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,8 @@ impl BasicWorkerBuilder {
         let resilience = self.resilience.unwrap_or_default();
 
         BasicWorker {
+            kv_engine_id: ArcSwapOption::new(metadata.spec.kv_engine_id.clone().map(Arc::new)),
+            kv_engine_id_unconfirmed: AtomicBool::new(false),
             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,10 @@ 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() && engine_id_may_be_stale(launched_status, worker.as_ref()) {
+                let timeout = Duration::from_secs(health_config.timeout_secs.max(1));
+                refresh_kv_engine_id_after_recovery(&worker, timeout).await;
+            }
             ProbeCompletion {
                 worker_id,
                 worker,
@@ -700,6 +704,87 @@ fn schedule_worker_at(
     );
 }
 
+/// Whether a probe launched at `launched_status` that succeeds should re-read
+/// the worker's KV engine id: after a failure (the engine may be a new
+/// process), while a worker that registered before its engine reported an id
+/// is still pending, or while a Ready worker's last recovery re-read has not
+/// been confirmed. A Draining worker, or a pending one that already has an
+/// id, is the same process registration read.
+fn engine_id_may_be_stale(launched_status: WorkerStatus, worker: &dyn Worker) -> bool {
+    match launched_status {
+        WorkerStatus::Failed | WorkerStatus::NotReady => true,
+        WorkerStatus::Pending => worker.kv_engine_id().is_none(),
+        WorkerStatus::Ready => !worker.kv_engine_id_confirmed(),
+        WorkerStatus::Draining => false,
+    }
+}
+
+/// A PD worker that answers its probe again may be a new engine process on
+/// the same address, with a new KV transfer engine id (#2491); a worker
+/// registered while its engine was still coming up may have none at all.
+/// 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. The read is bounded by the probe timeout: one that fails, expires
+/// or comes back without an id keeps the previous id and says so, since a
+/// stale id is a better outcome than a probe slot held forever or a handoff
+/// with no id at all. An id the engine does report wins over the spec's: it
+/// is the one a handoff must target.
+async fn refresh_kv_engine_id_after_recovery(worker: &Arc<dyn Worker>, timeout: Duration) {
+    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;
+    }
+    let read = tokio::time::timeout(
+        timeout,
+        discover_grpc_kv_engine_id(worker.url(), spec.runtime_type.as_str()),
+    )
+    .await;
+    // Anything but a confirmed id leaves the flag set, so the next probe
+    // retries instead of the worker serving with a possibly stale id until
+    // its next outage.
+    worker.set_kv_engine_id_confirmed(false);
+    match read {
+        Ok(Ok(Some(discovered))) => {
+            let previous = worker.kv_engine_id();
+            if worker.refresh_kv_engine_id(Some(discovered.clone())) {
+                info!(
+                    worker_url = %worker.url(),
+                    ?previous,
+                    discovered,
+                    "Recovered PD worker reports a new KV engine id"
+                );
+            }
+            worker.set_kv_engine_id_confirmed(true);
+        }
+        // A partial read (server info tolerated as missing) or an engine
+        // that reports no id: a known id must never be cleared by it.
+        Ok(Ok(None)) => match worker.kv_engine_id() {
+            Some(previous) => warn!(
+                worker_url = %worker.url(),
+                previous,
+                "Recovered PD worker reported no KV engine id; keeping the previous one until it does"
+            ),
+            None => warn!(
+                worker_url = %worker.url(),
+                "PD worker reports no KV engine id yet; handoffs carry none until it does"
+            ),
+        },
+        Ok(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 and retrying on the next probe"
+        ),
+        Err(_) => warn!(
+            worker_url = %worker.url(),
+            timeout_secs = timeout.as_secs(),
+            "Re-reading the KV engine id of a recovered PD worker timed out; keeping the previous one and retrying on the next probe"
+        ),
+    }
+}
+
 /// Apply the state machine to a probe outcome. Returns the next status if
 /// a transition is needed, `None` if the worker stays in its current state.
 ///
@@ -1170,14 +1255,100 @@ mod tests {
         },
     };
 
