From c638f9216a08bfb5644d8a266ddd35421e04118d Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 1 Jul 2026 22:28:21 +0800 Subject: [PATCH] [Rust Frontend] Split engine core DTOs into separate modules (#47265) Signed-off-by: Bugen Zhao --- rust/src/chat/src/lib.rs | 3 +- rust/src/chat/src/multimodal.rs | 2 +- rust/src/chat/src/multimodal/tensor.rs | 2 +- .../chat/src/output/default/structural_tag.rs | 8 +- rust/src/chat/tests/chat.rs | 5 +- .../examples/external_engine_logprobs.rs | 6 +- rust/src/engine-core-client/src/client.rs | 3 +- rust/src/engine-core-client/src/client/imp.rs | 7 +- .../engine-core-client/src/client/state.rs | 4 +- .../engine-core-client/src/client/stream.rs | 2 +- .../src/coordinator/inproc.rs | 7 +- .../src/engine-core-client/src/mock_engine.rs | 3 +- .../src/protocol/handshake.rs | 3 +- .../src/protocol/logprobs.rs | 49 +- .../src/protocol/logprobs/tests.rs | 4 +- .../engine-core-client/src/protocol/mod.rs | 642 +----------------- .../{classified_outputs.rs => output.rs} | 222 +++++- .../src/protocol/request.rs | 175 +++++ .../src/protocol/sampling.rs | 211 ++++++ .../src/protocol/structured_outputs.rs | 81 +++ .../engine-core-client/src/tests/client.rs | 9 +- rust/src/engine-core-client/src/transport.rs | 5 +- .../src/llm/examples/external_engine_smoke.rs | 2 +- rust/src/llm/src/output.rs | 2 +- rust/src/llm/src/request.rs | 8 +- rust/src/llm/src/request_metrics.rs | 10 +- rust/src/llm/tests/generate.rs | 8 +- rust/src/mock-engine/src/engine.rs | 7 +- rust/src/mock-engine/src/io.rs | 5 +- rust/src/mock-engine/src/tests.rs | 6 +- rust/src/parser/benches/utils/mod.rs | 3 +- rust/src/parser/python/src/lib.rs | 3 +- rust/src/parser/src/reasoning/seed_oss.rs | 6 +- rust/src/parser/src/reasoning/step3p5.rs | 6 +- rust/src/parser/src/tool/mod.rs | 4 +- rust/src/parser/src/unified/combined.rs | 3 +- rust/src/parser/src/unified/gemma4.rs | 3 +- rust/src/parser/src/unified/mod.rs | 5 +- rust/src/server/src/error.rs | 3 +- rust/src/server/src/grpc/convert.rs | 5 +- rust/src/server/src/grpc/tests.rs | 5 +- rust/src/server/src/listener.rs | 2 +- rust/src/server/src/middleware/offload.rs | 2 +- .../server/src/routes/http_client_tests.rs | 5 +- .../src/routes/openai/chat_completions.rs | 4 +- .../routes/openai/utils/structured_outputs.rs | 2 +- rust/src/server/src/routes/tests.rs | 9 +- rust/src/server/src/state.rs | 3 +- rust/src/text/src/lower.rs | 6 +- rust/src/text/src/lower/logprobs.rs | 3 +- rust/src/text/src/lower/token_ids.rs | 2 +- rust/src/text/src/output/decoded.rs | 2 +- rust/src/text/src/request.rs | 3 +- rust/src/tokenizer/src/hf/added_tokens.rs | 5 +- 54 files changed, 806 insertions(+), 789 deletions(-) rename rust/src/engine-core-client/src/protocol/{classified_outputs.rs => output.rs} (50%) create mode 100644 rust/src/engine-core-client/src/protocol/request.rs create mode 100644 rust/src/engine-core-client/src/protocol/sampling.rs create mode 100644 rust/src/engine-core-client/src/protocol/structured_outputs.rs diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index c16921ea758..8284ddd1285 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -50,7 +50,8 @@ mod request; mod stream; use vllm_engine_core_client::EngineCoreClient; -use vllm_engine_core_client::protocol::{ModelDtype, ReasoningParserKwargs}; +use vllm_engine_core_client::protocol::dtype::ModelDtype; +use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_llm::Llm; use vllm_text::{Prompt, TextLlm, TextRequest}; diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index 8fd44376f99..9ec67cdd32e 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -22,7 +22,7 @@ use llm_multimodal::{ TrackedMedia, }; use tracing::warn; -use vllm_engine_core_client::protocol::ModelDtype; +use vllm_engine_core_client::protocol::dtype::ModelDtype; use vllm_engine_core_client::protocol::multimodal::{ MmBatchedField, MmFeatureSpec, MmFeatures, MmField, MmFieldElem, MmFlatField, MmKwargsItem, MmSharedField, MmSlice, PlaceholderRange, SliceSpec, diff --git a/rust/src/chat/src/multimodal/tensor.rs b/rust/src/chat/src/multimodal/tensor.rs index eddf8f707e9..b5a5f78f264 100644 --- a/rust/src/chat/src/multimodal/tensor.rs +++ b/rust/src/chat/src/multimodal/tensor.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use half::{bf16, f16}; use llm_multimodal::{ModelSpecificValue, PreprocessedImages}; -use vllm_engine_core_client::protocol::ModelDtype; +use vllm_engine_core_client::protocol::dtype::ModelDtype; use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue; use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor}; diff --git a/rust/src/chat/src/output/default/structural_tag.rs b/rust/src/chat/src/output/default/structural_tag.rs index 6ba2458ca8d..4dbebd50fd2 100644 --- a/rust/src/chat/src/output/default/structural_tag.rs +++ b/rust/src/chat/src/output/default/structural_tag.rs @@ -1,7 +1,9 @@ //! Applies xgrammar structural-tag constraints for strict tool calling. use thiserror_ext::AsReport; -use vllm_engine_core_client::protocol::{StructuredOutputBackend, StructuredOutputsParams}; +use vllm_engine_core_client::protocol::structured_outputs::{ + StructuredOutputBackend, StructuredOutputsParams, +}; use vllm_parser::tool::StructuralTagModel; use xgrammar_structural_tag::{ FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam, @@ -76,7 +78,9 @@ fn structural_tag_tool_choice(request: &ChatRequest) -> Option>; @@ -452,7 +452,7 @@ mod tests { EngineLoadSnapshot, EngineRoutingState, RequestRegistry, UtilityRegistry, }; use crate::mock_engine::default_ready_response; - use crate::protocol::{ + use crate::protocol::output::{ EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput, }; use crate::transport::ConnectedEngine; diff --git a/rust/src/engine-core-client/src/client/stream.rs b/rust/src/engine-core-client/src/client/stream.rs index 3cbb215b0ef..56c6a7cb663 100644 --- a/rust/src/engine-core-client/src/client/stream.rs +++ b/rust/src/engine-core-client/src/client/stream.rs @@ -10,7 +10,7 @@ use tracing::{debug, error, warn}; use crate::client::AbortRequest; use crate::client::state::OutputReceiver; -use crate::protocol::{EngineCoreFinishReason, EngineCoreOutput}; +use crate::protocol::output::{EngineCoreFinishReason, EngineCoreOutput}; use crate::{AbortCause, Error, Result}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/rust/src/engine-core-client/src/coordinator/inproc.rs b/rust/src/engine-core-client/src/coordinator/inproc.rs index 54c9f810d03..526c640003e 100644 --- a/rust/src/engine-core-client/src/coordinator/inproc.rs +++ b/rust/src/engine-core-client/src/coordinator/inproc.rs @@ -10,10 +10,9 @@ use zeromq::{XPubSocket, ZmqMessage}; use crate::client::imp::ClientInner; use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState}; use crate::error::{Error, Result, bail_unexpected_coordinator_output}; -use crate::protocol::{ - ClassifiedEngineCoreOutputs, DpControlMessage, EngineCoreOutputs, EngineCoreRequestType, - encode_msgpack, -}; +use crate::protocol::encode_msgpack; +use crate::protocol::output::{ClassifiedEngineCoreOutputs, DpControlMessage, EngineCoreOutputs}; +use crate::protocol::request::EngineCoreRequestType; /// Coordinator-to-engine `START_DP_WAVE` control payload encoded on the /// engine-facing coordinator socket. diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index be6947bd45a..781b004c7b9 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -8,8 +8,9 @@ use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage}; use crate::EngineId; use crate::error::{Error, Result, bail_unexpected_handshake_message}; +use crate::protocol::dtype::ModelDtype; use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage}; -use crate::protocol::{ModelDtype, decode_msgpack, encode_msgpack}; +use crate::protocol::{decode_msgpack, encode_msgpack}; /// Default model length advertised by reusable mock engine helpers. pub const DEFAULT_MOCK_MAX_MODEL_LEN: u64 = 1024 * 1024; diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index 1eea6630446..7a295209613 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -2,7 +2,8 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; -use crate::protocol::{ModelDtype, OpaqueValue}; +use crate::protocol::OpaqueValue; +use crate::protocol::dtype::ModelDtype; /// Decoded engine startup-handshake payload sent on the handshake socket. /// diff --git a/rust/src/engine-core-client/src/protocol/logprobs.rs b/rust/src/engine-core-client/src/protocol/logprobs.rs index 00c01df671c..24e6ae2fee1 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs.rs @@ -9,8 +9,7 @@ use enum_as_inner::EnumAsInner; use serde::{Deserialize, Deserializer, Serialize}; use self::wire::*; -use super::{EngineCoreOutput, EngineCoreOutputs, decode_msgpack}; -use crate::error::{Error, Result, bail_ext_value_decode, ext_value_decode}; +use crate::error::{Error, Result, bail_ext_value_decode}; use crate::protocol::tensor::{WireArrayData, WireNdArray}; /// One token candidate and its logprob metadata for a single sequence position. @@ -160,7 +159,7 @@ impl Serialize for MaybeWireLogprobs { impl MaybeWireLogprobs { /// Resolve the wire representation into decoded logprobs by looking up aux /// frames and decoding raw views as needed. - fn resolve(self, frames: &[Frame], field_prefix: &str) -> Result + pub(super) fn resolve(self, frames: &[Frame], field_prefix: &str) -> Result where Frame: AsRef<[u8]>, { @@ -171,37 +170,6 @@ impl MaybeWireLogprobs { } } -impl EngineCoreOutputs { - /// Resolve all wire-format fields in-place by looking up aux frames and - /// decoding raw-view payloads as needed. - fn resolve_in_place(&mut self, frames: &[Frame]) -> Result<()> - where - Frame: AsRef<[u8]>, - { - for output in &mut self.outputs { - output.resolve_in_place(frames)?; - } - Ok(()) - } -} - -impl EngineCoreOutput { - /// Resolve all wire-format fields in-place by looking up aux frames and - /// decoding raw-view payloads as needed. - fn resolve_in_place(&mut self, frames: &[Frame]) -> Result<()> - where - Frame: AsRef<[u8]>, - { - self.new_logprobs = (self.new_logprobs.take()) - .map(|value| value.resolve(frames, "new_logprobs")) - .transpose()?; - self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take()) - .map(|value| value.resolve(frames, "new_prompt_logprobs_tensors")) - .transpose()?; - Ok(()) - } -} - impl WireLogprobs { /// Convert semantic per-position logprobs into the Python wire tuple shape. /// @@ -315,16 +283,3 @@ impl WireLogprobs { Ok(Logprobs { positions }) } } - -/// 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 -where - Frame: AsRef<[u8]>, -{ - let first_frame = frames.first().ok_or_else(|| ext_value_decode!("missing output frame"))?; - - let mut outputs: EngineCoreOutputs = decode_msgpack(first_frame.as_ref())?; - outputs.resolve_in_place(frames)?; - Ok(outputs) -} 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 7408b98f50c..6fbc57378a1 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs @@ -3,8 +3,8 @@ use std::collections::BTreeSet; use bytes::Bytes; use rmpv::Value; -use super::{Logprobs, PositionLogprobs, TokenLogprob, decode_engine_core_outputs}; -use crate::protocol::EngineCoreFinishReason; +use super::{Logprobs, PositionLogprobs, TokenLogprob}; +use crate::protocol::output::{EngineCoreFinishReason, decode_engine_core_outputs}; fn encode_value(value: &Value) -> Vec { let mut out = Vec::new(); diff --git a/rust/src/engine-core-client/src/protocol/mod.rs b/rust/src/engine-core-client/src/protocol/mod.rs index d7502615336..d434a4e3e94 100644 --- a/rust/src/engine-core-client/src/protocol/mod.rs +++ b/rust/src/engine-core-client/src/protocol/mod.rs @@ -1,28 +1,11 @@ use std::any::type_name; -use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::io::Cursor; -use bytes::Bytes; use rmpv::Value; use serde::{Deserialize, Serialize}; -use serde_default::DefaultFromSerde; -use serde_repr::{Deserialize_repr, Serialize_repr}; -use serde_tuple::{Deserialize_tuple, Serialize_tuple}; use thiserror_ext::AsReport; use crate::error::{Error, Result}; -use crate::protocol::logprobs::MaybeWireLogprobs; -use crate::protocol::multimodal::MmFeatures; -use crate::protocol::stats::{PrefillStats, SchedulerStats}; -use crate::protocol::utility::UtilityOutput; - -// TODO: This module currently mixes reusable frontend-facing semantic types -// (for example `FinishReason`, `StopReason`, `RequestOutputKind`, and future -// cleaned-up frontend sampling types) with engine-core-specific wire DTOs and -// handshake/control messages. While the Rust frontend is still evolving -// quickly, keep them co-located here for iteration speed. Once the higher-level -// API boundary stabilizes, move the truly reusable semantic types into a -// lower-level common crate and keep the engine transport/wire messages here. /// Dynamic msgpack value used for schema positions that are preserved but not /// yet strongly typed in the early-stage Rust client. @@ -36,499 +19,18 @@ fn is_false(v: &bool) -> bool { !v } -fn default_top_p() -> f32 { - 1.0 -} - -fn default_repetition_penalty() -> f32 { - 1.0 -} - -fn default_temperature() -> f32 { - 1.0 -} - -fn default_max_tokens() -> u32 { - 16 -} - -mod classified_outputs; pub mod dtype; pub mod handshake; pub mod logprobs; pub mod lora; pub mod multimodal; +pub mod output; +pub mod request; +pub mod sampling; pub mod stats; +pub mod structured_outputs; pub mod tensor; pub mod utility; -pub use classified_outputs::{ - ClassifiedEngineCoreOutputs, DpControlMessage, RequestBatchOutputs, UtilityCallOutput, -}; -pub use dtype::ModelDtype; -pub use logprobs::decode_engine_core_outputs; - -/// Request types are encoded as single-byte protocol constants so they can be -/// sent over the ZMQ socket without an extra encoding step. -/// -/// Original Python definition: -/// -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -#[repr(u8)] -pub enum EngineCoreRequestType { - Add = 0, - Abort = 1, - StartDpWave = 2, - Utility = 3, -} - -impl EngineCoreRequestType { - /// Decode the single-byte request type frame used on the engine input - /// socket. Returns `None` for unrecognized values. - pub fn from_frame(frame: &[u8]) -> Option { - let [value] = frame else { - return None; - }; - - match value { - 0 => Some(Self::Add), - 1 => Some(Self::Abort), - 2 => Some(Self::StartDpWave), - 3 => Some(Self::Utility), - _ => None, - } - } - - /// Encode the request type as the single-byte frame used on the engine - /// input socket. - pub fn to_frame(self) -> Bytes { - Bytes::from_static(match self { - Self::Add => b"\x00", - Self::Abort => b"\x01", - Self::StartDpWave => b"\x02", - Self::Utility => b"\x03", - }) - } -} - -/// Reason a request finished: stop, length, abort, error, or repetition. -/// -/// This mirrors the Python enum and uses integer encoding for compact wire -/// representation. -/// -/// Original Python definition: -/// -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)] -#[repr(u8)] -pub enum EngineCoreFinishReason { - /// A stop string was emitted. - Stop = 0, - /// `max_tokens` or `max_model_len` was reached. - Length = 1, - /// The request was aborted by the client. - Abort = 2, - /// A retryable request-level internal error occurred. - Error = 3, - /// A repetitive token pattern was detected. - Repetition = 4, -} - -/// Event types emitted by engine-core for one request. -/// -/// Original Python definition: -/// -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)] -#[repr(u8)] -pub enum EngineCoreEventType { - Queued = 1, - Scheduled = 2, - Preempted = 3, -} - -/// A timestamped engine-core event associated with one request. -/// -/// Original Python definition: -/// -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct EngineCoreEvent { - pub r#type: EngineCoreEventType, - pub timestamp: f64, -} - -/// Controls how intermediate outputs are returned to the frontend. -/// -/// `Cumulative = 0` is intentionally not supported in Rust frontend. -/// -/// Original Python definition: -/// -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize_repr, Deserialize_repr)] -#[repr(u8)] -pub enum RequestOutputKind { - /// Return only token deltas in each update. - #[default] - Delta = 1, - /// Suppress intermediate updates and return only the final output. - FinalOnly = 2, -} - -/// Structured-output backend selected for EngineCore grammar compilation. -/// -/// Python vLLM stores this in `StructuredOutputsParams._backend` after request -/// validation. The Rust frontend currently always lowers structured-output -/// requests to guidance, while ignoring any user-supplied `_backend` value. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "kebab-case")] -pub enum StructuredOutputBackend { - Xgrammar, - #[default] - Guidance, - Outlines, - LmFormatEnforcer, -} - -/// The stop reason associated with a finished output. -/// -/// Python models this as the union-typed `stop_reason: int | str | None` -/// field on `EngineCoreOutput`; the Rust client narrows it into a tagged enum. -/// -/// Original Python field: -/// -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(untagged)] -pub enum StopReason { - TokenId(u32), - Text(String), -} - -/// 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 { - /// JSON schema (as a dict/object or JSON string) constraining the output. - pub json: Option, - /// Regular expression the output must match. - pub regex: Option, - /// List of allowed output strings (the model must produce one of these). - pub choice: Option>, - /// 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, - /// 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, - /// 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. - #[serde( - default, - rename = "_backend", - deserialize_with = "serde_with::rust::deserialize_ignore_any" - )] - pub backend: StructuredOutputBackend, -} - -/// Engine-core-facing sampling parameters for text generation. -/// -/// This is the normalized southbound subset used by the Rust frontend when it -/// talks to Python engine-core over the wire. User-facing request semantics -/// such as `stop` strings, `n`, `ignore_eos`, and output aggregation mode are -/// intentionally handled by higher layers before values reach this DTO. -/// -/// Original Python definition: -/// -// Python's SamplingParams is `omit_defaults=True`, so msgpack drops -// default-valued keys; default the whole struct. Per-field fns cover the -// non-zero defaults. -#[serde_with::skip_serializing_none] -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, DefaultFromSerde)] -#[serde(default)] -pub struct EngineCoreSamplingParams { - /// Controls randomness. Lower values are more deterministic; zero means - /// greedy sampling. - #[serde(default = "default_temperature")] - pub temperature: f32, - /// Cumulative probability threshold for nucleus sampling. - #[serde(default = "default_top_p")] - pub top_p: f32, - /// Maximum number of top tokens to consider. `0` means all tokens. - pub top_k: u32, - /// Random seed used by the sampler when present. - pub seed: Option, - /// Maximum number of tokens to generate per output sequence. - #[serde(default = "default_max_tokens")] - pub max_tokens: u32, - /// Minimum number of tokens to generate before EOS or stop-token handling. - pub min_tokens: u32, - /// Maximum number of reasoning ("thinking") tokens to emit before the - /// reasoning section is force-closed. `None` means unlimited; the - /// user-facing `-1` sentinel is normalized to `None` by the frontend before - /// reaching this DTO, so only non-negative values are sent. Enforced - /// engine-side (and only when a reasoning parser is configured). - pub thinking_token_budget: Option, - /// Number of log probabilities to return per generated token. - /// - /// `None` disables sample logprobs. `-1` requests the full vocabulary. - pub logprobs: Option, - /// Number of log probabilities to return per prompt token. - /// - /// `None` disables prompt logprobs. `-1` requests the full vocabulary. - pub prompt_logprobs: Option, - /// Minimum probability threshold for token sampling. - pub min_p: f32, - /// Frequency penalty applied by the sampler. - pub frequency_penalty: f32, - /// Presence penalty applied by the sampler. - pub presence_penalty: f32, - /// Repetition penalty applied by the sampler. - #[serde(default = "default_repetition_penalty")] - pub repetition_penalty: f32, - /// Token IDs that stop generation. - pub stop_token_ids: Vec, - /// Primary EOS token ID used by engine-core's dedicated EOS stop path. - /// - /// This mirrors Python's internal `_eos_token_id` field and is derived by - /// the frontend from tokenizer/model metadata rather than supplied directly - /// by end users. - #[serde(rename = "_eos_token_id")] - pub eos_token_id: Option, - /// Complete stop-token set used by engine-core for `min_tokens` masking. - /// - /// This mirrors Python's internal `_all_stop_token_ids` field and should - /// contain explicit `stop_token_ids` plus any frontend-derived EOS token - /// IDs. - #[serde(rename = "_all_stop_token_ids")] - pub all_stop_token_ids: BTreeSet, - /// Logit biases to apply during sampling. - /// Keys are token IDs - pub logit_bias: Option>, - /// Restrict output to these token IDs only. - pub allowed_token_ids: Option>, - /// Tokenized bad words to avoid during generation. - #[serde(rename = "_bad_words_token_ids")] - pub bad_words_token_ids: Option>>, - /// Parameters for configuring structured outputs (guided decoding). - pub structured_outputs: Option, - /// Specific token IDs for which log probabilities should be returned at - /// each position. - /// - /// When set, the engine returns logprobs for exactly these tokens in - /// addition to the sampled/scored token. Mutually exclusive with the - /// `logprobs` count field in practice. - pub logprob_token_ids: Option>, - /// If `Some(true)`, the request will not attempt to read from the prefix - /// cache; newly computed blocks may still populate the cache. `None` - /// defers to engine-core defaults. - pub skip_reading_prefix_cache: Option, - /// Additional request parameters for custom extensions (from `vllm_xargs`). - pub extra_args: Option>, -} - -impl EngineCoreSamplingParams { - /// Constructs a default sampling params for testing purposes only. - pub fn for_test() -> Self { - Self { - temperature: 1.0, - top_p: 1.0, - top_k: 0, - seed: None, - max_tokens: 65536, - min_tokens: 0, - thinking_token_budget: None, - logprobs: None, - prompt_logprobs: None, - min_p: 0.0, - frequency_penalty: 0.0, - presence_penalty: 0.0, - repetition_penalty: 1.0, - stop_token_ids: Vec::new(), - eos_token_id: None, - all_stop_token_ids: BTreeSet::new(), - logit_bias: None, - allowed_token_ids: None, - bad_words_token_ids: None, - structured_outputs: None, - logprob_token_ids: None, - skip_reading_prefix_cache: None, - extra_args: None, - } - } -} - -/// Extra kwargs consumed by engine-side reasoning parsers. -/// -/// Original Python construction point: -/// -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct ReasoningParserKwargs { - /// Effective kwargs visible to the chat template for this request. - pub chat_template_kwargs: HashMap, -} - -/// Engine-core add-request payload sent from frontend to engine. -/// -/// Original Python definition: -/// -#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)] -pub struct EngineCoreRequest { - pub request_id: String, - pub prompt_token_ids: Option>, - /// Multimodal features attached to the request. - pub mm_features: Option, - pub sampling_params: Option, - /// Pooling parameters are preserved in the schema but not yet strongly - /// typed. - pub pooling_params: Option, - pub arrival_time: f64, - #[serde(default)] - pub lora_request: Option, - #[serde(default)] - pub cache_salt: Option, - #[serde(default)] - pub data_parallel_rank: Option, - /// Unsupported in the first-stage Rust client because Python uses a custom - /// tensor/aux-frame encoding path for this field. - #[serde(default)] - pub prompt_embeds: Option, - /// Per-position mask for mixed-mode inputs (e.g. chat completion with - /// `prompt_embeds` content parts). `Some(true)` means real token id; - /// `Some(false)` means the position uses a pre-computed entry from - /// `prompt_embeds`. `None` for pure-tokens and pure-embeds requests. - #[serde(default)] - pub prompt_is_token_ids: Option>, - /// Index of the client, used to ensure outputs are sent back to the same - /// client when scaling out the frontend. - #[serde(default)] - pub client_index: u32, - /// In DP mode, indicates which wave this request is expected to belong to. - #[serde(default)] - pub current_wave: u32, - #[serde(default)] - pub priority: i32, - #[serde(default)] - pub trace_headers: Option>, - #[serde(default)] - pub resumable: bool, - /// Original user-provided request ID, used for output reporting and aborts. - #[serde(default)] - pub external_req_id: Option, - #[serde(default)] - pub reasoning_ended: Option, - /// Reasoning-parser kwargs forwarded from the frontend to the - /// structured-output backend. - #[serde(default)] - pub reasoning_parser_kwargs: Option, - /// If `true`, the request should be added to the scheduler's waiting queue - /// and immediately aborted, so connector-side cleanup runs via the - /// standard `request_finished` hook. - #[serde(default)] - pub abort_immediately: bool, -} - -impl EngineCoreRequest { - /// Validate fields intentionally not supported in the first-stage client. - pub fn validate(&self) -> Result<()> { - if self.prompt_embeds.is_some() { - return Err(Error::UnsupportedField { - context: "EngineCoreRequest", - field: "prompt_embeds", - }); - } - Ok(()) - } -} - -/// Engine-core output for a single request. -/// -/// Original Python definition: -/// -#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)] -pub struct EngineCoreOutput { - pub request_id: String, - pub new_token_ids: Vec, - /// Decoded sample logprobs for the newly generated positions in this - /// output. - #[serde(default)] - pub new_logprobs: Option, - /// Decoded prompt logprobs for the scored prompt positions emitted in this - /// output. - #[serde(default)] - pub new_prompt_logprobs_tensors: Option, - #[serde(default)] - pub pooling_output: Option, - #[serde(default)] - pub finish_reason: Option, - #[serde(default)] - pub stop_reason: Option, - #[serde(default)] - pub events: Option>, - #[serde(default)] - pub kv_transfer_params: Option, - #[serde(default)] - pub trace_headers: Option, - /// Breakdown of the scheduled prefill computation, set on the first output - /// of a newly scheduled prefill and elided for subsequent decode outputs. - #[serde(default)] - pub prefill_stats: Option, - #[serde(default)] - pub routed_experts: Option, - /// Number of NaNs seen in logits. Values above zero indicate corruption. - #[serde(default)] - pub num_nans_in_logits: u32, -} - -impl EngineCoreOutput { - /// Returns whether this output is terminal for the request. - pub fn finished(&self) -> bool { - self.finish_reason.is_some() - } -} - -/// Batch of engine-core outputs returned to a frontend client. -/// -/// Original Python definition: -/// -#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)] -pub struct EngineCoreOutputs { - #[serde(default)] - pub engine_index: u32, - /// Outputs grouped for this client in the current engine tick. - #[serde(default)] - pub outputs: Vec, - #[serde(default)] - pub scheduler_stats: Option>, - #[serde(default)] - pub timestamp: f64, - #[serde(default)] - pub utility_output: Option, - #[serde(default)] - pub finished_requests: Option>, - /// In DP mode, signals that the current wave finished and engines are - /// paused. - #[serde(default)] - pub 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, -} /// Encode a Rust value into msgpack using the protocol crate's serde model. pub fn encode_msgpack(value: &T) -> Result> @@ -564,81 +66,17 @@ where }) } +/// Decode a msgpack payload into a dynamic value for diagnostics and tests. pub fn decode_value(bytes: &[u8]) -> Result { Ok(rmpv::decode::read_value(&mut Cursor::new(bytes))?) } #[cfg(test)] mod tests { - use std::collections::BTreeSet; + use std::collections::BTreeMap; use super::*; - #[test] - fn engine_core_request_serializes_as_full_array() { - let request = EngineCoreRequest { - request_id: "req-1".to_string(), - prompt_token_ids: Some(vec![1, 2, 3]), - sampling_params: Some(EngineCoreSamplingParams { - max_tokens: 8, - ..EngineCoreSamplingParams::for_test() - }), - arrival_time: 1234.5, - client_index: 7, - ..EngineCoreRequest::default() - }; - - let encoded = encode_msgpack(&request).unwrap(); - let value = decode_value(&encoded).unwrap(); - let array = match value { - Value::Array(array) => array, - other => panic!("expected array, got {other:?}"), - }; - - assert_eq!(array.len(), 20); - assert_eq!(array[0], Value::from("req-1")); - assert_eq!(array[2], Value::Nil); - assert_eq!(array[4], Value::Nil); - assert_eq!(array[10], Value::Nil); - assert_eq!(array[11], Value::from(7)); - } - - #[test] - fn engine_core_outputs_roundtrip_finished_fields() { - let outputs = EngineCoreOutputs { - outputs: vec![EngineCoreOutput { - request_id: "req-1".to_string(), - new_token_ids: vec![42], - new_logprobs: None, - new_prompt_logprobs_tensors: None, - pooling_output: None, - finish_reason: Some(EngineCoreFinishReason::Length), - stop_reason: Some(StopReason::Text("stop".to_string())), - events: None, - kv_transfer_params: None, - trace_headers: None, - prefill_stats: None, - routed_experts: None, - num_nans_in_logits: 0, - }], - finished_requests: Some(BTreeSet::from(["req-1".to_string()])), - ..Default::default() - }; - - let encoded = encode_msgpack(&outputs).unwrap(); - let decoded: EngineCoreOutputs = decode_msgpack(&encoded).unwrap(); - - assert_eq!(decoded.outputs.len(), 1); - assert_eq!( - decoded.outputs[0].finish_reason, - Some(EngineCoreFinishReason::Length) - ); - assert_eq!( - decoded.finished_requests, - Some(BTreeSet::from(["req-1".to_string()])) - ); - } - #[test] fn decode_msgpack_includes_type_name_and_value_fallback() { let error = decode_msgpack::( @@ -648,72 +86,4 @@ mod tests { expect_test::expect![[r#"messagepack decode failed for u64: wrong msgpack marker FixMap(1); value fallback: {"status": "READY"}"#]].assert_eq(&error.to_report_string()); } - - #[test] - fn structured_outputs_backend_ignores_deserialized_value() { - let params: StructuredOutputsParams = serde_json::from_value(serde_json::json!({ - "json_object": true, - "_backend": "xgrammar", - })) - .unwrap(); - - assert_eq!(params.backend, StructuredOutputBackend::Guidance); - - let value = serde_json::to_value(params).unwrap(); - assert_eq!(value["_backend"], "guidance"); - } - - /// A real `sampling_params` is a sparse `omit_defaults` map; absent fields - /// must fall back to defaults. `python_compat` can't catch this since Rust - /// encodes full maps (see `engine_core_request_serializes_as_full_array`). - #[test] - fn decodes_sampling_params_with_omitted_defaults() { - let sampling_params = Value::Map(vec![ - ( - Value::from("stop_token_ids"), - Value::Array(vec![Value::from(151643u32)]), - ), - (Value::from("skip_reading_prefix_cache"), Value::from(false)), - ]); - let request = Value::Array(vec![ - Value::from("req-omit-defaults"), - Value::Array(vec![ - Value::from(1u32), - Value::from(2u32), - Value::from(3u32), - ]), - Value::Nil, - sampling_params, - Value::Nil, - Value::from(1.0f64), - ]); - - let mut bytes = Vec::new(); - rmpv::encode::write_value(&mut bytes, &request).unwrap(); - - let decoded: EngineCoreRequest = decode_msgpack(&bytes) - .expect("a real omit_defaults request must decode (regression: missing field)"); - - assert_eq!(decoded.request_id, "req-omit-defaults"); - let sampling = decoded.sampling_params.expect("sampling params present"); - - assert_eq!(sampling.stop_token_ids, vec![151643]); - assert_eq!(sampling.skip_reading_prefix_cache, Some(false)); - - // Omitted fields -> Python defaults. - assert_eq!(sampling.temperature, 1.0); - assert_eq!(sampling.top_p, 1.0); - assert_eq!(sampling.top_k, 0); - assert_eq!(sampling.seed, None); - assert_eq!(sampling.max_tokens, 16); - assert_eq!(sampling.min_tokens, 0); - assert_eq!(sampling.min_p, 0.0); - assert_eq!(sampling.frequency_penalty, 0.0); - assert_eq!(sampling.presence_penalty, 0.0); - assert_eq!(sampling.repetition_penalty, 1.0); - assert_eq!(sampling.logprobs, None); - assert_eq!(sampling.prompt_logprobs, None); - assert_eq!(sampling.eos_token_id, None); - assert!(sampling.all_stop_token_ids.is_empty()); - } } diff --git a/rust/src/engine-core-client/src/protocol/classified_outputs.rs b/rust/src/engine-core-client/src/protocol/output.rs similarity index 50% rename from rust/src/engine-core-client/src/protocol/classified_outputs.rs rename to rust/src/engine-core-client/src/protocol/output.rs index d572f8f925b..c82ecfa2b8b 100644 --- a/rust/src/engine-core-client/src/protocol/classified_outputs.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -1,10 +1,164 @@ use std::collections::BTreeSet; use enum_as_inner::EnumAsInner; +use serde::{Deserialize, Serialize}; +use serde_default::DefaultFromSerde; +use serde_repr::{Deserialize_repr, Serialize_repr}; +use serde_tuple::{Deserialize_tuple, Serialize_tuple}; use super::utility::UtilityOutput; -use super::{EngineCoreOutput, EngineCoreOutputs}; -use crate::protocol::stats::SchedulerStats; +use crate::error::{Error, Result, ext_value_decode}; +use crate::protocol::logprobs::MaybeWireLogprobs; +use crate::protocol::stats::{PrefillStats, SchedulerStats}; +use crate::protocol::{OpaqueValue, decode_msgpack}; + +/// The stop reason associated with a finished output. +/// +/// Python models this as the union-typed `stop_reason: int | str | None` +/// field on `EngineCoreOutput`; the Rust client narrows it into a tagged enum. +/// +/// Original Python field: +/// +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum StopReason { + TokenId(u32), + Text(String), +} + +/// Reason a request finished: stop, length, abort, error, or repetition. +/// +/// This mirrors the Python enum and uses integer encoding for compact wire +/// representation. +/// +/// Original Python definition: +/// +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)] +#[repr(u8)] +pub enum EngineCoreFinishReason { + /// A stop string was emitted. + Stop = 0, + /// `max_tokens` or `max_model_len` was reached. + Length = 1, + /// The request was aborted by the client. + Abort = 2, + /// A retryable request-level internal error occurred. + Error = 3, + /// A repetitive token pattern was detected. + Repetition = 4, +} + +/// Event types emitted by engine-core for one request. +/// +/// Original Python definition: +/// +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)] +#[repr(u8)] +pub enum EngineCoreEventType { + Queued = 1, + Scheduled = 2, + Preempted = 3, +} + +/// A timestamped engine-core event associated with one request. +/// +/// Original Python definition: +/// +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct EngineCoreEvent { + pub r#type: EngineCoreEventType, + pub timestamp: f64, +} + +/// Engine-core output for a single request. +/// +/// Original Python definition: +/// +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)] +pub struct EngineCoreOutput { + pub request_id: String, + pub new_token_ids: Vec, + /// Decoded sample logprobs for the newly generated positions in this + /// output. + #[serde(default)] + pub new_logprobs: Option, + /// Decoded prompt logprobs for the scored prompt positions emitted in this + /// output. + #[serde(default)] + pub new_prompt_logprobs_tensors: Option, + #[serde(default)] + pub pooling_output: Option, + #[serde(default)] + pub finish_reason: Option, + #[serde(default)] + pub stop_reason: Option, + #[serde(default)] + pub events: Option>, + #[serde(default)] + pub kv_transfer_params: Option, + #[serde(default)] + pub trace_headers: Option, + /// Breakdown of the scheduled prefill computation, set on the first output + /// of a newly scheduled prefill and elided for subsequent decode outputs. + #[serde(default)] + pub prefill_stats: Option, + #[serde(default)] + pub routed_experts: Option, + /// Number of NaNs seen in logits. Values above zero indicate corruption. + #[serde(default)] + pub num_nans_in_logits: u32, +} + +impl EngineCoreOutput { + /// Returns whether this output is terminal for the request. + pub fn finished(&self) -> bool { + self.finish_reason.is_some() + } + + /// Resolve all wire-format fields in-place by looking up aux frames and + /// decoding raw-view payloads as needed. + fn resolve_in_place(&mut self, frames: &[Frame]) -> Result<()> + where + Frame: AsRef<[u8]>, + { + self.new_logprobs = (self.new_logprobs.take()) + .map(|value| value.resolve(frames, "new_logprobs")) + .transpose()?; + self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take()) + .map(|value| value.resolve(frames, "new_prompt_logprobs_tensors")) + .transpose()?; + Ok(()) + } +} + +/// Batch of engine-core outputs returned to a frontend client. +/// +/// Original Python definition: +/// +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)] +pub struct EngineCoreOutputs { + #[serde(default)] + pub engine_index: u32, + /// Outputs grouped for this client in the current engine tick. + #[serde(default)] + pub outputs: Vec, + #[serde(default)] + pub scheduler_stats: Option>, + #[serde(default)] + pub timestamp: f64, + #[serde(default)] + pub utility_output: Option, + #[serde(default)] + pub finished_requests: Option>, + /// In DP mode, signals that the current wave finished and engines are + /// paused. + #[serde(default)] + pub 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, +} /// Data-parallel control notifications multiplexed through `EngineCoreOutputs`. #[derive(Debug, Clone, PartialEq, Eq)] @@ -49,6 +203,18 @@ pub enum ClassifiedEngineCoreOutputs { } impl EngineCoreOutputs { + /// Resolve all wire-format fields in-place by looking up aux frames and + /// decoding raw-view payloads as needed. + fn resolve_in_place(&mut self, frames: &[Frame]) -> Result<()> + where + Frame: AsRef<[u8]>, + { + for output in &mut self.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() @@ -92,12 +258,62 @@ impl EngineCoreOutputs { } } +/// 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 +where + Frame: AsRef<[u8]>, +{ + let first_frame = frames.first().ok_or_else(|| ext_value_decode!("missing output frame"))?; + + let mut outputs: EngineCoreOutputs = decode_msgpack(first_frame.as_ref())?; + outputs.resolve_in_place(frames)?; + Ok(outputs) +} + #[cfg(test)] mod tests { use std::collections::BTreeSet; use super::*; - use crate::protocol::EngineCoreOutput; + use crate::protocol::output::EngineCoreOutput; + use crate::protocol::{decode_msgpack, encode_msgpack}; + + #[test] + fn engine_core_outputs_roundtrip_finished_fields() { + let outputs = EngineCoreOutputs { + outputs: vec![EngineCoreOutput { + request_id: "req-1".to_string(), + new_token_ids: vec![42], + new_logprobs: None, + new_prompt_logprobs_tensors: None, + pooling_output: None, + finish_reason: Some(EngineCoreFinishReason::Length), + stop_reason: Some(StopReason::Text("stop".to_string())), + events: None, + kv_transfer_params: None, + trace_headers: None, + prefill_stats: None, + routed_experts: None, + num_nans_in_logits: 0, + }], + finished_requests: Some(BTreeSet::from(["req-1".to_string()])), + ..Default::default() + }; + + let encoded = encode_msgpack(&outputs).unwrap(); + let decoded: EngineCoreOutputs = decode_msgpack(&encoded).unwrap(); + + assert_eq!(decoded.outputs.len(), 1); + assert_eq!( + decoded.outputs[0].finish_reason, + Some(EngineCoreFinishReason::Length) + ); + assert_eq!( + decoded.finished_requests, + Some(BTreeSet::from(["req-1".to_string()])) + ); + } #[test] fn engine_core_outputs_classify_request_batch() { diff --git a/rust/src/engine-core-client/src/protocol/request.rs b/rust/src/engine-core-client/src/protocol/request.rs new file mode 100644 index 00000000000..b7993a3f7c0 --- /dev/null +++ b/rust/src/engine-core-client/src/protocol/request.rs @@ -0,0 +1,175 @@ +use std::collections::{BTreeMap, HashMap}; + +use bytes::Bytes; +use serde::{Deserialize, Serialize}; +use serde_default::DefaultFromSerde; +use serde_tuple::{Deserialize_tuple, Serialize_tuple}; + +use crate::protocol::multimodal::MmFeatures; +use crate::protocol::sampling::EngineCoreSamplingParams; +use crate::protocol::{OpaqueValue, lora}; +use crate::{Error, Result}; + +/// Request types are encoded as single-byte protocol constants so they can be +/// sent over the ZMQ socket without an extra encoding step. +/// +/// Original Python definition: +/// +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +pub enum EngineCoreRequestType { + Add = 0, + Abort = 1, + StartDpWave = 2, + Utility = 3, +} + +impl EngineCoreRequestType { + /// Decode the single-byte request type frame used on the engine input + /// socket. Returns `None` for unrecognized values. + pub fn from_frame(frame: &[u8]) -> Option { + let [value] = frame else { + return None; + }; + + match value { + 0 => Some(Self::Add), + 1 => Some(Self::Abort), + 2 => Some(Self::StartDpWave), + 3 => Some(Self::Utility), + _ => None, + } + } + + /// Encode the request type as the single-byte frame used on the engine + /// input socket. + pub fn to_frame(self) -> Bytes { + Bytes::from_static(match self { + Self::Add => b"\x00", + Self::Abort => b"\x01", + Self::StartDpWave => b"\x02", + Self::Utility => b"\x03", + }) + } +} + +/// Extra kwargs consumed by engine-side reasoning parsers. +/// +/// Original Python construction point: +/// +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ReasoningParserKwargs { + /// Effective kwargs visible to the chat template for this request. + pub chat_template_kwargs: HashMap, +} + +/// Engine-core add-request payload sent from frontend to engine. +/// +/// Original Python definition: +/// +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)] +pub struct EngineCoreRequest { + pub request_id: String, + pub prompt_token_ids: Option>, + /// Multimodal features attached to the request. + pub mm_features: Option, + pub sampling_params: Option, + /// Pooling parameters are preserved in the schema but not yet strongly + /// typed. + pub pooling_params: Option, + pub arrival_time: f64, + #[serde(default)] + pub lora_request: Option, + #[serde(default)] + pub cache_salt: Option, + #[serde(default)] + pub data_parallel_rank: Option, + /// Unsupported in the first-stage Rust client because Python uses a custom + /// tensor/aux-frame encoding path for this field. + #[serde(default)] + pub prompt_embeds: Option, + /// Per-position mask for mixed-mode inputs (e.g. chat completion with + /// `prompt_embeds` content parts). `Some(true)` means real token id; + /// `Some(false)` means the position uses a pre-computed entry from + /// `prompt_embeds`. `None` for pure-tokens and pure-embeds requests. + #[serde(default)] + pub prompt_is_token_ids: Option>, + /// Index of the client, used to ensure outputs are sent back to the same + /// client when scaling out the frontend. + #[serde(default)] + pub client_index: u32, + /// In DP mode, indicates which wave this request is expected to belong to. + #[serde(default)] + pub current_wave: u32, + #[serde(default)] + pub priority: i32, + #[serde(default)] + pub trace_headers: Option>, + #[serde(default)] + pub resumable: bool, + /// Original user-provided request ID, used for output reporting and aborts. + #[serde(default)] + pub external_req_id: Option, + #[serde(default)] + pub reasoning_ended: Option, + /// Reasoning-parser kwargs forwarded from the frontend to the + /// structured-output backend. + #[serde(default)] + pub reasoning_parser_kwargs: Option, + /// If `true`, the request should be added to the scheduler's waiting queue + /// and immediately aborted, so connector-side cleanup runs via the + /// standard `request_finished` hook. + #[serde(default)] + pub abort_immediately: bool, +} + +impl EngineCoreRequest { + /// Validate fields intentionally not supported in the first-stage client. + pub fn validate(&self) -> Result<()> { + if self.prompt_embeds.is_some() { + return Err(Error::UnsupportedField { + context: "EngineCoreRequest", + field: "prompt_embeds", + }); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use rmpv::Value; + + use super::*; + use crate::protocol::sampling::EngineCoreSamplingParams; + use crate::protocol::{decode_value, encode_msgpack}; + + #[test] + fn engine_core_request_serializes_as_full_array() { + let request = EngineCoreRequest { + request_id: "req-1".to_string(), + prompt_token_ids: Some(vec![1, 2, 3]), + sampling_params: Some(EngineCoreSamplingParams { + max_tokens: 8, + ..EngineCoreSamplingParams::for_test() + }), + arrival_time: 1234.5, + client_index: 7, + ..EngineCoreRequest::default() + }; + + let encoded = encode_msgpack(&request).unwrap(); + let value = decode_value(&encoded).unwrap(); + let array = match value { + Value::Array(array) => array, + other => panic!("expected array, got {other:?}"), + }; + + assert_eq!(array.len(), 20); + assert_eq!(array[0], Value::from("req-1")); + assert_eq!(array[2], Value::Nil); + assert_eq!(array[4], Value::Nil); + assert_eq!(array[10], Value::Nil); + assert_eq!(array[11], Value::from(7)); + } +} diff --git a/rust/src/engine-core-client/src/protocol/sampling.rs b/rust/src/engine-core-client/src/protocol/sampling.rs new file mode 100644 index 00000000000..b724b36dc53 --- /dev/null +++ b/rust/src/engine-core-client/src/protocol/sampling.rs @@ -0,0 +1,211 @@ +use std::collections::{BTreeSet, HashMap}; + +use serde::{Deserialize, Serialize}; +use serde_default::DefaultFromSerde; + +use crate::protocol::structured_outputs::StructuredOutputsParams; + +fn default_top_p() -> f32 { + 1.0 +} + +fn default_repetition_penalty() -> f32 { + 1.0 +} + +fn default_temperature() -> f32 { + 1.0 +} + +fn default_max_tokens() -> u32 { + 16 +} + +/// Engine-core-facing sampling parameters for text generation. +/// +/// This is the normalized southbound subset used by the Rust frontend when it +/// talks to Python engine-core over the wire. User-facing request semantics +/// such as `stop` strings, `n`, `ignore_eos`, and output aggregation mode are +/// intentionally handled by higher layers before values reach this DTO. +/// +/// Original Python definition: +/// +// Python's SamplingParams is `omit_defaults=True`, so msgpack drops +// default-valued keys; default the whole struct. Per-field fns cover the +// non-zero defaults. +#[serde_with::skip_serializing_none] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, DefaultFromSerde)] +#[serde(default)] +pub struct EngineCoreSamplingParams { + /// Controls randomness. Lower values are more deterministic; zero means + /// greedy sampling. + #[serde(default = "default_temperature")] + pub temperature: f32, + /// Cumulative probability threshold for nucleus sampling. + #[serde(default = "default_top_p")] + pub top_p: f32, + /// Maximum number of top tokens to consider. `0` means all tokens. + pub top_k: u32, + /// Random seed used by the sampler when present. + pub seed: Option, + /// Maximum number of tokens to generate per output sequence. + #[serde(default = "default_max_tokens")] + pub max_tokens: u32, + /// Minimum number of tokens to generate before EOS or stop-token handling. + pub min_tokens: u32, + /// Maximum number of reasoning ("thinking") tokens to emit before the + /// reasoning section is force-closed. `None` means unlimited; the + /// user-facing `-1` sentinel is normalized to `None` by the frontend before + /// reaching this DTO, so only non-negative values are sent. Enforced + /// engine-side (and only when a reasoning parser is configured). + pub thinking_token_budget: Option, + /// Number of log probabilities to return per generated token. + /// + /// `None` disables sample logprobs. `-1` requests the full vocabulary. + pub logprobs: Option, + /// Number of log probabilities to return per prompt token. + /// + /// `None` disables prompt logprobs. `-1` requests the full vocabulary. + pub prompt_logprobs: Option, + /// Minimum probability threshold for token sampling. + pub min_p: f32, + /// Frequency penalty applied by the sampler. + pub frequency_penalty: f32, + /// Presence penalty applied by the sampler. + pub presence_penalty: f32, + /// Repetition penalty applied by the sampler. + #[serde(default = "default_repetition_penalty")] + pub repetition_penalty: f32, + /// Token IDs that stop generation. + pub stop_token_ids: Vec, + /// Primary EOS token ID used by engine-core's dedicated EOS stop path. + /// + /// This mirrors Python's internal `_eos_token_id` field and is derived by + /// the frontend from tokenizer/model metadata rather than supplied directly + /// by end users. + #[serde(rename = "_eos_token_id")] + pub eos_token_id: Option, + /// Complete stop-token set used by engine-core for `min_tokens` masking. + /// + /// This mirrors Python's internal `_all_stop_token_ids` field and should + /// contain explicit `stop_token_ids` plus any frontend-derived EOS token + /// IDs. + #[serde(rename = "_all_stop_token_ids")] + pub all_stop_token_ids: BTreeSet, + /// Logit biases to apply during sampling. + /// Keys are token IDs + pub logit_bias: Option>, + /// Restrict output to these token IDs only. + pub allowed_token_ids: Option>, + /// Tokenized bad words to avoid during generation. + #[serde(rename = "_bad_words_token_ids")] + pub bad_words_token_ids: Option>>, + /// Parameters for configuring structured outputs (guided decoding). + pub structured_outputs: Option, + /// Specific token IDs for which log probabilities should be returned at + /// each position. + /// + /// When set, the engine returns logprobs for exactly these tokens in + /// addition to the sampled/scored token. Mutually exclusive with the + /// `logprobs` count field in practice. + pub logprob_token_ids: Option>, + /// If `Some(true)`, the request will not attempt to read from the prefix + /// cache; newly computed blocks may still populate the cache. `None` + /// defers to engine-core defaults. + pub skip_reading_prefix_cache: Option, + /// Additional request parameters for custom extensions (from `vllm_xargs`). + pub extra_args: Option>, +} + +impl EngineCoreSamplingParams { + /// Constructs a default sampling params for testing purposes only. + pub fn for_test() -> Self { + Self { + temperature: 1.0, + top_p: 1.0, + top_k: 0, + seed: None, + max_tokens: 65536, + min_tokens: 0, + thinking_token_budget: None, + logprobs: None, + prompt_logprobs: None, + min_p: 0.0, + frequency_penalty: 0.0, + presence_penalty: 0.0, + repetition_penalty: 1.0, + stop_token_ids: Vec::new(), + eos_token_id: None, + all_stop_token_ids: BTreeSet::new(), + logit_bias: None, + allowed_token_ids: None, + bad_words_token_ids: None, + structured_outputs: None, + logprob_token_ids: None, + skip_reading_prefix_cache: None, + extra_args: None, + } + } +} + +#[cfg(test)] +mod tests { + use rmpv::Value; + + use crate::protocol::decode_msgpack; + use crate::protocol::request::EngineCoreRequest; + + /// A real `sampling_params` is a sparse `omit_defaults` map; absent fields + /// must fall back to defaults. `python_compat` can't catch this since Rust + /// encodes full maps (see `engine_core_request_serializes_as_full_array`). + #[test] + fn decodes_sampling_params_with_omitted_defaults() { + let sampling_params = Value::Map(vec![ + ( + Value::from("stop_token_ids"), + Value::Array(vec![Value::from(151643u32)]), + ), + (Value::from("skip_reading_prefix_cache"), Value::from(false)), + ]); + let request = Value::Array(vec![ + Value::from("req-omit-defaults"), + Value::Array(vec![ + Value::from(1u32), + Value::from(2u32), + Value::from(3u32), + ]), + Value::Nil, + sampling_params, + Value::Nil, + Value::from(1.0f64), + ]); + + let mut bytes = Vec::new(); + rmpv::encode::write_value(&mut bytes, &request).unwrap(); + + let decoded: EngineCoreRequest = decode_msgpack(&bytes) + .expect("a real omit_defaults request must decode (regression: missing field)"); + + assert_eq!(decoded.request_id, "req-omit-defaults"); + let sampling = decoded.sampling_params.expect("sampling params present"); + + assert_eq!(sampling.stop_token_ids, vec![151643]); + assert_eq!(sampling.skip_reading_prefix_cache, Some(false)); + + // Omitted fields -> Python defaults. + assert_eq!(sampling.temperature, 1.0); + assert_eq!(sampling.top_p, 1.0); + assert_eq!(sampling.top_k, 0); + assert_eq!(sampling.seed, None); + assert_eq!(sampling.max_tokens, 16); + assert_eq!(sampling.min_tokens, 0); + assert_eq!(sampling.min_p, 0.0); + assert_eq!(sampling.frequency_penalty, 0.0); + assert_eq!(sampling.presence_penalty, 0.0); + assert_eq!(sampling.repetition_penalty, 1.0); + assert_eq!(sampling.logprobs, None); + assert_eq!(sampling.prompt_logprobs, None); + assert_eq!(sampling.eos_token_id, None); + assert!(sampling.all_stop_token_ids.is_empty()); + } +} diff --git a/rust/src/engine-core-client/src/protocol/structured_outputs.rs b/rust/src/engine-core-client/src/protocol/structured_outputs.rs new file mode 100644 index 00000000000..9bf8102aa6b --- /dev/null +++ b/rust/src/engine-core-client/src/protocol/structured_outputs.rs @@ -0,0 +1,81 @@ +use serde::{Deserialize, Serialize}; + +/// Structured-output backend selected for EngineCore grammar compilation. +/// +/// Python vLLM stores this in `StructuredOutputsParams._backend` after request +/// validation. The Rust frontend currently always lowers structured-output +/// requests to guidance, while ignoring any user-supplied `_backend` value. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum StructuredOutputBackend { + Xgrammar, + #[default] + Guidance, + Outlines, + 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 { + /// JSON schema (as a dict/object or JSON string) constraining the output. + pub json: Option, + /// Regular expression the output must match. + pub regex: Option, + /// List of allowed output strings (the model must produce one of these). + pub choice: Option>, + /// 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, + /// 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, + /// 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. + #[serde( + default, + rename = "_backend", + deserialize_with = "serde_with::rust::deserialize_ignore_any" + )] + pub backend: StructuredOutputBackend, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn structured_outputs_backend_ignores_deserialized_value() { + let params: StructuredOutputsParams = serde_json::from_value(serde_json::json!({ + "json_object": true, + "_backend": "xgrammar", + })) + .unwrap(); + + assert_eq!(params.backend, StructuredOutputBackend::Guidance); + + let value = serde_json::to_value(params).unwrap(); + assert_eq!(value["_backend"], "guidance"); + } +} diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 11c403c5637..f60973d29b5 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -22,13 +22,14 @@ use crate::protocol::multimodal::{ MmFeatureSpec, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmSlice, PlaceholderRange, SliceSpec, }; +use crate::protocol::output::{ + EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, decode_engine_core_outputs, +}; +use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType}; +use crate::protocol::sampling::EngineCoreSamplingParams; use crate::protocol::stats::SchedulerStats; use crate::protocol::tensor::WireTensor; use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope}; -use crate::protocol::{ - EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, - EngineCoreRequestType, EngineCoreSamplingParams, decode_engine_core_outputs, -}; use crate::test_utils::{ IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets, setup_mock_engine_with_init, spawn_mock_engine_task, diff --git a/rust/src/engine-core-client/src/transport.rs b/rust/src/engine-core-client/src/transport.rs index d0d9b4efe39..aecf9625000 100644 --- a/rust/src/engine-core-client/src/transport.rs +++ b/rust/src/engine-core-client/src/transport.rs @@ -18,9 +18,8 @@ use crate::error::{Error, Result, bail_unexpected_handshake_message}; use crate::protocol::handshake::{ EngineCoreReadyResponse, HandshakeAddresses, HandshakeInitMessage, ReadyMessage, }; -use crate::protocol::{ - EngineCoreOutputs, decode_engine_core_outputs, decode_msgpack, encode_msgpack, -}; +use crate::protocol::output::{EngineCoreOutputs, decode_engine_core_outputs}; +use crate::protocol::{decode_msgpack, encode_msgpack}; /// Dedicated single-frame sentinel emitted by Python `EngineCoreProc` when the /// engine dies. diff --git a/rust/src/llm/examples/external_engine_smoke.rs b/rust/src/llm/examples/external_engine_smoke.rs index 83a22d7dbd4..e5e153f3347 100644 --- a/rust/src/llm/examples/external_engine_smoke.rs +++ b/rust/src/llm/examples/external_engine_smoke.rs @@ -5,7 +5,7 @@ use clap::Parser; use futures::StreamExt as _; use tokio::time::timeout; use tracing_subscriber::EnvFilter; -use vllm_engine_core_client::protocol::EngineCoreSamplingParams; +use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode}; use vllm_llm::{FinishReason, GenerateOutputStream, GenerateRequest, Llm}; diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 8cfc38d0bc9..7602df5b801 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -8,7 +8,7 @@ use futures::stream::FusedStream; use futures::{Stream, StreamExt as _, pin_mut}; use serde::{Deserialize, Serialize}; use vllm_engine_core_client::protocol::logprobs::Logprobs; -use vllm_engine_core_client::protocol::{EngineCoreFinishReason, StopReason}; +use vllm_engine_core_client::protocol::output::{EngineCoreFinishReason, StopReason}; use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream}; use crate::error::Result; diff --git a/rust/src/llm/src/request.rs b/rust/src/llm/src/request.rs index bbb1d60fc6d..159cf823de4 100644 --- a/rust/src/llm/src/request.rs +++ b/rust/src/llm/src/request.rs @@ -4,9 +4,8 @@ use std::time::{SystemTime, UNIX_EPOCH}; use uuid::Uuid; use vllm_engine_core_client::protocol::lora::LoraRequest; use vllm_engine_core_client::protocol::multimodal::MmFeatures; -use vllm_engine_core_client::protocol::{ - EngineCoreRequest, EngineCoreSamplingParams, ReasoningParserKwargs, -}; +use vllm_engine_core_client::protocol::request::{EngineCoreRequest, ReasoningParserKwargs}; +use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; use crate::error::{Error, Result}; @@ -134,7 +133,8 @@ fn current_unix_timestamp_secs() -> f64 { mod tests { use std::collections::BTreeMap; - use vllm_engine_core_client::protocol::{EngineCoreSamplingParams, ReasoningParserKwargs}; + use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; + use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; use super::GenerateRequest; use crate::error::Error; diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs index 6612fa3cc4f..38795a70928 100644 --- a/rust/src/llm/src/request_metrics.rs +++ b/rust/src/llm/src/request_metrics.rs @@ -1,7 +1,9 @@ use std::time::{SystemTime, UNIX_EPOCH}; +use vllm_engine_core_client::protocol::output::{ + EngineCoreEvent, EngineCoreEventType, EngineCoreOutput, +}; use vllm_engine_core_client::protocol::stats::PrefillStats; -use vllm_engine_core_client::protocol::{EngineCoreEvent, EngineCoreEventType, EngineCoreOutput}; use vllm_metrics::{ EngineLabels, FinishedReasonLabels, METRICS, PromptTokenSourceLabels, RequestMetrics, }; @@ -328,8 +330,8 @@ pub(crate) fn current_unix_timestamp_secs() -> f64 { #[cfg(test)] mod tests { + use vllm_engine_core_client::protocol::output::{EngineCoreEvent, EngineCoreEventType}; use vllm_engine_core_client::protocol::stats::PrefillStats; - use vllm_engine_core_client::protocol::{EngineCoreEvent, EngineCoreEventType}; use super::{RequestMetricsTracker, diff_or_zero}; @@ -341,7 +343,7 @@ mod tests { 2, 10.0, 100.2, - &vllm_engine_core_client::protocol::EngineCoreOutput { + &vllm_engine_core_client::protocol::output::EngineCoreOutput { request_id: "req-1".to_string(), new_token_ids: vec![1], finish_reason: None, @@ -369,7 +371,7 @@ mod tests { 2, 11.5, 100.4, - &vllm_engine_core_client::protocol::EngineCoreOutput { + &vllm_engine_core_client::protocol::output::EngineCoreOutput { request_id: "req-1".to_string(), new_token_ids: vec![2, 3], finish_reason: None, diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 98108334731..b7942910d44 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -9,11 +9,13 @@ use uuid::Uuid; use vllm_engine_core_client::protocol::logprobs::{ Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob, }; -use vllm_engine_core_client::protocol::stats::PrefillStats; -use vllm_engine_core_client::protocol::{ +use vllm_engine_core_client::protocol::output::{ EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput, - EngineCoreOutputs, EngineCoreRequest, EngineCoreSamplingParams, + EngineCoreOutputs, }; +use vllm_engine_core_client::protocol::request::EngineCoreRequest; +use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; +use vllm_engine_core_client::protocol::stats::PrefillStats; use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task}; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig}; use vllm_llm::{ diff --git a/rust/src/mock-engine/src/engine.rs b/rust/src/mock-engine/src/engine.rs index 2aa2f7bb397..42a8f0efa04 100644 --- a/rust/src/mock-engine/src/engine.rs +++ b/rust/src/mock-engine/src/engine.rs @@ -11,12 +11,13 @@ use tokio::sync::mpsc; 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, +}; +use vllm_engine_core_client::protocol::request::EngineCoreRequest; use vllm_engine_core_client::protocol::utility::{ EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope, }; -use vllm_engine_core_client::protocol::{ - EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, -}; use super::Opt; diff --git a/rust/src/mock-engine/src/io.rs b/rust/src/mock-engine/src/io.rs index 28d77639c77..77c0f14b57c 100644 --- a/rust/src/mock-engine/src/io.rs +++ b/rust/src/mock-engine/src/io.rs @@ -4,10 +4,9 @@ use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use tracing::warn; use vllm_engine_core_client::mock_engine::MockEngineDataSockets; +use vllm_engine_core_client::protocol::request::{EngineCoreRequest, EngineCoreRequestType}; use vllm_engine_core_client::protocol::utility::EngineCoreUtilityRequest; -use vllm_engine_core_client::protocol::{ - EngineCoreRequest, EngineCoreRequestType, decode_msgpack, encode_msgpack, -}; +use vllm_engine_core_client::protocol::{decode_msgpack, encode_msgpack}; use zeromq::{DealerSocket, PushSocket, SocketRecv as _, SocketSend as _, ZmqMessage}; use crate::engine::{EngineInput, EngineOutput}; diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs index a80aef40300..71a4e0a6575 100644 --- a/rust/src/mock-engine/src/tests.rs +++ b/rust/src/mock-engine/src/tests.rs @@ -5,9 +5,9 @@ use anyhow::Result; use futures::StreamExt as _; use tokio::time::timeout; use tokio_util::sync::CancellationToken; -use vllm_engine_core_client::protocol::{ - EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams, -}; +use vllm_engine_core_client::protocol::output::EngineCoreFinishReason; +use vllm_engine_core_client::protocol::request::EngineCoreRequest; +use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; use vllm_engine_core_client::test_utils::IpcNamespace; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode}; diff --git a/rust/src/parser/benches/utils/mod.rs b/rust/src/parser/benches/utils/mod.rs index 229f40a3681..bb674131718 100644 --- a/rust/src/parser/benches/utils/mod.rs +++ b/rust/src/parser/benches/utils/mod.rs @@ -5,14 +5,13 @@ mod adapter; +pub(super) use adapter::UnifiedToolParserAdapter; use futures::FutureExt as _; use openai_protocol::common::{Function as OpenAiFunction, Tool as OpenAiTool}; use tool_parser::traits::ToolParser as ExternalToolParser; use vllm_parser::tool::test_utils::collect_stream; use vllm_parser::tool::{Tool, ToolParser}; -pub(super) use adapter::UnifiedToolParserAdapter; - pub(super) fn openai_tools(tools: &[Tool]) -> Vec { tools .iter() diff --git a/rust/src/parser/python/src/lib.rs b/rust/src/parser/python/src/lib.rs index 4567348bcd9..8057930ade0 100644 --- a/rust/src/parser/python/src/lib.rs +++ b/rust/src/parser/python/src/lib.rs @@ -231,9 +231,10 @@ fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> { #[cfg(test)] mod tests { - use super::*; use serde_json::json; + use super::*; + fn with_python(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R { Python::initialize(); Python::attach(f) diff --git a/rust/src/parser/src/reasoning/seed_oss.rs b/rust/src/parser/src/reasoning/seed_oss.rs index 580e95e7957..61dde653779 100644 --- a/rust/src/parser/src/reasoning/seed_oss.rs +++ b/rust/src/parser/src/reasoning/seed_oss.rs @@ -49,10 +49,8 @@ mod tests { use std::sync::Arc; use super::SeedOssReasoningParser; - use crate::reasoning::{ - ReasoningParser, - tests::{SEED_THINK_END_ID, SEED_THINK_START_ID, fake_tokenizer}, - }; + use crate::reasoning::ReasoningParser; + use crate::reasoning::tests::{SEED_THINK_END_ID, SEED_THINK_START_ID, fake_tokenizer}; #[test] fn without_prompt_markers_expects_start_token() { diff --git a/rust/src/parser/src/reasoning/step3p5.rs b/rust/src/parser/src/reasoning/step3p5.rs index 1bcf54fcd54..79677e682f1 100644 --- a/rust/src/parser/src/reasoning/step3p5.rs +++ b/rust/src/parser/src/reasoning/step3p5.rs @@ -127,10 +127,8 @@ mod tests { use std::sync::Arc; use super::Step3p5ReasoningParser; - use crate::reasoning::{ - ReasoningParser, - tests::{THINK_START_ID, fake_tokenizer}, - }; + use crate::reasoning::ReasoningParser; + use crate::reasoning::tests::{THINK_START_ID, fake_tokenizer}; #[test] fn picks_up_prompt_start_boundary() { diff --git a/rust/src/parser/src/tool/mod.rs b/rust/src/parser/src/tool/mod.rs index dd4630b1c6b..5a06f2311ed 100644 --- a/rust/src/parser/src/tool/mod.rs +++ b/rust/src/parser/src/tool/mod.rs @@ -14,8 +14,6 @@ mod parameters; mod qwen_coder; #[cfg(any(test, feature = "test-util"))] pub mod test_utils; -use crate::utils; - use std::collections::{BTreeMap, btree_map}; pub use deepseek_dsml::{DeepSeekV4ToolParser, DeepSeekV32ToolParser}; @@ -35,6 +33,8 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; pub use xgrammar_structural_tag::Model as StructuralTagModel; +use crate::utils; + /// One function-style tool made available to the model. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Tool { diff --git a/rust/src/parser/src/unified/combined.rs b/rust/src/parser/src/unified/combined.rs index f06334c72ee..3b6abb4b932 100644 --- a/rust/src/parser/src/unified/combined.rs +++ b/rust/src/parser/src/unified/combined.rs @@ -2,11 +2,10 @@ use vllm_tokenizer::DynTokenizer; +use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput}; use crate::reasoning::ReasoningParser; use crate::tool::{StructuralTagModel, Tool, ToolParser, ToolParserOutput}; -use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput}; - /// Unified parser that composes existing reasoning and tool parsers. pub struct CombinedParser { reasoning: Option>, diff --git a/rust/src/parser/src/unified/gemma4.rs b/rust/src/parser/src/unified/gemma4.rs index 71541f94c69..51a1de12c2f 100644 --- a/rust/src/parser/src/unified/gemma4.rs +++ b/rust/src/parser/src/unified/gemma4.rs @@ -1,4 +1,5 @@ use serde_json::{Map, Number, Value}; +use vllm_tokenizer::DynTokenizer; use winnow::ascii::multispace0 as ws0; use winnow::combinator::{alt, delimited, eof, opt, separated, seq, terminated}; use winnow::error::{ContextError, ErrMode, ModalResult}; @@ -6,8 +7,6 @@ use winnow::prelude::*; use winnow::stream::{Partial, Stream}; use winnow::token::{literal, take_till, take_until}; -use vllm_tokenizer::DynTokenizer; - use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput}; use crate::reasoning::last_reasoning_boundary; use crate::tool::{Tool, ToolCallDelta}; diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index 6fe7d29b879..0410c568461 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -3,13 +3,12 @@ mod combined; mod gemma4; +pub use combined::CombinedParser; +pub use gemma4::Gemma4UnifiedParser; use thiserror::Error; use thiserror_ext::Macro; use vllm_tokenizer::DynTokenizer; -pub use combined::CombinedParser; -pub use gemma4::Gemma4UnifiedParser; - use crate::reasoning::ReasoningError; use crate::tool::{ StructuralTagModel, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput, diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index 3eba278267f..c566f0e2273 100644 --- a/rust/src/server/src/error.rs +++ b/rust/src/server/src/error.rs @@ -1,8 +1,7 @@ use axum::Json; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; -use thiserror_ext::AsReport as _; -use thiserror_ext::{Construct, Macro}; +use thiserror_ext::{AsReport as _, Construct, Macro}; use crate::routes::openai::utils::types::{ErrorDetail, ErrorResponse}; diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 3ebe6b31fd2..9836745dd60 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -3,7 +3,8 @@ use tonic::Status; use uuid::Uuid; -use vllm_engine_core_client::protocol::{StopReason, StructuredOutputsParams}; +use vllm_engine_core_client::protocol::output::StopReason; +use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; use vllm_text::{ DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, SamplingParams, TextDecodeOptions, TextRequest, @@ -502,7 +503,7 @@ impl ResponseOpts { #[cfg(test)] mod tests { - use vllm_engine_core_client::protocol::StopReason; + use vllm_engine_core_client::protocol::output::StopReason; use vllm_text::{FinishReason, Finished, Prompt}; use super::pb::finish_info::{FinishReason as PbFinishReason, StopReason as PbStopReason}; diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 48de928cfe7..e0d7ba635c3 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -18,9 +18,10 @@ use vllm_chat::{ ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt, }; -use vllm_engine_core_client::protocol::{ - EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, +use vllm_engine_core_client::protocol::output::{ + EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, }; +use vllm_engine_core_client::protocol::request::EngineCoreRequest; use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task}; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId}; use vllm_llm::Llm; diff --git a/rust/src/server/src/listener.rs b/rust/src/server/src/listener.rs index b1dcc919678..e61484398ae 100644 --- a/rust/src/server/src/listener.rs +++ b/rust/src/server/src/listener.rs @@ -152,8 +152,8 @@ impl axum::serve::Listener for Listener { /// Allow the unified listener to be adaptable to `tls_listener`. impl AsyncAccept for Listener { - type Connection = ListenerIo; type Address = ListenerAddr; + type Connection = ListenerIo; type Error = std::io::Error; fn poll_accept( diff --git a/rust/src/server/src/middleware/offload.rs b/rust/src/server/src/middleware/offload.rs index cad560eb7ae..28cde754c4c 100644 --- a/rust/src/server/src/middleware/offload.rs +++ b/rust/src/server/src/middleware/offload.rs @@ -57,9 +57,9 @@ where S::Error: Send + 'static, B: Send + 'static, { - type Response = S::Response; type Error = S::Error; type Future = BoxFuture<'static, Result>; + type Response = S::Response; fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_ready(cx) diff --git a/rust/src/server/src/routes/http_client_tests.rs b/rust/src/server/src/routes/http_client_tests.rs index b3c6977f5ec..9c04701e866 100644 --- a/rust/src/server/src/routes/http_client_tests.rs +++ b/rust/src/server/src/routes/http_client_tests.rs @@ -18,9 +18,10 @@ use vllm_chat::{ ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt, }; -use vllm_engine_core_client::protocol::{ - EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, +use vllm_engine_core_client::protocol::output::{ + EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, }; +use vllm_engine_core_client::protocol::request::EngineCoreRequest; use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task}; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId}; use vllm_llm::Llm; diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index 60cd14f9a81..5a879267c9f 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -21,7 +21,7 @@ use vllm_chat::{ AssistantBlockKind, AssistantMessageExt as _, ChatEvent, ChatEventStream, ChatEventStreamTrait, CollectedAssistantMessage, FinishReason, }; -use vllm_engine_core_client::protocol::StopReason; +use vllm_engine_core_client::protocol::output::StopReason; use self::convert::{ResponseOptions, prepare_chat_request}; use crate::config::ApiServerOptions; @@ -825,7 +825,7 @@ mod tests { use vllm_chat::{ AssistantBlockKind, AssistantContentBlock, AssistantToolCall, ChatEvent, FinishReason, }; - use vllm_engine_core_client::protocol::StopReason; + use vllm_engine_core_client::protocol::output::StopReason; use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob}; use super::{ 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 e974c836bb5..24d261314f1 100644 --- a/rust/src/server/src/routes/openai/utils/structured_outputs.rs +++ b/rust/src/server/src/routes/openai/utils/structured_outputs.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; -use vllm_engine_core_client::protocol::StructuredOutputsParams; +use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; use crate::error::ApiError; diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index b40542a7463..8e7953a263d 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -24,14 +24,15 @@ use vllm_chat::{ NewChatOutputProcessorOptions, }; use vllm_engine_core_client::mock_engine::default_ready_response; +use vllm_engine_core_client::protocol::decode_value; use vllm_engine_core_client::protocol::logprobs::{ Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob, }; -use vllm_engine_core_client::protocol::utility::{UtilityOutput, UtilityResultEnvelope}; -use vllm_engine_core_client::protocol::{ - EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason, - decode_value, +use vllm_engine_core_client::protocol::output::{ + EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, StopReason, }; +use vllm_engine_core_client::protocol::request::EngineCoreRequest; +use vllm_engine_core_client::protocol::utility::{UtilityOutput, UtilityResultEnvelope}; use vllm_engine_core_client::test_utils::{ IpcNamespace, spawn_mock_engine_task, spawn_mock_engine_task_with_ready, }; diff --git a/rust/src/server/src/state.rs b/rust/src/server/src/state.rs index eb37df40ea4..b56ca594871 100644 --- a/rust/src/server/src/state.rs +++ b/rust/src/server/src/state.rs @@ -1,6 +1,5 @@ -use std::sync::Arc; -use std::sync::OnceLock; use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, OnceLock}; use serde_json::Value; use sha2::{Digest, Sha256}; diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index a36480edf6b..31cf0058cdc 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -3,15 +3,15 @@ use std::collections::BTreeSet; pub(crate) mod logprobs; pub(crate) mod token_ids; -use vllm_engine_core_client::protocol::EngineCoreSamplingParams; +use logprobs::validate_logprobs; +use token_ids::{validate_prompt_token_ids, validate_vocab_range}; +use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; use vllm_llm::GenerateRequest; use vllm_tokenizer::Tokenizer; use crate::backend::{SamplingHints, SamplingLimits}; use crate::error::{Error, Result}; use crate::request::{SamplingParams, TextRequest}; -use logprobs::validate_logprobs; -use token_ids::{validate_prompt_token_ids, validate_vocab_range}; /// One text request after it has been lowered into the raw generate boundary. #[derive(Debug)] diff --git a/rust/src/text/src/lower/logprobs.rs b/rust/src/text/src/lower/logprobs.rs index 3c90f339107..116835560ce 100644 --- a/rust/src/text/src/lower/logprobs.rs +++ b/rust/src/text/src/lower/logprobs.rs @@ -3,9 +3,10 @@ //! `-1` is expanded only for bounds checks. The original request values are //! passed through to engine-core. -use crate::backend::SamplingLimits; use thiserror::Error; +use crate::backend::SamplingLimits; + #[derive(Debug, Error)] pub enum LogprobsError { #[error("{parameter} must be non-negative or -1, got {value}")] diff --git a/rust/src/text/src/lower/token_ids.rs b/rust/src/text/src/lower/token_ids.rs index 740329b8bde..c2371858837 100644 --- a/rust/src/text/src/lower/token_ids.rs +++ b/rust/src/text/src/lower/token_ids.rs @@ -1,7 +1,7 @@ use std::result::Result; use thiserror::Error; -use vllm_engine_core_client::protocol::EngineCoreSamplingParams; +use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; use crate::SamplingLimits; diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs index 203efab460c..a9444e459c4 100644 --- a/rust/src/text/src/output/decoded.rs +++ b/rust/src/text/src/output/decoded.rs @@ -5,7 +5,7 @@ use futures::{Stream, StreamExt}; use serde::{Deserialize, Serialize}; use tracing::{Level, debug, trace}; use vllm_engine_core_client::AbortCause; -use vllm_engine_core_client::protocol::StopReason; +use vllm_engine_core_client::protocol::output::StopReason; use vllm_llm::{FinishReason, GenerateOutput, TokenUsage}; use vllm_tokenizer::{DynTokenizer, IncrementalDecoder}; diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index 621da75ad51..3b59b687b5a 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -5,7 +5,8 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use vllm_engine_core_client::protocol::lora::LoraRequest; use vllm_engine_core_client::protocol::multimodal::MmFeatures; -use vllm_engine_core_client::protocol::{ReasoningParserKwargs, StructuredOutputsParams}; +use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; +use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; use crate::error::{Error, Result}; use crate::output::TextDecodeOptions; diff --git a/rust/src/tokenizer/src/hf/added_tokens.rs b/rust/src/tokenizer/src/hf/added_tokens.rs index d1d9fa8b4b4..ab0c0155598 100644 --- a/rust/src/tokenizer/src/hf/added_tokens.rs +++ b/rust/src/tokenizer/src/hf/added_tokens.rs @@ -1,11 +1,12 @@ +use std::fs; +use std::path::Path; + use serde::{Deserialize, Serialize}; use thiserror_ext::AsReport as _; use tracing::warn; use crate::Result; -use std::{fs, path::Path}; - /// Minimal `tokenizer.json` projection used to patch `added_tokens` while /// preserving the rest of the tokenizer definition verbatim. #[derive(Debug, Deserialize, Serialize)]