A high-throughput and memory-efficient inference and serving engine for LLMs
[Rust Frontend] [Refactor] Refine utility call interfaces (#44856)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Bugen Zhao committed
Jun 9, 2026 at 06:13 UTC
bc941f375de8c3b27973550a8791428ba7c773d4
5 files changed
+125
-82
rust/src/engine-core-client/examples/external_engine_utility_call.rs
+4
-3
@@ -3,6 +3,7 @@ use std::time::Duration;
3
use anyhow::{Context, Result, bail};
4
use clap::Parser;
5
use tracing_subscriber::EnvFilter;
6
+use vllm_engine_core_client::protocol::utility::PauseMode;
7
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
8
9
#[derive(Debug, Parser)]
@@ -32,8 +33,8 @@ struct Args {
33
reset_external: bool,
34
#[arg(long, default_value_t = 1)]
35
sleep_level: u32,
35
- #[arg(long, default_value = "abort")]
36
- sleep_mode: String,
36
+ #[arg(long, default_value_t = PauseMode::Abort)]
37
+ sleep_mode: PauseMode,
38
#[arg(
39
long,
40
default_value_t = false,
@@ -106,7 +107,7 @@ async fn main() -> Result<()> {
107
if args.skip_sleep_wake {
108
println!("sleep_wake=skipped");
109
} else {
109
- client.sleep(args.sleep_level, &args.sleep_mode).await.with_context(|| {
110
+ client.sleep(args.sleep_level, args.sleep_mode).await.with_context(|| {
111
format!(
112
"failed to call sleep utility with level={} mode={}",
113
args.sleep_level, args.sleep_mode
rust/src/engine-core-client/src/client.rs
+30
-51
@@ -2,6 +2,7 @@ use std::sync::Arc;
2
use std::time::Duration;
3
4
use futures::future::{join_all, try_join_all};
5
+use itertools::Itertools;
6
use serde::Serialize;
7
use tokio::sync::mpsc;
8
use tokio_util::task::AbortOnDropHandle;
@@ -12,7 +13,7 @@ use crate::coordinator::CoordinatorHandle;
13
use crate::error::{Error, Result};
14
use crate::protocol::handshake::EngineCoreReadyResponse;
15
use crate::protocol::lora::LoraRequest;
15
-use crate::protocol::utility::EngineCoreUtilityRequest;
16
+use crate::protocol::utility::{EngineCoreUtilityRequest, PauseMode};
17
use crate::protocol::{EngineCoreRequest, EngineCoreRequestType, ModelDtype};
18
use crate::transport::{self, ConnectedEngine};
19
@@ -571,6 +572,26 @@ impl EngineCoreClient {
572
try_join_all(futures).await
573
}
574
575
+ /// Call a utility method on all connected engines and return the shared
576
+ /// result if every engine agrees.
577
+ pub async fn call_utility_consensus<T, A>(&self, method: &str, args: A) -> Result<T>
578
+ where
579
+ T: serde::de::DeserializeOwned + std::fmt::Debug + PartialEq,
580
+ A: serde::Serialize + std::fmt::Debug,
581
+ {
582
+ let results: Vec<T> = self.call_utility(method, args).await?;
583
+
584
+ if results.iter().all_equal() {
585
+ // `engine_count >= 1` is enforced during startup handshake so `results` must be non-empty.
586
+ Ok(results.into_iter().next().unwrap())
587
+ } else {
588
+ Err(Error::InconsistentUtilityResults {
589
+ method: method.to_string(),
590
+ values: format!("{results:?}"),
591
+ })
592
+ }
593
+ }
594
+
595
/// Execute `collective_rpc` on all engines and flatten all engine results
596
/// into one list.
597
pub async fn collective_rpc<A, K>(
@@ -599,27 +620,8 @@ impl EngineCoreClient {
620
}
621
622
/// Return whether the engine is currently sleeping at any level.
602
- ///
603
- /// Under data parallel, all engines should agree on the sleep state: a
604
- /// divergence signals a control-plane bug. Returns
605
- /// `Error::InconsistentUtilityResults` if engines disagree.
623
pub async fn is_sleeping(&self) -> Result<bool> {
607
- let results: Vec<bool> = self.call_utility("is_sleeping", ()).await?;
608
- // `engine_count >= 1` is enforced during startup handshake, so `results`
609
- // is normally non-empty; fall back to a fail-loud error rather than
610
- // indexing in case that invariant is ever bypassed.
611
- let first = *results.first().ok_or_else(|| Error::InconsistentUtilityResults {
612
- method: "is_sleeping".to_string(),
613
- values: "[]".to_string(),
614
- })?;
615
- if results.iter().all(|&v| v == first) {
616
- Ok(first)
617
- } else {
618
- Err(Error::InconsistentUtilityResults {
619
- method: "is_sleeping".to_string(),
620
- values: format!("{results:?}"),
621
- })
622
- }
624
+ self.call_utility_consensus("is_sleeping", ()).await
625
}
626
627
/// Reset the multi-modal cache.
@@ -643,22 +645,14 @@ impl EngineCoreClient {
645
reset_running_requests: bool,
646
reset_connector: bool,
647
) -> Result<bool> {
646
- let results: Vec<bool> = self
648
+ Ok(self
649
.call_utility(
650
"reset_prefix_cache",
651
(reset_running_requests, reset_connector),
652
)
651
- .await?;
652
- // `engine_count >= 1` is enforced during startup handshake, so `results`
653
- // is normally non-empty; fail loud rather than reporting a vacuous
654
- // success (`[].all() == true`) in case that invariant is ever bypassed.
655
- if results.is_empty() {
656
- return Err(Error::InconsistentUtilityResults {
657
- method: "reset_prefix_cache".to_string(),
658
- values: "[]".to_string(),
659
- });
660
- }
661
- Ok(results.into_iter().all(|ok| ok))
653
+ .await?
654
+ .into_iter()
655
+ .all(|reset| reset))
656
}
657
658
/// Load or refresh one LoRA adapter on every connected engine.
@@ -680,7 +674,7 @@ impl EngineCoreClient {
674
}
675
676
/// Put the engine to sleep.
683
- pub async fn sleep(&self, level: u32, mode: &str) -> Result<()> {
677
+ pub async fn sleep(&self, level: u32, mode: PauseMode) -> Result<()> {
678
self.call_utility::<(), _>("sleep", (level, mode)).await?;
679
Ok(())
680
}
@@ -693,7 +687,7 @@ impl EngineCoreClient {
687
}
688
689
/// Pause the scheduler so generation can be halted
696
- pub async fn pause_scheduler(&self, mode: &str, clear_cache: bool) -> Result<()> {
690
+ pub async fn pause_scheduler(&self, mode: PauseMode, clear_cache: bool) -> Result<()> {
691
self.call_utility::<(), _>("pause_scheduler", (mode, clear_cache)).await?;
692
Ok(())
693
}
@@ -706,22 +700,7 @@ impl EngineCoreClient {
700
701
/// Return whether the scheduler is currently in any pause state.
702
pub async fn is_scheduler_paused(&self) -> Result<bool> {
709
- let results: Vec<bool> = self.call_utility("is_scheduler_paused", ()).await?;
710
- // `engine_count >= 1` is enforced during startup handshake, so `results`
711
- // is normally non-empty; fall back to a fail-loud error rather than
712
- // indexing in case that invariant is ever bypassed.
713
- let first = *results.first().ok_or_else(|| Error::InconsistentUtilityResults {
714
- method: "is_scheduler_paused".to_string(),
715
- values: "[]".to_string(),
716
- })?;
717
- if results.iter().all(|&v| v == first) {
718
- Ok(first)
719
- } else {
720
- Err(Error::InconsistentUtilityResults {
721
- method: "is_scheduler_paused".to_string(),
722
- values: format!("{results:?}"),
723
- })
724
- }
703
+ self.call_utility_consensus("is_scheduler_paused", ()).await
704
}
705
706
/// Shut down local client tasks and close transport state.
rust/src/engine-core-client/src/protocol/utility.rs
+70
-2
@@ -1,15 +1,63 @@
1
use std::any::type_name;
2
-use std::fmt;
2
+use std::{fmt, str::FromStr};
3
4
use rmpv::Value;
5
use serde::{Deserialize, Serialize};
6
use serde_default::DefaultFromSerde;
7
use serde_tuple::{Deserialize_tuple, Serialize_tuple};
8
+use serde_with::{DeserializeFromStr, SerializeDisplay};
9
use thiserror_ext::AsReport;
10
11
use super::{OpaqueValue, default_opaque_value_nil};
12
use crate::error::{Error, Result};
13
14
+/// How pause/sleep utility calls handle in-flight requests.
15
+///
16
+/// Use display/from-str serde so MessagePack utility args stay as Python
17
+/// literal strings instead of serde enum variant tuples.
18
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, SerializeDisplay, DeserializeFromStr)]
19
+pub enum PauseMode {
20
+ /// Abort all in-flight requests immediately.
21
+ #[default]
22
+ Abort,
23
+ /// Wait for in-flight requests to complete.
24
+ Wait,
25
+ /// Freeze queued requests so they can resume later.
26
+ Keep,
27
+}
28
+
29
+impl PauseMode {
30
+ /// Return the Python literal used on the utility-call wire.
31
+ pub fn as_str(self) -> &'static str {
32
+ match self {
33
+ Self::Abort => "abort",
34
+ Self::Wait => "wait",
35
+ Self::Keep => "keep",
36
+ }
37
+ }
38
+}
39
+
40
+impl fmt::Display for PauseMode {
41
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42
+ f.write_str(self.as_str())
43
+ }
44
+}
45
+
46
+impl FromStr for PauseMode {
47
+ type Err = String;
48
+
49
+ fn from_str(value: &str) -> std::result::Result<Self, Self::Err> {
50
+ match value {
51
+ "abort" => Ok(Self::Abort),
52
+ "wait" => Ok(Self::Wait),
53
+ "keep" => Ok(Self::Keep),
54
+ other => Err(format!(
55
+ "invalid pause mode `{other}`; expected one of: abort, wait, keep"
56
+ )),
57
+ }
58
+ }
59
+}
60
+
61
/// Utility call id as carried on the engine-core MessagePack wire.
62
///
63
/// Python emits utility ids as MessagePack integers, including values that may
@@ -212,7 +260,7 @@ mod tests {
260
use rmpv::Value;
261
use serde::Serialize;
262
215
- use super::{EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope};
263
+ use super::{EngineCoreUtilityRequest, PauseMode, UtilityOutput, UtilityResultEnvelope};
264
use crate::Error;
265
use crate::protocol::{decode_msgpack, decode_value, encode_msgpack};
266
@@ -241,6 +289,26 @@ mod tests {
289
assert_eq!(array[3], Value::Array(Vec::new()));
290
}
291
292
+ #[test]
293
+ fn pause_mode_serializes_as_python_literal() {
294
+ let request =
295
+ EngineCoreUtilityRequest::new(7, 42, "pause_scheduler", (PauseMode::Abort, true))
296
+ .unwrap();
297
+
298
+ let encoded = encode_msgpack(&request).unwrap();
299
+ let value = decode_value(&encoded).unwrap();
300
+ let array = match value {
301
+ Value::Array(array) => array,
302
+ other => panic!("expected utility request array, got {other:?}"),
303
+ };
304
+
305
+ assert_eq!(array[2], Value::from("pause_scheduler"));
306
+ assert_eq!(
307
+ array[3],
308
+ Value::Array(vec![Value::from("abort"), Value::from(true)])
309
+ );
310
+ }
311
+
312
#[test]
313
fn utility_output_decodes_typed_result() {
314
let output = UtilityOutput {
rust/src/server/src/routes/pause.rs
+11
-20
@@ -1,8 +1,10 @@
1
use std::sync::Arc;
2
3
use axum::Json;
4
+use axum::extract::rejection::QueryRejection;
5
use axum::extract::{Query, State};
6
use serde::{Deserialize, Serialize};
7
+use vllm_engine_core_client::protocol::utility::PauseMode;
8
9
use crate::error::ApiError;
10
use crate::state::AppState;
@@ -10,8 +12,8 @@ use crate::utils::utility_call_error;
12
13
#[derive(Debug, Deserialize)]
14
pub(crate) struct PauseParams {
13
- #[serde(default = "default_pause_mode")]
14
- mode: String,
15
+ #[serde(default)]
16
+ mode: PauseMode,
17
#[serde(default = "default_clear_cache")]
18
clear_cache: bool,
19
}
@@ -26,17 +28,14 @@ pub(crate) struct IsPausedResponse {
28
is_paused: bool,
29
}
30
29
-/// Pause modes accepted by the engine (`PauseMode` in the Python frontend).
30
-const VALID_PAUSE_MODES: [&str; 3] = ["abort", "wait", "keep"];
31
-
32
-fn default_pause_mode() -> String {
33
- "abort".to_string()
34
-}
35
-
31
const fn default_clear_cache() -> bool {
32
true
33
}
34
35
+fn invalid_query(error: QueryRejection) -> ApiError {
36
+ ApiError::invalid_request(error.body_text(), Some("mode"))
37
+}
38
+
39
// TODO: the Python frontend also accepts the deprecated
40
// `wait_for_inflight_requests` flag (equivalent to `mode="wait"`); it is
41
// intentionally omitted here in favor of the `mode` parameter.
@@ -44,21 +43,13 @@ const fn default_clear_cache() -> bool {
43
/// Pause the scheduler so generation can be halted (e.g. for weight updates).
44
pub async fn pause(
45
State(state): State<Arc<AppState>>,
47
- Query(params): Query<PauseParams>,
46
+ params: Result<Query<PauseParams>, QueryRejection>,
47
) -> Result<Json<StatusResponse>, ApiError> {
49
- if !VALID_PAUSE_MODES.contains(¶ms.mode.as_str()) {
50
- return Err(ApiError::invalid_request(
51
- format!(
52
- "Invalid pause mode '{}'; expected one of: abort, wait, keep",
53
- params.mode
54
- ),
55
- Some("mode"),
56
- ));
57
- }
48
+ let Query(params) = params.map_err(invalid_query)?;
49
50
state
51
.engine_core_client()
61
- .pause_scheduler(¶ms.mode, params.clear_cache)
52
+ .pause_scheduler(params.mode, params.clear_cache)
53
.await
54
.map_err(|error| utility_call_error("pause", error))?;
55
rust/src/server/src/routes/sleep.rs
+10
-6
@@ -1,9 +1,11 @@
1
use std::sync::Arc;
2
3
use axum::Json;
4
+use axum::extract::rejection::QueryRejection;
5
use axum::extract::{Query, State};
6
use axum::http::StatusCode;
7
use serde::{Deserialize, Serialize};
8
+use vllm_engine_core_client::protocol::utility::PauseMode;
9
10
use crate::error::ApiError;
11
use crate::state::AppState;
@@ -18,8 +20,8 @@ pub(crate) struct IsSleepingResponse {
20
pub(crate) struct SleepParams {
21
#[serde(default = "default_sleep_level")]
22
level: u32,
21
- #[serde(default = "default_sleep_mode")]
22
- mode: String,
23
+ #[serde(default)]
24
+ mode: PauseMode,
25
}
26
27
#[derive(Debug, Default, Deserialize)]
@@ -32,18 +34,20 @@ const fn default_sleep_level() -> u32 {
34
1
35
}
36
35
-fn default_sleep_mode() -> String {
36
- "abort".to_string()
37
+fn invalid_query(error: QueryRejection) -> ApiError {
38
+ ApiError::invalid_request(error.body_text(), Some("mode"))
39
}
40
41
/// Put the engine to sleep.
42
pub async fn sleep(
43
State(state): State<Arc<AppState>>,
42
- Query(params): Query<SleepParams>,
44
+ params: Result<Query<SleepParams>, QueryRejection>,
45
) -> Result<StatusCode, ApiError> {
46
+ let Query(params) = params.map_err(invalid_query)?;
47
+
48
state
49
.engine_core_client()
46
- .sleep(params.level, ¶ms.mode)
50
+ .sleep(params.level, params.mode)
51
.await
52
.map_err(|error| utility_call_error("sleep", error))?;
53