+    /// The re-read gate: after a failure always, while pending only without
+    /// an id, never on a Ready or Draining probe.
+    #[test]
+    fn the_engine_id_is_re_read_after_a_failure_or_while_pending_without_one() {
+        let worker = |kv_engine_id: Option<&str>| -> Arc<dyn Worker> {
+            let mut builder = BasicWorkerBuilder::new("grpc://p:1")
+                .worker_type(WorkerType::Prefill)
+                .connection_mode(ConnectionMode::Grpc)
+                .runtime_type(RuntimeType::Vllm)
+                .kv_connector("MooncakeConnector");
+            if let Some(id) = kv_engine_id {
+                builder = builder.kv_engine_id(id);
+            }
+            Arc::new(builder.build())
+        };
+        let with_id = worker(Some("eng"));
+        let without_id = worker(None);
+        let unconfirmed = worker(Some("eng"));
+        unconfirmed.set_kv_engine_id_confirmed(false);
+        let cases = [
+            (WorkerStatus::Failed, &with_id, true),
+            (WorkerStatus::NotReady, &with_id, true),
+            (WorkerStatus::Pending, &with_id, false),
+            (WorkerStatus::Pending, &without_id, true),
+            (WorkerStatus::Ready, &with_id, false),
+            (WorkerStatus::Ready, &without_id, false),
+            (WorkerStatus::Ready, &unconfirmed, true),
+            (WorkerStatus::Draining, &unconfirmed, false),
+        ];
+        for (status, worker, expected) in cases {
+            assert_eq!(
+                engine_id_may_be_stale(status, worker.as_ref()),
+                expected,
+                "{status:?} with id {:?}",
+                worker.kv_engine_id()
+            );
+        }
+    }
+
+    /// A recovered worker whose engine accepts the connection but never
+    /// answers the metadata read must not hold the probe slot: the read
+    /// expires and the previous engine id stays (#2491).
+    #[tokio::test]
+    #[expect(
+        clippy::disallowed_methods,
+        reason = "the held connections live for the duration of the test"
+    )]
+    async fn a_hanging_engine_id_read_expires_and_keeps_the_previous_id() {
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("bind a loopback listener");
+        let port = listener.local_addr().expect("listener address").port();
+        // Accept every connection and never answer on it.
+        let hold = tokio::spawn(async move {
+            let mut held = Vec::new();
+            loop {
+                if let Ok((socket, _)) = listener.accept().await {
+                    held.push(socket);
+                }
+            }
+        });
+        let worker: Arc<dyn Worker> = Arc::new(
+            BasicWorkerBuilder::new(format!("grpc://127.0.0.1:{port}"))
+                .worker_type(WorkerType::Prefill)
+                .connection_mode(ConnectionMode::Grpc)
+                .runtime_type(RuntimeType::Vllm)
+                .kv_connector("MooncakeConnector")
+                .kv_engine_id("eng-old")
+                .build(),
+        );
+
+        let started = std::time::Instant::now();
+        refresh_kv_engine_id_after_recovery(&worker, Duration::from_millis(300)).await;
+
+        assert!(
+            started.elapsed() < Duration::from_secs(5),
+            "the re-read must be bounded by the timeout"
+        );
+        assert_eq!(worker.kv_engine_id().as_deref(), Some("eng-old"));
+        // Not confirmed: a later Ready probe re-reads instead of giving up.
+        assert!(!worker.kv_engine_id_confirmed());
+        assert!(engine_id_may_be_stale(WorkerStatus::Ready, worker.as_ref()));
+        hold.abort();
+    }
+
     use openai_protocol::{
         model_card::ModelCard,
         worker::{HealthCheckConfig, WorkerStatus},
     };
 
     use super::*;
     use crate::worker::{
-        BasicWorkerBuilder, ConnectionMode, Worker, WorkerError, WorkerRegistry, WorkerType,
+        BasicWorkerBuilder, ConnectionMode, RuntimeType, Worker, WorkerError, WorkerRegistry,
+        WorkerType,
     };
 
     fn make_worker(url: &str, success_threshold: u32, failure_threshold: u32) -> Arc<dyn Worker> {

model_gateway/src/worker/worker.rs

@@ -616,6 +616,31 @@ 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
+    }
+
+    /// Whether the engine id in force has been confirmed by the engine since
+    /// the worker last recovered. `false` after a re-read that failed,
+    /// expired or returned no id, so the next probe tries again instead of
+    /// the worker serving with a possibly stale id until its next outage.
+    fn kv_engine_id_confirmed(&self) -> bool {
+        true
+    }
+
+    /// Record the outcome of a re-read (see [`Self::kv_engine_id_confirmed`]).
+    fn set_kv_engine_id_confirmed(&self, _confirmed: bool) {}
+
     /// Transform a request for DP-aware routing.
     ///
     /// When the worker has a `dp_rank`, injects `data_parallel_rank`
@@ -1282,6 +1307,16 @@ 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 worker built from a fresh discovery starts from its
+    /// own spec, and one rebuilt by a properties update starts from that
+    /// spec's (unrefreshed) id.
+    pub kv_engine_id: ArcSwapOption<String>,
+    /// Set while a recovery re-read of the engine id has not succeeded yet
+    /// (see [`Worker::kv_engine_id_confirmed`]).
+    pub kv_engine_id_unconfirmed: AtomicBool,
     /// Worker-directed HTTP client, shared across same-config workers, built
     /// on first use (see [`LazyHttpClient`]).
     pub http_client: Arc<LazyHttpClient>,
@@ -1300,6 +1335,10 @@ 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()),
+            kv_engine_id_unconfirmed: AtomicBool::new(
+                self.kv_engine_id_unconfirmed.load(Ordering::Relaxed),
+            ),
             http_client: Arc::clone(&self.http_client),
             resilience: self.resilience.clone(),
         }
@@ -1435,6 +1474,28 @@ 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 kv_engine_id_confirmed(&self) -> bool {
+        !self.kv_engine_id_unconfirmed.load(Ordering::Relaxed)
+    }
+
+    fn set_kv_engine_id_confirmed(&self, confirmed: bool) {
+        self.kv_engine_id_unconfirmed
+            .store(!confirmed, Ordering::Relaxed);
+    }
+
     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;