WIP - New cancellation mechanism works

Author: markVaykhansky

src/guidellm/benchmark/aggregator.py

@@ -638,7 +638,7 @@ def _calculate_error_rate(self) -> float:
         total_successful = self.requests_stats.totals.successful.total
         total_errored = self.requests_stats.totals.errored.total
         total_finished = total_errored + total_successful
-        return total_errored / total_finished
+        return total_errored / total_finished if total_finished > 0 else 0
 
     def _compile_results(
         self,

src/guidellm/scheduler/scheduler.py

@@ -175,11 +175,6 @@ async def run(
                             and not iter_result.request_info.canceled
                             and self._is_max_error_rate_reached(iter_result.run_info)
                         ):
-                            if shutdown_event is None:
-                                raise RuntimeError(
-                                    "We've reached max_error_rate "
-                                    "but shutdown_event is corrupt"
-                                )
                             shutdown_event.set()
                             max_error_rate_reached = True
                             logger.info(
@@ -199,7 +194,7 @@ async def run(
                 run_info=run_info,
             )
 
-            await self._stop_processes(futures, shutdown_event, requests_queue)
+            await self._stop_processes(futures, shutdown_event)
 
     def _validate_scheduler_params(
         self,
@@ -252,7 +247,6 @@ async def _start_processes(
             scheduling_strategy.processes_limit,
             scheduling_strategy.processing_requests_limit,
         )
-        num_processes = 1
         requests_limit_split = (
             scheduling_strategy.processing_requests_limit
             // scheduling_strategy.processes_limit
@@ -459,8 +453,8 @@ async def _stop_processes(
         self,
         futures: list[asyncio.Future],
         shutdown_event: MultiprocessingEvent,
-        requests_queue: multiprocessing.Queue,
     ):
-        shutdown_event.set()
+        if not shutdown_event.is_set():
+            shutdown_event.set()
         logger.debug("Waiting for futures to shut down")
         await asyncio.gather(*futures)

src/guidellm/scheduler/worker.py

@@ -42,6 +42,10 @@
 ]
 
 
+class ShutdownSignalReceived(Exception):
+    pass
+
+
 @dataclass
 class WorkerProcessRequest(Generic[RequestT]):
     request: RequestT
@@ -132,25 +136,24 @@ async def get_request(
         # if we simply use asyncio.to_thread(requests_queue.get)
         # the cancellation task doesn't propagate because the
         # asyncio.to_thread is blocking
-        return await asyncio.to_thread(requests_queue.get)
-        # def _get_queue_intermittently():
-        #     while True:
-        #         try:
-        #             return requests_queue.get(timeout=shutdown_poll_interval_seconds)
-        #         except queue.Empty:
-        #             logger.info("Checking shutdown even is set in get_request")
-        #             if shutdown_event.is_set():
-        #                 logger.info(f"Shutdown signal received in future {process_id}")
-        #                 raise asyncio.CancelledError()
-        #                 # return None
-        #
-        # try:
-        #     return await asyncio.to_thread(_get_queue_intermittently)  # type: ignore[attr-defined]
-        # except asyncio.CancelledError:
-        #     logger.info("kaki")
-        #     # return None
-        #     raise
-        #     # raise
+        def _get_queue_intermittently():
+            while True:
+                try:
+                    return requests_queue.get(timeout=shutdown_poll_interval_seconds)
+                except queue.Empty:
+                    logger.info("Checking shutdown even is set in get_request")
+                    if shutdown_event.is_set():
+                        logger.info(f"Shutdown signal received in future {process_id}")
+                        raise asyncio.CancelledError()
+                        # return None
+
+        try:
+            return await asyncio.to_thread(_get_queue_intermittently)  # type: ignore[attr-defined]
+        except asyncio.CancelledError:
+            logger.info("kaki")
+            # return None
+            raise
+            # raise
 
     async def send_result(
         self,
@@ -267,7 +270,7 @@ async def _process_runner():
                 ],
                 return_when=asyncio.FIRST_EXCEPTION,
             )
-            logger.info("First exception happened")
+            logger.info(f"First exception happened, done: [{[r.get_name() for r in done]}")
 
             for task in pending:
                 logger.debug(f"Cancelling task {task.get_name()}")
@@ -281,7 +284,7 @@ async def _process_runner():
 
             for task in done:
                 task_exception = task.exception()
-                if not isinstance(task_exception, asyncio.CancelledError):
+                if not isinstance(task_exception, ShutdownSignalReceived):
                     raise task_exception
         try:
             asyncio.run(_process_runner())
@@ -303,7 +306,8 @@ async def _wait_for_shutdown(
             await asyncio.sleep(shutdown_poll_interval)
 
         logger.debug("Shutdown signal received")
-        raise asyncio.CancelledError("Shutdown event set, cancelling process loop.")
+        raise ShutdownSignalReceived("Shutdown event set, cancelling process loop.")
+        # raise asyncio.CancelledError("Shutdown event set, cancelling process loop.")
 
     async def _process_synchronous_requests_loop(
             self,