From cc745a63f3d67411a562f0e68008fecaf17c48cc Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Fri, 10 Jul 2026 09:56:15 -0700 Subject: [PATCH 1/6] feat(grpc): extend the existing service protocol Signed-off-by: Connor Carpenter --- rust/proto/vllm_grpc.proto | 209 ++++++++++++++++++++++++++++++ rust/src/server/src/grpc/mod.rs | 70 ++++++++++ rust/src/server/src/grpc/tests.rs | 58 +++++++++ 3 files changed, 337 insertions(+) diff --git a/rust/proto/vllm_grpc.proto b/rust/proto/vllm_grpc.proto index 56c5f36442db..df833855f265 100644 --- a/rust/proto/vllm_grpc.proto +++ b/rust/proto/vllm_grpc.proto @@ -9,6 +9,19 @@ service Generate { rpc Generate (GenerateRequest) returns (GenerateResponse) {} // Generates text given a prompt, streaming the outputs rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {} + + rpc GetEngineInfo (GetEngineInfoRequest) returns (EngineInfo) {} + rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {} + rpc Health (HealthRequest) returns (HealthResponse) {} + rpc Abort (AbortRequest) returns (AbortResponse) {} + rpc Drain (DrainRequest) returns (DrainResponse) {} + + rpc LoadLora (LoadLoraRequest) returns (LoadLoraResponse) {} + rpc UnloadLora (UnloadLoraRequest) returns (UnloadLoraResponse) {} + rpc ListLoras (ListLorasRequest) returns (ListLorasResponse) {} + + rpc GetKvConnectorInfo (GetKvConnectorInfoRequest) returns (KvConnectorInfo) {} + rpc GetKvEventSources (GetKvEventSourcesRequest) returns (GetKvEventSourcesResponse) {} } // ====================================================================================== @@ -42,6 +55,13 @@ message GenerateRequest { uint32 truncate_prompt_tokens = 11; int32 priority = 12; + + // Multimodal inputs aligned with placeholder markers in token_ids. + repeated MediaItem media = 13; + // Loaded LoRA name; empty selects the base model. + string lora_name = 14; + // Router-selected data-parallel rank. + optional uint32 data_parallel_rank = 15; } message RandomSampling { @@ -194,3 +214,192 @@ message TokenIds { repeated uint32 ids = 1; } +// ====================================================================================== +// Media +// ====================================================================================== + +enum Modality { + MODALITY_UNSPECIFIED = 0; + MODALITY_IMAGE = 1; + MODALITY_VIDEO = 2; + MODALITY_AUDIO = 3; +} + +message MediaItem { + Modality modality = 1; + oneof source { + string url = 2; + string data_uri = 3; + bytes raw_bytes = 4; + } + string mime_type = 5; + string uuid = 6; +} + +// ====================================================================================== +// Discovery and lifecycle +// ====================================================================================== + +enum EngineRole { + ENGINE_ROLE_UNSPECIFIED = 0; + ENGINE_ROLE_AGGREGATED = 1; + ENGINE_ROLE_PREFILL = 2; + ENGINE_ROLE_DECODE = 3; +} + +message GetEngineInfoRequest {} + +message EngineInfo { + string engine_name = 1; + string engine_version = 2; + string api_version = 3; + EngineRole role = 4; + string instance_id = 5; + repeated string supported_models = 6; + ParallelismInfo parallelism = 7; + KvConnectorInfo kv_connector = 8; +} + +message ParallelismInfo { + uint32 tensor_parallel_size = 1; + uint32 pipeline_parallel_size = 2; + uint32 data_parallel_size = 3; + uint32 data_parallel_rank = 4; + uint32 data_parallel_start_rank = 5; +} + +message GetModelInfoRequest {} + +message ModelInfo { + string model_id = 1; + string served_model_name = 2; + repeated string served_model_aliases = 3; + uint32 max_context_length = 4; + uint32 max_output_tokens = 5; + uint32 kv_block_size = 6; + uint64 total_kv_blocks = 7; + uint64 max_running_requests = 8; + uint64 max_batched_tokens = 9; + repeated string tokenizer_modes = 10; + uint32 max_loras = 11; + + bool supports_text_input = 20; + bool supports_token_ids_input = 21; + bool supports_lora = 22; + bool supports_multimodal = 23; + string reasoning_parser = 24; + string tool_call_parser = 25; +} + +message HealthRequest { + bool include_inference_probe = 1; + string model = 2; +} + +message HealthResponse { + HealthState state = 1; + repeated HealthCheck checks = 2; +} + +enum HealthState { + HEALTH_STATE_UNSPECIFIED = 0; + HEALTH_STATE_STARTING = 1; + HEALTH_STATE_READY = 2; + HEALTH_STATE_DEGRADED = 3; + HEALTH_STATE_DRAINING = 4; + HEALTH_STATE_NOT_READY = 5; +} + +message HealthCheck { + string name = 1; + HealthState state = 2; + string message = 3; +} + +message AbortRequest { + string request_id = 1; + bool abort_all = 2; +} + +message AbortResponse { + AbortStatus status = 1; + string message = 2; +} + +enum AbortStatus { + ABORT_STATUS_UNSPECIFIED = 0; + ABORT_STATUS_ABORTED = 1; + ABORT_STATUS_NOT_FOUND = 2; + ABORT_STATUS_ALREADY_FINISHED = 3; + ABORT_STATUS_UNSUPPORTED = 4; +} + +message DrainRequest {} + +message DrainResponse { + DrainState state = 1; + uint32 in_flight_requests = 2; + string message = 3; +} + +enum DrainState { + DRAIN_STATE_UNSPECIFIED = 0; + DRAIN_STATE_IN_PROGRESS = 1; + DRAIN_STATE_COMPLETE = 2; +} + +// ====================================================================================== +// LoRA lifecycle +// ====================================================================================== + +message LoraAdapter { + int64 lora_id = 1; + string lora_name = 2; + string source_path = 3; +} + +message LoadLoraRequest { LoraAdapter adapter = 1; } +message LoadLoraResponse { + LoraAdapter adapter = 1; + bool already_loaded = 2; +} +message UnloadLoraRequest { string lora_name = 1; } +message UnloadLoraResponse { LoraAdapter adapter = 1; } +message ListLorasRequest {} +message ListLorasResponse { repeated LoraAdapter adapters = 1; } + +// ====================================================================================== +// KV discovery +// ====================================================================================== + +message KvEndpoint { + string host = 1; + uint32 port = 2; + string protocol = 3; +} + +message GetKvConnectorInfoRequest {} + +message KvConnectorInfo { + bool enabled = 1; + string transfer_backend = 2; + repeated KvEndpoint local_endpoints = 3; + repeated string supported_protocols = 4; + uint32 schema_version = 5; +} + +message GetKvEventSourcesRequest {} +message GetKvEventSourcesResponse { repeated KvEventSource sources = 1; } + +message KvEventSource { + string transport = 1; + KvEndpoint endpoint_addr = 2; + string topic = 3; + string replay_endpoint = 4; + uint32 data_parallel_rank = 5; + string encoding = 6; + uint32 schema_version = 7; + uint32 buffer_steps = 8; + uint32 hwm = 9; + uint32 max_queue_size = 10; +} diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs index 1fcb8674fee1..57549a3da90e 100644 --- a/rust/src/server/src/grpc/mod.rs +++ b/rust/src/server/src/grpc/mod.rs @@ -151,6 +151,76 @@ impl pb::generate_server::Generate for GenerateServiceImpl { let response_stream = ReceiverStream::new(rx); Ok(Response::new(Box::pin(response_stream))) } + + async fn get_engine_info( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("GetEngineInfo")) + } + + async fn get_model_info( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("GetModelInfo")) + } + + async fn health( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("Health")) + } + + async fn abort( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("Abort")) + } + + async fn drain( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("Drain")) + } + + async fn load_lora( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("LoadLora")) + } + + async fn unload_lora( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("UnloadLora")) + } + + async fn list_loras( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("ListLoras")) + } + + async fn get_kv_connector_info( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("GetKvConnectorInfo")) + } + + async fn get_kv_event_sources( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("GetKvEventSources")) + } } fn text_error_to_status(error: vllm_text::Error) -> Status { diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 83c4de440efc..d2fc85d46411 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1031,3 +1031,61 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() { server_task.abort(); } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn extension_methods_are_unimplemented() { + let (mut client, server_task, _engine_task) = + grpc_test_server(b"engine-grpc-stubs", default_stream_output_specs()).await; + + assert_eq!( + client.get_engine_info(pb::GetEngineInfoRequest {}).await.unwrap_err().code(), + tonic::Code::Unimplemented + ); + assert_eq!( + client.get_model_info(pb::GetModelInfoRequest {}).await.unwrap_err().code(), + tonic::Code::Unimplemented + ); + assert_eq!( + client.health(pb::HealthRequest::default()).await.unwrap_err().code(), + tonic::Code::Unimplemented + ); + assert_eq!( + client.abort(pb::AbortRequest::default()).await.unwrap_err().code(), + tonic::Code::Unimplemented + ); + assert_eq!( + client.drain(pb::DrainRequest {}).await.unwrap_err().code(), + tonic::Code::Unimplemented + ); + assert_eq!( + client.load_lora(pb::LoadLoraRequest::default()).await.unwrap_err().code(), + tonic::Code::Unimplemented + ); + assert_eq!( + client.unload_lora(pb::UnloadLoraRequest::default()).await.unwrap_err().code(), + tonic::Code::Unimplemented + ); + assert_eq!( + client.list_loras(pb::ListLorasRequest {}).await.unwrap_err().code(), + tonic::Code::Unimplemented + ); + assert_eq!( + client + .get_kv_connector_info(pb::GetKvConnectorInfoRequest {}) + .await + .unwrap_err() + .code(), + tonic::Code::Unimplemented + ); + assert_eq!( + client + .get_kv_event_sources(pb::GetKvEventSourcesRequest {}) + .await + .unwrap_err() + .code(), + tonic::Code::Unimplemented + ); + + server_task.abort(); +} From a5a701b832038b1ad4ceed3330b11f07a0a21afa Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Tue, 14 Jul 2026 15:06:18 -0700 Subject: [PATCH 2/6] fix(grpc): address foundation review feedback Signed-off-by: Connor Carpenter --- rust/Cargo.lock | 15 ++++ rust/Cargo.toml | 1 + rust/proto/vllm_grpc.proto | 108 +++++++-------------------- rust/src/server/Cargo.toml | 1 + rust/src/server/src/grpc/mod.rs | 36 ++++----- rust/src/server/src/grpc/tests.rs | 118 +++++++++++++++++++----------- rust/src/server/src/lib.rs | 16 +++- 7 files changed, 148 insertions(+), 147 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index cc055505dc4e..9c149facb241 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -4534,6 +4534,7 @@ dependencies = [ "futures-core", "pin-project-lite", "tokio", + "tokio-util", ] [[package]] @@ -4634,6 +4635,19 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "tonic-health" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4ff0636fef47afb3ec02818f5bceb4377b8abb9d6a386aeade18bd6212f8eb7" +dependencies = [ + "prost", + "tokio", + "tokio-stream", + "tonic", + "tonic-prost", +] + [[package]] name = "tonic-prost" version = "0.14.5" @@ -5312,6 +5326,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 435350c07117..811424102cf1 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -111,6 +111,7 @@ tokio-stream = "0.1" tokio-util = { version = "0.7.18", features = ["rt"] } tonic = "0.14.5" tonic-build = "0.14.5" +tonic-health = "0.14.5" tonic-prost = "0.14.5" tonic-prost-build = "0.14.5" tool-parser = "1.2.0" diff --git a/rust/proto/vllm_grpc.proto b/rust/proto/vllm_grpc.proto index df833855f265..f18d05530543 100644 --- a/rust/proto/vllm_grpc.proto +++ b/rust/proto/vllm_grpc.proto @@ -9,10 +9,11 @@ service Generate { rpc Generate (GenerateRequest) returns (GenerateResponse) {} // Generates text given a prompt, streaming the outputs rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {} +} - rpc GetEngineInfo (GetEngineInfoRequest) returns (EngineInfo) {} +service Engine { + rpc GetDeploymentInfo (GetDeploymentInfoRequest) returns (DeploymentInfo) {} rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {} - rpc Health (HealthRequest) returns (HealthResponse) {} rpc Abort (AbortRequest) returns (AbortResponse) {} rpc Drain (DrainRequest) returns (DrainResponse) {} @@ -20,7 +21,6 @@ service Generate { rpc UnloadLora (UnloadLoraRequest) returns (UnloadLoraResponse) {} rpc ListLoras (ListLorasRequest) returns (ListLorasResponse) {} - rpc GetKvConnectorInfo (GetKvConnectorInfoRequest) returns (KvConnectorInfo) {} rpc GetKvEventSources (GetKvEventSourcesRequest) returns (GetKvEventSourcesResponse) {} } @@ -60,8 +60,6 @@ message GenerateRequest { repeated MediaItem media = 13; // Loaded LoRA name; empty selects the base model. string lora_name = 14; - // Router-selected data-parallel rank. - optional uint32 data_parallel_rank = 15; } message RandomSampling { @@ -240,24 +238,18 @@ message MediaItem { // Discovery and lifecycle // ====================================================================================== -enum EngineRole { - ENGINE_ROLE_UNSPECIFIED = 0; - ENGINE_ROLE_AGGREGATED = 1; - ENGINE_ROLE_PREFILL = 2; - ENGINE_ROLE_DECODE = 3; -} +message GetDeploymentInfoRequest {} -message GetEngineInfoRequest {} - -message EngineInfo { - string engine_name = 1; - string engine_version = 2; - string api_version = 3; - EngineRole role = 4; - string instance_id = 5; - repeated string supported_models = 6; - ParallelismInfo parallelism = 7; - KvConnectorInfo kv_connector = 8; +message DeploymentInfo { + string engine_version = 1; + string api_version = 2; + string instance_id = 3; + repeated string supported_models = 4; + ParallelismInfo parallelism = 5; + uint32 kv_block_size = 6; + uint64 total_kv_blocks = 7; + uint64 max_running_requests = 8; + uint64 max_batched_tokens = 9; } message ParallelismInfo { @@ -266,6 +258,7 @@ message ParallelismInfo { uint32 data_parallel_size = 3; uint32 data_parallel_rank = 4; uint32 data_parallel_start_rank = 5; + uint32 decode_context_parallel_size = 6; } message GetModelInfoRequest {} @@ -276,12 +269,8 @@ message ModelInfo { repeated string served_model_aliases = 3; uint32 max_context_length = 4; uint32 max_output_tokens = 5; - uint32 kv_block_size = 6; - uint64 total_kv_blocks = 7; - uint64 max_running_requests = 8; - uint64 max_batched_tokens = 9; - repeated string tokenizer_modes = 10; - uint32 max_loras = 11; + repeated string tokenizer_modes = 6; + uint32 max_loras = 7; bool supports_text_input = 20; bool supports_token_ids_input = 21; @@ -291,48 +280,11 @@ message ModelInfo { string tool_call_parser = 25; } -message HealthRequest { - bool include_inference_probe = 1; - string model = 2; -} - -message HealthResponse { - HealthState state = 1; - repeated HealthCheck checks = 2; -} - -enum HealthState { - HEALTH_STATE_UNSPECIFIED = 0; - HEALTH_STATE_STARTING = 1; - HEALTH_STATE_READY = 2; - HEALTH_STATE_DEGRADED = 3; - HEALTH_STATE_DRAINING = 4; - HEALTH_STATE_NOT_READY = 5; -} - -message HealthCheck { - string name = 1; - HealthState state = 2; - string message = 3; -} - message AbortRequest { - string request_id = 1; - bool abort_all = 2; + repeated string request_ids = 1; } -message AbortResponse { - AbortStatus status = 1; - string message = 2; -} - -enum AbortStatus { - ABORT_STATUS_UNSPECIFIED = 0; - ABORT_STATUS_ABORTED = 1; - ABORT_STATUS_NOT_FOUND = 2; - ABORT_STATUS_ALREADY_FINISHED = 3; - ABORT_STATUS_UNSUPPORTED = 4; -} +message AbortResponse {} message DrainRequest {} @@ -372,31 +324,21 @@ message ListLorasResponse { repeated LoraAdapter adapters = 1; } // KV discovery // ====================================================================================== -message KvEndpoint { +message GetKvEventSourcesRequest {} +message GetKvEventSourcesResponse { repeated KvEventSource sources = 1; } + +message KvEventEndpoint { string host = 1; uint32 port = 2; string protocol = 3; } -message GetKvConnectorInfoRequest {} - -message KvConnectorInfo { - bool enabled = 1; - string transfer_backend = 2; - repeated KvEndpoint local_endpoints = 3; - repeated string supported_protocols = 4; - uint32 schema_version = 5; -} - -message GetKvEventSourcesRequest {} -message GetKvEventSourcesResponse { repeated KvEventSource sources = 1; } - message KvEventSource { string transport = 1; - KvEndpoint endpoint_addr = 2; + KvEventEndpoint endpoint_addr = 2; string topic = 3; string replay_endpoint = 4; - uint32 data_parallel_rank = 5; + optional uint32 data_parallel_rank = 5; string encoding = 6; uint32 schema_version = 7; uint32 buffer_steps = 8; diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index f3e03863d495..1ed397d589cc 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/mod.rs b/rust/src/server/src/grpc/mod.rs index 57549a3da90e..b8fcc7408774 100644 --- a/rust/src/server/src/grpc/mod.rs +++ b/rust/src/server/src/grpc/mod.rs @@ -21,6 +21,7 @@ pub mod pb { tonic::include_proto!("vllm"); } +pub use pb::engine_server::EngineServer; pub use pb::generate_server::GenerateServer; #[cfg(test)] @@ -37,6 +38,16 @@ impl GenerateServiceImpl { } } +/// Unimplemented control-plane service registered on the existing gRPC listener. +#[derive(Default)] +pub struct EngineServiceImpl; + +impl EngineServiceImpl { + pub fn new() -> Self { + Self + } +} + #[tonic::async_trait] impl pb::generate_server::Generate for GenerateServiceImpl { type GenerateStreamStream = @@ -151,12 +162,15 @@ impl pb::generate_server::Generate for GenerateServiceImpl { let response_stream = ReceiverStream::new(rx); Ok(Response::new(Box::pin(response_stream))) } +} - async fn get_engine_info( +#[tonic::async_trait] +impl pb::engine_server::Engine for EngineServiceImpl { + async fn get_deployment_info( &self, - _request: Request, - ) -> Result, Status> { - Err(Status::unimplemented("GetEngineInfo")) + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("GetDeploymentInfo")) } async fn get_model_info( @@ -166,13 +180,6 @@ impl pb::generate_server::Generate for GenerateServiceImpl { Err(Status::unimplemented("GetModelInfo")) } - async fn health( - &self, - _request: Request, - ) -> Result, Status> { - Err(Status::unimplemented("Health")) - } - async fn abort( &self, _request: Request, @@ -208,13 +215,6 @@ impl pb::generate_server::Generate for GenerateServiceImpl { Err(Status::unimplemented("ListLoras")) } - async fn get_kv_connector_info( - &self, - _request: Request, - ) -> Result, Status> { - Err(Status::unimplemented("GetKvConnectorInfo")) - } - async fn get_kv_event_sources( &self, _request: Request, diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index d2fc85d46411..fb14118d5889 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -13,6 +13,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, @@ -31,8 +35,9 @@ use vllm_tokenizer::test_utils::TestTokenizer; use zeromq::prelude::{SocketRecv, SocketSend}; use zeromq::{DealerSocket, PushSocket, ZmqMessage}; +use super::pb::engine_client::EngineClient; use super::pb::generate_client::GenerateClient; -use super::{GenerateServer, GenerateServiceImpl, pb}; +use super::{EngineServer, EngineServiceImpl, GenerateServer, GenerateServiceImpl, pb}; use crate::listener::{Listener, MaybeTlsListener}; use crate::state::AppState; use crate::tls; @@ -247,10 +252,16 @@ async fn grpc_test_server( output_specs: Vec<(Vec, Option)>, ) -> ( GenerateClient, + EngineClient, + HealthClient, tokio::task::JoinHandle<()>, MockEngineTask, ) { let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await; + let engine_service = EngineServer::new(EngineServiceImpl::new()); + let (health_reporter, health_service) = health_reporter(); + health_reporter.set_serving::>().await; + 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"); @@ -258,17 +269,30 @@ async fn grpc_test_server( let server_task = tokio::spawn(async move { let incoming = MaybeTlsListener::plain(Listener::Tcp(listener)); TonicServer::builder() + .add_service(health_service) .add_service(svc) + .add_service(engine_service) .serve_with_incoming(incoming) .await .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"); + let grpc_client = GenerateClient::new(channel.clone()); + let engine_client = EngineClient::new(channel.clone()); + let health_client = HealthClient::new(channel); - (grpc_client, server_task, engine_task) + ( + grpc_client, + engine_client, + health_client, + server_task, + engine_task, + ) } /// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode). @@ -430,7 +454,7 @@ async fn h2_unresponsive_peer_closed_within(addr: &str, wait: Duration) -> bool #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn unary_generate_returns_collected_text() { - let (mut client, server_task, engine_task) = + let (mut client, _control_client, _health_client, server_task, engine_task) = grpc_test_server(b"engine-grpc-unary", default_stream_output_specs()).await; let response = client @@ -473,7 +497,7 @@ async fn unary_generate_returns_collected_text() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn unary_generate_with_token_ids_prompt() { - let (mut client, server_task, engine_task) = + let (mut client, _control_client, _health_client, server_task, engine_task) = grpc_test_server(b"engine-grpc-token-ids", default_stream_output_specs()).await; let response = client @@ -507,7 +531,7 @@ async fn unary_generate_with_token_ids_prompt() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn unary_generate_returns_token_ids_when_requested() { - let (mut client, server_task, engine_task) = + let (mut client, _control_client, _health_client, server_task, engine_task) = grpc_test_server(b"engine-grpc-tok-resp", default_stream_output_specs()).await; let response = client @@ -547,7 +571,7 @@ async fn unary_generate_returns_token_ids_when_requested() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn unary_generate_missing_prompt_returns_invalid_argument() { - let (mut client, server_task, _engine_task) = + let (mut client, _control_client, _health_client, server_task, _engine_task) = grpc_test_server(b"engine-grpc-no-prompt", default_stream_output_specs()).await; let status = client @@ -569,7 +593,7 @@ async fn unary_generate_missing_prompt_returns_invalid_argument() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn unary_generate_min_tokens_above_max_tokens_returns_invalid_argument() { - let (mut client, server_task, _engine_task) = + let (mut client, _control_client, _health_client, server_task, _engine_task) = grpc_test_server(b"engine-grpc-min-above-max", default_stream_output_specs()).await; let status = client @@ -597,7 +621,7 @@ async fn unary_generate_min_tokens_above_max_tokens_returns_invalid_argument() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn streaming_generate_yields_incremental_responses() { - let (mut client, server_task, engine_task) = + let (mut client, _control_client, _health_client, server_task, engine_task) = grpc_test_server(b"engine-grpc-stream", default_stream_output_specs()).await; let stream = client @@ -661,11 +685,12 @@ async fn streaming_generate_yields_incremental_responses() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn streaming_generate_missing_prompt_returns_invalid_argument() { - let (mut client, server_task, _engine_task) = grpc_test_server( - b"engine-grpc-stream-no-prompt", - default_stream_output_specs(), - ) - .await; + let (mut client, _control_client, _health_client, server_task, _engine_task) = + grpc_test_server( + b"engine-grpc-stream-no-prompt", + default_stream_output_specs(), + ) + .await; let status = client .generate_stream(pb::GenerateRequest { @@ -685,11 +710,12 @@ async fn streaming_generate_missing_prompt_returns_invalid_argument() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn streaming_generate_min_tokens_above_max_tokens_returns_invalid_argument() { - let (mut client, server_task, _engine_task) = grpc_test_server( - b"engine-grpc-stream-min-above-max", - default_stream_output_specs(), - ) - .await; + let (mut client, _control_client, _health_client, server_task, _engine_task) = + grpc_test_server( + b"engine-grpc-stream-min-above-max", + default_stream_output_specs(), + ) + .await; let status = client .generate_stream(pb::GenerateRequest { @@ -716,7 +742,7 @@ async fn streaming_generate_min_tokens_above_max_tokens_returns_invalid_argument #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn unary_generate_with_sampling_params() { - let (mut client, server_task, engine_task) = + let (mut client, _control_client, _health_client, server_task, engine_task) = grpc_test_server(b"engine-grpc-sampling", default_stream_output_specs()).await; let response = client @@ -752,7 +778,7 @@ async fn unary_generate_with_sampling_params() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn unary_generate_rejects_wrong_model() { - let (mut client, server_task, _engine_task) = + let (mut client, _control_client, _health_client, server_task, _engine_task) = grpc_test_server(b"engine-grpc-wrong-model", default_stream_output_specs()).await; let status = client @@ -778,11 +804,12 @@ async fn unary_generate_rejects_wrong_model() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn streaming_generate_rejects_wrong_model() { - let (mut client, server_task, _engine_task) = grpc_test_server( - b"engine-grpc-stream-wrong-model", - default_stream_output_specs(), - ) - .await; + let (mut client, _control_client, _health_client, server_task, _engine_task) = + grpc_test_server( + b"engine-grpc-stream-wrong-model", + default_stream_output_specs(), + ) + .await; let status = client .generate_stream(pb::GenerateRequest { @@ -807,7 +834,7 @@ async fn streaming_generate_rejects_wrong_model() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn unary_generate_accepts_empty_model() { - let (mut client, server_task, engine_task) = + let (mut client, _control_client, _health_client, server_task, engine_task) = grpc_test_server(b"engine-grpc-empty-model", default_stream_output_specs()).await; // Empty `model` (proto3 default) is treated as "unset" and should be accepted. @@ -836,7 +863,7 @@ async fn unary_generate_accepts_empty_model() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn unary_generate_output_text_defaults_to_true() { - let (mut client, server_task, engine_task) = + let (mut client, _control_client, _health_client, server_task, engine_task) = grpc_test_server(b"engine-grpc-default-text", default_stream_output_specs()).await; // No response options at all — output_text should default to true. @@ -1034,22 +1061,33 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] -async fn extension_methods_are_unimplemented() { - let (mut client, server_task, _engine_task) = +async fn canonical_health_and_unimplemented_extensions_share_listener() { + let (_generate_client, mut client, mut health_client, server_task, _engine_task) = grpc_test_server(b"engine-grpc-stubs", default_stream_output_specs()).await; + for service in ["vllm.Generate", "vllm.Engine", ""] { + let health = health_client + .check(HealthCheckRequest { + service: service.to_string(), + }) + .await + .unwrap() + .into_inner(); + assert_eq!(health.status, HealthServingStatus::Serving as i32); + } + assert_eq!( - client.get_engine_info(pb::GetEngineInfoRequest {}).await.unwrap_err().code(), + client + .get_deployment_info(pb::GetDeploymentInfoRequest {}) + .await + .unwrap_err() + .code(), tonic::Code::Unimplemented ); assert_eq!( client.get_model_info(pb::GetModelInfoRequest {}).await.unwrap_err().code(), tonic::Code::Unimplemented ); - assert_eq!( - client.health(pb::HealthRequest::default()).await.unwrap_err().code(), - tonic::Code::Unimplemented - ); assert_eq!( client.abort(pb::AbortRequest::default()).await.unwrap_err().code(), tonic::Code::Unimplemented @@ -1070,14 +1108,6 @@ async fn extension_methods_are_unimplemented() { client.list_loras(pb::ListLorasRequest {}).await.unwrap_err().code(), tonic::Code::Unimplemented ); - assert_eq!( - client - .get_kv_connector_info(pb::GetKvConnectorInfoRequest {}) - .await - .unwrap_err() - .code(), - tonic::Code::Unimplemented - ); assert_eq!( client .get_kv_event_sources(pb::GetKvEventSourcesRequest {}) diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 008a29149e62..a8f7323bdc7d 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -36,6 +36,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}; @@ -200,12 +201,23 @@ 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(); + health_reporter + .set_serving::>() + .await; + health_reporter + .set_serving::>() + .await; + let generate_service = + grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone())); + let engine_service = grpc::EngineServer::new(grpc::EngineServiceImpl::new()); 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) + .add_service(engine_service); info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server"); Some((grpc_listener, svc, grpc_tls)) } else { From ce52c03619fa9aec3487d45ba955bd5a48563bdf Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Tue, 14 Jul 2026 16:56:17 -0700 Subject: [PATCH 3/6] refactor(grpc): clarify server metadata Signed-off-by: Connor Carpenter --- rust/proto/vllm_grpc.proto | 21 ++++++++++----------- rust/src/server/src/grpc/mod.rs | 8 ++++---- rust/src/server/src/grpc/tests.rs | 6 +----- 3 files changed, 15 insertions(+), 20 deletions(-) diff --git a/rust/proto/vllm_grpc.proto b/rust/proto/vllm_grpc.proto index f18d05530543..0ae50010fd40 100644 --- a/rust/proto/vllm_grpc.proto +++ b/rust/proto/vllm_grpc.proto @@ -12,7 +12,7 @@ service Generate { } service Engine { - rpc GetDeploymentInfo (GetDeploymentInfoRequest) returns (DeploymentInfo) {} + rpc GetServerInfo (GetServerInfoRequest) returns (ServerInfo) {} rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {} rpc Abort (AbortRequest) returns (AbortResponse) {} rpc Drain (DrainRequest) returns (DrainResponse) {} @@ -238,18 +238,19 @@ message MediaItem { // Discovery and lifecycle // ====================================================================================== -message GetDeploymentInfoRequest {} +message GetServerInfoRequest {} -message DeploymentInfo { +message ServerInfo { string engine_version = 1; string api_version = 2; string instance_id = 3; repeated string supported_models = 4; ParallelismInfo parallelism = 5; - uint32 kv_block_size = 6; - uint64 total_kv_blocks = 7; - uint64 max_running_requests = 8; - uint64 max_batched_tokens = 9; + uint32 max_model_len = 6; + uint32 kv_block_size = 7; + uint64 total_kv_blocks = 8; + uint64 max_running_requests = 9; + uint64 max_batched_tokens = 10; } message ParallelismInfo { @@ -267,10 +268,8 @@ message ModelInfo { string model_id = 1; string served_model_name = 2; repeated string served_model_aliases = 3; - uint32 max_context_length = 4; - uint32 max_output_tokens = 5; - repeated string tokenizer_modes = 6; - uint32 max_loras = 7; + repeated string tokenizer_modes = 4; + uint32 max_loras = 5; bool supports_text_input = 20; bool supports_token_ids_input = 21; diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs index b8fcc7408774..fb26dbf6343d 100644 --- a/rust/src/server/src/grpc/mod.rs +++ b/rust/src/server/src/grpc/mod.rs @@ -166,11 +166,11 @@ impl pb::generate_server::Generate for GenerateServiceImpl { #[tonic::async_trait] impl pb::engine_server::Engine for EngineServiceImpl { - async fn get_deployment_info( + async fn get_server_info( &self, - _request: Request, - ) -> Result, Status> { - Err(Status::unimplemented("GetDeploymentInfo")) + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented("GetServerInfo")) } async fn get_model_info( diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index fb14118d5889..8106a9f3b52e 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1077,11 +1077,7 @@ async fn canonical_health_and_unimplemented_extensions_share_listener() { } assert_eq!( - client - .get_deployment_info(pb::GetDeploymentInfoRequest {}) - .await - .unwrap_err() - .code(), + client.get_server_info(pb::GetServerInfoRequest {}).await.unwrap_err().code(), tonic::Code::Unimplemented ); assert_eq!( From 2e15e1ca9ff3b61cdf78c34de4edc97826f69477 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Tue, 14 Jul 2026 17:37:04 -0700 Subject: [PATCH 4/6] refactor(grpc): align model and server metadata Signed-off-by: Connor Carpenter --- rust/proto/vllm_grpc.proto | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/rust/proto/vllm_grpc.proto b/rust/proto/vllm_grpc.proto index 0ae50010fd40..8a6af960b2d9 100644 --- a/rust/proto/vllm_grpc.proto +++ b/rust/proto/vllm_grpc.proto @@ -244,13 +244,13 @@ message ServerInfo { string engine_version = 1; string api_version = 2; string instance_id = 3; - repeated string supported_models = 4; - ParallelismInfo parallelism = 5; - uint32 max_model_len = 6; - uint32 kv_block_size = 7; - uint64 total_kv_blocks = 8; - uint64 max_running_requests = 9; - uint64 max_batched_tokens = 10; + ParallelismInfo parallelism = 4; + uint32 max_model_len = 5; + uint32 kv_block_size = 6; + uint64 total_kv_blocks = 7; + uint64 max_running_requests = 8; + uint64 max_batched_tokens = 9; + uint32 max_loras = 10; } message ParallelismInfo { @@ -269,7 +269,6 @@ message ModelInfo { string served_model_name = 2; repeated string served_model_aliases = 3; repeated string tokenizer_modes = 4; - uint32 max_loras = 5; bool supports_text_input = 20; bool supports_token_ids_input = 21; From f0edef260eef43e9582f327932cc9b52d41f4426 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Wed, 15 Jul 2026 10:26:52 -0700 Subject: [PATCH 5/6] refactor(grpc): rename Engine service to Control Signed-off-by: Connor Carpenter --- rust/proto/vllm_grpc.proto | 2 +- rust/src/server/src/grpc/mod.rs | 8 ++++---- rust/src/server/src/grpc/tests.rs | 18 +++++++++--------- rust/src/server/src/lib.rs | 6 +++--- 4 files changed, 17 insertions(+), 17 deletions(-) diff --git a/rust/proto/vllm_grpc.proto b/rust/proto/vllm_grpc.proto index 8a6af960b2d9..545c42336350 100644 --- a/rust/proto/vllm_grpc.proto +++ b/rust/proto/vllm_grpc.proto @@ -11,7 +11,7 @@ service Generate { rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {} } -service Engine { +service Control { rpc GetServerInfo (GetServerInfoRequest) returns (ServerInfo) {} rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {} rpc Abort (AbortRequest) returns (AbortResponse) {} diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs index fb26dbf6343d..a52a7f79ab60 100644 --- a/rust/src/server/src/grpc/mod.rs +++ b/rust/src/server/src/grpc/mod.rs @@ -21,7 +21,7 @@ pub mod pb { tonic::include_proto!("vllm"); } -pub use pb::engine_server::EngineServer; +pub use pb::control_server::ControlServer; pub use pb::generate_server::GenerateServer; #[cfg(test)] @@ -40,9 +40,9 @@ impl GenerateServiceImpl { /// Unimplemented control-plane service registered on the existing gRPC listener. #[derive(Default)] -pub struct EngineServiceImpl; +pub struct ControlServiceImpl; -impl EngineServiceImpl { +impl ControlServiceImpl { pub fn new() -> Self { Self } @@ -165,7 +165,7 @@ impl pb::generate_server::Generate for GenerateServiceImpl { } #[tonic::async_trait] -impl pb::engine_server::Engine for EngineServiceImpl { +impl pb::control_server::Control for ControlServiceImpl { async fn get_server_info( &self, _request: Request, diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 8106a9f3b52e..4e8a5e580b03 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -35,9 +35,9 @@ use vllm_tokenizer::test_utils::TestTokenizer; use zeromq::prelude::{SocketRecv, SocketSend}; use zeromq::{DealerSocket, PushSocket, ZmqMessage}; -use super::pb::engine_client::EngineClient; +use super::pb::control_client::ControlClient; use super::pb::generate_client::GenerateClient; -use super::{EngineServer, EngineServiceImpl, GenerateServer, GenerateServiceImpl, pb}; +use super::{ControlServer, ControlServiceImpl, GenerateServer, GenerateServiceImpl, pb}; use crate::listener::{Listener, MaybeTlsListener}; use crate::state::AppState; use crate::tls; @@ -252,16 +252,16 @@ async fn grpc_test_server( output_specs: Vec<(Vec, Option)>, ) -> ( GenerateClient, - EngineClient, + ControlClient, HealthClient, tokio::task::JoinHandle<()>, MockEngineTask, ) { let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await; - let engine_service = EngineServer::new(EngineServiceImpl::new()); + let control_service = ControlServer::new(ControlServiceImpl::new()); let (health_reporter, health_service) = health_reporter(); health_reporter.set_serving::>().await; - health_reporter.set_serving::>().await; + 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"); @@ -271,7 +271,7 @@ async fn grpc_test_server( TonicServer::builder() .add_service(health_service) .add_service(svc) - .add_service(engine_service) + .add_service(control_service) .serve_with_incoming(incoming) .await .expect("grpc server"); @@ -283,12 +283,12 @@ async fn grpc_test_server( .await .expect("connect grpc channel"); let grpc_client = GenerateClient::new(channel.clone()); - let engine_client = EngineClient::new(channel.clone()); + let control_client = ControlClient::new(channel.clone()); let health_client = HealthClient::new(channel); ( grpc_client, - engine_client, + control_client, health_client, server_task, engine_task, @@ -1065,7 +1065,7 @@ async fn canonical_health_and_unimplemented_extensions_share_listener() { let (_generate_client, mut client, mut health_client, server_task, _engine_task) = grpc_test_server(b"engine-grpc-stubs", default_stream_output_specs()).await; - for service in ["vllm.Generate", "vllm.Engine", ""] { + for service in ["vllm.Generate", "vllm.Control", ""] { let health = health_client .check(HealthCheckRequest { service: service.to_string(), diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index a8f7323bdc7d..f7f818487b9a 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -206,18 +206,18 @@ where .set_serving::>() .await; health_reporter - .set_serving::>() + .set_serving::>() .await; let generate_service = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone())); - let engine_service = grpc::EngineServer::new(grpc::EngineServiceImpl::new()); + let control_service = grpc::ControlServer::new(grpc::ControlServiceImpl::new()); 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(health_service) .add_service(generate_service) - .add_service(engine_service); + .add_service(control_service); info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server"); Some((grpc_listener, svc, grpc_tls)) } else { From a4a4977940381ad42d05613b451780f9ff0e0c91 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Wed, 15 Jul 2026 11:45:24 -0700 Subject: [PATCH 6/6] feat(grpc): report engine failures through health Signed-off-by: Connor Carpenter --- rust/src/engine-core-client/src/client.rs | 6 + rust/src/engine-core-client/src/client/imp.rs | 40 +++-- rust/src/server/src/grpc/tests.rs | 141 +++++++++++++++--- rust/src/server/src/lib.rs | 59 +++++++- 4 files changed, 211 insertions(+), 35 deletions(-) diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index eb899e3e7843..caf4ddb4d900 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -457,6 +457,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 c91a93d27146..46117669c181 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -6,7 +6,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; @@ -33,6 +33,7 @@ pub(crate) struct ClientInner { request_reg: Mutex, utility_reg: Mutex, health_error: ArcSwapOption, + health_tx: watch::Sender, } impl ClientInner { @@ -54,6 +55,7 @@ impl ClientInner { request_reg: Mutex::new(RequestRegistry::new(engines)), utility_reg: Mutex::new(UtilityRegistry::default()), health_error: ArcSwapOption::empty(), + health_tx: watch::channel(true).0, } } @@ -188,6 +190,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 { @@ -267,14 +275,23 @@ impl ClientInner { /// recorded for this client. Later failures do not overwrite the first /// one so `/health` and post-close callers observe a stable cause. fn record_health_error(&self, error: Arc) -> Arc { - if let Some(existing) = self.health_error.load_full() { - return existing; - } - self.health_error - .rcu(|current| current.clone().unwrap_or_else(|| error.clone())); - self.health_error - .load_full() - .expect("health error must be recorded before registries close") + let persistent_error = if let Some(existing) = self.health_error.load_full() { + existing + } else { + self.health_error + .rcu(|current| current.clone().unwrap_or_else(|| error.clone())); + self.health_error + .load_full() + .expect("health error must be recorded before registries close") + }; + self.health_tx.send_if_modified(|healthy| { + if !*healthy { + return false; + } + *healthy = false; + true + }); + persistent_error } /// Assert there is a recorded health error and return a `Shared` variant @@ -458,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/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 4e8a5e580b03..cd97cfe6472d 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -27,7 +27,9 @@ use vllm_engine_core_client::protocol::output::{ }; use vllm_engine_core_client::protocol::request::EngineCoreRequest; use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task}; -use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId}; +use vllm_engine_core_client::{ + ENGINE_CORE_DEAD_SENTINEL, EngineCoreClient, EngineCoreClientConfig, EngineId, +}; use vllm_llm::Llm; use vllm_text::tokenizer::DynTokenizer; use vllm_text::{Prompt, TextBackend}; @@ -201,7 +203,37 @@ 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, +) { + setup_grpc_service_with_engine(engine_id, move |dealer, push| { + boxed_test_future(async move { + let add = recv_engine_message(dealer).await; + let request: EngineCoreRequest = + rmp_serde::from_slice(&add[1]).expect("decode request"); + send_outputs( + push, + engine_outputs_for_request(&request.request_id, output_specs), + ) + .await; + }) + }) + .await +} + +async fn setup_grpc_service_with_engine( + engine_id: impl Into, + run: F, +) -> ( + GenerateServer, + tokio::sync::watch::Receiver, + MockEngineTask, +) +where + F: for<'a> FnOnce(&'a mut DealerSocket, &'a mut PushSocket) -> TestFuture<'a> + Send + 'static, +{ let ipc = IpcNamespace::new().expect("create ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let engine_id = engine_id.into(); @@ -209,18 +241,7 @@ async fn setup_grpc_service( let engine_task = MockEngineTask::new(spawn_mock_engine_task( handshake_address.clone(), engine_id.clone(), - move |dealer, push| { - boxed_test_future(async move { - let add = recv_engine_message(dealer).await; - let request: EngineCoreRequest = - rmp_serde::from_slice(&add[1]).expect("decode request"); - send_outputs( - push, - engine_outputs_for_request(&request.request_id, output_specs), - ) - .await; - }) - }, + run, )); let client = EngineCoreClient::connect( @@ -233,6 +254,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), @@ -241,6 +263,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, ) } @@ -257,7 +280,21 @@ 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; + start_grpc_test_server(svc, engine_health, engine_task).await +} + +async fn start_grpc_test_server( + generate_service: GenerateServer, + engine_health: tokio::sync::watch::Receiver, + engine_task: MockEngineTask, +) -> ( + GenerateClient, + ControlClient, + HealthClient, + tokio::task::JoinHandle<()>, + MockEngineTask, +) { let control_service = ControlServer::new(ControlServiceImpl::new()); let (health_reporter, health_service) = health_reporter(); health_reporter.set_serving::>().await; @@ -268,13 +305,18 @@ async fn grpc_test_server( let server_task = tokio::spawn(async move { let incoming = MaybeTlsListener::plain(Listener::Tcp(listener)); - TonicServer::builder() + let server = TonicServer::builder() .add_service(health_service) - .add_service(svc) + .add_service(generate_service) .add_service(control_service) - .serve_with_incoming(incoming) - .await - .expect("grpc server"); + .serve_with_incoming(incoming); + let health_monitor = crate::monitor_grpc_health( + health_reporter, + engine_health, + tokio_util::sync::CancellationToken::new(), + ); + let (server_result, ()) = tokio::join!(server, health_monitor); + server_result.expect("grpc server"); }); let channel = Endpoint::from_shared(format!("http://{addr}")) @@ -303,7 +345,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"); @@ -393,7 +435,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(); @@ -1115,3 +1158,57 @@ async fn canonical_health_and_unimplemented_extensions_share_listener() { server_task.abort(); } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn grpc_health_transitions_to_not_serving_when_engine_fails() { + let (fail_tx, fail_rx) = tokio::sync::oneshot::channel(); + let (generate_service, engine_health, engine_task) = + setup_grpc_service_with_engine(b"engine-grpc-health-failure", move |_dealer, push| { + boxed_test_future(async move { + fail_rx.await.expect("trigger engine failure"); + push.send(ZmqMessage::from(ENGINE_CORE_DEAD_SENTINEL.to_vec())) + .await + .expect("send engine-dead sentinel"); + }) + }) + .await; + let (_generate_client, _control_client, mut health_client, server_task, engine_task) = + start_grpc_test_server(generate_service, engine_health, engine_task).await; + + let mut health_streams = Vec::new(); + for service in ["vllm.Generate", "vllm.Control", ""] { + let mut stream = health_client + .watch(HealthCheckRequest { + service: service.to_string(), + }) + .await + .expect("start health watch") + .into_inner(); + let initial = stream + .message() + .await + .expect("read initial health") + .expect("initial health response"); + assert_eq!(initial.status, HealthServingStatus::Serving as i32); + health_streams.push((service, stream)); + } + + fail_tx.send(()).expect("trigger engine failure"); + + for (service, mut stream) in health_streams { + let update = tokio::time::timeout(Duration::from_secs(2), stream.message()) + .await + .expect("health update timeout") + .expect("read health update") + .expect("health update response"); + assert_eq!( + update.status, + HealthServingStatus::NotServing as i32, + "unexpected health status for {service}" + ); + } + + server_task.abort(); + engine_task.await.expect("mock engine task"); +} diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index f7f818487b9a..06bdab848c2f 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -33,10 +33,12 @@ use hyper_util::rt::{TokioIo, TokioTimer}; use hyper_util::server::graceful::GracefulShutdown; use hyper_util::service::TowerToHyperService; use tokio::net::TcpListener; +use tokio::sync::watch; use tokio::time::{Instant, sleep_until}; use tokio_util::sync::CancellationToken; use tonic::transport::Server as TonicServer; -use tonic_health::server::health_reporter; +use tonic_health::ServingStatus; +use tonic_health::server::{HealthReporter, health_reporter}; use tower::ServiceExt as _; use tracing::{info, trace, warn}; use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends}; @@ -57,6 +59,36 @@ const GRPC_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(7200); /// connection. 20s matches the gRPC-core default. const GRPC_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(20); +async fn wait_until_engine_unhealthy(mut engine_health: watch::Receiver) { + loop { + if !*engine_health.borrow_and_update() { + return; + } + if engine_health.changed().await.is_err() { + return; + } + } +} + +async fn monitor_grpc_health( + health_reporter: HealthReporter, + engine_health: watch::Receiver, + shutdown: CancellationToken, +) { + tokio::select! { + _ = wait_until_engine_unhealthy(engine_health) => {} + _ = shutdown.cancelled() => {} + } + + health_reporter + .set_not_serving::>() + .await; + health_reporter + .set_not_serving::>() + .await; + health_reporter.set_service_status("", ServingStatus::NotServing).await; +} + /// Resolve the public model names accepted by the frontend. fn effective_served_model_names(model: &str, served_model_name: &[String]) -> Vec { if served_model_name.is_empty() { @@ -202,6 +234,7 @@ where .transpose() .context("invalid gRPC TLS configuration")?; let (health_reporter, health_service) = health_reporter(); + let engine_health = state.engine_core_client().subscribe_health(); health_reporter .set_serving::>() .await; @@ -219,7 +252,7 @@ where .add_service(generate_service) .add_service(control_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 }; @@ -239,6 +272,14 @@ where let force_shutdown = CancellationToken::new(); let shutdown_deadline = Arc::new(OnceLock::new()); + let (grpc_server_setup, grpc_health_setup) = match grpc_setup { + Some((listener, service, tls, reporter, engine_health)) => ( + Some((listener, service, tls)), + Some((reporter, engine_health)), + ), + None => (None, None), + }; + // Spawn a task to trigger `force_shutdown` after shutdown deadline elapses. tokio::spawn({ let shutdown = server_shutdown.clone(); @@ -303,7 +344,7 @@ 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)) = grpc_server_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; @@ -330,7 +371,17 @@ where } }; - let (http_res, grpc_res) = tokio::join!(http_fut, grpc_fut); + let grpc_health_fut = { + let shutdown = server_shutdown.child_token(); + async move { + let Some((health_reporter, engine_health)) = grpc_health_setup else { + return; + }; + monitor_grpc_health(health_reporter, engine_health, shutdown).await; + } + }; + + let (http_res, grpc_res, ()) = tokio::join!(http_fut, grpc_fut, grpc_health_fut); http_res.and(grpc_res)?; let shutdown_deadline = shutdown_deadline