From baba7d1e2ef9d326486493146b43157f7cd4f8d0 Mon Sep 17 00:00:00 2001 From: Yukim1 <121286183+zwhzzz0821@users.noreply.github.com> Date: Sun, 16 Aug 2026 11:22:28 +0800 Subject: [PATCH] [Refactor] Remove legacy stage_args YAML loader (#6200) Signed-off-by: zwhzzz0821 <2831474076@qq.com> Co-authored-by: Hongsheng Liu --- .../offline_inference/aura_omni/end2end.py | 2 - examples/offline_inference/lance/end2end.py | 3 - .../nemotron_voicechat/end2end.py | 6 +- .../text_to_speech/fish_speech/end2end.py | 14 +--- .../qwen3_tts/word_timestamps.py | 8 +- tests/config/test_config_factory.py | 10 ++- .../test_hunyuanimage3_i2t_expansion.py | 9 +-- .../test_async_omni_diffusion_config.py | 18 +++++ tests/entrypoints/test_utils.py | 73 +++-------------- vllm_omni/config/stage_config.py | 5 ++ vllm_omni/engine/async_omni_engine.py | 10 +-- vllm_omni/entrypoints/utils.py | 79 ++++--------------- 12 files changed, 73 insertions(+), 164 deletions(-) diff --git a/examples/offline_inference/aura_omni/end2end.py b/examples/offline_inference/aura_omni/end2end.py index bac385d8e..29ccb3539 100644 --- a/examples/offline_inference/aura_omni/end2end.py +++ b/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() diff --git a/examples/offline_inference/lance/end2end.py b/examples/offline_inference/lance/end2end.py index d26910161..2b70bd749 100644 --- a/examples/offline_inference/lance/end2end.py +++ b/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() diff --git a/examples/offline_inference/nemotron_voicechat/end2end.py b/examples/offline_inference/nemotron_voicechat/end2end.py index b020e5bad..ac5fac376 100644 --- a/examples/offline_inference/nemotron_voicechat/end2end.py +++ b/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) diff --git a/examples/offline_inference/text_to_speech/fish_speech/end2end.py b/examples/offline_inference/text_to_speech/fish_speech/end2end.py index 8ce4664b9..e2e5738fd 100644 --- a/examples/offline_inference/text_to_speech/fish_speech/end2end.py +++ b/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 diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py b/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py index 62f5e2caa..7b282c560 100644 --- a/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py +++ b/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() diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 5bfbb67e8..41fc15df5 100644 --- a/tests/config/test_config_factory.py +++ b/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): diff --git a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py b/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py index 84a9291b0..c97de0069 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py +++ b/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 diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 75a9de2a7..1daca1e99 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/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 = { diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 803097dcf..66c7d5c30 100644 --- a/tests/entrypoints/test_utils.py +++ b/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: diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 86a6fed1b..cdfd35ee2 100644 --- a/vllm_omni/config/stage_config.py +++ b/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", [])] diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index c457b7052..38eea768e 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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) diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 9e884db48..e20dd7ec2 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/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(