guidellm.scheduler

@SchedulingStrategy.register("constant")
class AsyncConstantStrategy(SchedulingStrategy):
    """
    Constant-rate scheduling for predictable load patterns.

    Schedules requests at a fixed rate distributed evenly across worker processes,
    providing predictable timing behavior for steady-state load simulation and
    consistent system performance measurement. Requests arrive at uniform intervals.
    """

    type_: Literal["constant"] = "constant"  # type: ignore[assignment]
    rate: float = Field(
        description="Request scheduling rate in requests per second",
        gt=0,
    )
    max_concurrency: PositiveInt | None = Field(
        default=None,
        description="Maximum number of concurrent requests to schedule",
    )
    rampup_duration: NonNegativeFloat = Field(
        default=0.0,
        description=(
            "Duration in seconds to linearly ramp up from 0 to target rate "
            "at the beginning of each strategy run"
        ),
    )

    def __str__(self) -> str:
        """
        :return: String identifier with rate value
        """
        return f"constant@{self.rate:.2f}"

    @property
    def processes_limit(self) -> PositiveInt | None:
        """
        :return: Max concurrency if set, otherwise None for unlimited
        """
        return self.max_concurrency

    @property
    def requests_limit(self) -> PositiveInt | None:
        """
        :return: Max concurrency if set, otherwise None for unlimited
        """
        return self.max_concurrency

    async def next_request_time(self, worker_index: PositiveInt) -> float:
        """
        Calculate next request time at fixed intervals with optional linear rampup.

        Schedules requests at uniform intervals determined by the configured rate,
        independent of request completion times. If rampup_duration is set, the rate
        increases linearly from 0 to the target rate during the rampup period, then
        continues at the constant rate.

        :param worker_index: Unused for constant strategy
        :return: Start time plus interval based on request index and
            rampup configuration
        """
        _ = worker_index  # unused
        current_index = self.next_request_index()
        start_time = await self.get_processes_start_time()

        if self.rampup_duration > 0:
            # Calculate number of requests that would be sent during rampup
            # Cumulative requests by time t during rampup:
            # n = rate * t² / (2 * rampup_duration)
            # At end of rampup (t = rampup_duration), n_rampup is calculated below
            n_rampup = self.rate * self.rampup_duration / 2.0

            if current_index == 1:
                # First request at start_time
                return start_time
            elif current_index <= n_rampup:
                # During rampup: solve for t where
                # n = rate * t² / (2 * rampup_duration)
                time_offset = math.sqrt(
                    2.0 * current_index * self.rampup_duration / self.rate
                )
                return start_time + time_offset
            else:
                # After rampup: continue at constant rate
                time_offset = (
                    self.rampup_duration + (current_index - n_rampup) / self.rate
                )
                return start_time + time_offset
        else:
            # No rampup: uniform intervals
            return start_time + current_index / self.rate

    def request_completed(self, request_info: RequestInfo):
        """
        Handle request completion (no-op for constant strategy).

        :param request_info: Completed request metadata (unused)
        """
        _ = request_info  # request_info unused for async constant strategy

