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

[Rust Frontend][gRPC] Add abort control RPC (#49255)

Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Connor Carpenter <connorc@nvidia.com>

Connor Carpenter committed Jul 22, 2026 at 04:31 UTC c79ad3ae213249760bc360b50eb733af1bf1f715
5 files changed +177 -29
rust/proto/vllm_grpc.proto
+13
@@ -14,6 +14,10 @@ service Generate {
14 rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {}
15 }
16
17 +service Control {
18 + rpc Abort (AbortRequest) returns (AbortResponse) {}
19 +}
20 +
21 // ======================================================================================
22 // Generate Request
23 // ======================================================================================
@@ -201,3 +205,12 @@ message TokenIds {
205 repeated uint32 ids = 1;
206 }
207
208 +// ======================================================================================
209 +// Control
210 +// ======================================================================================
211 +
212 +message AbortRequest {
213 + repeated string request_ids = 1;
214 +}
215 +
216 +message AbortResponse {}
rust/src/server/src/grpc/health.rs
+6 -11
@@ -8,7 +8,7 @@ use tonic_health::ServingStatus;
8 use tonic_health::server::HealthReporter;
9 use tracing::{info, warn};
10
11 -use super::GenerateGrpcService;
11 +use super::{ControlGrpcService, GenerateGrpcService};
12
13 pub(crate) async fn monitor_health(
14 mut health_reporter: HealthReporter,
@@ -16,21 +16,18 @@ pub(crate) async fn monitor_health(
16 shutdown: CancellationToken,
17 ) {
18 let generate_service = GenerateGrpcService::NAME;
19 + let control_service = ControlGrpcService::NAME;
20 let status = ServingStatus::NotServing;
21 let health_event_first = tokio::select! {
22 result = engine_health.wait_for(|healthy| !*healthy) => {
23 match result {
24 Ok(_) => warn!(
24 - generate_service,
25 - overall_service = true,
25 status = ?status,
26 reason = "engine_unhealthy",
27 "marking gRPC health services as not serving"
28 ),
29 Err(error) => warn!(
30 %error,
32 - generate_service,
33 - overall_service = true,
31 status = ?status,
32 reason = "health_channel_closed",
33 "engine health channel closed; marking gRPC health services as not serving"
@@ -40,8 +37,6 @@ pub(crate) async fn monitor_health(
37 }
38 _ = shutdown.cancelled() => {
39 info!(
43 - generate_service,
44 - overall_service = true,
40 status = ?status,
41 reason = "server_shutdown",
42 "server shutting down; marking gRPC health services as not serving"
@@ -51,20 +46,20 @@ pub(crate) async fn monitor_health(
46 };
47
48 health_reporter.set_not_serving::<GenerateGrpcService>().await;
54 - // Generate is currently the only engine-backed gRPC service, so overall
55 - // server health intentionally mirrors it.
49 + health_reporter.set_not_serving::<ControlGrpcService>().await;
50 + // Both gRPC services use the same engine client, so overall server health
51 + // mirrors their shared engine health.
52 health_reporter.set_service_status("", status).await;
53
54 if health_event_first {
55 shutdown.cancelled().await;
56 info!(
61 - generate_service,
62 - overall_service = true,
57 reason = "server_shutdown",
58 "server shutting down; closing gRPC health watches"
59 );
60 }
61
62 health_reporter.clear_service_status(generate_service).await;
63 + health_reporter.clear_service_status(control_service).await;
64 health_reporter.clear_service_status("").await;
65 }
rust/src/server/src/grpc/mod.rs
+32
@@ -26,8 +26,10 @@ pub mod pb {
26 }
27
28 pub(crate) use health::monitor_health;
29 +pub use pb::control_server::ControlServer;
30 pub use pb::generate_server::GenerateServer;
31
32 +pub(crate) type ControlGrpcService = ControlServer<ControlServiceImpl>;
33 pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>;
34
35 #[cfg(test)]
@@ -44,6 +46,36 @@ impl GenerateServiceImpl {
46 }
47 }
48
49 +/// gRPC control service backed by the shared application state.
50 +pub struct ControlServiceImpl {
51 + state: Arc<AppState>,
52 +}
53 +
54 +impl ControlServiceImpl {
55 + pub fn new(state: Arc<AppState>) -> Self {
56 + Self { state }
57 + }
58 +}
59 +
60 +#[tonic::async_trait]
61 +impl pb::control_server::Control for ControlServiceImpl {
62 + async fn abort(
63 + &self,
64 + request: Request<pb::AbortRequest>,
65 + ) -> Result<Response<pb::AbortResponse>, Status> {
66 + let request_ids = request.into_inner().request_ids;
67 + if request_ids.is_empty() {
68 + return Ok(Response::new(pb::AbortResponse {}));
69 + }
70 + self.state
71 + .chat
72 + .abort(&request_ids)
73 + .await
74 + .map_err(|error| Status::internal(error.to_report_string()))?;
75 + Ok(Response::new(pb::AbortResponse {}))
76 + }
77 +}
78 +
79 #[tonic::async_trait]
80 impl pb::generate_server::Generate for GenerateServiceImpl {
81 type GenerateStreamStream =
rust/src/server/src/grpc/tests.rs
+122 -18
@@ -38,8 +38,9 @@ use vllm_tokenizer::test_utils::TestTokenizer;
38 use zeromq::prelude::{SocketRecv, SocketSend};
39 use zeromq::{DealerSocket, PushSocket, ZmqMessage};
40
41 +use super::pb::control_client::ControlClient;
42 use super::pb::generate_client::GenerateClient;
42 -use super::{GenerateServer, GenerateServiceImpl, pb};
43 +use super::{ControlServer, ControlServiceImpl, GenerateServer, GenerateServiceImpl, pb};
44 use crate::listener::{Listener, MaybeTlsListener};
45 use crate::state::AppState;
46 use crate::tls;
@@ -153,10 +154,6 @@ async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec<bytes::Bytes> {
154 dealer.recv().await.expect("recv engine message").into_vec()
155 }
156
156 -fn test_llm(client: EngineCoreClient) -> Llm {
157 - Llm::new(client).with_request_id_randomization(false)
158 -}
159 -
157 #[derive(Clone, Debug)]
158 struct FakeTextBackend;
159
@@ -206,6 +203,7 @@ async fn setup_grpc_service(
203 output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
204 ) -> (
205 GenerateServer<GenerateServiceImpl>,
206 + ControlServer<ControlServiceImpl>,
207 tokio::sync::watch::Receiver<bool>,
208 MockEngineTask,
209 ) {
@@ -243,12 +241,13 @@ async fn setup_grpc_service(
241 let engine_health = client.subscribe_health();
242
243 let chat = ChatLlm::from_shared_backend(
246 - test_llm(client),
244 + Llm::new(client),
245 Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
246 );
247 let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
248 (
251 - GenerateServer::new(GenerateServiceImpl::new(state)),
249 + GenerateServer::new(GenerateServiceImpl::new(state.clone())),
250 + ControlServer::new(ControlServiceImpl::new(state)),
251 engine_health,
252 engine_task,
253 )
@@ -264,9 +263,11 @@ async fn grpc_test_server(
263 tokio::task::JoinHandle<()>,
264 MockEngineTask,
265 ) {
267 - let (svc, engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
266 + let (generate_service, control_service, engine_health, engine_task) =
267 + setup_grpc_service(engine_id, output_specs).await;
268 let (channel, server_task) = start_grpc_test_server(
269 - svc,
269 + generate_service,
270 + control_service,
271 engine_health,
272 tokio_util::sync::CancellationToken::new(),
273 )
@@ -276,11 +277,13 @@ async fn grpc_test_server(
277
278 async fn start_grpc_test_server(
279 generate_service: GenerateServer<GenerateServiceImpl>,
280 + control_service: ControlServer<ControlServiceImpl>,
281 engine_health: tokio::sync::watch::Receiver<bool>,
282 shutdown: tokio_util::sync::CancellationToken,
283 ) -> (Channel, tokio::task::JoinHandle<()>) {
284 let (health_reporter, health_service) = health_reporter();
285 health_reporter.set_serving::<GenerateServer<GenerateServiceImpl>>().await;
286 + health_reporter.set_serving::<ControlServer<ControlServiceImpl>>().await;
287
288 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
289 let addr = listener.local_addr().expect("local addr");
@@ -289,6 +292,7 @@ async fn start_grpc_test_server(
292 let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
293 let server = TonicServer::builder()
294 .add_service(health_service)
295 + .add_service(control_service)
296 .add_service(generate_service)
297 .serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
298 let health_monitor =
@@ -319,7 +323,8 @@ async fn grpc_tls_test_server(
323 certs: &TestCerts,
324 cert_reqs: i32,
325 ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
322 - let (svc, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
326 + let (generate_service, control_service, _engine_health, engine_task) =
327 + setup_grpc_service(engine_id, output_specs).await;
328 let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
329 .expect("build grpc tls config");
330
@@ -329,7 +334,8 @@ async fn grpc_tls_test_server(
334 let server_task = tokio::spawn(async move {
335 let incoming = MaybeTlsListener::tls(Listener::Tcp(listener), context);
336 TonicServer::builder()
332 - .add_service(svc)
337 + .add_service(control_service)
338 + .add_service(generate_service)
339 .serve_with_incoming(incoming)
340 .await
341 .expect("grpc tls server");
@@ -409,7 +415,7 @@ async fn grpc_server_with_keepalive(
415 engine_id: impl Into<EngineId>,
416 keepalive: Option<Duration>,
417 ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
412 - let (svc, _engine_health, engine_task) =
418 + let (generate_service, control_service, _engine_health, engine_task) =
419 setup_grpc_service(engine_id, default_stream_output_specs()).await;
420
421 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
@@ -425,7 +431,8 @@ async fn grpc_server_with_keepalive(
431 let server_task = tokio::spawn(async move {
432 let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
433 builder
428 - .add_service(svc)
434 + .add_service(control_service)
435 + .add_service(generate_service)
436 .serve_with_incoming(incoming)
437 .await
438 .expect("grpc server");
@@ -1073,14 +1080,106 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
1080 server_task.abort();
1081 }
1082
1083 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1084 +#[serial]
1085 +async fn control_abort_resolves_external_id_and_empty_is_noop() {
1086 + let (generate_service, control_service, engine_health, engine_task) =
1087 + setup_grpc_service(b"engine-grpc-abort-active", vec![(vec![b'h' as u32], None)]).await;
1088 + let (channel, server_task) = start_grpc_test_server(
1089 + generate_service,
1090 + control_service,
1091 + engine_health,
1092 + tokio_util::sync::CancellationToken::new(),
1093 + )
1094 + .await;
1095 + let mut generate_client = GenerateClient::new(channel.clone());
1096 + let mut control_client = ControlClient::new(channel);
1097 + let request_id = "test-abort-active";
1098 +
1099 + let mut stream = generate_client
1100 + .generate_stream(pb::GenerateRequest {
1101 + request_id: request_id.to_string(),
1102 + model: "test-model".to_string(),
1103 + prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())),
1104 + stopping: Some(pb::StoppingCriteria {
1105 + max_new_tokens: 10,
1106 + ..Default::default()
1107 + }),
1108 + ..Default::default()
1109 + })
1110 + .await
1111 + .expect("start generation")
1112 + .into_inner();
1113 +
1114 + loop {
1115 + let response = tokio::time::timeout(Duration::from_secs(2), stream.message())
1116 + .await
1117 + .expect("timed out waiting for active generation output")
1118 + .expect("read active generation output")
1119 + .expect("generation ended before producing output");
1120 + if let Some(output) = response.outputs {
1121 + assert!(
1122 + output.finish_info.is_none(),
1123 + "generation finished before abort behavior was exercised"
1124 + );
1125 + break;
1126 + }
1127 + }
1128 +
1129 + control_client
1130 + .abort(pb::AbortRequest::default())
1131 + .await
1132 + .expect("empty abort should be a no-op");
1133 + assert!(
1134 + tokio::time::timeout(Duration::from_millis(100), stream.message())
1135 + .await
1136 + .is_err(),
1137 + "empty abort unexpectedly ended the active generation"
1138 + );
1139 +
1140 + control_client
1141 + .abort(pb::AbortRequest {
1142 + request_ids: vec![
1143 + request_id.to_string(),
1144 + request_id.to_string(),
1145 + "unknown".to_string(),
1146 + ],
1147 + })
1148 + .await
1149 + .expect("abort active generation");
1150 +
1151 + let finish_reason = loop {
1152 + let response = tokio::time::timeout(Duration::from_secs(2), stream.message())
1153 + .await
1154 + .expect("timed out waiting for aborted generation")
1155 + .expect("read aborted generation")
1156 + .expect("generation ended without an aborted response");
1157 + if let Some(finish_info) = response.outputs.and_then(|output| output.finish_info) {
1158 + break finish_info.finish_reason;
1159 + }
1160 + };
1161 + assert_eq!(finish_reason, pb::finish_info::FinishReason::Aborted as i32);
1162 +
1163 + control_client
1164 + .abort(pb::AbortRequest {
1165 + request_ids: vec![request_id.to_string()],
1166 + })
1167 + .await
1168 + .expect("repeated abort should be idempotent");
1169 +
1170 + engine_task.await.expect("mock engine task");
1171 + server_task.abort();
1172 +}
1173 +
1174 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1175 #[serial]
1176 async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() {
1079 - let (generate_service, _connected_engine_health, _engine_task) =
1177 + let (generate_service, control_service, _connected_engine_health, _engine_task) =
1178 setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await;
1179 let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true);
1180 let (channel, server_task) = start_grpc_test_server(
1181 generate_service,
1182 + control_service,
1183 engine_health,
1184 tokio_util::sync::CancellationToken::new(),
1185 )
@@ -1088,7 +1187,7 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy()
1187 let mut health_client = HealthClient::new(channel);
1188
1189 let mut health_streams = Vec::new();
1091 - for service in ["vllm.Generate", ""] {
1190 + for service in ["vllm.Generate", "vllm.Control", ""] {
1191 let service_label = if service.is_empty() {
1192 "overall"
1193 } else {
@@ -1143,14 +1242,19 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy()
1242 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1243 #[serial]
1244 async fn grpc_health_watch_closes_on_graceful_shutdown() {
1146 - let (generate_service, engine_health, _engine_task) = setup_grpc_service(
1245 + let (generate_service, control_service, engine_health, _engine_task) = setup_grpc_service(
1246 b"engine-grpc-health-shutdown",
1247 default_stream_output_specs(),
1248 )
1249 .await;
1250 let shutdown = tokio_util::sync::CancellationToken::new();
1152 - let (channel, server_task) =
1153 - start_grpc_test_server(generate_service, engine_health, shutdown.clone()).await;
1251 + let (channel, server_task) = start_grpc_test_server(
1252 + generate_service,
1253 + control_service,
1254 + engine_health,
1255 + shutdown.clone(),
1256 + )
1257 + .await;
1258 let mut health_client = HealthClient::new(channel);
1259 let mut stream = health_client
1260 .watch(HealthCheckRequest {
rust/src/server/src/lib.rs
+4
@@ -207,6 +207,9 @@ where
207 let (health_reporter, health_service) = health_reporter();
208 let engine_health = state.engine_core_client().subscribe_health();
209 health_reporter.set_serving::<grpc::GenerateGrpcService>().await;
210 + health_reporter.set_serving::<grpc::ControlGrpcService>().await;
211 + let control_service =
212 + grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone()));
213 let generate_service =
214 grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone()));
215 let svc = TonicServer::builder()
@@ -214,6 +217,7 @@ where
217 .http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
218 .layer(middleware::request_runtime_layer(state.clone()))
219 .add_service(health_service)
220 + .add_service(control_service)
221 .add_service(generate_service);
222 info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
223 Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))