feat(mock-worker): --capture writes each Generate request as a JSON line

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": <the string as sent>}`, 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,

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<PathBuf>,
 }
 
 impl Config {
     /// Parse the configuration from `std::env::args`, falling back to defaults.
     pub fn from_args() -> Result<Self, String> {
+        Self::parse(std::env::args().skip(1))
+    }
+
+    /// Parse the configuration from command-line flags, without the program
+    /// name.
+    fn parse(args: impl IntoIterator<Item = String>) -> Result<Self, String> {
         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 <path>       tokenizer path for gRPC autoload (default = model)\n\
        --gen-ms <ms>            canned per-request latency (default 0)\n\
        --output-tokens <n>      output tokens per request when unspecified (default 8)\n\
+       --capture <path>         append each gRPC Generate request to <path> as a JSON line\n\
      \n\
      Realistic engine simulator (continuous batching; opt-in):\n\
        --engine <canned|realistic>  engine mode (default canned)\n\
@@ -171,3 +192,26 @@ fn usage() -> String {
        --prefix-cache <bool>    enable prefix caching + KV events (default true)"
         .to_string()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn parse(flags: &[&str]) -> Result<Config, String> {
+        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"))
+        );
+    }
+}

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<Config>, 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<Config>, 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<Config>,
     /// Present iff the worker runs the realistic engine simulator.
     engine: Option<Engine>,
+    /// Present iff `--capture` is set: every `Generate` request is recorded.
+    capture: Option<Arc<Capture>>,
 }
 
 type GenStream = Pin<Box<dyn Stream<Item = Result<ts::GenerateResponse, Status>> + Send>>;
@@ -79,9 +99,18 @@ impl TokenSpeedScheduler for MockScheduler {
         &self,
         request: Request<ts::GenerateRequest>,
     ) -> Result<Response<Self::GenerateStream>, 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;
         }

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;

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,