Просмотр исходного кода

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 1 неделю назад
Родитель
Сommit
4988a86654
1 измененных файлов: 45 добавлений и 14 удалений
  1. +45
    -14
      src/cli.rs

+ 45
- 14
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::<Vec<_>>()
.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::<Vec<_>>()
.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<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();
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::<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> {


Загрузка…
Отмена
Сохранить