From 83edc789df28bc35cb0c62f13c5a1e85e5bb3dba Mon Sep 17 00:00:00 2001 From: orangels Date: Wed, 6 May 2026 18:10:15 +0800 Subject: [PATCH] fix: make benchmark and rpm execution reliable --- src/benchmarks/aime.rs | 20 ++++++++ src/benchmarks/gpqa.rs | 49 ++++++++++++++---- src/benchmarks/mod.rs | 111 ++++++++++++++++++++++++++++++++++++++--- src/cli.rs | 53 ++++++++++++++++---- src/report.rs | 33 +++++++++++- 5 files changed, 240 insertions(+), 26 deletions(-) diff --git a/src/benchmarks/aime.rs b/src/benchmarks/aime.rs index addaa16..0d935fc 100644 --- a/src/benchmarks/aime.rs +++ b/src/benchmarks/aime.rs @@ -118,6 +118,26 @@ mod tests { assert_eq!(loaded.cases[0].answer, "4"); } + #[test] + fn loads_jsonl_normalized_from_hugging_face_rows() { + 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":"aime2026-7","problem":"Compute 6 times 7.","answer":"42"} +"#, + ) + .expect("write normalized fetch output"); + + let loaded = load_cases(temp_dir.path()).expect("load cases"); + + assert_eq!(loaded.cases.len(), 1); + assert_eq!(loaded.cases[0].id, "aime2026-7"); + assert_eq!(loaded.cases[0].problem, "Compute 6 times 7."); + assert_eq!(loaded.cases[0].answer, "42"); + } + #[test] fn parquet_only_returns_conversion_hint() { let temp_dir = tempfile::tempdir().expect("create temp dir"); diff --git a/src/benchmarks/gpqa.rs b/src/benchmarks/gpqa.rs index ba1fd38..294f199 100644 --- a/src/benchmarks/gpqa.rs +++ b/src/benchmarks/gpqa.rs @@ -54,16 +54,22 @@ pub fn load_cases(data_dir: &Path) -> Result { format!("missing Incorrect Answer 3 column in row {}", index + 2) })?; + let id = format!("gpqa-diamond-{}", index + 1); + let (choices, answer) = rotated_choices( + [ + correct.to_string(), + incorrect_1.to_string(), + incorrect_2.to_string(), + incorrect_3.to_string(), + ], + index % 4, + ); + Ok(GpqaCase { - id: format!("gpqa-diamond-{}", index + 1), + id, 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', + choices, + answer, }) }) .collect::>>()?; @@ -95,6 +101,19 @@ fn normalize_header(header: &str) -> String { .to_ascii_lowercase() } +fn rotated_choices(mut values: [String; 4], rotation: usize) -> ([(String, String); 4], char) { + let rotation = rotation % values.len(); + values.rotate_left(rotation); + let answer_index = (4 - rotation) % 4; + let labels = ["A", "B", "C", "D"]; + let answer = labels[answer_index] + .chars() + .next() + .expect("labels are non-empty"); + let choices = std::array::from_fn(|index| (labels[index].to_string(), values[index].clone())); + (choices, answer) +} + fn parse_csv(input: &str) -> Result>> { let mut rows = Vec::new(); let mut row = Vec::new(); @@ -187,19 +206,29 @@ mod tests { 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", + "Question,Correct Answer,Incorrect Answer 1,Incorrect Answer 2,Incorrect Answer 3\n\"Which is correct?\",\"Right\",\"Wrong, with comma\",\"Nope\",\"Never\"\n\"Which is second?\",\"Second right\",\"Second wrong 1\",\"Second wrong 2\",\"Second wrong 3\"\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.len(), 2); 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'); + assert_eq!(loaded.cases[1].answer, 'D'); + assert_eq!( + loaded.cases[1] + .choices + .iter() + .find(|(label, _)| label == "D") + .map(|(_, text)| text.as_str()), + Some("Second right") + ); + assert!(loaded.cases.iter().any(|case| case.answer != 'A')); } } diff --git a/src/benchmarks/mod.rs b/src/benchmarks/mod.rs index 8de2167..e6d2745 100644 --- a/src/benchmarks/mod.rs +++ b/src/benchmarks/mod.rs @@ -4,12 +4,13 @@ pub mod judge; use anyhow::{Context, Result, bail}; use chrono::Utc; -use serde::Serialize; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; use sha2::{Digest, Sha256}; use std::path::{Path, PathBuf}; use tokio::fs; -const AIME_2026_URL: &str = "https://huggingface.co/datasets/MathArena/aime_2026/resolve/main/data/train-00000-of-00001.parquet"; +const AIME_2026_URL: &str = "https://datasets-server.huggingface.co/rows?dataset=MathArena%2Faime_2026&config=default&split=train&offset=0&length=100"; const GPQA_DIAMOND_URL: &str = "https://huggingface.co/datasets/Idavidrein/gpqa/resolve/main/gpqa_diamond.csv"; @@ -58,7 +59,12 @@ pub async fn fetch_dataset(dataset: &str, data_dir: &Path) -> Result { .bytes() .await .with_context(|| format!("failed to read downloaded bytes for {dataset}"))?; - fs::write(&spec.file_path, &bytes) + let file_bytes = match spec.normalize { + Some(normalize) => normalize(&bytes)?, + None => bytes.to_vec(), + }; + + fs::write(&spec.file_path, &file_bytes) .await .with_context(|| format!("failed to write dataset file {}", spec.file_path.display()))?; @@ -66,8 +72,8 @@ pub async fn fetch_dataset(dataset: &str, data_dir: &Path) -> Result { name: dataset, source_url: spec.source_url, downloaded_at: Utc::now().to_rfc3339(), - sha256: sha256_hex(&bytes), - bytes: bytes.len(), + sha256: sha256_hex(&file_bytes), + bytes: file_bytes.len(), }; let metadata_yaml = serde_yaml::to_string(&metadata).context("failed to serialize dataset metadata")?; @@ -84,6 +90,7 @@ struct DatasetSpec<'a> { dir: PathBuf, file_path: PathBuf, use_hf_token: bool, + normalize: Option Result>>, } fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result> { @@ -92,9 +99,10 @@ fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result let dir = data_dir.join("aime2026"); Ok(DatasetSpec { source_url: AIME_2026_URL, - file_path: dir.join("train-00000-of-00001.parquet"), + file_path: dir.join("aime2026.jsonl"), dir, use_hf_token: false, + normalize: Some(normalize_aime_rows_json), }) } "gpqa-diamond" => { @@ -104,6 +112,7 @@ fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result file_path: dir.join("gpqa_diamond.csv"), dir, use_hf_token: true, + normalize: None, }) } other => bail!( @@ -112,6 +121,45 @@ fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result } } +#[derive(Debug, Deserialize)] +struct AimeRowsResponse { + rows: Vec, +} + +#[derive(Debug, Deserialize)] +struct AimeRowWrapper { + row: AimeSourceRow, +} + +#[derive(Debug, Deserialize)] +struct AimeSourceRow { + problem_idx: u64, + problem: String, + answer: Value, +} + +fn normalize_aime_rows_json(bytes: &[u8]) -> Result> { + let response: AimeRowsResponse = + serde_json::from_slice(bytes).context("failed to parse AIME rows response")?; + let mut output = Vec::new(); + for row in response.rows { + let answer = match row.row.answer { + Value::String(answer) => answer, + Value::Number(answer) => answer.to_string(), + other => bail!("unexpected AIME answer value: {other}"), + }; + let line = json!({ + "id": format!("aime2026-{}", row.row.problem_idx), + "problem": row.row.problem, + "answer": answer, + }); + serde_json::to_writer(&mut output, &line) + .context("failed to serialize normalized AIME JSONL")?; + output.push(b'\n'); + } + Ok(output) +} + fn sha256_hex(bytes: &[u8]) -> String { let digest = Sha256::digest(bytes); let mut hex = String::with_capacity(digest.len() * 2); @@ -120,3 +168,54 @@ fn sha256_hex(bytes: &[u8]) -> String { } hex } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn normalizes_aime_rows_response_to_loader_jsonl() { + let rows_response = br#"{ + "rows": [ + { + "row": { + "problem_idx": 1, + "problem": "What is 20 + 22?", + "answer": 42 + } + }, + { + "row": { + "problem_idx": 2, + "problem": "What is 10 squared?", + "answer": "100" + } + } + ] + }"#; + + let normalized = normalize_aime_rows_json(rows_response).expect("normalize rows"); + + let lines = std::str::from_utf8(&normalized) + .expect("utf8 jsonl") + .lines() + .map(|line| serde_json::from_str::(line).expect("json line")) + .collect::>(); + + assert_eq!( + lines, + vec![ + json!({ + "id": "aime2026-1", + "problem": "What is 20 + 22?", + "answer": "42", + }), + json!({ + "id": "aime2026-2", + "problem": "What is 10 squared?", + "answer": "100", + }), + ] + ); + } +} diff --git a/src/cli.rs b/src/cli.rs index 3184f8b..66e8436 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -14,7 +14,7 @@ use futures::{StreamExt, stream}; use regex::Regex; use std::path::{Path, PathBuf}; use std::time::{Duration, Instant}; -use tokio::time::sleep; +use tokio::time::{Instant as TokioInstant, sleep_until}; #[derive(Debug, Parser)] #[command( @@ -342,22 +342,32 @@ async fn run_rpm( 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(), + prompt, ..request_template(provider_config, &model, 0.0, 1024) }; - let delay = Duration::from_secs_f64(60.0 / rpm as f64); + let schedule = rpm_start_schedule(duration, rpm); let started_at = Utc::now(); let started = Instant::now(); + let tokio_started = TokioInstant::now(); + let in_flight_limit = schedule.len().max(1); let mut metrics = Metrics::new(); - while started.elapsed() < duration { - match run_model_request(provider_config.protocol, request.clone()).await { + let results = stream::iter(schedule.into_iter().map(|offset| { + let request = request.clone(); + async move { + sleep_until(tokio_started + offset).await; + run_model_request(provider_config.protocol, request).await + } + })) + .buffer_unordered(in_flight_limit) + .collect::>() + .await; + + for result in results { + match result { 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(); @@ -369,7 +379,6 @@ async fn run_rpm( started_at, duration_ms: started.elapsed().as_millis(), target_rpm: rpm, - prompt, temperature: 0.0, max_tokens: 1024, }, @@ -579,6 +588,22 @@ fn parse_duration(value: &str) -> Result { Ok(Duration::from_secs(seconds)) } +fn rpm_start_schedule(duration: Duration, rpm: u32) -> Vec { + let interval_nanos = 60_000_000_000u128 / u128::from(rpm); + if interval_nanos == 0 { + return Vec::new(); + } + + let duration_nanos = duration.as_nanos(); + let mut offsets = Vec::new(); + let mut offset = 0u128; + while offset < duration_nanos { + offsets.push(Duration::from_nanos(offset.min(u128::from(u64::MAX)) as u64)); + offset += interval_nanos; + } + offsets +} + fn dataset_data_dir(config_path: &Path) -> Result { if !config_path.exists() { return Ok(PathBuf::from("data/benchmarks")); @@ -619,6 +644,16 @@ mod tests { assert_eq!(parse_duration("5m").expect("minutes").as_secs(), 300); } + #[test] + fn computes_rpm_start_schedule_from_run_start() { + let schedule = rpm_start_schedule(Duration::from_secs(60), 120); + + assert_eq!(schedule.len(), 120); + assert_eq!(schedule[0], Duration::ZERO); + assert_eq!(schedule[1], Duration::from_millis(500)); + assert_eq!(schedule[119], Duration::from_millis(59_500)); + } + #[test] fn rejects_invalid_duration() { let error = parse_duration("one hour").expect_err("invalid duration"); diff --git a/src/report.rs b/src/report.rs index 64694f2..12a64c7 100644 --- a/src/report.rs +++ b/src/report.rs @@ -66,7 +66,6 @@ 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, } @@ -221,4 +220,36 @@ mod tests { assert!(path.ends_with("reports/aime2026-openai-gpt-test-20260506T010203Z.json")); assert!(path.exists()); } + + #[test] + fn rpm_report_does_not_serialize_prompt() { + let report = RpmReport { + benchmark: "rpm".to_string(), + provider: "openai".to_string(), + model: "gpt/test".to_string(), + run: RpmRunReport { + started_at: Utc.with_ymd_and_hms(2026, 5, 6, 1, 2, 3).unwrap(), + duration_ms: 1000, + target_rpm: 60, + temperature: 0.0, + max_tokens: 1024, + }, + summary: RpmSummaryReport { + actual_requests: 1, + success: 1, + failure: 0, + latency_ms: LatencyReport { + p50: Some(10), + p95: Some(10), + p99: Some(10), + }, + }, + errors: vec![], + }; + + let json = serde_json::to_string(&report).expect("serialize report"); + + assert!(!json.contains("sensitive prompt")); + assert!(!json.contains("\"prompt\"")); + } }