feat(model-serving): add preStop-based lossless upgrade for ModelServing

examples/model-serving/lossless-upgrade-pd-disaggregation.yaml

@@ -0,0 +1,51 @@
+# Lossless upgrade example: PD disaggregated ModelServing.
+#
+# The drain-timeout annotation on the ModelServing enables the lossless drain.
+# The drain-kv-buffer annotation on the prefill role's entryTemplate reserves
+# up to 30s after the prefill drains its own in-flight requests, so its paired
+# decodes can finish pulling the last KV (SGLang push / NIXL-Mooncake pull).
+# The decode role carries no buffer annotation: its metrics are reliable and it
+# exits as soon as its own in-flight requests finish.
+apiVersion: workload.serving.volcano.sh/v1alpha1
+kind: ModelServing
+metadata:
+  name: deepseek-v2-lite-lossless
+  namespace: dev
+  annotations:
+    modelserving.volcano.sh/drain-timeout: "300s"
+spec:
+  schedulerName: volcano
+  replicas: 1
+  recoveryPolicy: ServingGroupRecreate
+  template:
+    restartGracePeriodSeconds: 60
+    roles:
+      - name: prefill
+        replicas: 1
+        entryTemplate:
+          metadata:
+            annotations:
+              # PD prefill only: hold a KV buffer so paired decodes finish
+              # pulling before the prefill is terminated.
+              modelserving.volcano.sh/drain-kv-buffer: "30s"
+          spec:
+            containers:
+              - name: vllm
+                image: ghcr.io/volcano-sh/kthena-engine:vllm-ascend_v0.10.1rc1_mooncake_v0.3.5
+                ports:
+                  - containerPort: 8000
+                env:
+                  - name: HF_HUB_OFFLINE
+                    value: "1"
+      - name: decode
+        replicas: 1
+        entryTemplate:
+          spec:
+            containers:
+              - name: vllm
+                image: ghcr.io/volcano-sh/kthena-engine:vllm-ascend_v0.10.1rc1_mooncake_v0.3.5
+                ports:
+                  - containerPort: 8000
+                env:
+                  - name: HF_HUB_OFFLINE
+                    value: "1"

examples/model-serving/lossless-upgrade.yaml

@@ -0,0 +1,31 @@
+# Lossless upgrade example: aggregated (non-PD) ModelServing.
+#
+# The annotation below enables the lossless drain: on deletion the router stops
+# dispatching new requests to the Pod and its preStop hook polls the engine
+# metrics until in-flight reaches zero (bounded by min(drain-timeout, 70% of
+# the grace period)) before SIGTERM is sent. Remove the annotation to restore
+# the default immediate-delete behavior.
+# A single-role aggregated deployment needs no drain-kv-buffer annotation: the
+# engine serves both prefill and decode, and its own counters are reliable.
+apiVersion: workload.serving.volcano.sh/v1alpha1
+kind: ModelServing
+metadata:
+  name: sample-lossless
+  namespace: default
+  annotations:
+    modelserving.volcano.sh/drain-timeout: "300s"
+spec:
+  schedulerName: volcano
+  replicas: 1
+  template:
+    restartGracePeriodSeconds: 60
+    roles:
+      - name: inference          # single role: one Pod runs the full engine (aggregated)
+        replicas: 1
+        entryTemplate:
+          spec:
+            containers:
+              - name: vllm
+                image: ghcr.io/volcano-sh/kthena-engine:vllm-ascend_v0.10.1rc1_mooncake_v0.3.5
+                ports:
+                  - containerPort: 8000    # serving + metrics port, read by the preStop hook

pkg/apis/workload/v1alpha1/labels.go

@@ -35,4 +35,13 @@ const (
 
 	// InPlaceUpdateStateAnnotationKey is the pod annotation that records in-place update state.
 	InPlaceUpdateStateAnnotationKey = "modelserving.volcano.sh/inplace-update-state"
+
+	// DrainTimeoutAnnotationKey is the ModelServing annotation that enables lossless drain.
+	// The value is a duration (e.g. "300s") bounding how long a terminating Pod waits for
+	// its in-flight requests to finish. Missing, "0" or invalid values keep the default
+	// immediate-delete behavior.
+	DrainTimeoutAnnotationKey = "modelserving.volcano.sh/drain-timeout"
+	// DrainKVBufferAnnotationKey is the role entryTemplate annotation that reserves extra
+	// time after in-flight drains for paired decodes to pull the last KV (PD prefill Pods).
+	DrainKVBufferAnnotationKey = "modelserving.volcano.sh/drain-kv-buffer"
 )

