Skip to content

Commit fecea8c

Browse files
committed
feat: add progress tracking with ETA for benchmark runs
Add comprehensive progress tracking system: - ProgressStats dataclass with running accuracy, error rate, avg time, ETA - tqdm progress bar support with live accuracy and ETA display - Simple progress mode (prints periodic status updates) - Custom progress_callback for programmatic integration - CLI --progress/-p option: auto, tqdm, simple, none Progress shows: completed/total, current accuracy, errors, ETA
1 parent 5a076d8 commit fecea8c

4 files changed

Lines changed: 325 additions & 44 deletions

File tree

‎benchmarks/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
from benchmarks.base import Benchmark, BenchmarkResult, BenchmarkSample
2121
from benchmarks.metrics import Metrics
2222
from benchmarks.results import ExperimentConfig, ExperimentRecord, ResultsStore
23-
from benchmarks.runner import BenchmarkRunner, compare_methods
23+
from benchmarks.runner import BenchmarkRunner, ProgressStats, compare_methods
2424
from benchmarks.tasks import (
2525
BrowseCompPlusBenchmark,
2626
NIAHBenchmark,
@@ -35,6 +35,7 @@
3535
"BenchmarkSample",
3636
# Runner
3737
"BenchmarkRunner",
38+
"ProgressStats",
3839
"compare_methods",
3940
# Metrics
4041
"Metrics",

‎benchmarks/cli.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,15 @@ def main():
316316
default=1,
317317
help="Parallel workers (1=sequential, >1=parallel threads, default: 1)",
318318
)
319+
run_parser.add_argument(
320+
"--progress",
321+
"-p",
322+
type=str,
323+
default="auto",
324+
choices=["auto", "tqdm", "simple", "none"],
325+
help="Progress display mode: auto (uses tqdm if available), tqdm (progress bar), "
326+
"simple (periodic status), none (quiet). Default: auto",
327+
)
319328

320329
# Benchmark-specific options for run
321330
run_parser.add_argument("--context-length", type=int, default=100_000, help="NIAH context len")
@@ -390,6 +399,7 @@ def cmd_run(args: argparse.Namespace) -> int:
390399
verbose=args.verbose,
391400
log_dir=args.log_dir,
392401
max_workers=args.max_workers,
402+
progress=args.progress,
393403
)
394404

395405
if args.benchmark == "all":

‎benchmarks/runner.py‎