@SchedulingStrategy.register("poisson")
class AsyncPoissonStrategy(SchedulingStrategy):
    """
    Poisson-distributed scheduling for realistic load simulation.

    Schedules requests following a Poisson process with exponentially distributed
    inter-arrival times, providing realistic simulation of user behavior and network
    traffic patterns. Request arrivals have random variance around the target rate.
    """

    type_: Literal["poisson"] = "poisson"  # type: ignore[assignment]
    rate: float = Field(
        description="Request scheduling rate in requests per second",
        gt=0,
    )
    max_concurrency: PositiveInt | None = Field(
        default=None,
        description="Maximum number of concurrent requests to schedule",
    )
    random_seed: int = Field(
        default=42,
        description="Random seed for Poisson distribution reproducibility",
    )

    _random: random.Random | None = PrivateAttr(None)
    _offset: Synchronized[float] | None = PrivateAttr(None)

    def __str__(self) -> str:
        """
        :return: String identifier with rate value
        """
        return f"poisson@{self.rate:.2f}"

    @property
    def processes_limit(self) -> PositiveInt | None:
        """
        :return: Max concurrency if set, otherwise None for unlimited
        """
        return self.max_concurrency

    @property
    def requests_limit(self) -> PositiveInt | None:
        """
        :return: Max concurrency if set, otherwise None for unlimited
        """
        return self.max_concurrency

    def init_processes_timings(
        self,
        worker_count: PositiveInt,
        max_concurrency: PositiveInt,
        mp_context: BaseContext,
    ):
        """
        Initialize Poisson-specific timing state.

        Sets up shared offset value for coordinating exponentially distributed
        request timing across worker processes.

        :param worker_count: Number of worker processes to coordinate
        :param max_concurrency: Maximum number of concurrent requests allowed
        """
        self._offset = mp_context.Value("d", -1.0)
        # Call base implementation last to avoid
        # setting Event before offset is ready
        super().init_processes_timings(worker_count, max_concurrency, mp_context)

    def init_processes_start(self, start_time: float):
        """
        Initialize the offset time for Poisson timing calculations.

        Sets the initial timing offset from which exponentially distributed
        intervals are calculated.

        :param start_time: Unix timestamp when request processing should begin
        """
        ThroughputStrategy.init_processes_start(self, start_time)

        if self._offset is None:
            raise RuntimeError(
                "_offset is None in init_processes_start; was "
                "init_processes_timings not called?"
            )
        with self._offset.get_lock():
            self._offset.value = start_time

    async def next_request_time(self, worker_index: PositiveInt) -> float:
        """
        Calculate next request time using exponential distribution.

        Generates inter-arrival times following exponential distribution,
        accumulating delays to produce Poisson-distributed request arrivals.

        :param worker_index: Unused for Poisson strategy
        :return: Next arrival time based on Poisson process
        """
        _ = worker_index  # unused
        _ = await self.get_processes_start_time()  # ensure offset is initialized

        if self._random is None:
            self._random = random.Random(self.random_seed)

        next_delay = self._random.expovariate(self.rate)

        if self._offset is None:
            raise RuntimeError(
                "_offset is None in next_request_time; was "
                "init_processes_timings not called?"
            )
        with self._offset.get_lock():
            self._offset.value += next_delay

            return self._offset.value

    def request_completed(self, request_info: RequestInfo):
        """
        Handle request completion (no-op for Poisson strategy).

        :param request_info: Completed request metadata (unused)
        """
        _ = request_info  # request_info unused for async poisson strategy

class BackendInterface(Protocol, Generic[RequestT, ResponseT]):
    """
    Protocol defining the interface for request processing backends.

    Establishes the contract for backend implementations that process requests
    within the scheduler system. Backends manage initialization, validation,
    processing, and shutdown lifecycle. All properties must be pickleable before
    process_startup is called for multi-process environments.

    Example:
    ::
        class CustomBackend(BackendInterface):
            @property
            def processes_limit(self) -> int:
                return 4

            async def resolve(self, request, request_info, history=None):
                yield response, updated_request_info
    """

    @property
    def processes_limit(self) -> int | None:
        """
        :return: Maximum worker processes supported, or None if unlimited
        """

    @property
    def requests_limit(self) -> int | None:
        """
        :return: Maximum concurrent requests supported, or None if unlimited
        """

    @property
    def info(self) -> dict[str, Any]:
        """
        :return: Backend metadata including model initialization and configuration
        """

    async def process_startup(self) -> None:
        """
        Perform backend initialization and startup procedures.

        :raises Exception: Implementation-specific exceptions for startup failures
        """

    async def validate(self) -> None:
        """
        Validate backend configuration and operational status.

        :raises Exception: Implementation-specific exceptions for validation failures
        """

    async def process_shutdown(self) -> None:
        """
        Perform backend cleanup and shutdown procedures.

        :raises Exception: Implementation-specific exceptions for shutdown failures
        """

    async def resolve(
        self,
        request: RequestT,
        request_info: RequestInfo,
        history: HistoryT[RequestT, ResponseT] | None = None,
    ) -> AsyncIterator[tuple[ResponseT | None, RequestInfo]]:
        """
        Process a request and yield incremental response updates.

        :param request: The request object to process
        :param request_info: Scheduling metadata and timing information
        :param history: Conversation history for multi-turn requests
        :yield: Tuples of (response, updated_request_info) for each response chunk.
            Response may be None for intermediate updates (e.g., first token arrival).
        :raises Exception: Implementation-specific exceptions for processing failures
        """

