@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f28a508162

[Feature][Rust Frontend] Add Shutdown control RPC (#59316)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Julien Darve <jdarve@NVIDIA.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

JulienDarve committed Sep 30, 2026 at 05:49 UTC f28a5081629377c36cd219a35070555b72279109

10 files changed +129 -9

docs/usage/security.md

+1 -1

@@ -375,7 +375,7 @@ An attacker who can reach the gRPC port can:

375 1. **Run arbitrary inference** via the `Generate` and `GenerateStream` RPCs without any credentials

376 2. **Mutate engine state** by pausing generation, sleeping the engine, or initiating configured RL weight updates through the `Control` service

377 3. **Consume GPU and compute resources** by submitting unbounded generation requests

378 -4. **Cause Denial of Service** by exploiting bugs in the gRPC interface that can crash vLLM.

378 +4. **Stop a managed engine** through `Control.Shutdown`, or cause denial of service by exploiting bugs in the gRPC interface.

379

380 ### Recommendations

381

rust/proto/README.md

+1 -1

@@ -16,7 +16,7 @@ the schemas on pull requests, including on forks.

17 ```toml

18 [dependencies]

19 -vllm-proto = "0.3"

19 +vllm-proto = "0.4"

20 ```

22 For example, import `vllm_proto::inference_client::InferenceClient` or

rust/proto/control.proto

@@ -9,6 +9,10 @@ service Control {

9 rpc GetServerInfo (GetServerInfoRequest) returns (ServerInfo) {}

10 rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {}

12 + // Request managed-engine shutdown using the configured shutdown timeout.

13 + // Returns FAILED_PRECONDITION when this server does not manage the engine process.

14 + rpc Shutdown (ShutdownRequest) returns (ShutdownResponse) {}

15 +

16 // Request lifecycle.

17 rpc Abort (AbortRequest) returns (AbortResponse) {}

@@ -86,6 +90,10 @@ message ModelInfo {

90 string tool_call_parser = 25;

91 }

93 +message ShutdownRequest {}

94 +// Acknowledges acceptance, not completion; shutdown may close the connection first.

95 +message ShutdownResponse {}

96 +

97 message AbortRequest {

98 repeated string request_ids = 1;

99 }

rust/src/cmd/src/cli.rs

@@ -557,6 +557,8 @@ impl SharedRuntimeArgs {

557 disable_log_stats: self.disable_log_stats,

558 grpc_port: self.grpc_port,

559 shutdown_timeout,

560 + // The engine is launched and supervised by another process.

561 + manages_engine: false,

562 keep_alive_timeout,

563 profiler,

564 }

@@ -572,6 +574,7 @@ impl SharedRuntimeArgs {

574 engine_count: usize,

575 local_input_address: Option<String>,

576 local_output_address: Option<String>,

577 + manages_engine: bool,

578 ) -> Config {

579 let ready_timeout = self.ready_timeout();

580 let shutdown_timeout = self.shutdown_timeout();

@@ -615,6 +618,7 @@ impl SharedRuntimeArgs {

618 disable_log_stats: self.disable_log_stats,

619 grpc_port: self.grpc_port,

620 shutdown_timeout,

621 + manages_engine,

622 keep_alive_timeout,

623 profiler,

624 }

@@ -793,6 +797,8 @@ impl ServeArgs {

797 self.managed_engine.data_parallel_size,

798 local_input_address,

799 local_output_address,

800 + // `--data-parallel-size-local 0` runs the frontend without a local engine.

801 + self.managed_engine.data_parallel_size_local != Some(0),

802 )

803 }

804

rust/src/cmd/src/cli/tests.rs

+26

@@ -1889,6 +1889,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {

1889 disable_log_stats: false,

1890 grpc_port: None,

1891 shutdown_timeout: 0ns,

1892 + manages_engine: true,

1893 keep_alive_timeout: 5s,

1894 profiler: None,

1895 }

@@ -1983,6 +1984,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {

1984 disable_log_stats: false,

1985 grpc_port: None,

1986 shutdown_timeout: 0ns,

1987 + manages_engine: true,

1988 keep_alive_timeout: 5s,

1989 profiler: None,

1990 }

@@ -2011,6 +2013,29 @@ fn frontend_args_reject_legacy_handshake_flags() {

2013 assert!(error.to_string().contains("--handshake-address"));

2014 }

2015

2016 +#[test]

2017 +fn serve_frontend_config_does_not_manage_engine_without_local_engines() {

2018 + let cli = Cli::try_parse_from([

2019 + "vllm-rs",

2020 + "serve",

2021 + "Qwen/Qwen3-0.6B",

2022 + "--data-parallel-address",

2023 + "10.99.48.128",

2024 + "--data-parallel-size",

2025 + "2",

2026 + "--data-parallel-size-local",

2027 + "0",

2028 + ])

2029 + .unwrap();

2030 +

2031 + let Command::Serve(args) = cli.command else {

2032 + panic!("expected serve args");

2033 + };

2034 + let config = args.to_frontend_config("tcp://10.99.48.128:29550".to_string());

2035 +

2036 + assert!(!config.manages_engine);

2037 +}

2038 +

2039 #[test]

2040 fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present() {

2041 let cli = Cli::try_parse_from([

@@ -2098,6 +2123,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present

2123 disable_log_stats: false,

2124 grpc_port: None,

2125 shutdown_timeout: 0ns,

2126 + manages_engine: false,

2127 keep_alive_timeout: 5s,

2128 profiler: None,

2129 }

rust/src/server/examples/external_engine_openai_qwen.rs

@@ -87,6 +87,7 @@ async fn main() -> Result<()> {

87 disable_log_stats: false,

88 grpc_port: None,

89 shutdown_timeout: Duration::ZERO,

90 + manages_engine: false,

91 keep_alive_timeout: Duration::from_secs(5),

92 profiler: None,

93 };

rust/src/server/src/config.rs

@@ -277,6 +277,9 @@ pub struct Config {

277 pub grpc_port: Option<u16>,

278 /// Maximum time to wait for active HTTP/gRPC requests to drain on shutdown.

279 pub shutdown_timeout: Duration,

280 + /// Whether the caller manages the engine process and shuts it down when

281 + /// the shutdown token is cancelled. Enables the gRPC `Control.Shutdown` RPC.

282 + pub manages_engine: bool,

283 /// Maximum idle time on a keep-alive HTTP connection before the server

284 /// closes it (`VLLM_HTTP_TIMEOUT_KEEP_ALIVE`, default 5s).

285 pub keep_alive_timeout: Duration,

rust/src/server/src/grpc/control.rs

+20

@@ -6,6 +6,7 @@ use std::sync::Arc;

6 use serde_json::Value as JsonValue;

7 use thiserror_ext::AsReport as _;

8 use tokio::sync::Mutex;

9 +use tokio_util::sync::CancellationToken;

10 use tonic::{Code, Request, Response, Status};

11 use vllm_engine_core_client::EngineCoreClient;

12 use vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse;

@@ -23,6 +24,7 @@ pub(crate) type ControlGrpcService = ControlServer<ControlServiceImpl>;

24 pub struct ControlServiceImpl {

25 state: Arc<AppState>,

26 rl_lock: Mutex<()>,

27 + engine_shutdown: Option<CancellationToken>,

28 }

30 impl ControlServiceImpl {

@@ -30,9 +32,15 @@ impl ControlServiceImpl {

32 Self {

33 state,

34 rl_lock: Mutex::new(()),

35 + engine_shutdown: None,

36 }

37 }

39 + pub fn with_engine_shutdown(mut self, shutdown: Option<CancellationToken>) -> Self {

40 + self.engine_shutdown = shutdown;

41 + self

42 + }

43 +

44 fn ready(&self) -> &EngineCoreReadyResponse {

45 self.state.engine_core_client().ready_response()

46 }

@@ -179,6 +187,18 @@ fn list_loras_status(error: LoraDisabledError) -> Status {

187

188 #[tonic::async_trait]

189 impl pb::control_server::Control for ControlServiceImpl {

190 + async fn shutdown(

191 + &self,

192 + _request: Request<pb::ShutdownRequest>,

193 + ) -> Result<Response<pb::ShutdownResponse>, Status> {

194 + let shutdown = self.engine_shutdown.as_ref().ok_or_else(|| {

195 + Status::failed_precondition("this frontend does not manage the engine process")

196 + })?;

197 + tracing::info!("engine shutdown requested via gRPC");

198 + shutdown.cancel();

199 + Ok(Response::new(pb::ShutdownResponse {}))

200 + }

201 +

202 async fn get_server_info(

203 &self,

204 _request: Request<pb::GetServerInfoRequest>,

rust/src/server/src/grpc/tests.rs

+58 -4

@@ -366,6 +366,7 @@ where

366 engine_id,

367 default_ready_response(),

368 backend,

369 + None,

370 move |dealer, push| {

371 boxed_test_future(async move {

372 let add = recv_engine_message(dealer).await;

@@ -387,6 +388,7 @@ async fn setup_grpc_service_with_engine_script<F>(

388 engine_id: impl Into<EngineId>,

389 ready: EngineCoreReadyResponse,

390 backend: Arc<dyn ChatTextBackend>,

391 + engine_shutdown: Option<tokio_util::sync::CancellationToken>,

392 script: F,

393 ) -> (

394 InferenceServer<InferenceServiceImpl>,

@@ -425,7 +427,7 @@ where

427 (

428 InferenceServer::new(InferenceServiceImpl::new(state.clone()))

429 .max_decoding_message_size(crate::DEFAULT_REQUEST_BODY_LIMIT_BYTES),

428 - ControlServer::new(ControlServiceImpl::new(state)),

430 + ControlServer::new(ControlServiceImpl::new(state).with_engine_shutdown(engine_shutdown)),

431 engine_health,

432 engine_task,

433 )

@@ -1920,6 +1922,7 @@ async fn control_reports_server_and_model_info() {

1922 b"engine-grpc-info".to_vec(),

1923 ready,

1924 Arc::new(FakeTextBackend),

1925 + None,

1926 |_, _| boxed_test_future(async {}),

1927 )

1928 .await;

@@ -1987,6 +1990,7 @@ async fn control_lora_lifecycle_selects_adapter_for_generation() {

1990 b"engine-grpc-lora".to_vec(),

1991 ready,

1992 Arc::new(FakeTextBackend),

1993 + None,

1994 |dealer, push| {

1995 boxed_test_future(async move {

1996 reply_utility_bool(dealer, push, "add_lora", true).await;

@@ -2111,6 +2115,7 @@ async fn control_list_loras_requires_lora_enabled_engine() {

2115 b"engine-grpc-lora-disabled".to_vec(),

2116 default_ready_response(),

2117 Arc::new(FakeTextBackend),

2118 + None,

2119 |_, _| boxed_test_future(async move {}),

2120 )

2121 .await;

@@ -2143,6 +2148,7 @@ async fn control_forwards_weight_update_without_pause_guard() {

2148 b"engine-grpc-rl".to_vec(),

2149 ready,

2150 Arc::new(FakeTextBackend),

2151 + None,

2152 |dealer, push| {

2153 boxed_test_future(async move {

2154 let frames = recv_engine_message(dealer).await;

@@ -2392,9 +2398,9 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy()

2398

2399 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]

2400 #[serial]

2395 -async fn grpc_health_watch_closes_on_graceful_shutdown() {

2401 +async fn grpc_shutdown_requires_managed_engine() {

2402 let (inference_service, control_service, engine_health, _engine_task) = setup_grpc_service(

2397 - b"engine-grpc-health-shutdown",

2403 + b"engine-grpc-external-shutdown",

2404 default_stream_output_specs(),

2405 )

2406 .await;

@@ -2406,6 +2412,47 @@ async fn grpc_health_watch_closes_on_graceful_shutdown() {

2412 shutdown.clone(),

2413 )

2414 .await;

2415 + let mut control_client = ControlClient::new(channel.clone());

2416 +

2417 + let status = control_client

2418 + .shutdown(pb::ShutdownRequest {})

2419 + .await

2420 + .expect_err("external engine shutdown must be rejected");

2421 + assert_eq!(status.code(), tonic::Code::FailedPrecondition);

2422 + assert!(!shutdown.is_cancelled());

2423 +

2424 + let health = HealthClient::new(channel)

2425 + .check(HealthCheckRequest {

2426 + service: "vllm.Inference".to_string(),

2427 + })

2428 + .await

2429 + .expect("server remains healthy after rejected shutdown")

2430 + .into_inner();

2431 + assert_eq!(health.status, HealthServingStatus::Serving as i32);

2432 + server_task.abort();

2433 +}

2434 +

2435 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

2436 +#[serial]

2437 +async fn grpc_shutdown_closes_health_watch_and_server() {

2438 + let shutdown = tokio_util::sync::CancellationToken::new();

2439 + let (inference_service, control_service, engine_health, _engine_task) =

2440 + setup_grpc_service_with_engine_script(

2441 + b"engine-grpc-health-shutdown",

2442 + default_ready_response(),

2443 + Arc::new(FakeTextBackend),

2444 + Some(shutdown.clone()),

2445 + |_, _| boxed_test_future(async {}),

2446 + )

2447 + .await;

2448 + let (channel, server_task) = start_grpc_test_server(

2449 + inference_service,

2450 + control_service,

2451 + engine_health,

2452 + shutdown.child_token(),

2453 + )

2454 + .await;

2455 + let mut control_client = ControlClient::new(channel.clone());

2456 let mut health_client = HealthClient::new(channel);

2457 let mut stream = health_client

2458 .watch(HealthCheckRequest {

@@ -2426,7 +2473,14 @@ async fn grpc_health_watch_closes_on_graceful_shutdown() {

2473 "unexpected initial health status for vllm.Inference"

2474 );

2475

2429 - shutdown.cancel();

2476 + control_client

2477 + .shutdown(pb::ShutdownRequest {})

2478 + .await

2479 + .expect("accept managed engine shutdown");

2480 + assert!(

2481 + shutdown.is_cancelled(),

2482 + "engine owner must receive shutdown"

2483 + );

2484

2485 let update = tokio::time::timeout(Duration::from_secs(2), stream.message())

2486 .await

rust/src/server/src/lib.rs

+5 -3

@@ -247,9 +247,11 @@ where

247 let engine_health = state.engine_core_client().subscribe_health();

248 health_reporter.set_serving::<grpc::InferenceGrpcService>().await;

249 health_reporter.set_serving::<grpc::ControlGrpcService>().await;

250 - let control_service =

251 - grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone()))

252 - .max_decoding_message_size(DEFAULT_REQUEST_BODY_LIMIT_BYTES);

250 + let control_service = grpc::ControlGrpcService::new(

251 + grpc::ControlServiceImpl::new(state.clone())

252 + .with_engine_shutdown(config.manages_engine.then(|| shutdown.clone())),

253 + )

254 + .max_decoding_message_size(DEFAULT_REQUEST_BODY_LIMIT_BYTES);

255 let inference_service =

256 grpc::InferenceGrpcService::new(grpc::InferenceServiceImpl::new(state.clone()))

257 .max_decoding_message_size(DEFAULT_REQUEST_BODY_LIMIT_BYTES);