diff --git a/.gitignore b/.gitignore index 7245a4d86..9bd42de79 100644 --- a/.gitignore +++ b/.gitignore @@ -157,3 +157,4 @@ clients/java/target/ clients/java/build.sbt clients/java/gradle.properties openapitools.json +scripts/generate_sim/profiles/*.local.json diff --git a/Cargo.lock b/Cargo.lock index 08a978420..22bfcac14 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7311,6 +7311,16 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "sim-loadgen" +version = "0.1.0" +dependencies = [ + "futures", + "reqwest 0.13.4", + "serde_json", + "tokio", +] + [[package]] name = "simd-adler32" version = "0.3.10" diff --git a/Cargo.toml b/Cargo.toml index 17dee0caa..7e14ca487 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [workspace] members = ["model_gateway", "crates/protocols", "crates/reasoning_parser", "crates/tool_parser", "crates/workflow", "crates/tokenizer", "crates/auth", "crates/mcp", - "crates/external_router", "crates/kv_index", "crates/data_connector", "crates/multimodal", "crates/mm_rdma", "crates/wasm", "crates/mesh", "crates/grpc_client", "crates/engine_zmq_client", "bindings/python", "bindings/golang", "clients/rust", "clients/openapi-gen", "crates/mock_worker", "crates/rl", "crates/radix_tree"] + "crates/external_router", "crates/kv_index", "crates/data_connector", "crates/multimodal", "crates/mm_rdma", "crates/wasm", "crates/mesh", "crates/grpc_client", "crates/engine_zmq_client", "bindings/python", "bindings/golang", "clients/rust", "clients/openapi-gen", "crates/mock_worker", "crates/rl", "crates/radix_tree", "crates/sim_loadgen"] resolver = "2" [workspace.dependencies] diff --git a/crates/mock_worker/src/engine.rs b/crates/mock_worker/src/engine.rs index 9da7cd204..e1cc33e3f 100644 --- a/crates/mock_worker/src/engine.rs +++ b/crates/mock_worker/src/engine.rs @@ -438,7 +438,8 @@ impl SchedulerState { }); } - /// Tokens currently resident in KV. + /// Tokens physically resident in KV (admission control and eviction + /// watermarks); see [`Self::pinned_tokens`] for what is REPORTED. fn used_tokens(&self, p: &EngineParams) -> u64 { if p.prefix_cache { // KV holds the shared radix cache (blocks persist across requests @@ -460,8 +461,24 @@ impl SchedulerState { } } + /// Tokens pinned by running requests (prompt + generated so far): the + /// KV a scheduler cannot evict. This — not physical occupancy — is what + /// engines report as `num_used_tokens` / `token_usage` (SGLang: used = + /// total − available − evictable), so a warm radix cache that keeps + /// KV physically full does not read as an overloaded worker. Reporting + /// occupancy instead trips a gateway's token-usage overload gate on + /// every warm worker and makes its cache-aware routing avoid exactly + /// the workers holding the prefixes (measured: same-worker follow-ups + /// fell from 0.87 to 0.15 over a 100 s run at 0.9 threshold). + fn pinned_tokens(&self) -> u64 { + self.running + .iter() + .map(|r| u64::from(r.prompt_tokens + r.generated)) + .sum() + } + fn snapshot(&self, p: &EngineParams) -> LoadSnapshot { - let used = self.used_tokens(p); + let used = self.pinned_tokens().min(p.kv_capacity_tokens); let waiting_uncached: i64 = self.waiting.iter().map(|w| w.uncached_tokens as i64).sum(); LoadSnapshot { num_running_reqs: self.running.len() as i32, @@ -866,6 +883,53 @@ mod tests { panic!("no token produced"); } + /// A warm prefix cache keeps KV physically occupied, but the reported + /// usage must fall back to the running requests' pinned tokens once + /// they finish — otherwise every warm worker reads as overloaded. + #[test] + fn reported_usage_excludes_evictable_cache() { + let p = EngineParams { + prefix_cache: true, + block_size: 4, + kv_capacity_tokens: 4096, + prefill_chunk_tokens: 1_000_000, + ..Default::default() + }; + let mut st = SchedulerState::new(); + let (r, mut rx) = req("a", vec![3; 256], 8); + st.enqueue(r, &p); + let mut done = false; + for _ in 0..10_000 { + let step = st.step(&p); + for (tx, ev) in step.sends { + let _ = tx.send(ev); + } + while let Ok(ev) = rx.try_recv() { + if matches!(ev, GenEvent::Done { .. }) { + done = true; + } + } + if done { + break; + } + } + assert!(done, "request did not finish"); + // Drain the completed request out of the batch. + let _ = st.step(&p); + assert!( + st.cache.len() >= 256 / 4, + "prefix cache retained the prompt blocks" + ); + assert!( + st.used_tokens(&p) >= 256, + "KV is physically occupied by the cache" + ); + let snap = st.snapshot(&p); + assert_eq!(snap.num_running_reqs, 0); + assert_eq!(snap.num_used_tokens, 0, "no running request pins KV"); + assert_eq!(snap.token_usage, 0.0); + } + #[test] fn ttft_scales_with_uncached_prompt_length() { let p = EngineParams { diff --git a/crates/mock_worker/src/http.rs b/crates/mock_worker/src/http.rs index 128700636..db40fd0d0 100644 --- a/crates/mock_worker/src/http.rs +++ b/crates/mock_worker/src/http.rs @@ -7,6 +7,13 @@ //! synthetic token ids (one per whitespace word) so shared text prefixes still //! produce cache hits and prompt length still drives prefill latency. Token-id //! KV events (event-driven `cache_aware`) remain a gRPC-path feature. +//! +//! `/generate` is SGLang-native in realistic mode: it reads `input_ids` when +//! the body carries them (text is the fallback), and answers in the native +//! shape (`output_ids` + `meta_info`), so a load generator that speaks the +//! production `/generate` contract scores cached tokens and rebuilds +//! multi-turn context from the worker's own output. In canned mode it stays +//! the historical chat-shaped alias the existing rigs depend on. use std::{ convert::Infallible, @@ -39,6 +46,9 @@ use crate::{ pub struct AppState { cfg: Arc, engine: Option, + /// The listener's port, echoed in native `meta_info.worker_port` so a + /// client can attribute a response to a worker without trusting routing. + port: u16, } /// Build the router serving the mock HTTP worker contract. @@ -48,7 +58,7 @@ pub fn router(state: Arc) -> Router { .route("/v1/models", get(models)) .route("/v1/chat/completions", post(chat_completions)) .route("/v1/completions", post(completions)) - .route("/generate", post(chat_completions)) + .route("/generate", post(generate)) .route("/v1/loads", get(loads)) .with_state(state) } @@ -64,7 +74,7 @@ pub async fn serve(cfg: Arc, host: String, port: u16) { }; // One simulated engine per listener (i.e. per virtual worker). let engine = cfg.realistic.then(|| Engine::spawn(cfg.engine.clone())); - let state = Arc::new(AppState { cfg, engine }); + let state = Arc::new(AppState { cfg, engine, port }); if let Err(e) = axum::serve(listener, router(state)).await { tracing::error!("http worker {port} stopped: {e}"); } @@ -142,6 +152,197 @@ async fn completions(State(state): State>, body: Bytes) -> Respons handle(Endpoint::Completions, state, body).await } +/// `/generate`: SGLang-native in realistic mode, chat-shaped alias otherwise. +async fn generate(State(state): State>, body: Bytes) -> Response { + if state.engine.is_none() { + return handle(Endpoint::Chat, state, body).await; + } + let parsed: Value = serde_json::from_slice(&body).unwrap_or(Value::Null); + drop(body); + let stream_requested = parsed + .get("stream") + .and_then(Value::as_bool) + .unwrap_or(false); + let prompt_ids = extract_input_ids(&parsed) + .unwrap_or_else(|| synth_token_ids(&extract_prompt_text(&parsed))); + let max_new = extract_native_max_new(&parsed).unwrap_or(state.cfg.output_tokens); + drop(parsed); + let request_id = next_request_id(); + let (tx, rx) = mpsc::unbounded_channel(); + state + .engine + .as_ref() + .expect("checked above") + .submit(NewRequest { + request_id: request_id.clone(), + prompt_token_ids: prompt_ids, + max_new, + events: tx, + }); + if stream_requested { + native_sse(rx, request_id, state.port).into_response() + } else { + Json(native_completion(rx, request_id, state.port).await).into_response() + } +} + +/// Everything a native response reports, accumulated from the engine's events. +struct NativeProgress { + request_id: String, + worker_port: u16, + prompt_tokens: u32, + cached_tokens: u32, + output_ids: Vec, + finished: bool, +} + +impl NativeProgress { + fn new(request_id: String, worker_port: u16) -> Self { + Self { + request_id, + worker_port, + prompt_tokens: 0, + cached_tokens: 0, + output_ids: Vec::new(), + finished: false, + } + } + + fn absorb(&mut self, ev: engine::GenEvent) { + match ev { + engine::GenEvent::Token { + token_id, + prompt_tokens, + cached_tokens, + } => { + self.prompt_tokens = prompt_tokens; + self.cached_tokens = cached_tokens; + self.output_ids.push(token_id); + } + engine::GenEvent::Done { + prompt_tokens, + cached_tokens, + .. + } => { + self.prompt_tokens = prompt_tokens; + self.cached_tokens = cached_tokens; + self.finished = true; + } + } + } + + /// The SGLang-native frame: `output_ids` and completion accounting are + /// reported once, on the terminal frame; earlier frames carry the prompt + /// accounting only (first frame = time to first token). + fn frame(&self) -> Value { + let completion = if self.finished { + self.output_ids.len() + } else { + 0 + }; + json!({ + "text": if self.finished { "mock" } else { "" }, + "output_ids": if self.finished { Value::from(self.output_ids.clone()) } else { Value::from(Vec::::new()) }, + "meta_info": { + "id": self.request_id, + "prompt_tokens": self.prompt_tokens, + "completion_tokens": completion, + "cached_tokens": self.cached_tokens, + "finish_reason": if self.finished { + json!({"type": "length", "length": completion}) + } else { + Value::Null + }, + "worker_port": self.worker_port, + }, + }) + } +} + +/// Non-streaming native `/generate`: one JSON object after the last token. +async fn native_completion( + mut rx: mpsc::UnboundedReceiver, + request_id: String, + worker_port: u16, +) -> Value { + let mut progress = NativeProgress::new(request_id, worker_port); + while let Some(ev) = rx.recv().await { + progress.absorb(ev); + } + progress.frame() +} + +/// Streaming native `/generate`: a first frame at the first token (so the +/// client's TTFT is the engine's), the terminal frame with `output_ids` at +/// completion, then `[DONE]`. Intermediate tokens are accumulated, not +/// streamed one per frame: the production `/generate` clients this mock +/// serves read the first and last frame, and a frame per token through the +/// gateway would make the mock fleet's SSE volume the bottleneck. +fn native_sse( + rx: mpsc::UnboundedReceiver, + request_id: String, + worker_port: u16, +) -> Sse>> { + enum St { + Active { + rx: mpsc::UnboundedReceiver, + progress: NativeProgress, + first_sent: bool, + }, + Closing, + Ended, + } + + let body = stream::unfold( + St::Active { + rx, + progress: NativeProgress::new(request_id, worker_port), + first_sent: false, + }, + |st| async move { + match st { + St::Active { + mut rx, + mut progress, + mut first_sent, + } => loop { + match rx.recv().await { + Some(ev) => { + progress.absorb(ev); + if progress.finished { + let frame = progress.frame(); + return Some(( + Ok(Event::default().data(frame.to_string())), + St::Closing, + )); + } + if !first_sent { + first_sent = true; + let frame = progress.frame(); + return Some(( + Ok(Event::default().data(frame.to_string())), + St::Active { + rx, + progress, + first_sent, + }, + )); + } + } + // Engine dropped the request without a terminal + // event: end the stream so the client sees an + // incomplete response, not a hang. + None => return None, + } + }, + St::Closing => Some((Ok(Event::default().data("[DONE]")), St::Ended)), + St::Ended => None, + } + }, + ); + Sse::new(body) +} + async fn handle(endpoint: Endpoint, state: Arc, body: Bytes) -> Response { let parsed: Option = serde_json::from_slice(&body).ok(); let stream_requested = parsed @@ -410,6 +611,31 @@ fn hash_word(w: &str) -> u32 { h % 30_000 } +/// Native `input_ids`: a flat token list, or a batch of one (`[[...]]`). +fn extract_input_ids(v: &Value) -> Option> { + let ids = v.get("input_ids")?.as_array()?; + let seq = match ids.first() { + Some(Value::Array(inner)) => inner, + _ => ids, + }; + Some( + seq.iter() + .filter_map(Value::as_u64) + .map(|id| id as u32) + .collect(), + ) +} + +/// Native limit: `sampling_params.max_new_tokens` first, then the top-level +/// OpenAI-style keys. +fn extract_native_max_new(v: &Value) -> Option { + v.get("sampling_params") + .and_then(|sp| sp.get("max_new_tokens")) + .and_then(Value::as_u64) + .map(|n| n as u32) + .or_else(|| extract_max_tokens(v)) +} + fn extract_max_tokens(v: &Value) -> Option { for key in ["max_tokens", "max_new_tokens"] { if let Some(n) = v.get(key).and_then(Value::as_u64) { @@ -423,3 +649,59 @@ fn next_request_id() -> String { static COUNTER: AtomicU64 = AtomicU64::new(0); format!("mock-http-{}", COUNTER.fetch_add(1, Ordering::Relaxed)) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn native_input_ids_accepts_flat_and_batched_shapes() { + let flat = json!({"input_ids": [1, 2, 3]}); + let batched = json!({"input_ids": [[4, 5]]}); + let text_only = json!({"text": "a b"}); + assert_eq!(extract_input_ids(&flat), Some(vec![1, 2, 3])); + assert_eq!(extract_input_ids(&batched), Some(vec![4, 5])); + assert_eq!(extract_input_ids(&text_only), None); + } + + #[test] + fn native_max_new_prefers_sampling_params() { + let native = json!({"sampling_params": {"max_new_tokens": 7}, "max_tokens": 99}); + let openai = json!({"max_tokens": 99}); + assert_eq!(extract_native_max_new(&native), Some(7)); + assert_eq!(extract_native_max_new(&openai), Some(99)); + assert_eq!(extract_native_max_new(&json!({})), None); + } + + #[test] + fn native_frames_report_output_only_when_finished() { + let mut p = NativeProgress::new("r1".into(), 9007); + p.absorb(engine::GenEvent::Token { + token_id: 11, + prompt_tokens: 100, + cached_tokens: 64, + }); + let first = p.frame(); + assert_eq!(first["output_ids"].as_array().unwrap().len(), 0); + assert_eq!(first["meta_info"]["completion_tokens"], 0); + assert_eq!(first["meta_info"]["cached_tokens"], 64); + assert!(first["meta_info"]["finish_reason"].is_null()); + p.absorb(engine::GenEvent::Token { + token_id: 12, + prompt_tokens: 100, + cached_tokens: 64, + }); + p.absorb(engine::GenEvent::Done { + finish_reason: "length", + prompt_tokens: 100, + completion_tokens: 2, + cached_tokens: 64, + }); + let last = p.frame(); + assert_eq!(last["output_ids"], json!([11, 12])); + assert_eq!(last["meta_info"]["completion_tokens"], 2); + assert_eq!(last["meta_info"]["prompt_tokens"], 100); + assert_eq!(last["meta_info"]["worker_port"], 9007); + assert_eq!(last["meta_info"]["finish_reason"]["type"], "length"); + } +} diff --git a/crates/sim_loadgen/Cargo.toml b/crates/sim_loadgen/Cargo.toml new file mode 100644 index 000000000..8ae54ff18 --- /dev/null +++ b/crates/sim_loadgen/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "sim-loadgen" +version = "0.1.0" +edition = "2021" +publish = false +description = "Open-loop /generate load generator simulating production ingress for SMG scale tests" + +[[bin]] +name = "sim-loadgen" +path = "src/main.rs" + +[dependencies] +tokio = { workspace = true, features = ["full"] } +reqwest = { workspace = true, features = ["stream", "json", "rustls", "http2"] } +serde_json.workspace = true +futures.workspace = true diff --git a/crates/sim_loadgen/src/args.rs b/crates/sim_loadgen/src/args.rs new file mode 100644 index 000000000..8ac7ed873 --- /dev/null +++ b/crates/sim_loadgen/src/args.rs @@ -0,0 +1,486 @@ +//! Runtime configuration for the load generator, parsed from CLI flags. + +/// Low anchors of the length CDFs: the token count at cumulative 0.0. +/// Live here (not in main) so flag validation can reject a first user +/// anchor at or below them — otherwise the inverse CDF's first segment +/// would interpolate DOWNWARD and stop being an inverse CDF. +pub const PROMPT_LOW_TOKENS: u32 = 256; +pub const OUTPUT_LOW_TOKENS: u32 = 16; + +/// How a session picks its SMG for turn 1. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Ingress { + /// Consistent choice: splitmix64 of the routing key modulo the URL count. + Hash, + /// A fresh uniform choice per request. + Random, +} + +impl Ingress { + pub fn as_str(self) -> &'static str { + match self { + Self::Hash => "hash", + Self::Random => "random", + } + } +} + +/// How a session picks its SMG for turn 2. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Turn2Ingress { + /// Reuse the SMG turn 1 landed on. + Same, + /// Consistent hash of the routing key (matches turn 1 under `--ingress hash`). + Hash, + /// A fresh uniform choice. + Random, +} + +impl Turn2Ingress { + pub fn as_str(self) -> &'static str { + match self { + Self::Same => "same", + Self::Hash => "hash", + Self::Random => "random", + } + } +} + +/// Wire format for the prompt in the `/generate` body. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Payload { + /// Pre-tokenized `input_ids` — the gateway routes on its token tree + /// (or the hash index / sticky override when those are configured). + Ids, + /// Untokenized `text`: the token context space-joined as decimal words. + /// The gateway routes on its approximate string tree; the mock worker + /// re-derives one stable id per word, so prefix reuse is preserved. + Text, +} + +impl Payload { + pub fn as_str(self) -> &'static str { + match self { + Self::Ids => "ids", + Self::Text => "text", + } + } +} + +/// Configuration for one load-generation run. +#[derive(Debug, Clone)] +pub struct Args { + /// SMG base URLs the generator spreads sessions across (required). + pub smg_urls: Vec, + /// Length of the session-arrival window; in-flight sessions still finish. + pub duration_secs: u64, + /// Poisson session arrival rate. + pub session_rps: f64, + /// Per-turn continue probability: after each turn the session sends + /// another with this probability (name kept from the two-turn contract, + /// where it was exactly the turn-2 probability). + pub t2_ratio: f64, + /// Hard cap on turns per session; context growth also ends a session + /// when the next turn would exceed `--prompt-max` (the model's + /// context-window limit stands in for both). + pub max_turns: u32, + /// Per-request client timeout; a wedged stream records status 0 instead + /// of hanging the end-of-run drain. + pub request_timeout_secs: u64, + /// Give every turn a fresh routing key (models clients that do not + /// carry a stable session key): each turn re-pins under the sticky + /// override and, under hash ingress, may land on a different SMG. + pub key_per_turn: bool, + /// Mean of the exponential think time between turn 1 and turn 2. + pub think_secs: f64, + /// Request SSE streaming responses (TTFT is only measurable when true). + pub stream: bool, + /// Speak HTTP/2 prior knowledge to the SMGs (multiplexed streams). + pub http2: bool, + /// Independent connections per SMG origin (requests round-robin across + /// them). In `--http2` mode each client multiplexes ONE connection per + /// origin, so this bounds concurrent streams per SMG; without it a + /// small gateway count throttles the generator, not the gateway. + pub conns_per_origin: usize, + /// Turn-1 SMG choice. + pub ingress: Ingress, + /// Turn-2 SMG choice. + pub turn2_ingress: Turn2Ingress, + /// Fraction of sessions using one of 32 shared routing keys instead of a + /// unique per-session key. + pub routing_key_reuse: f64, + /// Shared warm prefix length; 0 = per-session unique prefix (cold). + pub system_prefix_tokens: u32, + /// Number of distinct shared system prefixes ("agents"): each session + /// picks one, so the population reuses `system_prefix_pool` large + /// prefixes. 1 = a single global shared prefix (byte-identical to the + /// pre-pool behavior). + pub system_prefix_pool: u32, + /// Images per session. + pub image_count: u32, + /// Base64 characters per image payload. + pub image_bytes: usize, + /// Token id marking image positions in `input_ids`. + pub image_placeholder_id: u32, + /// Placeholder ids emitted per image. + pub image_placeholder_run: u32, + /// Fresh ids appended after the echoed turn-1 output in turn 2. + pub t2_suffix_tokens: u32, + /// Prompt-length CDF anchors as (tokens, cumulative) pairs. + pub prompt_cdf: Vec<(u32, f64)>, + /// Prompt length at cumulative 1.0 (the CDF tail anchor). + pub prompt_max: u32, + /// Output-length CDF anchors as (tokens, cumulative) pairs. + pub output_cdf: Vec<(u32, f64)>, + /// Output length at cumulative 1.0 (the CDF tail anchor). + pub output_max: u32, + /// Send `x-smg-routing-tokens` (first <=512 input ids) with each request. + pub tokens_hint: bool, + /// Prompt wire format: `ids` (default) or `text`. + pub payload: Payload, + /// `model` field for the request body; empty omits it. IGW-mode + /// gateways (gRPC worker legs) reject /generate without a model. + pub model: String, + /// Global cap on in-flight requests (a permit per request, not session). + pub max_inflight: usize, + /// Requests finishing this early are excluded from summary stats. + pub warmup_secs: u64, + /// Base seed; every random stream in the run derives from it. + pub seed: u64, + /// Output directory for requests.jsonl and summary.json. + pub out: String, +} + +impl Args { + /// Flag defaults; the single source `from_args` mutates and tests build + /// synthetic configs from. + pub(crate) fn defaults() -> Self { + Self { + smg_urls: Vec::new(), + duration_secs: 60, + session_rps: 5.0, + t2_ratio: 0.5, + think_secs: 30.0, + stream: true, + http2: false, + conns_per_origin: 4, + max_turns: 2, + request_timeout_secs: 300, + key_per_turn: false, + ingress: Ingress::Hash, + turn2_ingress: Turn2Ingress::Same, + routing_key_reuse: 0.0, + system_prefix_tokens: 2048, + system_prefix_pool: 1, + image_count: 1, + image_bytes: 620_000, + image_placeholder_id: 151_655, + image_placeholder_run: 256, + t2_suffix_tokens: 64, + prompt_cdf: vec![(5000, 0.216), (10_000, 0.530), (20_000, 0.994)], + prompt_max: 32_000, + output_cdf: vec![(1000, 0.351), (2000, 0.513), (5000, 0.998)], + output_max: 8192, + tokens_hint: false, + payload: Payload::Ids, + model: String::new(), + max_inflight: 200_000, + warmup_secs: 0, + seed: 42, + out: "sim_out".to_string(), + } + } + + /// Parse the configuration from `std::env::args`, falling back to defaults. + pub fn from_args() -> Result { + let mut cfg = Self::defaults(); + + let mut args = std::env::args().skip(1); + while let Some(flag) = args.next() { + match flag.as_str() { + "--smg-urls" => { + cfg.smg_urls = value(&mut args, &flag)? + .split(',') + .map(|url| url.trim().trim_end_matches('/').to_string()) + .filter(|url| !url.is_empty()) + .collect(); + } + "--duration-secs" => cfg.duration_secs = parse(value(&mut args, &flag)?, &flag)?, + "--session-rps" => cfg.session_rps = parse(value(&mut args, &flag)?, &flag)?, + "--t2-ratio" => cfg.t2_ratio = parse(value(&mut args, &flag)?, &flag)?, + "--think-secs" => cfg.think_secs = parse(value(&mut args, &flag)?, &flag)?, + "--stream" => cfg.stream = parse(value(&mut args, &flag)?, &flag)?, + "--http2" => cfg.http2 = parse(value(&mut args, &flag)?, &flag)?, + "--conns-per-origin" => { + cfg.conns_per_origin = parse(value(&mut args, &flag)?, &flag)?; + } + "--max-turns" => cfg.max_turns = parse(value(&mut args, &flag)?, &flag)?, + "--request-timeout-secs" => { + cfg.request_timeout_secs = parse(value(&mut args, &flag)?, &flag)?; + } + "--key-per-turn" => cfg.key_per_turn = parse(value(&mut args, &flag)?, &flag)?, + "--ingress" => { + cfg.ingress = match value(&mut args, &flag)?.as_str() { + "hash" => Ingress::Hash, + "random" => Ingress::Random, + other => return Err(format!("--ingress must be hash|random, got {other}")), + } + } + "--turn2-ingress" => { + cfg.turn2_ingress = match value(&mut args, &flag)?.as_str() { + "same" => Turn2Ingress::Same, + "hash" => Turn2Ingress::Hash, + "random" => Turn2Ingress::Random, + other => { + return Err(format!( + "--turn2-ingress must be same|hash|random, got {other}" + )) + } + } + } + "--routing-key-reuse" => { + cfg.routing_key_reuse = parse(value(&mut args, &flag)?, &flag)?; + } + "--system-prefix-tokens" => { + cfg.system_prefix_tokens = parse(value(&mut args, &flag)?, &flag)?; + } + "--system-prefix-pool" => { + cfg.system_prefix_pool = parse(value(&mut args, &flag)?, &flag)?; + } + "--image-count" => cfg.image_count = parse(value(&mut args, &flag)?, &flag)?, + "--image-bytes" => cfg.image_bytes = parse(value(&mut args, &flag)?, &flag)?, + "--image-placeholder-id" => { + cfg.image_placeholder_id = parse(value(&mut args, &flag)?, &flag)?; + } + "--image-placeholder-run" => { + cfg.image_placeholder_run = parse(value(&mut args, &flag)?, &flag)?; + } + "--t2-suffix-tokens" => { + cfg.t2_suffix_tokens = parse(value(&mut args, &flag)?, &flag)?; + } + "--prompt-cdf" => cfg.prompt_cdf = parse_cdf(&value(&mut args, &flag)?, &flag)?, + "--prompt-max" => cfg.prompt_max = parse(value(&mut args, &flag)?, &flag)?, + "--output-cdf" => cfg.output_cdf = parse_cdf(&value(&mut args, &flag)?, &flag)?, + "--output-max" => cfg.output_max = parse(value(&mut args, &flag)?, &flag)?, + "--tokens-hint" => cfg.tokens_hint = parse(value(&mut args, &flag)?, &flag)?, + "--payload" => { + cfg.payload = match value(&mut args, &flag)?.as_str() { + "ids" => Payload::Ids, + "text" => Payload::Text, + other => return Err(format!("--payload must be ids|text, got {other}")), + } + } + "--model" => cfg.model = value(&mut args, &flag)?, + "--max-inflight" => cfg.max_inflight = parse(value(&mut args, &flag)?, &flag)?, + "--warmup-secs" => cfg.warmup_secs = parse(value(&mut args, &flag)?, &flag)?, + "--seed" => cfg.seed = parse(value(&mut args, &flag)?, &flag)?, + "--out" => cfg.out = value(&mut args, &flag)?, + "-h" | "--help" => return Err(usage()), + other => return Err(format!("unknown flag: {other}\n\n{}", usage())), + } + } + + if cfg.smg_urls.is_empty() { + return Err(format!("--smg-urls is required\n\n{}", usage())); + } + if cfg.session_rps <= 0.0 || !cfg.session_rps.is_finite() { + return Err("--session-rps must be a positive finite number".to_string()); + } + for (flag, ratio) in [ + ("--t2-ratio", cfg.t2_ratio), + ("--routing-key-reuse", cfg.routing_key_reuse), + ] { + if !(0.0..=1.0).contains(&ratio) { + return Err(format!("{flag} must be within [0, 1], got {ratio}")); + } + } + if cfg.think_secs < 0.0 || !cfg.think_secs.is_finite() { + return Err("--think-secs must be a non-negative finite number".to_string()); + } + if cfg.max_inflight == 0 { + return Err("--max-inflight must be at least 1".to_string()); + } + // `build_clients` would clamp 0 to one client while summary.json + // echoed 0: the reported topology must match the executed one. + if cfg.conns_per_origin == 0 { + return Err("--conns-per-origin must be at least 1".to_string()); + } + // A warmup at or past the arrival window leaves the steady-state + // window empty, so every statistic would report null instead of + // failing. + if cfg.warmup_secs >= cfg.duration_secs { + return Err(format!( + "--warmup-secs ({}) must be less than --duration-secs ({})", + cfg.warmup_secs, cfg.duration_secs + )); + } + // A session always sends turn 1 before the cap is consulted, so a + // cap of 0 would silently behave as 1 instead of honoring the + // documented hard limit. + if cfg.max_turns == 0 { + return Err("--max-turns must be at least 1".to_string()); + } + // The session code treats any pool <= 1 as one global prefix, so + // 0 would quietly measure shared-prefix behavior under a flag + // that claims otherwise. + if cfg.system_prefix_pool == 0 { + return Err("--system-prefix-pool must be at least 1".to_string()); + } + for (max_flag, cdf_flag, cdf, max, low) in [ + ( + "--prompt-max", + "--prompt-cdf", + &cfg.prompt_cdf, + cfg.prompt_max, + PROMPT_LOW_TOKENS, + ), + ( + "--output-max", + "--output-cdf", + &cfg.output_cdf, + cfg.output_max, + OUTPUT_LOW_TOKENS, + ), + ] { + if let Some(&(tokens, _)) = cdf.first() { + if tokens <= low { + return Err(format!( + "{cdf_flag} first anchor ({tokens}) must exceed the fixed low anchor ({low} tokens)" + )); + } + } + if let Some(&(tokens, _)) = cdf.last() { + if max < tokens { + return Err(format!( + "{max_flag} ({max}) must be at least the last {cdf_flag} anchor ({tokens})" + )); + } + } + } + Ok(cfg) + } +} + +fn value(args: &mut impl Iterator, flag: &str) -> Result { + args.next() + .ok_or_else(|| format!("missing value for {flag}")) +} + +fn parse(raw: String, flag: &str) -> Result { + raw.parse() + .map_err(|_| format!("invalid value for {flag}: {raw}")) +} + +fn parse_cdf(raw: &str, flag: &str) -> Result, String> { + let mut anchors = Vec::new(); + for part in raw.split(',') { + let part = part.trim(); + let Some((tokens, cum)) = part.split_once(':') else { + return Err(format!( + "invalid {flag} anchor (want tokens:cumulative): {part}" + )); + }; + let tokens: u32 = tokens + .parse() + .map_err(|_| format!("invalid {flag} token count: {part}"))?; + let cum: f64 = cum + .parse() + .map_err(|_| format!("invalid {flag} cumulative value: {part}"))?; + anchors.push((tokens, cum)); + } + if anchors.is_empty() { + return Err(format!( + "{flag} needs at least one tokens:cumulative anchor" + )); + } + // `f64::from_str` accepts "NaN", which passes every ordered compare. + if anchors + .iter() + .any(|&(_, cum)| !cum.is_finite() || cum <= 0.0 || cum > 1.0) + { + return Err(format!("{flag} cumulative values must be within (0, 1]")); + } + if anchors + .windows(2) + .any(|pair| pair[1].0 <= pair[0].0 || pair[1].1 <= pair[0].1) + { + return Err(format!( + "{flag} anchors must be strictly increasing in both tokens and cumulative" + )); + } + Ok(anchors) +} + +fn usage() -> String { + "sim-loadgen — open-loop /generate load generator for SMG scale simulation\n\n\ + Flags:\n\ + --smg-urls SMG base URLs (required)\n\ + --duration-secs session-arrival window length (default 60)\n\ + --session-rps Poisson session arrival rate (default 5.0)\n\ + --t2-ratio probability a session sends turn 2 (default 0.5)\n\ + --think-secs mean exponential think time before turn 2 (default 30)\n\ + --stream request SSE streaming responses (default true)\n\ + --http2 HTTP/2 prior knowledge to the SMGs (default false)\n\ + --conns-per-origin connections per SMG, round-robined (default 4)\n\ + --max-turns turn cap per session; each turn continues with\n\ + probability --t2-ratio (default 2)\n\ + --request-timeout-secs per-request client timeout (default 300)\n\ + --key-per-turn fresh routing key every turn (default false)\n\ + --ingress turn-1 SMG choice (default hash)\n\ + --turn2-ingress turn-2 SMG choice (default same)\n\ + --routing-key-reuse fraction of sessions sharing one of 32 keys (default 0.0)\n\ + --system-prefix-tokens shared warm prefix length; 0 = cold (default 2048)\n\ + --system-prefix-pool distinct shared prefixes; 1 = single global (default 1)\n\ + --payload prompt wire format (default ids)\n\ + --model `model` field in the body; empty omits it\n\ + (required by IGW-mode gateways)\n\ + --image-count images per session (default 1)\n\ + --image-bytes base64 chars per image (default 620000)\n\ + --image-placeholder-id token id marking image positions (default 151655)\n\ + --image-placeholder-run placeholder ids per image (default 256)\n\ + --t2-suffix-tokens fresh ids appended in turn 2 (default 64)\n\ + --prompt-cdf prompt-length CDF anchors\n\ + (default 5000:0.216,10000:0.530,20000:0.994)\n\ + --prompt-max prompt length at cumulative 1.0 (default 32000)\n\ + --output-cdf output-length CDF anchors\n\ + (default 1000:0.351,2000:0.513,5000:0.998)\n\ + --output-max output length at cumulative 1.0 (default 8192)\n\ + --tokens-hint send x-smg-routing-tokens, first <=512 ids (default false)\n\ + --max-inflight global in-flight request cap (default 200000)\n\ + --warmup-secs exclude requests finishing this early from stats (default 0)\n\ + --seed base seed for all randomness (default 42)\n\ + --out output directory (default sim_out)" + .to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_cdf_rejects_non_finite_and_out_of_range_cumulatives() { + assert!(parse_cdf("100:NaN", "--prompt-cdf").is_err()); + assert!(parse_cdf("100:inf", "--prompt-cdf").is_err()); + assert!(parse_cdf("100:0", "--prompt-cdf").is_err()); + assert!(parse_cdf("100:1.5", "--prompt-cdf").is_err()); + assert_eq!( + parse_cdf("100:0.5,200:1.0", "--prompt-cdf").unwrap(), + vec![(100, 0.5), (200, 1.0)] + ); + } + + #[test] + fn usage_lists_every_accepted_flag() { + let text = usage(); + for flag in [ + "--system-prefix-pool", + "--payload", + "--model", + "--max-turns", + "--tokens-hint", + "--out", + ] { + assert!(text.contains(flag), "usage() omits {flag}"); + } + } +} diff --git a/crates/sim_loadgen/src/dist.rs b/crates/sim_loadgen/src/dist.rs new file mode 100644 index 000000000..0a0494652 --- /dev/null +++ b/crates/sim_loadgen/src/dist.rs @@ -0,0 +1,204 @@ +//! Deterministic randomness for the whole generator: a splitmix64 generator +//! plus the helpers derived from it (exponential arrivals, piecewise-linear +//! inverse CDFs, token-id and base64 streams). One algorithm everywhere makes +//! every run reproducible from `--seed` alone. + +/// Salts naming the independent streams derived from `--seed`. Values are +/// arbitrary but fixed; [`mix`] spreads them apart. +pub const SALT_ARRIVAL: u64 = 1; +pub const SALT_SESSION: u64 = 2; +pub const SALT_PREFIX: u64 = 3; +pub const SALT_PAD: u64 = 4; +pub const SALT_SUFFIX: u64 = 5; +pub const SALT_IMAGE: u64 = 6; + +/// Generated token ids stay below this bound — under the default image +/// placeholder id (151655), so padding never fakes an image run. +const TOKEN_ID_SPACE: u64 = 150_000; + +const GOLDEN_GAMMA: u64 = 0x9e37_79b9_7f4a_7c15; + +/// The splitmix64 output mixer. +pub fn mix(mut z: u64) -> u64 { + z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + z ^ (z >> 31) +} + +/// Derive an independent stream seed from a parent seed and a salt. +pub fn sub_seed(seed: u64, salt: u64) -> u64 { + mix(seed ^ mix(salt)) +} + +/// FNV-1a over the key bytes, finished with the splitmix64 mixer so +/// consecutive keys spread across a small modulus. +pub fn hash_str(s: &str) -> u64 { + let mut h: u64 = 0xcbf2_9ce4_8422_2325; + for &b in s.as_bytes() { + h ^= u64::from(b); + h = h.wrapping_mul(0x100_0000_01b3); + } + mix(h) +} + +/// splitmix64 sequence generator. +pub struct Rng { + state: u64, +} + +impl Rng { + pub fn new(seed: u64) -> Self { + Self { state: seed } + } + + pub fn next_u64(&mut self) -> u64 { + self.state = self.state.wrapping_add(GOLDEN_GAMMA); + mix(self.state) + } + + /// Uniform in [0, 1): the top 53 bits as a double. + pub fn next_f64(&mut self) -> f64 { + (self.next_u64() >> 11) as f64 * (1.0 / (1u64 << 53) as f64) + } + + /// Exponential with the given mean via inverse CDF; `1 - u` keeps the + /// argument of `ln` in (0, 1]. + pub fn next_exp(&mut self, mean: f64) -> f64 { + -(1.0 - self.next_f64()).ln() * mean + } + + /// Uniform index in [0, n). `n` must be non-zero. + pub fn next_index(&mut self, n: usize) -> usize { + (self.next_u64() % n as u64) as usize + } +} + +/// Piecewise-linear inverse CDF through (tokens, cumulative) anchors, with a +/// fixed low anchor at cumulative 0.0 and a tail anchor at 1.0. +pub struct PiecewiseCdf { + /// (cumulative, tokens) points, strictly increasing in both. + points: Vec<(f64, f64)>, +} + +impl PiecewiseCdf { + pub fn new(low_tokens: u32, anchors: &[(u32, f64)], max_tokens: u32) -> Self { + // Flag validation rejects this first (a user-facing error); the + // assert keeps the "strictly increasing" invariant of `points` + // honest for every other caller. + assert!( + anchors + .first() + .is_none_or(|&(tokens, _)| tokens > low_tokens), + "first CDF anchor {:?} must exceed the low anchor ({low_tokens})", + anchors.first() + ); + let mut points = vec![(0.0, f64::from(low_tokens))]; + for &(tokens, cum) in anchors { + points.push((cum, f64::from(tokens))); + } + // The tail segment reaches max_tokens at cumulative 1.0 unless the + // caller's last anchor is already there. + if points[points.len() - 1].0 < 1.0 { + points.push((1.0, f64::from(max_tokens))); + } + Self { points } + } + + /// Sample with `u` uniform in [0, 1), linearly interpolating between the + /// two anchors bracketing `u`. + pub fn sample(&self, u: f64) -> u32 { + for pair in self.points.windows(2) { + let (c0, t0) = pair[0]; + let (c1, t1) = pair[1]; + if u <= c1 { + let width = c1 - c0; + let t = if width > 0.0 { + t0 + (t1 - t0) * (u - c0) / width + } else { + t1 + }; + return (t.round() as u32).max(1); + } + } + (self.points[self.points.len() - 1].1.round() as u32).max(1) + } +} + +/// A deterministic run of token ids for the given stream seed. +pub fn token_ids(seed: u64, count: usize) -> Vec { + let mut rng = Rng::new(seed); + (0..count) + .map(|_| (rng.next_u64() % TOKEN_ID_SPACE) as u32) + .collect() +} + +const BASE64_ALPHABET: &[u8; 64] = + b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; + +/// A deterministic base64-alphabet string of exactly `len` characters, so the +/// same (session, image) seed reproduces byte-identical payloads across turns. +pub fn base64_blob(seed: u64, len: usize) -> String { + let mut rng = Rng::new(seed); + let mut out = String::with_capacity(len); + 'fill: loop { + // Ten 6-bit symbols per 64-bit draw; the top 4 bits are discarded. + let mut word = rng.next_u64(); + for _ in 0..10 { + if out.len() == len { + break 'fill; + } + out.push(char::from(BASE64_ALPHABET[(word & 63) as usize])); + word >>= 6; + } + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cdf_hits_its_anchors_exactly_and_stays_monotonic() { + let cdf = PiecewiseCdf::new( + 256, + &[(5000, 0.216), (10_000, 0.530), (20_000, 0.994)], + 32_000, + ); + // Inverse-CDF at each anchor's cumulative probability returns the + // anchor's token count — the property that makes the sampled + // distribution match the production percentiles. + assert_eq!(cdf.sample(0.216), 5000); + assert_eq!(cdf.sample(0.530), 10_000); + assert_eq!(cdf.sample(0.994), 20_000); + assert_eq!(cdf.sample(0.0), 256); + assert_eq!(cdf.sample(1.0), 32_000); + + let mut prev = 0; + for i in 0..=1000 { + let v = cdf.sample(f64::from(i) / 1000.0); + assert!(v >= prev, "inverse CDF must be monotonic"); + assert!((256..=32_000).contains(&v)); + prev = v; + } + } + + #[test] + fn derived_streams_are_deterministic_and_independent() { + // Same seed → identical stream (turn-2 image regeneration and + // cross-run reproducibility depend on this). + assert_eq!(token_ids(7, 32), token_ids(7, 32)); + assert_eq!(base64_blob(7, 64), base64_blob(7, 64)); + // Different salts under one seed give unrelated streams. + assert_ne!( + token_ids(sub_seed(42, SALT_PREFIX), 32), + token_ids(sub_seed(42, SALT_PAD), 32) + ); + // Blob is exactly the requested length and base64-alphabet only. + let blob = base64_blob(9, 257); + assert_eq!(blob.len(), 257); + assert!(blob + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'+' || b == b'/')); + } +} diff --git a/crates/sim_loadgen/src/main.rs b/crates/sim_loadgen/src/main.rs new file mode 100644 index 000000000..d504b4365 --- /dev/null +++ b/crates/sim_loadgen/src/main.rs @@ -0,0 +1,243 @@ +//! sim-loadgen: open-loop `/generate` load generator simulating production +//! ingress in front of N SMG replicas — Poisson session arrivals, paired +//! turn-1/turn-2 requests sharing prefixes, routing keys, and multimodal +//! payloads, with client-side TTFT/E2E measurement. See +//! `.claude/generate-scale-sim/01-design.md` for the workload profile. + +mod args; +mod dist; +mod report; +mod session; + +use std::{ + fs, + io::{BufWriter, Write}, + path::Path, + process::ExitCode, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }, + time::{Duration, Instant}, +}; + +use tokio::{ + sync::{mpsc, Semaphore}, + task::JoinSet, +}; + +use crate::{ + args::{Args, OUTPUT_LOW_TOKENS, PROMPT_LOW_TOKENS}, + dist::PiecewiseCdf, + report::RequestRecord, + session::Ctx, +}; + +const PROGRESS_INTERVAL_SECS: u64 = 10; + +#[tokio::main] +async fn main() -> ExitCode { + let cli = match Args::from_args() { + Ok(cli) => Arc::new(cli), + Err(message) => { + eprintln!("{message}"); + return ExitCode::from(2); + } + }; + + if let Err(e) = fs::create_dir_all(&cli.out) { + eprintln!("failed to create output directory {}: {e}", cli.out); + return ExitCode::from(2); + } + let jsonl_path = Path::new(&cli.out).join("requests.jsonl"); + let jsonl = match fs::File::create(&jsonl_path) { + Ok(file) => file, + Err(e) => { + eprintln!("failed to create {}: {e}", jsonl_path.display()); + return ExitCode::from(2); + } + }; + let clients = match build_clients(&cli) { + Ok(clients) => clients, + Err(e) => { + eprintln!("failed to build HTTP client: {e}"); + return ExitCode::from(2); + } + }; + + // The collector solely owns the JSONL file and the in-memory records; it + // drains until every sender (one per Ctx handle) is gone. + let (records_tx, mut records_rx) = mpsc::unbounded_channel::(); + // The collector reports whether the JSONL write failed: a truncated + // requests.jsonl must fail the run, not exit 0 — the harness reads + // that file for cache-hit and balance figures and would otherwise + // compute them over a partial record set that looks valid. + let mut collector: JoinSet<(Vec, bool)> = JoinSet::new(); + collector.spawn(async move { + let mut out = BufWriter::new(jsonl); + let mut records = Vec::new(); + let mut write_failed = false; + while let Some(record) = records_rx.recv().await { + if !write_failed && writeln!(out, "{}", record.to_json()).is_err() { + eprintln!("requests.jsonl write failed; keeping records in memory only"); + write_failed = true; + } + records.push(record); + } + write_failed |= out.flush().is_err(); + (records, write_failed) + }); + + let ctx = Arc::new(Ctx { + args: cli.clone(), + clients, + next_client: AtomicU64::new(0), + limiter: Arc::new(Semaphore::new(cli.max_inflight.min(Semaphore::MAX_PERMITS))), + records: records_tx, + prompt_cdf: PiecewiseCdf::new(PROMPT_LOW_TOKENS, &cli.prompt_cdf, cli.prompt_max), + output_cdf: PiecewiseCdf::new(OUTPUT_LOW_TOKENS, &cli.output_cdf, cli.output_max), + sent: AtomicU64::new(0), + done: AtomicU64::new(0), + errors: AtomicU64::new(0), + }); + + let run_start = Instant::now(); + let run_start_ms = session::epoch_ms(); + let mut aux: JoinSet<()> = JoinSet::new(); + { + let ctx = ctx.clone(); + aux.spawn(progress(ctx, run_start)); + } + + // Open-loop Poisson arrivals on an ABSOLUTE schedule: each arrival time + // is the running sum of exponential gaps, and the spawner sleeps until + // that instant. Sleeping per-gap instead would add the timer's minimum + // resolution to every gap and undershoot high rates by 20-30%. + let deadline = run_start + Duration::from_secs(cli.duration_secs); + let mut arrivals = dist::Rng::new(dist::sub_seed(cli.seed, dist::SALT_ARRIVAL)); + let mut sessions: JoinSet<()> = JoinSet::new(); + let mut spawned: u64 = 0; + let mut next_arrival = run_start; + loop { + let gap = arrivals.next_exp(1.0 / cli.session_rps); + next_arrival += Duration::from_secs_f64(gap); + if next_arrival >= deadline { + break; + } + tokio::time::sleep_until(next_arrival.into()).await; + sessions.spawn(session::run(ctx.clone(), spawned)); + spawned += 1; + } + + // The arrival window is closed; in-flight sessions (think times included) + // still run to completion. + while let Some(joined) = sessions.join_next().await { + if let Err(e) = joined { + eprintln!("session task panicked: {e}"); + } + } + aux.abort_all(); + while aux.join_next().await.is_some() {} + // The last live sender drops here, letting the collector drain and exit. + drop(ctx); + let (records, write_failed) = match collector.join_next().await { + Some(Ok(outcome)) => outcome, + _ => { + eprintln!("record collector failed; summary will be empty"); + (Vec::new(), true) + } + }; + + let elapsed_secs = run_start.elapsed().as_secs_f64(); + let summary = report::summarize(&cli, &records, run_start_ms, elapsed_secs, spawned); + let summary_path = Path::new(&cli.out).join("summary.json"); + match serde_json::to_string_pretty(&summary) { + Ok(text) => { + if let Err(e) = fs::write(&summary_path, text) { + eprintln!("failed to write {}: {e}", summary_path.display()); + return ExitCode::FAILURE; + } + } + Err(e) => { + eprintln!("failed to serialize summary: {e}"); + return ExitCode::FAILURE; + } + } + + let total = records.len(); + let errors = records.iter().filter(|r| !r.is_ok()).count(); + eprintln!( + "[sim-loadgen] finished: {spawned} sessions, {total} requests, {errors} errors \ + in {elapsed_secs:.1}s; wrote {} and {}", + jsonl_path.display(), + summary_path.display(), + ); + if total > 0 && errors * 2 > total { + eprintln!("[sim-loadgen] error rate {errors}/{total} exceeds 50%"); + return ExitCode::FAILURE; + } + if write_failed { + eprintln!( + "[sim-loadgen] {} is incomplete (write failure); the run is not usable", + jsonl_path.display() + ); + return ExitCode::FAILURE; + } + ExitCode::SUCCESS +} + +/// `--conns-per-origin` independent clients, round-robined per request. Each +/// keeps a large idle pool so the h1 mode reuses per-stream connections; +/// `--http2` multiplexes each client's streams over ONE connection per SMG, +/// so the client count bounds concurrent streams per origin — without it a +/// small gateway count throttles the generator instead of the gateway. +fn build_clients(cli: &Args) -> Result, reqwest::Error> { + (0..cli.conns_per_origin.max(1)) + .map(|_| { + let mut builder = reqwest::Client::builder() + .pool_idle_timeout(Some(Duration::from_secs(90))) + .pool_max_idle_per_host(4096) + // A wedged stream must fail, not hang the end-of-run drain. + .timeout(Duration::from_secs(cli.request_timeout_secs)) + .tcp_nodelay(true); + if cli.http2 { + // Hundreds of concurrent multi-hundred-KB uploads share each + // connection; the h2 defaults (64 KiB stream / 1 MiB conn + // windows) throttle body upload long before the gateway is + // the bottleneck and show up as server-side 408s. Explicit + // windows, NOT hyper's adaptive window: adaptive overrides + // the configured sizes with the 64 KiB default, and h2 + // (0.4.19+) sizes its small-DATA-frame budget from the + // configured connection window — at 64 KiB about 128 + // sub-256-byte frames buffered on one connection (the + // first-token frame and `[DONE]` of ~60 streams) kill it + // with ENHANCE_YOUR_CALM "too_many_data_frames". + builder = builder + .http2_prior_knowledge() + .http2_initial_stream_window_size(4 * 1024 * 1024) + .http2_initial_connection_window_size(32 * 1024 * 1024); + } + builder.build() + }) + .collect() +} + +/// Progress line to stderr every 10 s with the instantaneous completion rate. +async fn progress(ctx: Arc, run_start: Instant) { + let mut ticks = tokio::time::interval(Duration::from_secs(PROGRESS_INTERVAL_SECS)); + // The first tick fires immediately; skip it. + ticks.tick().await; + let mut last_done: u64 = 0; + loop { + ticks.tick().await; + let sent = ctx.sent.load(Ordering::Relaxed); + let done = ctx.done.load(Ordering::Relaxed); + let errors = ctx.errors.load(Ordering::Relaxed); + let rps = (done - last_done) as f64 / PROGRESS_INTERVAL_SECS as f64; + eprintln!( + "[sim-loadgen] t={:.0}s sent={sent} done={done} errors={errors} rps={rps:.1}", + run_start.elapsed().as_secs_f64(), + ); + last_done = done; + } +} diff --git a/crates/sim_loadgen/src/report.rs b/crates/sim_loadgen/src/report.rs new file mode 100644 index 000000000..4a6902eb2 --- /dev/null +++ b/crates/sim_loadgen/src/report.rs @@ -0,0 +1,485 @@ +//! Per-request records and the end-of-run summary: latency percentiles, +//! cache-hit ratios, turn-2 worker affinity, and the turn-1 worker spread. + +use std::collections::{BTreeMap, HashMap}; + +use serde_json::{json, Value}; + +use crate::args::Args; + +/// A request's `cached_tokens / prompt_tokens` at or above this counts as a +/// cache hit (the design's hit definition). +const CACHE_HIT_RATIO: f64 = 0.3; + +/// One completed (or failed) `/generate` request. +#[derive(Debug)] +pub struct RequestRecord { + pub turn: u8, + pub session: u64, + pub key: String, + pub smg: usize, + pub worker_port: Option, + /// Local count of the input ids sent, not the server's echo. + pub prompt_tokens: usize, + pub cached_tokens: Option, + pub completion_tokens: Option, + pub max_new: u32, + /// Elapsed to the first SSE data frame; `None` for non-streaming requests. + pub ttft_ms: Option, + pub e2e_ms: f64, + /// HTTP status; 0 for transport failures (connect or mid-stream). + pub status: u16, + /// Request start, milliseconds since the Unix epoch. + pub start_ms: u64, + /// Gateway's remote-index echo (`x-smg-index-source`): what the + /// prefetch resolved for this decision. `None` when the gateway runs + /// without a remote index. + pub index_source: Option, + /// Gateway's predicted cached tokens for the served worker + /// (`x-smg-index-predicted-tokens`); comparing against the worker's + /// actual `cached_tokens` separates index error from policy spill. + pub index_predicted_tokens: Option, +} + +impl RequestRecord { + pub fn is_ok(&self) -> bool { + (200..300).contains(&self.status) + } + + fn finish_ms(&self) -> u64 { + self.start_ms.saturating_add(self.e2e_ms as u64) + } + + fn cached_ratio(&self) -> Option { + let cached = self.cached_tokens? as f64; + if self.prompt_tokens == 0 { + return None; + } + Some(cached / self.prompt_tokens as f64) + } + + /// One JSONL line. + pub fn to_json(&self) -> Value { + json!({ + "turn": self.turn, + "session": self.session, + "key": self.key, + "smg": self.smg, + "worker_port": self.worker_port, + "prompt_tokens": self.prompt_tokens, + "cached_tokens": self.cached_tokens, + "completion_tokens": self.completion_tokens, + "max_new": self.max_new, + "ttft_ms": self.ttft_ms, + "e2e_ms": self.e2e_ms, + "status": self.status, + "start_ms": self.start_ms, + "index_source": self.index_source, + "index_predicted_tokens": self.index_predicted_tokens, + }) + } +} + +/// Build the summary document. Statistics cover the STEADY-STATE window +/// only: requests completing after `--warmup-secs` and before the arrival +/// window closes (`--duration-secs`). The drain tail — turns that finish +/// after arrivals stop — is reported separately and never mixed into +/// cache/latency/throughput comparisons. +pub fn summarize( + args: &Args, + records: &[RequestRecord], + run_start_ms: u64, + elapsed_secs: f64, + sessions: u64, +) -> Value { + let warmup_end_ms = run_start_ms.saturating_add(args.warmup_secs.saturating_mul(1000)); + let arrival_end_ms = run_start_ms.saturating_add(args.duration_secs.saturating_mul(1000)); + let measured: Vec<&RequestRecord> = records + .iter() + .filter(|r| { + let finish = r.finish_ms(); + finish >= warmup_end_ms && finish < arrival_end_ms + }) + .collect(); + let drain: Vec<&RequestRecord> = records + .iter() + .filter(|r| r.finish_ms() >= arrival_end_ms) + .collect(); + let measured_secs = (args.duration_secs.saturating_sub(args.warmup_secs)).max(1) as f64; + + let mut errors: BTreeMap = BTreeMap::new(); + for record in records { + if !record.is_ok() { + *errors.entry(record.status.to_string()).or_insert(0) += 1; + } + } + + let ok: Vec<&RequestRecord> = measured.iter().copied().filter(|r| r.is_ok()).collect(); + let ttfts: Vec = ok.iter().filter_map(|r| r.ttft_ms).collect(); + let e2es: Vec = ok.iter().map(|r| r.e2e_ms).collect(); + + // Turn-1 worker per session, from the whole run: a warmup turn 1 still + // anchors its session's turn-2 affinity. + let mut t1_ports: HashMap = HashMap::new(); + for record in records { + if record.turn == 1 && record.is_ok() { + if let Some(port) = record.worker_port { + t1_ports.entry(record.session).or_insert(port); + } + } + } + let t2_matches: Vec = ok + .iter() + .filter(|r| r.turn == 2) + .filter_map(|r| { + let port = r.worker_port?; + Some(port == *t1_ports.get(&r.session)?) + }) + .collect(); + let same_worker_rate = if t2_matches.is_empty() { + None + } else { + Some(t2_matches.iter().filter(|&&same| same).count() as f64 / t2_matches.len() as f64) + }; + + // Consecutive-turn stickiness across ALL follow-up turns: each turn t is + // compared with its session's turn t-1 worker (from the whole run). + let mut turn_ports: HashMap<(u64, u8), u64> = HashMap::new(); + for record in records { + if record.is_ok() { + if let Some(port) = record.worker_port { + turn_ports + .entry((record.session, record.turn)) + .or_insert(port); + } + } + } + let followup_matches: Vec = ok + .iter() + .filter(|r| r.turn >= 2) + .filter_map(|r| { + let port = r.worker_port?; + Some(port == *turn_ports.get(&(r.session, r.turn - 1))?) + }) + .collect(); + let followup_same_worker_rate = if followup_matches.is_empty() { + None + } else { + Some( + followup_matches.iter().filter(|&&same| same).count() as f64 + / followup_matches.len() as f64, + ) + }; + + // Mean turns per session, over sessions with at least one recorded turn. + let mut max_turn: HashMap = HashMap::new(); + for record in records { + let entry = max_turn.entry(record.session).or_insert(0); + *entry = (*entry).max(record.turn); + } + let mean_turns = if max_turn.is_empty() { + None + } else { + Some(max_turn.values().map(|&t| f64::from(t)).sum::() / max_turn.len() as f64) + }; + + let mut worker_counts: BTreeMap = BTreeMap::new(); + for record in ok.iter().filter(|r| r.turn == 1) { + if let Some(port) = record.worker_port { + *worker_counts.entry(port).or_insert(0) += 1; + } + } + let t1_total: u64 = worker_counts.values().sum(); + let distinct = worker_counts.len(); + let max_share = if t1_total > 0 { + worker_counts + .values() + .max() + .map(|&m| m as f64 / t1_total as f64) + } else { + None + }; + // Normalized over the observed workers; a single worker is complete + // concentration, so it reports 0 rather than dividing by ln(1). + let normalized_entropy = if t1_total == 0 { + None + } else if distinct > 1 { + let h: f64 = worker_counts + .values() + .map(|&c| { + let p = c as f64 / t1_total as f64; + -p * p.ln() + }) + .sum(); + Some(h / (distinct as f64).ln()) + } else { + Some(0.0) + }; + + let mut per_smg: Vec = vec![0; args.smg_urls.len()]; + for record in records { + if let Some(slot) = per_smg.get_mut(record.smg) { + *slot += 1; + } + } + let per_smg_requests: Vec = args + .smg_urls + .iter() + .zip(&per_smg) + .map(|(url, &requests)| json!({"url": url, "requests": requests})) + .collect(); + + json!({ + "config": config_json(args), + "totals": { + "sessions": sessions, + "requests": records.len(), + "errors": errors, + }, + "elapsed_secs": elapsed_secs, + // Steady-state window: [warmup, arrival end). Offered rate and the + // drain tail are reported separately so scenarios stay comparable. + "window": { + "start_secs": args.warmup_secs, + "end_secs": args.duration_secs, + "measured_requests": measured.len(), + }, + "offered_session_rps": args.session_rps, + "drain": { + "requests": drain.len(), + "ok": drain.iter().filter(|r| r.is_ok()).count(), + }, + "achieved_rps": measured.len() as f64 / measured_secs, + "ttft_ms": stats(&ttfts), + "e2e_ms": stats(&e2es), + "turns": { + "turn1": turn_block(&measured, Some(1)), + "turn2": turn_block(&measured, Some(2)), + "followup": turn_block(&measured, None), + }, + // All turns together; `cached_token_ratio` here is THE number to + // compare with backend cached-token telemetry. + "overall": overall_block(&measured), + "turn2_same_worker_rate": same_worker_rate, + "followup_same_worker_rate": followup_same_worker_rate, + "mean_turns_per_session": mean_turns, + "turn1_workers": { + "distinct": distinct, + "max_share": max_share, + "normalized_entropy": normalized_entropy, + }, + "per_smg_requests": per_smg_requests, + }) +} + +fn config_json(args: &Args) -> Value { + json!({ + "smg_urls": args.smg_urls, + "duration_secs": args.duration_secs, + "session_rps": args.session_rps, + "t2_ratio": args.t2_ratio, + "think_secs": args.think_secs, + "stream": args.stream, + "http2": args.http2, + "conns_per_origin": args.conns_per_origin, + "max_turns": args.max_turns, + "request_timeout_secs": args.request_timeout_secs, + "key_per_turn": args.key_per_turn, + "ingress": args.ingress.as_str(), + "turn2_ingress": args.turn2_ingress.as_str(), + "routing_key_reuse": args.routing_key_reuse, + "system_prefix_tokens": args.system_prefix_tokens, + "system_prefix_pool": args.system_prefix_pool, + "image_count": args.image_count, + "image_bytes": args.image_bytes, + "image_placeholder_id": args.image_placeholder_id, + "image_placeholder_run": args.image_placeholder_run, + "t2_suffix_tokens": args.t2_suffix_tokens, + "prompt_cdf": cdf_json(&args.prompt_cdf), + "prompt_max": args.prompt_max, + "output_cdf": cdf_json(&args.output_cdf), + "output_max": args.output_max, + "tokens_hint": args.tokens_hint, + "payload": args.payload.as_str(), + "model": args.model, + "max_inflight": args.max_inflight, + "warmup_secs": args.warmup_secs, + "seed": args.seed, + "out": args.out, + }) +} + +fn cdf_json(anchors: &[(u32, f64)]) -> Vec { + anchors + .iter() + .map(|&(tokens, cum)| json!([tokens, cum])) + .collect() +} + +/// Stats over every measured request regardless of turn. +fn overall_block(measured: &[&RequestRecord]) -> Value { + let ok: Vec<&RequestRecord> = measured.iter().copied().filter(|r| r.is_ok()).collect(); + let ratios: Vec = ok.iter().filter_map(|r| r.cached_ratio()).collect(); + let prompt_sum: u64 = ok + .iter() + .filter(|r| r.cached_tokens.is_some()) + .map(|r| r.prompt_tokens as u64) + .sum(); + let cached_sum: u64 = ok.iter().filter_map(|r| r.cached_tokens).sum(); + json!({ + "ok": ok.len(), + "prompt_tokens_sum": prompt_sum, + "cached_tokens_sum": cached_sum, + "cached_token_ratio": if prompt_sum > 0 { + Value::from(cached_sum as f64 / prompt_sum as f64) + } else { + Value::Null + }, + "cached_ratio_request_mean": mean(&ratios), + }) +} + +/// Stats over one exact turn (`Some(n)`) or every follow-up turn (`None`, +/// i.e. turn >= 2). +fn turn_block(measured: &[&RequestRecord], turn: Option) -> Value { + let of_turn: Vec<&RequestRecord> = measured + .iter() + .copied() + .filter(|r| match turn { + Some(t) => r.turn == t, + None => r.turn >= 2, + }) + .collect(); + let ok: Vec<&RequestRecord> = of_turn.iter().copied().filter(|r| r.is_ok()).collect(); + let ratios: Vec = ok.iter().filter_map(|r| r.cached_ratio()).collect(); + let hit_rate = if ratios.is_empty() { + None + } else { + Some(ratios.iter().filter(|&&r| r >= CACHE_HIT_RATIO).count() as f64 / ratios.len() as f64) + }; + // Token-weighted ratio (Σcached/Σprompt) is the number comparable with + // backend cached-token telemetry; the per-request mean is kept under its + // own name because long prompts weigh the two differently. + let prompt_sum: u64 = ok + .iter() + .filter(|r| r.cached_tokens.is_some()) + .map(|r| r.prompt_tokens as u64) + .sum(); + let cached_sum: u64 = ok.iter().filter_map(|r| r.cached_tokens).sum(); + let ttfts: Vec = ok.iter().filter_map(|r| r.ttft_ms).collect(); + let e2es: Vec = ok.iter().map(|r| r.e2e_ms).collect(); + json!({ + "count": of_turn.len(), + "ok": ok.len(), + "prompt_tokens_sum": prompt_sum, + "cached_tokens_sum": cached_sum, + "cached_token_ratio": if prompt_sum > 0 { + Value::from(cached_sum as f64 / prompt_sum as f64) + } else { + Value::Null + }, + "cached_ratio_request_mean": mean(&ratios), + "hit_rate": hit_rate, + "ttft_ms": stats(&ttfts), + "e2e_ms": stats(&e2es), + }) +} + +fn mean(values: &[f64]) -> Option { + if values.is_empty() { + None + } else { + Some(values.iter().sum::() / values.len() as f64) + } +} + +/// {mean, p50, p90, p99} with nearest-rank percentiles from a sorted copy. +fn stats(values: &[f64]) -> Value { + if values.is_empty() { + return json!({"mean": null, "p50": null, "p90": null, "p99": null}); + } + let mut sorted = values.to_vec(); + sorted.sort_unstable_by(f64::total_cmp); + let pct = |q: f64| { + let rank = ((q * (sorted.len() - 1) as f64).round() as usize).min(sorted.len() - 1); + sorted[rank] + }; + json!({ + "mean": sorted.iter().sum::() / sorted.len() as f64, + "p50": pct(0.50), + "p90": pct(0.90), + "p99": pct(0.99), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rec(turn: u8, session: u64, worker: u64, prompt: usize, cached: u64) -> RequestRecord { + RequestRecord { + turn, + session, + key: format!("sess-{session}"), + smg: 0, + worker_port: Some(worker), + prompt_tokens: prompt, + cached_tokens: Some(cached), + completion_tokens: Some(8), + max_new: 8, + ttft_ms: Some(10.0), + e2e_ms: 100.0, + status: 200, + start_ms: 1000, + index_source: None, + index_predicted_tokens: None, + } + } + + fn args() -> Args { + let mut args = Args::defaults(); + args.smg_urls = vec!["http://127.0.0.1:30000".to_string()]; + args + } + + #[test] + fn token_weighted_ratio_is_not_the_request_mean() { + // One fully-cached short prompt + one uncached long prompt: the + // request mean says 0.5, the token-weighted ratio says 0.1. Backend + // telemetry is token-weighted, so the summary must carry both under + // distinct names. + let records = vec![rec(1, 1, 9001, 100, 100), rec(1, 2, 9002, 900, 0)]; + let summary = summarize(&args(), &records, 0, 10.0, 2); + let t1 = &summary["turns"]["turn1"]; + assert_eq!(t1["prompt_tokens_sum"], 1000); + assert_eq!(t1["cached_tokens_sum"], 100); + let token_weighted = t1["cached_token_ratio"].as_f64().unwrap(); + let request_mean = t1["cached_ratio_request_mean"].as_f64().unwrap(); + assert!((token_weighted - 0.1).abs() < 1e-9); + assert!((request_mean - 0.5).abs() < 1e-9); + let overall = &summary["overall"]; + assert!((overall["cached_token_ratio"].as_f64().unwrap() - 0.1).abs() < 1e-9); + } + + #[test] + fn consecutive_turn_stickiness_and_mean_turns() { + // Session 1: three turns all on 9001 (two sticky follow-ups). + // Session 2: turn 2 moves workers (one non-sticky follow-up). + let records = vec![ + rec(1, 1, 9001, 1000, 0), + rec(2, 1, 9001, 2000, 1900), + rec(3, 1, 9001, 3000, 2900), + rec(1, 2, 9001, 1000, 0), + rec(2, 2, 9002, 2000, 100), + ]; + let summary = summarize(&args(), &records, 0, 10.0, 2); + let followup_rate = summary["followup_same_worker_rate"].as_f64().unwrap(); + assert!((followup_rate - 2.0 / 3.0).abs() < 1e-9); + // turn2_same_worker_rate keeps its original turn-2-vs-turn-1 meaning. + assert!((summary["turn2_same_worker_rate"].as_f64().unwrap() - 0.5).abs() < 1e-9); + let mean_turns = summary["mean_turns_per_session"].as_f64().unwrap(); + assert!((mean_turns - 2.5).abs() < 1e-9); + // Follow-up block spans turns 2 and 3. + assert_eq!(summary["turns"]["followup"]["ok"], 3); + } +} diff --git a/crates/sim_loadgen/src/session.rs b/crates/sim_loadgen/src/session.rs new file mode 100644 index 000000000..f522a96e9 --- /dev/null +++ b/crates/sim_loadgen/src/session.rs @@ -0,0 +1,559 @@ +//! Session lifecycle: build paired turn-1/turn-2 `/generate` requests, send +//! them through the shared client, and parse the SGLang-native responses +//! (single JSON or SSE frames). + +use std::{ + fmt::Write as _, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use futures::StreamExt; +use serde_json::Value; +use tokio::sync::{mpsc::UnboundedSender, Semaphore}; + +use crate::{ + args::{Args, Ingress, Payload, Turn2Ingress}, + dist::{self, PiecewiseCdf, Rng}, + report::RequestRecord, +}; + +/// Number of routing keys shared by `--routing-key-reuse` sessions. +const SHARED_ROUTING_KEYS: usize = 32; + +/// Most input ids a routing-tokens hint carries (the gateway caps it there). +const TOKENS_HINT_CAP: usize = 512; + +/// Shared state every session task needs. +pub struct Ctx { + pub args: Arc, + /// One client per configured connection-per-origin; requests + /// round-robin so h2 streams spread over several connections. + pub clients: Vec, + pub next_client: AtomicU64, + pub limiter: Arc, + pub records: UnboundedSender, + pub prompt_cdf: PiecewiseCdf, + pub output_cdf: PiecewiseCdf, + pub sent: AtomicU64, + pub done: AtomicU64, + pub errors: AtomicU64, +} + +/// Turn-1 SMG index for a routing key under hash ingress — the stand-in for +/// ingress consistent hashing, so it must be a pure function of the key. +pub(crate) fn hash_smg(key: &str, smg_count: usize) -> usize { + (dist::hash_str(key) % smg_count.max(1) as u64) as usize +} + +/// Whether the session sends another turn after `turn`. +pub(crate) fn session_continues(cont_draw: bool, turn: u32, max_turns: u32) -> bool { + cont_draw && turn < max_turns +} + +/// Whether the next turn's context (current ⊕ output ⊕ suffix) still fits +/// the model window; a session that would exceed it ends, standing in for +/// the context-length limit. +pub(crate) fn next_context_fits( + context_len: usize, + output_len: usize, + suffix_len: usize, + prompt_max: u32, +) -> bool { + context_len + output_len + suffix_len <= prompt_max as usize +} + +/// Milliseconds since the Unix epoch. +pub fn epoch_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} + +/// Run one session: turn 1, then after each turn continue with probability +/// `--t2-ratio` (up to `--max-turns`), each new turn extending the context +/// with the previous turn's returned output plus a fresh suffix. +pub async fn run(ctx: Arc, sid: u64) { + let args = &ctx.args; + let session_seed = dist::sub_seed(dist::sub_seed(args.seed, dist::SALT_SESSION), sid); + let mut rng = Rng::new(session_seed); + + let reuse_draw = rng.next_f64(); + let mut key = if reuse_draw < args.routing_key_reuse { + format!("shared-{}", rng.next_index(SHARED_ROUTING_KEYS)) + } else { + format!("sess-{sid}") + }; + + let prompt_len = ctx.prompt_cdf.sample(rng.next_f64()) as usize; + let max_new_1 = ctx.output_cdf.sample(rng.next_f64()); + + // Warm runs share one prefix stream; cold runs give each session its own, + // so no cross-session prefix ever matches. + let prefix_seed = if args.system_prefix_tokens > 0 { + // A population of shared system prompts ("agents"): each session + // picks one of `system_prefix_pool`, so a small set of large + // prefixes is reused across many sessions — the agentic/chat + // sharing shape. pool <= 1 keeps the single-global-prefix + // behavior byte-identical. + let base = dist::sub_seed(args.seed, dist::SALT_PREFIX); + if args.system_prefix_pool > 1 { + let agent = session_seed % args.system_prefix_pool as u64; + dist::sub_seed(base, agent) + } else { + base + } + } else { + dist::sub_seed(session_seed, dist::SALT_PREFIX) + }; + let mut input_ids = dist::token_ids(prefix_seed, args.system_prefix_tokens as usize); + let placeholder_total = args.image_placeholder_run as usize * args.image_count as usize; + input_ids.resize( + input_ids.len() + placeholder_total, + args.image_placeholder_id, + ); + let pad = prompt_len.saturating_sub(input_ids.len()); + input_ids.extend(dist::token_ids( + dist::sub_seed(session_seed, dist::SALT_PAD), + pad, + )); + + let n = args.smg_urls.len(); + let t1_smg = match args.ingress { + Ingress::Hash => hash_smg(&key, n), + Ingress::Random => rng.next_index(n), + }; + + // Turn loop: each turn extends the context with the previous turn's + // returned output plus a fresh suffix; the session ends on the continue + // draw, the turn cap, a failed turn, or when the next context would + // exceed the model window (`--prompt-max`). + let mut context = input_ids; + let mut turn: u32 = 1; + let mut smg = t1_smg; + loop { + let max_new = if turn == 1 { + max_new_1 + } else { + ctx.output_cdf.sample(rng.next_f64()) + }; + let output_ids = send_turn( + &ctx, + &TurnRequest { + sid, + session_seed, + turn: turn.min(u32::from(u8::MAX)) as u8, + key: &key, + smg, + input_ids: &context, + max_new, + }, + ) + .await; + + let cont = rng.next_f64() < args.t2_ratio; + let think = rng.next_exp(args.think_secs); + if !session_continues(cont, turn, args.max_turns) { + return; + } + // A failed turn has no output to extend; the session ends there. + let Some(output_ids) = output_ids else { + return; + }; + let suffix = args.t2_suffix_tokens as usize; + if !next_context_fits(context.len(), output_ids.len(), suffix, args.prompt_max) { + return; + } + tokio::time::sleep(Duration::from_secs_f64(think)).await; + + context.extend(output_ids); + context.extend(dist::token_ids( + dist::sub_seed( + dist::sub_seed(session_seed, dist::SALT_SUFFIX), + u64::from(turn), + ), + suffix, + )); + turn += 1; + // Clients without a stable session key present a fresh key each + // turn: the sticky override re-pins, and hash ingress re-hashes. + if args.key_per_turn { + key = format!("sess-{sid}-t{turn}"); + } + smg = match args.turn2_ingress { + Turn2Ingress::Same => t1_smg, + Turn2Ingress::Hash => hash_smg(&key, n), + Turn2Ingress::Random => rng.next_index(n), + }; + } +} + +struct TurnRequest<'a> { + sid: u64, + session_seed: u64, + turn: u8, + key: &'a str, + smg: usize, + input_ids: &'a [u32], + max_new: u32, +} + +/// Send one `/generate` request and record its outcome. Returns the returned +/// `output_ids` on success (empty when the response carried none), `None` on +/// any error. +async fn send_turn(ctx: &Ctx, req: &TurnRequest<'_>) -> Option> { + let args = &ctx.args; + let permit = match ctx.limiter.clone().acquire_owned().await { + Ok(permit) => permit, + // The semaphore is never closed; a close means shutdown. + Err(_) => return None, + }; + + // Image payloads are regenerated from the session seed inside the permit, + // so waiting sessions do not hold hundreds of KB, and turn 2 reproduces + // byte-identical bytes without storing them across the think time. + let images: Vec = (0..args.image_count) + .map(|i| { + dist::base64_blob( + dist::sub_seed(req.session_seed, dist::SALT_IMAGE + u64::from(i)), + args.image_bytes, + ) + }) + .collect(); + let body = build_body( + req.input_ids, + &images, + req.max_new, + args.stream, + args.payload, + &args.model, + ); + drop(images); + + let hint = args.tokens_hint.then(|| { + let head = &req.input_ids[..req.input_ids.len().min(TOKENS_HINT_CAP)]; + let mut joined = String::with_capacity(head.len() * 7); + for (i, id) in head.iter().enumerate() { + if i > 0 { + joined.push(','); + } + let _ = write!(joined, "{id}"); + } + joined + }); + + ctx.sent.fetch_add(1, Ordering::Relaxed); + let url = format!("{}/generate", args.smg_urls[req.smg]); + let start_ms = epoch_ms(); + let started = Instant::now(); + let pick = ctx.next_client.fetch_add(1, Ordering::Relaxed) as usize % ctx.clients.len(); + let mut request = ctx.clients[pick] + .post(&url) + .header("content-type", "application/json") + .header("x-smg-routing-key", req.key); + if let Some(hint) = &hint { + request = request.header("x-smg-routing-tokens", hint.as_str()); + } + + let mut status: u16 = 0; + let mut ttft_ms: Option = None; + let mut response: Option = None; + let mut index_source: Option = None; + let mut index_predicted_tokens: Option = None; + if let Ok(resp) = request.body(body).send().await { + status = resp.status().as_u16(); + // Remote-index echo headers (absent unless the gateway runs with + // --kv-indexer-url); captured before the body consumes `resp`. + index_source = resp + .headers() + .get("x-smg-index-source") + .and_then(|v| v.to_str().ok()) + .map(str::to_string); + index_predicted_tokens = resp + .headers() + .get("x-smg-index-predicted-tokens") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse().ok()); + if resp.status().is_success() { + if args.stream { + match consume_sse(resp, started).await { + Ok((ttft, last)) => { + ttft_ms = ttft; + response = last; + } + // A mid-stream transport failure is an incomplete + // request, not a success at the original status. + Err(_) => status = 0, + } + } else { + match resp.bytes().await { + // The gRPC router returns an array of responses for + // non-streaming /generate; the HTTP mock returns one + // object. Normalize to the first (only) element. + // An unparsable body is a failed request, not a + // success with no worker/cached fields: recorded as + // ok it would vanish from the cache-ratio denominator + // and a malformed responder would improve the numbers. + Ok(bytes) => match serde_json::from_slice::(&bytes) { + Ok(Value::Array(mut items)) if !items.is_empty() => { + response = Some(items.swap_remove(0)); + } + Ok(other) => response = Some(other), + Err(_) => status = 0, + }, + Err(_) => status = 0, + } + } + } else { + // Drain the error body so the connection can be reused. + let _ = resp.bytes().await; + } + } + let e2e_ms = started.elapsed().as_secs_f64() * 1000.0; + drop(permit); + + let mut worker_port = None; + let mut cached_tokens = None; + let mut completion_tokens = None; + let mut output_ids: Option> = None; + if let Some(value) = &response { + let meta = &value["meta_info"]; + // The HTTP mock injects its port directly; the gRPC router carries + // no worker identity, so gRPC legs register each worker with a + // `weight_version` label of its port, relayed here verbatim. + worker_port = meta["worker_port"].as_u64().or_else(|| { + meta["weight_version"] + .as_str() + .and_then(|v| v.parse::().ok()) + }); + cached_tokens = meta["cached_tokens"].as_u64(); + completion_tokens = meta["completion_tokens"].as_u64(); + output_ids = value["output_ids"].as_array().map(|ids| { + ids.iter() + .filter_map(|id| id.as_u64().map(|id| id as u32)) + .collect() + }); + } + + let is_error = !(200..300).contains(&status); + if is_error { + ctx.errors.fetch_add(1, Ordering::Relaxed); + } + ctx.done.fetch_add(1, Ordering::Relaxed); + + let record = RequestRecord { + turn: req.turn, + session: req.sid, + key: req.key.to_string(), + smg: req.smg, + worker_port, + prompt_tokens: req.input_ids.len(), + cached_tokens, + completion_tokens, + max_new: req.max_new, + ttft_ms, + e2e_ms, + status, + start_ms, + index_source, + index_predicted_tokens, + }; + // A send error means the collector is gone (shutdown); nothing to do. + let _ = ctx.records.send(record); + + if is_error { + None + } else { + Some(output_ids.unwrap_or_default()) + } +} + +/// Serialize the request body by hand: the multi-hundred-KB image strings are +/// appended directly instead of being copied through an intermediate +/// `serde_json::Value` (the base64 alphabet needs no JSON escaping). +/// +/// `Payload::Text` sends the same token context as a `text` field of +/// space-joined decimal words instead of `input_ids`. Appending tokens +/// appends text, so a turn's context is a string prefix of the next turn's +/// and prefix matching survives the format change. The gateway then routes +/// on its approximate string tree; the mock re-derives one stable id per +/// word. Image placeholders are only expanded on the ids path, so text runs +/// should use `--image-count 0`. +fn build_body( + input_ids: &[u32], + images: &[String], + max_new: u32, + stream: bool, + payload: Payload, + model: &str, +) -> String { + let image_len: usize = images.iter().map(|image| image.len() + 3).sum(); + let mut body = String::with_capacity(image_len + input_ids.len() * 7 + 128); + body.push('{'); + if !model.is_empty() { + let _ = write!(body, "\"model\":\"{model}\","); + } + match payload { + Payload::Ids => { + body.push_str("\"input_ids\":["); + for (i, id) in input_ids.iter().enumerate() { + if i > 0 { + body.push(','); + } + let _ = write!(body, "{id}"); + } + body.push(']'); + } + Payload::Text => { + body.push_str("\"text\":\""); + for (i, id) in input_ids.iter().enumerate() { + if i > 0 { + body.push(' '); + } + let _ = write!(body, "{id}"); + } + body.push('"'); + } + } + if !images.is_empty() { + body.push_str(",\"image_data\":["); + for (i, image) in images.iter().enumerate() { + if i > 0 { + body.push(','); + } + body.push('"'); + body.push_str(image); + body.push('"'); + } + body.push(']'); + } + let _ = write!( + body, + ",\"sampling_params\":{{\"max_new_tokens\":{max_new}}},\"stream\":{stream}}}" + ); + body +} + +/// Drain an SSE response: TTFT is the elapsed time to the first data frame, +/// and the last data frame before `[DONE]` carries the same full JSON shape +/// as a non-streaming response. Frames may span chunks, so bytes are buffered +/// and consumed line by line. +async fn consume_sse( + resp: reqwest::Response, + started: Instant, +) -> Result<(Option, Option), reqwest::Error> { + let mut stream = resp.bytes_stream(); + let mut buf: Vec = Vec::new(); + let mut ttft_ms = None; + let mut last = None; + 'read: while let Some(chunk) = stream.next().await { + buf.extend_from_slice(&chunk?); + while let Some(newline) = buf.iter().position(|&b| b == b'\n') { + let line: Vec = buf.drain(..=newline).collect(); + let line = line.strip_suffix(b"\n").unwrap_or(&line); + let line = line.strip_suffix(b"\r").unwrap_or(line); + let Some(data) = line + .strip_prefix(b"data: ") + .or_else(|| line.strip_prefix(b"data:")) + else { + continue; + }; + if data == b"[DONE]".as_slice() { + break 'read; + } + if ttft_ms.is_none() { + ttft_ms = Some(started.elapsed().as_secs_f64() * 1000.0); + } + if let Ok(value) = serde_json::from_slice::(data) { + last = Some(value); + } + } + } + // Drain any trailing bytes after [DONE] so the connection can be reused. + while let Some(chunk) = stream.next().await { + if chunk.is_err() { + break; + } + } + Ok((ttft_ms, last)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn hash_ingress_is_deterministic_and_spreads() { + // Same key always maps to the same SMG (ingress stickiness)... + for key in ["sess-1", "sess-2", "shared-7"] { + assert_eq!(hash_smg(key, 8), hash_smg(key, 8)); + } + // ...and distinct keys reach every SMG (no degenerate hashing). + let mut seen = [false; 8]; + for sid in 0..256 { + seen[hash_smg(&format!("sess-{sid}"), 8)] = true; + } + assert!(seen.iter().all(|&s| s), "keys must spread over all SMGs"); + // A single SMG never divides by zero. + assert_eq!(hash_smg("sess-1", 1), 0); + } + + #[test] + fn session_continuation_respects_draw_and_turn_cap() { + assert!(session_continues(true, 1, 2), "draw passed, under cap"); + assert!( + !session_continues(false, 1, 8), + "failed draw ends the session" + ); + assert!(!session_continues(true, 2, 2), "turn cap ends the session"); + assert!(session_continues(true, 7, 8)); + assert!(!session_continues(true, 8, 8)); + } + + #[test] + fn context_window_caps_session_growth() { + assert!(next_context_fits(10_000, 2000, 256, 24_576)); + assert!(!next_context_fits(23_000, 2000, 256, 24_576)); + // Exactly at the window still fits. + assert!(next_context_fits(24_000, 500, 76, 24_576)); + } + + #[test] + fn model_field_prepends_only_when_set() { + let with: serde_json::Value = + serde_json::from_str(&build_body(&[1], &[], 2, false, Payload::Ids, "mock-model")) + .unwrap(); + assert_eq!(with["model"], "mock-model"); + let without: serde_json::Value = + serde_json::from_str(&build_body(&[1], &[], 2, false, Payload::Ids, "")).unwrap(); + assert!( + without.get("model").is_none(), + "empty model must be omitted" + ); + } + + #[test] + fn text_payload_is_valid_json_and_prefix_preserving() { + let turn1 = build_body(&[12, 3], &[], 8, false, Payload::Text, ""); + let turn2 = build_body(&[12, 3, 45], &[], 8, false, Payload::Text, ""); + let v1: serde_json::Value = serde_json::from_str(&turn1).unwrap(); + let v2: serde_json::Value = serde_json::from_str(&turn2).unwrap(); + assert_eq!(v1["text"], "12 3"); + assert_eq!(v2["text"], "12 3 45"); + assert!(v1.get("input_ids").is_none(), "text mode must omit ids"); + // Extending the token context extends the text — the string tree + // sees turn 1 as a prefix of turn 2. + assert!(v2["text"] + .as_str() + .unwrap() + .starts_with(v1["text"].as_str().unwrap())); + // The ids path is untouched by the payload switch. + let ids: serde_json::Value = + serde_json::from_str(&build_body(&[12, 3], &[], 8, false, Payload::Ids, "")).unwrap(); + assert_eq!(ids["input_ids"], serde_json::json!([12, 3])); + } +} diff --git a/scripts/generate_sim/README.md b/scripts/generate_sim/README.md new file mode 100644 index 000000000..285c0c599 --- /dev/null +++ b/scripts/generate_sim/README.md @@ -0,0 +1,248 @@ +# /generate scale simulation + +Local-only reproduction of the sanitized production `/generate` workload +(design: `.claude/generate-scale-sim/01-design.md`): K SMG replicas in front of +a mock worker fleet (`crates/mock_worker --engine realistic`, SGLang-native +`/generate`), driven by the +open-loop `sim-loadgen` crate. No Kubernetes, no production access. + +## Quick start + +```sh +# Smoke run on a laptop (24 workers, 2 SMGs, 40 s): +python3 scripts/generate_sim/sim.py run \ + --profile scripts/generate_sim/profiles/smoke.json + +# Laptop profile at production-like worker pressure (~450 req/s, 120 workers, 2.5 min): +python3 scripts/generate_sim/sim.py run \ + --profile scripts/generate_sim/profiles/local-small.json + +# Reuse existing binaries, tweak a knob without editing the profile: +python3 scripts/generate_sim/sim.py run --profile ... --skip-build \ + --override loadgen.ingress=random --override smg_count=1 + +# Rebuild the report for a finished (or aborted) run: +python3 scripts/generate_sim/sim.py report --run-dir target/generate-sim/ + +# Named comparisons (side-by-side compare.md): +python3 scripts/generate_sim/scenarios.py list +python3 scripts/generate_sim/scenarios.py compare \ + --scenario stable-key-vs-random-ingress \ + --profile scripts/generate_sim/profiles/local-small.json +``` + +A run builds `smg`, `mock-worker`, and `sim-loadgen` (release), launches the +fleet, registers every worker with every SMG (`POST /workers`, +`disable_health_check`, 64-way per SMG), gates on `GET /workers` reaching 99%, +warms up, drives the load, samples each SMG every 5 s (RSS/%CPU via `ps`, fds +via `lsof` or `/proc`, plus `/metrics` admission-queue and connection gauges), +then tears everything down and writes `report.json` / `report.md` into the run +dir (default `target/generate-sim/-/`, with per-process +logs under `logs/`). + +## Port plan + +| range | use | +|---|---| +| 9000 .. 9000+workers-1 | mock workers (template `workers_total` 10000: 9000–18999) | +| 30000 .. 30000+K-1 | SMG data ports | +| 39000 .. 39000+K-1 | SMG prometheus ports | +| 40000 .. 40000+replicas-1 | index replicas (profiles with `index_service`) | +| 40100 .. | index replica metrics / readiness | +| 40200 .. | severable peer proxies (`index_service.partitionable`) | +| 41000 .. 41000+K-1 | gateway mesh ports (`mesh_smgs`) | + +Everything binds 127.0.0.1. Runs `pkill` leftover `smg`/`mock-worker`/ +`sim-loadgen`/`radix-index-service`/`radix-index-bridge` processes started +from the same binary paths before and after, so ports are free across runs. + +## Profiles + +Plain JSON consumed by `sim.py`; `profiles/local-small.json` is the schema by +example. Every unknown production property is an explicit knob — never edit +the harness to change the workload: + +- top level: `smg_count`, `workers_total`, `mock_processes`, `duration_secs`, + `warmup_secs`, `sample_interval_secs`, `sample_fds`, `readiness_*`. +- `mock`: forwarded to `mock-worker` as `--key-with-hyphens value` + (`decode_base_ms`, `decode_per_req_ms`, `prefill_tps`, `max_running`, + `kv_tokens`, `block_size`, `prefix_cache`). Image bytes in request bodies + are payload only: the engine counts prompt tokens, not images. +- `smg_flags`: the literal gateway argv tail — the design doc's + production-equivalent cache_aware set. The harness adds only + `--host/--port/--prometheus-*`; it never passes `--enable-igw`. +- `loadgen`: forwarded to `sim-loadgen` the same way (`session_rps`, + `t2_ratio`, `think_secs`, `system_prefix_tokens`, `prompt_cdf`/`output_cdf`, + `image_bytes`, `image_count`, `routing_key_reuse`, `ingress`, + `turn2_ingress`, `tokens_hint`, `stream`, `http2`). Booleans are + value-style (`--stream true`), matching mock-worker's `--prefix-cache`. + The harness adds `--smg-urls`, `--duration-secs`, and `--out` (the run + dir, where `requests.jsonl` and `summary.json` land). +- `index_service` (optional): launches the shared radix-index service and, + by default, its event bridge — `replicas`, `bridge`, `inferred_ttl_secs`, + `event_ttl_secs`, `default_capacity_blocks`, `sweep_interval_secs`, + `apply_delay_stored_ms` / `apply_delay_removed_ms` (staleness injection), + `partitionable` (peer traffic via severable proxies), `deferred_replicas` + (listed as peers, launched by a drill). Needs the `smg-radix-index` + package in the tree; the build step adds it only when this block is set. +- `mesh_smgs`: gateway-to-gateway TreeSync mesh (one mesh port per SMG). + +## Fault drills + +Top-level profile keys, each firing once mid-run on its own thread; every +drill records epoch-ms timestamps of what it did in `meta.json` (surfaced in +`report.md` under "Drills"), and a drill that fails records `_error` +instead of dying silently. Any key that looks like a drill but is not one of +these fails the run at start, so a leg can never silently measure nothing. + +| key | shape | effect | +|---|---|---| +| `restart_smgs_at_secs` | seconds | kill and relaunch every SMG (sticky pins and placements lost) | +| `kill_index_replica` | `{at_secs, replica, relaunch_after_secs?}` | SIGKILL a replica; relaunch bootstrapping from a survivor | +| `flap_index_replica` | `{at_secs, replica, cycles, period_secs}` | kill + relaunch the replica `cycles` times | +| `hang_index_replica` | `{at_secs, replica, resume_after_secs?}` | SIGSTOP (TCP up, nothing drains), then SIGCONT | +| `start_deferred_replica` | `{at_secs, replica}` | launch a `deferred_replicas` member under load, bootstrapping from replica 0 | +| `partition_drill` | `{at_secs, heal_after_secs?}` | sever every inter-replica link (needs `partitionable`), then heal | +| `remove_workers_drill` | `{at_secs, count}` | deregister the last `count` workers from every gateway (DELETE /workers); listeners stay up | +| `add_workers_drill` | `{at_secs, count}` | start `count` NEW workers on fresh ports and register them with every gateway | +| `restart_one_smg_drill` | `{at_secs, smg}` | kill ONE gateway and relaunch it cold, re-registering the fleet with it | +| `rolling_replica_restart_drill` | `{at_secs, gap_secs}` | kill and relaunch every replica in turn, each bootstrapping from a live peer | +| `gateway_partition_drill` | `{at_secs, heal_after_secs?, scope: all\|half}` | sever the gateways' link to replica 0 (needs `gateway_proxy`; `half` = even-numbered gateways only) | + +`index_service.anti_entropy_secs` is forwarded to the service (peer digest +exchange period; 0 disables it). Two more `index_service` keys support the audits: `gateway_proxy` routes every +gateway's `--kv-indexer-url` through a severable proxy (even gateways one, +odd gateways another), and `dump_on_exit` pulls every live replica's state +with `radix-index-dump` before teardown and records per-holder divergence +(`index_divergence` in `meta.json` / the `replicas converged at end` row). +Differing holders are classified: event-fed holders are sequenced ground +truth and must be identical; placement-fed holders are copied, never agreed +on, so two replicas both holding a worker above its capacity (the mock's +`kv_tokens / block_size`) legitimately differ by when each ran its capacity +cut — counted as `in capacity band (cut timing)` — while a placement holder +differing under capacity is a lost update. `converged` is true only when no +event-fed holder differs, no holder is missing on a replica, and no placement +holder differs under capacity. + +Rows the compare table adds for index legs: gateway-side lookup latency +p50/p90/p99 (what the 2 ms deadline is measured against), service-side apply +and query engine time p50/p90/p99, index replica peak RSS and mean CPU, +prediction error bias / exact share / p50 / p90 / p95 / max, and the +follow-up cache ratio in the first and last full minute (drift). + +Scenarios `eval-matrix` (regimes × gateways × concurrency), `io-shapes`, +`chaos` and `soak` are the side-by-side evaluation; `production` in those +legs is the shipping configuration (hash placement index + sticky routing +key, HTTP workers). + +`failover_bins.py ` bins follow-up cache ratios around the +kill instant of a `kill_index_replica` run. + +Aggregate request rps = `session_rps × (1 + t2_ratio)`; the local profiles +compress time 10× (`decode_base_ms` 4.3, `prefill_tps` 80000 → ~8.9 s mean +lifetime; `prefill_chunk` 320 keeps one step's prefill under one decode +step, since the engine's step time is max(prefill chunk, decode) — a +larger chunk stretches every decode step while any prompt is prefilling) and bodies 10× (`image_bytes` 62000) together, per the design doc. + +## Mock engine fidelity notes + +Two properties of `crates/mock_worker`'s realistic engine decide whether a +cache-aware result means anything, and both are set by this harness: + +- `prefill_chunk` 320 (see the compression note below): the engine's step + time is max(prefill chunk, decode step), so a large chunk stretches every + decode step while any prompt is prefilling. +- Reported usage is the running requests' pinned tokens, not physical KV + occupancy (SGLang's definition: used = total − available − evictable). A + warm radix cache keeps KV physically full; reporting that as usage tripped + the gateway's `--worker-overload-token-usage 0.9` gate on every warm + worker and made cache-aware routing avoid the workers holding the + prefixes (same-worker follow-ups 0.87 → 0.15 over 100 s). + +## Load generator transport + +The profiles drive the gateways over h2 (`loadgen.http2: true`), one +multiplexed connection per client per gateway, so a high request rate does +not exhaust ephemeral ports the way HTTP/1.1 connection churn does (~16k on +macOS; 600+ sessions/s over h1 hit it). This needs the gateway fix in #2488: +before it, every streamed `/generate` response relayed by the gateway ended +with a zero-length non-terminal DATA frame, and h2 ≥ 0.4.16 clients close a +connection after 100 of those — a 60% transport-error rate under load that +had nothing to do with routing. Set `loadgen.http2: false` to measure +against a gateway without that fix. + +## File descriptors / ulimit + +`sim.py` raises `RLIMIT_NOFILE` to the hard limit before spawning children +(they inherit it). If the hard limit itself is low: + +- macOS: `sudo launchctl limit maxfiles 65536 1048576`, then a new shell; the + per-process cap is `kern.maxfilesperproc`. +- Linux: raise `nofile` in `/etc/security/limits.conf` or run under + `prlimit --nofile=1048576`. + +Budget: each mock process holds its listeners plus accepted upstream conns; +each SMG holds ~1 h2c conn per worker (`--upstream-http2`) plus client conns. +The full profile needs ≥200k fds system-wide; `local-*` fit default-raised +laptop limits. `lsof`-based fd sampling is slow at high fd counts — the full +profile sets `"sample_fds": false`. + +## Full-profile host sizing + +`profiles/full.template.json` carries generic round placeholders — copy it to +`profiles/full.local.json` (gitignored) and fill in your fleet's real worker +count, request rate, timing, and body sizes; never commit those values. A +full-scale run belongs on a large Linux host, not a laptop. Scaling rules of +thumb for K SMGs, W workers, R req/s, mean body B, mean lifetime L: + +- Aggregate body ingest ≈ R × B, all loopback. +- Concurrent client h2 streams ≈ R × L — run the loadgen with `http2: true` + and size `conns_per_origin`; h1 would need one socket per stream. +- SMG→worker connections ≈ K × W (one h2c conn per pair with + `--upstream-http2`). +- Registration is K × W `POST /workers` calls (64-way per SMG, SMGs in + parallel); allow a generous readiness timeout. +- Body memory depends on the routing regime: with the sticky override and a + valid routing key, bodies STREAM (`REASON_PURE_FORWARD`) and are not + buffered; without it (or without a key) the typed path buffers each body — + the `body path streamed share` report row verifies which regime a run was + actually in. +- Resource conclusions (CPU/RSS/fds) are only meaningful at full scale; + reduced-scale reports carry an explicit warning banner. + +## Policy A/B + +`scenarios.py compare --scenario policy-ab` runs the identical fleet and +workload twice, once per gateway binary. Build binary B from another checkout +with an **isolated** `CARGO_TARGET_DIR` so the two builds never trample each +other's artifacts (and A stays warm): + +```sh +cd /path/to/other-checkout +CARGO_TARGET_DIR=/tmp/smg-ab-b RUSTC_WRAPPER= cargo build --release -p smg + +python3 scripts/generate_sim/scenarios.py compare --scenario policy-ab \ + --profile scripts/generate_sim/profiles/local-medium.json \ + --smg-bin-a "$PWD/target/release/smg" \ + --smg-bin-b /tmp/smg-ab-b/release/smg +``` + +`--only LEG` (repeatable) runs a subset of a scenario's legs by label, for +reruns after a fix; an unknown label is an error rather than an empty run. + +`--smg-bin` also works on plain `sim.py run` for one-off runs against a +prebuilt gateway. Mock workers and the loadgen always come from this +checkout, so only the gateway varies. + +## What the report contains + +- `report.md` / `report.json`: loadgen summary (TTFT/E2E percentiles, + cache-hit rates), per-worker imbalance from `requests.jsonl` (fleet CoV, + max/mean, distinct workers; split per turn), turn-2 same-worker rate, + per-SMG cache-aware branch counts (`hash_hit`/`hash_spill`/fallbacks, + parsed from the scoped `RUST_LOG=warn,smg::policies::cache_aware=debug` + logs — branches are debug-log-only, there is no branch metric), and per-SMG + resource samples (RSS, CPU, fds, admission-queue depth, active connections, + rejected/selection counters). +- `samples.jsonl`: the raw 5 s samples; `meta.json`: pids, registration and + readiness counts, timings. diff --git a/scripts/generate_sim/failover_bins.py b/scripts/generate_sim/failover_bins.py new file mode 100755 index 000000000..fe477dd0a --- /dev/null +++ b/scripts/generate_sim/failover_bins.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +"""Time-sliced failover analysis: bin follow-up cache ratios around the +index-replica kill instant recorded in meta.json (`index_killed_at_ms`, +written by sim.py's kill_index_replica drill; epoch milliseconds, the same +clock as the loadgen's `start_ms`). + +Usage: failover_bins.py [bin_secs] + +Exits non-zero when no seed recorded a kill: an empty analysis of a drill +that never fired must not look like a clean result. +""" + +import glob +import json +import sys + + +def main(): + if len(sys.argv) < 2: + raise SystemExit(__doc__) + base = sys.argv[1] + width = int(sys.argv[2]) if len(sys.argv) > 2 else 10 + bins = {} + seeds = sorted(glob.glob(f"{base}/*/seed-*/")) + observed = 0 + errors = requests = 0 + for sd in seeds: + with open(sd + "meta.json") as f: + meta = json.load(f) + killed = meta.get("index_killed_at_ms") + # `is None`, not falsy: a legitimate 0 offset must not read as absent. + if killed is None: + print(f"WARN: kill not observed in {sd}") + continue + observed += 1 + with open(sd + "requests.jsonl") as f: + for line in f: + rec = json.loads(line) + requests += 1 + if not (200 <= rec.get("status", 0) < 300): + errors += 1 + if rec.get("turn", 1) < 2: + continue + prompt = rec.get("prompt_tokens") or 0 + cached = rec.get("cached_tokens") or 0 + if not prompt: + continue + offset = (rec["start_ms"] - killed) / 1000.0 + b = int(offset // width) * width + slot = bins.setdefault(b, {"p": 0, "c": 0, "src": {}}) + slot["p"] += prompt + slot["c"] += cached + src = rec.get("index_source") or "none" + slot["src"][src] = slot["src"].get(src, 0) + 1 + print(f"kill observed in {observed}/{len(seeds)} seeds; errors {errors}/{requests}") + if not observed: + raise SystemExit("no seed recorded index_killed_at_ms; did the kill drill fire?") + for b in sorted(bins): + slot = bins[b] + total = sum(slot["src"].values()) + top = ", ".join( + f"{k} {v / total:.0%}" + for k, v in sorted(slot["src"].items(), key=lambda kv: -kv[1])[:3] + ) + print(f"{b:+6d}s followup cached {slot['c'] / slot['p']:.4f} {top}") + + +if __name__ == "__main__": + main() diff --git a/scripts/generate_sim/profiles/agentic-small.json b/scripts/generate_sim/profiles/agentic-small.json new file mode 100644 index 000000000..986244692 --- /dev/null +++ b/scripts/generate_sim/profiles/agentic-small.json @@ -0,0 +1,82 @@ +{ + "name": "agentic-small", + "description": "Agentic workload: a population of large shared system prompts (system prompt + tool schemas) reused across many sessions, deep multi-turn tool cycles. 8 SMGs x 120 grpc workers, 10x compression. Cache semantics only. smg_count is swept by the gw-scaleout scenario.", + "smg_count": 8, + "workers_total": 120, + "mock_processes": 1, + "model_id": "mock-model", + "duration_secs": 180, + "warmup_secs": 20, + "sample_interval_secs": 5, + "sample_fds": true, + "readiness_fraction": 0.99, + "readiness_timeout_secs": 120, + "mock": { + "engine": "realistic", + "block_size": 256, + "prefix_cache": true, + "decode_base_ms": 4.3, + "decode_per_req_ms": 0, + "prefill_tps": 80000, + "prefill_chunk": 320, + "max_running": 80, + "kv_tokens": 1200000 + }, + "smg_flags": [ + "--policy", + "cache_aware", + "--cache-index", + "hash", + "--cache-threshold", + "0.60", + "--block-size", + "256", + "--cache-boundaries", + "3072,4096,6144,8192,12288,16384", + "--balance-abs-threshold", + "8", + "--balance-rel-threshold", + "1.2", + "--overlap-decay", + "1.0", + "--disable-retries", + "--upstream-http2", + "--max-concurrent-requests", + "180000", + "--queue-size", + "512", + "--queue-timeout-secs", + "5", + "--worker-overload-waiting-requests", + "64", + "--worker-overload-token-usage", + "0.9", + "--routing-key-override", + "--assignment-mode", + "delegate" + ], + "loadgen": { + "session_rps": 100, + "t2_ratio": 0.8, + "max_turns": 8, + "think_secs": 3, + "system_prefix_tokens": 6144, + "prompt_cdf": "6500:0.30,9000:0.70,16000:0.99", + "prompt_max": 32768, + "output_cdf": "400:0.40,900:0.75,2000:0.99", + "output_max": 4096, + "t2_suffix_tokens": 384, + "image_bytes": 62000, + "image_count": 0, + "ingress": "random", + "turn2_ingress": "random", + "routing_key_reuse": 0.0, + "tokens_hint": false, + "stream": false, + "http2": true, + "conns_per_origin": 16, + "max_inflight": 50000, + "seed": 42, + "system_prefix_pool": 16 + } +} diff --git a/scripts/generate_sim/profiles/conversational-small.json b/scripts/generate_sim/profiles/conversational-small.json new file mode 100644 index 000000000..f7f0a5af2 --- /dev/null +++ b/scripts/generate_sim/profiles/conversational-small.json @@ -0,0 +1,82 @@ +{ + "name": "conversational-small", + "description": "Conversational/chat workload: a few shared system prompts, 2-5 turn conversations growing context each turn, moderate cross-user reuse. 8 SMGs x 120 grpc workers, 10x compression. Cache semantics only. smg_count swept by the gw-scaleout scenario.", + "smg_count": 8, + "workers_total": 120, + "mock_processes": 1, + "model_id": "mock-model", + "duration_secs": 150, + "warmup_secs": 20, + "sample_interval_secs": 5, + "sample_fds": true, + "readiness_fraction": 0.99, + "readiness_timeout_secs": 120, + "mock": { + "engine": "realistic", + "block_size": 256, + "prefix_cache": true, + "decode_base_ms": 4.3, + "decode_per_req_ms": 0, + "prefill_tps": 80000, + "prefill_chunk": 320, + "max_running": 80, + "kv_tokens": 1200000 + }, + "smg_flags": [ + "--policy", + "cache_aware", + "--cache-index", + "hash", + "--cache-threshold", + "0.60", + "--block-size", + "256", + "--cache-boundaries", + "3072,4096,6144,8192,12288,16384", + "--balance-abs-threshold", + "8", + "--balance-rel-threshold", + "1.2", + "--overlap-decay", + "1.0", + "--disable-retries", + "--upstream-http2", + "--max-concurrent-requests", + "180000", + "--queue-size", + "512", + "--queue-timeout-secs", + "5", + "--worker-overload-waiting-requests", + "64", + "--worker-overload-token-usage", + "0.9", + "--routing-key-override", + "--assignment-mode", + "delegate" + ], + "loadgen": { + "session_rps": 200, + "t2_ratio": 0.6, + "max_turns": 5, + "think_secs": 2, + "system_prefix_tokens": 1024, + "prompt_cdf": "1500:0.35,4000:0.75,9000:0.99", + "prompt_max": 16384, + "output_cdf": "300:0.45,800:0.80,1800:0.99", + "output_max": 3072, + "t2_suffix_tokens": 256, + "image_bytes": 62000, + "image_count": 0, + "ingress": "random", + "turn2_ingress": "random", + "routing_key_reuse": 0.0, + "tokens_hint": false, + "stream": false, + "http2": true, + "conns_per_origin": 16, + "max_inflight": 50000, + "seed": 42, + "system_prefix_pool": 4 + } +} diff --git a/scripts/generate_sim/profiles/full.template.json b/scripts/generate_sim/profiles/full.template.json new file mode 100644 index 000000000..1c9fbb274 --- /dev/null +++ b/scripts/generate_sim/profiles/full.template.json @@ -0,0 +1,82 @@ +{ + "name": "full-template", + "description": "TEMPLATE for a full-scale run: copy to profiles/full.local.json (untracked) and fill in your fleet's real worker count, rates, timing, and body sizes. The committed values below are generic round placeholders, NOT any deployment's figures.", + "smg_count": 8, + "workers_total": 10000, + "mock_processes": 5, + "model_id": "mock-model", + "duration_secs": 600, + "warmup_secs": 30, + "sample_interval_secs": 5, + "sample_fds": false, + "readiness_fraction": 0.99, + "readiness_timeout_secs": 900, + "mock": { + "engine": "realistic", + "block_size": 256, + "prefix_cache": true, + "decode_base_ms": 40, + "decode_per_req_ms": 0, + "prefill_tps": 8000, + "prefill_chunk": 320, + "max_running": 80, + "kv_tokens": 1200000 + }, + "smg_flags": [ + "--policy", + "cache_aware", + "--cache-index", + "hash", + "--cache-threshold", + "0.60", + "--block-size", + "128", + "--cache-boundaries", + "3072,4096,6144,8192,12288,16384", + "--balance-abs-threshold", + "8", + "--balance-rel-threshold", + "1.2", + "--overlap-decay", + "1.0", + "--disable-retries", + "--upstream-http2", + "--max-concurrent-requests", + "180000", + "--queue-size", + "512", + "--queue-timeout-secs", + "5", + "--worker-overload-waiting-requests", + "64", + "--worker-overload-token-usage", + "0.9", + "--routing-key-override", + "--assignment-mode", + "delegate" + ], + "loadgen": { + "session_rps": 2000, + "t2_ratio": 0.5, + "think_secs": 30, + "system_prefix_tokens": 2048, + "prompt_cdf": "5000:0.216,10000:0.530,20000:0.994", + "prompt_max": 24576, + "output_cdf": "1000:0.351,2000:0.513,5000:0.998", + "output_max": 6144, + "t2_suffix_tokens": 512, + "image_bytes": 500000, + "image_count": 1, + "ingress": "hash", + "turn2_ingress": "same", + "routing_key_reuse": 0.0, + "tokens_hint": false, + "stream": true, + "http2": true, + "conns_per_origin": 16, + "max_turns": 2, + "seed": 42, + "max_inflight": 700000 + }, + "requires_large_linux_host": "Full scale needs a large Linux host; see README host sizing." +} diff --git a/scripts/generate_sim/profiles/local-medium.json b/scripts/generate_sim/profiles/local-medium.json new file mode 100644 index 000000000..cf7f90682 --- /dev/null +++ b/scripts/generate_sim/profiles/local-medium.json @@ -0,0 +1,81 @@ +{ + "name": "local-medium", + "description": "Workstation profile at production-like worker pressure: 8 SMGs x 240 workers, 10x compression. Cache semantics only — not for resource conclusions.", + "smg_count": 8, + "workers_total": 240, + "mock_processes": 2, + "model_id": "mock-model", + "duration_secs": 180, + "warmup_secs": 15, + "sample_interval_secs": 5, + "sample_fds": true, + "readiness_fraction": 0.99, + "readiness_timeout_secs": 300, + "mock": { + "engine": "realistic", + "block_size": 256, + "prefix_cache": true, + "decode_base_ms": 4.3, + "decode_per_req_ms": 0, + "prefill_tps": 80000, + "prefill_chunk": 320, + "max_running": 80, + "kv_tokens": 1200000 + }, + "smg_flags": [ + "--policy", + "cache_aware", + "--cache-index", + "hash", + "--cache-threshold", + "0.60", + "--block-size", + "128", + "--cache-boundaries", + "3072,4096,6144,8192,12288,16384", + "--balance-abs-threshold", + "8", + "--balance-rel-threshold", + "1.2", + "--overlap-decay", + "1.0", + "--disable-retries", + "--upstream-http2", + "--max-concurrent-requests", + "180000", + "--queue-size", + "512", + "--queue-timeout-secs", + "5", + "--worker-overload-waiting-requests", + "64", + "--worker-overload-token-usage", + "0.9", + "--routing-key-override", + "--assignment-mode", + "delegate" + ], + "loadgen": { + "session_rps": 611, + "t2_ratio": 0.5, + "think_secs": 3, + "system_prefix_tokens": 2048, + "prompt_cdf": "5000:0.216,10000:0.530,20000:0.994", + "prompt_max": 24576, + "output_cdf": "1000:0.351,2000:0.513,5000:0.998", + "output_max": 6144, + "t2_suffix_tokens": 512, + "image_bytes": 62000, + "image_count": 1, + "ingress": "hash", + "turn2_ingress": "same", + "routing_key_reuse": 0.0, + "tokens_hint": false, + "stream": true, + "http2": true, + "conns_per_origin": 16, + "max_turns": 2, + "seed": 42, + "max_inflight": 100000 + } +} diff --git a/scripts/generate_sim/profiles/local-small.json b/scripts/generate_sim/profiles/local-small.json new file mode 100644 index 000000000..0febb3bba --- /dev/null +++ b/scripts/generate_sim/profiles/local-small.json @@ -0,0 +1,81 @@ +{ + "name": "local-small", + "description": "Laptop profile at production-like worker pressure: 8 SMGs x 120 workers, ~30-38 concurrent requests per worker under 10x time compression and reduced body size. smg_count and every workload property are knobs; cache semantics only — not for resource conclusions.", + "smg_count": 8, + "workers_total": 120, + "mock_processes": 1, + "model_id": "mock-model", + "duration_secs": 150, + "warmup_secs": 20, + "sample_interval_secs": 5, + "sample_fds": true, + "readiness_fraction": 0.99, + "readiness_timeout_secs": 120, + "mock": { + "engine": "realistic", + "block_size": 256, + "prefix_cache": true, + "decode_base_ms": 4.3, + "decode_per_req_ms": 0, + "prefill_tps": 80000, + "prefill_chunk": 320, + "max_running": 80, + "kv_tokens": 1200000 + }, + "smg_flags": [ + "--policy", + "cache_aware", + "--cache-index", + "hash", + "--cache-threshold", + "0.60", + "--block-size", + "128", + "--cache-boundaries", + "3072,4096,6144,8192,12288,16384", + "--balance-abs-threshold", + "8", + "--balance-rel-threshold", + "1.2", + "--overlap-decay", + "1.0", + "--disable-retries", + "--upstream-http2", + "--max-concurrent-requests", + "180000", + "--queue-size", + "512", + "--queue-timeout-secs", + "5", + "--worker-overload-waiting-requests", + "64", + "--worker-overload-token-usage", + "0.9", + "--routing-key-override", + "--assignment-mode", + "delegate" + ], + "loadgen": { + "session_rps": 305, + "t2_ratio": 0.5, + "max_turns": 2, + "think_secs": 3, + "system_prefix_tokens": 2048, + "prompt_cdf": "5000:0.216,10000:0.530,20000:0.994", + "prompt_max": 24576, + "output_cdf": "1000:0.351,2000:0.513,5000:0.998", + "output_max": 6144, + "t2_suffix_tokens": 512, + "image_bytes": 62000, + "image_count": 1, + "ingress": "hash", + "turn2_ingress": "same", + "routing_key_reuse": 0.0, + "tokens_hint": false, + "stream": true, + "http2": true, + "conns_per_origin": 16, + "max_inflight": 50000, + "seed": 42 + } +} diff --git a/scripts/generate_sim/profiles/smoke.json b/scripts/generate_sim/profiles/smoke.json new file mode 100644 index 000000000..81a53a524 --- /dev/null +++ b/scripts/generate_sim/profiles/smoke.json @@ -0,0 +1,80 @@ +{ + "name": "smoke", + "description": "Fast wiring check: 2 SMGs x 24 workers, ~4 req/s, 40 s. Not a measurement profile.", + "smg_count": 2, + "workers_total": 24, + "mock_processes": 1, + "model_id": "mock-model", + "duration_secs": 40, + "warmup_secs": 5, + "sample_interval_secs": 5, + "sample_fds": true, + "readiness_fraction": 0.99, + "readiness_timeout_secs": 60, + "mock": { + "engine": "realistic", + "block_size": 256, + "prefix_cache": true, + "decode_base_ms": 2.0, + "decode_per_req_ms": 0, + "prefill_tps": 200000, + "prefill_chunk": 320, + "max_running": 80, + "kv_tokens": 1200000 + }, + "smg_flags": [ + "--policy", + "cache_aware", + "--cache-index", + "hash", + "--cache-threshold", + "0.60", + "--block-size", + "128", + "--cache-boundaries", + "3072,4096,6144,8192,12288,16384", + "--balance-abs-threshold", + "8", + "--balance-rel-threshold", + "1.2", + "--overlap-decay", + "1.0", + "--disable-retries", + "--upstream-http2", + "--max-concurrent-requests", + "180000", + "--queue-size", + "512", + "--queue-timeout-secs", + "5", + "--worker-overload-waiting-requests", + "64", + "--worker-overload-token-usage", + "0.9", + "--routing-key-override", + "--assignment-mode", + "delegate" + ], + "loadgen": { + "session_rps": 2.5, + "t2_ratio": 0.5, + "think_secs": 2, + "system_prefix_tokens": 2048, + "prompt_cdf": "5000:0.216,10000:0.530,20000:0.994", + "prompt_max": 24576, + "output_cdf": "1000:0.351,2000:0.513,5000:0.998", + "output_max": 6144, + "t2_suffix_tokens": 512, + "image_bytes": 20000, + "image_count": 1, + "ingress": "hash", + "turn2_ingress": "same", + "routing_key_reuse": 0.0, + "tokens_hint": false, + "stream": true, + "http2": true, + "conns_per_origin": 16, + "max_turns": 2, + "seed": 42 + } +} diff --git a/scripts/generate_sim/scenarios.py b/scripts/generate_sim/scenarios.py new file mode 100755 index 000000000..2d49d0505 --- /dev/null +++ b/scripts/generate_sim/scenarios.py @@ -0,0 +1,1337 @@ +#!/usr/bin/env python3 +"""Named comparison scenarios on top of sim.py. + +Each scenario is a list of legs; a leg is a profile-override set (dotted keys) +plus, for policy A/B, which prebuilt gateway binary to use. `compare` runs the +legs sequentially against a fresh fleet each and emits a side-by-side markdown +table built from every leg's report.json. Per-turn rows are always included, +so `turn-ab` (t2_ratio 1.0) needs only one leg. + +Usage: + scripts/generate_sim/scenarios.py list + scripts/generate_sim/scenarios.py compare --scenario smg1-vs-smg8 \ + --profile scripts/generate_sim/profiles/local-small.json + scripts/generate_sim/scenarios.py compare --scenario policy-ab \ + --profile ... --smg-bin-a /path/A/smg --smg-bin-b /path/B/smg +""" + +import argparse +import json +import sys +import time +from datetime import datetime +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import sim # noqa: E402 + +# Multi-turn conversational traffic at the SAME total request rate as the +# 1.5-turn baseline (305 sessions/s x ~1.5 turns ~= 110 x ~4.15 turns): +# turn-mix comparisons must hold request RPS constant, not session RPS. +MULTITURN = { + "loadgen.session_rps": 110, + "loadgen.t2_ratio": 0.85, + "loadgen.max_turns": 8, + "loadgen.t2_suffix_tokens": 256, +} + + +# 10x-compressed gateway clocks for timing-sensitive legs. The profiles +# compress request time 10x (itl/think/prefill), but the gateway's +# wall-clock timers keep production defaults (load monitor 10 s, +# eviction 120 s, cache TTL 180 s) — uncompressed they distort any leg +# whose behavior depends on load freshness, tree eviction, or TTL. +COMPRESSED_CLOCK_FLAGS = { + "--load-monitor-interval": "1", + "--eviction-interval": "12", + "--cache-ttl-secs": "18", +} + +# Shared gateway-flag patch for the radix-replica legs: route on the +# approximate radix tree with no sticky short-circuit. False removes a flag. +RADIX_TREE_FLAGS = { + "--cache-index": "tree", + "--routing-key-override": False, + "--assignment-mode": False, +} + +# kv-events legs: gRPC-mode workers streaming KV cache events, tree index, +# no sticky short-circuit. --enable-igw is required because workers are +# registered dynamically: without it the gateway's single router is chosen +# from --worker-urls schemes at startup (HTTP by default) and would never +# route to gRPC workers. Loadgen runs non-streaming: the gRPC router's +# final SSE frame carries only the tail output token, which would starve +# the multi-turn context build. +KV_EVENT_OVERRIDES = { + "worker_mode": "grpc", + "loadgen.image_count": 0, + "loadgen.stream": False, + "loadgen.model": "mock-model", + "smg_flag_overrides": {**RADIX_TREE_FLAGS, "--enable-igw": None}, +} + + +def patch_smg_flags(flags, patches): + """Return a copy of `flags` with each `--flag: value` replaced in place + (or appended); a value of False removes the flag (and its value when it + has one). Scenario legs that must differ in exactly one gateway setting + go through here, so the rest of the flag list is shared by + construction.""" + out = list(flags) + for flag, value in patches.items(): + if value is False: + if flag in out: + idx = out.index(flag) + span = 2 if idx + 1 < len(out) and not out[idx + 1].startswith("--") else 1 + del out[idx : idx + span] + continue + if flag in out: + idx = out.index(flag) + if value is None: + continue + out[idx + 1] = str(value) + else: + out.append(flag) + if value is not None: + out.append(str(value)) + return out + + +REMOTE_INDEX_FLAGS = { + "--enable-igw": None, + "--kv-indexer-url": "http://127.0.0.1:40000", + "--kv-indexer-block-size": "256", +} + + +def remote_leg(index_overrides, extra=None): + """A sprayed remote-index leg; sweep legs vary ONLY index_overrides.""" + leg = { + **KV_EVENT_OVERRIDES, + "loadgen.ingress": "random", + "loadgen.turn2_ingress": "random", + "index_service": { + "replicas": 2, + "bridge": True, + "inferred_ttl_secs": 18, + "sweep_interval_secs": 1, + "default_capacity_blocks": 4688, + **index_overrides, + }, + "smg_flag_overrides": { + **RADIX_TREE_FLAGS, + **COMPRESSED_CLOCK_FLAGS, + **REMOTE_INDEX_FLAGS, + }, + } + if extra: + leg.update(extra) + return leg + + +# Leg: (label, {dotted override: value}, smg_bin slot: None | "a" | "b"). +# The special override key "smg_flag_overrides" patches individual gateway +# flags via patch_smg_flags. smg1-vs-smg8 keeps the same session_rps: +# aggregate rps is a loadgen-side property, so halving the gateway count +# concentrates rather than shrinks load. +SCENARIOS = { + # Production sticky-session assignment A/B: delegate (pin the worker the + # policy chose for turn 1) vs min_group (pin the group-least-keys + # worker). Everything else identical. + "assignment-mode-ab": [ + ("delegate", {}, None), + ( + "min-group", + {"smg_flag_overrides": {"--assignment-mode": "min_group"}}, + None, + ), + ], + # Controlled TTL comparison: traffic identical (think fixed at 6 s + # compressed = 60 s production), ONLY --cache-ttl-secs differs. 18 s is + # the production 180 s under 10x time compression; 2 s puts the TTL + # below the think time. + "ttl-controlled": [ + ( + "ttl-18s", + dict( + MULTITURN, + **{ + "loadgen.think_secs": 6, + "smg_flag_overrides": {"--cache-ttl-secs": "18"}, + }, + ), + None, + ), + ( + "ttl-2s", + dict( + MULTITURN, + **{ + "loadgen.think_secs": 6, + "smg_flag_overrides": {"--cache-ttl-secs": "2"}, + }, + ), + None, + ), + ], + "smg1-vs-smg8": [ + ("smg1", {"smg_count": 1}, None), + ("smg8", {"smg_count": 8}, None), + ], + "stable-key-vs-random-ingress": [ + ("hash-ingress", {"loadgen.ingress": "hash"}, None), + ("random-ingress", {"loadgen.ingress": "random"}, None), + ], + # What does it take to sustain >= 0.80 aggregate cached tokens? The + # aggregate is (1-f)*prefix_share + f*followup_ratio where f is the + # follow-up share of requests — a traffic property. Request RPS is held + # constant across legs (see MULTITURN). TTL effects live in the separate + # ttl-controlled scenario so this one varies traffic shape only. + "hit-rate-calibration": [ + ("baseline-1p5turn", {}, None), + ("multiturn", dict(MULTITURN), None), + ( + "multiturn-prefix4k", + dict(MULTITURN, **{"loadgen.system_prefix_tokens": 4096}), + None, + ), + ], + # Gateway-revision comparison: the deployed revision (prebuilt binary, + # slot a — build it from the actual deployed SHA with an ISOLATED + # CARGO_TARGET_DIR) vs the branch's ambient latest-main build, plus + # min_group on latest main. Routing-regime differences show directly in + # the "body path streamed share" row. + "revision-ab": [ + ("deployed-rev", {}, "a"), + ("latest-main", {}, None), + ( + "latest-main-min-group", + {"smg_flag_overrides": {"--assignment-mode": "min_group"}}, + None, + ), + ], + # Routing-key stability: stable per-session keys (baseline), a fresh key + # every turn (sticky pin lost + ingress re-hash each turn), and everyone + # sharing 32 keys (concurrent same-key pressure on the sticky cap). + "key-stability": [ + ("stable-key", {}, None), + ( + "key-per-turn", + { + "loadgen.key_per_turn": True, + "loadgen.turn2_ingress": "hash", + }, + None, + ), + ("shared-keys", {"loadgen.routing_key_reuse": 1.0}, None), + ], + # Approximate radix trees (cache_index=tree) are PER-REPLICA state: each + # gateway learns only from its own placements — over HTTP there are no + # worker KV events and no gateway-to-gateway sync. These legs drop the + # sticky override so cache_aware actually consults the tree, and vary + # (a) the index input — pre-tokenized ids (token tree) vs raw text + # (string tree), (b) whether a session's turns stay on one SMG (hash + # ingress, like a consistent-hashing LB) or spray uniformly (random), + # and (c) the replica count. Prediction accuracy should collapse only + # when per-replica trees AND sprayed sessions combine. token-hint keeps + # the body streaming while still feeding the token tree via the + # x-smg-routing-tokens header. Images off: placeholder expansion is an + # ids-path feature, so this keeps ids and text legs byte-comparable. + "radix-replica": [ + ( + "token-affine", + { + "smg_flag_overrides": RADIX_TREE_FLAGS, + "loadgen.image_count": 0, + }, + None, + ), + ( + "token-random", + { + "smg_flag_overrides": RADIX_TREE_FLAGS, + "loadgen.image_count": 0, + "loadgen.ingress": "random", + "loadgen.turn2_ingress": "random", + }, + None, + ), + ( + "token-single", + { + "smg_flag_overrides": RADIX_TREE_FLAGS, + "loadgen.image_count": 0, + "loadgen.ingress": "random", + "loadgen.turn2_ingress": "random", + "smg_count": 1, + }, + None, + ), + ( + "text-affine", + { + "smg_flag_overrides": RADIX_TREE_FLAGS, + "loadgen.image_count": 0, + "loadgen.payload": "text", + }, + None, + ), + ( + "text-random", + { + "smg_flag_overrides": RADIX_TREE_FLAGS, + "loadgen.image_count": 0, + "loadgen.payload": "text", + "loadgen.ingress": "random", + "loadgen.turn2_ingress": "random", + }, + None, + ), + ( + "token-hint-streamed", + { + "smg_flag_overrides": RADIX_TREE_FLAGS, + "loadgen.image_count": 0, + "loadgen.tokens_hint": True, + }, + None, + ), + ], + # Event-driven cache-aware routing: gRPC workers broadcast their actual + # cache contents (SubscribeKvEvents -> PositionalIndexer), so every + # gateway replica independently converges on ground truth with no + # gateway-to-gateway sync. The paired approximate-tree legs in + # radix-replica collapse under sprayed ingress; if worker broadcast + # works, the sprayed leg here should NOT collapse — that is the whole + # comparison. buffered-control repeats the approximate token tree under + # the same gRPC fleet, isolating "events vs approximation" from any + # HTTP-vs-gRPC pipeline difference... except events cannot be disabled + # per leg without a mock flag, so the control instead sprays with the + # sticky override ON (events subscribed but placement sticky), pinning + # the event machinery's cost while removing its routing influence. + "kv-events": [ + ("event-affine", dict(KV_EVENT_OVERRIDES), None), + ( + "event-sprayed", + { + **KV_EVENT_OVERRIDES, + "loadgen.ingress": "random", + "loadgen.turn2_ingress": "random", + }, + None, + ), + ( + "sticky-control", + { + "worker_mode": "grpc", + "loadgen.image_count": 0, + "loadgen.stream": False, + "loadgen.model": "mock-model", + "smg_flag_overrides": {"--cache-index": "tree", "--enable-igw": None}, + }, + None, + ), + ], + # M2 core matrix (02-experiment-plan.md): sprayed ingress everywhere + # (the regime that breaks per-replica state), one variable per + # comparison. local-event = per-gateway event indexers (round-3 + # architecture); remote-event = same feed, ONE shared index (parity = + # location only); remote-placement = shared index fed ONLY by gateway + # placements (bridge off — the eventless-engine thesis, the number + # that does not exist yet); mesh-tree = the in-repo TreeSync + # alternative (approximate trees synced gateway-to-gateway). + # Compressed gateway clocks on every leg (held constant). + "remote-index": [ + ( + "local-event-sprayed", + { + **KV_EVENT_OVERRIDES, + "loadgen.ingress": "random", + "loadgen.turn2_ingress": "random", + "smg_flag_overrides": { + **RADIX_TREE_FLAGS, + "--enable-igw": None, + **COMPRESSED_CLOCK_FLAGS, + }, + }, + None, + ), + ( + "remote-event-sprayed", + { + **KV_EVENT_OVERRIDES, + "loadgen.ingress": "random", + "loadgen.turn2_ingress": "random", + "index_service": { + "replicas": 2, + "bridge": True, + "inferred_ttl_secs": 18, + "sweep_interval_secs": 1, + "default_capacity_blocks": 4688, + }, + "smg_flag_overrides": { + **RADIX_TREE_FLAGS, + "--enable-igw": None, + **COMPRESSED_CLOCK_FLAGS, + "--kv-indexer-url": "http://127.0.0.1:40000", + "--kv-indexer-block-size": "256", + }, + }, + None, + ), + ( + "remote-placement-sprayed", + { + **KV_EVENT_OVERRIDES, + "loadgen.ingress": "random", + "loadgen.turn2_ingress": "random", + "index_service": { + "replicas": 2, + "bridge": False, + "inferred_ttl_secs": 18, + "sweep_interval_secs": 1, + "default_capacity_blocks": 4688, + }, + "smg_flag_overrides": { + **RADIX_TREE_FLAGS, + "--enable-igw": None, + **COMPRESSED_CLOCK_FLAGS, + "--kv-indexer-url": "http://127.0.0.1:40000", + "--kv-indexer-block-size": "256", + }, + }, + None, + ), + ( + "mesh-tree-sprayed", + { + **KV_EVENT_OVERRIDES, + "loadgen.ingress": "random", + "loadgen.turn2_ingress": "random", + "mesh_smgs": True, + "smg_flag_overrides": { + **RADIX_TREE_FLAGS, + "--enable-igw": None, + **COMPRESSED_CLOCK_FLAGS, + }, + }, + None, + ), + ], + # Mesh TreeSync leg alone (the core matrix's first three legs already + # completed; peer-url format fix made this rerunnable separately). + "remote-index-mesh": [ + ( + "mesh-tree-sprayed", + { + **KV_EVENT_OVERRIDES, + "loadgen.ingress": "random", + "loadgen.turn2_ingress": "random", + "mesh_smgs": True, + "smg_flag_overrides": { + **RADIX_TREE_FLAGS, + "--enable-igw": None, + **COMPRESSED_CLOCK_FLAGS, + }, + }, + None, + ), + ], + # Placement-fed leg alone: rerunnable after the prompt-plus-output + # placement-chain fix to measure how much of the 1.6-point gap to the + # event feed it closes (the routing-time chain covered only the + # prompt; the worker's blocks span the generated tail too). + "remote-index-placement": [ + ( + "remote-placement-sprayed", + { + **KV_EVENT_OVERRIDES, + "loadgen.ingress": "random", + "loadgen.turn2_ingress": "random", + "index_service": { + "replicas": 2, + "bridge": False, + "inferred_ttl_secs": 18, + "sweep_interval_secs": 1, + "default_capacity_blocks": 4688, + }, + "smg_flag_overrides": { + **RADIX_TREE_FLAGS, + "--enable-igw": None, + **COMPRESSED_CLOCK_FLAGS, + "--kv-indexer-url": "http://127.0.0.1:40000", + "--kv-indexer-block-size": "256", + }, + }, + None, + ), + ], + # Fault injection on the shared index: a full inter-replica network + # partition (severable TCP proxies) healed mid-run, and a wedged + # replica (SIGSTOP: TCP up, nothing drains) resumed mid-run. Both + # on the placement feed with 2 replicas; the per-replica metrics + # timeline in meta records divergence and reconvergence. + "index-partition": [ + ( + "partition-45s", + remote_leg( + {"bridge": False, "partitionable": True}, + extra={"partition_drill": {"at_secs": 60, "heal_after_secs": 45}}, + ), + None, + ), + ( + "hang-45s", + remote_leg( + {"bridge": False}, + extra={ + "hang_index_replica": {"at_secs": 60, "replica": 1, "resume_after_secs": 45} + }, + ), + None, + ), + ], + # F5: scale-up — replica 2 is in everyone's peer list from the + # start but only launches at t=60s, bootstrapping from replica 0 + # under live load. F7: replica 1 flaps (kill+relaunch x3). + "index-scaleup-flap": [ + ( + "scaleup-r3", + remote_leg( + {"bridge": False, "replicas": 3, "deferred_replicas": [2]}, + extra={"start_deferred_replica": {"replica": 2, "at_secs": 60}}, + ), + None, + ), + ( + "flap-r1", + remote_leg( + {"bridge": False}, + extra={ + "flap_index_replica": { + "replica": 1, + "at_secs": 45, + "cycles": 3, + "period_secs": 20, + } + }, + ), + None, + ), + ], + # Staleness sweep (event feed; the placement feed has no Removed to + # delay): constant injected apply lag, reported against the 3 s + # compressed think time (= 30 s production). Stored and Removed are + # delayed in SEPARATE legs — they fail in opposite directions. + "index-staleness": [ + ("stored-30ms", remote_leg({"apply_delay_stored_ms": 30}), None), + ("stored-300ms", remote_leg({"apply_delay_stored_ms": 300}), None), + ("stored-3000ms", remote_leg({"apply_delay_stored_ms": 3000}), None), + ("removed-3000ms", remote_leg({"apply_delay_removed_ms": 3000}), None), + ], + # Capacity-model sensitivity for the inferred feed (placement-only): + # 0.5x / 2x of the 1x=4688 blocks used by the core matrix leg. + "index-capacity": [ + ( + "capacity-half", + remote_leg({"bridge": False, "default_capacity_blocks": 2344}), + None, + ), + ( + "capacity-double", + remote_leg({"bridge": False, "default_capacity_blocks": 9375}), + None, + ), + ], + # Failover drill: kill replica 0 (the endpoint every gateway and + # publisher dials) mid-window, relaunch 30 s later bootstrapping from + # the survivor. Placement-fed (the harder case: no replayable feed). + # Run with --seeds 6: timing nondeterminism needs the wider n. + "index-failover": [ + ( + "kill-replica0", + remote_leg( + {"bridge": False}, + { + "kill_index_replica": { + "at_secs": 60, + "replica": 0, + "relaunch_after_secs": 30, + } + }, + ), + None, + ), + ], + # Kill and relaunch every SMG mid-window: sticky pins and placements are + # process state, so affinity must rebuild; errors during the blackout + # are part of the result. + "router-restart": [ + ( + "restart-mid-run", + dict( + MULTITURN, + **{"restart_smgs_at_secs": 60}, + ), + None, + ), + ], + # The hash placement index is LOCAL to each SMG: turn-2 affinity only + # survives if turn 2 reaches the same SMG. This isolates that effect. + # Does a smaller SMG fleet raise cache hit rates? With sticky turns the + # placement index fragmentation shouldn't matter; with scattered turns + # the chance of landing on the SMG that holds the placement is 1/K. + "fleet-size-sweep": [ + ("smg2-sticky", {"smg_count": 2}, None), + ("smg8-sticky", {"smg_count": 8}, None), + ( + "smg2-t2random", + { + "smg_count": 2, + "loadgen.ingress": "random", + "loadgen.turn2_ingress": "random", + }, + None, + ), + ( + "smg8-t2random", + { + "smg_count": 8, + "loadgen.ingress": "random", + "loadgen.turn2_ingress": "random", + }, + None, + ), + ], + "turn2-same-vs-random-smg": [ + ("t2-same-smg", {"loadgen.turn2_ingress": "same"}, None), + ( + "t2-random-smg", + {"loadgen.ingress": "random", "loadgen.turn2_ingress": "random"}, + None, + ), + ], + "cold-vs-warm-prefix": [ + ("cold", {"loadgen.system_prefix_tokens": 0}, None), + ("warm", {"loadgen.system_prefix_tokens": 2048}, None), + ], + "turn-ab": [ + ("turn-ab", {"loadgen.t2_ratio": 1.0}, None), + ], + "policy-ab": [ + ("policy-a", {}, "a"), + ("policy-b", {}, "b"), + ], +} + + +# Gateway scale-out sweep: the DB thesis. Sessions spray across gateways +# (ingress + turn2 random), so a follow-up frequently lands on a gateway +# that did NOT route turn 1. Compare per-gateway LOCAL event indexers +# vs one SHARED remote index (placement-only feed, bridge off) as the +# gateway count grows. cache-hit row: "AGG cached (request mean)". +# local-event converges per gateway (full worker view) so it should hold +# roughly flat; the question is whether the eventless shared DB MATCHES +# it as gateways scale — and, once Mode 1 lands, the string case where +# per-gateway state genuinely fragments. +def _scaleout_legs(counts=(1, 2, 4, 8)): + """Scale-out sweep: four index regimes x gateway count. Sessions spray + across gateways (ingress + turn2 random), so a follow-up often lands on + a gateway that did not route turn 1 — the shared state's value shows + there. Cache-hit rows: "AGG cached (request mean)", "followup cached + (request mean)", "followup same-worker"; validity rows: "index + remote_hit share" (must be populated on the remote legs).""" + spray = {"loadgen.ingress": "random", "loadgen.turn2_ingress": "random"} + idx_flags = { + "--kv-indexer-url": "http://127.0.0.1:40000", + "--kv-indexer-block-size": "256", + } + + def db(bridge): + return { + "replicas": 1, + "bridge": bridge, + "inferred_ttl_secs": 18, + "sweep_interval_secs": 1, + "default_capacity_blocks": 4688, + } + + legs = [] + for count in counts: + # per-gateway event indexers: each gateway sees ALL worker events + # (full view, N x fan-in) — the well-informed local upper reference. + legs.append( + ( + f"local-event-smg{count}", + { + **KV_EVENT_OVERRIDES, + "smg_count": count, + **spray, + "smg_flag_overrides": { + **RADIX_TREE_FLAGS, + "--enable-igw": None, + **COMPRESSED_CLOCK_FLAGS, + }, + }, + None, + ) + ) + # shared DB fed by the event bridge — the real "DB replaces + # per-gateway indexing" leg (shared view, no per-gateway fan-in). + legs.append( + ( + f"remote-event-smg{count}", + { + **KV_EVENT_OVERRIDES, + "smg_count": count, + **spray, + "index_service": db(True), + "smg_flag_overrides": { + **RADIX_TREE_FLAGS, + "--enable-igw": None, + **COMPRESSED_CLOCK_FLAGS, + **idx_flags, + }, + }, + None, + ) + ) + # shared DB fed ONLY by gateway placements (eventless engine). + legs.append( + ( + f"remote-placement-smg{count}", + { + **KV_EVENT_OVERRIDES, + "smg_count": count, + **spray, + "index_service": db(False), + "smg_flag_overrides": { + **RADIX_TREE_FLAGS, + "--enable-igw": None, + **COMPRESSED_CLOCK_FLAGS, + **idx_flags, + }, + }, + None, + ) + ) + # fragmenting floor: HTTP workers, per-gateway tree, no events, no + # shared DB — a gateway knows only what it itself routed. + legs.append( + ( + f"local-nosharing-smg{count}", + { + "smg_count": count, + **spray, + "loadgen.image_count": 0, + "smg_flag_overrides": {**RADIX_TREE_FLAGS, **COMPRESSED_CLOCK_FLAGS}, + }, + None, + ) + ) + return legs + + +SCENARIOS["gw-scaleout"] = _scaleout_legs() + + +# ---- the side-by-side evaluation ------------------------------------------ +# Four regimes, identical workload, differing only in where prefix knowledge +# lives: +# production — the shipping configuration: hash placement index + +# sticky routing key, HTTP workers (the incumbent). +# local-events — per-gateway kv_index event trees (every gateway +# consumes every worker's KV events; the well-informed +# local upper reference). +# remote-events — the shared index fed by the event bridge. +# remote-placement— the shared index fed only by gateway placements +# (eventless engines). +# Every regime drops images from the bodies (the gRPC legs cannot carry +# them) and runs the compressed gateway clocks. +def _regime(name, count, extra=None): + spray = {"loadgen.ingress": "random", "loadgen.turn2_ingress": "random"} + idx_flags = {"--kv-indexer-url": "http://127.0.0.1:40000", "--kv-indexer-block-size": "256"} + db = { + "replicas": 1, + "inferred_ttl_secs": 18, + "sweep_interval_secs": 1, + "default_capacity_blocks": 4688, + } + if name == "production": + leg = { + "smg_count": count, + **spray, + "loadgen.image_count": 0, + "smg_flag_overrides": {**COMPRESSED_CLOCK_FLAGS}, + } + elif name == "local-events": + leg = { + **KV_EVENT_OVERRIDES, + "smg_count": count, + **spray, + "smg_flag_overrides": { + **RADIX_TREE_FLAGS, + "--enable-igw": None, + **COMPRESSED_CLOCK_FLAGS, + }, + } + elif name == "remote-events": + leg = { + **KV_EVENT_OVERRIDES, + "smg_count": count, + **spray, + "index_service": {**db, "bridge": True}, + "smg_flag_overrides": { + **RADIX_TREE_FLAGS, + "--enable-igw": None, + **COMPRESSED_CLOCK_FLAGS, + **idx_flags, + }, + } + elif name == "remote-placement": + leg = { + **KV_EVENT_OVERRIDES, + "smg_count": count, + **spray, + "index_service": {**db, "bridge": False}, + "smg_flag_overrides": { + **RADIX_TREE_FLAGS, + "--enable-igw": None, + **COMPRESSED_CLOCK_FLAGS, + **idx_flags, + }, + } + else: + raise ValueError(name) + if extra: + for k, v in extra.items(): + if k == "index_service": + leg.setdefault("index_service", {}).update(v) + else: + leg[k] = v + return leg + + +REGIMES = ("production", "local-events", "remote-events", "remote-placement") + + +def _eval_matrix(counts=(1, 8), rates=(305, 611, 900)): + """gateways × concurrency × regime at the profile's worker count.""" + legs = [] + for count in counts: + for rate in rates: + for regime in REGIMES: + legs.append( + ( + f"{regime}-smg{count}-rps{rate}", + _regime(regime, count, {"loadgen.session_rps": rate}), + None, + ) + ) + return legs + + +SCENARIOS["eval-matrix"] = _eval_matrix() + +# Input/output shape sweep at 8 gateways. Rates are scaled per shape so the +# per-worker concurrency stays near the calibrated ~35 (request lifetime +# tracks output length at 4.3 ms/token): short outputs take a higher +# session rate, long outputs a lower one. +IO_SHAPES = { + "short-short": { + "loadgen.prompt_cdf": "1000:0.50,2000:0.95,4000:1.0", + "loadgen.prompt_max": 4096, + "loadgen.output_cdf": "200:0.50,400:0.95,800:1.0", + "loadgen.output_max": 1024, + "loadgen.session_rps": 900, + }, + "long-short": { + "loadgen.prompt_cdf": "12000:0.30,20000:0.90,24000:1.0", + "loadgen.prompt_max": 24576, + "loadgen.output_cdf": "200:0.50,400:0.95,800:1.0", + "loadgen.output_max": 1024, + "loadgen.session_rps": 900, + }, + "long-long": { + "loadgen.prompt_cdf": "12000:0.30,20000:0.90,24000:1.0", + "loadgen.prompt_max": 24576, + "loadgen.output_cdf": "2000:0.50,3000:0.90,5000:1.0", + "loadgen.output_max": 6144, + "loadgen.session_rps": 400, + }, +} +SCENARIOS["io-shapes"] = [ + (f"{regime}-{shape}", _regime(regime, 8, overrides), None) + for shape, overrides in IO_SHAPES.items() + for regime in REGIMES +] + +# Chaos on the placement-fed shared index (8 gateways), each leg ending with a +# replica dump so replica convergence is measured, not assumed. The cold +# gateway restart is run on the local-events regime too: that comparison IS +# the thesis (a cold gateway routing on the fleet's knowledge vs. its own). +CHAOS_DB = {"index_service": {"replicas": 2, "dump_on_exit": True}} +SCENARIOS["chaos"] = [ + ( + "worker-churn", + _regime( + "remote-placement", + 8, + { + **CHAOS_DB, + "remove_workers_drill": {"at_secs": 60, "count": 20}, + "add_workers_drill": {"at_secs": 90, "count": 20}, + }, + ), + None, + ), + ( + "cold-gateway-remote", + _regime( + "remote-placement", 8, {**CHAOS_DB, "restart_one_smg_drill": {"at_secs": 60, "smg": 3}} + ), + None, + ), + ( + "cold-gateway-local", + _regime("local-events", 8, {"restart_one_smg_drill": {"at_secs": 60, "smg": 3}}), + None, + ), + ( + "rolling-replica-restart", + _regime( + "remote-placement", + 8, + {**CHAOS_DB, "rolling_replica_restart_drill": {"at_secs": 45, "gap_secs": 30}}, + ), + None, + ), + ( + "replica-removed-permanently", + _regime( + "remote-placement", 8, {**CHAOS_DB, "kill_index_replica": {"at_secs": 60, "replica": 1}} + ), + None, + ), + ( + "gateway-index-partition-all", + _regime( + "remote-placement", + 8, + { + "index_service": {"replicas": 2, "dump_on_exit": True, "gateway_proxy": True}, + "gateway_partition_drill": {"at_secs": 60, "heal_after_secs": 45, "scope": "all"}, + }, + ), + None, + ), + ( + "gateway-index-partition-half", + _regime( + "remote-placement", + 8, + { + "index_service": {"replicas": 2, "dump_on_exit": True, "gateway_proxy": True}, + "gateway_partition_drill": {"at_secs": 60, "heal_after_secs": 45, "scope": "half"}, + }, + ), + None, + ), + ( + "replica-partition-converge", + _regime( + "remote-events", + 8, + { + "index_service": {"replicas": 2, "dump_on_exit": True, "partitionable": True}, + "partition_drill": {"at_secs": 60, "heal_after_secs": 45}, + }, + ), + None, + ), +] + +# 30-minute soaks: drift over time (capacity, staleness, leaks) on both feeds. +SCENARIOS["soak"] = [ + ( + f"{regime}-soak", + _regime(regime, 8, {"duration_secs": 1800, "index_service": {"dump_on_exit": True}}), + None, + ) + for regime in ("remote-placement", "remote-events") +] +# One short remote leg to VALIDATE the port: DB launches, bridge feeds it, +# and index_sources populates (run with --override duration_secs=30). +SCENARIOS["idx-smoke"] = [_scaleout_legs(counts=(1,))[1]] # remote-event-smg1 + + +def _get(mapping, *keys): + node = mapping + for key in keys: + if not isinstance(node, dict): + return None + node = node.get(key) + return node + + +def extract_rows(report): + """Flatten one report.json into the comparison rows (all guarded).""" + summary = report.get("loadgen_summary", {}) + req = report.get("requests", {}) + samples = report.get("samples", []) + branches = report.get("cache_aware_branches", []) + + rss_peaks = [_get(s, "rss_kib", "peak") for s in samples] + rss_peaks = [v for v in rss_peaks if v is not None] + cpu_means = [_get(s, "cpu_pct", "mean") for s in samples] + cpu_means = [v for v in cpu_means if v is not None] + queue_peaks = [_get(s, "queue_depth", "peak") for s in samples] + queue_peaks = [v for v in queue_peaks if v is not None] + rejected = sum(s.get("rejected_total") or 0 for s in samples) + + branch_totals = {} + for entry in branches: + for name, count in entry.get("branches", {}).items(): + branch_totals[name] = branch_totals.get(name, 0) + count + total_decisions = sum(branch_totals.values()) + hash_hit = branch_totals.get("hash_hit", 0) + + rows = {} + totals = summary.get("totals", {}) + errors = totals.get("errors", {}) + err_count = sum(errors.values()) if isinstance(errors, dict) else errors + requests_total = totals.get("requests") + rows["ok"] = ( + requests_total - err_count + if isinstance(requests_total, int) and isinstance(err_count, int) + else None + ) + rows["err"] = err_count + rows["achieved_rps"] = summary.get("achieved_rps") + for metric in ("ttft_ms", "e2e_ms"): + for pct in ("p50", "p90", "p99"): + rows[f"{metric}_{pct}"] = _get(summary, metric, pct) + # Token-weighted (sum cached / sum prompt) is THE number comparable with + # backend cached-token telemetry; the per-request mean is a different + # statistic and is reported under its own name. + rows["AGG cached tokens (sum/sum)"] = _get(summary, "overall", "cached_token_ratio") + rows["AGG cached (request mean)"] = _get(summary, "overall", "cached_ratio_request_mean") + for turn in ("turn1", "followup"): + rows[turn + " cached tokens (sum/sum)"] = _get(summary, "turns", turn, "cached_token_ratio") + rows[turn + " cached (request mean)"] = _get( + summary, "turns", turn, "cached_ratio_request_mean" + ) + rows[turn + " prompt tokens sum"] = _get(summary, "turns", turn, "prompt_tokens_sum") + rows[turn + " cached tokens sum"] = _get(summary, "turns", turn, "cached_tokens_sum") + rows["mean turns/session"] = summary.get("mean_turns_per_session") + rows["t2 same-worker (loadgen)"] = summary.get("turn2_same_worker_rate") + rows["followup same-worker"] = summary.get("followup_same_worker_rate") + rows["t1 max worker share"] = _get(summary, "turn1_workers", "max_share") + rows["t1 entropy (norm)"] = _get(summary, "turn1_workers", "normalized_entropy") + for turn in ("turn1", "turn2"): + rows[turn + " cached/prompt"] = _get(req, "turns", turn, "cached_over_prompt") + rows[turn + " hit rate"] = _get(req, "turns", turn, "hit_rate") + rows[turn + " CoV (fleet)"] = _get(req, "turns", turn, "imbalance", "cov_fleet") + rows["t2 same-worker rate"] = req.get("t2_same_worker_rate") + imb = req.get("overall_imbalance", {}) + rows["overall CoV (fleet)"] = imb.get("cov_fleet") + rows["distinct workers"] = imb.get("distinct_workers") + rows["hash_hit share"] = round(hash_hit / total_decisions, 4) if total_decisions else None + # Sticky-session outcomes (--routing-key-override): occupied_hit is a + # follow-up landing on its pinned worker; vacant is a fresh key. + sticky = {} + for s in samples: + for name, count in (s.get("sticky_branches") or {}).items(): + sticky[name] = sticky.get(name, 0) + count + sticky_total = sum(sticky.values()) + rows["sticky occupied_hit share"] = ( + round(sticky.get("occupied_hit", 0) / sticky_total, 4) if sticky_total else None + ) + rows["sticky cap_respill count"] = sticky.get("cap_respill", 0) if sticky else None + # Direct verification of the request-body regime: share of requests + # routed via streaming (header-only selection) vs buffered (typed path). + paths = {} + for s in samples: + for name, count in (s.get("body_paths") or {}).items(): + paths[name] = paths.get(name, 0) + count + path_total = sum(paths.values()) + streamed = sum(v for k, v in paths.items() if k.startswith("streamed")) + # Remote-index rows (--kv-indexer-url legs): per-request echo shares + # and the direct accuracy signal (predicted-vs-actual cached tokens). + sources = req.get("index_sources") or {} + total_sourced = sum(sources.values()) + if total_sourced: + rows["index remote_hit share"] = round(sources.get("remote_hit", 0) / total_sourced, 4) + misses = total_sourced - sources.get("remote_hit", 0) - sources.get("remote_empty", 0) + rows["index degraded share (timeout+disconnect)"] = round(misses / total_sourced, 4) + pred = req.get("index_prediction_error_tokens") + if pred: + rows["index prediction error mean (tokens)"] = pred.get("mean") + rows["index prediction exact share"] = pred.get("exact_share") + rows["index prediction error p50 abs (tokens)"] = pred.get("p50_abs") + rows["index prediction error p90 abs (tokens)"] = pred.get("p90_abs") + rows["index prediction error p95 abs (tokens)"] = pred.get("p95_abs") + rows["index prediction error max abs (tokens)"] = pred.get("max_abs") + # Index latency: the gateway's view of a lookup (client, wire, service + # and queueing — what the 2 ms deadline is measured against) and the + # service's own apply/query engine time. + idx = report.get("index_samples") or {} + for label, key in ( + ("index lookup ms (gateway) p50", "p50"), + ("index lookup ms (gateway) p90", "p90"), + ("index lookup ms (gateway) p99", "p99"), + ): + if (idx.get("gateway_lookup_ms") or {}).get(key) is not None: + rows[label] = idx["gateway_lookup_ms"][key] + for label, key in ( + ("index apply ms (service) p50", "p50"), + ("index apply ms (service) p90", "p90"), + ("index apply ms (service) p99", "p99"), + ): + if (idx.get("apply_ms") or {}).get(key) is not None: + rows[label] = idx["apply_ms"][key] + for label, key in ( + ("index query ms (service) p50", "p50"), + ("index query ms (service) p90", "p90"), + ("index query ms (service) p99", "p99"), + ): + if (idx.get("query_ms") or {}).get(key) is not None: + rows[label] = idx["query_ms"][key] + reps = idx.get("replicas") or {} + if reps: + rss = [r["rss_peak_kib"] for r in reps.values() if r.get("rss_peak_kib") is not None] + cpu = [r["cpu_mean_pct"] for r in reps.values() if r.get("cpu_mean_pct") is not None] + if rss: + rows["index rss peak MiB (max replica)"] = round(max(rss) / 1024, 1) + if cpu: + rows["index cpu mean % (max replica)"] = max(cpu) + div = report.get("index_divergence") or {} + if div.get("holders_compared") is not None: + rows["replicas converged at end"] = div.get("converged") + rows["holders differing between replicas"] = div.get("holders_differing") + rows["event-fed holders differing"] = div.get("holders_differing_event_fed") + rows["placement holders differing, in capacity band (cut timing)"] = div.get( + "holders_differing_placement_in_band" + ) + rows["placement holders differing, under capacity"] = div.get( + "holders_differing_placement_out_of_band" + ) + rows["holders only on one replica"] = div.get("holders_only_in_one") + tl = req.get("followup_timeline") or [] + if len(tl) >= 4: + # First and last full minutes of follow-up cache ratio: a drift over + # the run reads straight off the table. + rows["followup cached minute 1"] = tl[1]["cached_over_prompt"] + rows["followup cached last minute"] = tl[-2]["cached_over_prompt"] + + rows["body path streamed share"] = round(streamed / path_total, 4) if path_total else None + rows["offered session rps"] = summary.get("offered_session_rps") + rows["drain requests (excluded)"] = _get(summary, "drain", "requests") + rows["rss peak MiB (max smg)"] = round(max(rss_peaks) / 1024, 1) if rss_peaks else None + rows["cpu mean % (max smg)"] = max(cpu_means) if cpu_means else None + rows["queue depth peak"] = max(queue_peaks) if queue_peaks else None + rows["rejected total"] = rejected + return rows + + +def write_compare_md(scenario, leg_results, path): + labels = [label for label, _ in leg_results] + row_keys = [] + for _, rows in leg_results: + for key in rows: + if key not in row_keys: + row_keys.append(key) + lines = [f"# generate-sim compare — {scenario}", ""] + lines.append("| metric | " + " | ".join(labels) + " |") + lines.append("|---|" + "---|" * len(labels)) + for key in row_keys: + cells = [sim._fmt(rows.get(key)) for _, rows in leg_results] + lines.append("| {} | {} |".format(key, " | ".join(cells))) + lines.append("") + for label, _ in leg_results: + lines.append(f"- {label}: see its run dir for report.md / report.json") + lines.append("") + text = "\n".join(lines) + with open(path, "w") as f: + f.write(text) + print(text) + + +def select_legs(legs, only): + """Legs whose label is in `only`; every leg when `only` is empty. + Unknown labels are an error so a typo cannot silently run nothing.""" + if not only: + return list(legs) + labels = [label for label, _, _ in legs] + unknown = sorted(set(only) - set(labels)) + if unknown: + raise SystemExit(f"--only: no such leg(s) {unknown}; legs are {labels}") + return [leg for leg in legs if leg[0] in only] + + +def cmd_compare(args): + legs = select_legs(SCENARIOS[args.scenario], args.only) + bins = {"a": args.smg_bin_a, "b": args.smg_bin_b, None: args.smg_bin} + for _, _, slot in legs: + if slot is not None and not bins[slot]: + raise SystemExit( + f"scenario {args.scenario} needs --smg-bin-{slot} (a prebuilt gateway binary)" + ) + + base = sim.load_profile(args.profile) + stamp = datetime.now().strftime("%Y%m%d-%H%M%S") + out_root = Path(args.out_root or sim.REPO_ROOT / "target" / "generate-sim") + scenario_dir = out_root / ("{}-{}-{}".format(args.scenario, base.get("name", "run"), stamp)) + + leg_results = [] + built = args.skip_build + index_built = args.skip_build + for idx, (label, overrides, slot) in enumerate(legs): + seed_rows = [] + for seed_idx in range(args.seeds): + profile = json.loads(json.dumps(base)) # deep copy: legs must not leak + flag_patches = None + for key, val in overrides.items(): + if key == "smg_flag_overrides": + flag_patches = val + else: + sim.apply_override(profile, key, val) + if flag_patches: + profile["smg_flags"] = patch_smg_flags(profile["smg_flags"], flag_patches) + for raw in args.override: + key, val = sim.parse_override_arg(raw) + sim.apply_override(profile, key, val) + base_seed = int(profile.get("loadgen", {}).get("seed", 42)) + sim.apply_override(profile, "loadgen.seed", base_seed + seed_idx) + sim.log( + f"scenario {args.scenario} leg {idx + 1}/{len(legs)}: {label} " + f"seed {seed_idx + 1}/{args.seeds}" + ) + run_dir = sim.run_profile( + profile, + scenario_dir / label / f"seed-{base_seed + seed_idx}", + smg_bin=bins[slot], + skip_build=built and (index_built or "index_service" not in overrides), + ) + # Only a slot-None run builds the gateway; slot legs use a + # prebuilt binary, so ambient target/release/smg may not exist. + # Index binaries are built per-profile, so the first leg that + # wants the index must still build even if an earlier leg ran. + built = built or slot is None + index_built = index_built or "index_service" in overrides + with open(Path(run_dir) / "report.json") as f: + seed_rows.append(extract_rows(json.load(f))) + time.sleep(2) + with open(scenario_dir / label / "seed-rows.json", "w") as f: + json.dump(seed_rows, f, indent=2) + leg_results.append((label, aggregate_seed_rows(seed_rows))) + + # Leg-to-leg binary identity: when no leg deliberately uses a prebuilt + # slot (revision A/B does), every leg must have run the SAME gateway + # binary — otherwise the comparison silently includes a build delta. + if all(slot is None for _, _, slot in legs): + smg_shas = set() + for label, _, _ in legs: + for meta_path in sorted((scenario_dir / label).glob("seed-*/meta.json")): + with open(meta_path) as f: + smg_shas.add(json.load(f)["binary_sha256"]["smg"]) + if len(smg_shas) > 1: + raise SystemExit( + f"scenario {args.scenario} legs ran different gateway binaries: {sorted(smg_shas)}" + ) + + write_compare_md(args.scenario, leg_results, scenario_dir / "compare.md") + sim.log("compare: %s" % (scenario_dir / "compare.md")) + + +# Two-sided 97.5% Student-t quantiles by degrees of freedom; small seed +# counts need these, not the normal 1.96 (n=3 would otherwise report +# intervals ~2.2x too narrow). +STUDENT_T_975 = { + 1: 12.706, + 2: 4.303, + 3: 3.182, + 4: 2.776, + 5: 2.571, + 6: 2.447, + 7: 2.365, + 8: 2.306, + 9: 2.262, +} + + +def aggregate_seed_rows(seed_rows): + """Mean ± 95% CI half-width (Student-t) across seeds for numeric rows; + single-seed runs and non-numeric rows pass through the first value.""" + if len(seed_rows) == 1: + return seed_rows[0] + keys = [] + for rows in seed_rows: + for key in rows: + if key not in keys: + keys.append(key) + out = {} + for key in keys: + vals = [r.get(key) for r in seed_rows] + nums = [v for v in vals if isinstance(v, (int, float))] + if len(nums) == len(seed_rows): + n = len(nums) + m = sum(nums) / n + var = sum((v - m) ** 2 for v in nums) / (n - 1) + t = STUDENT_T_975.get(n - 1, 1.96) + half = t * (var**0.5) / (n**0.5) + out[key] = f"{m:.4g} ±{half:.2g}" + elif nums: + # Partial coverage: show the value plus how many seeds had it, + # so a validity row cannot read like a full-seed aggregate. + out[key] = f"{sum(nums) / len(nums):.4g} ({len(nums)}/{len(seed_rows)} seeds)" + else: + out[key] = vals[0] + return out + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + sub = parser.add_subparsers(dest="cmd", required=True) + + sub.add_parser("list", help="list scenario names") + + cmp_p = sub.add_parser("compare", help="run a scenario's legs and emit compare.md") + cmp_p.add_argument("--scenario", required=True, choices=sorted(SCENARIOS)) + cmp_p.add_argument("--profile", required=True, help="base profile JSON") + cmp_p.add_argument("--skip-build", action="store_true") + cmp_p.add_argument("--smg-bin", help="gateway binary for non-A/B legs") + cmp_p.add_argument("--smg-bin-a", help="gateway binary A (policy-ab)") + cmp_p.add_argument("--smg-bin-b", help="gateway binary B (policy-ab)") + cmp_p.add_argument("--out-root", help="parent dir for the scenario run dirs") + cmp_p.add_argument( + "--seeds", + type=int, + default=3, + help="loadgen seeds per leg; rows report mean ±95%% CI (default 3)", + ) + cmp_p.add_argument( + "--override", + action="append", + default=[], + metavar="KEY=VALUE", + help="extra dotted profile override applied to every leg (repeatable)", + ) + cmp_p.add_argument( + "--only", + action="append", + default=[], + metavar="LEG", + help="run only this leg of the scenario, by label (repeatable; default all)", + ) + + args = parser.parse_args() + if args.cmd == "list": + for name, legs in sorted(SCENARIOS.items()): + print(f"{name:<30} {' vs '.join(label for label, _, _ in legs)}") + elif args.cmd == "compare": + cmd_compare(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/generate_sim/sim.py b/scripts/generate_sim/sim.py new file mode 100755 index 000000000..bfb430fac --- /dev/null +++ b/scripts/generate_sim/sim.py @@ -0,0 +1,2216 @@ +#!/usr/bin/env python3 +"""Local /generate scale-simulation orchestrator for SMG. + +Implements the harness from .claude/generate-scale-sim/01-design.md: build the +gateway, mock fleet, and load generator; launch K SMG replicas against mock +workers in `--engine realistic` mode; register every worker with every SMG; drive +the sim-loadgen workload; sample per-SMG resources and /metrics while it runs; +then merge loadgen output, samples, and per-SMG cache-aware branch logs into +report.json / report.md. + +Stdlib only. macOS-first (BSD ps/lsof invocations) with Linux fallbacks +(/proc). Profiles are plain JSON (see profiles/); every workload unknown is a +profile or CLI knob, never hardcoded here. + +Usage: + scripts/generate_sim/sim.py run --profile scripts/generate_sim/profiles/local-small.json + scripts/generate_sim/sim.py run --profile ... --skip-build --smg-bin /path/to/smg + scripts/generate_sim/sim.py report --run-dir target/generate-sim/ +""" + +import argparse +import hashlib +import json +import os +import platform +import re +import signal +import socket +import statistics +import subprocess +import threading +import time +import urllib.error +import urllib.request +from collections import Counter +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] + +MOCK_BASE_PORT = 9000 +SMG_BASE_PORT = 30000 +PROM_BASE_PORT = 39000 +INDEX_BASE_PORT = 40000 +INDEX_METRICS_BASE = 40100 +# Peer-facing severable proxies in front of each index replica (partition +# drill): replica j is reached by its peers via INDEX_PROXY_BASE + j. +INDEX_PROXY_BASE = 40200 +# Severable proxies between the GATEWAYS and replica 0 (gateway-side partition +# drills): even-numbered gateways dial +0, odd-numbered dial +1, so a drill can +# sever all gateways or only half of them. +INDEX_GW_PROXY_BASE = 40300 +# Gateway-to-gateway mesh (TreeSync) ports, one per SMG. +MESH_BASE_PORT = 41000 + +# Cache-aware decision branches are DEBUG logs only (no branch metric), scoped +# so the rest of the gateway stays at warn. +SMG_RUST_LOG = "warn,smg::policies::cache_aware=debug" + +METRIC_PREFIXES = ( + "smg_admission_queue_depth", + "smg_http_connections_active", + "smg_admission_queue_rejected_total", + "smg_worker_selection_total", + # With --routing-key-override, follow-up turns route through the sticky + # pin, which shows up here (occupied_hit/occupied_miss/vacant/...) — + # cache-aware debug lines then cover only delegated (turn-1) decisions. + "smg_manual_policy_branch_total", + "smg_routing_key_source_total", + # Buffered vs streamed body routing, per path/reason — the direct + # verification of which request-body regime a leg actually ran in. + "smg_router_request_body_path_total", + # Remote-index lookups as the gateway sees them: outcome counters and + # the latency histogram the read p50/p90/p99 rows are computed from. + "smg_remote_index_", +) + +BRANCH_RE = re.compile(r'branch="?([A-Za-z0-9_.-]+)"?') + +# Every `index_service` key the harness forwards or acts on, and every +# top-level fault-drill key it implements. A profile (or scenario leg) that +# sets anything else in these shapes fails at run start: a silently ignored +# knob would produce a plausible comparison table with no independent +# variable — the worst failure mode a measurement harness can have. +SUPPORTED_INDEX_KEYS = { + "replicas", + "bridge", + "inferred_ttl_secs", + "default_capacity_blocks", + "sweep_interval_secs", + "event_ttl_secs", + "apply_delay_stored_ms", + "apply_delay_removed_ms", + "partitionable", + "deferred_replicas", + # Gateways reach replica 0 through severable proxies (gateway_partition). + "gateway_proxy", + # Pull every live replica's state at the end of the run and record the + # per-holder divergence between replicas (consistency audit). + "dump_on_exit", + # Forwarded to the service: peer anti-entropy period (0 disables). + "anti_entropy_secs", +} +SUPPORTED_DRILLS = { + "restart_smgs_at_secs", + "kill_index_replica", + "flap_index_replica", + "hang_index_replica", + "start_deferred_replica", + "partition_drill", + # Worker churn: deregister N workers from every gateway, then bring N + # NEW workers (fresh ports) up and register them mid-run. + "remove_workers_drill", + "add_workers_drill", + # One gateway restarted cold while the others keep serving. + "restart_one_smg_drill", + # Every replica killed and relaunched in turn (rolling restart). + "rolling_replica_restart_drill", + # Gateways lose their link to the index (all of them, or half). + "gateway_partition_drill", +} +DRILL_KEY_RE = re.compile(r"(_drill|_replica|_at_secs)$") + + +# ---- small helpers ---------------------------------------------------------- + + +def log(msg): + print("==> " + msg, flush=True) + + +def epoch_ms(): + return int(time.time() * 1000) + + +def http_get(url, timeout): + with urllib.request.urlopen(url, timeout=timeout) as resp: + return resp.read().decode("utf-8", "replace") + + +def http_post_json(url, payload, timeout): + req = urllib.request.Request( + url, + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=timeout) as resp: + resp.read() + + +def raise_nofile_limit(): + # Thousands of mock ports + per-SMG upstream connections need plenty of + # file descriptors; children inherit the raised limit. + try: + import resource + except ImportError: + return + soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) + target = hard if hard != resource.RLIM_INFINITY else 1048576 + if soft >= target: + return + try: + resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard)) + log(f"raised RLIMIT_NOFILE {soft} -> {target}") + except (ValueError, OSError): + log(f"WARN: could not raise RLIMIT_NOFILE beyond {soft}") + + +def flags_from(params): + """dict -> CLI flags: {"decode_base_ms": 4.3} -> ["--decode-base-ms", "4.3"]. + + Booleans are value-style ("--stream true"), matching mock-worker's + --prefix-cache convention; None values are omitted. + """ + flags = [] + for key, val in params.items(): + if val is None: + continue + flag = "--" + key.replace("_", "-") + if isinstance(val, bool): + flags += [flag, "true" if val else "false"] + else: + flags += [flag, str(val)] + return flags + + +def load_profile(path): + with open(path) as f: + return json.load(f) + + +def apply_override(profile, dotted, value): + """Set a dotted path ("loadgen.ingress") in the profile dict.""" + node = profile + keys = dotted.split(".") + for key in keys[:-1]: + node = node.setdefault(key, {}) + node[keys[-1]] = value + + +def parse_override_arg(raw): + key, _, val = raw.partition("=") + if not _: + raise SystemExit("--override expects key=value, got: " + raw) + try: + return key, json.loads(val) + except ValueError: + return key, val + + +def validate_profile(profile): + """Reject knobs the harness would silently ignore (see SUPPORTED_*).""" + unknown_index = set(profile.get("index_service") or {}) - SUPPORTED_INDEX_KEYS + unknown_drills = { + key for key in profile if DRILL_KEY_RE.search(key) and key not in SUPPORTED_DRILLS + } + problems = [] + if unknown_index: + problems.append(f"index_service keys {sorted(unknown_index)} are not implemented") + if unknown_drills: + problems.append(f"drill keys {sorted(unknown_drills)} are not implemented") + index_cfg = profile.get("index_service") or {} + if not index_cfg: + needs_index = SUPPORTED_DRILLS - { + "restart_smgs_at_secs", + "remove_workers_drill", + "add_workers_drill", + "restart_one_smg_drill", + } + used = sorted(k for k in needs_index if profile.get(k)) + if used: + problems.append(f"drills {used} need an index_service block") + if profile.get("partition_drill") and not index_cfg.get("partitionable"): + problems.append("partition_drill needs index_service.partitionable = true") + if profile.get("gateway_partition_drill") and not index_cfg.get("gateway_proxy"): + problems.append("gateway_partition_drill needs index_service.gateway_proxy = true") + cfg = profile.get("restart_one_smg_drill") or {} + if cfg and not 0 <= int(cfg.get("smg", 0)) < int(profile["smg_count"]): + problems.append("restart_one_smg_drill.smg must be < smg_count") + cfg = profile.get("remove_workers_drill") or {} + if cfg and not 0 < int(cfg.get("count", 0)) < int(profile["workers_total"]): + problems.append("remove_workers_drill.count must be in (0, workers_total)") + if 0 in (index_cfg.get("deferred_replicas") or []): + problems.append("replica 0 (the publish endpoint) cannot be deferred") + if index_cfg: + # A drill naming a replica outside the fleet (or a non-deferred one + # for the deferred start) would fail only at drill time. + replicas = int(index_cfg.get("replicas", 1)) + deferred = {int(r) for r in index_cfg.get("deferred_replicas") or []} + for key in ("kill_index_replica", "flap_index_replica", "hang_index_replica"): + cfg = profile.get(key) or {} + if cfg and not 0 <= int(cfg.get("replica", 1)) < replicas: + problems.append(f"{key}.replica must be < index_service.replicas ({replicas})") + cfg = profile.get("start_deferred_replica") or {} + if cfg and int(cfg.get("replica", -1)) not in deferred: + problems.append( + "start_deferred_replica.replica must be in index_service.deferred_replicas" + ) + if problems: + raise SystemExit("profile invalid; the drill would not fire: " + "; ".join(problems)) + + +# ---- process management ----------------------------------------------------- + + +def spawn(name, cmd, log_path, env=None): + fh = open(log_path, "ab") + # Children must never inherit our stdout: an unread pipe fills and blocks + # the process (see scale_test.sh), so everything goes to per-process logs. + proc = subprocess.Popen(cmd, stdout=fh, stderr=subprocess.STDOUT, env=env) + return {"name": name, "proc": proc, "log": fh} + + +def teardown(children, bins): + for child in reversed(children): + if child["proc"].poll() is None: + # A SIGSTOPped child (hang drill) never handles SIGTERM. + try: + child["proc"].send_signal(signal.SIGCONT) + except OSError: + pass + child["proc"].terminate() + deadline = time.time() + 10 + for child in reversed(children): + remaining = max(0.1, deadline - time.time()) + try: + child["proc"].wait(timeout=remaining) + except subprocess.TimeoutExpired: + child["proc"].kill() + child["log"].close() + children.clear() + # Safety net for anything reparented or leaked from prior runs. Anchored to + # the start of the command line so it matches only processes exec'd from + # these binaries — not this orchestrator, whose own argv can mention the + # smg path (--smg-bin), and not unrelated smg processes. + for path in bins: + if path: + subprocess.run( + ["pkill", "-9", "-f", "^" + str(path)], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=False, + ) + + +def wait_health(url, timeout, what): + deadline = time.time() + timeout + while time.time() < deadline: + try: + http_get(url, timeout=5) + return + except OSError: + time.sleep(1) + raise RuntimeError(f"{what} never became healthy at {url}") + + +def wait_tcp(port, timeout, what): + """Liveness for listeners with no HTTP surface (gRPC workers).""" + deadline = time.time() + timeout + while time.time() < deadline: + try: + with socket.create_connection(("127.0.0.1", port), timeout=5): + return + except OSError: + time.sleep(1) + raise RuntimeError(f"{what} never accepted TCP on port {port}") + + +class TcpProxy: + """A severable TCP forwarder: `listen_port` -> 127.0.0.1:`target_port`. + + The partition drill puts one in front of each index replica for its + PEERS (never for gateways or the bridge): severing closes every relayed + connection and refuses new ones, so the replicas cannot reach each other + while every client still can — a true inter-replica network partition. + """ + + def __init__(self, listen_port, target_port): + self.target_port = target_port + self._open = threading.Event() + self._open.set() + self._conns = set() + self._lock = threading.Lock() + self._server = socket.create_server(("127.0.0.1", listen_port)) + # 0 = ephemeral (tests); read back the bound port either way. + self.listen_port = self._server.getsockname()[1] + self._server.settimeout(0.5) + self._closed = False + threading.Thread(target=self._accept_loop, daemon=True).start() + + def _accept_loop(self): + while not self._closed: + try: + client, _ = self._server.accept() + except TimeoutError: + # socket.timeout is a distinct class before Python 3.10. + continue + except OSError: + return + if not self._open.is_set(): + client.close() + continue + try: + upstream = socket.create_connection(("127.0.0.1", self.target_port), timeout=5) + except OSError: + client.close() + continue + with self._lock: + self._conns.update((client, upstream)) + for src, dst in ((client, upstream), (upstream, client)): + threading.Thread(target=self._pump, args=(src, dst), daemon=True).start() + + def _pump(self, src, dst): + try: + while True: + data = src.recv(65536) + if not data: + break + dst.sendall(data) + except OSError: + pass + finally: + self._drop(src, dst) + + def _drop(self, *socks): + with self._lock: + for s in socks: + self._conns.discard(s) + for s in socks: + try: + s.shutdown(socket.SHUT_RDWR) + except OSError: + pass + try: + s.close() + except OSError: + pass + + def sever(self): + """Refuse new connections and cut every relayed one.""" + self._open.clear() + with self._lock: + live = list(self._conns) + self._drop(*live) + + def heal(self): + self._open.set() + + def close(self): + self._closed = True + self.sever() + try: + self._server.close() + except OSError: + pass + + +def ensure_local_tokenizer(): + """Generate (once) a WordLevel tokenizer.json covering every token id + the sim can produce (loadgen ids < 150k, sim outputs < 30k), so the + gateway's gRPC pipeline can resolve and decode without any network + fetch. Returned path goes into each gRPC worker's tokenizer_path label. + """ + path = REPO_ROOT / "target" / "generate-sim" / "wordlevel-tokenizer.json" + if path.exists(): + return path + path.parent.mkdir(parents=True, exist_ok=True) + vocab = {f"t{i}": i for i in range(150_000)} + vocab[""] = 150_000 + tok = { + "version": "1.0", + "truncation": None, + "padding": None, + "added_tokens": [], + "normalizer": None, + "pre_tokenizer": {"type": "Whitespace"}, + "post_processor": None, + "decoder": None, + "model": {"type": "WordLevel", "vocab": vocab, "unk_token": ""}, + } + with open(path, "w") as f: + json.dump(tok, f) + log(f"generated local tokenizer: {path.name}") + return path + + +# ---- run steps -------------------------------------------------------------- + + +def build_binaries(target_dir, build_gateway, build_index): + """Release build of the harness binaries. The index service package + (`smg-radix-index`) is built only when the profile asks for an + `index_service`: it is optional on this branch, and a profile that + requests it on a tree without the crate fails here, loudly.""" + packages = ["mock-worker", "sim-loadgen"] + if build_gateway: + packages.append("smg") + if build_index: + packages.append("smg-radix-index") + cmd = ["cargo", "build", "--release"] + for pkg in packages: + cmd += ["-p", pkg] + log("building (release): " + " ".join(packages)) + env = dict(os.environ) + env["CARGO_TARGET_DIR"] = str(target_dir) + env["RUSTC_WRAPPER"] = "" + subprocess.run(cmd, cwd=str(REPO_ROOT), env=env, check=True) + + +def launch_mocks(profile, logs_dir, mock_bin): + total = int(profile["workers_total"]) + procs = int(profile["mock_processes"]) + grpc = profile.get("worker_mode", "http") == "grpc" + port_flags = ( + ("--grpc-base-port", "--grpc-count") if grpc else ("--http-base-port", "--http-count") + ) + per_proc = (total + procs - 1) // procs + children = [] + started = 0 + for j in range(procs): + count = min(per_proc, total - started) + base = MOCK_BASE_PORT + started + cmd = [ + str(mock_bin), + "--host", + "127.0.0.1", + port_flags[0], + str(base), + port_flags[1], + str(count), + "--model", + profile.get("model_id", "mock-model"), + ] + flags_from(profile.get("mock", {})) + children.append(spawn(f"mock-{j}", cmd, logs_dir / f"mock-{j}.log")) + started += count + mode = "grpc" if grpc else "http" + log( + f"mock fleet: {total} {mode} workers over {procs} processes " + f"(ports {MOCK_BASE_PORT}-{MOCK_BASE_PORT + total - 1})" + ) + time.sleep(2) + for child in children: + if child["proc"].poll() is not None: + raise RuntimeError("{} exited early; see its log".format(child["name"])) + if grpc: + wait_tcp(MOCK_BASE_PORT, 30, "mock fleet") + else: + wait_health(f"http://127.0.0.1:{MOCK_BASE_PORT}/health", 30, "mock fleet") + return children + + +def index_replica_cmd(profile, index_bin, replica, bootstrap_from=None): + """argv for index replica `replica`. Peers are every OTHER replica in the + profile (deferred ones included: they join later and must already be in + everyone's relay list), reached through the severable proxies when the + profile is `partitionable`.""" + cfg = profile.get("index_service") or {} + replicas = int(cfg.get("replicas", 1)) + peer_base = INDEX_PROXY_BASE if cfg.get("partitionable") else INDEX_BASE_PORT + cmd = [ + str(index_bin), + "--port", + str(INDEX_BASE_PORT + replica), + "--metrics-port", + str(INDEX_METRICS_BASE + replica), + ] + peers = ",".join(f"http://127.0.0.1:{peer_base + j}" for j in range(replicas) if j != replica) + if peers: + cmd += ["--peers", peers] + if bootstrap_from is not None: + cmd += ["--bootstrap-from", f"http://127.0.0.1:{INDEX_BASE_PORT + bootstrap_from}"] + for key in ( + "inferred_ttl_secs", + "default_capacity_blocks", + "sweep_interval_secs", + "event_ttl_secs", + "apply_delay_stored_ms", + "apply_delay_removed_ms", + "anti_entropy_secs", + ): + if key in cfg: + cmd += ["--" + key.replace("_", "-"), str(cfg[key])] + return cmd + + +def spawn_index_replica(profile, logs_dir, index_bin, replica, bootstrap_from=None, tag=""): + env = dict(os.environ) + env["RUST_LOG"] = "info" + cmd = index_replica_cmd(profile, index_bin, replica, bootstrap_from) + child = spawn(f"index-{replica}", cmd, logs_dir / f"index-{replica}{tag}.log", env=env) + wait_tcp(INDEX_BASE_PORT + replica, 30, f"index-{replica}") + return child + + +def launch_index_service(profile, logs_dir, index_bin, bridge_bin): + """Optional radix index service (+ event bridge) from the profile's + `index_service` block, e.g. + {"replicas": 2, "bridge": true, "inferred_ttl_secs": 18, + "sweep_interval_secs": 1, "default_capacity_blocks": N, + "apply_delay_stored_ms": 0, "partitionable": false, + "deferred_replicas": []} + Replicas relay to each other; the bridge (when enabled) subscribes to + every gRPC worker and publishes to replica 0. `partitionable` routes + peer traffic through severable proxies (partition drill); replicas in + `deferred_replicas` are listed as peers but launched later by the + `start_deferred_replica` drill. Returns (children, proxies).""" + cfg = profile.get("index_service") + if not cfg: + return [], [] + replicas = int(cfg.get("replicas", 1)) + deferred = {int(r) for r in cfg.get("deferred_replicas") or []} + proxies = [] + if cfg.get("partitionable"): + proxies = [TcpProxy(INDEX_PROXY_BASE + i, INDEX_BASE_PORT + i) for i in range(replicas)] + if cfg.get("gateway_proxy"): + # Two gateway-facing proxies to replica 0 (see INDEX_GW_PROXY_BASE); + # launch_smgs points each gateway at one of them. + proxies += [TcpProxy(INDEX_GW_PROXY_BASE + k, INDEX_BASE_PORT) for k in range(2)] + children = [] + for i in range(replicas): + if i in deferred: + continue + children.append(spawn_index_replica(profile, logs_dir, index_bin, i)) + bridged = 0 + if cfg.get("bridge", True): + total = int(profile["workers_total"]) + grpc = profile.get("worker_mode", "http") == "grpc" + grpc_workers = ( + [f"grpc://127.0.0.1:{port}" for port in range(MOCK_BASE_PORT, MOCK_BASE_PORT + total)] + if grpc + else [] + ) + if grpc_workers: + env = dict(os.environ) + env["RUST_LOG"] = "info" + cmd = [ + str(bridge_bin), + "--workers", + ",".join(grpc_workers), + "--index", + f"http://127.0.0.1:{INDEX_BASE_PORT}", + "--model", + profile.get("model_id", "mock-model"), + "--block-size", + str(profile.get("mock", {}).get("block_size", 128)), + ] + children.append(spawn("bridge", cmd, logs_dir / "bridge.log", env=env)) + bridged = len(grpc_workers) + log( + f"index service: {replicas - len(deferred)}/{replicas} replicas on ports " + f"{INDEX_BASE_PORT}.. (bridging {bridged} grpc workers" + f"{', partitionable' if proxies else ''})" + ) + return children, proxies + + +def smg_cmd(profile, smg_bin, i): + """argv for gateway `i`. With `index_service.gateway_proxy`, the + `--kv-indexer-url` flag is rewritten to this gateway's severable proxy + (even gateways +0, odd +1) so a drill can cut all or half of the fleet + off from the index.""" + cmd = [ + str(smg_bin), + "--host", + "127.0.0.1", + "--port", + str(SMG_BASE_PORT + i), + "--prometheus-host", + "127.0.0.1", + "--prometheus-port", + str(PROM_BASE_PORT + i), + ] + list(profile["smg_flags"]) + if (profile.get("index_service") or {}).get("gateway_proxy") and "--kv-indexer-url" in cmd: + cmd[cmd.index("--kv-indexer-url") + 1] = f"http://127.0.0.1:{INDEX_GW_PROXY_BASE + (i % 2)}" + if profile.get("mesh_smgs"): + # Gateway-to-gateway mesh (TreeSync of approximate-tree inserts): + # per-instance port. The gateway uses only mesh_peer_urls[0], as + # its ONE-SHOT gossip init peer (round 0 only), so every + # instance bootstraps from smg-0 — whose port is bound first + # and gated on below — rather than from a gateway this loop has + # not spawned yet. + peers = [f"127.0.0.1:{MESH_BASE_PORT}"] if i else [] + cmd += [ + "--enable-mesh", + "--mesh-host", + "127.0.0.1", + "--mesh-advertise-host", + "127.0.0.1", + "--mesh-port", + str(MESH_BASE_PORT + i), + "--mesh-server-name", + f"smg-{i}", + ] + if peers: + cmd += ["--mesh-peer-urls"] + peers + return cmd + + +def launch_one_smg(profile, logs_dir, smg_bin, i, tag=""): + env = dict(os.environ) + env["RUST_LOG"] = SMG_RUST_LOG + child = spawn(f"smg-{i}", smg_cmd(profile, smg_bin, i), logs_dir / f"smg-{i}{tag}.log", env=env) + if profile.get("mesh_smgs") and i == 0: + # smg-0 must be listening before any peer's single init round. + wait_tcp(MESH_BASE_PORT, 30, "smg-0 mesh") + return child + + +def launch_smgs(profile, logs_dir, smg_bin): + count = int(profile["smg_count"]) + children = [launch_one_smg(profile, logs_dir, smg_bin, i) for i in range(count)] + log(f"gateways: {count} on ports {SMG_BASE_PORT}.. (prometheus {PROM_BASE_PORT}..)") + for i in range(count): + wait_health(f"http://127.0.0.1:{SMG_BASE_PORT + i}/health", 60, f"smg-{i}") + return children + + +def register_workers(profile, worker_ports=None, smg_ports=None): + """POST every worker URL to every SMG: 64-way per SMG, SMGs in parallel. + + Same WorkerSpec shape scale_test.sh proved at 2k ports; health disabled so + workers are instantly routable and registration cost stays isolated from + the health-probe loop. `worker_ports` / `smg_ports` narrow the fan-out + (the add-workers and single-gateway-restart drills). + """ + total = int(profile["workers_total"]) + model_id = profile.get("model_id", "mock-model") + grpc = profile.get("worker_mode", "http") == "grpc" + if worker_ports is None: + worker_ports = list(range(MOCK_BASE_PORT, MOCK_BASE_PORT + total)) + if smg_ports is None: + smg_ports = [SMG_BASE_PORT + i for i in range(int(profile["smg_count"]))] + tokenizer_path = str(ensure_local_tokenizer()) if grpc else None + + def register_one(smg_port, worker_port): + if grpc: + # The URL scheme selects the connection mode; runtime picks the + # proto dialect the mock implements. weight_version is relayed + # verbatim in every response's meta_info — it is how the + # loadgen learns which worker served a request (the gRPC + # router exposes no other worker identity). + body = { + "url": f"grpc://127.0.0.1:{worker_port}", + "connection_mode": "grpc", + "runtime": "tokenspeed", + "models": [{"id": model_id}], + "kv_block_size": int(profile.get("mock", {}).get("block_size", 128)), + "labels": { + "tokenizer_path": tokenizer_path, + "weight_version": str(worker_port), + }, + "health": {"disable_health_check": True}, + } + else: + body = { + "url": f"http://127.0.0.1:{worker_port}", + "connection_mode": "http", + "runtime": "sglang", + "models": [{"id": model_id}], + "health": {"disable_health_check": True}, + } + try: + http_post_json(f"http://127.0.0.1:{smg_port}/workers", body, timeout=20) + return True + except OSError: + return False + + def register_all(smg_port): + ok = 0 + with ThreadPoolExecutor(max_workers=64) as pool: + for success in pool.map(lambda p: register_one(smg_port, p), worker_ports): + ok += 1 if success else 0 + return ok + + log(f"registering {len(worker_ports)} workers with {len(smg_ports)} SMGs via POST /workers") + registered = {} + with ThreadPoolExecutor(max_workers=len(smg_ports)) as pool: + for port, ok in zip(smg_ports, pool.map(register_all, smg_ports)): + registered[port] = ok + for port in smg_ports: + log(f" smg :{port} accepted {registered[port]}/{len(worker_ports)}") + return registered + + +def wait_ready(profile): + """Gate on every SMG reporting >= readiness_fraction of the fleet. + + A timeout FAILS the run: a report over a fleet that was never fully + routable would describe a topology that did not exist (imbalance + divides by workers_total, cache hit depends on the reachable set). + """ + total = int(profile["workers_total"]) + need = max(1, int(total * float(profile.get("readiness_fraction", 0.99)))) + timeout = float(profile.get("readiness_timeout_secs", 300)) + smg_ports = [SMG_BASE_PORT + i for i in range(int(profile["smg_count"]))] + log(f"waiting for >= {need}/{total} workers per SMG (timeout {timeout:.0f}s)") + counts = {port: 0 for port in smg_ports} + deadline = time.time() + timeout + while time.time() < deadline: + for port in smg_ports: + try: + body = http_get(f"http://127.0.0.1:{port}/workers", timeout=60) + counts[port] = body.count('"url"') + except OSError: + pass + if all(n >= need for n in counts.values()): + log(" ready: " + " ".join(f"{p}:{counts[p]}" for p in smg_ports)) + return counts + time.sleep(2) + seen = " ".join(f"{p}:{counts[p]}" for p in smg_ports) + raise RuntimeError(f"readiness timeout: need >= {need}/{total} workers per SMG, got {seen}") + + +# ---- sampling --------------------------------------------------------------- + + +def ps_stats(pids): + out = subprocess.run( + ["ps", "-o", "pid=,rss=,pcpu=", "-p", ",".join(str(p) for p in pids)], + capture_output=True, + text=True, + check=False, + ) + stats = {} + for line in out.stdout.splitlines(): + parts = line.split() + if len(parts) >= 3: + try: + stats[int(parts[0])] = (int(parts[1]), float(parts[2])) + except ValueError: + pass + return stats + + +def fd_count(pid): + proc_fd = f"/proc/{pid}/fd" + if os.path.isdir(proc_fd): + try: + return len(os.listdir(proc_fd)) + except OSError: + return None + # Darwin: no /proc; lsof is accurate but slow at very high fd counts, + # which is why the full profile sets sample_fds=false. + out = subprocess.run(["lsof", "-p", str(pid)], capture_output=True, text=True, check=False) + if out.returncode != 0: + return None + return max(0, len(out.stdout.splitlines()) - 1) + + +def scrape_metrics(prom_port): + try: + body = http_get(f"http://127.0.0.1:{prom_port}/metrics", timeout=4) + except OSError: + return {} + values = {} + for line in body.splitlines(): + if line.startswith("#") or not line.startswith(METRIC_PREFIXES): + continue + parts = line.rsplit(None, 1) + if len(parts) != 2: + continue + try: + values[parts[0]] = float(parts[1]) + except ValueError: + pass + return values + + +def scrape_index_metrics(metrics_port): + """radix_index_* lines from a replica's admin port (gauges, counters and + the apply/query latency histograms).""" + try: + body = http_get(f"http://127.0.0.1:{metrics_port}/metrics", timeout=2) + except OSError: + return {} + values = {} + for line in body.splitlines(): + if line.startswith("#") or not line.startswith("radix_index_"): + continue + parts = line.rsplit(None, 1) + if len(parts) != 2: + continue + try: + values[parts[0]] = float(parts[1]) + except ValueError: + pass + return values + + +def sampler_loop(stop, smg_pids, out_path, interval, sample_fds, index_pids=None): + """`index_pids`: callable returning [(replica, pid)] for the live index + replicas (they change across kill/relaunch drills), or None.""" + start = time.time() + with open(out_path, "a") as out: + while True: + index_live = index_pids() if index_pids else [] + stats = ps_stats(list(smg_pids) + [pid for _, pid in index_live]) + index_entries = [] + for replica, pid in index_live: + rss, cpu = stats.get(pid, (None, None)) + index_entries.append( + { + "replica": replica, + "pid": pid, + "rss_kib": rss, + "cpu_pct": cpu, + "metrics": scrape_index_metrics(INDEX_METRICS_BASE + replica), + } + ) + entries = [] + for idx, pid in enumerate(smg_pids): + rss, cpu = stats.get(pid, (None, None)) + entry = { + "idx": idx, + "pid": pid, + "rss_kib": rss, + "cpu_pct": cpu, + "fds": fd_count(pid) if sample_fds and rss is not None else None, + "metrics": scrape_metrics(PROM_BASE_PORT + idx), + } + entries.append(entry) + record = { + "ts": round(time.time(), 3), + "elapsed_s": round(time.time() - start, 1), + "smg": entries, + "index": index_entries, + } + out.write(json.dumps(record) + "\n") + out.flush() + if stop.wait(interval): + return + + +# ---- report ----------------------------------------------------------------- + + +ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") + + +def branch_counts(log_path): + counts = Counter() + if not log_path.exists(): + return counts + with open(log_path, errors="replace") as f: + for line in f: + if "Cache-aware selection" not in line: + continue + # tracing's fmt layer colors field names, leaving escape codes + # between `branch` and `=`; strip them before matching. + m = BRANCH_RE.search(ANSI_RE.sub("", line)) + if m: + counts[m.group(1)] += 1 + return counts + + +def imbalance(counter, workers_total): + if not counter: + return {"requests": 0, "distinct_workers": 0} + observed = list(counter.values()) + fleet = observed + [0] * max(0, workers_total - len(observed)) + mean = statistics.mean(fleet) + return { + "requests": sum(observed), + "distinct_workers": len(observed), + "workers_total": workers_total, + "cov_fleet": round(statistics.pstdev(fleet) / mean, 4) if mean else None, + "max_over_mean_fleet": round(max(fleet) / mean, 2) if mean else None, + "cov_observed": ( + round(statistics.pstdev(observed) / statistics.mean(observed), 4) + if statistics.mean(observed) + else None + ), + } + + +def analyze_requests(path, workers_total): + per_worker = Counter() + per_turn_worker = {} + turn_stats = {} + session_turn_worker = {} + index_sources = Counter() + prediction_errors = [] + # Per-minute timeline of follow-up cache ratio and index outcomes, so a + # drift over a long run (capacity, staleness, leaks) is visible as a + # curve rather than averaged away. + timeline = {} + t0 = None + if not path.exists(): + return {"error": "requests.jsonl missing"} + with open(path, errors="replace") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + rec = json.loads(line) + except ValueError: + continue + start_ms = rec.get("start_ms") + if start_ms is not None: + t0 = start_ms if t0 is None else min(t0, start_ms) + turn = int(rec.get("turn", 1)) + ts = turn_stats.setdefault(turn, {"n": 0, "prompt": 0, "cached": 0, "hits": 0}) + ts["n"] += 1 + prompt = rec.get("prompt_tokens") or 0 + cached = rec.get("cached_tokens") or 0 + ts["prompt"] += prompt + ts["cached"] += cached + # Request-level "hit" per the design doc: cached/prompt >= 0.3. + if prompt and cached / prompt >= 0.3: + ts["hits"] += 1 + port = rec.get("worker_port") + if port is not None: + per_worker[port] += 1 + per_turn_worker.setdefault(turn, Counter())[port] += 1 + session = rec.get("session") + if session is not None: + session_turn_worker.setdefault(session, {})[turn] = port + src = rec.get("index_source") + if src: + index_sources[src] += 1 + pred = rec.get("index_predicted_tokens") + if pred is not None and cached is not None: + prediction_errors.append(int(pred) - int(cached)) + if start_ms is not None and turn >= 2 and rec.get("status") == 200 and prompt: + minute = int(start_ms // 60_000) + bucket = timeline.setdefault( + minute, {"n": 0, "prompt": 0, "cached": 0, "timeouts": 0, "disconnected": 0} + ) + bucket["n"] += 1 + bucket["prompt"] += prompt + bucket["cached"] += cached + bucket["timeouts"] += src == "remote_timeout" + bucket["disconnected"] += src == "remote_disconnected" + both = [s for s in session_turn_worker.values() if 1 in s and 2 in s] + same = sum(1 for s in both if s[1] == s[2]) + turns = {} + for turn, ts in sorted(turn_stats.items()): + turns[f"turn{turn}"] = { + "requests": ts["n"], + # Raw sums so every ratio in the tables is verifiable. + "prompt_tokens_sum": ts["prompt"], + "cached_tokens_sum": ts["cached"], + "cached_over_prompt": round(ts["cached"] / ts["prompt"], 4) if ts["prompt"] else None, + "hit_rate": round(ts["hits"] / ts["n"], 4) if ts["n"] else None, + "imbalance": imbalance(per_turn_worker.get(turn, Counter()), workers_total), + } + return { + "overall_imbalance": imbalance(per_worker, workers_total), + "turns": turns, + "t2_sessions": len(both), + "t2_same_worker_rate": round(same / len(both), 4) if both else None, + "index_sources": dict(index_sources), + "index_prediction_error_tokens": prediction_error_summary(prediction_errors), + # Minute-by-minute follow-up cache ratio (token-weighted) and index + # outcome shares, keyed by minutes since the first request. + "followup_timeline": [ + { + "minute": minute - int(t0 // 60_000), + "requests": b["n"], + "cached_over_prompt": round(b["cached"] / b["prompt"], 4), + "timeout_share": round(b["timeouts"] / b["n"], 4), + "disconnected_share": round(b["disconnected"] / b["n"], 4), + } + for minute, b in sorted(timeline.items()) + ] + if timeline and t0 is not None + else [], + } + + +def prediction_error_summary(errors): + """predicted − actual cached tokens over every index-routed request: + signed mean (bias), absolute p50/p90/p95/max, and the share of exact + answers. An index that is wrong by a block on one request in twenty is + a different thing from one wrong by 4k tokens on every request, and the + p95 alone cannot tell them apart.""" + if not errors: + return None + abs_sorted = sorted(abs(e) for e in errors) + n = len(abs_sorted) + pick = lambda q: abs_sorted[min(n - 1, int(q * (n - 1)))] # noqa: E731 + return { + "requests": n, + "mean": round(statistics.fmean(errors), 2), + "exact_share": round(sum(1 for e in errors if e == 0) / n, 4), + "p50_abs": pick(0.50), + "p90_abs": pick(0.90), + "p95_abs": pick(0.95), + "max_abs": abs_sorted[-1], + } + + +HIST_BUCKET_RE = re.compile(r"^(?P[a-zA-Z_:]+)_bucket\{(?P[^}]*)\}$") + + +def histogram_deltas(first, last, name): + """Cumulative bucket counts of Prometheus histogram `name` accumulated + between two metric snapshots (dicts of metric-line -> value), summed over + any extra labels: {upper_bound_seconds: count}, plus the sample count.""" + buckets = {} + for key, val in last.items(): + m = HIST_BUCKET_RE.match(key) + if not m or m.group("name") != name: + continue + le = None + for label in m.group("labels").split(","): + label = label.strip() + if label.startswith("le="): + le = label[3:].strip('"') + if le is None: + continue + bound = float("inf") if le == "+Inf" else float(le) + buckets[bound] = buckets.get(bound, 0.0) + val - first.get(key, 0.0) + return buckets + + +def histogram_percentiles(buckets, quantiles=(0.5, 0.9, 0.99)): + """Percentiles (seconds) from cumulative buckets by linear interpolation + inside the bucket that crosses each quantile; the +Inf bucket reports + its lower edge (the last finite bound) — a floor, never an estimate.""" + if not buckets: + return {} + bounds = sorted(buckets) + total = buckets[bounds[-1]] + if total <= 0: + return {} + out = {} + for q in quantiles: + target = q * total + prev_bound, prev_count = 0.0, 0.0 + for bound in bounds: + count = buckets[bound] + if count >= target: + if bound == float("inf"): + out[f"p{int(q * 100)}"] = prev_bound + else: + span = count - prev_count + frac = (target - prev_count) / span if span > 0 else 1.0 + out[f"p{int(q * 100)}"] = prev_bound + frac * (bound - prev_bound) + break + prev_bound, prev_count = bound, count + return {k: round(v * 1000, 3) for k, v in out.items()} # milliseconds + + +def summarize_index_samples(path, window=None): + """Per-replica peak RSS / mean CPU inside `window`, and the service-side + apply and query latency percentiles over the window (histogram deltas, + summed across replicas); plus the gateway-side remote-index query + latency percentiles (summed across gateways).""" + per_replica = {} + first_index, last_index = {}, {} + first_gw, last_gw = {}, {} + if not path.exists(): + return {} + with open(path, errors="replace") as f: + for line in f: + try: + rec = json.loads(line) + except ValueError: + continue + if window is not None: + elapsed = rec.get("elapsed_s") + if elapsed is None or not (window[0] <= elapsed <= window[1]): + continue + for entry in rec.get("index", []): + r = entry.get("replica") + s = per_replica.setdefault(r, {"rss_kib": [], "cpu_pct": []}) + if entry.get("rss_kib") is not None: + s["rss_kib"].append(entry["rss_kib"]) + if entry.get("cpu_pct") is not None: + s["cpu_pct"].append(entry["cpu_pct"]) + metrics = entry.get("metrics") or {} + if metrics: + # Counters are monotonic per replica process; a relaunch + # resets them, so keep the first/last per replica and sum + # the deltas afterwards. + first_index.setdefault(r, metrics) + last_index[r] = metrics + for entry in rec.get("smg", []): + metrics = entry.get("metrics") or {} + if any(k.startswith("smg_remote_index_query_duration_seconds") for k in metrics): + first_gw.setdefault(entry.get("idx"), metrics) + last_gw[entry.get("idx")] = metrics + + def merged(firsts, lasts, name): + total = {} + for key in lasts: + for bound, count in histogram_deltas(firsts.get(key, {}), lasts[key], name).items(): + total[bound] = total.get(bound, 0.0) + count + return total + + out = { + "replicas": { + str(r): { + "rss_peak_kib": max(s["rss_kib"]) if s["rss_kib"] else None, + "cpu_mean_pct": round(statistics.mean(s["cpu_pct"]), 1) if s["cpu_pct"] else None, + "cpu_peak_pct": max(s["cpu_pct"]) if s["cpu_pct"] else None, + } + for r, s in sorted(per_replica.items()) + }, + "apply_ms": histogram_percentiles( + merged(first_index, last_index, "radix_index_apply_duration_seconds") + ), + "query_ms": histogram_percentiles( + merged(first_index, last_index, "radix_index_query_duration_seconds") + ), + "gateway_lookup_ms": histogram_percentiles( + merged(first_gw, last_gw, "smg_remote_index_query_duration_seconds") + ), + } + return out + + +def summarize_samples(path, smg_count, window=None): + """Aggregate sampler records; with `window` = (start_s, end_s), only + samples with `elapsed_s` inside it count. The sampler starts AFTER the + warmup sleep, so `elapsed_s == 0` is already the start of the load + period: the steady-state window is (0, duration_secs).""" + series = [ + {"rss_kib": [], "cpu_pct": [], "fds": [], "queue_depth": [], "conns": []} + for _ in range(smg_count) + ] + last_counters = [{"rejected": 0.0, "selections": 0.0} for _ in range(smg_count)] + sticky_branches = [{} for _ in range(smg_count)] + body_paths = [{} for _ in range(smg_count)] + if not path.exists(): + return [] + with open(path, errors="replace") as f: + for line in f: + try: + rec = json.loads(line) + except ValueError: + continue + if window is not None: + elapsed = rec.get("elapsed_s") + if elapsed is None or not (window[0] <= elapsed <= window[1]): + continue + for entry in rec.get("smg", []): + idx = entry.get("idx") + if idx is None or idx >= smg_count: + continue + s = series[idx] + if entry.get("rss_kib") is not None: + s["rss_kib"].append(entry["rss_kib"]) + if entry.get("cpu_pct") is not None: + s["cpu_pct"].append(entry["cpu_pct"]) + if entry.get("fds") is not None: + s["fds"].append(entry["fds"]) + metrics = entry.get("metrics", {}) + depth = conns = rejected = selections = 0.0 + for key, val in metrics.items(): + if key.startswith("smg_admission_queue_depth"): + depth += val + elif key.startswith("smg_http_connections_active"): + conns += val + elif key.startswith("smg_admission_queue_rejected_total"): + rejected += val + elif key.startswith("smg_worker_selection_total"): + selections += val + elif key.startswith("smg_manual_policy_branch_total"): + m = BRANCH_RE.search(key) + if m: + sticky_branches[idx][m.group(1)] = val + elif key.startswith("smg_router_request_body_path_total"): + m = re.search(r'path="?(\w+)"?.*?reason="?(\w+)"?', key) + if m: + body_paths[idx]["{}:{}".format(*m.groups())] = val + if metrics: + s["queue_depth"].append(depth) + s["conns"].append(conns) + last_counters[idx] = {"rejected": rejected, "selections": selections} + + def agg(values, as_int=True): + if not values: + return {"peak": None, "mean": None} + mean = statistics.mean(values) + return { + "peak": max(values), + "mean": int(mean) if as_int else round(mean, 1), + } + + out = [] + for idx in range(smg_count): + s = series[idx] + out.append( + { + "idx": idx, + "rss_kib": agg(s["rss_kib"]), + "cpu_pct": agg(s["cpu_pct"], as_int=False), + "fds": agg(s["fds"]), + "queue_depth": agg(s["queue_depth"]), + "http_connections_active": agg(s["conns"]), + "rejected_total": last_counters[idx]["rejected"], + "worker_selection_total": last_counters[idx]["selections"], + # Final counter values: sticky-session outcomes under + # --routing-key-override (occupied_hit = pinned follow-up). + "sticky_branches": sticky_branches[idx], + # Final "path:reason" counters — verifies buffered vs + # streamed routing directly. + "body_paths": body_paths[idx], + } + ) + return out + + +def build_report(run_dir): + run_dir = Path(run_dir) + profile = load_profile(run_dir / "profile.json") + smg_count = int(profile["smg_count"]) + workers_total = int(profile["workers_total"]) + + summary_path = run_dir / "summary.json" + loadgen_summary = {} + if summary_path.exists(): + loadgen_summary = load_profile(summary_path) + + meta_path = run_dir / "meta.json" + meta = load_profile(meta_path) if meta_path.exists() else {} + + per_smg_branches = [] + for i in range(smg_count): + counts = branch_counts(run_dir / "logs" / f"smg-{i}.log") + per_smg_branches.append({"idx": i, "branches": dict(counts)}) + + report = { + "profile": profile, + "run_dir": str(run_dir), + "loadgen_summary": loadgen_summary, + "requests": analyze_requests(run_dir / "requests.jsonl", workers_total), + "samples": summarize_samples( + run_dir / "samples.jsonl", + smg_count, + window=(0, int(profile["duration_secs"])), + ), + "index_samples": summarize_index_samples( + run_dir / "samples.jsonl", + window=(0, int(profile["duration_secs"])), + ), + "index_divergence": meta.get("index_divergence"), + "cache_aware_branches": per_smg_branches, + # Drill outcomes ride into the report so a drill that never fired + # (or failed half-way) cannot be mistaken for a measured effect. + "drills": { + k: v + for k, v in meta.items() + if k.startswith("index_") + or k.startswith("workers_") + or k.startswith("smg_") + or k.startswith("gateway_") + or "restart" in k + or k == "loadgen_exit" + or any(k.startswith(d) for d in SUPPORTED_DRILLS) + }, + } + with open(run_dir / "report.json", "w") as f: + json.dump(report, f, indent=2) + write_markdown(report, run_dir / "report.md") + return report + + +def _fmt(val): + if val is None: + return "n/a" + if isinstance(val, float): + return f"{val:.4g}" + return str(val) + + +def write_markdown(report, path): + profile = report["profile"] + lines = [] + if int(profile.get("workers_total", 0)) < 4000: + lines.append( + "> **Reduced-scale run.** Cache/affinity semantics are " + "meaningful; CPU/RSS/fd/connection figures are NOT " + "production-representative (fleet size, body size, concurrency, " + "sticky-map cardinality, and stream chunking are all scaled " + "down). Use profiles/full.local.json on a large host for " + "resource conclusions." + ) + lines.append("") + lines.append("# generate-sim report — {}".format(profile.get("name", "run"))) + lines.append("") + loadgen = profile.get("loadgen", {}) + derived_rps = None + if "session_rps" in loadgen: + derived_rps = float(loadgen["session_rps"]) * (1.0 + float(loadgen.get("t2_ratio", 0))) + lines.append("| key | value |") + lines.append("|---|---|") + for key, val in [ + ("run_dir", report["run_dir"]), + ("smg_count", profile.get("smg_count")), + ("workers_total", profile.get("workers_total")), + ("mock_processes", profile.get("mock_processes")), + ("duration_secs", profile.get("duration_secs")), + ("target aggregate rps", derived_rps), + ("ingress", loadgen.get("ingress")), + ("system_prefix_tokens", loadgen.get("system_prefix_tokens")), + ("t2_ratio", loadgen.get("t2_ratio")), + ]: + lines.append(f"| {key} | {_fmt(val)} |") + lines.append("") + + drills = report.get("drills") or {} + if drills: + lines.append("## Drills (from meta.json)") + lines.append("") + lines.append("| key | value |") + lines.append("|---|---|") + for key in sorted(drills): + lines.append(f"| {key} | {_fmt(drills[key])} |") + lines.append("") + + summary = report.get("loadgen_summary", {}) + scalars = {k: v for k, v in summary.items() if isinstance(v, (int, float, str, bool))} + if scalars: + lines.append("## Loadgen summary") + lines.append("") + lines.append("| key | value |") + lines.append("|---|---|") + for key in sorted(scalars): + lines.append(f"| {key} | {_fmt(scalars[key])} |") + lines.append("") + + req = report.get("requests", {}) + if "overall_imbalance" in req: + lines.append("## Worker balance (from requests.jsonl)") + lines.append("") + lines.append( + "| slice | requests | distinct | CoV (fleet) | max/mean | cached/prompt | hit rate |" + ) + lines.append("|---|---|---|---|---|---|---|") + imb = req["overall_imbalance"] + lines.append( + "| overall | {} | {}/{} | {} | {} | | |".format( + _fmt(imb.get("requests")), + _fmt(imb.get("distinct_workers")), + _fmt(imb.get("workers_total")), + _fmt(imb.get("cov_fleet")), + _fmt(imb.get("max_over_mean_fleet")), + ) + ) + for name, ts in sorted(req.get("turns", {}).items()): + timb = ts.get("imbalance", {}) + lines.append( + "| {} | {} | {}/{} | {} | {} | {} | {} |".format( + name, + _fmt(ts.get("requests")), + _fmt(timb.get("distinct_workers")), + _fmt(timb.get("workers_total")), + _fmt(timb.get("cov_fleet")), + _fmt(timb.get("max_over_mean_fleet")), + _fmt(ts.get("cached_over_prompt")), + _fmt(ts.get("hit_rate")), + ) + ) + lines.append("") + lines.append( + "turn-2 same-worker rate: {} (over {} two-turn sessions)".format( + _fmt(req.get("t2_same_worker_rate")), _fmt(req.get("t2_sessions")) + ) + ) + lines.append("") + + branches = report.get("cache_aware_branches", []) + all_names = sorted({name for e in branches for name in e["branches"]}) + if all_names: + lines.append("## Cache-aware branches (per SMG, from debug logs)") + lines.append("") + lines.append("| smg | " + " | ".join(all_names) + " | total |") + lines.append("|---|" + "---|" * (len(all_names) + 1)) + for entry in branches: + row = [str(entry["idx"])] + row += [str(entry["branches"].get(name, 0)) for name in all_names] + row.append(str(sum(entry["branches"].values()))) + lines.append("| " + " | ".join(row) + " |") + lines.append("") + + samples = report.get("samples", []) + if samples: + lines.append("## Gateway resources (per SMG, 5 s samples)") + lines.append("") + lines.append( + "| smg | rss peak MiB | rss mean MiB | cpu mean % | cpu peak % | fds peak " + "| queue peak | conns peak | rejected | selections |" + ) + lines.append("|---|---|---|---|---|---|---|---|---|---|") + + def mib(kib): + return _fmt(round(kib / 1024, 1)) if kib is not None else "n/a" + + for s in samples: + cells = [ + str(s["idx"]), + mib(s["rss_kib"]["peak"]), + mib(s["rss_kib"]["mean"]), + _fmt(s["cpu_pct"]["mean"]), + _fmt(s["cpu_pct"]["peak"]), + _fmt(s["fds"]["peak"]), + _fmt(s["queue_depth"]["peak"]), + _fmt(s["http_connections_active"]["peak"]), + _fmt(s["rejected_total"]), + _fmt(s["worker_selection_total"]), + ] + lines.append("| " + " | ".join(cells) + " |") + lines.append("") + + with open(path, "w") as f: + f.write("\n".join(lines)) + + +# ---- orchestration ---------------------------------------------------------- + + +def _repo_relative(path): + """Repo-relative rendering for provenance: committed artifacts must not + carry absolute local paths.""" + path = Path(path) + try: + return str(path.resolve().relative_to(REPO_ROOT.resolve())) + except ValueError: + return path.name + + +def _git(args): + try: + return ( + subprocess.run( + ["git"] + args, + cwd=str(REPO_ROOT), + capture_output=True, + text=True, + timeout=10, + check=False, + ).stdout.strip() + or None + ) + except OSError: + return None + + +def _sha256(path): + try: + digest = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + except OSError: + return None + + +class Drills: + """The mid-run fault drills. Each runs on its own daemon thread, records + what it did (epoch-ms timestamps) in `meta`, and on ANY failure records + `_error` instead of dying silently — a drill that half-happened + must be distinguishable from a measured blackout. `children` mutations + are serialized with teardown through `lock`, and every step checks + `stop` so a drill cannot relaunch anything into a run that is ending. + """ + + def __init__( + self, + profile, + meta, + children, + lock, + stop, + logs_dir, + index_bin, + proxies, + bins=None, + smg_pids=None, + ): + self.profile = profile + self.meta = meta + self.children = children + self.lock = lock + self.stop = stop + self.logs_dir = logs_dir + self.index_bin = index_bin + # Inter-replica proxies come first (one per replica when + # `partitionable`), then the two gateway-facing ones (`gateway_proxy`). + replicas = int((profile.get("index_service") or {}).get("replicas", 1)) + n_peer = replicas if (profile.get("index_service") or {}).get("partitionable") else 0 + self.proxies = proxies[:n_peer] + self.gateway_proxies = proxies[n_peer:] + self.bins = bins or {} + self.smg_pids = smg_pids + self._threads = [] + + def start_all(self): + for key, target in ( + ("kill_index_replica", self._kill_index_replica), + ("flap_index_replica", self._flap_index_replica), + ("hang_index_replica", self._hang_index_replica), + ("start_deferred_replica", self._start_deferred_replica), + ("partition_drill", self._partition), + ("remove_workers_drill", self._remove_workers), + ("add_workers_drill", self._add_workers), + ("restart_one_smg_drill", self._restart_one_smg), + ("rolling_replica_restart_drill", self._rolling_replica_restart), + ("gateway_partition_drill", self._gateway_partition), + ): + cfg = self.profile.get(key) + if cfg: + thread = threading.Thread( + target=self._guarded, args=(key, target, cfg), daemon=True + ) + thread.start() + self._threads.append(thread) + + def join(self, timeout): + """Let in-flight drills finish recording before meta is written: a + drill past its last sleep (e.g. mid-relaunch) would otherwise write + into a dict that has already been dumped, or make the dump raise.""" + deadline = time.time() + timeout + for thread in self._threads: + thread.join(timeout=max(0.0, deadline - time.time())) + + def _guarded(self, key, target, cfg): + try: + target(cfg) + except Exception as e: # noqa: BLE001 — the outcome IS the result + log(f"WARN: drill {key} failed: {e!r}") + self.meta[f"{key}_error"] = repr(e) + + def _sleep(self, secs): + """Sleep unless the run is ending; returns False when it is.""" + return not self.stop.wait(float(secs)) + + def _live_replica(self, replica): + name = f"index-{replica}" + with self.lock: + return [c for c in self.children if c["name"] == name and c["proc"].poll() is None] + + def _kill(self, replica): + victims = self._live_replica(replica) + if not victims: + raise RuntimeError(f"index-{replica} is not running; nothing to kill") + for child in victims: + child["proc"].kill() + return epoch_ms() + + def _relaunch(self, replica, tag): + # Bootstrap from the lowest live replica that is not this one. + replicas = int(self.profile["index_service"].get("replicas", 1)) + source = next( + (r for r in range(replicas) if r != replica and self._live_replica(r)), + None, + ) + child = spawn_index_replica( + self.profile, self.logs_dir, self.index_bin, replica, bootstrap_from=source, tag=tag + ) + with self.lock: + if self.stop.is_set(): + teardown([child], []) + raise RuntimeError("run ended during relaunch; replica torn down") + self.children.append(child) + return epoch_ms() + + def _kill_index_replica(self, cfg): + replica = int(cfg.get("replica", 1)) + if not self._sleep(cfg.get("at_secs", 60)): + return + log(f"drill: killing index-{replica}") + self.meta["index_killed_at_ms"] = self._kill(replica) + self.meta["index_killed_replica"] = replica + relaunch_after = cfg.get("relaunch_after_secs") + if relaunch_after is not None and self._sleep(relaunch_after): + log(f"drill: relaunching index-{replica}") + self.meta["index_relaunched_at_ms"] = self._relaunch(replica, "-relaunch") + + def _flap_index_replica(self, cfg): + replica = int(cfg.get("replica", 1)) + cycles = int(cfg.get("cycles", 3)) + period = float(cfg.get("period_secs", 20)) + if not self._sleep(cfg.get("at_secs", 45)): + return + events = self.meta.setdefault("index_flap_events", []) + for cycle in range(cycles): + log(f"drill: flap {cycle + 1}/{cycles} of index-{replica}") + killed = self._kill(replica) + if not self._sleep(period / 2): + return + relaunched = self._relaunch(replica, f"-flap{cycle + 1}") + events.append({"killed_at_ms": killed, "relaunched_at_ms": relaunched}) + if cycle + 1 < cycles and not self._sleep(period / 2): + return + self.meta["index_flap_replica"] = replica + + def _hang_index_replica(self, cfg): + replica = int(cfg.get("replica", 1)) + if not self._sleep(cfg.get("at_secs", 60)): + return + victims = self._live_replica(replica) + if not victims: + raise RuntimeError(f"index-{replica} is not running; nothing to hang") + log(f"drill: SIGSTOP index-{replica} (TCP up, nothing drains)") + for child in victims: + child["proc"].send_signal(signal.SIGSTOP) + self.meta["index_hung_at_ms"] = epoch_ms() + self.meta["index_hung_replica"] = replica + resume_after = cfg.get("resume_after_secs") + if resume_after is not None and self._sleep(resume_after): + log(f"drill: SIGCONT index-{replica}") + for child in victims: + child["proc"].send_signal(signal.SIGCONT) + self.meta["index_resumed_at_ms"] = epoch_ms() + + def _start_deferred_replica(self, cfg): + replica = int(cfg["replica"]) + deferred = {int(r) for r in self.profile["index_service"].get("deferred_replicas") or []} + if replica not in deferred: + raise RuntimeError(f"replica {replica} is not in index_service.deferred_replicas") + if not self._sleep(cfg.get("at_secs", 60)): + return + log(f"drill: starting deferred index-{replica} (bootstrap under load)") + self.meta["index_deferred_started_at_ms"] = self._relaunch(replica, "-deferred") + self.meta["index_deferred_replica"] = replica + + # ---- worker churn -------------------------------------------------------- + + def _smg_ports(self): + return [SMG_BASE_PORT + i for i in range(int(self.profile["smg_count"]))] + + def _remove_workers(self, cfg): + """Deregister the LAST `count` workers from every gateway (DELETE + /workers/{id}, resolved from each gateway's own listing). The mock + listeners stay up: this is the control-plane removal a drain or a + scale-down performs, and the index must learn it through the + gateways' lifecycle signals, not through the worker dying.""" + count = int(cfg["count"]) + total = int(self.profile["workers_total"]) + if not self._sleep(cfg.get("at_secs", 60)): + return + victims = {MOCK_BASE_PORT + total - 1 - k for k in range(count)} + log(f"drill: removing {count} workers from every gateway") + removed = {} + for smg_port in self._smg_ports(): + listing = json.loads(http_get(f"http://127.0.0.1:{smg_port}/workers", timeout=30)) + workers = listing["workers"] if isinstance(listing, dict) else listing + done = 0 + for w in workers: + port = int(str(w.get("url", "")).rsplit(":", 1)[-1] or 0) + if port in victims: + req = urllib.request.Request( + f"http://127.0.0.1:{smg_port}/workers/{w['id']}", method="DELETE" + ) + with urllib.request.urlopen(req, timeout=30): + done += 1 + removed[str(smg_port)] = done + self.meta["workers_removed_at_ms"] = epoch_ms() + self.meta["workers_removed"] = {"ports": sorted(victims), "per_gateway": removed} + + def _add_workers(self, cfg): + """Bring `count` NEW workers up (a fresh mock process on ports past + the fleet) and register them with every gateway mid-run.""" + count = int(cfg["count"]) + total = int(self.profile["workers_total"]) + if not self._sleep(cfg.get("at_secs", 60)): + return + base = MOCK_BASE_PORT + total + grpc = self.profile.get("worker_mode", "http") == "grpc" + port_flags = ( + ("--grpc-base-port", "--grpc-count") if grpc else ("--http-base-port", "--http-count") + ) + cmd = [ + str(self.bins["mock"]), + "--host", + "127.0.0.1", + port_flags[0], + str(base), + port_flags[1], + str(count), + "--model", + self.profile.get("model_id", "mock-model"), + ] + flags_from(self.profile.get("mock", {})) + env = dict(os.environ) + env["RUST_LOG"] = "info" + log(f"drill: adding {count} new workers on ports {base}-{base + count - 1}") + child = spawn("mock-added", cmd, self.logs_dir / "mock-added.log", env=env) + wait_tcp(base, 30, "added workers") + with self.lock: + if self.stop.is_set(): + teardown([child], []) + return + self.children.append(child) + registered = register_workers(self.profile, worker_ports=list(range(base, base + count))) + self.meta["workers_added_at_ms"] = epoch_ms() + self.meta["workers_added"] = {"ports": [base, base + count - 1], "per_gateway": registered} + + # ---- gateway churn ------------------------------------------------------- + + def _restart_one_smg(self, cfg): + """Kill ONE gateway and relaunch it cold (no local trees, no sticky + pins); re-register every worker with it. The shared-index thesis in + one drill: a cold gateway routes on the fleet's knowledge from its + first request, or it does not.""" + idx = int(cfg.get("smg", 0)) + if not self._sleep(cfg.get("at_secs", 60)): + return + log(f"drill: restarting smg-{idx} cold") + with self.lock: + old = [ + c for c in self.children if c["name"] == f"smg-{idx}" and c["proc"].poll() is None + ] + for child in old: + child["proc"].kill() + self.meta["smg_restarted_at_ms"] = epoch_ms() + child = launch_one_smg(self.profile, self.logs_dir, self.bins["smg"], idx, tag="-restart") + with self.lock: + if self.stop.is_set(): + teardown([child], []) + raise RuntimeError("run ended during gateway restart") + self.children.append(child) + wait_health(f"http://127.0.0.1:{SMG_BASE_PORT + idx}/health", 60, f"smg-{idx}") + register_workers(self.profile, smg_ports=[SMG_BASE_PORT + idx]) + if self.smg_pids is not None and idx < len(self.smg_pids): + self.smg_pids[idx] = child["proc"].pid + self.meta["smg_restarted"] = idx + self.meta["smg_restart_ready_at_ms"] = epoch_ms() + + # ---- replica churn ------------------------------------------------------- + + def _rolling_replica_restart(self, cfg): + """Kill and relaunch every replica in turn, `gap_secs` apart, each + bootstrapping from a live peer — the rolling-restart a deploy does.""" + replicas = int(self.profile["index_service"].get("replicas", 1)) + gap = float(cfg.get("gap_secs", 20)) + if not self._sleep(cfg.get("at_secs", 60)): + return + events = self.meta.setdefault("index_rolling_restart_events", []) + for replica in range(replicas): + log(f"drill: rolling restart — killing index-{replica}") + killed = self._kill(replica) + if not self._sleep(gap / 2): + return + relaunched = self._relaunch(replica, f"-rolling{replica}") + events.append( + {"replica": replica, "killed_at_ms": killed, "relaunched_at_ms": relaunched} + ) + if replica + 1 < replicas and not self._sleep(gap / 2): + return + + # ---- partitions ---------------------------------------------------------- + + def _gateway_partition(self, cfg): + """Sever the gateways' link to the index: `scope: all` cuts every + gateway (even- and odd-numbered proxies), `scope: half` only the + even-numbered ones — an asymmetric partition where half the fleet + routes on the index and half falls open to load-only.""" + if not self.gateway_proxies: + raise RuntimeError("gateway_partition_drill needs index_service.gateway_proxy") + scope = cfg.get("scope", "all") + targets = self.gateway_proxies if scope == "all" else self.gateway_proxies[:1] + if not self._sleep(cfg.get("at_secs", 60)): + return + log(f"drill: severing gateway→index links ({scope})") + for proxy in targets: + proxy.sever() + self.meta["gateway_partitioned_at_ms"] = epoch_ms() + self.meta["gateway_partition_scope"] = scope + heal_after = cfg.get("heal_after_secs") + if heal_after is not None and self._sleep(heal_after): + log("drill: healing gateway→index links") + for proxy in targets: + proxy.heal() + self.meta["gateway_partition_healed_at_ms"] = epoch_ms() + + def _partition(self, cfg): + if not self.proxies: + raise RuntimeError("partition_drill needs index_service.partitionable") + if not self._sleep(cfg.get("at_secs", 60)): + return + log("drill: severing every inter-replica link") + for proxy in self.proxies: + proxy.sever() + self.meta["index_partitioned_at_ms"] = epoch_ms() + heal_after = cfg.get("heal_after_secs") + if heal_after is not None and self._sleep(heal_after): + log("drill: healing the partition") + for proxy in self.proxies: + proxy.heal() + self.meta["index_healed_at_ms"] = epoch_ms() + + +def placement_capacity_blocks(profile): + """The per-worker capacity the gateway publishes for placement-fed + holders: the mock's KV budget in blocks (None when the profile has no + realistic mock section).""" + mock = profile.get("mock") or {} + kv_tokens = mock.get("kv_tokens") + block_size = mock.get("block_size") + if not kv_tokens or not block_size: + return None + return -(-int(kv_tokens) // int(block_size)) + + +def classify_divergence(base, other, capacity_blocks): + """Compare two replicas' per-holder dumps. A differing holder is one + whose block SET differs (digest). Event-fed holders are sequenced + ground truth and must be identical once anti-entropy has run. + Placement-fed holders are copied, never agreed on: each replica runs + its own capacity cut, so two replicas that both hold a worker above + its capacity legitimately differ by WHEN they cut (the 1x-2x + hysteresis band). Those are counted separately from a placement + holder that differs while under capacity on either side, which + would be a lost update.""" + differing_event, in_band, out_of_band, only_in_one = [], [], [], [] + for key in set(base) | set(other): + if key not in base or key not in other: + only_in_one.append(key) + continue + a, b = base[key], other[key] + if a["digest"] == b["digest"]: + continue + if a.get("event_fed") or b.get("event_fed"): + differing_event.append(key) + elif capacity_blocks is not None and min(a["blocks"], b["blocks"]) >= capacity_blocks: + in_band.append(key) + else: + out_of_band.append(key) + return { + "holders_differing": len(differing_event) + len(in_band) + len(out_of_band), + "holders_differing_event_fed": len(differing_event), + "holders_differing_placement_in_band": len(in_band), + "holders_differing_placement_out_of_band": len(out_of_band), + "holders_only_in_one": len(only_in_one), + "capacity_blocks": capacity_blocks, + # Converged = nothing a peer should have repaired is left: every + # event-fed holder identical, every holder on every replica, and + # no placement holder differing outside the capacity band. + "converged": not differing_event and not out_of_band and not only_in_one, + } + + +def dump_replicas(profile, run_dir, dump_bin, live_replicas): + """Pull every live replica's state with radix-index-dump and compare + their per-holder digests (see `classify_divergence`). Returns the + divergence summary recorded in meta.""" + dumps = {} + for replica in live_replicas: + out = subprocess.run( + [str(dump_bin), "--connect", f"http://127.0.0.1:{INDEX_BASE_PORT + replica}"], + capture_output=True, + text=True, + timeout=120, + check=False, + ) + if out.returncode != 0: + dumps[replica] = {"error": out.stderr.strip()[-400:]} + continue + try: + dumps[replica] = json.loads(out.stdout) + except ValueError as e: + dumps[replica] = {"error": f"unparsable dump: {e}"} + with open(run_dir / f"index-dump-{replica}.json", "w") as f: + f.write(out.stdout) + good = {r: d for r, d in dumps.items() if "holders" in d} + summary = { + "replicas": sorted(dumps), + "errors": {str(r): d["error"] for r, d in dumps.items() if "error" in d}, + "total_blocks": {str(r): d.get("total_blocks") for r, d in good.items()}, + } + if len(good) >= 2: + base_r = min(good) + capacity = placement_capacity_blocks(profile) + merged = None + for r, d in good.items(): + if r == base_r: + continue + part = classify_divergence(good[base_r]["holders"], d["holders"], capacity) + if merged is None: + merged = part + else: + for k, v in part.items(): + if isinstance(v, int) and not isinstance(v, bool): + merged[k] = max(merged[k], v) + elif isinstance(v, bool): + merged[k] = merged[k] and v + summary["holders_compared"] = len(good[base_r]["holders"]) + summary.update(merged) + return summary + + +def run_profile(profile, run_dir, smg_bin=None, skip_build=False): + """Full run: build -> mocks -> SMGs -> register -> loadgen -> report. + + Returns the run dir; report.json / report.md are inside it. + """ + validate_profile(profile) + run_dir = Path(run_dir) + logs_dir = run_dir / "logs" + logs_dir.mkdir(parents=True, exist_ok=True) + + # Pin the target dir so build and binary resolution always agree + # (see scale_test.sh); override by exporting CARGO_TARGET_DIR. + target_dir = Path(os.environ.get("CARGO_TARGET_DIR") or REPO_ROOT / "target") + wants_index = bool(profile.get("index_service")) + if not skip_build: + build_binaries(target_dir, build_gateway=smg_bin is None, build_index=wants_index) + smg_bin = Path(smg_bin) if smg_bin else target_dir / "release" / "smg" + mock_bin = target_dir / "release" / "mock-worker" + loadgen_bin = target_dir / "release" / "sim-loadgen" + index_bin = target_dir / "release" / "radix-index-service" + bridge_bin = target_dir / "release" / "radix-index-bridge" + dump_bin = target_dir / "release" / "radix-index-dump" + required = [smg_bin, mock_bin, loadgen_bin] + if wants_index: + required += [index_bin, bridge_bin] + if (profile.get("index_service") or {}).get("dump_on_exit"): + required.append(dump_bin) + for path in required: + if not os.access(str(path), os.X_OK): + raise SystemExit(f"binary missing: {path} (drop --skip-build?)") + + if profile.get("requires_large_linux_host"): + log("NOTE: " + str(profile["requires_large_linux_host"])) + + with open(run_dir / "profile.json", "w") as f: + json.dump(profile, f, indent=2) + + raise_nofile_limit() + # Every binary the run may spawn, so a crashed run cannot leave a stale + # index replica holding 40000/40100 for the next run to silently reuse. + all_bins = [smg_bin, mock_bin, loadgen_bin, index_bin, bridge_bin] + teardown([], all_bins) # clear leftovers from prior runs so ports are free + time.sleep(1) + + children = [] + children_lock = threading.Lock() + proxies = [] + meta = { + "smg_bin": _repo_relative(smg_bin), + "started_at": datetime.now().isoformat(), + # Provenance: enough to reproduce or audit any table built from this + # run — repo state, exact binaries, profile content, and seed. + "git_commit": _git(["rev-parse", "HEAD"]), + "git_dirty": bool(_git(["status", "--porcelain"])), + "binary_sha256": { + "smg": _sha256(smg_bin), + "mock-worker": _sha256(mock_bin), + "sim-loadgen": _sha256(loadgen_bin), + "radix-index-service": _sha256(index_bin) if index_bin.exists() else None, + "radix-index-bridge": _sha256(bridge_bin) if bridge_bin.exists() else None, + }, + "profile_sha256": hashlib.sha256(json.dumps(profile, sort_keys=True).encode()).hexdigest(), + "loadgen_seed": profile.get("loadgen", {}).get("seed"), + "host": platform.platform(), + } + stop = threading.Event() + sampler = None + drills = None + try: + children += launch_mocks(profile, logs_dir, mock_bin) + index_children, proxies = launch_index_service(profile, logs_dir, index_bin, bridge_bin) + children += index_children + children += launch_smgs(profile, logs_dir, smg_bin) + meta["registered"] = register_workers(profile) + meta["ready"] = wait_ready(profile) + + warmup = float(profile.get("warmup_secs", 10)) + log(f"warmup sleep {warmup:.0f}s") + time.sleep(warmup) + + smg_pids = [c["proc"].pid for c in children if c["name"].startswith("smg-")] + + def live_index_pids(): + with children_lock: + return [ + (int(c["name"].split("-", 1)[1]), c["proc"].pid) + for c in children + if c["name"].startswith("index-") and c["proc"].poll() is None + ] + + sampler = threading.Thread( + target=sampler_loop, + args=( + stop, + smg_pids, + run_dir / "samples.jsonl", + float(profile.get("sample_interval_secs", 5)), + bool(profile.get("sample_fds", True)), + live_index_pids if wants_index else None, + ), + daemon=True, + ) + sampler.start() + + duration = int(profile["duration_secs"]) + smg_urls = ",".join( + f"http://127.0.0.1:{SMG_BASE_PORT + i}" for i in range(int(profile["smg_count"])) + ) + cmd = [ + str(loadgen_bin), + "--smg-urls", + smg_urls, + "--duration-secs", + str(duration), + "--out", + str(run_dir), + ] + flags_from(profile.get("loadgen", {})) + log(f"loadgen: {duration}s run") + loadgen = spawn("loadgen", cmd, logs_dir / "loadgen.log") + with children_lock: + children.append(loadgen) + + drills = Drills( + profile, + meta, + children, + children_lock, + stop, + logs_dir, + index_bin, + proxies, + bins={"smg": smg_bin, "mock": mock_bin}, + smg_pids=smg_pids, + ) + drills.start_all() + + # Optional mid-window gateway restart: sticky pins and hash + # placements are process state, so affinity must rebuild from + # scratch; requests during the blackout fail and count as errors. + restart_at = profile.get("restart_smgs_at_secs") + if restart_at: + + def _restart_smgs(): + if stop.wait(float(restart_at)): + return + log("restarting all SMGs (sticky pins and placements lost)") + meta["restart_attempted_at_secs"] = restart_at + try: + with children_lock: + old = [c for c in children if c["name"].startswith("smg-")] + for child in old: + if child["proc"].poll() is None: + child["proc"].kill() + new_smgs = launch_smgs(profile, logs_dir, smg_bin) + with children_lock: + if stop.is_set(): + teardown(new_smgs, []) + return + children.extend(new_smgs) + register_workers(profile) + # Re-gate readiness: registration only waits for the + # POSTs, not for readiness_fraction, and metrics over a + # partial relaunched fleet would read as measured. + meta["restart_ready"] = wait_ready(profile) + # The sampler reads this list each tick; swap in the new pids. + smg_pids[:] = [c["proc"].pid for c in new_smgs] + meta["restarted_at_secs"] = restart_at + except Exception as e: # noqa: BLE001 — the outcome IS the result + log(f"WARN: SMG restart drill failed: {e!r}") + meta["restart_error"] = repr(e) + + threading.Thread(target=_restart_smgs, daemon=True).start() + try: + meta["loadgen_exit"] = loadgen["proc"].wait(timeout=duration * 3 + 300) + except subprocess.TimeoutExpired: + log("WARN: loadgen overran; killing") + loadgen["proc"].kill() + meta["loadgen_exit"] = "timeout" + finally: + stop.set() + if sampler is not None: + sampler.join(timeout=30) + if drills is not None: + drills.join(timeout=10) + if wants_index and (profile.get("index_service") or {}).get("dump_on_exit"): + # Consistency audit before the replicas go away: pull every live + # replica and diff their per-holder block sets. + with children_lock: + live = sorted( + int(c["name"].split("-", 1)[1]) + for c in children + if c["name"].startswith("index-") and c["proc"].poll() is None + ) + try: + meta["index_divergence"] = dump_replicas(profile, run_dir, dump_bin, live) + log(f"index divergence: {meta['index_divergence']}") + except Exception as e: # noqa: BLE001 — the outcome IS the result + meta["index_divergence"] = {"error": repr(e)} + with children_lock: + teardown(children, all_bins) + for proxy in proxies: + proxy.close() + meta["finished_at"] = datetime.now().isoformat() + with open(run_dir / "meta.json", "w") as f: + json.dump(meta, f, indent=2) + + build_report(run_dir) + if meta.get("loadgen_exit") not in (0, None): + # The loadgen exits non-zero only for run-invalidating conditions + # (>50% errors, truncated requests.jsonl, timeout). The report is + # still written for diagnosis, but the run must not read as measured. + raise SystemExit( + f"loadgen exited {meta['loadgen_exit']}; the run is not usable " + f"(see {run_dir / 'logs' / 'loadgen.log'})" + ) + log(f"report: {run_dir / 'report.md'}") + return run_dir + + +def default_run_dir(profile, tag=None): + stamp = datetime.now().strftime("%Y%m%d-%H%M%S") + name = profile.get("name", "run") + if tag: + name = f"{name}-{tag}" + return REPO_ROOT / "target" / "generate-sim" / (f"{name}-{stamp}") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + sub = parser.add_subparsers(dest="cmd", required=True) + + run_p = sub.add_parser("run", help="run one profile end to end") + run_p.add_argument("--profile", required=True, help="path to a profile JSON") + run_p.add_argument("--skip-build", action="store_true", help="use existing binaries") + run_p.add_argument( + "--smg-bin", + help="prebuilt gateway binary (e.g. from another checkout for policy A/B); " + "skips building the smg package", + ) + run_p.add_argument("--out", help="run directory (default target/generate-sim/-)") + run_p.add_argument("--tag", help="suffix for the default run dir name") + run_p.add_argument( + "--override", + action="append", + default=[], + metavar="KEY=VALUE", + help="dotted profile override, e.g. loadgen.ingress=random (repeatable)", + ) + + report_p = sub.add_parser("report", help="rebuild report.json/report.md for a run dir") + report_p.add_argument("--run-dir", required=True) + + args = parser.parse_args() + if args.cmd == "run": + profile = load_profile(args.profile) + for raw in args.override: + key, val = parse_override_arg(raw) + apply_override(profile, key, val) + run_dir = Path(args.out) if args.out else default_run_dir(profile, args.tag) + run_profile(profile, run_dir, smg_bin=args.smg_bin, skip_build=args.skip_build) + elif args.cmd == "report": + build_report(args.run_dir) + log(f"report: {Path(args.run_dir) / 'report.md'}") + + +if __name__ == "__main__": + main() diff --git a/scripts/generate_sim/test_generate_sim.py b/scripts/generate_sim/test_generate_sim.py new file mode 100755 index 000000000..a20d700cc --- /dev/null +++ b/scripts/generate_sim/test_generate_sim.py @@ -0,0 +1,575 @@ +#!/usr/bin/env python3 +"""Focused tests for the generate-sim harness: profile invariants the +benchmark's validity depends on, metric aggregation, scenario construction +(controlled comparisons differ in exactly the intended knob), and the +branch-log parser. Run with: + + python3 -m unittest discover -s scripts/generate_sim -p 'test_*.py' +""" + +import json +import tempfile +import unittest +from pathlib import Path + +import scenarios +import sim + +PROFILES = Path(__file__).resolve().parent / "profiles" +PRODUCTION_PROFILES = ["local-small.json", "local-medium.json", "full.template.json"] + + +def load(name): + with open(PROFILES / name) as f: + return json.load(f) + + +class ProfileInvariants(unittest.TestCase): + def test_production_profiles_enable_sticky_override(self): + # Production runs --routing-key-override; a profile without it + # measures hash-placement affinity instead of production behavior. + for name in PRODUCTION_PROFILES + ["smoke.json"]: + flags = load(name)["smg_flags"] + self.assertIn("--routing-key-override", flags, name) + self.assertIn("--assignment-mode", flags, name) + + def test_smg_routing_block_stays_128_while_engine_block_is_256(self): + for name in PRODUCTION_PROFILES: + profile = load(name) + flags = profile["smg_flags"] + self.assertEqual(flags[flags.index("--block-size") + 1], "128", name) + self.assertEqual(profile["mock"]["block_size"], 256, name) + self.assertEqual(profile["mock"]["max_running"], 80, name) + + def test_mock_blocks_use_only_flags_the_mock_worker_accepts(self): + # The engine is crates/mock_worker's realistic simulator; a key it + # does not know aborts every worker at launch (the fleet dies before + # the first request), so the profile schema is pinned to its flags. + accepted = { + "engine", + "prefill_tps", + "decode_base_ms", + "decode_per_req_ms", + "prefill_chunk", + "max_running", + "kv_tokens", + "block_size", + "prefix_cache", + } + for path in sorted(PROFILES.glob("*.json")): + mock = load(path.name)["mock"] + self.assertEqual(mock["engine"], "realistic", path.name) + self.assertTrue(mock.get("prefix_cache"), path.name) + self.assertLessEqual(set(mock), accepted, path.name) + + def test_full_profile_supports_production_concurrency(self): + # 6,000 rps x 89 s mean lifetime ~= 534k concurrent requests. + self.assertGreaterEqual(load("full.template.json")["loadgen"]["max_inflight"], 534_000) + + def test_local_profiles_target_production_worker_pressure(self): + # ~30-38 concurrent per worker: session_rps x mean_turns x lifetime + # / workers. Baseline mean turns ~1.5 at t2_ratio 0.5 / max 2. + for name in ["local-small.json", "local-medium.json"]: + profile = load(name) + lg = profile["loadgen"] + request_rps = lg["session_rps"] * 1.5 + concurrent_per_worker = request_rps * 8.9 / profile["workers_total"] + self.assertGreater(concurrent_per_worker, 25, name) + self.assertLess(concurrent_per_worker, 45, name) + + +class ScenarioConstruction(unittest.TestCase): + def test_ttl_scenario_differs_only_in_ttl(self): + base = load("local-small.json") + rendered = [] + for _, overrides, _ in scenarios.SCENARIOS["ttl-controlled"]: + profile = json.loads(json.dumps(base)) + patches = None + for key, val in overrides.items(): + if key == "smg_flag_overrides": + patches = val + else: + sim.apply_override(profile, key, val) + profile["smg_flags"] = scenarios.patch_smg_flags(profile["smg_flags"], patches) + rendered.append(profile) + a, b = rendered + self.assertEqual(a["loadgen"], b["loadgen"], "traffic must be identical") + diff = [(fa, fb) for fa, fb in zip(a["smg_flags"], b["smg_flags"]) if fa != fb] + self.assertEqual(len(a["smg_flags"]), len(b["smg_flags"])) + self.assertEqual(diff, [("18", "2")], "only the TTL value may differ") + + def test_assignment_ab_differs_only_in_mode(self): + base = load("local-small.json") + legs = scenarios.SCENARIOS["assignment-mode-ab"] + flags_a = base["smg_flags"] + flags_b = scenarios.patch_smg_flags(base["smg_flags"], legs[1][1]["smg_flag_overrides"]) + diff = [(fa, fb) for fa, fb in zip(flags_a, flags_b) if fa != fb] + self.assertEqual(diff, [("delegate", "min_group")]) + + def test_turn_mix_legs_hold_request_rps_constant(self): + # 305 sessions/s x ~1.5 turns == 110 x ~4.15 turns (within 10%). + base = load("local-small.json")["loadgen"]["session_rps"] + multi = scenarios.MULTITURN["loadgen.session_rps"] + self.assertAlmostEqual(base * 1.5, multi * 4.15, delta=base * 1.5 * 0.10) + + def test_patch_smg_flags_replaces_in_place_and_appends(self): + flags = ["--assignment-mode", "delegate", "--disable-retries"] + patched = scenarios.patch_smg_flags( + flags, {"--assignment-mode": "min_group", "--cache-ttl-secs": "18"} + ) + self.assertEqual( + patched, + [ + "--assignment-mode", + "min_group", + "--disable-retries", + "--cache-ttl-secs", + "18", + ], + ) + self.assertEqual(flags[1], "delegate", "input must not be mutated") + + def test_patch_smg_flags_false_removes_bare_and_valued_flags(self): + flags = [ + "--cache-index", + "hash", + "--routing-key-override", + "--assignment-mode", + "delegate", + "--disable-retries", + ] + patched = scenarios.patch_smg_flags(flags, scenarios.RADIX_TREE_FLAGS) + self.assertEqual(patched, ["--cache-index", "tree", "--disable-retries"]) + self.assertIn("--routing-key-override", flags, "input must not be mutated") + + def test_kv_event_legs_run_grpc_nonstreaming_with_igw(self): + for label, overrides, _ in scenarios.SCENARIOS["kv-events"]: + self.assertEqual(overrides["worker_mode"], "grpc", label) + self.assertIs( + overrides["loadgen.stream"], + False, + f"{label}: gRPC streaming's final frame carries only the tail " + "token; multi-turn context needs the full output", + ) + self.assertIn( + "--enable-igw", + overrides["smg_flag_overrides"], + f"{label}: dynamically registered gRPC workers are unreachable without IGW routing", + ) + # The event legs must drop the sticky short-circuit; the control + # must keep it (that is what makes it a control). + by_label = {label: o for label, o, _ in scenarios.SCENARIOS["kv-events"]} + self.assertIs( + by_label["event-affine"]["smg_flag_overrides"]["--routing-key-override"], + False, + ) + self.assertNotIn( + "--routing-key-override", + by_label["sticky-control"]["smg_flag_overrides"], + ) + + def test_radix_legs_share_flag_patch_and_disable_images(self): + for label, overrides, _ in scenarios.SCENARIOS["radix-replica"]: + self.assertEqual( + overrides["smg_flag_overrides"], + scenarios.RADIX_TREE_FLAGS, + f"leg {label} must route on the tree without the sticky override", + ) + self.assertEqual( + overrides["loadgen.image_count"], + 0, + f"leg {label}: placeholder expansion is ids-only; images must be off", + ) + + +class MetricAggregation(unittest.TestCase): + def _analyze(self, records): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "requests.jsonl" + with open(path, "w") as f: + for record in records: + f.write(json.dumps(record) + "\n") + return sim.analyze_requests(path, workers_total=4) + + def test_cached_over_prompt_is_token_weighted(self): + # A fully-cached short prompt plus an uncached long prompt: the + # token-weighted ratio is 0.1, NOT the 0.5 request mean. + records = [ + { + "turn": 1, + "session": 1, + "worker_port": 9001, + "prompt_tokens": 100, + "cached_tokens": 100, + "status": 200, + }, + { + "turn": 1, + "session": 2, + "worker_port": 9002, + "prompt_tokens": 900, + "cached_tokens": 0, + "status": 200, + }, + ] + report = self._analyze(records) + turn1 = report["turns"]["turn1"] + self.assertEqual(turn1["prompt_tokens_sum"], 1000) + self.assertEqual(turn1["cached_tokens_sum"], 100) + self.assertAlmostEqual(turn1["cached_over_prompt"], 0.1) + + def test_seed_aggregation_reports_mean_and_ci(self): + rows = [{"x": 1.0, "label": "a"}, {"x": 2.0, "label": "a"}, {"x": 3.0, "label": "a"}] + agg = scenarios.aggregate_seed_rows(rows) + self.assertTrue(agg["x"].startswith("2 ±"), agg["x"]) + self.assertEqual(agg["label"], "a") + + +class ArtifactHygiene(unittest.TestCase): + def test_committed_results_carry_no_absolute_local_paths(self): + # OSS hygiene: committed artifacts must not leak local usernames or + # machine paths (repo-relative provenance only). + results = Path(__file__).resolve().parent / "results" + if not results.exists(): + self.skipTest("no committed results") + offenders = [] + for f in results.rglob("*.json"): + text = f.read_text(errors="replace") + if "/Users/" in text or "/home/" in text: + offenders.append(str(f.relative_to(results))) + self.assertEqual(offenders, [], "absolute local paths in artifacts") + + +class BranchParsing(unittest.TestCase): + def test_branch_counts_strip_ansi_color(self): + line = ( + "\x1b[2m2026-08-27\x1b[0m \x1b[34mDEBUG\x1b[0m Cache-aware selection " + '\x1b[3mbranch\x1b[0m\x1b[2m=\x1b[0m"hash_hit" worker="http://w"\n' + ) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "smg.log" + with open(path, "w") as f: + f.write(line * 3) + f.write("unrelated line\n") + counts = sim.branch_counts(path) + self.assertEqual(counts["hash_hit"], 3) + + +class DrillValidation(unittest.TestCase): + def test_unknown_or_unfireable_drill_keys_fail_loudly(self): + # A knob the harness would silently ignore must fail at run start. + for bad in [ + {"index_service": {"replicas": 2}, "explode_index_replica": {"at_secs": 1}}, + {"index_service": {"replicas": 2, "bogus_knob": 1}}, + {"kill_index_replica": {"at_secs": 1}}, # no index_service to kill + {"index_service": {"replicas": 2}, "partition_drill": {"at_secs": 1}}, + {"index_service": {"replicas": 2, "deferred_replicas": [0]}}, + ]: + with self.assertRaises(SystemExit, msg=str(bad)): + sim.validate_profile(bad) + sim.validate_profile( + { + "index_service": {"replicas": 2, "partitionable": True}, + "partition_drill": {"at_secs": 1, "heal_after_secs": 1}, + } + ) + sim.validate_profile({"restart_smgs_at_secs": 60}) + + def test_every_committed_scenario_leg_validates(self): + # Every leg's knobs must be ones the harness implements, so no + # scenario can produce a comparison with no independent variable. + base = load("local-small.json") + for name, legs in scenarios.SCENARIOS.items(): + for label, overrides, _ in legs: + profile = json.loads(json.dumps(base)) + for key, val in overrides.items(): + if key != "smg_flag_overrides": + sim.apply_override(profile, key, val) + try: + sim.validate_profile(profile) + except SystemExit as e: + self.fail(f"{name}/{label}: {e}") + + def test_index_replica_cmd_forwards_delays_and_proxied_peers(self): + profile = { + "index_service": {"replicas": 3, "apply_delay_stored_ms": 300, "partitionable": True} + } + cmd = sim.index_replica_cmd(profile, "/bin/index", 1, bootstrap_from=0) + self.assertEqual(cmd[cmd.index("--apply-delay-stored-ms") + 1], "300") + self.assertEqual( + cmd[cmd.index("--peers") + 1], + f"http://127.0.0.1:{sim.INDEX_PROXY_BASE},http://127.0.0.1:{sim.INDEX_PROXY_BASE + 2}", + ) + self.assertEqual( + cmd[cmd.index("--bootstrap-from") + 1], f"http://127.0.0.1:{sim.INDEX_BASE_PORT}" + ) + plain = sim.index_replica_cmd({"index_service": {"replicas": 2}}, "/bin/index", 0) + self.assertEqual( + plain[plain.index("--peers") + 1], f"http://127.0.0.1:{sim.INDEX_BASE_PORT + 1}" + ) + self.assertNotIn("--bootstrap-from", plain) + + +class PartitionProxy(unittest.TestCase): + def test_sever_cuts_live_and_new_connections_and_heal_restores(self): + import socket + import threading + + server = socket.create_server(("127.0.0.1", 0)) + server.settimeout(10) + + def echo(): + while True: + try: + conn, _ = server.accept() + except OSError: + return + + def pump(c): + with c: + while True: + data = c.recv(4096) + if not data: + return + c.sendall(data) + + threading.Thread(target=pump, args=(conn,), daemon=True).start() + + threading.Thread(target=echo, daemon=True).start() + proxy = sim.TcpProxy(0, server.getsockname()[1]) + + def read_eof(sock): + sock.settimeout(5) + try: + return sock.recv(4) + except OSError: + return b"" + + try: + with socket.create_connection(("127.0.0.1", proxy.listen_port), timeout=5) as c: + c.sendall(b"ping") + self.assertEqual(c.recv(4), b"ping") + proxy.sever() + self.assertEqual(read_eof(c), b"", "a severed link must drop the live connection") + with socket.create_connection(("127.0.0.1", proxy.listen_port), timeout=5) as c2: + self.assertEqual(read_eof(c2), b"", "a severed link must refuse new connections") + proxy.heal() + with socket.create_connection(("127.0.0.1", proxy.listen_port), timeout=5) as c3: + c3.sendall(b"pong") + self.assertEqual(c3.recv(4), b"pong") + finally: + proxy.close() + server.close() + + +class SampleWindow(unittest.TestCase): + def test_steady_state_window_starts_at_zero(self): + # The sampler starts after the warmup sleep, so elapsed_s == 0 is + # already steady state; the first sample must count. + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "samples.jsonl" + with open(path, "w") as f: + for elapsed, rss in ((0.0, 100), (5.0, 200), (400.0, 999)): + rec = {"elapsed_s": elapsed, "smg": [{"idx": 0, "rss_kib": rss, "metrics": {}}]} + f.write(json.dumps(rec) + "\n") + out = sim.summarize_samples(path, 1, window=(0, 150)) + self.assertEqual(out[0]["rss_kib"]["peak"], 200) + self.assertEqual(out[0]["rss_kib"]["mean"], 150) + + +class FailoverBins(unittest.TestCase): + def _run(self, base): + import subprocess + import sys + + script = Path(__file__).resolve().parent / "failover_bins.py" + return subprocess.run( + [sys.executable, str(script), str(base)], capture_output=True, text=True, check=False + ) + + def test_empty_analysis_exits_non_zero_and_recorded_kill_bins(self): + with tempfile.TemporaryDirectory() as tmp: + seed = Path(tmp) / "leg" / "seed-42" + seed.mkdir(parents=True) + (seed / "meta.json").write_text(json.dumps({})) + (seed / "requests.jsonl").write_text("") + self.assertNotEqual(self._run(tmp).returncode, 0, "no kill observed must fail") + (seed / "meta.json").write_text(json.dumps({"index_killed_at_ms": 1_000_000})) + rows = [ + { + "turn": 2, + "status": 200, + "prompt_tokens": 100, + "cached_tokens": 80, + "start_ms": 1_000_000 - 5_000, + "index_source": "remote_hit", + }, + { + "turn": 2, + "status": 200, + "prompt_tokens": 100, + "cached_tokens": 20, + "start_ms": 1_000_000 + 5_000, + "index_source": "remote_timeout", + }, + ] + (seed / "requests.jsonl").write_text("".join(json.dumps(r) + "\n" for r in rows)) + done = self._run(tmp) + self.assertEqual(done.returncode, 0, done.stderr) + self.assertIn("kill observed in 1/1 seeds", done.stdout) + self.assertIn("-10s", done.stdout) + self.assertIn("+0s", done.stdout) + + +class SeedAggregation(unittest.TestCase): + def test_partial_coverage_is_marked_not_passed_through(self): + rows = [{"x": 1.0}, {"x": None}, {"x": 3.0}] + agg = scenarios.aggregate_seed_rows(rows) + self.assertEqual(agg["x"], "2 (2/3 seeds)") + + +class DivergenceClassificationTest(unittest.TestCase): + CAP = 100 + + @staticmethod + def holder(blocks, digest, event_fed=False): + return {"blocks": blocks, "digest": digest, "event_fed": event_fed, "dropped": False} + + def test_identical_replicas_converge(self): + base = {"w1": self.holder(50, "a", True), "w2": self.holder(150, "b")} + d = sim.classify_divergence(base, dict(base), self.CAP) + self.assertTrue(d["converged"]) + self.assertEqual(d["holders_differing"], 0) + + def test_event_fed_difference_is_never_converged(self): + base = {"w1": self.holder(50, "a", True)} + other = {"w1": self.holder(50, "b", True)} + d = sim.classify_divergence(base, other, self.CAP) + self.assertFalse(d["converged"]) + self.assertEqual(d["holders_differing_event_fed"], 1) + + def test_placement_difference_above_capacity_on_both_sides_is_cut_timing(self): + base = {"w1": self.holder(190, "a")} + other = {"w1": self.holder(101, "b")} + d = sim.classify_divergence(base, other, self.CAP) + self.assertTrue(d["converged"]) + self.assertEqual(d["holders_differing_placement_in_band"], 1) + self.assertEqual(d["holders_differing"], 1) + + def test_placement_difference_under_capacity_is_a_lost_update(self): + base = {"w1": self.holder(190, "a")} + other = {"w1": self.holder(99, "b")} + d = sim.classify_divergence(base, other, self.CAP) + self.assertFalse(d["converged"]) + self.assertEqual(d["holders_differing_placement_out_of_band"], 1) + + def test_unknown_capacity_treats_every_placement_difference_as_real(self): + base = {"w1": self.holder(190, "a")} + other = {"w1": self.holder(150, "b")} + d = sim.classify_divergence(base, other, None) + self.assertFalse(d["converged"]) + self.assertEqual(d["holders_differing_placement_out_of_band"], 1) + + def test_holder_missing_on_one_replica_is_not_converged(self): + base = {"w1": self.holder(10, "a"), "w2": self.holder(10, "c")} + other = {"w1": self.holder(10, "a")} + d = sim.classify_divergence(base, other, self.CAP) + self.assertFalse(d["converged"]) + self.assertEqual(d["holders_only_in_one"], 1) + + def test_capacity_comes_from_the_mock_kv_budget(self): + self.assertEqual( + sim.placement_capacity_blocks({"mock": {"kv_tokens": 1_200_000, "block_size": 256}}), + 4688, + ) + self.assertIsNone(sim.placement_capacity_blocks({"mock": {}})) + + +class SelectLegsTest(unittest.TestCase): + LEGS = [("a", {}, None), ("b", {}, None), ("c", {}, None)] + + def test_empty_filter_runs_every_leg_in_order(self): + self.assertEqual(scenarios.select_legs(self.LEGS, []), self.LEGS) + + def test_filter_keeps_scenario_order_not_flag_order(self): + self.assertEqual( + [leg[0] for leg in scenarios.select_legs(self.LEGS, ["c", "a"])], ["a", "c"] + ) + + def test_unknown_label_is_an_error(self): + with self.assertRaises(SystemExit): + scenarios.select_legs(self.LEGS, ["a", "nope"]) + + +if __name__ == "__main__": + unittest.main() + + +class IndexAnalysis(unittest.TestCase): + def test_histogram_percentiles_interpolate_and_floor_at_inf(self): + # 100 samples: 50 in (0, 1 ms], 40 in (1, 2 ms], 10 beyond 50 ms. + buckets = {0.001: 50.0, 0.002: 90.0, 0.05: 90.0, float("inf"): 100.0} + pct = sim.histogram_percentiles(buckets, quantiles=(0.5, 0.9, 0.99)) + self.assertAlmostEqual(pct["p50"], 1.0, places=3) # exactly at the first edge + self.assertAlmostEqual(pct["p90"], 2.0, places=3) + # p99 lands in +Inf: reported as the last finite bound (a floor). + self.assertEqual(pct["p99"], 50.0) + self.assertEqual(sim.histogram_percentiles({}), {}) + + def test_histogram_deltas_subtract_and_sum_labels(self): + first = {'h_bucket{le="0.001",gw="a"}': 5.0, 'h_bucket{le="+Inf",gw="a"}': 7.0} + last = { + 'h_bucket{le="0.001",gw="a"}': 15.0, + 'h_bucket{le="+Inf",gw="a"}': 20.0, + 'h_bucket{le="0.001",gw="b"}': 2.0, + 'h_bucket{le="+Inf",gw="b"}': 3.0, + 'other_bucket{le="0.001"}': 99.0, + } + deltas = sim.histogram_deltas(first, last, "h") + self.assertEqual(deltas, {0.001: 12.0, float("inf"): 16.0}) + + def test_prediction_error_summary_reports_bias_exactness_and_tails(self): + errors = [0] * 90 + [256] * 5 + [-4096] * 5 + s = sim.prediction_error_summary(errors) + self.assertEqual(s["requests"], 100) + self.assertEqual(s["exact_share"], 0.9) + self.assertEqual(s["p50_abs"], 0) + self.assertEqual(s["p90_abs"], 0) + self.assertEqual(s["p95_abs"], 256) + self.assertEqual(s["max_abs"], 4096) + self.assertLess(s["mean"], 0) + self.assertIsNone(sim.prediction_error_summary([])) + + def test_evaluation_scenarios_validate_and_differ_only_by_regime(self): + base = load("local-medium.json") + for name in ("eval-matrix", "io-shapes", "chaos", "soak"): + for label, overrides, _ in scenarios.SCENARIOS[name]: + profile = json.loads(json.dumps(base)) + for key, val in overrides.items(): + if key == "smg_flag_overrides": + profile["smg_flags"] = scenarios.patch_smg_flags(profile["smg_flags"], val) + else: + sim.apply_override(profile, key, val) + sim.validate_profile(profile) # raises SystemExit on a bad leg + self.assertEqual(profile["loadgen"]["image_count"], 0, label) + # Same gateway count and rate across the four regimes of one cell. + cell = [ + leg for leg in scenarios.SCENARIOS["eval-matrix"] if leg[0].endswith("-smg8-rps611") + ] + self.assertEqual(len(cell), 4) + self.assertEqual({leg[1]["smg_count"] for leg in cell}, {8}) + self.assertEqual({leg[1]["loadgen.session_rps"] for leg in cell}, {611}) + + def test_gateway_proxy_rewrites_each_gateways_index_url(self): + profile = load("local-small.json") + profile["index_service"] = {"replicas": 1, "gateway_proxy": True} + profile["smg_flags"] = profile["smg_flags"] + ["--kv-indexer-url", "http://127.0.0.1:40000"] + cmd0 = sim.smg_cmd(profile, "/bin/smg", 0) + cmd1 = sim.smg_cmd(profile, "/bin/smg", 1) + self.assertEqual( + cmd0[cmd0.index("--kv-indexer-url") + 1], f"http://127.0.0.1:{sim.INDEX_GW_PROXY_BASE}" + ) + self.assertEqual( + cmd1[cmd1.index("--kv-indexer-url") + 1], + f"http://127.0.0.1:{sim.INDEX_GW_PROXY_BASE + 1}", + )