Przeglądaj źródła

fix: make benchmark and rpm execution reliable

main
orangels 1 tydzień temu
rodzic
commit
83edc789df
5 zmienionych plików z 240 dodań i 26 usunięć
  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 Wyświetl plik

@@ -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 Wyświetl plik

@@ -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 Wyświetl plik

@@ -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 Wyświetl plik

@@ -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 Wyświetl plik

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

Ładowanie…
Anuluj
Zapisz