Sleep interminetly

Author: markVaykhansky

src/guidellm/scheduler/worker.py

@@ -165,6 +165,7 @@ async def resolve_scheduler_request(
         timeout_time: float,
         results_queue: multiprocessing.Queue,
         process_id: int,
+        shutdown_event: Optional[MultiprocessingEvent] = None,
     ):
         info = SchedulerRequestInfo(
             targeted_start_time=start_time,
@@ -183,7 +184,21 @@ async def resolve_scheduler_request(
         asyncio.create_task(self.send_result(results_queue, request_scheduled_result))
 
         if (wait_time := start_time - time.time()) > 0:
-            await asyncio.sleep(wait_time)
+            if shutdown_event is None:
+                await asyncio.sleep(wait_time)
+            else:
+                shutdown_signal_received = \
+                    await self._sleep_intermittently_until_timestamp_or_shutdown(
+                       sleep_until_timestamp=start_time,
+                       shutdown_event=shutdown_event,
+                    )
+                if shutdown_signal_received:
+                    logger.info(
+                        "Received shutdown signal "
+                        "while waiting to start "
+                        f"|| Process ID {process_id}"
+                    )
+                    return
 
         info.worker_start = time.time()
         request_start_result: WorkerProcessResult[RequestT, ResponseT] = \
@@ -211,6 +226,18 @@ async def resolve_scheduler_request(
         )
         asyncio.create_task(self.send_result(results_queue, result))
 
+    async def _sleep_intermittently_until_timestamp_or_shutdown(
+            self,
+            sleep_until_timestamp: float,
+            shutdown_event: MultiprocessingEvent,
+    ) -> bool:
+        delta = timedelta(seconds=10).total_seconds()
+        while time.time() < sleep_until_timestamp:
+            await asyncio.sleep(delta)
+            if shutdown_event.is_set():
+                return True
+        return False
+
     def process_loop_synchronous(
         self,
         requests_queue: multiprocessing.Queue,
@@ -240,6 +267,7 @@ async def _process_runner():
                     timeout_time=process_request.timeout_time,
                     results_queue=results_queue,
                     process_id=process_id,
+                    shutdown_event=shutdown_event,
                 )
 
         try:
@@ -271,10 +299,26 @@ async def _process_runner():
                     shutdown_event=shutdown_event,
                     process_id=process_id)
             ) is not None:
+                if shutdown_event and shutdown_event.is_set():
+                    logger.error("This shouldn't happen! "
+                                 "We should catch the "
+                                 "shutdown in the get wrapper")
+                    logger.info(f"Shutdown signal received"
+                                f" in future {process_id}")
+                    break
+
                 dequeued_time = time.time()
+                logger.debug(f"Dequeued Process ID {process_id} || "
+                             f"Timestamp {dequeued_time} || "
+                             f"Semaphore {pending._value}/{max_concurrency}")
 
                 await pending.acquire()
 
+                lock_acquired_at = time.time()
+                logger.debug(f"Lock acquired Process ID {process_id} ||"
+                             f" Timestamp {lock_acquired_at} ||"
+                             f" Semaphore {pending._value}/{max_concurrency}")
+
                 def _task_done(_: asyncio.Task):
                     nonlocal pending
                     pending.release()
@@ -292,6 +336,7 @@ def _task_done(_: asyncio.Task):
                         timeout_time=process_request.timeout_time,
                         results_queue=results_queue,
                         process_id=process_id,
+                        shutdown_event=shutdown_event,
                     )
                 )
                 task.add_done_callback(_task_done)