[BugFix] Make Stage Device Initialization Respect Visible GPUs (#2025)

Signed-off-by: Alex Brooks <albrooks@redhat.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>

tests/entrypoints/test_stage_utils.py

@@ -75,6 +75,26 @@ def test_set_stage_devices_respects_logical_ids(mocker: MockerFixture, monkeypat
     assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
 
 
+@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
+def test_set_stage_devices_handles_not_enough_devices(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    # Preserve an existing logical mapping and ensure devices "0,1" map through it.
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "6,7")
+    call_log: list[int] = []
+    dummy_torch = _make_dummy_torch(call_log)
+    monkeypatch.setitem(sys.modules, "torch", dummy_torch)
+
+    # Mock the platform at the source module where it's defined
+    mock_platform = _make_mock_platform(mocker, device_type="cuda", env_var="CUDA_VISIBLE_DEVICES")
+    monkeypatch.setattr(
+        "vllm_omni.platforms.current_omni_platform",
+        mock_platform,
+    )
+
+    # Raise since we need 4 GPUs, but we only have 2 visible
+    with pytest.raises(ValueError):
+        set_stage_devices(stage_id=0, devices="0,1,2,3")
+
+
 @pytest.mark.usefixtures("clean_gpu_memory_between_tests")
 def test_set_stage_devices_npu_platform(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
     """Test that set_stage_devices works correctly for NPU platform."""

vllm_omni/engine/async_omni_engine.py

@@ -64,6 +64,7 @@ from vllm_omni.engine.stage_init_utils import (
 from vllm_omni.entrypoints.utils import (
     load_and_resolve_stage_configs,
 )
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -297,8 +298,6 @@ class AsyncOmniEngine:
         omni_kv_connector: tuple[dict[str, Any] | None, str | None, str | None] = (None, None, None),
     ) -> StartedLlmStage:
         """Launch one LLM stage to READY state in a helper thread."""
-        from vllm_omni.platforms import current_omni_platform
-
         started_stage: StartedLlmStage | None = None
         lock_fds: list[int] = []
         device_control_env = current_omni_platform.device_control_env_var
@@ -429,6 +428,7 @@ class AsyncOmniEngine:
 
     def _initialize_stages(self, stage_init_timeout: int) -> None:
         """Initialize stage clients/processors in orchestrator thread and assign to self."""
+        device_control_env = current_omni_platform.device_control_env_var
 
         num_stages = self.num_stages
         stage_clients: list[Any | None] = [None] * num_stages
@@ -469,25 +469,32 @@ class AsyncOmniEngine:
                     omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id)
 
                     if metadata.stage_type == "diffusion":
-                        setup_stage_devices(stage_id, metadata.runtime_cfg)
-                        omni_conn_cfg, omni_from, omni_to = omni_kv_connector
-                        if omni_conn_cfg:
-                            from vllm_omni.entrypoints.utils import inject_omni_kv_config
+                        previous_visible_devices = os.environ.get(device_control_env)
+                        try:
+                            setup_stage_devices(stage_id, metadata.runtime_cfg)
+                            omni_conn_cfg, omni_from, omni_to = omni_kv_connector
+                            if omni_conn_cfg:
+                                from vllm_omni.entrypoints.utils import inject_omni_kv_config
 
-                            inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
-                        _inject_kv_stage_info(stage_cfg, stage_id)
-                        stage_clients[stage_id] = initialize_diffusion_stage(
-                            self.model,
-                            stage_cfg,
-                            metadata,
-                            batch_size=self.diffusion_batch_size,
-                        )
-                        logger.info(
-                            "[AsyncOmniEngine] Stage %s initialized (diffusion, batch_size=%d)",
-                            stage_id,
-                            self.diffusion_batch_size,
-                        )
-                        continue
+                                inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
+                            _inject_kv_stage_info(stage_cfg, stage_id)
+                            stage_clients[stage_id] = initialize_diffusion_stage(
+                                self.model,
+                                stage_cfg,
+                                metadata,
+                                batch_size=self.diffusion_batch_size,
+                            )
+                            logger.info(
+                                "[AsyncOmniEngine] Stage %s initialized (diffusion, batch_size=%d)",
+                                stage_id,
+                                self.diffusion_batch_size,
+                            )
+                            continue
+                        finally:
+                            if previous_visible_devices is None:
+                                current_omni_platform.unset_device_control_env_var()
+                            else:
+                                current_omni_platform.set_device_control_env_var(previous_visible_devices)
 
                     llm_stage_ids.append(stage_id)
                     llm_launch_futures[stage_id] = launch_executor.submit(

vllm_omni/engine/stage_init_utils.py

@@ -27,6 +27,7 @@ from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices
 from vllm_omni.entrypoints.utils import filter_dataclass_kwargs, resolve_model_config_path
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -83,8 +84,6 @@ def resolve_worker_cls(engine_args: dict[str, Any]) -> None:
     if worker_cls is not None and worker_cls != "auto":
         return
 
-    from vllm_omni.platforms import current_omni_platform
-
     worker_type = str(worker_type).lower()
     if worker_type == "ar":
         engine_args["worker_cls"] = current_omni_platform.get_omni_ar_worker_cls()
@@ -214,23 +213,17 @@ def prepare_engine_environment() -> None:
 
 def setup_stage_devices(stage_id: int, runtime_cfg: Any) -> None:
     """Device mapping via set_stage_devices for a single stage."""
-    try:
-        from vllm_omni.platforms import current_omni_platform
-
-        device_type = current_omni_platform.device_type
-        set_stage_devices(
-            stage_id,
-            runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else None,
-            device_type=device_type,
-        )
+    physical_devices = set_stage_devices(
+        stage_id,
+        runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else None,
+    )
+    # Only log if we actually set the env vars in the stage
+    if physical_devices:
         logger.info(
-            "[stage_init] Stage-%s set devices for %s, runtime devices: %s",
+            "[stage_init] Stage-%s set runtime devices: %s",
             stage_id,
-            device_type,
-            runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else None,
+            physical_devices,
         )
-    except Exception as e:
-        logger.warning("Device setup failed for stage %s: %s", stage_id, e)
 
 
 def build_engine_args_dict(
@@ -295,8 +288,6 @@ def acquire_device_locks(
     """
     lock_fds: list[int] = []
     try:
-        from vllm_omni.platforms import current_omni_platform
-
         # Get parallel sizes
         if "parallel_config" in engine_args_dict:
             pc = engine_args_dict["parallel_config"]

vllm_omni/entrypoints/stage_utils.py

@@ -6,6 +6,7 @@ from multiprocessing import shared_memory as _shm
 from typing import Any
 
 from vllm_omni.config.yaml_util import to_dict as _omega_to_dict
+from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
 
@@ -13,108 +14,95 @@ logger = logging.getLogger(__name__)
 def set_stage_devices(
     stage_id: int,
     devices: str | int | None,
-    device_type: str | None = None,
-) -> None:
+) -> str | None:
     """Configure per-stage device visibility and current device (CUDA or NPU).
 
     This function sets environment variables that control which devices are visible
-    to the process, and sets the current device. It must be called BEFORE worker
-    initialization so that workers see the correct devices.
+    to the process. It must be called BEFORE worker initialization so that workers
+    see the correct devices.
+
+
+    NOTE: This will set the control variable for the appropriate platform.
+        - CUDA: CUDA_VISIBLE_DEVICES
+        - NPU: ASCEND_RT_VISIBLE_DEVICES
 
     Args:
         stage_id: Stage identifier for logging
-        devices: Device specification:
-            - Comma-separated string (e.g. "2,5,7"): interpreted as logical
-              indices against the current device visibility env var (e.g.
-              CUDA_VISIBLE_DEVICES/ASCEND_RT_VISIBLE_DEVICES) when present;
-              falls back to physical IDs if no mapping exists. Logical index 0
-              is used as current device.
-            - Integer or digit-string: treat as logical index (0-based) into the
-              current device visibility mapping; map to physical device, then set
-              env var to this single device.
-            - None/"cpu": keep default visibility.
-            - Otherwise: set env var to the provided single device string.
-        device_type: Device type ("cuda" or "npu"). If None, auto-detects.
+        devices: Devices specified as either:
+            - None / "cpu"; uses the default visibility.
+            - An int or a str composed of one or more ints separated by commas,
+              which correspond to logical indices. If the control env var is
+              set, e.g., CUDA_VISIBLE_DEVICES, we will map the logical indices
+              to physical, e.g.,
+                    devices: [0,1,2,3]
+                    CUDA_VISIBLE_DEVICES -> [1, 3, 4, 5, 6]
+            will leverage [1, 3, 4, 5]
 
-    Behavior:
-        - CUDA: Sets CUDA_VISIBLE_DEVICES and calls torch.cuda.set_device()
-        - NPU: Sets ASCEND_RT_VISIBLE_DEVICES and calls torch.npu.set_device()
+    Returns:
+        The list of physical devices that were set for the given stage
+        or None if we have no passed devices / are using cpu.
     """
-    from vllm_omni.platforms import current_omni_platform
-
-    if device_type is None:
-        device_type = current_omni_platform.device_type
-
     env_var = current_omni_platform.device_control_env_var
+    vis = os.environ.get(env_var)
 
-    try:
-        selected_physical: int | None = None
-        logical_idx: int | None = None
+    if devices in (None, "cpu"):
+        logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices)
+        return None
 
-        if isinstance(devices, str) and "," in devices:
-            toks = [t.strip() for t in devices.split(",") if t.strip() != ""]
-            vis = os.environ.get(env_var)
-            mapped_devices: list[str] = []
-            mapping: list[int] = []
-            if vis:
-                try:
-                    mapping = [int(x) for x in vis.split(",") if x.strip() != ""]
-                except Exception as e:
-                    logger.debug("[Stage-%s] Failed to parse existing %s: %s", stage_id, env_var, e)
-            for tok in toks:
-                try:
-                    idx = int(tok)
-                except Exception:
-                    mapped_devices.append(tok)
-                    continue
-                if mapping and 0 <= idx < len(mapping):
-                    mapped_devices.append(str(mapping[idx]))
-                else:
-                    mapped_devices.append(str(idx))
-            mapped_devices_str = ",".join(mapped_devices)
-            current_omni_platform.set_device_control_env_var(mapped_devices_str)
-            if toks:
-                try:
-                    selected_physical = int(mapped_devices[0])
-                    logger.debug(
-                        "[Stage-%s] Set %s to %s; logical 0 -> physical %s",
-                        stage_id,
-                        env_var,
-                        mapped_devices_str,
-                        selected_physical,
-                    )
-                except Exception as e:
-                    logger.debug("[Stage-%s] Failed to parse first %s device: %s", stage_id, device_type, e)
-                    selected_physical = None
-        elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()):
-            logical_idx = max(0, int(devices))
-            vis = os.environ.get(env_var)
-            if vis:
-                try:
-                    mapping = [int(x) for x in vis.split(",") if x.strip() != ""]
-                    if 0 <= logical_idx < len(mapping):
-                        selected_physical = mapping[logical_idx]
-                except Exception as e:
-                    logger.debug("[Stage-%s] Failed to map logical index via %s: %s", stage_id, env_var, e)
-                    selected_physical = None
-            if selected_physical is None:
-                selected_physical = int(logical_idx)
-            current_omni_platform.set_device_control_env_var(str(selected_physical))
-            logger.debug(
-                "[Stage-%s] Logical index %d -> physical %s; set %s to single device",
-                stage_id,
-                logical_idx + 1,
-                selected_physical,
-                env_var,
-            )
-        elif devices in (None, "cpu"):
-            logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices)
-        else:
-            selected_physical = int(str(devices))
-            current_omni_platform.set_device_control_env_var(str(selected_physical))
-            logger.debug("[Stage-%s] Set %s to single device %s (fallback)", stage_id, env_var, selected_physical)
-    except Exception as e:
-        logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e)
+    elif isinstance(devices, (int, str)):
+        device_list = _parse_device_list(devices)
+        if vis is not None:
+            visible_device_list = _parse_device_list(vis)
+            device_list = _map_device_list(stage_id, device_list, visible_device_list)
+        device_str = ",".join(device_list)
+        current_omni_platform.set_device_control_env_var(device_str)
+        return device_str
+
+    raise TypeError(f"Expected str or int device IDs for stage initialization, got type {type(devices)}")
+
+
+def _parse_device_list(devices: str | int) -> list[str]:
+    """Given an int or a str representing one or more comma separated
+    non-negative IDs, coerce it to a list of strs.
+
+    Args:
+        devices: devices to be converted to a list of strs.
+    """
+    if isinstance(devices, int):
+        if devices < 0:
+            raise ValueError("Device IDs must be non-negative integers!")
+        return [str(devices)]
+    # Devices will usually be ints, but not always
+    # so we don't explicitly validate that here.
+    return [t.strip() for t in devices.split(",") if t.strip() != ""]
+
+
+def _map_device_list(stage_id: int, device_list: list[str], visible_device_list: list[str]) -> list[str]:
+    """Maps logical to physical devices if we have enough visible devices available.
+
+    Args:
+        stage_id: The stage ID currently configuring devices.
+        device_list: List of (logical) devices to be used, which are strings
+            holding non-negative nums counting from 0, 1, ..., n devices needed.
+        visible_device_list: List of physical devices available.
+    """
+    num_visible = len(visible_device_list)
+    num_logical = len(device_list)
+    if num_visible < num_logical:
+        raise ValueError(f"Stage {stage_id} requires {num_logical} devices, but only {num_visible} devices are visible")
+
+    # Ensure that the logical IDs are actually in range to avoid index errors;
+    # If the check above passes and those below fail, the logical devices are wrong,
+    # i.e., not actually 0, 1, ..., n
+    if not all(device.isdigit() for device in device_list):
+        raise ValueError("Logical devices must be non-negative integers")
+
+    logical_ids = [int(device) for device in device_list]
+    if max(logical_ids) >= num_visible:
+        raise ValueError(
+            f"Stage {stage_id} has logical IDs {device_list}, one or more of which exceed the number of visible devices"
+        )
+    return [visible_device_list[idx] for idx in logical_ids]
 
 
 def serialize_obj(obj: Any) -> bytes: