| @@ -1,3 +1,7 @@ | |||||
| use anyhow::{Context, Result, bail}; | |||||
| use serde::Deserialize; | |||||
| use std::path::{Path, PathBuf}; | |||||
| #[derive(Debug, Clone, PartialEq, Eq)] | #[derive(Debug, Clone, PartialEq, Eq)] | ||||
| pub struct AimeCase { | pub struct AimeCase { | ||||
| pub id: String, | pub id: String, | ||||
| @@ -5,6 +9,68 @@ pub struct AimeCase { | |||||
| pub answer: String, | 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 { | impl AimeCase { | ||||
| pub fn prompt(&self) -> String { | pub fn prompt(&self) -> String { | ||||
| format!( | format!( | ||||
| @@ -31,4 +97,41 @@ mod tests { | |||||
| assert!(prompt.contains("What is 20 + 22?")); | assert!(prompt.contains("What is 20 + 22?")); | ||||
| assert!(prompt.contains("final integer answer")); | 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)] | #[derive(Debug, Clone, PartialEq, Eq)] | ||||
| pub struct GpqaCase { | pub struct GpqaCase { | ||||
| pub id: String, | pub id: String, | ||||
| @@ -6,6 +10,132 @@ pub struct GpqaCase { | |||||
| pub answer: char, | 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 { | impl GpqaCase { | ||||
| pub fn prompt(&self) -> String { | pub fn prompt(&self) -> String { | ||||
| let choices = self | let choices = self | ||||
| @@ -49,4 +179,27 @@ mod tests { | |||||
| assert!(prompt.contains("D. Delta")); | assert!(prompt.contains("D. Delta")); | ||||
| assert!(prompt.contains("answer with exactly one letter")); | 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; | ||||
| use crate::benchmarks::judge; | |||||
| use crate::config::AppConfig; | 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 crate::runner::{ModelRequest, run_model_request}; | ||||
| use anyhow::Result; | |||||
| use anyhow::{Context, Result, bail}; | |||||
| use chrono::Utc; | |||||
| use clap::{Parser, Subcommand}; | use clap::{Parser, Subcommand}; | ||||
| use futures::{StreamExt, stream}; | |||||
| use regex::Regex; | |||||
| use std::path::{Path, PathBuf}; | use std::path::{Path, PathBuf}; | ||||
| use std::time::{Duration, Instant}; | |||||
| use tokio::time::sleep; | |||||
| #[derive(Debug, Parser)] | #[derive(Debug, Parser)] | ||||
| #[command( | #[command( | ||||
| @@ -119,9 +130,453 @@ pub async fn dispatch(cli: Cli) -> Result<()> { | |||||
| println!("{}", path.display()); | println!("{}", path.display()); | ||||
| Ok(()) | 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> { | fn dataset_data_dir(config_path: &Path) -> Result<PathBuf> { | ||||
| @@ -157,4 +612,17 @@ mod tests { | |||||
| assert!(error.to_string().contains("failed to parse config")); | 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 config; | ||||
| mod metrics; | mod metrics; | ||||
| mod protocols; | mod protocols; | ||||
| mod report; | |||||
| mod runner; | mod runner; | ||||
| use anyhow::Result; | 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()); | |||||
| } | |||||
| } | |||||