diff --git a/README.md b/README.md index 0d70b14..ed6dd50 100644 --- a/README.md +++ b/README.md @@ -134,6 +134,11 @@ independent between agents and turns; `--seed` makes the sampled workload reproducible. `--tokenizer-model` can be supplied when the endpoint's model name is not also a Hugging Face tokenizer identifier. +Every request includes a `user` field containing a UUID that remains stable for the +life of that agent and differs between agents. Use `--disable-user-tagging` to omit +the field, or `--user-prefix ` to use deterministic +`-` values instead of UUIDs. + Tool-call latency defaults to zero. Set a fixed delay with `--tool-call-latency-ms`, or sample milliseconds independently for every invocation with `--tool-call-latency-lognorm-median-ms` (or @@ -142,18 +147,25 @@ optional `--tool-call-latency-lognorm-max-ms`. After a model response succeeds, agent sleeps for the sampled duration before making the environment result available and submitting its next model request. -One tool invocation means one model request followed by one synthetic environment -response. The benchmark asks the model for an `environment` tool call and preserves -the returned assistant message in history. If an endpoint returns plain assistant -content instead, BatchBench wraps it in a valid synthetic tool call before appending -the environment response so the loop can continue. +One tool invocation means one unconstrained model request followed by one synthetic +environment response. BatchBench treats generated assistant output as opaque state: +it preserves content and reasoning output but ignores any model-generated tool calls. +It then adds its own valid synthetic `environment` tool call and appends the sampled +environment response. This keeps the trajectory protocol-valid without assuming +anything about the generated output. The final report includes: - total input tokens sent, from successful responses' `usage.prompt_tokens`; - total output tokens generated, from `usage.completion_tokens`; - estimated cached input tokens under perfect prefix caching; -- total simulated tool-call latency across all agents. +- total simulated tool-call latency across all agents; +- request-latency p50, p90, and p99 across successful requests; +- end-to-end p50, p90, and p99 across completed agents, measured from the start + of each agent loop through its final synthetic tool delay. + +Failed agents are excluded from the end-to-end latency distribution because their +lifetimes end early. Their request failures are still included in the failure report. 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). diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 70df5c0..c05457b 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -161,6 +161,7 @@ dependencies = [ "thiserror 1.0.69", "tokenizers", "tokio", + "uuid", ] [[package]] @@ -632,10 +633,21 @@ checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4" dependencies = [ "cfg-if", "libc", - "r-efi", + "r-efi 5.3.0", "wasi 0.14.7+wasi-0.2.4", ] +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + [[package]] name = "gimli" version = "0.32.3" @@ -1340,6 +1352,12 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + [[package]] name = "rand" version = "0.8.5" @@ -2103,6 +2121,17 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "uuid" +version = "1.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239" +dependencies = [ + "getrandom 0.4.3", + "js-sys", + "wasm-bindgen", +] + [[package]] name = "version_check" version = "0.9.5" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 7cf1274..2801589 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -25,4 +25,5 @@ rand_distr = "0.4" chrono = { version = "0.4", features = ["serde"] } csv = "1" tokenizers = { version = "0.22.2", features = ["http"] } +uuid = { version = "1", features = ["v4"] } pyo3 = { version = "0.22.6", features = ["extension-module", "abi3-py39"], optional = true } diff --git a/rust/src/agent.rs b/rust/src/agent.rs index c4701db..b489b64 100644 --- a/rust/src/agent.rs +++ b/rust/src/agent.rs @@ -8,6 +8,7 @@ use reqwest::{Client, Url}; use serde_json::{json, Map, Value}; use tokenizers::Tokenizer; use tokio::task::JoinSet; +use uuid::Uuid; /// A fixed value or a value independently sampled from a log-normal distribution. #[derive(Clone, Debug)] @@ -118,6 +119,8 @@ pub struct AgentLoopConfig { pub verbose: bool, pub dry_run: bool, pub sglang: bool, + pub ignore_eos: bool, + pub user_tagging: bool, pub seed: Option, } @@ -176,6 +179,8 @@ impl AgentLoopConfig { verbose: false, dry_run: false, sglang: false, + ignore_eos: false, + user_tagging: true, seed: None, }) } @@ -215,6 +220,16 @@ impl AgentLoopConfig { self } + pub fn with_ignore_eos(mut self, ignore_eos: bool) -> Self { + self.ignore_eos = ignore_eos; + self + } + + pub fn with_user_tagging(mut self, user_tagging: bool) -> Self { + self.user_tagging = user_tagging; + self + } + pub fn with_seed(mut self, seed: u64) -> Self { self.seed = Some(seed); self @@ -248,6 +263,7 @@ struct AgentTurnPlan { #[derive(Clone, Debug)] struct AgentPlan { agent_id: usize, + user_tag: Option, initial_prompt_tokens: usize, initial_prompt: String, turns: Vec, @@ -279,12 +295,16 @@ pub struct AgentBenchmarkReport { pub latency_p50: Option, pub latency_p90: Option, pub latency_p99: Option, + pub agent_end_to_end_latency_p50: Option, + pub agent_end_to_end_latency_p90: Option, + pub agent_end_to_end_latency_p99: Option, pub failures: Vec, } #[derive(Debug)] struct AgentWorkerReport { completed: bool, + end_to_end_latency: Duration, successful_requests: u64, failed_requests: u64, input_tokens: u64, @@ -300,6 +320,7 @@ impl AgentWorkerReport { fn new() -> Self { Self { completed: false, + end_to_end_latency: Duration::ZERO, successful_requests: 0, failed_requests: 0, input_tokens: 0, @@ -367,6 +388,7 @@ pub async fn run_agent_benchmark(config: AgentLoopConfig) -> Result Result Result Result Result> { plans.push(AgentPlan { agent_id, + user_tag: generate_user_tag( + config.user_tagging, + config.user_prefix.as_deref(), + agent_id, + ), initial_prompt_tokens, initial_prompt, turns, @@ -533,6 +565,7 @@ async fn run_agent( client: Client, config: std::sync::Arc, ) -> Result { + let agent_start = Instant::now(); let mut report = AgentWorkerReport::new(); let mut messages = vec![json!({ "role": "user", @@ -565,7 +598,12 @@ async fn run_agent( continue; } - let body = build_request_body(&config, &messages, turn.output_tokens, plan.agent_id); + let body = build_request_body( + &config, + &messages, + turn.output_tokens, + plan.user_tag.as_deref(), + ); match request_with_retries(&client, &config, &body).await { Ok(result) => { report.successful_requests += 1; @@ -604,54 +642,44 @@ async fn run_agent( invocation, error: err.to_string(), }); + report.end_to_end_latency = agent_start.elapsed(); return Ok(report); } } } report.completed = true; + report.end_to_end_latency = agent_start.elapsed(); Ok(report) } +fn generate_user_tag(enabled: bool, prefix: Option<&str>, agent_id: usize) -> Option { + enabled.then(|| match prefix { + Some(prefix) => format!("{prefix}-{agent_id}"), + None => Uuid::new_v4().to_string(), + }) +} + fn build_request_body( config: &AgentLoopConfig, messages: &[Value], output_tokens: usize, - agent_id: usize, + user_tag: Option<&str>, ) -> Value { let mut body = json!({ "model": config.model, - "messages": messages, - "tools": [{ - "type": "function", - "function": { - "name": "environment", - "description": "Interact with the benchmark's synthetic environment.", - "parameters": { - "type": "object", - "properties": { - "request": { - "type": "string", - "description": "The environment operation to perform." - } - }, - "required": ["request"] - } - } - }], - "tool_choice": { - "type": "function", - "function": {"name": "environment"} - }, - "parallel_tool_calls": false + "messages": messages }); let (max_key, min_key) = output_token_field_names(config.sglang); if let Some(map) = body.as_object_mut() { map.insert(max_key.to_string(), json!(output_tokens)); map.insert(min_key.to_string(), json!(output_tokens)); - if let Some(prefix) = &config.user_prefix { - map.insert("user".to_string(), json!(format!("{prefix}-{agent_id}"))); + if config.ignore_eos { + map.insert("ignore_eos".to_string(), json!(true)); + } + if let Some(user_tag) = user_tag { + map.insert("user".to_string(), json!(user_tag)); } } body @@ -752,43 +780,26 @@ fn append_model_and_environment( environment_content: &str, ) -> Result<()> { let mut assistant = normalize_assistant_message(assistant_message)?; - let mut tool_call_ids = assistant - .get("tool_calls") - .and_then(Value::as_array) - .map(|calls| { - calls - .iter() - .filter_map(|call| call.get("id").and_then(Value::as_str)) - .map(str::to_string) - .collect::>() - }) - .unwrap_or_default(); - - if tool_call_ids.is_empty() { - let tool_call_id = synthetic_tool_call_id(agent_id, invocation); - let tool_call = json!({ - "id": tool_call_id, - "type": "function", - "function": { - "name": "environment", - "arguments": "{\"request\":\"continue\"}" - } - }); - assistant - .as_object_mut() - .ok_or_else(|| anyhow!("normalized assistant message is not an object"))? - .insert("tool_calls".to_string(), json!([tool_call])); - tool_call_ids.push(synthetic_tool_call_id(agent_id, invocation)); - } + let tool_call_id = synthetic_tool_call_id(agent_id, invocation); + let tool_call = json!({ + "id": tool_call_id, + "type": "function", + "function": { + "name": "environment", + "arguments": "{\"request\":\"continue\"}" + } + }); + assistant + .as_object_mut() + .ok_or_else(|| anyhow!("normalized assistant message is not an object"))? + .insert("tool_calls".to_string(), json!([tool_call])); messages.push(assistant); - for (index, tool_call_id) in tool_call_ids.into_iter().enumerate() { - messages.push(json!({ - "role": "tool", - "tool_call_id": tool_call_id, - "content": if index == 0 { environment_content } else { "" }, - })); - } + messages.push(json!({ + "role": "tool", + "tool_call_id": tool_call_id, + "content": environment_content, + })); Ok(()) } @@ -825,12 +836,14 @@ fn normalize_assistant_message(message: Value) -> Result { .ok_or_else(|| anyhow!("choices[0].message is not an object"))?; let mut normalized = Map::new(); normalized.insert("role".to_string(), Value::String("assistant".to_string())); - for key in ["content", "name", "tool_calls", "function_call", "refusal"] { + // Treat generated output as opaque assistant state. Provider-generated tool calls are + // deliberately discarded; BatchBench adds its own valid synthetic tool-call envelope. + for key in ["content", "reasoning_content", "name", "refusal"] { if let Some(value) = source.get(key) { normalized.insert(key.to_string(), value.clone()); } } - if !normalized.contains_key("content") && !normalized.contains_key("tool_calls") { + if !normalized.contains_key("content") { normalized.insert("content".to_string(), Value::Null); } Ok(Value::Object(normalized)) @@ -960,7 +973,7 @@ mod tests { } #[test] - fn appended_turn_preserves_model_message_and_adds_tool_result() { + fn appended_turn_preserves_output_but_replaces_model_tool_calls() { let mut messages = vec![json!({"role": "user", "content": "start"})]; let assistant = json!({ "role": "assistant", @@ -968,17 +981,25 @@ mod tests { "tool_calls": [{ "id": "call_123", "type": "function", - "function": {"name": "environment", "arguments": "{\"request\":\"x\"}"} + "function": {"name": "unexpected", "arguments": "{\"unterminated\":"} }], "reasoning_content": "server-specific field" }); append_model_and_environment(&mut messages, assistant, 0, 1, "result").unwrap(); assert_eq!(messages.len(), 3); - assert_eq!(messages[1]["tool_calls"][0]["id"], "call_123"); - assert!(messages[1].get("reasoning_content").is_none()); + assert_eq!(messages[1]["tool_calls"][0]["id"], "call_batchbench_0_1"); + assert_eq!( + messages[1]["tool_calls"][0]["function"]["arguments"], + "{\"request\":\"continue\"}" + ); + assert_eq!( + messages[1]["tool_calls"][0]["function"]["name"], + "environment" + ); + assert_eq!(messages[1]["reasoning_content"], "server-specific field"); assert_eq!(messages[2]["role"], "tool"); - assert_eq!(messages[2]["tool_call_id"], "call_123"); + assert_eq!(messages[2]["tool_call_id"], "call_batchbench_0_1"); assert_eq!(messages[2]["content"], "result"); } @@ -1037,34 +1058,67 @@ mod tests { json!({"role": "tool", "tool_call_id": "call_1", "content": "result"}), ]; - let body = build_request_body(&config, &messages, 17, 0); + let body = build_request_body(&config, &messages, 17, Some("agent-user-tag")); assert_eq!(body["messages"], json!(messages)); assert_eq!(body["max_tokens"], 17); assert_eq!(body["min_tokens"], 17); + assert_eq!(body["user"], "agent-user-tag"); + assert!(body.get("tools").is_none()); + assert!(body.get("tool_choice").is_none()); + } + + #[test] + fn user_tags_are_uuid_v4_values_and_can_be_disabled() { + let first = generate_user_tag(true, None, 0).unwrap(); + let second = generate_user_tag(true, None, 1).unwrap(); + assert_ne!(first, second); + assert_eq!(Uuid::parse_str(&first).unwrap().get_version_num(), 4); + assert_eq!(generate_user_tag(false, None, 0), None); + } + + #[test] + fn user_prefix_stamps_a_deterministic_per_agent_user_field() { assert_eq!( - body["tool_choice"]["function"]["name"], - Value::String("environment".to_string()) + generate_user_tag(true, Some("loadtest"), 1).as_deref(), + Some("loadtest-1") ); + assert_eq!(generate_user_tag(false, Some("loadtest"), 1), None); + } + + #[test] + fn request_body_omits_disabled_user_tag() { + 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(); + let messages = vec![json!({"role": "user", "content": "start"})]; + let body = build_request_body(&config, &messages, 4, None); assert!(body.get("user").is_none()); } #[test] - fn user_prefix_stamps_a_per_agent_user_field() { + fn request_body_can_force_generation_past_eos() { let config = AgentLoopConfig::try_new( "http://localhost:8000/v1/chat/completions", None, "test-model", - 2, + 1, SampleSpec::fixed(8).unwrap(), SampleSpec::fixed(4).unwrap(), SampleSpec::fixed(6).unwrap(), SampleSpec::fixed(2).unwrap(), ) .unwrap() - .with_user_prefix("loadtest"); + .with_ignore_eos(true); let messages = vec![json!({"role": "user", "content": "start"})]; - - let body = build_request_body(&config, &messages, 5, 1); - assert_eq!(body["user"], Value::String("loadtest-1".to_string())); + let body = build_request_body(&config, &messages, 4, None); + assert_eq!(body["ignore_eos"], true); } } diff --git a/rust/src/agent_cli.rs b/rust/src/agent_cli.rs index 21a47e8..2bd49b1 100644 --- a/rust/src/agent_cli.rs +++ b/rust/src/agent_cli.rs @@ -172,6 +172,14 @@ struct Args { #[arg(long)] sglang: bool, + /// Ignore EOS so generation continues to the configured output-token limit + #[arg(long)] + ignore_eos: bool, + + /// Omit the persistent per-agent value from the request's user field + #[arg(long)] + disable_user_tagging: bool, + /// Print truncated request and response payloads #[arg(long, short)] verbose: bool, @@ -210,9 +218,14 @@ struct CsvResult { latency_p50_ms: Option, latency_p90_ms: Option, latency_p99_ms: Option, + agent_end_to_end_latency_p50_ms: Option, + agent_end_to_end_latency_p90_ms: Option, + agent_end_to_end_latency_p99_ms: Option, host: String, endpoint: String, sglang: bool, + ignore_eos: bool, + user_tagging: bool, seed: Option, tool_call_latency_ms: Option, tool_call_latency_lognorm_mu: Option, @@ -359,6 +372,15 @@ async fn run(args: Args) -> Result<()> { "default (min_tokens/max_tokens)" } ); + println!("Ignore EOS: {}", args.ignore_eos); + println!( + "Per-agent user tagging: {}", + if args.disable_user_tagging { + "disabled" + } else { + "enabled" + } + ); println!("==========================================\n"); let mut config = AgentLoopConfig::try_new( @@ -375,6 +397,8 @@ async fn run(args: Args) -> Result<()> { .with_request_timeout(Duration::from_secs(args.request_timeout_secs)) .with_retry(args.max_retries, Duration::from_millis(args.retry_delay_ms)) .with_sglang(args.sglang) + .with_ignore_eos(args.ignore_eos) + .with_user_tagging(!args.disable_user_tagging) .with_verbose(args.verbose) .with_dry_run(args.dry_run); if let Some(tool_call_latency_ms) = tool_call_latency_ms { @@ -412,9 +436,14 @@ async fn run(args: Args) -> Result<()> { latency_p50_ms: milliseconds(report.latency_p50), latency_p90_ms: milliseconds(report.latency_p90), latency_p99_ms: milliseconds(report.latency_p99), + agent_end_to_end_latency_p50_ms: milliseconds(report.agent_end_to_end_latency_p50), + agent_end_to_end_latency_p90_ms: milliseconds(report.agent_end_to_end_latency_p90), + agent_end_to_end_latency_p99_ms: milliseconds(report.agent_end_to_end_latency_p99), host: args.host, endpoint, sglang: args.sglang, + ignore_eos: args.ignore_eos, + user_tagging: !args.disable_user_tagging, seed: args.seed, tool_call_latency_ms: args.tool_call_latency_ms, tool_call_latency_lognorm_mu: args.tool_call_latency_lognorm_mu, @@ -613,11 +642,17 @@ fn print_summary(report: &AgentBenchmarkReport) { report.requests_per_second ); println!( - "Latency (ms): p50={} p90={} p99={}", + "Request latency (ms): p50={} p90={} p99={}", format_latency(report.latency_p50), format_latency(report.latency_p90), format_latency(report.latency_p99) ); + println!( + "Agent end-to-end latency (ms, completed agents): p50={} p90={} p99={}", + format_latency(report.agent_end_to_end_latency_p50), + format_latency(report.agent_end_to_end_latency_p90), + format_latency(report.agent_end_to_end_latency_p99) + ); for failure in &report.failures { println!( "Failure: agent {} invocation {}: {}",