Quellcode durchsuchen

feat: add progress bars to benchmark and RPM execution

Show real-time progress with elapsed time, completion count, and ETA
using indicatif. Progress bar clears on completion to keep final
output clean.
main
orangels vor 1 Woche
Ursprung
Commit
4988a86654
1 geänderte Dateien mit 45 neuen und 14 gelöschten Zeilen
  1. +45
    -14
      src/cli.rs

+ 45
- 14
src/cli.rs Datei anzeigen

@@ -17,6 +17,7 @@ use anyhow::{Context, Result, bail};
use chrono::Utc; use chrono::Utc;
use clap::{Parser, Subcommand}; use clap::{Parser, Subcommand};
use futures::{StreamExt, stream}; use futures::{StreamExt, stream};
use indicatif::{ProgressBar, ProgressStyle};
use regex::Regex; use regex::Regex;
use std::collections::BTreeMap; use std::collections::BTreeMap;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
@@ -239,13 +240,21 @@ async fn run_aime_benchmark(
&loaded.local_path, &loaded.local_path,
); );
let cases = apply_limit(loaded.cases, limit); let cases = apply_limit(loaded.cases, limit);
let total = cases.len() as u64;
let started_at = Utc::now(); let started_at = Utc::now();
let started = Instant::now(); let started = Instant::now();
let mut base_request = request_template(&provider_config, &model, 0.0, 1024); let mut base_request = request_template(&provider_config, &model, 0.0, 1024);
base_request.stream = stream.unwrap_or(provider_config.stream); base_request.stream = stream.unwrap_or(provider_config.stream);
let protocol = provider_config.protocol; let protocol = provider_config.protocol;


let results = stream::iter(cases)
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| { .map(|case| {
let mut request = base_request.clone(); let mut request = base_request.clone();
request.prompt = case.prompt(); request.prompt = case.prompt();
@@ -254,13 +263,12 @@ async fn run_aime_benchmark(
(case, result) (case, result)
} }
}) })
.buffer_unordered(nonzero_concurrency(concurrency))
.collect::<Vec<_>>()
.await;
.buffer_unordered(nonzero_concurrency(concurrency));


let mut metrics = Metrics::new(); let mut metrics = Metrics::new();
let mut wrong_cases = Vec::new(); let mut wrong_cases = Vec::new();
for (case, result) in results {
while let Some((case, result)) = results.next().await {
pb.inc(1);
match result { match result {
Ok(response) => { Ok(response) => {
metrics.record_success(response.status, response.elapsed_ms as u64, response.first_token_ms.map(|ms| ms as u64)); metrics.record_success(response.status, response.elapsed_ms as u64, response.first_token_ms.map(|ms| ms as u64));
@@ -281,6 +289,7 @@ async fn run_aime_benchmark(
Err(error) => metrics.record_failure(error_code(&error)), Err(error) => metrics.record_failure(error_code(&error)),
} }
} }
pb.finish_and_clear();


let summary = metrics.summary(); let summary = metrics.summary();
let report = benchmark_report(BenchmarkReportInput { let report = benchmark_report(BenchmarkReportInput {
@@ -323,13 +332,21 @@ async fn run_gpqa_benchmark(
&loaded.local_path, &loaded.local_path,
); );
let cases = apply_limit(loaded.cases, limit); let cases = apply_limit(loaded.cases, limit);
let total = cases.len() as u64;
let started_at = Utc::now(); let started_at = Utc::now();
let started = Instant::now(); let started = Instant::now();
let mut base_request = request_template(&provider_config, &model, 0.0, 1024); let mut base_request = request_template(&provider_config, &model, 0.0, 1024);
base_request.stream = stream.unwrap_or(provider_config.stream); base_request.stream = stream.unwrap_or(provider_config.stream);
let protocol = provider_config.protocol; let protocol = provider_config.protocol;


let results = stream::iter(cases)
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| { .map(|case| {
let mut request = base_request.clone(); let mut request = base_request.clone();
request.prompt = case.prompt(); request.prompt = case.prompt();
@@ -338,13 +355,12 @@ async fn run_gpqa_benchmark(
(case, result) (case, result)
} }
}) })
.buffer_unordered(nonzero_concurrency(concurrency))
.collect::<Vec<_>>()
.await;
.buffer_unordered(nonzero_concurrency(concurrency));


let mut metrics = Metrics::new(); let mut metrics = Metrics::new();
let mut wrong_cases = Vec::new(); let mut wrong_cases = Vec::new();
for (case, result) in results {
while let Some((case, result)) = results.next().await {
pb.inc(1);
match result { match result {
Ok(response) => { Ok(response) => {
metrics.record_success(response.status, response.elapsed_ms as u64, response.first_token_ms.map(|ms| ms as u64)); metrics.record_success(response.status, response.elapsed_ms as u64, response.first_token_ms.map(|ms| ms as u64));
@@ -367,6 +383,7 @@ async fn run_gpqa_benchmark(
Err(error) => metrics.record_failure(error_code(&error)), Err(error) => metrics.record_failure(error_code(&error)),
} }
} }
pb.finish_and_clear();


let summary = metrics.summary(); let summary = metrics.summary();
let report = benchmark_report(BenchmarkReportInput { let report = benchmark_report(BenchmarkReportInput {
@@ -796,8 +813,17 @@ async fn run_scheduled_requests(
starts: Vec<ScheduledProbe>, starts: Vec<ScheduledProbe>,
max_in_flight: usize, max_in_flight: usize,
) -> Vec<ScheduledResult> { ) -> 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 tokio_started = TokioInstant::now();
stream::iter(starts.into_iter().map(|start| {
let mut results = Vec::with_capacity(starts.len());
let mut s = stream::iter(starts.into_iter().map(|start| {
let request = request.clone(); let request = request.clone();
async move { async move {
sleep_until(tokio_started + start.offset).await; sleep_until(tokio_started + start.offset).await;
@@ -808,9 +834,14 @@ async fn run_scheduled_requests(
} }
} }
})) }))
.buffer_unordered(nonzero_concurrency(max_in_flight))
.collect::<Vec<_>>()
.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> { fn provider_name(config: &AppConfig, provider: Option<&str>) -> Result<String> {


Laden…
Abbrechen
Speichern