[Bugfix] Fix precedence between caller runtime args and default stage configs (#2076)

Signed-off-by: xiaohajiayou <923390377@qq.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>

tests/entrypoints/test_stage_utils.py

@@ -92,9 +92,10 @@ def test_set_stage_devices_handles_not_enough_devices(mocker: MockerFixture, mon
         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")
+    # Keep the logical mapping and resolve to the visible subset.
+    set_stage_devices(stage_id=0, devices="0,1,2,3")
+
+    assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
 
 
 @pytest.mark.usefixtures("clean_gpu_memory_between_tests")

vllm_omni/engine/stage_init_utils.py

@@ -336,7 +336,13 @@ def acquire_device_locks(
             num_devices = current_omni_platform.get_device_count()
             physical_devices = list(range(num_devices))
 
-        num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
+        if len(physical_devices) < num_devices_per_stage:
+            raise RuntimeError(
+                f"Stage {stage_id} requires {num_devices_per_stage} device(s) based on parallel_config, "
+                f"but only {len(physical_devices)} device(s) are available: {physical_devices}"
+            )
+
+        num_devices_to_lock = num_devices_per_stage
         devices_to_lock = sorted(physical_devices[:num_devices_to_lock])
 
         logger.debug(
@@ -462,6 +468,21 @@ def initialize_diffusion_stage(
         model=model,
         **_to_dict(stage_cfg.engine_args),
     )
+    num_devices_per_stage = od_config.parallel_config.world_size
+    device_control_env = current_omni_platform.device_control_env_var
+    visible_devices_str = os.environ.get(device_control_env)
+    if visible_devices_str:
+        physical_devices = [device.strip() for device in visible_devices_str.split(",") if device.strip()]
+    else:
+        physical_devices = list(range(current_omni_platform.get_device_count()))
+
+    if len(physical_devices) < num_devices_per_stage:
+        raise ValueError(
+            f"Stage {metadata.stage_id} requires {num_devices_per_stage} device(s) based on parallel_config, "
+            f"but {len(physical_devices)} device(s) are available: {physical_devices}"
+        )
+
+    od_config.num_gpus = num_devices_per_stage
     if metadata.cfg_kv_collect_func is not None:
         od_config.cfg_kv_collect_func = metadata.cfg_kv_collect_func
     return StageDiffusionClient(model, od_config, metadata, batch_size=batch_size)

vllm_omni/entrypoints/stage_utils.py

@@ -78,7 +78,7 @@ def _parse_device_list(devices: str | int) -> list[str]:
 
 
 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.
+    """Map logical stage devices onto the currently available device pool.
 
     Args:
         stage_id: The stage ID currently configuring devices.
@@ -87,22 +87,42 @@ def _map_device_list(stage_id: int, device_list: list[str], visible_device_list:
         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 some requested ids exceed the available pool, we will fall back to the
+    # subset that can be mapped and leave the final capacity check to the later
+    # parallel-config validation path.
     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:
+    mapped_devices = [visible_device_list[idx] for idx in logical_ids if idx < num_visible]
+    mapping_pairs = [
+        f"{logical_id}->{visible_device_list[logical_id]}" for logical_id in logical_ids if logical_id < num_visible
+    ]
+    if not mapped_devices:
         raise ValueError(
-            f"Stage {stage_id} has logical IDs {device_list}, one or more of which exceed the number of visible devices"
+            f"Stage {stage_id} has logical IDs {device_list}, none of which map to the visible devices "
+            f"{visible_device_list}"
         )
-    return [visible_device_list[idx] for idx in logical_ids]
+    if len(mapped_devices) < len(logical_ids):
+        logger.warning(
+            "Stage %s requested logical devices %s, but only %d device(s) are currently available: %s. "
+            "Resolved logical-to-physical mapping: %s. Falling back to mapped subset %s",
+            stage_id,
+            device_list,
+            num_visible,
+            visible_device_list,
+            ", ".join(mapping_pairs) if mapping_pairs else "(none)",
+            mapped_devices,
+        )
+    else:
+        logger.info(
+            "Stage %s logical-to-physical device mapping: %s",
+            stage_id,
+            ", ".join(mapping_pairs),
+        )
+    return mapped_devices
 
 
 def serialize_obj(obj: Any) -> bytes:

vllm_omni/entrypoints/utils.py

@@ -299,11 +299,19 @@ def load_stage_configs_from_model(model: str, base_engine_args: dict | None = No
     stage_config_path = resolve_model_config_path(model)
     if stage_config_path is None:
         return []
-    stage_configs = load_stage_configs_from_yaml(config_path=stage_config_path, base_engine_args=base_engine_args)
+    stage_configs = load_stage_configs_from_yaml(
+        config_path=stage_config_path,
+        base_engine_args=base_engine_args,
+        prefer_stage_engine_args=False,
+    )
     return stage_configs
 
 
-def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None = None) -> list:
+def load_stage_configs_from_yaml(
+    config_path: str,
+    base_engine_args: dict | None = None,
+    prefer_stage_engine_args: bool = True,
+) -> list:
     """Load stage configurations from a YAML file.
 
     .. deprecated::
@@ -311,6 +319,9 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None
 
     Args:
         config_path: Path to the YAML configuration file
+        base_engine_args: Engine args supplied by the caller.
+        prefer_stage_engine_args: When True, YAML stage args override caller
+            engine args. When False, caller engine args override YAML defaults.
 
     Returns:
         List of stage configuration dictionaries from the file's stage_args
@@ -327,7 +338,11 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None
         base_engine_args_tmp = base_engine_args.copy()
         # Update base_engine_args with stage-specific engine_args if they exist
         if hasattr(stage_arg, "engine_args") and stage_arg.engine_args is not None:
-            base_engine_args_tmp = create_config(merge_configs(base_engine_args_tmp, stage_arg.engine_args))
+            if prefer_stage_engine_args:
+                merged_engine_args = merge_configs(base_engine_args_tmp, stage_arg.engine_args)
+            else:
+                merged_engine_args = merge_configs(stage_arg.engine_args, base_engine_args_tmp)
+            base_engine_args_tmp = create_config(merged_engine_args)
         stage_type = getattr(stage_arg, "stage_type", "llm")
         if hasattr(stage_arg, "runtime") and stage_arg.runtime is not None and stage_type != "diffusion":
             base_engine_args_tmp.async_chunk = global_async_chunk