Drop times queue and handling timing in each worker

Author: sjmonson

src/guidellm/scheduler/result.py

@@ -17,7 +17,7 @@
     "SchedulerRequestResult",
     "SchedulerResult",
     "SchedulerRunInfo",
-    "WorkerProcessRequestTime",
+    "WorkerProcessRequest",
     "WorkerProcessResult",
 ]
 
@@ -147,8 +147,8 @@ class SchedulerRequestResult(
 
 
 @dataclass
-class WorkerProcessRequestTime:
-    start_time: float
+class WorkerProcessRequest(Generic[RequestT, ResponseT]):
+    session: RequestSession[RequestT, ResponseT]
     timeout_time: float
     queued_time: float
 
@@ -163,6 +163,5 @@ class WorkerProcessResult(Generic[RequestT, ResponseT]):
 
 @dataclass
 class MPQueues(Generic[RequestT, ResponseT]):
-    requests: Queue[RequestSession[RequestT, ResponseT]]
-    times: Queue[WorkerProcessRequestTime]
+    requests: Queue[WorkerProcessRequest[RequestT, ResponseT]]
     responses: Queue[WorkerProcessResult[RequestT, ResponseT]]

src/guidellm/scheduler/scheduler.py

@@ -17,7 +17,6 @@
 from loguru import logger
 
 from guidellm.config import settings
-from guidellm.request.session import RequestSession
 from guidellm.request.types import (
     RequestT,
     ResponseT,
@@ -27,7 +26,7 @@
     SchedulerRequestResult,
     SchedulerResult,
     SchedulerRunInfo,
-    WorkerProcessRequestTime,
+    WorkerProcessRequest,
     WorkerProcessResult,
 )
 from guidellm.scheduler.strategy import SchedulingStrategy
@@ -130,7 +129,7 @@ async def run(
             futures, queues, stop_event = await self._start_processes(
                 manager, executor, scheduling_strategy
             )
-            run_info, requests_iter, times_iter = self._run_setup(
+            run_info, requests_iter = self._run_setup(
                 futures, scheduling_strategy, max_number, max_duration
             )
             yield SchedulerResult(
@@ -155,9 +154,7 @@ async def run(
 
                     requests_iter = self._add_requests(
                         requests_iter,
-                        times_iter,
                         queues.requests,
-                        queues.times,
                         run_info,
                     )
                     await asyncio.sleep(0)  # enable requests to start
@@ -196,7 +193,6 @@ async def _start_processes(
             requests=manager.Queue(
                 maxsize=scheduling_strategy.processing_requests_limit
             ),
-            times=manager.Queue(maxsize=scheduling_strategy.processing_requests_limit),
             responses=manager.Queue(),
         )
         stop_event = manager.Event()
@@ -214,6 +210,7 @@ async def _start_processes(
             % scheduling_strategy.processes_limit
         )
         process_ids = (id_ for id_ in range(num_processes))
+        start_time = time.time()
         process_requests_limits = (
             requests_limit_split + 1
             if i < requests_limit_remain
@@ -229,10 +226,13 @@ async def _start_processes(
                     executor,
                     self.worker.process_loop_asynchronous,
                     queues,
+                    scheduling_strategy,
+                    start_time,
                     stop_event,
                     False,  # TODO: Make configurable
                     requests_limit,
                     id_,
+                    num_processes,
                 )
             )
 
@@ -246,10 +246,9 @@ def _run_setup(
         scheduling_strategy: SchedulingStrategy,
         max_number: Optional[int],
         max_duration: Optional[float],
-    ) -> tuple[SchedulerRunInfo, Iterator[Any], Iterator[float]]:
+    ) -> tuple[SchedulerRunInfo, Iterator[Any]]:
         requests_iter = iter(self.request_loader)
         start_time = time.time()
-        times_iter = iter(scheduling_strategy.request_times())
         end_time = time.time() + (max_duration or math.inf)
         end_number = max_number or math.inf
 
@@ -275,14 +274,12 @@ def _run_setup(
             strategy=scheduling_strategy,
         )
 
-        return info, requests_iter, times_iter
+        return info, requests_iter
 
     def _add_requests(
         self,
         requests_iter: Optional[Iterator[Any]],
-        times_iter: Iterator[float],
-        requests_queue: Queue[RequestSession[RequestT, ResponseT]],
-        times_queue: Queue[WorkerProcessRequestTime],
+        requests_queue: Queue[WorkerProcessRequest[RequestT, ResponseT]],
         run_info: SchedulerRunInfo,
     ) -> Optional[Iterator[Any]]:
         if requests_iter is not None:
@@ -296,24 +293,20 @@ def _add_requests(
                     if run_info.created_requests >= run_info.end_number:
                         raise StopIteration
 
+                    if time.time() >= run_info.end_time:
+                        raise StopIteration
+
                     session = next(requests_iter)
-                    requests_queue.put(session)
-                    for _ in range(len(session)):
-                        if (
-                            request_time := next(times_iter)
-                        ) >= run_info.end_time or time.time() >= run_info.end_time:
-                            raise StopIteration
-
-                        work_req = WorkerProcessRequestTime(
-                            start_time=request_time,
-                            timeout_time=run_info.end_time,
-                            queued_time=time.time(),
-                        )
-                        times_queue.put(work_req)
-
-                        run_info.created_requests += 1
-                        run_info.queued_requests += 1
-                        added_count += 1
+                    work_req = WorkerProcessRequest(
+                        session=session,
+                        timeout_time=run_info.end_time,
+                        queued_time=time.time(),
+                    )
+                    requests_queue.put(work_req)
+
+                    run_info.created_requests += len(session)
+                    run_info.queued_requests += len(session)
+                    added_count += len(session)
             except StopIteration:
                 # we've reached the limit number, limit time, or exhausted the requests
                 # set to None to stop adding more and tell the loop no more requests

src/guidellm/scheduler/strategy.py

@@ -94,7 +94,9 @@ def processing_requests_limit(self) -> int:
         """
         return settings.max_concurrency
 
-    def request_times(self) -> Generator[float, None, None]:
+    def request_times(
+        self, start_time: Optional[float] = None
+    ) -> Generator[float, None, None]:
         """
         A generator that yields timestamps for when requests should be sent.
         This method should be implemented by subclasses to provide specific
@@ -168,7 +170,9 @@ def processing_requests_limit(self) -> int:
         """
         return 1
 
-    def request_times(self) -> Generator[float, None, None]:
+    def request_times(
+        self, start_time: Optional[float] = None
+    ) -> Generator[float, None, None]:
         """
         A generator that yields time.time() so requests are sent immediately,
             while scheduling them synchronously.
@@ -254,7 +258,9 @@ def processing_requests_limit(self) -> int:
         """
         return self.streams
 
-    def request_times(self) -> Generator[float, None, None]:
+    def request_times(
+        self, start_time: Optional[float] = None
+    ) -> Generator[float, None, None]:
         """
         A generator that yields time.time() so requests are sent
         immediately, while scheduling them concurrently with the specified
@@ -328,15 +334,18 @@ def processing_requests_limit(self) -> int:
         """
         return self.max_concurrency or super().processing_requests_limit
 
-    def request_times(self) -> Generator[float, None, None]:
+    def request_times(
+        self, start_time: Optional[float] = None
+    ) -> Generator[float, None, None]:
         """
         A generator that yields the start time.time() so requests are sent
         immediately, while scheduling as many asynchronously as possible.
 
         :return: A generator that yields the start time.time()
             for immediate request scheduling.
         """
-        start_time = time.time()
+        if start_time is None:
+            start_time = time.time()
 
         while True:
             yield start_time
@@ -379,7 +388,9 @@ class AsyncConstantStrategy(ThroughputStrategy):
         ),
     )
 
-    def request_times(self) -> Generator[float, None, None]:
+    def request_times(
+        self, start_time: Optional[float] = None
+    ) -> Generator[float, None, None]:
         """
         A generator that yields timestamps for when requests should be sent.
         This method schedules requests asynchronously at a constant rate
@@ -391,7 +402,8 @@ def request_times(self) -> Generator[float, None, None]:
 
         :return: A generator that yields timestamps for request scheduling.
         """
-        start_time = time.time()
+        if start_time is None:
+            start_time = time.time()
         constant_increment = 1.0 / self.rate
 
         # handle bursts first to get to the desired rate
@@ -451,7 +463,9 @@ class AsyncPoissonStrategy(ThroughputStrategy):
         description=("The random seed to use for the Poisson distribution. "),
     )
 
-    def request_times(self) -> Generator[float, None, None]:
+    def request_times(
+        self, start_time: Optional[float] = None
+    ) -> Generator[float, None, None]:
         """
         A generator that yields timestamps for when requests should be sent.
         This method schedules requests asynchronously at a Poisson rate
@@ -461,7 +475,8 @@ def request_times(self) -> Generator[float, None, None]:
 
         :return: A generator that yields timestamps for request scheduling.
         """
-        start_time = time.time()
+        if start_time is None:
+            start_time = time.time()
 
         if self.initial_burst is not None:
             # send an initial burst equal to the rate