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

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 17:01 UTC

[Rust Frontend] Introduce mock engine for benchmark baseline (#43469)

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

Bugen Zhao committed May 28, 2026 at 09:40 UTC c87f62ccf8f698d4dff3b9d020edfc692443686f
19 files changed +1462 -204
rust/Cargo.lock
+38
@@ -2371,6 +2371,15 @@ version = "0.2.16"
2371 source = "registry+https://github.com/rust-lang/crates.io-index"
2372 checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981"
2373
2374 +[[package]]
2375 +name = "libmimalloc-sys"
2376 +version = "0.1.49"
2377 +source = "registry+https://github.com/rust-lang/crates.io-index"
2378 +checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9"
2379 +dependencies = [
2380 + "cc",
2381 +]
2382 +
2383 [[package]]
2384 name = "libredox"
2385 version = "0.1.14"
@@ -2569,6 +2578,15 @@ version = "0.3.3"
2578 source = "registry+https://github.com/rust-lang/crates.io-index"
2579 checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b"
2580
2581 +[[package]]
2582 +name = "mimalloc"
2583 +version = "0.1.52"
2584 +source = "registry+https://github.com/rust-lang/crates.io-index"
2585 +checksum = "2d4139bb28d14ad1facf21d5eb8825051b326e172d216b39f6d31df53cc97862"
2586 +dependencies = [
2587 + "libmimalloc-sys",
2588 +]
2589 +
2590 [[package]]
2591 name = "mime"
2592 version = "0.3.17"
@@ -5644,6 +5662,7 @@ dependencies = [
5662 "educe",
5663 "expect-test",
5664 "itertools 0.14.0",
5665 + "mimalloc",
5666 "native-tls",
5667 "serde",
5668 "serde_json",
@@ -5742,6 +5761,25 @@ dependencies = [
5761 "prometheus-client",
5762 ]
5763
5764 +[[package]]
5765 +name = "vllm-mock-engine"
5766 +version = "0.1.0"
5767 +dependencies = [
5768 + "anyhow",
5769 + "asynk-strim-attr",
5770 + "clap",
5771 + "futures",
5772 + "rand 0.9.2",
5773 + "rmpv",
5774 + "serde",
5775 + "tokio",
5776 + "tokio-util",
5777 + "tracing",
5778 + "tracing-subscriber",
5779 + "vllm-engine-core-client",
5780 + "zeromq",
5781 +]
5782 +
5783 [[package]]
5784 name = "vllm-reasoning-parser"
5785 version = "0.1.0"
rust/Cargo.toml
+3
@@ -6,6 +6,7 @@ members = [
6 "src/llm",
7 "src/managed-engine",
8 "src/metrics",
9 + "src/mock-engine",
10 "src/reasoning-parser",
11 "src/server",
12 "src/text",
@@ -45,6 +46,7 @@ http-body = "1.0.1"
46 itertools = "0.14.0"
47 libc = "0.2.177"
48 llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "5b558989844d1c7af3e43d0f604069ffd9c06320" }
49 +mimalloc = "0.1.52"
50 minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] }
51 minijinja-contrib = { version = "2.0", features = ["pycompat"] }
52 native-tls-vendored = { package = "native-tls", version = "0.2.18", features = ["vendored"] }
@@ -57,6 +59,7 @@ prometheus-client = "0.24.0"
59 prometheus-client-derive-encode = "0.5.0"
60 prost = "0.14.3"
61 prost-types = "0.14.3"
62 +rand = "0.9.2"
63 reasoning-parser = "1.2.2"
64 reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] }
65 riptoken = { version = "0.3.0", default-features = false }
rust/src/cmd/Cargo.toml
+1
@@ -17,6 +17,7 @@ anyhow.workspace = true
17 clap.workspace = true
18 educe.workspace = true
19 itertools.workspace = true
20 +mimalloc.workspace = true
21 native-tls-vendored = { workspace = true, optional = true }
22 serde.workspace = true
23 serde_json.workspace = true
rust/src/cmd/src/main.rs
+3
@@ -11,6 +11,9 @@ use vllm_managed_engine::ManagedEngineHandle;
11
12 use crate::cli::{Cli, Command};
13
14 +#[global_allocator]
15 +static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
16 +
17 const TOKIO_WORKER_THREADS_ENV: &str = "TOKIO_WORKER_THREADS";
18 const DEFAULT_MAX_TOKIO_WORKER_THREADS: usize = 32;
19
rust/src/engine-core-client/src/client/imp.rs
+9 -8
@@ -107,13 +107,13 @@ impl ClientInner {
107 Ok(registry.abortable_request_ids(request_ids))
108 }
109
110 - /// Obtain the stream sender for one output. If it indicates the request is
111 - /// finished, it will be removed from the registry.
112 - pub fn take_sender_for_output(
110 + /// Obtain stream senders for a whole engine output batch with one registry
111 + /// lock acquisition.
112 + pub fn take_senders_for_outputs<'a>(
113 &self,
114 - output: &EngineCoreOutput,
115 - ) -> Option<mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>> {
116 - self.request_reg.lock().sender_for_output(output)
114 + outputs: impl IntoIterator<Item = &'a EngineCoreOutput>,
115 + ) -> Vec<Option<mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>>> {
116 + self.request_reg.lock().senders_for_outputs(outputs)
117 }
118
119 /// Remove a batch of requests that have finished or aborted, returning
@@ -301,9 +301,10 @@ pub(crate) async fn run_output_dispatcher_loop(
301
302 match outputs.classify() {
303 ClassifiedEngineCoreOutputs::RequestBatch(batch) => {
304 - for output in batch.outputs {
304 + let senders = inner.take_senders_for_outputs(&batch.outputs);
305 + for (output, sender) in batch.outputs.into_iter().zip(senders) {
306 let request_id = output.request_id.clone();
306 - let Some(sender) = inner.take_sender_for_output(&output) else {
307 + let Some(sender) = sender else {
308 debug!(request_id, "dropping output for inactive request");
309 continue;
310 };
rust/src/engine-core-client/src/client/state.rs
+12 -3
@@ -1,4 +1,4 @@
1 -use std::collections::BTreeMap;
1 +use std::collections::{BTreeMap, HashMap};
2 use std::sync::atomic::{AtomicU64, Ordering};
3
4 use tokio::sync::{mpsc, oneshot};
@@ -80,7 +80,7 @@ impl EngineRoutingState {
80 #[derive(Debug)]
81 pub struct RequestRegistry {
82 closed: bool,
83 - requests: BTreeMap<String, TrackedRequest>,
83 + requests: HashMap<String, TrackedRequest>,
84 routing_per_engine: BTreeMap<EngineId, EngineRoutingState>,
85 }
86
@@ -88,7 +88,7 @@ impl RequestRegistry {
88 pub fn new(engines: &[ConnectedEngine]) -> Self {
89 Self {
90 closed: false,
91 - requests: BTreeMap::default(),
91 + requests: HashMap::default(),
92 routing_per_engine: engines
93 .iter()
94 .map(|engine| (engine.engine_id.clone(), EngineRoutingState::default()))
@@ -180,6 +180,15 @@ impl RequestRegistry {
180 }
181 }
182
183 + /// Obtain stream senders for a whole engine output batch under one
184 + /// registry lock. Finished outputs are removed before returning.
185 + pub fn senders_for_outputs<'a>(
186 + &mut self,
187 + outputs: impl IntoIterator<Item = &'a EngineCoreOutput>,
188 + ) -> Vec<Option<OutputSender>> {
189 + outputs.into_iter().map(|output| self.sender_for_output(output)).collect()
190 + }
191 +
192 /// Remove a batch of requests that have finished or aborted, returning
193 /// their stream senders.
194 pub fn finish_many<'a>(
rust/src/engine-core-client/src/lib.rs
+1
@@ -2,6 +2,7 @@ mod client;
2 mod coordinator;
3 mod error;
4 mod metrics;
5 +pub mod mock_engine;
6 pub mod protocol;
7 #[cfg(any(test, feature = "test-util"))]
8 pub mod test_utils;
rust/src/engine-core-client/src/mock_engine.rs new
+265
@@ -0,0 +1,265 @@
1 +use std::path::Path;
2 +use std::time::Duration;
3 +
4 +use tokio::time::timeout;
5 +use zeromq::prelude::{Socket, SocketRecv, SocketSend};
6 +use zeromq::util::PeerIdentity;
7 +use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage};
8 +
9 +use crate::EngineId;
10 +use crate::error::{Error, Result, bail_unexpected_handshake_message};
11 +use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};
12 +use crate::protocol::{ModelDtype, decode_msgpack, encode_msgpack};
13 +
14 +/// Default model length advertised by reusable mock engine helpers.
15 +pub const DEFAULT_MOCK_MAX_MODEL_LEN: u64 = 1024 * 1024;
16 +/// Default KV block count advertised by reusable mock engine helpers.
17 +pub const DEFAULT_MOCK_NUM_GPU_BLOCKS: u64 = 0;
18 +
19 +/// Startup behavior for one mock engine joining a frontend.
20 +#[derive(Debug, Clone)]
21 +pub struct MockEngineConfig {
22 + /// Whether the engine should advertise itself as local to the frontend.
23 + pub local: bool,
24 + /// Whether the engine should advertise itself as headless.
25 + pub headless: bool,
26 + /// Engine-ready payload reported after INIT, including max model length,
27 + /// KV block count, and dtype.
28 + pub ready_response: EngineCoreReadyResponse,
29 + /// Maximum time to wait for IPC endpoints to appear before connecting.
30 + pub connect_timeout: Duration,
31 +}
32 +
33 +impl Default for MockEngineConfig {
34 + fn default() -> Self {
35 + Self {
36 + local: false,
37 + headless: true,
38 + ready_response: default_ready_response(),
39 + connect_timeout: Duration::from_secs(5),
40 + }
41 + }
42 +}
43 +
44 +/// Construct the ready response used by the standalone mock engine CLI.
45 +pub fn default_ready_response() -> EngineCoreReadyResponse {
46 + EngineCoreReadyResponse {
47 + max_model_len: DEFAULT_MOCK_MAX_MODEL_LEN,
48 + num_gpu_blocks: DEFAULT_MOCK_NUM_GPU_BLOCKS,
49 + dp_stats_address: None,
50 + dtype: Some(ModelDtype::Float32),
51 + }
52 +}
53 +
54 +/// Coordinator-side sockets used by one mock engine when coordinator mode
55 +/// is enabled.
56 +pub struct MockCoordinatorSockets {
57 + /// Subscription socket that receives coordinator broadcasts such as
58 + /// `START_DP_WAVE`.
59 + pub input_sub: SubSocket,
60 + /// Push socket used to send coordinator-only `EngineCoreOutputs` back to
61 + /// the frontend.
62 + pub output_push: PushSocket,
63 +}
64 +
65 +/// One mock engine's connection to one frontend client.
66 +///
67 +/// vLLM launches one engine-client pair per API server process. A remote
68 +/// engine connects to every advertised input/output pair and uses the request's
69 +/// `client_index` to route outputs back to the originating API server.
70 +pub struct MockEngineDataSockets {
71 + /// Socket used to receive frontend requests.
72 + pub dealer: DealerSocket,
73 + /// Socket used to publish normal request outputs back to the frontend.
74 + pub push: PushSocket,
75 +}
76 +
77 +/// Frontend-facing sockets owned by one mock engine.
78 +pub struct MockEngineSockets {
79 + /// Decoded INIT message sent by the frontend during handshake.
80 + pub init: HandshakeInitMessage,
81 + /// Data sockets for all frontend clients in client-index order.
82 + ///
83 + /// For Rust frontend this will always be one socket, while for Python frontend
84 + /// this may be multiple sockets if there are multiple API server processes.
85 + pub data_sockets: Vec<MockEngineDataSockets>,
86 + /// Optional coordinator sockets when the client enabled the in-process
87 + /// coordinator.
88 + pub coordinator: Option<MockCoordinatorSockets>,
89 +}
90 +
91 +/// Build a HELLO or READY handshake status payload.
92 +fn ready_message(status: &str, config: &MockEngineConfig) -> ReadyMessage {
93 + ReadyMessage {
94 + status: Some(status.to_string()),
95 + local: Some(config.local),
96 + headless: Some(config.headless),
97 + parallel_config_hash: None,
98 + }
99 +}
100 +
101 +/// Convert an engine id into a ZMQ DEALER identity.
102 +fn peer_identity(engine_id: impl Into<EngineId>) -> Result<PeerIdentity> {
103 + let engine_id = engine_id.into();
104 + PeerIdentity::try_from(engine_id.clone()).map_err(|error| Error::UnexpectedHandshakeMessage {
105 + message: format!(
106 + "invalid mock engine identity {:?}: {error}",
107 + engine_id.to_vec()
108 + ),
109 + })
110 +}
111 +
112 +/// Wait for an IPC endpoint path to appear before attempting to connect.
113 +async fn wait_for_ipc_endpoint(endpoint: &str, connect_timeout: Duration) -> Result<()> {
114 + let Some(socket_path) = endpoint.strip_prefix("ipc://") else {
115 + return Ok(());
116 + };
117 +
118 + timeout(connect_timeout, async {
119 + while !Path::new(socket_path).exists() {
120 + tokio::time::sleep(Duration::from_millis(20)).await;
121 + }
122 + })
123 + .await
124 + .map_err(|_| Error::HandshakeTimeout {
125 + stage: "mock engine IPC endpoint",
126 + timeout: connect_timeout,
127 + })
128 +}
129 +
130 +/// Encode the engine-ready response sent on input socket registration.
131 +fn ready_response_payload(config: &MockEngineConfig) -> Result<Vec<u8>> {
132 + encode_msgpack(&config.ready_response)
133 +}
134 +
135 +/// Join a frontend-owned handshake endpoint and open mock engine sockets.
136 +pub async fn connect_to_frontend(
137 + engine_handshake: impl AsRef<str>,
138 + engine_id: impl Into<EngineId>,
139 + config: MockEngineConfig,
140 +) -> Result<MockEngineSockets> {
141 + let engine_handshake = engine_handshake.as_ref();
142 + wait_for_ipc_endpoint(engine_handshake, config.connect_timeout).await?;
143 +
144 + let peer_identity = peer_identity(engine_id)?;
145 + let mut options = SocketOptions::default();
146 + options.peer_identity(peer_identity.clone());
147 + let mut handshake = DealerSocket::with_options(options);
148 + handshake.connect(engine_handshake).await?;
149 + handshake
150 + .send(ZmqMessage::from(encode_msgpack(&ready_message(
151 + "HELLO", &config,
152 + ))?))
153 + .await?;
154 +
155 + let init_frames = handshake.recv().await?.into_vec();
156 + if init_frames.len() != 1 {
157 + bail_unexpected_handshake_message!(
158 + "expected one INIT frame from frontend, got {}",
159 + init_frames.len()
160 + );
161 + }
162 + let init: HandshakeInitMessage = decode_msgpack(init_frames[0].as_ref())?;
163 +
164 + if init.addresses.inputs.is_empty() {
165 + return Err(Error::UnexpectedHandshakeMessage {
166 + message: "frontend INIT did not include an input address".to_string(),
167 + });
168 + }
169 + if init.addresses.inputs.len() != init.addresses.outputs.len() {
170 + return Err(Error::UnexpectedHandshakeMessage {
171 + message: format!(
172 + "frontend INIT input/output address count mismatch: {} inputs, {} outputs",
173 + init.addresses.inputs.len(),
174 + init.addresses.outputs.len()
175 + ),
176 + });
177 + }
178 +
179 + let mut data_sockets = Vec::with_capacity(init.addresses.inputs.len());
180 + for (input_address, output_address) in
181 + init.addresses.inputs.iter().zip(init.addresses.outputs.iter())
182 + {
183 + wait_for_ipc_endpoint(input_address, config.connect_timeout).await?;
184 + wait_for_ipc_endpoint(output_address, config.connect_timeout).await?;
185 +
186 + let mut input_options = SocketOptions::default();
187 + input_options.peer_identity(peer_identity.clone());
188 + let mut dealer = DealerSocket::with_options(input_options);
189 + dealer.connect(input_address).await?;
190 + dealer.send(ZmqMessage::from(ready_response_payload(&config)?)).await?;
191 +
192 + let mut push = PushSocket::new();
193 + push.connect(output_address).await?;
194 +
195 + data_sockets.push(MockEngineDataSockets { dealer, push });
196 + }
197 +
198 + let coordinator = match (
199 + init.addresses.coordinator_input.as_deref(),
200 + init.addresses.coordinator_output.as_deref(),
201 + ) {
202 + (Some(coordinator_input), Some(coordinator_output)) => {
203 + let mut input_sub = SubSocket::new();
204 + input_sub.connect(coordinator_input).await?;
205 + input_sub.subscribe("").await?;
206 +
207 + let mut output_push = PushSocket::new();
208 + output_push.connect(coordinator_output).await?;
209 +
210 + let ready = input_sub.recv().await?.into_vec();
211 + if ready.len() != 1 || ready[0].as_ref() != b"READY" {
212 + bail_unexpected_handshake_message!(
213 + "expected coordinator READY marker, got {:?}",
214 + ready
215 + );
216 + }
217 +
218 + Some(MockCoordinatorSockets {
219 + input_sub,
220 + output_push,
221 + })
222 + }
223 + (None, None) => None,
224 + _ => bail_unexpected_handshake_message!(
225 + "coordinator handshake addresses must be both present or both absent"
226 + ),
227 + };
228 +
229 + handshake
230 + .send(ZmqMessage::from(encode_msgpack(&ready_message(
231 + "READY", &config,
232 + ))?))
233 + .await?;
234 +
235 + Ok(MockEngineSockets {
236 + init,
237 + data_sockets,
238 + coordinator,
239 + })
240 +}
241 +
242 +/// Join already-bootstrapped frontend input/output sockets directly.
243 +pub async fn connect_to_bootstrapped_frontend(
244 + input_address: impl AsRef<str>,
245 + output_address: impl AsRef<str>,
246 + engine_id: impl Into<EngineId>,
247 + config: MockEngineConfig,
248 +) -> Result<(DealerSocket, PushSocket)> {
249 + let input_address = input_address.as_ref();
250 + let output_address = output_address.as_ref();
251 + wait_for_ipc_endpoint(input_address, config.connect_timeout).await?;
252 + wait_for_ipc_endpoint(output_address, config.connect_timeout).await?;
253 +
254 + let peer_identity = peer_identity(engine_id)?;
255 + let mut input_options = SocketOptions::default();
256 + input_options.peer_identity(peer_identity);
257 + let mut dealer = DealerSocket::with_options(input_options);
258 + dealer.connect(input_address).await?;
259 + dealer.send(ZmqMessage::from(ready_response_payload(&config)?)).await?;
260 +
261 + let mut push = PushSocket::new();
262 + push.connect(output_address).await?;
263 +
264 + Ok((dealer, push))
265 +}
rust/src/engine-core-client/src/protocol/mod.rs
+31
@@ -36,6 +36,14 @@ fn is_false(v: &bool) -> bool {
36 !v
37 }
38
39 +fn default_top_p() -> f32 {
40 + 1.0
41 +}
42 +
43 +fn default_repetition_penalty() -> f32 {
44 + 1.0
45 +}
46 +
47 mod classified_outputs;
48 pub mod dtype;
49 pub mod handshake;
@@ -65,6 +73,24 @@ pub enum EngineCoreRequestType {
73 }
74
75 impl EngineCoreRequestType {
76 + /// Decode the single-byte request type frame used on the engine input
77 + /// socket. Returns `None` for unrecognized values.
78 + pub fn from_frame(frame: &[u8]) -> Option<Self> {
79 + let [value] = frame else {
80 + return None;
81 + };
82 +
83 + match value {
84 + 0 => Some(Self::Add),
85 + 1 => Some(Self::Abort),
86 + 2 => Some(Self::StartDpWave),
87 + 3 => Some(Self::Utility),
88 + _ => None,
89 + }
90 + }
91 +
92 + /// Encode the request type as the single-byte frame used on the engine
93 + /// input socket.
94 pub fn to_frame(self) -> Bytes {
95 Bytes::from_static(match self {
96 Self::Add => b"\x00",
@@ -200,14 +226,17 @@ pub struct EngineCoreSamplingParams {
226 /// greedy sampling.
227 pub temperature: f32,
228 /// Cumulative probability threshold for nucleus sampling.
229 + #[serde(default = "default_top_p")]
230 pub top_p: f32,
231 /// Maximum number of top tokens to consider. `0` means all tokens.
232 + #[serde(default)]
233 pub top_k: u32,
234 /// Random seed used by the sampler when present.
235 pub seed: Option<i64>,
236 /// Maximum number of tokens to generate per output sequence.
237 pub max_tokens: u32,
238 /// Minimum number of tokens to generate before EOS or stop-token handling.
239 + #[serde(default)]
240 pub min_tokens: u32,
241 /// Number of log probabilities to return per generated token.
242 ///
@@ -218,12 +247,14 @@ pub struct EngineCoreSamplingParams {
247 /// `None` disables prompt logprobs. `-1` requests the full vocabulary.
248 pub prompt_logprobs: Option<i32>,
249 /// Minimum probability threshold for token sampling.
250 + #[serde(default)]
251 pub min_p: f32,
252 /// Frequency penalty applied by the sampler.
253 pub frequency_penalty: f32,
254 /// Presence penalty applied by the sampler.
255 pub presence_penalty: f32,
256 /// Repetition penalty applied by the sampler.
257 + #[serde(default = "default_repetition_penalty")]
258 pub repetition_penalty: f32,
259 /// Token IDs that stop generation.
260 pub stop_token_ids: Vec<u32>,
rust/src/engine-core-client/src/protocol/utility.rs
+1 -1
@@ -102,7 +102,7 @@ impl<'de> Deserialize<'de> for UtilityCallId {
102 ///
103 /// Original Python payload shape:
104 /// `(client_index, call_id, method_name, args)`
105 -#[derive(Debug, Clone, PartialEq, Serialize_tuple)]
105 +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
106 pub struct EngineCoreUtilityRequest {
107 pub client_index: u32,
108 pub call_id: UtilityCallId,
rust/src/engine-core-client/src/test_utils.rs
+34 -175
@@ -1,17 +1,18 @@
1 use std::future::Future;
2 use std::path::Path;
3 use std::pin::Pin;
4 -use std::time::Duration;
4
5 use tempfile::TempDir;
6 use tokio::sync::oneshot;
8 -use zeromq::prelude::{Socket, SocketRecv, SocketSend};
9 -use zeromq::util::PeerIdentity;
10 -use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage};
7 +use zeromq::{DealerSocket, PushSocket};
8
9 use crate::EngineId;
10 +pub use crate::mock_engine::{MockCoordinatorSockets, MockEngineSockets};
11 +use crate::mock_engine::{
12 + MockEngineConfig, MockEngineDataSockets, connect_to_bootstrapped_frontend, connect_to_frontend,
13 +};
14 use crate::protocol::ModelDtype;
14 -use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};
15 +use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage};
16
17 /// Per-test IPC endpoint namespace backed by a unique temporary directory.
18 ///
@@ -52,156 +53,29 @@ impl IpcNamespace {
53 }
54 }
55
55 -/// Construct a standard local READY message used by mock engines in tests.
56 -fn ready_message(status: &str) -> ReadyMessage {
57 - ReadyMessage {
58 - status: Some(status.to_string()),
59 - local: Some(true),
60 - headless: Some(true),
61 - parallel_config_hash: None,
56 +fn test_mock_engine_config() -> MockEngineConfig {
57 + MockEngineConfig {
58 + local: true,
59 + headless: true,
60 + ready_response: EngineCoreReadyResponse {
61 + max_model_len: 4096,
62 + num_gpu_blocks: 0,
63 + dp_stats_address: None,
64 + dtype: Some(ModelDtype::Float32),
65 + },
66 + ..Default::default()
67 }
68 }
69
65 -/// Construct a default ready response payload for mock engine input
66 -/// registration.
67 -fn ready_response_payload() -> Vec<u8> {
68 - rmp_serde::to_vec_named(&EngineCoreReadyResponse {
69 - max_model_len: 4096,
70 - num_gpu_blocks: 0,
71 - dp_stats_address: None,
72 - dtype: Some(ModelDtype::Float32),
73 - })
74 - .expect("encode ready response payload")
75 -}
76 -
77 -/// Coordinator-side sockets connected by one mock engine when coordinator mode
78 -/// is enabled.
79 -pub struct MockCoordinatorConnections {
80 - /// Subscription socket that receives coordinator broadcasts such as
81 - /// `START_DP_WAVE`.
82 - pub input_sub: SubSocket,
83 - /// Push socket used to send coordinator-only `EngineCoreOutputs` back to
84 - /// the frontend.
85 - pub output_push: PushSocket,
86 -}
87 -
88 -/// Fully connected mock engine transport state used by tests.
89 -pub struct MockEngineConnections {
90 - /// Decoded INIT message sent by the frontend during handshake.
91 - pub init: HandshakeInitMessage,
92 - /// Socket used to receive frontend requests.
93 - pub dealer: DealerSocket,
94 - /// Socket used to publish normal request outputs back to the frontend.
95 - pub push: PushSocket,
96 - /// Optional coordinator sockets when the client enabled the in-process
97 - /// coordinator.
98 - pub coordinator: Option<MockCoordinatorConnections>,
99 -}
100 -
70 /// Complete the engine-core handshake and connect mock input/output sockets
71 /// plus optional coordinator sockets.
103 -pub async fn setup_mock_engine_connections(
72 +pub async fn setup_mock_engine_sockets(
73 engine_handshake: String,
74 engine_id: impl Into<EngineId>,
106 -) -> MockEngineConnections {
107 - // Wait for the client to bind the handshake socket before connecting.
108 - // A fixed sleep is racy under CI load; instead poll for the socket file.
109 - let socket_path = engine_handshake
110 - .strip_prefix("ipc://")
111 - .expect("handshake address must be ipc://");
112 - for _ in 0..100 {
113 - if Path::new(socket_path).exists() {
114 - break;
115 - }
116 - tokio::time::sleep(Duration::from_millis(20)).await;
117 - }
118 -
119 - let peer_identity = PeerIdentity::try_from(engine_id.into()).expect("peer id");
120 -
121 - let mut options = SocketOptions::default();
122 - options.peer_identity(peer_identity.clone());
123 - let mut handshake = DealerSocket::with_options(options);
124 - handshake
125 - .connect(&engine_handshake)
126 - .await
127 - .expect("connect mock engine handshake socket");
128 - handshake
129 - .send(ZmqMessage::from(
130 - rmp_serde::to_vec_named(&ready_message("HELLO")).expect("encode HELLO ready message"),
131 - ))
132 - .await
133 - .expect("send HELLO ready message");
134 -
135 - let init_frames = handshake.recv().await.expect("receive handshake init message").into_vec();
136 - assert_eq!(init_frames.len(), 1);
137 - let init: HandshakeInitMessage =
138 - rmp_serde::from_slice(init_frames[0].as_ref()).expect("decode handshake init message");
139 -
140 - let mut input_options = SocketOptions::default();
141 - input_options.peer_identity(peer_identity);
142 - let mut dealer = DealerSocket::with_options(input_options);
143 - dealer
144 - .connect(&init.addresses.inputs[0])
145 - .await
146 - .expect("connect mock engine input socket");
147 - dealer
148 - .send(ZmqMessage::from(ready_response_payload()))
149 - .await
150 - .expect("send mock engine input ready frame");
151 -
152 - let mut push = PushSocket::new();
153 - push.connect(&init.addresses.outputs[0])
154 - .await
155 - .expect("connect mock engine output socket");
156 -
157 - let coordinator = match (
158 - init.addresses.coordinator_input.as_deref(),
159 - init.addresses.coordinator_output.as_deref(),
160 - ) {
161 - (Some(coordinator_input), Some(coordinator_output)) => {
162 - let mut input_sub = SubSocket::new();
163 - input_sub
164 - .connect(coordinator_input)
165 - .await
166 - .expect("connect mock engine coordinator input socket");
167 - input_sub
168 - .subscribe("")
169 - .await
170 - .expect("subscribe mock engine coordinator input socket");
171 -
172 - let mut output_push = PushSocket::new();
173 - output_push
174 - .connect(coordinator_output)
175 - .await
176 - .expect("connect mock engine coordinator output socket");
177 -
178 - let ready =
179 - input_sub.recv().await.expect("receive coordinator READY marker").into_vec();
180 - assert_eq!(ready.len(), 1);
181 - assert_eq!(ready[0].as_ref(), b"READY");
182 -
183 - Some(MockCoordinatorConnections {
184 - input_sub,
185 - output_push,
186 - })
187 - }
188 - (None, None) => None,
189 - _ => panic!("coordinator handshake addresses must be both present or both absent"),
190 - };
191 -
192 - handshake
193 - .send(ZmqMessage::from(
194 - rmp_serde::to_vec_named(&ready_message("READY")).expect("encode READY ready message"),
195 - ))
75 +) -> MockEngineSockets {
76 + connect_to_frontend(engine_handshake, engine_id, test_mock_engine_config())
77 .await
197 - .expect("send READY ready message");
198 -
199 - MockEngineConnections {
200 - init,
201 - dealer,
202 - push,
203 - coordinator,
204 - }
78 + .expect("connect mock engine")
79 }
80
81 /// Connect one mock engine directly to already-bootstrapped frontend
@@ -211,31 +85,14 @@ pub async fn setup_bootstrapped_mock_engine(
85 output_address: String,
86 engine_id: impl Into<EngineId>,
87 ) -> (DealerSocket, PushSocket) {
214 - for endpoint in [&input_address, &output_address] {
215 - if let Some(socket_path) = endpoint.strip_prefix("ipc://") {
216 - for _ in 0..100 {
217 - if Path::new(socket_path).exists() {
218 - break;
219 - }
220 - tokio::time::sleep(Duration::from_millis(20)).await;
221 - }
222 - }
223 - }
224 -
225 - let peer_identity = PeerIdentity::try_from(engine_id.into()).expect("peer id");
226 - let mut input_options = SocketOptions::default();
227 - input_options.peer_identity(peer_identity);
228 - let mut dealer = DealerSocket::with_options(input_options);
229 - dealer.connect(&input_address).await.expect("connect mock engine input socket");
230 - dealer
231 - .send(ZmqMessage::from(ready_response_payload()))
232 - .await
233 - .expect("send mock engine input ready frame");
234 -
235 - let mut push = PushSocket::new();
236 - push.connect(&output_address).await.expect("connect mock engine output socket");
237 -
238 - (dealer, push)
88 + connect_to_bootstrapped_frontend(
89 + input_address,
90 + output_address,
91 + engine_id,
92 + test_mock_engine_config(),
93 + )
94 + .await
95 + .expect("connect bootstrapped mock engine")
96 }
97
98 /// Complete the engine-core handshake and connect mock input/output sockets.
@@ -247,9 +104,11 @@ pub async fn setup_mock_engine_with_init(
104 engine_handshake: String,
105 engine_id: impl Into<EngineId>,
106 ) -> (HandshakeInitMessage, DealerSocket, PushSocket) {
250 - let MockEngineConnections {
251 - init, dealer, push, ..
252 - } = setup_mock_engine_connections(engine_handshake, engine_id).await;
107 + let MockEngineSockets {
108 + init, data_sockets, ..
109 + } = setup_mock_engine_sockets(engine_handshake, engine_id).await;
110 + let MockEngineDataSockets { dealer, push } =
111 + data_sockets.into_iter().next().expect("mock engine data socket");
112 (init, dealer, push)
113 }
114
rust/src/engine-core-client/src/tests/client.rs
+20 -17
@@ -30,7 +30,7 @@ use crate::protocol::{
30 EngineCoreRequestType, EngineCoreSamplingParams, decode_engine_core_outputs,
31 };
32 use crate::test_utils::{
33 - IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_connections,
33 + IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets,
34 setup_mock_engine_with_init, spawn_mock_engine_task,
35 };
36 use crate::{
@@ -477,8 +477,8 @@ async fn coordinator_handshake_includes_engine_control_addresses() {
477 let (init_tx, init_rx) = oneshot::channel();
478 let (shutdown_tx, shutdown_rx) = oneshot::channel();
479 let engine_task = tokio::spawn(async move {
480 - let connections = setup_mock_engine_connections(handshake_address, &engine_id).await;
481 - let _ = init_tx.send(connections.init.clone());
480 + let sockets = setup_mock_engine_sockets(handshake_address, &engine_id).await;
481 + let _ = init_tx.send(sockets.init.clone());
482 let _ = shutdown_rx.await;
483 });
484
@@ -515,14 +515,15 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
515 let engine0_task = tokio::spawn({
516 let handshake_address = handshake_address.clone();
517 async move {
518 - let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await;
518 + let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await;
519 let mut coordinator =
520 engine.coordinator.take().expect("coordinator sockets should be present");
521 + let data_socket = engine.data_sockets.first_mut().expect("data socket");
522
523 let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
524 assert_eq!((wave, exclude_engine), (0, 0));
525
525 - let add = recv_engine_message(&mut engine.dealer).await;
526 + let add = recv_engine_message(&mut data_socket.dealer).await;
527 assert_eq!(add[0].as_ref(), &[0x00]);
528 let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
529 assert_eq!(request.request_id, "req-1");
@@ -538,7 +539,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
539 );
540
541 send_outputs(
541 - &mut engine.push,
542 + &mut data_socket.push,
543 EngineCoreOutputs {
544 engine_index: 0,
545 outputs: vec![request_output(
@@ -565,14 +566,14 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
566 let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
567 assert_eq!((wave, exclude_engine), (1, 0));
568
568 - let add = recv_engine_message(&mut engine.dealer).await;
569 + let add = recv_engine_message(&mut data_socket.dealer).await;
570 assert_eq!(add[0].as_ref(), &[0x00]);
571 let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
572 assert_eq!(request.request_id, "req-3");
573 assert_eq!(request.current_wave, 1);
574
575 send_outputs(
575 - &mut engine.push,
576 + &mut data_socket.push,
577 EngineCoreOutputs {
578 engine_index: 0,
579 outputs: vec![request_output(
@@ -594,14 +595,15 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
595 let engine1_task = tokio::spawn({
596 let handshake_address = handshake_address.clone();
597 async move {
597 - let mut engine = setup_mock_engine_connections(handshake_address, &[0x01, 0x00]).await;
598 + let mut engine = setup_mock_engine_sockets(handshake_address, &[0x01, 0x00]).await;
599 let mut coordinator =
600 engine.coordinator.take().expect("coordinator sockets should be present");
601 + let data_socket = engine.data_sockets.first_mut().expect("data socket");
602
603 let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
604 assert_eq!((wave, exclude_engine), (0, 0));
605
604 - let add = recv_engine_message(&mut engine.dealer).await;
606 + let add = recv_engine_message(&mut data_socket.dealer).await;
607 assert_eq!(add[0].as_ref(), &[0x00]);
608 let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
609 assert_eq!(request.request_id, "req-2");
@@ -617,7 +619,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
619 );
620
621 send_outputs(
620 - &mut engine.push,
622 + &mut data_socket.push,
623 EngineCoreOutputs {
624 engine_index: 1,
625 outputs: vec![request_output(
@@ -637,7 +639,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
639 assert!(
640 timeout(
641 Duration::from_millis(200),
640 - recv_engine_message(&mut engine.dealer)
642 + recv_engine_message(&mut data_socket.dealer)
643 )
644 .await
645 .is_err()
@@ -712,7 +714,7 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() {
714 let engine0_task = tokio::spawn({
715 let handshake_address = handshake_address.clone();
716 async move {
715 - let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await;
717 + let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await;
718 let mut coordinator =
719 engine.coordinator.take().expect("coordinator sockets should be present");
720
@@ -727,7 +729,7 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() {
729 let engine1_task = tokio::spawn({
730 let handshake_address = handshake_address.clone();
731 async move {
730 - let mut engine = setup_mock_engine_connections(handshake_address, &[0x01, 0x00]).await;
732 + let mut engine = setup_mock_engine_sockets(handshake_address, &[0x01, 0x00]).await;
733 let mut coordinator =
734 engine.coordinator.take().expect("coordinator sockets should be present");
735
@@ -778,9 +780,10 @@ async fn coordinator_accepts_stats_only_outputs() {
780
781 let (shutdown_tx, shutdown_rx) = oneshot::channel();
782 let engine_task = tokio::spawn(async move {
781 - let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await;
783 + let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await;
784 let mut coordinator =
785 engine.coordinator.take().expect("coordinator sockets should be present");
786 + let data_socket = engine.data_sockets.first_mut().expect("data socket");
787
788 let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
789 assert_eq!((wave, exclude_engine), (0, 0));
@@ -799,13 +802,13 @@ async fn coordinator_accepts_stats_only_outputs() {
802 )
803 .await;
804
802 - let add = recv_engine_message(&mut engine.dealer).await;
805 + let add = recv_engine_message(&mut data_socket.dealer).await;
806 assert_eq!(add[0].as_ref(), &[0x00]);
807 let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
808 assert_eq!(request.request_id, "req-stats");
809
810 send_outputs(
808 - &mut engine.push,
811 + &mut data_socket.push,
812 EngineCoreOutputs {
813 engine_index: 0,
814 outputs: vec![request_output(
rust/src/mock-engine/Cargo.toml new
+30
@@ -0,0 +1,30 @@
1 +[package]
2 +name = "vllm-mock-engine"
3 +version.workspace = true
4 +edition.workspace = true
5 +license.workspace = true
6 +
7 +[[bin]]
8 +name = "vllm-mock-engine"
9 +path = "src/main.rs"
10 +
11 +[dependencies]
12 +anyhow.workspace = true
13 +asynk-strim-attr.workspace = true
14 +clap.workspace = true
15 +futures.workspace = true
16 +rand.workspace = true
17 +rmpv.workspace = true
18 +serde.workspace = true
19 +tokio = { workspace = true, features = ["signal"] }
20 +tokio-util.workspace = true
21 +tracing.workspace = true
22 +tracing-subscriber.workspace = true
23 +vllm-engine-core-client.workspace = true
24 +zeromq.workspace = true
25 +
26 +[dev-dependencies]
27 +vllm-engine-core-client = { workspace = true, features = ["test-util"] }
28 +
29 +[lints]
30 +workspace = true
rust/src/mock-engine/README.md new
+104
@@ -0,0 +1,104 @@
1 +# vLLM Mock Engine
2 +
3 +`vllm-mock-engine` is a small engine-side process for frontend stress testing. It
4 +joins a frontend-owned startup handshake, reports a large ready response, treats
5 +prefill as instant, and emits random decode tokens until each request reaches
6 +its `max_tokens`.
7 +
8 +The frontend must own the handshake socket. Start the frontend first, then start
9 +the mock engine with the same handshake address.
10 +
11 +## Start the mock engine
12 +
13 +```bash
14 +cargo run -p vllm-mock-engine -- \
15 + --handshake-address tcp://127.0.0.1:29550 \
16 + --engine-count 1 \
17 + --output-token-chunk-size 1 \
18 + --vocab-size 32000 \
19 + --seed 0 \
20 + --log-requests
21 +```
22 +
23 +Useful knobs:
24 +
25 +- `--engine-count` must match the frontend's expected data-parallel engine
26 + count.
27 +- `--output-token-chunk-size` controls how many token IDs appear in one
28 + `EngineCoreOutput`; values greater than 1 are useful for MTP/spec-decode
29 + shaped frontend tests.
30 +- `--vocab-size` should stay within the tokenizer vocabulary of the model used
31 + by the frontend.
32 +
33 +Stop it with Ctrl-C.
34 +
35 +## Rust Frontend
36 +
37 +Terminal 1:
38 +
39 +```bash
40 +cargo run --bin vllm-rs -- serve Qwen/Qwen3-0.6B \
41 + --data-parallel-size 1 \
42 + --data-parallel-size-local 0 \
43 + --handshake-port 29550
44 +```
45 +
46 +Terminal 2:
47 +
48 +```bash
49 +cargo run -p vllm-mock-engine -- \
50 + --handshake-address tcp://127.0.0.1:29550
51 +```
52 +
53 +For multiple mock engines, set both sides to the same count:
54 +
55 +```bash
56 +cargo run --bin vllm-rs -- serve Qwen/Qwen3-0.6B \
57 + --data-parallel-size 4 \
58 + --data-parallel-size-local 0 \
59 + --handshake-port 29550
60 +
61 +cargo run -p vllm-mock-engine -- \
62 + --handshake-address tcp://127.0.0.1:29550 \
63 + --engine-count 4
64 +```
65 +
66 +## Python Frontend
67 +
68 +Use `vllm serve` with `--data-parallel-size-local 0` so the Python process runs
69 +as a frontend/API server and waits for external engines on
70 +`--data-parallel-rpc-port`.
71 +
72 +Terminal 1:
73 +
74 +```bash
75 +vllm serve Qwen/Qwen3-0.6B \
76 + --data-parallel-address 127.0.0.1 \
77 + --data-parallel-rpc-port 29550 \
78 + --data-parallel-size 1 \
79 + --data-parallel-size-local 0
80 +```
81 +
82 +Terminal 2:
83 +
84 +```bash
85 +cargo run -p vllm-mock-engine -- \
86 + --handshake-address tcp://127.0.0.1:29550
87 +```
88 +
89 +## Smoke Request
90 +
91 +After either frontend is ready:
92 +
93 +```bash
94 +curl http://127.0.0.1:8000/v1/chat/completions \
95 + -H "Content-Type: application/json" \
96 + -d '{
97 + "model": "Qwen/Qwen3-0.6B",
98 + "messages": [{"role": "user", "content": "hello"}],
99 + "max_tokens": 16,
100 + "stream": true
101 + }'
102 +```
103 +
104 +Always pass `max_tokens`; the mock engine stops by length.
rust/src/mock-engine/src/engine.rs new
+416
@@ -0,0 +1,416 @@
1 +use std::collections::{BTreeMap, BTreeSet, HashMap};
2 +use std::hash::{Hash as _, Hasher as _};
3 +use std::time::{SystemTime, UNIX_EPOCH};
4 +
5 +use anyhow::{Result, anyhow};
6 +use rand::rngs::StdRng;
7 +use rand::{Rng as _, SeedableRng as _};
8 +use rmpv::Value;
9 +use serde::Serialize;
10 +use tokio::sync::mpsc;
11 +use tokio::task::yield_now;
12 +use tokio_util::sync::CancellationToken;
13 +use tracing::{debug, info, warn};
14 +use vllm_engine_core_client::protocol::utility::{
15 + EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope,
16 +};
17 +use vllm_engine_core_client::protocol::{
18 + EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
19 +};
20 +
21 +use super::Opt;
22 +
23 +/// Derive a stable per-request seed from the CLI seed, engine, and request id.
24 +fn request_seed(base_seed: u64, engine_index: u32, request_id: &str) -> u64 {
25 + let mut hasher = std::hash::DefaultHasher::new();
26 + base_seed.hash(&mut hasher);
27 + engine_index.hash(&mut hasher);
28 + request_id.hash(&mut hasher);
29 + hasher.finish()
30 +}
31 +
32 +/// Current UNIX timestamp in seconds for engine-core output envelopes.
33 +fn now_secs() -> f64 {
34 + SystemTime::now()
35 + .duration_since(UNIX_EPOCH)
36 + .map(|duration| duration.as_secs_f64())
37 + .unwrap_or_default()
38 +}
39 +
40 +/// Build one request output with only token IDs and terminal status populated.
41 +fn request_output(
42 + request_id: String,
43 + new_token_ids: Vec<u32>,
44 + finish_reason: Option<EngineCoreFinishReason>,
45 +) -> EngineCoreOutput {
46 + EngineCoreOutput {
47 + request_id,
48 + new_token_ids,
49 + finish_reason,
50 + ..Default::default()
51 + }
52 +}
53 +
54 +/// Produce an empty output with a terminal finish reason for an invalid request.
55 +fn empty_finish_outputs(
56 + engine_index: u32,
57 + request_id: String,
58 + finish_reason: EngineCoreFinishReason,
59 +) -> EngineCoreOutputs {
60 + let output = request_output(request_id, Vec::new(), Some(finish_reason));
61 + let finished_requests = BTreeSet::from([output.request_id.clone()]);
62 +
63 + EngineCoreOutputs {
64 + engine_index,
65 + outputs: vec![output],
66 + timestamp: now_secs(),
67 + finished_requests: Some(finished_requests),
68 + ..Default::default()
69 + }
70 +}
71 +
72 +/// Encode a utility result into the protocol's msgpack value envelope.
73 +fn utility_envelope<T>(value: T) -> Result<UtilityResultEnvelope>
74 +where
75 + T: Serialize,
76 +{
77 + Ok(UtilityResultEnvelope::without_type_info(
78 + rmpv::ext::to_value(value)?,
79 + ))
80 +}
81 +
82 +/// Produce the minimal utility responses needed by the Rust frontend.
83 +fn utility_response(
84 + engine_index: u32,
85 + request: EngineCoreUtilityRequest,
86 +) -> Result<EngineCoreOutputs> {
87 + let result = match request.method_name.as_str() {
88 + "get_supported_tasks" => utility_envelope(vec!["generate"]),
89 + "is_sleeping" => utility_envelope(false),
90 + "reset_prefix_cache" => utility_envelope(true),
91 + "reset_mm_cache"
92 + | "reset_encoder_cache"
93 + | "profile"
94 + | "sleep"
95 + | "wake_up"
96 + | "execute_dummy_batch" => utility_envelope(()),
97 + _ => utility_envelope(Value::Nil),
98 + }?;
99 +
100 + Ok(EngineCoreOutputs {
101 + engine_index,
102 + utility_output: Some(UtilityOutput {
103 + call_id: request.call_id,
104 + failure_message: None,
105 + result: Some(result),
106 + }),
107 + timestamp: now_secs(),
108 + ..Default::default()
109 + })
110 +}
111 +
112 +/// Message sent from the frontend to the mock engine task to drive the engine loop.
113 +pub(crate) enum EngineInput {
114 + Request(Box<EngineCoreRequest>),
115 + Abort(Vec<String>),
116 + Utility(EngineCoreUtilityRequest),
117 + StartDpWave,
118 +}
119 +
120 +/// Message sent from the mock engine task to the frontend for one engine output batch.
121 +pub(crate) struct EngineOutput {
122 + pub client_index: u32,
123 + pub outputs: EngineCoreOutputs,
124 +}
125 +
126 +/// Per-request decode state owned by one mock engine.
127 +#[derive(Debug)]
128 +struct ActiveRequest {
129 + request_id: String,
130 + client_index: u32,
131 + prompt_len: usize,
132 + max_tokens: usize,
133 + generated: usize,
134 + rng: StdRng,
135 +}
136 +
137 +impl ActiveRequest {
138 + /// Create a new active request from an incoming EngineCoreRequest, or return an immediate
139 + /// finish reason if the request is invalid.
140 + fn new(
141 + engine_index: u32,
142 + request: Box<EngineCoreRequest>,
143 + opt: &Opt,
144 + ) -> Result<Self, EngineCoreFinishReason> {
145 + let request_id = request.request_id;
146 + let client_index = request.client_index;
147 + let prompt_len = request.prompt_token_ids.as_ref().map(Vec::len).unwrap_or_default();
148 +
149 + let Some(sampling_params) = request.sampling_params else {
150 + warn!(
151 + request_id,
152 + "request has no sampling params; returning engine error"
153 + );
154 + return Err(EngineCoreFinishReason::Error);
155 + };
156 + let max_tokens = sampling_params.max_tokens as usize;
157 +
158 + if opt.log_requests {
159 + info!(
160 + request_id,
161 + prompt_len,
162 + max_tokens,
163 + chunk_size = opt.output_token_chunk_size,
164 + "mock request started"
165 + );
166 + }
167 +
168 + if max_tokens == 0 {
169 + return Err(EngineCoreFinishReason::Length);
170 + }
171 +
172 + Ok(ActiveRequest {
173 + rng: StdRng::seed_from_u64(request_seed(opt.seed, engine_index, &request_id)),
174 + request_id,
175 + client_index,
176 + prompt_len,
177 + max_tokens,
178 + generated: 0,
179 + })
180 + }
181 +
182 + /// Advance this request by one mock engine step.
183 + fn step(&mut self, opt: &Opt) -> EngineCoreOutput {
184 + let remaining = self.max_tokens - self.generated;
185 + let chunk_len = remaining.min(opt.output_token_chunk_size);
186 + let mut new_token_ids = Vec::with_capacity(chunk_len);
187 + for _ in 0..chunk_len {
188 + new_token_ids.push(self.rng.random_range(0..opt.vocab_size));
189 + }
190 + self.generated += chunk_len;
191 +
192 + let finished = self.generated >= self.max_tokens;
193 + request_output(
194 + self.request_id.clone(),
195 + new_token_ids,
196 + finished.then_some(EngineCoreFinishReason::Length),
197 + )
198 + }
199 +}
200 +
201 +/// Internal state for one mock engine instance, owned by the engine loop task.
202 +struct Engine {
203 + engine_index: u32,
204 + opt: Opt,
205 + active_requests: HashMap<String, ActiveRequest>,
206 +}
207 +
208 +impl Engine {
209 + /// Drain one frontend request message received on the input DEALER socket.
210 + fn handle_input(&mut self, input: EngineInput) -> Result<Vec<EngineOutput>> {
211 + let mut outputs = Vec::new();
212 +
213 + match input {
214 + EngineInput::Request(request) => {
215 + let request_id = request.request_id.clone();
216 + let client_index = request.client_index;
217 +
218 + if self.active_requests.contains_key(&request_id) {
219 + warn!(
220 + engine_index = self.engine_index,
221 + request_id, "duplicate mock request id"
222 + );
223 + return Ok(vec![EngineOutput {
224 + client_index,
225 + outputs: empty_finish_outputs(
226 + self.engine_index,
227 + request_id,
228 + EngineCoreFinishReason::Error,
229 + ),
230 + }]);
231 + }
232 +
233 + match ActiveRequest::new(self.engine_index, request, &self.opt) {
234 + Ok(request) => {
235 + self.active_requests.insert(request_id, request);
236 + }
237 + Err(finish_reason) => {
238 + return Ok(vec![EngineOutput {
239 + client_index,
240 + outputs: empty_finish_outputs(
241 + self.engine_index,
242 + request_id,
243 + finish_reason,
244 + ),
245 + }]);
246 + }
247 + }
248 + }
249 +
250 + EngineInput::Abort(request_ids) => {
251 + let mut outputs_by_client =
252 + BTreeMap::<u32, (Vec<EngineCoreOutput>, BTreeSet<String>)>::new();
253 + for request_id in request_ids {
254 + if let Some(request) = self.active_requests.remove(&request_id) {
255 + let output = request_output(
256 + request_id.clone(),
257 + Vec::new(),
258 + Some(EngineCoreFinishReason::Abort),
259 + );
260 + let (outputs, finished_requests) = outputs_by_client
261 + .entry(request.client_index)
262 + .or_insert_with(|| (Vec::new(), BTreeSet::new()));
263 + outputs.push(output);
264 + finished_requests.insert(request_id.clone());
265 + if self.opt.log_requests {
266 + info!(request_id, finish_reason = "abort", "mock request aborted");
267 + }
268 + }
269 + }
270 + for (client_index, (client_outputs, finished_requests)) in outputs_by_client {
271 + outputs.push({
272 + let outputs = EngineCoreOutputs {
273 + engine_index: self.engine_index,
274 + outputs: client_outputs,
275 + timestamp: now_secs(),
276 + finished_requests: Some(finished_requests),
277 + ..Default::default()
278 + };
279 + EngineOutput {
280 + client_index,
281 + outputs,
282 + }
283 + });
284 + }
285 + }
286 +
287 + EngineInput::Utility(request) => {
288 + debug!(
289 + engine_index = self.engine_index,
290 + call_id = %request.call_id,
291 + method = request.method_name,
292 + "mock utility request"
293 + );
294 + let client_index = request.client_index;
295 + outputs.push({
296 + let outputs = utility_response(self.engine_index, request)?;
297 + EngineOutput {
298 + client_index,
299 + outputs,
300 + }
301 + });
302 + }
303 +
304 + EngineInput::StartDpWave => {
305 + debug!(
306 + engine_index = self.engine_index,
307 + "ignoring START_DP_WAVE in mock engine"
308 + );
309 + }
310 + }
311 +
312 + Ok(outputs)
313 + }
314 +
315 + /// Advance active requests once and return one batched engine output.
316 + fn step(&mut self) -> Vec<EngineOutput> {
317 + if self.active_requests.is_empty() {
318 + return Vec::new();
319 + }
320 +
321 + let mut outputs_by_client =
322 + BTreeMap::<u32, (Vec<EngineCoreOutput>, BTreeSet<String>)>::new();
323 + let mut all_finished_requests = BTreeSet::new();
324 +
325 + for request in self.active_requests.values_mut() {
326 + let client_index = request.client_index;
327 + let output = request.step(&self.opt);
328 + let request_id = request.request_id.clone();
329 + let finished = output.finished();
330 + if output.finished() {
331 + all_finished_requests.insert(request_id.clone());
332 + if self.opt.log_requests {
333 + info!(
334 + request_id,
335 + prompt_len = request.prompt_len,
336 + output_tokens = request.generated,
337 + finish_reason = "length",
338 + "mock request finished"
339 + );
340 + }
341 + }
342 + let (outputs, finished_requests) = outputs_by_client
343 + .entry(client_index)
344 + .or_insert_with(|| (Vec::new(), BTreeSet::new()));
345 + if finished {
346 + finished_requests.insert(request_id.clone());
347 + }
348 + outputs.push(output);
349 + }
350 +
351 + for request_id in &all_finished_requests {
352 + self.active_requests.remove(request_id);
353 + }
354 +
355 + outputs_by_client
356 + .into_iter()
357 + .filter_map(|(client_index, (outputs, finished_requests))| {
358 + (!outputs.is_empty()).then(|| EngineOutput {
359 + client_index,
360 + outputs: EngineCoreOutputs {
361 + engine_index: self.engine_index,
362 + outputs,
363 + timestamp: now_secs(),
364 + finished_requests: (!finished_requests.is_empty())
365 + .then_some(finished_requests),
366 + ..Default::default()
367 + },
368 + })
369 + })
370 + .collect()
371 + }
372 +}
373 +
374 +/// Run the main loop for the mock engine, receiving `EngineInput` from `input_rx`
375 +/// and sending `EngineOutput` to `output_tx` until `shutdown` is cancelled.
376 +pub(crate) async fn run_engine_loop(
377 + engine_index: u32,
378 + opt: Opt,
379 + mut input_rx: mpsc::UnboundedReceiver<EngineInput>,
380 + output_tx: mpsc::Sender<EngineOutput>,
381 + shutdown: CancellationToken,
382 +) -> Result<()> {
383 + let mut engine = Engine {
384 + engine_index,
385 + opt,
386 + active_requests: HashMap::new(),
387 + };
388 +
389 + loop {
390 + let outputs = tokio::select! {
391 + biased;
392 + _ = shutdown.cancelled() => break,
393 +
394 + input = input_rx.recv() => {
395 + let input = input
396 + .ok_or_else(|| anyhow!("mock engine input channel closed"))?;
397 + engine.handle_input(input)?
398 + }
399 +
400 + // If there are active requests, step them once after yielding to the scheduler to
401 + // avoid blocking the engine loop while still making steady progress on request outputs.
402 + _ = yield_now(), if !engine.active_requests.is_empty() => {
403 + engine.step()
404 + }
405 + };
406 +
407 + for output in outputs {
408 + output_tx
409 + .send(output)
410 + .await
411 + .map_err(|_| anyhow!("mock engine IO task shut down"))?;
412 + }
413 + }
414 +
415 + Ok(())
416 +}
rust/src/mock-engine/src/io.rs new
+121
@@ -0,0 +1,121 @@
1 +use anyhow::{Context as _, Result, anyhow, bail};
2 +use futures::{Stream, StreamExt as _, stream};
3 +use tokio::sync::mpsc;
4 +use tokio_util::sync::CancellationToken;
5 +use tracing::warn;
6 +use vllm_engine_core_client::mock_engine::MockEngineDataSockets;
7 +use vllm_engine_core_client::protocol::utility::EngineCoreUtilityRequest;
8 +use vllm_engine_core_client::protocol::{
9 + EngineCoreRequest, EngineCoreRequestType, decode_msgpack, encode_msgpack,
10 +};
11 +use zeromq::{DealerSocket, PushSocket, SocketRecv as _, SocketSend as _, ZmqMessage};
12 +
13 +use crate::engine::{EngineInput, EngineOutput};
14 +
15 +/// Send one engine output batch to the client over the appropriate push socket.
16 +async fn send_engine_outputs_to_client(
17 + push_sockets: &mut [PushSocket],
18 + EngineOutput {
19 + client_index,
20 + outputs,
21 + }: EngineOutput,
22 +) -> Result<()> {
23 + let message = ZmqMessage::from(encode_msgpack(&outputs)?);
24 + push_sockets[client_index as usize].send(message).await?;
25 + Ok(())
26 +}
27 +
28 +/// Create a stream of `EngineInput` by continuously receiving messages from the given dealer socket
29 +/// and decoding them into `EngineInput`.
30 +fn dealer_input_stream(dealer: DealerSocket) -> impl Stream<Item = Result<EngineInput>> {
31 + stream::unfold(dealer, |mut dealer| async {
32 + let input = loop {
33 + let message =
34 + match dealer.recv().await.context("failed to receive message from dealer socket") {
35 + Ok(message) => message,
36 + Err(err) => break Err(err),
37 + };
38 +
39 + match decode_request(message) {
40 + Ok(input) => break Ok(input),
41 + Err(err) => {
42 + warn!(%err, "failed to decode engine request message; ignoring");
43 + }
44 + }
45 + };
46 +
47 + Some((input, dealer))
48 + })
49 +}
50 +
51 +/// Decode a `ZmqMessage` into an `EngineInput`. Returns an error if the message is malformed or
52 +/// contains an unknown/unsupported request type.
53 +fn decode_request(message: ZmqMessage) -> Result<EngineInput> {
54 + let frames = message.into_vec();
55 + if frames.is_empty() {
56 + bail!("empty engine request message");
57 + }
58 + if frames.len() != 2 {
59 + bail!("invalid frame count for engine request: {}", frames.len());
60 + }
61 +
62 + let request_type_frame = frames[0].as_ref();
63 + let Some(request_type) = EngineCoreRequestType::from_frame(request_type_frame) else {
64 + bail!("unknown engine request type: {:?}", request_type_frame);
65 + };
66 +
67 + let input = match request_type {
68 + EngineCoreRequestType::Add => {
69 + let request: Box<EngineCoreRequest> = decode_msgpack(frames[1].as_ref())?;
70 + EngineInput::Request(request)
71 + }
72 + EngineCoreRequestType::Abort => {
73 + let request_ids: Vec<String> = decode_msgpack(frames[1].as_ref())?;
74 + EngineInput::Abort(request_ids)
75 + }
76 + EngineCoreRequestType::Utility => {
77 + let request: EngineCoreUtilityRequest = decode_msgpack(frames[1].as_ref())?;
78 + EngineInput::Utility(request)
79 + }
80 + EngineCoreRequestType::StartDpWave => EngineInput::StartDpWave,
81 + };
82 +
83 + Ok(input)
84 +}
85 +
86 +/// Run the main IO loop for the mock engine, continuously receiving and decoding raw messages from
87 +/// the dealer sockets, sending them to the engine loop task via `input_tx`, and receiving
88 +/// `EngineOutput` from the engine loop task via `output_rx` and sending them to the client over the
89 +/// appropriate push socket, until `shutdown` is cancelled.
90 +pub(crate) async fn run_io_loop(
91 + data_sockets: Vec<MockEngineDataSockets>,
92 + input_tx: mpsc::UnboundedSender<EngineInput>,
93 + mut output_rx: mpsc::Receiver<EngineOutput>,
94 + shutdown: CancellationToken,
95 +) -> Result<()> {
96 + let (dealers, mut push_sockets): (Vec<_>, Vec<_>) =
97 + data_sockets.into_iter().map(|sockets| (sockets.dealer, sockets.push)).unzip();
98 + let mut input_streams =
99 + stream::select_all(dealers.into_iter().map(dealer_input_stream).map(Box::pin));
100 +
101 + loop {
102 + tokio::select! {
103 + biased;
104 + _ = shutdown.cancelled() => return Ok(()),
105 +
106 + output = output_rx.recv() => {
107 + let output = output
108 + .ok_or_else(|| anyhow!("mock engine output channel closed"))?;
109 + send_engine_outputs_to_client(&mut push_sockets, output).await?;
110 + }
111 +
112 + input = input_streams.next() => {
113 + let input = input
114 + .ok_or_else(|| anyhow!("mock engine input streams closed"))??;
115 + input_tx
116 + .send(input)
117 + .map_err(|_| anyhow!("mock engine state task shut down"))?;
118 + }
119 + }
120 + }
121 +}
rust/src/mock-engine/src/lib.rs new
+138
@@ -0,0 +1,138 @@
1 +use anyhow::{Context, Result, bail};
2 +use clap::Parser;
3 +use tokio::sync::mpsc;
4 +use tokio::task::JoinSet;
5 +use tokio_util::sync::CancellationToken;
6 +use tracing::{error, info};
7 +use vllm_engine_core_client::EngineId;
8 +use vllm_engine_core_client::mock_engine::{
9 + MockEngineConfig, MockEngineSockets, connect_to_frontend,
10 +};
11 +
12 +pub mod engine;
13 +pub mod io;
14 +
15 +/// Standalone engine-core protocol emulator for frontend stress testing.
16 +#[derive(Debug, Clone, Parser)]
17 +#[command(
18 + name = "vllm-mock-engine",
19 + about = "Run a mock vLLM headless engine for Rust frontend stress testing."
20 +)]
21 +pub struct Opt {
22 + /// Frontend-owned ZMQ handshake address.
23 + #[arg(long, default_value = "tcp://127.0.0.1:29550")]
24 + pub handshake_address: String,
25 +
26 + /// Number of mock engine identities to register with the frontend.
27 + #[arg(long, default_value_t = 1)]
28 + pub engine_count: usize,
29 +
30 + /// Number of accepted output tokens included in each EngineCoreOutput.
31 + #[arg(long, default_value_t = 1)]
32 + pub output_token_chunk_size: usize,
33 +
34 + /// Random token IDs are sampled uniformly from 0..vocab_size.
35 + #[arg(long, default_value_t = 32_000)]
36 + pub vocab_size: u32,
37 +
38 + /// Base seed for deterministic random token generation.
39 + #[arg(long, default_value_t = 0)]
40 + pub seed: u64,
41 +
42 + /// Log a summary line for each request.
43 + #[arg(long)]
44 + pub log_requests: bool,
45 +}
46 +
47 +/// Run one mock engine until shutdown or transport failure.
48 +async fn run_engine(engine_index: u32, opt: Opt, shutdown: CancellationToken) -> Result<()> {
49 + let MockEngineSockets { data_sockets, .. } = connect_to_frontend(
50 + &opt.handshake_address,
51 + EngineId::from_engine_index(engine_index),
52 + MockEngineConfig::default(),
53 + )
54 + .await
55 + .with_context(|| format!("mock engine {engine_index} failed to connect to frontend"))?;
56 +
57 + info!(engine_index, "mock engine connected to frontend");
58 +
59 + let (input_tx, input_rx) = mpsc::unbounded_channel();
60 + let (output_tx, output_rx) = mpsc::channel(64);
61 +
62 + // IO loop: dealer -> input_tx, output_rx -> push
63 + let mut io_loop = tokio::spawn(io::run_io_loop(
64 + data_sockets,
65 + input_tx,
66 + output_rx,
67 + shutdown.clone(),
68 + ));
69 + // Engine loop: input_rx -> engine logic -> output_tx
70 + let mut engine_loop = tokio::spawn(engine::run_engine_loop(
71 + engine_index,
72 + opt,
73 + input_rx,
74 + output_tx,
75 + shutdown.clone(),
76 + ));
77 +
78 + tokio::select! {
79 + biased;
80 + _ = shutdown.cancelled() => {
81 + io_loop.abort();
82 + engine_loop.abort();
83 + io_loop.await.ok();
84 + engine_loop.await.ok();
85 + }
86 +
87 + result = &mut io_loop => {
88 + error!(engine_index, "mock engine IO loop exited unexpectedly");
89 + engine_loop.abort();
90 + engine_loop.await.ok();
91 + result??;
92 + }
93 + result = &mut engine_loop => {
94 + error!(engine_index, "mock engine loop exited unexpectedly");
95 + io_loop.abort();
96 + io_loop.await.ok();
97 + result??;
98 + }
99 + }
100 +
101 + info!(engine_index, "mock engine shut down");
102 + Ok(())
103 +}
104 +
105 +/// Run all requested mock engines until cancellation or one engine task fails.
106 +pub async fn run(opt: Opt, shutdown: CancellationToken) -> Result<()> {
107 + info!(?opt, "starting mock engine");
108 +
109 + let mut engines = JoinSet::new();
110 + for engine_index in 0..opt.engine_count {
111 + engines.spawn(run_engine(
112 + engine_index as u32,
113 + opt.clone(),
114 + shutdown.clone(),
115 + ));
116 + }
117 +
118 + tokio::select! {
119 + biased;
120 + _ = shutdown.cancelled() => {
121 + engines.abort_all();
122 + while engines.join_next().await.is_some() {}
123 + Ok(())
124 + }
125 +
126 + joined = engines.join_next() => {
127 + match joined {
128 + Some(Ok(Ok(()))) => bail!("mock engine exited unexpectedly"),
129 + Some(Ok(Err(error))) => Err(error),
130 + Some(Err(error)) => Err(error).context("mock engine task join failed"),
131 + None => Ok(()),
132 + }
133 + }
134 + }
135 +}
136 +
137 +#[cfg(test)]
138 +mod tests;
rust/src/mock-engine/src/main.rs new
+38
@@ -0,0 +1,38 @@
1 +use anyhow::{Context, Result};
2 +use clap::Parser as _;
3 +use tokio_util::sync::CancellationToken;
4 +use tracing::{Level, info};
5 +use vllm_mock_engine::Opt;
6 +
7 +fn init_tracing() {
8 + tracing_subscriber::fmt().with_max_level(Level::INFO).init();
9 +}
10 +
11 +/// Create a cancellation token that is triggered by Ctrl-C.
12 +fn shutdown_signal() -> CancellationToken {
13 + let token = CancellationToken::new();
14 + let shutdown = token.clone();
15 +
16 + tokio::spawn(async move {
17 + tokio::signal::ctrl_c().await.expect("failed to install Ctrl-C signal handler");
18 + info!("received shutdown signal (Ctrl-C), shutting down...");
19 + shutdown.cancel();
20 + });
21 +
22 + token
23 +}
24 +
25 +fn main() -> Result<()> {
26 + init_tracing();
27 + let opt = Opt::parse();
28 +
29 + let runtime = tokio::runtime::Builder::new_multi_thread()
30 + .enable_all()
31 + .build()
32 + .context("failed to build Tokio runtime")?;
33 +
34 + runtime.block_on(async move {
35 + let shutdown = shutdown_signal();
36 + vllm_mock_engine::run(opt, shutdown).await
37 + })
38 +}
rust/src/mock-engine/src/tests.rs new
+197
@@ -0,0 +1,197 @@
1 +use std::net::TcpListener;
2 +use std::time::Duration;
3 +
4 +use anyhow::Result;
5 +use futures::StreamExt as _;
6 +use tokio::time::timeout;
7 +use tokio_util::sync::CancellationToken;
8 +use vllm_engine_core_client::protocol::{
9 + EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams,
10 +};
11 +use vllm_engine_core_client::test_utils::IpcNamespace;
12 +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
13 +
14 +use crate::{Opt, run};
15 +
16 +fn free_tcp_address() -> String {
17 + let listener = TcpListener::bind("127.0.0.1:0").expect("bind free port");
18 + let port = listener.local_addr().expect("local addr").port();
19 + drop(listener);
20 + format!("tcp://127.0.0.1:{port}")
21 +}
22 +
23 +fn client_config(handshake_address: String, engine_count: usize) -> EngineCoreClientConfig {
24 + EngineCoreClientConfig {
25 + transport_mode: TransportMode::HandshakeOwner {
26 + handshake_address,
27 + advertised_host: "127.0.0.1".to_string(),
28 + engine_count,
29 + ready_timeout: Duration::from_secs(5),
30 + local_input_address: None,
31 + local_output_address: None,
32 + },
33 + coordinator_mode: None,
34 + model_name: "mock-model".to_string(),
35 + client_index: 0,
36 + }
37 +}
38 +
39 +async fn connect_with_mock(
40 + handshake_address: String,
41 + engine_count: usize,
42 + output_token_chunk_size: usize,
43 +) -> (
44 + EngineCoreClient,
45 + CancellationToken,
46 + tokio::task::JoinHandle<Result<()>>,
47 +) {
48 + let shutdown = CancellationToken::new();
49 + let task = tokio::spawn(run(
50 + Opt {
51 + handshake_address: handshake_address.clone(),
52 + engine_count,
53 + output_token_chunk_size,
54 + vocab_size: 32_000,
55 + seed: 0,
56 + log_requests: false,
57 + },
58 + shutdown.clone(),
59 + ));
60 +
61 + let client = timeout(
62 + Duration::from_secs(5),
63 + EngineCoreClient::connect(client_config(handshake_address, engine_count)),
64 + )
65 + .await
66 + .expect("client connect timeout")
67 + .expect("connect client");
68 +
69 + (client, shutdown, task)
70 +}
71 +
72 +fn sample_request(request_id: &str, max_tokens: u32) -> EngineCoreRequest {
73 + EngineCoreRequest {
74 + request_id: request_id.to_string(),
75 + prompt_token_ids: Some(vec![1, 2, 3]),
76 + sampling_params: Some(EngineCoreSamplingParams {
77 + max_tokens,
78 + ..EngineCoreSamplingParams::for_test()
79 + }),
80 + arrival_time: 0.0,
81 + ..Default::default()
82 + }
83 +}
84 +
85 +async fn shutdown_mock(
86 + client: EngineCoreClient,
87 + shutdown: CancellationToken,
88 + task: tokio::task::JoinHandle<Result<()>>,
89 +) {
90 + client.shutdown().await.expect("client shutdown");
91 + shutdown.cancel();
92 + task.await.expect("mock join").expect("mock run");
93 +}
94 +
95 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
96 +async fn mock_engine_connects_over_tcp() {
97 + let handshake_address = free_tcp_address();
98 + let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
99 + assert_eq!(client.engine_count(), 1);
100 + assert_eq!(client.engine_identities()[0], &[0, 0]);
101 + assert_eq!(client.max_model_len(), Some(1024 * 1024));
102 + shutdown_mock(client, shutdown, task).await;
103 +}
104 +
105 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
106 +async fn mock_engine_connects_over_ipc() {
107 + let ipc = IpcNamespace::new().expect("ipc namespace");
108 + let handshake_address = ipc.handshake_endpoint();
109 + let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
110 + assert_eq!(client.engine_count(), 1);
111 + assert_eq!(client.engine_identities()[0], &[0, 0]);
112 + shutdown_mock(client, shutdown, task).await;
113 +}
114 +
115 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
116 +async fn mock_engine_registers_multiple_identities() {
117 + let handshake_address = free_tcp_address();
118 + let (client, shutdown, task) = connect_with_mock(handshake_address, 2, 1).await;
119 + assert_eq!(client.engine_count(), 2);
120 + assert_eq!(client.engine_identities(), vec![&[0, 0][..], &[1, 0][..]]);
121 + shutdown_mock(client, shutdown, task).await;
122 +}
123 +
124 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
125 +async fn chunk_size_one_outputs_one_token_per_update() {
126 + let handshake_address = free_tcp_address();
127 + let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
128 + let mut stream = client.call(sample_request("req-1", 3)).await.expect("call");
129 +
130 + let first = stream.next().await.expect("first").expect("first ok");
131 + assert_eq!(first.new_token_ids.len(), 1);
132 + assert_eq!(first.finish_reason, None);
133 + let second = stream.next().await.expect("second").expect("second ok");
134 + assert_eq!(second.new_token_ids.len(), 1);
135 + assert_eq!(second.finish_reason, None);
136 + let third = stream.next().await.expect("third").expect("third ok");
137 + assert_eq!(third.new_token_ids.len(), 1);
138 + assert_eq!(third.finish_reason, Some(EngineCoreFinishReason::Length));
139 + assert!(stream.next().await.is_none());
140 +
141 + shutdown_mock(client, shutdown, task).await;
142 +}
143 +
144 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
145 +async fn chunk_size_clips_final_output_to_max_tokens() {
146 + let handshake_address = free_tcp_address();
147 + let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 4).await;
148 + let mut stream = client.call(sample_request("req-clip", 6)).await.expect("call");
149 +
150 + let first = stream.next().await.expect("first").expect("first ok");
151 + assert_eq!(first.new_token_ids.len(), 4);
152 + assert_eq!(first.finish_reason, None);
153 + let second = stream.next().await.expect("second").expect("second ok");
154 + assert_eq!(second.new_token_ids.len(), 2);
155 + assert_eq!(second.finish_reason, Some(EngineCoreFinishReason::Length));
156 + assert!(stream.next().await.is_none());
157 +
158 + shutdown_mock(client, shutdown, task).await;
159 +}
160 +
161 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
162 +async fn abort_cancels_active_request_and_emits_terminal_output() {
163 + let handshake_address = free_tcp_address();
164 + let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
165 + let mut stream = client.call(sample_request("req-abort", 1_000_000)).await.expect("call");
166 + let first = stream.next().await.expect("first").expect("first ok");
167 + assert_eq!(first.finish_reason, None);
168 +
169 + client.abort(&["req-abort".to_string()]).await.expect("abort");
170 +
171 + loop {
172 + let output = timeout(Duration::from_secs(5), stream.next())
173 + .await
174 + .expect("stream timeout")
175 + .expect("terminal output")
176 + .expect("output ok");
177 + if output.finish_reason.is_some() {
178 + assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Abort));
179 + break;
180 + }
181 + }
182 +
183 + shutdown_mock(client, shutdown, task).await;
184 +}
185 +
186 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
187 +async fn utility_requests_return_minimal_success_responses() {
188 + let handshake_address = free_tcp_address();
189 + let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
190 +
191 + assert!(!client.is_sleeping().await.expect("is sleeping"));
192 + assert!(client.reset_prefix_cache(false, false).await.expect("reset prefix cache"));
193 + client.reset_mm_cache().await.expect("reset mm cache");
194 + client.reset_encoder_cache().await.expect("reset encoder cache");
195 +
196 + shutdown_mock(client, shutdown, task).await;
197 +}