diff --git a/manifests/charts/base/crds/runtime.agentcube.volcano.sh_agentruntimes.yaml b/manifests/charts/base/crds/runtime.agentcube.volcano.sh_agentruntimes.yaml index de219e65d..620451df5 100644 --- a/manifests/charts/base/crds/runtime.agentcube.volcano.sh_agentruntimes.yaml +++ b/manifests/charts/base/crds/runtime.agentcube.volcano.sh_agentruntimes.yaml @@ -49,6 +49,18 @@ spec: MaxSessionDuration describes the maximum duration for a session. After this duration, the session will be terminated no matter active or inactive. type: string + nodeStickiness: + description: |- + NodeStickiness configures soft node-affinity stickiness across sessions. + Disabled unless explicitly enabled. + properties: + enabled: + default: false + description: |- + Enabled turns on node stickiness for this workload's sessions. + Defaults to false. + type: boolean + type: object podTemplate: description: PodTemplate describes the template that will be used to create an agent sandbox. diff --git a/manifests/charts/base/crds/runtime.agentcube.volcano.sh_codeinterpreters.yaml b/manifests/charts/base/crds/runtime.agentcube.volcano.sh_codeinterpreters.yaml index a8cb4e3a9..f703cc361 100644 --- a/manifests/charts/base/crds/runtime.agentcube.volcano.sh_codeinterpreters.yaml +++ b/manifests/charts/base/crds/runtime.agentcube.volcano.sh_codeinterpreters.yaml @@ -65,6 +65,19 @@ spec: After this duration, the session will be terminated regardless of activity, to prevent long-lived sandboxes from accumulating unbounded state. type: string + nodeStickiness: + description: |- + NodeStickiness configures soft node-affinity stickiness across sessions. + Disabled unless explicitly enabled. Has no effect when WarmPoolSize > 0, + since warm-pool sandboxes are pre-scheduled before a session claims them. + properties: + enabled: + default: false + description: |- + Enabled turns on node stickiness for this workload's sessions. + Defaults to false. + type: boolean + type: object ports: description: |- Ports is a list of ports that the code interpreter runtime will expose. diff --git a/manifests/charts/base/templates/rbac/workloadmanager.yaml b/manifests/charts/base/templates/rbac/workloadmanager.yaml index 73f6e3b0a..18bae4206 100644 --- a/manifests/charts/base/templates/rbac/workloadmanager.yaml +++ b/manifests/charts/base/templates/rbac/workloadmanager.yaml @@ -36,7 +36,7 @@ rules: verbs: ["update"] - apiGroups: ["runtime.agentcube.volcano.sh"] resources: ["agentruntimes"] - verbs: ["get", "list", "watch"] + verbs: ["get", "list", "watch", "patch"] - apiGroups: ["runtime.agentcube.volcano.sh"] resources: ["agentruntimes/status"] verbs: ["update", "patch"] diff --git a/pkg/apis/runtime/v1alpha1/agent_type.go b/pkg/apis/runtime/v1alpha1/agent_type.go index fe5413faa..85a4325b8 100644 --- a/pkg/apis/runtime/v1alpha1/agent_type.go +++ b/pkg/apis/runtime/v1alpha1/agent_type.go @@ -56,6 +56,11 @@ type AgentRuntimeSpec struct { // +kubebuilder:validation:Required // +kubebuilder:default="8h" MaxSessionDuration *metav1.Duration `json:"maxSessionDuration,omitempty" protobuf:"bytes,3,opt,name=maxSessionDuration"` + + // NodeStickiness configures soft node-affinity stickiness across sessions. + // Disabled unless explicitly enabled. + // +optional + NodeStickiness *NodeStickinessSpec `json:"nodeStickiness,omitempty"` } // AgentRuntimeStatus represents the observed state of an AgentRuntime. diff --git a/pkg/apis/runtime/v1alpha1/codeinterpreter_types.go b/pkg/apis/runtime/v1alpha1/codeinterpreter_types.go index c12154f3e..9dcefaf24 100644 --- a/pkg/apis/runtime/v1alpha1/codeinterpreter_types.go +++ b/pkg/apis/runtime/v1alpha1/codeinterpreter_types.go @@ -81,6 +81,12 @@ type CodeInterpreterSpec struct { // +kubebuilder:validation:Enum=picod;none // +optional AuthMode AuthModeType `json:"authMode,omitempty"` + + // NodeStickiness configures soft node-affinity stickiness across sessions. + // Disabled unless explicitly enabled. Has no effect when WarmPoolSize > 0, + // since warm-pool sandboxes are pre-scheduled before a session claims them. + // +optional + NodeStickiness *NodeStickinessSpec `json:"nodeStickiness,omitempty"` } // CodeInterpreterStatus represents the observed state of a CodeInterpreter. @@ -170,6 +176,17 @@ type TargetPort struct { Protocol ProtocolType `json:"protocol"` } +// NodeStickinessSpec configures soft node-affinity stickiness: when enabled, a new +// session prefers (soft / PreferredDuringScheduling) the node the previous session +// for the same workload landed on. It is disabled unless explicitly turned on. +type NodeStickinessSpec struct { + // Enabled turns on node stickiness for this workload's sessions. + // Defaults to false. + // +kubebuilder:default=false + // +optional + Enabled bool `json:"enabled,omitempty"` +} + // AuthModeType defines the authentication mode for the sandbox runtime. type AuthModeType string diff --git a/pkg/apis/runtime/v1alpha1/zz_generated.deepcopy.go b/pkg/apis/runtime/v1alpha1/zz_generated.deepcopy.go index a7dbb4d41..429019fd7 100644 --- a/pkg/apis/runtime/v1alpha1/zz_generated.deepcopy.go +++ b/pkg/apis/runtime/v1alpha1/zz_generated.deepcopy.go @@ -108,6 +108,11 @@ func (in *AgentRuntimeSpec) DeepCopyInto(out *AgentRuntimeSpec) { *out = new(v1.Duration) **out = **in } + if in.NodeStickiness != nil { + in, out := &in.NodeStickiness, &out.NodeStickiness + *out = new(NodeStickinessSpec) + **out = **in + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new AgentRuntimeSpec. @@ -286,6 +291,11 @@ func (in *CodeInterpreterSpec) DeepCopyInto(out *CodeInterpreterSpec) { *out = new(int32) **out = **in } + if in.NodeStickiness != nil { + in, out := &in.NodeStickiness, &out.NodeStickiness + *out = new(NodeStickinessSpec) + **out = **in + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new CodeInterpreterSpec. @@ -320,6 +330,21 @@ func (in *CodeInterpreterStatus) DeepCopy() *CodeInterpreterStatus { return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *NodeStickinessSpec) DeepCopyInto(out *NodeStickinessSpec) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new NodeStickinessSpec. +func (in *NodeStickinessSpec) DeepCopy() *NodeStickinessSpec { + if in == nil { + return nil + } + out := new(NodeStickinessSpec) + in.DeepCopyInto(out) + return out +} + // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *SandboxTemplate) DeepCopyInto(out *SandboxTemplate) { *out = *in diff --git a/pkg/workloadmanager/handlers.go b/pkg/workloadmanager/handlers.go index 52f85eb2e..926b548de 100644 --- a/pkg/workloadmanager/handlers.go +++ b/pkg/workloadmanager/handlers.go @@ -425,12 +425,12 @@ func (s *Server) createSandbox(ctx context.Context, dynamicClient dynamic.Interf sandboxPodName = podName } - podIP, err := s.k8sClient.GetSandboxPodIP(ctx, createdSandbox.Namespace, sandboxNameForPod, sandboxPodName) + podIP, nodeName, err := s.k8sClient.GetSandboxPodInfo(ctx, createdSandbox.Namespace, sandboxNameForPod, sandboxPodName) if err != nil { if isContextError(err) { return nil, err } - return nil, api.NewInternalError(fmt.Errorf("failed to get sandbox %s/%s pod IP: %w", createdSandbox.Namespace, sandboxNameForPod, err)) + return nil, api.NewInternalError(fmt.Errorf("failed to get sandbox %s/%s pod info: %w", createdSandbox.Namespace, sandboxNameForPod, err)) } if err := s.waitForSandboxEntryPointsReady(ctx, podIP, sandboxEntry); err != nil { if isContextError(err) { @@ -464,11 +464,50 @@ func (s *Server) createSandbox(ctx context.Context, dynamicClient dynamic.Interf } needRollbackSandbox = false + + // Only record the node after the store update succeeds, so that a rollback + // (which deletes the sandbox) does not race with a goroutine that has already + // recorded a node for a session that was never fully established. + s.recordStickyNode(createdSandbox.Namespace, sandboxEntry, nodeName) klog.V(2).Infof("init sandbox %s/%s successfully, kind: %s, sessionID: %s", createdSandbox.Namespace, createdSandbox.Name, createdSandbox.Kind, sandboxEntry.SessionID) return response, nil } +// recordStickyNode records the scheduled node on the owning workload so the next +// session prefers the same node. Best-effort: a failure here must not fail the +// create, since the session is already usable. The patch runs in a detached +// goroutine with its own short timeout so that a slow API server or a client +// disconnect does not delay the sandbox creation response. +func (s *Server) recordStickyNode(namespace string, sandboxEntry *sandboxEntry, nodeName string) { + if sandboxEntry.StickyWorkloadName == "" || nodeName == "" { + return + } + kind := sandboxEntry.StickyWorkloadKind + name := sandboxEntry.StickyWorkloadName + + // Hold a read lock while checking the flag and adding to the wait group. + // Concurrent recordStickyNode calls do not block each other (read lock), + // while Shutdown() takes the write lock so that no recordStickyNode can + // observe shuttingDown=false and then call wg.Add after wg.Wait has begun. + s.shutdownMu.RLock() + if s.shuttingDown { + s.shutdownMu.RUnlock() + return + } + s.wg.Add(1) + s.shutdownMu.RUnlock() + + go func() { + defer s.wg.Done() + patchCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := s.k8sClient.PatchWorkloadLastNode(patchCtx, namespace, kind, name, nodeName); err != nil { + klog.Warningf("failed to patch %s %s/%s last-node=%s: %v", kind, namespace, name, nodeName, err) + } + }() +} + // rollbackSandboxCreation deletes the sandbox (or sandbox claim) and its store // placeholder when creation fails. It runs in a fresh context so that a // canceled request context does not prevent cleanup. diff --git a/pkg/workloadmanager/handlers_test.go b/pkg/workloadmanager/handlers_test.go index 2105967d5..513ff83ef 100644 --- a/pkg/workloadmanager/handlers_test.go +++ b/pkg/workloadmanager/handlers_test.go @@ -33,6 +33,7 @@ import ( "github.com/agiledragon/gomonkey/v2" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" + cubefake "github.com/volcano-sh/agentcube/client-go/clientset/versioned/fake" "github.com/volcano-sh/agentcube/pkg/api" runtimev1alpha1 "github.com/volcano-sh/agentcube/pkg/apis/runtime/v1alpha1" "github.com/volcano-sh/agentcube/pkg/common/types" @@ -243,11 +244,11 @@ func TestServerCreateSandbox(t *testing.T) { deleteCalls++ return nil }) - patches.ApplyMethod(reflect.TypeOf((*K8sClient)(nil)), "GetSandboxPodIP", func(_ *K8sClient, _ context.Context, _, _, _ string) (string, error) { + patches.ApplyMethod(reflect.TypeOf((*K8sClient)(nil)), "GetSandboxPodInfo", func(_ *K8sClient, _ context.Context, _, _, _ string) (string, string, error) { if cur.podIPErr != nil { - return "", cur.podIPErr + return "", "", cur.podIPErr } - return "10.0.0.9", nil + return "10.0.0.9", "node-1", nil }) patches.ApplyPrivateMethod(reflect.TypeOf(server), "waitForSandboxEntryPointsReady", func(_ *Server, _ context.Context, _ string, _ *sandboxEntry) error { return cur.readyErr @@ -386,11 +387,11 @@ func TestServerCreateSandboxClaimUsesAdoptedSandboxButStoresClaimName(t *testing patches.ApplyFunc(deleteSandboxClaim, func(_ context.Context, _ dynamic.Interface, _, _ string) error { return nil }) - patches.ApplyMethod(reflect.TypeOf((*K8sClient)(nil)), "GetSandboxPodIP", func(_ *K8sClient, _ context.Context, namespace, sandboxName, podName string) (string, error) { + patches.ApplyMethod(reflect.TypeOf((*K8sClient)(nil)), "GetSandboxPodInfo", func(_ *K8sClient, _ context.Context, namespace, sandboxName, podName string) (string, string, error) { gotNamespace = namespace gotSandboxName = sandboxName gotPodName = podName - return "10.0.0.10", nil + return "10.0.0.10", "node-1", nil }) patches.ApplyPrivateMethod(reflect.TypeOf(server), "waitForSandboxEntryPointsReady", func(_ *Server, _ context.Context, _ string, _ *sandboxEntry) error { return nil @@ -728,21 +729,23 @@ func makeSandbox(kind, ns, name string) (*sandboxv1alpha1.Sandbox, *sandboxEntry }, entry } +type testCaseHSC struct { + name string + kind string + body string + buildErr error + buildNotFound bool + createErr error + createResp *types.CreateSandboxResponse + expectStatus int + expectMessage string + expectCreateCalls int +} + func TestHandleSandboxCreate(t *testing.T) { gin.SetMode(gin.TestMode) - tests := []struct { - name string - kind string - body string - buildErr error - buildNotFound bool - createErr error - createResp *types.CreateSandboxResponse - expectStatus int - expectMessage string - expectCreateCalls int - }{ + tests := []testCaseHSC{ { name: "invalid json", kind: types.AgentRuntimeKind, @@ -836,10 +839,55 @@ func TestHandleSandboxCreate(t *testing.T) { }, } - for _, tt := range tests { - tc := tt + // Patch ONCE at the outer level and read per-case state through a shared + // pointer. Re-patching the same function per-subtest on arm64 causes gomonkey + // to silently fail on the second apply, because PC-relative branch + // instructions don't recalculate correctly after a reset. + type fixture struct { + tc *testCaseHSC + sb *sandboxv1alpha1.Sandbox + entry *sandboxEntry + claim *extensionsv1alpha1.SandboxClaim + createObj int + } + var cur *fixture + fakeServer := newFakeServer() + + patches := gomonkey.NewPatches() + defer patches.Reset() + + patches.ApplyFunc(buildSandboxByAgentRuntime, func(_, _, _ string, _ *Informers) (*sandboxv1alpha1.Sandbox, *sandboxEntry, error) { + if cur.tc.kind != types.AgentRuntimeKind { + return nil, nil, errors.New("unexpected kind") + } + if cur.tc.buildErr != nil { + return nil, nil, cur.tc.buildErr + } + return cur.sb, cur.entry, nil + }) + patches.ApplyFunc(buildSandboxByCodeInterpreter, func(_, _, _ string, _ *Informers) (*sandboxv1alpha1.Sandbox, *extensionsv1alpha1.SandboxClaim, *sandboxEntry, error) { + if cur.tc.kind != types.CodeInterpreterKind { + return nil, nil, nil, errors.New("unexpected kind") + } + if cur.tc.buildErr != nil { + return nil, nil, nil, cur.tc.buildErr + } + return cur.sb, cur.claim, cur.entry, nil + }) + patches.ApplyPrivateMethod(reflect.TypeOf(fakeServer), "createSandbox", func(_ *Server, _ context.Context, _ dynamic.Interface, _ *sandboxv1alpha1.Sandbox, _ *extensionsv1alpha1.SandboxClaim, _ *sandboxEntry, _ <-chan SandboxStatusUpdate) (*types.CreateSandboxResponse, error) { + cur.createObj++ + if cur.tc.createErr != nil { + return nil, cur.tc.createErr + } + if cur.tc.createResp != nil { + return cur.tc.createResp, nil + } + return nil, nil + }) + + for i := range tests { + tc := &tests[i] t.Run(tc.name, func(t *testing.T) { - fakeServer := newFakeServer() w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) @@ -849,45 +897,11 @@ func TestHandleSandboxCreate(t *testing.T) { sb, entry := makeSandbox(tc.kind, "ns", "sandbox-1") claim := &extensionsv1alpha1.SandboxClaim{ObjectMeta: metav1.ObjectMeta{Name: sb.Name, Namespace: sb.Namespace}} - - patches := gomonkey.NewPatches() - defer patches.Reset() - - patches.ApplyFunc(buildSandboxByAgentRuntime, func(_, _, _ string, _ *Informers) (*sandboxv1alpha1.Sandbox, *sandboxEntry, error) { - if tc.kind != types.AgentRuntimeKind { - return nil, nil, errors.New("unexpected kind") - } - if tc.buildErr != nil { - return nil, nil, tc.buildErr - } - return sb, entry, nil - }) - - patches.ApplyFunc(buildSandboxByCodeInterpreter, func(_, _, _ string, _ *Informers) (*sandboxv1alpha1.Sandbox, *extensionsv1alpha1.SandboxClaim, *sandboxEntry, error) { - if tc.kind != types.CodeInterpreterKind { - return nil, nil, nil, errors.New("unexpected kind") - } - if tc.buildErr != nil { - return nil, nil, nil, tc.buildErr - } - return sb, claim, entry, nil - }) - - createCalls := 0 - patches.ApplyPrivateMethod(reflect.TypeOf(fakeServer), "createSandbox", func(_ *Server, _ context.Context, _ dynamic.Interface, _ *sandboxv1alpha1.Sandbox, _ *extensionsv1alpha1.SandboxClaim, _ *sandboxEntry, _ <-chan SandboxStatusUpdate) (*types.CreateSandboxResponse, error) { - createCalls++ - if tc.createErr != nil { - return nil, tc.createErr - } - if tc.createResp != nil { - return tc.createResp, nil - } - return nil, nil - }) + cur = &fixture{tc: tc, sb: sb, entry: entry, claim: claim} fakeServer.handleSandboxCreate(c, tc.kind) - require.Equal(t, tc.expectCreateCalls, createCalls, "createSandbox call count") + require.Equal(t, tc.expectCreateCalls, cur.createObj, "createSandbox call count") require.Equal(t, tc.expectStatus, w.Code) if tc.expectStatus != http.StatusOK { @@ -1011,3 +1025,48 @@ func TestHandleSandboxCreate_IdentityErrors(t *testing.T) { }) } } + +// TestRecordStickyNode verifies the best-effort write-back guard clauses. +func TestRecordStickyNode(t *testing.T) { + t.Run("no-op when StickyWorkloadName is empty", func(_ *testing.T) { + server := &Server{k8sClient: &K8sClient{}} + entry := &sandboxEntry{StickyWorkloadName: "", StickyWorkloadKind: runtimev1alpha1.AgentRuntimeKind} + // Must not panic; the guard clause prevents calling PatchWorkloadLastNode. + server.recordStickyNode("ns", entry, "node-1") + }) + + t.Run("no-op when nodeName is empty", func(_ *testing.T) { + server := &Server{k8sClient: &K8sClient{}} + entry := &sandboxEntry{StickyWorkloadName: "ar-1", StickyWorkloadKind: runtimev1alpha1.AgentRuntimeKind} + // Must not panic; the guard clause prevents calling PatchWorkloadLastNode. + server.recordStickyNode("ns", entry, "") + }) + + t.Run("no-op when shutting down", func(_ *testing.T) { + server := &Server{k8sClient: &K8sClient{}} + server.shuttingDown = true + entry := &sandboxEntry{StickyWorkloadName: "ar-1", StickyWorkloadKind: runtimev1alpha1.AgentRuntimeKind} + // Must not spawn a goroutine (no wg.Add) and must not panic. + server.recordStickyNode("ns", entry, "node-1") + }) + + t.Run("concurrent recordStickyNode does not block", func(_ *testing.T) { + // RWMutex with RLock allows concurrent readers, so multiple + // recordStickyNode calls should not block each other. + // Provide a fake cubeClientset so the async goroutine does not panic. + cs := cubefake.NewSimpleClientset() + server := &Server{k8sClient: &K8sClient{cubeClientset: cs}} + entry := &sandboxEntry{StickyWorkloadName: "ar-1", StickyWorkloadKind: runtimev1alpha1.AgentRuntimeKind} + + done := make(chan struct{}, 10) + for i := 0; i < 10; i++ { + go func() { + server.recordStickyNode("ns", entry, "node-1") + done <- struct{}{} + }() + } + for i := 0; i < 10; i++ { + <-done + } + }) +} diff --git a/pkg/workloadmanager/k8s_client.go b/pkg/workloadmanager/k8s_client.go index 1d241693d..2fae8bdc1 100644 --- a/pkg/workloadmanager/k8s_client.go +++ b/pkg/workloadmanager/k8s_client.go @@ -30,6 +30,7 @@ import ( "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" "k8s.io/apimachinery/pkg/labels" "k8s.io/apimachinery/pkg/runtime" + k8stypes "k8s.io/apimachinery/pkg/types" "k8s.io/client-go/dynamic" "k8s.io/client-go/dynamic/dynamicinformer" "k8s.io/client-go/informers" @@ -58,6 +59,9 @@ var ( LastActivityAnnotationKey = "last-activity-time" // IdleTimeoutAnnotationKey key for idle timeout IdleTimeoutAnnotationKey = "runtime.agentcube.io/idle-timeout" + // LastNodeAnnotationKey records the node the workload's previous session landed + // on, so the next session can prefer the same node (soft node stickiness). + LastNodeAnnotationKey = "runtime.agentcube.io/last-node" ) // K8sClient encapsulates the Kubernetes client @@ -81,6 +85,12 @@ type sandboxEntry struct { OwnerID string Ports []runtimev1alpha1.TargetPort IdleTimeout time.Duration + // StickyWorkloadName/StickyWorkloadKind identify the owning workload whose + // last-node annotation should be patched for node stickiness. They are set + // only on the direct-create path (no warm pool) and are not persisted to the + // store. An empty StickyWorkloadName disables the write-back. + StickyWorkloadName string + StickyWorkloadKind string } // NewK8sClient creates a new Kubernetes client @@ -336,50 +346,77 @@ func (u *UserK8sClient) DeleteSandboxClaim(ctx context.Context, namespace, sandb return deleteSandboxClaim(ctx, u.dynamicClient, namespace, sandboxClaimName) } -// GetSandboxPodIP gets the IP address of the pod corresponding to the Sandbox -func (c *K8sClient) GetSandboxPodIP(ctx context.Context, namespace, sandboxName, podName string) (string, error) { +// GetSandboxPodInfo gets the IP address and scheduled node of the pod corresponding to the Sandbox. +func (c *K8sClient) GetSandboxPodInfo(ctx context.Context, namespace, sandboxName, podName string) (podIP string, nodeName string, err error) { // An explicit Pod name comes from the live Sandbox object. Read the Pod from // the API too, so a lagging informer cannot reject an otherwise ready Sandbox. if podName != "" { pod, err := c.clientset.CoreV1().Pods(namespace).Get(ctx, podName, metav1.GetOptions{}) if err != nil { - return "", fmt.Errorf("failed to get sandbox pod %s/%s: %w", namespace, podName, err) + return "", "", fmt.Errorf("failed to get sandbox pod %s/%s: %w", namespace, podName, err) } - return validateAndGetPodIP(pod) + return validateAndGetPodInfo(pod) } // Find pod through label selector (sandbox-name label we set) pods, err := c.podLister.Pods(namespace).List(labels.SelectorFromSet(map[string]string{SandboxNameLabelKey: sandboxName})) if err != nil { - return "", fmt.Errorf("failed to list pods from cache: %w", err) + return "", "", fmt.Errorf("failed to list pods from cache: %w", err) } // Find the pod that belongs to this sandbox by checking ownerReferences for _, pod := range pods { for _, ownerRef := range pod.OwnerReferences { if ownerRef.Kind == "Sandbox" && ownerRef.Name == sandboxName { if ownerRef.Controller == nil || *ownerRef.Controller { - return validateAndGetPodIP(pod) + return validateAndGetPodInfo(pod) } } } } - return "", fmt.Errorf("no pod found for sandbox %s", sandboxName) + return "", "", fmt.Errorf("no pod found for sandbox %s", sandboxName) } -// validateAndGetPodIP validates pod status and returns IP -func validateAndGetPodIP(pod *corev1.Pod) (string, error) { +// validateAndGetPodInfo validates pod status and returns the pod IP and scheduled node +func validateAndGetPodInfo(pod *corev1.Pod) (podIP string, nodeName string, err error) { // Check if Pod is running if pod.Status.Phase != corev1.PodRunning { - return "", fmt.Errorf("pod not running yet, status: %s", pod.Status.Phase) + return "", "", fmt.Errorf("pod not running yet, status: %s", pod.Status.Phase) } // Check if Pod IP is assigned if pod.Status.PodIP == "" { - return "", fmt.Errorf("pod IP not assigned yet") + return "", "", fmt.Errorf("pod IP not assigned yet") } - return pod.Status.PodIP, nil + return pod.Status.PodIP, pod.Spec.NodeName, nil +} + +// PatchWorkloadLastNode records the node the latest session landed on as an +// annotation on the owning workload (AgentRuntime or CodeInterpreter), so the +// next session can prefer the same node. +func (c *K8sClient) PatchWorkloadLastNode(ctx context.Context, namespace, kind, name, nodeName string) error { + patch := []byte(fmt.Sprintf( + `{"metadata":{"annotations":{%q:%q}}}`, + LastNodeAnnotationKey, nodeName, + )) + var err error + switch kind { + case runtimev1alpha1.AgentRuntimeKind: + _, err = c.cubeClientset.RuntimeV1alpha1().AgentRuntimes(namespace).Patch( + ctx, name, k8stypes.MergePatchType, patch, metav1.PatchOptions{}, + ) + case runtimev1alpha1.CodeInterpreterKind: + _, err = c.cubeClientset.RuntimeV1alpha1().CodeInterpreters(namespace).Patch( + ctx, name, k8stypes.MergePatchType, patch, metav1.PatchOptions{}, + ) + default: + return fmt.Errorf("patch last-node: unsupported workload kind %q", kind) + } + if err != nil { + return fmt.Errorf("patch %s %s/%s last-node: %w", kind, namespace, name, err) + } + return nil } // WaitForSandboxReady waits for the Sandbox to be ready diff --git a/pkg/workloadmanager/k8s_client_test.go b/pkg/workloadmanager/k8s_client_test.go index 5875bc559..f67f98f06 100644 --- a/pkg/workloadmanager/k8s_client_test.go +++ b/pkg/workloadmanager/k8s_client_test.go @@ -24,16 +24,20 @@ import ( "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/labels" "k8s.io/client-go/kubernetes" listersv1 "k8s.io/client-go/listers/core/v1" "k8s.io/client-go/rest" + + cubefake "github.com/volcano-sh/agentcube/client-go/clientset/versioned/fake" + runtimev1alpha1 "github.com/volcano-sh/agentcube/pkg/apis/runtime/v1alpha1" ) // Helper function to create a pod with owner reference -func createPodWithOwner(name, namespace, sandboxName string, phase corev1.PodPhase, podIP string) *corev1.Pod { +func createPodWithOwner(name, namespace, sandboxName string, phase corev1.PodPhase, podIP, nodeName string) *corev1.Pod { controller := true return &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{ @@ -50,6 +54,9 @@ func createPodWithOwner(name, namespace, sandboxName string, phase corev1.PodPha }, }, }, + Spec: corev1.PodSpec{ + NodeName: nodeName, + }, Status: corev1.PodStatus{ Phase: phase, PodIP: podIP, @@ -120,10 +127,10 @@ func (m *mockPodLister) addPod(pod *corev1.Pod) { m.podsByNamespace[pod.Namespace] = append(m.podsByNamespace[pod.Namespace], pod) } -// TestGetSandboxPodIP_Success verifies GetSandboxPodIP returns IP when pod is present and valid -func TestGetSandboxPodIP_Success(t *testing.T) { +// TestGetSandboxPodInfo_Success verifies GetSandboxPodInfo returns IP and nodeName when pod is present and valid +func TestGetSandboxPodInfo_Success(t *testing.T) { // Setup: Create a mock pod lister with a valid running pod - pod := createPodWithOwner("test-pod", "test-namespace", "test-sandbox", corev1.PodRunning, "10.0.0.1") + pod := createPodWithOwner("test-pod", "test-namespace", "test-sandbox", corev1.PodRunning, "10.0.0.1", "node-1") mockPodLister := newMockPodLister() mockPodLister.addPod(pod) @@ -132,15 +139,19 @@ func TestGetSandboxPodIP_Success(t *testing.T) { } // Execute - ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "") + ip, nodeName, err := client.GetSandboxPodInfo(context.Background(), "test-namespace", "test-sandbox", "") // Verify assert.NoError(t, err, "Expected no error for valid pod") assert.Equal(t, "10.0.0.1", ip, "Expected IP to match pod IP") + assert.Equal(t, "node-1", nodeName, "Expected nodeName to match pod spec") } -func TestGetSandboxPodIPReadsNamedPodFromLiveAPI(t *testing.T) { - pod := createPodWithOwner("warm-pool-pod", "test-namespace", "warm-pool-sandbox", corev1.PodRunning, "10.0.0.2") +// TestGetSandboxPodInfoReadsNamedPodFromLiveAPI verifies that when a podName is +// provided, GetSandboxPodInfo reads the Pod from the live API server rather than +// the informer cache, so that a stale cache entry cannot reject a ready Sandbox. +func TestGetSandboxPodInfoReadsNamedPodFromLiveAPI(t *testing.T) { + pod := createPodWithOwner("warm-pool-pod", "test-namespace", "warm-pool-sandbox", corev1.PodRunning, "10.0.0.2", "node-warm") pod.TypeMeta = metav1.TypeMeta{APIVersion: "v1", Kind: "Pod"} apiServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -162,18 +173,20 @@ func TestGetSandboxPodIPReadsNamedPodFromLiveAPI(t *testing.T) { stalePod := pod.DeepCopy() stalePod.Status.Phase = corev1.PodPending stalePod.Status.PodIP = "" + stalePod.Spec.NodeName = "" staleLister := newMockPodLister() staleLister.addPod(stalePod) client := &K8sClient{clientset: clientset, podLister: staleLister} - ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "warm-pool-sandbox", "warm-pool-pod") + ip, nodeName, err := client.GetSandboxPodInfo(context.Background(), "test-namespace", "warm-pool-sandbox", "warm-pool-pod") assert.NoError(t, err) assert.Equal(t, "10.0.0.2", ip) + assert.Equal(t, "node-warm", nodeName) } -// TestGetSandboxPodIP_PodNotFound verifies GetSandboxPodIP returns error when pod is not found -func TestGetSandboxPodIP_PodNotFound(t *testing.T) { +// TestGetSandboxPodInfo_PodNotFound verifies GetSandboxPodInfo returns error when pod is not found +func TestGetSandboxPodInfo_PodNotFound(t *testing.T) { // Setup: Create a mock pod lister with pod that has wrong label mockPodLister := newMockPodLister() pod := &corev1.Pod{ @@ -196,41 +209,45 @@ func TestGetSandboxPodIP_PodNotFound(t *testing.T) { } // Execute - ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "") + ip, nodeName, err := client.GetSandboxPodInfo(context.Background(), "test-namespace", "test-sandbox", "") // Verify assert.Error(t, err, "Expected error when pod not found") assert.Empty(t, ip, "Expected empty IP when error occurs") + assert.Empty(t, nodeName, "Expected empty nodeName when error occurs") assert.Contains(t, err.Error(), "no pod found for sandbox test-sandbox", "Error message should indicate pod not found") } -// TestGetSandboxPodIP_InvalidPodStatus verifies GetSandboxPodIP returns error when pod status is invalid -func TestGetSandboxPodIP_InvalidPodStatus(t *testing.T) { +// TestGetSandboxPodInfo_InvalidPodStatus verifies GetSandboxPodInfo returns error when pod status is invalid +func TestGetSandboxPodInfo_InvalidPodStatus(t *testing.T) { // Test cases: pod not running or pod without IP testCases := []struct { - name string - phase corev1.PodPhase - podIP string - errMsg string + name string + phase corev1.PodPhase + podIP string + nodeName string + errMsg string }{ { - name: "pod not running", - phase: corev1.PodPending, - podIP: "10.0.0.1", - errMsg: "pod not running yet", + name: "pod not running", + phase: corev1.PodPending, + podIP: "10.0.0.1", + nodeName: "node-1", + errMsg: "pod not running yet", }, { - name: "pod without IP", - phase: corev1.PodRunning, - podIP: "", - errMsg: "pod IP not assigned yet", + name: "pod without IP", + phase: corev1.PodRunning, + podIP: "", + nodeName: "node-2", + errMsg: "pod IP not assigned yet", }, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { // Setup: Create a mock pod lister with invalid pod status - pod := createPodWithOwner("test-pod", "test-namespace", "test-sandbox", tc.phase, tc.podIP) + pod := createPodWithOwner("test-pod", "test-namespace", "test-sandbox", tc.phase, tc.podIP, tc.nodeName) mockPodLister := newMockPodLister() mockPodLister.addPod(pod) @@ -239,12 +256,69 @@ func TestGetSandboxPodIP_InvalidPodStatus(t *testing.T) { } // Execute - ip, err := client.GetSandboxPodIP(context.Background(), "test-namespace", "test-sandbox", "") + ip, nodeName, err := client.GetSandboxPodInfo(context.Background(), "test-namespace", "test-sandbox", "") // Verify assert.Error(t, err, "Expected error for invalid pod status") assert.Empty(t, ip, "Expected empty IP when error occurs") + assert.Empty(t, nodeName, "Expected empty nodeName when error occurs") assert.Contains(t, err.Error(), tc.errMsg, "Error message should indicate the issue") }) } } + +// TestPatchWorkloadLastNode verifies the last-node annotation is written to the +// correct workload kind, and that an unsupported kind is rejected. +func TestPatchWorkloadLastNode(t *testing.T) { + const ( + namespace = "test-namespace" + nodeName = "node-a" + ) + + t.Run("AgentRuntime", func(t *testing.T) { + ar := &runtimev1alpha1.AgentRuntime{ + ObjectMeta: metav1.ObjectMeta{Name: "ar-1", Namespace: namespace}, + } + cs := cubefake.NewSimpleClientset(ar) + client := &K8sClient{cubeClientset: cs} + + err := client.PatchWorkloadLastNode(context.Background(), namespace, runtimev1alpha1.AgentRuntimeKind, "ar-1", nodeName) + require.NoError(t, err) + + got, err := cs.RuntimeV1alpha1().AgentRuntimes(namespace).Get(context.Background(), "ar-1", metav1.GetOptions{}) + require.NoError(t, err) + assert.Equal(t, nodeName, got.Annotations[LastNodeAnnotationKey]) + }) + + t.Run("CodeInterpreter", func(t *testing.T) { + ci := &runtimev1alpha1.CodeInterpreter{ + ObjectMeta: metav1.ObjectMeta{Name: "ci-1", Namespace: namespace}, + } + cs := cubefake.NewSimpleClientset(ci) + client := &K8sClient{cubeClientset: cs} + + err := client.PatchWorkloadLastNode(context.Background(), namespace, runtimev1alpha1.CodeInterpreterKind, "ci-1", nodeName) + require.NoError(t, err) + + got, err := cs.RuntimeV1alpha1().CodeInterpreters(namespace).Get(context.Background(), "ci-1", metav1.GetOptions{}) + require.NoError(t, err) + assert.Equal(t, nodeName, got.Annotations[LastNodeAnnotationKey]) + }) + + t.Run("unsupported kind", func(t *testing.T) { + client := &K8sClient{cubeClientset: cubefake.NewSimpleClientset()} + err := client.PatchWorkloadLastNode(context.Background(), namespace, "Sandbox", "x", nodeName) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported workload kind") + }) + + t.Run("non-existent resource returns wrapped error", func(t *testing.T) { + // Patch on a missing AgentRuntime produces an API error that is wrapped. + client := &K8sClient{cubeClientset: cubefake.NewSimpleClientset()} + err := client.PatchWorkloadLastNode(context.Background(), namespace, runtimev1alpha1.AgentRuntimeKind, "missing", nodeName) + require.Error(t, err) + assert.Contains(t, err.Error(), "patch AgentRuntime") + assert.Contains(t, err.Error(), namespace) + assert.Contains(t, err.Error(), "missing") + }) +} diff --git a/pkg/workloadmanager/server.go b/pkg/workloadmanager/server.go index f77c157e4..45ee6c864 100644 --- a/pkg/workloadmanager/server.go +++ b/pkg/workloadmanager/server.go @@ -45,6 +45,8 @@ type Server struct { informers *Informers storeClient store.Store wg sync.WaitGroup + shutdownMu sync.RWMutex + shuttingDown bool certWatcher *mtls.CertWatcher // mTLS cert watcher for graceful cleanup } @@ -226,6 +228,14 @@ func configureHTTP2TLSServer(server *http.Server, tlsConfig *tls.Config) error { // Shutdown performs graceful shutdown of the HTTP server. func (s *Server) Shutdown(ctx context.Context) error { + // Signal that the server is shutting down so that new best-effort + // background work (e.g. recordStickyNode) is not accepted. Use write lock + // so that no concurrent recordStickyNode (which holds a read lock) can + // observe shuttingDown=false and then call wg.Add after wg.Wait has begun. + s.shutdownMu.Lock() + s.shuttingDown = true + s.shutdownMu.Unlock() + if s.httpServer != nil { klog.Info("Shutting down HTTP server...") if err := s.httpServer.Shutdown(ctx); err != nil { diff --git a/pkg/workloadmanager/workload_builder.go b/pkg/workloadmanager/workload_builder.go index 2a42e4f07..a936f7ab2 100644 --- a/pkg/workloadmanager/workload_builder.go +++ b/pkg/workloadmanager/workload_builder.go @@ -274,6 +274,14 @@ func buildSandboxByAgentRuntime(namespace string, name string, ownerID string, i podSpec.RuntimeClassName = nil } + // Node stickiness: if enabled and a previous session recorded a node, prefer it (soft). + stickinessOn := nodeStickinessEnabled(agentRuntimeObj.Spec.NodeStickiness) + if stickinessOn { + if lastNode := agentRuntimeObj.Annotations[LastNodeAnnotationKey]; lastNode != "" { + injectPreferredNodeAffinity(podSpec, lastNode) + } + } + buildParams := &buildSandboxParams{ namespace: namespace, workloadName: name, @@ -305,9 +313,78 @@ func buildSandboxByAgentRuntime(namespace string, name string, ownerID string, i SessionID: sessionID, IdleTimeout: idleTimeout, } + if stickinessOn { + entry.StickyWorkloadName = name + entry.StickyWorkloadKind = runtimev1alpha1.AgentRuntimeKind + } return sandbox, entry, nil } +// nodeStickinessEnabled reports whether soft node-affinity stickiness is turned on. +// A nil spec (field omitted) means disabled. +func nodeStickinessEnabled(spec *runtimev1alpha1.NodeStickinessSpec) bool { + return spec != nil && spec.Enabled +} + +// injectPreferredNodeAffinity appends a soft (preferred) node-affinity term that +// favors scheduling onto nodeName, without overwriting any affinity the user +// already defined on the pod spec. It is a no-op when nodeName is empty or an +// equivalent hostname preference already exists, to avoid duplicate terms that +// would amplify scheduler weighting and bloat the pod spec. +func injectPreferredNodeAffinity(podSpec *corev1.PodSpec, nodeName string) { + if nodeName == "" { + return + } + if podSpec.Affinity != nil && podSpec.Affinity.NodeAffinity != nil { + for _, existing := range podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution { + if isEquivalentHostnamePreference(existing.Preference, nodeName) { + return + } + } + } + term := corev1.PreferredSchedulingTerm{ + Weight: 100, + Preference: corev1.NodeSelectorTerm{ + MatchExpressions: []corev1.NodeSelectorRequirement{ + { + Key: corev1.LabelHostname, + Operator: corev1.NodeSelectorOpIn, + Values: []string{nodeName}, + }, + }, + }, + } + if podSpec.Affinity == nil { + podSpec.Affinity = &corev1.Affinity{} + } + if podSpec.Affinity.NodeAffinity == nil { + podSpec.Affinity.NodeAffinity = &corev1.NodeAffinity{} + } + podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append( + podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution, + term, + ) +} + +// isEquivalentHostnamePreference reports whether term is a standalone hostname +// preference that already favors nodeName exclusively. A NodeSelectorTerm ANDs +// its MatchExpressions and MatchFields, so a term carrying any additional +// constraint is strictly more restrictive than the soft hostname term we inject +// and must not suppress injection. We also only treat a term as equivalent when +// the values list is exactly [nodeName]; a multi-value hostname preference +// (e.g., hostname In [nodeA, nodeB]) does not strongly prefer nodeName and must +// not prevent injection of the stickiness term. +func isEquivalentHostnamePreference(term corev1.NodeSelectorTerm, nodeName string) bool { + if len(term.MatchFields) != 0 || len(term.MatchExpressions) != 1 { + return false + } + expr := term.MatchExpressions[0] + if expr.Key != corev1.LabelHostname || expr.Operator != corev1.NodeSelectorOpIn { + return false + } + return len(expr.Values) == 1 && expr.Values[0] == nodeName +} + // buildCodeInterpreterEnvVars copies the template env vars and injects the // public key when authMode is picod. func buildCodeInterpreterEnvVars(templateEnv []corev1.EnvVar, authMode runtimev1alpha1.AuthModeType) []corev1.EnvVar { @@ -418,6 +495,15 @@ func buildSandboxByCodeInterpreter(namespace string, codeInterpreterName string, }, } + // Node stickiness: if enabled and a previous session recorded a node, prefer it (soft). + if nodeStickinessEnabled(codeInterpreterObj.Spec.NodeStickiness) { + if lastNode := codeInterpreterObj.Annotations[LastNodeAnnotationKey]; lastNode != "" { + injectPreferredNodeAffinity(&podSpec, lastNode) + } + sandboxEntry.StickyWorkloadName = codeInterpreterName + sandboxEntry.StickyWorkloadKind = runtimev1alpha1.CodeInterpreterKind + } + buildParams := &buildSandboxParams{ sandboxName: sandboxName, namespace: namespace, diff --git a/pkg/workloadmanager/workload_builder_test.go b/pkg/workloadmanager/workload_builder_test.go index 035dd5c6b..3ca7444a7 100644 --- a/pkg/workloadmanager/workload_builder_test.go +++ b/pkg/workloadmanager/workload_builder_test.go @@ -349,6 +349,332 @@ func TestBuildCodeInterpreterEnvVars(t *testing.T) { } } +func TestInjectPreferredNodeAffinity(t *testing.T) { + hostnameTerm := func(node string) corev1.PreferredSchedulingTerm { + return corev1.PreferredSchedulingTerm{ + Weight: 100, + Preference: corev1.NodeSelectorTerm{ + MatchExpressions: []corev1.NodeSelectorRequirement{ + { + Key: corev1.LabelHostname, + Operator: corev1.NodeSelectorOpIn, + Values: []string{node}, + }, + }, + }, + } + } + + t.Run("nil affinity is initialized", func(t *testing.T) { + podSpec := &corev1.PodSpec{} + injectPreferredNodeAffinity(podSpec, "node-a") + + preferred := podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution + if len(preferred) != 1 { + t.Fatalf("expected 1 preferred term, got %d", len(preferred)) + } + assert.Equal(t, hostnameTerm("node-a"), preferred[0]) + }) + + t.Run("appends without overwriting existing terms", func(t *testing.T) { + existing := corev1.PreferredSchedulingTerm{ + Weight: 10, + Preference: corev1.NodeSelectorTerm{ + MatchExpressions: []corev1.NodeSelectorRequirement{ + {Key: "zone", Operator: corev1.NodeSelectorOpIn, Values: []string{"z1"}}, + }, + }, + } + podSpec := &corev1.PodSpec{ + Affinity: &corev1.Affinity{ + NodeAffinity: &corev1.NodeAffinity{ + PreferredDuringSchedulingIgnoredDuringExecution: []corev1.PreferredSchedulingTerm{existing}, + }, + }, + } + injectPreferredNodeAffinity(podSpec, "node-b") + + preferred := podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution + if len(preferred) != 2 { + t.Fatalf("expected 2 preferred terms, got %d", len(preferred)) + } + assert.Equal(t, existing, preferred[0], "existing term must be preserved") + assert.Equal(t, hostnameTerm("node-b"), preferred[1]) + }) + + t.Run("preserves existing required affinity and pod affinity", func(t *testing.T) { + required := &corev1.NodeSelector{ + NodeSelectorTerms: []corev1.NodeSelectorTerm{ + {MatchExpressions: []corev1.NodeSelectorRequirement{ + {Key: "disk", Operator: corev1.NodeSelectorOpIn, Values: []string{"ssd"}}, + }}, + }, + } + podSpec := &corev1.PodSpec{ + Affinity: &corev1.Affinity{ + NodeAffinity: &corev1.NodeAffinity{ + RequiredDuringSchedulingIgnoredDuringExecution: required, + }, + PodAffinity: &corev1.PodAffinity{}, + }, + } + injectPreferredNodeAffinity(podSpec, "node-c") + + assert.Equal(t, required, podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution) + assert.NotNil(t, podSpec.Affinity.PodAffinity) + assert.Len(t, podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution, 1) + }) + + t.Run("skips when an equivalent standalone hostname preference exists", func(t *testing.T) { + podSpec := &corev1.PodSpec{ + Affinity: &corev1.Affinity{ + NodeAffinity: &corev1.NodeAffinity{ + PreferredDuringSchedulingIgnoredDuringExecution: []corev1.PreferredSchedulingTerm{ + hostnameTerm("node-d"), + }, + }, + }, + } + injectPreferredNodeAffinity(podSpec, "node-d") + + preferred := podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution + assert.Len(t, preferred, 1, "duplicate hostname preference must not be appended") + }) + + t.Run("still injects when hostname term carries extra match expressions", func(t *testing.T) { + // A NodeSelectorTerm ANDs its expressions, so this term only matches + // node-e when it is also in zone z1 — strictly more restrictive than the + // soft hostname term we inject, so injection must NOT be suppressed. + restrictive := corev1.PreferredSchedulingTerm{ + Weight: 100, + Preference: corev1.NodeSelectorTerm{ + MatchExpressions: []corev1.NodeSelectorRequirement{ + {Key: corev1.LabelHostname, Operator: corev1.NodeSelectorOpIn, Values: []string{"node-e"}}, + {Key: "zone", Operator: corev1.NodeSelectorOpIn, Values: []string{"z1"}}, + }, + }, + } + podSpec := &corev1.PodSpec{ + Affinity: &corev1.Affinity{ + NodeAffinity: &corev1.NodeAffinity{ + PreferredDuringSchedulingIgnoredDuringExecution: []corev1.PreferredSchedulingTerm{restrictive}, + }, + }, + } + injectPreferredNodeAffinity(podSpec, "node-e") + + preferred := podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution + if len(preferred) != 2 { + t.Fatalf("expected 2 preferred terms, got %d", len(preferred)) + } + assert.Equal(t, restrictive, preferred[0], "restrictive term must be preserved") + assert.Equal(t, hostnameTerm("node-e"), preferred[1]) + }) + + t.Run("still injects when hostname term carries match fields", func(t *testing.T) { + withFields := corev1.PreferredSchedulingTerm{ + Weight: 100, + Preference: corev1.NodeSelectorTerm{ + MatchExpressions: []corev1.NodeSelectorRequirement{ + {Key: corev1.LabelHostname, Operator: corev1.NodeSelectorOpIn, Values: []string{"node-f"}}, + }, + MatchFields: []corev1.NodeSelectorRequirement{ + {Key: "metadata.name", Operator: corev1.NodeSelectorOpIn, Values: []string{"node-f"}}, + }, + }, + } + podSpec := &corev1.PodSpec{ + Affinity: &corev1.Affinity{ + NodeAffinity: &corev1.NodeAffinity{ + PreferredDuringSchedulingIgnoredDuringExecution: []corev1.PreferredSchedulingTerm{withFields}, + }, + }, + } + injectPreferredNodeAffinity(podSpec, "node-f") + + preferred := podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution + assert.Len(t, preferred, 2, "term with match fields is more restrictive; injection must proceed") + }) + + t.Run("no-op when nodeName is empty", func(t *testing.T) { + podSpec := &corev1.PodSpec{} + injectPreferredNodeAffinity(podSpec, "") + assert.Nil(t, podSpec.Affinity, "affinity must remain nil when nodeName is empty") + }) + + t.Run("still injects when hostname expression uses different key", func(t *testing.T) { + // A term with "zone" In [...] is not a hostname preference, so injection must proceed. + zoneTerm := corev1.PreferredSchedulingTerm{ + Weight: 100, + Preference: corev1.NodeSelectorTerm{ + MatchExpressions: []corev1.NodeSelectorRequirement{ + {Key: "zone", Operator: corev1.NodeSelectorOpIn, Values: []string{"node-g"}}, + }, + }, + } + podSpec := &corev1.PodSpec{ + Affinity: &corev1.Affinity{ + NodeAffinity: &corev1.NodeAffinity{ + PreferredDuringSchedulingIgnoredDuringExecution: []corev1.PreferredSchedulingTerm{zoneTerm}, + }, + }, + } + injectPreferredNodeAffinity(podSpec, "node-g") + preferred := podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution + assert.Len(t, preferred, 2, "non-hostname key must not suppress injection") + }) + + t.Run("still injects when hostname expression uses different operator", func(t *testing.T) { + // A term with "hostname NotIn [...]" is not a soft preference for the same node. + notInTerm := corev1.PreferredSchedulingTerm{ + Weight: 100, + Preference: corev1.NodeSelectorTerm{ + MatchExpressions: []corev1.NodeSelectorRequirement{ + {Key: corev1.LabelHostname, Operator: corev1.NodeSelectorOpNotIn, Values: []string{"node-h"}}, + }, + }, + } + podSpec := &corev1.PodSpec{ + Affinity: &corev1.Affinity{ + NodeAffinity: &corev1.NodeAffinity{ + PreferredDuringSchedulingIgnoredDuringExecution: []corev1.PreferredSchedulingTerm{notInTerm}, + }, + }, + } + injectPreferredNodeAffinity(podSpec, "node-h") + preferred := podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution + assert.Len(t, preferred, 2, "NotIn operator must not suppress injection") + }) + + t.Run("still injects when hostname preference exists for a different node", func(t *testing.T) { + // A hostname In preference for node-i must not suppress injection for node-j. + otherTerm := corev1.PreferredSchedulingTerm{ + Weight: 100, + Preference: corev1.NodeSelectorTerm{ + MatchExpressions: []corev1.NodeSelectorRequirement{ + {Key: corev1.LabelHostname, Operator: corev1.NodeSelectorOpIn, Values: []string{"node-i"}}, + }, + }, + } + podSpec := &corev1.PodSpec{ + Affinity: &corev1.Affinity{ + NodeAffinity: &corev1.NodeAffinity{ + PreferredDuringSchedulingIgnoredDuringExecution: []corev1.PreferredSchedulingTerm{otherTerm}, + }, + }, + } + injectPreferredNodeAffinity(podSpec, "node-j") + preferred := podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution + assert.Len(t, preferred, 2, "hostname preference for a different node must not suppress injection") + }) + + t.Run("still injects when hostname preference contains multiple values", func(t *testing.T) { + // hostname In [node-k, node-l] is not an exclusive preference for node-k. + multiTerm := corev1.PreferredSchedulingTerm{ + Weight: 100, + Preference: corev1.NodeSelectorTerm{ + MatchExpressions: []corev1.NodeSelectorRequirement{ + {Key: corev1.LabelHostname, Operator: corev1.NodeSelectorOpIn, Values: []string{"node-k", "node-l"}}, + }, + }, + } + podSpec := &corev1.PodSpec{ + Affinity: &corev1.Affinity{ + NodeAffinity: &corev1.NodeAffinity{ + PreferredDuringSchedulingIgnoredDuringExecution: []corev1.PreferredSchedulingTerm{multiTerm}, + }, + }, + } + injectPreferredNodeAffinity(podSpec, "node-k") + preferred := podSpec.Affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution + assert.Len(t, preferred, 2, "multi-value hostname preference must not suppress injection") + }) +} + +func TestBuildSandboxByAgentRuntime_NodeStickiness(t *testing.T) { + newIfm := func(ar *runtimev1alpha1.AgentRuntime) *Informers { + fakeClient := cubefake.NewSimpleClientset(ar) + factory := cubeinformers.NewSharedInformerFactory(fakeClient, 0) + agentRuntimeInformer := factory.Runtime().V1alpha1().AgentRuntimes() + if err := agentRuntimeInformer.Informer().GetStore().Add(ar); err != nil { + t.Fatalf("failed to add agent runtime to informer store: %v", err) + } + return &Informers{ + AgentRuntimeLister: agentRuntimeInformer.Lister(), + AgentRuntimeInformer: agentRuntimeInformer.Informer(), + informerFactory: newFactory(), + cubeInformerFactory: factory, + } + } + + baseSpec := runtimev1alpha1.AgentRuntimeSpec{ + Template: &runtimev1alpha1.SandboxTemplate{ + Spec: corev1.PodSpec{ + Containers: []corev1.Container{{Name: "main", Image: "nginx"}}, + }, + }, + } + enabledSpec := func() runtimev1alpha1.AgentRuntimeSpec { + s := baseSpec + s.NodeStickiness = &runtimev1alpha1.NodeStickinessSpec{Enabled: true} + return s + } + + t.Run("disabled by default does not inject affinity or set sticky fields", func(t *testing.T) { + ar := &runtimev1alpha1.AgentRuntime{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: testNamespace, + Name: testAgentRuntimeName, + Annotations: map[string]string{LastNodeAnnotationKey: "sticky-node"}, + }, + Spec: baseSpec, + } + sandbox, entry, err := buildSandboxByAgentRuntime(testNamespace, testAgentRuntimeName, "", newIfm(ar)) + assert.NoError(t, err) + assert.Nil(t, sandbox.Spec.PodTemplate.Spec.Affinity, "affinity must not be injected when stickiness is off") + assert.Empty(t, entry.StickyWorkloadName, "write-back must be disabled when stickiness is off") + }) + + t.Run("enabled with no annotation sets sticky fields but injects no affinity", func(t *testing.T) { + ar := &runtimev1alpha1.AgentRuntime{ + ObjectMeta: metav1.ObjectMeta{Namespace: testNamespace, Name: testAgentRuntimeName}, + Spec: enabledSpec(), + } + sandbox, entry, err := buildSandboxByAgentRuntime(testNamespace, testAgentRuntimeName, "", newIfm(ar)) + assert.NoError(t, err) + assert.Nil(t, sandbox.Spec.PodTemplate.Spec.Affinity) + assert.Equal(t, testAgentRuntimeName, entry.StickyWorkloadName) + assert.Equal(t, runtimev1alpha1.AgentRuntimeKind, entry.StickyWorkloadKind) + }) + + t.Run("enabled with annotation injects preferred hostname affinity", func(t *testing.T) { + ar := &runtimev1alpha1.AgentRuntime{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: testNamespace, + Name: testAgentRuntimeName, + Annotations: map[string]string{LastNodeAnnotationKey: "sticky-node"}, + }, + Spec: enabledSpec(), + } + sandbox, entry, err := buildSandboxByAgentRuntime(testNamespace, testAgentRuntimeName, "", newIfm(ar)) + assert.NoError(t, err) + assert.Equal(t, testAgentRuntimeName, entry.StickyWorkloadName) + + affinity := sandbox.Spec.PodTemplate.Spec.Affinity + if affinity == nil || affinity.NodeAffinity == nil { + t.Fatal("expected node affinity to be injected") + } + preferred := affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution + if len(preferred) != 1 { + t.Fatalf("expected 1 preferred term, got %d", len(preferred)) + } + expr := preferred[0].Preference.MatchExpressions + if len(expr) != 1 || expr[0].Key != corev1.LabelHostname || expr[0].Values[0] != "sticky-node" { + t.Errorf("unexpected match expression: %+v", expr) + } + }) +} + func TestBuildSandboxByAgentRuntime_NotFound(t *testing.T) { fakeClient := cubefake.NewSimpleClientset() factory := cubeinformers.NewSharedInformerFactory(fakeClient, 0) @@ -686,6 +1012,110 @@ func TestBuildSandboxByCodeInterpreter_SuccessWithWarmPool(t *testing.T) { assertOwnerReference(t, claim.OwnerReferences[0]) } +func TestBuildSandboxByCodeInterpreter_NodeStickiness(t *testing.T) { + const ciName = "ci-sticky" + newIfm := func(ci *runtimev1alpha1.CodeInterpreter) *Informers { + fakeClient := cubefake.NewSimpleClientset(ci) + factory := cubeinformers.NewSharedInformerFactory(fakeClient, 0) + codeInterpreterInformer := factory.Runtime().V1alpha1().CodeInterpreters() + if err := codeInterpreterInformer.Informer().GetStore().Add(ci); err != nil { + t.Fatalf("failed to add code interpreter to informer store: %v", err) + } + return &Informers{ + CodeInterpreterLister: codeInterpreterInformer.Lister(), + CodeInterpreterInformer: codeInterpreterInformer.Informer(), + informerFactory: newFactory(), + cubeInformerFactory: factory, + } + } + + baseSpec := runtimev1alpha1.CodeInterpreterSpec{ + AuthMode: runtimev1alpha1.AuthModeNone, + Template: &runtimev1alpha1.CodeInterpreterSandboxTemplate{Image: "my-ci-image:latest"}, + } + enabledSpec := func() runtimev1alpha1.CodeInterpreterSpec { + s := baseSpec + s.NodeStickiness = &runtimev1alpha1.NodeStickinessSpec{Enabled: true} + return s + } + + t.Run("disabled by default does not inject affinity or set sticky fields", func(t *testing.T) { + ci := &runtimev1alpha1.CodeInterpreter{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: testNamespace, + Name: ciName, + Annotations: map[string]string{LastNodeAnnotationKey: "sticky-node"}, + }, + Spec: baseSpec, + } + sandbox, _, entry, err := buildSandboxByCodeInterpreter(testNamespace, ciName, "", newIfm(ci)) + assert.NoError(t, err) + assert.Nil(t, sandbox.Spec.PodTemplate.Spec.Affinity, "affinity must not be injected when stickiness is off") + assert.Empty(t, entry.StickyWorkloadName, "write-back must be disabled when stickiness is off") + }) + + t.Run("enabled with no annotation sets sticky fields but injects no affinity", func(t *testing.T) { + ci := &runtimev1alpha1.CodeInterpreter{ + ObjectMeta: metav1.ObjectMeta{Namespace: testNamespace, Name: ciName}, + Spec: enabledSpec(), + } + sandbox, _, entry, err := buildSandboxByCodeInterpreter(testNamespace, ciName, "", newIfm(ci)) + assert.NoError(t, err) + assert.Nil(t, sandbox.Spec.PodTemplate.Spec.Affinity) + assert.Equal(t, ciName, entry.StickyWorkloadName) + assert.Equal(t, runtimev1alpha1.CodeInterpreterKind, entry.StickyWorkloadKind) + }) + + t.Run("enabled with annotation injects preferred hostname affinity", func(t *testing.T) { + ci := &runtimev1alpha1.CodeInterpreter{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: testNamespace, + Name: ciName, + Annotations: map[string]string{LastNodeAnnotationKey: "sticky-node"}, + }, + Spec: enabledSpec(), + } + sandbox, _, entry, err := buildSandboxByCodeInterpreter(testNamespace, ciName, "", newIfm(ci)) + assert.NoError(t, err) + assert.Equal(t, ciName, entry.StickyWorkloadName) + assert.Equal(t, runtimev1alpha1.CodeInterpreterKind, entry.StickyWorkloadKind) + + affinity := sandbox.Spec.PodTemplate.Spec.Affinity + if affinity == nil || affinity.NodeAffinity == nil { + t.Fatal("expected node affinity to be injected") + } + preferred := affinity.NodeAffinity.PreferredDuringSchedulingIgnoredDuringExecution + if len(preferred) != 1 { + t.Fatalf("expected 1 preferred term, got %d", len(preferred)) + } + expr := preferred[0].Preference.MatchExpressions + if len(expr) != 1 || expr[0].Key != corev1.LabelHostname || expr[0].Values[0] != "sticky-node" { + t.Errorf("unexpected match expression: %+v", expr) + } + }) + + t.Run("warm pool path does not set sticky fields even when enabled", func(t *testing.T) { + ci := &runtimev1alpha1.CodeInterpreter{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: testNamespace, + Name: ciName, + Annotations: map[string]string{LastNodeAnnotationKey: "sticky-node"}, + }, + Spec: runtimev1alpha1.CodeInterpreterSpec{ + AuthMode: runtimev1alpha1.AuthModeNone, + WarmPoolSize: ptr.To(int32(3)), + Template: &runtimev1alpha1.CodeInterpreterSandboxTemplate{Image: "my-ci-image:latest"}, + NodeStickiness: &runtimev1alpha1.NodeStickinessSpec{Enabled: true}, + }, + } + sandbox, claim, entry, err := buildSandboxByCodeInterpreter(testNamespace, ciName, "", newIfm(ci)) + assert.NoError(t, err) + assert.NotNil(t, claim, "warm pool path should return a claim") + assert.Nil(t, sandbox.Spec.PodTemplate.Spec.Affinity, "warm pool sandbox has no inline pod spec to attach affinity") + assert.Empty(t, entry.StickyWorkloadName, "warm pool path cannot influence scheduling, so no write-back") + }) +} + func TestBuildSandboxObject_OwnershipLabels(t *testing.T) { tests := []struct { name string