Type cleanup

Signed-off-by: Samuel Monson <smonson@redhat.com>

Author: sjmonson

src/guidellm/scheduler/__init__.py

@@ -15,10 +15,10 @@
 from .objects import (
     BackendInterface,
     BackendT,
+    DatasetIterT,
     HistoryT,
     MeasuredRequestTimings,
-    MultiTurnRequestT,
-    MultiTurnT,
+    RequestDataT,
     RequestSchedulerTimings,
     RequestT,
     ResponseT,
@@ -28,7 +28,6 @@
     SchedulerState,
     SchedulerUpdateAction,
     SchedulerUpdateActionProgress,
-    TurnT,
 )
 from .scheduler import Scheduler
 from .strategies import (
@@ -59,6 +58,7 @@
     "Constraint",
     "ConstraintInitializer",
     "ConstraintsInitializerFactory",
+    "DatasetIterT",
     "Environment",
     "HistoryT",
     "LastCompletionRequestTimings",
@@ -68,12 +68,11 @@
     "MaxGlobalErrorRateConstraint",
     "MaxNumberConstraint",
     "MeasuredRequestTimings",
-    "MultiTurnRequestT",
-    "MultiTurnT",
     "NoDelayRequestTimings",
     "NonDistributedEnvironment",
     "PoissonRateRequestTimings",
     "PydanticConstraintInitializer",
+    "RequestDataT",
     "RequestSchedulerTimings",
     "RequestT",
     "ResponseT",
@@ -91,7 +90,6 @@
     "StrategyType",
     "SynchronousStrategy",
     "ThroughputStrategy",
-    "TurnT",
     "UnserializableConstraintInitializer",
     "WorkerProcess",
     "WorkerProcessGroup",

src/guidellm/scheduler/environments.py

@@ -19,14 +19,14 @@
 
 import time
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Iterable
+from collections.abc import AsyncIterator
 from typing import (
     Generic,
 )
 
 from guidellm.scheduler.constraints import Constraint
 from guidellm.scheduler.objects import (
-    MultiTurnRequestT,
+    DatasetIterT,
     RequestT,
     ResponseT,
     ScheduledRequestInfo,
@@ -52,11 +52,11 @@ class Environment(ABC, Generic[RequestT, ResponseT], InfoMixin):
     @abstractmethod
     async def sync_run_params(
         self,
-        requests: Iterable[RequestT | MultiTurnRequestT[RequestT]],
+        requests: DatasetIterT[RequestT],
         strategy: SchedulingStrategy,
         constraints: dict[str, Constraint],
     ) -> tuple[
-        Iterable[RequestT | MultiTurnRequestT[RequestT]],
+        DatasetIterT[RequestT],
         SchedulingStrategy,
         dict[str, Constraint],
     ]:
@@ -130,7 +130,7 @@ async def sync_run_end(
     ) -> AsyncIterator[
         tuple[
             ResponseT,
-            RequestT | MultiTurnRequestT[RequestT],
+            RequestT,
             ScheduledRequestInfo,
             SchedulerState,
         ]
@@ -194,11 +194,11 @@ def __init__(self):
 
     async def sync_run_params(
         self,
-        requests: Iterable[RequestT | MultiTurnRequestT[RequestT]],
+        requests: DatasetIterT[RequestT],
         strategy: SchedulingStrategy,
         constraints: dict[str, Constraint],
     ) -> tuple[
-        Iterable[RequestT | MultiTurnRequestT[RequestT]],
+        DatasetIterT[RequestT],
         SchedulingStrategy,
         dict[str, Constraint],
     ]:
@@ -250,7 +250,7 @@ async def sync_run_end(
     ) -> AsyncIterator[
         tuple[
             ResponseT,
-            RequestT | MultiTurnRequestT[RequestT],
+            RequestT,
             ScheduledRequestInfo,
             SchedulerState,
         ]

src/guidellm/scheduler/objects.py

@@ -11,7 +11,7 @@
 
 import time
 import uuid
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Iterable
 from typing import (
     Any,
     ClassVar,
@@ -34,10 +34,10 @@
 __all__ = [
     "BackendInterface",
     "BackendT",
+    "DatasetIterT",
     "HistoryT",
     "MeasuredRequestTimings",
-    "MultiTurnRequestT",
-    "MultiTurnT",
+    "RequestDataT",
     "RequestSchedulerTimings",
     "RequestT",
     "ResponseT",
@@ -47,36 +47,32 @@
     "SchedulerState",
     "SchedulerUpdateAction",
     "SchedulerUpdateActionProgress",
-    "TurnT",
 ]
 
 RequestT = TypeVar("RequestT")
 """Generic request object type for scheduler processing."""
 
-# TODO: Remove
-MultiTurnRequestT = RequestT
-
 ResponseT = TypeVar("ResponseT")
 """Generic response object type returned by backend processing."""
 
-TurnT = TypeAliasType(
-    "TurnT",
+RequestDataT = TypeAliasType(
+    "RequestDataT",
     tuple[RequestT, "ScheduledRequestAugmentation", "ScheduledRequestInfo"],
     type_params=(RequestT,),
 )
-
-MultiTurnT = TypeAliasType(
-    "MultiTurnT",
-    list[TurnT[RequestT]],
-    type_params=(RequestT,),
-)
-"""Multi-turn request structure supporting conversation history with optional delays."""
+"""Request including external metadata and scheduling config."""
 
 HistoryT = TypeAliasType(
     "HistoryT",
     list[tuple[RequestT, ResponseT]],
     type_params=(RequestT, ResponseT),
 )
+"""Record of requests + responses in conversation."""
+
+
+DatasetIterT = TypeAliasType(
+    "DatasetIterT", Iterable[Iterable[tuple[RequestT, float]]], type_params=(RequestT,)
+)
 
 
 class SchedulerMessagingPydanticRegistry(RegistryMixin[RegistryObjT]):

src/guidellm/scheduler/scheduler.py

@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-from collections.abc import AsyncIterator, Iterable
+from collections.abc import AsyncIterator
 from typing import Any, Generic
 
 from guidellm.scheduler.constraints import (
@@ -20,7 +20,7 @@
 from guidellm.scheduler.environments import Environment, NonDistributedEnvironment
 from guidellm.scheduler.objects import (
     BackendInterface,
-    MultiTurnRequestT,
+    DatasetIterT,
     RequestT,
     ResponseT,
     ScheduledRequestInfo,
@@ -66,7 +66,7 @@ class Scheduler(
 
     async def run(
         self,
-        requests: Iterable[RequestT | MultiTurnRequestT[RequestT]],
+        requests: DatasetIterT[RequestT],
         backend: BackendInterface[RequestT, ResponseT],
         strategy: SchedulingStrategy,
         env: Environment | None,

src/guidellm/scheduler/worker.py

@@ -32,8 +32,7 @@
 from guidellm.scheduler.objects import (
     BackendInterface,
     HistoryT,
-    MultiTurnRequestT,
-    MultiTurnT,
+    RequestDataT,
     RequestT,
     ResponseT,
     ScheduledRequestAugmentation,
@@ -54,7 +53,7 @@
     "ProcessRequestT",
     tuple[
         HistoryT[RequestT, ResponseT],
-        MultiTurnT[RequestT],
+        list[RequestDataT[RequestT]],
         ScheduledRequestAugmentation,
     ],
     type_params=(RequestT, ResponseT),
@@ -87,11 +86,8 @@ class WorkerProcess(Generic[RequestT, ResponseT]):
     def __init__(
         self,
         messaging: InterProcessMessaging[
-            tuple[
-                ResponseT | None,
-                RequestT | MultiTurnRequestT[RequestT],
-                ScheduledRequestInfo,
-            ],
+            tuple[ResponseT | None, RequestT, ScheduledRequestInfo],
+            list[RequestDataT[RequestT]],
         ],
         backend: BackendInterface[RequestT, ResponseT],
         request_timings: ScheduledRequestTimings,
@@ -132,7 +128,7 @@ def __init__(
         self.backend_started = False
         self.messaging_started = False
         self.turns_queue: list[
-            tuple[HistoryT[RequestT, ResponseT], MultiTurnT[RequestT]]
+            tuple[HistoryT[RequestT, ResponseT], list[RequestDataT[RequestT]]]
         ] = []
 
     def run(self):
@@ -332,7 +328,7 @@ async def _cancel_requests_loop(self):
                 self._send_update("cancelled", None, request, request_info)
 
     async def _process_next_request(self) -> ProcessRequestT[RequestT, ResponseT]:
-        conversation: MultiTurnT[RequestT] = []
+        conversation: list[RequestDataT[RequestT]] = []
         history: HistoryT[RequestT, ResponseT] = []
         request: RequestT | None = None
         request_info: ScheduledRequestInfo | None = None
@@ -409,7 +405,7 @@ async def _process_next_request(self) -> ProcessRequestT[RequestT, ResponseT]:
     async def _wait_then_requeue(
         self,
         history: HistoryT[RequestT, ResponseT],
-        conversation: MultiTurnT[RequestT],
+        conversation: list[RequestDataT[RequestT]],
         aug: ScheduledRequestAugmentation,
     ):
         try:
@@ -427,7 +423,7 @@ def _send_update(
             "pending", "in_progress", "completed", "errored", "cancelled"
         ],
         response: ResponseT | None,
-        request: RequestT | MultiTurnRequestT[RequestT],
+        request: RequestT,
         request_info: ScheduledRequestInfo,
     ):
         prev_status = request_info.status

src/guidellm/scheduler/worker_group.py

@@ -25,8 +25,8 @@
 from guidellm.scheduler.constraints import Constraint, RequestsExhaustedConstraint
 from guidellm.scheduler.objects import (
     BackendInterface,
-    MultiTurnRequestT,
-    MultiTurnT,
+    DatasetIterT,
+    RequestDataT,
     RequestT,
     ResponseT,
     ScheduledRequestAugmentation,
@@ -83,8 +83,8 @@ class WorkerProcessGroup(Generic[RequestT, ResponseT]):
 
     def __init__(
         self,
-        requests: Iterable[RequestT | MultiTurnRequestT[RequestT]] | None,
-        cycle_requests: Iterable[RequestT | MultiTurnRequestT[RequestT]] | None,
+        requests: DatasetIterT[RequestT] | None,
+        cycle_requests: DatasetIterT[RequestT] | None,
         backend: BackendInterface[RequestT, ResponseT],
         strategy: SchedulingStrategy,
         constraints: dict[str, Constraint],
@@ -131,16 +131,8 @@ def __init__(
         # Scheduler and messaging state, created in start
         self.state: WorkerGroupState[ResponseT, RequestT] = None
         self.messaging: InterProcessMessaging[
-            tuple[
-                RequestT | MultiTurnRequestT[RequestT],
-                ScheduledRequestInfo,
-            ],
-            tuple[
-                ResponseT | None,
-                RequestT | MultiTurnRequestT[RequestT],
-                ScheduledRequestInfo,
-                SchedulerState,
-            ],
+            list[RequestDataT[RequestT]],
+            tuple[ResponseT | None, RequestT, ScheduledRequestInfo, SchedulerState],
         ] = None
 
     async def create_processes(self):
@@ -473,9 +465,9 @@ def __init__(
 
     def requests_generator(
         self,
-        requests: Iterable[Iterable[tuple[RequestT, float]]] | None,
-        cycle_requests: Iterable[Iterable[tuple[RequestT, float]]] | None,
-    ) -> Generator[MultiTurnT[RequestT], None, None]:
+        requests: DatasetIterT[RequestT] | None,
+        cycle_requests: DatasetIterT[RequestT] | None,
+    ) -> Generator[list[RequestDataT[RequestT]], None, None]:
         """
         Generate request-info pairs for worker processing with constraint evaluation.
 
@@ -544,12 +536,12 @@ def received_callback(
         self,
         update: tuple[
             ResponseT | None,
-            RequestT | MultiTurnRequestT,
+            RequestT,
             ScheduledRequestInfo,
         ],
     ) -> tuple[
         ResponseT | None,
-        RequestT | MultiTurnRequestT,
+        RequestT,
         ScheduledRequestInfo,
         SchedulerState,
     ]:

tests/unit/scheduler/test_objects.py

@@ -7,13 +7,11 @@
 
 import pytest
 from pydantic import ValidationError
-from typing_extensions import TypeAliasType
 
 from guidellm.scheduler import (
     BackendInterface,
     BackendT,
     MeasuredRequestTimings,
-    MultiTurnRequestT,
     RequestSchedulerTimings,
     RequestT,
     ResponseT,
@@ -49,20 +47,6 @@ def test_backend_t():
     assert BackendT.__constraints__ == ()
 
 
-def test_multi_turn_request_t():
-    """Validate MultiTurnRequestT is a TypeAliasType for multi-turn requests."""
-    assert isinstance(MultiTurnRequestT, TypeAliasType)
-    assert MultiTurnRequestT.__name__ == "MultiTurnRequestT"
-
-    value = MultiTurnRequestT.__value__
-    assert hasattr(value, "__origin__")
-    assert value.__origin__ is Union
-
-    type_params = getattr(MultiTurnRequestT, "__type_params__", ())
-    assert len(type_params) == 1
-    assert type_params[0].__name__ == "RequestT"
-
-
 class TestBackendInterface:
     """Test the BackendInterface abstract base class."""