From 5d1163d635c56fa7418ba56d39a03d5330dda88c Mon Sep 17 00:00:00 2001 From: orangels Date: Wed, 6 May 2026 17:59:16 +0800 Subject: [PATCH] feat: run benchmark and rpm tests --- src/benchmarks/aime.rs | 103 +++++++++ src/benchmarks/gpqa.rs | 153 +++++++++++++ src/cli.rs | 474 ++++++++++++++++++++++++++++++++++++++++- src/main.rs | 1 + src/metrics.rs | 180 ++++++++++++++++ src/report.rs | 224 +++++++++++++++++++ 6 files changed, 1132 insertions(+), 3 deletions(-) create mode 100644 src/report.rs diff --git a/src/benchmarks/aime.rs b/src/benchmarks/aime.rs index 50e9a21..addaa16 100644 --- a/src/benchmarks/aime.rs +++ b/src/benchmarks/aime.rs @@ -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, + pub local_path: PathBuf, +} + +#[derive(Debug, Deserialize)] +struct NormalizedAimeCase { + id: String, + problem: String, + answer: String, +} + +pub fn load_cases(data_dir: &Path) -> Result { + 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::>>()?; + 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")); + } } diff --git a/src/benchmarks/gpqa.rs b/src/benchmarks/gpqa.rs index bc11004..ba1fd38 100644 --- a/src/benchmarks/gpqa.rs +++ b/src/benchmarks/gpqa.rs @@ -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, + pub local_path: PathBuf, +} + +pub fn load_cases(data_dir: &Path) -> Result { + 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::>(); + + 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::>>()?; + + Ok(LoadedGpqaCases { + cases, + local_path: csv_path, + }) +} + +fn csv_value<'a>( + record: &'a [String], + header_index: &HashMap, + 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::() + .to_ascii_lowercase() +} + +fn parse_csv(input: &str) -> Result>> { + 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'); + } } diff --git a/src/cli.rs b/src/cli.rs index 8ca2f57..3184f8b 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -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, + model: Option, + concurrency: usize, + limit: Option, +) -> 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::>() + .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, + model: Option, + concurrency: usize, + limit: Option, +) -> 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::>() + .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, + model: Option, + 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 { + 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(cases: Vec, limit: Option) -> Vec { + 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, + duration_ms: u128, + concurrency: usize, + limit: Option, + summary: MetricsSummary, + wrong_cases: Vec, +) -> 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) -> 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 { + let value = value.trim(); + let Some(number) = value.strip_suffix('s') else { + if let Some(number) = value.strip_suffix('m') { + let minutes = number + .parse::() + .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::() + .with_context(|| format!("invalid duration: {value}"))?; + Ok(Duration::from_secs(seconds)) } fn dataset_data_dir(config_path: &Path) -> Result { @@ -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")); + } } diff --git a/src/main.rs b/src/main.rs index 7c28447..19ecb40 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,7 @@ mod cli; mod config; mod metrics; mod protocols; +mod report; mod runner; use anyhow::Result; diff --git a/src/metrics.rs b/src/metrics.rs index 8b13789..2fb3164 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -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, + latency: Histogram, +} + +#[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, + pub latency_ms: LatencySummary, + pub errors: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct LatencySummary { + pub p50: Option, + pub p95: Option, + pub p99: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct ErrorCount { + pub code: String, + pub count: u64, +} + +impl ErrorCount { + pub fn new(code: impl Into, 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) { + 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)); + } +} diff --git a/src/report.rs b/src/report.rs new file mode 100644 index 0000000..64694f2 --- /dev/null +++ b/src/report.rs @@ -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, + pub local_path: String, +} + +#[derive(Debug, Clone, Serialize)] +pub struct RunReport { + pub started_at: DateTime, + pub duration_ms: u128, + pub concurrency: usize, + pub limit: Option, + pub temperature: f32, + pub max_tokens: u32, +} + +#[derive(Debug, Clone, Serialize)] +pub struct LatencyReport { + pub p50: Option, + pub p95: Option, + pub p99: Option, +} + +#[derive(Debug, Clone, Serialize)] +pub struct BenchmarkSummaryReport { + pub accuracy: Option, + 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, + pub wrong_cases: Vec, +} + +#[derive(Debug, Clone, Serialize)] +pub struct RpmRunReport { + pub started_at: DateTime, + 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, +} + +pub fn write_benchmark_report(root: &Path, report: &BenchmarkReport) -> Result { + write_report( + root, + &report.benchmark, + &report.provider, + &report.model, + report.run.started_at, + report, + ) +} + +pub fn write_rpm_report(root: &Path, report: &RpmReport) -> Result { + write_report( + root, + &report.benchmark, + &report.provider, + &report.model, + report.run.started_at, + report, + ) +} + +fn write_report( + root: &Path, + benchmark: &str, + provider: &str, + model: &str, + started_at: DateTime, + report: &T, +) -> Result { + 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::(); + let collapsed = sanitized + .split('-') + .filter(|part| !part.is_empty()) + .collect::>() + .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()); + } +}