Add stop event

Signed-off-by: Samuel Monson <[email protected]>

Author: sjmonson

src/guidellm/scheduler/scheduler.py

@@ -6,6 +6,7 @@
 from multiprocessing import Manager
 from queue import Empty as QueueEmpty
 from queue import Queue
+from threading import Event
 from typing import (
     Any,
     Generic,
@@ -126,7 +127,7 @@ async def run(
             ) as executor,
         ):
             requests_iter: Optional[Iterator[Any]] = None
-            futures, queues = await self._start_processes(
+            futures, queues, stop_event = await self._start_processes(
                 manager, executor, scheduling_strategy
             )
             run_info, requests_iter, times_iter = self._run_setup(
@@ -178,7 +179,7 @@ async def run(
                 run_info=run_info,
             )
 
-            await self._stop_processes(futures, queues.requests)
+            await self._stop_processes(futures, stop_event)
 
     async def _start_processes(
         self,
@@ -188,6 +189,7 @@ async def _start_processes(
     ) -> tuple[
         list[asyncio.Future],
         MPQueues[RequestT, ResponseT],
+        Event,
     ]:
         await self.worker.prepare_multiprocessing()
         queues: MPQueues[RequestT, ResponseT] = MPQueues(
@@ -197,6 +199,7 @@ async def _start_processes(
             times=manager.Queue(maxsize=scheduling_strategy.processing_requests_limit),
             responses=manager.Queue(),
         )
+        stop_event = manager.Event()
 
         num_processes = min(
             scheduling_strategy.processes_limit,
@@ -226,6 +229,7 @@ async def _start_processes(
                     executor,
                     self.worker.process_loop_asynchronous,
                     queues,
+                    stop_event,
                     False,  # TODO: Make configurable
                     requests_limit,
                     id_,
@@ -234,7 +238,7 @@ async def _start_processes(
 
         await asyncio.sleep(0.1)  # give time for processes to start
 
-        return futures, queues
+        return futures, queues, stop_event
 
     def _run_setup(
         self,
@@ -369,10 +373,9 @@ def _check_result_ready(
     async def _stop_processes(
         self,
         futures: list[asyncio.Future],
-        requests_queue: Queue[RequestSession[RequestT, ResponseT]],
+        stop_event: Event,
     ):
-        # FIXME: Need new method for stopping workers
-        for _ in futures:
-            requests_queue.put(None)
+        # stop all processes
+        stop_event.set()
 
         await asyncio.gather(*futures)

src/guidellm/scheduler/worker.py

@@ -6,6 +6,7 @@
 from dataclasses import dataclass
 from queue import Empty as QueueEmpty
 from queue import Queue
+from threading import Event
 from typing import (
     Any,
     Generic,
@@ -109,9 +110,11 @@ async def resolve(
         ...
 
     async def get_request_time(
-        self, times_queue: Queue[WorkerProcessRequestTime]
+        self,
+        times_queue: Queue[WorkerProcessRequestTime],
+        timeout: Optional[int] = None,
     ) -> WorkerProcessRequestTime:
-        return await asyncio.to_thread(times_queue.get)  # type: ignore[attr-defined]
+        return await asyncio.to_thread(times_queue.get, timeout=timeout)  # type: ignore[attr-defined]
 
     async def send_result(
         self,
@@ -181,6 +184,7 @@ async def resolve_scheduler_request(
     def process_loop_asynchronous(
         self,
         queues: MPQueues[RequestT, ResponseT],
+        stop_event: Event,
         prioritize_sessions: bool,
         max_concurrency: int,
         process_id: int,
@@ -189,7 +193,7 @@ async def _process_runner():
             lock = asyncio.Semaphore(max_concurrency)
             pending_sessions: list[RequestSession[RequestT, ResponseT]] = []
 
-            while True:  # TODO: Exit condition
+            while True:
                 await asyncio.sleep(0)  # Yield control to the event loop
                 await lock.acquire()
 
@@ -201,13 +205,16 @@ async def _process_runner():
                         else queues.requests.get_nowait()
                     )
                     dequeued_time = time.time()
-                    request_times = await self.get_request_time(queues.times)
+                    request_times = await self.get_request_time(queues.times, 5)
                 except (QueueEmpty, IndexError):
                     # Requeue the session if we don't have a next time yet
                     if request_session is not None:
                         pending_sessions.append(request_session)
                     lock.release()
-                    continue
+                    if stop_event.is_set():
+                        return  # Exit if stop event is set
+                    else:
+                        continue
 
                 async def wait_then_requeue(
                     session: RequestSession[RequestT, ResponseT],
@@ -309,13 +316,15 @@ async def prepare_multiprocessing(self):
     def process_loop_asynchronous(
         self,
         queues: MPQueues[GenerationRequest, ResponseSummary],
+        stop_event: Event,
         prioritize_sessions: bool,
         max_concurrency: int,
         process_id: int,
     ):
         asyncio.run(self.backend.validate())
         super().process_loop_asynchronous(
             queues=queues,
+            stop_event=stop_event,
             prioritize_sessions=prioritize_sessions,
             max_concurrency=max_concurrency,
             process_id=process_id,