Updates to remove the measured request timings generic and replace it with pydantic polymorphism

Author: markurtz

research/multiprocesssing_communication_perf/utils.py

@@ -154,7 +154,7 @@ def create_test_objects(
                 GenerationRequest(
                     content=generate_str(objects_size),
                 ),
-                ScheduledRequestInfo[GenerationRequestTimings](
+                ScheduledRequestInfo(
                     scheduler_timings=RequestSchedulerTimings(
                         targeted_start=time.time(),
                         queued=time.time(),
@@ -173,7 +173,7 @@ def create_test_objects(
                 ),
             )
             for _ in range(num_objects)
-        ], [GenerationRequest, ScheduledRequestInfo[GenerationRequestTimings]]
+        ], [GenerationRequest, ScheduledRequestInfo]
 
     if type_ == "tuple[GenerationResponse]":
         return [
@@ -186,7 +186,7 @@ def create_test_objects(
                 GenerationRequest(
                     content=generate_str(objects_size // 2),
                 ),
-                ScheduledRequestInfo[GenerationRequestTimings](
+                ScheduledRequestInfo(
                     scheduler_timings=RequestSchedulerTimings(
                         targeted_start=time.time(),
                         queued=time.time(),
@@ -208,7 +208,7 @@ def create_test_objects(
         ], [
             GenerationResponse,
             GenerationRequest,
-            ScheduledRequestInfo[GenerationRequestTimings],
+            ScheduledRequestInfo,
         ]
 
     raise ValueError(f"Unknown type_: {type_}")

src/guidellm/backend/backend.py

@@ -18,7 +18,6 @@
 
 from guidellm.backend.objects import (
     GenerationRequest,
-    GenerationRequestTimings,
     GenerationResponse,
 )
 from guidellm.scheduler import BackendInterface
@@ -35,7 +34,7 @@
 
 class Backend(
     RegistryMixin["type[Backend]"],
-    BackendInterface[GenerationRequest, GenerationRequestTimings, GenerationResponse],
+    BackendInterface[GenerationRequest, GenerationResponse],
 ):
     """
     Base class for generative AI backends with registry and lifecycle.

src/guidellm/backend/objects.py

@@ -11,7 +11,10 @@
 
 from pydantic import Field
 
-from guidellm.scheduler import MeasuredRequestTimings
+from guidellm.scheduler import (
+    MeasuredRequestTimings,
+    SchedulerMessagingPydanticRegistry,
+)
 from guidellm.utils import StandardBaseModel
 
 __all__ = [
@@ -21,6 +24,7 @@
 ]
 
 
+@SchedulerMessagingPydanticRegistry.register()
 class GenerationRequest(StandardBaseModel):
     """Request model for backend generation operations."""
 
@@ -59,6 +63,7 @@ class GenerationRequest(StandardBaseModel):
     )
 
 
+@SchedulerMessagingPydanticRegistry.register()
 class GenerationResponse(StandardBaseModel):
     """Response model for backend generation operations."""
 
@@ -135,9 +140,11 @@ def preferred_output_tokens(
             return self.response_output_tokens or self.request_output_tokens
 
 
+@MeasuredRequestTimings.register()
 class GenerationRequestTimings(MeasuredRequestTimings):
     """Timing model for tracking generation request lifecycle events."""
 
+    timings_type: Literal["generation_request_timings"] = "generation_request_timings"
     first_iteration: Optional[float] = Field(
         default=None,
         description="Unix timestamp when the first generation iteration began.",
@@ -146,3 +153,6 @@ class GenerationRequestTimings(MeasuredRequestTimings):
         default=None,
         description="Unix timestamp when the last generation iteration completed.",
     )
+
+
+SchedulerMessagingPydanticRegistry.register_decorator(GenerationRequestTimings)

src/guidellm/backend/openai.py

@@ -279,11 +279,9 @@ async def default_model(self) -> Optional[str]:
     async def resolve(
         self,
         request: GenerationRequest,
-        request_info: ScheduledRequestInfo[GenerationRequestTimings],
+        request_info: ScheduledRequestInfo,
         history: Optional[list[tuple[GenerationRequest, GenerationResponse]]] = None,
-    ) -> AsyncIterator[
-        tuple[GenerationResponse, ScheduledRequestInfo[GenerationRequestTimings]]
-    ]:
+    ) -> AsyncIterator[tuple[GenerationResponse, ScheduledRequestInfo]]:
         """
         Process a generation request and yield progressive responses.
 

src/guidellm/benchmark/__init__.py

@@ -40,9 +40,6 @@
     BenchmarkerProgressGroup,
     GenerativeConsoleBenchmarkerProgress,
 )
-from .scheduler_registry import scheduler_register_benchmark_objects
-
-scheduler_register_benchmark_objects()
 
 __all__ = [
     "Aggregator",

src/guidellm/benchmark/aggregator.py

@@ -38,7 +38,6 @@
 
 from guidellm.backend import (
     GenerationRequest,
-    GenerationRequestTimings,
     GenerationResponse,
 )
 from guidellm.benchmark.objects import (
@@ -47,7 +46,6 @@
     GenerativeRequestStats,
 )
 from guidellm.scheduler import (
-    MeasuredRequestTimingsT,
     RequestT,
     ResponseT,
     ScheduledRequestInfo,
@@ -153,7 +151,7 @@ def get_metric(
 
 
 @runtime_checkable
-class Aggregator(Protocol[ResponseT, RequestT, MeasuredRequestTimingsT]):
+class Aggregator(Protocol[ResponseT, RequestT]):
     """
     Protocol for processing benchmark data updates during execution.
 
@@ -167,7 +165,7 @@ def __call__(
         state: AggregatorState,
         response: ResponseT | None,
         request: RequestT,
-        request_info: ScheduledRequestInfo[MeasuredRequestTimingsT],
+        request_info: ScheduledRequestInfo,
         scheduler_state: SchedulerState,
     ) -> dict[str, Any] | None:
         """
@@ -183,7 +181,7 @@ def __call__(
 
 
 @runtime_checkable
-class CompilableAggregator(Protocol[ResponseT, RequestT, MeasuredRequestTimingsT]):
+class CompilableAggregator(Protocol[ResponseT, RequestT]):
     """
     Protocol for aggregators that compile final results from aggregated state.
 
@@ -196,7 +194,7 @@ def __call__(
         state: AggregatorState,
         response: ResponseT | None,
         request: RequestT,
-        request_info: ScheduledRequestInfo[MeasuredRequestTimingsT],
+        request_info: ScheduledRequestInfo,
         scheduler_state: SchedulerState,
     ) -> dict[str, Any] | None:
         """
@@ -225,7 +223,7 @@ def compile(
 class SerializableAggregator(
     PydanticClassRegistryMixin[type["SerializableAggregator"]],
     ABC,
-    Generic[ResponseT, RequestT, MeasuredRequestTimingsT],
+    Generic[ResponseT, RequestT],
 ):
     schema_discriminator: ClassVar[str] = "type_"
 
@@ -286,7 +284,7 @@ def __call__(
         state: AggregatorState,
         response: ResponseT | None,
         request: RequestT,
-        request_info: ScheduledRequestInfo[MeasuredRequestTimingsT],
+        request_info: ScheduledRequestInfo,
         scheduler_state: SchedulerState,
     ) -> dict[str, Any] | None:
         """
@@ -314,9 +312,7 @@ def compile(
 
 
 @SerializableAggregator.register("inject_extras")
-class InjectExtrasAggregator(
-    SerializableAggregator[ResponseT, RequestT, MeasuredRequestTimingsT], InfoMixin
-):
+class InjectExtrasAggregator(SerializableAggregator[ResponseT, RequestT], InfoMixin):
     """
     Aggregator for injecting extra metadata into the output.
     """
@@ -333,7 +329,7 @@ def __call__(
         state: AggregatorState,
         response: ResponseT | None,
         request: RequestT,
-        request_info: ScheduledRequestInfo[MeasuredRequestTimingsT],
+        request_info: ScheduledRequestInfo,
         scheduler_state: SchedulerState,
     ) -> dict[str, Any] | None:
         """
@@ -355,9 +351,7 @@ def compile(
 
 
 @SerializableAggregator.register("scheduler_stats")
-class SchedulerStatsAggregator(
-    SerializableAggregator[ResponseT, RequestT, MeasuredRequestTimingsT], InfoMixin
-):
+class SchedulerStatsAggregator(SerializableAggregator[ResponseT, RequestT], InfoMixin):
     """
     Aggregates scheduler timing and performance metrics.
 
@@ -376,7 +370,7 @@ def __call__(
         state: AggregatorState,
         response: ResponseT | None,
         request: RequestT,
-        request_info: ScheduledRequestInfo[MeasuredRequestTimingsT],
+        request_info: ScheduledRequestInfo,
         scheduler_state: SchedulerState,
     ) -> dict[str, Any] | None:
         """
@@ -499,9 +493,7 @@ def compile(
 
 @SerializableAggregator.register("generative_stats_progress")
 class GenerativeStatsProgressAggregator(
-    SerializableAggregator[
-        GenerationResponse, GenerationRequest, GenerationRequestTimings
-    ]
+    SerializableAggregator[GenerationResponse, GenerationRequest]
 ):
     """
     Tracks generative model metrics during benchmark execution.
@@ -523,7 +515,7 @@ def __call__(
         state: AggregatorState,
         response: GenerationResponse | None,
         request: GenerationRequest,
-        request_info: ScheduledRequestInfo[GenerationRequestTimings],
+        request_info: ScheduledRequestInfo,
         scheduler_state: SchedulerState,
     ) -> dict[str, Any] | None:
         """
@@ -667,9 +659,7 @@ def compile(
 
 @SerializableAggregator.register("generative_requests")
 class GenerativeRequestsAggregator(
-    SerializableAggregator[
-        GenerationResponse, GenerationRequest, GenerationRequestTimings
-    ],
+    SerializableAggregator[GenerationResponse, GenerationRequest],
 ):
     """
     Compiles complete generative benchmark results with warmup/cooldown filtering.
@@ -712,7 +702,7 @@ def __call__(
         state: AggregatorState,
         response: GenerationResponse | None,
         request: GenerationRequest,
-        request_info: ScheduledRequestInfo[GenerationRequestTimings],
+        request_info: ScheduledRequestInfo,
         scheduler_state: SchedulerState,
     ) -> dict[str, Any] | None:
         """
@@ -875,7 +865,7 @@ def compile(
 
     def _is_in_warmup(
         self,
-        request_info: ScheduledRequestInfo[GenerationRequestTimings],
+        request_info: ScheduledRequestInfo,
         scheduler_state: SchedulerState,
     ) -> bool:
         """Check if the current request is within the warmup period."""
@@ -902,7 +892,7 @@ def _is_in_warmup(
 
     def _is_in_cooldown(
         self,
-        request_info: ScheduledRequestInfo[GenerationRequestTimings],
+        request_info: ScheduledRequestInfo,
         scheduler_state: SchedulerState,
     ) -> bool:
         """Check if the current request is within the cooldown period."""
@@ -936,7 +926,7 @@ def _create_generative_request_stats(
         cls,
         response: GenerationResponse,
         request: GenerationRequest,
-        request_info: ScheduledRequestInfo[GenerationRequestTimings],
+        request_info: ScheduledRequestInfo,
     ) -> GenerativeRequestStats:
         prompt_tokens = response.preferred_prompt_tokens(
             settings.preferred_prompt_tokens_source

src/guidellm/benchmark/benchmarker.py

@@ -36,7 +36,6 @@
     BackendInterface,
     Constraint,
     Environment,
-    MeasuredRequestTimingsT,
     NonDistributedEnvironment,
     RequestT,
     ResponseT,
@@ -51,7 +50,7 @@
 
 
 class Benchmarker(
-    Generic[BenchmarkT, RequestT, MeasuredRequestTimingsT, ResponseT],
+    Generic[BenchmarkT, RequestT, ResponseT],
     ABC,
     ThreadSafeSingletonMixin,
 ):
@@ -69,13 +68,12 @@ class Benchmarker(
     async def run(
         self,
         requests: Iterable[RequestT | Iterable[RequestT | tuple[RequestT, float]]],
-        backend: BackendInterface[RequestT, MeasuredRequestTimingsT, ResponseT],
+        backend: BackendInterface[RequestT, ResponseT],
         profile: Profile,
         benchmark_class: type[BenchmarkT],
         benchmark_aggregators: dict[
             str,
-            Aggregator[ResponseT, RequestT, MeasuredRequestTimingsT]
-            | CompilableAggregator[ResponseT, RequestT, MeasuredRequestTimingsT],
+            Aggregator[ResponseT, RequestT] | CompilableAggregator[ResponseT, RequestT],
         ],
         environment: Environment | None = None,
     ) -> AsyncIterator[
@@ -121,7 +119,7 @@ async def run(
                     request,
                     request_info,
                     scheduler_state,
-                ) in Scheduler[RequestT, MeasuredRequestTimingsT, ResponseT]().run(
+                ) in Scheduler[RequestT, ResponseT]().run(
                     requests=requests,
                     backend=backend,
                     strategy=strategy,
@@ -170,12 +168,11 @@ def _compile_benchmark_kwargs(
         run_index: int,
         profile: Profile,
         requests: Iterable[RequestT | Iterable[RequestT | tuple[RequestT, float]]],
-        backend: BackendInterface[RequestT, MeasuredRequestTimingsT, ResponseT],
+        backend: BackendInterface[RequestT, ResponseT],
         environment: Environment,
         aggregators: dict[
             str,
-            Aggregator[ResponseT, RequestT, MeasuredRequestTimingsT]
-            | CompilableAggregator[ResponseT, RequestT, MeasuredRequestTimingsT],
+            Aggregator[ResponseT, RequestT] | CompilableAggregator[ResponseT, RequestT],
         ],
         aggregators_state: dict[str, dict[str, Any]],
         strategy: SchedulingStrategy,

src/guidellm/benchmark/entrypoints.py

@@ -13,7 +13,6 @@
     Backend,
     BackendType,
     GenerationRequest,
-    GenerationRequestTimings,
     GenerationResponse,
 )
 from guidellm.benchmark.aggregator import (
@@ -266,7 +265,6 @@ async def benchmark_generative_text(  # noqa: C901
         Benchmarker[
             GenerativeBenchmark,
             GenerationRequest,
-            GenerationRequestTimings,
             GenerationResponse,
         ]().run(
             requests=request_loader,

src/guidellm/benchmark/objects.py

@@ -31,7 +31,6 @@
 import yaml
 from pydantic import Field, computed_field
 
-from guidellm.backend import GenerationRequestTimings
 from guidellm.benchmark.profile import (
     Profile,
 )
@@ -134,7 +133,7 @@ class BenchmarkMetrics(StandardBaseDict):
 class BenchmarkRequestStats(StandardBaseDict):
     """Individual request processing statistics and scheduling metadata."""
 
-    scheduler_info: ScheduledRequestInfo[GenerationRequestTimings] = Field(
+    scheduler_info: ScheduledRequestInfo = Field(
         description="Scheduler metadata and timing information for the request"
     )