[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>

rust/src/engine-core-client/src/client.rs

@@ -5,7 +5,7 @@ use std::collections::BTreeMap;
 use std::sync::Arc;
 use std::time::Duration;
 
-use futures::future::{join_all, try_join_all};
+use futures::future::join_all;
 use itertools::Itertools;
 use serde::Serialize;
 use serde_json::Value as JsonValue;
@@ -450,6 +450,11 @@ impl EngineCoreClient {
         self.config.transport_mode.data_parallel_size()
     }
 
+    #[cfg(test)]
+    pub(crate) fn pending_utility_call_count(&self) -> usize {
+        self.inner.pending_utility_call_count()
+    }
+
     /// Return the engine-side indices connected to this client.
     pub fn engine_indices(&self) -> Vec<u32> {
         self.engines
@@ -596,6 +601,16 @@ impl EngineCoreClient {
         let (engine_id, rx) =
             self.inner.register_request(request_id.clone(), lora_name, data_parallel_rank)?;
 
+        // Construct the output stream first before actually sending the request to the engine.
+        // This ensures that cancelling the future will properly clean up the request registration
+        // via `Drop` of the stream.
+        let stream = EngineCoreOutputStream::new(
+            request_id.clone(),
+            engine_id.engine_index().unwrap_or(0),
+            self.abort_tx.clone(),
+            rx,
+        );
+
         let result: Result<()> = async {
             if let Some(coordinator) = self.coordinator.as_ref() {
                 let snapshot = coordinator.snapshot();
@@ -622,12 +637,7 @@ impl EngineCoreClient {
             return Err(error);
         }
 
-        Ok(EngineCoreOutputStream::new(
-            request_id,
-            engine_id.engine_index().unwrap_or(0),
-            self.abort_tx.clone(),
-            rx,
-        ))
+        Ok(stream)
     }
 
     /// Abort currently in-flight requests by request ID.
@@ -651,8 +661,9 @@ impl EngineCoreClient {
     }
 
     /// Call a typed utility method on all connected engines, returning one
-    /// decoded result per connected engine if all calls succeed or an error
-    /// if any call fails.
+    /// decoded result per connected engine if all calls succeed. The client
+    /// waits for every engine outcome before returning an error so callers can
+    /// safely compensate partially applied mutations.
     ///
     /// Callers should pass utility arguments using Rust tuple semantics so the
     /// encoded payload matches Python's `(client_index, call_id,
@@ -669,62 +680,47 @@ impl EngineCoreClient {
             "sending utility request"
         );
 
-        // Phase 1: allocate one call id per engine and build the per-engine
-        // request payloads up-front. Any failure here (registry closed, encode
-        // error) must roll back the call ids already allocated so they do not
-        // leak in the utility registry until shutdown.
-        let mut pending_calls = Vec::with_capacity(self.engines.len());
-        let mut prepared_sends = Vec::with_capacity(self.engines.len());
+        /// Removes utility waiters if a call future is cancelled before completion.
+        struct UtilityCallGuard<'a> {
+            inner: &'a ClientInner,
+            call_ids: Vec<u64>,
+        }
+
+        impl Drop for UtilityCallGuard<'_> {
+            fn drop(&mut self) {
+                self.inner.unregister_utility_calls(self.call_ids.drain(..));
+            }
+        }
+
+        let mut call_guard = UtilityCallGuard {
+            inner: self.inner.as_ref(),
+            call_ids: Vec::with_capacity(self.engines.len()),
+        };
+
+        let mut prepared_calls = Vec::with_capacity(self.engines.len());
         for engine in &self.engines {
-            let (call_id, rx) = match self.inner.allocate_and_register_utility_call() {
-                Ok(pair) => pair,
-                Err(err) => {
-                    self.inner.unregister_utility_calls(pending_calls.iter().map(|(id, _)| *id));
-                    return Err(err);
-                }
-            };
-            let request = match EngineCoreUtilityRequest::new(
-                self.config.client_index,
-                call_id,
-                method,
-                &args,
-            ) {
-                Ok(request) => request,
-                Err(err) => {
-                    self.inner.unregister_utility_calls(
-                        pending_calls.iter().map(|(id, _)| *id).chain(std::iter::once(call_id)),
-                    );
-                    return Err(err);
-                }
-            };
-            pending_calls.push((call_id, rx));
-            prepared_sends.push((&engine.engine_id, request));
+            let (call_id, rx) = self.inner.allocate_and_register_utility_call()?;
+            call_guard.call_ids.push(call_id);
+            let request =
+                EngineCoreUtilityRequest::new(self.config.client_index, call_id, method, &args)?;
+            prepared_calls.push((&engine.engine_id, call_id, rx, request));
         }
 
-        // Phase 2: dispatch every utility request concurrently. `try_join_all`
-        // fails fast on the first transport error and drops the remaining send
-        // futures; any engines that already received the request will reply,
-        // but those replies are simply dropped because we roll back the call
-        // ids below.
-        let send_futures = prepared_sends.iter().map(|(engine_id, request)| {
-            self.inner.send_to_engine(engine_id, EngineCoreRequestType::Utility, request)
-        });
-        if let Err(err) = try_join_all(send_futures).await {
-            self.inner.unregister_utility_calls(pending_calls.iter().map(|(id, _)| *id));
-            return Err(err);
-        }
-
-        // Phase 3: wait for all engines to respond and preserve the per-engine
-        // result list.
-        let futures = pending_calls.into_iter().map(|(call_id, rx)| async move {
-            rx.await
-                .map_err(|_| Error::UtilityCallClosed {
-                    method: method.to_string(),
-                    call_id,
-                })??
-                .into_typed_result(method)
-        });
-        try_join_all(futures).await
+        let outcomes = join_all(prepared_calls.into_iter().map(
+            |(engine_id, call_id, rx, request)| async move {
+                self.inner
+                    .send_to_engine(engine_id, EngineCoreRequestType::Utility, &request)
+                    .await?;
+                rx.await
+                    .map_err(|_| Error::UtilityCallClosed {
+                        method: method.to_string(),
+                        call_id,
+                    })??
+                    .into_typed_result(method)
+            },
+        ))
+        .await;
+        outcomes.into_iter().collect()
     }
 
     /// Call a utility method on all connected engines and return the shared

rust/src/engine-core-client/src/client/imp.rs

@@ -118,6 +118,11 @@ impl ClientInner {
         self.utility_reg.lock().unregister_many(call_ids);
     }
 
+    #[cfg(test)]
+    pub fn pending_utility_call_count(&self) -> usize {
+        self.utility_reg.lock().len()
+    }
+
     /// Undo a request registration when `add_request()` fails.
     pub fn rollback_request(&self, request_id: &str) {
         let _ = self.request_reg.lock().remove(request_id);

rust/src/engine-core-client/src/client/state.rs

@@ -441,6 +441,11 @@ impl UtilityRegistry {
         self.utility_calls.contains_key(&call_id)
     }
 
+    #[cfg(test)]
+    pub fn len(&self) -> usize {
+        self.utility_calls.len()
+    }
+
     pub fn is_closed(&self) -> bool {
         self.closed
     }

rust/src/engine-core-client/src/tests/client.rs

@@ -1431,23 +1431,20 @@ async fn is_sleeping_wrapper_sends_typed_request_and_returns_typed_response() {
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn call_utility_failure_message_surfaces_as_error() {
+async fn call_utility_waits_for_all_engines_before_returning_error() {
     init_tracing();
     let ipc = IpcNamespace::new().unwrap();
     let handshake_address = ipc.handshake_endpoint();
-    let engine_id = b"engine-utility-fail".to_vec();
-
-    let (shutdown_tx, engine_task) = spawn_mock_engine_task(
+    let (failure_sent_tx, failure_sent_rx) = oneshot::channel();
+    let (shutdown_tx_0, engine_task_0) = spawn_mock_engine_task(
         handshake_address.clone(),
-        engine_id.clone(),
-        |dealer, push| {
+        EngineId::from_engine_index(0).into_frame().to_vec(),
+        move |dealer, push| {
             Box::pin(async move {
                 let utility = recv_engine_message(dealer).await;
-                assert_eq!(utility[0].as_ref(), &[0x03]);
                 let payload = decode_value(&utility[1]);
                 let call_id =
                     payload.as_array().and_then(|array| array[1].as_u64()).expect("call_id");
-
                 send_outputs(
                     push,
                     UtilityCallOutput {
@@ -1462,35 +1459,134 @@ async fn call_utility_failure_message_surfaces_as_error() {
                     .into(),
                 )
                 .await;
+                let _ = failure_sent_tx.send(());
+            })
+        },
+    );
+    let (second_received_tx, second_received_rx) = oneshot::channel();
+    let (release_second_tx, release_second_rx) = oneshot::channel();
+    let (shutdown_tx_1, engine_task_1) = spawn_mock_engine_task(
+        handshake_address.clone(),
+        EngineId::from_engine_index(1).into_frame().to_vec(),
+        move |dealer, push| {
+            Box::pin(async move {
+                let utility = recv_engine_message(dealer).await;
+                let payload = decode_value(&utility[1]);
+                let call_id =
+                    payload.as_array().and_then(|array| array[1].as_u64()).expect("call_id");
+                let _ = second_received_tx.send(());
+                let _ = release_second_rx.await;
+                send_outputs(
+                    push,
+                    UtilityCallOutput {
+                        engine_index: 1,
+                        timestamp: 0.0,
+                        output: UtilityOutput {
+                            call_id: call_id.into(),
+                            failure_message: None,
+                            result: Some(utility_result_value(true)),
+                        },
+                    }
+                    .into(),
+                )
+                .await;
             })
         },
     );
 
-    let client = connect_client_with_ipc(
-        handshake_test_config(
-            handshake_address,
-            1,
-            "test-model",
-            Duration::from_secs(2),
-            0,
-            None,
-        ),
-        &ipc,
-    )
-    .await;
+    let client = std::sync::Arc::new(
+        connect_client_with_ipc(
+            handshake_test_config(
+                handshake_address,
+                2,
+                "test-model",
+                Duration::from_secs(2),
+                0,
+                None,
+            ),
+            &ipc,
+        )
+        .await,
+    );
+    let call_client = client.clone();
+    let call =
+        tokio::spawn(async move { call_client.call_utility::<bool, _>("test_mutation", ()).await });
 
-    let error = client.call_utility::<bool, _>("is_sleeping", ()).await.unwrap_err();
+    failure_sent_rx.await.unwrap();
+    second_received_rx.await.unwrap();
+    timeout(Duration::from_secs(2), async {
+        while client.pending_utility_call_count() != 1 {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .expect("wait for first engine failure");
+    assert!(!call.is_finished());
+
+    let _ = release_second_tx.send(());
+    let error = call.await.unwrap().unwrap_err();
     assert!(matches!(
         error,
         Error::UtilityCallFailed {
             method,
             message,
             ..
-        } if method == "is_sleeping" && message == "boom"
+        } if method == "test_mutation" && message == "boom"
     ));
 
-    let _ = shutdown_tx.send(());
-    engine_task.await.unwrap();
+    let _ = shutdown_tx_0.send(());
+    let _ = shutdown_tx_1.send(());
+    engine_task_0.await.unwrap();
+    engine_task_1.await.unwrap();
+    let client = std::sync::Arc::try_unwrap(client)
+        .unwrap_or_else(|_| panic!("utility task retained client after completion"));
+    client.shutdown().await.unwrap();
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn cancelled_utility_call_unregisters_waiter() {
+    init_tracing();
+    let ipc = IpcNamespace::new().unwrap();
+    let handshake_address = ipc.handshake_endpoint();
+    let (received_tx, received_rx) = oneshot::channel();
+    let (_shutdown, engine_task) = spawn_mock_engine_task(
+        handshake_address.clone(),
+        vec![0x00, 0x00],
+        |dealer, _push| {
+            Box::pin(async move {
+                let _utility = recv_engine_message(dealer).await;
+                let _ = received_tx.send(());
+                std::future::pending::<()>().await;
+            })
+        },
+    );
+    let client = std::sync::Arc::new(
+        connect_client_with_ipc(
+            handshake_test_config(
+                handshake_address,
+                1,
+                "test-model",
+                Duration::from_secs(2),
+                0,
+                None,
+            ),
+            &ipc,
+        )
+        .await,
+    );
+    let call_client = client.clone();
+    let call =
+        tokio::spawn(async move { call_client.call_utility::<bool, _>("add_lora", ()).await });
+
+    received_rx.await.unwrap();
+    assert_eq!(client.pending_utility_call_count(), 1);
+    call.abort();
+    assert!(call.await.unwrap_err().is_cancelled());
+    assert_eq!(client.pending_utility_call_count(), 0);
+
+    engine_task.abort();
+    let client = std::sync::Arc::try_unwrap(client)
+        .unwrap_or_else(|_| panic!("utility task retained client after cancellation"));
     client.shutdown().await.unwrap();
 }