@SchedulingStrategy.register("concurrent")
class ConcurrentStrategy(SchedulingStrategy):
    """
    Parallel request processing with fixed concurrency limits.

    Enables concurrent request processing up to a specified number of streams,
    providing balanced throughput while maintaining predictable resource usage.
    Requests are distributed across streams with completion-based timing coordination.
    """

    type_: Literal["concurrent"] = "concurrent"  # type: ignore[assignment]
    streams: PositiveInt = Field(
        description="Number of concurrent streams for scheduling requests",
    )
    rampup_duration: NonNegativeFloat = Field(
        default=0.0,
        description=(
            "Duration in seconds to spread initial requests up to max_concurrency "
            "at the beginning of each strategy run"
        ),
    )

    _process_last_request_time: float | None = PrivateAttr(None)

    def __str__(self) -> str:
        """
        :return: String identifier with stream count
        """
        return f"concurrent@{self.streams}"

    @property
    def processes_limit(self) -> PositiveInt:
        """
        :return: Number of streams as maximum worker processes
        """
        return self.streams

    @property
    def requests_limit(self) -> PositiveInt:
        """
        :return: Number of streams as maximum concurrent requests
        """
        return self.streams

    async def next_request_time(self, worker_index: PositiveInt) -> float:
        """
        Calculate next request time with stream-based distribution.

        Initial requests are staggered across streams during rampup, subsequent
        requests scheduled after previous completion within each stream.

        :param worker_index: Worker process index for distributing initial requests
        :return: Time of last completion or staggered start time if first request
        """
        _ = worker_index  # unused
        current_index = self.next_request_index()
        start_time = await self.get_processes_start_time()

        if current_index < self.streams and self.rampup_duration > 0:
            # linearly spread start times for first concurrent requests across rampup
            return start_time + self.rampup_duration * (current_index / self.streams)

        if self._process_last_request_time is not None:
            return self._process_last_request_time

        return start_time

    def request_completed(self, request_info: RequestInfo):
        """
        Update timing state with completed request information.

        Tracks completion time to schedule next request in the same stream.

        :param request_info: Completed request metadata including timing
        """
        if request_info.completed_at is not None:
            self._process_last_request_time = request_info.completed_at

@runtime_checkable
class Constraint(Protocol):
    """
    Protocol for constraint evaluation functions that control scheduler behavior.

    Defines the interface that all constraint implementations must follow. Constraints
    are callable objects that evaluate scheduler state and request information to
    determine whether processing should continue or stop. The protocol enables type
    checking and runtime validation of constraint implementations while allowing
    flexible implementation approaches (functions, classes, closures).

    Example:
    ::
        def my_constraint(
            state: SchedulerState, request: RequestInfo
        ) -> SchedulerUpdateAction:
            if state.processing_requests > 100:
                return SchedulerUpdateAction(request_queuing="stop")
            return SchedulerUpdateAction(request_queuing="continue")
    """

    def __call__(
        self, state: SchedulerState, request: RequestInfo
    ) -> SchedulerUpdateAction:
        """
        Evaluate constraint against scheduler state and request information.

        :param state: Current scheduler state with metrics and timing information
        :param request: Individual request information and metadata
        :return: Action indicating whether to continue or stop scheduler operations
        """

class ConstraintArgs(PydanticClassRegistryMixin["ConstraintArgs"]):
    """
    Base class for constraint configuration arguments.

    Uses ``PydanticClassRegistryMixin`` to enable polymorphic deserialization
    based on the ``kind`` field. Each registered subclass represents a specific
    constraint type with its own parameters.

    :cvar schema_discriminator: Field name for polymorphic deserialization
    """

    model_config = ConfigDict(
        extra="forbid",
        serialize_by_alias=True,
        ser_json_bytes="base64",
        val_json_bytes="base64",
    )

    schema_discriminator: ClassVar[str] = "kind"

    @classmethod
    def __pydantic_schema_base_type__(cls) -> type[ConstraintArgs]:
        """
        Return base type for polymorphic validation hierarchy.

        :return: Base ConstraintArgs class for schema validation
        """
        if cls.__name__ == "ConstraintArgs":
            return cls

        return ConstraintArgs

    kind: str = Field(
        description="Constraint type discriminator for polymorphic serialization",
    )

    @property
    def constraint_key(self) -> str:
        """
        The key to use when inserting into the constraints dict.

        Defaults to ``kind``, but subclasses may override if the factory
        registry key differs from the args kind.

        :return: Registry key for this constraint type
        """
        return self.kind

