Parcourir la source

feat: run benchmark and rpm tests

main
orangels il y a 1 mois
Parent
révision
5d1163d635
6 fichiers modifiés avec 1132 ajouts et 3 suppressions
  1. +103
    -0
      src/benchmarks/aime.rs
  2. +153
    -0
      src/benchmarks/gpqa.rs
  3. +471
    -3
      src/cli.rs
  4. +1
    -0
      src/main.rs
  5. +180
    -0
      src/metrics.rs
  6. +224
    -0
      src/report.rs

+ 103
- 0
src/benchmarks/aime.rs Voir le fichier

@@ -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"));
}
}

+ 153
- 0
src/benchmarks/gpqa.rs Voir le fichier

@@ -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');
}
}

+ 471
- 3
src/cli.rs Voir le fichier

@@ -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"));
}
}

+ 1
- 0
src/main.rs Voir le fichier

@@ -3,6 +3,7 @@ mod cli;
mod config;
mod metrics;
mod protocols;
mod report;
mod runner;

use anyhow::Result;


+ 180
- 0
src/metrics.rs Voir le fichier

@@ -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));
}
}

+ 224
- 0
src/report.rs Voir le fichier

@@ -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());
}
}

Chargement…
Annuler
Enregistrer