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

[Rust Frontend] Wait for all utility calls to finish (#52671)

Signed-off-by: Connor Carpenter <connorc@nvidia.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Bugen Zhao <i@bugenzhao.com>

Connor Carpenter committed Aug 18, 2026 at 11:23 UTC d75136c030cc62973dc470d1981199be8de47d62
4 files changed +187 -85
rust/src/engine-core-client/src/client.rs
+57 -61
@@ -5,7 +5,7 @@ use std::collections::BTreeMap;
5 use std::sync::Arc;
6 use std::time::Duration;
7
8 -use futures::future::{join_all, try_join_all};
8 +use futures::future::join_all;
9 use itertools::Itertools;
10 use serde::Serialize;
11 use serde_json::Value as JsonValue;
@@ -450,6 +450,11 @@ impl EngineCoreClient {
450 self.config.transport_mode.data_parallel_size()
451 }
452
453 + #[cfg(test)]
454 + pub(crate) fn pending_utility_call_count(&self) -> usize {
455 + self.inner.pending_utility_call_count()
456 + }
457 +
458 /// Return the engine-side indices connected to this client.
459 pub fn engine_indices(&self) -> Vec<u32> {
460 self.engines
@@ -596,6 +601,16 @@ impl EngineCoreClient {
601 let (engine_id, rx) =
602 self.inner.register_request(request_id.clone(), lora_name, data_parallel_rank)?;
603
604 + // Construct the output stream first before actually sending the request to the engine.
605 + // This ensures that cancelling the future will properly clean up the request registration
606 + // via `Drop` of the stream.
607 + let stream = EngineCoreOutputStream::new(
608 + request_id.clone(),
609 + engine_id.engine_index().unwrap_or(0),
610 + self.abort_tx.clone(),
611 + rx,
612 + );
613 +
614 let result: Result<()> = async {
615 if let Some(coordinator) = self.coordinator.as_ref() {
616 let snapshot = coordinator.snapshot();
@@ -622,12 +637,7 @@ impl EngineCoreClient {
637 return Err(error);
638 }
639
625 - Ok(EngineCoreOutputStream::new(
626 - request_id,
627 - engine_id.engine_index().unwrap_or(0),
628 - self.abort_tx.clone(),
629 - rx,
630 - ))
640 + Ok(stream)
641 }
642
643 /// Abort currently in-flight requests by request ID.
@@ -651,8 +661,9 @@ impl EngineCoreClient {
661 }
662
663 /// Call a typed utility method on all connected engines, returning one
654 - /// decoded result per connected engine if all calls succeed or an error
655 - /// if any call fails.
664 + /// decoded result per connected engine if all calls succeed. The client
665 + /// waits for every engine outcome before returning an error so callers can
666 + /// safely compensate partially applied mutations.
667 ///
668 /// Callers should pass utility arguments using Rust tuple semantics so the
669 /// encoded payload matches Python's `(client_index, call_id,
@@ -669,62 +680,47 @@ impl EngineCoreClient {
680 "sending utility request"
681 );
682
672 - // Phase 1: allocate one call id per engine and build the per-engine
673 - // request payloads up-front. Any failure here (registry closed, encode
674 - // error) must roll back the call ids already allocated so they do not
675 - // leak in the utility registry until shutdown.
676 - let mut pending_calls = Vec::with_capacity(self.engines.len());
677 - let mut prepared_sends = Vec::with_capacity(self.engines.len());
678 - for engine in &self.engines {
679 - let (call_id, rx) = match self.inner.allocate_and_register_utility_call() {
680 - Ok(pair) => pair,
681 - Err(err) => {
682 - self.inner.unregister_utility_calls(pending_calls.iter().map(|(id, _)| *id));
683 - return Err(err);
684 - }
685 - };
686 - let request = match EngineCoreUtilityRequest::new(
687 - self.config.client_index,
688 - call_id,
689 - method,
690 - &args,
691 - ) {
692 - Ok(request) => request,
693 - Err(err) => {
694 - self.inner.unregister_utility_calls(
695 - pending_calls.iter().map(|(id, _)| *id).chain(std::iter::once(call_id)),
696 - );
697 - return Err(err);
698 - }
699 - };
700 - pending_calls.push((call_id, rx));
701 - prepared_sends.push((&engine.engine_id, request));
683 + /// Removes utility waiters if a call future is cancelled before completion.
684 + struct UtilityCallGuard<'a> {
685 + inner: &'a ClientInner,
686 + call_ids: Vec<u64>,
687 }
688
704 - // Phase 2: dispatch every utility request concurrently. `try_join_all`
705 - // fails fast on the first transport error and drops the remaining send
706 - // futures; any engines that already received the request will reply,
707 - // but those replies are simply dropped because we roll back the call
708 - // ids below.
709 - let send_futures = prepared_sends.iter().map(|(engine_id, request)| {
710 - self.inner.send_to_engine(engine_id, EngineCoreRequestType::Utility, request)
711 - });
712 - if let Err(err) = try_join_all(send_futures).await {
713 - self.inner.unregister_utility_calls(pending_calls.iter().map(|(id, _)| *id));
714 - return Err(err);
689 + impl Drop for UtilityCallGuard<'_> {
690 + fn drop(&mut self) {
691 + self.inner.unregister_utility_calls(self.call_ids.drain(..));
692 + }
693 }
694
717 - // Phase 3: wait for all engines to respond and preserve the per-engine
718 - // result list.
719 - let futures = pending_calls.into_iter().map(|(call_id, rx)| async move {
720 - rx.await
721 - .map_err(|_| Error::UtilityCallClosed {
722 - method: method.to_string(),
723 - call_id,
724 - })??
725 - .into_typed_result(method)
726 - });
727 - try_join_all(futures).await
695 + let mut call_guard = UtilityCallGuard {
696 + inner: self.inner.as_ref(),
697 + call_ids: Vec::with_capacity(self.engines.len()),
698 + };
699 +
700 + let mut prepared_calls = Vec::with_capacity(self.engines.len());
701 + for engine in &self.engines {
702 + let (call_id, rx) = self.inner.allocate_and_register_utility_call()?;
703 + call_guard.call_ids.push(call_id);
704 + let request =
705 + EngineCoreUtilityRequest::new(self.config.client_index, call_id, method, &args)?;
706 + prepared_calls.push((&engine.engine_id, call_id, rx, request));
707 + }
708 +
709 + let outcomes = join_all(prepared_calls.into_iter().map(
710 + |(engine_id, call_id, rx, request)| async move {
711 + self.inner
712 + .send_to_engine(engine_id, EngineCoreRequestType::Utility, &request)
713 + .await?;
714 + rx.await
715 + .map_err(|_| Error::UtilityCallClosed {
716 + method: method.to_string(),
717 + call_id,
718 + })??
719 + .into_typed_result(method)
720 + },
721 + ))
722 + .await;
723 + outcomes.into_iter().collect()
724 }
725
726 /// Call a utility method on all connected engines and return the shared
rust/src/engine-core-client/src/client/imp.rs
+5
@@ -118,6 +118,11 @@ impl ClientInner {
118 self.utility_reg.lock().unregister_many(call_ids);
119 }
120
121 + #[cfg(test)]
122 + pub fn pending_utility_call_count(&self) -> usize {
123 + self.utility_reg.lock().len()
124 + }
125 +
126 /// Undo a request registration when `add_request()` fails.
127 pub fn rollback_request(&self, request_id: &str) {
128 let _ = self.request_reg.lock().remove(request_id);
rust/src/engine-core-client/src/client/state.rs
+5
@@ -441,6 +441,11 @@ impl UtilityRegistry {
441 self.utility_calls.contains_key(&call_id)
442 }
443
444 + #[cfg(test)]
445 + pub fn len(&self) -> usize {
446 + self.utility_calls.len()
447 + }
448 +
449 pub fn is_closed(&self) -> bool {
450 self.closed
451 }
rust/src/engine-core-client/src/tests/client.rs
+120 -24
@@ -1431,23 +1431,20 @@ async fn is_sleeping_wrapper_sends_typed_request_and_returns_typed_response() {
1431 }
1432
1433 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1434 -async fn call_utility_failure_message_surfaces_as_error() {
1434 +async fn call_utility_waits_for_all_engines_before_returning_error() {
1435 init_tracing();
1436 let ipc = IpcNamespace::new().unwrap();
1437 let handshake_address = ipc.handshake_endpoint();
1438 - let engine_id = b"engine-utility-fail".to_vec();
1439 -
1440 - let (shutdown_tx, engine_task) = spawn_mock_engine_task(
1438 + let (failure_sent_tx, failure_sent_rx) = oneshot::channel();
1439 + let (shutdown_tx_0, engine_task_0) = spawn_mock_engine_task(
1440 handshake_address.clone(),
1442 - engine_id.clone(),
1443 - |dealer, push| {
1441 + EngineId::from_engine_index(0).into_frame().to_vec(),
1442 + move |dealer, push| {
1443 Box::pin(async move {
1444 let utility = recv_engine_message(dealer).await;
1446 - assert_eq!(utility[0].as_ref(), &[0x03]);
1445 let payload = decode_value(&utility[1]);
1446 let call_id =
1447 payload.as_array().and_then(|array| array[1].as_u64()).expect("call_id");
1450 -
1448 send_outputs(
1449 push,
1450 UtilityCallOutput {
@@ -1462,35 +1459,134 @@ async fn call_utility_failure_message_surfaces_as_error() {
1459 .into(),
1460 )
1461 .await;
1462 + let _ = failure_sent_tx.send(());
1463 + })
1464 + },
1465 + );
1466 + let (second_received_tx, second_received_rx) = oneshot::channel();
1467 + let (release_second_tx, release_second_rx) = oneshot::channel();
1468 + let (shutdown_tx_1, engine_task_1) = spawn_mock_engine_task(
1469 + handshake_address.clone(),
1470 + EngineId::from_engine_index(1).into_frame().to_vec(),
1471 + move |dealer, push| {
1472 + Box::pin(async move {
1473 + let utility = recv_engine_message(dealer).await;
1474 + let payload = decode_value(&utility[1]);
1475 + let call_id =
1476 + payload.as_array().and_then(|array| array[1].as_u64()).expect("call_id");
1477 + let _ = second_received_tx.send(());
1478 + let _ = release_second_rx.await;
1479 + send_outputs(
1480 + push,
1481 + UtilityCallOutput {
1482 + engine_index: 1,
1483 + timestamp: 0.0,
1484 + output: UtilityOutput {
1485 + call_id: call_id.into(),
1486 + failure_message: None,
1487 + result: Some(utility_result_value(true)),
1488 + },
1489 + }
1490 + .into(),
1491 + )
1492 + .await;
1493 })
1494 },
1495 );
1496
1469 - let client = connect_client_with_ipc(
1470 - handshake_test_config(
1471 - handshake_address,
1472 - 1,
1473 - "test-model",
1474 - Duration::from_secs(2),
1475 - 0,
1476 - None,
1477 - ),
1478 - &ipc,
1479 - )
1480 - .await;
1497 + let client = std::sync::Arc::new(
1498 + connect_client_with_ipc(
1499 + handshake_test_config(
1500 + handshake_address,
1501 + 2,
1502 + "test-model",
1503 + Duration::from_secs(2),
1504 + 0,
1505 + None,
1506 + ),
1507 + &ipc,
1508 + )
1509 + .await,
1510 + );
1511 + let call_client = client.clone();
1512 + let call =
1513 + tokio::spawn(async move { call_client.call_utility::<bool, _>("test_mutation", ()).await });
1514
1482 - let error = client.call_utility::<bool, _>("is_sleeping", ()).await.unwrap_err();
1515 + failure_sent_rx.await.unwrap();
1516 + second_received_rx.await.unwrap();
1517 + timeout(Duration::from_secs(2), async {
1518 + while client.pending_utility_call_count() != 1 {
1519 + tokio::task::yield_now().await;
1520 + }
1521 + })
1522 + .await
1523 + .expect("wait for first engine failure");
1524 + assert!(!call.is_finished());
1525 +
1526 + let _ = release_second_tx.send(());
1527 + let error = call.await.unwrap().unwrap_err();
1528 assert!(matches!(
1529 error,
1530 Error::UtilityCallFailed {
1531 method,
1532 message,
1533 ..
1489 - } if method == "is_sleeping" && message == "boom"
1534 + } if method == "test_mutation" && message == "boom"
1535 ));
1536
1492 - let _ = shutdown_tx.send(());
1493 - engine_task.await.unwrap();
1537 + let _ = shutdown_tx_0.send(());
1538 + let _ = shutdown_tx_1.send(());
1539 + engine_task_0.await.unwrap();
1540 + engine_task_1.await.unwrap();
1541 + let client = std::sync::Arc::try_unwrap(client)
1542 + .unwrap_or_else(|_| panic!("utility task retained client after completion"));
1543 + client.shutdown().await.unwrap();
1544 +}
1545 +
1546 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1547 +async fn cancelled_utility_call_unregisters_waiter() {
1548 + init_tracing();
1549 + let ipc = IpcNamespace::new().unwrap();
1550 + let handshake_address = ipc.handshake_endpoint();
1551 + let (received_tx, received_rx) = oneshot::channel();
1552 + let (_shutdown, engine_task) = spawn_mock_engine_task(
1553 + handshake_address.clone(),
1554 + vec![0x00, 0x00],
1555 + |dealer, _push| {
1556 + Box::pin(async move {
1557 + let _utility = recv_engine_message(dealer).await;
1558 + let _ = received_tx.send(());
1559 + std::future::pending::<()>().await;
1560 + })
1561 + },
1562 + );
1563 + let client = std::sync::Arc::new(
1564 + connect_client_with_ipc(
1565 + handshake_test_config(
1566 + handshake_address,
1567 + 1,
1568 + "test-model",
1569 + Duration::from_secs(2),
1570 + 0,
1571 + None,
1572 + ),
1573 + &ipc,
1574 + )
1575 + .await,
1576 + );
1577 + let call_client = client.clone();
1578 + let call =
1579 + tokio::spawn(async move { call_client.call_utility::<bool, _>("add_lora", ()).await });
1580 +
1581 + received_rx.await.unwrap();
1582 + assert_eq!(client.pending_utility_call_count(), 1);
1583 + call.abort();
1584 + assert!(call.await.unwrap_err().is_cancelled());
1585 + assert_eq!(client.pending_utility_call_count(), 0);
1586 +
1587 + engine_task.abort();
1588 + let client = std::sync::Arc::try_unwrap(client)
1589 + .unwrap_or_else(|_| panic!("utility task retained client after cancellation"));
1590 client.shutdown().await.unwrap();
1591 }
1592