guidellm 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

guidellm/benchmark/progress.py

@@ -0,0 +1,720 @@
+import math
+import time
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Generic, Optional, TypeVar, Union
+
+from rich.console import Group
+from rich.live import Live
+from rich.panel import Panel
+from rich.progress import (
+    BarColumn,
+    Progress,
+    ProgressColumn,
+    SpinnerColumn,
+    TaskID,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+
+from guidellm.benchmark.aggregator import (
+    BenchmarkAggregator,
+    GenerativeBenchmarkAggregator,
+)
+from guidellm.benchmark.benchmark import Benchmark, GenerativeBenchmark
+from guidellm.benchmark.benchmarker import BenchmarkerResult
+from guidellm.scheduler import (
+    SchedulingStrategy,
+    StrategyType,
+    strategy_display_str,
+)
+from guidellm.utils import Colors
+
+__all__ = [
+    "BenchmarkerTaskProgressState",
+    "BenchmarkerProgressDisplay",
+    "GenerativeTextBenchmarkerTaskProgressState",
+    "GenerativeTextBenchmarkerProgressDisplay",
+]
+
+
+@dataclass
+class BenchmarkerTaskProgressState:
+    display_scheduler_stats: bool
+
+    task_id: TaskID
+    strategy: Union[StrategyType, SchedulingStrategy]
+    started: bool = False
+    compiling: bool = False
+    ended: bool = False
+
+    start_time: Optional[float] = None
+    max_number: Optional[float] = None
+    max_duration: Optional[float] = None
+    in_warmup: bool = False
+    in_cooldown: bool = False
+
+    requests_rate: float = 0
+    request_latency: float = 0
+    requests_processing: float = 0
+    requests_successful: float = 0
+    requests_incomplete: float = 0
+    requests_errored: float = 0
+
+    worker_overheads_time_ms: float = 0.0
+    backend_overheads_time_ms: float = 0.0
+    requests_sleep_time_ms: float = 0.0
+    requests_targeted_start_time_delay_ms: float = 0.0
+
+    @property
+    def description(self) -> str:
+        return strategy_display_str(self.strategy)
+
+    @property
+    def total(self) -> Optional[float]:
+        if self.max_number is None and self.max_duration is None:
+            return None
+
+        return 1000
+
+    @property
+    def completed(self) -> int:
+        if self.ended:
+            return 1000
+
+        if self.max_number is None and self.max_duration is None:
+            return 0
+
+        number = self.requests_successful + self.requests_errored
+        number_percent = (
+            number / float(self.max_number) * 1000 if self.max_number else -math.inf
+        )
+        duration_percent = (
+            (time.time() - self.start_time) / self.max_duration * 1000
+            if self.max_duration and self.start_time
+            else -math.inf
+        )
+
+        return min(int(max(number_percent, duration_percent)), 1000)
+
+    @property
+    def fields(self) -> dict[str, str]:
+        fields = {
+            "start_time": self.formatted_start_time,
+            "progress_status": self.formatted_progress_status,
+            "requests_summary": self.formatted_requests_summary,
+        }
+
+        if self.display_scheduler_stats:
+            fields["scheduler_stats"] = self.formatted_scheduler_stats
+
+        return fields
+
+    @property
+    def formatted_start_time(self) -> str:
+        if self.start_time is None:
+            return "--:--:--"
+
+        return datetime.fromtimestamp(self.start_time).strftime("%H:%M:%S")
+
+    @property
+    def formatted_progress_status(self) -> str:
+        if self.ended:
+            status = "complete"
+            color = Colors.SUCCESS
+        elif self.compiling:
+            status = "compiling"
+            color = Colors.PROGRESS
+        elif self.started and self.in_warmup:
+            status = "warmup"
+            color = Colors.PROGRESS
+        elif self.started and self.in_cooldown:
+            status = "cooldown"
+            color = Colors.PROGRESS
+        elif self.started:
+            status = "running"
+            color = Colors.PROGRESS
+        else:
+            status = "pending"
+            color = Colors.INFO
+
+        return f"[{color}]{status.ljust(8)}[/{color}]"
+
+    @property
+    def formatted_requests_summary(self) -> str:
+        if not self.started:
+            return " "
+
+        return (
+            f"[{Colors.INFO}]Req:[/{Colors.INFO}] "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.requests_rate,
+                label="req/s",
+                total_characters=12,
+                digits_places=4,
+                decimal_places=1,
+            )
+            + ", "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.request_latency,
+                label="Lat",
+                units="s",
+                total_characters=12,
+                digits_places=4,
+                decimal_places=2,
+            )
+            + ", "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.requests_processing,
+                label="Conc",
+                total_characters=12,
+                digits_places=4,
+                decimal_places=1,
+            )
+            + ", "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.requests_successful,
+                label="Comp",
+                total_characters=12,
+                digits_places=5,
+                decimal_places=0,
+            )
+            + ", "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.requests_incomplete,
+                label="Inc",
+                total_characters=12,
+                digits_places=5,
+                decimal_places=0,
+            )
+            + ", "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.requests_errored,
+                label="Err",
+                total_characters=12,
+                digits_places=5,
+                decimal_places=0,
+            )
+        )
+
+    @property
+    def formatted_scheduler_stats(self) -> str:
+        if not self.started:
+            return " "
+
+        return (
+            f"[{Colors.INFO}]Sys:[/{Colors.INFO}] "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.worker_overheads_time_ms,
+                label="Work OH",
+                units="ms",
+                total_characters=18,
+                digits_places=3,
+                decimal_places=1,
+            )
+            + ", "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.backend_overheads_time_ms,
+                label="Back OH",
+                units="ms",
+                total_characters=18,
+                digits_places=3,
+                decimal_places=1,
+            )
+            + ", "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.requests_sleep_time_ms,
+                label="Req Sleep",
+                units="ms",
+                total_characters=18,
+                digits_places=5,
+                decimal_places=0,
+            )
+            + ", "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.requests_targeted_start_time_delay_ms,
+                label="Start Del",
+                units="ms",
+                total_characters=18,
+                digits_places=5,
+                decimal_places=0,
+            )
+        )
+
+    @staticmethod
+    def format_progress_display(
+        value: float,
+        label: str,
+        units: str = "",
+        total_characters: Optional[int] = None,
+        digits_places: Optional[int] = None,
+        decimal_places: Optional[int] = None,
+    ) -> str:
+        if decimal_places is None and digits_places is None:
+            formatted_number = f"{value}:.0f"
+        elif digits_places is None:
+            formatted_number = f"{value:.{decimal_places}f}"
+        elif decimal_places is None:
+            formatted_number = f"{value:>{digits_places}f}"
+        else:
+            formatted_number = f"{value:>{digits_places}.{decimal_places}f}"
+
+        result = f"{formatted_number}{units} [{Colors.INFO}]{label}[/{Colors.INFO}]"
+
+        if total_characters is not None:
+            total_characters += len(Colors.INFO) * 2 + 5
+
+            if len(result) < total_characters:
+                result = result.rjust(total_characters)
+
+        return result
+
+
+class GenerativeTextBenchmarkerTaskProgressState(BenchmarkerTaskProgressState):
+    output_tokens: float = 0
+    prompt_tokens: float = 0
+    output_tokens_rate: float = 0
+    total_tokens_rate: float = 0
+    tokens_ttft: float = 0
+    tokens_itl: float = 0
+
+    @property
+    def fields(self) -> dict[str, str]:
+        fields = super().fields
+        fields["tokens_summary"] = self.formatted_tokens_summary
+        return fields
+
+    @property
+    def formatted_tokens_summary(self) -> str:
+        if not self.started:
+            return " "
+
+        return (
+            f"[{Colors.INFO}]Tok:[/{Colors.INFO}] "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.output_tokens_rate,
+                label="gen/s",
+                total_characters=12,
+                digits_places=4,
+                decimal_places=1,
+            )
+            + ", "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.total_tokens_rate,
+                label="tot/s",
+                total_characters=12,
+                digits_places=4,
+                decimal_places=1,
+            )
+            + ", "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.tokens_ttft,
+                label="TTFT",
+                units="ms",
+                total_characters=12,
+                digits_places=3,
+                decimal_places=1,
+            )
+            + ", "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.tokens_itl,
+                label="ITL",
+                units="ms",
+                total_characters=12,
+                digits_places=3,
+                decimal_places=1,
+            )
+            + ", "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.prompt_tokens,
+                label="Prompt",
+                total_characters=12,
+                digits_places=4,
+                decimal_places=0,
+            )
+            + ", "
+            + BenchmarkerTaskProgressState.format_progress_display(
+                value=self.output_tokens,
+                label="Gen",
+                total_characters=12,
+                digits_places=4,
+                decimal_places=0,
+            )
+        )
+
+
+BTPS = TypeVar("BTPS", bound=BenchmarkerTaskProgressState)
+
+
+class BenchmarkerProgressDisplay(Generic[BTPS]):
+    def __init__(self, display_scheduler_stats: bool):
+        self.display_scheduler_stats = display_scheduler_stats
+        self.started = False
+        self.benchmarker_tasks_progress = Progress(*self.create_task_progress_columns())
+        self.benchmarker_tasks_panel = Panel(
+            self.benchmarker_tasks_progress,
+            title="Benchmarks",
+            title_align="left",
+            expand=True,
+        )
+        self.benchmarker_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.benchmarker_live = Live(
+            Group(
+                self.benchmarker_tasks_panel,
+                self.benchmarker_progress,
+            ),
+            redirect_stdout=True,
+            redirect_stderr=True,
+        )
+        self.active_task: Optional[TaskID] = None
+        self.benchmarker_tasks: list[BTPS] = []
+        self.progress_task: Optional[TaskID] = None
+
+    def update(self, result: BenchmarkerResult):
+        if result.type_ == "run_start":
+            if self.started:
+                raise RuntimeError("Progress display already started.")
+
+            self.handle_start(result)
+            self.started = True
+        elif result.type_ == "run_complete":
+            if not self.started:
+                raise RuntimeError("Progress display not started.")
+
+            self.handle_end(result)
+            self.started = False
+        else:
+            if not self.started:
+                raise RuntimeError("Progress display not started.")
+
+            self.handle_update(result)
+
+    def handle_start(self, result: BenchmarkerResult):
+        self.benchmarker_live.start()
+
+        for index, strategy_type in enumerate(result.profile.strategy_types):
+            task_id = self.benchmarker_tasks_progress.add_task(
+                description=strategy_type,
+                start=False,
+                total=None,
+                completed=0,
+                visible=False,
+            )
+            task_progress_state = self.create_task_progress_state(
+                task_id=task_id,
+                index=index,
+                strategy_type=strategy_type,
+                result=result,
+            )
+            self.benchmarker_tasks.append(task_progress_state)
+            self.benchmarker_tasks_progress.update(
+                task_id,
+                description=task_progress_state.description,
+                visible=True,
+                **task_progress_state.fields,  # type: ignore[arg-type]
+            )
+
+        self.progress_task = self.benchmarker_progress.add_task(
+            "",
+            total=len(self.benchmarker_tasks) * 1000,
+            completed_benchmarks=0,
+            total_benchmarks=len(self.benchmarker_tasks),
+        )
+
+    def handle_update(self, result: BenchmarkerResult):
+        current_state: BTPS = self.benchmarker_tasks[result.current_index]
+
+        if result.type_ == "scheduler_start":
+            self.handle_update_scheduler_start(current_state, result)
+            self.active_task = current_state.task_id
+        elif result.type_ == "scheduler_update":
+            self.handle_update_scheduler_update(current_state, result)
+        elif result.type_ == "scheduler_complete":
+            self.handle_update_scheduler_complete(current_state, result)
+        elif result.type_ == "benchmark_compiled":
+            self.handle_update_benchmark_compiled(current_state, result)
+        else:
+            raise ValueError(f"Unknown result type: {result.type_}")
+
+        if self.progress_task is None:
+            raise RuntimeError("Progress task not set.")
+
+        self.benchmarker_tasks_progress.update(
+            current_state.task_id,
+            description=current_state.description,
+            completed=current_state.completed,
+            total=current_state.total,
+            **current_state.fields,  # type: ignore[arg-type]
+        )
+        self.benchmarker_progress.update(
+            self.progress_task,
+            completed=(result.current_index * 1000) + current_state.completed,
+            total=1000 * len(self.benchmarker_tasks),
+            completed_benchmarks=(
+                result.current_index + (1 if current_state.ended else 0)
+            ),
+            total_benchmarks=len(self.benchmarker_tasks),
+        )
+
+        if current_state.ended:
+            self.benchmarker_tasks_progress.stop_task(current_state.task_id)
+            self.active_task = None
+
+    def handle_update_scheduler_start(
+        self, progress_state: BTPS, result: BenchmarkerResult
+    ):
+        if self.active_task is not None:
+            raise RuntimeError("Active task already set.")
+
+        progress_state.strategy = result.current_strategy  # type: ignore[assignment]
+        progress_state.started = True
+        current_aggregator: BenchmarkAggregator = result.current_aggregator  # type: ignore[assignment]
+        progress_state.start_time = (
+            current_aggregator.requests_stats.totals.total.start_time
+        )
+        progress_state.max_number = current_aggregator.args.max_number
+        progress_state.max_duration = current_aggregator.args.max_duration
+
+    def handle_update_scheduler_update(
+        self, progress_state: BTPS, result: BenchmarkerResult
+    ):
+        if self.active_task is None:
+            raise RuntimeError("Active task not set.")
+
+        if self.active_task != progress_state.task_id:
+            raise RuntimeError("Active task does not match current task.")
+
+        current_aggregator: BenchmarkAggregator = result.current_aggregator  # type: ignore[assignment]
+        progress_state.in_warmup = current_aggregator.in_warmup
+        progress_state.in_cooldown = current_aggregator.in_cooldown
+        progress_state.requests_rate = (
+            current_aggregator.requests_stats.totals.successful.rate
+        )
+        progress_state.request_latency = (
+            current_aggregator.requests_stats.request_time.mean
+        )
+        progress_state.requests_processing = (
+            current_aggregator.scheduler_stats.processing_requests.last
+        )
+        progress_state.requests_successful = (
+            current_aggregator.requests_stats.totals.successful.total
+        )
+        progress_state.requests_incomplete = (
+            current_aggregator.requests_stats.totals.incomplete.total
+        )
+        progress_state.requests_errored = (
+            current_aggregator.requests_stats.totals.errored.total
+        )
+        progress_state.worker_overheads_time_ms = (
+            current_aggregator.requests_stats.scheduled_time_delay.mean_ms
+            + current_aggregator.requests_stats.worker_start_delay.mean_ms
+        )
+        progress_state.backend_overheads_time_ms = (
+            current_aggregator.requests_stats.request_time_delay.mean_ms
+        )
+        progress_state.requests_sleep_time_ms = (
+            current_aggregator.requests_stats.scheduled_time_sleep.mean_ms
+        )
+        progress_state.requests_targeted_start_time_delay_ms = (
+            current_aggregator.requests_stats.request_start_time_targeted_delay.mean_ms
+        )
+
+    def handle_update_scheduler_complete(
+        self,
+        progress_state: BTPS,
+        result: BenchmarkerResult,  # noqa: ARG002
+    ):
+        if self.active_task is None:
+            raise RuntimeError("Active task not set.")
+
+        if self.active_task != progress_state.task_id:
+            raise RuntimeError("Active task does not match current task.")
+
+        progress_state.in_warmup = False
+        progress_state.in_cooldown = False
+        progress_state.compiling = True
+
+    def handle_update_benchmark_compiled(
+        self, progress_state: BTPS, result: BenchmarkerResult
+    ):
+        if self.active_task is None:
+            raise RuntimeError("Active task not set.")
+
+        if self.active_task != progress_state.task_id:
+            raise RuntimeError("Active task does not match current task.")
+
+        current_benchmark: Benchmark = result.current_benchmark  # type: ignore[assignment]
+        progress_state.compiling = False
+        progress_state.ended = True
+        progress_state.requests_rate = (
+            current_benchmark.metrics.requests_per_second.successful.mean
+        )
+        progress_state.requests_processing = (
+            current_benchmark.metrics.request_concurrency.successful.mean
+        )
+
+    def handle_end(self, result: BenchmarkerResult):  # noqa: ARG002
+        if self.progress_task is None:
+            raise RuntimeError("Progress task not set.")
+
+        self.benchmarker_progress.update(
+            self.progress_task,
+            completed=len(self.benchmarker_tasks) * 1000,
+            total=len(self.benchmarker_tasks) * 1000,
+            completed_benchmarks=len(self.benchmarker_tasks),
+            total_benchmarks=len(self.benchmarker_tasks),
+        )
+        self.benchmarker_progress.stop_task(self.progress_task)
+        self.benchmarker_live.stop()
+        self.active_task = None
+        self.benchmarker_tasks = []
+        self.progress_task = None
+
+    def create_task_progress_columns(self) -> list[ProgressColumn]:
+        columns = [
+            TextColumn("[{task.fields[start_time]}]"),
+            SpinnerColumn(style=Colors.PROGRESS),
+            TaskProgressColumn(style=Colors.PROGRESS),
+            TextColumn("{task.description}"),
+            TextColumn("({task.fields[progress_status]})"),
+            TextColumn(" "),
+        ]
+
+        if not self.display_scheduler_stats:
+            columns += [
+                TextColumn("{task.fields[requests_summary]}\n"),
+            ]
+        else:
+            columns += [
+                TextColumn(
+                    "{task.fields[requests_summary]}\n{task.fields[scheduler_stats]}\n"
+                ),
+            ]
+
+        return columns
+
+    def create_task_progress_state(
+        self,
+        task_id: TaskID,
+        index: int,  # noqa: ARG002
+        strategy_type: StrategyType,
+        result: BenchmarkerResult,  # noqa: ARG002
+    ) -> BTPS:
+        return BenchmarkerTaskProgressState(  # type: ignore[return-value]
+            display_scheduler_stats=self.display_scheduler_stats,
+            task_id=task_id,
+            strategy=strategy_type,
+        )
+
+
+class GenerativeTextBenchmarkerProgressDisplay(
+    BenchmarkerProgressDisplay[GenerativeTextBenchmarkerTaskProgressState]
+):
+    def handle_update_scheduler_update(
+        self,
+        progress_state: GenerativeTextBenchmarkerTaskProgressState,
+        result: BenchmarkerResult,
+    ):
+        super().handle_update_scheduler_update(progress_state, result)
+        current_aggregator: GenerativeBenchmarkAggregator = result.current_aggregator  # type: ignore[assignment]
+        progress_state.output_tokens = (
+            current_aggregator.requests_stats.output_tokens.mean
+        )
+        progress_state.prompt_tokens = (
+            current_aggregator.requests_stats.prompt_tokens.mean
+        )
+        progress_state.output_tokens_rate = (
+            current_aggregator.requests_stats.output_tokens.rate
+        )
+        progress_state.total_tokens_rate = (
+            current_aggregator.requests_stats.total_tokens.rate
+        )
+        progress_state.tokens_ttft = (
+            current_aggregator.requests_stats.time_to_first_token.mean_ms
+        )
+        progress_state.tokens_itl = (
+            current_aggregator.requests_stats.inter_token_latency.mean_ms
+        )
+
+    def handle_update_benchmark_compiled(
+        self,
+        progress_state: GenerativeTextBenchmarkerTaskProgressState,
+        result: BenchmarkerResult,
+    ):
+        super().handle_update_benchmark_compiled(progress_state, result)
+
+        current_benchmark: GenerativeBenchmark = result.current_benchmark  # type: ignore[assignment]
+        progress_state.request_latency = (
+            current_benchmark.metrics.request_latency.successful.mean
+        )
+        progress_state.requests_successful = current_benchmark.request_totals.successful
+        progress_state.requests_errored = current_benchmark.request_totals.errored
+        progress_state.requests_incomplete = current_benchmark.request_totals.incomplete
+        progress_state.output_tokens = (
+            current_benchmark.metrics.output_token_count.successful.mean
+        )
+        progress_state.prompt_tokens = (
+            current_benchmark.metrics.prompt_token_count.successful.mean
+        )
+        progress_state.output_tokens_rate = (
+            current_benchmark.metrics.output_tokens_per_second.successful.mean
+        )
+        progress_state.total_tokens_rate = (
+            current_benchmark.metrics.tokens_per_second.successful.mean
+        )
+        progress_state.tokens_ttft = (
+            current_benchmark.metrics.time_to_first_token_ms.successful.mean
+        )
+        progress_state.tokens_itl = (
+            current_benchmark.metrics.inter_token_latency_ms.successful.mean
+        )
+
+    def create_task_progress_state(
+        self,
+        task_id: TaskID,
+        index: int,  # noqa: ARG002
+        strategy_type: StrategyType,
+        result: BenchmarkerResult,  # noqa: ARG002
+    ) -> GenerativeTextBenchmarkerTaskProgressState:
+        return GenerativeTextBenchmarkerTaskProgressState(
+            display_scheduler_stats=self.display_scheduler_stats,
+            task_id=task_id,
+            strategy=strategy_type,
+        )
+
+    def create_task_progress_columns(self) -> list[ProgressColumn]:
+        columns = super().create_task_progress_columns()
+        columns = columns[:-1]  # remove the last display info column
+
+        if not self.display_scheduler_stats:
+            columns += [
+                TextColumn(
+                    "{task.fields[requests_summary]}\n{task.fields[tokens_summary]}",
+                ),
+            ]
+        else:
+            columns += [
+                TextColumn(
+                    "{task.fields[requests_summary]}\n{task.fields[tokens_summary]}\n{task.fields[scheduler_stats]}",
+                ),
+            ]
+
+        return columns