@runtime_checkable
class ConstraintInitializer(Protocol):
    """
    Protocol for constraint initializer factory functions that create constraints.

    Defines the interface for factory objects that create constraint instances from
    configuration parameters. Constraint initializers enable dynamic constraint
    creation and configuration, supporting both simple boolean flags and complex
    parameter dictionaries. The protocol allows type checking while maintaining
    flexibility for different initialization patterns.

    Example:
    ::
        class MaxRequestsInitializer:
            def __init__(self, max_requests: int):
                self.max_requests = max_requests

            def create_constraint(self) -> Constraint:
                def evaluate(state, request):
                    if state.total_requests >= self.max_requests:
                        return SchedulerUpdateAction(request_queuing="stop")
                    return SchedulerUpdateAction(request_queuing="continue")
                return evaluate
    """

    def create_constraint(self, **kwargs) -> Constraint:
        """
        Create a constraint instance from configuration parameters.

        :param kwargs: Configuration parameters for constraint creation
        :return: Configured constraint evaluation function
        """

class ConstraintsInitializerFactory(RegistryMixin[ConstraintInitializer]):
    """
    Registry factory for creating and managing constraint initializers.

    Provides centralized access to registered constraint types with support for
    creating constraints from ``ConstraintArgs`` instances or pre-configured
    initializer instances. Handles constraint resolution and type validation
    for the scheduler constraint system.

    Example:
    ::
        from guidellm.scheduler import ConstraintsInitializerFactory

        # Register new constraint type
        @ConstraintsInitializerFactory.register("new_constraint")
        class NewConstraint:
            def create_constraint(self, **kwargs) -> Constraint:
                return lambda state, request: SchedulerUpdateAction()

        # Create and use constraint
        args = NewConstraintArgs(kind="new_constraint")
        initializer = ConstraintsInitializerFactory.create(args)
        constraint = initializer.create_constraint()
    """

    @classmethod
    def create(cls, args: ConstraintArgs) -> ConstraintInitializer:
        """
        Create a constraint initializer from a ``ConstraintArgs`` instance.

        :param args: Validated constraint arguments with kind discriminator
        :return: Configured constraint initializer instance
        :raises ValueError: If args.kind is not registered in the factory
        """
        if cls.registry is None or args.kind not in cls.registry:
            raise ValueError(f"Unknown constraint discriminator: {args.kind}")

        initializer_class = cls.registry[args.kind]
        return initializer_class(args=args)  # type: ignore[operator]

    @classmethod
    def deserialize(
        cls, initializer_dict: dict[str, Any]
    ) -> SerializableConstraintInitializer | UnserializableConstraintInitializer:
        """
        Deserialize constraint initializer from dictionary format.

        :param initializer_dict: Dictionary representation of constraint initializer
        :return: Reconstructed constraint initializer instance
        :raises ValueError: If constraint type is unknown or cannot be deserialized
        """
        if initializer_dict.get("type_") == "unserializable":
            return UnserializableConstraintInitializer.model_validate(initializer_dict)

        if (
            cls.registry is not None
            and initializer_dict.get("type_")
            and initializer_dict["type_"] in cls.registry
        ):
            initializer_class = cls.registry[initializer_dict["type_"]]
            if hasattr(initializer_class, "model_validate"):
                return initializer_class.model_validate(initializer_dict)  # type: ignore[return-value]
            else:
                return initializer_class(**initializer_dict)  # type: ignore[return-value,operator]

        raise ValueError(
            f"Cannot deserialize unknown constraint initializer: "
            f"{initializer_dict.get('type_', 'unknown')}"
        )

    @classmethod
    def resolve(
        cls,
        initializers: dict[
            str,
            Constraint | ConstraintInitializer,
        ],
    ) -> dict[str, Constraint]:
        """
        Resolve constraint initializers to callable constraints.

        :param initializers: Dictionary mapping constraint keys to specifications.
            Values must be Constraint instances or ConstraintInitializer instances.
        :return: Dictionary mapping constraint keys to callable functions
        :raises TypeError: If a value is not a supported type
        """
        constraints = {}

        for key, val in initializers.items():
            if isinstance(val, Constraint):
                constraints[key] = val
            elif isinstance(val, ConstraintInitializer):
                constraints[key] = val.create_constraint()
            else:
                raise TypeError(
                    f"Constraint '{key}' has unsupported value type "
                    f"{type(val).__name__}. Expected a Constraint instance or "
                    f"ConstraintInitializer instance."
                )

        return constraints