diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 42b3c230cc1..3da176cd1c6 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2371,6 +2371,15 @@ version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" +[[package]] +name = "libmimalloc-sys" +version = "0.1.49" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9" +dependencies = [ + "cc", +] + [[package]] name = "libredox" version = "0.1.14" @@ -2569,6 +2578,15 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b" +[[package]] +name = "mimalloc" +version = "0.1.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d4139bb28d14ad1facf21d5eb8825051b326e172d216b39f6d31df53cc97862" +dependencies = [ + "libmimalloc-sys", +] + [[package]] name = "mime" version = "0.3.17" @@ -5644,6 +5662,7 @@ dependencies = [ "educe", "expect-test", "itertools 0.14.0", + "mimalloc", "native-tls", "serde", "serde_json", @@ -5742,6 +5761,25 @@ dependencies = [ "prometheus-client", ] +[[package]] +name = "vllm-mock-engine" +version = "0.1.0" +dependencies = [ + "anyhow", + "asynk-strim-attr", + "clap", + "futures", + "rand 0.9.2", + "rmpv", + "serde", + "tokio", + "tokio-util", + "tracing", + "tracing-subscriber", + "vllm-engine-core-client", + "zeromq", +] + [[package]] name = "vllm-reasoning-parser" version = "0.1.0" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index a43aa53862c..e742b68b2ad 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -6,6 +6,7 @@ members = [ "src/llm", "src/managed-engine", "src/metrics", + "src/mock-engine", "src/reasoning-parser", "src/server", "src/text", @@ -45,6 +46,7 @@ http-body = "1.0.1" itertools = "0.14.0" libc = "0.2.177" llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "5b558989844d1c7af3e43d0f604069ffd9c06320" } +mimalloc = "0.1.52" minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } minijinja-contrib = { version = "2.0", features = ["pycompat"] } native-tls-vendored = { package = "native-tls", version = "0.2.18", features = ["vendored"] } @@ -57,6 +59,7 @@ prometheus-client = "0.24.0" prometheus-client-derive-encode = "0.5.0" prost = "0.14.3" prost-types = "0.14.3" +rand = "0.9.2" reasoning-parser = "1.2.2" reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] } riptoken = { version = "0.3.0", default-features = false } diff --git a/rust/src/cmd/Cargo.toml b/rust/src/cmd/Cargo.toml index b684d072202..b0caa65b4e8 100644 --- a/rust/src/cmd/Cargo.toml +++ b/rust/src/cmd/Cargo.toml @@ -17,6 +17,7 @@ anyhow.workspace = true clap.workspace = true educe.workspace = true itertools.workspace = true +mimalloc.workspace = true native-tls-vendored = { workspace = true, optional = true } serde.workspace = true serde_json.workspace = true diff --git a/rust/src/cmd/src/main.rs b/rust/src/cmd/src/main.rs index ce4e37e09bd..c015319ed18 100644 --- a/rust/src/cmd/src/main.rs +++ b/rust/src/cmd/src/main.rs @@ -11,6 +11,9 @@ use vllm_managed_engine::ManagedEngineHandle; use crate::cli::{Cli, Command}; +#[global_allocator] +static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; + const TOKIO_WORKER_THREADS_ENV: &str = "TOKIO_WORKER_THREADS"; const DEFAULT_MAX_TOKIO_WORKER_THREADS: usize = 32; diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 03805388243..e432638f350 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -107,13 +107,13 @@ impl ClientInner { Ok(registry.abortable_request_ids(request_ids)) } - /// Obtain the stream sender for one output. If it indicates the request is - /// finished, it will be removed from the registry. - pub fn take_sender_for_output( + /// Obtain stream senders for a whole engine output batch with one registry + /// lock acquisition. + pub fn take_senders_for_outputs<'a>( &self, - output: &EngineCoreOutput, - ) -> Option>> { - self.request_reg.lock().sender_for_output(output) + outputs: impl IntoIterator, + ) -> Vec>>> { + self.request_reg.lock().senders_for_outputs(outputs) } /// Remove a batch of requests that have finished or aborted, returning @@ -301,9 +301,10 @@ pub(crate) async fn run_output_dispatcher_loop( match outputs.classify() { ClassifiedEngineCoreOutputs::RequestBatch(batch) => { - for output in batch.outputs { + 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(); - let Some(sender) = inner.take_sender_for_output(&output) else { + let Some(sender) = sender else { debug!(request_id, "dropping output for inactive request"); continue; }; diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index 0dd0bd3e968..d47c5a80719 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -1,4 +1,4 @@ -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap}; use std::sync::atomic::{AtomicU64, Ordering}; use tokio::sync::{mpsc, oneshot}; @@ -80,7 +80,7 @@ impl EngineRoutingState { #[derive(Debug)] pub struct RequestRegistry { closed: bool, - requests: BTreeMap, + requests: HashMap, routing_per_engine: BTreeMap, } @@ -88,7 +88,7 @@ impl RequestRegistry { pub fn new(engines: &[ConnectedEngine]) -> Self { Self { closed: false, - requests: BTreeMap::default(), + requests: HashMap::default(), routing_per_engine: engines .iter() .map(|engine| (engine.engine_id.clone(), EngineRoutingState::default())) @@ -180,6 +180,15 @@ impl RequestRegistry { } } + /// Obtain stream senders for a whole engine output batch under one + /// registry lock. Finished outputs are removed before returning. + pub fn senders_for_outputs<'a>( + &mut self, + outputs: impl IntoIterator, + ) -> Vec> { + outputs.into_iter().map(|output| self.sender_for_output(output)).collect() + } + /// Remove a batch of requests that have finished or aborted, returning /// their stream senders. pub fn finish_many<'a>( diff --git a/rust/src/engine-core-client/src/lib.rs b/rust/src/engine-core-client/src/lib.rs index 914ce874e9e..e39e29c4e5a 100644 --- a/rust/src/engine-core-client/src/lib.rs +++ b/rust/src/engine-core-client/src/lib.rs @@ -2,6 +2,7 @@ mod client; mod coordinator; mod error; mod metrics; +pub mod mock_engine; pub mod protocol; #[cfg(any(test, feature = "test-util"))] pub mod test_utils; diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs new file mode 100644 index 00000000000..aa1cfecaee4 --- /dev/null +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -0,0 +1,265 @@ +use std::path::Path; +use std::time::Duration; + +use tokio::time::timeout; +use zeromq::prelude::{Socket, SocketRecv, SocketSend}; +use zeromq::util::PeerIdentity; +use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage}; + +use crate::EngineId; +use crate::error::{Error, Result, bail_unexpected_handshake_message}; +use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage}; +use crate::protocol::{ModelDtype, decode_msgpack, encode_msgpack}; + +/// Default model length advertised by reusable mock engine helpers. +pub const DEFAULT_MOCK_MAX_MODEL_LEN: u64 = 1024 * 1024; +/// Default KV block count advertised by reusable mock engine helpers. +pub const DEFAULT_MOCK_NUM_GPU_BLOCKS: u64 = 0; + +/// Startup behavior for one mock engine joining a frontend. +#[derive(Debug, Clone)] +pub struct MockEngineConfig { + /// Whether the engine should advertise itself as local to the frontend. + pub local: bool, + /// Whether the engine should advertise itself as headless. + pub headless: bool, + /// Engine-ready payload reported after INIT, including max model length, + /// KV block count, and dtype. + pub ready_response: EngineCoreReadyResponse, + /// Maximum time to wait for IPC endpoints to appear before connecting. + pub connect_timeout: Duration, +} + +impl Default for MockEngineConfig { + fn default() -> Self { + Self { + local: false, + headless: true, + ready_response: default_ready_response(), + connect_timeout: Duration::from_secs(5), + } + } +} + +/// Construct the ready response used by the standalone mock engine CLI. +pub fn default_ready_response() -> EngineCoreReadyResponse { + EngineCoreReadyResponse { + max_model_len: DEFAULT_MOCK_MAX_MODEL_LEN, + num_gpu_blocks: DEFAULT_MOCK_NUM_GPU_BLOCKS, + dp_stats_address: None, + dtype: Some(ModelDtype::Float32), + } +} + +/// Coordinator-side sockets used by one mock engine when coordinator mode +/// is enabled. +pub struct MockCoordinatorSockets { + /// Subscription socket that receives coordinator broadcasts such as + /// `START_DP_WAVE`. + pub input_sub: SubSocket, + /// Push socket used to send coordinator-only `EngineCoreOutputs` back to + /// the frontend. + pub output_push: PushSocket, +} + +/// One mock engine's connection to one frontend client. +/// +/// vLLM launches one engine-client pair per API server process. A remote +/// engine connects to every advertised input/output pair and uses the request's +/// `client_index` to route outputs back to the originating API server. +pub struct MockEngineDataSockets { + /// Socket used to receive frontend requests. + pub dealer: DealerSocket, + /// Socket used to publish normal request outputs back to the frontend. + pub push: PushSocket, +} + +/// Frontend-facing sockets owned by one mock engine. +pub struct MockEngineSockets { + /// Decoded INIT message sent by the frontend during handshake. + pub init: HandshakeInitMessage, + /// Data sockets for all frontend clients in client-index order. + /// + /// For Rust frontend this will always be one socket, while for Python frontend + /// this may be multiple sockets if there are multiple API server processes. + pub data_sockets: Vec, + /// Optional coordinator sockets when the client enabled the in-process + /// coordinator. + pub coordinator: Option, +} + +/// Build a HELLO or READY handshake status payload. +fn ready_message(status: &str, config: &MockEngineConfig) -> ReadyMessage { + ReadyMessage { + status: Some(status.to_string()), + local: Some(config.local), + headless: Some(config.headless), + parallel_config_hash: None, + } +} + +/// Convert an engine id into a ZMQ DEALER identity. +fn peer_identity(engine_id: impl Into) -> Result { + let engine_id = engine_id.into(); + PeerIdentity::try_from(engine_id.clone()).map_err(|error| Error::UnexpectedHandshakeMessage { + message: format!( + "invalid mock engine identity {:?}: {error}", + engine_id.to_vec() + ), + }) +} + +/// Wait for an IPC endpoint path to appear before attempting to connect. +async fn wait_for_ipc_endpoint(endpoint: &str, connect_timeout: Duration) -> Result<()> { + let Some(socket_path) = endpoint.strip_prefix("ipc://") else { + return Ok(()); + }; + + timeout(connect_timeout, async { + while !Path::new(socket_path).exists() { + tokio::time::sleep(Duration::from_millis(20)).await; + } + }) + .await + .map_err(|_| Error::HandshakeTimeout { + stage: "mock engine IPC endpoint", + timeout: connect_timeout, + }) +} + +/// Encode the engine-ready response sent on input socket registration. +fn ready_response_payload(config: &MockEngineConfig) -> Result> { + encode_msgpack(&config.ready_response) +} + +/// Join a frontend-owned handshake endpoint and open mock engine sockets. +pub async fn connect_to_frontend( + engine_handshake: impl AsRef, + engine_id: impl Into, + config: MockEngineConfig, +) -> Result { + let engine_handshake = engine_handshake.as_ref(); + wait_for_ipc_endpoint(engine_handshake, config.connect_timeout).await?; + + let peer_identity = peer_identity(engine_id)?; + let mut options = SocketOptions::default(); + options.peer_identity(peer_identity.clone()); + let mut handshake = DealerSocket::with_options(options); + handshake.connect(engine_handshake).await?; + handshake + .send(ZmqMessage::from(encode_msgpack(&ready_message( + "HELLO", &config, + ))?)) + .await?; + + let init_frames = handshake.recv().await?.into_vec(); + if init_frames.len() != 1 { + bail_unexpected_handshake_message!( + "expected one INIT frame from frontend, got {}", + init_frames.len() + ); + } + let init: HandshakeInitMessage = decode_msgpack(init_frames[0].as_ref())?; + + if init.addresses.inputs.is_empty() { + return Err(Error::UnexpectedHandshakeMessage { + message: "frontend INIT did not include an input address".to_string(), + }); + } + if init.addresses.inputs.len() != init.addresses.outputs.len() { + return Err(Error::UnexpectedHandshakeMessage { + message: format!( + "frontend INIT input/output address count mismatch: {} inputs, {} outputs", + init.addresses.inputs.len(), + init.addresses.outputs.len() + ), + }); + } + + let mut data_sockets = Vec::with_capacity(init.addresses.inputs.len()); + for (input_address, output_address) in + init.addresses.inputs.iter().zip(init.addresses.outputs.iter()) + { + wait_for_ipc_endpoint(input_address, config.connect_timeout).await?; + wait_for_ipc_endpoint(output_address, config.connect_timeout).await?; + + let mut input_options = SocketOptions::default(); + input_options.peer_identity(peer_identity.clone()); + let mut dealer = DealerSocket::with_options(input_options); + dealer.connect(input_address).await?; + dealer.send(ZmqMessage::from(ready_response_payload(&config)?)).await?; + + let mut push = PushSocket::new(); + push.connect(output_address).await?; + + data_sockets.push(MockEngineDataSockets { dealer, push }); + } + + let coordinator = match ( + init.addresses.coordinator_input.as_deref(), + init.addresses.coordinator_output.as_deref(), + ) { + (Some(coordinator_input), Some(coordinator_output)) => { + let mut input_sub = SubSocket::new(); + input_sub.connect(coordinator_input).await?; + input_sub.subscribe("").await?; + + let mut output_push = PushSocket::new(); + output_push.connect(coordinator_output).await?; + + let ready = input_sub.recv().await?.into_vec(); + if ready.len() != 1 || ready[0].as_ref() != b"READY" { + bail_unexpected_handshake_message!( + "expected coordinator READY marker, got {:?}", + ready + ); + } + + Some(MockCoordinatorSockets { + input_sub, + output_push, + }) + } + (None, None) => None, + _ => bail_unexpected_handshake_message!( + "coordinator handshake addresses must be both present or both absent" + ), + }; + + handshake + .send(ZmqMessage::from(encode_msgpack(&ready_message( + "READY", &config, + ))?)) + .await?; + + Ok(MockEngineSockets { + init, + data_sockets, + coordinator, + }) +} + +/// Join already-bootstrapped frontend input/output sockets directly. +pub async fn connect_to_bootstrapped_frontend( + input_address: impl AsRef, + output_address: impl AsRef, + engine_id: impl Into, + config: MockEngineConfig, +) -> Result<(DealerSocket, PushSocket)> { + let input_address = input_address.as_ref(); + let output_address = output_address.as_ref(); + wait_for_ipc_endpoint(input_address, config.connect_timeout).await?; + wait_for_ipc_endpoint(output_address, config.connect_timeout).await?; + + let peer_identity = peer_identity(engine_id)?; + let mut input_options = SocketOptions::default(); + input_options.peer_identity(peer_identity); + let mut dealer = DealerSocket::with_options(input_options); + dealer.connect(input_address).await?; + dealer.send(ZmqMessage::from(ready_response_payload(&config)?)).await?; + + let mut push = PushSocket::new(); + push.connect(output_address).await?; + + Ok((dealer, push)) +} diff --git a/rust/src/engine-core-client/src/protocol/mod.rs b/rust/src/engine-core-client/src/protocol/mod.rs index 009311eae41..4a00d9d31c5 100644 --- a/rust/src/engine-core-client/src/protocol/mod.rs +++ b/rust/src/engine-core-client/src/protocol/mod.rs @@ -36,6 +36,14 @@ fn is_false(v: &bool) -> bool { !v } +fn default_top_p() -> f32 { + 1.0 +} + +fn default_repetition_penalty() -> f32 { + 1.0 +} + mod classified_outputs; pub mod dtype; pub mod handshake; @@ -65,6 +73,24 @@ pub enum EngineCoreRequestType { } 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", @@ -200,14 +226,17 @@ pub struct EngineCoreSamplingParams { /// greedy sampling. 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. + #[serde(default)] pub top_k: u32, /// Random seed used by the sampler when present. pub seed: Option, /// Maximum number of tokens to generate per output sequence. pub max_tokens: u32, /// Minimum number of tokens to generate before EOS or stop-token handling. + #[serde(default)] pub min_tokens: u32, /// Number of log probabilities to return per generated token. /// @@ -218,12 +247,14 @@ pub struct EngineCoreSamplingParams { /// `None` disables prompt logprobs. `-1` requests the full vocabulary. pub prompt_logprobs: Option, /// Minimum probability threshold for token sampling. + #[serde(default)] 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, diff --git a/rust/src/engine-core-client/src/protocol/utility.rs b/rust/src/engine-core-client/src/protocol/utility.rs index b46caef49cc..ef7e862d517 100644 --- a/rust/src/engine-core-client/src/protocol/utility.rs +++ b/rust/src/engine-core-client/src/protocol/utility.rs @@ -102,7 +102,7 @@ impl<'de> Deserialize<'de> for UtilityCallId { /// /// Original Python payload shape: /// `(client_index, call_id, method_name, args)` -#[derive(Debug, Clone, PartialEq, Serialize_tuple)] +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)] pub struct EngineCoreUtilityRequest { pub client_index: u32, pub call_id: UtilityCallId, diff --git a/rust/src/engine-core-client/src/test_utils.rs b/rust/src/engine-core-client/src/test_utils.rs index 99c266dc626..f1c5c65503f 100644 --- a/rust/src/engine-core-client/src/test_utils.rs +++ b/rust/src/engine-core-client/src/test_utils.rs @@ -1,17 +1,18 @@ use std::future::Future; use std::path::Path; use std::pin::Pin; -use std::time::Duration; use tempfile::TempDir; use tokio::sync::oneshot; -use zeromq::prelude::{Socket, SocketRecv, SocketSend}; -use zeromq::util::PeerIdentity; -use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage}; +use zeromq::{DealerSocket, PushSocket}; use crate::EngineId; +pub use crate::mock_engine::{MockCoordinatorSockets, MockEngineSockets}; +use crate::mock_engine::{ + MockEngineConfig, MockEngineDataSockets, connect_to_bootstrapped_frontend, connect_to_frontend, +}; use crate::protocol::ModelDtype; -use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage}; +use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage}; /// Per-test IPC endpoint namespace backed by a unique temporary directory. /// @@ -52,156 +53,29 @@ impl IpcNamespace { } } -/// Construct a standard local READY message used by mock engines in tests. -fn ready_message(status: &str) -> ReadyMessage { - ReadyMessage { - status: Some(status.to_string()), - local: Some(true), - headless: Some(true), - parallel_config_hash: None, +fn test_mock_engine_config() -> MockEngineConfig { + MockEngineConfig { + local: true, + headless: true, + ready_response: EngineCoreReadyResponse { + max_model_len: 4096, + num_gpu_blocks: 0, + dp_stats_address: None, + dtype: Some(ModelDtype::Float32), + }, + ..Default::default() } } -/// Construct a default ready response payload for mock engine input -/// registration. -fn ready_response_payload() -> Vec { - rmp_serde::to_vec_named(&EngineCoreReadyResponse { - max_model_len: 4096, - num_gpu_blocks: 0, - dp_stats_address: None, - dtype: Some(ModelDtype::Float32), - }) - .expect("encode ready response payload") -} - -/// Coordinator-side sockets connected by one mock engine when coordinator mode -/// is enabled. -pub struct MockCoordinatorConnections { - /// Subscription socket that receives coordinator broadcasts such as - /// `START_DP_WAVE`. - pub input_sub: SubSocket, - /// Push socket used to send coordinator-only `EngineCoreOutputs` back to - /// the frontend. - pub output_push: PushSocket, -} - -/// Fully connected mock engine transport state used by tests. -pub struct MockEngineConnections { - /// Decoded INIT message sent by the frontend during handshake. - pub init: HandshakeInitMessage, - /// Socket used to receive frontend requests. - pub dealer: DealerSocket, - /// Socket used to publish normal request outputs back to the frontend. - pub push: PushSocket, - /// Optional coordinator sockets when the client enabled the in-process - /// coordinator. - pub coordinator: Option, -} - /// Complete the engine-core handshake and connect mock input/output sockets /// plus optional coordinator sockets. -pub async fn setup_mock_engine_connections( +pub async fn setup_mock_engine_sockets( engine_handshake: String, engine_id: impl Into, -) -> MockEngineConnections { - // Wait for the client to bind the handshake socket before connecting. - // A fixed sleep is racy under CI load; instead poll for the socket file. - let socket_path = engine_handshake - .strip_prefix("ipc://") - .expect("handshake address must be ipc://"); - for _ in 0..100 { - if Path::new(socket_path).exists() { - break; - } - tokio::time::sleep(Duration::from_millis(20)).await; - } - - let peer_identity = PeerIdentity::try_from(engine_id.into()).expect("peer id"); - - let mut options = SocketOptions::default(); - options.peer_identity(peer_identity.clone()); - let mut handshake = DealerSocket::with_options(options); - handshake - .connect(&engine_handshake) +) -> MockEngineSockets { + connect_to_frontend(engine_handshake, engine_id, test_mock_engine_config()) .await - .expect("connect mock engine handshake socket"); - handshake - .send(ZmqMessage::from( - rmp_serde::to_vec_named(&ready_message("HELLO")).expect("encode HELLO ready message"), - )) - .await - .expect("send HELLO ready message"); - - let init_frames = handshake.recv().await.expect("receive handshake init message").into_vec(); - assert_eq!(init_frames.len(), 1); - let init: HandshakeInitMessage = - rmp_serde::from_slice(init_frames[0].as_ref()).expect("decode handshake init message"); - - let mut input_options = SocketOptions::default(); - input_options.peer_identity(peer_identity); - let mut dealer = DealerSocket::with_options(input_options); - dealer - .connect(&init.addresses.inputs[0]) - .await - .expect("connect mock engine input socket"); - dealer - .send(ZmqMessage::from(ready_response_payload())) - .await - .expect("send mock engine input ready frame"); - - let mut push = PushSocket::new(); - push.connect(&init.addresses.outputs[0]) - .await - .expect("connect mock engine output socket"); - - let coordinator = match ( - init.addresses.coordinator_input.as_deref(), - init.addresses.coordinator_output.as_deref(), - ) { - (Some(coordinator_input), Some(coordinator_output)) => { - let mut input_sub = SubSocket::new(); - input_sub - .connect(coordinator_input) - .await - .expect("connect mock engine coordinator input socket"); - input_sub - .subscribe("") - .await - .expect("subscribe mock engine coordinator input socket"); - - let mut output_push = PushSocket::new(); - output_push - .connect(coordinator_output) - .await - .expect("connect mock engine coordinator output socket"); - - let ready = - input_sub.recv().await.expect("receive coordinator READY marker").into_vec(); - assert_eq!(ready.len(), 1); - assert_eq!(ready[0].as_ref(), b"READY"); - - Some(MockCoordinatorConnections { - input_sub, - output_push, - }) - } - (None, None) => None, - _ => panic!("coordinator handshake addresses must be both present or both absent"), - }; - - handshake - .send(ZmqMessage::from( - rmp_serde::to_vec_named(&ready_message("READY")).expect("encode READY ready message"), - )) - .await - .expect("send READY ready message"); - - MockEngineConnections { - init, - dealer, - push, - coordinator, - } + .expect("connect mock engine") } /// Connect one mock engine directly to already-bootstrapped frontend @@ -211,31 +85,14 @@ pub async fn setup_bootstrapped_mock_engine( output_address: String, engine_id: impl Into, ) -> (DealerSocket, PushSocket) { - for endpoint in [&input_address, &output_address] { - if let Some(socket_path) = endpoint.strip_prefix("ipc://") { - for _ in 0..100 { - if Path::new(socket_path).exists() { - break; - } - tokio::time::sleep(Duration::from_millis(20)).await; - } - } - } - - let peer_identity = PeerIdentity::try_from(engine_id.into()).expect("peer id"); - let mut input_options = SocketOptions::default(); - input_options.peer_identity(peer_identity); - let mut dealer = DealerSocket::with_options(input_options); - dealer.connect(&input_address).await.expect("connect mock engine input socket"); - dealer - .send(ZmqMessage::from(ready_response_payload())) - .await - .expect("send mock engine input ready frame"); - - let mut push = PushSocket::new(); - push.connect(&output_address).await.expect("connect mock engine output socket"); - - (dealer, push) + connect_to_bootstrapped_frontend( + input_address, + output_address, + engine_id, + test_mock_engine_config(), + ) + .await + .expect("connect bootstrapped mock engine") } /// Complete the engine-core handshake and connect mock input/output sockets. @@ -247,9 +104,11 @@ pub async fn setup_mock_engine_with_init( engine_handshake: String, engine_id: impl Into, ) -> (HandshakeInitMessage, DealerSocket, PushSocket) { - let MockEngineConnections { - init, dealer, push, .. - } = setup_mock_engine_connections(engine_handshake, engine_id).await; + let MockEngineSockets { + init, data_sockets, .. + } = setup_mock_engine_sockets(engine_handshake, engine_id).await; + let MockEngineDataSockets { dealer, push } = + data_sockets.into_iter().next().expect("mock engine data socket"); (init, dealer, push) } diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 54128330bb0..d1e0b03cc7d 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -30,7 +30,7 @@ use crate::protocol::{ EngineCoreRequestType, EngineCoreSamplingParams, decode_engine_core_outputs, }; use crate::test_utils::{ - IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_connections, + IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets, setup_mock_engine_with_init, spawn_mock_engine_task, }; use crate::{ @@ -477,8 +477,8 @@ async fn coordinator_handshake_includes_engine_control_addresses() { let (init_tx, init_rx) = oneshot::channel(); let (shutdown_tx, shutdown_rx) = oneshot::channel(); let engine_task = tokio::spawn(async move { - let connections = setup_mock_engine_connections(handshake_address, &engine_id).await; - let _ = init_tx.send(connections.init.clone()); + let sockets = setup_mock_engine_sockets(handshake_address, &engine_id).await; + let _ = init_tx.send(sockets.init.clone()); let _ = shutdown_rx.await; }); @@ -515,14 +515,15 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { let engine0_task = tokio::spawn({ let handshake_address = handshake_address.clone(); async move { - let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await; + let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await; let mut coordinator = engine.coordinator.take().expect("coordinator sockets should be present"); + let data_socket = engine.data_sockets.first_mut().expect("data socket"); let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; assert_eq!((wave, exclude_engine), (0, 0)); - let add = recv_engine_message(&mut engine.dealer).await; + let add = recv_engine_message(&mut data_socket.dealer).await; assert_eq!(add[0].as_ref(), &[0x00]); let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); assert_eq!(request.request_id, "req-1"); @@ -538,7 +539,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { ); send_outputs( - &mut engine.push, + &mut data_socket.push, EngineCoreOutputs { engine_index: 0, outputs: vec![request_output( @@ -565,14 +566,14 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; assert_eq!((wave, exclude_engine), (1, 0)); - let add = recv_engine_message(&mut engine.dealer).await; + let add = recv_engine_message(&mut data_socket.dealer).await; assert_eq!(add[0].as_ref(), &[0x00]); let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); assert_eq!(request.request_id, "req-3"); assert_eq!(request.current_wave, 1); send_outputs( - &mut engine.push, + &mut data_socket.push, EngineCoreOutputs { engine_index: 0, outputs: vec![request_output( @@ -594,14 +595,15 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { let engine1_task = tokio::spawn({ let handshake_address = handshake_address.clone(); async move { - let mut engine = setup_mock_engine_connections(handshake_address, &[0x01, 0x00]).await; + let mut engine = setup_mock_engine_sockets(handshake_address, &[0x01, 0x00]).await; let mut coordinator = engine.coordinator.take().expect("coordinator sockets should be present"); + let data_socket = engine.data_sockets.first_mut().expect("data socket"); let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; assert_eq!((wave, exclude_engine), (0, 0)); - let add = recv_engine_message(&mut engine.dealer).await; + let add = recv_engine_message(&mut data_socket.dealer).await; assert_eq!(add[0].as_ref(), &[0x00]); let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); assert_eq!(request.request_id, "req-2"); @@ -617,7 +619,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { ); send_outputs( - &mut engine.push, + &mut data_socket.push, EngineCoreOutputs { engine_index: 1, outputs: vec![request_output( @@ -637,7 +639,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { assert!( timeout( Duration::from_millis(200), - recv_engine_message(&mut engine.dealer) + recv_engine_message(&mut data_socket.dealer) ) .await .is_err() @@ -712,7 +714,7 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() { let engine0_task = tokio::spawn({ let handshake_address = handshake_address.clone(); async move { - let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await; + let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await; let mut coordinator = engine.coordinator.take().expect("coordinator sockets should be present"); @@ -727,7 +729,7 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() { let engine1_task = tokio::spawn({ let handshake_address = handshake_address.clone(); async move { - let mut engine = setup_mock_engine_connections(handshake_address, &[0x01, 0x00]).await; + let mut engine = setup_mock_engine_sockets(handshake_address, &[0x01, 0x00]).await; let mut coordinator = engine.coordinator.take().expect("coordinator sockets should be present"); @@ -778,9 +780,10 @@ async fn coordinator_accepts_stats_only_outputs() { let (shutdown_tx, shutdown_rx) = oneshot::channel(); let engine_task = tokio::spawn(async move { - let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await; + let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await; let mut coordinator = engine.coordinator.take().expect("coordinator sockets should be present"); + let data_socket = engine.data_sockets.first_mut().expect("data socket"); let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; assert_eq!((wave, exclude_engine), (0, 0)); @@ -799,13 +802,13 @@ async fn coordinator_accepts_stats_only_outputs() { ) .await; - let add = recv_engine_message(&mut engine.dealer).await; + let add = recv_engine_message(&mut data_socket.dealer).await; assert_eq!(add[0].as_ref(), &[0x00]); let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); assert_eq!(request.request_id, "req-stats"); send_outputs( - &mut engine.push, + &mut data_socket.push, EngineCoreOutputs { engine_index: 0, outputs: vec![request_output( diff --git a/rust/src/mock-engine/Cargo.toml b/rust/src/mock-engine/Cargo.toml new file mode 100644 index 00000000000..c6ba9f499e6 --- /dev/null +++ b/rust/src/mock-engine/Cargo.toml @@ -0,0 +1,30 @@ +[package] +name = "vllm-mock-engine" +version.workspace = true +edition.workspace = true +license.workspace = true + +[[bin]] +name = "vllm-mock-engine" +path = "src/main.rs" + +[dependencies] +anyhow.workspace = true +asynk-strim-attr.workspace = true +clap.workspace = true +futures.workspace = true +rand.workspace = true +rmpv.workspace = true +serde.workspace = true +tokio = { workspace = true, features = ["signal"] } +tokio-util.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true +vllm-engine-core-client.workspace = true +zeromq.workspace = true + +[dev-dependencies] +vllm-engine-core-client = { workspace = true, features = ["test-util"] } + +[lints] +workspace = true diff --git a/rust/src/mock-engine/README.md b/rust/src/mock-engine/README.md new file mode 100644 index 00000000000..6c48ce1e758 --- /dev/null +++ b/rust/src/mock-engine/README.md @@ -0,0 +1,104 @@ +# vLLM Mock Engine + +`vllm-mock-engine` is a small engine-side process for frontend stress testing. It +joins a frontend-owned startup handshake, reports a large ready response, treats +prefill as instant, and emits random decode tokens until each request reaches +its `max_tokens`. + +The frontend must own the handshake socket. Start the frontend first, then start +the mock engine with the same handshake address. + +## Start the mock engine + +```bash +cargo run -p vllm-mock-engine -- \ + --handshake-address tcp://127.0.0.1:29550 \ + --engine-count 1 \ + --output-token-chunk-size 1 \ + --vocab-size 32000 \ + --seed 0 \ + --log-requests +``` + +Useful knobs: + +- `--engine-count` must match the frontend's expected data-parallel engine + count. +- `--output-token-chunk-size` controls how many token IDs appear in one + `EngineCoreOutput`; values greater than 1 are useful for MTP/spec-decode + shaped frontend tests. +- `--vocab-size` should stay within the tokenizer vocabulary of the model used + by the frontend. + +Stop it with Ctrl-C. + +## Rust Frontend + +Terminal 1: + +```bash +cargo run --bin vllm-rs -- serve Qwen/Qwen3-0.6B \ + --data-parallel-size 1 \ + --data-parallel-size-local 0 \ + --handshake-port 29550 +``` + +Terminal 2: + +```bash +cargo run -p vllm-mock-engine -- \ + --handshake-address tcp://127.0.0.1:29550 +``` + +For multiple mock engines, set both sides to the same count: + +```bash +cargo run --bin vllm-rs -- serve Qwen/Qwen3-0.6B \ + --data-parallel-size 4 \ + --data-parallel-size-local 0 \ + --handshake-port 29550 + +cargo run -p vllm-mock-engine -- \ + --handshake-address tcp://127.0.0.1:29550 \ + --engine-count 4 +``` + +## Python Frontend + +Use `vllm serve` with `--data-parallel-size-local 0` so the Python process runs +as a frontend/API server and waits for external engines on +`--data-parallel-rpc-port`. + +Terminal 1: + +```bash +vllm serve Qwen/Qwen3-0.6B \ + --data-parallel-address 127.0.0.1 \ + --data-parallel-rpc-port 29550 \ + --data-parallel-size 1 \ + --data-parallel-size-local 0 +``` + +Terminal 2: + +```bash +cargo run -p vllm-mock-engine -- \ + --handshake-address tcp://127.0.0.1:29550 +``` + +## Smoke Request + +After either frontend is ready: + +```bash +curl http://127.0.0.1:8000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Qwen/Qwen3-0.6B", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 16, + "stream": true + }' +``` + +Always pass `max_tokens`; the mock engine stops by length. diff --git a/rust/src/mock-engine/src/engine.rs b/rust/src/mock-engine/src/engine.rs new file mode 100644 index 00000000000..2aa2f7bb397 --- /dev/null +++ b/rust/src/mock-engine/src/engine.rs @@ -0,0 +1,416 @@ +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::hash::{Hash as _, Hasher as _}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use anyhow::{Result, anyhow}; +use rand::rngs::StdRng; +use rand::{Rng as _, SeedableRng as _}; +use rmpv::Value; +use serde::Serialize; +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::utility::{ + EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope, +}; +use vllm_engine_core_client::protocol::{ + EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, +}; + +use super::Opt; + +/// Derive a stable per-request seed from the CLI seed, engine, and request id. +fn request_seed(base_seed: u64, engine_index: u32, request_id: &str) -> u64 { + let mut hasher = std::hash::DefaultHasher::new(); + base_seed.hash(&mut hasher); + engine_index.hash(&mut hasher); + request_id.hash(&mut hasher); + hasher.finish() +} + +/// Current UNIX timestamp in seconds for engine-core output envelopes. +fn now_secs() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or_default() +} + +/// Build one request output with only token IDs and terminal status populated. +fn request_output( + request_id: String, + new_token_ids: Vec, + finish_reason: Option, +) -> EngineCoreOutput { + EngineCoreOutput { + request_id, + new_token_ids, + finish_reason, + ..Default::default() + } +} + +/// Produce an empty output with a terminal finish reason for an invalid request. +fn empty_finish_outputs( + engine_index: u32, + request_id: String, + finish_reason: EngineCoreFinishReason, +) -> EngineCoreOutputs { + let output = request_output(request_id, Vec::new(), Some(finish_reason)); + let finished_requests = BTreeSet::from([output.request_id.clone()]); + + EngineCoreOutputs { + engine_index, + outputs: vec![output], + timestamp: now_secs(), + finished_requests: Some(finished_requests), + ..Default::default() + } +} + +/// Encode a utility result into the protocol's msgpack value envelope. +fn utility_envelope(value: T) -> Result +where + T: Serialize, +{ + Ok(UtilityResultEnvelope::without_type_info( + rmpv::ext::to_value(value)?, + )) +} + +/// Produce the minimal utility responses needed by the Rust frontend. +fn utility_response( + engine_index: u32, + request: EngineCoreUtilityRequest, +) -> Result { + let result = match request.method_name.as_str() { + "get_supported_tasks" => utility_envelope(vec!["generate"]), + "is_sleeping" => utility_envelope(false), + "reset_prefix_cache" => utility_envelope(true), + "reset_mm_cache" + | "reset_encoder_cache" + | "profile" + | "sleep" + | "wake_up" + | "execute_dummy_batch" => utility_envelope(()), + _ => utility_envelope(Value::Nil), + }?; + + Ok(EngineCoreOutputs { + engine_index, + utility_output: Some(UtilityOutput { + call_id: request.call_id, + failure_message: None, + result: Some(result), + }), + timestamp: now_secs(), + ..Default::default() + }) +} + +/// Message sent from the frontend to the mock engine task to drive the engine loop. +pub(crate) enum EngineInput { + Request(Box), + Abort(Vec), + Utility(EngineCoreUtilityRequest), + StartDpWave, +} + +/// Message sent from the mock engine task to the frontend for one engine output batch. +pub(crate) struct EngineOutput { + pub client_index: u32, + pub outputs: EngineCoreOutputs, +} + +/// Per-request decode state owned by one mock engine. +#[derive(Debug)] +struct ActiveRequest { + request_id: String, + client_index: u32, + prompt_len: usize, + max_tokens: usize, + generated: usize, + rng: StdRng, +} + +impl ActiveRequest { + /// Create a new active request from an incoming EngineCoreRequest, or return an immediate + /// finish reason if the request is invalid. + fn new( + engine_index: u32, + request: Box, + opt: &Opt, + ) -> Result { + let request_id = request.request_id; + let client_index = request.client_index; + let prompt_len = request.prompt_token_ids.as_ref().map(Vec::len).unwrap_or_default(); + + let Some(sampling_params) = request.sampling_params else { + warn!( + request_id, + "request has no sampling params; returning engine error" + ); + return Err(EngineCoreFinishReason::Error); + }; + let max_tokens = sampling_params.max_tokens as usize; + + if opt.log_requests { + info!( + request_id, + prompt_len, + max_tokens, + chunk_size = opt.output_token_chunk_size, + "mock request started" + ); + } + + if max_tokens == 0 { + return Err(EngineCoreFinishReason::Length); + } + + Ok(ActiveRequest { + rng: StdRng::seed_from_u64(request_seed(opt.seed, engine_index, &request_id)), + request_id, + client_index, + prompt_len, + max_tokens, + generated: 0, + }) + } + + /// Advance this request by one mock engine step. + fn step(&mut self, opt: &Opt) -> EngineCoreOutput { + let remaining = self.max_tokens - self.generated; + let chunk_len = remaining.min(opt.output_token_chunk_size); + let mut new_token_ids = Vec::with_capacity(chunk_len); + for _ in 0..chunk_len { + new_token_ids.push(self.rng.random_range(0..opt.vocab_size)); + } + self.generated += chunk_len; + + let finished = self.generated >= self.max_tokens; + request_output( + self.request_id.clone(), + new_token_ids, + finished.then_some(EngineCoreFinishReason::Length), + ) + } +} + +/// Internal state for one mock engine instance, owned by the engine loop task. +struct Engine { + engine_index: u32, + opt: Opt, + active_requests: HashMap, +} + +impl Engine { + /// Drain one frontend request message received on the input DEALER socket. + fn handle_input(&mut self, input: EngineInput) -> Result> { + let mut outputs = Vec::new(); + + match input { + EngineInput::Request(request) => { + let request_id = request.request_id.clone(); + let client_index = request.client_index; + + if self.active_requests.contains_key(&request_id) { + warn!( + engine_index = self.engine_index, + request_id, "duplicate mock request id" + ); + return Ok(vec![EngineOutput { + client_index, + outputs: empty_finish_outputs( + self.engine_index, + request_id, + EngineCoreFinishReason::Error, + ), + }]); + } + + match ActiveRequest::new(self.engine_index, request, &self.opt) { + Ok(request) => { + self.active_requests.insert(request_id, request); + } + Err(finish_reason) => { + return Ok(vec![EngineOutput { + client_index, + outputs: empty_finish_outputs( + self.engine_index, + request_id, + finish_reason, + ), + }]); + } + } + } + + EngineInput::Abort(request_ids) => { + let mut outputs_by_client = + BTreeMap::, BTreeSet)>::new(); + for request_id in request_ids { + if let Some(request) = self.active_requests.remove(&request_id) { + let output = request_output( + request_id.clone(), + Vec::new(), + Some(EngineCoreFinishReason::Abort), + ); + let (outputs, finished_requests) = outputs_by_client + .entry(request.client_index) + .or_insert_with(|| (Vec::new(), BTreeSet::new())); + outputs.push(output); + finished_requests.insert(request_id.clone()); + if self.opt.log_requests { + info!(request_id, finish_reason = "abort", "mock request aborted"); + } + } + } + for (client_index, (client_outputs, finished_requests)) in outputs_by_client { + outputs.push({ + let outputs = EngineCoreOutputs { + engine_index: self.engine_index, + outputs: client_outputs, + timestamp: now_secs(), + finished_requests: Some(finished_requests), + ..Default::default() + }; + EngineOutput { + client_index, + outputs, + } + }); + } + } + + EngineInput::Utility(request) => { + debug!( + engine_index = self.engine_index, + call_id = %request.call_id, + method = request.method_name, + "mock utility request" + ); + let client_index = request.client_index; + outputs.push({ + let outputs = utility_response(self.engine_index, request)?; + EngineOutput { + client_index, + outputs, + } + }); + } + + EngineInput::StartDpWave => { + debug!( + engine_index = self.engine_index, + "ignoring START_DP_WAVE in mock engine" + ); + } + } + + Ok(outputs) + } + + /// Advance active requests once and return one batched engine output. + fn step(&mut self) -> Vec { + if self.active_requests.is_empty() { + return Vec::new(); + } + + let mut outputs_by_client = + BTreeMap::, BTreeSet)>::new(); + let mut all_finished_requests = BTreeSet::new(); + + for request in self.active_requests.values_mut() { + let client_index = request.client_index; + let output = request.step(&self.opt); + let request_id = request.request_id.clone(); + let finished = output.finished(); + if output.finished() { + all_finished_requests.insert(request_id.clone()); + if self.opt.log_requests { + info!( + request_id, + prompt_len = request.prompt_len, + output_tokens = request.generated, + finish_reason = "length", + "mock request finished" + ); + } + } + let (outputs, finished_requests) = outputs_by_client + .entry(client_index) + .or_insert_with(|| (Vec::new(), BTreeSet::new())); + if finished { + finished_requests.insert(request_id.clone()); + } + outputs.push(output); + } + + for request_id in &all_finished_requests { + self.active_requests.remove(request_id); + } + + outputs_by_client + .into_iter() + .filter_map(|(client_index, (outputs, finished_requests))| { + (!outputs.is_empty()).then(|| EngineOutput { + client_index, + outputs: EngineCoreOutputs { + engine_index: self.engine_index, + outputs, + timestamp: now_secs(), + finished_requests: (!finished_requests.is_empty()) + .then_some(finished_requests), + ..Default::default() + }, + }) + }) + .collect() + } +} + +/// Run the main loop for the mock engine, receiving `EngineInput` from `input_rx` +/// and sending `EngineOutput` to `output_tx` until `shutdown` is cancelled. +pub(crate) async fn run_engine_loop( + engine_index: u32, + opt: Opt, + mut input_rx: mpsc::UnboundedReceiver, + output_tx: mpsc::Sender, + shutdown: CancellationToken, +) -> Result<()> { + let mut engine = Engine { + engine_index, + opt, + active_requests: HashMap::new(), + }; + + loop { + let outputs = tokio::select! { + biased; + _ = shutdown.cancelled() => break, + + input = input_rx.recv() => { + let input = input + .ok_or_else(|| anyhow!("mock engine input channel closed"))?; + engine.handle_input(input)? + } + + // If there are active requests, step them once after yielding to the scheduler to + // avoid blocking the engine loop while still making steady progress on request outputs. + _ = yield_now(), if !engine.active_requests.is_empty() => { + engine.step() + } + }; + + for output in outputs { + output_tx + .send(output) + .await + .map_err(|_| anyhow!("mock engine IO task shut down"))?; + } + } + + Ok(()) +} diff --git a/rust/src/mock-engine/src/io.rs b/rust/src/mock-engine/src/io.rs new file mode 100644 index 00000000000..28d77639c77 --- /dev/null +++ b/rust/src/mock-engine/src/io.rs @@ -0,0 +1,121 @@ +use anyhow::{Context as _, Result, anyhow, bail}; +use futures::{Stream, StreamExt as _, stream}; +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::utility::EngineCoreUtilityRequest; +use vllm_engine_core_client::protocol::{ + EngineCoreRequest, EngineCoreRequestType, decode_msgpack, encode_msgpack, +}; +use zeromq::{DealerSocket, PushSocket, SocketRecv as _, SocketSend as _, ZmqMessage}; + +use crate::engine::{EngineInput, EngineOutput}; + +/// Send one engine output batch to the client over the appropriate push socket. +async fn send_engine_outputs_to_client( + push_sockets: &mut [PushSocket], + EngineOutput { + client_index, + outputs, + }: EngineOutput, +) -> Result<()> { + let message = ZmqMessage::from(encode_msgpack(&outputs)?); + push_sockets[client_index as usize].send(message).await?; + Ok(()) +} + +/// Create a stream of `EngineInput` by continuously receiving messages from the given dealer socket +/// and decoding them into `EngineInput`. +fn dealer_input_stream(dealer: DealerSocket) -> impl Stream> { + stream::unfold(dealer, |mut dealer| async { + let input = loop { + let message = + match dealer.recv().await.context("failed to receive message from dealer socket") { + Ok(message) => message, + Err(err) => break Err(err), + }; + + match decode_request(message) { + Ok(input) => break Ok(input), + Err(err) => { + warn!(%err, "failed to decode engine request message; ignoring"); + } + } + }; + + Some((input, dealer)) + }) +} + +/// Decode a `ZmqMessage` into an `EngineInput`. Returns an error if the message is malformed or +/// contains an unknown/unsupported request type. +fn decode_request(message: ZmqMessage) -> Result { + let frames = message.into_vec(); + if frames.is_empty() { + bail!("empty engine request message"); + } + if frames.len() != 2 { + bail!("invalid frame count for engine request: {}", frames.len()); + } + + let request_type_frame = frames[0].as_ref(); + let Some(request_type) = EngineCoreRequestType::from_frame(request_type_frame) else { + bail!("unknown engine request type: {:?}", request_type_frame); + }; + + let input = match request_type { + EngineCoreRequestType::Add => { + let request: Box = decode_msgpack(frames[1].as_ref())?; + EngineInput::Request(request) + } + EngineCoreRequestType::Abort => { + let request_ids: Vec = decode_msgpack(frames[1].as_ref())?; + EngineInput::Abort(request_ids) + } + EngineCoreRequestType::Utility => { + let request: EngineCoreUtilityRequest = decode_msgpack(frames[1].as_ref())?; + EngineInput::Utility(request) + } + EngineCoreRequestType::StartDpWave => EngineInput::StartDpWave, + }; + + Ok(input) +} + +/// Run the main IO loop for the mock engine, continuously receiving and decoding raw messages from +/// the dealer sockets, sending them to the engine loop task via `input_tx`, and receiving +/// `EngineOutput` from the engine loop task via `output_rx` and sending them to the client over the +/// appropriate push socket, until `shutdown` is cancelled. +pub(crate) async fn run_io_loop( + data_sockets: Vec, + input_tx: mpsc::UnboundedSender, + mut output_rx: mpsc::Receiver, + shutdown: CancellationToken, +) -> Result<()> { + let (dealers, mut push_sockets): (Vec<_>, Vec<_>) = + data_sockets.into_iter().map(|sockets| (sockets.dealer, sockets.push)).unzip(); + let mut input_streams = + stream::select_all(dealers.into_iter().map(dealer_input_stream).map(Box::pin)); + + loop { + tokio::select! { + biased; + _ = shutdown.cancelled() => return Ok(()), + + output = output_rx.recv() => { + let output = output + .ok_or_else(|| anyhow!("mock engine output channel closed"))?; + send_engine_outputs_to_client(&mut push_sockets, output).await?; + } + + input = input_streams.next() => { + let input = input + .ok_or_else(|| anyhow!("mock engine input streams closed"))??; + input_tx + .send(input) + .map_err(|_| anyhow!("mock engine state task shut down"))?; + } + } + } +} diff --git a/rust/src/mock-engine/src/lib.rs b/rust/src/mock-engine/src/lib.rs new file mode 100644 index 00000000000..98cf54c7eaa --- /dev/null +++ b/rust/src/mock-engine/src/lib.rs @@ -0,0 +1,138 @@ +use anyhow::{Context, Result, bail}; +use clap::Parser; +use tokio::sync::mpsc; +use tokio::task::JoinSet; +use tokio_util::sync::CancellationToken; +use tracing::{error, info}; +use vllm_engine_core_client::EngineId; +use vllm_engine_core_client::mock_engine::{ + MockEngineConfig, MockEngineSockets, connect_to_frontend, +}; + +pub mod engine; +pub mod io; + +/// Standalone engine-core protocol emulator for frontend stress testing. +#[derive(Debug, Clone, Parser)] +#[command( + name = "vllm-mock-engine", + about = "Run a mock vLLM headless engine for Rust frontend stress testing." +)] +pub struct Opt { + /// Frontend-owned ZMQ handshake address. + #[arg(long, default_value = "tcp://127.0.0.1:29550")] + pub handshake_address: String, + + /// Number of mock engine identities to register with the frontend. + #[arg(long, default_value_t = 1)] + pub engine_count: usize, + + /// Number of accepted output tokens included in each EngineCoreOutput. + #[arg(long, default_value_t = 1)] + pub output_token_chunk_size: usize, + + /// Random token IDs are sampled uniformly from 0..vocab_size. + #[arg(long, default_value_t = 32_000)] + pub vocab_size: u32, + + /// Base seed for deterministic random token generation. + #[arg(long, default_value_t = 0)] + pub seed: u64, + + /// Log a summary line for each request. + #[arg(long)] + pub log_requests: bool, +} + +/// Run one mock engine until shutdown or transport failure. +async fn run_engine(engine_index: u32, opt: Opt, shutdown: CancellationToken) -> Result<()> { + let MockEngineSockets { data_sockets, .. } = connect_to_frontend( + &opt.handshake_address, + EngineId::from_engine_index(engine_index), + MockEngineConfig::default(), + ) + .await + .with_context(|| format!("mock engine {engine_index} failed to connect to frontend"))?; + + info!(engine_index, "mock engine connected to frontend"); + + let (input_tx, input_rx) = mpsc::unbounded_channel(); + let (output_tx, output_rx) = mpsc::channel(64); + + // IO loop: dealer -> input_tx, output_rx -> push + let mut io_loop = tokio::spawn(io::run_io_loop( + data_sockets, + input_tx, + output_rx, + shutdown.clone(), + )); + // Engine loop: input_rx -> engine logic -> output_tx + let mut engine_loop = tokio::spawn(engine::run_engine_loop( + engine_index, + opt, + input_rx, + output_tx, + shutdown.clone(), + )); + + tokio::select! { + biased; + _ = shutdown.cancelled() => { + io_loop.abort(); + engine_loop.abort(); + io_loop.await.ok(); + engine_loop.await.ok(); + } + + result = &mut io_loop => { + error!(engine_index, "mock engine IO loop exited unexpectedly"); + engine_loop.abort(); + engine_loop.await.ok(); + result??; + } + result = &mut engine_loop => { + error!(engine_index, "mock engine loop exited unexpectedly"); + io_loop.abort(); + io_loop.await.ok(); + result??; + } + } + + info!(engine_index, "mock engine shut down"); + Ok(()) +} + +/// Run all requested mock engines until cancellation or one engine task fails. +pub async fn run(opt: Opt, shutdown: CancellationToken) -> Result<()> { + info!(?opt, "starting mock engine"); + + let mut engines = JoinSet::new(); + for engine_index in 0..opt.engine_count { + engines.spawn(run_engine( + engine_index as u32, + opt.clone(), + shutdown.clone(), + )); + } + + tokio::select! { + biased; + _ = shutdown.cancelled() => { + engines.abort_all(); + while engines.join_next().await.is_some() {} + Ok(()) + } + + joined = engines.join_next() => { + match joined { + Some(Ok(Ok(()))) => bail!("mock engine exited unexpectedly"), + Some(Ok(Err(error))) => Err(error), + Some(Err(error)) => Err(error).context("mock engine task join failed"), + None => Ok(()), + } + } + } +} + +#[cfg(test)] +mod tests; diff --git a/rust/src/mock-engine/src/main.rs b/rust/src/mock-engine/src/main.rs new file mode 100644 index 00000000000..d28574fa8b5 --- /dev/null +++ b/rust/src/mock-engine/src/main.rs @@ -0,0 +1,38 @@ +use anyhow::{Context, Result}; +use clap::Parser as _; +use tokio_util::sync::CancellationToken; +use tracing::{Level, info}; +use vllm_mock_engine::Opt; + +fn init_tracing() { + tracing_subscriber::fmt().with_max_level(Level::INFO).init(); +} + +/// Create a cancellation token that is triggered by Ctrl-C. +fn shutdown_signal() -> CancellationToken { + let token = CancellationToken::new(); + let shutdown = token.clone(); + + tokio::spawn(async move { + tokio::signal::ctrl_c().await.expect("failed to install Ctrl-C signal handler"); + info!("received shutdown signal (Ctrl-C), shutting down..."); + shutdown.cancel(); + }); + + token +} + +fn main() -> Result<()> { + init_tracing(); + let opt = Opt::parse(); + + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .context("failed to build Tokio runtime")?; + + runtime.block_on(async move { + let shutdown = shutdown_signal(); + vllm_mock_engine::run(opt, shutdown).await + }) +} diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs new file mode 100644 index 00000000000..fd04761090a --- /dev/null +++ b/rust/src/mock-engine/src/tests.rs @@ -0,0 +1,197 @@ +use std::net::TcpListener; +use std::time::Duration; + +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::test_utils::IpcNamespace; +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode}; + +use crate::{Opt, run}; + +fn free_tcp_address() -> String { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind free port"); + let port = listener.local_addr().expect("local addr").port(); + drop(listener); + format!("tcp://127.0.0.1:{port}") +} + +fn client_config(handshake_address: String, engine_count: usize) -> EngineCoreClientConfig { + EngineCoreClientConfig { + transport_mode: TransportMode::HandshakeOwner { + handshake_address, + advertised_host: "127.0.0.1".to_string(), + engine_count, + ready_timeout: Duration::from_secs(5), + local_input_address: None, + local_output_address: None, + }, + coordinator_mode: None, + model_name: "mock-model".to_string(), + client_index: 0, + } +} + +async fn connect_with_mock( + handshake_address: String, + engine_count: usize, + output_token_chunk_size: usize, +) -> ( + EngineCoreClient, + CancellationToken, + tokio::task::JoinHandle>, +) { + let shutdown = CancellationToken::new(); + let task = tokio::spawn(run( + Opt { + handshake_address: handshake_address.clone(), + engine_count, + output_token_chunk_size, + vocab_size: 32_000, + seed: 0, + log_requests: false, + }, + shutdown.clone(), + )); + + let client = timeout( + Duration::from_secs(5), + EngineCoreClient::connect(client_config(handshake_address, engine_count)), + ) + .await + .expect("client connect timeout") + .expect("connect client"); + + (client, shutdown, task) +} + +fn sample_request(request_id: &str, max_tokens: u32) -> EngineCoreRequest { + EngineCoreRequest { + request_id: request_id.to_string(), + prompt_token_ids: Some(vec![1, 2, 3]), + sampling_params: Some(EngineCoreSamplingParams { + max_tokens, + ..EngineCoreSamplingParams::for_test() + }), + arrival_time: 0.0, + ..Default::default() + } +} + +async fn shutdown_mock( + client: EngineCoreClient, + shutdown: CancellationToken, + task: tokio::task::JoinHandle>, +) { + client.shutdown().await.expect("client shutdown"); + shutdown.cancel(); + task.await.expect("mock join").expect("mock run"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mock_engine_connects_over_tcp() { + let handshake_address = free_tcp_address(); + let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await; + assert_eq!(client.engine_count(), 1); + assert_eq!(client.engine_identities()[0], &[0, 0]); + assert_eq!(client.max_model_len(), Some(1024 * 1024)); + shutdown_mock(client, shutdown, task).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mock_engine_connects_over_ipc() { + let ipc = IpcNamespace::new().expect("ipc namespace"); + let handshake_address = ipc.handshake_endpoint(); + let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await; + assert_eq!(client.engine_count(), 1); + assert_eq!(client.engine_identities()[0], &[0, 0]); + shutdown_mock(client, shutdown, task).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mock_engine_registers_multiple_identities() { + let handshake_address = free_tcp_address(); + let (client, shutdown, task) = connect_with_mock(handshake_address, 2, 1).await; + assert_eq!(client.engine_count(), 2); + assert_eq!(client.engine_identities(), vec![&[0, 0][..], &[1, 0][..]]); + shutdown_mock(client, shutdown, task).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn chunk_size_one_outputs_one_token_per_update() { + let handshake_address = free_tcp_address(); + let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await; + let mut stream = client.call(sample_request("req-1", 3)).await.expect("call"); + + let first = stream.next().await.expect("first").expect("first ok"); + assert_eq!(first.new_token_ids.len(), 1); + assert_eq!(first.finish_reason, None); + let second = stream.next().await.expect("second").expect("second ok"); + assert_eq!(second.new_token_ids.len(), 1); + assert_eq!(second.finish_reason, None); + let third = stream.next().await.expect("third").expect("third ok"); + assert_eq!(third.new_token_ids.len(), 1); + assert_eq!(third.finish_reason, Some(EngineCoreFinishReason::Length)); + assert!(stream.next().await.is_none()); + + shutdown_mock(client, shutdown, task).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn chunk_size_clips_final_output_to_max_tokens() { + let handshake_address = free_tcp_address(); + let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 4).await; + let mut stream = client.call(sample_request("req-clip", 6)).await.expect("call"); + + let first = stream.next().await.expect("first").expect("first ok"); + assert_eq!(first.new_token_ids.len(), 4); + assert_eq!(first.finish_reason, None); + let second = stream.next().await.expect("second").expect("second ok"); + assert_eq!(second.new_token_ids.len(), 2); + assert_eq!(second.finish_reason, Some(EngineCoreFinishReason::Length)); + assert!(stream.next().await.is_none()); + + shutdown_mock(client, shutdown, task).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn abort_cancels_active_request_and_emits_terminal_output() { + let handshake_address = free_tcp_address(); + let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await; + let mut stream = client.call(sample_request("req-abort", 1_000_000)).await.expect("call"); + let first = stream.next().await.expect("first").expect("first ok"); + assert_eq!(first.finish_reason, None); + + client.abort(&["req-abort".to_string()]).await.expect("abort"); + + loop { + let output = timeout(Duration::from_secs(5), stream.next()) + .await + .expect("stream timeout") + .expect("terminal output") + .expect("output ok"); + if output.finish_reason.is_some() { + assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Abort)); + break; + } + } + + shutdown_mock(client, shutdown, task).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn utility_requests_return_minimal_success_responses() { + let handshake_address = free_tcp_address(); + let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await; + + assert!(!client.is_sleeping().await.expect("is sleeping")); + assert!(client.reset_prefix_cache(false, false).await.expect("reset prefix cache")); + client.reset_mm_cache().await.expect("reset mm cache"); + client.reset_encoder_cache().await.expect("reset encoder cache"); + + shutdown_mock(client, shutdown, task).await; +}