@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 91f5b92438
Fork 0

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 3, 00:30 UTC

[Rust Frontend] [Refactor] Extract a newtype for utility call ID (#43405)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed May 22, 2026 at 23:22 UTC 91f5b92438a568c89e8b9d6c2c55de5a552291f6
9 files changed +392 -237
rust/src/engine-core-client/src/client.rs
+2 -3
@@ -10,9 +10,8 @@ use crate::client::imp::{ClientInner, run_abort_loop, run_output_dispatcher_loop
10 use crate::coordinator::CoordinatorHandle;
11 use crate::error::{Error, Result};
12 use crate::protocol::handshake::EngineCoreReadyResponse;
13 -use crate::protocol::{
14 - EngineCoreRequest, EngineCoreRequestType, EngineCoreUtilityRequest, ModelDtype,
15 -};
13 +use crate::protocol::utility::EngineCoreUtilityRequest;
14 +use crate::protocol::{EngineCoreRequest, EngineCoreRequestType, ModelDtype};
15 use crate::transport::{self, ConnectedEngine};
16
17 pub(crate) mod imp;
rust/src/engine-core-client/src/client/imp.rs
+12 -6
@@ -16,9 +16,10 @@ use crate::client::{AbortCause, AbortRequest};
16 use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output};
17 use crate::metrics::record_scheduler_stats;
18 use crate::protocol::stats::SchedulerStats;
19 +use crate::protocol::utility::UtilityOutput;
20 use crate::protocol::{
21 ClassifiedEngineCoreOutputs, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequestType,
21 - UtilityOutput, encode_msgpack,
22 + encode_msgpack,
23 };
24 use crate::transport::{ConnectedEngine, EngineId};
25 use crate::{Error, Result, transport};
@@ -72,7 +73,7 @@ impl ClientInner {
73 }
74
75 /// Allocate the next utility `call_id` and register its waiting receiver.
75 - pub fn allocate_and_register_utility_call(&self) -> Result<(i64, UtilityReceiver)> {
76 + pub fn allocate_and_register_utility_call(&self) -> Result<(u64, UtilityReceiver)> {
77 let mut registry = self.utility_reg.lock();
78 if registry.is_closed() {
79 return Err(self.closed_error());
@@ -83,7 +84,7 @@ impl ClientInner {
84 /// Undo a batch of utility call allocations when the fan-out send fails
85 /// partway through. Silently ignores unknown call ids so callers can pass
86 /// the full set without first filtering successful sends.
86 - pub fn unregister_utility_calls(&self, call_ids: impl IntoIterator<Item = i64>) {
87 + pub fn unregister_utility_calls(&self, call_ids: impl IntoIterator<Item = u64>) {
88 self.utility_reg.lock().unregister_many(call_ids);
89 }
90
@@ -160,7 +161,12 @@ impl ClientInner {
161 /// Resolve one utility output to the waiting caller. Returns `true` if a
162 /// waiting caller existed.
163 pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool {
163 - match self.utility_reg.lock().resolve(&output.call_id) {
164 + let Some(call_id) = output.call_id.as_u64() else {
165 + // Currently, all utility call issued by the client should have unsigned call IDs.
166 + return false;
167 + };
168 +
169 + match self.utility_reg.lock().resolve(&call_id) {
170 Some(sender) => {
171 sender.send(Ok(output)).unwrap_or_default();
172 true
@@ -342,13 +348,13 @@ pub(crate) async fn run_output_dispatcher_loop(
348 let call_id = utility.output.call_id;
349 if inner.resolve_utility_output(utility.output) {
350 trace!(
345 - call_id,
351 + %call_id,
352 engine_index = utility.engine_index,
353 "resolved utility output"
354 );
355 } else {
356 warn!(
351 - call_id,
357 + %call_id,
358 engine_index = utility.engine_index,
359 "dropping output for unexpected utility call"
360 );
rust/src/engine-core-client/src/client/state.rs
+10 -9
@@ -1,5 +1,5 @@
1 use std::collections::BTreeMap;
2 -use std::sync::atomic::{AtomicI64, Ordering};
2 +use std::sync::atomic::{AtomicU64, Ordering};
3
4 use tokio::sync::{mpsc, oneshot};
5 use tracing::trace;
@@ -7,8 +7,9 @@ use tracing::trace;
7 use crate::EngineId;
8 use crate::client::stream::EngineCoreStreamOutput;
9 use crate::error::{Error, Result};
10 +use crate::protocol::EngineCoreOutput;
11 use crate::protocol::stats::SchedulerStats;
11 -use crate::protocol::{EngineCoreOutput, UtilityOutput};
12 +use crate::protocol::utility::UtilityOutput;
13 use crate::transport::ConnectedEngine;
14
15 pub type OutputSender = mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>;
@@ -266,15 +267,15 @@ impl RequestRegistry {
267 #[derive(Debug)]
268 pub struct UtilityRegistry {
269 closed: bool,
269 - next_call_id: AtomicI64,
270 - utility_calls: BTreeMap<i64, UtilitySender>,
270 + next_call_id: AtomicU64,
271 + utility_calls: BTreeMap<u64, UtilitySender>,
272 }
273
274 impl Default for UtilityRegistry {
275 fn default() -> Self {
276 Self {
277 closed: false,
277 - next_call_id: AtomicI64::new(1),
278 + next_call_id: AtomicU64::new(1),
279 utility_calls: BTreeMap::default(),
280 }
281 }
@@ -283,7 +284,7 @@ impl Default for UtilityRegistry {
284 impl UtilityRegistry {
285 /// Allocate the next utility `call_id` and register a newly added utility
286 /// call.
286 - pub fn allocate_and_register(&mut self) -> (i64, UtilityReceiver) {
287 + pub fn allocate_and_register(&mut self) -> (u64, UtilityReceiver) {
288 let call_id = self.next_call_id.fetch_add(1, Ordering::Relaxed);
289 let (tx, rx) = oneshot::channel();
290 self.utility_calls.insert(call_id, tx);
@@ -291,14 +292,14 @@ impl UtilityRegistry {
292 }
293
294 /// Resolve a utility output to its waiting receiver.
294 - pub fn resolve(&mut self, call_id: &i64) -> Option<UtilitySender> {
295 + pub fn resolve(&mut self, call_id: &u64) -> Option<UtilitySender> {
296 self.utility_calls.remove(call_id)
297 }
298
299 /// Drop a batch of registered utility calls without delivering a result.
300 /// Used to roll back allocations when the dispatch fan-out fails before
301 /// every engine could accept the request.
301 - pub fn unregister_many(&mut self, call_ids: impl IntoIterator<Item = i64>) {
302 + pub fn unregister_many(&mut self, call_ids: impl IntoIterator<Item = u64>) {
303 for call_id in call_ids {
304 self.utility_calls.remove(&call_id);
305 }
@@ -315,7 +316,7 @@ impl UtilityRegistry {
316 }
317
318 #[cfg(test)]
318 - pub fn contains(&self, call_id: i64) -> bool {
319 + pub fn contains(&self, call_id: u64) -> bool {
320 self.utility_calls.contains_key(&call_id)
321 }
322
rust/src/engine-core-client/src/error.rs
+5 -3
@@ -4,6 +4,8 @@ use std::time::Duration;
4 use thiserror::Error;
5 use thiserror_ext::Macro;
6
7 +use crate::protocol::utility::UtilityCallId;
8 +
9 pub type Result<T> = std::result::Result<T, Error>;
10
11 /// Public error type for the Rust engine-core client.
@@ -70,17 +72,17 @@ pub enum Error {
72 #[error("utility call `{method}` failed (call_id={call_id}): {message}")]
73 UtilityCallFailed {
74 method: String,
73 - call_id: i64,
75 + call_id: UtilityCallId,
76 message: String,
77 },
78 #[error("utility call `{method}` returned an invalid result (call_id={call_id}): {message}")]
79 UtilityResultDecode {
80 method: String,
79 - call_id: i64,
81 + call_id: UtilityCallId,
82 message: String,
83 },
84 #[error("utility call `{method}` closed unexpectedly (call_id={call_id})")]
83 - UtilityCallClosed { method: String, call_id: i64 },
85 + UtilityCallClosed { method: String, call_id: u64 },
86
87 /// A special variant to allow cloning the same error.
88 #[error(transparent)]
rust/src/engine-core-client/src/protocol/classified_outputs.rs
+4 -3
@@ -2,7 +2,8 @@ use std::collections::BTreeSet;
2
3 use enum_as_inner::EnumAsInner;
4
5 -use super::{EngineCoreOutput, EngineCoreOutputs, UtilityOutput};
5 +use super::utility::UtilityOutput;
6 +use super::{EngineCoreOutput, EngineCoreOutputs};
7 use crate::protocol::stats::SchedulerStats;
8
9 /// Data-parallel control notifications multiplexed through `EngineCoreOutputs`.
@@ -150,7 +151,7 @@ mod tests {
151 fn engine_core_outputs_classify_utility() {
152 let outputs = EngineCoreOutputs {
153 utility_output: Some(UtilityOutput {
153 - call_id: 42,
154 + call_id: 42_u64.into(),
155 failure_message: None,
156 result: None,
157 }),
@@ -201,7 +202,7 @@ mod tests {
202 ..Default::default()
203 }],
204 utility_output: Some(UtilityOutput {
204 - call_id: 1,
205 + call_id: 1_u64.into(),
206 failure_message: None,
207 result: None,
208 }),
rust/src/engine-core-client/src/protocol/mod.rs
+2 -187
@@ -14,6 +14,7 @@ use crate::error::{Error, Result};
14 use crate::protocol::logprobs::MaybeWireLogprobs;
15 use crate::protocol::multimodal::MmFeatures;
16 use crate::protocol::stats::{PrefillStats, SchedulerStats};
17 +use crate::protocol::utility::UtilityOutput;
18
19 // TODO: This module currently mixes reusable frontend-facing semantic types
20 // (for example `FinishReason`, `StopReason`, `RequestOutputKind`, and future
@@ -42,6 +43,7 @@ pub mod logprobs;
43 pub mod multimodal;
44 pub mod stats;
45 pub mod tensor;
46 +pub mod utility;
47 pub use classified_outputs::{
48 ClassifiedEngineCoreOutputs, DpControlMessage, RequestBatchOutputs, UtilityCallOutput,
49 };
@@ -373,51 +375,6 @@ impl EngineCoreRequest {
375 }
376 }
377
376 -/// Engine-core utility call payload sent from frontend to engine.
377 -///
378 -/// Original Python payload shape:
379 -/// `(client_index, call_id, method_name, args)`
380 -#[derive(Debug, Clone, PartialEq, Serialize_tuple)]
381 -pub struct EngineCoreUtilityRequest {
382 - pub client_index: u32,
383 - pub call_id: i64,
384 - pub method_name: String,
385 - pub args: OpaqueValue,
386 -}
387 -
388 -impl EngineCoreUtilityRequest {
389 - /// Create a new utility request with the given strongly typed arguments,
390 - /// encoding them into the expected msgpack value format.
391 - pub fn new<T>(
392 - client_index: u32,
393 - call_id: i64,
394 - method_name: impl Into<String>,
395 - args: T,
396 - ) -> Result<Self>
397 - where
398 - T: Serialize + std::fmt::Debug,
399 - {
400 - let args = rmpv::ext::to_value(&args).map_err(|error| Error::Encode {
401 - target_type: type_name::<T>(),
402 - message: format!(
403 - "failed to encode utility args `{args:?}`: {}",
404 - error.to_report_string()
405 - ),
406 - })?;
407 - let args = match args {
408 - Value::Nil => Value::Array(Vec::new()),
409 - other => other,
410 - };
411 -
412 - Ok(Self {
413 - client_index,
414 - call_id,
415 - method_name: method_name.into(),
416 - args,
417 - })
418 - }
419 -}
420 -
378 /// Engine-core output for a single request.
379 ///
380 /// Original Python definition:
@@ -464,70 +421,6 @@ impl EngineCoreOutput {
421 }
422 }
423
467 -/// Result of a utility call.
468 -///
469 -/// Original Python definition:
470 -/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L174-L183>
471 -#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
472 -pub struct UtilityOutput {
473 - pub call_id: i64,
474 - /// Non-`None` implies the call failed and `result` should be ignored.
475 - #[serde(default)]
476 - pub failure_message: Option<String>,
477 - #[serde(default)]
478 - pub result: Option<UtilityResultEnvelope>,
479 -}
480 -
481 -/// Python `UtilityResult` wrapper carried inside `UtilityOutput.result`.
482 -///
483 -/// Upstream reference:
484 -/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/serial_utils.py#L178-L185>
485 -#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
486 -pub struct UtilityResultEnvelope {
487 - /// Recursive type information encoded on Python side, serving as the hint
488 - /// for deserialization. We don't care it here as in Rust frontend all
489 - /// utility calls are strongly-typed.
490 - #[serde(default)]
491 - type_info: Option<OpaqueValue>,
492 - /// The actual utility result.
493 - #[serde(default = "default_opaque_value_nil")]
494 - result: OpaqueValue,
495 -}
496 -
497 -impl UtilityResultEnvelope {
498 - /// Create a utility result envelope without type information.
499 - pub fn without_type_info(result: OpaqueValue) -> Self {
500 - Self {
501 - type_info: None,
502 - result,
503 - }
504 - }
505 -}
506 -
507 -impl UtilityOutput {
508 - /// Decode the typed result of a utility call.
509 - pub fn into_typed_result<T>(self, method: &str) -> Result<T>
510 - where
511 - T: serde::de::DeserializeOwned,
512 - {
513 - if let Some(message) = self.failure_message {
514 - return Err(Error::UtilityCallFailed {
515 - method: method.to_string(),
516 - call_id: self.call_id,
517 - message,
518 - });
519 - }
520 -
521 - let result = self.result.map(|e| e.result).unwrap_or(Value::Nil);
522 -
523 - rmpv::ext::from_value(result).map_err(|error| Error::UtilityResultDecode {
524 - method: method.to_string(),
525 - call_id: self.call_id,
526 - message: error.to_report_string(),
527 - })
528 - }
529 -}
530 -
424 /// Batch of engine-core outputs returned to a frontend client.
425 ///
426 /// Original Python definition:
@@ -601,13 +494,6 @@ mod tests {
494
495 use super::*;
496
604 - fn utility_result_value<T>(value: T) -> UtilityResultEnvelope
605 - where
606 - T: Serialize,
607 - {
608 - UtilityResultEnvelope::without_type_info(rmpv::ext::to_value(value).unwrap())
609 - }
610 -
497 #[test]
498 fn engine_core_request_serializes_as_full_array() {
499 let request = EngineCoreRequest {
@@ -673,77 +559,6 @@ mod tests {
559 );
560 }
561
676 - #[test]
677 - fn utility_request_serializes_as_tuple_payload() {
678 - let request = EngineCoreUtilityRequest::new(7, 42, "is_sleeping", ()).unwrap();
679 -
680 - let encoded = encode_msgpack(&request).unwrap();
681 - let value = decode_value(&encoded).unwrap();
682 - let array = match value {
683 - Value::Array(array) => array,
684 - other => panic!("expected utility request array, got {other:?}"),
685 - };
686 -
687 - assert_eq!(array.len(), 4);
688 - assert_eq!(array[0], Value::from(7));
689 - assert_eq!(array[1], Value::from(42));
690 - assert_eq!(array[2], Value::from("is_sleeping"));
691 - assert_eq!(array[3], Value::Array(Vec::new()));
692 - }
693 -
694 - #[test]
695 - fn utility_output_decodes_typed_result() {
696 - let output = UtilityOutput {
697 - call_id: 9,
698 - failure_message: None,
699 - result: Some(utility_result_value(true)),
700 - };
701 -
702 - assert!(output.into_typed_result::<bool>("is_sleeping").unwrap());
703 - }
704 -
705 - #[test]
706 - fn utility_output_reports_failure_message() {
707 - let error = UtilityOutput {
708 - call_id: 9,
709 - failure_message: Some("boom".to_string()),
710 - result: None,
711 - }
712 - .into_typed_result::<bool>("is_sleeping")
713 - .unwrap_err();
714 -
715 - assert!(matches!(
716 - error,
717 - Error::UtilityCallFailed {
718 - method,
719 - call_id,
720 - message
721 - } if method == "is_sleeping" && call_id == 9 && message == "boom"
722 - ));
723 - }
724 -
725 - #[test]
726 - fn utility_output_decodes_missing_result_as_unit() {
727 - UtilityOutput {
728 - call_id: 3,
729 - failure_message: None,
730 - result: None,
731 - }
732 - .into_typed_result::<()>("reset_mm_cache")
733 - .unwrap();
734 - }
735 -
736 - #[test]
737 - fn utility_output_decodes_nil_result_as_unit() {
738 - UtilityOutput {
739 - call_id: 4,
740 - failure_message: None,
741 - result: Some(UtilityResultEnvelope::without_type_info(Value::Nil)),
742 - }
743 - .into_typed_result::<()>("sleep")
744 - .unwrap();
745 - }
746 -
562 #[test]
563 fn decode_msgpack_includes_type_name_and_value_fallback() {
564 let error = decode_msgpack::<u64>(
rust/src/engine-core-client/src/protocol/utility.rs new
+330
@@ -0,0 +1,330 @@
1 +use std::any::type_name;
2 +use std::fmt;
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 thiserror_ext::AsReport;
9 +
10 +use super::{OpaqueValue, default_opaque_value_nil};
11 +use crate::error::{Error, Result};
12 +
13 +/// Utility call id as carried on the engine-core MessagePack wire.
14 +///
15 +/// Python emits utility ids as MessagePack integers, including values that may
16 +/// require unsigned 64-bit encoding. Keep MessagePack's signed/unsigned
17 +/// integer distinction instead of flattening to `i64` or `u64` at decode time.
18 +#[derive(Clone, Copy, PartialEq)]
19 +pub struct UtilityCallId(rmpv::Integer);
20 +
21 +impl UtilityCallId {
22 + /// Returns the integer represented as `u64` if possible, or else `None`.
23 + /// This is the typical case for utility calls.
24 + pub fn as_u64(self) -> Option<u64> {
25 + self.0.as_u64()
26 + }
27 +
28 + /// Returns the integer represented as `i64` if possible, or else `None`.
29 + pub fn as_i64(self) -> Option<i64> {
30 + self.0.as_i64()
31 + }
32 +}
33 +
34 +impl Default for UtilityCallId {
35 + fn default() -> Self {
36 + Self(0_u64.into())
37 + }
38 +}
39 +
40 +impl From<u64> for UtilityCallId {
41 + fn from(value: u64) -> Self {
42 + Self(value.into())
43 + }
44 +}
45 +
46 +impl TryFrom<Value> for UtilityCallId {
47 + type Error = String;
48 +
49 + fn try_from(value: Value) -> std::result::Result<Self, Self::Error> {
50 + match value {
51 + Value::Integer(value) => Ok(UtilityCallId(value)),
52 + other => Err(format!(
53 + "expected a MessagePack integer utility call id, got {other}"
54 + )),
55 + }
56 + }
57 +}
58 +
59 +impl PartialEq<u64> for UtilityCallId {
60 + fn eq(&self, other: &u64) -> bool {
61 + self.as_u64() == Some(*other)
62 + }
63 +}
64 +
65 +impl fmt::Debug for UtilityCallId {
66 + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
67 + fmt::Display::fmt(&self.0, f)
68 + }
69 +}
70 +
71 +impl fmt::Display for UtilityCallId {
72 + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
73 + fmt::Display::fmt(&self.0, f)
74 + }
75 +}
76 +
77 +impl Serialize for UtilityCallId {
78 + fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
79 + where
80 + S: serde::Serializer,
81 + {
82 + Value::Integer(self.0).serialize(serializer)
83 + }
84 +}
85 +
86 +impl<'de> Deserialize<'de> for UtilityCallId {
87 + fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
88 + where
89 + D: serde::Deserializer<'de>,
90 + {
91 + let value = Value::deserialize(deserializer)?;
92 + match value {
93 + Value::Integer(value) => Ok(UtilityCallId(value)),
94 + other => Err(serde::de::Error::custom(format!(
95 + "expected a MessagePack integer utility call id, got {other}"
96 + ))),
97 + }
98 + }
99 +}
100 +
101 +/// Engine-core utility call payload sent from frontend to engine.
102 +///
103 +/// Original Python payload shape:
104 +/// `(client_index, call_id, method_name, args)`
105 +#[derive(Debug, Clone, PartialEq, Serialize_tuple)]
106 +pub struct EngineCoreUtilityRequest {
107 + pub client_index: u32,
108 + pub call_id: UtilityCallId,
109 + pub method_name: String,
110 + pub args: OpaqueValue,
111 +}
112 +
113 +impl EngineCoreUtilityRequest {
114 + /// Create a new utility request with the given strongly typed arguments,
115 + /// encoding them into the expected msgpack value format.
116 + pub fn new<T>(
117 + client_index: u32,
118 + call_id: u64,
119 + method_name: impl Into<String>,
120 + args: T,
121 + ) -> Result<Self>
122 + where
123 + T: Serialize + std::fmt::Debug,
124 + {
125 + let args = rmpv::ext::to_value(&args).map_err(|error| Error::Encode {
126 + target_type: type_name::<T>(),
127 + message: format!(
128 + "failed to encode utility args `{args:?}`: {}",
129 + error.to_report_string()
130 + ),
131 + })?;
132 + let args = match args {
133 + Value::Nil => Value::Array(Vec::new()),
134 + other => other,
135 + };
136 +
137 + Ok(Self {
138 + client_index,
139 + call_id: UtilityCallId::from(call_id),
140 + method_name: method_name.into(),
141 + args,
142 + })
143 + }
144 +}
145 +
146 +/// Result of a utility call.
147 +///
148 +/// Original Python definition:
149 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L174-L183>
150 +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
151 +pub struct UtilityOutput {
152 + pub call_id: UtilityCallId,
153 + /// Non-`None` implies the call failed and `result` should be ignored.
154 + #[serde(default)]
155 + pub failure_message: Option<String>,
156 + #[serde(default)]
157 + pub result: Option<UtilityResultEnvelope>,
158 +}
159 +
160 +/// Python `UtilityResult` wrapper carried inside `UtilityOutput.result`.
161 +///
162 +/// Upstream reference:
163 +/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/serial_utils.py#L178-L185>
164 +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
165 +pub struct UtilityResultEnvelope {
166 + /// Recursive type information encoded on Python side, serving as the hint
167 + /// for deserialization. We don't care it here as in Rust frontend all
168 + /// utility calls are strongly-typed.
169 + #[serde(default)]
170 + type_info: Option<OpaqueValue>,
171 + /// The actual utility result.
172 + #[serde(default = "default_opaque_value_nil")]
173 + result: OpaqueValue,
174 +}
175 +
176 +impl UtilityResultEnvelope {
177 + /// Create a utility result envelope without type information.
178 + pub fn without_type_info(result: OpaqueValue) -> Self {
179 + Self {
180 + type_info: None,
181 + result,
182 + }
183 + }
184 +}
185 +
186 +impl UtilityOutput {
187 + /// Decode the typed result of a utility call.
188 + pub fn into_typed_result<T>(self, method: &str) -> Result<T>
189 + where
190 + T: serde::de::DeserializeOwned,
191 + {
192 + if let Some(message) = self.failure_message {
193 + return Err(Error::UtilityCallFailed {
194 + method: method.to_string(),
195 + call_id: self.call_id,
196 + message,
197 + });
198 + }
199 +
200 + let result = self.result.map(|e| e.result).unwrap_or(Value::Nil);
201 +
202 + rmpv::ext::from_value(result).map_err(|error| Error::UtilityResultDecode {
203 + method: method.to_string(),
204 + call_id: self.call_id,
205 + message: error.to_report_string(),
206 + })
207 + }
208 +}
209 +
210 +#[cfg(test)]
211 +mod tests {
212 + use rmpv::Value;
213 + use serde::Serialize;
214 +
215 + use super::{EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope};
216 + use crate::Error;
217 + use crate::protocol::{decode_msgpack, decode_value, encode_msgpack};
218 +
219 + fn utility_result_value<T>(value: T) -> UtilityResultEnvelope
220 + where
221 + T: Serialize,
222 + {
223 + UtilityResultEnvelope::without_type_info(rmpv::ext::to_value(value).unwrap())
224 + }
225 +
226 + #[test]
227 + fn utility_request_serializes_as_tuple_payload() {
228 + let request = EngineCoreUtilityRequest::new(7, 42, "is_sleeping", ()).unwrap();
229 +
230 + let encoded = encode_msgpack(&request).unwrap();
231 + let value = decode_value(&encoded).unwrap();
232 + let array = match value {
233 + Value::Array(array) => array,
234 + other => panic!("expected utility request array, got {other:?}"),
235 + };
236 +
237 + assert_eq!(array.len(), 4);
238 + assert_eq!(array[0], Value::from(7));
239 + assert_eq!(array[1], Value::from(42));
240 + assert_eq!(array[2], Value::from("is_sleeping"));
241 + assert_eq!(array[3], Value::Array(Vec::new()));
242 + }
243 +
244 + #[test]
245 + fn utility_output_decodes_typed_result() {
246 + let output = UtilityOutput {
247 + call_id: 9_u64.into(),
248 + failure_message: None,
249 + result: Some(utility_result_value(true)),
250 + };
251 +
252 + assert!(output.into_typed_result::<bool>("is_sleeping").unwrap());
253 + }
254 +
255 + #[test]
256 + fn utility_output_decodes_unsigned_64_bit_call_id() {
257 + let value = Value::Array(vec![Value::from(u64::MAX), Value::Nil, Value::Nil]);
258 + let mut encoded = Vec::new();
259 + rmpv::encode::write_value(&mut encoded, &value).unwrap();
260 +
261 + let output: UtilityOutput = decode_msgpack(&encoded).unwrap();
262 +
263 + assert_eq!(output.call_id.as_u64(), Some(u64::MAX));
264 + }
265 +
266 + #[test]
267 + fn utility_output_decodes_signed_negative_call_id() {
268 + let value = Value::Array(vec![Value::from(-1), Value::Nil, Value::Nil]);
269 + let mut encoded = Vec::new();
270 + rmpv::encode::write_value(&mut encoded, &value).unwrap();
271 +
272 + let output: UtilityOutput = decode_msgpack(&encoded).unwrap();
273 +
274 + assert_eq!(output.call_id.as_u64(), None);
275 + }
276 +
277 + #[test]
278 + fn utility_output_decodes_other_negative_call_id() {
279 + let value = Value::Array(vec![Value::from(-2), Value::Nil, Value::Nil]);
280 + let mut encoded = Vec::new();
281 + rmpv::encode::write_value(&mut encoded, &value).unwrap();
282 +
283 + let output: UtilityOutput = decode_msgpack(&encoded).unwrap();
284 +
285 + assert_eq!(output.call_id.as_u64(), None);
286 + assert_eq!(output.call_id.as_i64(), Some(-2));
287 + }
288 +
289 + #[test]
290 + fn utility_output_reports_failure_message() {
291 + let error = UtilityOutput {
292 + call_id: 9_u64.into(),
293 + failure_message: Some("boom".to_string()),
294 + result: None,
295 + }
296 + .into_typed_result::<bool>("is_sleeping")
297 + .unwrap_err();
298 +
299 + assert!(matches!(
300 + error,
301 + Error::UtilityCallFailed {
302 + method,
303 + call_id,
304 + message
305 + } if method == "is_sleeping" && call_id == 9 && message == "boom"
306 + ));
307 + }
308 +
309 + #[test]
310 + fn utility_output_decodes_missing_result_as_unit() {
311 + UtilityOutput {
312 + call_id: 3_u64.into(),
313 + failure_message: None,
314 + result: None,
315 + }
316 + .into_typed_result::<()>("reset_mm_cache")
317 + .unwrap();
318 + }
319 +
320 + #[test]
321 + fn utility_output_decodes_nil_result_as_unit() {
322 + UtilityOutput {
323 + call_id: 4_u64.into(),
324 + failure_message: None,
325 + result: Some(UtilityResultEnvelope::without_type_info(Value::Nil)),
326 + }
327 + .into_typed_result::<()>("sleep")
328 + .unwrap();
329 + }
330 +}
rust/src/engine-core-client/src/tests/client.rs
+16 -16
@@ -24,10 +24,10 @@ use crate::protocol::multimodal::{
24 };
25 use crate::protocol::stats::SchedulerStats;
26 use crate::protocol::tensor::WireTensor;
27 +use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
28 use crate::protocol::{
29 EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
29 - EngineCoreRequestType, EngineCoreSamplingParams, UtilityOutput, UtilityResultEnvelope,
30 - decode_engine_core_outputs,
30 + EngineCoreRequestType, EngineCoreSamplingParams, decode_engine_core_outputs,
31 };
32 use crate::test_utils::{
33 IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_connections,
@@ -878,7 +878,7 @@ async fn client_fail_closes_when_main_output_path_receives_dp_control() {
878 push,
879 EngineCoreOutputs {
880 utility_output: Some(UtilityOutput {
881 - call_id: 1,
881 + call_id: 1_u64.into(),
882 failure_message: None,
883 result: None,
884 }),
@@ -978,7 +978,7 @@ async fn client_fail_closes_when_main_output_path_receives_mixed_shape_output()
978 push,
979 EngineCoreOutputs {
980 utility_output: Some(UtilityOutput {
981 - call_id: 1,
981 + call_id: 1_u64.into(),
982 failure_message: None,
983 result: None,
984 }),
@@ -1309,7 +1309,7 @@ async fn is_sleeping_wrapper_sends_typed_request_and_returns_typed_response() {
1309 };
1310 assert_eq!(array.len(), 4);
1311 assert_eq!(array[0], Value::from(5));
1312 - let call_id = array[1].as_i64().expect("call_id");
1312 + let call_id = array[1].as_u64().expect("call_id");
1313 assert_eq!(array[2], Value::from("is_sleeping"));
1314 assert_eq!(array[3], Value::Array(Vec::new()));
1315
@@ -1317,7 +1317,7 @@ async fn is_sleeping_wrapper_sends_typed_request_and_returns_typed_response() {
1317 push,
1318 EngineCoreOutputs {
1319 utility_output: Some(UtilityOutput {
1320 - call_id,
1320 + call_id: call_id.into(),
1321 failure_message: None,
1322 result: Some(utility_result_value(true)),
1323 }),
@@ -1366,13 +1366,13 @@ async fn call_utility_failure_message_surfaces_as_error() {
1366 assert_eq!(utility[0].as_ref(), &[0x03]);
1367 let payload = decode_value(&utility[1]);
1368 let call_id =
1369 - payload.as_array().and_then(|array| array[1].as_i64()).expect("call_id");
1369 + payload.as_array().and_then(|array| array[1].as_u64()).expect("call_id");
1370
1371 send_outputs(
1372 push,
1373 EngineCoreOutputs {
1374 utility_output: Some(UtilityOutput {
1375 - call_id,
1375 + call_id: call_id.into(),
1376 failure_message: Some("boom".to_string()),
1377 result: None,
1378 }),
@@ -1865,13 +1865,13 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
1865 Value::Array(array) => array,
1866 other => panic!("expected utility payload array, got {other:?}"),
1867 };
1868 - let call_id = array[1].as_i64().expect("call_id");
1868 + let call_id = array[1].as_u64().expect("call_id");
1869 assert_eq!(array[2], Value::from("is_sleeping"));
1870 send_outputs(
1871 push,
1872 EngineCoreOutputs {
1873 utility_output: Some(UtilityOutput {
1874 - call_id,
1874 + call_id: call_id.into(),
1875 failure_message: None,
1876 result: Some(utility_result_value(true)),
1877 }),
@@ -1919,13 +1919,13 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
1919 Value::Array(array) => array,
1920 other => panic!("expected utility payload array, got {other:?}"),
1921 };
1922 - let call_id = array[1].as_i64().expect("call_id");
1922 + let call_id = array[1].as_u64().expect("call_id");
1923 assert_eq!(array[2], Value::from("is_sleeping"));
1924 send_outputs(
1925 push,
1926 EngineCoreOutputs {
1927 utility_output: Some(UtilityOutput {
1928 - call_id,
1928 + call_id: call_id.into(),
1929 failure_message: None,
1930 result: Some(utility_result_value(true)),
1931 }),
@@ -2029,14 +2029,14 @@ async fn collective_rpc_flattens_results_from_all_engines() {
2029 Value::Array(array) => array,
2030 other => panic!("expected utility payload array, got {other:?}"),
2031 };
2032 - let call_id = array[1].as_i64().expect("call_id");
2032 + let call_id = array[1].as_u64().expect("call_id");
2033 assert_eq!(array[2], Value::from("collective_rpc"));
2034
2035 send_outputs(
2036 push,
2037 EngineCoreOutputs {
2038 utility_output: Some(UtilityOutput {
2039 - call_id,
2039 + call_id: call_id.into(),
2040 failure_message: None,
2041 result: Some(utility_result_value(vec!["engine-0-worker"])),
2042 }),
@@ -2060,14 +2060,14 @@ async fn collective_rpc_flattens_results_from_all_engines() {
2060 Value::Array(array) => array,
2061 other => panic!("expected utility payload array, got {other:?}"),
2062 };
2063 - let call_id = array[1].as_i64().expect("call_id");
2063 + let call_id = array[1].as_u64().expect("call_id");
2064 assert_eq!(array[2], Value::from("collective_rpc"));
2065
2066 send_outputs(
2067 push,
2068 EngineCoreOutputs {
2069 utility_output: Some(UtilityOutput {
2070 - call_id,
2070 + call_id: call_id.into(),
2071 failure_message: None,
2072 result: Some(utility_result_value(vec!["engine-1-worker"])),
2073 }),
rust/src/server/src/routes/tests.rs
+11 -10
@@ -27,9 +27,10 @@ use vllm_chat::{
27 use vllm_engine_core_client::protocol::logprobs::{
28 Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
29 };
30 +use vllm_engine_core_client::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
31 use vllm_engine_core_client::protocol::{
32 EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason,
32 - UtilityOutput, UtilityResultEnvelope, decode_value,
33 + decode_value,
34 };
35 use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
36 use vllm_engine_core_client::{
@@ -369,10 +370,10 @@ fn utility_none_result() -> UtilityResultEnvelope {
370 UtilityResultEnvelope::without_type_info(Value::Nil)
371 }
372
372 -fn utility_outputs(call_id: i64, result: UtilityResultEnvelope) -> EngineCoreOutputs {
373 +fn utility_outputs(call_id: u64, result: UtilityResultEnvelope) -> EngineCoreOutputs {
374 EngineCoreOutputs {
375 utility_output: Some(UtilityOutput {
375 - call_id,
376 + call_id: call_id.into(),
377 failure_message: None,
378 result: Some(result),
379 }),
@@ -2989,7 +2990,7 @@ async fn reset_prefix_cache_route_sends_expected_utility_call() {
2990
2991 let payload = decode_value(&utility[1]).expect("decode utility payload");
2992 let array = payload.as_array().expect("utility payload array");
2992 - let call_id = array[1].as_i64().expect("call id");
2993 + let call_id = array[1].as_u64().expect("call id");
2994
2995 assert_eq!(array[2], Value::from("reset_prefix_cache"));
2996 assert_eq!(
@@ -3031,7 +3032,7 @@ async fn reset_mm_cache_route_sends_expected_utility_call() {
3032
3033 let payload = decode_value(&utility[1]).expect("decode utility payload");
3034 let array = payload.as_array().expect("utility payload array");
3034 - let call_id = array[1].as_i64().expect("call id");
3035 + let call_id = array[1].as_u64().expect("call id");
3036
3037 assert_eq!(array[2], Value::from("reset_mm_cache"));
3038 assert_eq!(array[3], Value::Array(Vec::new()));
@@ -3070,7 +3071,7 @@ async fn reset_encoder_cache_route_sends_expected_utility_call() {
3071
3072 let payload = decode_value(&utility[1]).expect("decode utility payload");
3073 let array = payload.as_array().expect("utility payload array");
3073 - let call_id = array[1].as_i64().expect("call id");
3074 + let call_id = array[1].as_u64().expect("call id");
3075
3076 assert_eq!(array[2], Value::from("reset_encoder_cache"));
3077 assert_eq!(array[3], Value::Array(Vec::new()));
@@ -3109,7 +3110,7 @@ async fn collective_rpc_route_sends_expected_utility_call_and_returns_results()
3110
3111 let payload = decode_value(&utility[1]).expect("decode utility payload");
3112 let array = payload.as_array().expect("utility payload array");
3112 - let call_id = array[1].as_i64().expect("call id");
3113 + let call_id = array[1].as_u64().expect("call id");
3114
3115 assert_eq!(array[2], Value::from("collective_rpc"));
3116 assert_eq!(
@@ -3194,7 +3195,7 @@ async fn sleep_route_uses_python_compatible_default_query_values() {
3195
3196 let payload = decode_value(&utility[1]).expect("decode utility payload");
3197 let array = payload.as_array().expect("utility payload array");
3197 - let call_id = array[1].as_i64().expect("call id");
3198 + let call_id = array[1].as_u64().expect("call id");
3199
3200 assert_eq!(array[2], Value::from("sleep"));
3201 assert_eq!(
@@ -3236,7 +3237,7 @@ async fn wake_up_route_without_tags_sends_none() {
3237
3238 let payload = decode_value(&utility[1]).expect("decode utility payload");
3239 let array = payload.as_array().expect("utility payload array");
3239 - let call_id = array[1].as_i64().expect("call id");
3240 + let call_id = array[1].as_u64().expect("call id");
3241
3242 assert_eq!(array[2], Value::from("wake_up"));
3243 assert_eq!(array[3], Value::Array(vec![Value::Nil]));
@@ -3275,7 +3276,7 @@ async fn is_sleeping_route_returns_json_payload() {
3276
3277 let payload = decode_value(&utility[1]).expect("decode utility payload");
3278 let array = payload.as_array().expect("utility payload array");
3278 - let call_id = array[1].as_i64().expect("call id");
3279 + let call_id = array[1].as_u64().expect("call id");
3280
3281 assert_eq!(array[2], Value::from("is_sleeping"));
3282 assert_eq!(array[3], Value::Array(Vec::new()));