feat: add soft node-affinity stickiness for AgentRuntime and CodeInte…

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.

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.

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"]

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.

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
 

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,49 @@ 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 the shutdown lock while checking the flag and adding to the wait
+	// group. This guarantees that Shutdown() cannot observe a false flag and
+	// then begin wg.Wait between the check and the wg.Add.
+	s.shutdownMu.Lock()
+	if s.shuttingDown {
+		s.shutdownMu.Unlock()
+		return
+	}
+	s.wg.Add(1)
+	s.shutdownMu.Unlock()
+
+	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.

pkg/workloadmanager/handlers_test.go

@@ -243,11 +243,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 +386,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 +728,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 +838,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 +896,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 +1024,23 @@ func TestHandleSandboxCreate_IdentityErrors(t *testing.T) {
 		})
 	}
 }
+
+// TestRecordStickyNode verifies the best-effort write-back guard clauses.
+// recordStickyNode now runs the patch in a goroutine (best-effort async), so
+// these tests verify that the guard clauses prevent a goroutine from being
+// spawned when preconditions are not met.
+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, "")
+	})
+}