From 2d980133fee0c9d6c2ec09e839366368f8f555e3 Mon Sep 17 00:00:00 2001 From: Haco <75477391+xiaohajiayou@users.noreply.github.com> Date: Thu, 9 Apr 2026 16:44:31 +0800 Subject: [PATCH] [Bugfix] Fix precedence between caller runtime args and default stage configs (#2076) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: xiaohajiayou <923390377@qq.com> Co-authored-by: 汪志鹏 Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com> --- tests/entrypoints/test_stage_utils.py | 7 ++--- vllm_omni/engine/stage_init_utils.py | 23 +++++++++++++++- vllm_omni/entrypoints/stage_utils.py | 38 ++++++++++++++++++++------- vllm_omni/entrypoints/utils.py | 21 ++++++++++++--- 4 files changed, 73 insertions(+), 16 deletions(-) diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py index 3afc6f12f..15ee9c32a 100644 --- a/tests/entrypoints/test_stage_utils.py +++ b/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") diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index f71afad83..e6f603d2a 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/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) diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 8674d3c33..7b725f469 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/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: diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 0e1000ec9..c5e49a933 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/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