Parcourir la source

fix: make benchmark and rpm execution reliable

main
orangels il y a 1 semaine
Parent
révision
83edc789df
5 fichiers modifiés avec 240 ajouts et 26 suppressions
  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 Voir le fichier

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


+ 39
- 10
src/benchmarks/gpqa.rs Voir le fichier

@@ -54,16 +54,22 @@ pub fn load_cases(data_dir: &Path) -> Result<LoadedGpqaCases> {
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::<Result<Vec<_>>>()?;
@@ -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<Vec<Vec<String>>> {
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'));
}
}

+ 105
- 6
src/benchmarks/mod.rs Voir le fichier

@@ -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<PathBuf> {
.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<PathBuf> {
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<fn(&[u8]) -> Result<Vec<u8>>>,
}

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");
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<DatasetSpec<'a>
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<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 {
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::<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 Voir le fichier

@@ -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::<Vec<_>>()
.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<Duration> {
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> {
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");


+ 32
- 1
src/report.rs Voir le fichier

@@ -66,7 +66,6 @@ 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,
}
@@ -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\""));
}
}

Chargement…
Annuler
Enregistrer