[Rust Frontend][gRPC] Separate liveness from inference readiness

docs/deployment/k8s.md

@@ -431,6 +431,43 @@ You can also verify the health service manually with `grpcurl`:
 grpcurl -plaintext localhost:50051 grpc.health.v1.Health/Check
 ```
 
+### Rust frontend gRPC health checks
+
+The Rust frontend is a separate server from `--grpc` above. Start it with `VLLM_USE_RUST_FRONTEND=1 vllm serve <model> --grpc-port 50051` or `vllm-rs serve <model> --grpc-port 50051`. Its gRPC port serves `grpc.health.v1.Health` and reports engine liveness and inference readiness under separate service names, so a paused engine can leave the load balancer without being restarted:
+
+| Service | Reports `NOT_SERVING` when |
+| --- | --- |
+| `liveness` | The engine has failed. Pause, sleep, weight updates, and orderly shutdown do not affect it until the listener closes. |
+| `readiness` and `vllm.Inference` | The engine has failed, shutdown has started, the engine is paused (including for a process checkpoint), asleep, or partially awake, or a weight update has not finished successfully. |
+| `vllm.Control` and `""` | The engine has failed or shutdown has started. |
+
+Point the startup and liveness probes at `liveness` and the readiness probe at `readiness`. A liveness probe on `readiness` or `vllm.Inference` restarts an engine that stays paused longer than the probe's failure threshold. The listener stops accepting connections when shutdown starts, so probes see the endpoint as unavailable while `--shutdown-timeout` drains in-flight requests.
+
+```yaml
+startupProbe:
+  grpc:
+    port: 50051
+    service: liveness
+  periodSeconds: 10
+  failureThreshold: 120 # Allow for model loading.
+livenessProbe:
+  grpc:
+    port: 50051
+    service: liveness
+readinessProbe:
+  grpc:
+    port: 50051
+    service: readiness
+```
+
+A paused or sleeping pod leaves the endpoints of its Services, so send resume, wake-up, and weight-update calls to the pod directly, for example through a headless Service with `publishNotReadyAddresses: true`. After a process checkpoint fails, the engine stays paused and `liveness` stays `SERVING`, so whatever drives the checkpoint must restart the pod when `/checkpoint/status` reports the failure.
+
+This server does not enable gRPC reflection, so check a service manually with [`grpc_health_probe`](https://github.com/grpc-ecosystem/grpc-health-probe):
+
+```bash
+grpc_health_probe -addr=localhost:50051 -service=readiness
+```
+
 ## Troubleshooting
 
 ### Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated"

rust/README.md

@@ -64,6 +64,12 @@ SMG `VllmEngine` gRPC server in place of HTTP. The Rust listener exposes the
 `vllm.Inference` and `vllm.Control` services alongside HTTP. These are different
 gRPC APIs; `--grpc` and `--grpc-port` are mutually exclusive.
 
+### gRPC health checks
+
+The gRPC listener also serves `grpc.health.v1.Health`. Kubernetes probes should
+use the `liveness` service for liveness and `readiness` for readiness; see
+[Rust frontend gRPC health checks](../docs/deployment/k8s.md#rust-frontend-grpc-health-checks).
+
 ### RL weight synchronization
 
 With `VLLM_SERVER_DEV_MODE=1`, the Rust frontend supports the HTTP weight-transfer

rust/src/engine-core-client/src/client.rs

@@ -15,6 +15,7 @@ use tokio_util::task::AbortOnDropHandle;
 use tracing::{debug, info, trace};
 
 use crate::client::imp::{ClientInner, run_abort_loop, run_output_dispatcher_loop};
+use crate::client::readiness::{RecheckGuard, run_readiness_loop};
 use crate::coordinator::CoordinatorHandle;
 use crate::error::{Error, Result, bail_invalid_client_config};
 use crate::metrics::EngineSleepState;
@@ -29,6 +30,7 @@ use crate::transport::{self, ConnectedEngine};
 
 pub(crate) mod imp;
 mod output_channel;
+mod readiness;
 mod state;
 mod stream;
 
@@ -328,6 +330,7 @@ pub struct EngineCoreClient {
     output_task: AbortOnDropHandle<()>,
     dispatcher_task: AbortOnDropHandle<()>,
     abort_task: AbortOnDropHandle<()>,
+    readiness_task: AbortOnDropHandle<()>,
     coordinator_output_task: Option<AbortOnDropHandle<()>>,
     coordinator_task: Option<AbortOnDropHandle<()>>,
 }
@@ -440,6 +443,8 @@ impl EngineCoreClient {
         );
         let abort_task =
             AbortOnDropHandle::new(runtime.spawn(run_abort_loop(inner.clone(), abort_rx)));
+        let readiness_task =
+            AbortOnDropHandle::new(runtime.spawn(run_readiness_loop(inner.clone())));
 
         // If any engine reported a dp_stats_address in its ready response, use it
         // as the external coordinator address.
@@ -480,6 +485,14 @@ impl EngineCoreClient {
                 (None, None, None)
             };
 
+        // Only engines with sleep mode can start paused, prepared for a
+        // checkpoint. Asking here also fails startup with a clear error when
+        // the engines predate the `is_inference_ready` utility.
+        if engines.iter().all(|engine| engine.ready_response.enable_sleep_mode) {
+            let ready = inner.is_inference_ready().await?;
+            inner.readiness().start(ready);
+        }
+
         Ok(Self {
             config,
             input_address: connected.input_address,
@@ -494,6 +507,7 @@ impl EngineCoreClient {
             output_task,
             dispatcher_task,
             abort_task,
+            readiness_task,
             coordinator_output_task,
             coordinator_task,
         })
@@ -527,6 +541,11 @@ impl EngineCoreClient {
         self.inner.pending_utility_call_count()
     }
 
+    #[cfg(test)]
+    pub(crate) fn readiness_recheck_pending(&self) -> bool {
+        self.inner.readiness().pending_ask().is_some()
+    }
+
     /// Return the engine-side indices connected to this client, as stamped on
     /// request output streams: an engine identity that does not encode an
     /// index maps to 0, as in [`Self::call`].
@@ -640,6 +659,13 @@ impl EngineCoreClient {
     pub fn health_error(&self) -> Option<Arc<Error>> {
         self.inner.health_error()
     }
+
+    /// Subscribe to whether the engines can accept and process new requests.
+    /// It is `false` while they are paused, asleep, or prepared for a
+    /// checkpoint, or while a weight update has not finished successfully.
+    pub fn subscribe_readiness(&self) -> tokio::sync::watch::Receiver<bool> {
+        self.inner.readiness().subscribe()
+    }
 }
 
 /// Build the prompt block hasher if every engine reports the same supported
@@ -798,7 +824,7 @@ impl EngineCoreClient {
         T: serde::de::DeserializeOwned,
         A: serde::Serialize + std::fmt::Debug,
     {
-        self.inner.call_utility(method, args).await
+        self.inner.call_utility(method, args, None).await
     }
 
     /// Call a utility method on all connected engines and return the shared
@@ -831,12 +857,35 @@ impl EngineCoreClient {
         args: A,
         kwargs: K,
     ) -> Result<Vec<rmpv::Value>>
+    where
+        A: serde::Serialize + std::fmt::Debug,
+        K: serde::Serialize + std::fmt::Debug,
+    {
+        // Worker methods can change readiness, as a weight update does.
+        let recheck = self.inner.readiness().recheck_after_call();
+        self.collective_rpc_with(recheck, method, timeout, args, kwargs).await
+    }
+
+    /// [`Self::collective_rpc`] as a control call that passes `recheck`.
+    async fn collective_rpc_with<A, K>(
+        &self,
+        recheck: RecheckGuard,
+        method: &str,
+        timeout: Option<f64>,
+        args: A,
+        kwargs: K,
+    ) -> Result<Vec<rmpv::Value>>
     where
         A: serde::Serialize + std::fmt::Debug,
         K: serde::Serialize + std::fmt::Debug,
     {
         let results = self
-            .call_utility::<rmpv::Value, _>("collective_rpc", (method, timeout, args, kwargs))
+            .inner
+            .call_utility::<rmpv::Value, _>(
+                "collective_rpc",
+                (method, timeout, args, kwargs),
+                Some(recheck),
+            )
             .await?;
 
         Ok(results
@@ -875,7 +924,8 @@ impl EngineCoreClient {
 
     /// Start a weight update for the base model.
     pub async fn start_weight_update(&self) -> Result<()> {
-        self.collective_rpc(
+        self.collective_rpc_with(
+            self.inner.readiness().withdraw_until_recheck(),
             "start_weight_update",
             None,
             Vec::<JsonValue>::new(),
@@ -887,7 +937,8 @@ impl EngineCoreClient {
 
     /// Start a weight update for the speculative draft model.
     pub async fn start_draft_weight_update(&self) -> Result<()> {
-        self.collective_rpc(
+        self.collective_rpc_with(
+            self.inner.readiness().withdraw_until_recheck(),
             "start_draft_weight_update",
             None,
             Vec::<JsonValue>::new(),
@@ -988,7 +1039,8 @@ impl EngineCoreClient {
 
     /// Put the engine to sleep.
     pub async fn sleep(&self, level: SleepLevel, mode: PauseMode) -> Result<()> {
-        self.call_utility::<(), _>("sleep", (level, mode)).await?;
+        let recheck = self.inner.readiness().withdraw_until_recheck();
+        self.inner.call_utility::<(), _>("sleep", (level, mode), Some(recheck)).await?;
         self.inner.record_sleep_state(EngineSleepState::Asleep(level));
         Ok(())
     }
@@ -1003,8 +1055,10 @@ impl EngineCoreClient {
     /// Wake the engine from sleep, optionally limiting the wake-up to specific
     /// tags, and return whether every engine is fully awake.
     pub async fn wake_up(&self, tags: Option<Vec<String>>) -> Result<bool> {
+        let recheck = self.inner.readiness().recheck_after_call();
         let fully_awake = self
-            .call_utility::<bool, _>("wake_up", (tags,))
+            .inner
+            .call_utility::<bool, _>("wake_up", (tags,), Some(recheck))
             .await?
             .into_iter()
             .all(|fully_awake| fully_awake);
@@ -1018,13 +1072,17 @@ impl EngineCoreClient {
     /// Pause generation (e.g. for weight updates). In `abort` and `wait` mode
     /// the engine rejects new requests with a retryable error until resume.
     pub async fn pause_generation(&self, mode: PauseMode, clear_cache: bool) -> Result<()> {
-        self.call_utility::<(), _>("pause_generation", (mode, clear_cache)).await?;
+        let recheck = self.inner.readiness().withdraw_until_recheck();
+        self.inner
+            .call_utility::<(), _>("pause_generation", (mode, clear_cache), Some(recheck))
+            .await?;
         Ok(())
     }
 
     /// Resume the scheduler after a pause
     pub async fn resume_scheduler(&self) -> Result<()> {
-        self.call_utility::<(), _>("resume_scheduler", ()).await?;
+        let recheck = self.inner.readiness().recheck_after_call();
+        self.inner.call_utility::<(), _>("resume_scheduler", (), Some(recheck)).await?;
         Ok(())
     }
 
@@ -1033,14 +1091,18 @@ impl EngineCoreClient {
     /// Returns the first refusal from any engine. An engine that refuses
     /// changes nothing, but with data parallelism the others may have started.
     pub async fn checkpoint_prepare(&self, mode: PauseMode) -> Result<Option<String>> {
+        let recheck = self.inner.readiness().withdraw_until_recheck();
         let refusals: Vec<Option<String>> =
-            self.call_utility("checkpoint_prepare", (mode,)).await?;
+            self.inner.call_utility("checkpoint_prepare", (mode,), Some(recheck)).await?;
         Ok(refusals.into_iter().flatten().next())
     }
 
     /// Bring back what `checkpoint_prepare` released and serve again.
     pub async fn checkpoint_restore(&self) -> Result<()> {
-        self.call_utility::<(), _>("checkpoint_restore", ()).await?;
+        let recheck = self.inner.readiness().recheck_after_call();
+        self.inner
+            .call_utility::<(), _>("checkpoint_restore", (), Some(recheck))
+            .await?;
         Ok(())
     }
 
@@ -1122,6 +1184,7 @@ impl EngineCoreClient {
             output_task,
             dispatcher_task,
             abort_task,
+            readiness_task,
             coordinator_output_task,
             coordinator_task,
             ..
@@ -1133,7 +1196,7 @@ impl EngineCoreClient {
 
         // Abort all client tasks first, then await them.
         // Note the aborting orders here.
-        let mut tasks = vec![abort_task, dispatcher_task, output_task];
+        let mut tasks = vec![readiness_task, abort_task, dispatcher_task, output_task];
         tasks.extend(coordinator_task);
         tasks.extend(coordinator_output_task);
 

rust/src/engine-core-client/src/client/imp.rs

@@ -17,6 +17,7 @@ use tracing::{debug, info, trace, warn};
 use vllm_metrics::METRICS;
 use zeromq::RouterSendHalf;
 
+use crate::client::readiness::{Readiness, RecheckGuard};
 use crate::client::state::{
     OutputReceiver, OutputSender, RequestRegistry, UtilityReceiver, UtilityRegistry,
 };
@@ -66,6 +67,9 @@ pub(crate) struct ClientInner {
     utility_reg: Mutex<UtilityRegistry>,
     health_error: ArcSwapOption<Error>,
     health_tx: watch::Sender<bool>,
+    /// Whether the engines can accept and process new requests, maintained
+    /// around the control calls this client sends.
+    readiness: Arc<Readiness>,
 }
 
 impl ClientInner {
@@ -103,6 +107,7 @@ impl ClientInner {
             utility_reg: Mutex::new(UtilityRegistry::default()),
             health_error: ArcSwapOption::empty(),
             health_tx: watch::Sender::new(true),
+            readiness: Arc::new(Readiness::new()),
         }
     }
 
@@ -260,9 +265,21 @@ impl ClientInner {
         self.health_tx.subscribe()
     }
 
+    pub fn readiness(&self) -> &Arc<Readiness> {
+        &self.readiness
+    }
+
     /// Call a typed utility method on every engine. See
     /// [`crate::EngineCoreClient::call_utility`].
-    pub async fn call_utility<T, A>(self: &Arc<Self>, method: &str, args: A) -> Result<Vec<T>>
+    ///
+    /// A control call passes its `recheck` guard, which the call holds until
+    /// the engines reply.
+    pub async fn call_utility<T, A>(
+        self: &Arc<Self>,
+        method: &str,
+        args: A,
+        recheck: Option<RecheckGuard>,
+    ) -> Result<Vec<T>>
     where
         T: serde::de::DeserializeOwned,
         A: serde::Serialize + std::fmt::Debug,
@@ -295,6 +312,7 @@ impl ClientInner {
         let inner = self.clone();
         let method_name = method.to_string();
         let call = self.handle.spawn(async move {
+            let _recheck = recheck;
             let outputs = join_all(calls.into_iter().map(|(engine_id, call_id, rx, payload)| {
                 let mut input_send = inner.input_send.clone();
                 let method_name = &method_name;
@@ -330,6 +348,12 @@ impl ClientInner {
         outputs.into_iter().map(|output| output?.into_typed_result(method)).collect()
     }
 
+    /// Ask every engine whether it can accept and process new requests.
+    pub async fn is_inference_ready(self: &Arc<Self>) -> Result<bool> {
+        let ready: Vec<bool> = self.call_utility("is_inference_ready", (), None).await?;
+        Ok(ready.into_iter().all(|ready| ready))
+    }
+
     /// Resolve one utility output to the waiting caller. Returns `true` if a
     /// waiting caller existed.
     pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool {