[Bugfix] Release stage launch lock before handshake (#2717)

Signed-off-by: Chenguang ZHENG <645327136@qq.com>

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.
 

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: