Almost working e2e

Author: markurtz

src/guidellm/benchmark/__init__.py

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

src/guidellm/benchmark/scheduler_registry.py

@@ -0,0 +1,21 @@
+from guidellm.backend import (
+    GenerationRequest,
+    GenerationRequestTimings,
+    GenerationResponse,
+)
+from guidellm.scheduler import ScheduledRequestInfo, SchedulerMessagingPydanticRegistry
+
+__all__ = ["scheduler_register_benchmark_objects"]
+
+
+def scheduler_register_benchmark_objects():
+    SchedulerMessagingPydanticRegistry.register("GenerationRequest")(GenerationRequest)
+    SchedulerMessagingPydanticRegistry.register("GenerationResponse")(
+        GenerationResponse
+    )
+    SchedulerMessagingPydanticRegistry.register("GenerationRequestTimings")(
+        GenerationRequestTimings
+    )
+    SchedulerMessagingPydanticRegistry.register("ScheduledRequestInfo")(
+        ScheduledRequestInfo
+    )

src/guidellm/scheduler/scheduler.py

@@ -129,6 +129,7 @@ async def run(
                 worker_group = WorkerProcessGroup[
                     RequestT, MeasuredRequestTimingsT, ResponseT
                 ](
+                    requests=None,
                     cycle_requests=local_requests,
                     backend=backend,
                     strategy=local_strategy,

src/guidellm/scheduler/worker_group.py

@@ -578,6 +578,7 @@ def _locked_update(
                 self._internal_constraints.update(add_constraints)
             if update_constraints:
                 self._update_with_constraints(info)
+            self._state.end_time = time.time()
             state_copy: SchedulerState = self._state.model_copy()
 
         return (
@@ -588,23 +589,6 @@ def _locked_update(
             ),
         )
 
-    def _locked_cancel_request(
-        self, info: ScheduledRequestInfo[MeasuredRequestTimingsT]
-    ):
-        if info.status != "queued":
-            raise ValueError(f"Cannot cancel request in {info.status} state")
-
-        with self._update_lock:
-            self._state.queued_requests -= 1
-            self._state.processed_requests += 1
-            self._state.cancelled_requests += 1
-
-            info.status = "cancelled"
-            info.scheduler_timings.resolve_end = time.time()
-            state_copy: SchedulerState = self._state.model_copy()
-
-        return state_copy
-
     def _update_new_request(self):
         self._state.created_requests += 1
         self._state.queued_requests += 1