pkg/kthena-router/controller/modelserver_controller.go

@@ -273,6 +273,13 @@ func (c *ModelServerController) syncPodHandler(key string) error {
 		return err
 	}
 
+	// Terminating Pods stay tracked (in-flight metrics keep being scraped) and
+	// are excluded from candidates by the store queries; the store is cleared
+	// by the object-delete event above.
+	if pod.DeletionTimestamp != nil {
+		return c.addOrUpdatePod(pod)
+	}
+
 	if !isPodReady(pod) {
 		_ = c.store.DeletePod(types.NamespacedName{Namespace: namespace, Name: name})
 		return nil

pkg/kthena-router/controller/modelserver_controller_test.go

@@ -1189,3 +1189,84 @@ func waitForObjectInCache(t *testing.T, timeout time.Duration, checkFunc func()
 		}
 	}
 }
+
+// TestModelServerController_TerminatingPodTracked verifies that a Pod with a
+// deletionTimestamp stays in the store (tracked for the drain) but is excluded
+// from scheduling candidates, and is cleared on object deletion.
+func TestModelServerController_TerminatingPodTracked(t *testing.T) {
+	kubeClient := kubefake.NewSimpleClientset()
+	kthenaClient := kthenafake.NewSimpleClientset()
+
+	ms := &aiv1alpha1.ModelServer{
+		ObjectMeta: metav1.ObjectMeta{
+			Namespace: "default",
+			Name:      "test-modelserver-terminating",
+		},
+		Spec: aiv1alpha1.ModelServerSpec{
+			InferenceEngine: aiv1alpha1.VLLM,
+			WorkloadSelector: &aiv1alpha1.WorkloadSelector{
+				MatchLabels: map[string]string{"app": "terminating-test"},
+			},
+		},
+	}
+	_, err := kthenaClient.NetworkingV1alpha1().ModelServers("default").Create(
+		context.Background(), ms, metav1.CreateOptions{})
+	require.NoError(t, err)
+
+	kubeInformerFactory := informers.NewSharedInformerFactory(kubeClient, 0)
+	kthenaInformerFactory := informersv1alpha1.NewSharedInformerFactory(kthenaClient, 0)
+	store := newStoreWithMockBackend()
+	controller, err := NewModelServerController(
+		kthenaInformerFactory, kubeInformerFactory, store, common.NewTransportRegistry())
+	require.NoError(t, err)
+
+	stop := make(chan struct{})
+	defer close(stop)
+	go controller.Run(stop)
+	kthenaInformerFactory.Start(stop)
+	kubeInformerFactory.Start(stop)
+	waitForCacheSync(t, 5*time.Second, controller.modelServerSynced, controller.podSynced)
+
+	pod := &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Namespace: "default",
+			Name:      "terminating-pod",
+			Labels:    map[string]string{"app": "terminating-test"},
+		},
+		Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "engine"}}},
+		Status: corev1.PodStatus{Phase: corev1.PodRunning,
+			Conditions: []corev1.PodCondition{{Type: corev1.PodReady, Status: corev1.ConditionTrue}}},
+	}
+	_, err = kubeClient.CoreV1().Pods("default").Create(context.Background(), pod, metav1.CreateOptions{})
+	require.NoError(t, err)
+
+	msName := types.NamespacedName{Namespace: "default", Name: ms.Name}
+	podName := types.NamespacedName{Namespace: "default", Name: pod.Name}
+	assert.True(t, waitForObjectInCache(t, 5*time.Second, func() bool {
+		pods, err := store.GetPodsByModelServer(msName)
+		return err == nil && len(pods) == 1
+	}), "ready pod should be a scheduling candidate")
+
+	// Mark the Pod terminating: it must stay tracked but leave the candidates.
+	now := metav1.Now()
+	pod.DeletionTimestamp = &now
+	_, err = kubeClient.CoreV1().Pods("default").Update(context.Background(), pod, metav1.UpdateOptions{})
+	require.NoError(t, err)
+
+	assert.True(t, waitForObjectInCache(t, 5*time.Second, func() bool {
+		pods, err := store.GetPodsByModelServer(msName)
+		return err == nil && len(pods) == 0
+	}), "terminating pod must be excluded from scheduling candidates")
+
+	assert.True(t, waitForObjectInCache(t, 5*time.Second, func() bool {
+		info := store.GetPodInfo(podName)
+		return info != nil && info.IsDraining()
+	}), "terminating pod must stay tracked in the store for the drain")
+
+	// Object deletion clears the store entry.
+	err = kubeClient.CoreV1().Pods("default").Delete(context.Background(), pod.Name, metav1.DeleteOptions{})
+	require.NoError(t, err)
+	assert.True(t, waitForObjectInCache(t, 5*time.Second, func() bool {
+		return store.GetPodInfo(podName) == nil
+	}), "deleted pod must be removed from the store")
+}

