[Refactor] Remove legacy stage_args YAML loader (#6200)

Signed-off-by: zwhzzz0821 <2831474076@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

examples/offline_inference/aura_omni/end2end.py

@@ -24,7 +24,6 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.media.audio import load_audio
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.stage_input_processors.aura_omni import (
     DEFAULT_QWEN3_TTS_REF_TEXT,
@@ -228,7 +227,6 @@ def parse_args():
     parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.")
     parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.")
     parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.")
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/lance/end2end.py

@@ -82,9 +82,6 @@ def parse_args():
     # ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
     # and ``create_default_diffusion`` materializes the stage config.
 
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(p)
     return p.parse_args()
 
 

examples/offline_inference/nemotron_voicechat/end2end.py

@@ -67,7 +67,7 @@ def main() -> None:
     parser.add_argument("--output-dir", default="results/nemotron_voicechat")
     parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
     parser.add_argument(
-        "--stage-configs-path",
+        "--deploy-config",
         default=None,
         help=(
             "Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml "
@@ -134,8 +134,8 @@ def main() -> None:
     ]
 
     omni_kwargs: dict = {}
-    if args.stage_configs_path:
-        omni_kwargs["stage_configs_path"] = args.stage_configs_path
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
     omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
     outputs = omni.generate(inputs, sampling_params_list)
 

examples/offline_inference/text_to_speech/fish_speech/end2end.py

@@ -38,17 +38,6 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 logger = logging.getLogger(__name__)
 
 DEFAULT_MODEL = "fishaudio/s2-pro"
-DEFAULT_DEPLOY_CONFIG = os.path.join(
-    os.path.dirname(__file__),
-    "..",
-    "..",
-    "..",
-    "..",
-    "vllm_omni",
-    "model_executor",
-    "stage_configs",
-    "fish_speech_s2_pro.yaml",
-)
 
 
 def build_prompt(
@@ -152,7 +141,8 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )
-    omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
     return omni_kwargs
 
 

examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py

@@ -32,11 +32,10 @@ from end2end import _estimate_prompt_len
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni import Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
 
 
-def _default_stage_config() -> str:
+def _default_deploy_config() -> str:
     repo_root = Path(__file__).resolve().parents[4]
     return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
 
@@ -109,7 +108,7 @@ def main(args: Any) -> None:
         "output_dir",
     ):
         omni_kwargs.pop(key, None)
-    omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
+    omni_kwargs["deploy_config"] = args.deploy_config or _default_deploy_config()
     omni_kwargs["log_stats"] = args.log_stats
     omni = Omni(**omni_kwargs)
 
@@ -156,7 +155,7 @@ def parse_args() -> Any:
         default=None,
         help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
     )
-    parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
+    parser.add_argument("--deploy-config", default=None, help="Qwen3-TTS deploy YAML")
     parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
     parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
     parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
@@ -164,7 +163,6 @@ def parse_args() -> Any:
     parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
     parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
     parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

tests/config/test_config_factory.py

@@ -1102,6 +1102,13 @@ class TestResolveScheduler:
 
 
 class TestDeployConfigLoading:
+    def test_rejects_legacy_stage_args_schema(self, tmp_path):
+        deploy_path = tmp_path / "legacy.yaml"
+        deploy_path.write_text("stage_args:\n  - stage_id: 0\n", encoding="utf-8")
+
+        with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"):
+            load_deploy_config(deploy_path)
+
     def test_load_minicpmo_duplex_deploy_config(self):
         deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml"))
 
@@ -1377,7 +1384,8 @@ stages:
         assert async_stages[1].custom_process_input_func is None
 
     def test_no_bundled_legacy_stage_config_yamls(self):
-        stage_config_dir = Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "stage_configs"
+        repo_root = Path(__file__).resolve().parents[2]
+        stage_config_dir = repo_root / "vllm_omni" / "model_executor" / "stage_configs"
         assert not list(stage_config_dir.glob("*.yaml"))
 
     def test_merge_pipeline_deploy(self):

tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py

@@ -3,19 +3,18 @@
 """Smoke test for HunyuanImage-3.0 Image-to-Text (I2T) pipeline."""
 
 from collections.abc import Generator
-from pathlib import Path
 
 import pytest
 import torch
 from PIL import Image
 
 from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni import Omni
 from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt
 
 MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
-REPO_ROOT = Path(__file__).resolve().parents[3]
-STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image3_i2t.yaml"
+DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image3_ar.yaml")
 
 # First 20 generated token IDs from the HF greedy reference on this input.
 # vllm-omni AR output matches this prefix bitwise; the two implementations
@@ -45,14 +44,14 @@ EXPECTED_PREFIX_TOKEN_IDS: list[int] = [
 # Decoded form, kept only for human-readable assertion messages.
 EXPECTED_PREFIX_TEXT = "The image is a solid, uniform green color with no variations, objects, or details present. It"
 
-pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion, pytest.mark.cuda]
 
 
 @pytest.fixture(scope="module")
 def omni() -> Generator[Omni, None, None]:
     with OmniRunner(
         MODEL_NAME,
-        deploy_config=str(STAGE_CONFIG_PATH),
+        deploy_config=DEPLOY_CONFIG_PATH,
         trust_remote_code=True,
     ) as runner:
         yield runner.omni

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -492,6 +492,24 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
     assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}}
 
 
