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 CompositeBenchmarkerProgress(
    BenchmarkerProgress[BenchmarkAccumulatorT, BenchmarkT]
):
    """
    Composite progress tracker that aggregates multiple progress trackers.

    Delegates lifecycle events to all registered progress trackers, enabling
    simultaneous monitoring and display of benchmark execution across different
    interfaces or formats.
    """

    def __init__(
        self, trackers: list[BenchmarkerProgress[BenchmarkAccumulatorT, BenchmarkT]]
    ):
        """
        Initialize composite progress tracker with a list of individual trackers.

        :param progress_trackers: Individual progress trackers to aggregate
        """
        super().__init__()
        self.trackers = trackers

    @staticmethod
    async def _notify_progress(*callbacks: Awaitable) -> None:
        """Finish every callback before propagating the first lifecycle failure."""
        results = await asyncio.gather(*callbacks, return_exceptions=True)
        for result in results:
            if isinstance(result, BaseException):
                raise result

    async def on_initialize(self, profile: Profile):
        await self._notify_progress(
            *(tracker.on_initialize(profile) for tracker in self.trackers)
        )

    async def on_benchmark_start(self, strategy: SchedulingStrategy):
        await self._notify_progress(
            *(tracker.on_benchmark_start(strategy) for tracker in self.trackers)
        )

    async def on_benchmark_update(
        self, accumulator: BenchmarkAccumulatorT, scheduler_state: SchedulerState
    ):
        await self._notify_progress(
            *(
                tracker.on_benchmark_update(accumulator, scheduler_state)
                for tracker in self.trackers
            )
        )

    async def on_benchmark_complete(self, benchmark: BenchmarkT):
        await self._notify_progress(
            *(tracker.on_benchmark_complete(benchmark) for tracker in self.trackers)
        )

    async def on_finalize(self):
        await self._notify_progress(
            *(tracker.on_finalize() for tracker in self.trackers)
        )

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, cleanup: bool = True):
        """
        Initialize console progress display with rendering configuration.

        :param display_scheduler_stats: Whether to display scheduler timing statistics
        :param cleanup: Whether to clear the console display after completion
        """
        super().__init__()
        Live.__init__(
            self,
            refresh_per_second=4,
            auto_refresh=True,
            redirect_stdout=True,
            # Only redirect stderr if it would interfere with stdout display
            redirect_stderr=stderr_eq_stdout(),
            transient=cleanup,
        )
        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("]"),
            auto_refresh=False,
        )
        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,
            )

class GenerativeLoggingBenchmarkerProgress(
    BenchmarkerProgress[GenerativeBenchmarkAccumulator, GenerativeBenchmark]
):
    """Log interval-limited progress independently of other progress trackers."""

    def __init__(self, interval: float = 10.0):
        """
        Initialize independent progress logging.

        :param interval: Positive finite minimum seconds between periodic records
        :raises ValueError: If the interval is not positive and finite
        """
        super().__init__()
        if not isfinite(interval) or interval <= 0:
            raise ValueError("Progress log interval must be positive and finite")
        self.interval = interval
        self._state: _GenerativeProgressTaskState | None = None
        self._index = 0
        self._started_at = 0.0
        self._last_update = 0.0

    async def on_initialize(self, profile: Profile):
        """
        Reset logging for a new run.

        :param profile: Benchmark profile
        """
        self.profile = profile
        self._index = 0
        self._state = None

    async def on_benchmark_start(self, strategy: SchedulingStrategy):
        """
        Log strategy start immediately.

        :param strategy: Strategy being executed
        """
        self._index += 1
        self._state = _GenerativeProgressTaskState(strategy_type=strategy.type_)
        self._state.start(strategy)
        self._started_at = monotonic()
        self._log_update("started")

    async def on_benchmark_update(
        self,
        accumulator: GenerativeBenchmarkAccumulator,
        scheduler_state: SchedulerState,
    ):
        """
        Periodically log current metrics.

        :param accumulator: Accumulated benchmark metrics
        :param scheduler_state: Scheduler counters and progress
        """
        if self._state and monotonic() - self._last_update >= self.interval:
            self._state.update(accumulator, scheduler_state)
            self._log_update(self._state.benchmark_status)

    async def on_benchmark_complete(self, benchmark: GenerativeBenchmark):
        """
        Log final metrics regardless of the interval.

        :param benchmark: Completed result
        """
        if self._state:
            self._state.complete(benchmark)
            self._log_update("completed")
            self._state = None

    async def on_finalize(self):
        """Release progress state."""
        self._state = None

    def _log_update(self, status: str):
        if self._state is None:
            return
        state = self._state
        now = monotonic()
        logger.bind(benchmark_index=self._index, progress_status=status).info(
            "Benchmark {} ({}): {} | elapsed={:.1f}s | "
            "successful={} errored={} incomplete={} | "
            "requests/s={:.2f} output_tokens/s={:.2f}",
            self._index,
            state.strategy,
            status,
            now - self._started_at,
            state.successful_requests,
            state.errored_requests,
            state.cancelled_requests,
            state.requests_per_second,
            state.output_tokens_rate,
        )
        self._last_update = now