pkg/kthena-router/datastore/store.go

@@ -1020,7 +1020,11 @@ func (s *store) GetPodsByModelServer(name types.NamespacedName) ([]*PodInfo, err
 
 	for _, podName := range podNames {
 		if value, ok := s.pods.Load(podName); ok {
-			pods = append(pods, value.(*PodInfo))
+			pi := value.(*PodInfo)
+			if pi.IsDraining() {
+				continue
+			}
+			pods = append(pods, pi)
 		}
 	}
 
@@ -1085,7 +1089,11 @@ func (s *store) GetDecodePods(modelServerName types.NamespacedName) ([]*PodInfo,
 
 	for _, podName := range decodePodNames {
 		if value, ok := s.pods.Load(podName); ok {
-			decodePods = append(decodePods, value.(*PodInfo))
+			pi := value.(*PodInfo)
+			if pi.IsDraining() {
+				continue
+			}
+			decodePods = append(decodePods, pi)
 		}
 	}
 
@@ -1105,7 +1113,11 @@ func (s *store) GetPrefillPods(modelServerName types.NamespacedName) ([]*PodInfo
 
 	for _, podName := range prefillPodNames {
 		if value, ok := s.pods.Load(podName); ok {
-			prefillPods = append(prefillPods, value.(*PodInfo))
+			pi := value.(*PodInfo)
+			if pi.IsDraining() {
+				continue
+			}
+			prefillPods = append(prefillPods, pi)
 		}
 	}
 
@@ -1130,7 +1142,11 @@ func (s *store) GetPrefillPodsForDecodeGroup(modelServerName types.NamespacedNam
 	prefillPods := make([]*PodInfo, 0, len(prefillPodNames))
 	for _, podName := range prefillPodNames {
 		if value, ok := s.pods.Load(podName); ok {
-			prefillPods = append(prefillPods, value.(*PodInfo))
+			pi := value.(*PodInfo)
+			if pi.IsDraining() {
+				continue
+			}
+			prefillPods = append(prefillPods, pi)
 		}
 	}
 
@@ -2128,6 +2144,13 @@ func (p *PodInfo) GetRequestRunningNum() float64 {
 	return p.RequestRunningNum
 }
 
+// IsDraining reports whether the Pod is terminating; draining Pods stay
+// tracked but are excluded from scheduling candidates.
+func (p *PodInfo) IsDraining() bool {
+	pod := p.GetPod()
+	return pod != nil && pod.DeletionTimestamp != nil
+}
+
 // IncrOnFlightRequests atomically increments the local in-flight counter and
 // returns the new value.
 func (p *PodInfo) IncrOnFlightRequests() int64 {
@@ -2455,7 +2478,7 @@ func (s *store) GetPodsByInferencePool(name types.NamespacedName) ([]*PodInfo, e
 	s.pods.Range(func(key, value interface{}) bool {
 		podInfo := value.(*PodInfo)
 		pod := podInfo.GetPod()
-		if pod != nil && pod.Namespace == name.Namespace && selector.Matches(labels.Set(podInfo.GetPodLabels())) {
+		if pod != nil && pod.Namespace == name.Namespace && !podInfo.IsDraining() && selector.Matches(labels.Set(podInfo.GetPodLabels())) {
 			pods = append(pods, podInfo)
 		}
 		return true

pkg/model-serving-controller/controller/model_serving_controller.go

@@ -2993,6 +2993,9 @@ func (c *ModelServingController) runRoleDeletePlugins(ctx context.Context, ms *w
 
 func (c *ModelServingController) CreatePodsForServingGroup(ctx context.Context, ms *workloadv1alpha1.ModelServing, servingGroupIndex int, revision string, roles []workloadv1alpha1.Role) error {
 	servingGroupName := utils.GenerateServingGroupName(ms.Name, servingGroupIndex)
+	if cfg := utils.ParseDrainConfig(ms); cfg.Invalid {
+		c.recorder.Event(ms, corev1.EventTypeWarning, "InvalidDrainTimeout", cfg.Reason)
+	}
 	for _, role := range roles {
 		roleTemplateHash := utils.CalRoleTemplateHash(role)
 		replicas := int(*role.Replicas)