Refactor shutdown logic in OmniLLM and stage worker to ensure proper engine closure. The shutdown method now checks for the presence of engine_core before invoking shutdown, enhancing resource management. Additionally, the stage worker now explicitly calls stage_engine.close() upon receiving a shutdown signal. This improves the overall stability and reliability of the shutdown process.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

vllm_omni/entrypoints/omni_llm.py

@@ -217,10 +217,10 @@ class OmniLLM(LLM):
         Note: Stage management is now handled by Omni class.
         This method closes the LLM engine but not stages.
         """
-        # Close the LLM engine if it exists
         if hasattr(self, "llm_engine") and self.llm_engine is not None:
-            if hasattr(self.llm_engine, "shutdown"):
-                self.llm_engine.shutdown()
+            engine_core = getattr(self.llm_engine, "engine_core", None)
+            if engine_core is not None and hasattr(engine_core, "shutdown"):
+                engine_core.shutdown()
 
     def __del__(self) -> None:  # best-effort
         try:

vllm_omni/entrypoints/omni_stage.py

@@ -968,6 +968,7 @@ def _stage_worker(
         task_type = task.get("type", OmniStageTaskType.GENERATE)
         if task_type == OmniStageTaskType.SHUTDOWN:
             logger.info("Received shutdown signal")
+            stage_engine.close()
             break
 
         # Handle profiler control commands