From cc4ade4ba797b689721fbce539b30a24a895bc9b Mon Sep 17 00:00:00 2001 From: tzhouam Date: Tue, 17 Mar 2026 08:19:56 +0000 Subject: [PATCH] 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 --- vllm_omni/entrypoints/omni_llm.py | 6 +++--- vllm_omni/entrypoints/omni_stage.py | 1 + 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index e06e50c83..f7880eb8a 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/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: diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index b40311d19..1f14d0ab4 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/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