From 51c14eeaede4fb428f456dc785467ce995356518 Mon Sep 17 00:00:00 2001 From: CAICAII <3360776475@qq.com> Date: Sat, 4 Jul 2026 12:16:21 +0800 Subject: [PATCH] [BugFix] Preserve explicit deploy stage engine_extras (#4255) Signed-off-by: CAICAIIs <3360776475@qq.com> --- tests/helpers/stage_config.py | 13 +++++++ tests/test_config_factory.py | 59 +++++++++++++++++++++++++++++++- vllm_omni/config/stage_config.py | 3 +- 3 files changed, 73 insertions(+), 2 deletions(-) diff --git a/tests/helpers/stage_config.py b/tests/helpers/stage_config.py index 4351eb27d..8706bb1bc 100644 --- a/tests/helpers/stage_config.py +++ b/tests/helpers/stage_config.py @@ -633,6 +633,18 @@ def get_deploy_config_path(rel_path: str) -> str: return str(_DEPLOY_DIR / rel_path) +def get_deploy_config_stage(rel_path: str, stage_id: int) -> dict[str, Any]: + """Return one stage entry from a deploy yaml by ``stage_id``.""" + with open(get_deploy_config_path(rel_path), encoding="utf-8") as f: + cfg = yaml.safe_load(f) or {} + + stage_key = "stages" if "stages" in cfg else "stage_args" + for stage in cfg.get(stage_key, []): + if stage.get("stage_id") == stage_id: + return stage + raise KeyError(f"No stage_id={stage_id} in deploy config {rel_path!r}") + + def _get_config_value_by_path(config_dict: dict, path: str) -> Any: """Read a dot-separated path from a nested dict (e.g. ``engine_args.load_format``).""" current: Any = config_dict @@ -707,6 +719,7 @@ def stage_config_path_for_run_level(stage_config_path: str | None, run_level: st __all__ = [ "get_deploy_config_path", + "get_deploy_config_stage", "modify_stage_config", "stage_config_path_for_run_level", ] diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 23d36c338..4b956dfcd 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -14,7 +14,7 @@ from unittest.mock import patch import pytest from transformers import PretrainedConfig, Qwen3OmniMoeConfig -from tests.helpers.stage_config import get_deploy_config_path +from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline from vllm_omni.config.stage_config import ( @@ -739,6 +739,15 @@ stages: assert deploy.stages[0].compilation_config == {"pass_config": {"fuse_allreduce_rms": False}} assert "compilation_config" not in deploy.stages[0].engine_extras + def test_load_voxcpm2_deploy_config_preserves_engine_extras(self): + deploy_path = get_deploy_config_path("voxcpm2.yaml") + raw_stage = get_deploy_config_stage("voxcpm2.yaml", 0) + expected_runtime_config = raw_stage["engine_extras"]["hf_overrides"]["voxcpm2_runtime_config"] + + deploy = load_deploy_config(deploy_path) + runtime_config = deploy.stages[0].engine_extras["hf_overrides"]["voxcpm2_runtime_config"] + assert runtime_config == expected_runtime_config + def test_merge_pipeline_deploy(self): deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): @@ -928,6 +937,54 @@ stages: assert "foo" in stage.engine_extras assert stage.engine_extras["foo"] == {"a": 111, "b": {"d": 9, "e": 199}, "c": 3} + def test_explicit_engine_extras_merge_with_flat_and_engine_args(self): + """Explicit engine_extras should be preserved with existing pass-through styles.""" + fake_config = { + "stages": [ + { + "stage_id": 0, + "engine_extras": { + "hf_overrides": { + "runtime_config": { + "explicit_only": True, + "shared": "explicit", + "nested": {"a": 1}, + } + } + }, + "hf_overrides": { + "runtime_config": { + "flat_only": True, + "shared": "flat", + "nested": {"b": 2}, + } + }, + "engine_args": { + "hf_overrides": { + "runtime_config": { + "engine_args_only": True, + "shared": "engine_args", + "nested": {"c": 3}, + } + } + }, + }, + ] + } + + with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): + deploy = load_deploy_config("dummy.yaml") + + assert deploy.stages[0].engine_extras["hf_overrides"] == { + "runtime_config": { + "engine_args_only": True, + "explicit_only": True, + "flat_only": True, + "nested": {"a": 1, "b": 2, "c": 3}, + "shared": "engine_args", + } + } + def test_deep_merge_does_not_mutate_inputs(self): """Merging engine_args must not mutate the base stage dict.""" fake_config = { diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index f66023bc4..858e9fa15 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -468,6 +468,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: """Parse a single stage entry from deploy YAML into StageDeployConfig.""" # Get the non-reserved keys for this stage flat_args = {k: v for k, v in stage_data.items() if k not in _STAGE_RESERVED_KEYS} + explicit_engine_extras = dict(stage_data.get("engine_extras") or {}) runtime_cfg = dict(stage_data.get("runtime", {})) devices = runtime_cfg.get("devices", stage_data.get("devices")) num_replicas = runtime_cfg.get("num_replicas", stage_data.get("num_replicas", 1)) @@ -495,7 +496,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: kwargs["output_connectors"] = stage_data.get("output_connectors") kwargs["input_connectors"] = stage_data.get("input_connectors") kwargs["default_sampling_params"] = stage_data.get("default_sampling_params") - kwargs["engine_extras"] = flat_args + kwargs["engine_extras"] = _get_recursively_merged_dict(explicit_engine_extras, flat_args) return StageDeployConfig(**kwargs)