From c8c60e19709e3be7326e0bcd9124f935517dd598 Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Thu, 12 Mar 2026 21:09:34 +0800 Subject: [PATCH] [Config Refactor][1/2] Model Pipeline Configuration System (#1115) Signed-off-by: lishunyang --- tests/test_config_factory.py | 570 +++++++++++++++++++++++++++ vllm_omni/config/__init__.py | 20 + vllm_omni/config/stage_config.py | 487 +++++++++++++++++++++++ vllm_omni/config/yaml_util.py | 65 +++ vllm_omni/entrypoints/omni.py | 60 ++- vllm_omni/entrypoints/omni_stage.py | 5 +- vllm_omni/entrypoints/stage_utils.py | 4 +- vllm_omni/entrypoints/utils.py | 23 +- 8 files changed, 1188 insertions(+), 46 deletions(-) create mode 100644 tests/test_config_factory.py create mode 100644 vllm_omni/config/stage_config.py create mode 100644 vllm_omni/config/yaml_util.py diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py new file mode 100644 index 000000000..21ef22523 --- /dev/null +++ b/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 diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py index e2db6f427..2aa236e69 100644 --- a/vllm_omni/config/__init__.py +++ b/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", ] diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py new file mode 100644 index 000000000..ad298aa9d --- /dev/null +++ b/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// +_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 diff --git a/vllm_omni/config/yaml_util.py b/vllm_omni/config/yaml_util.py new file mode 100644 index 000000000..09b27ca02 --- /dev/null +++ b/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] diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index c414b972b..488b986d8 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/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"] diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 3f5a14e97..129832267 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/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) diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index de4a71719..99f2042e2 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/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) diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 69ce73fc4..0e31bfa7c 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/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; \