diff --git a/src/cli.rs b/src/cli.rs index 40992f6..195ed48 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -17,6 +17,7 @@ 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}; @@ -239,13 +240,21 @@ async fn run_aime_benchmark( &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, 1024); base_request.stream = stream.unwrap_or(provider_config.stream); 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| { let mut request = base_request.clone(); request.prompt = case.prompt(); @@ -254,13 +263,12 @@ async fn run_aime_benchmark( (case, result) } }) - .buffer_unordered(nonzero_concurrency(concurrency)) - .collect::>() - .await; + .buffer_unordered(nonzero_concurrency(concurrency)); let mut metrics = Metrics::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 { Ok(response) => { 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)), } } + pb.finish_and_clear(); let summary = metrics.summary(); let report = benchmark_report(BenchmarkReportInput { @@ -323,13 +332,21 @@ async fn run_gpqa_benchmark( &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, 1024); base_request.stream = stream.unwrap_or(provider_config.stream); 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| { let mut request = base_request.clone(); request.prompt = case.prompt(); @@ -338,13 +355,12 @@ async fn run_gpqa_benchmark( (case, result) } }) - .buffer_unordered(nonzero_concurrency(concurrency)) - .collect::>() - .await; + .buffer_unordered(nonzero_concurrency(concurrency)); let mut metrics = Metrics::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 { Ok(response) => { 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)), } } + pb.finish_and_clear(); let summary = metrics.summary(); let report = benchmark_report(BenchmarkReportInput { @@ -796,8 +813,17 @@ async fn run_scheduled_requests( starts: Vec, max_in_flight: usize, ) -> Vec { + 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(); - 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(); async move { sleep_until(tokio_started + start.offset).await; @@ -808,9 +834,14 @@ async fn run_scheduled_requests( } } })) - .buffer_unordered(nonzero_concurrency(max_in_flight)) - .collect::>() - .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 {