From d38566ae80afd80ae362327d2e4c1cbc4011859a Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Thu, 24 Sep 2026 01:03:06 +0800 Subject: [PATCH] [Bugfix][Core] Fix stage initialization spawn-lock/device-lock inversion (#7926) Signed-off-by: tzhouam --- .../test_async_omni_engine_stage_init.py | 187 ++++++++++++++++-- vllm_omni/engine/stage_engine_startup.py | 50 ++--- vllm_omni/engine/stage_init_utils.py | 6 +- vllm_omni/engine/stage_runtime.py | 44 ++--- 4 files changed, 215 insertions(+), 72 deletions(-) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index a8255f81a..c836ec152 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -6,6 +6,7 @@ import contextlib import importlib import json import os +import threading import time import types from dataclasses import dataclass, field @@ -429,6 +430,9 @@ def test_initialize_local_diffusion_replica_scopes_runtime_env(monkeypatch): def _capture_launch_diffusion_stage_replica(**_kwargs): captured["runtime_env"] = os.environ.get(runtime_env_var) captured["device_env"] = os.environ.get(device_env_var) + captured["visible_devices"] = _kwargs["stage_visible_devices"] + assert _kwargs["spawn_device_lock"] is runtime._spawn_device_lock + assert not runtime._spawn_device_lock.locked() raise RuntimeError("stop after capturing launch environment") monkeypatch.setattr( @@ -442,7 +446,8 @@ def test_initialize_local_diffusion_replica_scopes_runtime_env(monkeypatch): assert captured == { "runtime_env": "stage-value", - "device_env": "0", + "device_env": "0,1", + "visible_devices": "0", } assert runtime_env_var not in os.environ assert os.environ[device_env_var] == "0,1" @@ -634,14 +639,28 @@ def test_initialize_diffusion_stage_preserves_configured_max_num_seqs(monkeypatc } -def test_launch_diffusion_stage_replica_preserves_configured_max_num_seqs(monkeypatch): +@pytest.mark.parametrize("spawn_failure", [False, True]) +def test_launch_diffusion_stage_replica_preserves_configured_max_num_seqs(monkeypatch, spawn_failure): import vllm_omni.diffusion.stage_diffusion_client as client_mod import vllm_omni.diffusion.stage_diffusion_proc as proc_mod import vllm_omni.engine.stage_engine_startup as startup_mod + from vllm_omni.platforms import current_omni_platform + + spawn_lock = threading.Lock() + device_env = current_omni_platform.device_control_env_var + monkeypatch.setenv(device_env, "0,1") + released = [] + monkeypatch.setattr(startup_mod, "release_device_locks", lambda fds: released.extend(fds)) + + def acquire(*args, **kwargs): + assert not spawn_lock.locked() + assert os.environ[device_env] == "0,1" + assert kwargs["visible_devices"] == "1" + return [42] od_config = types.SimpleNamespace(max_num_seqs=4, parallel_config=types.SimpleNamespace(world_size=1)) monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config) - monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: []) + monkeypatch.setattr(startup_mod, "acquire_device_locks", acquire) monkeypatch.setattr( startup_mod, "register_stage_with_omni_master", @@ -663,7 +682,15 @@ def test_launch_diffusion_stage_replica_preserves_configured_max_num_seqs(monkey outputs=["tcp://127.0.0.1:26003"], ) ) - monkeypatch.setattr(proc_mod, "StageDiffusionProcManager", lambda **kwargs: proc_manager) + + def spawn(**kwargs): + assert spawn_lock.locked() + assert os.environ[device_env] == "1" + if spawn_failure: + raise RuntimeError("spawn failed") + return proc_manager + + monkeypatch.setattr(proc_mod, "StageDiffusionProcManager", spawn) sentinel_client = object() monkeypatch.setattr( client_mod.StageDiffusionClient, @@ -671,15 +698,26 @@ def test_launch_diffusion_stage_replica_preserves_configured_max_num_seqs(monkey lambda metadata, **kwargs: sentinel_client, ) - result, resources = startup_mod.launch_diffusion_stage_replica( - model="dummy-model", - stage_config=types.SimpleNamespace(), - metadata=types.SimpleNamespace(stage_id=0), - stage_init_timeout=12, - use_inline=False, - omni_master_server=omni_master_server, - ) + expectation = pytest.raises(RuntimeError, match="spawn failed") if spawn_failure else contextlib.nullcontext() + with expectation: + result, resources = startup_mod.launch_diffusion_stage_replica( + model="dummy-model", + stage_config=types.SimpleNamespace(), + metadata=types.SimpleNamespace(stage_id=0), + stage_init_timeout=12, + use_inline=False, + omni_master_server=omni_master_server, + stage_visible_devices="1", + spawn_device_lock=spawn_lock, + ) + assert not spawn_lock.locked() + assert os.environ[device_env] == "0,1" + if spawn_failure: + assert released == [42] + return + assert resources.lock_fds == [42] + assert released == [] assert result is sentinel_client assert od_config.max_num_seqs == 4 assert resources.manager is proc_manager @@ -717,7 +755,7 @@ def test_launch_diffusion_stage_replica_preserves_step_execution_max_num_seqs(mo parallel_config=types.SimpleNamespace(world_size=1), ) monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config) - monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: []) + monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args, **kwargs: []) monkeypatch.setattr( startup_mod, "register_stage_with_omni_master", @@ -930,7 +968,7 @@ def test_stage_runtime_launches_shared_engines_with_per_client_addresses(monkeyp monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans) monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda stage_id, _cfg: str(stage_id)) - monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: []) + monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: []) events: list[tuple[str, int]] = [] captured_launch_kwargs: list[dict[str, object]] = [] @@ -1067,7 +1105,7 @@ def test_stage_runtime_multi_api_maps_stage_devices_from_launcher_visibility(mon device_env = runtime_mod.current_omni_platform.device_control_env_var monkeypatch.setenv(device_env, "5") monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan]) - monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: []) + monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: []) captured_devices: list[str | None] = [] @contextlib.contextmanager @@ -1112,7 +1150,7 @@ def test_stage_runtime_multi_api_failure_shuts_down_before_exceptional_context_e stage_plan.replicas[0].engine_args_dict = {} monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan]) monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_args: None) - monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: []) + monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: []) events: list[str] = [] @@ -1406,7 +1444,7 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag prev_device_env = os.environ.get(device_env_var) os.environ[device_env_var] = "0" - def _capture_acquire_device_locks(*_args): + def _capture_acquire_device_locks(*_args, **_kwargs): nonlocal captured_timeout captured_timeout = _args[2] return [] @@ -2245,7 +2283,9 @@ def test_single_api_common_launch_ownership_and_rollback(monkeypatch, parallel, addresses = EngineZmqAddresses(inputs=["ipc://input"], outputs=["ipc://output"]) monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0") - def acquire(*_): + def acquire(*_, **kwargs): + assert kwargs["visible_devices"] == "0" + assert not runtime._spawn_device_lock.locked() events.append("lock") return [42] @@ -2400,7 +2440,6 @@ def test_typed_diffusion_replicas_share_one_config_between_planning_and_launch(m ) mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={}) mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None)) - mocker.patch.object(runtime, "_stage_device_scope", side_effect=lambda *_: contextlib.nullcontext()) client = mocker.Mock() launch = mocker.patch.object(runtime_module, "launch_diffusion_stage_replica", return_value=(client, None)) @@ -2529,3 +2568,113 @@ def test_native_kv_topology_rejects_silent_legacy_fallback(roles, sources, async else: with pytest.raises(ValueError, match="two-stage"): runtime._validate_native_kv_topology(plans) + + +@pytest.mark.parametrize("explicit", [None, "1"]) +def test_acquire_device_locks_visible_devices_precedence(monkeypatch, tmp_path, explicit): + import vllm_omni.engine.stage_init_utils as init_mod + from vllm_omni.platforms import current_omni_platform + + env_var = current_omni_platform.device_control_env_var + monkeypatch.setenv(env_var, "0") + monkeypatch.setattr(init_mod, "device_init_lock_path", lambda device: str(tmp_path / f"gpu-{device}")) + locked: set[int] = set() + fds = init_mod.acquire_device_locks(0, {}, 1, locked, visible_devices=explicit) + try: + assert len(fds) == 1 + assert locked == {int(explicit or "0")} + assert os.environ[env_var] == "0" + finally: + init_mod.release_device_locks(fds) + + +def test_llm_device_lock_wait_allows_holder_to_spawn(monkeypatch, tmp_path): + """A peer holding the GPU flock must be able to spawn to finish init.""" + import vllm_omni.engine.stage_init_utils as init_mod + import vllm_omni.engine.stage_runtime as runtime_mod + from vllm_omni.platforms import current_omni_platform + + runtime = _make_stage_runtime() + plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig()).replicas[0] + plan.engine_args_dict = {} + monkeypatch.setenv(current_omni_platform.device_control_env_var, "0") + monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0") + monkeypatch.setattr(init_mod, "device_init_lock_path", lambda device: str(tmp_path / f"gpu-{device}")) + holder_fds = init_mod.acquire_device_locks(1, {}, 1) + assert len(holder_fds) == 1 + waiting = threading.Event() + spawned = threading.Event() + + def holder(): + try: + if waiting.wait(5): + acquired = runtime._spawn_device_lock.acquire(timeout=2) + if acquired: + runtime._spawn_device_lock.release() + spawned.set() + finally: + init_mod.release_device_locks(holder_fds) + + def acquire(*args, **kwargs): + waiting.set() + # Ensure the peer has an opportunity to request the spawn lock before + # flock acquisition can return (including via its timeout escape). + assert spawned.wait(3), "GPU-lock holder blocked by the waiting stage's spawn lock" + return init_mod.acquire_device_locks(*args, **kwargs) + + monkeypatch.setattr(runtime_mod, "acquire_device_locks", acquire) + + @contextlib.contextmanager + def launch(**kwargs): + yield StageReplicaResources(addresses=EngineZmqAddresses(inputs=["ipc://in"], outputs=["ipc://out"])) + + monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch) + monkeypatch.setattr(runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", lambda **_: object()) + thread = threading.Thread(target=holder, daemon=True) + thread.start() + try: + runtime._initialize_local_llm_replica(plan, 5) + assert spawned.is_set() + finally: + waiting.set() + thread.join(6) + assert not thread.is_alive() + + +@pytest.mark.parametrize("fail", [False, True]) +def test_colocated_diffusion_scopes_devices(monkeypatch, fail): + import vllm_omni.engine.stage_engine_startup as startup_mod + from vllm_omni.platforms import current_omni_platform + + env_var = current_omni_platform.device_control_env_var + monkeypatch.setenv(env_var, "0,1") + lock = threading.Lock() + client = object() + plan = _make_diffusion_plan(0, stage_id=0).replicas[0] + + def initialize(*args, **kwargs): + assert lock.locked() + assert os.environ[env_var] == "1" + if fail: + raise RuntimeError("init failed") + return client + + def forbidden(*args, **kwargs): + pytest.fail("colocated diffusion must not acquire distributed device locks") + + monkeypatch.setattr(startup_mod, "initialize_diffusion_stage", initialize) + monkeypatch.setattr(startup_mod, "acquire_device_locks", forbidden) + expectation = pytest.raises(RuntimeError, match="init failed") if fail else contextlib.nullcontext() + with expectation: + result, _ = startup_mod.launch_diffusion_stage_replica( + model="dummy", + stage_config=plan.stage_cfg, + metadata=plan.metadata, + stage_init_timeout=1, + use_inline=True, + stage_visible_devices="1", + spawn_device_lock=lock, + ) + assert result is client + assert os.environ[env_var] == "0,1" + assert not lock.locked() diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 1aaebc82b..81af45d58 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -1588,6 +1588,8 @@ def launch_diffusion_stage_replica( replica_id: int = 0, omni_master_server: OmniMasterServer | None = None, omni_coordinator_address: str | None = None, + stage_visible_devices: str | None = None, + spawn_device_lock: threading.Lock | None = None, ) -> tuple[Any, StageReplicaResources]: """Launch a local diffusion stage replica. @@ -1596,20 +1598,22 @@ def launch_diffusion_stage_replica( ``StageDiffusionProc`` that heartbeats to ``OmniCoordinator``. """ if omni_master_server is None: - client = initialize_diffusion_stage( - metadata.stage_id, - model, - stage_config, - metadata, - stage_init_timeout=stage_init_timeout, - use_inline=use_inline, - ) + with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock): + client = initialize_diffusion_stage( + metadata.stage_id, + model, + stage_config, + metadata, + stage_init_timeout=stage_init_timeout, + use_inline=use_inline, + ) return client, StageReplicaResources() from vllm_omni.diffusion import stage_diffusion_proc from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient - od_config = build_diffusion_config(model, stage_config, metadata) + with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock): + od_config = build_diffusion_config(model, stage_config, metadata) parallel_config = getattr(od_config, "parallel_config", None) world_size = getattr(parallel_config, "world_size", 1) try: @@ -1620,6 +1624,7 @@ def launch_diffusion_stage_replica( metadata.stage_id, {"tensor_parallel_size": world_size}, stage_init_timeout, + visible_devices=stage_visible_devices, ) proc_manager = None try: @@ -1636,19 +1641,20 @@ def launch_diffusion_stage_replica( handshake=True, data=False, ) - proc_manager = stage_diffusion_proc.StageDiffusionProcManager( - model=model, - od_config=od_config, - stage_init_timeout=stage_init_timeout, - handshake_address=registration.handshake_address, - addresses=EngineZmqAddresses( - inputs=[registration.input_address], - outputs=[registration.output_address], - ), - omni_coordinator_address=omni_coordinator_address, - omni_stage_id=metadata.stage_id, - omni_replica_id=replica_id, - ) + with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock): + proc_manager = stage_diffusion_proc.StageDiffusionProcManager( + model=model, + od_config=od_config, + stage_init_timeout=stage_init_timeout, + handshake_address=registration.handshake_address, + addresses=EngineZmqAddresses( + inputs=[registration.input_address], + outputs=[registration.output_address], + ), + omni_coordinator_address=omni_coordinator_address, + omni_stage_id=metadata.stage_id, + omni_replica_id=replica_id, + ) omni_master_server.release_route_port_reservations( metadata.stage_id, replica_id, diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index faaf6b707..cd38d0e2e 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1755,9 +1755,13 @@ def acquire_device_locks( engine_args_dict: dict[str, Any], stage_init_timeout: int, locked_devices: set[int] | None = None, + *, + visible_devices: str | None = None, ) -> list[int]: """Acquire exclusive file locks on devices needed by this stage. + Pass resolved physical ``visible_devices`` to avoid changing the process + environment while waiting. When omitted, use the device-control environment. Returns list of lock file descriptors that must be released after init. """ lock_fds: list[int] = [] @@ -1790,7 +1794,7 @@ def acquire_device_locks( # Get physical device IDs device_control_env = current_omni_platform.device_control_env_var - visible_devices_str = os.environ.get(device_control_env) + visible_devices_str = visible_devices if visible_devices is not None else os.environ.get(device_control_env) physical_devices: list[int] = [] if visible_devices_str: diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index ed22bfa87..2f3e2e742 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -411,15 +411,17 @@ class StageRuntime: replica.metadata.runtime_cfg, ) if not self._parallel_stage_init: - with self._scoped_spawn_device_env(physical_devices): - lock_fds.extend( - acquire_device_locks( - replica.metadata.stage_id, - replica.engine_args_dict, - timeout, - locked_devices, - ) + # Never hold the spawn-env lock while waiting for a GPU: + # its current holder may need that lock to spawn. + lock_fds.extend( + acquire_device_locks( + replica.metadata.stage_id, + replica.engine_args_dict, + timeout, + locked_devices, + visible_devices=physical_devices, ) + ) launch_context = launch_stage_replica( vllm_config=replica.stage_vllm_config, @@ -557,17 +559,6 @@ class StageRuntime: """Hook for runtimes that own extra infrastructure during init.""" return None - @contextmanager - def _scoped_spawn_device_env(self, physical_devices: str | None) -> Iterator[None]: - """Briefly scope device visibility for spawn-sensitive setup steps.""" - from vllm_omni.engine.stage_engine_startup import scoped_spawn_device_env - - with scoped_spawn_device_env( - physical_devices, - self._spawn_device_lock, - ): - yield - def _resolve_replica_physical_devices(self, stage_id: int, runtime_cfg: Any) -> str | None: if runtime_cfg is None: runtime_cfg = {} @@ -578,13 +569,6 @@ class StageRuntime: visible_baseline=self._init_visible_devices_baseline, ) - @contextmanager - def _stage_device_scope(self, stage_id: int, runtime_cfg: Any) -> Iterator[None]: - """Temporarily apply the stage device env while launching a replica.""" - physical_devices = self._resolve_replica_physical_devices(stage_id, runtime_cfg) - with self._scoped_spawn_device_env(physical_devices): - yield - # ---- Internal methods ---- def _validate_native_kv_topology(self, plans: list[LogicalStageInitPlan]) -> None: @@ -1182,10 +1166,8 @@ class StageRuntime: client = None resources = None try: - with ( - stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg), - self._stage_device_scope(plan.metadata.stage_id, plan.metadata.runtime_cfg), - ): + physical_devices = self._resolve_replica_physical_devices(plan.metadata.stage_id, plan.metadata.runtime_cfg) + with stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg): omni_conn_cfg, omni_from, omni_to = plan.omni_kv_connector if omni_conn_cfg: if omni_from is None or omni_to is None: @@ -1221,6 +1203,8 @@ class StageRuntime: replica_id=plan.replica_id, omni_master_server=self._get_omni_master_server(), omni_coordinator_address=self._get_coordinator_address(), + stage_visible_devices=physical_devices, + spawn_device_lock=self._spawn_device_lock, ) logger.info(