Browse Source

fix: make benchmark and rpm execution reliable

main
orangels 1 week ago
parent
commit
83edc789df
5 changed files with 240 additions and 26 deletions
  1. +20
    -0
      src/benchmarks/aime.rs
  2. +39
    -10
      src/benchmarks/gpqa.rs
  3. +105
    -6
      src/benchmarks/mod.rs
  4. +44
    -9
      src/cli.rs
  5. +32
    -1
      src/report.rs

+ 20
- 0
src/benchmarks/aime.rs View File

@@ -118,6 +118,26 @@ mod tests {
assert_eq!(loaded.cases[0].answer, "4"); 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] #[test]
fn parquet_only_returns_conversion_hint() { fn parquet_only_returns_conversion_hint() {
let temp_dir = tempfile::tempdir().expect("create temp dir"); let temp_dir = tempfile::tempdir().expect("create temp dir");


+ 39
- 10
src/benchmarks/gpqa.rs View File

@@ -54,16 +54,22 @@ pub fn load_cases(data_dir: &Path) -> Result<LoadedGpqaCases> {
format!("missing Incorrect Answer 3 column in row {}", index + 2) 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 { Ok(GpqaCase {
id: format!("gpqa-diamond-{}", index + 1),
id,
question: question.to_string(), 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::<Result<Vec<_>>>()?; .collect::<Result<Vec<_>>>()?;
@@ -95,6 +101,19 @@ fn normalize_header(header: &str) -> String {
.to_ascii_lowercase() .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<Vec<Vec<String>>> { fn parse_csv(input: &str) -> Result<Vec<Vec<String>>> {
let mut rows = Vec::new(); let mut rows = Vec::new();
let mut row = 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::create_dir_all(&dataset_dir).expect("create dataset dir");
std::fs::write( std::fs::write(
dataset_dir.join("gpqa_diamond.csv"), 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"); .expect("write csv");


let loaded = load_cases(temp_dir.path()).expect("load cases"); let loaded = load_cases(temp_dir.path()).expect("load cases");


assert_eq!(loaded.local_path, dataset_dir.join("gpqa_diamond.csv")); 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].question, "Which is correct?");
assert_eq!( assert_eq!(
loaded.cases[0].choices[0], loaded.cases[0].choices[0],
("A".to_string(), "Right".to_string()) ("A".to_string(), "Right".to_string())
); );
assert_eq!(loaded.cases[0].answer, 'A'); 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'));
} }
} }

+ 105
- 6
src/benchmarks/mod.rs View File

@@ -4,12 +4,13 @@ pub mod judge;


use anyhow::{Context, Result, bail}; use anyhow::{Context, Result, bail};
use chrono::Utc; use chrono::Utc;
use serde::Serialize;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use tokio::fs; 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 = const GPQA_DIAMOND_URL: &str =
"https://huggingface.co/datasets/Idavidrein/gpqa/resolve/main/gpqa_diamond.csv"; "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<PathBuf> {
.bytes() .bytes()
.await .await
.with_context(|| format!("failed to read downloaded bytes for {dataset}"))?; .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 .await
.with_context(|| format!("failed to write dataset file {}", spec.file_path.display()))?; .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<PathBuf> {
name: dataset, name: dataset,
source_url: spec.source_url, source_url: spec.source_url,
downloaded_at: Utc::now().to_rfc3339(), 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 = let metadata_yaml =
serde_yaml::to_string(&metadata).context("failed to serialize dataset metadata")?; serde_yaml::to_string(&metadata).context("failed to serialize dataset metadata")?;
@@ -84,6 +90,7 @@ struct DatasetSpec<'a> {
dir: PathBuf, dir: PathBuf,
file_path: PathBuf, file_path: PathBuf,
use_hf_token: bool, use_hf_token: bool,
normalize: Option<fn(&[u8]) -> Result<Vec<u8>>>,
} }


fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result<DatasetSpec<'a>> { fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result<DatasetSpec<'a>> {
@@ -92,9 +99,10 @@ fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result<DatasetSpec<'a>
let dir = data_dir.join("aime2026"); let dir = data_dir.join("aime2026");
Ok(DatasetSpec { Ok(DatasetSpec {
source_url: AIME_2026_URL, source_url: AIME_2026_URL,
file_path: dir.join("train-00000-of-00001.parquet"),
file_path: dir.join("aime2026.jsonl"),
dir, dir,
use_hf_token: false, use_hf_token: false,
normalize: Some(normalize_aime_rows_json),
}) })
} }
"gpqa-diamond" => { "gpqa-diamond" => {
@@ -104,6 +112,7 @@ fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result<DatasetSpec<'a>
file_path: dir.join("gpqa_diamond.csv"), file_path: dir.join("gpqa_diamond.csv"),
dir, dir,
use_hf_token: true, use_hf_token: true,
normalize: None,
}) })
} }
other => bail!( other => bail!(
@@ -112,6 +121,45 @@ fn dataset_spec<'a>(dataset: &str, data_dir: &'a Path) -> Result<DatasetSpec<'a>
} }
} }


#[derive(Debug, Deserialize)]
struct AimeRowsResponse {
rows: Vec<AimeRowWrapper>,
}

#[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<Vec<u8>> {
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 { fn sha256_hex(bytes: &[u8]) -> String {
let digest = Sha256::digest(bytes); let digest = Sha256::digest(bytes);
let mut hex = String::with_capacity(digest.len() * 2); let mut hex = String::with_capacity(digest.len() * 2);
@@ -120,3 +168,54 @@ fn sha256_hex(bytes: &[u8]) -> String {
} }
hex 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::<Value>(line).expect("json line"))
.collect::<Vec<_>>();

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

+ 44
- 9
src/cli.rs View File

@@ -14,7 +14,7 @@ use futures::{StreamExt, stream};
use regex::Regex; use regex::Regex;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use tokio::time::sleep;
use tokio::time::{Instant as TokioInstant, sleep_until};


#[derive(Debug, Parser)] #[derive(Debug, Parser)]
#[command( #[command(
@@ -342,22 +342,32 @@ async fn run_rpm(
let provider_config = config.provider(Some(&provider_name))?; let provider_config = config.provider(Some(&provider_name))?;
let model = model.unwrap_or_else(|| provider_config.default_model.clone()); let model = model.unwrap_or_else(|| provider_config.default_model.clone());
let request = ModelRequest { let request = ModelRequest {
prompt: prompt.clone(),
prompt,
..request_template(provider_config, &model, 0.0, 1024) ..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_at = Utc::now();
let started = Instant::now(); let started = Instant::now();
let tokio_started = TokioInstant::now();
let in_flight_limit = schedule.len().max(1);
let mut metrics = Metrics::new(); 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::<Vec<_>>()
.await;

for result in results {
match result {
Ok(response) => metrics.record_success(response.status, response.elapsed_ms as u64), Ok(response) => metrics.record_success(response.status, response.elapsed_ms as u64),
Err(error) => metrics.record_failure(error_code(&error)), Err(error) => metrics.record_failure(error_code(&error)),
} }
if started.elapsed() + delay <= duration {
sleep(delay).await;
}
} }


let summary = metrics.summary(); let summary = metrics.summary();
@@ -369,7 +379,6 @@ async fn run_rpm(
started_at, started_at,
duration_ms: started.elapsed().as_millis(), duration_ms: started.elapsed().as_millis(),
target_rpm: rpm, target_rpm: rpm,
prompt,
temperature: 0.0, temperature: 0.0,
max_tokens: 1024, max_tokens: 1024,
}, },
@@ -579,6 +588,22 @@ fn parse_duration(value: &str) -> Result<Duration> {
Ok(Duration::from_secs(seconds)) Ok(Duration::from_secs(seconds))
} }


fn rpm_start_schedule(duration: Duration, rpm: u32) -> Vec<Duration> {
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<PathBuf> { fn dataset_data_dir(config_path: &Path) -> Result<PathBuf> {
if !config_path.exists() { if !config_path.exists() {
return Ok(PathBuf::from("data/benchmarks")); return Ok(PathBuf::from("data/benchmarks"));
@@ -619,6 +644,16 @@ mod tests {
assert_eq!(parse_duration("5m").expect("minutes").as_secs(), 300); 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] #[test]
fn rejects_invalid_duration() { fn rejects_invalid_duration() {
let error = parse_duration("one hour").expect_err("invalid duration"); let error = parse_duration("one hour").expect_err("invalid duration");


+ 32
- 1
src/report.rs View File

@@ -66,7 +66,6 @@ pub struct RpmRunReport {
pub started_at: DateTime<Utc>, pub started_at: DateTime<Utc>,
pub duration_ms: u128, pub duration_ms: u128,
pub target_rpm: u32, pub target_rpm: u32,
pub prompt: String,
pub temperature: f32, pub temperature: f32,
pub max_tokens: u32, pub max_tokens: u32,
} }
@@ -221,4 +220,36 @@ mod tests {
assert!(path.ends_with("reports/aime2026-openai-gpt-test-20260506T010203Z.json")); assert!(path.ends_with("reports/aime2026-openai-gpt-test-20260506T010203Z.json"));
assert!(path.exists()); 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\""));
}
} }

Loading…
Cancel
Save