diff --git a/rust/src/chat/src/output/default/structural_tag.rs b/rust/src/chat/src/output/default/structural_tag.rs index 4dbebd50fd..45a508ffee 100644 --- a/rust/src/chat/src/output/default/structural_tag.rs +++ b/rust/src/chat/src/output/default/structural_tag.rs @@ -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()); } } diff --git a/rust/src/chat/tests/chat.rs b/rust/src/chat/tests/chat.rs index b209b7b769..77e93ac5c5 100644 --- a/rust/src/chat/tests/chat.rs +++ b/rust/src/chat/tests/chat.rs @@ -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"innerouter"), @@ -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; } diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 50965cbb3f..3465817a45 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -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:?}" ))?; diff --git a/rust/src/engine-core-client/src/coordinator/inproc.rs b/rust/src/engine-core-client/src/coordinator/inproc.rs index 526c640003..aaa6d8171b 100644 --- a/rust/src/engine-core-client/src/coordinator/inproc.rs +++ b/rust/src/engine-core-client/src/coordinator/inproc.rs @@ -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) => { diff --git a/rust/src/engine-core-client/src/error.rs b/rust/src/engine-core-client/src/error.rs index 2172c98374..f4fc5c4823 100644 --- a/rust/src/engine-core-client/src/error.rs +++ b/rust/src/engine-core-client/src/error.rs @@ -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")] diff --git a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs index 6fbc57378a..17310c7a30 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs @@ -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, diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index c82ecfa2b8..157343c20b 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -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: /// #[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, + outputs: Vec, #[serde(default)] - pub scheduler_stats: Option>, + scheduler_stats: Option>, #[serde(default)] - pub timestamp: f64, + timestamp: f64, #[serde(default)] - pub utility_output: Option, + utility_output: Option, #[serde(default)] - pub finished_requests: Option>, + finished_requests: Option>, /// In DP mode, signals that the current wave finished and engines are /// paused. #[serde(default)] - pub wave_complete: Option, + wave_complete: Option, /// 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, + start_wave: Option, } /// 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, @@ -176,30 +176,48 @@ pub struct RequestBatchOutputs { pub finished_requests: Option>, } -#[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 for EngineCoreOutputs { + fn from(outputs: RequestBatchOutputs) -> Self { + Self::RequestBatch(outputs) + } +} + +impl From for EngineCoreOutputs { + fn from(output: UtilityCallOutput) -> Self { + Self::Utility(output) + } +} + +impl From 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 for EngineCoreOutputs { + type Error = Error; + + fn try_from(value: WireEngineCoreOutputs) -> Result { + 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 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(&self, serializer: S) -> std::result::Result + where + S: Serializer, + { + WireEngineCoreOutputs::from(self.clone()).serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for EngineCoreOutputs { + fn deserialize(deserializer: D) -> std::result::Result + 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(frames: &[Frame]) -> Result @@ -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()); } } diff --git a/rust/src/engine-core-client/src/protocol/structured_outputs.rs b/rust/src/engine-core-client/src/protocol/structured_outputs.rs index 9bf8102aa6..ed82d3e558 100644 --- a/rust/src/engine-core-client/src/protocol/structured_outputs.rs +++ b/rust/src/engine-core-client/src/protocol/structured_outputs.rs @@ -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: -/// -#[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, + Json(Value), /// Regular expression the output must match. - pub regex: Option, + Regex(String), /// List of allowed output strings (the model must produce one of these). - pub choice: Option>, + Choice(Vec), /// Context-free grammar (in EBNF-like notation) the output must conform to. - pub grammar: Option, - /// When `true`, output must be valid JSON (free-form, no schema). - pub json_object: Option, + 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, - /// Structural tag configuration (JSON-encoded string). - pub structural_tag: Option, +} + +/// 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: +/// +#[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) -> Self { + Self::from_constraint(StructuredOutputConstraint::Regex(regex.into())) + } + + pub fn choice(choice: Vec) -> Self { + Self::from_constraint(StructuredOutputConstraint::Choice(choice)) + } + + pub fn grammar(grammar: impl Into) -> 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) -> 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, + regex: Option, + choice: Option>, + grammar: Option, + json_object: Option, + #[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, + structural_tag: Option, #[serde( default, rename = "_backend", deserialize_with = "serde_with::rust::deserialize_ignore_any" )] - pub backend: StructuredOutputBackend, + backend: StructuredOutputBackend, +} + +impl TryFrom for StructuredOutputsParams { + type Error = Error; + + fn try_from(raw: WireStructuredOutputsParams) -> Result { + 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 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(&self, serializer: S) -> std::result::Result + where + S: Serializer, + { + WireStructuredOutputsParams::from(self.clone()).serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for StructuredOutputsParams { + fn deserialize(deserializer: D) -> std::result::Result + 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::(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::(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::(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()); + } } diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index f60973d29b..00aae6fe0c 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -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; diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index b7942910d4..9911841ead 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -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; diff --git a/rust/src/mock-engine/src/engine.rs b/rust/src/mock-engine/src/engine.rs index 42a8f0efa0..73d60b5597 100644 --- a/rust/src/mock-engine/src/engine.rs +++ b/rust/src/mock-engine/src/engine.rs @@ -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() diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 9836745dd6..5dfeba08a5 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -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)) } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index e0d7ba635c..83c4de440e 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -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, Option)>, ) -> 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, Option)> { diff --git a/rust/src/server/src/routes/http_client_tests.rs b/rust/src/server/src/routes/http_client_tests.rs index 9c04701e86..6d479854a0 100644 --- a/rust/src/server/src/routes/http_client_tests.rs +++ b/rust/src/server/src/routes/http_client_tests.rs @@ -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, Option)>, ) -> 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, Option)> { diff --git a/rust/src/server/src/routes/openai/utils/structured_outputs.rs b/rust/src/server/src/routes/openai/utils/structured_outputs.rs index 24d261314f..e717765c51 100644 --- a/rust/src/server/src/routes/openai/utils/structured_outputs.rs +++ b/rust/src/server/src/routes/openai/utils/structured_outputs.rs @@ -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))) } } } diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 6b70f51413..4590afef31 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -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, Option)>, ) -> 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"Need tool."), @@ -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"\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; })