From 6e63ea875e9a54f640b52478a567b9393df95892 Mon Sep 17 00:00:00 2001 From: Gao Han Date: Thu, 10 Sep 2026 09:54:20 +0800 Subject: [PATCH] [Bugfix] Allow LLM replicas on different GPUs to initialize concurrently (#7292) Signed-off-by: Gao Han Co-authored-by: amy-why-3459 --- .../test_stage_runtime_launch_concurrency.py | 149 ++++++++++++++++++ vllm_omni/engine/stage_runtime.py | 63 +++----- 2 files changed, 170 insertions(+), 42 deletions(-) create mode 100644 tests/engine/test_stage_runtime_launch_concurrency.py diff --git a/tests/engine/test_stage_runtime_launch_concurrency.py b/tests/engine/test_stage_runtime_launch_concurrency.py new file mode 100644 index 000000000..f880b9386 --- /dev/null +++ b/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) diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 5d7d07653..53400df5e 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/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: