fix: fix vllm graceful shutdown (#5818) (#5835)

Co-authored-by: dagil-nvidia <[email protected]>

Author: jh-nv

components/src/dynamo/vllm/engine_monitor.py

@@ -25,7 +25,12 @@ class VllmEngineMonitor:
     Monitors the health of the vLLM engine and initiates a shutdown if the engine is dead.
     """
 
-    def __init__(self, runtime: DistributedRuntime, engine_client: AsyncLLM):
+    def __init__(
+        self,
+        runtime: DistributedRuntime,
+        engine_client: AsyncLLM,
+        shutdown_event: asyncio.Event = None,
+    ):
         if not isinstance(runtime, DistributedRuntime):
             raise ValueError(
                 f"{self.__class__.__name__} requires an instance of DistributedRuntime."
@@ -37,6 +42,7 @@ def __init__(self, runtime: DistributedRuntime, engine_client: AsyncLLM):
 
         self.runtime = runtime
         self.engine_client = engine_client
+        self.shutdown_event = shutdown_event
         self._monitor_task = asyncio.create_task(self._check_engine_health())
 
         logger.info(
@@ -66,10 +72,41 @@ def timeout_handler(signum, frame):
             signal.alarm(0)
 
     async def _check_engine_health(self):
+        """
+        Continuously check engine health until:
+        1. Engine dies (EngineDeadError) - initiate shutdown
+        2. Shutdown event is triggered - stop monitoring gracefully
+        3. Task is cancelled - cleanup
+        """
         while True:
             try:
+                # Check if shutdown event was triggered - stop monitoring
+                if self.shutdown_event and self.shutdown_event.is_set():
+                    logger.info(
+                        f"{self.__class__.__name__}: Shutdown event detected, stopping engine health monitoring."
+                    )
+                    break
+
                 await self.engine_client.check_health()
-                await asyncio.sleep(HEALTH_CHECK_INTERVAL)
+
+                # Sleep with shutdown event awareness for faster response
+                if self.shutdown_event:
+                    try:
+                        await asyncio.wait_for(
+                            self.shutdown_event.wait(), timeout=HEALTH_CHECK_INTERVAL
+                        )
+                        # Shutdown event was set during sleep
+                        logger.info(
+                            f"{self.__class__.__name__}: Shutdown event detected, stopping engine health monitoring."
+                        )
+                        break
+                    except asyncio.TimeoutError:
+                        # Normal timeout, continue monitoring
+                        pass
+                else:
+                    # No shutdown event, just sleep normally
+                    await asyncio.sleep(HEALTH_CHECK_INTERVAL)
+
             except EngineDeadError as e:
                 logger.error(f"Traceback: {traceback.format_exc()}")
                 logger.error(f"vLLM AsyncLLM health check failed: {e}")
@@ -78,4 +115,5 @@ async def _check_engine_health(self):
                 self.runtime.shutdown()
                 os._exit(1)
             except asyncio.CancelledError:
-                pass
+                logger.debug(f"{self.__class__.__name__}: Health check task cancelled.")
+                break