diff --git a/crates/mock_worker/README.md b/crates/mock_worker/README.md index d24c9a27d0..def21937fc 100644 --- a/crates/mock_worker/README.md +++ b/crates/mock_worker/README.md @@ -80,6 +80,43 @@ label, e.g. `"labels":{"tokenizer_path":"gpt2"}`, and a `"kv_block_size":16`, an do **not** pass `--disable-tokenizer-autoload`. (HTTP workers need no tokenizer but cannot drive event-driven `cache_aware`, which requires token ids.) +## Capturing requests + +`--capture PATH` appends every gRPC `Generate` request a worker receives to +`PATH`, one JSON object per line, so a test can check exactly what the gateway +put on the wire. Responses are unchanged, in canned and realistic mode alike. +HTTP and ZMQ workers do not capture. + +```bash +cargo run --release -p mock-worker -- \ + --grpc-base-port 19000 --grpc-count 1 --model mock-model --capture generate.jsonl +``` + +Keys are the field names in +[`tokenspeed_scheduler.proto`](../grpc_client/proto/tokenspeed_scheduler.proto): + +- `request_id`, `input_ids`, `original_text` and `stream`. The gateway sends a + client `rid` as `request_id` (with a suffix under PD), so lines join to + responses on it. +- Every `SamplingParams` scalar, and `logit_bias` as an object with sorted + keys, so the same request always gives the same bytes. An optional the + gateway left unset is `null`. A float is the shortest decimal that reads back + as the same `f32`, so a request's `0.7` shows as `0.7`; NaN and infinities, + which JSON cannot hold, are strings such as `"NaN"`. +- `constraint`: `{"kind": "regex" | "json_schema" | "ebnf_grammar" | + "structural_tag", "value": }`, or `null`. +- `return_logprob`, `logprob_start_len`, `top_logprobs_num` and + `token_ids_logprob`. +- `has_custom_params`, `has_mm_inputs`, `has_encode_bootstrap_info`, + `has_kv_bootstrap_info` and `has_data_parallel_rank`: those fields are + recorded only as present or absent. + +Each line is in the file before the worker sends the first frame of its +response, so killing the worker loses no line. A capture path that cannot be +opened fails at startup with exit code 2, before any worker starts. A new +capture file is readable by its owner only (mode 0o600); an existing file keeps +its mode. + ## Scale-test rig (gateway CPU) `scripts/scale_test.sh` launches an IGW gateway, starts a canned mock fleet, diff --git a/crates/mock_worker/src/config.rs b/crates/mock_worker/src/config.rs index 0725093df4..8140db7ba4 100644 --- a/crates/mock_worker/src/config.rs +++ b/crates/mock_worker/src/config.rs @@ -1,6 +1,6 @@ //! Runtime configuration for the mock worker fleet, parsed from CLI flags. -use std::time::Duration; +use std::{path::PathBuf, time::Duration}; use crate::engine::EngineParams; @@ -38,11 +38,29 @@ pub struct Config { pub realistic: bool, /// Engine-simulator parameters (only used when `realistic`). pub engine: EngineParams, + /// Settings for replay testing (gRPC workers only). + pub replay: ReplayConfig, +} + +/// Settings for replaying recorded requests through the gateway against +/// this mock. Kept in one struct so that adding a replay flag does not +/// touch every `Config` literal. +#[derive(Debug, Clone, Default)] +pub struct ReplayConfig { + /// `--capture PATH`: append every gRPC `Generate` request to `PATH`, + /// one JSON object per line. + pub capture: Option, } impl Config { /// Parse the configuration from `std::env::args`, falling back to defaults. pub fn from_args() -> Result { + Self::parse(std::env::args().skip(1)) + } + + /// Parse the configuration from command-line flags, without the program + /// name. + fn parse(args: impl IntoIterator) -> Result { let mut cfg = Self { host: "127.0.0.1".to_string(), http_base_port: 9000, @@ -58,9 +76,10 @@ impl Config { output_tokens: 8, realistic: false, engine: EngineParams::default(), + replay: ReplayConfig::default(), }; - let mut args = std::env::args().skip(1); + let mut args = args.into_iter(); while let Some(flag) = args.next() { match flag.as_str() { "--host" => cfg.host = value(&mut args, &flag)?, @@ -79,6 +98,7 @@ impl Config { cfg.gen_delay = Duration::from_millis(parse(value(&mut args, &flag)?, &flag)?); } "--output-tokens" => cfg.output_tokens = parse(value(&mut args, &flag)?, &flag)?, + "--capture" => cfg.replay.capture = Some(value(&mut args, &flag)?.into()), "--engine" => { cfg.realistic = match value(&mut args, &flag)?.as_str() { "realistic" => true, @@ -158,6 +178,7 @@ fn usage() -> String { --tokenizer tokenizer path for gRPC autoload (default = model)\n\ --gen-ms canned per-request latency (default 0)\n\ --output-tokens output tokens per request when unspecified (default 8)\n\ + --capture append each gRPC Generate request to as a JSON line\n\ \n\ Realistic engine simulator (continuous batching; opt-in):\n\ --engine engine mode (default canned)\n\ @@ -171,3 +192,26 @@ fn usage() -> String { --prefix-cache enable prefix caching + KV events (default true)" .to_string() } + +#[cfg(test)] +mod tests { + use super::*; + + fn parse(flags: &[&str]) -> Result { + Config::parse(flags.iter().map(|flag| (*flag).to_string())) + } + + #[test] + fn capture_flag_sets_the_capture_path() { + let grpc = ["--grpc-base-port", "19000", "--grpc-count", "1"]; + let cfg = parse(&grpc).expect("gRPC flags parse"); + assert_eq!(cfg.replay.capture, None, "capture is off by default"); + + let cfg = parse(&[&grpc[..], &["--capture", "/tmp/generate.jsonl"]].concat()) + .expect("--capture parses"); + assert_eq!( + cfg.replay.capture, + Some(PathBuf::from("/tmp/generate.jsonl")) + ); + } +} diff --git a/crates/mock_worker/src/grpc.rs b/crates/mock_worker/src/grpc.rs index 2d6e72499e..187bd60fe7 100644 --- a/crates/mock_worker/src/grpc.rs +++ b/crates/mock_worker/src/grpc.rs @@ -20,6 +20,7 @@ use ts::{ use crate::{ config::Config, engine::{self, Engine, NewRequest}, + replay::Capture, }; /// Serve the mock TokenSpeed gRPC service on `port` until the process exits. @@ -45,9 +46,26 @@ pub async fn serve(cfg: Arc, host: String, port: u16) { /// port 0 and read it back instead of picking one and binding it later. pub async fn serve_with_listener(cfg: Arc, listener: TcpListener) { let addr = listener.local_addr().ok(); + let capture = match &cfg.replay.capture { + Some(path) => match Capture::open(path) { + Ok(capture) => Some(Arc::new(capture)), + Err(e) => { + tracing::error!( + "grpc worker {addr:?} cannot open capture file {}: {e}", + path.display() + ); + return; + } + }, + None => None, + }; // One simulated engine per listener (i.e. per virtual worker). let engine = cfg.realistic.then(|| Engine::spawn(cfg.engine.clone())); - let service = MockScheduler { cfg, engine }; + let service = MockScheduler { + cfg, + engine, + capture, + }; if let Err(e) = Server::builder() .add_service(TokenSpeedSchedulerServer::new(service)) .serve_with_incoming(TcpListenerStream::new(listener)) @@ -62,6 +80,8 @@ struct MockScheduler { cfg: Arc, /// Present iff the worker runs the realistic engine simulator. engine: Option, + /// Present iff `--capture` is set: every `Generate` request is recorded. + capture: Option>, } type GenStream = Pin> + Send>>; @@ -79,9 +99,18 @@ impl TokenSpeedScheduler for MockScheduler { &self, request: Request, ) -> Result, Status> { + let req = request.into_inner(); + // Record before anything else: the line is in the file before the + // stream is returned, so before the gateway can read any frame. + if let Some(capture) = &self.capture { + if let Err(e) = capture.record(&req) { + tracing::error!("capturing request {} failed: {e}", req.request_id); + return Err(Status::internal(format!("mock-worker capture failed: {e}"))); + } + } + // Realistic mode: submit to the engine simulator and stream its output. if let Some(engine) = &self.engine { - let req = request.into_inner(); let request_id = req.request_id; let prompt_token_ids = req.tokenized.map(|t| t.input_ids).unwrap_or_default(); // Omitted limit falls back to the worker default, matching the HTTP @@ -107,7 +136,7 @@ impl TokenSpeedScheduler for MockScheduler { } // Canned mode: a single up-front delay, then synthetic token ids. - let request_id = request.into_inner().request_id; + let request_id = req.request_id; if !self.cfg.gen_delay.is_zero() { tokio::time::sleep(self.cfg.gen_delay).await; } diff --git a/crates/mock_worker/src/lib.rs b/crates/mock_worker/src/lib.rs index 0a355e5129..c76d153c1b 100644 --- a/crates/mock_worker/src/lib.rs +++ b/crates/mock_worker/src/lib.rs @@ -5,4 +5,5 @@ pub mod config; pub mod engine; pub mod grpc; pub mod http; +pub mod replay; pub mod zmq; diff --git a/crates/mock_worker/src/main.rs b/crates/mock_worker/src/main.rs index 4e2d254579..1cacf5191e 100644 --- a/crates/mock_worker/src/main.rs +++ b/crates/mock_worker/src/main.rs @@ -8,7 +8,7 @@ use std::{process::ExitCode, sync::Arc}; -use mock_worker::{config::Config, grpc, http, zmq}; +use mock_worker::{config::Config, grpc, http, replay::Capture, zmq}; #[tokio::main] async fn main() -> ExitCode { @@ -24,6 +24,19 @@ async fn main() -> ExitCode { } }; + // A capture file that cannot be opened is a config error. Each gRPC worker + // opens it again, and one that cannot would stop while the process kept + // running, so check it once before any worker starts. + if let Some(path) = &cfg.replay.capture { + if let Err(e) = Capture::open(path) { + eprintln!( + "mock-worker: cannot open capture file {}: {e}", + path.display() + ); + return ExitCode::from(2); + } + } + tracing::info!( "mock-worker: {} http from :{}, {} grpc from :{}, model={}", cfg.http_count, diff --git a/crates/mock_worker/src/replay.rs b/crates/mock_worker/src/replay.rs new file mode 100644 index 0000000000..1e0c79b27f --- /dev/null +++ b/crates/mock_worker/src/replay.rs @@ -0,0 +1,178 @@ +//! Replay testing: `--capture` records every gRPC `Generate` request the +//! worker receives, so a test can check what the gateway put on the wire. + +#[cfg(unix)] +use std::os::unix::fs::OpenOptionsExt; +use std::{ + collections::BTreeMap, + fs::{File, OpenOptions}, + io::{self, Write}, + path::Path, + sync::{Mutex, PoisonError}, +}; + +use serde_json::{json, Map, Number, Value}; +use smg_grpc_client::tokenspeed_scheduler::tokenspeed_proto as ts; +use ts::sampling_params::Constraint; + +/// An append-mode capture file holding one JSON object per `Generate` request. +pub struct Capture { + file: Mutex, +} + +impl Capture { + /// Open `path` for appending, creating it if it does not exist. A capture + /// holds prompts, so a file this creates is readable by its owner only + /// (mode 0o600 on Unix). A file that already exists keeps its mode: it + /// was made by someone who chose who may read it. + pub fn open(path: &Path) -> io::Result { + let mut options = OpenOptions::new(); + options.create(true).append(true); + #[cfg(unix)] + options.mode(0o600); + let file = options.open(path)?; + Ok(Self { + file: Mutex::new(file), + }) + } + + /// Append `req` as one line. The line is in the file when this returns: + /// `File` keeps no buffer of its own, so a reader that opens the file + /// afterwards sees the line, and killing the process cannot lose it. + pub fn record(&self, req: &ts::GenerateRequest) -> io::Result<()> { + let mut line = serde_json::to_vec(&capture_json(req))?; + line.push(b'\n'); + // One `write_all` under the lock keeps this worker's lines whole and + // in order. Append mode puts each write at the current end of the + // file, so workers that share the file add to it, never overwrite it. + let mut file = self.file.lock().unwrap_or_else(PoisonError::into_inner); + file.write_all(&line)?; + file.flush() + } +} + +/// The capture line for `req`, written by hand because the prost types do not +/// derive `Serialize`. Keys are the proto field names. An optional the +/// client left unset is `null`; an absent `tokenized` or `sampling_params` +/// message reads as its defaults, as it does for an engine. Fields that +/// carry payloads (`custom_params`, `mm_inputs`, the bootstrap infos) and +/// `data_parallel_rank` are recorded only as present or absent. +fn capture_json(req: &ts::GenerateRequest) -> Value { + let (input_ids, original_text) = match &req.tokenized { + Some(tokenized) => ( + tokenized.input_ids.as_slice(), + tokenized.original_text.as_str(), + ), + None => (&[][..], ""), + }; + let unset = ts::SamplingParams::default(); + let sampling = req.sampling_params.as_ref().unwrap_or(&unset); + // Decoding puts `logit_bias` in a `HashMap`, whose order changes from one + // decode to the next; sorted keys give the same request the same bytes. + let logit_bias: Map = sampling + .logit_bias + .iter() + .collect::>() + .into_iter() + .map(|(token, bias)| (token.clone(), f32_json(*bias))) + .collect(); + json!({ + "request_id": req.request_id, + "input_ids": input_ids, + "original_text": original_text, + "stream": req.stream, + "temperature": sampling.temperature.map(f32_json), + "top_p": sampling.top_p.map(f32_json), + "top_k": sampling.top_k, + "min_p": sampling.min_p.map(f32_json), + "repetition_penalty": sampling.repetition_penalty.map(f32_json), + "frequency_penalty": sampling.frequency_penalty.map(f32_json), + "presence_penalty": sampling.presence_penalty.map(f32_json), + "max_new_tokens": sampling.max_new_tokens, + "min_new_tokens": sampling.min_new_tokens, + "stop": sampling.stop, + "stop_token_ids": sampling.stop_token_ids, + "ignore_eos": sampling.ignore_eos, + "no_stop_trim": sampling.no_stop_trim, + "skip_special_tokens": sampling.skip_special_tokens, + "spaces_between_special_tokens": sampling.spaces_between_special_tokens, + "n": sampling.n, + "sampling_seed": sampling.sampling_seed, + "logit_bias": logit_bias, + "constraint": constraint_json(sampling.constraint.as_ref()), + "return_logprob": req.return_logprob, + "logprob_start_len": req.logprob_start_len, + "top_logprobs_num": req.top_logprobs_num, + "token_ids_logprob": req.token_ids_logprob, + "has_custom_params": sampling.custom_params.is_some(), + "has_mm_inputs": req.mm_inputs.is_some(), + "has_encode_bootstrap_info": req.encode_bootstrap_info.is_some(), + "has_kv_bootstrap_info": req.kv_bootstrap_info.is_some(), + "has_data_parallel_rank": req.data_parallel_rank.is_some(), + }) +} + +/// An `f32` as the shortest decimal that reads back as the same `f32`, so a +/// request's `0.7` is captured as `0.7`, not as the `0.699999988079071` that +/// widening it to `f64` prints. JSON has no NaN or infinity, and `null` +/// means unset, so a non-finite value is written as its name, e.g. `"NaN"`. +fn f32_json(value: f32) -> Value { + let shortest = value.to_string(); + match shortest.parse::().ok().and_then(Number::from_f64) { + Some(number) => Value::Number(number), + None => Value::String(shortest), + } +} + +/// `{"kind", "value"}` for the constraint oneof, with the value as sent. +fn constraint_json(constraint: Option<&Constraint>) -> Value { + let (kind, value) = match constraint { + None => return Value::Null, + Some(Constraint::Regex(value)) => ("regex", value), + Some(Constraint::JsonSchema(value)) => ("json_schema", value), + Some(Constraint::EbnfGrammar(value)) => ("ebnf_grammar", value), + Some(Constraint::StructuralTag(value)) => ("structural_tag", value), + }; + json!({ "kind": kind, "value": value }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// JSON has no NaN or infinity, and `null` means unset, so a set + /// non-finite value must come back as something other than `null`. + #[test] + fn non_finite_floats_are_not_captured_as_unset() { + assert_eq!(f32_json(f32::NAN), json!("NaN")); + assert_eq!(f32_json(f32::INFINITY), json!("inf")); + assert_eq!(f32_json(f32::NEG_INFINITY), json!("-inf")); + } + + /// A capture holds prompts, so a file the mock creates is readable by + /// its owner only. A file that already exists keeps the mode it has. + #[cfg(unix)] + #[test] + fn a_new_capture_file_is_owner_only() { + use std::{fs, os::unix::fs::PermissionsExt}; + + let dir = tempfile::tempdir().expect("temp dir"); + let mode = |path: &Path| { + let mode = fs::metadata(path) + .expect("capture file metadata") + .permissions() + .mode(); + format!("{:o}", mode & 0o777) + }; + + let created = dir.path().join("created.jsonl"); + Capture::open(&created).expect("create a capture file"); + assert_eq!(mode(&created), "600", "a new file is owner-only"); + + let existing = dir.path().join("existing.jsonl"); + fs::write(&existing, b"").expect("create a file"); + fs::set_permissions(&existing, fs::Permissions::from_mode(0o644)).expect("chmod 644"); + Capture::open(&existing).expect("open an existing capture file"); + assert_eq!(mode(&existing), "644", "an existing file keeps its mode"); + } +} diff --git a/crates/mock_worker/src/zmq.rs b/crates/mock_worker/src/zmq.rs index ecff4f842d..c58d9079c7 100644 --- a/crates/mock_worker/src/zmq.rs +++ b/crates/mock_worker/src/zmq.rs @@ -308,6 +308,7 @@ mod tests { output_tokens: 4, realistic: false, engine: EngineParams::default(), + replay: Default::default(), } } diff --git a/crates/mock_worker/tests/capture.rs b/crates/mock_worker/tests/capture.rs new file mode 100644 index 0000000000..73dc652e7b --- /dev/null +++ b/crates/mock_worker/tests/capture.rs @@ -0,0 +1,456 @@ +//! `--capture` writes every gRPC `Generate` request a worker receives to a +//! file, one JSON object per line, before the worker sends any frame. These +//! tests serve the real TokenSpeed service on a free port and drive it with +//! the gateway's own client. + +use std::{ + path::{Path, PathBuf}, + sync::Arc, + time::Duration, +}; + +use futures::StreamExt; +use mock_worker::{ + config::{Config, ReplayConfig}, + engine::EngineParams, + grpc::serve_with_listener, +}; +use serde_json::{json, Value}; +use smg_grpc_client::{tokenspeed_proto as ts, TokenSpeedSchedulerClient}; +use tokio::net::TcpListener; +use ts::generate_response::Response as GenResp; + +/// A canned worker that captures to `capture` when it is set. +fn config(capture: Option) -> Config { + Config { + host: "127.0.0.1".to_string(), + http_base_port: 0, + http_count: 0, + grpc_base_port: 0, + grpc_count: 1, + zmq_handshake: None, + zmq_count: 0, + zmq_start_index: 0, + model_id: "capture-test".to_string(), + tokenizer_path: "capture-test".to_string(), + gen_delay: Duration::ZERO, + output_tokens: 3, + realistic: false, + engine: EngineParams::default(), + replay: ReplayConfig { capture }, + } +} + +/// Serve `cfg` on a free port and connect the gateway's client to it. +async fn start(cfg: Config) -> TokenSpeedSchedulerClient { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a free port"); + let port = listener.local_addr().expect("listener address").port(); + #[expect( + clippy::disallowed_methods, + reason = "the server serves until the test process exits" + )] + let _server = tokio::spawn(serve_with_listener(Arc::new(cfg), listener)); + TokenSpeedSchedulerClient::connect(&format!("http://127.0.0.1:{port}")) + .await + .expect("connect to the mock worker") +} + +/// A streamed request with every sampling field unset. +fn request(request_id: &str, input_ids: &[u32], text: &str) -> ts::GenerateRequest { + ts::GenerateRequest { + request_id: request_id.to_string(), + tokenized: Some(ts::TokenizedInput { + input_ids: input_ids.to_vec(), + original_text: text.to_string(), + }), + sampling_params: Some(ts::SamplingParams::default()), + stream: true, + ..Default::default() + } +} + +/// Send `req` and read every frame of the response. +async fn generate( + client: &TokenSpeedSchedulerClient, + req: ts::GenerateRequest, +) -> Vec { + let mut stream = client.generate(req).await.expect("generate"); + let mut frames = Vec::new(); + while let Some(frame) = stream.next().await { + frames.push(frame.expect("a response frame")); + } + stream.mark_completed(); + frames +} + +/// The capture file's lines, each parsed as JSON. +fn captured(path: &Path) -> Vec { + let text = std::fs::read_to_string(path).expect("read the capture file"); + assert!( + text.is_empty() || text.ends_with('\n'), + "every line ends with a newline: {text:?}" + ); + text.lines() + .map(|line| serde_json::from_str(line).expect("each line is one JSON object")) + .collect() +} + +#[tokio::test] +async fn writes_one_line_per_request_in_order() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("generate.jsonl"); + let client = start(config(Some(path.clone()))).await; + + // A newline inside the text must not split the line; ids span the u32 range. + let sent: [(&str, Vec, &str, bool); 3] = [ + ( + "first", + vec![151644, 872, 198], + "<|im_start|>user\nhi", + true, + ), + ("second", vec![], "", false), + ( + "third", + vec![0, 4_294_967_295], + "a \"quoted\" word, a \\ backslash,\ta tab and 日本語\n", + true, + ), + ]; + for (request_id, input_ids, text, stream) in &sent { + let mut req = request(request_id, input_ids, text); + req.stream = *stream; + generate(&client, req).await; + } + + let lines = captured(&path); + assert_eq!(lines.len(), sent.len(), "one line per request: {lines:?}"); + for (line, (request_id, input_ids, text, stream)) in lines.iter().zip(&sent) { + assert_eq!(line["request_id"], *request_id); + assert_eq!(line["input_ids"], json!(input_ids), "{request_id}"); + assert_eq!(line["original_text"], *text, "{request_id}"); + assert_eq!(line["stream"], *stream, "{request_id}"); + } +} + +/// The line is in the file once the client holds the response stream, +/// before it reads a frame. The realistic engine's slow decode step holds the +/// first frame back by 200 ms, so a line written while streaming would be +/// missing when the file is read. +#[tokio::test] +async fn line_is_on_disk_before_the_first_frame() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("generate.jsonl"); + let mut cfg = config(Some(path.clone())); + cfg.realistic = true; + cfg.engine.decode_base_ms = 200.0; + let client = start(cfg).await; + + let mut req = request("early", &[1, 2, 3], "early"); + if let Some(sampling) = req.sampling_params.as_mut() { + sampling.max_new_tokens = Some(1); + } + let mut stream = client.generate(req).await.expect("generate"); + + let lines = captured(&path); + assert_eq!(lines.len(), 1, "the line is written before any frame"); + assert_eq!(lines[0]["request_id"], "early"); + + let first = stream.next().await.expect("a first frame"); + assert_eq!(first.expect("a response frame").request_id, "early"); + while let Some(frame) = stream.next().await { + frame.expect("a response frame"); + } + stream.mark_completed(); +} + +/// With `--capture` on, a canned worker sends exactly the frames it sends +/// without it. +#[tokio::test] +async fn canned_output_is_unchanged_with_capture_on() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("generate.jsonl"); + let with_capture = start(config(Some(path.clone()))).await; + let without_capture = start(config(None)).await; + + for stream in [true, false] { + let mut req = request("canned", &[1, 2, 3], "canned"); + req.stream = stream; + let frames = generate(&with_capture, req.clone()).await; + assert!( + matches!( + frames.last().and_then(|frame| frame.response.as_ref()), + Some(GenResp::Complete(_)) + ), + "stream={stream}: the response ends with a complete: {frames:?}" + ); + assert_eq!( + frames, + generate(&without_capture, req).await, + "stream={stream}" + ); + } + assert_eq!(captured(&path).len(), 2, "both requests were captured"); +} + +/// An optional the client leaves unset is `null`; a field it sets comes back +/// as sent, zeros included. Whole lines are compared, so a missing key fails. +#[tokio::test] +async fn unset_optionals_are_null_and_set_values_come_back_as_sent() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("generate.jsonl"); + let client = start(config(Some(path.clone()))).await; + + generate(&client, request("unset", &[1], "unset")).await; + + let mut set = request("set", &[2], "set"); + set.sampling_params = Some(ts::SamplingParams { + temperature: Some(0.7), + top_p: Some(0.95), + top_k: Some(20), + min_p: Some(0.0), + frequency_penalty: Some(0.5), + presence_penalty: Some(-0.25), + repetition_penalty: Some(1.1), + max_new_tokens: Some(128), + min_new_tokens: 4, + stop: vec!["".to_string(), "\n\nUser:".to_string()], + stop_token_ids: vec![151643, 151645], + ignore_eos: true, + skip_special_tokens: true, + spaces_between_special_tokens: true, + n: 2, + logit_bias: [("151643".to_string(), -100.0), ("42".to_string(), 0.1)].into(), + constraint: None, + no_stop_trim: true, + custom_params: Some(Default::default()), + sampling_seed: Some(u64::MAX), + }); + set.return_logprob = true; + set.logprob_start_len = Some(0); + set.top_logprobs_num = 5; + set.token_ids_logprob = vec![7, 8]; + set.mm_inputs = Some(ts::MultimodalInputs::default()); + set.encode_bootstrap_info = Some(ts::EncodeBootstrapInfo::default()); + set.kv_bootstrap_info = Some(ts::KvBootstrapInfo::default()); + set.data_parallel_rank = Some(0); + generate(&client, set).await; + + let lines = captured(&path); + assert_eq!(lines.len(), 2, "{lines:?}"); + assert_eq!( + lines[0], + json!({ + "request_id": "unset", + "input_ids": [1], + "original_text": "unset", + "stream": true, + "temperature": null, + "top_p": null, + "top_k": null, + "min_p": null, + "repetition_penalty": null, + "frequency_penalty": null, + "presence_penalty": null, + "max_new_tokens": null, + "min_new_tokens": 0, + "stop": [], + "stop_token_ids": [], + "ignore_eos": false, + "no_stop_trim": false, + "skip_special_tokens": false, + "spaces_between_special_tokens": false, + "n": 0, + "sampling_seed": null, + "logit_bias": {}, + "constraint": null, + "return_logprob": false, + "logprob_start_len": null, + "top_logprobs_num": 0, + "token_ids_logprob": [], + "has_custom_params": false, + "has_mm_inputs": false, + "has_encode_bootstrap_info": false, + "has_kv_bootstrap_info": false, + "has_data_parallel_rank": false, + }) + ); + assert_eq!( + lines[1], + json!({ + "request_id": "set", + "input_ids": [2], + "original_text": "set", + "stream": true, + "temperature": 0.7, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "repetition_penalty": 1.1, + "frequency_penalty": 0.5, + "presence_penalty": -0.25, + "max_new_tokens": 128, + "min_new_tokens": 4, + "stop": ["", "\n\nUser:"], + "stop_token_ids": [151643, 151645], + "ignore_eos": true, + "no_stop_trim": true, + "skip_special_tokens": true, + "spaces_between_special_tokens": true, + "n": 2, + "sampling_seed": u64::MAX, + "logit_bias": {"151643": -100.0, "42": 0.1}, + "constraint": null, + "return_logprob": true, + "logprob_start_len": 0, + "top_logprobs_num": 5, + "token_ids_logprob": [7, 8], + "has_custom_params": true, + "has_mm_inputs": true, + "has_encode_bootstrap_info": true, + "has_kv_bootstrap_info": true, + "has_data_parallel_rank": true, + }) + ); +} + +/// The constraint is captured as its oneof arm and the string as sent. +#[tokio::test] +async fn constraint_is_captured_with_its_kind_and_value() { + use ts::sampling_params::Constraint; + + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("generate.jsonl"); + let client = start(config(Some(path.clone()))).await; + + /// Builds one arm of the constraint oneof from its value. + type Arm = fn(String) -> Constraint; + // The spaces in the JSON values show they are kept as sent, not re-encoded. + let cases: [(&str, Arm, &str); 4] = [ + ("regex", Constraint::Regex, "[a-z]+"), + ( + "json_schema", + Constraint::JsonSchema, + r#"{"type": "array", "minItems": 1}"#, + ), + ( + "ebnf_grammar", + Constraint::EbnfGrammar, + r#"root ::= "yes" | "no""#, + ), + ( + "structural_tag", + Constraint::StructuralTag, + r#"{"type": "structural_tag", "format": {"type": "const_string", "value": ""}}"#, + ), + ]; + for (kind, constraint, value) in &cases { + let mut req = request(kind, &[1], kind); + if let Some(sampling) = req.sampling_params.as_mut() { + sampling.constraint = Some(constraint(value.to_string())); + } + generate(&client, req).await; + } + + let lines = captured(&path); + assert_eq!(lines.len(), cases.len(), "{lines:?}"); + for (line, (kind, _, value)) in lines.iter().zip(&cases) { + assert_eq!( + line["constraint"], + json!({"kind": kind, "value": value}), + "{kind}" + ); + } +} + +/// A capture file that cannot be opened stops the worker instead of letting +/// it serve requests it cannot record. +#[tokio::test] +async fn unopenable_capture_file_stops_the_worker() { + let dir = tempfile::tempdir().expect("temp dir"); + // The parent directory does not exist, so the file cannot be created. + let path = dir.path().join("missing").join("generate.jsonl"); + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a free port"); + let serving = serve_with_listener(Arc::new(config(Some(path))), listener); + tokio::time::timeout(Duration::from_secs(5), serving) + .await + .expect("the worker stops instead of serving"); +} + +/// The binary checks the capture file before any worker starts: a path that +/// cannot be opened is a config error (exit 2), not a live process that +/// serves nothing. +#[tokio::test] +async fn unopenable_capture_file_fails_at_startup() { + let dir = tempfile::tempdir().expect("temp dir"); + // The parent directory does not exist, so the file cannot be created. + let path = dir.path().join("missing").join("generate.jsonl"); + let port = std::net::TcpListener::bind("127.0.0.1:0") + .and_then(|listener| listener.local_addr()) + .expect("a free port") + .port(); + let port_arg = port.to_string(); + + let mut mock_worker = tokio::process::Command::new(env!("CARGO_BIN_EXE_mock-worker")); + mock_worker + .args(["--grpc-count", "1", "--grpc-base-port", port_arg.as_str()]) + .arg("--capture") + .arg(&path) + .kill_on_drop(true); + let output = tokio::time::timeout(Duration::from_secs(10), mock_worker.output()) + .await + .expect("mock-worker exits instead of running on") + .expect("run mock-worker"); + + assert_eq!(output.status.code(), Some(2), "{output:?}"); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + stderr.contains(&format!("cannot open capture file {}", path.display())), + "stderr names the path: {stderr}" + ); + assert!( + std::net::TcpStream::connect(("127.0.0.1", port)).is_err(), + "nothing is left listening on port {port}" + ); +} + +/// The same request gives the same bytes. Decoding puts `logit_bias` in a +/// `HashMap`, whose order changes from one decode to the next, so its keys +/// must be written in a fixed order. +#[tokio::test] +async fn the_same_request_gives_byte_identical_lines() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("generate.jsonl"); + let client = start(config(Some(path.clone()))).await; + + let tokens = [ + 0, 1, 5, 7, 9, 13, 31, 42, 64, 77, 100, 256, 2000, 128000, 151643, 151645, + ]; + let mut req = request("biased", &[1], "biased"); + if let Some(sampling) = req.sampling_params.as_mut() { + sampling.logit_bias = tokens + .iter() + .map(|token| (token.to_string(), -1.5)) + .collect(); + } + generate(&client, req.clone()).await; + generate(&client, req).await; + + let text = std::fs::read_to_string(&path).expect("read the capture file"); + let lines: Vec<&str> = text.lines().collect(); + assert_eq!(lines.len(), 2, "{text}"); + assert_eq!(lines[0], lines[1], "the same request gives the same bytes"); + let line: Value = serde_json::from_str(lines[0]).expect("one JSON object"); + let keys: Vec<&String> = line["logit_bias"] + .as_object() + .expect("logit_bias is an object") + .keys() + .collect(); + let mut sorted = keys.clone(); + sorted.sort(); + assert_eq!(keys, sorted, "logit_bias keys are sorted"); +} diff --git a/model_gateway/src/routers/grpc/regular/streaming/eof_tests.rs b/model_gateway/src/routers/grpc/regular/streaming/eof_tests.rs index 6321ee1e75..984e6821f9 100644 --- a/model_gateway/src/routers/grpc/regular/streaming/eof_tests.rs +++ b/model_gateway/src/routers/grpc/regular/streaming/eof_tests.rs @@ -132,6 +132,7 @@ async fn scripted_stream( output_tokens: 0, realistic: false, engine: mock_worker::engine::EngineParams::default(), + replay: Default::default(), }); let server = tokio::spawn(mock_worker::grpc::serve_with_listener(config, listener)); let client = VllmEngineClient::connect(&format!("http://127.0.0.1:{port}")) diff --git a/model_gateway/tests/grpc_context_length_test.rs b/model_gateway/tests/grpc_context_length_test.rs index a489039b0c..6fad2b926e 100644 --- a/model_gateway/tests/grpc_context_length_test.rs +++ b/model_gateway/tests/grpc_context_length_test.rs @@ -68,6 +68,7 @@ async fn start_mock_grpc_worker() -> u16 { output_tokens: 2, realistic: false, engine: mock_worker::engine::EngineParams::default(), + replay: Default::default(), }); tokio::spawn(mock_worker::grpc::serve_with_listener(cfg, listener)); port diff --git a/model_gateway/tests/grpc_pd_fanout_test.rs b/model_gateway/tests/grpc_pd_fanout_test.rs index 439b60796e..5a5954678e 100644 --- a/model_gateway/tests/grpc_pd_fanout_test.rs +++ b/model_gateway/tests/grpc_pd_fanout_test.rs @@ -129,6 +129,7 @@ async fn start_mock_grpc_worker() -> u16 { output_tokens: OUTPUT_TOKENS, realistic: false, engine: mock_worker::engine::EngineParams::default(), + replay: Default::default(), }); tokio::spawn(mock_worker::grpc::serve_with_listener(cfg, listener)); port diff --git a/model_gateway/tests/tenant_rate_limiting_grpc_test.rs b/model_gateway/tests/tenant_rate_limiting_grpc_test.rs index 73907742a9..d5775eca98 100644 --- a/model_gateway/tests/tenant_rate_limiting_grpc_test.rs +++ b/model_gateway/tests/tenant_rate_limiting_grpc_test.rs @@ -78,6 +78,7 @@ async fn start_mock_grpc_worker(output_tokens: u32) -> u16 { output_tokens, realistic: false, engine: mock_worker::engine::EngineParams::default(), + replay: Default::default(), }); tokio::spawn(mock_worker::grpc::serve_with_listener(cfg, listener)); port diff --git a/model_gateway/tests/zmq_backend_test.rs b/model_gateway/tests/zmq_backend_test.rs index f5f9d4b03e..24679d19c2 100644 --- a/model_gateway/tests/zmq_backend_test.rs +++ b/model_gateway/tests/zmq_backend_test.rs @@ -87,6 +87,7 @@ fn start_mock_zmq_engines(handshake: &str, count: u16) { output_tokens: OUTPUT_TOKENS, realistic: false, engine: mock_worker::engine::EngineParams::default(), + replay: Default::default(), }); for rank in 0..u32::from(count) { tokio::spawn(mock_worker::zmq::serve(