From 142151fef540d167d3b139fd2845cadb01fc71db Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Sat, 5 Sep 2026 22:22:41 +0800 Subject: [PATCH] Parallel stage initialization (admission + SH/EX device locks) (Worker refactor 1/N) (#5224) Signed-off-by: Taichang Zhou Signed-off-by: tzhouam Co-authored-by: Claude Fable 5 --- .../test_async_omni_engine_stage_init.py | 1 + tests/engine/test_parallel_stage_init.py | 842 ++++++++++++++++++ tests/helpers/clean.py | 32 - tests/worker/test_worker_base.py | 57 +- vllm_omni/config/omni_config.py | 6 + vllm_omni/engine/arg_utils.py | 4 + vllm_omni/engine/async_omni_engine.py | 5 + vllm_omni/engine/stage_admission.py | 297 ++++++ vllm_omni/engine/stage_engine_core_proc.py | 39 + .../engine/stage_engine_core_proc_manager.py | 2 + vllm_omni/engine/stage_engine_startup.py | 16 + vllm_omni/engine/stage_init_utils.py | 161 +++- vllm_omni/engine/stage_phase_lock.py | 242 +++++ vllm_omni/engine/stage_runtime.py | 278 +++++- vllm_omni/worker/base.py | 100 +-- 15 files changed, 1888 insertions(+), 194 deletions(-) create mode 100644 tests/engine/test_parallel_stage_init.py create mode 100644 vllm_omni/engine/stage_admission.py create mode 100644 vllm_omni/engine/stage_phase_lock.py diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index c31009b76..528e46f3a 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -205,6 +205,7 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp engine._omni_lb_policy = "random" engine.request_queue = types.SimpleNamespace() engine._log_stats = True + engine._parallel_stage_init = False captured: dict[str, object] = {} runtime = types.SimpleNamespace(stage_pools=[], initialize=lambda: None) diff --git a/tests/engine/test_parallel_stage_init.py b/tests/engine/test_parallel_stage_init.py new file mode 100644 index 000000000..8a2dd3c09 --- /dev/null +++ b/tests/engine/test_parallel_stage_init.py @@ -0,0 +1,842 @@ +"""Unit tests for parallel stage initialization (CPU-only). + +Covers the pure/lock logic added for ``parallel_stage_init``: + * init-group device keying (resolved physical set; parallel unique keys), + * admission control arithmetic + graph reserve, + * SH/EX device phase locks (real flock on a temp dir), + * the phase-locked executor wrapper ordering, + * a regression guard that the NVML process-scope path stayed removed. + +The GPU-only same-device concurrency (actual overlapping init) requires a +hardware soak and is out of scope for these unit tests. +""" + +from __future__ import annotations + +import contextlib +import inspect +import os +import subprocess +import threading +import types +from pathlib import Path + +import pytest + +from vllm_omni.engine.stage_admission import ( + ADMISSION_EXEMPT, + DeviceLedger, + StageAdmissionError, + StageDemand, + check_admission, + evaluate, + graph_reserve_bytes, +) +from vllm_omni.engine.stage_engine_core_proc import _install_phase_locks +from vllm_omni.engine.stage_init_utils import ( + LogicalStageInitPlan, + ReplicaInitPlan, + device_init_lock_path, + open_device_lock_file, +) +from vllm_omni.engine.stage_phase_lock import ( + DeviceLockTimeoutError, + DevicePhaseLock, + resolve_driver_device_ids, + wrap_executor_with_phase_locks, +) +from vllm_omni.engine.stage_runtime import StageRuntime + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# --------------------------------------------------------------------------- # +# helpers +# --------------------------------------------------------------------------- # +def _llm_replica(stage_id: int, replica_id: int, devices: str | None, *, vllm_config=None) -> ReplicaInitPlan: + metadata = types.SimpleNamespace( + stage_id=stage_id, + stage_type="llm", + runtime_cfg={"devices": devices} if devices is not None else {}, + replica_id=replica_id, + ) + stage_cfg = types.SimpleNamespace( + stage_id=stage_id, + stage_type="llm", + runtime=types.SimpleNamespace(devices=devices), + engine_args={}, + ) + return ReplicaInitPlan( + replica_id=replica_id, + num_replicas=1, + launch_mode="local", + stage_cfg=stage_cfg, + metadata=metadata, + stage_connector_spec={}, + omni_kv_connector=(None, None, None), + stage_vllm_config=vllm_config, + executor_class=object, + ) + + +def _runtime(parallel_stage_init: bool = False) -> StageRuntime: + return StageRuntime( + stage_configs=[], + model="dummy", + config_path="dummy", + stage_init_timeout=5, + diffusion_batch_size=1, + async_chunk=False, + parallel_stage_init=parallel_stage_init, + ) + + +def _fake_vllm_config(util: float, *, model="m", dtype="bf16", cudagraph_mode="FULL", capture_sizes=(1, 2, 4)): + return types.SimpleNamespace( + cache_config=types.SimpleNamespace(gpu_memory_utilization=util), + model_config=types.SimpleNamespace(model=model, dtype=dtype, enforce_eager=False), + compilation_config=types.SimpleNamespace( + cudagraph_mode=cudagraph_mode, + cudagraph_capture_sizes=list(capture_sizes), + ), + ) + + +# --------------------------------------------------------------------------- # +# A1: init-group device key +# --------------------------------------------------------------------------- # +def test_init_group_key_resolves_physical_devices_serial(): + runtime = _runtime(parallel_stage_init=False) + runtime._init_visible_devices_baseline = None # no CUDA_VISIBLE_DEVICES mapping + r0 = _llm_replica(0, 0, "0") + r1 = _llm_replica(1, 0, "1") + r0b = _llm_replica(2, 0, "0") + + k0 = runtime._replica_init_group_key(r0) + k1 = runtime._replica_init_group_key(r1) + k0b = runtime._replica_init_group_key(r0b) + + assert k0 == "device:0" + assert k1 == "device:1" + # Different physical devices -> different groups (parallel). + assert k0 != k1 + # Same physical device -> same group (serialized under LOCK_EX). + assert k0 == k0b + + +def test_init_group_key_parallel_is_unique_per_replica(): + runtime = _runtime(parallel_stage_init=True) + runtime._init_visible_devices_baseline = None + # Both on the SAME device, but parallel mode must give distinct keys so they + # run concurrently, coordinated by the child SH/EX locks + admission. + keys = [ + runtime._replica_init_group_key(_llm_replica(0, 0, "0")), + runtime._replica_init_group_key(_llm_replica(0, 1, "0")), + runtime._replica_init_group_key(_llm_replica(1, 0, "0")), + ] + assert keys == ["parallel:0:0", "parallel:0:1", "parallel:1:0"] + assert len(set(keys)) == 3 + + +def test_init_group_key_remote_and_diffusion_unchanged(): + runtime = _runtime(parallel_stage_init=True) + remote = _llm_replica(1, 0, "0") + remote.launch_mode = "remote" + remote.metadata.runtime_cfg = None + assert runtime._replica_init_group_key(remote) == "remote:1:0" + + diffusion = _llm_replica(2, 0, "0") + diffusion.metadata.stage_type = "diffusion" + assert runtime._replica_init_group_key(diffusion) == "inline:diffusion" + + +# --------------------------------------------------------------------------- # +# B1: admission +# --------------------------------------------------------------------------- # +def test_evaluate_fits(): + cap = 40 * 1024**3 + demands = [StageDemand(0, 0, [0], utilization=0.5, graph_reserve_bytes=1 * 1024**3)] + ledgers = evaluate(demands, {0: cap}, external_reserve_bytes=0, safety_margin_bytes=0) + assert ledgers[0].fits + assert ledgers[0].kv_budget_bytes == int(cap * 0.5) + + +def test_evaluate_oversubscription_raises(): + cap = 40 * 1024**3 + demands = [ + StageDemand(0, 0, [0], utilization=0.7, graph_reserve_bytes=0), + StageDemand(1, 0, [0], utilization=0.7, graph_reserve_bytes=0), + ] + with pytest.raises(StageAdmissionError): + evaluate(demands, {0: cap}, external_reserve_bytes=0, safety_margin_bytes=0) + + +def test_evaluate_unknown_capacity_raises(): + with pytest.raises(StageAdmissionError): + evaluate([StageDemand(0, 0, [3], 0.5, 0)], {0: 1}) + + +def test_graph_reserve_positive_when_enabled_zero_when_disabled(): + # cudagraph enabled -> conservative positive reserve + reserve = graph_reserve_bytes(_fake_vllm_config(0.9, cudagraph_mode="FULL")) + assert reserve > 0 + # cudagraph disabled -> exactly 0 + assert graph_reserve_bytes(_fake_vllm_config(0.9, cudagraph_mode="NONE")) == 0 + eager = _fake_vllm_config(0.9) + eager.model_config.enforce_eager = True + assert graph_reserve_bytes(eager) == 0 + + +def test_check_admission_multi_device_and_grouping(): + cap = 80 * 1024**3 + plan_a = LogicalStageInitPlan( + stage_idx=0, + stage_id=0, + replicas=[_llm_replica(0, 0, "0,1", vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE"))], + ) + plan_b = LogicalStageInitPlan( + stage_idx=1, + stage_id=1, + replicas=[_llm_replica(1, 0, "1", vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE"))], + ) + ledgers = check_admission( + [plan_a, plan_b], + resolve_physical_devices=lambda r: [int(x) for x in r.metadata.runtime_cfg["devices"].split(",")], + device_total_memory=lambda _d: cap, + external_reserve_bytes=0, + safety_margin_bytes=0, + ) + # device 0: only stage 0 (0.4). device 1: stage 0 + stage 1 (0.8). + assert ledgers[0].kv_budget_bytes == int(cap * 0.4) + assert ledgers[1].kv_budget_bytes == int(cap * 0.4) + int(cap * 0.4) + assert ledgers[1].fits + + +def _local_diffusion_replica(stage_id: int, replica_id: int, devices: str, *, util: float | None) -> ReplicaInitPlan: + replica = _llm_replica(stage_id, replica_id, devices) + replica.metadata.stage_type = "diffusion" + replica.stage_vllm_config = None + replica.stage_cfg.engine_args = {} if util is None else {"gpu_memory_utilization": util} + return replica + + +def test_check_admission_admits_diffusion_with_explicit_util(): + cap = 80 * 1024**3 + diff = _local_diffusion_replica(0, 0, "0", util=0.3) + plan = LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[diff]) + ledgers = check_admission( + [plan], + resolve_physical_devices=lambda _r: [0], + device_total_memory=lambda _d: cap, + external_reserve_bytes=0, + safety_margin_bytes=0, + ) + assert set(ledgers) == {0} + assert ledgers[0].kv_budget_bytes == int(cap * 0.3) + + +def test_check_admission_unresolved_local_raises(): + """A local replica invisible to the ledger must fail admission, not skip it: + it would otherwise initialize concurrently with unbounded demand.""" + unresolved = _llm_replica(1, 0, None, vllm_config=_fake_vllm_config(0.9, cudagraph_mode="NONE")) + plan = LogicalStageInitPlan(stage_idx=0, stage_id=1, replicas=[unresolved]) + with pytest.raises(StageAdmissionError, match="stage1/replica0.*unresolved"): + check_admission( + [plan], + resolve_physical_devices=lambda _r: None, + device_total_memory=lambda _d: 80 * 1024**3, + ) + + +def test_check_admission_local_diffusion_without_util_raises(): + """The reviewer's repro: local LLM + local diffusion with no explicit util + used to leave the diffusion stage out of the ledger silently.""" + llm = _llm_replica(0, 0, "0", vllm_config=_fake_vllm_config(0.45, cudagraph_mode="NONE")) + diff = _local_diffusion_replica(1, 0, "0", util=None) + plans = [ + LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[llm]), + LogicalStageInitPlan(stage_idx=1, stage_id=1, replicas=[diff]), + ] + with pytest.raises(StageAdmissionError, match="diffusion, no gpu_memory_utilization"): + check_admission( + plans, + resolve_physical_devices=lambda _r: [0], + device_total_memory=lambda _d: 80 * 1024**3, + ) + + +def test_check_admission_exempt_replicas_skip_without_error(): + """Only ADMISSION_EXEMPT (remote/operator-isolated) may skip the ledger.""" + cap = 80 * 1024**3 + local = _llm_replica(0, 0, "0", vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE")) + remote = _llm_replica(1, 0, None, vllm_config=_fake_vllm_config(0.9, cudagraph_mode="NONE")) + remote.launch_mode = "remote" + plan = LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[local, remote]) + ledgers = check_admission( + [plan], + resolve_physical_devices=lambda r: ADMISSION_EXEMPT if r.launch_mode == "remote" else [0], + device_total_memory=lambda _d: cap, + external_reserve_bytes=0, + safety_margin_bytes=0, + ) + assert set(ledgers) == {0} + assert ledgers[0].contributors == ["stage0/replica0"] + assert ledgers[0].kv_budget_bytes == int(cap * 0.4) + + +def test_run_stage_admission_excludes_remote_replicas(monkeypatch): + """Remote replicas consume a remote node's memory — never the local ledger. + + Remote LLM plans already carry runtime_cfg=None, but remote diffusion + replicas keep their cfg + utilization, so without an explicit launch_mode + guard they would be counted against local devices. They must resolve to + ADMISSION_EXEMPT (not None): None now means an unaccountable LOCAL replica + and fails admission. + """ + import vllm_omni.engine.stage_admission as admission_mod + + runtime = _runtime(parallel_stage_init=True) + runtime._init_visible_devices_baseline = None + + local = _llm_replica(0, 0, "0", vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE")) + remote_diff = _llm_replica(1, 0, "0") + remote_diff.launch_mode = "remote" + remote_diff.metadata.stage_type = "diffusion" + remote_diff.stage_vllm_config = None + remote_diff.stage_cfg.engine_args = {"gpu_memory_utilization": 0.3} + + captured = {} + + def _fake_check_admission(stage_plans, *, resolve_physical_devices, device_total_memory, **kw): + captured["resolve"] = resolve_physical_devices + return {} + + monkeypatch.setattr(admission_mod, "check_admission", _fake_check_admission) + runtime._run_stage_admission([LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[local, remote_diff])]) + + assert captured["resolve"](local) == [0] + assert captured["resolve"](remote_diff) is ADMISSION_EXEMPT + + +# --------------------------------------------------------------------------- # +# B1b: admission capacity is queried in the visible-ordinal namespace +# --------------------------------------------------------------------------- # +def _capture_admission(monkeypatch, runtime, replicas): + """Run _run_stage_admission with check_admission stubbed; return its callables.""" + import vllm_omni.engine.stage_admission as admission_mod + + captured: dict = {} + + def _fake_check_admission(stage_plans, *, resolve_physical_devices, device_total_memory, **kw): + captured["resolve"] = resolve_physical_devices + captured["capacity"] = device_total_memory + return {} + + monkeypatch.setattr(admission_mod, "check_admission", _fake_check_admission) + runtime._run_stage_admission([LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=replicas)]) + return captured + + +def _stub_platform(monkeypatch, seen: list[int], total: int = 80 * 1024**3): + import vllm_omni.engine.stage_runtime as runtime_mod + + monkeypatch.setattr( + runtime_mod, + "current_omni_platform", + types.SimpleNamespace( + device_control_env_var="CUDA_VISIBLE_DEVICES", + get_device_total_memory=lambda d: (seen.append(d), total)[1], + ), + ) + + +@pytest.mark.parametrize( + ("baseline", "devices", "physical", "ordinal"), + [ + ("3,4,5", "1", 4, 1), # restricted visibility + ("1,0", "0", 1, 0), # reordered visibility + ("5,GPU-deadbeef", "0", 5, 0), # mixed identifiers, numeric first + ("GPU-deadbeef,5", "1", 5, 1), # mixed identifiers, numeric second + ("0,1,2,3", "2", 2, 2), # identity + ], + ids=["restricted", "reordered", "mixed-numeric-first", "mixed-numeric-second", "identity"], +) +def test_admission_capacity_uses_visible_ordinal(monkeypatch, baseline, devices, physical, ordinal): + """The ledger stays keyed by physical id, but the capacity query is translated. + + ``get_device_total_memory``/``torch.cuda.get_device_properties`` read their + argument as an ordinal in the *current* process's visible-device namespace. + With CUDA_VISIBLE_DEVICES=3,4,5 the orchestrator only has ordinals 0..2, so + passing physical id 4 straight through would raise — or, when visibility is + reordered, silently measure a different GPU. + """ + runtime = _runtime(parallel_stage_init=True) + runtime._init_visible_devices_baseline = baseline + replica = _llm_replica(0, 0, devices, vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE")) + + seen: list[int] = [] + _stub_platform(monkeypatch, seen) + captured = _capture_admission(monkeypatch, runtime, [replica]) + + assert captured["resolve"](replica) == [physical], "ledger key must stay physical" + assert captured["capacity"](physical) == 80 * 1024**3 + assert seen == [ordinal], f"expected the capacity query on ordinal {ordinal}, got {seen}" + + +def test_admission_capacity_rejects_device_outside_baseline(monkeypatch): + """A physical id absent from the captured baseline cannot be measured here.""" + runtime = _runtime(parallel_stage_init=True) + runtime._init_visible_devices_baseline = "3,4,5" + replica = _llm_replica(0, 0, "0", vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE")) + + seen: list[int] = [] + _stub_platform(monkeypatch, seen) + captured = _capture_admission(monkeypatch, runtime, [replica]) + + with pytest.raises(StageAdmissionError, match="not in the visibility baseline"): + captured["capacity"](7) + assert seen == [], "must not query the platform with an unmappable id" + + +def test_admission_capacity_identity_when_no_baseline(monkeypatch): + """With no device-control env captured, ordinals and physical ids coincide.""" + runtime = _runtime(parallel_stage_init=True) + runtime._init_visible_devices_baseline = None + replica = _llm_replica(0, 0, "2", vllm_config=_fake_vllm_config(0.4, cudagraph_mode="NONE")) + + seen: list[int] = [] + _stub_platform(monkeypatch, seen) + captured = _capture_admission(monkeypatch, runtime, [replica]) + + assert captured["capacity"](2) == 80 * 1024**3 + assert seen == [2] + + +# --------------------------------------------------------------------------- # +# B1c: fail closed for executors the phase locks cannot guard +# --------------------------------------------------------------------------- # +class _RayBackedExecutor: + uses_ray = True + + +class _LocalExecutor: + uses_ray = False + + +def test_parallel_init_refuses_ray_backed_executor(): + """Phase locks are node-local flocks and the ledger only sees this host. + + A Ray executor may place the workers that load weights, profile, allocate KV + and capture graphs on other nodes, which participate in neither mechanism — + and the parent skips its own device lock when parallel_stage_init is on. + """ + runtime = _runtime(parallel_stage_init=True) + replica = _llm_replica(0, 0, "0") + replica.executor_class = _RayBackedExecutor + plans = [LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[replica])] + + with pytest.raises(RuntimeError, match="cannot guard Ray-backed or multi-node executors"): + runtime._reject_unguardable_executors(plans) + + +@pytest.mark.parametrize( + ("use_ray", "nnodes", "reason"), + [ + (True, 1, "Ray-backed"), + (False, 2, r"multi-node \(nnodes=2\)"), + ], + ids=["parallel-config-ray", "multi-node-mp"], +) +def test_parallel_init_refuses_unguardable_parallel_config(use_ray, nnodes, reason): + """The config guard also covers non-Ray executors spread across nodes.""" + runtime = _runtime(parallel_stage_init=True) + vllm_config = types.SimpleNamespace( + parallel_config=types.SimpleNamespace(use_ray=use_ray, nnodes=nnodes), + ) + replica = _llm_replica(0, 0, "0", vllm_config=vllm_config) + replica.executor_class = _LocalExecutor + plans = [LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[replica])] + + with pytest.raises(RuntimeError, match=reason): + runtime._reject_unguardable_executors(plans) + + +def test_parallel_init_allows_non_ray_executor(): + runtime = _runtime(parallel_stage_init=True) + replicas = [_llm_replica(0, 0, "0"), _llm_replica(1, 0, "1")] + replicas[0].executor_class = _LocalExecutor + replicas[1].executor_class = object # no uses_ray attribute at all + runtime._reject_unguardable_executors([LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=replicas)]) + + +def test_parallel_init_ignores_remote_ray_replicas(): + """Remote replicas are launched by the runtime owning their node.""" + runtime = _runtime(parallel_stage_init=True) + replica = _llm_replica(0, 0, "0") + replica.executor_class = _RayBackedExecutor + replica.launch_mode = "remote" + runtime._reject_unguardable_executors([LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[replica])]) + + +# --------------------------------------------------------------------------- # +# B2: SH/EX device phase locks (real flock) +# --------------------------------------------------------------------------- # +def _lock(tmp: Path, **kw) -> DevicePhaseLock: + return DevicePhaseLock([0], lock_dir=str(tmp), timeout_s=kw.pop("timeout_s", 0.3), **kw) + + +def test_shared_locks_coexist(tmp_path): + a, b = _lock(tmp_path), _lock(tmp_path) + with a.shared(): + with b.shared(): + pass # no timeout + + +def test_exclusive_excludes_shared_fail_closed(tmp_path): + a, b = _lock(tmp_path, timeout_s=0.2), _lock(tmp_path, timeout_s=0.2) + with a.exclusive(): + with pytest.raises(DeviceLockTimeoutError): + with b.shared(): + pass + + +def test_exclusive_excludes_exclusive(tmp_path): + a, b = _lock(tmp_path, timeout_s=0.2), _lock(tmp_path, timeout_s=0.2) + with a.exclusive(): + with pytest.raises(DeviceLockTimeoutError): + with b.exclusive(): + pass + + +def test_empty_device_set_is_noop(tmp_path): + lock = DevicePhaseLock([], lock_dir=str(tmp_path)) + with lock.shared(): + pass + with lock.exclusive(): + pass + + +def _dead_pid() -> int: + """A PID that has certainly exited and been reaped.""" + proc = subprocess.Popen(["/bin/true"]) + proc.wait() + return proc.pid + + +def test_legacy_device_lock_never_unlinks_the_path(): + """The stale-PID unlink path is gone from the legacy acquirer. + + ``flock`` is released by the kernel when the holder exits (SIGKILL included), + so a dead holder never keeps the lock and unlink-based cleanup is never + needed. It was also unsafe: the cleanup only ran on ``BlockingIOError``, i.e. + only while the lock *was* held, so every case it fired in was a case where + unlinking broke a live holder. + """ + import vllm_omni.engine.stage_init_utils as siu + + assert not hasattr(siu, "_cleanup_stale_lock_if_dead") + assert "unlink" not in inspect.getsource(siu.acquire_device_locks) + + # Test cleanup runs in many autouse/runtime fixtures and must not unlink the + # shared production lock path either. Persistent unlocked flock files are + # harmless; deleting one can split live holders across different inodes. + helper_source = (Path(__file__).parents[1] / "helpers" / "clean.py").read_text() + assert "_cleanup_stale_device_locks" not in helper_source + assert "vllm_omni_device_*_init.lock" not in helper_source + + +def test_legacy_lock_cannot_break_a_live_shared_holder(tmp_path, monkeypatch): + """A live LOCK_SH holder must keep excluding the legacy LOCK_EX acquirer + even when the PID recorded in the lock file belongs to a dead process. + + Several LOCK_SH holders overwrite that PID, so the recorded value can name a + process that has exited while another holder still owns the inode. Unlinking + on that signal let the legacy acquirer create a fresh file under the same + name and take LOCK_EX on a *different* inode — conflicting with nobody. + """ + import vllm_omni.engine.stage_init_utils as siu + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + monkeypatch.setattr( + siu, + "device_init_lock_path", + lambda device_id, lock_dir="/tmp": device_init_lock_path(device_id, str(tmp_path)), + ) + + path = Path(device_init_lock_path(0, str(tmp_path))) + holder = DevicePhaseLock([0], lock_dir=str(tmp_path), timeout_s=5) + + with holder.shared(): + path.write_text(f"{_dead_pid()}\n") # peer SH holder exited after recording its PID + inode_before = path.stat().st_ino + + lock_fds = siu.acquire_device_locks(0, {"tensor_parallel_size": 1}, stage_init_timeout=1) + + assert lock_fds == [], "legacy LOCK_EX must not be granted while a shared holder is live" + assert path.exists(), "the lock path must never be unlinked" + assert path.stat().st_ino == inode_before, "the lock inode must stay stable" + + siu.release_device_locks(lock_fds) + + +def test_phase_lock_reuses_the_existing_lock_inode(tmp_path): + """Repeated acquisitions coordinate on one stable inode per device.""" + path = Path(device_init_lock_path(0, str(tmp_path))) + lock = DevicePhaseLock([0], lock_dir=str(tmp_path), timeout_s=1) + + with lock.shared(): + first = path.stat().st_ino + with lock.exclusive(): + assert path.stat().st_ino == first + assert os.path.exists(path) + + +def test_device_lock_file_never_visible_at_restrictive_mode(tmp_path): + """Concurrent creators must never observe the lock at a umask-filtered mode. + + Creating in place with O_EXCL and widening afterwards publishes the file at + 0600 under a restrictive umask; a second user opening it in that window is + locked out and parallel init aborts. Publishing atomically closes the window, + so every concurrent creator either wins the link or opens a readable file. + """ + path = device_init_lock_path(0, str(tmp_path)) + results: list[tuple[int, bool]] = [] + errors: list[BaseException] = [] + modes: list[int] = [] + barrier = threading.Barrier(8) + + def _worker() -> None: + barrier.wait() + try: + fd, writable = open_device_lock_file(path) + except BaseException as exc: # noqa: BLE001 - surfaced by the assert below + errors.append(exc) + return + try: + modes.append(os.stat(path).st_mode & 0o777) + results.append((fd, writable)) + finally: + os.close(fd) + + previous = os.umask(0o077) + try: + threads = [threading.Thread(target=_worker) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + finally: + os.umask(previous) + + assert not errors, f"concurrent creators were locked out: {errors!r}" + assert len(results) == 8 + assert all(m & 0o044 == 0o044 for m in modes), f"lock was visible at {[oct(m) for m in modes]}" + assert len({os.stat(path).st_ino}) == 1 + + +@pytest.mark.parametrize("umask_value", [0o022, 0o027, 0o077], ids=["022", "027", "077"]) +def test_device_lock_file_is_readable_by_other_users(tmp_path, umask_value): + """The lock file must stay group/other readable whatever the creator's umask. + + These files coordinate GPU init *across* users, so a second user has to be + able to open the first user's file. The mode passed to ``os.open`` is filtered + by umask, so a restrictive umask (0027, 0077) would produce 0640 or 0600 and + lock every other user out -- fatal for parallel_stage_init, and a silent loss + of serialization in the legacy path. + """ + path = device_init_lock_path(0, str(tmp_path)) + previous = os.umask(umask_value) + try: + fd, writable = open_device_lock_file(path) + finally: + os.umask(previous) + try: + assert writable is True + mode = os.stat(path).st_mode & 0o777 + assert mode & 0o044 == 0o044, f"lock file mode {mode:04o} is not readable by other users" + finally: + os.close(fd) + + +@pytest.mark.skipif(os.geteuid() == 0, reason="root bypasses file permission bits") +def test_phase_lock_tolerates_lock_file_owned_by_another_user(tmp_path): + """A lock file created by a different user must not break locking. + + These files coordinate *across* users -- two people on the same physical GPU + must contend on the same path -- so whoever creates one owns it at 0644 and + every later user's O_RDWR open fails with EACCES. That made parallel_stage_init + die at startup on any shared machine. flock attaches to the open file + description and needs no write access, so the read-only fallback keeps full + mutual exclusion and only forfeits the diagnostic PID write. + """ + path = Path(device_init_lock_path(0, str(tmp_path))) + path.write_text("999999\n") + path.chmod(0o444) + + fd, writable = open_device_lock_file(str(path)) + try: + assert writable is False, "expected the read-only fallback on an unwritable file" + finally: + os.close(fd) + + holder = DevicePhaseLock([0], lock_dir=str(tmp_path), timeout_s=0.3) + contender = DevicePhaseLock([0], lock_dir=str(tmp_path), timeout_s=0.2) + with holder.exclusive(): + with pytest.raises(DeviceLockTimeoutError): + with contender.shared(): + pass + + assert path.read_text() == "999999\n", "must not clobber a file it cannot own" + + +def test_resolve_driver_device_ids_dp_slice(monkeypatch): + from vllm_omni.platforms import current_omni_platform + + env_var = current_omni_platform.device_control_env_var + monkeypatch.setenv(env_var, "4,5,6,7") + pc = types.SimpleNamespace( + tensor_parallel_size=2, + pipeline_parallel_size=1, + prefill_context_parallel_size=1, + sequence_parallel_size=1, + cfg_parallel_size=1, + ) + vc = types.SimpleNamespace(parallel_config=pc) + assert resolve_driver_device_ids(vc, local_dp_rank=0) == [4, 5] + assert resolve_driver_device_ids(vc, local_dp_rank=1) == [6, 7] + + +# --------------------------------------------------------------------------- # +# B2: executor wrapper ordering + retry +# --------------------------------------------------------------------------- # +class _RecordingLocker: + def __init__(self): + self.events: list[str] = [] + + @contextlib.contextmanager + def shared(self): + self.events.append("SH-enter") + try: + yield + finally: + self.events.append("SH-exit") + + @contextlib.contextmanager + def exclusive(self): + self.events.append("EX-enter") + try: + yield + finally: + self.events.append("EX-exit") + + +class _FakeExecutor: + instances: list[object] = [] + + def __init__(self, vllm_config): + self.vllm_config = vllm_config + self.dam_calls = 0 + _FakeExecutor.instances.append(self) + + def determine_available_memory(self): + self.dam_calls += 1 + return 123 + + def initialize_from_config(self, kv_cache_configs): + return "ok" + + def compile_or_warm_up_model(self): + return "warm" + + +def test_wrapper_phase_ordering(): + locker = _RecordingLocker() + wrapped_cls = wrap_executor_with_phase_locks(_FakeExecutor, locker) + exec_ = wrapped_cls(vllm_config=object()) + assert exec_.determine_available_memory() == 123 + assert exec_.initialize_from_config([]) == "ok" + assert exec_.compile_or_warm_up_model() == "warm" + assert locker.events == [ + "SH-enter", + "SH-exit", # load + "EX-enter", + "EX-exit", # profile + "SH-enter", + "SH-exit", # KV cache allocation + "SH-enter", + "SH-exit", # kernel warmup + CUDA-graph capture + ] + + +def test_wrapper_locks_warmup_and_capture_separately(): + """Regression guard for the upstream split. + + ``Executor.initialize_from_config`` used to fuse ``compile_or_warm_up_model`` + into itself; since vLLM v0.20.0 ``EngineCore._initialize_kv_caches`` calls the + two separately. If the wrapper only brackets ``initialize_from_config``, one + stage captures CUDA graphs — mutating device memory — while a peer holds the + exclusive lock for its profiling measurement. + """ + locker = _RecordingLocker() + wrapped_cls = wrap_executor_with_phase_locks(_FakeExecutor, locker) + exec_ = wrapped_cls(vllm_config=object()) + locker.events.clear() + + assert exec_.compile_or_warm_up_model() == "warm" + + assert locker.events == ["SH-enter", "SH-exit"], "compile_or_warm_up_model must run under the shared device lock" + assert "compile_or_warm_up_model" in vars(wrapped_cls), ( + "the wrapper must override compile_or_warm_up_model, not inherit it unguarded" + ) + + +@pytest.mark.parametrize( + "kwargs", + [{}, {"vllm_config": object()}, {"executor_class": object}], + ids=["both-missing", "executor-missing", "config-missing"], +) +def test_install_phase_locks_fails_closed_on_missing_kwargs(kwargs): + """parallel_stage_init must never proceed with the phase-lock guard + uninstalled — a missing vllm_config/executor_class aborts the child.""" + with pytest.raises(RuntimeError, match="phase-lock guard cannot be installed"): + _install_phase_locks(kwargs, local_dp_rank=0) + + +def test_install_phase_locks_wraps_executor(monkeypatch): + import vllm_omni.engine.stage_phase_lock as phase_lock_mod + + locker = _RecordingLocker() + locker.device_ids = [0] + monkeypatch.setattr( + phase_lock_mod.DevicePhaseLock, + "from_child", + classmethod(lambda _cls, _cfg, _rank: locker), + ) + kwargs = {"vllm_config": object(), "executor_class": _FakeExecutor} + + _install_phase_locks(kwargs, local_dp_rank=0) + + exec_ = kwargs["executor_class"](vllm_config=object()) + assert exec_.determine_available_memory() == 123 + assert locker.events[:2] == ["SH-enter", "SH-exit"] + + +# --------------------------------------------------------------------------- # +# A2: NVML process-scope removal regression guard +# --------------------------------------------------------------------------- # +def test_base_worker_dropped_nvml_process_scope(): + import vllm_omni.worker.base as base + + src = Path(base.__file__).read_text() + assert "get_process_gpu_memory" not in src + assert "is_process_scoped_memory_available" not in src + + +def test_device_ledger_required_and_fits(): + led = DeviceLedger(device_id=0, capacity_bytes=100, kv_budget_bytes=40, graph_reserve_bytes=20) + led.external_reserve_bytes = 10 + led.safety_margin_bytes = 10 + assert led.required_bytes == 80 + assert led.fits + led.kv_budget_bytes = 90 + assert not led.fits diff --git a/tests/helpers/clean.py b/tests/helpers/clean.py index 17fef825d..5cd3f5b39 100644 --- a/tests/helpers/clean.py +++ b/tests/helpers/clean.py @@ -170,36 +170,6 @@ def _print_device_processes() -> None: print("=" * 80) -def _cleanup_stale_device_locks() -> None: - """Remove stale device-initialization lock files whose recorded PID is dead. - - Lock files at ``/tmp/vllm_omni_device_*_init.lock`` may persist after a - crashed / killed test run and block subsequent orchestrator startups. - """ - import glob as _glob - - for lock_file in _glob.glob("/tmp/vllm_omni_device_*_init.lock"): - try: - with open(lock_file) as fh: - content = fh.read().strip() - if not content: - continue - pid = int(content) - except (OSError, ValueError): - continue - - try: - os.kill(pid, 0) - except ProcessLookupError: - print(f" Removing stale device lock {lock_file} (PID {pid} is dead)") - try: - os.unlink(lock_file) - except OSError: - pass - except PermissionError: - pass - - def cleanup_test_environment(*, shutdown_ray: bool = False) -> None: """Tear down distributed state and reset device memory for tests.""" from vllm import envs @@ -224,8 +194,6 @@ def cleanup_test_environment(*, shutdown_ray: bool = False) -> None: ray.shutdown() print("Pre-test device status:") - _cleanup_stale_device_locks() - num_devices = current_omni_platform.device_count() if num_devices > 0: try: diff --git a/tests/worker/test_worker_base.py b/tests/worker/test_worker_base.py index 305b0778f..18131a31c 100644 --- a/tests/worker/test_worker_base.py +++ b/tests/worker/test_worker_base.py @@ -2,9 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Characterization tests for ``vllm_omni.worker.base.OmniGPUWorkerBase``. -Pins the CURRENT behaviour of ``determine_available_memory`` (BOTH the NVML / -process-scoped arm AND the profiling-fallback arm — this base still carries the -NVML branch), plus the +Pins the CURRENT behaviour of ``determine_available_memory`` (device-level +profiling is the only path — the NVML / process-scoped arm was removed with +parallel stage init, which coordinates concurrent same-device measurement via +admission + SH/EX device locks instead), plus the memory-pool / sleep / wake-up plumbing that a later change may touch. Workers are constructed via ``object.__new__`` with only the attributes each @@ -57,27 +58,20 @@ def _make_worker(*, requested_memory: int, kv_cache_memory_bytes: int = 0, model # --------------------------------------------------------------------------- # # determine_available_memory # # --------------------------------------------------------------------------- # -def test_determine_available_memory_nvml_arm(monkeypatch): - """NVML available -> available = requested - per-process memory.""" +def test_no_process_scoped_collaborators_remain(): + """The NVML arm is gone from base.py, so the method CANNOT take a + process-scoped path on any host: its collaborators are not even imported. + (The removal itself happened in the parallel-stage-init feature commit; + tests/engine/test_parallel_stage_init.py pins the same invariant.)""" + assert not hasattr(base, "is_process_scoped_memory_available") + assert not hasattr(base, "detect_pid_host") + assert not hasattr(base, "get_process_gpu_memory") + + +def test_determine_available_memory_profiling_path(monkeypatch): + """available = requested - (weights + peak + non_torch); the only path.""" worker = _make_worker(requested_memory=30 * GIB, model_memory_usage=10 * GIB) monkeypatch.setattr(base, "memory_profiling", _fake_memory_profiling(non_torch=1 * GIB, torch_peak=2 * GIB)) - monkeypatch.setattr(base, "is_process_scoped_memory_available", lambda: True) - monkeypatch.setattr(base, "detect_pid_host", lambda: True) - monkeypatch.setattr(base, "get_process_gpu_memory", lambda local_rank: 8 * GIB) - - out = OmniGPUWorkerBase.determine_available_memory(worker) - - # NVML arm ignores profiled usage; uses process_memory directly. - assert out == 30 * GIB - 8 * GIB - - -def test_determine_available_memory_profiling_fallback_arm(monkeypatch): - """NVML unavailable -> available = requested - (weights + peak + non_torch).""" - worker = _make_worker(requested_memory=30 * GIB, model_memory_usage=10 * GIB) - monkeypatch.setattr(base, "memory_profiling", _fake_memory_profiling(non_torch=1 * GIB, torch_peak=2 * GIB)) - monkeypatch.setattr(base, "is_process_scoped_memory_available", lambda: False) - # detect_pid_host is short-circuited by the `and`, but pin it anyway. - monkeypatch.setattr(base, "detect_pid_host", lambda: True) out = OmniGPUWorkerBase.determine_available_memory(worker) @@ -85,25 +79,20 @@ def test_determine_available_memory_profiling_fallback_arm(monkeypatch): assert out == 30 * GIB - profiled -def test_determine_available_memory_falls_back_when_not_pid_host(monkeypatch): - """detect_pid_host False also routes to the profiling fallback.""" - worker = _make_worker(requested_memory=20 * GIB, model_memory_usage=5 * GIB) - monkeypatch.setattr(base, "memory_profiling", _fake_memory_profiling(non_torch=0, torch_peak=0)) - monkeypatch.setattr(base, "is_process_scoped_memory_available", lambda: True) - monkeypatch.setattr(base, "detect_pid_host", lambda: False) +def test_determine_available_memory_populates_total_consumed(monkeypatch): + """total_consumed mirrors upstream's MemoryProfilingResult field.""" + worker = _make_worker(requested_memory=30 * GIB, model_memory_usage=10 * GIB) + monkeypatch.setattr(base, "memory_profiling", _fake_memory_profiling(non_torch=1 * GIB, torch_peak=2 * GIB)) - out = OmniGPUWorkerBase.determine_available_memory(worker) + OmniGPUWorkerBase.determine_available_memory(worker) - assert out == 20 * GIB - 5 * GIB + assert worker.total_consumed == 3 * GIB def test_determine_available_memory_clamps_to_zero(monkeypatch): """Over-subscription clamps the KV budget to 0, never negative.""" - worker = _make_worker(requested_memory=5 * GIB, model_memory_usage=0) + worker = _make_worker(requested_memory=5 * GIB, model_memory_usage=8 * GIB) monkeypatch.setattr(base, "memory_profiling", _fake_memory_profiling(non_torch=0, torch_peak=0)) - monkeypatch.setattr(base, "is_process_scoped_memory_available", lambda: True) - monkeypatch.setattr(base, "detect_pid_host", lambda: True) - monkeypatch.setattr(base, "get_process_gpu_memory", lambda local_rank: 8 * GIB) out = OmniGPUWorkerBase.determine_available_memory(worker) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 88efc8d60..dce775cd3 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -1299,6 +1299,12 @@ class VllmOmniOrchestratorConfig: omni_lb_policy: str = "random" omni_heartbeat_timeout: float = Field(default=30.0, gt=0.0) batch_timeout: int = Field(default=10, ge=0) + # When True, stages sharing a physical GPU initialize concurrently, guarded + # by pre-launch admission control + engine-core-held SH/EX device locks + # (see stage_admission / stage_phase_lock). Default False keeps the legacy + # per-device LOCK_EX serialization. Enable only when the GPU is dedicated to + # this deployment (see rfc_parallel_stage_init). + parallel_stage_init: bool = False @config(config=ConfigDict(arbitrary_types_allowed=True)) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 0eb18c686..22790fbd4 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -504,6 +504,10 @@ class OrchestratorArgs: # === Lifecycle === stage_init_timeout: int = 300 init_timeout: int = 600 + # Initialize stages sharing a physical GPU concurrently, guarded by + # pre-launch admission + engine-core SH/EX device locks. Off by default; + # enable only when the GPU is dedicated to this deployment. + parallel_stage_init: bool = False # === Cross-stage Communication === batch_timeout: int = 10 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 49393122b..8501d8730 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -185,6 +185,10 @@ class AsyncOmniEngine: self._omni_heartbeat_timeout: float = float(kwargs.get("omni_heartbeat_timeout") or 30.0) if self._omni_heartbeat_timeout <= 0: raise ValueError(f"--omni-heartbeat-timeout must be > 0, got {self._omni_heartbeat_timeout}") + # Concurrent same-device stage init (admission + SH/EX phase locks). + # Sourced from the parallel_stage_init orchestrator/CLI arg (config, + # not an env var); default False preserves serial init. + self._parallel_stage_init: bool = bool(kwargs.get("parallel_stage_init") or False) if single_stage_mode: logger.info( @@ -322,6 +326,7 @@ class AsyncOmniEngine: stage_init_timeout=stage_init_timeout, async_chunk=self.async_chunk, tokenizer=self.tokenizer, + parallel_stage_init=self._parallel_stage_init, single_stage_id_filter=self._single_stage_id_filter, omni_master_address=self._omni_master_address, omni_master_port=self._omni_master_port, diff --git a/vllm_omni/engine/stage_admission.py b/vllm_omni/engine/stage_admission.py new file mode 100644 index 000000000..469787bca --- /dev/null +++ b/vllm_omni/engine/stage_admission.py @@ -0,0 +1,297 @@ +"""Pre-launch admission control for parallel stage initialization. + +When ``VllmOmniOrchestratorConfig.parallel_stage_init`` is enabled, several stage +engine cores initialize concurrently and allocate KV cache / capture CUDA graphs +independently. The SH/EX phase locks (``stage_phase_lock``) keep each *measurement* +clean, but they cannot bound the *sum* of what all stages will allocate — two +stages can each profile against a mostly-empty GPU, each compute a large KV +budget, and then both allocate → OOM. + +Admission is the hard backstop for that: **before any stage launches**, prove for +every physical device *g*:: + + Σ_{s on g} capacity(g)·utilization(s) + Σ graph_reserve(s,g) + + external_reserve + safety_margin ≤ capacity(g) + +If it does not hold, fail fast (``StageAdmissionError``) rather than OOM at +runtime. Because the budgets are proven to fit before anyone allocates, any +profile/allocate interleaving is safe. + +The arithmetic (``evaluate``) is pure and unit-testable. Plan-walking +(``check_admission``) takes injectable callables for device resolution and total +memory so it can also be exercised without a GPU. +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from typing import Any + +from vllm.logger import init_logger +from vllm.utils.mem_utils import format_gib + +logger = init_logger(__name__) + +# Tunable reserves (bytes). Conservative defaults; operators can pass overrides +# to ``check_admission``. Kept as constants (not env vars) per project policy. +_DEFAULT_EXTERNAL_RESERVE_BYTES = 1 * 1024**3 # headroom for unmanaged consumers +_DEFAULT_SAFETY_MARGIN_BYTES = 1 * 1024**3 # fragmentation / allocator slack + +# Graph-pool reserve. A single conservative constant used whenever CUDA-graph +# capture is enabled (never 0 in that case); zero when capture is disabled. +# Kept as a constant (not env var / not a measured table) per project policy. +_DEFAULT_GRAPH_RESERVE_BYTES = 2 * 1024**3 + + +class StageAdmissionError(RuntimeError): + """Raised when the planned stage budgets do not fit the available devices.""" + + +class AdmissionExempt: + """Type of :data:`ADMISSION_EXEMPT`; not meant to be instantiated elsewhere.""" + + __slots__ = () + + def __repr__(self) -> str: + return "ADMISSION_EXEMPT" + + +ADMISSION_EXEMPT = AdmissionExempt() +"""Resolver sentinel: the replica is deliberately outside local admission +(e.g. it runs on a remote node and consumes that node's memory). Only replicas +marked exempt may skip the ledger — a local replica the resolver cannot account +for fails admission instead (see :func:`check_admission`).""" + + +@dataclass +class StageDemand: + """One replica's claim on a set of physical devices.""" + + stage_id: int + replica_id: int + device_ids: list[int] + utilization: float + graph_reserve_bytes: int + is_diffusion: bool = False + + +@dataclass +class DeviceLedger: + """Per-device admission accounting (also the integration-run record).""" + + device_id: int + capacity_bytes: int + kv_budget_bytes: int = 0 + graph_reserve_bytes: int = 0 + external_reserve_bytes: int = 0 + safety_margin_bytes: int = 0 + contributors: list[str] = field(default_factory=list) + + @property + def required_bytes(self) -> int: + return self.kv_budget_bytes + self.graph_reserve_bytes + self.external_reserve_bytes + self.safety_margin_bytes + + @property + def fits(self) -> bool: + return self.required_bytes <= self.capacity_bytes + + +# ---- graph reserve --------------------------------------------------------- + + +def _cudagraph_disabled(vllm_config: Any) -> bool: + model_config = getattr(vllm_config, "model_config", None) + if getattr(model_config, "enforce_eager", False): + return True + comp = getattr(vllm_config, "compilation_config", None) + mode = getattr(comp, "cudagraph_mode", None) + if mode is not None and str(mode).upper().endswith("NONE"): + return True + return False + + +def graph_reserve_bytes(vllm_config: Any) -> int: + """Reserve for this stage's CUDA-graph pool (bytes). + + A single conservative constant when capture is enabled (never 0 in that + case), 0 when capture is disabled. This intentionally over-reserves rather + than measure per-model — admission stays a safe upper bound. + """ + if _cudagraph_disabled(vllm_config): + return 0 + return _DEFAULT_GRAPH_RESERVE_BYTES + + +# ---- pure evaluation ------------------------------------------------------- + + +def evaluate( + demands: Sequence[StageDemand], + capacities: dict[int, int], + *, + external_reserve_bytes: int = _DEFAULT_EXTERNAL_RESERVE_BYTES, + safety_margin_bytes: int = _DEFAULT_SAFETY_MARGIN_BYTES, +) -> dict[int, DeviceLedger]: + """Build the per-device ledger and raise ``StageAdmissionError`` if any device + is over-subscribed. Pure: no config objects, no GPU.""" + ledgers: dict[int, DeviceLedger] = {} + + def _ledger(dev: int) -> DeviceLedger: + if dev not in ledgers: + if dev not in capacities: + raise StageAdmissionError(f"No capacity known for physical device {dev}") + ledgers[dev] = DeviceLedger( + device_id=dev, + capacity_bytes=capacities[dev], + external_reserve_bytes=external_reserve_bytes, + safety_margin_bytes=safety_margin_bytes, + ) + return ledgers[dev] + + for d in demands: + for dev in d.device_ids: + ledger = _ledger(dev) + ledger.kv_budget_bytes += int(ledger.capacity_bytes * d.utilization) + ledger.graph_reserve_bytes += d.graph_reserve_bytes + ledger.contributors.append(f"stage{d.stage_id}/replica{d.replica_id}") + + over = [led for led in ledgers.values() if not led.fits] + for led in ledgers.values(): + logger.info( + "[admission] device %d: capacity=%s required=%s (kv=%s graph=%s ext=%s margin=%s) " + "headroom=%s contributors=%s%s", + led.device_id, + format_gib(led.capacity_bytes), + format_gib(led.required_bytes), + format_gib(led.kv_budget_bytes), + format_gib(led.graph_reserve_bytes), + format_gib(led.external_reserve_bytes), + format_gib(led.safety_margin_bytes), + format_gib(led.capacity_bytes - led.required_bytes), + led.contributors, + " <<< OVER" if not led.fits else "", + ) + if over: + detail = "; ".join( + f"device {led.device_id}: need {format_gib(led.required_bytes)} GiB > " + f"{format_gib(led.capacity_bytes)} GiB (contributors {led.contributors})" + for led in over + ) + raise StageAdmissionError( + "parallel_stage_init admission failed — per-device budget exceeds capacity. " + "Lower gpu_memory_utilization, reduce co-located stages, or disable " + f"parallel_stage_init. {detail}" + ) + return ledgers + + +# ---- plan walking ---------------------------------------------------------- + + +def check_admission( + stage_plans: Sequence[Any], + *, + resolve_physical_devices: Callable[[Any], list[int] | AdmissionExempt | None], + device_total_memory: Callable[[int], int], + external_reserve_bytes: int = _DEFAULT_EXTERNAL_RESERVE_BYTES, + safety_margin_bytes: int = _DEFAULT_SAFETY_MARGIN_BYTES, + graph_reserve: Callable[[Any], int] = graph_reserve_bytes, +) -> dict[int, DeviceLedger]: + """Extract demands from the orchestrator's stage plans and admit them. + + ``resolve_physical_devices(replica)`` returns the physical GPU ids a replica + occupies, or :data:`ADMISSION_EXEMPT` for replicas deliberately outside + local admission (e.g. running on a remote node). Any *other* replica that + cannot be accounted — unresolved or non-integer devices, or a diffusion + stage without an explicit ``gpu_memory_utilization`` — raises + ``StageAdmissionError``: under parallel init every local replica gets its + own init group and no parent holds a whole-init exclusive lock, so a + consumer invisible to the ledger would initialize concurrently with + unbounded demand, defeating admission (fail-closed). + ``device_total_memory(id)`` returns a device's total bytes. + """ + demands: list[StageDemand] = [] + exempt: list[str] = [] + unaccounted: list[str] = [] + for plan in stage_plans: + for replica in getattr(plan, "replicas", []): + metadata = replica.metadata + label = f"stage{metadata.stage_id}/replica{replica.replica_id}" + device_ids = resolve_physical_devices(replica) + if isinstance(device_ids, AdmissionExempt): + exempt.append(label) + continue + if not device_ids: + unaccounted.append(f"{label} (unresolved devices)") + continue + vllm_config = replica.stage_vllm_config + if vllm_config is None: + # Diffusion stages have no resolved vllm_config pre-launch and + # skip profile/capture; admit them with their raw util. A local + # diffusion stage without one has unknowable demand and must + # fail admission below, not bypass it. + util = _diffusion_utilization(replica) + if util is None: + unaccounted.append(f"{label} (diffusion, no gpu_memory_utilization)") + continue + demands.append( + StageDemand( + stage_id=metadata.stage_id, + replica_id=replica.replica_id, + device_ids=list(device_ids), + utilization=util, + graph_reserve_bytes=0, + is_diffusion=True, + ) + ) + continue + demands.append( + StageDemand( + stage_id=metadata.stage_id, + replica_id=replica.replica_id, + device_ids=list(device_ids), + utilization=float(vllm_config.cache_config.gpu_memory_utilization), + graph_reserve_bytes=int(graph_reserve(vllm_config)), + ) + ) + + if unaccounted: + raise StageAdmissionError( + "parallel_stage_init admission cannot account for local replicas: " + f"{unaccounted}. Every local stage must resolve to integer physical " + "device ids and declare a memory budget (diffusion stages: set " + "gpu_memory_utilization in engine_args); otherwise it would " + "initialize concurrently without bounding its demand. Fix the " + "stage config or disable parallel_stage_init." + ) + if exempt: + logger.info( + "[admission] exempt from local admission (remote/operator-isolated): %s. " + "These replicas consume no local device memory.", + exempt, + ) + + capacities: dict[int, int] = {} + for d in demands: + for dev in d.device_ids: + if dev not in capacities: + capacities[dev] = int(device_total_memory(dev)) + + return evaluate( + demands, + capacities, + external_reserve_bytes=external_reserve_bytes, + safety_margin_bytes=safety_margin_bytes, + ) + + +def _diffusion_utilization(replica: Any) -> float | None: + """Best-effort gpu_memory_utilization for a diffusion replica from raw args.""" + stage_cfg = getattr(replica, "stage_cfg", None) + engine_args = getattr(stage_cfg, "engine_args", None) + if isinstance(engine_args, dict): + util = engine_args.get("gpu_memory_utilization") + if util is not None: + return float(util) + return None diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 4bcf47d3b..e483e1bd4 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -41,6 +41,35 @@ logger = init_logger(__name__) _SIGNAL_EXIT_BASE = 128 +def _install_phase_locks(kwargs: dict[str, Any], local_dp_rank: int) -> None: + """Wrap ``kwargs["executor_class"]`` with the SH/EX phase-lock guard. + + Fail-closed: ``parallel_stage_init`` promises that every memory-mutating + init phase in this child runs under the per-device locks, so a missing + ``vllm_config`` or ``executor_class`` must abort the launch rather than + silently proceed with an unguarded parallel initialization. + """ + from vllm_omni.engine.stage_phase_lock import ( + DevicePhaseLock, + wrap_executor_with_phase_locks, + ) + + missing = [key for key in ("vllm_config", "executor_class") if kwargs.get(key) is None] + if missing: + raise RuntimeError( + f"parallel_stage_init is enabled but EngineCore kwargs are missing {missing}, " + "so the SH/EX phase-lock guard cannot be installed. Refusing to run an " + "unguarded parallel initialization; fix the launch plumbing or disable " + "parallel_stage_init." + ) + locker = DevicePhaseLock.from_child(kwargs["vllm_config"], local_dp_rank) + kwargs["executor_class"] = wrap_executor_with_phase_locks(kwargs["executor_class"], locker) + logger.info( + "[StageEngineCoreProc] parallel_stage_init: SH/EX phase locks on devices %s", + locker.device_ids, + ) + + def _signal_exit_code(signum: int) -> int: """Return the conventional process exit code for signal-driven exits.""" return _SIGNAL_EXIT_BASE + signum @@ -69,6 +98,7 @@ class StageEngineCoreProc(EngineCoreProc): omni_coordinator_address: str | None = None, omni_stage_id: int | None = None, omni_replica_id: int = 0, + omni_parallel_stage_init: bool = False, **kwargs: Any, ) -> None: """Launch StageEngineCoreProc busy loop in background process. @@ -139,6 +169,15 @@ class StageEngineCoreProc(EngineCoreProc): # default sampling extra_args. maybe_apply_cfg_scheduler_patches(kwargs.get("vllm_config")) + # When parallel stage init is enabled, wrap this driver's executor + # so its memory-mutating phases (load / KV alloc / capture) hold a + # per-device LOCK_SH and its profiling measurement holds LOCK_EX. + # The wrapper must be installed here (in the engine-core child): + # phase boundaries live inside EngineCore.__init__, invisible to the + # orchestrator. See vllm_omni.engine.stage_phase_lock. + if omni_parallel_stage_init: + _install_phase_locks(kwargs, local_dp_rank) + engine_core = StageEngineCoreProc( *args, engine_index=dp_rank, diff --git a/vllm_omni/engine/stage_engine_core_proc_manager.py b/vllm_omni/engine/stage_engine_core_proc_manager.py index 6b613bf8e..b14c62d8f 100644 --- a/vllm_omni/engine/stage_engine_core_proc_manager.py +++ b/vllm_omni/engine/stage_engine_core_proc_manager.py @@ -75,6 +75,7 @@ class StageEngineCoreProcManager(CoreEngineProcManager): omni_replica_base_id: int = 0, client_handshake_address: str | None = None, tensor_queue: Queue | None = None, + omni_parallel_stage_init: bool = False, ) -> None: # NOTE: we intentionally do not call ``super().__init__`` — the # parent's body hardcodes the wrong target. We re-implement it here @@ -92,6 +93,7 @@ class StageEngineCoreProcManager(CoreEngineProcManager): "tensor_queue": tensor_queue, "omni_stage_id": int(omni_stage_id), "omni_coordinator_address": omni_coordinator_address, + "omni_parallel_stage_init": bool(omni_parallel_stage_init), } if client_handshake_address: diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 5abcbd249..71ba17c7c 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -974,6 +974,7 @@ def _launch_omni_core_engines( omni_coordinator_address: str | None = None, stage_visible_devices: str | None = None, spawn_device_lock: threading.Lock | None = None, + omni_parallel_stage_init: bool = False, ) -> Iterator[tuple[CoreEngineProcManager, DPCoordinator | None, EngineZmqAddresses]]: """Launch local engine cores using the omni registration flow. @@ -1068,8 +1069,20 @@ def _launch_omni_core_engines( omni_stage_id=stage_id, omni_coordinator_address=omni_coordinator_address, omni_replica_base_id=replica_id, + omni_parallel_stage_init=omni_parallel_stage_init, ) else: + if omni_parallel_stage_init: + # The upstream CoreEngineProcManager targets EngineCoreProc.run_engine_core, + # which does NOT install the SH/EX phase-lock executor wrapper. Since the + # orchestrator skips its own device lock when parallel_stage_init is on, + # running here would leave same-device init unguarded. Fail closed. + raise RuntimeError( + f"parallel_stage_init is enabled but stage {stage_id} would launch via the " + "upstream CoreEngineProcManager (no omni coordinator), which cannot install " + "the phase-lock guard. Disable parallel_stage_init for this deployment or " + "run with the omni coordinator." + ) with scoped_spawn_device_env( stage_visible_devices, spawn_device_lock, stage_id=stage_id, replica_id=replica_id ): @@ -1120,6 +1133,7 @@ def launch_stage_replica( omni_coordinator_address: str | None = None, stage_visible_devices: str | None = None, spawn_device_lock: threading.Lock | None = None, + omni_parallel_stage_init: bool = False, ) -> Iterator[StageReplicaResources]: """Launch a local LLM stage replica. @@ -1141,6 +1155,7 @@ def launch_stage_replica( omni_coordinator_address=omni_coordinator_address, stage_visible_devices=stage_visible_devices, spawn_device_lock=spawn_device_lock, + omni_parallel_stage_init=omni_parallel_stage_init, ) as resources: engine_manager, coordinator, addresses = resources yield StageReplicaResources( @@ -1170,6 +1185,7 @@ def launch_stage_replica( omni_stage_id=stage_id, omni_coordinator_address=omni_coordinator_address, omni_replica_base_id=replica_id, + omni_parallel_stage_init=omni_parallel_stage_init, ) with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 8efc4152a..f2bc42ab1 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -17,6 +17,7 @@ import importlib import json import multiprocessing as mp import os +import tempfile import time from collections.abc import Callable, Collection, Generator, Mapping, Sequence from contextlib import contextmanager @@ -1562,43 +1563,108 @@ def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor: return input_processor -def _cleanup_stale_lock_if_dead(lock_file: str) -> bool: - """If *lock_file* exists and its recorded PID is dead, unlink the file. +def device_init_lock_path(device_id: int, lock_dir: str = "/tmp") -> str: + """Return the per-physical-device initialization lock file path. - Returns ``True`` if the stale lock was cleaned up (caller should retry), - ``False`` otherwise (lock holder appears alive, or file could not be read). + Shared by the orchestrator-side ``acquire_device_locks`` (legacy full-init + ``LOCK_EX``) and the engine-core-side ``DevicePhaseLock`` (parallel-stage-init + SH/EX phase locks) so both coordinate on the *same* file per device. + + That coordination only holds while the **inode** is stable, so nothing may + unlink this path: two holders of the same pathname on different inodes do not + conflict. The PID written into the file is diagnostic only. + """ + return os.path.join(lock_dir, f"vllm_omni_device_{device_id}_init.lock") + + +def _open_existing_lock_file(lock_file: str) -> tuple[int, bool] | None: + """Open an already-existing lock file; ``None`` if it does not exist yet. + + Prefers ``O_RDWR`` so the holder can stamp its PID, but falls back to + ``O_RDONLY``: ``flock`` locks attach to the open file description and, unlike + POSIX ``fcntl`` record locks, require no write access, so a read-only + descriptor still provides full mutual exclusion. """ try: - with open(lock_file) as fh: - content = fh.read().strip() - if not content: - return False - pid = int(content) - except (OSError, ValueError): - return False - - # Check whether the PID is still alive. - try: - os.kill(pid, 0) - except ProcessLookupError: - # PID does not exist — stale lock. - logger.info( - "Removing stale device lock %s (PID %s is dead)", - lock_file, - pid, - ) - try: - os.unlink(lock_file) - return True - except OSError: - logger.debug("Failed to unlink stale lock %s", lock_file) - return False + return os.open(lock_file, os.O_RDWR), True + except FileNotFoundError: + return None except PermissionError: - # PID exists but we cannot signal it (different user) — treat as alive. - return False + # Created by another user on this shared machine; read-only still locks. + pass + try: + return os.open(lock_file, os.O_RDONLY), False + except FileNotFoundError: + return None + except PermissionError as exc: + raise PermissionError( + f"Device init lock {lock_file} exists but is not readable by this user " + f"({exc}). It coordinates GPU initialization across users, so it must stay " + "readable by all of them; have its owner remove it or chmod it to 0644." + ) from exc - # PID is alive — legitimate lock holder. - return False + +def open_device_lock_file(lock_file: str) -> tuple[int, bool]: + """Open a per-device lock file, tolerating one created by another user. + + Returns ``(fd, writable)``. + + These lock files coordinate *across* users -- two people running on the same + physical GPU must contend on the same file -- so every user has to be able to + open whichever one exists. Two things make that awkward: + + * whoever creates it owns it, so later users may only get read access, which + ``flock`` is perfectly happy with; and + * the mode passed to ``os.open`` is filtered by the creator's ``umask``, so + under 0027 or 0077 the file would land 0640 or 0600 and lock everyone else + out entirely. + + Creating it therefore stages a temporary file, widens it with ``fchmod``, and + publishes it with an atomic ``os.link``. Creating in place with ``O_EXCL`` and + widening afterwards would briefly expose the file at the umask-filtered mode, + and a concurrent user opening it in that window would be locked out -- the + very failure this avoids. ``link`` also fails cleanly if another process wins + the race, which keeps the "first creator wins" inode stable. + """ + existing = _open_existing_lock_file(lock_file) + if existing is not None: + return existing + + lock_dir = os.path.dirname(lock_file) or "." + tmp_fd, tmp_path = tempfile.mkstemp(dir=lock_dir, prefix=".vllm_omni_device_lock_") + published = False + try: + os.fchmod(tmp_fd, 0o644) + try: + os.link(tmp_path, lock_file) + except FileExistsError: + pass # another process published first; fall through and open theirs + else: + published = True + return tmp_fd, True + finally: + try: + os.unlink(tmp_path) + except OSError: + logger.debug("Could not remove staged device lock %s", tmp_path) + if not published: + os.close(tmp_fd) + + existing = _open_existing_lock_file(lock_file) + if existing is not None: + return existing + raise PermissionError(f"Device init lock {lock_file} could not be created or opened") + + +def record_lock_holder_pid(fd: int, writable: bool) -> None: + """Stamp the holder's PID into an open lock file (diagnostic only).""" + if not writable: + return + try: + os.ftruncate(fd, 0) + os.write(fd, f"{os.getpid()}\n".encode()) + except OSError: + pass def acquire_device_locks( @@ -1677,28 +1743,31 @@ def acquire_device_locks( # Acquire locks wait_start = time.time() for device_id in devices_to_lock: - lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock" + lock_file = device_init_lock_path(device_id) lock_acquired = False - already_cleaned_stale = False # only try stale cleanup once per device while not lock_acquired: try: - lock_fd = os.open(lock_file, os.O_CREAT | os.O_RDWR, 0o644) + lock_fd, lock_writable = open_device_lock_file(lock_file) try: fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - os.ftruncate(lock_fd, 0) - os.write(lock_fd, f"{os.getpid()}\n".encode()) - os.fsync(lock_fd) + record_lock_holder_pid(lock_fd, lock_writable) lock_acquired = True lock_fds.append(lock_fd) logger.debug("Acquired exclusive lock for device %s", device_id) except BlockingIOError: os.close(lock_fd) - # Detect and clean stale locks from dead processes. - if not already_cleaned_stale: - already_cleaned_stale = True - if _cleanup_stale_lock_if_dead(lock_file): - continue # retry flock immediately + # NOTE: no stale-lock cleanup here. ``flock`` is released + # by the kernel when the holder exits (SIGKILL included), + # so a dead holder never keeps this lock. Unlinking the + # path on a dead *recorded* PID was actively unsafe: the + # PID is written by every holder, and under the + # parallel-stage-init SH/EX protocol several LOCK_SH + # holders overwrite it. Unlinking while one still held the + # old inode let a contender create a fresh file with the + # same name and take LOCK_EX on it, which then conflicted + # with nobody -- breaking the interoperability between the + # legacy full-init lock and the phase locks. if time.time() - wait_start > stage_init_timeout: logger.warning( "Timeout waiting for device %s initialization lock, proceeding anyway", @@ -1707,8 +1776,10 @@ def acquire_device_locks( break time.sleep(0.01) except OSError as e: - logger.debug( - "Failed to acquire lock for device %s: %s, continuing anyway", + logger.warning( + "Failed to acquire lock for device %s: %s. Continuing WITHOUT " + "device-init serialization; concurrent initialization on this " + "device may incorrectly measure available memory.", device_id, e, ) diff --git a/vllm_omni/engine/stage_phase_lock.py b/vllm_omni/engine/stage_phase_lock.py new file mode 100644 index 000000000..2a723fdd0 --- /dev/null +++ b/vllm_omni/engine/stage_phase_lock.py @@ -0,0 +1,242 @@ +"""Per-device shared/exclusive init phase locks for parallel stage init. + +When ``VllmOmniOrchestratorConfig.parallel_stage_init`` is enabled, several +stage engine-core processes may initialize on the same physical GPU at the same +time. Correctness then rests on two mechanisms (see +``rfc_parallel_stage_init.zh.md``): + + * **Admission** (``stage_admission``) proves, before any stage launches, that + every physical device's summed budget fits — this is the hard OOM backstop. + * **SH/EX phase locks** (this module) make each memory *measurement* quiescent: + memory-mutating phases (weight load, KV allocation, CUDA-graph capture) take + a per-device shared (``LOCK_SH``) lock; the profiling measurement inside + ``determine_available_memory`` takes an exclusive (``LOCK_EX``) lock. The + exclusive holder waits for in-flight mutations to drain and blocks new ones + for its (seconds-long) window, so the device-level profiling number each + stage computes stays within its admitted budget. + +The lock is held by the engine-core **driver** process (one per stage / per local +DP rank) and wraps the driver-side collective calls, so a single acquisition +covers every worker rank with no per-rank locking (which would deadlock on +inverted lock order). +Acquisition is always over the driver's whole device slice in **sorted global +order** so different drivers cannot deadlock, and it is **fail-closed**: a lock +timeout raises rather than silently proceeding. + +The core locking logic here intentionally depends only on ``os``/``fcntl`` so it +is importable and unit-testable without a GPU; torch/platform imports are lazy. +""" + +from __future__ import annotations + +import contextlib +import fcntl +import os +import time +from collections.abc import Iterator + +from vllm.logger import init_logger + +from vllm_omni.engine.stage_init_utils import ( + device_init_lock_path, + open_device_lock_file, + record_lock_holder_pid, +) + +logger = init_logger(__name__) + +# Default upper bound on how long a single phase will wait for its device locks +# before failing closed. Generous: a peer's exclusive profile window is seconds, +# but a peer's whole shared load/capture span can be much longer. +_DEFAULT_LOCK_TIMEOUT_S = 900.0 + + +class DeviceLockTimeoutError(RuntimeError): + """Raised (fail-closed) when a phase cannot acquire its device locks in time.""" + + +def resolve_driver_device_ids(vllm_config, local_dp_rank: int = 0) -> list[int]: + """Physical GPU ids this engine-core driver owns, as a sorted list. + + A driver (one local vLLM DP rank) owns a contiguous slice of the visible + devices of size ``TP*PP*PCP*SP*CFG`` (DP excluded — each DP rank is its own + driver). Visible ids come from the device-control env var the parent set on + the child (e.g. ``CUDA_VISIBLE_DEVICES``); when unset, all devices are visible. + """ + from vllm_omni.platforms import current_omni_platform + + env_var = current_omni_platform.device_control_env_var + visible = os.environ.get(env_var) if env_var else None + physical: list[int] + if visible: + physical = [] + for tok in visible.split(","): + tok = tok.strip() + if not tok: + continue + try: + physical.append(int(tok)) + except ValueError: + # Non-integer visibility (UUID / MIG) can't be mapped to a lock + # id; fall back to positional ordinals below. + physical = [] + break + else: + physical = [] + if not physical: + physical = list(range(current_omni_platform.get_device_count())) + + pc = vllm_config.parallel_config + devices_per_dp = ( + int(getattr(pc, "tensor_parallel_size", 1)) + * int(getattr(pc, "pipeline_parallel_size", 1)) + * int(getattr(pc, "prefill_context_parallel_size", 1)) + * int(getattr(pc, "sequence_parallel_size", 1)) + * int(getattr(pc, "cfg_parallel_size", 1)) + ) + devices_per_dp = max(1, devices_per_dp) + start = int(local_dp_rank) * devices_per_dp + slice_ = physical[start : start + devices_per_dp] + if not slice_: + # Defensive: fall back to the whole visible set rather than lock nothing. + slice_ = physical + return sorted(slice_) + + +class DevicePhaseLock: + """Sorted, fail-closed, whole-set SH/EX flock over a driver's physical devices.""" + + def __init__( + self, + device_ids: list[int], + *, + timeout_s: float = _DEFAULT_LOCK_TIMEOUT_S, + lock_dir: str = "/tmp", + ) -> None: + self._device_ids = sorted(device_ids) + self._timeout_s = timeout_s + self._lock_dir = lock_dir + + @classmethod + def from_child( + cls, + vllm_config, + local_dp_rank: int = 0, + *, + timeout_s: float = _DEFAULT_LOCK_TIMEOUT_S, + ) -> DevicePhaseLock: + """Build the lock for the current engine-core child from its config/env.""" + device_ids = resolve_driver_device_ids(vllm_config, local_dp_rank) + return cls(device_ids, timeout_s=timeout_s) + + @property + def device_ids(self) -> list[int]: + return list(self._device_ids) + + # ---- acquisition ------------------------------------------------------- + + def _acquire(self, mode: int) -> list[int]: + """Acquire *mode* (LOCK_SH|LOCK_EX) on every device in sorted order. + + Fail-closed: on timeout, releases whatever was acquired and raises + ``DeviceLockTimeoutError``. Sorted global order makes cross-driver deadlock + impossible. + """ + acquired: list[int] = [] + deadline = time.monotonic() + self._timeout_s + try: + for device_id in self._device_ids: + fd, writable = open_device_lock_file(device_init_lock_path(device_id, self._lock_dir)) + while True: + try: + fcntl.flock(fd, mode | fcntl.LOCK_NB) + break + except BlockingIOError: + if time.monotonic() > deadline: + os.close(fd) + raise DeviceLockTimeoutError( + f"Timed out after {self._timeout_s:.0f}s acquiring " + f"{'EX' if mode == fcntl.LOCK_EX else 'SH'} lock on device {device_id}" + ) + time.sleep(0.01) + record_lock_holder_pid(fd, writable) + acquired.append(fd) + return acquired + except BaseException: + _release_fds(acquired) + raise + + @contextlib.contextmanager + def shared(self) -> Iterator[None]: + """Hold ``LOCK_SH`` on the device slice for a memory-mutating phase.""" + if not self._device_ids: + yield + return + fds = self._acquire(fcntl.LOCK_SH) + try: + yield + finally: + _release_fds(fds) + + @contextlib.contextmanager + def exclusive(self) -> Iterator[None]: + """Hold ``LOCK_EX`` on the device slice for the profiling measurement.""" + if not self._device_ids: + yield + return + fds = self._acquire(fcntl.LOCK_EX) + try: + yield + finally: + _release_fds(fds) + + +def _release_fds(fds: list[int]) -> None: + for fd in fds: + try: + fcntl.flock(fd, fcntl.LOCK_UN) + os.close(fd) + except (OSError, ValueError): + pass + + +def wrap_executor_with_phase_locks(executor_class, locker: DevicePhaseLock): + """Return a subclass of *executor_class* that brackets init phases with locks. + + * ``__init__`` (executor construction → weight load) under ``shared()``. + * ``determine_available_memory`` (profiling) under ``exclusive()``. + * ``initialize_from_config`` (KV cache allocation) under ``shared()``. + * ``compile_or_warm_up_model`` (kernel warmup + CUDA-graph capture) under + ``shared()``. + + ``compile_or_warm_up_model`` needs its own bracket: upstream used to fuse it + into ``Executor.initialize_from_config``, but since vLLM v0.20.0 (vllm#39240) + ``EngineCore._initialize_kv_caches`` calls the two separately. Leaving it + unwrapped would let one stage capture graphs — mutating device memory — while + a peer holds the exclusive lock for its profiling measurement, which is + exactly the interference this protocol exists to prevent. + + Works uniformly for any executor (UniProc/MultiProc): the lock is taken in + the driver process that instantiates the executor and issues its collectives. + """ + + class _PhaseLockedExecutor(executor_class): + def __init__(self, vllm_config): + with locker.shared(): + super().__init__(vllm_config) + + def determine_available_memory(self): + with locker.exclusive(): + return super().determine_available_memory() + + def initialize_from_config(self, kv_cache_configs): + with locker.shared(): + return super().initialize_from_config(kv_cache_configs) + + def compile_or_warm_up_model(self): + with locker.shared(): + return super().compile_or_warm_up_model() + + _PhaseLockedExecutor.__name__ = f"PhaseLocked{executor_class.__name__}" + _PhaseLockedExecutor.__qualname__ = _PhaseLockedExecutor.__name__ + return _PhaseLockedExecutor diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 049d09688..93e482b42 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -9,6 +9,7 @@ import concurrent.futures import copy import os import threading +import time from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass @@ -125,6 +126,7 @@ class StageRuntime: stage_init_timeout: int, async_chunk: bool, tokenizer: str | None = None, + parallel_stage_init: bool = False, log_stats: bool = False, ) -> None: self._stage_configs = stage_configs @@ -133,6 +135,11 @@ class StageRuntime: self._stage_init_timeout = stage_init_timeout self._async_chunk = async_chunk self._tokenizer = tokenizer + # When True, same-device stages initialize concurrently, coordinated by + # pre-launch admission control + the engine-core-held SH/EX device locks + # (see VllmOmniOrchestratorConfig.parallel_stage_init). Default False + # keeps the legacy per-device LOCK_EX serialization. + self._parallel_stage_init = parallel_stage_init self._log_stats = log_stats self._num_stages = len(stage_configs) @@ -450,6 +457,13 @@ class StageRuntime: init_state_lock = threading.Lock() self._init_visible_devices_baseline = os.environ.get(current_omni_platform.device_control_env_var) + if self._parallel_stage_init: + # Refuse executor backends the phase locks cannot cover, then prove + # every physical device's summed budget fits BEFORE any stage spawns. + # Both fail fast here rather than OOM at allocation time. + self._reject_unguardable_executors(stage_plans) + self._run_stage_admission(stage_plans) + init_groups: dict[str, list[tuple[int, ReplicaInitPlan]]] = {} for plan in stage_plans: for replica in plan.replicas: @@ -506,8 +520,154 @@ class StageRuntime: return initialized_clients_by_stage + def _reject_unguardable_executors(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None: + """Fail closed for executors the SH/EX phase locks cannot actually guard. + + ``DevicePhaseLock`` is a node-local ``flock`` over files on the engine-core + driver's host, and the admission ledger is built only from this host's + devices. A Ray executor may place the workers that load weights, profile, + allocate KV cache and capture graphs on **other** nodes, which participate + in neither mechanism. Running there would take the parallel path — the + parent skips its own device lock when ``parallel_stage_init`` is on — + while only advertising the protection, so refuse before anything spawns. + + Ray detection uses both the executor's ``uses_ray`` class attribute and + ``ParallelConfig.use_ray``. Multi-node multiprocessing/external-launcher + configurations are detected through ``ParallelConfig.nnodes``. Both are + refused because some workers can live outside the driver's host. + + Remote replicas are skipped: they are launched by the runtime owning their + node and are already ``ADMISSION_EXEMPT`` in this host's ledger. + """ + offenders: list[str] = [] + for plan in stage_plans: + for replica in plan.replicas: + if replica.launch_mode == "remote": + continue + executor_class = replica.executor_class + parallel_config = getattr(replica.stage_vllm_config, "parallel_config", None) + uses_ray = bool(getattr(executor_class, "uses_ray", False)) or bool( + getattr(parallel_config, "use_ray", False) + ) + nnodes = int(getattr(parallel_config, "nnodes", 1) or 1) + reasons: list[str] = [] + if uses_ray: + reasons.append("Ray-backed") + if nnodes > 1: + reasons.append(f"multi-node (nnodes={nnodes})") + if reasons: + executor_name = getattr(executor_class, "__name__", "unknown executor") + offenders.append( + f"stage{replica.metadata.stage_id}/replica{replica.replica_id} " + f"({executor_name}; {', '.join(reasons)})" + ) + if offenders: + raise RuntimeError( + "parallel_stage_init cannot guard Ray-backed or multi-node executors: the SH/EX " + "phase locks are node-local file locks and the admission ledger only accounts " + "for this host's devices, so workers placed on other nodes would initialize unprotected. " + f"Offending replicas: {', '.join(offenders)}. Disable parallel_stage_init for " + "this deployment, or use a single-node non-Ray executor backend." + ) + + def _run_stage_admission(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None: + """Pre-launch per-device admission for parallel stage init (fail-fast).""" + from vllm_omni.engine.stage_admission import ( + ADMISSION_EXEMPT, + AdmissionExempt, + StageAdmissionError, + check_admission, + ) + + def _resolve(replica: ReplicaInitPlan) -> list[int] | AdmissionExempt | None: + if replica.launch_mode == "remote": + # Remote replicas consume a remote node's memory and the ledger + # is per-LOCAL-device. Remote LLM plans already carry + # runtime_cfg=None, but remote diffusion replicas keep their + # cfg — mark both exempt explicitly. Only this exemption may + # skip the ledger: a LOCAL replica that resolves to None below + # fails admission (fail-closed in check_admission). + return ADMISSION_EXEMPT + resolved = self._resolve_replica_physical_devices(replica.metadata.stage_id, replica.metadata.runtime_cfg) + if not resolved: + return None + ids: list[int] = [] + for tok in str(resolved).split(","): + tok = tok.strip() + if not tok: + continue + try: + ids.append(int(tok)) + except ValueError: + # Non-integer visibility (UUID / MIG) can't be reasoned about. + return None + return ids or None + + def _visible_ordinal(physical_id: int) -> int: + """Translate a physical device id to this process's visible ordinal. + + The ledger is keyed by *physical* id (resolved against the visibility + baseline captured before any stage narrowed the env), but + ``get_device_total_memory`` and ``torch.cuda.get_device_properties`` + interpret their argument as an ordinal in the **current** process's + visible-device namespace. Passing the physical id straight through + raises an invalid-ordinal error under a restricted + ``CUDA_VISIBLE_DEVICES`` and, when visibility is reordered, silently + reads a different GPU — which yields a wrong capacity, and therefore + an unsafe admission result, on heterogeneous devices. + """ + baseline = self._init_visible_devices_baseline + if not baseline: + # No device-control env captured: ordinals are physical ids. + return physical_id + for ordinal, tok in enumerate(baseline.split(",")): + tok = tok.strip() + if not tok: + continue + try: + candidate = int(tok) + except ValueError: + # UUID/MIG entries cannot be converted to physical integer + # ids, but they still occupy an ordinal in the visible list. + continue + if candidate == physical_id: + return ordinal + raise StageAdmissionError( + f"Physical device {physical_id} is not in the visibility baseline " + f"{baseline!r} captured at init, so its capacity cannot be measured " + "from this process. Refusing to admit against an unknown device." + ) + + def _total_memory(device_id: int) -> int: + ordinal = _visible_ordinal(device_id) + try: + mem = current_omni_platform.get_device_total_memory(ordinal) + if mem: + return int(mem) + except (NotImplementedError, AttributeError, TypeError): + pass + import torch + + return int(torch.cuda.get_device_properties(ordinal).total_memory) + + check_admission( + stage_plans, + resolve_physical_devices=_resolve, + device_total_memory=_total_memory, + ) + def _replica_init_group_key(self, replica: ReplicaInitPlan) -> str: - """Return the scheduling group used during replica initialization.""" + """Return the scheduling group used during replica initialization. + + Replicas sharing a group initialize sequentially; different groups + initialize in parallel threads. Local LLM replicas are keyed by their + **resolved canonical physical device set** so stages on different + physical GPUs land in different groups (parallel) while stages sharing a + device stay in one group (serialized, then further guarded by the + per-device ``LOCK_EX`` file lock). Keying on the raw ``runtime.devices`` + config value instead would fail to parallelize logically-distinct stages + that map to different physical GPUs. + """ if replica.launch_mode == "local" and replica.metadata.stage_type == "diffusion": # Local diffusion process spawning must stay on the orchestrator # thread. Keep all local diffusion replicas in one sequential group. @@ -515,9 +675,28 @@ class StageRuntime: if replica.launch_mode == "remote": return f"remote:{replica.metadata.stage_id}:{replica.replica_id}" + physical_devices = self._resolve_replica_physical_devices( + replica.metadata.stage_id, + replica.metadata.runtime_cfg, + ) + if self._parallel_stage_init: + # Same-device concurrency is coordinated by the engine-core SH/EX + # device locks + pre-launch admission, so give every replica its own + # group to let them all initialize in parallel. + return f"parallel:{replica.metadata.stage_id}:{replica.replica_id}" + runtime_cfg = replica.metadata.runtime_cfg or {} - devices = runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) - return f"device:{devices}" + raw_devices = ( + runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) + ) + if str(raw_devices) != str(physical_devices): + logger.debug( + "[stage_init] Stage-%s init-group key: raw devices=%s -> resolved physical=%s", + replica.metadata.stage_id, + raw_devices, + physical_devices, + ) + return f"device:{physical_devices}" def _initialize_replica( self, @@ -566,28 +745,76 @@ class StageRuntime: raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing executor_class") if plan.engine_args_dict is None: raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing engine args") - with self._scoped_spawn_device_env(physical_devices): - lock_fds = acquire_device_locks( + # G3 device locks: in the default (serial) path the parent holds a + # per-device LOCK_EX across the whole child init. When parallel stage + # init is enabled the engine-core child takes phased SH/EX locks + # itself (see stage_phase_lock), so the parent must NOT also hold the + # lock here — it would self-deadlock while waiting for the child's + # READY handshake. + if not self._parallel_stage_init: + g3_start = time.perf_counter() + with self._scoped_spawn_device_env(physical_devices): + lock_fds = acquire_device_locks( + plan.metadata.stage_id, + plan.engine_args_dict, + stage_init_timeout, + ) + logger.debug( + "[stage_init] Stage-%s G3 device-lock acquire took %.3fs", plan.metadata.stage_id, - plan.engine_args_dict, - stage_init_timeout, + time.perf_counter() - g3_start, + ) + + launch_cm = launch_stage_replica( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=self._log_stats, + stage_id=plan.metadata.stage_id, + replica_id=plan.replica_id, + stage_config=plan.stage_cfg, + 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, + 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() + 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, ) - # Serialize engine-core spawning across all LLM replicas to avoid - # ZMQ port-allocation races and simultaneous CUDA context init. - with self._replica_launch_lock, stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg): - with launch_stage_replica( - vllm_config=vllm_config, - executor_class=executor_class, - log_stats=self._log_stats, - stage_id=plan.metadata.stage_id, - replica_id=plan.replica_id, - stage_config=plan.stage_cfg, - 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, - ) as resources: - pass logger.info("[StageRuntime] Stage %s engine startup completed", plan.metadata.stage_id) if resources is None: @@ -773,6 +1000,7 @@ class DistStageRuntime(StageRuntime): omni_heartbeat_timeout: float = 30.0, omni_lb_policy: str = "random", request_queue: janus.Queue[EngineQueueMessage] | None = None, + parallel_stage_init: bool = False, ) -> None: super().__init__( stage_configs=stage_configs, @@ -781,6 +1009,7 @@ class DistStageRuntime(StageRuntime): stage_init_timeout=stage_init_timeout, async_chunk=async_chunk, tokenizer=tokenizer, + parallel_stage_init=parallel_stage_init, log_stats=log_stats, ) self._single_stage_id_filter = single_stage_id_filter @@ -1101,6 +1330,7 @@ def create_stage_runtime( stage_init_timeout: int, async_chunk: bool, tokenizer: str | None = None, + parallel_stage_init: bool = False, # Distributed-only params: single_stage_id_filter: int | None = None, omni_master_address: str | None = None, @@ -1122,6 +1352,7 @@ def create_stage_runtime( stage_init_timeout=stage_init_timeout, async_chunk=async_chunk, tokenizer=tokenizer, + parallel_stage_init=parallel_stage_init, log_stats=log_stats, single_stage_id_filter=single_stage_id_filter, omni_master_address=omni_master_address, @@ -1138,5 +1369,6 @@ def create_stage_runtime( stage_init_timeout=stage_init_timeout, async_chunk=async_chunk, tokenizer=tokenizer, + parallel_stage_init=parallel_stage_init, log_stats=log_stats, ) diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py index d616347d8..2a7e9c876 100644 --- a/vllm_omni/worker/base.py +++ b/vllm_omni/worker/base.py @@ -1,4 +1,4 @@ -"""Base worker class for vLLM-Omni with process-scoped GPU memory accounting.""" +"""Base worker class for vLLM-Omni with device-level GPU memory profiling.""" from __future__ import annotations @@ -17,22 +17,20 @@ from vllm_omni.diffusion.data import ( OmniSleepTask, OmniWakeTask, ) -from vllm_omni.entrypoints.utils import detect_pid_host from vllm_omni.platforms import current_omni_platform -from vllm_omni.worker.gpu_memory_utils import ( - get_process_gpu_memory, - is_process_scoped_memory_available, -) logger = init_logger(__name__) class OmniGPUWorkerBase(GPUWorker): - """Base GPU worker for vLLM-Omni with process-scoped memory accounting. + """Base GPU worker for vLLM-Omni. - This class overrides determine_available_memory() to use per-process GPU - memory tracking via pynvml, allowing multiple stages to initialize - concurrently on the same GPU without memory accounting interference. + This class overrides determine_available_memory() to compute the KV cache + budget from device-level profiling data. Concurrent same-device stage + initialization is coordinated externally by the parallel-stage-init + admission control + SH/EX device locks (see + ``VllmOmniOrchestratorConfig.parallel_stage_init``), which keep each + measurement quiescent instead of relying on per-process NVML accounting. It also replaces vLLM's TorchProfilerWrapper with OmniTorchProfilerWrapper for custom trace naming, background gzip, and trace path collection. @@ -91,21 +89,27 @@ class OmniGPUWorkerBase(GPUWorker): @torch.inference_mode() def determine_available_memory(self) -> int: - """Process-scoped GPU memory profiling for concurrent stage initialization. + """Device-level GPU memory profiling for the KV cache budget. Algorithm: 1. requested_memory = total_gpu_memory * gpu_memory_utilization (computed in init_device from cache_config) - 2. process_memory = memory used by THIS process only (via pynvml) - - Uses nvmlDeviceGetComputeRunningProcesses to get per-PID memory - - Supports CUDA_VISIBLE_DEVICES with indices, UUIDs, or MIG IDs + 2. profiled_usage = weights + peak_activation + non_torch_increase + (measured by ``memory_profiling`` around ``profile_run()``; + ``non_torch_increase`` is device-level, so it reflects whatever + else is resident on the GPU at profiling time) - 3. available_kv_cache = requested_memory - process_memory + 3. available_kv_cache = requested_memory - profiled_usage - Fallback: - If NVML is unavailable, falls back to profiling data: - available = requested - (weights + activations + non_torch) + Note: + Process-scoped NVML estimation was removed in favour of the + device-level profiling fallback plus the parallel-stage-init + admission + SH/EX device locks (see + ``VllmOmniOrchestratorConfig.parallel_stage_init``): the EX lock + makes each measurement quiescent, so the device-level number is the + correct, conservative budget. The NVML helpers in + ``gpu_memory_utils`` are retained for diffusion memory reporting. """ if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes: self.model_runner.profile_run() @@ -127,49 +131,25 @@ class OmniGPUWorkerBase(GPUWorker): # true here). Mirror upstream so the omni override keeps it populated. self.total_consumed = profile_result.total_consumed - process_memory = ( - get_process_gpu_memory(self.local_rank) - if is_process_scoped_memory_available() and detect_pid_host() - else None + profiled_usage = ( + int(self.model_runner.model_memory_usage) + + profile_result.torch_peak_increase + + profile_result.non_torch_increase + ) + self.available_kv_cache_memory_bytes = max(0, self.requested_memory - profiled_usage) + logger.debug( + "Profiling KV budget (PID %d, GPU %d): requested=%s, profiled=%s, available=%s", + os.getpid(), + self.local_rank, + format_gib(self.requested_memory), + format_gib(profiled_usage), + format_gib(self.available_kv_cache_memory_bytes), + ) + logger.info_once( + "Available KV cache memory: %s GiB (device-level profiling)", + format_gib(self.available_kv_cache_memory_bytes), + scope="local", ) - - if process_memory is not None: - # NVML available: use per-process memory - self.available_kv_cache_memory_bytes = max(0, self.requested_memory - process_memory) - logger.debug( - "Process-scoped memory (PID %d, GPU %d): requested=%s, used=%s, available=%s", - os.getpid(), - self.local_rank, - format_gib(self.requested_memory), - format_gib(process_memory), - format_gib(self.available_kv_cache_memory_bytes), - ) - logger.info_once( - "Available KV cache memory: %s GiB (process-scoped)", - format_gib(self.available_kv_cache_memory_bytes), - scope="local", - ) - else: - # NVML unavailable: use profiling data as conservative fallback - profiled_usage = ( - int(self.model_runner.model_memory_usage) - + profile_result.torch_peak_increase - + profile_result.non_torch_increase - ) - self.available_kv_cache_memory_bytes = max(0, self.requested_memory - profiled_usage) - logger.debug( - "Profiling fallback (PID %d, GPU %d): requested=%s, profiled=%s, available=%s", - os.getpid(), - self.local_rank, - format_gib(self.requested_memory), - format_gib(profiled_usage), - format_gib(self.available_kv_cache_memory_bytes), - ) - logger.info_once( - "Available KV cache memory: %s GiB (profiling fallback)", - format_gib(self.available_kv_cache_memory_bytes), - scope="local", - ) return int(self.available_kv_cache_memory_bytes)