diff --git a/rust/Cargo.lock b/rust/Cargo.lock index e61b16e6d9c..0902898f364 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -4876,6 +4876,7 @@ dependencies = [ "futures-core", "pin-project-lite", "tokio", + "tokio-util", ] [[package]] @@ -4937,9 +4938,9 @@ dependencies = [ [[package]] name = "tonic" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" dependencies = [ "async-trait", "axum", @@ -4966,9 +4967,9 @@ dependencies = [ [[package]] name = "tonic-build" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1882ac3bf5ef12877d7ed57aad87e75154c11931c2ba7e6cde5e22d63522c734" +checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322" dependencies = [ "prettyplease", "proc-macro2", @@ -4977,10 +4978,23 @@ dependencies = [ ] [[package]] -name = "tonic-prost" -version = "0.14.5" +name = "tonic-health" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309" +checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082" +dependencies = [ + "prost", + "tokio", + "tokio-stream", + "tonic", + "tonic-prost", +] + +[[package]] +name = "tonic-prost" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" dependencies = [ "bytes", "prost", @@ -4989,9 +5003,9 @@ dependencies = [ [[package]] name = "tonic-prost-build" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3144df636917574672e93d0f56d7edec49f90305749c668df5101751bb8f95a" +checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27" dependencies = [ "prettyplease", "proc-macro2", @@ -5730,6 +5744,7 @@ dependencies = [ "tokio-stream", "tokio-util", "tonic", + "tonic-health", "tonic-prost", "tonic-prost-build", "tower", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index d42e51c649a..a33a13ff69b 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -118,10 +118,11 @@ tokio = { version = "1.47.1", features = [ tokio-openssl = "0.6" tokio-stream = "0.1" tokio-util = { version = "0.7.18", features = ["rt"] } -tonic = "0.14.5" -tonic-build = "0.14.5" -tonic-prost = "0.14.5" -tonic-prost-build = "0.14.5" +tonic = "0.14.6" +tonic-build = "0.14.6" +tonic-health = "0.14.6" +tonic-prost = "0.14.6" +tonic-prost-build = "0.14.6" tool-parser = "1.2.0" tower = { version = "0.5.3", features = ["util"] } tower-http = { version = "0.6.8", features = ["cors", "trace"] } diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 5b9cd500d39..9ea7b703671 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -460,6 +460,12 @@ impl EngineCoreClient { self.inner.is_healthy() } + /// Subscribe to engine health changes. The current value is `true` while + /// the client is healthy and changes permanently to `false` on failure. + pub fn subscribe_health(&self) -> tokio::sync::watch::Receiver { + self.inner.subscribe_health() + } + /// Return the first persistent health error observed by the client, if any. pub fn health_error(&self) -> Option> { self.inner.health_error() diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index a699f85c26a..e2a1b1b3d82 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -9,7 +9,7 @@ use arc_swap::ArcSwapOption; use parking_lot::Mutex; use thiserror_ext::AsReport as _; use tokio::runtime::Handle; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, watch}; use tracing::{debug, info, trace, warn}; use vllm_metrics::METRICS; use zeromq::RouterSendHalf; @@ -36,6 +36,7 @@ pub(crate) struct ClientInner { request_reg: Mutex, utility_reg: Mutex, health_error: ArcSwapOption, + health_tx: watch::Sender, } impl ClientInner { @@ -57,6 +58,7 @@ impl ClientInner { request_reg: Mutex::new(RequestRegistry::new(engines)), utility_reg: Mutex::new(UtilityRegistry::default()), health_error: ArcSwapOption::empty(), + health_tx: watch::Sender::new(true), } } @@ -169,6 +171,7 @@ impl ClientInner { /// persistent health error. pub fn close_registries(&self, error: Arc) { let persistent_error = self.record_health_error(error); + self.publish_unhealthy(); let request_senders = self.request_reg.lock().close(); let utility_senders = self.utility_reg.lock().close(); @@ -191,6 +194,12 @@ impl ClientInner { self.health_error.load().is_none() } + /// Subscribe to engine health changes. The current value is `true` while + /// the client is healthy and changes permanently to `false` on failure. + pub fn subscribe_health(&self) -> watch::Receiver { + self.health_tx.subscribe() + } + /// Resolve one utility output to the waiting caller. Returns `true` if a /// waiting caller existed. pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool { @@ -280,6 +289,11 @@ impl ClientInner { .expect("health error must be recorded before registries close") } + /// Publish the sticky healthy-to-unhealthy transition. + fn publish_unhealthy(&self) { + self.health_tx.send_if_modified(|healthy| std::mem::replace(healthy, false)); + } + /// Assert there is a recorded health error and return a `Shared` variant /// wrapping it for error returns when the client is already closed. fn closed_error(&self) -> Error { @@ -461,13 +475,18 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn close_registries_records_first_health_error_only() { let inner = test_inner().await; + let mut health = inner.subscribe_health(); + assert!(*health.borrow()); inner.close_registries(Arc::new(Error::EngineCoreDead)); + health.changed().await.expect("health sender remains open"); assert!(!inner.is_healthy()); + assert!(!*health.borrow()); assert!(matches!( inner.health_error().as_deref(), Some(Error::EngineCoreDead) )); + assert!(!*inner.subscribe_health().borrow()); inner.close_registries(Arc::new(client_closed!("shutdown"))); assert!(matches!( diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index f3e03863d49..1ed397d589c 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -35,6 +35,7 @@ tokio-openssl.workspace = true tokio-stream.workspace = true tokio-util.workspace = true tonic.workspace = true +tonic-health.workspace = true tonic-prost.workspace = true tower.workspace = true tower-http.workspace = true diff --git a/rust/src/server/src/grpc/health.rs b/rust/src/server/src/grpc/health.rs new file mode 100644 index 00000000000..70bb81da9f9 --- /dev/null +++ b/rust/src/server/src/grpc/health.rs @@ -0,0 +1,70 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use tokio::sync::watch; +use tokio_util::sync::CancellationToken; +use tonic::server::NamedService; +use tonic_health::ServingStatus; +use tonic_health::server::HealthReporter; +use tracing::{info, warn}; + +use super::GenerateGrpcService; + +pub(crate) async fn monitor_health( + mut health_reporter: HealthReporter, + mut engine_health: watch::Receiver, + shutdown: CancellationToken, +) { + let generate_service = GenerateGrpcService::NAME; + let status = ServingStatus::NotServing; + let health_event_first = tokio::select! { + result = engine_health.wait_for(|healthy| !*healthy) => { + match result { + Ok(_) => warn!( + generate_service, + overall_service = true, + status = ?status, + reason = "engine_unhealthy", + "marking gRPC health services as not serving" + ), + Err(error) => warn!( + %error, + generate_service, + overall_service = true, + status = ?status, + reason = "health_channel_closed", + "engine health channel closed; marking gRPC health services as not serving" + ), + } + true + } + _ = shutdown.cancelled() => { + info!( + generate_service, + overall_service = true, + status = ?status, + reason = "server_shutdown", + "server shutting down; marking gRPC health services as not serving" + ); + false + } + }; + + health_reporter.set_not_serving::().await; + // Generate is currently the only engine-backed gRPC service, so overall + // server health intentionally mirrors it. + health_reporter.set_service_status("", status).await; + + if health_event_first { + shutdown.cancelled().await; + info!( + generate_service, + overall_service = true, + reason = "server_shutdown", + "server shutting down; closing gRPC health watches" + ); + } + + health_reporter.clear_service_status(generate_service).await; + health_reporter.clear_service_status("").await; +} diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs index 023aae38ac0..4170e5bb895 100644 --- a/rust/src/server/src/grpc/mod.rs +++ b/rust/src/server/src/grpc/mod.rs @@ -4,6 +4,7 @@ //! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade. mod convert; +mod health; use std::pin::Pin; use std::sync::Arc; @@ -24,8 +25,11 @@ pub mod pb { tonic::include_proto!("vllm"); } +pub(crate) use health::monitor_health; pub use pb::generate_server::GenerateServer; +pub(crate) type GenerateGrpcService = GenerateServer; + #[cfg(test)] mod tests; diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 56d0c1c2f5a..32412dd9ab7 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -16,6 +16,10 @@ use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; use tokio::net::TcpStream; use tokio_openssl::SslStream; use tonic::transport::{Channel, Endpoint, Server as TonicServer, Uri}; +use tonic_health::pb::HealthCheckRequest; +use tonic_health::pb::health_check_response::ServingStatus as HealthServingStatus; +use tonic_health::pb::health_client::HealthClient; +use tonic_health::server::health_reporter; use tower::service_fn; use vllm_chat::{ ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor, @@ -200,7 +204,11 @@ impl ChatRenderer for FakeTextBackend { async fn setup_grpc_service( engine_id: impl Into, output_specs: Vec<(Vec, Option)>, -) -> (GenerateServer, MockEngineTask) { +) -> ( + GenerateServer, + tokio::sync::watch::Receiver, + MockEngineTask, +) { let ipc = IpcNamespace::new().expect("create ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let engine_id = engine_id.into(); @@ -232,6 +240,7 @@ async fn setup_grpc_service( ) .await .expect("connect client"); + let engine_health = client.subscribe_health(); let chat = ChatLlm::from_shared_backend( test_llm(client), @@ -240,6 +249,7 @@ async fn setup_grpc_service( let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat)); ( GenerateServer::new(GenerateServiceImpl::new(state)), + engine_health, engine_task, ) } @@ -254,25 +264,51 @@ async fn grpc_test_server( tokio::task::JoinHandle<()>, MockEngineTask, ) { - let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await; + let (svc, engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await; + let (channel, server_task) = start_grpc_test_server( + svc, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + (GenerateClient::new(channel), server_task, engine_task) +} + +async fn start_grpc_test_server( + generate_service: GenerateServer, + engine_health: tokio::sync::watch::Receiver, + shutdown: tokio_util::sync::CancellationToken, +) -> (Channel, tokio::task::JoinHandle<()>) { + let (health_reporter, health_service) = health_reporter(); + health_reporter.set_serving::>().await; let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener"); let addr = listener.local_addr().expect("local addr"); let server_task = tokio::spawn(async move { let incoming = MaybeTlsListener::plain(Listener::Tcp(listener)); - TonicServer::builder() - .add_service(svc) - .serve_with_incoming(incoming) - .await - .expect("grpc server"); + let server = TonicServer::builder() + .add_service(health_service) + .add_service(generate_service) + .serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned()); + let health_monitor = + super::monitor_health(health_reporter, engine_health, shutdown.clone()); + let server = async move { + let result = server.await; + shutdown.cancel(); + result + }; + let (server_result, ()) = tokio::join!(server, health_monitor); + server_result.expect("grpc server"); }); - let grpc_client = GenerateClient::connect(format!("http://{addr}")) + let channel = Endpoint::from_shared(format!("http://{addr}")) + .expect("grpc endpoint") + .connect() .await - .expect("connect grpc client"); + .expect("connect grpc channel"); - (grpc_client, server_task, engine_task) + (channel, server_task) } /// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode). @@ -283,7 +319,7 @@ async fn grpc_tls_test_server( certs: &TestCerts, cert_reqs: i32, ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) { - let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await; + let (svc, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await; let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs)) .expect("build grpc tls config"); @@ -373,7 +409,8 @@ async fn grpc_server_with_keepalive( engine_id: impl Into, keepalive: Option, ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) { - let (svc, engine_task) = setup_grpc_service(engine_id, default_stream_output_specs()).await; + let (svc, _engine_health, engine_task) = + setup_grpc_service(engine_id, default_stream_output_specs()).await; let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener"); let addr = listener.local_addr().expect("local addr").to_string(); @@ -1035,3 +1072,129 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() { server_task.abort(); } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() { + let (generate_service, _connected_engine_health, _engine_task) = + setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await; + let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true); + let (channel, server_task) = start_grpc_test_server( + generate_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut health_client = HealthClient::new(channel); + + let mut health_streams = Vec::new(); + for service in ["vllm.Generate", ""] { + let service_label = if service.is_empty() { + "overall" + } else { + service + }; + let mut stream = health_client + .watch(HealthCheckRequest { + service: service.to_string(), + }) + .await + .unwrap_or_else(|error| { + panic!("failed to start health watch for {service_label}: {error}") + }) + .into_inner(); + let initial = stream + .message() + .await + .unwrap_or_else(|error| { + panic!("failed to read initial health status for {service_label}: {error}") + }) + .unwrap_or_else(|| { + panic!("health watch for {service_label} ended before its initial status") + }); + assert_eq!( + initial.status, + HealthServingStatus::Serving as i32, + "unexpected initial health status for {service_label}" + ); + health_streams.push((service_label, stream)); + } + + engine_health_tx.send(false).expect("publish unhealthy engine state"); + + for (service_label, mut stream) in health_streams { + let update = tokio::time::timeout(Duration::from_secs(2), stream.message()) + .await + .unwrap_or_else(|_| panic!("timed out waiting for health update for {service_label}")) + .unwrap_or_else(|error| { + panic!("failed to read health update for {service_label}: {error}") + }) + .unwrap_or_else(|| panic!("health watch for {service_label} ended before its update")); + assert_eq!( + update.status, + HealthServingStatus::NotServing as i32, + "unexpected health status for {service_label}" + ); + } + + server_task.abort(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn grpc_health_watch_closes_on_graceful_shutdown() { + let (generate_service, engine_health, _engine_task) = setup_grpc_service( + b"engine-grpc-health-shutdown", + default_stream_output_specs(), + ) + .await; + let shutdown = tokio_util::sync::CancellationToken::new(); + let (channel, server_task) = + start_grpc_test_server(generate_service, engine_health, shutdown.clone()).await; + let mut health_client = HealthClient::new(channel); + let mut stream = health_client + .watch(HealthCheckRequest { + service: "vllm.Generate".to_string(), + }) + .await + .expect("start health watch for vllm.Generate") + .into_inner(); + + let initial = stream + .message() + .await + .expect("read initial health status for vllm.Generate") + .expect("health watch ended before its initial status"); + assert_eq!( + initial.status, + HealthServingStatus::Serving as i32, + "unexpected initial health status for vllm.Generate" + ); + + shutdown.cancel(); + + let update = tokio::time::timeout(Duration::from_secs(2), stream.message()) + .await + .expect("timed out waiting for shutdown health update for vllm.Generate") + .expect("failed to read shutdown health update for vllm.Generate") + .expect("health watch ended before its shutdown update"); + assert_eq!( + update.status, + HealthServingStatus::NotServing as i32, + "unexpected shutdown health status for vllm.Generate" + ); + + let stream_end = tokio::time::timeout(Duration::from_secs(2), stream.message()) + .await + .expect("timed out waiting for vllm.Generate health watch to close") + .expect("failed while closing vllm.Generate health watch"); + assert!( + stream_end.is_none(), + "vllm.Generate health watch remained open" + ); + + tokio::time::timeout(Duration::from_secs(2), server_task) + .await + .expect("timed out waiting for gRPC server shutdown") + .expect("gRPC server task failed"); +} diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 7f1ec91c6f2..220da778af8 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -39,6 +39,7 @@ use tokio::net::TcpListener; use tokio::time::{Instant, sleep_until}; use tokio_util::sync::CancellationToken; use tonic::transport::Server as TonicServer; +use tonic_health::server::health_reporter; use tower::ServiceExt as _; use tracing::{info, trace, warn}; use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends}; @@ -203,14 +204,19 @@ where .map(tls::build_grpc_server_config) .transpose() .context("invalid gRPC TLS configuration")?; - let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone())); + let (health_reporter, health_service) = health_reporter(); + let engine_health = state.engine_core_client().subscribe_health(); + health_reporter.set_serving::().await; + let generate_service = + grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone())); let svc = TonicServer::builder() .http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL)) .http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT)) .layer(middleware::request_runtime_layer(state.clone())) - .add_service(svc); + .add_service(health_service) + .add_service(generate_service); info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server"); - Some((grpc_listener, svc, grpc_tls)) + Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) } else { None }; @@ -294,7 +300,8 @@ where let server_shutdown = server_shutdown.clone(); let force_shutdown = force_shutdown.clone(); async move { - let Some((grpc_listener, svc, grpc_tls)) = grpc_setup else { + let Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) = grpc_setup + else { // No gRPC configured: just wait for shutdown so we do not race the // join! by resolving early and tripping the cancellation token. shutdown.cancelled().await; @@ -304,19 +311,26 @@ where Some(context) => MaybeTlsListener::tls(grpc_listener, context), None => MaybeTlsListener::plain(grpc_listener), }; - let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned()); + let server = + svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned()); + let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown); - let result = tokio::select! { - result = server => { - result.context("gRPC server failed") - } - _ = force_shutdown.cancelled() => { - warn!("gRPC graceful shutdown deadline elapsed; aborting server"); - Ok(()) - } + let server = async move { + let result = tokio::select! { + result = server => { + result.context("gRPC server failed") + } + _ = force_shutdown.cancelled() => { + warn!("gRPC graceful shutdown deadline elapsed; aborting server"); + Ok(()) + } + }; + + server_shutdown.cancel(); + result }; - server_shutdown.cancel(); + let (result, ()) = tokio::join!(server, health_monitor); result } };