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 2, 13:30 UTC
[rust] fix: aggregate `is_sleeping` and `reset_prefix_cache` across DP engines (#43429)
Signed-off-by: Will.hou <1205157517@qq.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Will.hou committed
May 28, 2026 at 22:56 UTC
577d69383810b6d83753b7030c56bad71f0765fd
3 files changed
+257
-5
rust/src/engine-core-client/src/client.rs
+35
-5
@@ -595,9 +595,27 @@ impl EngineCoreClient {
595
}
596
597
/// Return whether the engine is currently sleeping at any level.
598
+ ///
599
+ /// Under data parallel, all engines should agree on the sleep state: a
600
+ /// divergence signals a control-plane bug. Returns
601
+ /// `Error::InconsistentUtilityResults` if engines disagree.
602
pub async fn is_sleeping(&self) -> Result<bool> {
599
- // TODO: we only return the result of the first engine here.
600
- Ok(self.call_utility("is_sleeping", ()).await?[0])
603
+ let results: Vec<bool> = self.call_utility("is_sleeping", ()).await?;
604
+ // `engine_count >= 1` is enforced during startup handshake, so `results`
605
+ // is normally non-empty; fall back to a fail-loud error rather than
606
+ // indexing in case that invariant is ever bypassed.
607
+ let first = *results.first().ok_or_else(|| Error::InconsistentUtilityResults {
608
+ method: "is_sleeping".to_string(),
609
+ values: "[]".to_string(),
610
+ })?;
611
+ if results.iter().all(|&v| v == first) {
612
+ Ok(first)
613
+ } else {
614
+ Err(Error::InconsistentUtilityResults {
615
+ method: "is_sleeping".to_string(),
616
+ values: format!("{results:?}"),
617
+ })
618
+ }
619
}
620
621
/// Reset the multi-modal cache.
@@ -613,18 +631,30 @@ impl EngineCoreClient {
631
}
632
633
/// Reset the prefix cache and optionally the external connector cache.
634
+ ///
635
+ /// Under data parallel, returns `true` only when every engine confirms the
636
+ /// reset (AND aggregation).
637
pub async fn reset_prefix_cache(
638
&self,
639
reset_running_requests: bool,
640
reset_connector: bool,
641
) -> Result<bool> {
621
- // TODO: we only return the result of the first engine here.
622
- Ok(self
642
+ let results: Vec<bool> = self
643
.call_utility(
644
"reset_prefix_cache",
645
(reset_running_requests, reset_connector),
646
)
627
- .await?[0])
647
+ .await?;
648
+ // `engine_count >= 1` is enforced during startup handshake, so `results`
649
+ // is normally non-empty; fail loud rather than reporting a vacuous
650
+ // success (`[].all() == true`) in case that invariant is ever bypassed.
651
+ if results.is_empty() {
652
+ return Err(Error::InconsistentUtilityResults {
653
+ method: "reset_prefix_cache".to_string(),
654
+ values: "[]".to_string(),
655
+ });
656
+ }
657
+ Ok(results.into_iter().all(|ok| ok))
658
}
659
660
/// Put the engine to sleep.
rust/src/engine-core-client/src/error.rs
+2
@@ -83,6 +83,8 @@ pub enum Error {
83
},
84
#[error("utility call `{method}` closed unexpectedly (call_id={call_id})")]
85
UtilityCallClosed { method: String, call_id: u64 },
86
+ #[error("utility call `{method}` returned inconsistent results across engines: {values}")]
87
+ InconsistentUtilityResults { method: String, values: String },
88
89
/// A special variant to allow cloning the same error.
90
#[error(transparent)]
rust/src/engine-core-client/src/tests/client.rs
+220
@@ -2119,6 +2119,226 @@ async fn collective_rpc_flattens_results_from_all_engines() {
2119
client.shutdown().await.unwrap();
2120
}
2121
2122
+/// Spawn a mock engine that handles a single utility call, asserts the method
2123
+/// name and serialized args match, and replies with `result`.
2124
+fn spawn_mock_utility_engine(
2125
+ handshake_address: String,
2126
+ engine_id: Vec<u8>,
2127
+ expected_method: &'static str,
2128
+ expected_args: Value,
2129
+ result: bool,
2130
+) -> (
2131
+ tokio::sync::oneshot::Sender<()>,
2132
+ tokio::task::JoinHandle<()>,
2133
+) {
2134
+ spawn_mock_engine_task(handshake_address, engine_id, move |dealer, push| {
2135
+ Box::pin(async move {
2136
+ let utility = recv_engine_message(dealer).await;
2137
+ assert_eq!(utility[0].as_ref(), &[0x03]);
2138
+ let payload = decode_value(&utility[1]);
2139
+ let array = match payload {
2140
+ Value::Array(array) => array,
2141
+ other => panic!("expected utility payload array, got {other:?}"),
2142
+ };
2143
+ // Utility requests serialize as `(client_index, call_id, method, args)`.
2144
+ let call_id = array[1].as_u64().expect("call_id");
2145
+ assert_eq!(array[2], Value::from(expected_method));
2146
+ assert_eq!(array[3], expected_args, "unexpected utility args");
2147
+ send_outputs(
2148
+ push,
2149
+ EngineCoreOutputs {
2150
+ utility_output: Some(UtilityOutput {
2151
+ call_id: call_id.into(),
2152
+ failure_message: None,
2153
+ result: Some(utility_result_value(result)),
2154
+ }),
2155
+ ..Default::default()
2156
+ },
2157
+ )
2158
+ .await;
2159
+ })
2160
+ })
2161
+}
2162
+
2163
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2164
+async fn is_sleeping_returns_error_when_engines_disagree() {
2165
+ init_tracing();
2166
+ let ipc = IpcNamespace::new().unwrap();
2167
+ let handshake_address = ipc.handshake_endpoint();
2168
+
2169
+ let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine(
2170
+ handshake_address.clone(),
2171
+ b"engine-0".to_vec(),
2172
+ "is_sleeping",
2173
+ Value::Array(vec![]),
2174
+ true,
2175
+ );
2176
+ let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine(
2177
+ handshake_address.clone(),
2178
+ b"engine-1".to_vec(),
2179
+ "is_sleeping",
2180
+ Value::Array(vec![]),
2181
+ false,
2182
+ );
2183
+
2184
+ let client = connect_client_with_ipc(
2185
+ handshake_test_config(
2186
+ handshake_address,
2187
+ 2,
2188
+ "test-model",
2189
+ Duration::from_secs(2),
2190
+ 5,
2191
+ None,
2192
+ ),
2193
+ &ipc,
2194
+ )
2195
+ .await;
2196
+
2197
+ let error = client.is_sleeping().await.unwrap_err();
2198
+ assert!(
2199
+ matches!(
2200
+ &error,
2201
+ Error::InconsistentUtilityResults { method, .. } if method == "is_sleeping"
2202
+ ),
2203
+ "expected InconsistentUtilityResults, got {error:?}",
2204
+ );
2205
+
2206
+ let _ = shutdown_tx_0.send(());
2207
+ let _ = shutdown_tx_1.send(());
2208
+ engine_task_0.await.unwrap();
2209
+ engine_task_1.await.unwrap();
2210
+ client.shutdown().await.unwrap();
2211
+}
2212
+
2213
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2214
+async fn is_sleeping_returns_value_when_all_engines_agree() {
2215
+ init_tracing();
2216
+ let ipc = IpcNamespace::new().unwrap();
2217
+ let handshake_address = ipc.handshake_endpoint();
2218
+
2219
+ let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine(
2220
+ handshake_address.clone(),
2221
+ b"engine-0".to_vec(),
2222
+ "is_sleeping",
2223
+ Value::Array(vec![]),
2224
+ true,
2225
+ );
2226
+ let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine(
2227
+ handshake_address.clone(),
2228
+ b"engine-1".to_vec(),
2229
+ "is_sleeping",
2230
+ Value::Array(vec![]),
2231
+ true,
2232
+ );
2233
+
2234
+ let client = connect_client_with_ipc(
2235
+ handshake_test_config(
2236
+ handshake_address,
2237
+ 2,
2238
+ "test-model",
2239
+ Duration::from_secs(2),
2240
+ 5,
2241
+ None,
2242
+ ),
2243
+ &ipc,
2244
+ )
2245
+ .await;
2246
+
2247
+ assert!(client.is_sleeping().await.unwrap());
2248
+
2249
+ let _ = shutdown_tx_0.send(());
2250
+ let _ = shutdown_tx_1.send(());
2251
+ engine_task_0.await.unwrap();
2252
+ engine_task_1.await.unwrap();
2253
+ client.shutdown().await.unwrap();
2254
+}
2255
+
2256
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2257
+async fn reset_prefix_cache_returns_true_when_all_engines_succeed() {
2258
+ init_tracing();
2259
+ let ipc = IpcNamespace::new().unwrap();
2260
+ let handshake_address = ipc.handshake_endpoint();
2261
+
2262
+ let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine(
2263
+ handshake_address.clone(),
2264
+ b"engine-0".to_vec(),
2265
+ "reset_prefix_cache",
2266
+ Value::Array(vec![Value::from(false), Value::from(false)]),
2267
+ true,
2268
+ );
2269
+ let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine(
2270
+ handshake_address.clone(),
2271
+ b"engine-1".to_vec(),
2272
+ "reset_prefix_cache",
2273
+ Value::Array(vec![Value::from(false), Value::from(false)]),
2274
+ true,
2275
+ );
2276
+
2277
+ let client = connect_client_with_ipc(
2278
+ handshake_test_config(
2279
+ handshake_address,
2280
+ 2,
2281
+ "test-model",
2282
+ Duration::from_secs(2),
2283
+ 5,
2284
+ None,
2285
+ ),
2286
+ &ipc,
2287
+ )
2288
+ .await;
2289
+
2290
+ assert!(client.reset_prefix_cache(false, false).await.unwrap());
2291
+
2292
+ let _ = shutdown_tx_0.send(());
2293
+ let _ = shutdown_tx_1.send(());
2294
+ engine_task_0.await.unwrap();
2295
+ engine_task_1.await.unwrap();
2296
+ client.shutdown().await.unwrap();
2297
+}
2298
+
2299
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2300
+async fn reset_prefix_cache_returns_false_when_any_engine_fails() {
2301
+ init_tracing();
2302
+ let ipc = IpcNamespace::new().unwrap();
2303
+ let handshake_address = ipc.handshake_endpoint();
2304
+
2305
+ let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine(
2306
+ handshake_address.clone(),
2307
+ b"engine-0".to_vec(),
2308
+ "reset_prefix_cache",
2309
+ Value::Array(vec![Value::from(false), Value::from(false)]),
2310
+ true,
2311
+ );
2312
+ let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine(
2313
+ handshake_address.clone(),
2314
+ b"engine-1".to_vec(),
2315
+ "reset_prefix_cache",
2316
+ Value::Array(vec![Value::from(false), Value::from(false)]),
2317
+ false,
2318
+ );
2319
+
2320
+ let client = connect_client_with_ipc(
2321
+ handshake_test_config(
2322
+ handshake_address,
2323
+ 2,
2324
+ "test-model",
2325
+ Duration::from_secs(2),
2326
+ 5,
2327
+ None,
2328
+ ),
2329
+ &ipc,
2330
+ )
2331
+ .await;
2332
+
2333
+ assert!(!client.reset_prefix_cache(false, false).await.unwrap());
2334
+
2335
+ let _ = shutdown_tx_0.send(());
2336
+ let _ = shutdown_tx_1.send(());
2337
+ engine_task_0.await.unwrap();
2338
+ engine_task_1.await.unwrap();
2339
+ client.shutdown().await.unwrap();
2340
+}
2341
+
2342
#[test]
2343
fn python_msgpack_fixtures_match_rust_encoding() {
2344
init_tracing();