|
- use crate::benchmarks;
- use crate::benchmarks::judge;
- use crate::config::AppConfig;
- use crate::metrics::{LatencySummary, Metrics, MetricsSummary};
- use crate::report::{
- BenchmarkParamsReport, BenchmarkReport, BenchmarkSummaryReport, CorrectCaseReport,
- DatasetReport, LatencyReport, LimiterInferenceKind, LimiterInferenceReport,
- PhaseSummaryReport, ProbeSecondReport, RpmModeDetailReport, RpmParamsReport, RpmReport,
- RpmRunReport, RpmSummaryReport, RunReport, WindowBoundaryReport, WrongCaseReport,
- write_benchmark_report, write_rpm_report,
- };
- use crate::rpm_modes::{
- ProbePhase, RpmMode, ScheduledProbe, burst_schedule, sliding_window_schedule,
- sustained_schedule, token_bucket_schedule, window_boundary_plan,
- };
- use crate::runner::{ModelRequest, run_model_request};
- use anyhow::{Context, Result, bail};
- use chrono::Utc;
- use clap::{Parser, Subcommand};
- use futures::{StreamExt, stream};
- use indicatif::{ProgressBar, ProgressStyle};
- use regex::Regex;
- use std::collections::BTreeMap;
- use std::path::{Path, PathBuf};
- use std::time::{Duration, Instant};
- use tokio::time::{Instant as TokioInstant, sleep_until};
-
- #[derive(Debug, Parser)]
- #[command(
- name = "lq_token_test",
- version,
- about = "Test LLM relay protocols, RPM, and benchmark accuracy"
- )]
- pub struct Cli {
- #[command(subcommand)]
- pub command: Command,
- }
-
- #[derive(Debug, Subcommand)]
- pub enum Command {
- Check {
- #[arg(long, default_value = "config.yaml")]
- config: PathBuf,
- #[arg(long)]
- provider: Option<String>,
- #[arg(long)]
- model: Option<String>,
- #[arg(long)]
- prompt: String,
- #[arg(long)]
- stream: Option<bool>,
- },
- Dataset {
- #[command(subcommand)]
- command: DatasetCommand,
- },
- Bench {
- #[command(subcommand)]
- command: BenchCommand,
- },
- Rpm {
- #[arg(long, default_value = "config.yaml")]
- config: PathBuf,
- #[arg(long, value_enum, default_value_t = RpmMode::Sustained)]
- mode: RpmMode,
- #[arg(long)]
- provider: Option<String>,
- #[arg(long)]
- model: Option<String>,
- #[arg(long)]
- rpm: Option<u32>,
- #[arg(long, default_value = "60s")]
- duration: String,
- #[arg(long)]
- burst: Option<u32>,
- #[arg(long)]
- probe_seconds: Option<u64>,
- #[arg(long, default_value_t = 500)]
- window_offset_ms: u64,
- #[arg(long)]
- concurrency: Option<usize>,
- #[arg(long)]
- prompt: String,
- #[arg(long)]
- stream: Option<bool>,
- },
- }
-
- #[derive(Debug, Subcommand)]
- pub enum DatasetCommand {
- Fetch { dataset: String },
- }
-
- #[derive(Debug, Subcommand)]
- pub enum BenchCommand {
- Aime2026 {
- #[arg(long, default_value = "config.yaml")]
- config: PathBuf,
- #[arg(long)]
- provider: Option<String>,
- #[arg(long)]
- model: Option<String>,
- #[arg(long, default_value_t = 4)]
- concurrency: usize,
- #[arg(long)]
- limit: Option<usize>,
- #[arg(long)]
- stream: Option<bool>,
- #[arg(long, default_value_t = 32768)]
- max_tokens: u32,
- },
- GpqaDiamond {
- #[arg(long, default_value = "config.yaml")]
- config: PathBuf,
- #[arg(long)]
- provider: Option<String>,
- #[arg(long)]
- model: Option<String>,
- #[arg(long, default_value_t = 4)]
- concurrency: usize,
- #[arg(long)]
- limit: Option<usize>,
- #[arg(long)]
- stream: Option<bool>,
- #[arg(long, default_value_t = 32768)]
- max_tokens: u32,
- },
- }
-
- pub async fn dispatch(cli: Cli) -> Result<()> {
- match cli.command {
- Command::Check {
- config,
- provider,
- model,
- prompt,
- stream,
- } => {
- let config = AppConfig::load(&config)?;
- let provider = config.resolved_provider(provider.as_deref())?;
- let request = ModelRequest {
- base_url: provider.base_url.clone(),
- api_token: provider.api_token.clone(),
- model: model.unwrap_or_else(|| provider.default_model.clone()),
- prompt,
- temperature: 0.0,
- max_tokens: 1024,
- stream: stream.unwrap_or(provider.stream),
- };
- let response = run_model_request(provider.protocol, request).await?;
-
- println!("status: {}", response.status);
- println!("elapsed_ms: {}", response.elapsed_ms);
- if let Some(ttft) = response.first_token_ms {
- println!("first_token_ms: {}", ttft);
- }
- println!("{}", response.text);
-
- Ok(())
- }
- Command::Dataset {
- command: DatasetCommand::Fetch { dataset },
- } => {
- let data_dir = dataset_data_dir(Path::new("config.yaml"))?;
- let path = benchmarks::fetch_dataset(&dataset, &data_dir).await?;
- println!("{}", path.display());
- Ok(())
- }
- Command::Bench { command } => dispatch_bench(command).await,
- Command::Rpm {
- config,
- mode,
- provider,
- model,
- rpm,
- duration,
- burst,
- probe_seconds,
- window_offset_ms,
- concurrency,
- prompt,
- stream,
- } => {
- run_rpm(
- config,
- RpmCommandOptions {
- mode,
- provider,
- model,
- rpm,
- duration,
- burst,
- probe_seconds,
- window_offset_ms,
- concurrency,
- prompt,
- stream,
- },
- )
- .await
- }
- }
- }
-
- async fn dispatch_bench(command: BenchCommand) -> Result<()> {
- match command {
- BenchCommand::Aime2026 {
- config,
- provider,
- model,
- concurrency,
- limit,
- stream,
- max_tokens,
- } => run_aime_benchmark(config, provider, model, concurrency, limit, stream, max_tokens).await,
- BenchCommand::GpqaDiamond {
- config,
- provider,
- model,
- concurrency,
- limit,
- stream,
- max_tokens,
- } => run_gpqa_benchmark(config, provider, model, concurrency, limit, stream, max_tokens).await,
- }
- }
-
- async fn run_aime_benchmark(
- config_path: PathBuf,
- provider: Option<String>,
- model: Option<String>,
- concurrency: usize,
- limit: Option<usize>,
- stream: Option<bool>,
- max_tokens: u32,
- ) -> Result<()> {
- let config = AppConfig::load(&config_path)?;
- let provider_name = provider_name(&config, provider.as_deref())?;
- let provider_config = config.resolved_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 total = cases.len() as u64;
- let started_at = Utc::now();
- let started = Instant::now();
- let mut base_request = request_template(&provider_config, &model, 0.0, max_tokens);
- base_request.stream = stream.unwrap_or(provider_config.stream);
- let protocol = provider_config.protocol;
-
- let pb = ProgressBar::new(total);
- pb.set_style(
- ProgressStyle::default_bar()
- .template("[{elapsed_precise}] {bar:40} {pos}/{len} ({eta})")
- .expect("valid template"),
- );
-
- let mut results = stream::iter(cases)
- .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));
-
- let mut metrics = Metrics::new();
- let mut wrong_cases = Vec::new();
- let mut correct_samples = Vec::new();
- while let Some((case, result)) = results.next().await {
- pb.inc(1);
- match result {
- Ok(response) => {
- metrics.record_success(response.status, response.elapsed_ms as u64, response.first_token_ms.map(|ms| ms as u64));
- let actual = judge::extract_final_integer(&response.text)
- .unwrap_or_else(|| "no_answer".to_string());
- let correct = judge::judge_integer(&response.text, &case.answer);
- metrics.record_judgement(correct);
- if correct {
- if correct_samples.len() < 5 {
- correct_samples.push(CorrectCaseReport {
- id: case.id,
- question: case.problem,
- expected: case.answer,
- raw_output: response.text,
- });
- }
- } else {
- 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)),
- }
- }
- pb.finish_and_clear();
-
- let summary = metrics.summary();
- let report = benchmark_report(BenchmarkReportInput {
- benchmark: "aime2026",
- provider: provider_name,
- model,
- stream: base_request.stream,
- dataset,
- started_at,
- duration_ms: started.elapsed().as_millis(),
- concurrency,
- limit,
- max_tokens,
- summary,
- correct_samples,
- 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>,
- stream: Option<bool>,
- max_tokens: u32,
- ) -> Result<()> {
- let config = AppConfig::load(&config_path)?;
- let provider_name = provider_name(&config, provider.as_deref())?;
- let provider_config = config.resolved_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 total = cases.len() as u64;
- let started_at = Utc::now();
- let started = Instant::now();
- let mut base_request = request_template(&provider_config, &model, 0.0, max_tokens);
- base_request.stream = stream.unwrap_or(provider_config.stream);
- let protocol = provider_config.protocol;
-
- let pb = ProgressBar::new(total);
- pb.set_style(
- ProgressStyle::default_bar()
- .template("[{elapsed_precise}] {bar:40} {pos}/{len} ({eta})")
- .expect("valid template"),
- );
-
- let mut results = stream::iter(cases)
- .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));
-
- let mut metrics = Metrics::new();
- let mut wrong_cases = Vec::new();
- let mut correct_samples = Vec::new();
- while let Some((case, result)) = results.next().await {
- pb.inc(1);
- match result {
- Ok(response) => {
- metrics.record_success(response.status, response.elapsed_ms as u64, response.first_token_ms.map(|ms| 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 = judge::judge_choice(&response.text, case.answer);
- metrics.record_judgement(correct);
- if correct {
- if correct_samples.len() < 5 {
- correct_samples.push(CorrectCaseReport {
- id: case.id,
- question: case.question,
- expected,
- raw_output: response.text,
- });
- }
- } else {
- wrong_cases.push(WrongCaseReport {
- id: case.id,
- question: case.question,
- expected,
- actual,
- raw_output: response.text,
- });
- }
- }
- Err(error) => metrics.record_failure(error_code(&error)),
- }
- }
- pb.finish_and_clear();
-
- let summary = metrics.summary();
- let report = benchmark_report(BenchmarkReportInput {
- benchmark: "gpqa-diamond",
- provider: provider_name,
- model,
- stream: base_request.stream,
- dataset,
- started_at,
- duration_ms: started.elapsed().as_millis(),
- concurrency,
- limit,
- max_tokens,
- summary,
- correct_samples,
- wrong_cases,
- });
- let report_path = write_benchmark_report(Path::new("."), &report)?;
- print_benchmark_report(&report, &report_path);
- Ok(())
- }
-
- struct RpmCommandOptions {
- mode: RpmMode,
- provider: Option<String>,
- model: Option<String>,
- rpm: Option<u32>,
- duration: String,
- burst: Option<u32>,
- probe_seconds: Option<u64>,
- window_offset_ms: u64,
- concurrency: Option<usize>,
- prompt: String,
- stream: Option<bool>,
- }
-
- async fn run_rpm(config_path: PathBuf, options: RpmCommandOptions) -> Result<()> {
- let config = AppConfig::load(&config_path)?;
- let mode_plan = build_rpm_mode_plan(
- options.mode,
- options.rpm,
- &options.duration,
- options.burst,
- options.probe_seconds,
- options.window_offset_ms,
- )?;
- let provider_name = provider_name(&config, options.provider.as_deref())?;
- let provider_config = config.resolved_provider(Some(&provider_name))?;
- let model = options
- .model
- .unwrap_or_else(|| provider_config.default_model.clone());
- let stream_enabled = options.stream.unwrap_or(provider_config.stream);
- let concurrency = options.concurrency.unwrap_or(mode_plan.default_concurrency);
- let request = ModelRequest {
- prompt: options.prompt.clone(),
- stream: stream_enabled,
- ..request_template(&provider_config, &model, 0.0, 1024)
- };
- let started_at = Utc::now();
- let started = Instant::now();
- let mut metrics = Metrics::new();
- let mut mode_summary = RpmModeSummaryBuilder::default();
-
- let results = run_scheduled_requests(
- provider_config.protocol,
- request,
- mode_plan.probes,
- concurrency,
- )
- .await;
-
- for result in results {
- let success = result.result.is_ok();
- mode_summary.record(result.phase, result.second, success);
- match result.result {
- Ok(response) => metrics.record_success(response.status, response.elapsed_ms as u64, response.first_token_ms.map(|ms| ms as u64)),
- Err(error) => metrics.record_failure(error_code(&error)),
- }
- }
-
- let summary = metrics.summary();
- let report = RpmReport {
- benchmark: "rpm".to_string(),
- provider: provider_name,
- model,
- params: RpmParamsReport {
- prompt: options.prompt,
- stream: stream_enabled,
- duration: options.duration,
- burst: options.burst,
- concurrency,
- probe_seconds: options.probe_seconds,
- window_offset_ms: options.window_offset_ms,
- },
- run: RpmRunReport {
- started_at,
- duration_ms: started.elapsed().as_millis(),
- target_rpm: mode_plan.target_rpm,
- actual_rpm: actual_rpm(summary.total, started.elapsed()),
- 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),
- ttft_ms: latency_report(&summary.ttft_ms),
- },
- mode: mode_plan.mode_name.to_string(),
- mode_detail: mode_summary.into_report(options.mode),
- errors: summary.errors,
- };
- let report_path = write_rpm_report(Path::new("."), &report)?;
- print_rpm_report(&report, &report_path);
- Ok(())
- }
-
- #[derive(Debug)]
- struct RpmModePlan {
- mode_name: &'static str,
- target_rpm: u32,
- probes: Vec<ScheduledProbe>,
- default_concurrency: usize,
- }
-
- #[derive(Debug)]
- struct ScheduledResult {
- phase: ProbePhase,
- second: Option<u64>,
- result: Result<crate::runner::ModelResponse>,
- }
-
- fn build_rpm_mode_plan(
- mode: RpmMode,
- rpm: Option<u32>,
- duration: &str,
- burst: Option<u32>,
- probe_seconds: Option<u64>,
- window_offset_ms: u64,
- ) -> Result<RpmModePlan> {
- let target_rpm = rpm.unwrap_or(0);
- let burst = burst.unwrap_or(target_rpm);
- let probes = match mode {
- RpmMode::Sustained => {
- let rpm = require_positive("rpm", rpm)?;
- let duration = parse_duration(duration)?;
- sustained_schedule(duration, rpm)
- .into_iter()
- .map(|offset| ScheduledProbe {
- offset,
- phase: ProbePhase::RefillProbe,
- second: Some(offset.as_secs()),
- })
- .collect()
- }
- RpmMode::Burst => {
- let burst = require_positive_value("burst", burst)?;
- burst_schedule(burst)
- .into_iter()
- .map(|offset| ScheduledProbe {
- offset,
- phase: ProbePhase::Burst,
- second: Some(0),
- })
- .collect()
- }
- RpmMode::TokenBucket => {
- let rpm = require_positive("rpm", rpm)?;
- let burst = require_positive_value("burst", burst)?;
- token_bucket_schedule(rpm, burst, probe_seconds.unwrap_or(30))
- }
- RpmMode::SlidingWindow => {
- let burst = require_positive_value("burst", burst)?;
- sliding_window_schedule(burst, probe_seconds.unwrap_or(90))
- }
- RpmMode::WindowBoundary => {
- let burst = require_positive_value("burst", burst)?;
- window_boundary_plan(Utc::now(), burst, window_offset_ms).probes
- }
- RpmMode::Diagnose => {
- let rpm = require_positive("rpm", rpm)?;
- let burst = require_positive_value("burst", burst)?;
- let mut probes = token_bucket_schedule(rpm, burst, probe_seconds.unwrap_or(90));
- probes.extend(sliding_window_schedule(burst, probe_seconds.unwrap_or(90)));
- probes.extend(window_boundary_plan(Utc::now(), burst, window_offset_ms).probes);
- probes
- }
- };
-
- let default_concurrency = probes.len().max(1);
- Ok(RpmModePlan {
- mode_name: mode_name(mode),
- target_rpm,
- default_concurrency,
- probes,
- })
- }
-
- fn mode_name(mode: RpmMode) -> &'static str {
- match mode {
- RpmMode::Sustained => "sustained",
- RpmMode::Burst => "burst",
- RpmMode::TokenBucket => "token-bucket",
- RpmMode::SlidingWindow => "sliding-window",
- RpmMode::WindowBoundary => "window-boundary",
- RpmMode::Diagnose => "diagnose",
- }
- }
-
- #[derive(Default)]
- struct RpmModeSummaryBuilder {
- phases: BTreeMap<&'static str, PhaseAccumulator>,
- refill_seconds: BTreeMap<u64, PhaseAccumulator>,
- sliding_seconds: BTreeMap<u64, PhaseAccumulator>,
- }
-
- impl RpmModeSummaryBuilder {
- fn record(&mut self, phase: ProbePhase, second: Option<u64>, success: bool) {
- match phase {
- ProbePhase::Burst => self.phase("burst").record(success),
- ProbePhase::RefillProbe => {
- self.phase("refill_probe").record(success);
- self.refill_seconds
- .entry(second.unwrap_or(0))
- .or_default()
- .record(success);
- }
- ProbePhase::SlidingProbe => {
- self.phase("sliding_probe").record(success);
- self.sliding_seconds
- .entry(second.unwrap_or(0))
- .or_default()
- .record(success);
- }
- ProbePhase::BeforeBoundary => self.phase("before_boundary").record(success),
- ProbePhase::AfterBoundary => self.phase("after_boundary").record(success),
- }
- }
-
- fn into_report(self, mode: RpmMode) -> Option<RpmModeDetailReport> {
- if mode == RpmMode::Sustained {
- return None;
- }
-
- let burst = self.phases.get("burst").map(PhaseAccumulator::to_report);
- let refill_probe = probe_seconds_report(self.refill_seconds);
- let sliding_probe = probe_seconds_report(self.sliding_seconds);
- let window_boundary = match (
- self.phases.get("before_boundary"),
- self.phases.get("after_boundary"),
- ) {
- (Some(before), Some(after)) => Some(WindowBoundaryReport {
- before: before.to_report(),
- after: after.to_report(),
- }),
- _ => None,
- };
- let inference = if mode == RpmMode::Diagnose {
- Some(infer_limiter(
- burst.as_ref(),
- &refill_probe,
- &sliding_probe,
- window_boundary.as_ref(),
- ))
- } else {
- None
- };
-
- Some(RpmModeDetailReport {
- burst,
- refill_probe,
- sliding_probe,
- window_boundary,
- inference,
- })
- }
-
- fn phase(&mut self, name: &'static str) -> &mut PhaseAccumulator {
- self.phases.entry(name).or_default()
- }
- }
-
- #[derive(Default)]
- struct PhaseAccumulator {
- sent: u64,
- success: u64,
- failure: u64,
- }
-
- impl PhaseAccumulator {
- fn record(&mut self, success: bool) {
- self.sent += 1;
- if success {
- self.success += 1;
- } else {
- self.failure += 1;
- }
- }
-
- fn to_report(&self) -> PhaseSummaryReport {
- PhaseSummaryReport {
- sent: self.sent,
- success: self.success,
- failure: self.failure,
- }
- }
-
- fn success_rate(&self) -> f64 {
- if self.sent == 0 {
- 0.0
- } else {
- self.success as f64 / self.sent as f64
- }
- }
- }
-
- fn probe_seconds_report(seconds: BTreeMap<u64, PhaseAccumulator>) -> Vec<ProbeSecondReport> {
- seconds
- .into_iter()
- .map(|(second, accumulator)| ProbeSecondReport {
- second,
- sent: accumulator.sent,
- success: accumulator.success,
- failure: accumulator.failure,
- })
- .collect()
- }
-
- fn infer_limiter(
- _burst: Option<&PhaseSummaryReport>,
- refill_probe: &[ProbeSecondReport],
- sliding_probe: &[ProbeSecondReport],
- window_boundary: Option<&WindowBoundaryReport>,
- ) -> LimiterInferenceReport {
- let mut signals = Vec::new();
-
- if let Some(boundary) = window_boundary {
- let before_rate = phase_success_rate(&boundary.before);
- let after_rate = phase_success_rate(&boundary.after);
- if after_rate > before_rate + 0.3 {
- signals.push("after-boundary success rate was much higher than before-boundary".into());
- return LimiterInferenceReport {
- likely_limiter: LimiterInferenceKind::FixedWindow,
- confidence: "medium".to_string(),
- signals,
- };
- }
- }
-
- let refill_sent: u64 = refill_probe.iter().map(|probe| probe.sent).sum();
- let refill_success: u64 = refill_probe.iter().map(|probe| probe.success).sum();
- if refill_sent > 0 && refill_success as f64 / refill_sent as f64 >= 0.5 {
- signals.push("refill probes recovered at a steady rate".into());
- return LimiterInferenceReport {
- likely_limiter: LimiterInferenceKind::TokenBucket,
- confidence: "medium".to_string(),
- signals,
- };
- }
-
- let early = sliding_probe
- .iter()
- .filter(|probe| probe.second <= 30)
- .fold(PhaseAccumulator::default(), |mut acc, probe| {
- acc.sent += probe.sent;
- acc.success += probe.success;
- acc.failure += probe.failure;
- acc
- });
- let late = sliding_probe
- .iter()
- .filter(|probe| probe.second >= 60)
- .fold(PhaseAccumulator::default(), |mut acc, probe| {
- acc.sent += probe.sent;
- acc.success += probe.success;
- acc.failure += probe.failure;
- acc
- });
- if late.sent > 0 && late.success_rate() > early.success_rate() + 0.3 {
- signals.push("probe recovery improved near the 60 second rolling window".into());
- return LimiterInferenceReport {
- likely_limiter: LimiterInferenceKind::SlidingWindow,
- confidence: "medium".to_string(),
- signals,
- };
- }
-
- signals.push("signals did not clearly match a limiter model".into());
- LimiterInferenceReport {
- likely_limiter: LimiterInferenceKind::Unknown,
- confidence: "low".to_string(),
- signals,
- }
- }
-
- fn phase_success_rate(phase: &PhaseSummaryReport) -> f64 {
- if phase.sent == 0 {
- 0.0
- } else {
- phase.success as f64 / phase.sent as f64
- }
- }
-
- fn actual_rpm(total_requests: u64, elapsed: Duration) -> Option<f64> {
- let elapsed_seconds = elapsed.as_secs_f64();
- if elapsed_seconds == 0.0 {
- None
- } else {
- Some(total_requests as f64 / elapsed_seconds * 60.0)
- }
- }
-
- fn require_positive(name: &str, value: Option<u32>) -> Result<u32> {
- let value = value.with_context(|| format!("{name} is required for this rpm mode"))?;
- require_positive_value(name, value)
- }
-
- fn require_positive_value(name: &str, value: u32) -> Result<u32> {
- if value == 0 {
- bail!("{name} must be greater than 0");
- }
- Ok(value)
- }
-
- async fn run_scheduled_requests(
- protocol: crate::config::ProtocolKind,
- request: ModelRequest,
- starts: Vec<ScheduledProbe>,
- max_in_flight: usize,
- ) -> Vec<ScheduledResult> {
- let total = starts.len() as u64;
- let pb = ProgressBar::new(total);
- pb.set_style(
- ProgressStyle::default_bar()
- .template("[{elapsed_precise}] {bar:40} {pos}/{len} ({eta})")
- .expect("valid template"),
- );
-
- let tokio_started = TokioInstant::now();
- let mut results = Vec::with_capacity(starts.len());
- let mut s = stream::iter(starts.into_iter().map(|start| {
- let request = request.clone();
- async move {
- sleep_until(tokio_started + start.offset).await;
- ScheduledResult {
- phase: start.phase,
- second: start.second,
- result: run_model_request(protocol, request).await,
- }
- }
- }))
- .buffer_unordered(nonzero_concurrency(max_in_flight));
-
- while let Some(result) = s.next().await {
- pb.inc(1);
- results.push(result);
- }
- pb.finish_and_clear();
- results
- }
-
- 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,
- stream: provider.stream,
- }
- }
-
- 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(),
- }
- }
-
- struct BenchmarkReportInput {
- benchmark: &'static str,
- provider: String,
- model: String,
- stream: bool,
- dataset: DatasetReport,
- started_at: chrono::DateTime<Utc>,
- duration_ms: u128,
- concurrency: usize,
- limit: Option<usize>,
- max_tokens: u32,
- summary: MetricsSummary,
- correct_samples: Vec<CorrectCaseReport>,
- wrong_cases: Vec<WrongCaseReport>,
- }
-
- fn benchmark_report(input: BenchmarkReportInput) -> BenchmarkReport {
- BenchmarkReport {
- benchmark: input.benchmark.to_string(),
- provider: input.provider,
- model: input.model,
- params: BenchmarkParamsReport {
- stream: input.stream,
- },
- dataset: input.dataset,
- run: RunReport {
- started_at: input.started_at,
- duration_ms: input.duration_ms,
- concurrency: input.concurrency,
- limit: input.limit,
- temperature: 0.0,
- max_tokens: input.max_tokens,
- },
- summary: BenchmarkSummaryReport {
- accuracy: input.summary.accuracy,
- success: input.summary.success,
- total: input.summary.total,
- correct: input.summary.correct,
- wrong: input.summary.wrong,
- failed: input.summary.failed,
- latency_ms: latency_report(&input.summary.latency_ms),
- ttft_ms: latency_report(&input.summary.ttft_ms),
- },
- errors: input.summary.errors,
- correct_samples: input.correct_samples,
- wrong_cases: input.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)
- );
- if report.summary.ttft_ms.p50.is_some() {
- println!(
- "ttft_ms: p50={} p95={} p99={}",
- format_optional_latency(report.summary.ttft_ms.p50),
- format_optional_latency(report.summary.ttft_ms.p95),
- format_optional_latency(report.summary.ttft_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!("mode: {}", report.mode);
- println!("target_rpm: {}", report.run.target_rpm);
- println!(
- "actual_rpm: {}",
- report
- .run
- .actual_rpm
- .map(|rpm| format!("{rpm:.2}"))
- .unwrap_or_else(|| "n/a".to_string())
- );
- 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)
- );
- if report.summary.ttft_ms.p50.is_some() {
- println!(
- "ttft_ms: p50={} p95={} p99={}",
- format_optional_latency(report.summary.ttft_ms.p50),
- format_optional_latency(report.summary.ttft_ms.p95),
- format_optional_latency(report.summary.ttft_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> {
- if !config_path.exists() {
- return Ok(PathBuf::from("data/benchmarks"));
- }
-
- let config = AppConfig::load(config_path)?;
- Ok(PathBuf::from(config.benchmarks.data_dir))
- }
-
- #[cfg(test)]
- mod tests {
- use super::*;
-
- #[test]
- fn dataset_data_dir_defaults_when_config_is_missing() {
- let temp_dir = tempfile::tempdir().expect("create temp dir");
- let missing_config = temp_dir.path().join("missing-config.yaml");
-
- let data_dir = dataset_data_dir(&missing_config).expect("default data dir");
-
- assert_eq!(data_dir, PathBuf::from("data/benchmarks"));
- }
-
- #[test]
- fn dataset_data_dir_propagates_invalid_existing_config() {
- let temp_dir = tempfile::tempdir().expect("create temp dir");
- let config_path = temp_dir.path().join("config.yaml");
- std::fs::write(&config_path, "providers: [").expect("write invalid config");
-
- let error = dataset_data_dir(&config_path).expect_err("invalid config should fail");
-
- 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 computes_rpm_start_schedule_from_run_start() {
- let schedule = sustained_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");
-
- assert!(error.to_string().contains("duration"));
- }
-
- #[test]
- fn rpm_command_defaults_to_sustained_mode() {
- let cli = Cli::try_parse_from([
- "lq_token_test",
- "rpm",
- "--provider",
- "anthropic",
- "--rpm",
- "120",
- "--prompt",
- "hello",
- ])
- .expect("parse rpm");
-
- let Command::Rpm { mode, rpm, .. } = cli.command else {
- panic!("expected rpm command");
- };
-
- assert_eq!(mode, RpmMode::Sustained);
- assert_eq!(rpm, Some(120));
- }
-
- #[test]
- fn rpm_command_parses_token_bucket_mode() {
- let cli = Cli::try_parse_from([
- "lq_token_test",
- "rpm",
- "--mode",
- "token-bucket",
- "--provider",
- "anthropic",
- "--rpm",
- "120",
- "--burst",
- "120",
- "--probe-seconds",
- "30",
- "--prompt",
- "hello",
- ])
- .expect("parse token bucket rpm");
-
- let Command::Rpm {
- mode,
- burst,
- probe_seconds,
- ..
- } = cli.command
- else {
- panic!("expected rpm command");
- };
-
- assert_eq!(mode, RpmMode::TokenBucket);
- assert_eq!(burst, Some(120));
- assert_eq!(probe_seconds, Some(30));
- }
-
- #[test]
- fn rpm_command_parses_window_boundary_offset() {
- let cli = Cli::try_parse_from([
- "lq_token_test",
- "rpm",
- "--mode",
- "window-boundary",
- "--provider",
- "anthropic",
- "--burst",
- "10",
- "--window-offset-ms",
- "250",
- "--prompt",
- "hello",
- ])
- .expect("parse window boundary rpm");
-
- let Command::Rpm {
- mode,
- window_offset_ms,
- ..
- } = cli.command
- else {
- panic!("expected rpm command");
- };
-
- assert_eq!(mode, RpmMode::WindowBoundary);
- assert_eq!(window_offset_ms, 250);
- }
-
- #[test]
- fn rpm_mode_validation_rejects_zero_values() {
- let zero_rpm = build_rpm_mode_plan(RpmMode::Sustained, Some(0), "60s", None, None, 500)
- .expect_err("zero rpm should fail");
- let zero_burst = build_rpm_mode_plan(RpmMode::Burst, None, "60s", Some(0), None, 500)
- .expect_err("zero burst should fail");
-
- assert!(zero_rpm.to_string().contains("rpm"));
- assert!(zero_burst.to_string().contains("burst"));
- }
- }
|