diff --git a/.gitignore b/.gitignore index b8e528e..2d90889 100644 --- a/.gitignore +++ b/.gitignore @@ -23,3 +23,5 @@ src/batchbench/_core*.pyd *.png pricing/ examples/* +!examples/trajectory-replay/ +!examples/trajectory-replay/plans.jsonl diff --git a/README.md b/README.md index 5cbd409..f6ae2aa 100644 --- a/README.md +++ b/README.md @@ -135,6 +135,73 @@ reproducible. `--tokenizer-model` can be supplied when the endpoint's model name not also a Hugging Face tokenizer identifier. It accepts a Hugging Face model ID, a local `tokenizer.json` file, or a local directory containing `tokenizer.json`. +### Trajectory replay and rolling admission + +Use `--agent-plans-jsonl` to replay complete empirical token shapes instead of +sampling initial prompts, outputs, environment results, and invocation counts +independently. Each non-empty JSONL line is one trajectory, and file order is the +FIFO admission order: + +```json +{"schema_version":1,"trajectory_id":"agent-001","requests":[{"prompt_tokens":13381,"output_tokens":288},{"prompt_tokens":14800,"output_tokens":512,"delay_after_ms":25},{"prompt_tokens":7000,"output_tokens":128,"reset_before":true}]} +``` + +`prompt_tokens` and `output_tokens` are required and must be positive. +`delay_after_ms` defaults to zero. For a normal transition, BatchBench infers the +synthetic environment growth as: + +```text +next prompt - current prompt - current output +``` + +If a trajectory compacts or resets such that this value would be negative, mark the +next request with `"reset_before": true`. BatchBench then starts that request from a +fresh synthetic prompt while preserving the trajectory's user/routing identity. A +reset deliberately claims no prefix-cache reuse from the preceding request. + +Absolute prompt counts commonly include chat-template and tool-envelope tokens that +are not part of the synthetic message content. After calibrating those values against +the target backend, use `--replay-initial-overhead-tokens` for a first request or +reset, and `--replay-turn-overhead-tokens` for each normal appended turn. BatchBench +subtracts these values when generating content while retaining the manifest's +absolute prompt targets for reporting. + +`--max-active-agents` bounds simultaneous trajectories while retaining every plan +in the manifest. BatchBench initially admits up to that limit and immediately +admits the next queued trajectory whenever any active trajectory terminates. The +replacement inherits the freed scheduler/routing slot, keeping active load balanced +across data-parallel ranks when perfect routing is enabled. + +```bash +batchbench-agent \ + --model Qwen/Qwen3-8B \ + --host http://localhost:8000 \ + --agent-plans-jsonl examples/trajectory-replay/plans.jsonl \ + --max-active-agents 2 \ + --replay-initial-overhead-tokens 0 \ + --replay-turn-overhead-tokens 0 \ + --agent-events-jsonl agent-events.jsonl \ + --seed 42 \ + --dry-run +``` + +Do not combine `--agent-plans-jsonl` with `--agents` or the synthetic workload-shape +flags. `--seed` remains useful because it makes generated placeholder text +reproducible. + +`--agent-events-jsonl` writes each trajectory's queue/admission time, reusable +routing slot, optional DP rank, finish time, runtime, and completion status. The +summary separately reports the final admission time and final drain duration. + +Trajectory replay reproduces the joint request-count and nominal token-shape plan, +not the original text. Prompt targets include workload-specific framing, while +BatchBench generates its own chat/tool framing and exact-length synthetic message +content. Uncalibrated framing can therefore make live prompts drift systematically +from their targets over a long trajectory. Live `usage.prompt_tokens` is +authoritative and should be compared with the plan before a capacity result is +accepted. Likewise, exact output lengths require a backend that honors the requested +minimum/maximum token constraints. + Every request includes a `user` field and an `X-SMG-Routing-Key` header containing the same UUID, which remains stable for the life of that agent and differs between agents. Use `--disable-user-tagging` to omit both, or `--user-prefix ` to @@ -142,9 +209,10 @@ use deterministic `-` values instead of UUIDs. For deterministic data-parallel routing, pass both `--dp-rank-perfect-routing` and `--dp-rank-perfect-routing-num `. -Each request then includes `X-SMG-Target-Worker: `, distributing -consecutive agents evenly across the configured ranks. This routing is independent -of user tagging. +Each request then includes `X-SMG-Target-Worker: `. Initial +agents occupy consecutive slots; under rolling admission, a replacement inherits +the slot freed by the trajectory that just terminated. This preserves the active +per-rank balance and is independent of user tagging. Tool-call latency defaults to zero. Set a fixed delay with `--tool-call-latency-ms`, or sample milliseconds independently for every invocation @@ -178,7 +246,7 @@ For each successful request after an agent's first, the cache estimate adds that same agent's preceding prompt-token count (capped by the current prompt count). Retries and failed requests are excluded because they do not provide reliable usage data. Use `--results-csv ` to persist the same totals, or `--dry-run` to inspect -the independently sampled workload without sending requests. +the sampled or replayed plans without sending requests. ## Rust CLI diff --git a/examples/trajectory-replay/plans.jsonl b/examples/trajectory-replay/plans.jsonl new file mode 100644 index 0000000..abcf405 --- /dev/null +++ b/examples/trajectory-replay/plans.jsonl @@ -0,0 +1,3 @@ +{"schema_version":1,"trajectory_id":"short-agent","requests":[{"prompt_tokens":256,"output_tokens":64},{"prompt_tokens":448,"output_tokens":96}]} +{"schema_version":1,"trajectory_id":"long-agent","requests":[{"prompt_tokens":512,"output_tokens":128},{"prompt_tokens":896,"output_tokens":128,"delay_after_ms":25},{"prompt_tokens":1280,"output_tokens":64}]} +{"schema_version":1,"trajectory_id":"compacting-agent","requests":[{"prompt_tokens":768,"output_tokens":128},{"prompt_tokens":384,"output_tokens":64,"reset_before":true}]} diff --git a/rust/Cargo.lock b/rust/Cargo.lock index c05457b..cce7f13 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -158,6 +158,7 @@ dependencies = [ "reqwest", "serde", "serde_json", + "sha2", "thiserror 1.0.69", "tokenizers", "tokio", @@ -176,6 +177,15 @@ version = "2.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2261d10cca569e4643e526d8dc2e62e433cc8aba21ab764233731f8d369bf394" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "bumpalo" version = "3.19.0" @@ -336,6 +346,15 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -370,6 +389,16 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + [[package]] name = "csv" version = "1.4.0" @@ -466,6 +495,16 @@ dependencies = [ "syn", ] +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + [[package]] name = "dirs" version = "6.0.0" @@ -614,6 +653,16 @@ dependencies = [ "pin-utils", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "getrandom" version = "0.2.16" @@ -1702,6 +1751,17 @@ dependencies = [ "serde", ] +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "shlex" version = "1.3.0" @@ -2027,6 +2087,12 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + [[package]] name = "unicode-ident" version = "1.0.19" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 2801589..93fb799 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -16,6 +16,7 @@ python = ["dep:pyo3"] anyhow = "1" serde = { version = "1", features = ["derive"] } serde_json = "1" +sha2 = "0.10" tokio = { version = "1", features = ["macros", "rt-multi-thread", "time", "fs", "signal"] } reqwest = { version = "0.11", default-features = false, features = ["json", "rustls-tls"] } thiserror = "1" diff --git a/rust/README.md b/rust/README.md index a57e93b..d7df38e 100644 --- a/rust/README.md +++ b/rust/README.md @@ -139,3 +139,11 @@ tokens expected to hit a perfect per-agent prefix cache. Tool latency can likewi be fixed with `--tool-call-latency-ms` or sampled with the `--tool-call-latency-lognorm-*` flags; each agent sleeps for that duration before continuing after a successful tool call. + +For empirical token-shape replay, pass `--agent-plans-jsonl ` instead of the +synthetic shape flags. Each JSONL record contains one trajectory's ordered +`prompt_tokens`/`output_tokens` requests. Add `--max-active-agents ` to queue all +plans while keeping at most `n` trajectories active; a queued plan immediately +replaces whichever active trajectory finishes. See the root README and +`examples/trajectory-replay/plans.jsonl` for the versioned schema and reset +semantics. diff --git a/rust/src/agent.rs b/rust/src/agent.rs index 0e36313..4c88b2b 100644 --- a/rust/src/agent.rs +++ b/rust/src/agent.rs @@ -1,3 +1,7 @@ +use std::collections::{HashSet, VecDeque}; +use std::fs::File; +use std::io::{BufRead, BufReader}; +use std::path::{Path, PathBuf}; use std::time::{Duration, Instant}; use anyhow::{anyhow, Context, Result}; @@ -5,6 +9,7 @@ use rand::{Rng, SeedableRng}; use rand_distr::{Distribution, LogNormal}; use reqwest::header::{HeaderMap, HeaderName, HeaderValue, AUTHORIZATION, CONTENT_TYPE}; use reqwest::{Client, Url}; +use serde::Deserialize; use serde_json::{json, Map, Value}; use tokenizers::Tokenizer; use tokio::task::JoinSet; @@ -128,6 +133,10 @@ pub struct AgentLoopConfig { pub ignore_eos: bool, pub user_tagging: bool, pub seed: Option, + pub agent_plans_jsonl: Option, + pub max_active_agents: Option, + pub replay_initial_overhead_tokens: usize, + pub replay_turn_overhead_tokens: usize, } impl AgentLoopConfig { @@ -189,6 +198,10 @@ impl AgentLoopConfig { ignore_eos: false, user_tagging: true, seed: None, + agent_plans_jsonl: None, + max_active_agents: None, + replay_initial_overhead_tokens: 0, + replay_turn_overhead_tokens: 0, }) } @@ -263,6 +276,29 @@ impl AgentLoopConfig { Ok(self) } + pub fn with_agent_plans_jsonl(mut self, path: impl Into) -> Self { + self.agent_plans_jsonl = Some(path.into()); + self + } + + pub fn with_max_active_agents(mut self, max_active_agents: usize) -> Result { + if max_active_agents == 0 { + return Err(anyhow!("max active agents must be greater than zero")); + } + self.max_active_agents = Some(max_active_agents); + Ok(self) + } + + pub fn with_replay_prompt_overhead( + mut self, + initial_overhead_tokens: usize, + turn_overhead_tokens: usize, + ) -> Self { + self.replay_initial_overhead_tokens = initial_overhead_tokens; + self.replay_turn_overhead_tokens = turn_overhead_tokens; + self + } + pub fn add_header(mut self, name: HeaderName, value: HeaderValue) -> Self { self.headers.insert(name, value); self @@ -271,17 +307,20 @@ impl AgentLoopConfig { #[derive(Clone, Debug)] struct AgentTurnPlan { + input_content_tokens: usize, + target_prompt_tokens: usize, output_tokens: usize, environment_tokens: usize, environment_content: String, tool_call_latency: Duration, + reset_prompt: Option, } #[derive(Clone, Debug)] struct AgentPlan { agent_id: usize, + trajectory_id: String, user_tag: Option, - initial_prompt_tokens: usize, initial_prompt: String, turns: Vec, } @@ -289,13 +328,25 @@ struct AgentPlan { #[derive(Clone, Debug)] pub struct AgentFailureRecord { pub agent_id: usize, + pub trajectory_id: String, pub invocation: usize, pub error: String, } +#[derive(Clone, Debug)] +pub struct AgentLifecycleRecord { + pub agent_id: usize, + pub trajectory_id: String, + pub routing_slot: usize, + pub admitted_at: Duration, + pub finished_at: Duration, + pub completed: bool, +} + #[derive(Clone, Debug)] pub struct AgentBenchmarkReport { pub total_agents: usize, + pub max_active_agents: usize, pub completed_agents: usize, pub planned_tool_invocations: u64, pub total_requests: u64, @@ -306,6 +357,8 @@ pub struct AgentBenchmarkReport { pub estimated_cached_input_tokens: u64, pub total_tool_call_latency: Duration, pub total_duration: Duration, + pub last_agent_admitted_at: Duration, + pub final_drain_duration: Duration, pub input_tokens_per_second: f64, pub output_tokens_per_second: f64, pub requests_per_second: f64, @@ -316,10 +369,15 @@ pub struct AgentBenchmarkReport { pub agent_end_to_end_latency_p90: Option, pub agent_end_to_end_latency_p99: Option, pub failures: Vec, + pub agent_lifecycles: Vec, } #[derive(Debug)] struct AgentWorkerReport { + agent_id: usize, + trajectory_id: String, + routing_slot: usize, + admitted_at: Duration, completed: bool, end_to_end_latency: Duration, successful_requests: u64, @@ -334,8 +392,17 @@ struct AgentWorkerReport { } impl AgentWorkerReport { - fn new() -> Self { + fn new( + agent_id: usize, + trajectory_id: String, + routing_slot: usize, + admitted_at: Duration, + ) -> Self { Self { + agent_id, + trajectory_id, + routing_slot, + admitted_at, completed: false, end_to_end_latency: Duration::ZERO, successful_requests: 0, @@ -354,11 +421,39 @@ impl AgentWorkerReport { #[derive(Clone, Debug)] struct DryRunRecord { agent_id: usize, + trajectory_id: String, + routing_slot: usize, + admitted_at: Duration, invocation: usize, input_content_tokens: usize, + target_prompt_tokens: usize, output_tokens: usize, environment_tokens: usize, tool_call_latency_ms: usize, + reset_before: bool, +} + +pub(crate) const TRAJECTORY_PLAN_SCHEMA_VERSION: u32 = 1; + +#[derive(Clone, Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct TrajectoryPlanSpec { + schema_version: u32, + trajectory_id: String, + requests: Vec, + #[serde(default, rename = "metadata")] + _metadata: Option, +} + +#[derive(Clone, Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct TrajectoryRequestSpec { + prompt_tokens: usize, + output_tokens: usize, + #[serde(default)] + reset_before: bool, + #[serde(default)] + delay_after_ms: u64, } #[derive(Debug)] @@ -369,6 +464,41 @@ struct RequestResult { latency: Duration, } +#[derive(Debug)] +struct RollingAdmission { + pending: VecDeque, + max_active: usize, +} + +impl RollingAdmission { + fn new(plans: Vec, max_active: usize) -> Result { + if max_active == 0 { + return Err(anyhow!("max active agents must be greater than zero")); + } + let effective_max_active = max_active.min(plans.len()); + Ok(Self { + pending: VecDeque::from(plans), + max_active: effective_max_active, + }) + } + + fn initial_admissions(&mut self) -> Vec<(usize, T)> { + (0..self.max_active) + .map(|routing_slot| { + let plan = self + .pending + .pop_front() + .expect("validated admission limit cannot exceed plan count"); + (routing_slot, plan) + }) + .collect() + } + + fn replacement(&mut self, routing_slot: usize) -> Option<(usize, T)> { + self.pending.pop_front().map(|plan| (routing_slot, plan)) + } +} + /// Run independent, stateful agent loops concurrently. /// /// Each agent begins with a separately generated user message. Every model response and synthetic @@ -376,6 +506,8 @@ struct RequestResult { /// previous request as a prefix for server-side KV-cache testing. pub async fn run_agent_benchmark(config: AgentLoopConfig) -> Result { let plans = build_agent_plans(&config)?; + let total_agents = plans.len(); + let requested_max_active_agents = config.max_active_agents.unwrap_or(total_agents); let planned_tool_invocations = plans .iter() .try_fold(0u64, |total, plan| { @@ -391,11 +523,19 @@ pub async fn run_agent_benchmark(config: AgentLoopConfig) -> Result Result Result Result Result>, + plan: AgentPlan, + routing_slot: usize, + admitted_at: Duration, + client: &Client, + config: &std::sync::Arc, +) { + let client = client.clone(); + let config = std::sync::Arc::clone(config); + join_set.spawn(async move { run_agent(plan, routing_slot, admitted_at, client, config).await }); +} + fn build_agent_plans(config: &AgentLoopConfig) -> Result> { let tokenizer = load_tokenizer(&config.tokenizer_model)?; let root_seed = match config.seed { @@ -481,6 +669,11 @@ fn build_agent_plans(config: &AgentLoopConfig) -> Result> { None => rand::thread_rng().gen(), }; + if let Some(path) = config.agent_plans_jsonl.as_deref() { + let specs = load_trajectory_plan_specs(path)?; + return build_replay_agent_plans(config, &tokenizer, root_seed, specs); + } + let mut plans = Vec::with_capacity(config.agent_count); for agent_id in 0..config.agent_count { let mut input_rng = stream_rng(root_seed, agent_id, 0); @@ -498,6 +691,7 @@ fn build_agent_plans(config: &AgentLoopConfig) -> Result> { generate_synthetic_text(&tokenizer, initial_prompt_tokens, &mut text_rng)?; let mut turns = Vec::with_capacity(invocation_count); + let mut target_prompt_tokens = initial_prompt_tokens; for _ in 0..invocation_count { let output_tokens = config .output_tokens @@ -514,21 +708,247 @@ fn build_agent_plans(config: &AgentLoopConfig) -> Result> { let environment_content = generate_synthetic_text(&tokenizer, environment_tokens, &mut text_rng)?; turns.push(AgentTurnPlan { + input_content_tokens: target_prompt_tokens, + target_prompt_tokens, output_tokens, environment_tokens, environment_content, tool_call_latency: Duration::from_millis(tool_call_latency_ms), + reset_prompt: None, + }); + target_prompt_tokens = target_prompt_tokens + .checked_add(output_tokens) + .and_then(|tokens| tokens.checked_add(environment_tokens)) + .ok_or_else(|| anyhow!("sampled agent prompt-token count overflowed usize"))?; + } + + plans.push(AgentPlan { + agent_id, + trajectory_id: format!("sampled-{agent_id}"), + user_tag: generate_user_tag( + config.user_tagging, + config.user_prefix.as_deref(), + agent_id, + ), + initial_prompt, + turns, + }); + } + Ok(plans) +} + +fn load_trajectory_plan_specs(path: &Path) -> Result> { + let file = File::open(path) + .with_context(|| format!("failed to open trajectory plan file {}", path.display()))?; + parse_trajectory_plan_specs(BufReader::new(file), &path.display().to_string()) +} + +fn parse_trajectory_plan_specs( + reader: R, + source: &str, +) -> Result> { + let mut specs = Vec::new(); + let mut trajectory_ids = HashSet::new(); + + for (line_index, line) in reader.lines().enumerate() { + let line_number = line_index + 1; + let line = line.with_context(|| format!("failed to read {source}:{line_number}"))?; + if line.trim().is_empty() { + continue; + } + let spec: TrajectoryPlanSpec = serde_json::from_str(&line) + .with_context(|| format!("invalid trajectory plan at {source}:{line_number}"))?; + validate_trajectory_plan_spec(&spec) + .with_context(|| format!("invalid trajectory plan at {source}:{line_number}"))?; + if !trajectory_ids.insert(spec.trajectory_id.clone()) { + return Err(anyhow!( + "invalid trajectory plan at {source}:{line_number}: duplicate trajectory_id {:?}", + spec.trajectory_id + )); + } + specs.push(spec); + } + + if specs.is_empty() { + return Err(anyhow!( + "trajectory plan file {source} contains no trajectory plans" + )); + } + Ok(specs) +} + +fn validate_trajectory_plan_spec(spec: &TrajectoryPlanSpec) -> Result<()> { + if spec.schema_version != TRAJECTORY_PLAN_SCHEMA_VERSION { + return Err(anyhow!( + "unsupported schema_version {}; expected {}", + spec.schema_version, + TRAJECTORY_PLAN_SCHEMA_VERSION + )); + } + if spec.trajectory_id.trim().is_empty() { + return Err(anyhow!("trajectory_id must not be empty")); + } + if spec.requests.is_empty() { + return Err(anyhow!("requests must contain at least one request")); + } + + for (request_index, request) in spec.requests.iter().enumerate() { + if request.prompt_tokens == 0 { + return Err(anyhow!( + "request {} prompt_tokens must be greater than zero", + request_index + 1 + )); + } + if request.output_tokens == 0 { + return Err(anyhow!( + "request {} output_tokens must be greater than zero", + request_index + 1 + )); + } + if request_index == 0 && request.reset_before { + return Err(anyhow!("the first request cannot set reset_before")); + } + if request_index > 0 { + let previous = &spec.requests[request_index - 1]; + inferred_environment_tokens(previous, request, 0).with_context(|| { + format!( + "request {} prompt_tokens ({}) cannot follow the preceding prompt and output; set reset_before=true to model compaction or reset", + request_index + 1, + request.prompt_tokens + ) + })?; + } + } + Ok(()) +} + +fn inferred_environment_tokens( + current: &TrajectoryRequestSpec, + next: &TrajectoryRequestSpec, + transition_overhead_tokens: usize, +) -> Result { + if next.reset_before { + return Ok(0); + } + let current_context = current + .prompt_tokens + .checked_add(current.output_tokens) + .ok_or_else(|| anyhow!("trajectory prompt-token count overflowed usize"))?; + let current_context = current_context + .checked_add(transition_overhead_tokens) + .ok_or_else(|| anyhow!("trajectory prompt-token count overflowed usize"))?; + next.prompt_tokens + .checked_sub(current_context) + .ok_or_else(|| { + anyhow!( + "next prompt ({}) is smaller than current prompt plus output ({current_context})", + next.prompt_tokens + ) + }) +} + +fn build_replay_agent_plans( + config: &AgentLoopConfig, + tokenizer: &Tokenizer, + root_seed: u64, + specs: Vec, +) -> Result> { + let mut plans = Vec::with_capacity(specs.len()); + for (agent_id, spec) in specs.into_iter().enumerate() { + let mut text_rng = stream_rng(root_seed, agent_id, 4); + let initial_prompt_tokens = spec.requests[0] + .prompt_tokens + .checked_sub(config.replay_initial_overhead_tokens) + .filter(|tokens| *tokens > 0) + .ok_or_else(|| { + anyhow!( + "trajectory {:?} first prompt ({}) must exceed replay initial overhead ({})", + spec.trajectory_id, + spec.requests[0].prompt_tokens, + config.replay_initial_overhead_tokens + ) + })?; + let initial_prompt = + generate_synthetic_text(tokenizer, initial_prompt_tokens, &mut text_rng)?; + let mut turns = Vec::with_capacity(spec.requests.len()); + let mut current_input_content_tokens = initial_prompt_tokens; + + for (request_index, request) in spec.requests.iter().enumerate() { + let reset_prompt = if request.reset_before { + let reset_prompt_tokens = request + .prompt_tokens + .checked_sub(config.replay_initial_overhead_tokens) + .filter(|tokens| *tokens > 0) + .ok_or_else(|| { + anyhow!( + "trajectory {:?} request {} prompt ({}) must exceed replay initial overhead ({})", + spec.trajectory_id, + request_index + 1, + request.prompt_tokens, + config.replay_initial_overhead_tokens + ) + })?; + current_input_content_tokens = reset_prompt_tokens; + Some(generate_synthetic_text( + tokenizer, + reset_prompt_tokens, + &mut text_rng, + )?) + } else { + None + }; + let environment_tokens = match spec.requests.get(request_index + 1) { + Some(next) => inferred_environment_tokens( + request, + next, + config.replay_turn_overhead_tokens, + ) + .with_context(|| { + format!( + "trajectory {:?} request {} cannot satisfy replay turn overhead {}; lower the overhead or set reset_before=true", + spec.trajectory_id, + request_index + 2, + config.replay_turn_overhead_tokens + ) + })?, + None => 0, + }; + let environment_content = + generate_synthetic_text(tokenizer, environment_tokens, &mut text_rng)?; + turns.push(AgentTurnPlan { + input_content_tokens: current_input_content_tokens, + target_prompt_tokens: request.prompt_tokens, + output_tokens: request.output_tokens, + environment_tokens, + environment_content, + tool_call_latency: Duration::from_millis(request.delay_after_ms), + reset_prompt, }); + if spec + .requests + .get(request_index + 1) + .is_some_and(|next| !next.reset_before) + { + current_input_content_tokens = current_input_content_tokens + .checked_add(request.output_tokens) + .and_then(|tokens| tokens.checked_add(environment_tokens)) + .ok_or_else(|| { + anyhow!( + "trajectory {:?} generated content-token count overflowed usize", + spec.trajectory_id + ) + })?; + } } plans.push(AgentPlan { agent_id, + trajectory_id: spec.trajectory_id, user_tag: generate_user_tag( config.user_tagging, config.user_prefix.as_deref(), agent_id, ), - initial_prompt_tokens, initial_prompt, turns, }); @@ -548,36 +968,100 @@ fn generate_synthetic_text( target_tokens: usize, rng: &mut R, ) -> Result { + if target_tokens == 0 { + return Ok(String::new()); + } let vocab_size = tokenizer.get_vocab_size(false); if vocab_size == 0 { return Err(anyhow!("tokenizer vocabulary is empty")); } - for _ in 0..8 { - let ids: Vec = (0..target_tokens) - .map(|_| rng.gen_range(0..vocab_size) as u32) - .collect(); - let decoded = tokenizer - .decode(&ids, true) - .map_err(|err| anyhow!("failed to decode synthetic tokens: {}", err))?; - if !decoded.trim().is_empty() { - return Ok(decoded); + let special_ids: HashSet = tokenizer + .get_added_tokens_decoder() + .into_iter() + .filter_map(|(id, token)| token.special.then_some(id)) + .collect(); + let special_base_ids = special_ids + .iter() + .filter(|id| (**id as usize) < vocab_size) + .count(); + if special_base_ids == vocab_size { + return Err(anyhow!( + "tokenizer vocabulary contains no non-special tokens" + )); + } + + let random_ids = |count: usize, rng: &mut R| -> Vec { + let mut ids = Vec::with_capacity(count); + while ids.len() < count { + let id = rng.gen_range(0..vocab_size) as u32; + if !special_ids.contains(&id) { + ids.push(id); + } + } + ids + }; + let mut text = tokenizer + .decode(&random_ids(target_tokens, rng), false) + .map_err(|err| anyhow!("failed to decode synthetic tokens: {}", err))?; + let mut actual_tokens = 0usize; + + for _ in 0..32 { + let encoding = tokenizer + .encode(text.as_str(), false) + .map_err(|err| anyhow!("failed to re-encode synthetic text: {}", err))?; + actual_tokens = encoding.len(); + let encoded_special = encoding.get_ids().iter().any(|id| special_ids.contains(id)); + if encoded_special { + text = tokenizer + .decode(&random_ids(target_tokens, rng), false) + .map_err(|err| anyhow!("failed to replace synthetic special tokens: {}", err))?; + continue; + } + if actual_tokens == target_tokens && !text.is_empty() { + return Ok(text); + } + + if actual_tokens > target_tokens { + let end_offset = encoding.get_offsets()[target_tokens - 1].1; + if end_offset > 0 && text.is_char_boundary(end_offset) { + text.truncate(end_offset); + } else { + text = tokenizer + .decode(&encoding.get_ids()[..target_tokens], false) + .map_err(|err| anyhow!("failed to trim synthetic text: {}", err))?; + } + } else { + let missing = target_tokens - actual_tokens; + let extra = tokenizer + .decode(&random_ids(missing.max(1), rng), false) + .map_err(|err| anyhow!("failed to extend synthetic text: {}", err))?; + if !text.ends_with(char::is_whitespace) { + text.push(' '); + } + text.push_str(&extra); } } - Ok((0..target_tokens) - .map(|_| rng.gen_range(1u32..=10_000u32).to_string()) - .collect::>() - .join(" ")) + Err(anyhow!( + "failed to generate synthetic text with exactly {target_tokens} tokens after 32 attempts (last count: {actual_tokens})" + )) } async fn run_agent( plan: AgentPlan, + routing_slot: usize, + admitted_at: Duration, client: Client, config: std::sync::Arc, ) -> Result { let agent_start = Instant::now(); - let mut report = AgentWorkerReport::new(); + let mut report = AgentWorkerReport::new( + plan.agent_id, + plan.trajectory_id.clone(), + routing_slot, + admitted_at, + ); let mut messages = vec![json!({ "role": "user", "content": plan.initial_prompt, @@ -586,18 +1070,26 @@ async fn run_agent( for (turn_index, turn) in plan.turns.iter().enumerate() { let invocation = turn_index + 1; + if let Some(reset_prompt) = turn.reset_prompt.as_deref() { + messages = vec![json!({ + "role": "user", + "content": reset_prompt, + })]; + previous_prompt_tokens = 0; + } if config.dry_run { report.dry_run_records.push(DryRunRecord { agent_id: plan.agent_id, + trajectory_id: plan.trajectory_id.clone(), + routing_slot, + admitted_at, invocation, - input_content_tokens: plan.initial_prompt_tokens - + plan.turns[..turn_index] - .iter() - .map(|prior| prior.output_tokens + prior.environment_tokens) - .sum::(), + input_content_tokens: turn.input_content_tokens, + target_prompt_tokens: turn.target_prompt_tokens, output_tokens: turn.output_tokens, environment_tokens: turn.environment_tokens, tool_call_latency_ms: turn.tool_call_latency.as_millis() as usize, + reset_before: turn.reset_prompt.is_some(), }); append_synthetic_turn( &mut messages, @@ -620,7 +1112,7 @@ async fn run_agent( &config, &body, plan.user_tag.as_deref(), - plan.agent_id, + routing_slot, ) .await { @@ -646,18 +1138,28 @@ async fn run_agent( .saturating_add(turn.tool_call_latency); } - append_model_and_environment( + if let Err(err) = append_model_and_environment( &mut messages, result.assistant_message, plan.agent_id, invocation, &turn.environment_content, - )?; + ) { + report.failures.push(AgentFailureRecord { + agent_id: plan.agent_id, + trajectory_id: plan.trajectory_id.clone(), + invocation, + error: format!("failed to append model response: {err:#}"), + }); + report.end_to_end_latency = agent_start.elapsed(); + return Ok(report); + } } Err(err) => { report.failed_requests += 1; report.failures.push(AgentFailureRecord { agent_id: plan.agent_id, + trajectory_id: plan.trajectory_id.clone(), invocation, error: err.to_string(), }); @@ -717,13 +1219,13 @@ async fn request_with_retries( config: &AgentLoopConfig, body: &Value, user_tag: Option<&str>, - agent_id: usize, + routing_slot: usize, ) -> Result { let start = Instant::now(); let mut last_error = None; for attempt in 0..=config.max_retries { - match single_attempt(client, config, body, user_tag, agent_id).await { + match single_attempt(client, config, body, user_tag, routing_slot).await { Ok(mut result) => { result.latency = start.elapsed(); return Ok(result); @@ -746,7 +1248,7 @@ async fn single_attempt( config: &AgentLoopConfig, body: &Value, user_tag: Option<&str>, - agent_id: usize, + routing_slot: usize, ) -> Result { if config.verbose { println!("[AGENT REQUEST] {}", sanitize_request(body)); @@ -754,7 +1256,7 @@ async fn single_attempt( let request = client .post(config.endpoint.clone()) - .headers(build_request_headers(config, user_tag, agent_id)?); + .headers(build_request_headers(config, user_tag, routing_slot)?); let response = request.json(body).send().await?; let status = response.status(); let bytes = response.bytes().await?; @@ -797,7 +1299,7 @@ async fn single_attempt( fn build_request_headers( config: &AgentLoopConfig, user_tag: Option<&str>, - agent_id: usize, + routing_slot: usize, ) -> Result { let mut headers = config.headers.clone(); if let Some(user_tag) = user_tag { @@ -806,7 +1308,7 @@ fn build_request_headers( headers.insert(SMG_ROUTING_KEY, value); } if let Some(num_ranks) = config.dp_rank_perfect_routing_num { - let target_worker = (agent_id % num_ranks).to_string(); + let target_worker = (routing_slot % num_ranks).to_string(); let value = HeaderValue::from_str(&target_worker) .context("failed to build X-SMG-Target-Worker header from agent rank")?; headers.insert(SMG_TARGET_WORKER, value); @@ -973,6 +1475,44 @@ mod tests { } } + #[test] + fn synthetic_text_reencodes_to_the_requested_token_count() { + use tokenizers::models::wordlevel::WordLevel; + use tokenizers::pre_tokenizers::whitespace::Whitespace; + use tokenizers::AddedToken; + + let vocab = [ + ("[UNK]".to_string(), 0), + ("alpha".to_string(), 1), + ("beta".to_string(), 2), + ("gamma".to_string(), 3), + ] + .into_iter() + .collect(); + let model = WordLevel::builder() + .vocab(vocab) + .unk_token("[UNK]".to_string()) + .build() + .unwrap(); + let mut tokenizer = Tokenizer::new(model); + tokenizer.with_pre_tokenizer(Some(Whitespace)); + tokenizer.add_special_tokens(&[AddedToken::from("beta", true)]); + let mut rng = rand::rngs::StdRng::seed_from_u64(17); + + let text = generate_synthetic_text(&tokenizer, 16, &mut rng).unwrap(); + let encoding = tokenizer.encode(text, false).unwrap(); + assert_eq!(encoding.len(), 16); + let special_ids: HashSet = tokenizer + .get_added_tokens_decoder() + .into_iter() + .filter_map(|(id, token)| token.special.then_some(id)) + .collect(); + assert!(encoding + .get_ids() + .iter() + .all(|id| !special_ids.contains(id))); + } + #[test] fn tool_call_latency_is_optional_and_configurable() { let config = AgentLoopConfig::try_new( @@ -1166,7 +1706,7 @@ mod tests { } #[test] - fn perfect_routing_targets_agent_id_modulo_rank_count() { + fn perfect_routing_targets_routing_slot_modulo_rank_count() { let config = AgentLoopConfig::try_new( "http://localhost:8000/v1/chat/completions", None, @@ -1245,4 +1785,180 @@ mod tests { let body = build_request_body(&config, &messages, 4, None); assert_eq!(body["ignore_eos"], true); } + + #[test] + fn trajectory_jsonl_preserves_request_order_and_blank_lines() { + let input = r#" +{"schema_version":1,"trajectory_id":"alpha","requests":[{"prompt_tokens":100,"output_tokens":20},{"prompt_tokens":150,"output_tokens":10,"delay_after_ms":25}]} + +{"schema_version":1,"trajectory_id":"beta","requests":[{"prompt_tokens":80,"output_tokens":5}]} +"#; + let specs = parse_trajectory_plan_specs(std::io::Cursor::new(input), "fixture").unwrap(); + + assert_eq!(specs.len(), 2); + assert_eq!(specs[0].trajectory_id, "alpha"); + assert_eq!(specs[0].requests[0].prompt_tokens, 100); + assert_eq!(specs[0].requests[1].prompt_tokens, 150); + assert_eq!(specs[0].requests[1].delay_after_ms, 25); + assert_eq!(specs[1].trajectory_id, "beta"); + } + + #[test] + fn trajectory_jsonl_accepts_explicit_compaction_reset() { + let input = r#"{"schema_version":1,"trajectory_id":"reset","requests":[{"prompt_tokens":100,"output_tokens":20},{"prompt_tokens":60,"output_tokens":10,"reset_before":true}]}"#; + let specs = parse_trajectory_plan_specs(std::io::Cursor::new(input), "fixture").unwrap(); + assert!(specs[0].requests[1].reset_before); + } + + #[test] + fn trajectory_environment_growth_is_derived_from_adjacent_requests() { + let current = TrajectoryRequestSpec { + prompt_tokens: 100, + output_tokens: 20, + reset_before: false, + delay_after_ms: 0, + }; + let next = TrajectoryRequestSpec { + prompt_tokens: 150, + output_tokens: 10, + reset_before: false, + delay_after_ms: 0, + }; + assert_eq!(inferred_environment_tokens(¤t, &next, 0).unwrap(), 30); + assert_eq!( + inferred_environment_tokens(¤t, &next, 10).unwrap(), + 20 + ); + + let reset = TrajectoryRequestSpec { + prompt_tokens: 60, + output_tokens: 10, + reset_before: true, + delay_after_ms: 0, + }; + assert_eq!( + inferred_environment_tokens(¤t, &reset, 10).unwrap(), + 0 + ); + } + + #[test] + fn replay_compiler_applies_overhead_and_reset_boundaries() { + use tokenizers::models::wordlevel::WordLevel; + use tokenizers::pre_tokenizers::whitespace::Whitespace; + + let vocab = [ + ("[UNK]".to_string(), 0), + ("alpha".to_string(), 1), + ("beta".to_string(), 2), + ("gamma".to_string(), 3), + ] + .into_iter() + .collect(); + let model = WordLevel::builder() + .vocab(vocab) + .unk_token("[UNK]".to_string()) + .build() + .unwrap(); + let mut tokenizer = Tokenizer::new(model); + tokenizer.with_pre_tokenizer(Some(Whitespace)); + + let input = r#"{"schema_version":1,"trajectory_id":"shape","requests":[{"prompt_tokens":100,"output_tokens":20},{"prompt_tokens":150,"output_tokens":10,"delay_after_ms":25},{"prompt_tokens":60,"output_tokens":5,"reset_before":true}]}"#; + let specs = parse_trajectory_plan_specs(std::io::Cursor::new(input), "fixture").unwrap(); + let config = AgentLoopConfig::try_new( + "http://localhost:8000/v1/chat/completions", + None, + "test-model", + 1, + SampleSpec::fixed(8).unwrap(), + SampleSpec::fixed(4).unwrap(), + SampleSpec::fixed(6).unwrap(), + SampleSpec::fixed(2).unwrap(), + ) + .unwrap() + .with_replay_prompt_overhead(5, 10); + + let plans = build_replay_agent_plans(&config, &tokenizer, 42, specs).unwrap(); + let plan = &plans[0]; + assert_eq!(plan.trajectory_id, "shape"); + assert_eq!( + tokenizer + .encode(plan.initial_prompt.as_str(), false) + .unwrap() + .len(), + 95 + ); + assert_eq!(plan.turns.len(), 3); + assert_eq!(plan.turns[0].input_content_tokens, 95); + assert_eq!(plan.turns[0].target_prompt_tokens, 100); + assert_eq!(plan.turns[0].environment_tokens, 20); + assert_eq!(plan.turns[1].input_content_tokens, 135); + assert_eq!(plan.turns[1].target_prompt_tokens, 150); + assert_eq!(plan.turns[1].environment_tokens, 0); + assert_eq!(plan.turns[1].tool_call_latency, Duration::from_millis(25)); + assert!(plan.turns[1].reset_prompt.is_none()); + assert_eq!(plan.turns[2].input_content_tokens, 55); + assert_eq!(plan.turns[2].target_prompt_tokens, 60); + let reset_prompt = plan.turns[2].reset_prompt.as_ref().unwrap(); + assert_eq!( + tokenizer + .encode(reset_prompt.as_str(), false) + .unwrap() + .len(), + 55 + ); + } + + #[test] + fn trajectory_jsonl_rejects_implicit_prompt_reduction_with_line_context() { + let input = r#"{"schema_version":1,"trajectory_id":"bad","requests":[{"prompt_tokens":100,"output_tokens":20},{"prompt_tokens":110,"output_tokens":10}]}"#; + let error = + parse_trajectory_plan_specs(std::io::Cursor::new(input), "fixture").unwrap_err(); + let message = format!("{error:#}"); + assert!(message.contains("fixture:1")); + assert!(message.contains("set reset_before=true")); + } + + #[test] + fn trajectory_jsonl_rejects_duplicate_ids_and_empty_input() { + let duplicated = concat!( + "{\"schema_version\":1,\"trajectory_id\":\"same\",\"requests\":[{\"prompt_tokens\":10,\"output_tokens\":1}]}\n", + "{\"schema_version\":1,\"trajectory_id\":\"same\",\"requests\":[{\"prompt_tokens\":20,\"output_tokens\":1}]}\n" + ); + let error = + parse_trajectory_plan_specs(std::io::Cursor::new(duplicated), "fixture").unwrap_err(); + assert!(error.to_string().contains("duplicate trajectory_id")); + + let error = parse_trajectory_plan_specs(std::io::Cursor::new("\n"), "fixture").unwrap_err(); + assert!(error.to_string().contains("contains no trajectory plans")); + } + + #[test] + fn trajectory_jsonl_rejects_unknown_schema_and_fields() { + let wrong_version = r#"{"schema_version":2,"trajectory_id":"v2","requests":[{"prompt_tokens":10,"output_tokens":1}]}"#; + let error = parse_trajectory_plan_specs(std::io::Cursor::new(wrong_version), "fixture") + .unwrap_err(); + assert!(format!("{error:#}").contains("unsupported schema_version")); + + let unknown = r#"{"schema_version":1,"trajectory_id":"extra","requests":[{"prompt_tokens":10,"output_tokens":1,"typo":true}]}"#; + let error = + parse_trajectory_plan_specs(std::io::Cursor::new(unknown), "fixture").unwrap_err(); + assert!(format!("{error:#}").contains("unknown field")); + } + + #[test] + fn rolling_admission_refills_the_freed_routing_slot_in_fifo_order() { + let mut admission = RollingAdmission::new(vec!["a", "b", "c", "d"], 2).unwrap(); + assert_eq!(admission.initial_admissions(), vec![(0, "a"), (1, "b")]); + assert_eq!(admission.replacement(1), Some((1, "c"))); + assert_eq!(admission.replacement(0), Some((0, "d"))); + assert_eq!(admission.replacement(1), None); + } + + #[test] + fn rolling_admission_validates_its_limit() { + assert!(RollingAdmission::new(vec![1], 0).is_err()); + let mut admission = RollingAdmission::new(vec![1], 2).unwrap(); + assert_eq!(admission.initial_admissions(), vec![(0, 1)]); + } } diff --git a/rust/src/agent_cli.rs b/rust/src/agent_cli.rs index 9b78e67..3f3ef9b 100644 --- a/rust/src/agent_cli.rs +++ b/rust/src/agent_cli.rs @@ -1,12 +1,15 @@ use std::ffi::OsString; use std::fs; +use std::io::{BufReader, BufWriter, Read, Write}; use std::path::{Path, PathBuf}; use std::time::Duration; use anyhow::{anyhow, Context, Result}; use clap::Parser; use serde::Serialize; +use sha2::{Digest, Sha256}; +use crate::agent::TRAJECTORY_PLAN_SCHEMA_VERSION; use crate::{run_agent_benchmark, AgentBenchmarkReport, AgentLoopConfig, SampleSpec}; const DEFAULT_MODEL: &str = "Qwen/Qwen3-VL-235B-A22B-Instruct-FP8"; @@ -17,9 +20,25 @@ const DEFAULT_MODEL: &str = "Qwen/Qwen3-VL-235B-A22B-Instruct-FP8"; about = "Benchmark concurrent stateful agent loops and server-side prefix caching" )] struct Args { - /// Number of independent agent loops to run concurrently - #[arg(long, default_value_t = 1)] - agents: usize, + /// Total number of synthetic agent loops (defaults to 1 without --agent-plans-jsonl) + #[arg(long)] + agents: Option, + + /// JSONL file containing complete empirical trajectory plans + #[arg(long, value_name = "PATH")] + agent_plans_jsonl: Option, + + /// Prompt tokens already accounted for by first-request/reset chat framing + #[arg(long, requires = "agent_plans_jsonl")] + replay_initial_overhead_tokens: Option, + + /// Prompt tokens already accounted for by each appended assistant/tool turn + #[arg(long, requires = "agent_plans_jsonl")] + replay_turn_overhead_tokens: Option, + + /// Maximum simultaneously active agents; completed agents are replaced from the pending queue + #[arg(long, value_parser = parse_positive_usize)] + max_active_agents: Option, /// OpenAI-style model identifier #[arg(long, default_value = DEFAULT_MODEL)] @@ -207,13 +226,23 @@ struct Args { /// Optional path to write a CSV report #[arg(long)] results_csv: Option, + + /// Optional path to write one JSONL lifecycle event per agent + #[arg(long, value_name = "PATH")] + agent_events_jsonl: Option, } #[derive(Debug, Serialize)] struct CsvResult { timestamp: String, model: String, + workload_mode: String, + agent_plans_jsonl: Option, + agent_plans_sha256: Option, + trajectory_schema_version: Option, agents: usize, + requested_max_active_agents: Option, + max_active_agents: usize, completed_agents: usize, planned_tool_invocations: u64, total_requests: u64, @@ -224,6 +253,8 @@ struct CsvResult { estimated_cached_input_tokens: u64, total_tool_call_latency_seconds: f64, total_duration_seconds: f64, + last_agent_admitted_seconds: f64, + final_drain_seconds: f64, requests_per_second: f64, input_tokens_per_second: f64, output_tokens_per_second: f64, @@ -241,6 +272,9 @@ struct CsvResult { dp_rank_perfect_routing: bool, dp_rank_perfect_routing_num: Option, seed: Option, + replay_initial_overhead_tokens: Option, + replay_turn_overhead_tokens: Option, + agent_events_jsonl: Option, tool_call_latency_ms: Option, tool_call_latency_lognorm_mu: Option, tool_call_latency_lognorm_median_ms: Option, @@ -248,6 +282,19 @@ struct CsvResult { tool_call_latency_lognorm_max_ms: Option, } +#[derive(Debug, Serialize)] +struct AgentLifecycleJson<'a> { + agent_id: usize, + trajectory_id: &'a str, + routing_slot: usize, + dp_rank: Option, + admitted_at_seconds: f64, + queue_wait_seconds: f64, + finished_at_seconds: f64, + runtime_seconds: f64, + completed: bool, +} + enum ParsedArgs { Ready(Box), Displayed, @@ -308,7 +355,24 @@ fn run_with_runtime(args: Args) -> Result<()> { } async fn run(args: Args) -> Result<()> { - if args.agents == 0 { + let replay_mode = args.agent_plans_jsonl.is_some(); + if replay_mode && args.agents.is_some() { + return Err(anyhow!( + "--agents cannot be combined with --agent-plans-jsonl; the manifest defines the total agent count" + )); + } + if replay_mode && has_synthetic_shape_args(&args) { + return Err(anyhow!( + "synthetic workload-shape flags cannot be combined with --agent-plans-jsonl" + )); + } + let agent_plans_sha256 = args + .agent_plans_jsonl + .as_deref() + .map(sha256_file) + .transpose()?; + let agent_count = args.agents.unwrap_or(1); + if agent_count == 0 { return Err(anyhow!("agents must be greater than zero")); } if args.request_timeout_secs == 0 { @@ -373,21 +437,43 @@ async fn run(args: Args) -> Result<()> { println!("Endpoint: {}", endpoint); println!("Model: {}", args.model); println!("Tokenizer: {}", tokenizer_model); - println!("Parallel agents: {}", args.agents); - println!("Initial prompt: {}", describe_spec(&input_tokens)); - println!("Model response: {}", describe_spec(&output_tokens)); - println!( - "Environment response: {}", - describe_spec(&environment_tokens) - ); - println!( - "Tool invocations per agent: {}", - describe_spec(&tool_invocations) - ); - println!( - "Tool-call latency: {}", - describe_latency_spec(tool_call_latency_ms.as_ref()) - ); + if let Some(path) = args.agent_plans_jsonl.as_deref() { + println!("Workload: trajectory replay from {}", path.display()); + if let Some(hash) = agent_plans_sha256.as_deref() { + println!("Trajectory manifest SHA-256: {hash}"); + } + println!( + "Replay prompt overhead: initial/reset {} tokens; appended turn {} tokens", + args.replay_initial_overhead_tokens.unwrap_or(0), + args.replay_turn_overhead_tokens.unwrap_or(0) + ); + println!( + "Maximum active agents: {}", + args.max_active_agents + .map(|limit| limit.to_string()) + .unwrap_or_else(|| "all manifest trajectories".to_string()) + ); + } else { + println!("Total synthetic agents: {}", agent_count); + println!( + "Maximum active agents: {}", + args.max_active_agents.unwrap_or(agent_count) + ); + println!("Initial prompt: {}", describe_spec(&input_tokens)); + println!("Model response: {}", describe_spec(&output_tokens)); + println!( + "Environment response: {}", + describe_spec(&environment_tokens) + ); + println!( + "Tool invocations per agent: {}", + describe_spec(&tool_invocations) + ); + println!( + "Tool-call latency: {}", + describe_latency_spec(tool_call_latency_ms.as_ref()) + ); + } println!( "Output token params: {}", if args.sglang { @@ -417,7 +503,7 @@ async fn run(args: Args) -> Result<()> { &endpoint, api_key, args.model.clone(), - args.agents, + agent_count, input_tokens, output_tokens, environment_tokens, @@ -431,6 +517,16 @@ async fn run(args: Args) -> Result<()> { .with_user_tagging(!args.disable_user_tagging) .with_verbose(args.verbose) .with_dry_run(args.dry_run); + if let Some(path) = args.agent_plans_jsonl.as_deref() { + config = config.with_agent_plans_jsonl(path); + config = config.with_replay_prompt_overhead( + args.replay_initial_overhead_tokens.unwrap_or(0), + args.replay_turn_overhead_tokens.unwrap_or(0), + ); + } + if let Some(max_active_agents) = args.max_active_agents { + config = config.with_max_active_agents(max_active_agents)?; + } if let Some(tool_call_latency_ms) = tool_call_latency_ms { config = config.with_tool_call_latency_ms(tool_call_latency_ms)?; } @@ -448,11 +544,29 @@ async fn run(args: Args) -> Result<()> { let report = run_agent_benchmark(config).await?; print_summary(&report); + if let Some(events_path) = args.agent_events_jsonl.as_deref() { + write_agent_events_jsonl(events_path, &report, args.dp_rank_perfect_routing_num)?; + println!("Wrote agent lifecycle events to {}", events_path.display()); + } + if let Some(csv_path) = args.results_csv.as_deref() { let record = CsvResult { timestamp: start_time.to_rfc3339(), model: args.model, - agents: args.agents, + workload_mode: if replay_mode { + "trajectory_replay".to_string() + } else { + "synthetic".to_string() + }, + agent_plans_jsonl: args + .agent_plans_jsonl + .as_ref() + .map(|path| path.display().to_string()), + agent_plans_sha256: agent_plans_sha256.clone(), + trajectory_schema_version: replay_mode.then_some(TRAJECTORY_PLAN_SCHEMA_VERSION), + agents: report.total_agents, + requested_max_active_agents: args.max_active_agents, + max_active_agents: report.max_active_agents, completed_agents: report.completed_agents, planned_tool_invocations: report.planned_tool_invocations, total_requests: report.total_requests, @@ -463,6 +577,8 @@ async fn run(args: Args) -> Result<()> { estimated_cached_input_tokens: report.estimated_cached_input_tokens, total_tool_call_latency_seconds: report.total_tool_call_latency.as_secs_f64(), total_duration_seconds: report.total_duration.as_secs_f64(), + last_agent_admitted_seconds: report.last_agent_admitted_at.as_secs_f64(), + final_drain_seconds: report.final_drain_duration.as_secs_f64(), requests_per_second: report.requests_per_second, input_tokens_per_second: report.input_tokens_per_second, output_tokens_per_second: report.output_tokens_per_second, @@ -480,6 +596,12 @@ async fn run(args: Args) -> Result<()> { dp_rank_perfect_routing: args.dp_rank_perfect_routing, dp_rank_perfect_routing_num: args.dp_rank_perfect_routing_num, seed: args.seed, + replay_initial_overhead_tokens: args.replay_initial_overhead_tokens, + replay_turn_overhead_tokens: args.replay_turn_overhead_tokens, + agent_events_jsonl: args + .agent_events_jsonl + .as_ref() + .map(|path| path.display().to_string()), tool_call_latency_ms: args.tool_call_latency_ms, tool_call_latency_lognorm_mu: args.tool_call_latency_lognorm_mu, tool_call_latency_lognorm_median_ms: args.tool_call_latency_lognorm_median_ms, @@ -493,6 +615,34 @@ async fn run(args: Args) -> Result<()> { Ok(()) } +fn has_synthetic_shape_args(args: &Args) -> bool { + args.input_tokens.is_some() + || args.input_lognorm_mu.is_some() + || args.input_lognorm_median.is_some() + || args.input_lognorm_sigma.is_some() + || args.input_lognorm_max.is_some() + || args.output_tokens.is_some() + || args.output_lognorm_mu.is_some() + || args.output_lognorm_median.is_some() + || args.output_lognorm_sigma.is_some() + || args.output_lognorm_max.is_some() + || args.environment_tokens.is_some() + || args.environment_lognorm_mu.is_some() + || args.environment_lognorm_median.is_some() + || args.environment_lognorm_sigma.is_some() + || args.environment_lognorm_max.is_some() + || args.tool_invocations.is_some() + || args.tool_invocations_lognorm_mu.is_some() + || args.tool_invocations_lognorm_median.is_some() + || args.tool_invocations_lognorm_sigma.is_some() + || args.tool_invocations_lognorm_max.is_some() + || args.tool_call_latency_ms.is_some() + || args.tool_call_latency_lognorm_mu.is_some() + || args.tool_call_latency_lognorm_median_ms.is_some() + || args.tool_call_latency_lognorm_sigma.is_some() + || args.tool_call_latency_lognorm_max_ms.is_some() +} + #[allow(clippy::too_many_arguments)] fn resolve_sample_spec( label: &str, @@ -649,6 +799,12 @@ fn print_summary(report: &AgentBenchmarkReport) { "Agents: {} completed / {} total", report.completed_agents, report.total_agents ); + println!( + "Rolling admission: {} maximum active; last admission at {:.2}s; final drain {:.2}s", + report.max_active_agents, + report.last_agent_admitted_at.as_secs_f64(), + report.final_drain_duration.as_secs_f64() + ); println!( "Tool invocations: {} planned; requests {} (success {}, failure {})", report.planned_tool_invocations, @@ -690,8 +846,8 @@ fn print_summary(report: &AgentBenchmarkReport) { ); for failure in &report.failures { println!( - "Failure: agent {} invocation {}: {}", - failure.agent_id, failure.invocation, failure.error + "Failure: agent {} trajectory {:?} invocation {}: {}", + failure.agent_id, failure.trajectory_id, failure.invocation, failure.error ); } } @@ -706,6 +862,24 @@ fn format_latency(duration: Option) -> String { .unwrap_or_else(|| "n/a".to_string()) } +fn sha256_file(path: &Path) -> Result { + let file = fs::File::open(path) + .with_context(|| format!("failed to open {} for hashing", path.display()))?; + let mut reader = BufReader::new(file); + let mut hasher = Sha256::new(); + let mut buffer = [0u8; 64 * 1024]; + loop { + let bytes_read = reader + .read(&mut buffer) + .with_context(|| format!("failed to hash {}", path.display()))?; + if bytes_read == 0 { + break; + } + hasher.update(&buffer[..bytes_read]); + } + Ok(format!("{:x}", hasher.finalize())) +} + fn write_results_csv(path: &Path, record: &CsvResult) -> Result<()> { if let Some(parent) = path.parent() { if !parent.as_os_str().is_empty() { @@ -724,6 +898,45 @@ fn write_results_csv(path: &Path, record: &CsvResult) -> Result<()> { Ok(()) } +fn write_agent_events_jsonl( + path: &Path, + report: &AgentBenchmarkReport, + dp_rank_count: Option, +) -> Result<()> { + if let Some(parent) = path.parent() { + if !parent.as_os_str().is_empty() { + fs::create_dir_all(parent) + .with_context(|| format!("failed to create directory {}", parent.display()))?; + } + } + let file = + fs::File::create(path).with_context(|| format!("failed to create {}", path.display()))?; + let mut writer = BufWriter::new(file); + for record in &report.agent_lifecycles { + let runtime = record.finished_at.saturating_sub(record.admitted_at); + let event = AgentLifecycleJson { + agent_id: record.agent_id, + trajectory_id: &record.trajectory_id, + routing_slot: record.routing_slot, + dp_rank: dp_rank_count.map(|ranks| record.routing_slot % ranks), + admitted_at_seconds: record.admitted_at.as_secs_f64(), + queue_wait_seconds: record.admitted_at.as_secs_f64(), + finished_at_seconds: record.finished_at.as_secs_f64(), + runtime_seconds: runtime.as_secs_f64(), + completed: record.completed, + }; + serde_json::to_writer(&mut writer, &event) + .with_context(|| format!("failed to serialize {}", path.display()))?; + writer + .write_all(b"\n") + .with_context(|| format!("failed to write {}", path.display()))?; + } + writer + .flush() + .with_context(|| format!("failed to flush {}", path.display()))?; + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -824,4 +1037,85 @@ mod tests { .to_string() .contains("value must be greater than zero")); } + + #[test] + fn trajectory_replay_flags_parse_without_synthetic_agent_count() { + let args = Args::try_parse_from([ + "batchbench-agent", + "--agent-plans-jsonl", + "plans.jsonl", + "--max-active-agents", + "128", + "--replay-initial-overhead-tokens", + "1113", + "--replay-turn-overhead-tokens", + "104", + ]) + .unwrap(); + + assert_eq!(args.agent_plans_jsonl, Some(PathBuf::from("plans.jsonl"))); + assert_eq!(args.max_active_agents, Some(128)); + assert_eq!(args.replay_initial_overhead_tokens, Some(1113)); + assert_eq!(args.replay_turn_overhead_tokens, Some(104)); + assert_eq!(args.agents, None); + assert!(!has_synthetic_shape_args(&args)); + } + + #[test] + fn trajectory_replay_detects_synthetic_shape_flags() { + let args = Args::try_parse_from([ + "batchbench-agent", + "--agent-plans-jsonl", + "plans.jsonl", + "--output-tokens", + "64", + ]) + .unwrap(); + assert!(has_synthetic_shape_args(&args)); + } + + #[tokio::test] + async fn trajectory_replay_rejects_explicit_agents_at_runtime() { + let args = Args::try_parse_from([ + "batchbench-agent", + "--agent-plans-jsonl", + "plans.jsonl", + "--agents", + "2", + ]) + .unwrap(); + let error = run(args).await.unwrap_err(); + assert!(error.to_string().contains("--agents cannot be combined")); + } + + #[tokio::test] + async fn trajectory_replay_rejects_shape_flags_at_runtime() { + let args = Args::try_parse_from([ + "batchbench-agent", + "--agent-plans-jsonl", + "plans.jsonl", + "--output-tokens", + "64", + ]) + .unwrap(); + let error = run(args).await.unwrap_err(); + assert!(error.to_string().contains("synthetic workload-shape flags")); + } + + #[test] + fn max_active_agents_rejects_zero_at_parse_time() { + let error = + Args::try_parse_from(["batchbench-agent", "--max-active-agents", "0"]).unwrap_err(); + assert!(error + .to_string() + .contains("value must be greater than zero")); + } + + #[test] + fn replay_overhead_requires_a_manifest() { + let error = + Args::try_parse_from(["batchbench-agent", "--replay-turn-overhead-tokens", "104"]) + .unwrap_err(); + assert!(error.to_string().contains("--agent-plans-jsonl")); + } } diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 8df7e47..a0c6a75 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -10,7 +10,8 @@ mod runner; mod tokenizer_loader; pub use agent::{ - run_agent_benchmark, AgentBenchmarkReport, AgentFailureRecord, AgentLoopConfig, SampleSpec, + run_agent_benchmark, AgentBenchmarkReport, AgentFailureRecord, AgentLifecycleRecord, + AgentLoopConfig, SampleSpec, }; pub use agent_cli::{run_from_argv as run_agent_from_argv, run_from_env as run_agent_from_env}; pub use cli::{run_from_argv, run_from_env};