diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 5fae5f1b2..0f7d4dab6 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -15,6 +15,7 @@ from vllm_omni.engine.stage_init_utils import ( ReplicaInitPlan, build_stage0_input_processor, compute_replica_layout, + split_devices_for_replicas, ) from vllm_omni.engine.stage_runtime import StageRuntime @@ -232,6 +233,76 @@ def test_compute_replica_layout_splits_diffusion_devices_by_world_size(): assert replica_devices_map == {0: ["0,1", "2,3"]} +@pytest.mark.parametrize("num_replicas", [1, 2, 4]) +def test_split_devices_for_replicas_returns_one_slot_per_replica_when_devices_unset(num_replicas): + """``devices`` is optional, and the result is indexed by replica id. + + A stage that declares ``num_replicas`` without ``devices`` lets every + replica inherit the launcher's CUDA_VISIBLE_DEVICES, so each one still + needs its own (empty) slot. + """ + assert split_devices_for_replicas(None, num_replicas, 1, 0) == [None] * num_replicas + + +def test_compute_replica_layout_covers_every_replica_when_devices_unset(): + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={}, + runtime=types.SimpleNamespace(num_replicas=3), + ) + + replicas_per_stage, replica_devices_map = compute_replica_layout([stage_cfg]) + + assert replicas_per_stage == [3] + assert replica_devices_map == {0: [None, None, None]} + + +def test_build_logical_stage_init_plans_handles_stage_without_devices(monkeypatch): + """Regression: a multi-replica stage with no ``devices`` must still plan. + + ``_build_logical_stage_init_plans`` indexes ``replica_devices_map`` by + replica id, so a short list raised ``IndexError`` for every replica after + the first. + """ + import vllm_omni.engine.stage_runtime as runtime_mod + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={}, + runtime=types.SimpleNamespace(num_replicas=3), + ) + runtime = StageRuntime( + stage_configs=[stage_cfg], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + ) + + monkeypatch.setattr( + runtime_mod, + "extract_legacy_stage_metadata", + lambda cfg: _make_llm_metadata(cfg.stage_id), + ) + monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(runtime_mod, "build_engine_args_dict", lambda *_, **__: {}) + monkeypatch.setattr(runtime_mod, "build_vllm_config", lambda *_args, **_kwargs: (types.SimpleNamespace(), object)) + + replicas_per_stage, replica_devices_map = compute_replica_layout([stage_cfg]) + stage_plans = runtime._build_logical_stage_init_plans( + omni_transfer_config=None, + replicas_per_stage=replicas_per_stage, + replica_devices_map=replica_devices_map, + ) + + assert [replica.replica_id for replica in stage_plans[0].replicas] == [0, 1, 2] + assert [replica.stage_cfg.runtime.devices for replica in stage_plans[0].replicas] == [None, None, None] + + def test_collect_initialized_clients_for_cleanup_deduplicates_clients(): shared = types.SimpleNamespace(name="shared") extra = types.SimpleNamespace(name="extra") diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 683726d28..a4c3e1650 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -567,9 +567,14 @@ def split_devices_for_replicas( num_replicas: int, tp_size: int, stage_id: int, -) -> list[str]: +) -> list[str | None]: """Split a devices string into per-replica subsets. + The result always has one entry per replica, because callers index it by + replica id. When ``devices_str`` is ``None`` the stage declares no explicit + placement, so every replica gets ``None`` and inherits the launcher's + ``CUDA_VISIBLE_DEVICES``. + When ``num_replicas`` is 1, returns ``[devices_str]`` unchanged. Otherwise, two YAML shapes are accepted: @@ -597,8 +602,14 @@ def split_devices_for_replicas( Any other length raises ``ValueError`` (the two modes are length-disjoint for ``num_replicas > 1``). """ - if num_replicas <= 1 or devices_str is None: - return [devices_str] if devices_str is not None else [devices_str] + if devices_str is None: + # No explicit placement: hand back one empty slot per replica, matching + # ``get_headless_replica_devices``. Returning a single-element list here + # made callers index past the end for every replica after the first. + return [None] * max(1, num_replicas) + + if num_replicas <= 1: + return [devices_str] device_list = [d.strip() for d in devices_str.split(",") if d.strip()] @@ -688,7 +699,7 @@ def compute_replica_layout( raise ValueError(f"num_replicas must be >= 0, got {num_replicas}") replicas_per_stage.append(num_replicas if allow_zero else max(1, num_replicas)) - replica_devices_map: dict[int, list[str]] = {} + replica_devices_map: dict[int, list[str | None]] = {} for stage_id, stage_cfg in enumerate(stage_configs): num_replicas = replicas_per_stage[stage_id] if num_replicas <= 1: diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index c0aafc86a..dea5242d1 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -337,7 +337,7 @@ class StageRuntime: self, omni_transfer_config: Any, replicas_per_stage: Sequence[int], - replica_devices_map: Mapping[int, Sequence[str]], + replica_devices_map: Mapping[int, Sequence[str | None]], ) -> list[LogicalStageInitPlan]: """Build startup plans for every logical stage and replica.""" stage_plans: list[LogicalStageInitPlan] = []