diff --git a/README.md b/README.md index 0d70b14..90b1d46 100644 --- a/README.md +++ b/README.md @@ -132,7 +132,13 @@ Each log-normal family accepts either `*-lognorm-median` or `*-lognorm-mu`, requires `*-lognorm-sigma`, and optionally accepts `*-lognorm-max`. Samples are 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. +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`. + +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 +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 @@ -142,18 +148,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..184c5cf 100644 --- a/rust/src/agent.rs +++ b/rust/src/agent.rs @@ -8,6 +8,11 @@ use reqwest::{Client, Url}; use serde_json::{json, Map, Value}; use tokenizers::Tokenizer; use tokio::task::JoinSet; +use uuid::Uuid; + +use crate::tokenizer_loader::load_tokenizer; + +const SMG_ROUTING_KEY: HeaderName = HeaderName::from_static("x-smg-routing-key"); /// A fixed value or a value independently sampled from a log-normal distribution. #[derive(Clone, Debug)] @@ -118,6 +123,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 +183,8 @@ impl AgentLoopConfig { verbose: false, dry_run: false, sglang: false, + ignore_eos: false, + user_tagging: true, seed: None, }) } @@ -215,6 +224,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 +267,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 +299,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 +324,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 +392,7 @@ pub async fn run_agent_benchmark(config: AgentLoopConfig) -> Result Result Result Result Result> { - let tokenizer = Tokenizer::from_pretrained(&config.tokenizer_model, None).map_err(|err| { - anyhow!( - "failed to load tokenizer {}: {}", - config.tokenizer_model, - err - ) - })?; + let tokenizer = load_tokenizer(&config.tokenizer_model)?; let root_seed = match config.seed { Some(seed) => seed, None => rand::thread_rng().gen(), @@ -485,6 +510,11 @@ fn build_agent_plans(config: &AgentLoopConfig) -> 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 +563,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,8 +596,13 @@ async fn run_agent( continue; } - let body = build_request_body(&config, &messages, turn.output_tokens, plan.agent_id); - match request_with_retries(&client, &config, &body).await { + let body = build_request_body( + &config, + &messages, + turn.output_tokens, + plan.user_tag.as_deref(), + ); + match request_with_retries(&client, &config, &body, plan.user_tag.as_deref()).await { Ok(result) => { report.successful_requests += 1; report.input_tokens = report.input_tokens.saturating_add(result.prompt_tokens); @@ -604,54 +640,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 @@ -669,12 +695,13 @@ async fn request_with_retries( client: &Client, config: &AgentLoopConfig, body: &Value, + user_tag: Option<&str>, ) -> Result { let start = Instant::now(); let mut last_error = None; for attempt in 0..=config.max_retries { - match single_attempt(client, config, body).await { + match single_attempt(client, config, body, user_tag).await { Ok(mut result) => { result.latency = start.elapsed(); return Ok(result); @@ -696,15 +723,15 @@ async fn single_attempt( client: &Client, config: &AgentLoopConfig, body: &Value, + user_tag: Option<&str>, ) -> Result { if config.verbose { println!("[AGENT REQUEST] {}", sanitize_request(body)); } - let mut request = client.post(config.endpoint.clone()); - for (name, value) in config.headers.iter() { - request = request.header(name, value); - } + let request = client + .post(config.endpoint.clone()) + .headers(build_request_headers(config, user_tag)?); let response = request.json(body).send().await?; let status = response.status(); let bytes = response.bytes().await?; @@ -744,6 +771,16 @@ async fn single_attempt( }) } +fn build_request_headers(config: &AgentLoopConfig, user_tag: Option<&str>) -> Result { + let mut headers = config.headers.clone(); + if let Some(user_tag) = user_tag { + let value = HeaderValue::from_str(user_tag) + .context("failed to build X-SMG-Routing-Key header from user tag")?; + headers.insert(SMG_ROUTING_KEY, value); + } + Ok(headers) +} + fn append_model_and_environment( messages: &mut Vec, assistant_message: Value, @@ -752,43 +789,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 +845,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 +982,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 +990,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 +1067,105 @@ 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 user_tag_stamps_the_same_per_agent_routing_header() { + 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 user_tag = generate_user_tag(true, Some("user"), 123).unwrap(); + let headers = build_request_headers(&config, Some(&user_tag)).unwrap(); + + assert_eq!(user_tag, "user-123"); + assert_eq!(headers.get(SMG_ROUTING_KEY).unwrap(), user_tag.as_str()); + } + + #[test] + fn routing_header_is_omitted_when_user_tagging_is_disabled() { + 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 headers = build_request_headers(&config, None).unwrap(); + + assert!(!headers.contains_key(SMG_ROUTING_KEY)); + } + + #[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..2c2eb28 100644 --- a/rust/src/agent_cli.rs +++ b/rust/src/agent_cli.rs @@ -143,8 +143,8 @@ struct Args { )] tool_call_latency_lognorm_max_ms: Option, - /// Stamp each request with an OpenAI `user` field of "-", - /// so gateways and routers can key session-sticky routing per agent + /// Stamp each request's OpenAI `user` field and X-SMG-Routing-Key header + /// with "-" for session-sticky routing per agent #[arg(long)] user_prefix: Option, @@ -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 user field and routing header + #[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 {}: {}", diff --git a/rust/src/generator.rs b/rust/src/generator.rs index def46fa..51e833b 100644 --- a/rust/src/generator.rs +++ b/rust/src/generator.rs @@ -2,9 +2,9 @@ use anyhow::{anyhow, Context, Result}; use rand::prelude::*; use rand::SeedableRng; use rand_distr::LogNormal; -use tokenizers::Tokenizer; use crate::config::RequestEntry; +use crate::tokenizer_loader::load_tokenizer; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum DistMode { @@ -45,8 +45,7 @@ pub fn generate_requests(opts: &GenerateOptions, model: &str) -> Result = Vec::with_capacity(opts.count); diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 6842e63..8df7e47 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -7,6 +7,7 @@ mod generator; mod py_bindings; mod report; mod runner; +mod tokenizer_loader; pub use agent::{ run_agent_benchmark, AgentBenchmarkReport, AgentFailureRecord, AgentLoopConfig, SampleSpec, diff --git a/rust/src/tokenizer_loader.rs b/rust/src/tokenizer_loader.rs new file mode 100644 index 0000000..5b36883 --- /dev/null +++ b/rust/src/tokenizer_loader.rs @@ -0,0 +1,73 @@ +use std::path::{Path, PathBuf}; + +use anyhow::{anyhow, Result}; +use tokenizers::Tokenizer; + +#[derive(Debug, PartialEq, Eq)] +enum TokenizerSource { + Local(PathBuf), + Hub(String), +} + +fn resolve_tokenizer_source(model_or_path: &str) -> Result { + let path = Path::new(model_or_path); + if path.is_dir() { + let tokenizer_json = path.join("tokenizer.json"); + if !tokenizer_json.is_file() { + return Err(anyhow!( + "local tokenizer directory {} does not contain tokenizer.json", + path.display() + )); + } + return Ok(TokenizerSource::Local(tokenizer_json)); + } + if path.is_file() { + return Ok(TokenizerSource::Local(path.to_path_buf())); + } + Ok(TokenizerSource::Hub(model_or_path.to_string())) +} + +pub(crate) fn load_tokenizer(model_or_path: &str) -> Result { + match resolve_tokenizer_source(model_or_path)? { + TokenizerSource::Local(path) => Tokenizer::from_file(&path).map_err(|error| { + anyhow!( + "failed to load local tokenizer {}: {}", + path.display(), + error + ) + }), + TokenizerSource::Hub(model) => Tokenizer::from_pretrained(&model, None) + .map_err(|error| anyhow!("failed to load tokenizer {}: {}", model, error)), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + + #[test] + fn local_directory_resolves_to_tokenizer_json() { + let directory = std::env::temp_dir().join(format!( + "batchbench-tokenizer-loader-{}", + std::process::id() + )); + fs::create_dir_all(&directory).unwrap(); + let tokenizer_json = directory.join("tokenizer.json"); + fs::write(&tokenizer_json, "{}").unwrap(); + + assert_eq!( + resolve_tokenizer_source(directory.to_str().unwrap()).unwrap(), + TokenizerSource::Local(tokenizer_json) + ); + fs::remove_dir_all(directory).unwrap(); + } + + #[test] + fn hub_identifier_remains_remote() { + assert_eq!( + resolve_tokenizer_source("zai-org/GLM-5.2-FP8").unwrap(), + TokenizerSource::Hub("zai-org/GLM-5.2-FP8".to_string()) + ); + } +}