From 155583f49f9a20477ea95a0119a7abfddbf0c646 Mon Sep 17 00:00:00 2001 From: Chenguang Zheng <645327136@qq.com> Date: Mon, 13 Apr 2026 18:35:59 +0800 Subject: [PATCH] [Bugfix] Release stage launch lock before handshake (#2717) Signed-off-by: Chenguang ZHENG <645327136@qq.com> --- .../test_async_omni_engine_stage_init.py | 89 +++++++++++++++++++ vllm_omni/engine/async_omni_engine.py | 23 ++--- 2 files changed, 101 insertions(+), 11 deletions(-) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 6993f391e..7b995fe70 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -227,6 +227,95 @@ def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake( assert captured_timeout == 302 +def test_launch_llm_stage_releases_launch_lock_before_complete_stage_handshake(monkeypatch): + """Regression test for parallel LLM stage startup during handshake wait.""" + import vllm_omni.engine.async_omni_engine as engine_mod + from vllm_omni.platforms import current_omni_platform + + engine = object.__new__(AsyncOmniEngine) + engine.log_stats = False + engine.model = "dummy-model" + engine.single_stage_mode = False + engine._omni_master_server = None + + fake_vllm_config = types.SimpleNamespace() + fake_addresses = types.SimpleNamespace() + shared_launch_lock = threading.Lock() + counter_lock = threading.Lock() + first_handshake_started = threading.Event() + second_stage_spawned = threading.Event() + allow_first_handshake_to_finish = threading.Event() + launch_errors: list[BaseException] = [] + spawn_count = 0 + + device_env_var = current_omni_platform.device_control_env_var + prev_device_env = os.environ.get(device_env_var) + os.environ[device_env_var] = "0" + + monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None) + monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) + monkeypatch.setattr(engine_mod, "build_vllm_config", lambda *_, **__: (fake_vllm_config, object)) + monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: []) + + def _spawn_stage_core(**_): + nonlocal spawn_count + with counter_lock: + spawn_count += 1 + call_idx = spawn_count + if call_idx == 2: + second_stage_spawned.set() + return fake_addresses, types.SimpleNamespace(), f"ipc://handshake-{call_idx}" + + def _complete_stage_handshake(_proc, handshake_address, _addresses, _vllm_cfg, _timeout): + if handshake_address == "ipc://handshake-1": + first_handshake_started.set() + assert second_stage_spawned.wait(timeout=1), ( + "second stage did not reach spawn_stage_core while first stage waited in handshake" + ) + assert allow_first_handshake_to_finish.wait(timeout=1), ( + "second stage did not enter handshake while first stage was still waiting" + ) + else: + allow_first_handshake_to_finish.set() + + monkeypatch.setattr(engine_mod, "spawn_stage_core", _spawn_stage_core) + monkeypatch.setattr(engine_mod, "complete_stage_handshake", _complete_stage_handshake) + + def _launch_stage(stage_id: int) -> None: + metadata = types.SimpleNamespace(stage_id=stage_id, runtime_cfg={"devices": str(stage_id)}) + try: + engine._launch_llm_stage( + stage_cfg=types.SimpleNamespace(engine_args={}), + metadata=metadata, + stage_connector_spec={}, + stage_init_timeout=302, + llm_stage_launch_lock=shared_launch_lock, + ) + except BaseException as exc: # pragma: no cover - surfaced through assertion below + launch_errors.append(exc) + + try: + first_thread = threading.Thread(target=_launch_stage, args=(0,)) + first_thread.start() + assert first_handshake_started.wait(timeout=1), "first stage never entered handshake" + + second_thread = threading.Thread(target=_launch_stage, args=(1,)) + second_thread.start() + + first_thread.join(timeout=3) + second_thread.join(timeout=3) + finally: + if prev_device_env is None: + os.environ.pop(device_env_var, None) + else: + os.environ[device_env_var] = prev_device_env + + assert not first_thread.is_alive() + assert not second_thread.is_alive() + assert second_stage_spawned.is_set() + assert not launch_errors + + def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch): """Regression test for GLM-Image t2i preprocessing path. diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 0a2e02d66..9609cf6e2 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -424,23 +424,24 @@ class AsyncOmniEngine: proc=proc, ) logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id) - # Keep the stage-specific device visibility until vLLM - # finishes starting all child processes. - if self.single_stage_mode and self._omni_master_server is not None: - launch_stack.close() - else: - assert proc is not None - assert handshake_address is not None - complete_stage_handshake( - proc, handshake_address, addresses, vllm_config, stage_init_timeout - ) - logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id) finally: if previous_visible_devices is None: current_omni_platform.unset_device_control_env_var() else: current_omni_platform.set_device_control_env_var(previous_visible_devices) + # After StageEngineCoreProc has been spawned it carries its + # stage-specific device visibility into descendants, so the + # slow HELLO/READY handshake can run without holding the + # process-wide launch lock. + if self.single_stage_mode and self._omni_master_server is not None: + launch_stack.close() + else: + assert proc is not None + assert handshake_address is not None + complete_stage_handshake(proc, handshake_address, addresses, vllm_config, stage_init_timeout) + logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id) + assert started_stage is not None return started_stage except Exception: