# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ Unit tests for StageConfigFactory and related classes. """ import importlib import inspect import warnings from dataclasses import dataclass, fields from pathlib import Path from unittest.mock import patch import pytest from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.config.stage_config import ( _PIPELINE_REGISTRY, DeployConfig, ModelPipeline, PipelineConfig, StageConfig, StageConfigFactory, StageDeployConfig, StageExecutionType, StagePipelineConfig, StageType, _apply_platform_overrides, _deep_merge_stage, _resolve_scheduler, build_stage_runtime_overrides, load_deploy_config, merge_pipeline_deploy, register_pipeline, strip_parent_engine_args, ) from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys pytestmark = [pytest.mark.core_model, pytest.mark.cpu] pytestmark = [pytest.mark.core_model, pytest.mark.cpu] 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" assert "max_num_seqs" not in omega_config.engine_args # 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, }, ) with warnings.catch_warnings(): warnings.simplefilter("ignore", FutureWarning) 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" # max_batch_size is migrated to engine_args.max_num_seqs assert omega_config.engine_args.max_num_seqs == 64 def test_to_omegaconf_max_batch_size_deprecation(self): """Test that runtime.max_batch_size emits a FutureWarning.""" config = StageConfig( stage_id=0, model_stage="thinker", runtime_overrides={"max_batch_size": 8}, ) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") config.to_omegaconf() deprecation_warnings = [x for x in w if issubclass(x.category, FutureWarning)] assert len(deprecation_warnings) == 1 assert "max_batch_size" in str(deprecation_warnings[0].message) def test_to_omegaconf_max_num_seqs_in_engine_args(self): """Test that max_num_seqs in yaml_engine_args takes precedence.""" config = StageConfig( stage_id=0, model_stage="thinker", yaml_engine_args={"max_num_seqs": 32}, ) omega_config = config.to_omegaconf() assert omega_config.engine_args.max_num_seqs == 32 def test_to_omegaconf_diffusion_parallel_overrides_replace_nested_values(self): config = StageConfig( stage_id=1, model_stage="diffusion", stage_type=StageType.DIFFUSION, yaml_engine_args={ "parallel_config": { "pipeline_parallel_size": 1, "data_parallel_size": 1, "tensor_parallel_size": 4, "enable_expert_parallel": False, "ulysses_degree": 1, "ring_degree": 1, "ulysses_mode": "strict", "sequence_parallel_size": 1, "cfg_parallel_size": 1, "vae_patch_parallel_size": 1, "use_hsdp": False, "hsdp_shard_size": -1, "hsdp_replicate_size": 1, } }, runtime_overrides={ "pipeline_parallel_size": 2, "data_parallel_size": 3, "tensor_parallel_size": 8, "enable_expert_parallel": True, "ulysses_degree": 2, "ring_degree": 4, "ulysses_mode": "advanced_uaa", "sequence_parallel_size": 8, "cfg_parallel_size": 2, "vae_patch_parallel_size": 2, "use_hsdp": True, "hsdp_shard_size": 8, "hsdp_replicate_size": 2, }, ) omega_config = config.to_omegaconf() assert omega_config.engine_args.parallel_config.pipeline_parallel_size == 2 assert omega_config.engine_args.parallel_config.data_parallel_size == 3 assert omega_config.engine_args.parallel_config.tensor_parallel_size == 8 assert omega_config.engine_args.parallel_config.enable_expert_parallel is True assert omega_config.engine_args.parallel_config.ulysses_degree == 2 assert omega_config.engine_args.parallel_config.ring_degree == 4 assert omega_config.engine_args.parallel_config.ulysses_mode == "advanced_uaa" assert omega_config.engine_args.parallel_config.sequence_parallel_size == 8 assert omega_config.engine_args.parallel_config.cfg_parallel_size == 2 assert omega_config.engine_args.parallel_config.vae_patch_parallel_size == 2 assert omega_config.engine_args.parallel_config.use_hsdp is True assert omega_config.engine_args.parallel_config.hsdp_shard_size == 8 assert omega_config.engine_args.parallel_config.hsdp_replicate_size == 2 assert "pipeline_parallel_size" not in omega_config.engine_args assert "data_parallel_size" not in omega_config.engine_args assert "tensor_parallel_size" not in omega_config.engine_args assert "enable_expert_parallel" not in omega_config.engine_args assert "ulysses_degree" not in omega_config.engine_args assert "ring_degree" not in omega_config.engine_args assert "ulysses_mode" not in omega_config.engine_args assert "sequence_parallel_size" not in omega_config.engine_args assert "cfg_parallel_size" not in omega_config.engine_args assert "vae_patch_parallel_size" not in omega_config.engine_args assert "use_hsdp" not in omega_config.engine_args assert "hsdp_shard_size" not in omega_config.engine_args assert "hsdp_replicate_size" not in omega_config.engine_args def test_to_omegaconf_diffusion_parallel_overrides_create_parallel_config(self): config = StageConfig( stage_id=1, model_stage="diffusion", stage_type=StageType.DIFFUSION, runtime_overrides={ "pipeline_parallel_size": 2, "data_parallel_size": 3, "tensor_parallel_size": 8, "enable_expert_parallel": True, "ulysses_degree": 2, "ring_degree": 4, "ulysses_mode": "advanced_uaa", "sequence_parallel_size": 8, "cfg_parallel_size": 2, "vae_patch_parallel_size": 2, "use_hsdp": True, "hsdp_shard_size": 8, "hsdp_replicate_size": 2, }, ) omega_config = config.to_omegaconf() assert omega_config.engine_args.parallel_config.pipeline_parallel_size == 2 assert omega_config.engine_args.parallel_config.data_parallel_size == 3 assert omega_config.engine_args.parallel_config.tensor_parallel_size == 8 assert omega_config.engine_args.parallel_config.enable_expert_parallel is True assert omega_config.engine_args.parallel_config.ulysses_degree == 2 assert omega_config.engine_args.parallel_config.ring_degree == 4 assert omega_config.engine_args.parallel_config.ulysses_mode == "advanced_uaa" assert omega_config.engine_args.parallel_config.sequence_parallel_size == 8 assert omega_config.engine_args.parallel_config.cfg_parallel_size == 2 assert omega_config.engine_args.parallel_config.vae_patch_parallel_size == 2 assert omega_config.engine_args.parallel_config.use_hsdp is True assert omega_config.engine_args.parallel_config.hsdp_shard_size == 8 assert omega_config.engine_args.parallel_config.hsdp_replicate_size == 2 assert "pipeline_parallel_size" not in omega_config.engine_args assert "data_parallel_size" not in omega_config.engine_args assert "tensor_parallel_size" not in omega_config.engine_args assert "enable_expert_parallel" not in omega_config.engine_args assert "ulysses_degree" not in omega_config.engine_args assert "ring_degree" not in omega_config.engine_args assert "ulysses_mode" not in omega_config.engine_args assert "sequence_parallel_size" not in omega_config.engine_args assert "cfg_parallel_size" not in omega_config.engine_args assert "vae_patch_parallel_size" not in omega_config.engine_args assert "use_hsdp" not in omega_config.engine_args assert "hsdp_shard_size" not in omega_config.engine_args assert "hsdp_replicate_size" not in omega_config.engine_args def test_to_omegaconf_diffusion_parallel_degree_overrides_recompute_sequence_parallel_size(self): config = StageConfig( stage_id=1, model_stage="diffusion", stage_type=StageType.DIFFUSION, yaml_engine_args={ "parallel_config": { "sequence_parallel_size": 1, "ulysses_degree": 1, "ring_degree": 1, } }, runtime_overrides={ "ulysses_degree": 2, "ring_degree": 4, }, ) omega_config = config.to_omegaconf() assert omega_config.engine_args.parallel_config.ulysses_degree == 2 assert omega_config.engine_args.parallel_config.ring_degree == 4 assert omega_config.engine_args.parallel_config.sequence_parallel_size == 8 assert "ulysses_degree" not in omega_config.engine_args assert "ring_degree" not in omega_config.engine_args assert "sequence_parallel_size" not in omega_config.engine_args def test_to_omegaconf_diffusion_parallel_explicit_sequence_parallel_size_is_preserved(self): config = StageConfig( stage_id=1, model_stage="diffusion", stage_type=StageType.DIFFUSION, yaml_engine_args={ "parallel_config": { "sequence_parallel_size": 1, "ulysses_degree": 1, "ring_degree": 1, } }, runtime_overrides={ "ulysses_degree": 2, "ring_degree": 4, "sequence_parallel_size": 16, }, ) omega_config = config.to_omegaconf() assert omega_config.engine_args.parallel_config.ulysses_degree == 2 assert omega_config.engine_args.parallel_config.ring_degree == 4 assert omega_config.engine_args.parallel_config.sequence_parallel_size == 16 def test_to_omegaconf_llm_parallel_overrides_remain_top_level(self): config = StageConfig( stage_id=0, model_stage="thinker", stage_type=StageType.LLM, runtime_overrides={ "pipeline_parallel_size": 2, "data_parallel_size": 3, "tensor_parallel_size": 8, }, ) omega_config = config.to_omegaconf() assert omega_config.engine_args.pipeline_parallel_size == 2 assert omega_config.engine_args.data_parallel_size == 3 assert omega_config.engine_args.tensor_parallel_size == 8 assert "pipeline_parallel_size" in omega_config.engine_args assert "data_parallel_size" in omega_config.engine_args assert "tensor_parallel_size" in omega_config.engine_args assert "parallel_config" not in omega_config.engine_args 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.""" @dataclass class MockParallelConfig: world_size: int = 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 orchestrator-owned, 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 TestStageResolutionHelpers: """Tests for shared stage override / filtering helpers.""" def test_build_stage_runtime_overrides_ignores_other_stage_and_internal_keys(self): # Pass the same filter set the function uses by default # (orchestrator-only fields plus SHARED_FIELDS so ``model`` is # treated as not-per-stage-overridable). overrides = build_stage_runtime_overrides( 0, { "gpu_memory_utilization": 0.5, "stage_0_gpu_memory_utilization": 0.9, "stage_1_gpu_memory_utilization": 0.1, "stage_0_model": "should_be_ignored", "parallel_config": {"world_size": 2}, }, internal_keys=internal_blacklist_keys() | SHARED_FIELDS, ) assert overrides["gpu_memory_utilization"] == 0.9 assert "model" not in overrides assert "parallel_config" not in overrides def test_strip_parent_engine_args_reports_only_surprising_parent_overrides(self): parent_fields = {f.name: f for f in fields(EngineArgs)} filtered, overridden = strip_parent_engine_args( { "model": "some/model", "stage_configs_path": "/tmp/stages.yaml", "tensor_parallel_size": 4, "worker_extension_cls": "some.Extension", "custom_pipeline_args": {"pipeline_class": "demo.Pipeline"}, }, parent_fields=parent_fields, keep_keys={"worker_extension_cls"}, strip_keys={"stage_configs_path"}, no_warn_keys={"model"}, ) assert filtered == { "worker_extension_cls": "some.Extension", "custom_pipeline_args": {"pipeline_class": "demo.Pipeline"}, } assert overridden == ["tensor_parallel_size"] def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self): parent_fields = {f.name: f for f in fields(EngineArgs)} filtered, overridden = strip_parent_engine_args( { "model": "some/model", "stage_configs_path": "/tmp/stages.yaml", "allowed_local_media_path": "/data/qwentts", "allowed_media_domains": ["example.com"], }, parent_fields=parent_fields, keep_keys={"allowed_local_media_path", "allowed_media_domains"}, strip_keys={"stage_configs_path"}, no_warn_keys={"model"}, ) assert filtered == { "allowed_local_media_path": "/data/qwentts", "allowed_media_domains": ["example.com"], } assert overridden == [] 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_yaml_migrates_runtime_max_batch_size(self, tmp_path): """Test that runtime.max_batch_size is migrated to engine_args.max_num_seqs.""" yaml_content = """\ model_type: legacy_batch stages: - stage_id: 0 model_stage: entry stage_type: llm input_sources: [] runtime: devices: "0" max_batch_size: 16 engine_args: model_arch: SomeModel """ yaml_file = tmp_path / "legacy_batch.yaml" yaml_file.write_text(yaml_content) pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_batch") s0 = pipeline.stages[0] assert "max_batch_size" not in s0.yaml_runtime assert s0.yaml_engine_args.get("max_num_seqs") == 16 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 TestPipelineDiscovery: """Tests for the central pipeline registry (``pipeline_registry._OMNI_PIPELINES``).""" def test_registry_has_known_models(self): """Built-in pipelines are lazy-loaded from the central declaration on first access; no eager import or discovery walk needed.""" # ``in`` triggers the lazy-map lookup without forcing a load. assert "qwen2_5_omni" in _PIPELINE_REGISTRY assert "qwen3_omni_moe" in _PIPELINE_REGISTRY assert "qwen3_tts" in _PIPELINE_REGISTRY def test_registry_loads_pipeline_on_getitem(self): """Looking up a registered model_type returns the matching PipelineConfig.""" pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] assert pipeline.model_type == "qwen3_omni_moe" assert len(pipeline.stages) == 3 # thinker + talker + code2wav def test_registry_returns_none_for_unknown(self): """Unknown model_types aren't found; ``get()`` returns None.""" assert "definitely_not_a_real_model" not in _PIPELINE_REGISTRY assert _PIPELINE_REGISTRY.get("definitely_not_a_real_model") is None def test_pipeline_config_supports_hf_architectures(self): """PipelineConfig accepts hf_architectures for HF-arch fallback (replaces the old _ARCHITECTURE_MODELS dict).""" p = PipelineConfig( model_type="custom_collide", hf_architectures=("SomeCollidingArch",), ) assert p.hf_architectures == ("SomeCollidingArch",) class TestStagePipelineConfig: def test_frozen(self): s = StagePipelineConfig(stage_id=0, model_stage="a") with pytest.raises(AttributeError): s.model_stage = "changed" def test_defaults(self): s = StagePipelineConfig(stage_id=0, model_stage="a") assert s.execution_type == StageExecutionType.LLM_AR assert s.input_sources == () assert s.final_output is False assert s.sampling_constraints == {} assert s.engine_output_type is None class TestPipelineConfigNew: def test_frozen(self): p = PipelineConfig(model_type="t", model_arch="A") with pytest.raises(AttributeError): p.model_type = "changed" def test_validate_valid(self): p = PipelineConfig( model_type="t", model_arch="A", stages=( StagePipelineConfig(stage_id=0, model_stage="a"), StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)), ), ) assert p.validate() == [] def test_validate_no_stages(self): p = PipelineConfig(model_type="t", model_arch="A") assert any("no stages" in e.lower() for e in p.validate()) class TestResolveScheduler: def test_all_execution_types_handled(self): for et in StageExecutionType: _resolve_scheduler(et) def test_ar_sync_when_false(self): cls = _resolve_scheduler(StageExecutionType.LLM_AR, async_scheduling=False) assert cls is not None assert "Async" not in cls.__name__ def test_ar_async_when_true(self): cls = _resolve_scheduler(StageExecutionType.LLM_AR, async_scheduling=True) assert cls is not None assert "Async" in cls.__name__ def test_generation(self): cls = _resolve_scheduler(StageExecutionType.LLM_GENERATION) assert cls is not None assert "Generation" in cls.__name__ def test_diffusion_returns_none(self): assert _resolve_scheduler(StageExecutionType.DIFFUSION) is None class TestPipelineRegistry: def test_register_and_lookup(self): p = PipelineConfig( model_type="__test_only__", model_arch="A", stages=(StagePipelineConfig(stage_id=0, model_stage="a"),), ) register_pipeline(p) assert _PIPELINE_REGISTRY["__test_only__"] is p del _PIPELINE_REGISTRY["__test_only__"] class TestDeployConfigLoading: def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path): deploy_path = tmp_path / "qwen3_tts_custom_voice.yaml" custom_voice_dir = tmp_path / "voices" deploy_path.write_text( f""" async_chunk: true custom_voice_dir: {custom_voice_dir} stages: - stage_id: 0 devices: "0" - stage_id: 1 devices: "0" """, encoding="utf-8", ) deploy = load_deploy_config(deploy_path) pipeline = _PIPELINE_REGISTRY["qwen3_tts"] stages = merge_pipeline_deploy(pipeline, deploy) assert deploy.custom_voice_dir == str(custom_voice_dir) assert {s.yaml_engine_args.get("custom_voice_dir") for s in stages} == {str(custom_voice_dir)} def test_load_qwen3_omni_moe_deploy_config(self): deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" deploy = load_deploy_config(deploy_path) assert len(deploy.stages) == 3 assert deploy.async_chunk is True assert deploy.connectors is not None assert deploy.platforms is not None def test_load_voxtral_tts_deploy_config_schema_fields(self): deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "voxtral_tts.yaml" deploy = load_deploy_config(deploy_path) assert deploy.stages[0].config_format == "mistral" assert deploy.stages[0].load_format == "mistral" assert deploy.stages[0].tokenizer_mode == "mistral" assert not any( name in deploy.stages[0].engine_extras for name in ("config_format", "load_format", "tokenizer_mode") ) def test_load_ming_flash_omni_deploy_config_schema_fields(self): deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml" deploy = load_deploy_config(deploy_path) assert deploy.stages[0].compilation_config == {"pass_config": {"fuse_allreduce_rms": False}} assert "compilation_config" not in deploy.stages[0].engine_extras def test_merge_pipeline_deploy(self): pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 3 s0 = stages[0] assert s0.model_stage == "thinker" assert s0.yaml_engine_args["model_arch"] == "Qwen3OmniMoeForConditionalGeneration" assert s0.yaml_engine_args["engine_output_type"] == "latent" assert s0.yaml_extras["default_sampling_params"]["detokenize"] is True def test_merge_pipeline_deploy_preserves_num_replicas(self, tmp_path): pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not base.exists(): pytest.skip("Deploy config not found") overlay = tmp_path / "multi_replicas.yaml" overlay.write_text(f'base_config: {base}\nstages:\n - stage_id: 1\n devices: "1,2"\n num_replicas: 2\n') deploy = load_deploy_config(overlay) assert deploy.stages[1].num_replicas == 2 stages = merge_pipeline_deploy(pipeline, deploy) assert stages[1].yaml_runtime["devices"] == "1,2" assert stages[1].yaml_runtime["num_replicas"] == 2 def test_merge_pipeline_deploy_preserves_requires_multimodal_data(self): pipeline = PipelineConfig( model_type="test_mm", model_arch="TestModel", stages=( StagePipelineConfig( stage_id=0, model_stage="ar", execution_type=StageExecutionType.LLM_AR, requires_multimodal_data=True, ), ), ) deploy = DeployConfig(async_chunk=False, stages=[StageDeployConfig(stage_id=0)]) stages = merge_pipeline_deploy(pipeline, deploy) assert stages[0].yaml_runtime["requires_multimodal_data"] is True def test_mixed_schema_preserves_flat_fields(self): """Ensure flat fields are not dropped when engine_args are present.""" fake_config = { "stages": [ { "stage_id": 0, "gpu_memory_utilization": 0.7, "max_num_seqs": 8, "engine_args": { "tensor_parallel_size": 4, "enforce_eager": True, }, }, ] } with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): deploy = load_deploy_config("dummy.yaml") assert len(deploy.stages) == 1 stage = deploy.stages[0] # Check that the engine args are set assert stage.tensor_parallel_size == 4 assert stage.enforce_eager is True # Check that the other settings are also preserved assert stage.gpu_memory_utilization == 0.7 assert stage.max_num_seqs == 8 def test_engine_parse_engine_fields(self): """Test that we correctly parse & recursively merge stage deploy fields.""" fake_config = { "stages": [ { "stage_id": 0, "compilation_config": { "encoder_cudagraph_token_budgets": [1024, 2048], "pass_config": {"fuse_norm_quant": True}, }, "engine_args": { "compilation_config": { "cudagraph_mm_encoder": True, "pass_config": {"fuse_allreduce_rms": False}, }, }, }, ] } with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): deploy = load_deploy_config("dummy.yaml") assert len(deploy.stages) == 1 stage = deploy.stages[0] assert stage.compilation_config == { "encoder_cudagraph_token_budgets": [1024, 2048], "pass_config": { "fuse_norm_quant": True, "fuse_allreduce_rms": False, }, "cudagraph_mm_encoder": True, } def test_engine_extras_deep_merges_dicts_simple(self): """Ensure dictionary valued keys merge properly for top level dicts.""" fake_config = { "stages": [ { "stage_id": 0, "foo": {"a": 111, "b": 1}, "engine_args": { "foo": {"b": 2, "c": 3}, }, }, ] } with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): deploy = load_deploy_config("dummy.yaml") assert len(deploy.stages) == 1 stage = deploy.stages[0] assert "foo" in stage.engine_extras assert stage.engine_extras["foo"] == {"a": 111, "b": 2, "c": 3} def test_engine_extras_dict_type_mismatch(self): """Ensure that we handle type mismatches with nested dicts correctly.""" fake_config = { "stages": [ { "stage_id": 0, "foo": {"b": {1: 1}}, "engine_args": { "foo": {"b": 2}, }, }, ] } with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): deploy = load_deploy_config("dummy.yaml") assert len(deploy.stages) == 1 stage = deploy.stages[0] assert "foo" in stage.engine_extras assert stage.engine_extras["foo"] == {"b": 2} def test_mixed_engine_extras_deep_merges_dicts(self): """Ensure dictionary valued keys merge properly for nested dicts.""" fake_config = { "stages": [ { "stage_id": 0, "foo": {"a": 111, "b": {"e": 199}}, "engine_args": { "foo": {"b": {"d": 9}, "c": 3}, }, }, ] } with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): deploy = load_deploy_config("dummy.yaml") assert len(deploy.stages) == 1 stage = deploy.stages[0] assert "foo" in stage.engine_extras assert stage.engine_extras["foo"] == {"a": 111, "b": {"d": 9, "e": 199}, "c": 3} def test_deep_merge_does_not_mutate_inputs(self): """Merging engine_args must not mutate the base stage dict.""" fake_config = { "stages": [ { "stage_id": 0, "foo": {"a": 1, "b": {"x": 10, "y": 20}}, "engine_args": { "foo": {"b": {"y": 99, "z": 30}, "c": 3}, }, }, ] } original_foo = fake_config["stages"][0]["foo"] original_b = original_foo["b"] with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): deploy = load_deploy_config("dummy.yaml") # Values should merge into the new dict correctly assert deploy.stages[0].engine_extras["foo"] == { "a": 1, "b": {"x": 10, "y": 99, "z": 30}, "c": 3, } # But the original nested dicts foo / b should be unchanged, since # we recursively shallow copy to avoid mutating in place. assert original_foo == {"a": 1, "b": {"x": 10, "y": 20}} assert original_b == {"x": 10, "y": 20} def test_mixed_schema_engine_args_wins_scalars(self): """engine_args takes precedence over flat fields for scalar conflicts.""" fake_config = { "stages": [ { "stage_id": 0, "gpu_memory_utilization": 0.7, "engine_args": { "gpu_memory_utilization": 0.5, }, }, ] } with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): deploy = load_deploy_config("dummy.yaml") assert deploy.stages[0].gpu_memory_utilization == 0.5 class TestQwen3OmniPipeline: def test_registered(self): p = _PIPELINE_REGISTRY.get("qwen3_omni_moe") assert p is not None assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration" assert len(p.stages) == 3 assert p.validate() == [] def test_thinker(self): s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(0) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.owns_tokenizer is True assert s.engine_output_type == "latent" assert s.sampling_constraints["detokenize"] is True def test_talker(self): s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(1) assert s.input_sources == (0,) assert s.sampling_constraints["stop_token_ids"] == [2150] assert s.custom_process_input_func is not None assert s.custom_process_next_stage_input_func is not None def test_code2wav(self): s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(2) assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.final_output_type == "audio" assert s.custom_process_input_func is not None class TestQwen2_5OmniPipeline: def test_registered(self): p = _PIPELINE_REGISTRY.get("qwen2_5_omni") assert p is not None assert p.model_arch == "Qwen2_5OmniForConditionalGeneration" assert len(p.stages) == 3 assert p.validate() == [] def test_thinker(self): s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(0) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.owns_tokenizer is True assert s.engine_output_type == "latent" assert s.requires_multimodal_data is True def test_talker(self): s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(1) assert s.input_sources == (0,) assert s.sampling_constraints["stop_token_ids"] == [8294] assert s.custom_process_input_func is not None def test_code2wav(self): s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(2) assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.final_output_type == "audio" assert s.engine_output_type == "audio" class TestQwen3TTSPipeline: def test_registered(self): p = _PIPELINE_REGISTRY.get("qwen3_tts") assert p is not None assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration" assert len(p.stages) == 2 assert p.validate() == [] def test_talker_stage(self): s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(0) assert s.model_stage == "qwen3_tts" assert s.execution_type == StageExecutionType.LLM_AR assert s.owns_tokenizer is True assert s.engine_output_type == "latent" assert s.sampling_constraints["stop_token_ids"] == [2150] # Stage 0 inherits the pipeline-level model_arch assert s.model_arch is None def test_code2wav_stage_has_per_stage_model_arch(self): s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(1) assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.final_output_type == "audio" assert s.engine_output_type == "audio" # Per-stage model_arch override (different from pipeline-level talker) assert s.model_arch == "Qwen3TTSCode2Wav" # tts_args is passed through via extras assert s.extras["tts_args"]["max_instructions_length"] == 500 def test_per_stage_model_arch_flows_through_merge(self, tmp_path): """Verify the new ps.model_arch override survives merge_pipeline_deploy.""" deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_tts.yaml" if not deploy_path.exists(): pytest.skip("qwen3_tts deploy yaml not found") deploy = load_deploy_config(deploy_path) pipeline = _PIPELINE_REGISTRY["qwen3_tts"] stages = merge_pipeline_deploy(pipeline, deploy) # Stage 0 inherits pipeline-level model_arch assert stages[0].yaml_engine_args["model_arch"] == "Qwen3TTSTalkerForConditionalGeneration" # Stage 1 uses its per-stage override assert stages[1].yaml_engine_args["model_arch"] == "Qwen3TTSCode2Wav" def test_subtalker_sampling_params_deep_merge_preserves_base_keys(self): """Verify subtalker sampling params participate in stage deep-merge.""" base = { "stage_id": 0, "subtalker_sampling_params": { "do_sample": True, "temperature": 0.9, "top_k": 50, "top_p": 1.0, }, } overlay = { "stage_id": 0, "subtalker_sampling_params": { "temperature": 0.7, "top_k": 32, }, } merged = _deep_merge_stage(base, overlay) assert merged["subtalker_sampling_params"] == { "do_sample": True, "temperature": 0.7, "top_k": 32, "top_p": 1.0, } class TestMingFlashOmniPipeline: def test_registered(self): p = _PIPELINE_REGISTRY.get("ming_flash_omni") assert p is not None assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 2 assert p.validate() == [] def test_thinker_stage(self): s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(0) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.owns_tokenizer is True assert s.requires_multimodal_data is True assert s.engine_output_type == "text" assert s.hf_config_name == "llm_config" assert s.sampling_constraints["detokenize"] is True def test_talker_stage(self): s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1) assert s.model_stage == "ming_tts" assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.input_sources == (0,) assert s.final_output_type == "audio" assert s.engine_output_type == "audio" assert s.hf_config_name == "talker_config" # Per-stage model_arch override (Ming talker has its own self-contained LLM) assert s.model_arch == "MingFlashOmniTalkerForConditionalGeneration" assert s.tokenizer_subdir == "talker/llm" assert s.custom_process_input_func is not None def test_talker_stage_processor_wiring_resolves(self): """The custom_process_input_func string must point to a real callable. Lazy string references only fail at first inference otherwise — this catches typos in the pipeline declaration at import / registration time. """ s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1) module_path, _, attr = s.custom_process_input_func.rpartition(".") module = importlib.import_module(module_path) assert callable(getattr(module, attr)) def test_tts_pipeline_registered(self): p = _PIPELINE_REGISTRY.get("ming_flash_omni_tts") assert p is not None assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration" assert len(p.stages) == 1 assert p.validate() == [] def test_tts_stage(self): s = _PIPELINE_REGISTRY["ming_flash_omni_tts"].get_stage(0) assert s.model_stage == "ming_tts" assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.input_sources == () assert s.owns_tokenizer is True assert s.final_output_type == "audio" assert s.engine_output_type == "audio" assert s.hf_config_name == "talker_config" assert s.tokenizer_subdir == "talker/llm" def test_full_yaml_loads_and_merges(self): """deploy/ming_flash_omni.yaml parses and merges with the registered pipeline.""" deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml" if not deploy_path.exists(): pytest.skip("ming_flash_omni deploy yaml not found") deploy = load_deploy_config(deploy_path) assert len(deploy.stages) == 2 assert deploy.async_chunk is False assert deploy.pipeline == "ming_flash_omni" # We won't test stage 0/1 colocation contract here, # as there could exist more variant of custom device setup pipeline = _PIPELINE_REGISTRY["ming_flash_omni"] stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 2 assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration" assert stages[1].yaml_engine_args["model_arch"] == "MingFlashOmniTalkerForConditionalGeneration" def test_tts_yaml_loads_and_merges(self): """deploy/ming_flash_omni_tts.yaml parses and routes to the TTS-only pipeline.""" deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_tts.yaml" if not deploy_path.exists(): pytest.skip("ming_flash_omni_tts deploy yaml not found") deploy = load_deploy_config(deploy_path) assert len(deploy.stages) == 1 assert deploy.pipeline == "ming_flash_omni_tts" pipeline = _PIPELINE_REGISTRY["ming_flash_omni_tts"] stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 1 assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniTalkerForConditionalGeneration" def test_thinker_only_pipeline_registered(self): p = _PIPELINE_REGISTRY.get("ming_flash_omni_thinker_only") assert p is not None assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 1 assert p.validate() == [] def test_thinker_only_stage(self): s = _PIPELINE_REGISTRY["ming_flash_omni_thinker_only"].get_stage(0) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.input_sources == () assert s.owns_tokenizer is True assert s.requires_multimodal_data is True assert s.final_output_type == "text" assert s.engine_output_type == "text" assert s.hf_config_name == "llm_config" assert s.sampling_constraints["detokenize"] is True def test_thinker_only_yaml_loads_and_merges(self): """deploy/ming_flash_omni_thinker_only.yaml parses and routes to the thinker-only pipeline.""" deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_thinker_only.yaml" if not deploy_path.exists(): pytest.skip("ming_flash_omni_thinker_only deploy yaml not found") deploy = load_deploy_config(deploy_path) assert len(deploy.stages) == 1 assert deploy.pipeline == "ming_flash_omni_thinker_only" pipeline = _PIPELINE_REGISTRY["ming_flash_omni_thinker_only"] stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 1 assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration" def test_image_pipeline_registered(self): p = _PIPELINE_REGISTRY.get("ming_flash_omni_image") assert p is not None assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 2 assert p.validate() == [] def test_image_thinker_stage(self): s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.input_sources == () assert s.final_output is False assert s.owns_tokenizer is True assert s.requires_multimodal_data is True # Image variant exports hidden states for the diffusion stage. assert s.engine_output_type == "latent" assert s.hf_config_name == "thinker_config" assert s.sampling_constraints["detokenize"] is False assert s.prompt_expand_func is not None def test_image_dit_stage(self): s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1) assert s.model_stage == "dit" assert s.execution_type == StageExecutionType.DIFFUSION assert s.input_sources == (0,) assert s.final_output is True assert s.final_output_type == "image" assert s.hf_config_name == "image_gen_config" assert s.model_arch == "MingImagePipeline" assert s.custom_process_input_func is not None def test_image_processor_wiring_resolves(self): """The prompt_expand_func and custom_process_input_func strings must point to real callables.""" thinker = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0) dit = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1) for ref in (thinker.prompt_expand_func, dit.custom_process_input_func): module_path, _, attr = ref.rpartition(".") module = importlib.import_module(module_path) assert callable(getattr(module, attr)) def test_image_yaml_loads_and_merges(self): """deploy/ming_flash_omni_image.yaml parses and routes to the image pipeline.""" deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_image.yaml" if not deploy_path.exists(): pytest.skip("ming_flash_omni_image deploy yaml not found") deploy = load_deploy_config(deploy_path) assert len(deploy.stages) == 2 assert deploy.async_chunk is False assert deploy.pipeline == "ming_flash_omni_image" assert deploy.connectors is not None assert "shared_memory_connector" in deploy.connectors pipeline = _PIPELINE_REGISTRY["ming_flash_omni_image"] stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 2 # Stage 0 thinker: AR worker that emits latents. assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration" assert stages[0].yaml_engine_args["engine_output_type"] == "latent" assert stages[0].yaml_extras["default_sampling_params"]["detokenize"] is False assert stages[0].yaml_extras["prompt_expand_func"] is not None # Stage 1 dit: diffusion stage with MingImagePipeline. assert stages[1].yaml_engine_args["model_arch"] == "MingImagePipeline" assert stages[1].custom_process_input_func is not None assert stages[1].final_output is True assert stages[1].final_output_type == "image" class TestBaseConfigInheritance: """Test deploy YAML base_config inheritance.""" def test_ci_inherits_from_main(self): ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe.yaml")) if not ci_path.exists(): pytest.skip("CI deploy config not found") deploy = load_deploy_config(ci_path) assert len(deploy.stages) == 3 # CI overrides assert deploy.stages[0].load_format is None assert "load_format" not in deploy.stages[0].engine_extras assert deploy.stages[0].max_num_seqs == 5 # Inherited from base assert deploy.stages[0].gpu_memory_utilization == 0.9 assert deploy.connectors is not None assert "connector_of_shared_memory" in deploy.connectors # CI overlay explicitly sets async_chunk: False (see # tests.helpers.stage_config._CI_OVERLAYS and PR #2383 discussion). Overlay # bool overrides base even when the base yaml has async_chunk: true. assert deploy.async_chunk is False def test_ci_sampling_merge(self): ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe.yaml")) if not ci_path.exists(): pytest.skip("CI deploy config not found") deploy = load_deploy_config(ci_path) s0 = deploy.stages[0].default_sampling_params # CI overrides max_tokens assert s0["max_tokens"] == 150 def test_pure_inheritance_overlay(self, tmp_path): """An overlay with only ``base_config`` inherits everything.""" base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not base.exists(): pytest.skip("Base deploy config not found") overlay = tmp_path / "overlay.yaml" overlay.write_text(f"base_config: {base}\n") deploy = load_deploy_config(overlay) assert len(deploy.stages) == 3 assert deploy.stages[0].gpu_memory_utilization == 0.9 def test_single_field_overlay(self, tmp_path): """An overlay overriding one stage field merges with the base.""" base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not base.exists(): pytest.skip("Base deploy config not found") overlay = tmp_path / "overlay.yaml" overlay.write_text(f"base_config: {base}\nstages:\n - stage_id: 2\n max_num_batched_tokens: 1000000\n") deploy = load_deploy_config(overlay) assert deploy.stages[2].max_num_batched_tokens == 1000000 # Rest inherited assert deploy.stages[0].gpu_memory_utilization == 0.9 class TestPlatformOverrides: """Test platform-specific deploy config overrides.""" def test_npu_overrides(self): deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) deploy = _apply_platform_overrides(deploy, platform="npu") assert deploy.stages[0].gpu_memory_utilization == 0.6 assert deploy.stages[0].tensor_parallel_size == 2 assert deploy.stages[0].devices == "0,1" # Stage 2 unaffected fields stay at base assert deploy.stages[2].enforce_eager is False def test_xpu_overrides(self): deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) deploy = _apply_platform_overrides(deploy, platform="xpu") assert deploy.stages[0].tensor_parallel_size == 4 assert deploy.stages[0].devices == "0,1,2,3" assert deploy.stages[0].engine_extras.get("max_cudagraph_capture_size") == 0 def test_unknown_platform_noop(self): deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) original_mem = deploy.stages[0].gpu_memory_utilization deploy = _apply_platform_overrides(deploy, platform="unknown_hw") assert deploy.stages[0].gpu_memory_utilization == original_mem def test_platforms_deep_merge_inheritance(self, tmp_path): """Overlay's platforms: block layers onto base's, per-stage.""" base = tmp_path / "base.yaml" base.write_text( "stages:\n" " - stage_id: 0\n" " gpu_memory_utilization: 0.9\n" "platforms:\n" " rocm:\n" " stages:\n" " - stage_id: 0\n" " enforce_eager: true\n" ) overlay = tmp_path / "overlay.yaml" overlay.write_text( f"base_config: {base.name}\n" "platforms:\n" " rocm:\n" " stages:\n" " - stage_id: 0\n" " max_num_seqs: 1\n" ) deploy = load_deploy_config(overlay) deploy = _apply_platform_overrides(deploy, platform="rocm") # Both base's enforce_eager and overlay's max_num_seqs should apply. assert deploy.stages[0].enforce_eager is True assert deploy.stages[0].max_num_seqs == 1 # Inherited stage default not touched by overlay platforms section. assert deploy.stages[0].gpu_memory_utilization == 0.9 class TestCLIOverrideFlow: """Test --stage-overrides JSON merge into StageConfig.""" def test_stage_overrides_merge(self): pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) stages = merge_pipeline_deploy(pipeline, deploy) # Simulate --stage-overrides '{"0": {"gpu_memory_utilization": 0.5}}' overrides = {"stage_0_gpu_memory_utilization": 0.5} stages[0].runtime_overrides = StageConfigFactory._merge_cli_overrides(stages[0], overrides) assert stages[0].runtime_overrides["gpu_memory_utilization"] == 0.5 def test_global_override_applies_to_all(self): pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) stages = merge_pipeline_deploy(pipeline, deploy) overrides = {"enforce_eager": True} for s in stages: s.runtime_overrides = StageConfigFactory._merge_cli_overrides(s, overrides) assert s.runtime_overrides["enforce_eager"] is True class TestSentinelDefaultPrecedence: """Caller-typed (non-None) values win over YAML; None values fall through to YAML / dataclass defaults (#3035).""" def _stages(self, cli_overrides): return StageConfigFactory._create_from_registry( "qwen3_omni_moe", cli_overrides=cli_overrides, ) def test_typed_kwarg_overrides_yaml(self): stages = self._stages({"max_num_seqs": 999}) assert stages[2].runtime_overrides.get("max_num_seqs") == 999 def test_none_value_skipped_yaml_wins(self): stages = self._stages({"max_num_seqs": None}) assert stages[2].runtime_overrides.get("max_num_seqs") is None assert stages[2].yaml_engine_args.get("max_num_seqs") == 64 def test_empty_kwargs_yaml_only(self): stages = self._stages({}) for stage in stages: assert stage.runtime_overrides == {} def test_typed_kwarg_equal_to_dataclass_default_still_overrides(self): # Caller intent honored regardless of value coincidence (no heuristic). stages = self._stages({"gpu_memory_utilization": 0.9}) assert stages[2].runtime_overrides.get("gpu_memory_utilization") == 0.9 def test_per_stage_kwarg_routed_to_correct_stage(self): stages = self._stages({"stage_0_gpu_memory_utilization": 0.42}) assert stages[0].runtime_overrides.get("gpu_memory_utilization") == 0.42 assert stages[2].runtime_overrides.get("gpu_memory_utilization") is None def test_async_chunk_false_overrides_yaml_true(self): stages = self._stages({"async_chunk": False}) for stage in stages: assert stage.yaml_engine_args.get("async_chunk") is not True def test_async_chunk_none_keeps_yaml_true(self): stages = self._stages({"async_chunk": None}) for stage in stages: assert stage.yaml_engine_args.get("async_chunk") is True def test_enable_prefix_caching_typed_overrides_yaml(self): stages = self._stages({"enable_prefix_caching": True}) for stage in stages: assert stage.runtime_overrides.get("enable_prefix_caching") is True def test_omni_with_vars_args_anti_pattern_is_safe(self): # Omni(**vars(args)) with mostly-None namespace must not clobber YAML. simulated_vars_args = { "gpu_memory_utilization": None, "max_num_seqs": None, "async_chunk": None, "enable_prefix_caching": None, "dtype": None, } stages = self._stages(simulated_vars_args) for stage in stages: assert stage.runtime_overrides == {} def test_create_from_registry_no_cli_explicit_keys_param(self): sig = inspect.signature(StageConfigFactory._create_from_registry) named = [p for p in sig.parameters.values() if p.kind != p.VAR_KEYWORD] assert "cli_explicit_keys" not in {p.name for p in named} def test_async_chunk_dispatches_processors(self): """A single ``qwen3_tts`` pipeline picks per-chunk vs end-to-end processors based on ``deploy.async_chunk``, without needing a separate variant pipeline registration.""" pipeline = _PIPELINE_REGISTRY["qwen3_tts"] # async_chunk=True → stage 0's per-chunk processor wires up, stage 1 # has no sync input processor. async_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=True)) assert ( async_stages[0] .yaml_engine_args.get("custom_process_next_stage_input_func", "") .endswith("talker2code2wav_async_chunk") ) assert async_stages[1].custom_process_input_func is None # async_chunk=False → stage 0 ships the bulk codec via the # worker-connector full-payload producer; stage 1 wires the # ``_token_only`` placeholder so the orchestrator emits no # legacy ``additional_information``-shaped input (PR3 sync- # via-connector data plane). sync_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=False)) assert ( sync_stages[0] .yaml_engine_args["custom_process_next_stage_input_func"] .endswith("talker2code2wav_full_payload") ) assert sync_stages[1].custom_process_input_func is not None assert sync_stages[1].custom_process_input_func.endswith("talker2code2wav_token_only") def test_async_chunk_dispatches_qwen3_omni_processors(self): import runpy from pathlib import Path from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy pipeline_path = ( Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "models" / "qwen3_omni" / "pipeline.py" ) pipeline = runpy.run_path(str(pipeline_path))["QWEN3_OMNI_PIPELINE"] async_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=True)) assert ( async_stages[0] .yaml_engine_args["custom_process_next_stage_input_func"] .endswith("thinker2talker_async_chunk") ) assert ( async_stages[1] .yaml_engine_args["custom_process_next_stage_input_func"] .endswith("talker2code2wav_async_chunk") ) sync_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=False)) assert ( sync_stages[0] .yaml_engine_args["custom_process_next_stage_input_func"] .endswith("thinker2talker_full_payload") ) assert ( sync_stages[1] .yaml_engine_args["custom_process_next_stage_input_func"] .endswith("talker2code2wav_full_payload") ) def test_ming_flash_omni_topology(self): """Guard ming_flash_omni's PR3 cleanup: stage 0 has no full-payload producer hook (the connector path was removed as fake -- arch is not in ``_FULL_PAYLOAD_INPUT_STAGES``), and stage 1 still wires the legacy ``thinker2talker`` (custom_process_input_func) plus the ``thinker2talker_token_only`` placeholder (sync_process_input_func). Merge under either async_chunk mode must not re-introduce a stage-0 full-payload hook.""" from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy pipeline = _PIPELINE_REGISTRY["ming_flash_omni"] stage0, stage1 = pipeline.stages assert stage0.custom_process_next_stage_input_func is None, ( "ming_flash_omni stage 0 must not declare a full-payload producer " "(connector path is not active for this arch)." ) assert stage1.custom_process_input_func is not None assert stage1.custom_process_input_func.endswith("thinker2talker") assert stage1.sync_process_input_func is not None assert stage1.sync_process_input_func.endswith("thinker2talker_token_only") # async_chunk=True must now be rejected: removing the fake hook means # there is no next-stage input processor for the validator to accept. # (Positive consequence -- users can't accidentally enable async_chunk # on an arch that doesn't actually support it.) import pytest as _pytest with _pytest.raises(ValueError, match="async_chunk=True"): merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=True)) # async_chunk=False merges cleanly and stage-0 yaml_engine_args carries # no spurious full-payload hook. merged = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=False)) assert "custom_process_next_stage_input_func" not in merged[0].yaml_engine_args, ( "stage-0 full-payload hook unexpectedly re-appeared in yaml_engine_args" ) class TestSamplingConstraintsPrecedence: """Test that pipeline sampling_constraints override deploy defaults.""" def test_constraints_win(self): pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) stages = merge_pipeline_deploy(pipeline, deploy) # Pipeline says detokenize=True for thinker, deploy can't override assert stages[0].yaml_extras["default_sampling_params"]["detokenize"] is True # Pipeline says stop_token_ids=[2150] for talker assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [2150]