| @@ -1,3 +1,7 @@ | |||
| use anyhow::{Context, Result, bail}; | |||
| use serde::Deserialize; | |||
| use std::path::{Path, PathBuf}; | |||
| #[derive(Debug, Clone, PartialEq, Eq)] | |||
| pub struct AimeCase { | |||
| pub id: String, | |||
| @@ -5,6 +9,68 @@ pub struct AimeCase { | |||
| pub answer: String, | |||
| } | |||
| #[derive(Debug, Clone, PartialEq, Eq)] | |||
| pub struct LoadedAimeCases { | |||
| pub cases: Vec<AimeCase>, | |||
| pub local_path: PathBuf, | |||
| } | |||
| #[derive(Debug, Deserialize)] | |||
| struct NormalizedAimeCase { | |||
| id: String, | |||
| problem: String, | |||
| answer: String, | |||
| } | |||
| pub fn load_cases(data_dir: &Path) -> Result<LoadedAimeCases> { | |||
| let dataset_dir = data_dir.join("aime2026"); | |||
| if !dataset_dir.exists() { | |||
| bail!("missing local dataset; run lq_token_test dataset fetch aime2026"); | |||
| } | |||
| let jsonl_path = dataset_dir.join("aime2026.jsonl"); | |||
| if jsonl_path.exists() { | |||
| let contents = std::fs::read_to_string(&jsonl_path) | |||
| .with_context(|| format!("failed to read {}", jsonl_path.display()))?; | |||
| let cases = contents | |||
| .lines() | |||
| .enumerate() | |||
| .filter(|(_, line)| !line.trim().is_empty()) | |||
| .map(|(index, line)| { | |||
| let parsed: NormalizedAimeCase = serde_json::from_str(line).with_context(|| { | |||
| format!( | |||
| "failed to parse normalized AIME JSONL line {} in {}", | |||
| index + 1, | |||
| jsonl_path.display() | |||
| ) | |||
| })?; | |||
| Ok(AimeCase { | |||
| id: parsed.id, | |||
| problem: parsed.problem, | |||
| answer: parsed.answer, | |||
| }) | |||
| }) | |||
| .collect::<Result<Vec<_>>>()?; | |||
| return Ok(LoadedAimeCases { | |||
| cases, | |||
| local_path: jsonl_path, | |||
| }); | |||
| } | |||
| let has_parquet = std::fs::read_dir(&dataset_dir) | |||
| .with_context(|| format!("failed to read {}", dataset_dir.display()))? | |||
| .filter_map(Result::ok) | |||
| .any(|entry| entry.path().extension().is_some_and(|ext| ext == "parquet")); | |||
| if has_parquet { | |||
| bail!( | |||
| "AIME 2026 parquet conversion is not implemented yet; provide normalized aime2026.jsonl with fields id, problem, answer in {}", | |||
| dataset_dir.display() | |||
| ); | |||
| } | |||
| bail!("missing local dataset; run lq_token_test dataset fetch aime2026"); | |||
| } | |||
| impl AimeCase { | |||
| pub fn prompt(&self) -> String { | |||
| format!( | |||
| @@ -31,4 +97,41 @@ mod tests { | |||
| assert!(prompt.contains("What is 20 + 22?")); | |||
| assert!(prompt.contains("final integer answer")); | |||
| } | |||
| #[test] | |||
| fn loads_normalized_jsonl_cases() { | |||
| let temp_dir = tempfile::tempdir().expect("create temp dir"); | |||
| let dataset_dir = temp_dir.path().join("aime2026"); | |||
| std::fs::create_dir_all(&dataset_dir).expect("create dataset dir"); | |||
| std::fs::write( | |||
| dataset_dir.join("aime2026.jsonl"), | |||
| r#"{"id":"aime-1","problem":"What is 2+2?","answer":"4"} | |||
| {"id":"aime-2","problem":"What is 3+5?","answer":"8"} | |||
| "#, | |||
| ) | |||
| .expect("write jsonl"); | |||
| let loaded = load_cases(temp_dir.path()).expect("load cases"); | |||
| assert_eq!(loaded.local_path, dataset_dir.join("aime2026.jsonl")); | |||
| assert_eq!(loaded.cases.len(), 2); | |||
| assert_eq!(loaded.cases[0].answer, "4"); | |||
| } | |||
| #[test] | |||
| fn parquet_only_returns_conversion_hint() { | |||
| let temp_dir = tempfile::tempdir().expect("create temp dir"); | |||
| let dataset_dir = temp_dir.path().join("aime2026"); | |||
| std::fs::create_dir_all(&dataset_dir).expect("create dataset dir"); | |||
| std::fs::write(dataset_dir.join("train.parquet"), b"parquet").expect("write parquet"); | |||
| let error = load_cases(temp_dir.path()).expect_err("parquet should not load"); | |||
| assert!( | |||
| error | |||
| .to_string() | |||
| .contains("parquet conversion is not implemented yet") | |||
| ); | |||
| assert!(error.to_string().contains("aime2026.jsonl")); | |||
| } | |||
| } | |||
| @@ -1,3 +1,7 @@ | |||
| use anyhow::{Context, Result, bail}; | |||
| use std::collections::HashMap; | |||
| use std::path::{Path, PathBuf}; | |||
| #[derive(Debug, Clone, PartialEq, Eq)] | |||
| pub struct GpqaCase { | |||
| pub id: String, | |||
| @@ -6,6 +10,132 @@ pub struct GpqaCase { | |||
| pub answer: char, | |||
| } | |||
| #[derive(Debug, Clone, PartialEq, Eq)] | |||
| pub struct LoadedGpqaCases { | |||
| pub cases: Vec<GpqaCase>, | |||
| pub local_path: PathBuf, | |||
| } | |||
| pub fn load_cases(data_dir: &Path) -> Result<LoadedGpqaCases> { | |||
| let dataset_dir = data_dir.join("gpqa_diamond"); | |||
| let csv_path = dataset_dir.join("gpqa_diamond.csv"); | |||
| if !csv_path.exists() { | |||
| bail!("missing local dataset; run lq_token_test dataset fetch gpqa-diamond"); | |||
| } | |||
| let contents = std::fs::read_to_string(&csv_path) | |||
| .with_context(|| format!("failed to read {}", csv_path.display()))?; | |||
| let rows = parse_csv(&contents).context("failed to parse GPQA CSV")?; | |||
| let (headers, records) = rows.split_first().context("GPQA CSV is empty")?; | |||
| let header_index = headers | |||
| .iter() | |||
| .enumerate() | |||
| .map(|(index, header)| (normalize_header(header), index)) | |||
| .collect::<HashMap<_, _>>(); | |||
| let cases = records | |||
| .iter() | |||
| .enumerate() | |||
| .map(|(index, record)| { | |||
| let question = csv_value(record, &header_index, &["question"]) | |||
| .with_context(|| format!("missing Question column in row {}", index + 2))?; | |||
| let correct = csv_value(record, &header_index, &["correctanswer", "correct"]) | |||
| .with_context(|| format!("missing Correct Answer column in row {}", index + 2))?; | |||
| let incorrect_1 = csv_value(record, &header_index, &["incorrectanswer1", "incorrect1"]) | |||
| .with_context(|| { | |||
| format!("missing Incorrect Answer 1 column in row {}", index + 2) | |||
| })?; | |||
| let incorrect_2 = csv_value(record, &header_index, &["incorrectanswer2", "incorrect2"]) | |||
| .with_context(|| { | |||
| format!("missing Incorrect Answer 2 column in row {}", index + 2) | |||
| })?; | |||
| let incorrect_3 = csv_value(record, &header_index, &["incorrectanswer3", "incorrect3"]) | |||
| .with_context(|| { | |||
| format!("missing Incorrect Answer 3 column in row {}", index + 2) | |||
| })?; | |||
| Ok(GpqaCase { | |||
| id: format!("gpqa-diamond-{}", index + 1), | |||
| question: question.to_string(), | |||
| choices: [ | |||
| ("A".to_string(), correct.to_string()), | |||
| ("B".to_string(), incorrect_1.to_string()), | |||
| ("C".to_string(), incorrect_2.to_string()), | |||
| ("D".to_string(), incorrect_3.to_string()), | |||
| ], | |||
| answer: 'A', | |||
| }) | |||
| }) | |||
| .collect::<Result<Vec<_>>>()?; | |||
| Ok(LoadedGpqaCases { | |||
| cases, | |||
| local_path: csv_path, | |||
| }) | |||
| } | |||
| fn csv_value<'a>( | |||
| record: &'a [String], | |||
| header_index: &HashMap<String, usize>, | |||
| names: &[&str], | |||
| ) -> Option<&'a str> { | |||
| names | |||
| .iter() | |||
| .find_map(|name| header_index.get(*name)) | |||
| .and_then(|index| record.get(*index)) | |||
| .map(String::as_str) | |||
| .filter(|value| !value.trim().is_empty()) | |||
| } | |||
| fn normalize_header(header: &str) -> String { | |||
| header | |||
| .chars() | |||
| .filter(|ch| ch.is_ascii_alphanumeric()) | |||
| .collect::<String>() | |||
| .to_ascii_lowercase() | |||
| } | |||
| fn parse_csv(input: &str) -> Result<Vec<Vec<String>>> { | |||
| let mut rows = Vec::new(); | |||
| let mut row = Vec::new(); | |||
| let mut field = String::new(); | |||
| let mut chars = input.chars().peekable(); | |||
| let mut in_quotes = false; | |||
| while let Some(ch) = chars.next() { | |||
| match ch { | |||
| '"' if in_quotes && chars.peek() == Some(&'"') => { | |||
| field.push('"'); | |||
| chars.next(); | |||
| } | |||
| '"' => in_quotes = !in_quotes, | |||
| ',' if !in_quotes => { | |||
| row.push(std::mem::take(&mut field)); | |||
| } | |||
| '\n' if !in_quotes => { | |||
| row.push(std::mem::take(&mut field)); | |||
| if !row.iter().all(|value| value.is_empty()) { | |||
| rows.push(std::mem::take(&mut row)); | |||
| } else { | |||
| row.clear(); | |||
| } | |||
| } | |||
| '\r' if !in_quotes => {} | |||
| _ => field.push(ch), | |||
| } | |||
| } | |||
| if in_quotes { | |||
| bail!("unterminated quoted field"); | |||
| } | |||
| if !field.is_empty() || !row.is_empty() { | |||
| row.push(field); | |||
| rows.push(row); | |||
| } | |||
| Ok(rows) | |||
| } | |||
| impl GpqaCase { | |||
| pub fn prompt(&self) -> String { | |||
| let choices = self | |||
| @@ -49,4 +179,27 @@ mod tests { | |||
| assert!(prompt.contains("D. Delta")); | |||
| assert!(prompt.contains("answer with exactly one letter")); | |||
| } | |||
| #[test] | |||
| fn loads_gpqa_csv_with_common_columns() { | |||
| let temp_dir = tempfile::tempdir().expect("create temp dir"); | |||
| let dataset_dir = temp_dir.path().join("gpqa_diamond"); | |||
| std::fs::create_dir_all(&dataset_dir).expect("create dataset dir"); | |||
| std::fs::write( | |||
| dataset_dir.join("gpqa_diamond.csv"), | |||
| "Question,Correct Answer,Incorrect Answer 1,Incorrect Answer 2,Incorrect Answer 3\n\"Which is correct?\",\"Right\",\"Wrong, with comma\",\"Nope\",\"Never\"\n", | |||
| ) | |||
| .expect("write csv"); | |||
| let loaded = load_cases(temp_dir.path()).expect("load cases"); | |||
| assert_eq!(loaded.local_path, dataset_dir.join("gpqa_diamond.csv")); | |||
| assert_eq!(loaded.cases.len(), 1); | |||
| assert_eq!(loaded.cases[0].question, "Which is correct?"); | |||
| assert_eq!( | |||
| loaded.cases[0].choices[0], | |||
| ("A".to_string(), "Right".to_string()) | |||
| ); | |||
| assert_eq!(loaded.cases[0].answer, 'A'); | |||
| } | |||
| } | |||
| @@ -1,9 +1,20 @@ | |||
| use crate::benchmarks; | |||
| use crate::benchmarks::judge; | |||
| use crate::config::AppConfig; | |||
| use crate::metrics::{LatencySummary, Metrics, MetricsSummary}; | |||
| use crate::report::{ | |||
| BenchmarkReport, BenchmarkSummaryReport, DatasetReport, LatencyReport, RpmReport, RpmRunReport, | |||
| RpmSummaryReport, RunReport, WrongCaseReport, write_benchmark_report, write_rpm_report, | |||
| }; | |||
| use crate::runner::{ModelRequest, run_model_request}; | |||
| use anyhow::Result; | |||
| use anyhow::{Context, Result, bail}; | |||
| use chrono::Utc; | |||
| use clap::{Parser, Subcommand}; | |||
| use futures::{StreamExt, stream}; | |||
| use regex::Regex; | |||
| use std::path::{Path, PathBuf}; | |||
| use std::time::{Duration, Instant}; | |||
| use tokio::time::sleep; | |||
| #[derive(Debug, Parser)] | |||
| #[command( | |||
| @@ -119,9 +130,453 @@ pub async fn dispatch(cli: Cli) -> Result<()> { | |||
| println!("{}", path.display()); | |||
| Ok(()) | |||
| } | |||
| Command::Bench { .. } => anyhow::bail!("bench is not implemented yet"), | |||
| Command::Rpm { .. } => anyhow::bail!("rpm is not implemented yet"), | |||
| Command::Bench { command } => dispatch_bench(command).await, | |||
| Command::Rpm { | |||
| config, | |||
| provider, | |||
| model, | |||
| rpm, | |||
| duration, | |||
| prompt, | |||
| } => run_rpm(config, provider, model, rpm, duration, prompt).await, | |||
| } | |||
| } | |||
| async fn dispatch_bench(command: BenchCommand) -> Result<()> { | |||
| match command { | |||
| BenchCommand::Aime2026 { | |||
| config, | |||
| provider, | |||
| model, | |||
| concurrency, | |||
| limit, | |||
| } => run_aime_benchmark(config, provider, model, concurrency, limit).await, | |||
| BenchCommand::GpqaDiamond { | |||
| config, | |||
| provider, | |||
| model, | |||
| concurrency, | |||
| limit, | |||
| } => run_gpqa_benchmark(config, provider, model, concurrency, limit).await, | |||
| } | |||
| } | |||
| async fn run_aime_benchmark( | |||
| config_path: PathBuf, | |||
| provider: Option<String>, | |||
| model: Option<String>, | |||
| concurrency: usize, | |||
| limit: Option<usize>, | |||
| ) -> Result<()> { | |||
| let config = AppConfig::load(&config_path)?; | |||
| let provider_name = provider_name(&config, provider.as_deref())?; | |||
| let provider_config = config.provider(Some(&provider_name))?; | |||
| let model = model.unwrap_or_else(|| provider_config.default_model.clone()); | |||
| let loaded = benchmarks::aime::load_cases(Path::new(&config.benchmarks.data_dir))?; | |||
| let dataset = dataset_report( | |||
| config | |||
| .benchmarks | |||
| .aime2026 | |||
| .as_ref() | |||
| .map(|dataset| (dataset.source.as_str(), dataset.split.as_str())), | |||
| &loaded.local_path, | |||
| ); | |||
| let cases = apply_limit(loaded.cases, limit); | |||
| let started_at = Utc::now(); | |||
| let started = Instant::now(); | |||
| let base_request = request_template(provider_config, &model, 0.0, 1024); | |||
| let protocol = provider_config.protocol; | |||
| let results = stream::iter(cases.into_iter()) | |||
| .map(|case| { | |||
| let mut request = base_request.clone(); | |||
| request.prompt = case.prompt(); | |||
| async move { | |||
| let result = run_model_request(protocol, request).await; | |||
| (case, result) | |||
| } | |||
| }) | |||
| .buffer_unordered(nonzero_concurrency(concurrency)) | |||
| .collect::<Vec<_>>() | |||
| .await; | |||
| let mut metrics = Metrics::new(); | |||
| let mut wrong_cases = Vec::new(); | |||
| for (case, result) in results { | |||
| match result { | |||
| Ok(response) => { | |||
| metrics.record_success(response.status, response.elapsed_ms as u64); | |||
| let actual = judge::extract_final_integer(&response.text) | |||
| .unwrap_or_else(|| "no_answer".to_string()); | |||
| let correct = actual == case.answer.trim(); | |||
| metrics.record_judgement(correct); | |||
| if !correct { | |||
| wrong_cases.push(WrongCaseReport { | |||
| id: case.id, | |||
| question: case.problem, | |||
| expected: case.answer, | |||
| actual, | |||
| raw_output: response.text, | |||
| }); | |||
| } | |||
| } | |||
| Err(error) => metrics.record_failure(error_code(&error)), | |||
| } | |||
| } | |||
| let summary = metrics.summary(); | |||
| let report = benchmark_report( | |||
| "aime2026", | |||
| provider_name, | |||
| model, | |||
| dataset, | |||
| started_at, | |||
| started.elapsed().as_millis(), | |||
| concurrency, | |||
| limit, | |||
| summary, | |||
| wrong_cases, | |||
| ); | |||
| let report_path = write_benchmark_report(Path::new("."), &report)?; | |||
| print_benchmark_report(&report, &report_path); | |||
| Ok(()) | |||
| } | |||
| async fn run_gpqa_benchmark( | |||
| config_path: PathBuf, | |||
| provider: Option<String>, | |||
| model: Option<String>, | |||
| concurrency: usize, | |||
| limit: Option<usize>, | |||
| ) -> Result<()> { | |||
| let config = AppConfig::load(&config_path)?; | |||
| let provider_name = provider_name(&config, provider.as_deref())?; | |||
| let provider_config = config.provider(Some(&provider_name))?; | |||
| let model = model.unwrap_or_else(|| provider_config.default_model.clone()); | |||
| let loaded = benchmarks::gpqa::load_cases(Path::new(&config.benchmarks.data_dir))?; | |||
| let dataset = dataset_report( | |||
| config | |||
| .benchmarks | |||
| .gpqa_diamond | |||
| .as_ref() | |||
| .map(|dataset| (dataset.source.as_str(), dataset.split.as_str())), | |||
| &loaded.local_path, | |||
| ); | |||
| let cases = apply_limit(loaded.cases, limit); | |||
| let started_at = Utc::now(); | |||
| let started = Instant::now(); | |||
| let base_request = request_template(provider_config, &model, 0.0, 1024); | |||
| let protocol = provider_config.protocol; | |||
| let results = stream::iter(cases.into_iter()) | |||
| .map(|case| { | |||
| let mut request = base_request.clone(); | |||
| request.prompt = case.prompt(); | |||
| async move { | |||
| let result = run_model_request(protocol, request).await; | |||
| (case, result) | |||
| } | |||
| }) | |||
| .buffer_unordered(nonzero_concurrency(concurrency)) | |||
| .collect::<Vec<_>>() | |||
| .await; | |||
| let mut metrics = Metrics::new(); | |||
| let mut wrong_cases = Vec::new(); | |||
| for (case, result) in results { | |||
| match result { | |||
| Ok(response) => { | |||
| metrics.record_success(response.status, response.elapsed_ms as u64); | |||
| let actual = judge::extract_choice(&response.text) | |||
| .map(|choice| choice.to_string()) | |||
| .unwrap_or_else(|| "no_answer".to_string()); | |||
| let expected = case.answer.to_string(); | |||
| let correct = actual == expected; | |||
| metrics.record_judgement(correct); | |||
| if !correct { | |||
| wrong_cases.push(WrongCaseReport { | |||
| id: case.id, | |||
| question: case.question, | |||
| expected, | |||
| actual, | |||
| raw_output: response.text, | |||
| }); | |||
| } | |||
| } | |||
| Err(error) => metrics.record_failure(error_code(&error)), | |||
| } | |||
| } | |||
| let summary = metrics.summary(); | |||
| let report = benchmark_report( | |||
| "gpqa-diamond", | |||
| provider_name, | |||
| model, | |||
| dataset, | |||
| started_at, | |||
| started.elapsed().as_millis(), | |||
| concurrency, | |||
| limit, | |||
| summary, | |||
| wrong_cases, | |||
| ); | |||
| let report_path = write_benchmark_report(Path::new("."), &report)?; | |||
| print_benchmark_report(&report, &report_path); | |||
| Ok(()) | |||
| } | |||
| async fn run_rpm( | |||
| config_path: PathBuf, | |||
| provider: Option<String>, | |||
| model: Option<String>, | |||
| rpm: u32, | |||
| duration: String, | |||
| prompt: String, | |||
| ) -> Result<()> { | |||
| if rpm == 0 { | |||
| bail!("rpm must be greater than 0"); | |||
| } | |||
| let duration = parse_duration(&duration)?; | |||
| let config = AppConfig::load(&config_path)?; | |||
| let provider_name = provider_name(&config, provider.as_deref())?; | |||
| let provider_config = config.provider(Some(&provider_name))?; | |||
| let model = model.unwrap_or_else(|| provider_config.default_model.clone()); | |||
| let request = ModelRequest { | |||
| prompt: prompt.clone(), | |||
| ..request_template(provider_config, &model, 0.0, 1024) | |||
| }; | |||
| let delay = Duration::from_secs_f64(60.0 / rpm as f64); | |||
| let started_at = Utc::now(); | |||
| let started = Instant::now(); | |||
| let mut metrics = Metrics::new(); | |||
| while started.elapsed() < duration { | |||
| match run_model_request(provider_config.protocol, request.clone()).await { | |||
| Ok(response) => metrics.record_success(response.status, response.elapsed_ms as u64), | |||
| Err(error) => metrics.record_failure(error_code(&error)), | |||
| } | |||
| if started.elapsed() + delay <= duration { | |||
| sleep(delay).await; | |||
| } | |||
| } | |||
| let summary = metrics.summary(); | |||
| let report = RpmReport { | |||
| benchmark: "rpm".to_string(), | |||
| provider: provider_name, | |||
| model, | |||
| run: RpmRunReport { | |||
| started_at, | |||
| duration_ms: started.elapsed().as_millis(), | |||
| target_rpm: rpm, | |||
| prompt, | |||
| temperature: 0.0, | |||
| max_tokens: 1024, | |||
| }, | |||
| summary: RpmSummaryReport { | |||
| actual_requests: summary.total, | |||
| success: summary.success, | |||
| failure: summary.failed, | |||
| latency_ms: latency_report(&summary.latency_ms), | |||
| }, | |||
| errors: summary.errors, | |||
| }; | |||
| let report_path = write_rpm_report(Path::new("."), &report)?; | |||
| print_rpm_report(&report, &report_path); | |||
| Ok(()) | |||
| } | |||
| fn provider_name(config: &AppConfig, provider: Option<&str>) -> Result<String> { | |||
| match provider { | |||
| Some(provider) => Ok(provider.to_string()), | |||
| None => config | |||
| .default_provider | |||
| .clone() | |||
| .context("no provider specified and config has no default_provider"), | |||
| } | |||
| } | |||
| fn request_template( | |||
| provider: &crate::config::ProviderConfig, | |||
| model: &str, | |||
| temperature: f32, | |||
| max_tokens: u32, | |||
| ) -> ModelRequest { | |||
| ModelRequest { | |||
| base_url: provider.base_url.clone(), | |||
| api_token: provider.api_token.clone(), | |||
| model: model.to_string(), | |||
| prompt: String::new(), | |||
| temperature, | |||
| max_tokens, | |||
| } | |||
| } | |||
| fn apply_limit<T>(cases: Vec<T>, limit: Option<usize>) -> Vec<T> { | |||
| match limit { | |||
| Some(limit) => cases.into_iter().take(limit).collect(), | |||
| None => cases, | |||
| } | |||
| } | |||
| fn dataset_report(config: Option<(&str, &str)>, local_path: &Path) -> DatasetReport { | |||
| let (source, split) = config.unwrap_or(("local", "train")); | |||
| DatasetReport { | |||
| source: source.to_string(), | |||
| split: split.to_string(), | |||
| revision: None, | |||
| local_path: local_path.display().to_string(), | |||
| } | |||
| } | |||
| fn benchmark_report( | |||
| benchmark: &str, | |||
| provider: String, | |||
| model: String, | |||
| dataset: DatasetReport, | |||
| started_at: chrono::DateTime<Utc>, | |||
| duration_ms: u128, | |||
| concurrency: usize, | |||
| limit: Option<usize>, | |||
| summary: MetricsSummary, | |||
| wrong_cases: Vec<WrongCaseReport>, | |||
| ) -> BenchmarkReport { | |||
| BenchmarkReport { | |||
| benchmark: benchmark.to_string(), | |||
| provider, | |||
| model, | |||
| dataset, | |||
| run: RunReport { | |||
| started_at, | |||
| duration_ms, | |||
| concurrency, | |||
| limit, | |||
| temperature: 0.0, | |||
| max_tokens: 1024, | |||
| }, | |||
| summary: BenchmarkSummaryReport { | |||
| accuracy: summary.accuracy, | |||
| success: summary.success, | |||
| total: summary.total, | |||
| correct: summary.correct, | |||
| wrong: summary.wrong, | |||
| failed: summary.failed, | |||
| latency_ms: latency_report(&summary.latency_ms), | |||
| }, | |||
| errors: summary.errors, | |||
| wrong_cases, | |||
| } | |||
| } | |||
| fn latency_report(summary: &LatencySummary) -> LatencyReport { | |||
| LatencyReport { | |||
| p50: summary.p50, | |||
| p95: summary.p95, | |||
| p99: summary.p99, | |||
| } | |||
| } | |||
| fn print_benchmark_report(report: &BenchmarkReport, report_path: &Path) { | |||
| println!("benchmark: {}", report.benchmark); | |||
| println!( | |||
| "accuracy: {}", | |||
| report | |||
| .summary | |||
| .accuracy | |||
| .map(|accuracy| format!("{:.2}%", accuracy * 100.0)) | |||
| .unwrap_or_else(|| "n/a".to_string()) | |||
| ); | |||
| println!( | |||
| "success: {}/{} (failed: {})", | |||
| report.summary.success, report.summary.total, report.summary.failed | |||
| ); | |||
| println!( | |||
| "latency_ms: p50={} p95={} p99={}", | |||
| format_optional_latency(report.summary.latency_ms.p50), | |||
| format_optional_latency(report.summary.latency_ms.p95), | |||
| format_optional_latency(report.summary.latency_ms.p99) | |||
| ); | |||
| println!("errors:"); | |||
| if report.errors.is_empty() { | |||
| println!(" none"); | |||
| } else { | |||
| for error in &report.errors { | |||
| println!(" {}: {}", error.code, error.count); | |||
| } | |||
| } | |||
| println!("wrong_cases:"); | |||
| if report.wrong_cases.is_empty() { | |||
| println!(" none"); | |||
| } else { | |||
| for case in &report.wrong_cases { | |||
| println!( | |||
| " {} expected={} actual={}", | |||
| case.id, case.expected, case.actual | |||
| ); | |||
| } | |||
| } | |||
| println!("report: {}", report_path.display()); | |||
| } | |||
| fn print_rpm_report(report: &RpmReport, report_path: &Path) { | |||
| println!("target_rpm: {}", report.run.target_rpm); | |||
| println!("actual_requests: {}", report.summary.actual_requests); | |||
| println!( | |||
| "success: {} failed: {}", | |||
| report.summary.success, report.summary.failure | |||
| ); | |||
| println!( | |||
| "latency_ms: p50={} p95={} p99={}", | |||
| format_optional_latency(report.summary.latency_ms.p50), | |||
| format_optional_latency(report.summary.latency_ms.p95), | |||
| format_optional_latency(report.summary.latency_ms.p99) | |||
| ); | |||
| println!("errors:"); | |||
| if report.errors.is_empty() { | |||
| println!(" none"); | |||
| } else { | |||
| for error in &report.errors { | |||
| println!(" {}: {}", error.code, error.count); | |||
| } | |||
| } | |||
| println!("report: {}", report_path.display()); | |||
| } | |||
| fn format_optional_latency(value: Option<u64>) -> String { | |||
| value | |||
| .map(|value| value.to_string()) | |||
| .unwrap_or_else(|| "n/a".to_string()) | |||
| } | |||
| fn error_code(error: &anyhow::Error) -> String { | |||
| let message = error.to_string(); | |||
| let status_regex = Regex::new(r"status\s+(\d{3})").expect("valid status regex"); | |||
| status_regex | |||
| .captures(&message) | |||
| .and_then(|captures| captures.get(1)) | |||
| .map(|code| code.as_str().to_string()) | |||
| .unwrap_or_else(|| "request_error".to_string()) | |||
| } | |||
| fn nonzero_concurrency(concurrency: usize) -> usize { | |||
| concurrency.max(1) | |||
| } | |||
| fn parse_duration(value: &str) -> Result<Duration> { | |||
| let value = value.trim(); | |||
| let Some(number) = value.strip_suffix('s') else { | |||
| if let Some(number) = value.strip_suffix('m') { | |||
| let minutes = number | |||
| .parse::<u64>() | |||
| .with_context(|| format!("invalid duration: {value}"))?; | |||
| return Ok(Duration::from_secs(minutes * 60)); | |||
| } | |||
| bail!("invalid duration: expected values like 60s or 5m"); | |||
| }; | |||
| let seconds = number | |||
| .parse::<u64>() | |||
| .with_context(|| format!("invalid duration: {value}"))?; | |||
| Ok(Duration::from_secs(seconds)) | |||
| } | |||
| fn dataset_data_dir(config_path: &Path) -> Result<PathBuf> { | |||
| @@ -157,4 +612,17 @@ mod tests { | |||
| assert!(error.to_string().contains("failed to parse config")); | |||
| } | |||
| #[test] | |||
| fn parses_duration_seconds_and_minutes() { | |||
| assert_eq!(parse_duration("60s").expect("seconds").as_secs(), 60); | |||
| assert_eq!(parse_duration("5m").expect("minutes").as_secs(), 300); | |||
| } | |||
| #[test] | |||
| fn rejects_invalid_duration() { | |||
| let error = parse_duration("one hour").expect_err("invalid duration"); | |||
| assert!(error.to_string().contains("duration")); | |||
| } | |||
| } | |||
| @@ -3,6 +3,7 @@ mod cli; | |||
| mod config; | |||
| mod metrics; | |||
| mod protocols; | |||
| mod report; | |||
| mod runner; | |||
| use anyhow::Result; | |||
| @@ -1 +1,181 @@ | |||
| use hdrhistogram::Histogram; | |||
| use serde::Serialize; | |||
| use std::collections::BTreeMap; | |||
| #[derive(Debug)] | |||
| pub struct Metrics { | |||
| success: u64, | |||
| failed: u64, | |||
| correct: u64, | |||
| total_judged: u64, | |||
| errors: BTreeMap<String, u64>, | |||
| latency: Histogram<u64>, | |||
| } | |||
| #[derive(Debug, Clone, PartialEq, Serialize)] | |||
| pub struct MetricsSummary { | |||
| pub success: u64, | |||
| pub failed: u64, | |||
| pub total: u64, | |||
| pub correct: u64, | |||
| pub wrong: u64, | |||
| pub total_judged: u64, | |||
| pub accuracy: Option<f64>, | |||
| pub latency_ms: LatencySummary, | |||
| pub errors: Vec<ErrorCount>, | |||
| } | |||
| #[derive(Debug, Clone, PartialEq, Eq, Serialize)] | |||
| pub struct LatencySummary { | |||
| pub p50: Option<u64>, | |||
| pub p95: Option<u64>, | |||
| pub p99: Option<u64>, | |||
| } | |||
| #[derive(Debug, Clone, PartialEq, Eq, Serialize)] | |||
| pub struct ErrorCount { | |||
| pub code: String, | |||
| pub count: u64, | |||
| } | |||
| impl ErrorCount { | |||
| pub fn new(code: impl Into<String>, count: u64) -> Self { | |||
| Self { | |||
| code: code.into(), | |||
| count, | |||
| } | |||
| } | |||
| } | |||
| impl Metrics { | |||
| pub fn new() -> Self { | |||
| Self { | |||
| success: 0, | |||
| failed: 0, | |||
| correct: 0, | |||
| total_judged: 0, | |||
| errors: BTreeMap::new(), | |||
| latency: Histogram::new(3).expect("valid histogram precision"), | |||
| } | |||
| } | |||
| pub fn record_success(&mut self, _status: u16, latency_ms: u64) { | |||
| self.success += 1; | |||
| let _ = self.latency.record(latency_ms); | |||
| } | |||
| pub fn record_failure(&mut self, code: impl Into<String>) { | |||
| self.failed += 1; | |||
| *self.errors.entry(code.into()).or_insert(0) += 1; | |||
| } | |||
| pub fn record_judgement(&mut self, correct: bool) { | |||
| self.total_judged += 1; | |||
| if correct { | |||
| self.correct += 1; | |||
| } | |||
| } | |||
| pub fn summary(&self) -> MetricsSummary { | |||
| let wrong = self.total_judged.saturating_sub(self.correct); | |||
| let accuracy = if self.total_judged == 0 { | |||
| None | |||
| } else { | |||
| Some(self.correct as f64 / self.total_judged as f64) | |||
| }; | |||
| let latency_ms = if self.latency.is_empty() { | |||
| LatencySummary { | |||
| p50: None, | |||
| p95: None, | |||
| p99: None, | |||
| } | |||
| } else { | |||
| LatencySummary { | |||
| p50: Some(self.latency.value_at_quantile(0.50)), | |||
| p95: Some(self.latency.value_at_quantile(0.95)), | |||
| p99: Some(self.latency.value_at_quantile(0.99)), | |||
| } | |||
| }; | |||
| MetricsSummary { | |||
| success: self.success, | |||
| failed: self.failed, | |||
| total: self.success + self.failed, | |||
| correct: self.correct, | |||
| wrong, | |||
| total_judged: self.total_judged, | |||
| accuracy, | |||
| latency_ms, | |||
| errors: self | |||
| .errors | |||
| .iter() | |||
| .map(|(code, count)| ErrorCount::new(code.clone(), *count)) | |||
| .collect(), | |||
| } | |||
| } | |||
| } | |||
| impl Default for Metrics { | |||
| fn default() -> Self { | |||
| Self::new() | |||
| } | |||
| } | |||
| #[cfg(test)] | |||
| mod tests { | |||
| use super::*; | |||
| #[test] | |||
| fn tracks_success_failure_and_error_counts() { | |||
| let mut metrics = Metrics::new(); | |||
| metrics.record_success(200, 120); | |||
| metrics.record_failure("request_error"); | |||
| metrics.record_failure("429"); | |||
| metrics.record_failure("429"); | |||
| let summary = metrics.summary(); | |||
| assert_eq!(summary.success, 1); | |||
| assert_eq!(summary.failed, 3); | |||
| assert_eq!(summary.total, 4); | |||
| assert_eq!( | |||
| summary.errors, | |||
| vec![ | |||
| ErrorCount::new("429", 2), | |||
| ErrorCount::new("request_error", 1) | |||
| ] | |||
| ); | |||
| } | |||
| #[test] | |||
| fn computes_accuracy_from_judged_cases() { | |||
| let mut metrics = Metrics::new(); | |||
| metrics.record_judgement(true); | |||
| metrics.record_judgement(false); | |||
| metrics.record_judgement(true); | |||
| let summary = metrics.summary(); | |||
| assert_eq!(summary.correct, 2); | |||
| assert_eq!(summary.wrong, 1); | |||
| assert_eq!(summary.total_judged, 3); | |||
| assert_eq!(summary.accuracy, Some(2.0 / 3.0)); | |||
| } | |||
| #[test] | |||
| fn computes_latency_percentiles() { | |||
| let mut metrics = Metrics::new(); | |||
| for latency in 1..=100 { | |||
| metrics.record_success(200, latency); | |||
| } | |||
| let summary = metrics.summary(); | |||
| assert_eq!(summary.latency_ms.p50, Some(50)); | |||
| assert_eq!(summary.latency_ms.p95, Some(95)); | |||
| assert_eq!(summary.latency_ms.p99, Some(99)); | |||
| } | |||
| } | |||
| @@ -0,0 +1,224 @@ | |||
| use crate::metrics::ErrorCount; | |||
| use anyhow::{Context, Result}; | |||
| use chrono::{DateTime, Utc}; | |||
| use serde::Serialize; | |||
| use std::path::{Path, PathBuf}; | |||
| #[derive(Debug, Clone, Serialize)] | |||
| pub struct DatasetReport { | |||
| pub source: String, | |||
| pub split: String, | |||
| pub revision: Option<String>, | |||
| pub local_path: String, | |||
| } | |||
| #[derive(Debug, Clone, Serialize)] | |||
| pub struct RunReport { | |||
| pub started_at: DateTime<Utc>, | |||
| pub duration_ms: u128, | |||
| pub concurrency: usize, | |||
| pub limit: Option<usize>, | |||
| pub temperature: f32, | |||
| pub max_tokens: u32, | |||
| } | |||
| #[derive(Debug, Clone, Serialize)] | |||
| pub struct LatencyReport { | |||
| pub p50: Option<u64>, | |||
| pub p95: Option<u64>, | |||
| pub p99: Option<u64>, | |||
| } | |||
| #[derive(Debug, Clone, Serialize)] | |||
| pub struct BenchmarkSummaryReport { | |||
| pub accuracy: Option<f64>, | |||
| pub success: u64, | |||
| pub total: u64, | |||
| pub correct: u64, | |||
| pub wrong: u64, | |||
| pub failed: u64, | |||
| pub latency_ms: LatencyReport, | |||
| } | |||
| #[derive(Debug, Clone, Serialize)] | |||
| pub struct WrongCaseReport { | |||
| pub id: String, | |||
| pub question: String, | |||
| pub expected: String, | |||
| pub actual: String, | |||
| pub raw_output: String, | |||
| } | |||
| #[derive(Debug, Clone, Serialize)] | |||
| pub struct BenchmarkReport { | |||
| pub benchmark: String, | |||
| pub provider: String, | |||
| pub model: String, | |||
| pub dataset: DatasetReport, | |||
| pub run: RunReport, | |||
| pub summary: BenchmarkSummaryReport, | |||
| pub errors: Vec<ErrorCount>, | |||
| pub wrong_cases: Vec<WrongCaseReport>, | |||
| } | |||
| #[derive(Debug, Clone, Serialize)] | |||
| pub struct RpmRunReport { | |||
| pub started_at: DateTime<Utc>, | |||
| pub duration_ms: u128, | |||
| pub target_rpm: u32, | |||
| pub prompt: String, | |||
| pub temperature: f32, | |||
| pub max_tokens: u32, | |||
| } | |||
| #[derive(Debug, Clone, Serialize)] | |||
| pub struct RpmSummaryReport { | |||
| pub actual_requests: u64, | |||
| pub success: u64, | |||
| pub failure: u64, | |||
| pub latency_ms: LatencyReport, | |||
| } | |||
| #[derive(Debug, Clone, Serialize)] | |||
| pub struct RpmReport { | |||
| pub benchmark: String, | |||
| pub provider: String, | |||
| pub model: String, | |||
| pub run: RpmRunReport, | |||
| pub summary: RpmSummaryReport, | |||
| pub errors: Vec<ErrorCount>, | |||
| } | |||
| pub fn write_benchmark_report(root: &Path, report: &BenchmarkReport) -> Result<PathBuf> { | |||
| write_report( | |||
| root, | |||
| &report.benchmark, | |||
| &report.provider, | |||
| &report.model, | |||
| report.run.started_at, | |||
| report, | |||
| ) | |||
| } | |||
| pub fn write_rpm_report(root: &Path, report: &RpmReport) -> Result<PathBuf> { | |||
| write_report( | |||
| root, | |||
| &report.benchmark, | |||
| &report.provider, | |||
| &report.model, | |||
| report.run.started_at, | |||
| report, | |||
| ) | |||
| } | |||
| fn write_report<T: Serialize>( | |||
| root: &Path, | |||
| benchmark: &str, | |||
| provider: &str, | |||
| model: &str, | |||
| started_at: DateTime<Utc>, | |||
| report: &T, | |||
| ) -> Result<PathBuf> { | |||
| let reports_dir = root.join("reports"); | |||
| std::fs::create_dir_all(&reports_dir).with_context(|| { | |||
| format!( | |||
| "failed to create reports directory {}", | |||
| reports_dir.display() | |||
| ) | |||
| })?; | |||
| let timestamp = started_at.format("%Y%m%dT%H%M%SZ"); | |||
| let filename = format!( | |||
| "{}-{}-{}-{timestamp}.json", | |||
| sanitize_filename_component(benchmark), | |||
| sanitize_filename_component(provider), | |||
| sanitize_filename_component(model), | |||
| ); | |||
| let path = reports_dir.join(filename); | |||
| let contents = serde_json::to_string_pretty(report).context("failed to serialize report")?; | |||
| std::fs::write(&path, contents) | |||
| .with_context(|| format!("failed to write report {}", path.display()))?; | |||
| Ok(path) | |||
| } | |||
| pub fn sanitize_filename_component(value: &str) -> String { | |||
| let sanitized = value | |||
| .trim() | |||
| .chars() | |||
| .map(|ch| { | |||
| if ch.is_ascii_alphanumeric() || ch == '.' || ch == '_' || ch == '-' { | |||
| ch | |||
| } else { | |||
| '-' | |||
| } | |||
| }) | |||
| .collect::<String>(); | |||
| let collapsed = sanitized | |||
| .split('-') | |||
| .filter(|part| !part.is_empty()) | |||
| .collect::<Vec<_>>() | |||
| .join("-"); | |||
| if collapsed.is_empty() { | |||
| "unknown".to_string() | |||
| } else { | |||
| collapsed | |||
| } | |||
| } | |||
| #[cfg(test)] | |||
| mod tests { | |||
| use super::*; | |||
| use chrono::TimeZone; | |||
| #[test] | |||
| fn sanitizes_model_names_for_filenames() { | |||
| assert_eq!( | |||
| sanitize_filename_component("openai/gpt-4.1:mini"), | |||
| "openai-gpt-4.1-mini" | |||
| ); | |||
| assert_eq!(sanitize_filename_component(" weird name "), "weird-name"); | |||
| } | |||
| #[test] | |||
| fn writes_benchmark_report_under_reports_dir() { | |||
| let temp_dir = tempfile::tempdir().expect("create temp dir"); | |||
| let report = BenchmarkReport { | |||
| benchmark: "aime2026".to_string(), | |||
| provider: "openai".to_string(), | |||
| model: "gpt/test".to_string(), | |||
| dataset: DatasetReport { | |||
| source: "local".to_string(), | |||
| split: "train".to_string(), | |||
| revision: None, | |||
| local_path: "data/benchmarks/aime2026/aime2026.jsonl".to_string(), | |||
| }, | |||
| run: RunReport { | |||
| started_at: Utc.with_ymd_and_hms(2026, 5, 6, 1, 2, 3).unwrap(), | |||
| duration_ms: 123, | |||
| concurrency: 2, | |||
| limit: Some(1), | |||
| temperature: 0.0, | |||
| max_tokens: 1024, | |||
| }, | |||
| summary: BenchmarkSummaryReport { | |||
| accuracy: Some(1.0), | |||
| success: 1, | |||
| total: 1, | |||
| correct: 1, | |||
| wrong: 0, | |||
| failed: 0, | |||
| latency_ms: LatencyReport { | |||
| p50: Some(10), | |||
| p95: Some(10), | |||
| p99: Some(10), | |||
| }, | |||
| }, | |||
| errors: vec![], | |||
| wrong_cases: vec![], | |||
| }; | |||
| let path = write_benchmark_report(temp_dir.path(), &report).expect("write report"); | |||
| assert!(path.ends_with("reports/aime2026-openai-gpt-test-20260506T010203Z.json")); | |||
| assert!(path.exists()); | |||
| } | |||
| } | |||