guidellm.benchmark.schemas

class Benchmark(StandardBaseDict, ABC, Generic[BenchmarkAccumulatorT]):
    """
    Compile and expose final benchmark execution metrics.

    Defines the interface for benchmark result implementations capturing
    comprehensive performance metrics including latency distributions, throughput
    measurements, and concurrency patterns. Subclasses implement compilation
    logic to transform accumulated metrics and scheduler state into structured
    results with statistical summaries.
    """

    @property
    @abstractmethod
    def start_time(self) -> float:
        """
        :return: Benchmark start timestamp in seconds since epoch
        """

    @property
    @abstractmethod
    def end_time(self) -> float:
        """
        :return: Benchmark completion timestamp in seconds since epoch
        """

    @property
    @abstractmethod
    def duration(self) -> float:
        """
        :return: Benchmark execution duration in seconds
        """

    @property
    @abstractmethod
    def request_latency(self) -> StatusDistributionSummary:
        """
        :return: Statistical distribution of request latencies
        """

    @property
    @abstractmethod
    def request_throughput(self) -> StatusDistributionSummary:
        """
        :return: Statistical distribution of throughput measurements
        """

    @property
    @abstractmethod
    def request_concurrency(self) -> StatusDistributionSummary:
        """
        :return: Statistical distribution of concurrent request counts
        """

    @classmethod
    @abstractmethod
    def compile(
        cls, accumulator: BenchmarkAccumulatorT, scheduler_state: SchedulerState
    ) -> Any:
        """
        Transform accumulated metrics into final benchmark results.

        :param accumulator: Accumulator instance with collected metrics and state
        :param scheduler_state: Scheduler's final state after execution completion
        :return: Compiled benchmark instance with complete statistical results
        """

class BenchmarkAccumulator(StandardBaseDict, ABC, Generic[RequestT, ResponseT]):
    """
    Track and accumulate benchmark metrics during scheduler execution.

    Maintains incremental metric estimates as requests are processed, enabling
    real-time progress monitoring and efficient metric compilation. Subclasses
    implement specific metric calculation strategies based on request/response
    characteristics and scheduler state evolution.
    """

    config: BenchmarkConfig = Field(
        description="Benchmark execution configuration and constraints",
    )

    @abstractmethod
    def update_estimate(
        self,
        response: ResponseT | None,
        request: RequestT,
        info: RequestInfo,
        scheduler_state: SchedulerState,
    ):
        """
        Incrementally update metrics with completed request data.

        :param response: Backend response data if request succeeded
        :param request: Request instance submitted to backend
        :param info: Request timing, status, and execution metadata
        :param scheduler_state: Current scheduler state with queue and concurrency info
        """
        ...

class BenchmarkConfig(StandardBaseDict):
    """
    Encapsulate execution parameters and constraints for benchmark runs.

    Defines comprehensive configuration including scheduler strategy, constraint
    sets, transient phase handling, metric sampling preferences, and execution
    metadata. Coordinates profile, request, backend, and environment configurations
    to enable reproducible benchmark execution with precise control over metric
    collection.
    """

    id_: str = Field(
        default_factory=lambda: str(uuid.uuid4()),
        description="Unique identifier for this benchmark execution",
    )
    run_id: str = Field(
        description="Identifier grouping related benchmark runs in a series",
    )
    run_index: int = Field(
        description="Zero-based index of this run within the benchmark series",
    )
    strategy: SchedulingStrategy = Field(
        description="Scheduler strategy controlling request execution patterns",
    )
    constraints: dict[str, dict[str, Any]] = Field(
        description="Constraint definitions applied to scheduler strategy execution",
    )
    sample_size: int | None = Field(
        default=None,
        description=(
            "Maximum number of requests per status group to retain full data for. "
            "None keeps all, 0 strips all, N > 0 uses reservoir sampling."
        ),
    )
    warmup: TransientPhaseConfig = Field(
        default_factory=TransientPhaseConfig,
        description="Warmup phase configuration excluding initial transient period",
    )
    cooldown: TransientPhaseConfig = Field(
        default_factory=TransientPhaseConfig,
        description="Cooldown phase configuration excluding final transient period",
    )
    prefer_response_metrics: bool = Field(
        default=True,
        description="Prioritize response-based metrics over request-based metrics",
    )
    slo: GoodputSLO | None = Field(
        default=None,
        description=(
            "Per-request latency objectives defining which requests count "
            "toward goodput. None disables goodput measurement"
        ),
    )
    confidence: float | None = Field(
        default=0.95,
        gt=0.0,
        lt=1.0,
        description=(
            "Two-sided confidence level for the intervals reported alongside "
            "request-level metrics. None reports those metrics without intervals"
        ),
    )
    warnings: list[WarningRuleArgs] = Field(
        default_factory=default_warning_rules,
        description="Post-benchmark warning checks and their thresholds",
    )
    profile: dict[str, Any] = Field(
        description="Profile instance coordinating multi-strategy execution",
    )
    requests: dict[str, Any] = Field(
        description="Request generation configuration and dataset metadata",
    )
    backend: dict[str, Any] = Field(
        description="Backend connection parameters and service configuration",
    )
    environment: dict[str, Any] = Field(
        description="Execution environment details and system metadata",
    )

