193 lines
7.6 KiB
Python
193 lines
7.6 KiB
Python
"""Rich terminal progress reporter for per-sample model generation."""
|
|
|
|
import asyncio
|
|
import time
|
|
|
|
from rich.console import Console
|
|
from rich.progress import (
|
|
BarColumn,
|
|
Progress,
|
|
SpinnerColumn,
|
|
TaskProgressColumn,
|
|
TextColumn,
|
|
TimeElapsedColumn,
|
|
TimeRemainingColumn,
|
|
)
|
|
|
|
|
|
class RichTerminalProgress:
|
|
def __init__(self, console=None):
|
|
# accept an EXTERNAL console: CLI phase messages and the live bar must
|
|
# share one console, or the two writers interleave and repaint wrongly
|
|
self.console = console or Console()
|
|
self.overall_id = None
|
|
self.overall_total = 0
|
|
self.progress = Progress(
|
|
SpinnerColumn(),
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(complete_style="green", finished_style="bold green"),
|
|
TaskProgressColumn(),
|
|
TextColumn("• Completed {task.completed}/{task.total}"),
|
|
TextColumn("• in-flight {task.fields[inflight]}"),
|
|
TextColumn("• {task.fields[rate]}/s"),
|
|
TextColumn("•"),
|
|
TimeElapsedColumn(),
|
|
TextColumn("•"),
|
|
TimeRemainingColumn(),
|
|
console=self.console,
|
|
refresh_per_second=4,
|
|
)
|
|
self.task_id = None
|
|
self.bench_name = ''
|
|
self.bench_tag = '' # e.g. '[1/6]': persistent benchmark counter
|
|
self._last_phase = ''
|
|
self.started = 0.0
|
|
self.current_started = 0.0
|
|
self.inflight = 0
|
|
self.heartbeat_task = None
|
|
|
|
def set_overall(self, total: int, done: int, label: str = 'benches'):
|
|
"""Multi-benchmark runs: a bar above the sample bar showing which
|
|
benchmark we are on (also covers loading/scoring phases)."""
|
|
self.overall_total = total
|
|
if self.overall_id is None:
|
|
self.progress.start()
|
|
# the shared column set reads fields from EVERY task: give the
|
|
# overall task the same fields or rendering raises KeyError
|
|
self.overall_id = self.progress.add_task(
|
|
f'[cyan]{label}[/cyan]', total=total, completed=min(done, total),
|
|
inflight=0, rate='0.00', success=done, failed=0,
|
|
waiting='00:00', last_result='')
|
|
else:
|
|
self.progress.update(self.overall_id, completed=min(done, total))
|
|
|
|
def advance_overall(self):
|
|
if self.overall_id is not None:
|
|
self.progress.advance(self.overall_id)
|
|
|
|
def start(self, total: int, description: str, completed: int = 0):
|
|
self.started = time.monotonic()
|
|
if self.task_id is not None:
|
|
return # one live reporter at a time; reuse across benchmarks
|
|
self.progress.start()
|
|
self.task_id = self.progress.add_task(
|
|
f"[green]{description}",
|
|
total=total,
|
|
completed=min(completed, total),
|
|
success=completed,
|
|
failed=0,
|
|
rate="0.00",
|
|
inflight=0,
|
|
waiting="00:00",
|
|
last_result="restored",
|
|
)
|
|
self.heartbeat_task = asyncio.create_task(self._heartbeat())
|
|
|
|
def reset_samples(self, total: int, description: str, completed: int = 0):
|
|
"""Start (or re-target) the per-sample task for the next benchmark."""
|
|
self.started = time.monotonic()
|
|
self.inflight = 0
|
|
self.bench_name = description
|
|
desc = f'[green]{self.bench_tag}{description} · generating[/green]'
|
|
if self.task_id is None:
|
|
self.progress.start()
|
|
self.task_id = self.progress.add_task(
|
|
desc, total=total, completed=min(completed, total),
|
|
success=completed, failed=0, rate='0.00', inflight=0,
|
|
waiting='00:00', last_result='restored')
|
|
self.heartbeat_task = asyncio.create_task(self._heartbeat())
|
|
else:
|
|
self.progress.update(self.task_id, description=desc,
|
|
total=total, completed=min(completed, total),
|
|
success=completed, failed=0, rate='0.00',
|
|
inflight=0, last_result='restored')
|
|
|
|
def set_bench_tag(self, tag: str):
|
|
"""Persistent counter shown on the sample bar, e.g. '[1/6]'."""
|
|
self.bench_tag = tag + ' ' if tag else ''
|
|
if self._last_phase:
|
|
self.set_phase(self._last_phase)
|
|
|
|
def set_phase(self, phase: str):
|
|
self._last_phase = phase
|
|
"""Retag the sample bar with what is happening (generating/scoring/
|
|
writing) -- the bar alone does not say which stage we are in."""
|
|
if self.task_id is not None:
|
|
self.progress.update(
|
|
self.task_id,
|
|
description=f'[green]{self.bench_tag}{self.bench_name} · {phase}[/green]')
|
|
|
|
def begin_sample(self, label: str):
|
|
if self.task_id is None:
|
|
return
|
|
self.inflight += 1
|
|
self.current_started = time.monotonic()
|
|
self.progress.update(self.task_id, inflight=self.inflight,
|
|
waiting="00:00", last_result=f"waiting {label}")
|
|
|
|
def rollback(self):
|
|
"""Pair a begin_sample that will NOT reach advance (retry path):
|
|
just decrement the in-flight count, no success/fail bookkeeping."""
|
|
self.inflight = max(0, self.inflight - 1)
|
|
if self.task_id is not None:
|
|
self.progress.update(self.task_id, inflight=self.inflight)
|
|
|
|
def advance(self, success: bool = True):
|
|
if self.task_id is None:
|
|
return
|
|
task = self.progress.tasks[self.task_id]
|
|
completed = task.completed + 1
|
|
ok = task.fields["success"] + (1 if success else 0)
|
|
failed = task.fields["failed"] + (0 if success else 1)
|
|
self.inflight = max(0, self.inflight - 1)
|
|
elapsed = max(time.monotonic() - self.started, 1e-6)
|
|
self.progress.update(
|
|
self.task_id,
|
|
advance=1,
|
|
success=ok,
|
|
failed=failed,
|
|
rate=f"{completed / elapsed:.2f}",
|
|
inflight=self.inflight,
|
|
waiting="00:00",
|
|
last_result="success" if success else "failed",
|
|
)
|
|
|
|
async def _heartbeat(self):
|
|
while self.task_id is not None:
|
|
if self.task_id is not None and self.inflight:
|
|
waiting = int(time.monotonic() - self.current_started)
|
|
self.progress.update(self.task_id, waiting=f"{waiting // 60:02d}:{waiting % 60:02d}")
|
|
await asyncio.sleep(1)
|
|
|
|
def log(self, message: str):
|
|
"""Print a status line ABOVE the live bar (safe during live display).
|
|
|
|
Plain console.print from another writer while Progress is live causes
|
|
interleaved/repainted output; Progress.print routes through the live
|
|
region correctly.
|
|
"""
|
|
# highlight=False: rich auto-colorizes numbers, which made narration
|
|
# lines look inconsistent (some numbers cyan, some plain)
|
|
if self.task_id is not None:
|
|
self.progress.print(message, highlight=False)
|
|
else:
|
|
self.console.print(message, highlight=False)
|
|
|
|
def pause(self):
|
|
"""Temporarily stop the live display (e.g. while hub downloads print
|
|
their own tqdm); task state is kept and resume() restores the bars."""
|
|
if self.task_id is not None or self.overall_id is not None:
|
|
self.progress.stop()
|
|
|
|
def resume(self):
|
|
if self.task_id is not None or self.overall_id is not None:
|
|
self.progress.start()
|
|
|
|
def close(self):
|
|
if self.task_id is not None:
|
|
if self.heartbeat_task is not None:
|
|
self.heartbeat_task.cancel()
|
|
self.heartbeat_task = None
|
|
self.progress.stop()
|
|
self.task_id = None
|