iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Rust Frontend] Use enum-backed domain types for engine outputs and structured outputs (#47283)

Browse Source 
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
Bugen Zhao
2026-07-02 10:41:46 +01:00
committed by GitHub
parent 08a8a4af3f
commit 25fcb65d51
16 changed files with 762 additions and 605 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/chat/src/output/default/structural_tag.rs
+12 -18
View File
@@ -47,9 +47,8 @@ pub(super) fn apply_structural_tag_constraint(
// Overwrite any existing structured output settings with the structural tag constraint.
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
structural_tag: Some(structural_tag),
backend: StructuredOutputBackend::Xgrammar,
..Default::default()
..StructuredOutputsParams::structural_tag(structural_tag)
});
Ok(())
@@ -119,10 +118,11 @@ mod tests {
.as_ref()
.expect("structured outputs should be set");
assert_eq!(params.backend, StructuredOutputBackend::Xgrammar);
serde_json::from_str(
params.structural_tag.as_deref().expect("structural_tag should be set"),
)
.expect("structural_tag should be valid JSON")
let structural_tag = params
.constraint
.as_structural_tag()
.expect("structured output constraint should be structural_tag");
serde_json::from_str(structural_tag).expect("structural_tag should be valid JSON")
}
fn structured_outputs(request: &ChatRequest) -> &StructuredOutputsParams {
@@ -161,9 +161,8 @@ mod tests {
fn auto_strict_tool_choice_overwrites_existing_json_guidance() {
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
json: Some(json!({"type": "object"})),
backend: StructuredOutputBackend::Xgrammar,
..Default::default()
..StructuredOutputsParams::json(json!({"type": "object"}))
});
let parser = qwen3_coder_parser(&request.tools);
@@ -171,8 +170,7 @@ mod tests {
.expect("structural tag should build");
let params = structured_outputs(&request);
assert!(params.json.is_none());
assert!(params.structural_tag.is_some());
assert!(params.constraint.is_structural_tag());
let tag = structural_tag_value(&request);
assert_eq!(tag["type"], "structural_tag");
assert!(tag.to_string().contains("search"));
@@ -195,9 +193,8 @@ mod tests {
fn required_tool_choice_overwrites_existing_json_object_guidance() {
let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
json_object: Some(true),
backend: StructuredOutputBackend::Xgrammar,
..Default::default()
..StructuredOutputsParams::json_object()
});
let parser = qwen3_coder_parser(&request.tools);
@@ -205,8 +202,7 @@ mod tests {
.expect("structural tag should build");
let params = structured_outputs(&request);
assert!(params.json_object.is_none());
assert!(params.structural_tag.is_some());
assert!(params.constraint.is_structural_tag());
let tag = structural_tag_value(&request);
assert_eq!(tag["type"], "structural_tag");
assert!(tag.to_string().contains("search"));
@@ -245,9 +241,8 @@ mod tests {
fn none_tool_choice_preserves_existing_json_object_guidance() {
let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
json_object: Some(true),
backend: StructuredOutputBackend::Xgrammar,
..Default::default()
..StructuredOutputsParams::json_object()
});
let parser = qwen3_coder_parser(&request.tools);
@@ -255,7 +250,6 @@ mod tests {
.expect("structural tag decision should succeed");
let params = structured_outputs(&request);
assert_eq!(params.json_object, Some(true));
assert!(params.structural_tag.is_none());
assert!(params.constraint.is_json_object());
}
}
rust/src/chat/tests/chat.rs
+34 -23
View File
@@ -16,7 +16,7 @@ use vllm_engine_core_client::protocol::logprobs::{
Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
};
use vllm_engine_core_client::protocol::output::{
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, StopReason,
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs, StopReason,
};
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
@@ -342,7 +342,7 @@ async fn chat_streams_text_events() {
);
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![
request_output("chat-1", vec![b'H' as u32], None, None),
request_output(
@@ -354,7 +354,8 @@ async fn chat_streams_text_events() {
],
finished_requests: Some(BTreeSet::from(["chat-1".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -454,7 +455,7 @@ async fn chat_stream_waits_for_complete_utf8_before_emitting() {
let _ = recv_engine_message(dealer).await;
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![
request_output("chat-utf8", bytes_to_token_ids(&[0xe4]), None, None),
request_output(
@@ -466,7 +467,8 @@ async fn chat_stream_waits_for_complete_utf8_before_emitting() {
],
finished_requests: Some(BTreeSet::from(["chat-utf8".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -542,7 +544,7 @@ async fn chat_stream_flushes_held_text_on_finish() {
let _ = recv_engine_message(dealer).await;
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![request_output(
"chat-final-flush",
bytes_to_token_ids(b"ok st"),
@@ -551,7 +553,8 @@ async fn chat_stream_flushes_held_text_on_finish() {
)],
finished_requests: Some(BTreeSet::from(["chat-final-flush".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -662,7 +665,7 @@ async fn chat_stream_reports_decode_failure_as_error_event() {
let _ = recv_engine_message(dealer).await;
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![request_output(
"chat-4",
vec![UNKNOWN_DECODE_TOKEN_ID],
@@ -670,7 +673,8 @@ async fn chat_stream_reports_decode_failure_as_error_event() {
None,
)],
..Default::default()
},
}
.into(),
)
.await;
})
@@ -725,7 +729,7 @@ async fn chat_stream_preserves_terminal_stop_token_when_requested() {
let _ = recv_engine_message(dealer).await;
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![request_output(
"chat-include-stop",
vec![b'H' as u32, b'i' as u32, b'!' as u32],
@@ -734,7 +738,8 @@ async fn chat_stream_preserves_terminal_stop_token_when_requested() {
)],
finished_requests: Some(BTreeSet::from(["chat-include-stop".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -812,7 +817,7 @@ async fn chat_stream_separates_reasoning_blocks_automatically() {
let _ = recv_engine_message(dealer).await;
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![
request_output(
"chat-reasoning",
@@ -841,7 +846,8 @@ async fn chat_stream_separates_reasoning_blocks_automatically() {
],
finished_requests: Some(BTreeSet::from(["chat-reasoning".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -955,7 +961,7 @@ async fn chat_collectors_return_structured_message_and_visible_text() {
let _ = recv_engine_message(dealer).await;
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![request_output(
"chat-collect",
bytes_to_token_ids(b"<think>inner</think>outer"),
@@ -964,7 +970,8 @@ async fn chat_collectors_return_structured_message_and_visible_text() {
)],
finished_requests: Some(BTreeSet::from(["chat-collect".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -1018,7 +1025,7 @@ async fn chat_explicitly_disables_reasoning_parser() {
let _ = recv_engine_message(dealer).await;
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![
request_output(
"chat-reasoning-disabled",
@@ -1049,7 +1056,8 @@ async fn chat_explicitly_disables_reasoning_parser() {
"chat-reasoning-disabled".to_string()
])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -1096,7 +1104,7 @@ async fn chat_stream_parses_tool_calls_automatically() {
let _ = recv_engine_message(dealer).await;
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![
request_output(
"chat-tool",
@@ -1121,7 +1129,8 @@ async fn chat_stream_parses_tool_calls_automatically() {
],
finished_requests: Some(BTreeSet::from(["chat-tool".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -1204,7 +1213,7 @@ async fn chat_collect_message_preserves_tool_call_arguments_in_final_only_mode()
let _ = recv_engine_message(dealer).await;
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![
request_output(
"chat-final-only-tool",
@@ -1231,7 +1240,8 @@ async fn chat_collect_message_preserves_tool_call_arguments_in_final_only_mode()
"chat-final-only-tool".to_string()
])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -1282,7 +1292,7 @@ async fn chat_stream_and_collect_preserve_prompt_and_sample_logprobs() {
let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![
request_output_with_logprobs(
&request.request_id,
@@ -1303,7 +1313,8 @@ async fn chat_stream_and_collect_preserve_prompt_and_sample_logprobs() {
],
finished_requests: Some(BTreeSet::from([request.request_id])),
..Default::default()
},
}
.into(),
)
.await;
}
rust/src/engine-core-client/src/client/imp.rs
+5 -6
View File
@@ -17,7 +17,7 @@ use crate::client::{AbortCause, AbortRequest};
use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output};
use crate::metrics::{LoraInfoExporter, record_scheduler_stats};
use crate::protocol::encode_msgpack;
use crate::protocol::output::{ClassifiedEngineCoreOutputs, EngineCoreOutput, EngineCoreOutputs};
use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs};
use crate::protocol::request::EngineCoreRequestType;
use crate::protocol::stats::SchedulerStats;
use crate::protocol::utility::UtilityOutput;
@@ -351,8 +351,8 @@ pub(crate) async fn run_output_dispatcher_loop(
)),
}?;
match outputs.classify() {
ClassifiedEngineCoreOutputs::RequestBatch(batch) => {
match outputs {
EngineCoreOutputs::RequestBatch(batch) => {
let senders = inner.take_senders_for_outputs(&batch.outputs);
for (output, sender) in batch.outputs.into_iter().zip(senders) {
let request_id = output.request_id.clone();
@@ -403,7 +403,7 @@ pub(crate) async fn run_output_dispatcher_loop(
let (running, waiting) = inner.lora_adapter_states();
lora_info.update(&METRICS.scheduler, running, waiting);
}
ClassifiedEngineCoreOutputs::Utility(utility) => {
EngineCoreOutputs::Utility(utility) => {
let call_id = utility.output.call_id;
if inner.resolve_utility_output(utility.output) {
trace!(
@@ -419,8 +419,7 @@ pub(crate) async fn run_output_dispatcher_loop(
);
}
}
other @ (ClassifiedEngineCoreOutputs::DpControl { .. }
| ClassifiedEngineCoreOutputs::Other(_)) => {
other => {
Err::<(), _>(unexpected_dispatcher_output!(
"received unexpected output on main dispatcher path: {other:?}"
))?;
rust/src/engine-core-client/src/coordinator/inproc.rs
+5 -5
View File
@@ -11,7 +11,7 @@ use crate::client::imp::ClientInner;
use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState};
use crate::error::{Error, Result, bail_unexpected_coordinator_output};
use crate::protocol::encode_msgpack;
use crate::protocol::output::{ClassifiedEngineCoreOutputs, DpControlMessage, EngineCoreOutputs};
use crate::protocol::output::{DpControlMessage, DpControlOutput, EngineCoreOutputs};
use crate::protocol::request::EngineCoreRequestType;
/// Coordinator-to-engine `START_DP_WAVE` control payload encoded on the
@@ -109,19 +109,19 @@ impl InProcCoordinatorRunner {
/// Apply one engine-originated control output to the coordinator state
/// machine.
async fn handle_outputs(&mut self, outputs: EngineCoreOutputs) -> Result<()> {
match outputs.classify() {
ClassifiedEngineCoreOutputs::RequestBatch(batch)
match outputs {
EngineCoreOutputs::RequestBatch(batch)
if batch.outputs.is_empty() && batch.finished_requests.is_none() =>
{
// Stats-only output for coordinator.
// Ignore since the Rust coordinator doesn't track stats for
// routing decisions.
}
ClassifiedEngineCoreOutputs::DpControl {
EngineCoreOutputs::DpControl(DpControlOutput {
engine_index,
control,
..
} => match control {
}) => match control {
// The engines signals they completed the current wave and are now paused.
// Advance the current wave and mark the state as paused.
DpControlMessage::WaveComplete(wave) => {
rust/src/engine-core-client/src/error.rs
+2
View File
@@ -25,6 +25,8 @@ pub enum Error {
ValueDecode(#[from] rmpv::decode::Error),
#[error("messagepack ext value decode failed: {message}")]
ExtValueDecode { message: String },
#[error("invalid structured outputs params: {message}")]
InvalidStructuredOutputsParams { message: String },
#[error("io error")]
Io(#[from] std::io::Error),
#[error("transport error")]
rust/src/engine-core-client/src/protocol/logprobs/tests.rs
+4 -4
View File
@@ -183,7 +183,7 @@ fn decodes_inline_new_logprobs() {
Some(inline_logprobs_value()),
None,
)))];
let decoded = decode_engine_core_outputs(&frames).unwrap();
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
let logprobs = decoded.outputs[0].new_logprobs.clone().unwrap().into_direct().unwrap();
assert_eq!(logprobs, expected_sample_logprobs());
@@ -214,7 +214,7 @@ fn decodes_multipart_new_logprobs() {
]),
Bytes::from_static(&[1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0]),
];
let decoded = decode_engine_core_outputs(&frames).unwrap();
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
let logprobs = decoded.outputs[0].new_logprobs.clone().unwrap().into_direct().unwrap();
assert_eq!(logprobs, expected_sample_logprobs());
@@ -226,7 +226,7 @@ fn decodes_inline_prompt_logprobs() {
None,
Some(inline_prompt_logprobs_value()),
)))];
let decoded = decode_engine_core_outputs(&frames).unwrap();
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
let logprobs = decoded.outputs[0]
.new_prompt_logprobs_tensors
@@ -252,7 +252,7 @@ fn decodes_big_endian_payloads() {
])),
None,
)))];
let decoded = decode_engine_core_outputs(&frames).unwrap();
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
let logprobs = decoded.outputs[0].new_logprobs.clone().unwrap().into_direct().unwrap();
assert_eq!(
logprobs,
rust/src/engine-core-client/src/protocol/output.rs
+167 -117
View File
@@ -1,7 +1,7 @@
use std::collections::BTreeSet;
use enum_as_inner::EnumAsInner;
use serde::{Deserialize, Serialize};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_default::DefaultFromSerde;
use serde_repr::{Deserialize_repr, Serialize_repr};
use serde_tuple::{Deserialize_tuple, Serialize_tuple};
@@ -131,33 +131,33 @@ impl EngineCoreOutput {
}
}
/// Batch of engine-core outputs returned to a frontend client.
/// Raw Python/msgpack engine-core output envelope.
///
/// Original Python definition:
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L186-L214>
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
pub struct EngineCoreOutputs {
struct WireEngineCoreOutputs {
#[serde(default)]
pub engine_index: u32,
engine_index: u32,
/// Outputs grouped for this client in the current engine tick.
#[serde(default)]
pub outputs: Vec<EngineCoreOutput>,
outputs: Vec<EngineCoreOutput>,
#[serde(default)]
pub scheduler_stats: Option<Box<SchedulerStats>>,
scheduler_stats: Option<Box<SchedulerStats>>,
#[serde(default)]
pub timestamp: f64,
timestamp: f64,
#[serde(default)]
pub utility_output: Option<UtilityOutput>,
utility_output: Option<UtilityOutput>,
#[serde(default)]
pub finished_requests: Option<BTreeSet<String>>,
finished_requests: Option<BTreeSet<String>>,
/// In DP mode, signals that the current wave finished and engines are
/// paused.
#[serde(default)]
pub wave_complete: Option<u32>,
wave_complete: Option<u32>,
/// In DP mode, signals that a request arrived for an old wave and the next
/// wave needs to start in other engines.
#[serde(default)]
pub start_wave: Option<u32>,
start_wave: Option<u32>,
}
/// Data-parallel control notifications multiplexed through `EngineCoreOutputs`.
@@ -167,7 +167,7 @@ pub enum DpControlMessage {
StartWave(u32),
}
#[derive(Debug, Clone, PartialEq)]
#[derive(Debug, Clone, PartialEq, Default)]
pub struct RequestBatchOutputs {
pub engine_index: u32,
pub outputs: Vec<EngineCoreOutput>,
@@ -176,30 +176,48 @@ pub struct RequestBatchOutputs {
pub finished_requests: Option<BTreeSet<String>>,
}
#[derive(Debug, Clone, PartialEq)]
#[derive(Debug, Clone, PartialEq, Default)]
pub struct UtilityCallOutput {
pub engine_index: u32,
pub timestamp: f64,
pub output: UtilityOutput,
}
/// Semantic classification of a raw `EngineCoreOutputs` message.
#[derive(Debug, Clone, PartialEq)]
pub struct DpControlOutput {
pub engine_index: u32,
pub timestamp: f64,
pub control: DpControlMessage,
}
/// Semantic engine-core output families.
///
/// Python currently uses one product-shaped wire struct for several distinct
/// output families. This enum exposes those families more explicitly without
/// changing the wire format.
/// Python currently uses one product-shaped wire struct. The Rust protocol
/// exposes the finite semantic families while preserving the same msgpack shape
/// for serialization.
#[derive(Debug, Clone, PartialEq, EnumAsInner)]
pub enum ClassifiedEngineCoreOutputs {
pub enum EngineCoreOutputs {
RequestBatch(RequestBatchOutputs),
Utility(UtilityCallOutput),
DpControl {
engine_index: u32,
timestamp: f64,
control: DpControlMessage,
},
/// Fallback for wire-shape combinations that do not map cleanly onto the
/// current semantic families.
Other(EngineCoreOutputs),
DpControl(DpControlOutput),
}
impl From<RequestBatchOutputs> for EngineCoreOutputs {
fn from(outputs: RequestBatchOutputs) -> Self {
Self::RequestBatch(outputs)
}
}
impl From<UtilityCallOutput> for EngineCoreOutputs {
fn from(output: UtilityCallOutput) -> Self {
Self::Utility(output)
}
}
impl From<DpControlOutput> for EngineCoreOutputs {
fn from(output: DpControlOutput) -> Self {
Self::DpControl(output)
}
}
impl EngineCoreOutputs {
@@ -209,55 +227,119 @@ impl EngineCoreOutputs {
where
Frame: AsRef<[u8]>,
{
for output in &mut self.outputs {
output.resolve_in_place(frames)?;
if let Self::RequestBatch(batch) = self {
for output in &mut batch.outputs {
output.resolve_in_place(frames)?;
}
}
Ok(())
}
}
/// Classify the raw wire message into a more semantic Rust enum.
pub fn classify(self) -> ClassifiedEngineCoreOutputs {
let has_request_payload = !self.outputs.is_empty()
|| self.scheduler_stats.is_some()
|| self.finished_requests.is_some();
/// Classify the raw wire message into a more semantic Rust enum.
impl TryFrom<WireEngineCoreOutputs> for EngineCoreOutputs {
type Error = Error;
fn try_from(value: WireEngineCoreOutputs) -> Result<Self> {
let has_request_payload = !value.outputs.is_empty()
|| value.scheduler_stats.is_some()
|| value.finished_requests.is_some();
match (
has_request_payload,
&self.utility_output,
&self.wave_complete,
&self.start_wave,
&value.utility_output,
&value.wave_complete,
&value.start_wave,
) {
(true, None, None, None) => {
ClassifiedEngineCoreOutputs::RequestBatch(RequestBatchOutputs {
engine_index: self.engine_index,
outputs: self.outputs,
scheduler_stats: self.scheduler_stats,
timestamp: self.timestamp,
finished_requests: self.finished_requests,
})
(true, None, None, None) => Ok(RequestBatchOutputs {
engine_index: value.engine_index,
outputs: value.outputs,
scheduler_stats: value.scheduler_stats,
timestamp: value.timestamp,
finished_requests: value.finished_requests,
}
(false, Some(_), None, None) => {
ClassifiedEngineCoreOutputs::Utility(UtilityCallOutput {
engine_index: self.engine_index,
timestamp: self.timestamp,
output: self.utility_output.unwrap(),
})
.into()),
(false, Some(_), None, None) => Ok(UtilityCallOutput {
engine_index: value.engine_index,
timestamp: value.timestamp,
output: value.utility_output.unwrap(),
}
(false, None, Some(_), None) => ClassifiedEngineCoreOutputs::DpControl {
engine_index: self.engine_index,
timestamp: self.timestamp,
control: DpControlMessage::WaveComplete(self.wave_complete.unwrap()),
},
(false, None, None, Some(_)) => ClassifiedEngineCoreOutputs::DpControl {
engine_index: self.engine_index,
timestamp: self.timestamp,
control: DpControlMessage::StartWave(self.start_wave.unwrap()),
},
_ => ClassifiedEngineCoreOutputs::Other(self),
.into()),
(false, None, Some(_), None) => Ok(DpControlOutput {
engine_index: value.engine_index,
timestamp: value.timestamp,
control: DpControlMessage::WaveComplete(value.wave_complete.unwrap()),
}
.into()),
(false, None, None, Some(_)) => Ok(DpControlOutput {
engine_index: value.engine_index,
timestamp: value.timestamp,
control: DpControlMessage::StartWave(value.start_wave.unwrap()),
}
.into()),
_ => Err(Error::Decode {
target_type: "EngineCoreOutputs",
message: "invalid wire shape".to_string(),
}),
}
}
}
impl From<EngineCoreOutputs> for WireEngineCoreOutputs {
fn from(value: EngineCoreOutputs) -> Self {
match value {
EngineCoreOutputs::RequestBatch(batch) => Self {
engine_index: batch.engine_index,
outputs: batch.outputs,
scheduler_stats: batch.scheduler_stats,
timestamp: batch.timestamp,
finished_requests: batch.finished_requests,
..Default::default()
},
EngineCoreOutputs::Utility(utility) => Self {
engine_index: utility.engine_index,
timestamp: utility.timestamp,
utility_output: Some(utility.output),
..Default::default()
},
EngineCoreOutputs::DpControl(control) => {
let (wave_complete, start_wave) = match control.control {
DpControlMessage::WaveComplete(wave) => (Some(wave), None),
DpControlMessage::StartWave(wave) => (None, Some(wave)),
};
Self {
engine_index: control.engine_index,
timestamp: control.timestamp,
wave_complete,
start_wave,
..Default::default()
}
}
}
}
}
impl Serialize for EngineCoreOutputs {
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
where
S: Serializer,
{
WireEngineCoreOutputs::from(self.clone()).serialize(serializer)
}
}
impl<'de> Deserialize<'de> for EngineCoreOutputs {
fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
where
D: Deserializer<'de>,
{
WireEngineCoreOutputs::deserialize(deserializer)?
.try_into()
.map_err(serde::de::Error::custom)
}
}
/// Decode one ordinary or multipart engine-core output message into the strong
/// typed public protocol shape.
pub fn decode_engine_core_outputs<Frame>(frames: &[Frame]) -> Result<EngineCoreOutputs>
@@ -281,7 +363,7 @@ mod tests {
#[test]
fn engine_core_outputs_roundtrip_finished_fields() {
let outputs = EngineCoreOutputs {
let outputs = WireEngineCoreOutputs {
outputs: vec![EngineCoreOutput {
request_id: "req-1".to_string(),
new_token_ids: vec![42],
@@ -302,7 +384,7 @@ mod tests {
};
let encoded = encode_msgpack(&outputs).unwrap();
let decoded: EngineCoreOutputs = decode_msgpack(&encoded).unwrap();
let decoded: WireEngineCoreOutputs = decode_msgpack(&encoded).unwrap();
assert_eq!(decoded.outputs.len(), 1);
assert_eq!(
@@ -317,7 +399,7 @@ mod tests {
#[test]
fn engine_core_outputs_classify_request_batch() {
let outputs = EngineCoreOutputs {
let outputs = WireEngineCoreOutputs {
outputs: vec![EngineCoreOutput {
request_id: "req-1".to_string(),
new_token_ids: vec![7],
@@ -360,12 +442,12 @@ mod tests {
},
)
"#]]
.assert_debug_eq(&outputs.classify());
.assert_debug_eq(&EngineCoreOutputs::try_from(outputs).unwrap());
}
#[test]
fn engine_core_outputs_classify_utility() {
let outputs = EngineCoreOutputs {
let outputs = WireEngineCoreOutputs {
utility_output: Some(UtilityOutput {
call_id: 42_u64.into(),
failure_message: None,
@@ -387,31 +469,33 @@ mod tests {
},
)
"#]]
.assert_debug_eq(&outputs.classify());
.assert_debug_eq(&EngineCoreOutputs::try_from(outputs).unwrap());
}
#[test]
fn engine_core_outputs_classify_control() {
let outputs = EngineCoreOutputs {
let outputs = WireEngineCoreOutputs {
start_wave: Some(3),
..Default::default()
};
expect_test::expect![[r#"
DpControl {
engine_index: 0,
timestamp: 0.0,
control: StartWave(
3,
),
}
DpControl(
DpControlOutput {
engine_index: 0,
timestamp: 0.0,
control: StartWave(
3,
),
},
)
"#]]
.assert_debug_eq(&outputs.classify());
.assert_debug_eq(&EngineCoreOutputs::try_from(outputs).unwrap());
}
#[test]
fn engine_core_outputs_classify_mixed_shape_as_raw() {
let outputs = EngineCoreOutputs {
fn engine_core_outputs_rejects_mixed_shape() {
let outputs = WireEngineCoreOutputs {
outputs: vec![EngineCoreOutput {
request_id: "req-1".to_string(),
new_token_ids: vec![7],
@@ -425,44 +509,10 @@ mod tests {
..Default::default()
};
expect_test::expect![[r#"
Other(
EngineCoreOutputs {
engine_index: 0,
outputs: [
EngineCoreOutput {
request_id: "req-1",
new_token_ids: [
7,
],
new_logprobs: None,
new_prompt_logprobs_tensors: None,
pooling_output: None,
finish_reason: None,
stop_reason: None,
events: None,
kv_transfer_params: None,
trace_headers: None,
prefill_stats: None,
routed_experts: None,
num_nans_in_logits: 0,
},
],
scheduler_stats: None,
timestamp: 0.0,
utility_output: Some(
UtilityOutput {
call_id: 1,
failure_message: None,
result: None,
},
),
finished_requests: None,
wave_complete: None,
start_wave: None,
},
)
"#]]
.assert_debug_eq(&outputs.classify());
let error = EngineCoreOutputs::try_from(outputs).unwrap_err();
expect_test::expect![[
r#"messagepack decode failed for EngineCoreOutputs: invalid wire shape"#
]]
.assert_eq(&error.to_string());
}
}
rust/src/engine-core-client/src/protocol/structured_outputs.rs
+249 -25
View File
@@ -1,4 +1,8 @@
use serde::{Deserialize, Serialize};
use enum_as_inner::EnumAsInner;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_json::Value;
use crate::error::{Error, Result};
/// Structured-output backend selected for EngineCore grammar compilation.
///
@@ -15,50 +19,219 @@ pub enum StructuredOutputBackend {
LmFormatEnforcer,
}
/// Parameters for configuring structured outputs (guided decoding).
///
/// Exactly one constraint field (`json`, `regex`, `choice`, `grammar`,
/// `json_object`, or `structural_tag`) should be set. The engine-core
/// backend selects the appropriate grammar compiler based on which field
/// is present.
///
/// Original Python definition:
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L36-L107>
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(default)]
pub struct StructuredOutputsParams {
/// The single structured-output constraint selected for a request.
#[derive(Debug, Clone, PartialEq, EnumAsInner)]
pub enum StructuredOutputConstraint {
/// JSON schema (as a dict/object or JSON string) constraining the output.
pub json: Option<serde_json::Value>,
Json(Value),
/// Regular expression the output must match.
pub regex: Option<String>,
Regex(String),
/// List of allowed output strings (the model must produce one of these).
pub choice: Option<Vec<String>>,
Choice(Vec<String>),
/// Context-free grammar (in EBNF-like notation) the output must conform to.
pub grammar: Option<String>,
/// When `true`, output must be valid JSON (free-form, no schema).
pub json_object: Option<bool>,
Grammar(String),
/// Output must be valid JSON (free-form, no schema).
JsonObject,
/// Structural tag configuration (JSON-encoded string).
StructuralTag(String),
}
/// Additional structured-output options that do not select the constraint mode.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct StructuredOutputOptions {
/// Disable any additional whitespace in guided JSON output.
#[serde(skip_serializing_if = "crate::protocol::is_false")]
pub disable_any_whitespace: bool,
/// Disable `additionalProperties` in JSON schema output.
#[serde(skip_serializing_if = "crate::protocol::is_false")]
pub disable_additional_properties: bool,
/// Custom whitespace pattern for guided JSON output.
pub whitespace_pattern: Option<String>,
/// Structural tag configuration (JSON-encoded string).
pub structural_tag: Option<String>,
}
/// Parameters for configuring structured outputs (guided decoding).
///
/// This is the semantic Rust representation: exactly one constraint mode is
/// always selected. The Python/msgpack product-shaped representation is kept in
/// the private wire type below and used only at serde boundaries.
///
/// Original Python definition:
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L36-L107>
#[derive(Debug, Clone, PartialEq)]
pub struct StructuredOutputsParams {
pub constraint: StructuredOutputConstraint,
pub options: StructuredOutputOptions,
/// Structured-output backend, mirroring Python's internal `_backend`.
///
/// User-supplied values are ignored during deserialization. This matches
/// Python's request boundary, where `_backend` is set by validation rather
/// than accepted as a request-level backend selector.
pub backend: StructuredOutputBackend,
}
impl StructuredOutputsParams {
pub fn json(json: Value) -> Self {
Self::from_constraint(StructuredOutputConstraint::Json(json))
}
pub fn regex(regex: impl Into<String>) -> Self {
Self::from_constraint(StructuredOutputConstraint::Regex(regex.into()))
}
pub fn choice(choice: Vec<String>) -> Self {
Self::from_constraint(StructuredOutputConstraint::Choice(choice))
}
pub fn grammar(grammar: impl Into<String>) -> Self {
Self::from_constraint(StructuredOutputConstraint::Grammar(grammar.into()))
}
pub fn json_object() -> Self {
Self::from_constraint(StructuredOutputConstraint::JsonObject)
}
pub fn structural_tag(structural_tag: impl Into<String>) -> Self {
Self::from_constraint(StructuredOutputConstraint::StructuralTag(
structural_tag.into(),
))
}
fn from_constraint(constraint: StructuredOutputConstraint) -> Self {
Self {
constraint,
options: StructuredOutputOptions::default(),
backend: StructuredOutputBackend::default(),
}
}
}
/// Wire-compatible structured-output payload used by Python engine-core.
///
/// Python models `StructuredOutputsParams` as a product-shaped dataclass with
/// several optional constraint fields, then validates that exactly one of those
/// fields is present. Rust exposes [`StructuredOutputsParams`] as an enum-backed
/// domain type instead, while using this private wire type for ser/de.
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(default)]
struct WireStructuredOutputsParams {
json: Option<Value>,
regex: Option<String>,
choice: Option<Vec<String>>,
grammar: Option<String>,
json_object: Option<bool>,
#[serde(skip_serializing_if = "crate::protocol::is_false")]
disable_any_whitespace: bool,
#[serde(skip_serializing_if = "crate::protocol::is_false")]
disable_additional_properties: bool,
whitespace_pattern: Option<String>,
structural_tag: Option<String>,
#[serde(
default,
rename = "_backend",
deserialize_with = "serde_with::rust::deserialize_ignore_any"
)]
pub backend: StructuredOutputBackend,
backend: StructuredOutputBackend,
}
impl TryFrom<WireStructuredOutputsParams> for StructuredOutputsParams {
type Error = Error;
fn try_from(raw: WireStructuredOutputsParams) -> Result<Self> {
use StructuredOutputConstraint::*;
let mut constraint = None;
macro_rules! insert_constraint {
($name:literal, $value:expr) => {
if let Some(value) = $value {
if let Some((existing, _)) = constraint {
return Err(Error::InvalidStructuredOutputsParams {
message: format!(
"multiple structured output constraints specified: {existing}, {}",
$name
),
});
}
constraint = Some(($name, value));
}
};
}
insert_constraint!("json", raw.json.map(Json));
insert_constraint!("regex", raw.regex.map(Regex));
insert_constraint!("choice", raw.choice.map(Choice));
insert_constraint!("grammar", raw.grammar.map(Grammar));
match raw.json_object {
Some(true) => {
insert_constraint!("json_object", Some(JsonObject))
}
Some(false) => {
return Err(Error::InvalidStructuredOutputsParams {
message: "structured_outputs.json_object must be true if set; omit structured_outputs to disable structured outputs".to_string(),
});
}
None => {}
}
insert_constraint!("structural_tag", raw.structural_tag.map(StructuralTag));
Ok(Self {
constraint: constraint.map(|(_, c)| c).ok_or_else(|| {
Error::InvalidStructuredOutputsParams {
message: "missing structured output constraint".to_string(),
}
})?,
options: StructuredOutputOptions {
disable_any_whitespace: raw.disable_any_whitespace,
disable_additional_properties: raw.disable_additional_properties,
whitespace_pattern: raw.whitespace_pattern,
},
backend: raw.backend,
})
}
}
impl From<StructuredOutputsParams> for WireStructuredOutputsParams {
fn from(params: StructuredOutputsParams) -> Self {
let mut raw = Self {
disable_any_whitespace: params.options.disable_any_whitespace,
disable_additional_properties: params.options.disable_additional_properties,
whitespace_pattern: params.options.whitespace_pattern,
backend: params.backend,
..Self::default()
};
match params.constraint {
StructuredOutputConstraint::Json(json) => raw.json = Some(json),
StructuredOutputConstraint::Regex(regex) => raw.regex = Some(regex),
StructuredOutputConstraint::Choice(choice) => raw.choice = Some(choice),
StructuredOutputConstraint::Grammar(grammar) => raw.grammar = Some(grammar),
StructuredOutputConstraint::JsonObject => raw.json_object = Some(true),
StructuredOutputConstraint::StructuralTag(structural_tag) => {
raw.structural_tag = Some(structural_tag);
}
}
raw
}
}
impl Serialize for StructuredOutputsParams {
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
where
S: Serializer,
{
WireStructuredOutputsParams::from(self.clone()).serialize(serializer)
}
}
impl<'de> Deserialize<'de> for StructuredOutputsParams {
fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
where
D: Deserializer<'de>,
{
WireStructuredOutputsParams::deserialize(deserializer)?
.try_into()
.map_err(serde::de::Error::custom)
}
}
#[cfg(test)]
@@ -74,8 +247,59 @@ mod tests {
.unwrap();
assert_eq!(params.backend, StructuredOutputBackend::Guidance);
assert_eq!(params.constraint, StructuredOutputConstraint::JsonObject);
let value = serde_json::to_value(params).unwrap();
assert_eq!(value["_backend"], "guidance");
}
#[test]
fn structured_outputs_rejects_missing_constraint() {
let error =
serde_json::from_value::<StructuredOutputsParams>(serde_json::json!({})).unwrap_err();
assert!(error.to_string().contains("missing structured output constraint"));
}
#[test]
fn structured_outputs_rejects_multiple_constraints() {
let error = serde_json::from_value::<StructuredOutputsParams>(serde_json::json!({
"json": {"type": "object"},
"regex": ".*",
}))
.unwrap_err();
assert!(
error
.to_string()
.contains("multiple structured output constraints specified: json, regex")
);
}
#[test]
fn structured_outputs_rejects_json_object_false() {
let error = serde_json::from_value::<StructuredOutputsParams>(serde_json::json!({
"json_object": false,
}))
.unwrap_err();
assert!(error.to_string().contains("json_object must be true"));
}
#[test]
fn structured_outputs_serializes_through_raw_shape() {
let params = StructuredOutputsParams {
constraint: StructuredOutputConstraint::StructuralTag(
r#"{"type":"structural_tag"}"#.to_string(),
),
options: StructuredOutputOptions::default(),
backend: StructuredOutputBackend::Xgrammar,
};
let value = serde_json::to_value(params).unwrap();
assert_eq!(value["structural_tag"], r#"{"type":"structural_tag"}"#);
assert_eq!(value["_backend"], "xgrammar");
assert!(value.get("json").is_none());
}
}
rust/src/engine-core-client/src/tests/client.rs
+163 -219
View File
@@ -23,7 +23,8 @@ use crate::protocol::multimodal::{
SliceSpec,
};
use crate::protocol::output::{
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, decode_engine_core_outputs,
DpControlMessage, DpControlOutput, EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs,
RequestBatchOutputs, UtilityCallOutput, decode_engine_core_outputs,
};
use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
use crate::protocol::sampling::EngineCoreSamplingParams;
@@ -571,8 +572,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
send_outputs(
&mut data_socket.push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![request_output(
"req-1",
vec![],
@@ -580,17 +580,19 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
)],
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
send_outputs(
&mut coordinator.output_push,
EngineCoreOutputs {
DpControlOutput {
engine_index: 0,
wave_complete: Some(0),
..Default::default()
},
timestamp: 0.0,
control: DpControlMessage::WaveComplete(0),
}
.into(),
)
.await;
@@ -605,8 +607,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
send_outputs(
&mut data_socket.push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![request_output(
"req-3",
vec![],
@@ -614,7 +615,8 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
)],
finished_requests: Some(BTreeSet::from(["req-3".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
@@ -651,7 +653,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
send_outputs(
&mut data_socket.push,
EngineCoreOutputs {
RequestBatchOutputs {
engine_index: 1,
outputs: vec![request_output(
"req-2",
@@ -660,7 +662,8 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
)],
finished_requests: Some(BTreeSet::from(["req-2".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
@@ -766,11 +769,12 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() {
send_outputs(
&mut coordinator.output_push,
EngineCoreOutputs {
DpControlOutput {
engine_index: 1,
start_wave: Some(4),
..Default::default()
},
timestamp: 0.0,
control: DpControlMessage::StartWave(4),
}
.into(),
)
.await;
@@ -821,15 +825,16 @@ async fn coordinator_accepts_stats_only_outputs() {
send_outputs(
&mut coordinator.output_push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: Vec::new(),
scheduler_stats: Some(Box::new(SchedulerStats {
num_running_reqs: 1,
current_wave: 0,
..Default::default()
})),
..Default::default()
},
}
.into(),
)
.await;
@@ -840,8 +845,7 @@ async fn coordinator_accepts_stats_only_outputs() {
send_outputs(
&mut data_socket.push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![request_output(
"req-stats",
vec![],
@@ -849,7 +853,8 @@ async fn coordinator_accepts_stats_only_outputs() {
)],
finished_requests: Some(BTreeSet::from(["req-stats".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
@@ -910,30 +915,34 @@ async fn client_fail_closes_when_main_output_path_receives_dp_control() {
send_outputs(
push,
EngineCoreOutputs {
utility_output: Some(UtilityOutput {
UtilityCallOutput {
output: UtilityOutput {
call_id: 1_u64.into(),
failure_message: None,
result: None,
}),
},
..Default::default()
},
}
.into(),
)
.await;
send_outputs(
push,
EngineCoreOutputs {
start_wave: Some(3),
..Default::default()
},
DpControlOutput {
engine_index: 0,
timestamp: 0.0,
control: DpControlMessage::StartWave(3),
}
.into(),
)
.await;
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![request_output("req-1", vec![999], None)],
..Default::default()
},
}
.into(),
)
.await;
@@ -985,91 +994,6 @@ async fn client_fail_closes_when_main_output_path_receives_dp_control() {
client.shutdown().await.unwrap();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn client_fail_closes_when_main_output_path_receives_mixed_shape_output() {
init_tracing();
let ipc = IpcNamespace::new().unwrap();
let handshake_address = ipc.handshake_endpoint();
let engine_id = b"engine-0".to_vec();
let (shutdown_tx, engine_task) = spawn_mock_engine_task(
handshake_address.clone(),
engine_id.clone(),
|dealer, push| {
Box::pin(async move {
let add_1 = recv_engine_message(dealer).await;
assert_eq!(add_1[0].as_ref(), &[0x00]);
let request_1: EngineCoreRequest = rmp_serde::from_slice(&add_1[1]).unwrap();
assert_eq!(request_1.client_index, 7);
assert_eq!(request_1.request_id, "req-1");
let add_2 = recv_engine_message(dealer).await;
assert_eq!(add_2[0].as_ref(), &[0x00]);
let request_2: EngineCoreRequest = rmp_serde::from_slice(&add_2[1]).unwrap();
assert_eq!(request_2.client_index, 7);
assert_eq!(request_2.request_id, "req-2");
send_outputs(
push,
EngineCoreOutputs {
utility_output: Some(UtilityOutput {
call_id: 1_u64.into(),
failure_message: None,
result: None,
}),
outputs: vec![request_output("req-1", vec![999], None)],
..Default::default()
},
)
.await;
tokio::time::sleep(Duration::from_millis(50)).await;
})
},
);
let client = connect_client_with_ipc(
handshake_test_config(
handshake_address,
1,
"test-model",
Duration::from_secs(2),
7,
None,
),
&ipc,
)
.await;
assert_eq!(client.engine_identities()[0], b"engine-0");
assert!(client.ready_responses()[0].max_model_len > 0);
let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap();
let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap();
let error_2 = timeout(Duration::from_secs(1), stream_2.next())
.await
.unwrap()
.unwrap()
.unwrap_err();
assert!(is_unexpected_dispatcher_output(&error_2));
let error_1 = timeout(Duration::from_secs(1), stream_1.next())
.await
.unwrap()
.unwrap()
.unwrap_err();
assert!(is_unexpected_dispatcher_output(&error_1));
assert!(matches!(
client.health_error().as_deref(),
Some(error) if is_unexpected_dispatcher_output(error)
));
let _ = shutdown_tx.send(());
engine_task.await.unwrap();
client.shutdown().await.unwrap();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn duplicate_request_ids_are_rejected_without_sending_a_second_add() {
init_tracing();
@@ -1091,7 +1015,7 @@ async fn duplicate_request_ids_are_rejected_without_sending_a_second_add() {
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![request_output(
"req-1",
vec![],
@@ -1099,7 +1023,8 @@ async fn duplicate_request_ids_are_rejected_without_sending_a_second_add() {
)],
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -1158,10 +1083,12 @@ async fn finished_requests_without_final_output_is_treated_as_unexpected_close()
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: Vec::new(),
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
@@ -1217,10 +1144,11 @@ async fn dropping_a_live_stream_triggers_abort() {
assert_eq!(add[0].as_ref(), &[0x00]);
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![request_output("req-1", vec![99], None)],
..Default::default()
},
}
.into(),
)
.await;
@@ -1275,14 +1203,15 @@ async fn dropping_multiple_live_streams_aborts_all_in_a_burst() {
}
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![
request_output("req-1", vec![99], None),
request_output("req-2", vec![99], None),
request_output("req-3", vec![99], None),
],
..Default::default()
},
}
.into(),
)
.await;
@@ -1436,14 +1365,16 @@ async fn is_sleeping_wrapper_sends_typed_request_and_returns_typed_response() {
send_outputs(
push,
EngineCoreOutputs {
utility_output: Some(UtilityOutput {
UtilityCallOutput {
engine_index: 0,
timestamp: 0.0,
output: UtilityOutput {
call_id: call_id.into(),
failure_message: None,
result: Some(utility_result_value(true)),
}),
..Default::default()
},
},
}
.into(),
)
.await;
})
@@ -1491,14 +1422,16 @@ async fn call_utility_failure_message_surfaces_as_error() {
send_outputs(
push,
EngineCoreOutputs {
utility_output: Some(UtilityOutput {
UtilityCallOutput {
engine_index: 0,
timestamp: 0.0,
output: UtilityOutput {
call_id: call_id.into(),
failure_message: Some("boom".to_string()),
result: None,
}),
..Default::default()
},
},
}
.into(),
)
.await;
})
@@ -1815,8 +1748,7 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
finish_req_1_rx.await.unwrap();
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![request_output(
&request_1.request_id,
vec![10],
@@ -1824,7 +1756,8 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
)],
finished_requests: Some(BTreeSet::from([request_1.request_id.clone()])),
..Default::default()
},
}
.into(),
)
.await;
@@ -1836,8 +1769,7 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
finish_req_3_rx.await.unwrap();
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![request_output(
&request_3.request_id,
vec![30],
@@ -1845,7 +1777,8 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
)],
finished_requests: Some(BTreeSet::from([request_3.request_id.clone()])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -1864,7 +1797,7 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
finish_req_2_rx.await.unwrap();
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
engine_index: 1,
outputs: vec![request_output(
&request_2.request_id,
@@ -1873,7 +1806,8 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
)],
finished_requests: Some(BTreeSet::from([request_2.request_id.clone()])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -1990,14 +1924,16 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
assert_eq!(array[2], Value::from("is_sleeping"));
send_outputs(
push,
EngineCoreOutputs {
utility_output: Some(UtilityOutput {
UtilityCallOutput {
engine_index: 0,
timestamp: 0.0,
output: UtilityOutput {
call_id: call_id.into(),
failure_message: None,
result: Some(utility_result_value(true)),
}),
..Default::default()
},
},
}
.into(),
)
.await;
@@ -2012,8 +1948,7 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
assert_eq!(aborted_ids, vec!["req-1".to_string()]);
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![request_output(
"req-1",
vec![],
@@ -2021,7 +1956,8 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
)],
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -2044,14 +1980,16 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
assert_eq!(array[2], Value::from("is_sleeping"));
send_outputs(
push,
EngineCoreOutputs {
utility_output: Some(UtilityOutput {
UtilityCallOutput {
engine_index: 0,
timestamp: 0.0,
output: UtilityOutput {
call_id: call_id.into(),
failure_message: None,
result: Some(utility_result_value(true)),
}),
..Default::default()
},
},
}
.into(),
)
.await;
@@ -2066,7 +2004,7 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
assert_eq!(aborted_ids, vec!["req-2".to_string()]);
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
engine_index: 1,
outputs: vec![request_output(
"req-2",
@@ -2075,7 +2013,8 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
)],
finished_requests: Some(BTreeSet::from(["req-2".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -2155,14 +2094,16 @@ async fn collective_rpc_flattens_results_from_all_engines() {
send_outputs(
push,
EngineCoreOutputs {
utility_output: Some(UtilityOutput {
UtilityCallOutput {
engine_index: 0,
timestamp: 0.0,
output: UtilityOutput {
call_id: call_id.into(),
failure_message: None,
result: Some(utility_result_value(vec!["engine-0-worker"])),
}),
..Default::default()
},
},
}
.into(),
)
.await;
})
@@ -2186,14 +2127,16 @@ async fn collective_rpc_flattens_results_from_all_engines() {
send_outputs(
push,
EngineCoreOutputs {
utility_output: Some(UtilityOutput {
UtilityCallOutput {
engine_index: 0,
timestamp: 0.0,
output: UtilityOutput {
call_id: call_id.into(),
failure_message: None,
result: Some(utility_result_value(vec!["engine-1-worker"])),
}),
..Default::default()
},
},
}
.into(),
)
.await;
})
@@ -2264,14 +2207,16 @@ fn spawn_mock_utility_engine(
assert_eq!(array[3], expected_args, "unexpected utility args");
send_outputs(
push,
EngineCoreOutputs {
utility_output: Some(UtilityOutput {
UtilityCallOutput {
engine_index: 0,
timestamp: 0.0,
output: UtilityOutput {
call_id: call_id.into(),
failure_message: None,
result: Some(utility_result_value(result)),
}),
..Default::default()
},
},
}
.into(),
)
.await;
})
@@ -2552,41 +2497,40 @@ fn python_msgpack_fixtures_match_rust_encoding() {
let decoded_outputs: EngineCoreOutputs = rmp_serde::from_slice(&outputs_bytes).unwrap();
expect_test::expect![[r#"
EngineCoreOutputs {
engine_index: 0,
outputs: [
EngineCoreOutput {
request_id: "req-1",
new_token_ids: [
7,
8,
],
new_logprobs: None,
new_prompt_logprobs_tensors: None,
pooling_output: None,
finish_reason: Some(
Length,
),
stop_reason: None,
events: None,
kv_transfer_params: None,
trace_headers: None,
prefill_stats: None,
routed_experts: None,
num_nans_in_logits: 0,
},
],
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: Some(
{
"req-1",
},
),
wave_complete: None,
start_wave: None,
}
RequestBatch(
RequestBatchOutputs {
engine_index: 0,
outputs: [
EngineCoreOutput {
request_id: "req-1",
new_token_ids: [
7,
8,
],
new_logprobs: None,
new_prompt_logprobs_tensors: None,
pooling_output: None,
finish_reason: Some(
Length,
),
stop_reason: None,
events: None,
kv_transfer_params: None,
trace_headers: None,
prefill_stats: None,
routed_experts: None,
num_nans_in_logits: 0,
},
],
scheduler_stats: None,
timestamp: 0.0,
finished_requests: Some(
{
"req-1",
},
),
},
)
"#]]
.assert_debug_eq(&decoded_outputs);
@@ -2599,7 +2543,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
let inline_logprobs =
decode_engine_core_outputs(&decode_frames(inline_logprobs_frames)).unwrap();
expect_sample_logprobs(
inline_logprobs.outputs[0]
inline_logprobs.as_request_batch().unwrap().outputs[0]
.new_logprobs
.as_ref()
.expect("inline logprobs decoded"),
@@ -2608,7 +2552,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
let multipart_logprobs =
decode_engine_core_outputs(&decode_frames(multipart_logprobs_frames)).unwrap();
expect_sample_logprobs(
multipart_logprobs.outputs[0]
multipart_logprobs.as_request_batch().unwrap().outputs[0]
.new_logprobs
.as_ref()
.expect("multipart logprobs decoded"),
@@ -2616,7 +2560,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
let inline_prompt = decode_engine_core_outputs(&decode_frames(inline_prompt_frames)).unwrap();
expect_prompt_logprobs(
inline_prompt.outputs[0]
inline_prompt.as_request_batch().unwrap().outputs[0]
.new_prompt_logprobs_tensors
.as_ref()
.expect("inline prompt logprobs decoded"),
@@ -2625,7 +2569,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
let multipart_prompt =
decode_engine_core_outputs(&decode_frames(multipart_prompt_frames)).unwrap();
expect_prompt_logprobs(
multipart_prompt.outputs[0]
multipart_prompt.as_request_batch().unwrap().outputs[0]
.new_prompt_logprobs_tensors
.as_ref()
.expect("multipart prompt logprobs decoded"),
@@ -2945,8 +2889,7 @@ async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends
send_outputs(
&mut push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![request_output(
"req-1",
vec![],
@@ -2954,7 +2897,8 @@ async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends
)],
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
@@ -3020,8 +2964,7 @@ async fn bootstrapped_external_coordinator_running_state_suppresses_wakeup() {
send_outputs(
&mut push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![request_output(
"req-1",
vec![],
@@ -3029,7 +2972,8 @@ async fn bootstrapped_external_coordinator_running_state_suppresses_wakeup() {
)],
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
..Default::default()
},
}
.into(),
)
.await;
rust/src/llm/tests/generate.rs
+36 -30
View File
@@ -11,7 +11,7 @@ use vllm_engine_core_client::protocol::logprobs::{
};
use vllm_engine_core_client::protocol::output::{
EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput,
EngineCoreOutputs,
EngineCoreOutputs, RequestBatchOutputs,
};
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
@@ -251,7 +251,7 @@ async fn generate_streams_outputs() {
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![
request_output_with_logprobs(
&request.request_id,
@@ -270,7 +270,8 @@ async fn generate_streams_outputs() {
],
finished_requests: Some(BTreeSet::from([request.request_id.clone()])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -331,8 +332,7 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![
EngineCoreOutput {
prefill_stats: Some(PrefillStats {
@@ -358,13 +358,9 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() {
Some(serde_json::json!({"connector": "x"})),
),
],
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: None,
wave_complete: None,
start_wave: None,
},
..Default::default()
}
.into(),
)
.await;
})
@@ -412,10 +408,12 @@ async fn generate_propagates_unexpected_close_errors() {
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: Vec::new(),
finished_requests: Some(BTreeSet::from([request.request_id])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -459,10 +457,11 @@ async fn dropping_a_live_generate_stream_triggers_abort() {
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![request_output(&request.request_id, vec![99], None)],
..Default::default()
},
}
.into(),
)
.await;
@@ -513,7 +512,7 @@ async fn duplicate_external_request_ids_are_randomized_before_reaching_engine_co
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![request_output(
&request_1.request_id,
vec![],
@@ -521,13 +520,14 @@ async fn duplicate_external_request_ids_are_randomized_before_reaching_engine_co
)],
finished_requests: Some(BTreeSet::from([request_1.request_id.clone()])),
..Default::default()
},
}
.into(),
)
.await;
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![request_output(
&request_2.request_id,
vec![],
@@ -535,7 +535,8 @@ async fn duplicate_external_request_ids_are_randomized_before_reaching_engine_co
)],
finished_requests: Some(BTreeSet::from([request_2.request_id])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -576,10 +577,11 @@ async fn abort_resolves_external_request_id_to_internal_before_reaching_engine()
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![request_output(&request.request_id, vec![7], None)],
..Default::default()
},
}
.into(),
)
.await;
@@ -643,13 +645,14 @@ async fn abort_by_external_id_aborts_all_internal_requests() {
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
outputs: vec![
request_output(&request_1.request_id, vec![7], None),
request_output(&request_2.request_id, vec![8], None),
],
..Default::default()
},
}
.into(),
)
.await;
@@ -710,7 +713,7 @@ async fn generate_records_request_metrics_in_prometheus_output() {
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
engine_index: 4,
timestamp: 10.0,
outputs: vec![EngineCoreOutput {
@@ -736,13 +739,14 @@ async fn generate_records_request_metrics_in_prometheus_output() {
)
}],
..Default::default()
},
}
.into(),
)
.await;
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
engine_index: 4,
timestamp: 11.5,
outputs: vec![request_output_with_events(
@@ -756,7 +760,8 @@ async fn generate_records_request_metrics_in_prometheus_output() {
)],
finished_requests: Some(BTreeSet::from([request.request_id])),
..Default::default()
},
}
.into(),
)
.await;
})
@@ -843,7 +848,7 @@ async fn dropping_stream_records_abort_terminal_request_metrics() {
send_outputs(
push,
EngineCoreOutputs {
RequestBatchOutputs {
engine_index: 5,
timestamp: 10.0,
outputs: vec![request_output_with_events(
@@ -862,7 +867,8 @@ async fn dropping_stream_records_abort_terminal_request_metrics() {
]),
)],
..Default::default()
},
}
.into(),
)
.await;
rust/src/mock-engine/src/engine.rs
+16 -12
View File
@@ -12,7 +12,8 @@ use tokio::task::yield_now;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, warn};
use vllm_engine_core_client::protocol::output::{
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs,
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs,
UtilityCallOutput,
};
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
use vllm_engine_core_client::protocol::utility::{
@@ -61,13 +62,14 @@ fn empty_finish_outputs(
let output = request_output(request_id, Vec::new(), Some(finish_reason));
let finished_requests = BTreeSet::from([output.request_id.clone()]);
EngineCoreOutputs {
RequestBatchOutputs {
engine_index,
outputs: vec![output],
timestamp: now_secs(),
finished_requests: Some(finished_requests),
..Default::default()
}
.into()
}
/// Encode a utility result into the protocol's msgpack value envelope.
@@ -98,16 +100,16 @@ fn utility_response(
_ => utility_envelope(Value::Nil),
}?;
Ok(EngineCoreOutputs {
Ok(UtilityCallOutput {
engine_index,
utility_output: Some(UtilityOutput {
timestamp: now_secs(),
output: UtilityOutput {
call_id: request.call_id,
failure_message: None,
result: Some(result),
}),
timestamp: now_secs(),
..Default::default()
})
},
}
.into())
}
/// Message sent from the frontend to the mock engine task to drive the engine loop.
@@ -270,13 +272,14 @@ impl Engine {
}
for (client_index, (client_outputs, finished_requests)) in outputs_by_client {
outputs.push({
let outputs = EngineCoreOutputs {
let outputs = RequestBatchOutputs {
engine_index: self.engine_index,
outputs: client_outputs,
timestamp: now_secs(),
finished_requests: Some(finished_requests),
..Default::default()
};
}
.into();
EngineOutput {
client_index,
outputs,
@@ -358,14 +361,15 @@ impl Engine {
.filter_map(|(client_index, (outputs, finished_requests))| {
(!outputs.is_empty()).then(|| EngineOutput {
client_index,
outputs: EngineCoreOutputs {
outputs: RequestBatchOutputs {
engine_index: self.engine_index,
outputs,
timestamp: now_secs(),
finished_requests: (!finished_requests.is_empty())
.then_some(finished_requests),
..Default::default()
},
}
.into(),
})
})
.collect()
rust/src/server/src/grpc/convert.rs
+10 -24
View File
@@ -231,32 +231,18 @@ fn convert_structured_output(
StructuredOutput::Json(schema) => {
let json: serde_json::Value = serde_json::from_str(schema)
.map_err(|e| Status::invalid_argument(format!("invalid json schema: {e}")))?;
StructuredOutputsParams {
json: Some(json),
..Default::default()
}
StructuredOutputsParams::json(json)
}
StructuredOutput::Regex(regex) => StructuredOutputsParams {
regex: Some(regex.clone()),
..Default::default()
},
StructuredOutput::Choice(choices) => StructuredOutputsParams {
choice: Some(choices.choices.clone()),
..Default::default()
},
StructuredOutput::Grammar(grammar) => StructuredOutputsParams {
grammar: Some(grammar.clone()),
..Default::default()
},
StructuredOutput::JsonObject(true) => StructuredOutputsParams {
json_object: Some(true),
..Default::default()
},
StructuredOutput::Regex(regex) => StructuredOutputsParams::regex(regex.clone()),
StructuredOutput::Choice(choices) => {
StructuredOutputsParams::choice(choices.choices.clone())
}
StructuredOutput::Grammar(grammar) => StructuredOutputsParams::grammar(grammar.clone()),
StructuredOutput::JsonObject(true) => StructuredOutputsParams::json_object(),
StructuredOutput::JsonObject(false) => return Ok(None),
StructuredOutput::StructuralTag(tag) => StructuredOutputsParams {
structural_tag: Some(tag.clone()),
..Default::default()
},
StructuredOutput::StructuralTag(tag) => {
StructuredOutputsParams::structural_tag(tag.clone())
}
};
Ok(Some(params))
}
rust/src/server/src/grpc/tests.rs
+4 -9
View File
@@ -19,7 +19,7 @@ use vllm_chat::{
DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt,
};
use vllm_engine_core_client::protocol::output::{
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs,
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs,
};
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
@@ -115,19 +115,14 @@ fn engine_outputs_for_request(
request_id: &str,
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
) -> EngineCoreOutputs {
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: output_specs
.into_iter()
.map(|(token_ids, finish_reason)| request_output(request_id, token_ids, finish_reason))
.collect(),
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: None,
wave_complete: None,
start_wave: None,
..Default::default()
}
.into()
}
fn default_stream_output_specs() -> Vec<(Vec<u32>, Option<EngineCoreFinishReason>)> {
rust/src/server/src/routes/http_client_tests.rs
+4 -9
View File
@@ -19,7 +19,7 @@ use vllm_chat::{
DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt,
};
use vllm_engine_core_client::protocol::output::{
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs,
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs,
};
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
@@ -111,19 +111,14 @@ fn engine_outputs_for_request(
request_id: &str,
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
) -> EngineCoreOutputs {
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: output_specs
.into_iter()
.map(|(token_ids, finish_reason)| request_output(request_id, token_ids, finish_reason))
.collect(),
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: None,
wave_complete: None,
start_wave: None,
..Default::default()
}
.into()
}
fn default_stream_output_specs() -> Vec<(Vec<u32>, Option<EngineCoreFinishReason>)> {
rust/src/server/src/routes/openai/utils/structured_outputs.rs
+6 -13
View File
@@ -37,7 +37,7 @@ pub enum ResponseFormat {
},
/// vLLM-specific structural tag format. The entire object (including the
/// `type` field) is JSON-serialized and passed as
/// `StructuredOutputsParams.structural_tag`.
/// `StructuredOutputConstraint::StructuralTag`.
///
/// We capture the payload as a catch-all map so both the legacy
/// (`structures`/`triggers`) and current (`format`) shapes are
@@ -79,14 +79,10 @@ pub fn convert_from_response_format(
};
match fmt {
ResponseFormat::Text => Ok(None),
ResponseFormat::JsonObject => Ok(Some(StructuredOutputsParams {
json_object: Some(true),
..Default::default()
})),
ResponseFormat::JsonSchema { json_schema } => Ok(Some(StructuredOutputsParams {
json: Some(json_schema.schema.clone()),
..Default::default()
})),
ResponseFormat::JsonObject => Ok(Some(StructuredOutputsParams::json_object())),
ResponseFormat::JsonSchema { json_schema } => Ok(Some(StructuredOutputsParams::json(
json_schema.schema.clone(),
))),
ResponseFormat::StructuralTag { .. } => {
// The Python frontend dumps the entire response_format object (including the
// `type` field) as a JSON string for the engine-core backend.
@@ -96,10 +92,7 @@ pub fn convert_from_response_format(
Some("response_format"),
)
})?;
Ok(Some(StructuredOutputsParams {
structural_tag: Some(tag_json),
..Default::default()
}))
Ok(Some(StructuredOutputsParams::structural_tag(tag_json)))
}
}
}
rust/src/server/src/routes/tests.rs
+45 -91
View File
@@ -29,7 +29,8 @@ use vllm_engine_core_client::protocol::logprobs::{
Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
};
use vllm_engine_core_client::protocol::output::{
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, StopReason,
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs, StopReason,
UtilityCallOutput,
};
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
use vllm_engine_core_client::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
@@ -232,19 +233,14 @@ fn engine_outputs_for_request(
request_id: &str,
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
) -> EngineCoreOutputs {
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: output_specs
.into_iter()
.map(|(token_ids, finish_reason)| request_output(request_id, token_ids, finish_reason))
.collect(),
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: None,
wave_complete: None,
start_wave: None,
..Default::default()
}
.into()
}
fn test_llm(client: EngineCoreClient) -> Llm {
@@ -390,14 +386,15 @@ fn utility_none_result() -> UtilityResultEnvelope {
}
fn utility_outputs(call_id: u64, result: UtilityResultEnvelope) -> EngineCoreOutputs {
EngineCoreOutputs {
utility_output: Some(UtilityOutput {
UtilityCallOutput {
output: UtilityOutput {
call_id: call_id.into(),
failure_message: None,
result: Some(result),
}),
},
..Default::default()
}
.into()
}
async fn send_outputs(push: &mut PushSocket, outputs: EngineCoreOutputs) {
@@ -2281,8 +2278,7 @@ async fn non_stream_chat_includes_logprobs_and_prompt_logprobs() {
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![request_output_with_logprobs(
&request.request_id,
bytes_to_token_ids(b"hi"),
@@ -2291,13 +2287,9 @@ async fn non_stream_chat_includes_logprobs_and_prompt_logprobs() {
Some(sample_logprobs_for_tokens(&bytes_to_token_ids(b"hi"))),
Some(prompt_logprobs_for_tokens(&prompt_token_ids)),
)],
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: None,
wave_complete: None,
start_wave: None,
},
..Default::default()
}
.into(),
)
.await;
})
@@ -3121,8 +3113,7 @@ async fn non_stream_completions_include_logprobs() {
rmp_serde::from_slice(&add[1]).expect("decode request");
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![
request_output_with_logprobs(
&request.request_id,
@@ -3141,13 +3132,10 @@ async fn non_stream_completions_include_logprobs() {
None,
),
],
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: Some(BTreeSet::from([request.request_id.clone()])),
wave_complete: None,
start_wave: None,
},
..Default::default()
}
.into(),
)
.await;
})
@@ -3224,8 +3212,7 @@ async fn non_stream_completions_include_prompt_logprobs() {
rmp_serde::from_slice(&add[1]).expect("decode request");
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![request_output_with_logprobs(
&request.request_id,
vec![b'h' as u32, b'i' as u32, b'!' as u32],
@@ -3246,13 +3233,9 @@ async fn non_stream_completions_include_prompt_logprobs() {
}),
Some(prompt_logprobs_for_hello()),
)],
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: None,
wave_complete: None,
start_wave: None,
},
..Default::default()
}
.into(),
)
.await;
})
@@ -3588,8 +3571,7 @@ async fn non_stream_raw_generate_returns_token_output_envelope() {
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![
request_output_with_logprobs(
&request.request_id,
@@ -3609,13 +3591,9 @@ async fn non_stream_raw_generate_returns_token_output_envelope() {
Some(json!({"connector": "x"})),
),
],
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: None,
wave_complete: None,
start_wave: None,
},
..Default::default()
}
.into(),
)
.await;
})
@@ -3708,8 +3686,7 @@ async fn stream_raw_generate_returns_sse_chunks_and_usage() {
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![
request_output_with_logprobs(
&request.request_id,
@@ -3728,13 +3705,9 @@ async fn stream_raw_generate_returns_sse_chunks_and_usage() {
None,
),
],
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: None,
wave_complete: None,
start_wave: None,
},
..Default::default()
}
.into(),
)
.await;
})
@@ -4426,8 +4399,7 @@ async fn include_reasoning_false_suppresses_non_stream_output_metadata() {
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![
request_output_with_logprobs(
&request.request_id,
@@ -4446,13 +4418,9 @@ async fn include_reasoning_false_suppresses_non_stream_output_metadata() {
None,
),
],
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: None,
wave_complete: None,
start_wave: None,
},
..Default::default()
}
.into(),
)
.await;
})
@@ -4601,8 +4569,7 @@ async fn tool_call_sse_chunks_can_carry_logprobs() {
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![request_output_with_logprobs(
&request.request_id,
bytes_to_token_ids(b"<think>Need tool.</think>"),
@@ -4613,19 +4580,14 @@ async fn tool_call_sse_chunks_can_carry_logprobs() {
))),
None,
)],
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: None,
wave_complete: None,
start_wave: None,
},
..Default::default()
}
.into(),
)
.await;
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![request_output_with_logprobs(
&request.request_id,
bytes_to_token_ids(b"<tool_call>\n{\"name\":\"get_weather\", "),
@@ -4636,19 +4598,14 @@ async fn tool_call_sse_chunks_can_carry_logprobs() {
))),
None,
)],
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: None,
wave_complete: None,
start_wave: None,
},
..Default::default()
}
.into(),
)
.await;
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
RequestBatchOutputs {
outputs: vec![request_output_with_logprobs(
&request.request_id,
bytes_to_token_ids(
@@ -4661,13 +4618,10 @@ async fn tool_call_sse_chunks_can_carry_logprobs() {
))),
None,
)],
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: Some(BTreeSet::from([request.request_id.clone()])),
wave_complete: None,
start_wave: None,
},
..Default::default()
}
.into(),
)
.await;
})