class GenerativeAudioMetricsSummary(StandardBaseDict):
    """
    Audio-specific metric summaries for generative benchmarks.

    Tracks token, sample count, duration, and byte-level metrics across input,
    output, and total usage for audio generation workloads.
    """

    tokens: GenerativeMetricsSummary | None = Field(
        description="Audio token count metrics and distributions"
    )
    samples: GenerativeMetricsSummary | None = Field(
        description="Sample count metrics and distributions"
    )
    seconds: GenerativeMetricsSummary | None = Field(
        description="Duration metrics in seconds and distributions"
    )
    bytes: GenerativeMetricsSummary | None = Field(
        description="Byte size metrics and distributions"
    )

    @classmethod
    def compile(
        cls,
        successful: list[GenerativeRequestStats],
        incomplete: list[GenerativeRequestStats],
        errored: list[GenerativeRequestStats],
    ) -> GenerativeAudioMetricsSummary:
        """
        Compile audio metrics summary from request statistics.

        :param successful: Successfully completed request statistics
        :param incomplete: Incomplete/cancelled request statistics
        :param errored: Failed request statistics
        :return: Compiled audio metrics summary
        """
        return GenerativeAudioMetricsSummary(
            tokens=GenerativeMetricsSummary.compile(
                property_name="audio_tokens",
                successful=successful,
                incomplete=incomplete,
                errored=errored,
            ),
            samples=GenerativeMetricsSummary.compile(
                property_name="audio_samples",
                successful=successful,
                incomplete=incomplete,
                errored=errored,
            ),
            seconds=GenerativeMetricsSummary.compile(
                property_name="audio_seconds",
                successful=successful,
                incomplete=incomplete,
                errored=errored,
            ),
            bytes=GenerativeMetricsSummary.compile(
                property_name="audio_bytes",
                successful=successful,
                incomplete=incomplete,
                errored=errored,
            ),
        )

