fix(pd): re-read a recovered PD worker's KV engine id before it is promoted

model_gateway/src/routers/common/kv_transfer.rs

@@ -100,7 +100,7 @@ pub(crate) fn connector_mode_for_worker(worker: &dyn Worker) -> KvConnectorMode
         None
     } else {
         let dp_size = worker.dp_size().or(label_dp);
-        effective_kv_engine_id(meta.spec.kv_engine_id.as_deref(), dp_size, worker.dp_rank())
+        effective_kv_engine_id(worker.kv_engine_id().as_deref(), dp_size, worker.dp_rank())
     };
     kv_connector_mode(
         meta.spec.kv_connector.as_deref(),
@@ -256,4 +256,28 @@ mod tests {
         assert_eq!(effective_kv_engine_id(Some(""), None, None), None);
         assert_eq!(effective_kv_engine_id(None, Some(2), Some(0)), None);
     }
+
+    #[test]
+    fn a_refreshed_engine_id_is_what_the_handoff_is_minted_for() {
+        use crate::worker::{BasicWorkerBuilder, ConnectionMode, RuntimeType, WorkerType};
+
+        let worker = BasicWorkerBuilder::new("grpc://p:1")
+            .worker_type(WorkerType::Prefill)
+            .connection_mode(ConnectionMode::Grpc)
+            .runtime_type(RuntimeType::Vllm)
+            .kv_connector("MooncakeConnector")
+            .kv_engine_id("eng-old")
+            .bootstrap_port(Some(8998))
+            .build();
+        assert!(matches!(
+            connector_mode_for_worker(&worker),
+            KvConnectorMode::Mooncake { engine_id: Some(ref id), .. } if id == "eng-old"
+        ));
+        assert!(worker.refresh_kv_engine_id(Some("eng-new".to_string())));
+        assert!(!worker.refresh_kv_engine_id(Some("eng-new".to_string())));
+        assert!(matches!(
+            connector_mode_for_worker(&worker),
+            KvConnectorMode::Mooncake { engine_id: Some(ref id), .. } if id == "eng-new"
+        ));
+    }
 }

model_gateway/src/worker/builder.rs

