From 194fa159a3d269022a451b20fc228f487321d9b9 Mon Sep 17 00:00:00 2001 From: Gao Han Date: Wed, 31 Dec 2025 03:23:54 +0800 Subject: [PATCH] [Bugfix]fix proces close problem for omni class (#554) Signed-off-by: Gaohan123 --- vllm_omni/entrypoints/omni.py | 8 +++++++- vllm_omni/entrypoints/omni_stage.py | 2 +- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index b9db56c78..e1af5cf6c 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -430,7 +430,13 @@ class Omni(OmniBase): per_stage_params.append(self.default_sampling_params_list[stage_id]) sampling_params_list = per_stage_params - return self._run_generation(prompts, sampling_params_list) + try: + return self._run_generation(prompts, sampling_params_list) + except Exception as e: + logger.exception("[Orchestrator] Failed to run generation: %s", e) + raise e + finally: + self.close() def _run_generation( self, diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index b0c4be08e..f5534290f 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -630,7 +630,7 @@ def _stage_worker( _recv_dequeue_ts = _time.time() if task is None: - logger.error("Received shutdown signal") + logger.info("Received shutdown signal") break batch_tasks: list[dict[str, Any]] = [task]