class GenerativeBenchmark(Benchmark[GenerativeBenchmarkAccumulator]):
    """Complete generative AI benchmark results with specialized metrics.

    Encapsulates comprehensive performance data from scheduler-driven generative
    workload executions including request-level statistics, token/latency distributions,
    throughput analysis, and concurrency patterns. Provides computed fields for temporal
    analysis and status-grouped request details for detailed post-execution reporting.
    """

    type_: Literal["generative_benchmark"] = "generative_benchmark"  # type: ignore[assignment]

    config: BenchmarkConfig = Field(
        description="Configuration parameters for this benchmark execution",
    )
    scheduler_state: SchedulerState = Field(
        description="Final state of the scheduler after benchmark completion",
    )
    scheduler_metrics: SchedulerMetrics = Field(
        description="Scheduler timing and performance statistics",
    )
    metrics: GenerativeMetrics = Field(
        description="Performance metrics and statistical distributions",
    )
    requests: StatusBreakdown[
        list[GenerativeRequestStats],
        list[GenerativeRequestStats],
        list[GenerativeRequestStats],
        None,
    ] = Field(
        description=(
            "Request details grouped by status: successful, incomplete, errored"
        ),
    )
    warnings: list[BenchmarkWarning] = Field(
        default_factory=list,
        description="Post-benchmark warnings for a missed arrival schedule",
    )

    @computed_field  # type: ignore[prop-decorator]
    @property
    def start_time(self) -> float:
        """
        :return: Benchmark start time in seconds since epoch
        """
        return self.scheduler_metrics.measure_start_time

    @computed_field  # type: ignore[prop-decorator]
    @property
    def end_time(self) -> float:
        """
        :return: Benchmark end time in seconds since epoch
        """
        return self.scheduler_metrics.measure_end_time

    @computed_field  # type: ignore[prop-decorator]
    @property
    def duration(self) -> float:
        """
        :return: Total benchmark execution duration in seconds
        """
        return self.end_time - self.start_time

    @computed_field  # type: ignore[prop-decorator]
    @property
    def warmup_duration(self) -> float:
        """
        :return: Warmup phase duration in seconds
        """
        return (
            self.scheduler_metrics.measure_start_time
            - self.scheduler_metrics.request_start_time
        )

    @computed_field  # type: ignore[prop-decorator]
    @property
    def cooldown_duration(self) -> float:
        """
        :return: Cooldown phase duration in seconds
        """
        return (
            self.scheduler_metrics.request_end_time
            - self.scheduler_metrics.measure_end_time
        )

    @property
    def request_latency(self) -> StatusDistributionSummary:
        """
        :return: Statistical distribution of request latencies across all requests
        """
        return self.metrics.request_latency

    @property
    def request_throughput(self) -> StatusDistributionSummary:
        """
        :return: Statistical distribution of throughput measured in requests per second
        """
        return self.metrics.requests_per_second

    @property
    def request_concurrency(self) -> StatusDistributionSummary:
        """
        :return: Statistical distribution of concurrent requests throughout execution
        """
        return self.metrics.request_concurrency

    @classmethod
    def compile(
        cls,
        accumulator: GenerativeBenchmarkAccumulator,
        scheduler_state: SchedulerState,
    ) -> GenerativeBenchmark:
        """
        Compile final benchmark results from accumulated execution state.

        :param accumulator: Accumulated benchmark state with request statistics
        :param scheduler_state: Final scheduler state after execution completion
        :return: Compiled generative benchmark instance with complete metrics
        """
        metrics = GenerativeMetrics.compile(accumulator)
        scheduler_metrics = SchedulerMetrics.compile(accumulator, scheduler_state)
        warnings = BenchmarkWarningAnalyzer(accumulator.config.warnings).analyze(
            scheduler_metrics=scheduler_metrics,
            metrics=metrics,
        )
        return GenerativeBenchmark(
            config=accumulator.config,
            scheduler_state=scheduler_state,
            scheduler_metrics=scheduler_metrics,
            metrics=metrics,
            requests=StatusBreakdown(
                successful=accumulator.completed.get_sampled(),
                incomplete=accumulator.incomplete.get_sampled(),
                errored=accumulator.errored.get_sampled(),
                total=None,
            ),
            warnings=warnings,
        )

