[Bugfix] Allow LLM replicas on different GPUs to initialize concurrently (#7292)

Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/engine/test_stage_runtime_launch_concurrency.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Check replica startup ordering without model-loading timing assumptions."""
+
+import contextlib
+import fcntl
+import os
+import threading
+from collections.abc import Iterator
+from pathlib import Path
+
+import pytest
+from pytest_mock import MockerFixture
+from vllm.config import VllmConfig
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine.utils import CoreEngineProcManager, EngineZmqAddresses
+
+from vllm_omni.config.stage_config import StageConfig
+from vllm_omni.engine import stage_init_utils, stage_runtime
+from vllm_omni.engine.stage_engine_startup import StageReplicaResources
+from vllm_omni.engine.stage_init_utils import LogicalStageInitPlan, ReplicaInitPlan, StageMetadata
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def launch_runtime(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path, mocker: MockerFixture
+) -> Iterator[stage_runtime.StageRuntime]:
+    runtime = stage_runtime.StageRuntime(
+        stage_configs=[],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=5,
+        async_chunk=False,
+    )
+    # These are virtual device IDs: only the real per-device file locking and
+    # scheduling run. No model, CUDA context, or engine process is created.
+    monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda _stage, cfg: cfg["devices"])
+    monkeypatch.setattr(
+        stage_init_utils, "device_init_lock_path", lambda device: str(tmp_path / f"device-{device}.lock")
+    )
+    mocker.patch.object(
+        stage_runtime.StageEngineCoreClientBase,
+        "make_async_mp_client",
+        return_value=mocker.Mock(spec=stage_runtime.StageClient),
+    )
+    monkeypatch.setenv("VLLM_OMNI_TEST_LAUNCH_ENV", "parent")
+    yield runtime
+    if runtime._stage_init_executor is not None:
+        runtime._stage_init_executor.shutdown(wait=True)
+
+
+def _plan(stage_id: int, device: str) -> LogicalStageInitPlan:
+    metadata = StageMetadata(
+        stage_id=stage_id,
+        stage_type="llm",
+        engine_output_type="text",
+        is_comprehension=False,
+        requires_multimodal_data=False,
+        engine_input_source=[],
+        final_output=False,
+        final_output_type=None,
+        default_sampling_params=SamplingParams(),
+        custom_process_input_func=None,
+        model_stage="dummy",
+        runtime_cfg={"devices": device, "env": {"VLLM_OMNI_TEST_LAUNCH_ENV": str(stage_id)}},
+    )
+    replica = ReplicaInitPlan(
+        replica_id=0,
+        num_replicas=1,
+        launch_mode="local",
+        stage_cfg=StageConfig(stage_id=stage_id, model_stage="dummy", yaml_runtime=metadata.runtime_cfg),
+        metadata=metadata,
+        stage_connector_spec={},
+        omni_kv_connector=(None, None, None),
+        stage_vllm_config=VllmConfig(),
+        executor_class=object,
+        engine_args_dict={},
+    )
+    return LogicalStageInitPlan(stage_idx=stage_id, stage_id=stage_id, replicas=[replica])
+
+
+def _resources(mocker: MockerFixture) -> StageReplicaResources:
+    return StageReplicaResources(
+        manager=mocker.Mock(spec=CoreEngineProcManager),
+        addresses=EngineZmqAddresses(inputs=["unused-input"], outputs=["unused-output"]),
+    )
+
+
+def test_distinct_devices_spawn_before_either_is_ready(
+    launch_runtime: stage_runtime.StageRuntime, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture
+) -> None:
+    spawned = [threading.Event(), threading.Event()]
+
+    @contextlib.contextmanager
+    def launch(*, stage_id: int, stage_visible_devices: str, **_kwargs: object) -> Iterator[StageReplicaResources]:
+        assert os.environ["VLLM_OMNI_TEST_LAUNCH_ENV"] == str(stage_id)
+        spawned[stage_id].set()
+        yield _resources(mocker)
+        # The device remains exclusively locked throughout loading/profiling,
+        # even though another device's replica must be able to start.
+        with open(stage_init_utils.device_init_lock_path(int(stage_visible_devices)), "r+") as lock:
+            with pytest.raises(BlockingIOError):
+                fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        assert spawned[1 - stage_id].wait(5), "A different device's spawn was blocked by the READY wait"
+
+    monkeypatch.setattr(stage_runtime, "launch_stage_replica", launch)
+    clients = launch_runtime._initialize_stage_replicas([_plan(0, "0"), _plan(1, "1")], 5)
+    assert all(clients[stage][0] is not None for stage in (0, 1))
+    assert os.environ["VLLM_OMNI_TEST_LAUNCH_ENV"] == "parent"
+
+
+def test_same_device_waits_for_ready_before_next_spawn(
+    launch_runtime: stage_runtime.StageRuntime, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture
+) -> None:
+    events: list[tuple[int, str]] = []
+
+    @contextlib.contextmanager
+    def launch(*, stage_id: int, **_kwargs: object) -> Iterator[StageReplicaResources]:
+        events.append((stage_id, "spawn"))
+        yield _resources(mocker)
+        events.append((stage_id, "ready"))
+
+    monkeypatch.setattr(stage_runtime, "launch_stage_replica", launch)
+    launch_runtime._initialize_stage_replicas([_plan(0, "0"), _plan(1, "0")], 5)
+    assert events == [(0, "spawn"), (0, "ready"), (1, "spawn"), (1, "ready")]
+    assert os.environ["VLLM_OMNI_TEST_LAUNCH_ENV"] == "parent"
+
+
+def test_ready_failure_cleans_up_resources_and_releases_device(
+    launch_runtime: stage_runtime.StageRuntime, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture
+) -> None:
+    manager = mocker.Mock(spec=CoreEngineProcManager)
+    resources = StageReplicaResources(manager=manager, addresses=EngineZmqAddresses(inputs=[], outputs=[]))
+
+    @contextlib.contextmanager
+    def launch(**_kwargs: object) -> Iterator[StageReplicaResources]:
+        yield resources
+        raise RuntimeError("READY handshake failed")
+
+    monkeypatch.setattr(stage_runtime, "launch_stage_replica", launch)
+    with pytest.raises(RuntimeError, match="READY handshake failed"):
+        launch_runtime._initialize_local_llm_replica(_plan(0, "0").replicas[0], 5)
+
+    manager.shutdown.assert_called_once_with()
+    assert os.environ["VLLM_OMNI_TEST_LAUNCH_ENV"] == "parent"
+    with open(stage_init_utils.device_init_lock_path(0), "r+") as lock:
+        fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)

vllm_omni/engine/stage_runtime.py

@@ -11,7 +11,7 @@ import os
 import threading
 import time
 from collections.abc import Callable, Iterator, Mapping, Sequence
-from contextlib import contextmanager
+from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass
 from typing import Any, cast
 
@@ -149,12 +149,9 @@ class StageRuntime:
         self.stage_pools: list[StagePool] = []
         self._stage_init_executor: concurrent.futures.ThreadPoolExecutor | None = None
         self._spawn_device_lock = threading.Lock()
-        # Serialize all LLM replica spawning + handshake across device groups
-        # to prevent ZMQ port-allocation races (get_engine_zmq_addresses) and
-        # CUDA-context conflicts when multiple engine core subprocesses
-        # initialize simultaneously on different GPUs.  Matches the old
-        # AsyncOmniEngine._initialize_llm_replica pattern which used a single
-        # ``llm_stage_launch_lock`` for all replicas.
+        # Serialize process spawning and process-global environment overlays.
+        # Readiness waits use per-device initialization protection so different
+        # device groups can initialize concurrently.
         self._replica_launch_lock = threading.Lock()
         self._init_visible_devices_baseline: str | None = None
 
@@ -778,43 +775,25 @@ class StageRuntime:
                 spawn_device_lock=self._spawn_device_lock,
                 omni_parallel_stage_init=self._parallel_stage_init,
             )
-            # G2 launch lock serializes engine-core *spawning* across replicas
-            # (ZMQ port-allocation races + simultaneous CUDA context init).
-            #   * Default path: hold it across the whole launch context manager,
-            #     whose __exit__ waits for READY — this serializes the full child
-            #     init (spawn + load + profile + KV + capture).
-            #   * Parallel path: hold it only around the spawn (__enter__); run
-            #     the READY-wait (__exit__) outside the lock so replicas init
-            #     concurrently, coordinated by the child SH/EX device locks.
-            if self._parallel_stage_init:
-                # The per-stage runtime.env overlay mutates os.environ
-                # (process-global), so it must be applied under the launch lock
-                # and only needs to cover the spawn (__enter__) — children
-                # inherit the env at spawn time; the READY-wait needs no env.
-                g2_start = time.perf_counter()
+            # Only process spawning and the runtime.env overlay need the global
+            # launch lock. The launch context's exit waits for READY: holding the
+            # lock there serializes model loading and compilation even across
+            # different GPUs. Default initialization keeps its per-device EX
+            # locks until READY; parallel_stage_init uses child phase locks.
+            g2_start = time.perf_counter()
+            with ExitStack() as launch_stack:
                 with self._replica_launch_lock:
-                    with stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg):
-                        resources = launch_cm.__enter__()
-                g2_spawned = time.perf_counter()
-                launch_cm.__exit__(None, None, None)
-                logger.debug(
-                    "[stage_init] Stage-%s G2 spawn(locked)=%.3fs, READY(unlocked)=%.3fs",
-                    plan.metadata.stage_id,
-                    g2_spawned - g2_start,
-                    time.perf_counter() - g2_spawned,
-                )
-            else:
-                g2_start = time.perf_counter()
-                with self._replica_launch_lock, stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg):
                     g2_locked = time.perf_counter()
-                    with launch_cm as resources:
-                        pass
-                logger.debug(
-                    "[stage_init] Stage-%s G2 launch-lock wait=%.3fs, spawn+READY=%.3fs",
-                    plan.metadata.stage_id,
-                    g2_locked - g2_start,
-                    time.perf_counter() - g2_locked,
-                )
+                    with stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg):
+                        resources = launch_stack.enter_context(launch_cm)
+                g2_spawned = time.perf_counter()
+            logger.debug(
+                "[stage_init] Stage-%s G2 launch-lock wait=%.3fs, spawn=%.3fs, READY=%.3fs",
+                plan.metadata.stage_id,
+                g2_locked - g2_start,
+                g2_spawned - g2_locked,
+                time.perf_counter() - g2_spawned,
+            )
 
             logger.info("[StageRuntime] Stage %s engine startup completed", plan.metadata.stage_id)
             if resources is None: