A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 1, 21:15 UTC
[Rust Frontend] Batch auto-abort requests by engine (#44591)
Signed-off-by: Hugh Ryan <197298026+HueCodes@users.noreply.github.com> Co-authored-by: Bugen Zhao <i@bugenzhao.com>
HueCodes committed
Jun 5, 2026 at 02:59 UTC
8a83e6f2d752660dfe598395fb13809bb55fb53f
2 files changed
+118
-25
rust/src/engine-core-client/src/client/imp.rs
+38
-25
@@ -1,5 +1,4 @@
1
use std::collections::BTreeMap;
2
-use std::slice;
2
use std::sync::Arc;
3
4
use arc_swap::ArcSwapOption;
@@ -253,33 +252,47 @@ pub(crate) async fn run_abort_loop(
252
inner: Arc<ClientInner>,
253
mut abort_rx: mpsc::UnboundedReceiver<AbortRequest>,
254
) {
256
- // TODO: receive and abort requests in batch
257
- while let Some(AbortRequest { request_id, cause }) = abort_rx.recv().await {
258
- let Some(engine_id) = inner.take_auto_abort_target(&request_id) else {
259
- debug!(request_id, "skip auto-abort for inactive request");
260
- continue;
261
- };
262
-
263
- match cause {
264
- AbortCause::DroppedStream => {
265
- info!(request_id, "auto-aborting request due to dropped stream")
266
- }
267
- AbortCause::StopStringMatched => {
268
- debug!(
269
- request_id,
270
- "auto-aborting request due to stop string matched"
271
- )
255
+ // Coalesce bursts of auto-aborts into a single Abort message per engine.
256
+ // A dropped-stream storm (e.g. many clients disconnecting at once under
257
+ // high concurrency) would otherwise issue one engine round-trip per
258
+ // request. `recv_many` returns as soon as at least one item is ready, so a
259
+ // lone abort is still forwarded promptly.
260
+ const MAX_DRAIN: usize = 1024;
261
+ let mut batch: Vec<AbortRequest> = Vec::new();
262
+
263
+ while abort_rx.recv_many(&mut batch, MAX_DRAIN).await > 0 {
264
+ let mut by_engine: BTreeMap<EngineId, Vec<String>> = BTreeMap::new();
265
+
266
+ for AbortRequest { request_id, cause } in batch.drain(..) {
267
+ let Some(engine_id) = inner.take_auto_abort_target(&request_id) else {
268
+ debug!(request_id, "skip auto-abort for inactive request");
269
+ continue;
270
+ };
271
+
272
+ match cause {
273
+ AbortCause::DroppedStream => {
274
+ info!(request_id, "auto-aborting request due to dropped stream")
275
+ }
276
+ AbortCause::StopStringMatched => {
277
+ debug!(
278
+ request_id,
279
+ "auto-aborting request due to stop string matched"
280
+ )
281
+ }
282
}
283
+
284
+ by_engine.entry(engine_id).or_default().push(request_id);
285
}
286
275
- if let Err(error) = inner.do_abort_requests(&engine_id, slice::from_ref(&request_id)).await
276
- {
277
- warn!(
278
- request_id,
279
- ?engine_id,
280
- error = %error.as_report(),
281
- "failed to auto-abort dropped request stream"
282
- );
287
+ for (engine_id, request_ids) in by_engine {
288
+ if let Err(error) = inner.do_abort_requests(&engine_id, &request_ids).await {
289
+ warn!(
290
+ ?engine_id,
291
+ ?request_ids,
292
+ error = %error.as_report(),
293
+ "failed to auto-abort request streams"
294
+ );
295
+ }
296
}
297
}
298
}
rust/src/engine-core-client/src/tests/client.rs
+80
@@ -1225,6 +1225,86 @@ async fn dropping_a_live_stream_triggers_abort() {
1225
client.shutdown().await.unwrap();
1226
}
1227
1228
+#[tokio::test]
1229
+async fn dropping_multiple_live_streams_aborts_all_in_a_burst() {
1230
+ init_tracing();
1231
+ let ipc = IpcNamespace::new().unwrap();
1232
+ let handshake_address = ipc.handshake_endpoint();
1233
+ let engine_id = b"engine-burst".to_vec();
1234
+ let request_ids = ["req-1", "req-2", "req-3"];
1235
+
1236
+ let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1237
+ handshake_address.clone(),
1238
+ engine_id.clone(),
1239
+ |dealer, push| {
1240
+ Box::pin(async move {
1241
+ for _ in 0..3 {
1242
+ let add = recv_engine_message(dealer).await;
1243
+ assert_eq!(add[0].as_ref(), &[0x00]);
1244
+ }
1245
+ send_outputs(
1246
+ push,
1247
+ EngineCoreOutputs {
1248
+ outputs: vec![
1249
+ request_output("req-1", vec![99], None),
1250
+ request_output("req-2", vec![99], None),
1251
+ request_output("req-3", vec![99], None),
1252
+ ],
1253
+ ..Default::default()
1254
+ },
1255
+ )
1256
+ .await;
1257
+
1258
+ let abort =
1259
+ timeout(Duration::from_secs(1), recv_engine_message(dealer)).await.unwrap();
1260
+ assert_eq!(abort[0].as_ref(), &[0x01]);
1261
+ let ids: Vec<String> = rmp_serde::from_slice(&abort[1]).unwrap();
1262
+ assert_eq!(
1263
+ ids,
1264
+ vec![
1265
+ "req-1".to_string(),
1266
+ "req-2".to_string(),
1267
+ "req-3".to_string()
1268
+ ]
1269
+ );
1270
+ assert!(
1271
+ timeout(Duration::from_millis(100), recv_engine_message(dealer)).await.is_err()
1272
+ );
1273
+ })
1274
+ },
1275
+ );
1276
+
1277
+ let client = connect_client_with_ipc(
1278
+ handshake_test_config(
1279
+ handshake_address,
1280
+ 1,
1281
+ "test-model",
1282
+ Duration::from_secs(2),
1283
+ 0,
1284
+ None,
1285
+ ),
1286
+ &ipc,
1287
+ )
1288
+ .await;
1289
+
1290
+ // Open every request first so all three adds reach the engine before it
1291
+ // emits outputs, then drain the first token from each stream.
1292
+ let mut streams = Vec::new();
1293
+ for id in request_ids {
1294
+ streams.push(client.call(sample_request_with_id(id)).await.unwrap());
1295
+ }
1296
+ for stream in streams.iter_mut() {
1297
+ let first = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap();
1298
+ assert_eq!(first.new_token_ids, vec![99]);
1299
+ }
1300
+ // Drop the whole burst back-to-back so the abort worker can batch them.
1301
+ drop(streams);
1302
+
1303
+ let _ = shutdown_tx.send(());
1304
+ engine_task.await.unwrap();
1305
+ client.shutdown().await.unwrap();
1306
+}
1307
+
1308
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1309
async fn dispatcher_failure_propagates_to_streams_and_future_calls() {
1310
init_tracing();