diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index b1771da28..e6f7f8a8b 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -926,10 +926,7 @@ def test_remote_replicas_use_distinct_init_group_keys(): replica.launch_mode = "remote" replica.metadata.runtime_cfg = None - assert [runtime._replica_init_group_key(replica) for replica in plan.replicas] == [ - "remote:1:0", - "remote:1:1", - ] + assert runtime._init_group_keys(plan.replicas) == ["remote:1:0", "remote:1:1"] def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_replica_failure(monkeypatch): diff --git a/tests/engine/test_parallel_stage_init.py b/tests/engine/test_parallel_stage_init.py index 5952d710e..d88d49825 100644 --- a/tests/engine/test_parallel_stage_init.py +++ b/tests/engine/test_parallel_stage_init.py @@ -1,7 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """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), + * init-group device keying (overlap components; 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, @@ -18,6 +21,7 @@ import inspect import os import subprocess import threading +import time import types from pathlib import Path @@ -37,7 +41,9 @@ from vllm_omni.engine.stage_init_utils import ( LogicalStageInitPlan, ReplicaInitPlan, device_init_lock_path, + device_overlap_group_keys, open_device_lock_file, + parse_physical_device_ids, ) from vllm_omni.engine.stage_phase_lock import ( DeviceLockTimeoutError, @@ -46,6 +52,7 @@ from vllm_omni.engine.stage_phase_lock import ( wrap_executor_with_phase_locks, ) from vllm_omni.engine.stage_runtime import StageRuntime +from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -104,49 +111,114 @@ def _fake_vllm_config(util: float, *, model="m", dtype="bf16", cudagraph_mode="F # --------------------------------------------------------------------------- # # A1: init-group device key # --------------------------------------------------------------------------- # -def test_init_group_key_resolves_physical_devices_serial(): +@pytest.mark.parametrize( + ("devices", "expected"), + [ + ("0", {0}), + ("1,0,1", {0, 1}), + (" 2, 3 ", {2, 3}), + (None, None), + ("", None), + ("GPU-uuid", None), + ], +) +def test_parse_physical_device_ids(devices, expected): + assert parse_physical_device_ids(devices) == expected + + +@pytest.mark.parametrize( + ("device_sets", "expected"), + [ + # equal / disjoint + ([{0}, {1}, {0}], ["device-group:0", "device-group:1", "device-group:0"]), + # unequal overlap, canonical order + ([{0, 1}, {0}, {2}], ["device-group:0,1", "device-group:0,1", "device-group:2"]), + # transitive: {0,1} ~ {1,2} ~ {2} + ([{0, 1}, {1, 2}, {2}], ["device-group:0,1,2"] * 3), + # unresolved may touch any GPU -> everything serial + ([{0}, None, {1}], ["device-group:*"] * 3), + ([], []), + ], +) +def test_device_overlap_group_keys(device_sets, expected): + sets = [None if s is None else frozenset(s) for s in device_sets] + assert device_overlap_group_keys(sets) == expected + + +def test_init_group_keys_serial_groups_by_device_overlap(): 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) + keys = runtime._init_group_keys( + [_llm_replica(0, 0, "0,1"), _llm_replica(1, 0, "0"), _llm_replica(2, 0, "2"), _llm_replica(3, 0, "1,0")] + ) - 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 + assert keys == ["device-group:0,1", "device-group:0,1", "device-group:2", "device-group:0,1"] -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(): +def test_init_group_keys_fixed_for_parallel_remote_and_diffusion(): 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" + + # Same device everywhere; parallel mode still keys per replica (child SH/EX locks coordinate). + keys = runtime._init_group_keys([_llm_replica(0, 0, "0"), _llm_replica(0, 1, "0"), remote, diffusion]) + + assert keys == ["parallel:0:0", "parallel:0:1", "remote:1:0", "inline:diffusion"] + + +def test_unresolved_devices_collapse_serial_groups_and_warn(caplog): + runtime = _runtime(parallel_stage_init=False) + with caplog.at_level("WARNING"): + keys = runtime._init_group_keys([_llm_replica(0, 0, "0"), _llm_replica(1, 0, "GPU-uuid")]) + + assert keys == ["device-group:*", "device-group:*"] + assert "Stage-1 replica 0" in caplog.text + assert "GPU-uuid" in caplog.text + + +def _stage_plans(*devices: str) -> list[LogicalStageInitPlan]: + return [ + LogicalStageInitPlan(stage_idx=i, stage_id=i, replicas=[_llm_replica(i, 0, dev)]) + for i, dev in enumerate(devices) + ] + + +def test_overlapping_device_stages_initialize_sequentially(monkeypatch): + monkeypatch.delenv(current_omni_platform.device_control_env_var, raising=False) + runtime = _runtime(parallel_stage_init=False) + in_flight, peak = 0, 0 + guard = threading.Lock() + + def _initialize_replica(plan, _timeout): + nonlocal in_flight, peak + with guard: + in_flight += 1 + peak = max(peak, in_flight) + time.sleep(0.05) + with guard: + in_flight -= 1 + return object() + + monkeypatch.setattr(runtime, "_initialize_replica", _initialize_replica) + runtime._initialize_stage_replicas(_stage_plans("0,1", "0"), stage_init_timeout=5) + + assert peak == 1 + + +def test_disjoint_device_stages_initialize_in_parallel(monkeypatch): + monkeypatch.delenv(current_omni_platform.device_control_env_var, raising=False) + runtime = _runtime(parallel_stage_init=False) + barrier = threading.Barrier(2, timeout=2) # breaks unless both inits are in flight together + + def _initialize_replica(_plan, _timeout): + barrier.wait() + return object() + + monkeypatch.setattr(runtime, "_initialize_replica", _initialize_replica) + runtime._initialize_stage_replicas(_stage_plans("0", "1"), stage_init_timeout=5) # --------------------------------------------------------------------------- # diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 1e57036a0..0562d2b08 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1668,6 +1668,41 @@ def record_lock_holder_pid(fd: int, writable: bool) -> None: pass +def parse_physical_device_ids(devices: str | None) -> frozenset[int] | None: + """Parse ``"0,1"`` into ``{0, 1}``; ``None`` if missing, empty or non-integer (UUID/MIG).""" + tokens = [tok.strip() for tok in str(devices or "").split(",") if tok.strip()] + if not tokens or not all(tok.isdigit() for tok in tokens): + return None + return frozenset(int(tok) for tok in tokens) + + +def device_overlap_group_keys(device_sets: Sequence[frozenset[int] | None]) -> list[str]: + """Key each device set by the connected component of sets it shares a GPU with. + + Transitively overlapping sets get one key, the sorted union of the component + (``{0,1}`` and ``{0}`` -> ``device-group:0,1``); disjoint sets get distinct + keys. ``None`` (unresolved: may touch any GPU) overlaps everything, so one + ``None`` collapses all sets into a single group. + """ + resolved = [devices for devices in device_sets if devices is not None] + if len(resolved) != len(device_sets): + return ["device-group:*"] * len(device_sets) + + components: list[set[int]] = [] + for devices in resolved: + merged = set(devices) + disjoint = [] + for comp in components: + if comp & merged: + merged |= comp + else: + disjoint.append(comp) + components = [*disjoint, merged] + + key_of = {device: "device-group:" + ",".join(map(str, sorted(comp))) for comp in components for device in comp} + return [key_of[next(iter(devices))] for devices in resolved] + + def acquire_device_locks( stage_id: int, engine_args_dict: dict[str, Any], diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 93e482b42..5d7d07653 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -52,11 +52,13 @@ from vllm_omni.engine.stage_init_utils import ( build_llm_stage_output_processor, build_vllm_config, compute_replica_layout, + device_overlap_group_keys, extract_legacy_stage_metadata, get_stage_connector_spec, inject_kv_stage_info, inject_omni_kv_connector_config, load_omni_transfer_config_for_model, + parse_physical_device_ids, prepare_engine_environment, release_device_locks, stage_runtime_env, @@ -446,9 +448,10 @@ class StageRuntime: ) -> dict[int, list[StagePoolClient | None]]: """Initialize all stage replicas. - Stages sharing the same GPU are initialized sequentially to avoid - memory profiling interference. Stages on different GPUs are - initialized in parallel. + Stages that share any physical GPU — including overlapping-but-unequal + sets such as ``{0,1}`` and ``{0}`` — initialize sequentially to avoid + memory profiling interference and same-process ``flock`` self-contention. + Stages on disjoint GPUs initialize in parallel. """ initialized_clients_by_stage: dict[int, list[StagePoolClient | None]] = { plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans @@ -464,10 +467,7 @@ class StageRuntime: 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: - init_groups.setdefault(self._replica_init_group_key(replica), []).append((plan.stage_idx, replica)) + init_groups = self._build_init_groups(stage_plans) def _init_group(group: list[tuple[int, ReplicaInitPlan]]) -> None: """Initialize replicas in one scheduling group sequentially.""" @@ -589,19 +589,8 @@ class StageRuntime: # 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 + parsed = parse_physical_device_ids(resolved) + return sorted(parsed) if parsed else None def _visible_ordinal(physical_id: int) -> int: """Translate a physical device id to this process's visible ordinal. @@ -656,47 +645,58 @@ class StageRuntime: device_total_memory=_total_memory, ) - def _replica_init_group_key(self, replica: ReplicaInitPlan) -> str: - """Return the scheduling group used during replica initialization. + def _build_init_groups( + self, + stage_plans: Sequence[LogicalStageInitPlan], + ) -> dict[str, list[tuple[int, ReplicaInitPlan]]]: + items = [(plan.stage_idx, replica) for plan in stage_plans for replica in plan.replicas] + groups: dict[str, list[tuple[int, ReplicaInitPlan]]] = {} + for key, item in zip(self._init_group_keys([replica for _, replica in items]), items, strict=True): + groups.setdefault(key, []).append(item) + return groups - 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. + def _init_group_keys(self, replicas: Sequence[ReplicaInitPlan]) -> list[str]: + """One init-group key per replica: same key -> sequential, distinct -> parallel. + + Diffusion, remote and parallel-stage-init replicas keep fixed per-replica + keys. Serial local LLM replicas are grouped by connected components of + their resolved physical-device overlap, so ``{0,1}`` and ``{0}`` share a + group and never re-``flock`` a device this process already holds. """ + keys = [self._init_group_key_override(replica) for replica in replicas] + pending = [i for i, key in enumerate(keys) if key is None] + device_sets: list[frozenset[int] | None] = [] + for i in pending: + replica = replicas[i] + physical = self._resolve_replica_physical_devices(replica.metadata.stage_id, replica.metadata.runtime_cfg) + parsed = parse_physical_device_ids(physical) + if parsed is None: + logger.warning( + "[stage_init] Stage-%s replica %s physical devices %r are not integer GPU ids; " + "all serial LLM replicas will share one init group", + replica.metadata.stage_id, + replica.replica_id, + physical, + ) + device_sets.append(parsed) + for i, key in zip(pending, device_overlap_group_keys(device_sets), strict=True): + keys[i] = key + return cast(list[str], keys) + + def _init_group_key_override(self, replica: ReplicaInitPlan) -> str | None: + """Key that bypasses device-overlap grouping, or ``None`` to use it.""" 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. return "inline:diffusion" 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 {} - 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}" + return None def _initialize_replica( self,