+@pytest.mark.parametrize(
+    ("legacy_arg", "value"),
+    [
+        ("stage_configs_path", "legacy.yaml"),
+        ("stage_configs", [{"stage_id": 0}]),
+    ],
+)
+def test_resolve_stage_configs_rejects_legacy_config_arguments(legacy_arg, value):
+    engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
+
+    with pytest.raises(ValueError, match=rf"`{legacy_arg}`.*`deploy_config`"):
+        engine._resolve_stage_configs(
+            "dummy-model",
+            {legacy_arg: value},
+            trust_remote_code=False,
+        )
+
+
 def test_default_stage_config_includes_quantization_config():
     """Ensure structured quantization_config survives default diffusion-stage creation."""
     quantization_config = {

tests/entrypoints/test_utils.py

@@ -27,7 +27,7 @@ from vllm_omni.entrypoints.utils import (
     filter_dataclass_kwargs,
     filter_stages,
     load_and_resolve_stage_configs,
-    load_stage_configs_from_yaml,
+    load_stage_configs_from_model,
     resolve_model_config_path,
 )
 
@@ -520,71 +520,20 @@ class TestLoadAndResolveStageConfigs:
         assert filtered[0].final_output_type is None
 
 
-class TestLoadStageConfigsFromYaml:
-    """Regression tests for stage-config loading and merging."""
-
-    def test_deep_merges_stage_engine_args(self, mocker: MockerFixture):
-        yaml_config = create_config(
-            {
-                "async_chunk": True,
-                "stage_args": [
-                    {
-                        "stage_id": 0,
-                        "runtime": {"device": 0},
-                        "engine_args": {
-                            "parallel_config": {"tensor_parallel_size": 4},
-                        },
-                    }
-                ],
-            }
-        )
+class TestLoadStageConfigsFromModel:
+    def test_unresolved_model_does_not_fall_back_to_yaml(self, mocker: MockerFixture, caplog):
         mocker.patch(
-            "vllm_omni.entrypoints.utils.load_yaml_config",
-            return_value=yaml_config,
+            "vllm_omni.entrypoints.utils.StageConfigFactory.create_legacy_stage_configs_from_model",
+            return_value=(None, None),
         )
+        resolve_path = mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path")
 
-        stages = load_stage_configs_from_yaml(
-            "fake.yaml",
-            base_engine_args={
-                "parallel_config": {
-                    "tensor_parallel_size": 1,
-                    "pipeline_parallel_size": 2,
-                },
-                "model": "base-model",
-            },
-        )
+        result = load_stage_configs_from_model("unregistered-model", trust_remote_code=False)
 
-        merged_engine_args = stages[0]["engine_args"]
-        assert merged_engine_args["parallel_config"]["tensor_parallel_size"] == 4
-        assert merged_engine_args["parallel_config"]["pipeline_parallel_size"] == 2
-        assert merged_engine_args["model"] == "base-model"
-        assert merged_engine_args["async_chunk"] is True
-
-    def test_merges_nested_stage_engine_args(self, mocker: MockerFixture):
-        yaml_config = create_config(
-            {
-                "stage_args": [
-                    {
-                        "stage_id": 0,
-                        "engine_args": {
-                            "nested": {"override": 2},
-                        },
-                    }
-                ],
-            }
-        )
-        mocker.patch(
-            "vllm_omni.entrypoints.utils.load_yaml_config",
-            return_value=yaml_config,
-        )
-
-        stages = load_stage_configs_from_yaml(
-            "fake.yaml",
-            base_engine_args={"nested": {"base": 1}},
-        )
-
-        assert stages[0]["engine_args"]["nested"]["base"] == 1
-        assert stages[0]["engine_args"]["nested"]["override"] == 2
+        assert result == ([], None)
+        resolve_path.assert_not_called()
+        assert "No registered PipelineConfig resolved" in caplog.text
+        assert "deploy_config" in caplog.text
 
 
 class TestCumulativeStreamingCoercion:

vllm_omni/config/stage_config.py

@@ -632,6 +632,11 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]:
 def load_deploy_config(path: str | Path) -> DeployConfig:
     """Load a deploy YAML (with optional base_config inheritance)."""
     raw_dict = resolve_deploy_yaml(path)
+    if "stage_args" in raw_dict:
+        raise ValueError(
+            f"Deploy config {path} uses the removed `stage_args` schema; "
+            "define topology in PipelineConfig and deployment overrides under `stages`."
+        )
 
     stages = [_parse_stage_deploy(s) for s in raw_dict.get("stages", [])]
 

