[Bugfix][Engine] Group overlapping device stages into one sequential init component (#7328)

Signed-off-by: ZhengWG <zwg0606@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

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):

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)
 
 
 # --------------------------------------------------------------------------- #

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],

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,