Lines changed: 234 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
- CodeAct (code-generation agents)
1010
1111
Supports parallel execution for faster evaluation.
12+
Includes progress tracking with ETA via tqdm or custom callbacks.
1213
"""
1314

1415
import time
@@ -19,6 +20,49 @@
1920

2021
from benchmarks.base import Benchmark, BenchmarkResult, BenchmarkSample, SampleResult
2122

23+
# Type alias for progress callback
24+
ProgressCallback = Callable[[int, int, "SampleResult | None", "ProgressStats"], None]
25+
26+
27+
@dataclass
28+
class ProgressStats:
29+
"""Running statistics for progress reporting."""
30+
31+
completed: int = 0
32+
total: int = 0
33+
correct: int = 0
34+
errors: int = 0
35+
total_time_ms: float = 0.0
36+
37+
@property
38+
def accuracy(self) -> float:
39+
"""Current running accuracy."""
40+
if self.completed == 0:
41+
return 0.0
42+
return self.correct / self.completed
43+
44+
@property
45+
def error_rate(self) -> float:
46+
"""Current error rate."""
47+
if self.completed == 0:
48+
return 0.0
49+
return self.errors / self.completed
50+
51+
@property
52+
def avg_time_ms(self) -> float:
53+
"""Average time per sample in milliseconds."""
54+
if self.completed == 0:
55+
return 0.0
56+
return self.total_time_ms / self.completed
57+
58+
@property
59+
def eta_seconds(self) -> float:
60+
"""Estimated time remaining in seconds."""
61+
remaining = self.total - self.completed
62+
if remaining <= 0 or self.avg_time_ms == 0:
63+
return 0.0
64+
return (remaining * self.avg_time_ms) / 1000.0
65+
2266

2367
@dataclass
2468
class RunnerConfig:
@@ -31,6 +75,8 @@ class RunnerConfig:
3175
verbose: bool = False
3276
log_dir: str | None = None
3377
max_workers: int = 1 # Number of parallel workers (1 = sequential)
78+
progress: str = "auto" # Progress display: "auto", "tqdm", "simple", "none"
79+
progress_callback: ProgressCallback | None = None # Custom progress callback
3480
backend_kwargs: dict[str, Any] = field(default_factory=dict)
3581
environment_kwargs: dict[str, Any] = field(default_factory=dict)
3682

@@ -67,6 +113,8 @@ def __init__(
67113
verbose: bool = False,
68114
log_dir: str | None = None,
69115
max_workers: int = 1,
116+
progress: str = "auto",
117+
progress_callback: ProgressCallback | None = None,
70118
**kwargs,
71119
):
72120
"""Initialize runner with configuration.
@@ -79,6 +127,13 @@ def __init__(
79127
verbose: Enable verbose output.
80128
log_dir: Directory for logging trajectories.
81129
max_workers: Number of parallel workers (1 = sequential, >1 = parallel).
130+
progress: Progress display mode:
131+
- "auto": Use tqdm if available, else simple
132+
- "tqdm": Force tqdm progress bar
133+
- "simple": Print periodic status updates
134+
- "none": No progress output
135+
progress_callback: Custom callback for progress updates.
136+
Signature: (completed, total, sample_result, stats) -> None
82137
**kwargs: Additional backend or environment kwargs.
83138
"""
84139
self.config = RunnerConfig(
@@ -89,9 +144,46 @@ def __init__(
89144
verbose=verbose,
90145
log_dir=log_dir,
91146
max_workers=max_workers,
147+
progress=progress,
148+
progress_callback=progress_callback,
92149
backend_kwargs={"model_name": model, **kwargs.get("backend_kwargs", {})},
93150
environment_kwargs=kwargs.get("environment_kwargs", {}),
94151
)
152+
self._tqdm_available: bool | None = None
153+
154+
def _check_tqdm(self) -> bool:
155+
"""Check if tqdm is available."""
156+
if self._tqdm_available is None:
157+
try:
158+
import tqdm # noqa: F401
159+
160+
self._tqdm_available = True
161+
except ImportError:
162+
self._tqdm_available = False
163+
return self._tqdm_available
164+
165+
def _get_progress_mode(self) -> str:
166+
"""Determine effective progress mode."""
167+
mode = self.config.progress
168+
if mode == "auto":
169+
return "tqdm" if self._check_tqdm() else "simple"
170+
if mode == "tqdm" and not self._check_tqdm():
171+
return "simple"
172+
return mode
173+
174+
def _format_eta(self, seconds: float) -> str:
175+
"""Format seconds as human-readable duration."""
176+
if seconds <= 0:
177+
return "--:--"
178+
if seconds < 60:
179+
return f"{seconds:.0f}s"
180+
if seconds < 3600:
181+
mins = int(seconds // 60)
182+
secs = int(seconds % 60)
183+
return f"{mins}:{secs:02d}"
184+
hours = int(seconds // 3600)
185+
mins = int((seconds % 3600) // 60)
186+
return f"{hours}h{mins:02d}m"
95187

96188
def run(
97189
self,
@@ -135,76 +227,175 @@ def run(
135227

136228
# Collect samples first (needed for parallel execution)
137229
samples = list(benchmark.load_samples(num_samples=num_samples, seed=seed))
230+
total = len(samples)
231+
232+
# Initialize progress tracking
233+
stats = ProgressStats(total=total)
234+
progress_mode = self._get_progress_mode()
138235

139236
if workers <= 1:
140-
# Sequential execution
237+
# Sequential execution with progress
238+
result.sample_results = self._run_sequential(
239+
samples, inference_fn, benchmark, stats, progress_mode
240+
)
241+
else:
242+
# Parallel execution with progress
243+
result.sample_results = self._run_parallel(
244+
samples, inference_fn, benchmark, workers, stats, progress_mode
245+
)
246+
247+
return result
248+
249+
def _update_progress(
250+
self,
251+
sample_result: SampleResult,
252+
stats: ProgressStats,
253+
progress_mode: str,
254+
pbar: Any = None,
255+
) -> None:
256+
"""Update progress statistics and display."""
257+
stats.completed += 1
258+
if sample_result.is_correct:
259+
stats.correct += 1
260+
if sample_result.error:
261+
stats.errors += 1
262+
stats.total_time_ms += sample_result.execution_time_ms
263+
264+
# Call custom callback if provided
265+
if self.config.progress_callback:
266+
self.config.progress_callback(
267+
stats.completed, stats.total, sample_result, stats
268+
)
269+
270+
# Update display based on mode
271+
if progress_mode == "tqdm" and pbar is not None:
272+
pbar.set_postfix(
273+
acc=f"{stats.accuracy:.1%}",
274+
err=stats.errors,
275+
eta=self._format_eta(stats.eta_seconds),
276+
refresh=False,
277+
)
278+
pbar.update(1)
279+
elif progress_mode == "simple":
280+
# Print every 10% or every sample for small runs
281+
interval = max(1, stats.total // 10)
282+
if stats.completed % interval == 0 or stats.completed == stats.total:
283+
print(
284+
f" Progress: {stats.completed}/{stats.total} "
285+
f"({stats.completed/stats.total:.0%}) | "
286+
f"Acc: {stats.accuracy:.1%} | "
287+
f"Errors: {stats.errors} | "
288+
f"ETA: {self._format_eta(stats.eta_seconds)}"
289+
)
290+
291+
# Verbose per-sample output
292+
if self.config.verbose:
293+
status = "✓" if sample_result.is_correct else "✗"
294+
print(f" [{status}] Sample {sample_result.sample_id}: {sample_result.metrics}")
295+
296+
def _run_sequential(
297+
self,
298+
samples: list[BenchmarkSample],
299+
inference_fn: Callable[[BenchmarkSample], tuple[str, dict[str, Any]]],
300+
benchmark: Benchmark,
301+
stats: ProgressStats,
302+
progress_mode: str,
303+
) -> list[SampleResult]:
304+
"""Run samples sequentially with progress tracking."""
305+
results: list[SampleResult] = []
306+
pbar = None
307+
308+
try:
309+
if progress_mode == "tqdm":
310+
from tqdm import tqdm
311+
312+
pbar = tqdm(
313+
total=stats.total,
314+
desc=f"{benchmark.name}",
315+
unit="sample",
316+
ncols=100,
317+
)
318+
141319
for sample in samples:
142320
sample_result = self._run_sample(sample, inference_fn, benchmark)
143-
result.sample_results.append(sample_result)
321+
results.append(sample_result)
322+
self._update_progress(sample_result, stats, progress_mode, pbar)
144323

145-
if self.config.verbose:
146-
status = "✓" if sample_result.is_correct else "✗"
147-
print(f" [{status}] Sample {sample.id}: {sample_result.metrics}")
148-
else:
149-
# Parallel execution
150-
result.sample_results = self._run_parallel(samples, inference_fn, benchmark, workers)
324+
finally:
325+
if pbar is not None:
326+
pbar.close()
151327

152-
return result
328+
return results
153329

154330
def _run_parallel(
155331
self,
156332
samples: list[BenchmarkSample],
157333
inference_fn: Callable[[BenchmarkSample], tuple[str, dict[str, Any]]],
158334
benchmark: Benchmark,
159335
max_workers: int,
336+
stats: ProgressStats,
337+
progress_mode: str,
160338
) -> list[SampleResult]:
161-
"""Run samples in parallel using thread pool.
339+
"""Run samples in parallel using thread pool with progress tracking.
162340
163341
Args:
164342
samples: List of samples to process.
165343
inference_fn: Inference function to apply.
166344
benchmark: Benchmark for evaluation.
167345
max_workers: Number of parallel threads.
346+
stats: Progress statistics to update.
347+
progress_mode: Progress display mode.
168348
169349
Returns:
170350
List of SampleResult in original sample order.
171351
"""
352+
import threading
353+
172354
results: dict[str, SampleResult] = {}
173-
completed = 0
174-
total = len(samples)
355+
lock = threading.Lock()
356+
pbar = None
175357

176-
with ThreadPoolExecutor(max_workers=max_workers) as executor:
177-
# Submit all tasks
178-
future_to_sample = {
179-
executor.submit(self._run_sample, sample, inference_fn, benchmark): sample
180-
for sample in samples
181-
}
182-
183-
# Collect results as they complete
184-
for future in as_completed(future_to_sample):
185-
sample = future_to_sample[future]
186-
try:
187-
sample_result = future.result()
188-
results[sample.id] = sample_result
189-
completed += 1
190-
191-
if self.config.verbose:
192-
status = "✓" if sample_result.is_correct else "✗"
193-
print(
194-
f" [{status}] ({completed}/{total}) "
195-
f"Sample {sample.id}: {sample_result.metrics}"
358+
try:
359+
if progress_mode == "tqdm":
360+
from tqdm import tqdm
361+
362+
pbar = tqdm(
363+
total=stats.total,
364+
desc=f"{benchmark.name}",
365+
unit="sample",
366+
ncols=100,
367+
)
368+
369+
with ThreadPoolExecutor(max_workers=max_workers) as executor:
370+
# Submit all tasks
371+
future_to_sample = {
372+
executor.submit(self._run_sample, sample, inference_fn, benchmark): sample
373+
for sample in samples
374+
}
375+
376+
# Collect results as they complete
377+
for future in as_completed(future_to_sample):
378+
sample = future_to_sample[future]
379+
try:
380+
sample_result = future.result()
381+
except Exception as e:
382+
# Handle unexpected errors
383+
sample_result = SampleResult(
384+
sample_id=sample.id,
385+
prediction="",
386+
expected=sample.expected_answer,
387+
is_correct=False,
388+
metrics={m: 0.0 for m in benchmark.default_metrics()},
389+
error=f"Parallel execution error: {e}",
196390
)
197-
except Exception as e:
198-
# Handle unexpected errors
199-
results[sample.id] = SampleResult(
200-
sample_id=sample.id,
201-
prediction="",
202-
expected=sample.expected_answer,
203-
is_correct=False,
204-
metrics={m: 0.0 for m in benchmark.default_metrics()},
205-
error=f"Parallel execution error: {e}",
206-
)
207-
completed += 1
391+
392+
with lock:
393+
results[sample.id] = sample_result
394+
self._update_progress(sample_result, stats, progress_mode, pbar)
395+
396+
finally:
397+
if pbar is not None:
398+
pbar.close()
208399

209400
# Return results in original sample order
210401
return [results[sample.id] for sample in samples]

0 commit comments

Comments
 (0)