class GenerativeBenchmarkAccumulator(
    BenchmarkAccumulator[GenerationRequest, GenerationResponse]
):
    """
    Primary accumulator for generative benchmark execution metrics and statistics.

    Orchestrates real-time metric collection across timing, scheduler, concurrency, and
    generative performance dimensions. Maintains separate accumulators for completed,
    errored, and incomplete requests while tracking overall metrics. Integrates with
    scheduler state to monitor warmup/cooldown phases and compute time-weighted
    statistics for throughput and latency analysis.
    """

    timings: GenerativeBenchmarkTimings = Field(
        default_factory=GenerativeBenchmarkTimings,
        description="Timing phases and transitions during benchmark execution",
    )
    completed: GenerativeRequestsAccumulator = Field(
        default_factory=GenerativeRequestsAccumulator,
        description="Accumulator for completed requests",
    )
    errored: GenerativeRequestsAccumulator = Field(
        default_factory=GenerativeRequestsAccumulator,
        description="Accumulator for errored requests",
    )
    incomplete: GenerativeRequestsAccumulator = Field(
        default_factory=GenerativeRequestsAccumulator,
        description="Accumulator for incomplete requests",
    )
    scheduler_metrics: SchedulerMetricsAccumulator = Field(
        default_factory=SchedulerMetricsAccumulator,
        description="Running metrics for scheduler state",
    )
    concurrency_metric: RunningMetricStats = Field(
        default_factory=RunningMetricStats,
        description="Accumulated request concurrency statistics",
    )
    total_metrics: GenerativeMetricsAccumulator = Field(
        default_factory=GenerativeMetricsAccumulator,
        description="Running metrics for all requests",
    )
    completed_metrics: GenerativeMetricsAccumulator = Field(
        default_factory=GenerativeMetricsAccumulator,
        description="Running metrics for completed requests",
    )
    errored_metrics: GenerativeMetricsAccumulator = Field(
        default_factory=GenerativeMetricsAccumulator,
        description="Running metrics for errored requests",
    )
    incomplete_metrics: GenerativeMetricsAccumulator = Field(
        default_factory=GenerativeMetricsAccumulator,
        description="Running metrics for incomplete requests",
    )

    def model_post_init(self, __context):
        """
        Initialize child accumulators with config values after model construction.

        Propagates sample_size from config to child request accumulators to ensure
        consistent sampling behavior across completed, errored, and incomplete request
        collections. This ensures the --metrics-sample-size option functions correctly.
        """
        super().model_post_init(__context)

        # Propagate sample_size from config to child accumulators
        self.completed.sample_size = self.config.sample_size
        self.errored.sample_size = self.config.sample_size
        self.incomplete.sample_size = self.config.sample_size

    def update_estimate(
        self,
        response: GenerationResponse | None,
        request: GenerationRequest,
        info: RequestInfo,
        scheduler_state: SchedulerState,
    ):
        """
        Update all benchmark metrics with a completed request.

        Processes request completion by updating timing phases, concurrency metrics,
        scheduler statistics, and generative performance metrics. Routes request to
        appropriate status-specific accumulator (completed/errored/incomplete) and
        updates aggregate totals. Cancelled requests that never started are ignored.

        :param response: Generation response with output and metrics, or None
        :param request: Original generation request with input data
        :param info: Request execution information and timing
        :param scheduler_state: Current scheduler state for phase tracking
        """
        self.timings.update_estimate(info, scheduler_state, self.config)
        duration = self.timings.duration
        elapsed_time_last_update = self.timings.elapsed_time_last_update
        self.concurrency_metric.update_estimate(
            value=scheduler_state.processing_requests,
            duration=duration,
            elapsed=elapsed_time_last_update,
        )

        requests_accumulator: GenerativeRequestsAccumulator
        metrics_accumulator: GenerativeMetricsAccumulator

        if info.status == "completed":
            requests_accumulator = self.completed
            metrics_accumulator = self.completed_metrics
        elif info.status == "errored":
            requests_accumulator = self.errored
            metrics_accumulator = self.errored_metrics
        elif info.status == "cancelled" and info.timings.resolve_start is not None:
            requests_accumulator = self.incomplete
            metrics_accumulator = self.incomplete_metrics
        else:
            # Not a terminal status or cancelled before starting
            # Do not include in requests or metrics
            return

        stats = requests_accumulator.update_estimate(
            response, request, info, self.config.prefer_response_metrics
        )
        metrics_accumulator.update_estimate(stats, duration)
        self.total_metrics.update_estimate(stats, duration)
        self.scheduler_metrics.update_estimate(scheduler_state, stats)

class GenerativeBenchmarkMetadata(StandardBaseModel):
    """
    Versioning and environment metadata for generative benchmark reports.
    """

    # Make sure to update version when making breaking changes to report schema
    version: Literal[2] = Field(
        description=(
            "Version of the benchmark report schema, increments "
            "whenever there is a breaking change to the output format"
        ),
        default=2,
    )
    guidellm_version: str = Field(
        description="Version of the guidellm package used for the benchmark",
        default_factory=lambda: version("guidellm"),
    )
    python_version: str = Field(
        description="Version of Python interpreter used during the benchmark",
        default_factory=platform.python_version,
    )
    platform: str = Field(
        description="Operating system platform where the benchmark was executed",
        default_factory=platform.platform,
    )

