From ac174317455a4e1a3ee24e9ccd6d66c2ddfd8dd4 Mon Sep 17 00:00:00 2001 From: Daniel Huang Date: Wed, 11 Mar 2026 19:40:42 -0700 Subject: [PATCH] [Enhancement] Patch OmniStage.try_collect() with _proc alive checks (#1560) Signed-off-by: Daniel Huang --- tests/entrypoints/test_omni_diffusion.py | 2 +- vllm_omni/entrypoints/omni_stage.py | 10 ++++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/test_omni_diffusion.py b/tests/entrypoints/test_omni_diffusion.py index 90b0fd053..9e555fa85 100644 --- a/tests/entrypoints/test_omni_diffusion.py +++ b/tests/entrypoints/test_omni_diffusion.py @@ -345,7 +345,7 @@ def _setup_multiprocessing_mocks(monkeypatch: pytest.MonkeyPatch, mocker: Mocker fake_process_instance = mocker.MagicMock() fake_process_instance.start = mocker.MagicMock() fake_process_instance.join = mocker.MagicMock() - fake_process_instance.is_alive = mocker.MagicMock(return_value=False) + fake_process_instance.is_alive = mocker.MagicMock(return_value=True) fake_process_instance.terminate = mocker.MagicMock() fake_process_class.return_value = fake_process_instance diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 3fa722b03..19091034a 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -693,8 +693,14 @@ class OmniStage: assert self._out_q is not None try: return self._out_q.get_nowait() - except Exception: - return None + except queue.Empty: + pass + except Exception as e: + logger.error("Unexpected error when collecting OmniStage output queue:", exc_info=e) + self.stop_stage_worker() + raise + if self._proc is not None and not self._proc.is_alive(): + raise RuntimeError(f"OmniStage Worker process died unexpectedly with exit code {self._proc.exitcode}") def process_engine_inputs( self,