@@ -363,6 +363,7 @@ impl BasicWorkerBuilder {
         let resilience = self.resilience.unwrap_or_default();
 
         BasicWorker {
+            kv_engine_id: ArcSwapOption::new(metadata.spec.kv_engine_id.clone().map(Arc::new)),
             runtime: ArcSwap::from_pointee(WorkerRuntime::new(&metadata.spec.url, initial_status)),
             circuit_breaker: ArcSwap::from_pointee(CircuitBreaker::with_config_and_label(
                 self.circuit_breaker_config,

model_gateway/src/worker/manager.rs

@@ -32,9 +32,9 @@ use crate::{
         metrics_aggregator::{self, MetricPack},
         registry::{WorkerDescriptor, WorkerId},
         worker::WorkerTypeExt,
-        ConnectionMode, Worker, WorkerOrigin, WorkerRegistry, WorkerResult,
+        ConnectionMode, Worker, WorkerOrigin, WorkerRegistry, WorkerResult, WorkerType,
     },
-    workflow::{Job, JobQueue},
+    workflow::{steps::local::discover_grpc_kv_engine_id, Job, JobQueue},
 };
 
 const REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
@@ -492,6 +492,14 @@ fn queue_due_probes(
         in_flight.insert(worker_id.clone());
         probes.push(Box::pin(async move {
             let probe_result = worker.check_health_async().await;
+            if probe_result.is_ok()
+                && matches!(
+                    launched_status,
+                    WorkerStatus::Failed | WorkerStatus::NotReady
+                )
+            {
+                refresh_kv_engine_id_after_recovery(&worker).await;
+            }
             ProbeCompletion {
                 worker_id,
                 worker,
@@ -713,6 +721,39 @@ fn schedule_worker_at(
 ///     came back on the same address (a restart), and without
 ///     `--remove-unhealthy-workers` this is the only way it rejoins
 ///   - Failed stays Failed on failure; Draining never transitions here
+/// A PD worker that answers its probe again after failing may be a new engine
+/// process on the same address, with a new KV transfer engine id (#2491).
+/// Re-read the id before the worker is promoted, so the next handoff is
+/// minted for the engine that is actually there. Only gRPC engines report
+/// the id; a read that fails keeps the previous one and says so.
+async fn refresh_kv_engine_id_after_recovery(worker: &Arc<dyn Worker>) {
+    let spec = &worker.metadata().spec;
+    if !matches!(spec.worker_type, WorkerType::Prefill | WorkerType::Decode)
+        || *worker.connection_mode() != ConnectionMode::Grpc
+        || spec.kv_connector.is_none()
+    {
+        return;
+    }
+    match discover_grpc_kv_engine_id(worker.url(), spec.runtime_type.as_str()).await {
+        Ok(discovered) => {
+            let previous = worker.kv_engine_id();
+            if worker.refresh_kv_engine_id(discovered.clone()) {
+                info!(
+                    worker_url = %worker.url(),
+                    ?previous,
+                    ?discovered,
+                    "Recovered PD worker reports a new KV engine id"
+                );
+            }
+        }
+        Err(error) => warn!(
+            worker_url = %worker.url(),
+            %error,
+            "Could not re-read the KV engine id of a recovered PD worker; keeping the previous one"
+        ),
+    }
+}
+
 fn compute_next_status(
     worker: &Arc<dyn Worker>,
     probe_ok: bool,

model_gateway/src/worker/worker.rs

@@ -616,6 +616,20 @@ pub trait Worker: Send + Sync + fmt::Debug + 'static {
         self.metadata().dp_size()
     }
 
+    /// The KV transfer engine id the worker's engine currently reports: the
+    /// spec's, unless a recovery re-read it (see
+    /// [`Self::refresh_kv_engine_id`]). A PD handoff must be minted for the
+    /// engine process that is there now, not the one discovered at
+    /// registration (#2491).
+    fn kv_engine_id(&self) -> Option<String> {
+        self.metadata().spec.kv_engine_id.clone()
+    }
+
+    /// Replace the engine id after a re-read. Returns `true` when it changed.
+    fn refresh_kv_engine_id(&self, _kv_engine_id: Option<String>) -> bool {
+        false
+    }
+
     /// Transform a request for DP-aware routing.
     ///
     /// When the worker has a `dp_rank`, injects `data_parallel_rank`
@@ -1282,6 +1296,11 @@ pub struct BasicWorker {
     /// When not `Wildcard`, overrides metadata.models for routing decisions.
     /// Uses `ArcSwap` for lock-free reads on the hot path (`supports_model`).
     pub models_override: Arc<ArcSwap<WorkerModels>>,
+    /// The KV transfer engine id in force, seeded from the spec and replaced
+    /// when a recovered engine reports a new one (see
+    /// [`Worker::refresh_kv_engine_id`]). Not shared across same-URL
+    /// replacements: a replacement is built from a fresh discovery.
+    pub kv_engine_id: ArcSwapOption<String>,
     /// Worker-directed HTTP client, shared across same-config workers, built
     /// on first use (see [`LazyHttpClient`]).
     pub http_client: Arc<LazyHttpClient>,
@@ -1300,6 +1319,7 @@ impl Clone for BasicWorker {
             zmq_connect_abort: Arc::clone(&self.zmq_connect_abort),
             connect_signal_tx: self.connect_signal_tx.clone(),
             models_override: Arc::clone(&self.models_override),
+            kv_engine_id: ArcSwapOption::new(self.kv_engine_id.load_full()),
             http_client: Arc::clone(&self.http_client),
             resilience: self.resilience.clone(),
         }
@@ -1435,6 +1455,19 @@ impl BasicWorker {
 
 #[async_trait]
 impl Worker for BasicWorker {
+    fn kv_engine_id(&self) -> Option<String> {
+        self.kv_engine_id.load_full().map(|id| (*id).clone())
+    }
+
+    fn refresh_kv_engine_id(&self, kv_engine_id: Option<String>) -> bool {
+        let previous = self.kv_engine_id.load_full();
+        if previous.as_deref() == kv_engine_id.as_ref() {
+            return false;
+        }
+        self.kv_engine_id.store(kv_engine_id.map(Arc::new));
+        true
+    }
+
     fn as_any(&self) -> &dyn Any {
         self
     }

model_gateway/src/workflow/steps/local/discover_metadata.rs

@@ -256,6 +256,17 @@ async fn fetch_vllm_http_metadata(
     labels
 }
 
+/// Re-read the KV transfer engine id a gRPC engine reports, for a PD worker
+/// that came back on the same address (#2491): a restarted engine process
+/// carries a new id, and a handoff minted for the old one strands the decode.
+pub(crate) async fn discover_grpc_kv_engine_id(
+    url: &str,
+    runtime_type: &str,
+) -> Result<Option<String>, String> {
+    let (mut labels, _) = fetch_grpc_metadata(url, runtime_type).await?;
+    Ok(labels.remove("kv_engine_id").filter(|id| !id.is_empty()))
+}
+
 async fn fetch_grpc_metadata(
     url: &str,
     runtime_type: &str,

model_gateway/src/workflow/steps/local/mod.rs

@@ -3,6 +3,7 @@ mod detect_backend;
 mod detect_connection;
 mod discover_dp;
 mod discover_metadata;
+pub(crate) use discover_metadata::discover_grpc_kv_engine_id;
 mod drain_workers;
 mod ensure_harmony_encoding;
 mod find_worker_to_update;