class GenerativeBenchmarkTimings(StandardBaseModel):
    """
    Tracks timing phases and transitions during benchmark execution.

    Monitors timestamps throughout benchmark execution including request submission,
    measurement period boundaries (warmup/active/cooldown), and completion events.
    Provides duration calculations and phase status determination based on configured
    warmup and cooldown periods.
    """

    request_start: float | None = Field(
        description="Timestamp when the first request was sent", default=None
    )
    measure_start: float | None = Field(
        description="Timestamp when measurement period started", default=None
    )
    measure_end: float | None = Field(
        description="Timestamp when measurement period ended", default=None
    )
    request_end: float | None = Field(
        description="Timestamp when the last request was completed", default=None
    )
    current_update: float | None = Field(
        description="Most recent timestamp observed during execution", default=None
    )
    current_request: float | None = Field(
        description="Most recent request completion timestamp observed", default=None
    )
    last_update: float | None = Field(
        description="Previous timestamp observed before the current one", default=None
    )
    last_request: float | None = Field(
        description="Previous request completion timestamp before the current one",
        default=None,
    )

    @property
    def status(self) -> Literal["pending", "warmup", "active", "cooldown"]:
        """
        :return: Current execution phase based on timing thresholds
        """
        if self.request_start is None or self.current_update is None:
            return "pending"

        if self.measure_start is None or self.current_update <= self.measure_start:
            return "warmup"

        if self.measure_end is not None and self.current_update >= self.measure_end:
            return "cooldown"

        return "active"

    @property
    def duration(self) -> float:
        """
        :return: Elapsed time since measurement or request start in seconds
        """
        if self.request_start is None or self.current_update is None:
            return 0.0

        return self.current_update - self.request_start

    @property
    def elapsed_time_last_update(self) -> float:
        """
        :return: Time elapsed between the last two update timestamps in seconds
        """
        if self.current_update is None or self.last_update is None:
            return 0.0

        return self.current_update - self.last_update

    @property
    def elapsed_time_last_request(self) -> float:
        """
        :return: Time elapsed between the last two request completions in seconds
        """
        if self.current_request is None or self.last_request is None:
            return 0.0

        return self.current_request - self.last_request

    @property
    def finalized_request_start(self) -> float:
        """
        :return: Finalized timestamp from the current state for when requests started
        """
        return self.request_start or -1.0

    @property
    def finalized_measure_start(self) -> float:
        """
        :return: Finalized timestamp from the current state for when measurement started
        """
        return self.measure_start or self.finalized_request_start

    @property
    def finalized_measure_end(self) -> float:
        """
        :return: Finalized timestamp from the current state for when measurement ended
        """
        return self.measure_end or self.finalized_request_end

    @property
    def finalized_request_end(self) -> float:
        """
        :return: Finalized timestamp from the current state for when requests ended
        """
        return self.request_end or self.current_request or -1.0

    def update_estimate(
        self,
        info: RequestInfo,
        scheduler_state: SchedulerState,
        config: BenchmarkConfig,
    ):
        """
        Update timing estimates based on request info and scheduler state.

        Advances timing markers through benchmark phases (warmup to active to cooldown)
        based on configured thresholds. Updates current/last timestamps for updates and
        request completions, determining measurement period boundaries.

        :param info: Request information containing timing data
        :param scheduler_state: Current scheduler state with progress metrics
        :param config: Benchmark configuration with warmup/cooldown settings
        """
        # First update non terminal timestamps
        self.request_start = scheduler_state.start_requests_time
        self.last_update = self.current_update
        if (current_time := info.timings.last_reported) is not None:
            self.current_update = (
                current_time
                if self.current_update is None
                else max(self.current_update, current_time)
            )

        # Next update measurement period timestamps, if available and possible
        warmup_active, measure_start = config.warmup.compute_transition_time(
            info=info, state=scheduler_state, period="start"
        )
        if not warmup_active:
            # No warmup, set measure_start to first request start
            self.measure_start = self.request_start
        elif measure_start is not None:
            self.measure_start = measure_start
        cooldown_active, measure_end = config.cooldown.compute_transition_time(
            info=info, state=scheduler_state, period="end"
        )
        if cooldown_active and measure_end is not None:
            self.measure_end = measure_end

        # Update last request terminal timestamps, if request is terminal
        if info.status in {"completed", "errored", "cancelled"}:
            self.last_request = self.current_request
            if info.completed_at is not None and (
                self.current_request is None or info.completed_at > self.current_request
            ):
                self.current_request = info.completed_at

        # Finally, update request stop timestamps, if at that stage and available
        if scheduler_state.end_processing_time is not None and self.request_end is None:
            self.request_end = (
                scheduler_state.progress.stop_time
                or self.current_request
                or scheduler_state.end_processing_time
            )
            if self.measure_end is None:
                # No cooldown triggered, set measure_end to request_end
                self.measure_end = self.request_end