vllm_omni/engine/async_omni_engine.py

@@ -1118,15 +1118,13 @@ class AsyncOmniEngine:
     ) -> tuple[str, list[Any]]:
         """Resolve stage configs and inject defaults shared by orchestrator/headless."""
 
+        for legacy_arg in ("stage_configs_path", "stage_configs"):
+            if legacy_arg in kwargs:
+                raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.")
+
         deploy_config_path = kwargs.pop("deploy_config", None)
         strategy_config_path = kwargs.pop("strategy_config", None)
         stage_overrides_json = kwargs.pop("stage_overrides", None)
-        explicit_stage_configs = kwargs.pop("stage_configs", None)
-        if explicit_stage_configs is not None:
-            logger.warning(
-                "`stage_configs` is not part of the public API. "
-                "Ignoring it and resolving stages from deploy_config/model factory."
-            )
 
         # Parse --stage-overrides JSON string if provided
         stage_overrides = parse_stage_overrides(stage_overrides_json)

vllm_omni/entrypoints/utils.py

@@ -19,7 +19,7 @@ from vllm_omni.config.config_factory import (
 )
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 from vllm_omni.config.stage_config import _DEPLOY_DIR
-from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs
+from vllm_omni.config.yaml_util import create_config, load_yaml_config
 from vllm_omni.diffusion.utils.hf_utils import (
     _looks_like_dreamzero,
     get_diffusion_model_index,
@@ -369,11 +369,12 @@ def load_stage_configs_from_model(
 ) -> tuple[list, str | None]:
     """Load stage configurations from model's default config file.
 
-    For models registered in the pipeline registry (new path), uses
+    For models registered in the pipeline registry, uses
     ``StageConfigFactory.create_legacy_stage_configs_from_model()`` which merges
     PipelineConfig + DeployConfig + CLI overrides.
 
-    For other models (legacy path), loads stage configs from YAML.
+    Models that cannot be resolved through the registry return no stages so the
+    caller can apply its default stage configuration, when available.
 
     Args:
         model: Model name or path (used to determine model_type)
@@ -383,7 +384,7 @@ def load_stage_configs_from_model(
         stage_overrides: Per-stage overrides from --stage-overrides.
         strategy_config_path: Optional path to a composable-parallel
             ``strategy.yaml`` whose derived sizing is overlaid onto the
-            registry-merged stages (opt-in; ignored on the legacy YAML path).
+            registry-merged stages.
 
     Returns:
         ``(stage_configs, omni_lb_policy)``: the list of stage configuration
@@ -427,69 +428,17 @@ def load_stage_configs_from_model(
         # Convert StageConfig objects to OmegaConf for backward compat
         return [stage.to_omegaconf() for stage in stages], omni_lb_policy
 
-    # Legacy fallback: load from YAML. A composable-parallel strategy cannot be
-    # applied here (it overlays onto registry-merged stages), so warn rather than
-    # silently dropping the operator's --strategy-config.
+    strategy_note = ""
     if strategy_config_path is not None:
-        logger.warning(
-            "--strategy-config (%s) was provided but model %r resolves via the "
-            "legacy stage_configs YAML path, which does not support "
-            "composable-parallel strategies; the strategy is ignored. Use a "
-            "registry-based model to apply it.",
-            strategy_config_path,
-            model,
-        )
-    stage_config_path = resolve_model_config_path(model)
-    if stage_config_path is None:
-        return [], None
-    stage_configs = load_stage_configs_from_yaml(
-        config_path=stage_config_path,
-        base_engine_args=base_engine_args,
-        prefer_stage_engine_args=True,
+        strategy_note = f" Strategy config {strategy_config_path!r} was not applied."
+    logger.warning(
+        "No registered PipelineConfig resolved for model %r. Legacy `stage_args` "
+        "YAMLs are no longer supported; register the pipeline and provide deployment "
+        "overrides through `deploy_config`.%s",
+        model,
+        strategy_note,
     )
-    return stage_configs, None
-
-
-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 (legacy OmegaConf path).
-
-    TODO(@lishunyang12): remove once all models use PipelineConfig + DeployConfig.
-
-    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
-    """
-    if base_engine_args is None:
-        base_engine_args = {}
-    config_data = load_yaml_config(config_path)
-    stage_args = config_data.stage_args
-    global_async_chunk = config_data.get("async_chunk", False)
-    # Convert any nested dataclass objects to dicts before creating DictConfig
-    base_engine_args = _convert_dataclasses_to_dict(base_engine_args)
-    base_engine_args = create_config(base_engine_args)
-    for stage_arg in stage_args:
-        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:
-            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
-        stage_arg.engine_args = base_engine_args_tmp
-    return stage_args
+    return [], None
 
 
 def filter_stages(