[Config Refactor][1/2] Model Pipeline Configuration System (#1115)

Signed-off-by: lishunyang <lishunyang12@163.com>

tests/test_config_factory.py

@@ -0,0 +1,570 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Unit tests for StageConfigFactory and related classes.
+"""
+
+from vllm_omni.config.stage_config import (
+    ModelPipeline,
+    StageConfig,
+    StageConfigFactory,
+    StageType,
+)
+
+
+class TestStageType:
+    """Tests for StageType enum."""
+
+    def test_stage_type_values(self):
+        """Test StageType enum values."""
+        assert StageType.LLM.value == "llm"
+        assert StageType.DIFFUSION.value == "diffusion"
+
+    def test_stage_type_from_string(self):
+        """Test creating StageType from string."""
+        assert StageType("llm") == StageType.LLM
+        assert StageType("diffusion") == StageType.DIFFUSION
+
+
+class TestStageConfig:
+    """Tests for StageConfig dataclass."""
+
+    def test_minimal_config(self):
+        """Test creating StageConfig with minimal required fields."""
+        config = StageConfig(stage_id=0, model_stage="thinker")
+        assert config.stage_id == 0
+        assert config.model_stage == "thinker"
+        assert config.stage_type == StageType.LLM
+        assert config.input_sources == []
+        assert config.final_output is False
+        assert config.worker_type is None
+
+    def test_full_config(self):
+        """Test creating StageConfig with all fields."""
+        config = StageConfig(
+            stage_id=1,
+            model_stage="talker",
+            stage_type=StageType.LLM,
+            input_sources=[0],
+            custom_process_input_func="module.path.func",
+            final_output=True,
+            final_output_type="audio",
+            worker_type="ar",
+            scheduler_cls="path.to.Scheduler",
+            hf_config_name="talker_config",
+            is_comprehension=False,
+        )
+        assert config.stage_id == 1
+        assert config.model_stage == "talker"
+        assert config.input_sources == [0]
+        assert config.final_output_type == "audio"
+        assert config.worker_type == "ar"
+
+    def test_to_omegaconf_basic(self):
+        """Test converting StageConfig to OmegaConf format."""
+        config = StageConfig(
+            stage_id=0,
+            model_stage="thinker",
+            stage_type=StageType.LLM,
+            worker_type="ar",
+            final_output=True,
+            final_output_type="text",
+        )
+        omega_config = config.to_omegaconf()
+
+        assert omega_config.stage_id == 0
+        assert omega_config.stage_type == "llm"
+        assert omega_config.engine_args.model_stage == "thinker"
+        assert omega_config.engine_args.worker_type == "ar"
+        assert omega_config.final_output is True
+        assert omega_config.final_output_type == "text"
+        # Legacy field name for backward compatibility
+        assert omega_config.engine_input_source == []
+
+    def test_to_omegaconf_with_runtime_overrides(self):
+        """Test that runtime overrides are applied to OmegaConf output."""
+        config = StageConfig(
+            stage_id=0,
+            model_stage="thinker",
+            runtime_overrides={
+                "gpu_memory_utilization": 0.9,
+                "tensor_parallel_size": 2,
+                "devices": "0,1",
+                "max_batch_size": 64,
+            },
+        )
+        omega_config = config.to_omegaconf()
+
+        assert omega_config.engine_args.gpu_memory_utilization == 0.9
+        assert omega_config.engine_args.tensor_parallel_size == 2
+        assert omega_config.runtime.devices == "0,1"
+        assert omega_config.runtime.max_batch_size == 64
+
+
+class TestModelPipeline:
+    """Tests for ModelPipeline class."""
+
+    def test_valid_linear_dag(self):
+        """Test validation of a valid linear DAG."""
+        stages = [
+            StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
+            StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
+            StageConfig(stage_id=2, model_stage="code2wav", input_sources=[1]),
+        ]
+        pipeline = ModelPipeline(model_type="test", stages=stages)
+        errors = pipeline.validate_pipeline()
+        assert errors == [], f"Unexpected errors: {errors}"
+
+    def test_valid_branching_dag(self):
+        """Test validation of a valid branching DAG."""
+        stages = [
+            StageConfig(stage_id=0, model_stage="input", input_sources=[]),
+            StageConfig(stage_id=1, model_stage="branch_a", input_sources=[0]),
+            StageConfig(stage_id=2, model_stage="branch_b", input_sources=[0]),
+        ]
+        pipeline = ModelPipeline(model_type="test", stages=stages)
+        errors = pipeline.validate_pipeline()
+        assert errors == [], f"Unexpected errors: {errors}"
+
+    def test_missing_entry_point(self):
+        """Test that missing entry point is detected."""
+        stages = [
+            StageConfig(stage_id=0, model_stage="stage_a", input_sources=[1]),
+            StageConfig(stage_id=1, model_stage="stage_b", input_sources=[0]),
+        ]
+        pipeline = ModelPipeline(model_type="test", stages=stages)
+        errors = pipeline.validate_pipeline()
+        assert any("entry point" in e.lower() for e in errors)
+
+    def test_missing_dependency(self):
+        """Test that missing stage reference is detected."""
+        stages = [
+            StageConfig(stage_id=0, model_stage="input", input_sources=[]),
+            StageConfig(stage_id=1, model_stage="output", input_sources=[99]),  # Invalid
+        ]
+        pipeline = ModelPipeline(model_type="test", stages=stages)
+        errors = pipeline.validate_pipeline()
+        assert any("non-existent" in e.lower() for e in errors)
+
+    def test_duplicate_stage_ids(self):
+        """Test that duplicate stage IDs are detected."""
+        stages = [
+            StageConfig(stage_id=0, model_stage="stage_a", input_sources=[]),
+            StageConfig(stage_id=0, model_stage="stage_b", input_sources=[]),  # Duplicate
+        ]
+        pipeline = ModelPipeline(model_type="test", stages=stages)
+        errors = pipeline.validate_pipeline()
+        assert any("duplicate" in e.lower() for e in errors)
+
+    def test_self_reference(self):
+        """Test that self-references are detected."""
+        stages = [
+            StageConfig(stage_id=0, model_stage="entry", input_sources=[]),
+            StageConfig(stage_id=1, model_stage="self_ref", input_sources=[1]),  # Self
+        ]
+        pipeline = ModelPipeline(model_type="test", stages=stages)
+        errors = pipeline.validate_pipeline()
+        assert any("itself" in e.lower() for e in errors)
+
+    def test_get_stage_by_id(self):
+        """Test getting stage by ID."""
+        stages = [
+            StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
+            StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
+        ]
+        pipeline = ModelPipeline(model_type="test", stages=stages)
+
+        stage = pipeline.get_stage(1)
+        assert stage is not None
+        assert stage.model_stage == "talker"
+
+        missing = pipeline.get_stage(99)
+        assert missing is None
+
+    def test_empty_pipeline(self):
+        """Test validation of empty pipeline."""
+        pipeline = ModelPipeline(model_type="test", stages=[])
+        errors = pipeline.validate_pipeline()
+        assert any("no stages" in e.lower() for e in errors)
+
+
+class TestStageConfigFactory:
+    """Tests for StageConfigFactory class."""
+
+    def test_default_diffusion_no_yaml(self):
+        """Test single-stage diffusion works without YAML config (@ZJY0516)."""
+        kwargs = {
+            "cache_backend": "none",
+            "cache_config": None,
+            "dtype": "bfloat16",
+        }
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+
+        assert len(configs) == 1
+        cfg = configs[0]
+        assert cfg["stage_id"] == 0
+        assert cfg["stage_type"] == "diffusion"
+        assert cfg["final_output"] is True
+        assert cfg["final_output_type"] == "image"
+
+    def test_default_diffusion_with_parallel_config(self):
+        """Test diffusion config calculates devices from parallel_config."""
+
+        class MockParallelConfig:
+            world_size = 4
+
+        kwargs = {
+            "parallel_config": MockParallelConfig(),
+            "cache_backend": "tea_cache",
+        }
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+
+        assert configs[0]["runtime"]["devices"] == "0,1,2,3"
+
+    def test_per_stage_override_precedence(self):
+        """Test that --stage-0-gpu-memory-utilization overrides global."""
+        stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[])
+        cli_overrides = {
+            "gpu_memory_utilization": 0.5,  # Global
+            "stage_0_gpu_memory_utilization": 0.9,  # Per-stage override
+        }
+
+        overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides)
+
+        # Per-stage should override global
+        assert overrides["gpu_memory_utilization"] == 0.9
+
+    def test_cli_override_forwards_engine_registered_args(self):
+        """Test that any engine-registered CLI arg is forwarded (@wuhang2014)."""
+        stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[])
+        cli_overrides = {
+            "gpu_memory_utilization": 0.9,  # Well-known param
+            "custom_engine_flag": True,  # Not in _INTERNAL_KEYS, so forwarded
+        }
+
+        overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides)
+
+        assert overrides["gpu_memory_utilization"] == 0.9
+        assert overrides["custom_engine_flag"] is True
+
+    def test_cli_override_excludes_internal_keys(self):
+        """Test that internal/orchestrator keys are not forwarded."""
+        stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[])
+        cli_overrides = {
+            "gpu_memory_utilization": 0.9,
+            "model": "some_model",  # Internal
+            "stage_configs_path": "/path",  # Internal
+            "batch_timeout": 10,  # Internal
+        }
+
+        overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides)
+
+        assert overrides["gpu_memory_utilization"] == 0.9
+        assert "model" not in overrides
+        assert "stage_configs_path" not in overrides
+        assert "batch_timeout" not in overrides
+
+    def test_per_stage_override_excludes_internal_keys(self):
+        """Test that per-stage overrides also skip internal keys."""
+        stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[])
+        cli_overrides = {
+            "stage_0_gpu_memory_utilization": 0.9,
+            "stage_0_model": "override_model",  # Internal, should be skipped
+            "stage_0_batch_timeout": 5,  # Internal, should be skipped
+        }
+
+        overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides)
+
+        assert overrides["gpu_memory_utilization"] == 0.9
+        assert "model" not in overrides
+        assert "batch_timeout" not in overrides
+
+
+class TestPipelineYamlParsing:
+    """Tests for pipeline YAML file parsing (@ZJY0516)."""
+
+    def test_parse_qwen3_omni_moe_yaml(self, tmp_path):
+        """Test parsing the qwen3_omni_moe pipeline YAML."""
+        yaml_content = """\
+model_type: qwen3_omni_moe
+async_chunk: false
+
+stages:
+  - stage_id: 0
+    model_stage: thinker
+    stage_type: llm
+    input_sources: []
+    worker_type: ar
+    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+    hf_config_name: thinker_config
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+
+  - stage_id: 1
+    model_stage: talker
+    stage_type: llm
+    input_sources: [0]
+    worker_type: ar
+    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+    hf_config_name: talker_config
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+
+  - stage_id: 2
+    model_stage: code2wav
+    stage_type: llm
+    input_sources: [1]
+    worker_type: generation
+    scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+    hf_config_name: thinker_config
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+    final_output: true
+    final_output_type: audio
+"""
+        yaml_file = tmp_path / "qwen3_omni_moe.yaml"
+        yaml_file.write_text(yaml_content)
+
+        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "qwen3_omni_moe")
+
+        assert pipeline.model_type == "qwen3_omni_moe"
+        assert len(pipeline.stages) == 3
+
+        # Stage 0: thinker
+        s0 = pipeline.stages[0]
+        assert s0.stage_id == 0
+        assert s0.model_stage == "thinker"
+        assert s0.stage_type == StageType.LLM
+        assert s0.input_sources == []
+        assert s0.worker_type == "ar"
+        assert s0.final_output is True
+        assert s0.final_output_type == "text"
+        assert s0.is_comprehension is True
+
+        # Stage 1: talker
+        s1 = pipeline.stages[1]
+        assert s1.stage_id == 1
+        assert s1.input_sources == [0]
+        assert s1.custom_process_input_func == (
+            "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker"
+        )
+        assert s1.final_output is False
+
+        # Stage 2: code2wav
+        s2 = pipeline.stages[2]
+        assert s2.stage_id == 2
+        assert s2.input_sources == [1]
+        assert s2.worker_type == "generation"
+        assert s2.final_output_type == "audio"
+
+    def test_parse_yaml_with_legacy_engine_input_source(self, tmp_path):
+        """Test backward compatibility with engine_input_source field."""
+        yaml_content = """\
+model_type: legacy_model
+
+stages:
+  - stage_id: 0
+    model_stage: entry
+    stage_type: llm
+  - stage_id: 1
+    model_stage: downstream
+    stage_type: llm
+    engine_input_source: [0]
+"""
+        yaml_file = tmp_path / "legacy.yaml"
+        yaml_file.write_text(yaml_content)
+
+        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_model")
+        assert pipeline.stages[1].input_sources == [0]
+
+    def test_parse_yaml_with_connectors_and_edges(self, tmp_path):
+        """Test parsing pipeline with optional connectors and edges."""
+        yaml_content = """\
+model_type: test_model
+
+stages:
+  - stage_id: 0
+    model_stage: entry
+    stage_type: llm
+    input_sources: []
+
+connectors:
+  type: ray
+
+edges:
+  - from: 0
+    to: 1
+"""
+        yaml_file = tmp_path / "with_connectors.yaml"
+        yaml_file.write_text(yaml_content)
+
+        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_model")
+        assert pipeline.connectors == {"type": "ray"}
+        assert pipeline.edges == [{"from": 0, "to": 1}]
+
+    def test_parsed_pipeline_passes_validation(self, tmp_path):
+        """Test that a well-formed YAML produces a valid pipeline."""
+        yaml_content = """\
+model_type: valid_model
+
+stages:
+  - stage_id: 0
+    model_stage: entry
+    stage_type: llm
+    input_sources: []
+    final_output: true
+    final_output_type: text
+  - stage_id: 1
+    model_stage: next
+    stage_type: llm
+    input_sources: [0]
+"""
+        yaml_file = tmp_path / "valid.yaml"
+        yaml_file.write_text(yaml_content)
+
+        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "valid_model")
+        errors = pipeline.validate_pipeline()
+        assert errors == [], f"Unexpected validation errors: {errors}"
+
+    def test_parse_diffusion_stage_type(self, tmp_path):
+        """Test parsing a diffusion stage type from YAML."""
+        yaml_content = """\
+model_type: diff_model
+
+stages:
+  - stage_id: 0
+    model_stage: dit
+    stage_type: diffusion
+    input_sources: []
+    final_output: true
+    final_output_type: image
+"""
+        yaml_file = tmp_path / "diffusion.yaml"
+        yaml_file.write_text(yaml_content)
+
+        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "diff_model")
+        assert pipeline.stages[0].stage_type == StageType.DIFFUSION
+
+    def test_parse_mimo_audio_yaml(self, tmp_path):
+        """Test parsing the MiMo Audio pipeline YAML."""
+        yaml_content = """\
+model_type: mimo_audio
+
+stages:
+  - stage_id: 0
+    model_stage: fused_thinker_talker
+    stage_type: llm
+    input_sources: []
+    worker_type: ar
+    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+
+  - stage_id: 1
+    model_stage: code2wav
+    stage_type: llm
+    input_sources: [0]
+    worker_type: generation
+    scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav
+    final_output: true
+    final_output_type: audio
+"""
+        yaml_file = tmp_path / "mimo_audio.yaml"
+        yaml_file.write_text(yaml_content)
+
+        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "mimo_audio")
+
+        assert pipeline.model_type == "mimo_audio"
+        assert len(pipeline.stages) == 2
+
+        # Stage 0: fused_thinker_talker
+        s0 = pipeline.stages[0]
+        assert s0.stage_id == 0
+        assert s0.model_stage == "fused_thinker_talker"
+        assert s0.input_sources == []
+        assert s0.worker_type == "ar"
+        assert s0.is_comprehension is True
+        assert s0.final_output is True
+        assert s0.final_output_type == "text"
+
+        # Stage 1: code2wav
+        s1 = pipeline.stages[1]
+        assert s1.stage_id == 1
+        assert s1.input_sources == [0]
+        assert s1.worker_type == "generation"
+        assert s1.custom_process_input_func == (
+            "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav"
+        )
+        assert s1.final_output is True
+        assert s1.final_output_type == "audio"
+
+        # Pipeline validation
+        errors = pipeline.validate_pipeline()
+        assert errors == [], f"Unexpected errors: {errors}"
+
+
+class TestAsyncChunk:
+    """Tests for async_chunk pipeline flag."""
+
+    def test_model_pipeline_async_chunk_default(self):
+        """Test ModelPipeline defaults async_chunk to False."""
+        stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
+        pipeline = ModelPipeline(model_type="test", stages=stages)
+        assert pipeline.async_chunk is False
+
+    def test_model_pipeline_async_chunk_set(self):
+        """Test ModelPipeline with async_chunk=True."""
+        stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
+        pipeline = ModelPipeline(model_type="test", stages=stages, async_chunk=True)
+        assert pipeline.async_chunk is True
+
+    def test_parse_async_chunk_from_yaml(self, tmp_path):
+        """Test that async_chunk is parsed from pipeline YAML."""
+        yaml_content = """\
+model_type: test_async
+async_chunk: true
+
+stages:
+  - stage_id: 0
+    model_stage: entry
+    stage_type: llm
+    input_sources: []
+"""
+        yaml_file = tmp_path / "async.yaml"
+        yaml_file.write_text(yaml_content)
+
+        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_async")
+        assert pipeline.async_chunk is True
+
+    def test_parse_missing_async_chunk_defaults_false(self, tmp_path):
+        """Test that missing async_chunk defaults to False."""
+        yaml_content = """\
+model_type: test_no_async
+
+stages:
+  - stage_id: 0
+    model_stage: entry
+    stage_type: llm
+    input_sources: []
+"""
+        yaml_file = tmp_path / "no_async.yaml"
+        yaml_file.write_text(yaml_content)
+
+        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_no_async")
+        assert pipeline.async_chunk is False
+
+
+class TestArchitectureFallback:
+    """Tests for architecture-based model detection fallback."""
+
+    def test_architecture_models_mapping_exists(self):
+        """Test that _ARCHITECTURE_MODELS contains expected entries."""
+        assert "MiMoAudioForConditionalGeneration" in StageConfigFactory._ARCHITECTURE_MODELS
+        assert StageConfigFactory._ARCHITECTURE_MODELS["MiMoAudioForConditionalGeneration"] == "mimo_audio"
+        assert "HunyuanImage3ForCausalMM" in StageConfigFactory._ARCHITECTURE_MODELS
+        assert StageConfigFactory._ARCHITECTURE_MODELS["HunyuanImage3ForCausalMM"] == "hunyuan_image3"
+
+    def test_mimo_audio_in_pipeline_models(self):
+        """Test that mimo_audio is registered in PIPELINE_MODELS."""
+        assert "mimo_audio" in StageConfigFactory.PIPELINE_MODELS

vllm_omni/config/__init__.py

@@ -4,8 +4,28 @@ Configuration module for vLLM-Omni.
 
 from vllm_omni.config.lora import LoRAConfig
 from vllm_omni.config.model import OmniModelConfig
+from vllm_omni.config.stage_config import (
+    ModelPipeline,
+    StageConfig,
+    StageConfigFactory,
+    StageType,
+)
+from vllm_omni.config.yaml_util import (
+    create_config,
+    load_yaml_config,
+    merge_configs,
+    to_dict,
+)
 
 __all__ = [
     "OmniModelConfig",
     "LoRAConfig",
+    "StageConfig",
+    "StageConfigFactory",
+    "ModelPipeline",
+    "StageType",
+    "create_config",
+    "load_yaml_config",
+    "merge_configs",
+    "to_dict",
 ]

vllm_omni/config/stage_config.py

@@ -0,0 +1,487 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Stage Configuration System for vLLM-Omni.
+
+Pipeline structure (stages, types, data-flow) is defined in per-model YAML
+files and is set by model developers at integration time.
+Runtime parameters (gpu_memory_utilization, tp_size, etc.) come from CLI.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import Any
+
+from vllm.logger import init_logger
+
+from vllm_omni.config.yaml_util import create_config, load_yaml_config, to_dict
+
+# Pipeline YAMLs live alongside model code in model_executor/models/<model>/
+_MODELS_DIR = Path(__file__).resolve().parent.parent / "model_executor" / "models"
+
+
+def get_pipeline_path(model_dir: str, filename: str) -> Path:
+    """Return the full path to a pipeline YAML file.
+
+    Args:
+        model_dir: Model subdirectory name (e.g., "qwen3_omni").
+        filename: Name of the YAML file (e.g., "pipeline.yaml").
+
+    Returns:
+        Absolute path to the file.
+    """
+    return _MODELS_DIR / model_dir / filename
+
+
+logger = init_logger(__name__)
+
+
+class StageType(str, Enum):
+    """Type of processing stage in the Omni pipeline."""
+
+    LLM = "llm"
+    DIFFUSION = "diffusion"
+
+
+@dataclass
+class StageConfig:
+    """Per-stage configuration — pipeline-structure fields only.
+
+    Engine params (gpu_memory_utilization, tp_size, etc.) come from CLI,
+    NOT from this class.
+    """
+
+    # Identity
+    stage_id: int
+    model_stage: str
+
+    # Stage type
+    stage_type: StageType = StageType.LLM
+
+    input_sources: list[int] = field(default_factory=list)
+    custom_process_input_func: str | None = None
+    final_output: bool = False
+    final_output_type: str | None = None  # "text", "audio", "image"
+    worker_type: str | None = None  # "ar" or "generation"
+    scheduler_cls: str | None = None
+    hf_config_name: str | None = None
+    is_comprehension: bool = False
+
+    # Runtime overrides (populated from CLI, not from pipeline YAML)
+    runtime_overrides: dict[str, Any] = field(default_factory=dict)
+
+    def to_omegaconf(self) -> Any:
+        """Convert to OmegaConf for backward compatibility with OmniStage.
+
+        Returns:
+            OmegaConf DictConfig with stage configuration in legacy format.
+        """
+        # Build engine_args dict with required fields
+        engine_args: dict[str, Any] = {
+            "model_stage": self.model_stage,
+        }
+
+        if self.worker_type:
+            engine_args["worker_type"] = self.worker_type
+        if self.scheduler_cls:
+            engine_args["scheduler_cls"] = self.scheduler_cls
+        if self.hf_config_name:
+            engine_args["hf_config_name"] = self.hf_config_name
+
+        # Apply runtime overrides (CLI args)
+        for key, value in self.runtime_overrides.items():
+            if key not in ("devices", "max_batch_size"):
+                engine_args[key] = value
+
+        # Build runtime config
+        runtime: dict[str, Any] = {
+            "process": True,
+            "max_batch_size": self.runtime_overrides.get("max_batch_size", 1),
+        }
+        if "devices" in self.runtime_overrides:
+            runtime["devices"] = self.runtime_overrides["devices"]
+
+        # Build full config dict
+        config_dict: dict[str, Any] = {
+            "stage_id": self.stage_id,
+            "stage_type": StageType(self.stage_type).value,
+            "engine_args": create_config(engine_args),
+            "runtime": create_config(runtime),
+            "engine_input_source": self.input_sources,  # Legacy field name
+            "final_output": self.final_output,
+            "final_output_type": self.final_output_type,
+            "is_comprehension": self.is_comprehension,
+        }
+
+        if self.custom_process_input_func:
+            config_dict["custom_process_input_func"] = self.custom_process_input_func
+
+        return create_config(config_dict)
+
+
+@dataclass
+class ModelPipeline:
+    """Complete pipeline definition for a multi-stage model.
+
+    Defined by model developers, bundled with the model, not user-editable.
+    """
+
+    model_type: str
+    stages: list[StageConfig]
+
+    # Pipeline-wide behavior flags
+    async_chunk: bool = False
+
+    # Optional distributed configuration
+    connectors: dict[str, Any] | None = None
+    edges: list[dict[str, Any]] | None = None
+
+    def get_stage(self, stage_id: int) -> StageConfig | None:
+        """Look up a stage by its ID.
+
+        Args:
+            stage_id: The stage ID to search for.
+
+        Returns:
+            The matching StageConfig, or None if not found.
+        """
+        for stage in self.stages:
+            if stage.stage_id == stage_id:
+                return stage
+        return None
+
+    def validate_pipeline(self) -> list[str]:
+        """Validate pipeline topology at model integration time (not runtime).
+
+        Checks:
+        - All stage IDs are unique
+        - All input_sources reference valid stage IDs
+        - At least one entry point (stage with empty input_sources)
+
+        Returns:
+            List of validation error messages. Empty list if valid.
+        """
+        errors: list[str] = []
+
+        if not self.stages:
+            errors.append("Topology has no stages defined")
+            return errors
+
+        # Check for unique stage IDs
+        stage_ids = [s.stage_id for s in self.stages]
+        if len(stage_ids) != len(set(stage_ids)):
+            errors.append("Duplicate stage IDs found")
+
+        stage_id_set = set(stage_ids)
+
+        # Check input_sources reference valid stages
+        for stage in self.stages:
+            for source_id in stage.input_sources:
+                if source_id not in stage_id_set:
+                    errors.append(f"Stage {stage.stage_id} references non-existent input source {source_id}")
+                if source_id == stage.stage_id:
+                    errors.append(f"Stage {stage.stage_id} references itself as input source")
+
+        # Check for at least one entry point
+        entry_points = [s for s in self.stages if not s.input_sources]
+        if not entry_points:
+            errors.append("No entry point found (stage with empty input_sources)")
+
+        return errors
+
+
+class StageConfigFactory:
+    """Factory that loads pipeline YAML and merges CLI overrides.
+
+    Handles both single-stage and multi-stage models.
+    """
+
+    # Mapping of model types to directories under model_executor/models/.
+    PIPELINE_MODELS: dict[str, str] = {
+        "qwen3_omni_moe": "qwen3_omni",
+        "qwen2_5_omni": "qwen2_5_omni",
+        "bagel": "bagel",
+        "qwen3_tts": "qwen3_tts",
+        "mimo_audio": "mimo_audio",
+        "glm-image": "glm_image",
+        "cosyvoice3": "cosyvoice3",
+        "mammothmoda2": "mammoth_moda2",
+    }
+
+    # Fallback: map HF architecture class names to pipeline dirs.
+    # Used when model_type collides with another model (e.g. MiMo Audio
+    # reports model_type="qwen2" which matches plain Qwen2, not our pipeline).
+    _ARCHITECTURE_MODELS: dict[str, str] = {
+        "MiMoAudioForConditionalGeneration": "mimo_audio",
+        "HunyuanImage3ForCausalMM": "hunyuan_image3",
+    }
+
+    @classmethod
+    def create_from_model(
+        cls,
+        model: str,
+        cli_overrides: dict[str, Any] | None = None,
+    ) -> list[StageConfig] | None:
+        """Load pipeline YAML, merge with CLI overrides.
+
+        Args:
+            model: Model name or path.
+            cli_overrides: CLI overrides from VllmConfig/OmniDiffusionConfig.
+
+        Returns:
+            List of StageConfig objects with CLI overrides applied,
+            or None if no pipeline definition was found for this model.
+        """
+        if cli_overrides is None:
+            cli_overrides = {}
+
+        trust_remote_code = cli_overrides.get("trust_remote_code", True)
+        pipeline = cls._load_pipeline(model, trust_remote_code=trust_remote_code)
+
+        if pipeline is None:
+            return None
+
+        errors = pipeline.validate_pipeline()
+        if errors:
+            logger.warning(f"Pipeline validation warnings for {model}: {errors}")
+
+        # Apply CLI overrides
+        result: list[StageConfig] = []
+        for stage in pipeline.stages:
+            # Merge global CLI overrides
+            stage.runtime_overrides = cls._merge_cli_overrides(stage, cli_overrides)
+            result.append(stage)
+
+        return result
+
+    @classmethod
+    def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
+        """Single-stage diffusion - no YAML needed.
+
+        Creates a default diffusion stage configuration for single-stage
+        diffusion models. Returns a legacy OmegaConf-compatible dict for
+        backward compatibility with OmniStage.
+
+        Args:
+            kwargs: Engine arguments from CLI/API.
+
+        Returns:
+            List containing a single config dict for the diffusion stage.
+        """
+        # Calculate devices based on parallel config
+        devices = "0"
+        if "parallel_config" in kwargs:
+            num_devices = kwargs["parallel_config"].world_size
+            for i in range(1, num_devices):
+                devices += f",{i}"
+
+        # Collect engine args – skip non-serializable objects
+        engine_args: dict[str, Any] = {}
+        for key, value in kwargs.items():
+            if key in ("parallel_config",):
+                continue
+            engine_args[key] = value
+
+        engine_args.setdefault("cache_backend", "none")
+        engine_args["model_stage"] = "diffusion"
+
+        # Convert dtype to string for OmegaConf
+        if "dtype" in engine_args:
+            engine_args["dtype"] = str(engine_args["dtype"])
+
+        config_dict: dict[str, Any] = {
+            "stage_id": 0,
+            "stage_type": StageType.DIFFUSION.value,
+            "runtime": {
+                "process": True,
+                "devices": devices,
+                "max_batch_size": 1,
+            },
+            "engine_args": create_config(engine_args),
+            "final_output": True,
+            "final_output_type": "image",
+        }
+
+        return [config_dict]
+
+    @classmethod
+    def _load_pipeline(cls, model: str, trust_remote_code: bool = True) -> ModelPipeline | None:
+        """Load pipeline YAML for the model.
+
+        Args:
+            model: Model name or path.
+            trust_remote_code: Whether to trust remote code for HF config loading.
+
+        Returns:
+            ModelPipeline if found, None otherwise.
+        """
+        model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code)
+        if model_type is None:
+            return None
+
+        pipeline_dir = cls.PIPELINE_MODELS.get(model_type)
+
+        # Fallback: check HF architectures when model_type doesn't match
+        if pipeline_dir is None and hf_config is not None:
+            for arch in getattr(hf_config, "architectures", []) or []:
+                pipeline_dir = cls._ARCHITECTURE_MODELS.get(arch)
+                if pipeline_dir is not None:
+                    model_type = pipeline_dir
+                    break
+
+        if pipeline_dir is None:
+            logger.debug(f"No pipeline mapping for model_type: {model_type}")
+            return None
+
+        pipeline_path = get_pipeline_path(pipeline_dir, "pipeline.yaml")
+
+        if not pipeline_path.exists():
+            logger.debug(f"Pipeline file not found: {pipeline_path}")
+            return None
+
+        return cls._parse_pipeline_yaml(pipeline_path, model_type)
+
+    @classmethod
+    def _parse_pipeline_yaml(cls, path: Path, model_type: str) -> ModelPipeline:
+        """Parse a pipeline YAML file.
+
+        Args:
+            path: Path to the YAML file.
+            model_type: Model type identifier.
+
+        Returns:
+            ModelPipeline object.
+        """
+        config_data = load_yaml_config(path)
+
+        stages: list[StageConfig] = []
+        for stage_data in config_data.stages:
+            # Use .get() for optional fields — idiomatic for OmegaConf DictConfig
+            stage_type_str = stage_data.get("stage_type", "llm")
+            stage_type = StageType(stage_type_str) if stage_type_str else StageType.LLM
+
+            # Handle both 'input_sources' (new) and 'engine_input_source' (legacy)
+            input_sources = stage_data.get("input_sources", None)
+            if input_sources is None:
+                input_sources = stage_data.get("engine_input_source", [])
+            if input_sources is None:
+                input_sources = []
+            input_sources = list(input_sources)
+
+            stage = StageConfig(
+                stage_id=stage_data.stage_id,
+                model_stage=stage_data.model_stage,
+                stage_type=stage_type,
+                input_sources=input_sources,
+                custom_process_input_func=stage_data.get("custom_process_input_func", None),
+                final_output=stage_data.get("final_output", False),
+                final_output_type=stage_data.get("final_output_type", None),
+                worker_type=stage_data.get("worker_type", None),
+                scheduler_cls=stage_data.get("scheduler_cls", None),
+                hf_config_name=stage_data.get("hf_config_name", None),
+                is_comprehension=stage_data.get("is_comprehension", False),
+            )
+            stages.append(stage)
+
+        # Get pipeline-wide flags
+        async_chunk = config_data.get("async_chunk", False)
+
+        # Get optional connector config
+        connectors = to_dict(config_data.connectors) if hasattr(config_data, "connectors") else None
+        edges = to_dict(config_data.edges) if hasattr(config_data, "edges") else None
+
+        return ModelPipeline(
+            model_type=getattr(config_data, "model_type", model_type),
+            stages=stages,
+            async_chunk=async_chunk,
+            connectors=connectors,
+            edges=edges,
+        )
+
+    @classmethod
+    def _auto_detect_model_type(cls, model: str, trust_remote_code: bool = True) -> tuple[str | None, Any]:
+        """Auto-detect model_type from model directory.
+
+        Args:
+            model: Model name or path.
+            trust_remote_code: Whether to trust remote code for HF config loading.
+
+        Returns:
+            Tuple of (model_type, hf_config). Both may be None on failure.
+        """
+        try:
+            from vllm.transformers_utils.config import get_config
+
+            hf_config = get_config(model, trust_remote_code=trust_remote_code)
+            return hf_config.model_type, hf_config
+        except Exception as e:
+            logger.debug(f"Failed to auto-detect model type for {model}: {e}")
+            return None, None
+
+    # Keys that should never be forwarded as engine overrides (internal /
+    # orchestrator-only knobs, complex objects, etc.).
+    _INTERNAL_KEYS: set[str] = {
+        "model",
+        "stage_configs_path",
+        "stage_id",
+        "stage_init_timeout",
+        "init_timeout",
+        "shm_threshold_bytes",
+        "worker_backend",
+        "ray_address",
+        "batch_timeout",
+        "log_stats",
+        "tokenizer",
+        "parallel_config",
+    }
+
+    @classmethod
+    def _merge_cli_overrides(
+        cls,
+        stage: StageConfig,
+        cli_overrides: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Merge CLI overrides into stage runtime config.
+
+        All CLI arguments registered by engine config classes (e.g.
+        EngineArgs / OmniDiffusionConfig) are accepted as overrides
+        unless they appear in ``_INTERNAL_KEYS``.
+
+        Handles:
+        - Global overrides (apply to all stages)
+        - Per-stage overrides (--stage-N-* format, take precedence)
+
+        Args:
+            stage: The stage to merge overrides into.
+            cli_overrides: CLI arguments from VllmConfig/OmniDiffusionConfig.
+
+        Returns:
+            Dict of runtime overrides for this stage.
+        """
+        result: dict[str, Any] = {}
+
+        # Apply global overrides – any key not in the internal blocklist
+        # is forwarded so that engine-registered params work out of the box.
+        for key, value in cli_overrides.items():
+            if key in cls._INTERNAL_KEYS:
+                continue
+            if re.match(r"stage_\d+_", key):
+                # Per-stage keys handled below
+                continue
+            if value is not None:
+                result[key] = value
+
+        # Apply per-stage overrides (--stage-N-* format, take precedence)
+        stage_prefix = f"stage_{stage.stage_id}_"
+        for key, value in cli_overrides.items():
+            if key.startswith(stage_prefix) and value is not None:
+                param_name = key[len(stage_prefix) :]
+                if param_name in cls._INTERNAL_KEYS:
+                    continue
+                result[param_name] = value
+
+        return result

vllm_omni/config/yaml_util.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Centralized OmegaConf wrapper for vLLM-Omni.
+
+All OmegaConf usage in the project MUST go through this module.
+Other modules should import these wrapper functions instead of
+using OmegaConf directly.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from omegaconf import DictConfig, OmegaConf
+
+
+def load_yaml_config(path: str | Any) -> DictConfig:
+    """Load a YAML file and return it as a DictConfig.
+
+    Args:
+        path: Path to the YAML file.
+
+    Returns:
+        OmegaConf DictConfig with attribute-style access.
+    """
+    return OmegaConf.load(path)  # type: ignore[return-value]
+
+
+def create_config(data: Any) -> DictConfig:
+    """Wrap a dict (or list) into a DictConfig.
+
+    Args:
+        data: Dict, list, or other structure to wrap.
+
+    Returns:
+        OmegaConf DictConfig / ListConfig.
+    """
+    return OmegaConf.create(data)
+
+
+def merge_configs(*cfgs: Any) -> dict:
+    """Deep-merge multiple configs and return a plain dict.
+
+    Args:
+        *cfgs: DictConfig or dict objects to merge (left to right).
+
+    Returns:
+        Plain dict with merged, resolved values.
+    """
+    merged = OmegaConf.merge(*cfgs)
+    return OmegaConf.to_container(merged, resolve=True)  # type: ignore[return-value]
+
+
+def to_dict(obj: Any, *, resolve: bool = True) -> Any:
+    """Convert a DictConfig (or similar) to a plain dict.
+
+    Args:
+        obj: OmegaConf container to convert.
+        resolve: Whether to resolve interpolations (default True).
+
+    Returns:
+        Plain dict.
+    """
+    return OmegaConf.to_container(obj, resolve=resolve)  # type: ignore[return-value]

vllm_omni/entrypoints/omni.py

@@ -15,13 +15,14 @@ import huggingface_hub
 import msgspec.msgpack
 import torch
 import zmq
-from omegaconf import OmegaConf
 from tqdm.auto import tqdm
 from vllm import SamplingParams
 from vllm.logger import init_logger
 from vllm.utils.network_utils import make_zmq_socket
 from vllm.v1.utils import get_engine_client_zmq_addr
 
+from vllm_omni.config.stage_config import StageConfigFactory
+from vllm_omni.config.yaml_util import create_config
 from vllm_omni.distributed.omni_connectors import (
     get_stage_connector_config,
     initialize_orchestrator_connectors,
@@ -224,46 +225,35 @@ class OmniBase:
             cache_config = self._get_default_cache_config(cache_backend)
         return cache_config
 
-    def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]:
-        """Create default diffusion stage configuration."""
-        # We temporally create a default config for diffusion stage.
-        # In the future, we should merge the default config with the user-provided config.
-        # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error.
+    def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
+        """Create default diffusion stage configuration.
+
+        Uses StageConfigFactory for typed configuration creation while
+        maintaining backward compatibility with the legacy format.
+
+        Args:
+            kwargs: Engine arguments from CLI/API.
+
+        Returns:
+            List containing a single OmegaConf config for the diffusion stage.
+        """
+        # Normalize dtype
         if "dtype" in kwargs and not isinstance(kwargs["dtype"], str):
             if not isinstance(kwargs["dtype"], torch.dtype):
                 raise TypeError(f"Provided dtype must be a string or torch.dtype, got {type(kwargs['dtype']).__name__}")
             kwargs["dtype"] = str(kwargs["dtype"]).removeprefix("torch.")
 
+        # Normalize cache config before passing to factory
         cache_backend = kwargs.get("cache_backend", "none")
         cache_config = self._normalize_cache_config(cache_backend, kwargs.get("cache_config", None))
-        # TODO: hack, calculate devices based on parallel config.
-        devices = "0"
-        if "parallel_config" in kwargs:
-            num_devices = kwargs["parallel_config"].world_size
-            for i in range(1, num_devices):
-                devices += f",{i}"
-        default_stage_cfg = [
-            {
-                "stage_id": 0,
-                "stage_type": "diffusion",
-                "runtime": {
-                    "process": True,
-                    "devices": devices,
-                    "max_batch_size": 1,
-                },
-                "engine_args": OmegaConf.create(
-                    {
-                        **kwargs,
-                        "cache_backend": cache_backend,
-                        "cache_config": cache_config,
-                    }
-                ),
-                "final_output": True,
-                "final_output_type": "image",
-            }
-        ]
-        default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion"
-        return default_stage_cfg
+
+        # Update kwargs with normalized values
+        kwargs_copy = dict(kwargs)
+        kwargs_copy["cache_backend"] = cache_backend
+        kwargs_copy["cache_config"] = cache_config
+
+        # Use the factory to create default diffusion config
+        return StageConfigFactory.create_default_diffusion(kwargs_copy)
 
     def _resolve_stage_configs(self, model: str, kwargs: dict[str, Any]) -> tuple[str, list[Any]]:
         """Resolve stage configs and inject defaults shared by orchestrator/headless."""
@@ -290,7 +280,7 @@ class OmniBase:
                 if getattr(cfg, "stage_type", None) != "diffusion":
                     continue
                 if not hasattr(cfg, "engine_args") or cfg.engine_args is None:
-                    cfg.engine_args = OmegaConf.create({})
+                    cfg.engine_args = create_config({})
                 if kwargs.get("lora_path") is not None:
                     if not hasattr(cfg.engine_args, "lora_path") or cfg.engine_args.lora_path is None:
                         cfg.engine_args.lora_path = kwargs["lora_path"]

vllm_omni/entrypoints/omni_stage.py

@@ -269,7 +269,10 @@ class OmniStage:
         self.engine_args = stage_config.engine_args
         self.model_stage = stage_config.engine_args.model_stage
         self.requires_multimodal_data = getattr(stage_config.runtime, "requires_multimodal_data", False)
-        self.engine_input_source = getattr(stage_config, "engine_input_source", [])
+        # Support both 'input_sources' (new format) and 'engine_input_source' (legacy)
+        self.engine_input_source = (
+            getattr(stage_config, "input_sources", None) or getattr(stage_config, "engine_input_source", []) or []
+        )
         self.engine_output_type = getattr(stage_config.engine_args, "engine_output_type", None)
         self.engine_outputs = None
         self.is_comprehension = getattr(stage_config, "is_comprehension", False)

vllm_omni/entrypoints/stage_utils.py

@@ -9,7 +9,7 @@ from collections.abc import Callable
 from multiprocessing import shared_memory as _shm
 from typing import Any
 
-from omegaconf import OmegaConf
+from vllm_omni.config.yaml_util import to_dict as _omega_to_dict
 
 logger = logging.getLogger(__name__)
 
@@ -315,7 +315,7 @@ def _to_dict(x: Any) -> dict[str, Any]:
     try:
         if isinstance(x, dict):
             return dict(x)
-        return OmegaConf.to_container(x, resolve=True)  # type: ignore[arg-type]
+        return _omega_to_dict(x)
     except Exception:
         try:
             return dict(x)

vllm_omni/entrypoints/utils.py

@@ -5,11 +5,11 @@ from dataclasses import asdict, fields, is_dataclass
 from pathlib import Path
 from typing import Any, get_args, get_origin
 
-from omegaconf import OmegaConf
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_config, get_hf_file_to_dict
 from vllm.transformers_utils.repo_utils import file_or_path_exists
 
+from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs
 from vllm_omni.entrypoints.stage_utils import _to_dict
 from vllm_omni.platforms import current_omni_platform
 
@@ -234,6 +234,11 @@ def resolve_model_config_path(model: str) -> str:
 def load_stage_configs_from_model(model: str, base_engine_args: dict | None = None) -> list:
     """Load stage configurations from model's default config file.
 
+    .. deprecated::
+        This is the legacy OmegaConf-based loading path. New code should use
+        ``StageConfigFactory.create_from_model()`` instead. This function will
+        be removed once all callers are migrated (see PR series [2/N]).
+
     Loads stage configurations based on the model type and device type.
     First tries to load a device-specific YAML file from stage_configs/{device_type}/
     directory. If not found, falls back to the default config file.
@@ -259,6 +264,9 @@ def load_stage_configs_from_model(model: str, base_engine_args: dict | None = No
 def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None = None) -> list:
     """Load stage configurations from a YAML file.
 
+    .. deprecated::
+        Legacy OmegaConf-based loader. Will be removed in PR series [2/N].
+
     Args:
         config_path: Path to the YAML configuration file
 
@@ -267,17 +275,17 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None
     """
     if base_engine_args is None:
         base_engine_args = {}
-    config_data = OmegaConf.load(config_path)
+    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 OmegaConf
+    # Convert any nested dataclass objects to dicts before creating DictConfig
     base_engine_args = _convert_dataclasses_to_dict(base_engine_args)
-    base_engine_args = OmegaConf.create(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:
-            base_engine_args_tmp = OmegaConf.merge(base_engine_args_tmp, stage_arg.engine_args)
+            base_engine_args_tmp = create_config(merge_configs(base_engine_args_tmp, stage_arg.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":
             runtime_cfg = stage_arg.runtime
@@ -312,7 +320,7 @@ def load_and_resolve_stage_configs(
         if not stage_configs:
             if default_stage_cfg_factory is not None:
                 default_stage_cfg = default_stage_cfg_factory()
-                stage_configs = OmegaConf.create(default_stage_cfg)
+                stage_configs = create_config(default_stage_cfg)
             else:
                 stage_configs = []
     else:
@@ -347,6 +355,7 @@ def get_final_stage_id_for_e2e(
         output_modalities = default_modalities
 
     try:
+        final_stage_id_for_e2e = last_stage_id
         for _sid in range(last_stage_id, -1, -1):
             if (
                 getattr(stage_list[_sid], "final_output", False)
@@ -354,8 +363,6 @@ def get_final_stage_id_for_e2e(
             ):
                 final_stage_id_for_e2e = _sid
                 break
-        if final_stage_id_for_e2e < 0:
-            final_stage_id_for_e2e = last_stage_id
     except Exception as e:
         logger.debug(
             "[Orchestrator] Failed to determine final stage for E2E; \