From c09b039f3776a7bb73e81349ec52fd694dc3e6c6 Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Wed, 25 Mar 2026 03:56:13 -0600 Subject: [PATCH] [BugFix] Make Stage Device Initialization Respect Visible GPUs (#2025) Signed-off-by: Alex Brooks Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com> --- tests/entrypoints/test_stage_utils.py | 20 +++ vllm_omni/engine/async_omni_engine.py | 47 ++++--- vllm_omni/engine/stage_init_utils.py | 27 ++-- vllm_omni/entrypoints/stage_utils.py | 170 ++++++++++++-------------- 4 files changed, 135 insertions(+), 129 deletions(-) diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py index 5eb6154f0..2bb2231cc 100644 --- a/tests/entrypoints/test_stage_utils.py +++ b/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.""" diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 47b97295c..dc62bda5b 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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( diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index cd1331dfb..9c246ce6e 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/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"] diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 1f45ff904..8674d3c33 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/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: