guidellm.benchmark.progress

class BenchmarkerProgress(Generic[BenchmarkAccumulatorT, BenchmarkT], ABC):
    """
    Abstract interface for tracking and displaying benchmark execution progress.

    Provides lifecycle hooks for monitoring benchmark stages including initialization,
    execution start, progress updates, completion, and finalization. Implementations
    handle display updates, progress tracking, and resource management for benchmark
    monitoring.
    """

    def __init__(self):
        """Initialize progress tracker with default state."""
        self.profile: Profile | None = None
        self.current_strategy: SchedulingStrategy | None = None

    @abstractmethod
    async def on_initialize(self, profile: Profile):
        """
        Initialize progress tracking for the given benchmark profile.

        :param profile: Benchmark profile configuration defining execution parameters
        """

    @abstractmethod
    async def on_benchmark_start(self, strategy: SchedulingStrategy):
        """
        Handle benchmark strategy execution start event.

        :param strategy: Scheduling strategy configuration being executed
        """

    @abstractmethod
    async def on_benchmark_update(
        self, accumulator: BenchmarkAccumulatorT, scheduler_state: SchedulerState
    ):
        """
        Handle benchmark execution progress update with current metrics.

        :param accumulator: Current accumulated benchmark metrics and statistics
        :param scheduler_state: Current scheduler execution state and counters
        """

    @abstractmethod
    async def on_benchmark_complete(self, benchmark: BenchmarkT):
        """
        Handle benchmark strategy execution completion event.

        :param benchmark: Completed benchmark results with final metrics
        """

    @abstractmethod
    async def on_finalize(self):
        """Finalize progress tracking and release associated resources."""

class GenerativeConsoleBenchmarkerProgress(
    BenchmarkerProgress[GenerativeBenchmarkAccumulator, GenerativeBenchmark], Live
):
    """
    Console-based real-time progress display for generative benchmarks.

    Renders live benchmark execution statistics using Rich library components with
    structured progress bars, timing information, request/token metrics, and optional
    scheduler statistics. Updates refresh automatically during benchmark execution.

    :cvar display_scheduler_stats: Whether to include scheduler statistics in display
    """

    def __init__(self, display_scheduler_stats: bool = False):
        """
        Initialize console progress display with rendering configuration.

        :param display_scheduler_stats: Whether to display scheduler timing statistics
        """
        super().__init__()
        Live.__init__(
            self,
            refresh_per_second=4,
            auto_refresh=True,
            redirect_stdout=True,
            redirect_stderr=True,
        )
        self.display_scheduler_stats: bool = display_scheduler_stats
        self.run_progress: Progress | None = None
        self.run_progress_task: TaskID | None = None
        self.tasks_progress: _GenerativeProgressTasks | None = None

    async def on_initialize(self, profile: Profile):
        """
        Initialize console display components and begin live rendering.

        :param profile: Benchmark profile configuration defining execution parameters
        """
        self.tasks_progress = _GenerativeProgressTasks(
            profile=profile, display_scheduler_stats=self.display_scheduler_stats
        )
        self.run_progress = Progress(
            TextColumn("Generating...", style=f"italic {Colors.progress}"),
            BarColumn(
                bar_width=None,
                complete_style=Colors.progress,
                finished_style=Colors.success,
            ),
            TextColumn(
                "({task.fields[completed_benchmarks]}/{task.fields[total_benchmarks]})",
                style=Colors.progress,
            ),
            TextColumn("["),
            TimeElapsedColumn(),
            TextColumn("<"),
            TimeRemainingColumn(),
            TextColumn("]"),
        )
        self.run_progress_task = self.run_progress.add_task("")
        self._sync_run_progress()
        self.update(
            Group(
                Panel(
                    self.tasks_progress,
                    title="Benchmarks",
                    title_align="left",
                    expand=True,
                ),
                self.run_progress,
            )
        )
        self.start()

    async def on_benchmark_start(self, strategy: SchedulingStrategy):
        """
        Update display for benchmark strategy execution start.

        :param strategy: Scheduling strategy configuration being executed
        """
        if self.tasks_progress is not None:
            self.tasks_progress.start_benchmark(strategy)
            self._sync_run_progress()

    async def on_benchmark_update(
        self,
        accumulator: GenerativeBenchmarkAccumulator,
        scheduler_state: SchedulerState,
    ):
        """
        Update display with current benchmark progress and metrics.

        :param accumulator: Current accumulated benchmark metrics and statistics
        :param scheduler_state: Current scheduler execution state and counters
        """
        if self.tasks_progress is not None:
            self.tasks_progress.update_benchmark(accumulator, scheduler_state)
            self._sync_run_progress()

    async def on_benchmark_complete(self, benchmark: GenerativeBenchmark):
        """
        Update display for completed benchmark strategy.

        :param benchmark: Completed benchmark results with final metrics
        """
        if self.tasks_progress is not None:
            self.tasks_progress.complete_benchmark(benchmark)
            self._sync_run_progress()

    async def on_finalize(self):
        """Stop display rendering and release resources."""
        if self.tasks_progress is not None:
            self.tasks_progress.finalize()
            self._sync_run_progress()
        if self.run_progress is not None and self.run_progress_task is not None:
            self.run_progress.stop_task(self.run_progress_task)
        self.stop()
        self.run_progress = None
        self.run_progress_task = None
        self.tasks_progress = None

    def _sync_run_progress(self):
        """Synchronize overall progress display with task progress."""
        if (
            self.run_progress is not None
            and self.run_progress_task is not None
            and self.tasks_progress is not None
        ):
            self.run_progress.update(
                self.run_progress_task,
                total=self.tasks_progress.steps_total,
                completed=self.tasks_progress.steps_progress,
                completed_benchmarks=self.tasks_progress.tasks_progress,
                total_benchmarks=self.tasks_progress.tasks_total,
            )