guidellm.scheduler.schemas.state

class SchedulerProgress(StandardBaseModel):
    """
    Progress tracking data for scheduler operations.

    Provides estimates for remaining work in scheduler operations, including
    fraction complete, request counts, and duration. Used by constraints and
    monitoring systems to track execution progress and make termination decisions.
    """

    remaining_requests: float | None = Field(
        description="Estimated number of remaining requests to process", default=None
    )
    total_requests: float | None = Field(
        description="Total number of requests to process", default=None
    )
    remaining_duration: float | None = Field(
        description="Estimated remaining duration in seconds", default=None
    )
    total_duration: float | None = Field(
        description="Total duration in seconds to process for", default=None
    )
    stop_time: float | None = Field(
        description="The timestamp the processing stopped at", default=None
    )

    @property
    def remaining_fraction(self) -> float | None:
        """
        :return: Estimated fraction of remaining progress, if known
        """
        fraction: float | None = None

        if (requests_fraction := self.remaining_requests_fraction) is not None:
            fraction = requests_fraction

        if (duration_fraction := self.remaining_duration_fraction) is not None:
            fraction = (
                duration_fraction
                if fraction is None
                else min(fraction, duration_fraction)
            )

        return fraction

    @property
    def remaining_requests_fraction(self) -> float | None:
        """
        :return: Estimated fraction of remaining requests, if known
        """
        return (
            self.remaining_requests / float(self.total_requests)
            if self.remaining_requests is not None
            and self.total_requests is not None
            and self.total_requests > 0
            else None
        )

    @property
    def remaining_duration_fraction(self) -> float | None:
        """
        :return: Estimated fraction of remaining duration, if known
        """
        return (
            self.remaining_duration / float(self.total_duration)
            if self.remaining_duration is not None
            and self.total_duration is not None
            and self.total_duration > 0
            else None
        )

    def combine(self, other: SchedulerProgress) -> SchedulerProgress:
        """
        Combine two progress instances, taking the minimum remaining estimates.

        :param other: Another progress instance to combine with
        :return: New progress instance with combined estimates
        """
        if (other_req_fraction := other.remaining_requests_fraction) is not None and (
            (cur_req_fraction := self.remaining_requests_fraction) is None
            or other_req_fraction < cur_req_fraction
        ):
            # Only update if the other is more advanced (lower fraction)
            self.remaining_requests = other.remaining_requests
            self.total_requests = other.total_requests

        if (other_dur_fraction := other.remaining_duration_fraction) is not None and (
            (cur_dur_fraction := self.remaining_duration_fraction) is None
            or other_dur_fraction < cur_dur_fraction
        ):
            # Only update if the other is more advanced (lower fraction)
            self.remaining_duration = other.remaining_duration
            self.total_duration = other.total_duration

        if other.stop_time is not None and (
            self.stop_time is None or other.stop_time < self.stop_time
        ):
            # Only update if the other has an earlier stop time
            self.stop_time = other.stop_time

        return self

class SchedulerState(StandardBaseModel):
    """
    Comprehensive state tracking for scheduler execution.

    Tracks scheduler execution progress, request counts, timing information,
    and constraint enforcement. Central to scheduler coordination, providing
    real-time metrics for monitoring and decision-making across distributed
    worker processes.

    Example:
    ::
        state = SchedulerState(node_id=0, num_processes=4)
        state.created_requests += 1
        state.queued_requests += 1
        completion_rate = state.processed_requests / state.created_requests
    """

    node_id: int = Field(
        description="Unique identifier for this scheduler node", default=-1
    )
    num_processes: int = Field(
        description="Number of worker processes in this scheduler", default=-1
    )
    start_time: float = Field(
        description="Unix timestamp when the scheduler started",
        default_factory=time.time,
    )
    end_time: float | None = Field(
        default=None, description="Unix timestamp when the scheduler stopped"
    )
    start_requests_time: float | None = Field(
        default=None, description="Unix timestamp of the first sent request"
    )
    end_requests_time: float | None = Field(
        default=None, description="Unix timestamp of the last finalized request"
    )
    end_queuing_time: float | None = Field(
        default=None, description="Unix timestamp when request queuing stopped"
    )
    end_queuing_constraints: dict[str, SchedulerUpdateAction] = Field(
        default_factory=dict,
        description="Constraints that triggered queuing termination",
    )
    end_processing_time: float | None = Field(
        default=None, description="Unix timestamp when request processing stopped"
    )
    end_processing_constraints: dict[str, SchedulerUpdateAction] = Field(
        default_factory=dict,
        description="Constraints that triggered processing termination",
    )
    scheduler_constraints: dict[str, SchedulerUpdateAction] = Field(
        default_factory=dict,
        description="Latest state from all constraints applied during scheduler run",
    )

    progress: SchedulerProgress = Field(
        default_factory=SchedulerProgress,
        description="Overall progress information for the scheduler run",
    )

    created_requests: int = Field(
        default=0, description="Total number of requests created"
    )
    queued_requests: int = Field(
        default=0, description="Total number of requests queued for processing"
    )
    pending_requests: int = Field(
        default=0,
        description="Number of requests pending processing within a worker",
    )
    processing_requests: int = Field(
        default=0, description="Number of requests currently being processed"
    )
    processed_requests: int = Field(
        default=0, description="Number of requests that completed processing"
    )
    successful_requests: int = Field(
        default=0, description="Number of requests that completed successfully"
    )
    errored_requests: int = Field(
        default=0, description="Number of requests that failed with errors"
    )
    cancelled_requests: int = Field(
        default=0, description="Number of requests that were cancelled"
    )
    # TODO: We should consider replacing this with per-conversation metrics
    generation_delay_samples: list[float] = Field(
        default_factory=list,
        exclude=True,  # Avoid including this potentially large list in output
        description="Delays in yielding conversations from the data generator",
    )

class SchedulerUpdateAction(StandardBaseModel):
    """
    Control directives for scheduler behavior and operations.

    Encapsulates control signals for scheduler operations including request
    queuing and processing directives. Used by constraints to communicate
    termination conditions and progress to scheduler components.

    Example:
    ::
        action = SchedulerUpdateAction(
            request_queuing="stop",
            request_processing="continue",
            metadata={"reason": "max_requests_reached"}
        )
    """

    request_queuing: Literal["continue", "stop"] = Field(
        default="continue", description="Action to take for request queuing operations"
    )
    request_processing: Literal["continue", "stop_local", "stop_all"] = Field(
        default="continue",
        description="Action to take for request processing operations",
    )
    stopping_scope: Literal["current", "all"] = Field(
        default="current",
        description=(
            "Whether this constraint's stop signal should halt only the current "
            "benchmark ('current') or also prevent escalation to subsequent "
            "rates/streams ('all')."
        ),
    )
    metadata: dict[str, Any] = Field(
        default_factory=dict,
        description="Additional context and data for the scheduler action",
    )
    progress: SchedulerProgress = Field(
        default_factory=SchedulerProgress,
        description="Progress information for the scheduler action",
    )