# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Unit tests for StageConfigFactory and related classes. """ import importlib import warnings from pathlib import Path from unittest.mock import patch import pytest from transformers import PretrainedConfig, Qwen3OmniMoeConfig from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage from vllm_omni.config import config_factory as config_factory_module from vllm_omni.config.config_factory import StageConfigFactory, _materialize_object_storage_configs from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability from vllm_omni.config.omni_config import VllmOmniConfig from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline, resolve_pipeline_config from vllm_omni.config.stage_config import ( _DEPLOY_DIR, DeployConfig, PipelineConfig, StageConfig, StageDeployConfig, StageExecutionType, StagePipelineConfig, StageType, _apply_platform_overrides, _deep_merge_stage, _resolve_scheduler, build_stage_runtime_overrides, load_deploy_config, merge_pipeline_deploy, normalize_pipeline_cli_overrides, pipeline_cfg_resolver, ) from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys pytestmark = [pytest.mark.core_model, pytest.mark.cpu] SINGLE_STAGE_PIPE_CFG = (StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),) # Validation strings to match against in post init NO_STAGES_MATCH_STR = "no stages" NO_TERMINAL_STAGE_MATCH_STR = "No terminal stage" @pytest.fixture(autouse=True) def _stable_test_platform(monkeypatch): from vllm_omni import platforms platform = platforms.current_omni_platform monkeypatch.setattr(platform, "device_name", "cpu", raising=False) monkeypatch.setattr(platform, "device_type", "cpu", raising=False) @pytest.fixture(autouse=True) def clear_config_factory_caches(): """Clear cached classmethods from the StageConfigFactory to prevent test pollution.""" yield StageConfigFactory.get_hf_config.cache_clear() StageConfigFactory.try_infer_model_type.cache_clear() _materialize_object_storage_configs.cache_clear() Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True) Q3_OMNI_THINKER_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=False) 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 == [] assert omega_config.session_mode == "turn" def test_to_omegaconf_duplex_session_mode(self): """Test that session mode is preserved as stage metadata.""" config = StageConfig( stage_id=0, model_stage="thinker", session_mode="duplex", ) omega_config = config.to_omegaconf() assert omega_config.session_mode == "duplex" 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_dict_override_deep_merges_yaml_dict(self): """A partial dict CLI override (e.g. --no-guardrails riding on model_config) must layer onto the YAML dict, not clobber siblings.""" config = StageConfig( stage_id=0, model_stage="diffusion", stage_type=StageType.DIFFUSION, yaml_engine_args={ "model_config": { "guardrails": True, "policy_server_config": {"action_space": "joint_position"}, } }, runtime_overrides={"model_config": {"guardrails": False}}, ) omega_config = config.to_omegaconf() assert omega_config.engine_args.model_config.guardrails is False assert omega_config.engine_args.model_config.policy_server_config.action_space == "joint_position" 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_allgather_override_recomputes_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, "allgather_degree": 1, } }, runtime_overrides={ "allgather_degree": 2, }, ) omega_config = config.to_omegaconf() assert omega_config.engine_args.parallel_config.allgather_degree == 2 assert omega_config.engine_args.parallel_config.sequence_parallel_size == 2 assert "allgather_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 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" @pytest.mark.parametrize( "model_class_name", [ "MiniMaxH3Pipeline", "WanPipeline", "LTX2Pipeline", "Cosmos3OmniDiffusersPipeline", ], ) def test_default_video_diffusion_output_type(self, model_class_name): configs = StageConfigFactory.create_default_diffusion({"model_class_name": model_class_name}) assert configs[0]["final_output_type"] == "video" def test_default_diffusion_with_parallel_config(self): """Test diffusion config calculates devices from parallel_config.""" kwargs = { "parallel_config": DiffusionParallelConfig(tensor_parallel_size=4), "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 "deploy_config": "/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 "deploy_config" 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 class TestPipelineDiscovery: """Tests for the central pipeline registry (``OMNI_PIPELINES``).""" def test_registry_has_known_models(self): """Check that specific models are in OMNI_PIPELINES.""" assert "qwen2_5_omni" in OMNI_PIPELINES assert "qwen3_omni_moe" in OMNI_PIPELINES assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES assert "qwen3_tts" in OMNI_PIPELINES def test_registry_resolver_qwen3_omni_all_stages(self): """Test that providing the HF config for qwen3 omni with audio enabled uses all stages.""" pipeline = resolve_pipeline_config( "qwen3_omni_moe", Q3_OMNI_ALL_STAGES_HF_CONFIG, ) assert isinstance(pipeline, PipelineConfig) assert pipeline.model_type == "qwen3_omni_moe" assert pipeline.default_deploy_config_name == "qwen3_omni_moe.yaml" assert len(pipeline.stages) == 3 # thinker + talker + code2wav def test_registry_resolver_qwen3_omni_thinker_only(self): """Test that providing the HF config for qwen3 omni without audio is thinker only.""" pipeline = resolve_pipeline_config( "qwen3_omni_moe", Q3_OMNI_THINKER_HF_CONFIG, ) assert isinstance(pipeline, PipelineConfig) assert pipeline.model_type == "qwen3_omni_moe_thinker_only" assert pipeline.default_deploy_config_name is None assert len(pipeline.stages) == 1 # thinker only def test_registry_qwen3_omni_thinker_only_key_is_static(self): """Explicit registry key must not re-enter the HF-config resolver.""" pipeline = resolve_pipeline_config( "qwen3_omni_moe_thinker_only", Q3_OMNI_ALL_STAGES_HF_CONFIG, ) assert pipeline is OMNI_PIPELINES["qwen3_omni_moe_thinker_only"] assert pipeline.model_type == "qwen3_omni_moe_thinker_only" assert len(pipeline.stages) == 1 assert pipeline.stages[0].engine_output_type == "text" assert pipeline.stages[0].final_output_type == "text" @pytest.mark.parametrize( "pipeline", [registered for registered in OMNI_PIPELINES.values() if isinstance(registered, PipelineConfig)], ids=lambda pipeline: pipeline.model_type, ) def test_registered_default_deploy_config_exists(self, pipeline): if pipeline.default_deploy_config_name is None: return assert Path(get_deploy_config_path(pipeline.default_deploy_config_name)).is_file() @pytest.mark.parametrize( "deploy_path", sorted(_DEPLOY_DIR.glob("*.yaml")), ids=lambda deploy_path: deploy_path.name, ) def test_explicit_deploy_pipeline_is_registered(self, deploy_path): deploy = load_deploy_config(deploy_path) if deploy.pipeline is not None: assert deploy.pipeline in OMNI_PIPELINES def test_registry_returns_none_for_unknown(self): """Unknown model_types aren't found and resolve to `None`.""" assert "definitely_not_a_real_model" not in OMNI_PIPELINES assert OMNI_PIPELINES.get("definitely_not_a_real_model") is None assert resolve_pipeline_config("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",), stages=SINGLE_STAGE_PIPE_CFG, ) assert p.hf_architectures == ("SomeCollidingArch",) class TestCosmos3PolicyPipeline: """Cosmos3 policy serving resolves via deploy yaml selection (π0 precedent).""" def test_registered_without_capturing_base_cosmos3_checkpoints(self): assert "cosmos3_policy" in OMNI_PIPELINES # T2I/video Cosmos3 checkpoints report model_type=cosmos3_omni and the # same model_index.json _class_name as policy checkpoints; they must # keep resolving through the single-stage diffusion fallback, so the # policy pipeline must not be reachable by auto-detection. assert "cosmos3_omni" not in OMNI_PIPELINES pipeline = OMNI_PIPELINES["cosmos3_policy"] assert pipeline.hf_architectures == () assert pipeline.diffusers_class_name is None def test_droid_deploy_yaml_carries_policy_server_config(self): deploy = load_deploy_config(get_deploy_config_path("cosmos3_policy_droid.yaml")) assert deploy.pipeline == "cosmos3_policy" stages = merge_pipeline_deploy(OMNI_PIPELINES["cosmos3_policy"], deploy) assert len(stages) == 1 stage = stages[0].to_omegaconf() assert stage.stage_type == "diffusion" assert stage.final_output_type == "action" assert stage.engine_args.model_class_name == "Cosmos3OmniDiffusersPipeline" # OpenPI websocket handshake metadata must reach the serving layer via # engine_args.model_config (ServingRealtimeRobotOpenPI reads it there). psc = stage.engine_args.model_config.policy_server_config assert list(psc.image_resolution) == [540, 640] assert psc.n_external_cameras == 2 assert psc.needs_wrist_camera is True assert psc.needs_stereo_camera is False assert psc.needs_session_id is True assert psc.action_space == "joint_position" # The DROID checkpoint's training recipe expects JSON-formatted prompts. assert stage.default_sampling_params.extra_args.format_prompt_as_json is True def test_no_guardrails_cli_flag_keeps_policy_server_config(self): """--no-guardrails becomes a partial model_config CLI override; it must deep-merge with the deploy yaml's model_config, not clobber the policy_server_config the OpenPI handshake depends on.""" deploy = load_deploy_config(get_deploy_config_path("cosmos3_policy_droid.yaml")) stages, _ = StageConfigFactory._create_legacy_from_registry( OMNI_PIPELINES["cosmos3_policy"], {"model_config": {"guardrails": False}}, user_deploy_config=deploy, ) stage = stages[0].to_omegaconf() assert stage.engine_args.model_config.guardrails is False assert stage.engine_args.model_config.policy_server_config.action_space == "joint_position" 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 assert s.scheduler_cls is None class TestPipelineConfigNew: def test_simple_init(self): # Ensure pipeline config validates in post init. PipelineConfig( model_type="t", model_arch="A", stages=SINGLE_STAGE_PIPE_CFG, ) def test_frozen(self): p = PipelineConfig( model_type="t", model_arch="A", stages=SINGLE_STAGE_PIPE_CFG, ) with pytest.raises(AttributeError): p.model_type = "changed" def test_validate_no_stages(self): with pytest.raises(ValueError, match=NO_STAGES_MATCH_STR): PipelineConfig(model_type="t", model_arch="A") def test_validate_no_terminal_stage(self): """A pipeline with no ``final_output`` stage can never emit a result.""" with pytest.raises(ValueError, match=NO_TERMINAL_STAGE_MATCH_STR): 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,)), ), ) def test_pipeline_can_have_final_output_in_any_stage(self): """Ensure any stage may carry ``final_output``.""" PipelineConfig( model_type="t", model_arch="A", stages=( StagePipelineConfig(stage_id=0, model_stage="a", final_output=True), StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)), ), ) class TestPipelineRegistration: def test_resolve_pipeline_prefers_deploy_pipeline_key(self, clean_pipeline_registry, tmp_path): deploy_key = "deploy_selected_pipeline" model_type_key = "hf_model_type_pipeline" deploy_pipe = PipelineConfig( model_type=deploy_key, stages=SINGLE_STAGE_PIPE_CFG, ) model_type_pipe = PipelineConfig( model_type=model_type_key, stages=SINGLE_STAGE_PIPE_CFG, ) register_pipeline(deploy_pipe) register_pipeline(model_type_pipe) deploy_path = tmp_path / "deploy.yaml" deploy_path.write_text(f"pipeline: {deploy_key}\n", encoding="utf-8") class FakeConfig(PretrainedConfig): model_type = model_type_key fake_config = FakeConfig() with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config): pipeline_cfg = StageConfigFactory.get_pipeline_config( "fake/model", trust_remote_code=True, deploy_config_path=str(deploy_path), ) assert pipeline_cfg is deploy_pipe def test_resolve_pipeline_rejects_invalid_deploy_pipeline_key( self, clean_pipeline_registry, tmp_path, ): model_type_key = "registered_model_type_pipeline" pipeline_cfg = PipelineConfig( model_type=model_type_key, stages=SINGLE_STAGE_PIPE_CFG, ) register_pipeline(pipeline_cfg) deploy_path = tmp_path / "deploy.yaml" deploy_path.write_text("pipeline: missing_pipeline\n", encoding="utf-8") class FakeConfig(PretrainedConfig): model_type = model_type_key with patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()): with pytest.raises(KeyError, match="missing_pipeline"): StageConfigFactory.get_pipeline_config( "fake/model", trust_remote_code=True, deploy_config_path=str(deploy_path), ) with pytest.raises(KeyError, match="missing_pipeline"): StageConfigFactory.create_from_model( "fake/model", trust_remote_code=False, cli_overrides={}, deploy_config_path=str(deploy_path), ) with pytest.raises(KeyError, match="missing_pipeline"): StageConfigFactory.create_legacy_stage_configs_from_model( "fake/model", trust_remote_code=False, cli_overrides={}, deploy_config_path=str(deploy_path), ) def test_resolve_pipeline_matches_hf_architecture_fallback(self, clean_pipeline_registry): pipeline_key = "architecture_fallback_pipeline" pipe_cfg = PipelineConfig( model_type=pipeline_key, hf_architectures=("ArchitectureFallbackForTest",), stages=SINGLE_STAGE_PIPE_CFG, ) register_pipeline(pipe_cfg) class FakeConfig(PretrainedConfig): model_type = "unregistered_model_type" fake_config = FakeConfig() fake_config.architectures = ["ArchitectureFallbackForTest"] with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config): pipeline_cfg = StageConfigFactory.get_pipeline_config( "fake/model", trust_remote_code=True, deploy_config_path=None, ) assert pipeline_cfg is pipe_cfg def test_resolve_pipeline_architecture_fallback_respects_predicates(self, clean_pipeline_registry): shared_arch = "SharedPredicateArchitectureForTest" def rejecting_predicate(_hf_config): return False def raising_predicate(_hf_config): raise ValueError("predicate failure") reject_cfg = PipelineConfig( model_type="rejected_by_predicate", hf_architectures=(shared_arch,), hf_config_predicate=rejecting_predicate, stages=SINGLE_STAGE_PIPE_CFG, ) raises_cfg = PipelineConfig( model_type="raises_in_predicate", hf_architectures=(shared_arch,), hf_config_predicate=raising_predicate, stages=SINGLE_STAGE_PIPE_CFG, ) accept_cfg = PipelineConfig( model_type="accepted_by_predicate", hf_architectures=(shared_arch,), stages=SINGLE_STAGE_PIPE_CFG, ) register_pipeline(reject_cfg) register_pipeline(raises_cfg) register_pipeline(accept_cfg) class FakeConfig(PretrainedConfig): model_type = "unregistered_model_type" fake_config = FakeConfig() fake_config.architectures = [shared_arch] with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config): pipeline_cfg = StageConfigFactory.get_pipeline_config( "fake/model", trust_remote_code=True, deploy_config_path=None, ) assert pipeline_cfg is accept_cfg def test_resolve_pipeline_architecture_fallback_rejects_non_matching_architecture( self, clean_pipeline_registry, ): predicate_calls = [] def predicate(_hf_config): predicate_calls.append(_hf_config) return True pipe_cfg = PipelineConfig( model_type="predicate_without_arch_match", hf_architectures=("DifferentArchitectureForTest",), hf_config_predicate=predicate, stages=SINGLE_STAGE_PIPE_CFG, ) register_pipeline(pipe_cfg) class FakeConfig(PretrainedConfig): model_type = "unregistered_model_type" fake_config = FakeConfig() fake_config.architectures = ["RequestedArchitectureForTest"] with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config): pipeline_cfg = StageConfigFactory.get_pipeline_config( "fake/model", trust_remote_code=True, deploy_config_path=None, ) assert pipeline_cfg is None assert predicate_calls == [fake_config] def test_resolve_pipeline_architecture_fallback_resolves_callable_pipeline(self, clean_pipeline_registry): pipeline_key = "callable_architecture_fallback" resolved_cfg = PipelineConfig( model_type="callable_resolved_pipeline", hf_architectures=("CallableArchitectureForTest",), stages=SINGLE_STAGE_PIPE_CFG, ) seen_hf_configs = [] def resolver(hf_config): seen_hf_configs.append(hf_config) return resolved_cfg register_pipeline(resolver, model_type=pipeline_key) class FakeConfig(PretrainedConfig): model_type = "unregistered_model_type" fake_config = FakeConfig() fake_config.architectures = ["CallableArchitectureForTest"] with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config): pipeline_cfg = StageConfigFactory.get_pipeline_config( "fake/model", trust_remote_code=True, deploy_config_path=None, ) assert pipeline_cfg is resolved_cfg assert seen_hf_configs == [fake_config] def test_resolve_pipeline_returns_none_without_registry_or_architecture_match(self, clean_pipeline_registry): class FakeConfig(PretrainedConfig): model_type = "unregistered_model_type" fake_config = FakeConfig() fake_config.architectures = ["UnmatchedArchitectureForTest"] with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config): pipeline_cfg = StageConfigFactory.get_pipeline_config( "fake/model", trust_remote_code=True, deploy_config_path=None, ) assert pipeline_cfg is None def test_create_from_model_returns_structured_omni_config(self): class FakeConfig(PretrainedConfig): model_type = "qwen3_tts" with patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()): omni_config = StageConfigFactory.create_from_model( "fake/model", trust_remote_code=False, cli_overrides={}, deploy_config_path=None, ) assert isinstance(omni_config, VllmOmniConfig) assert omni_config.pipeline_config is OMNI_PIPELINES["qwen3_tts"] assert len(omni_config.stage_configs) == 2 def test_create_from_model_preserves_model_on_structured_diffusion_stage(self): class FakeConfig(PretrainedConfig): model_type = "dreamzero" with patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()): omni_config = StageConfigFactory.create_from_model( "fake/model", trust_remote_code=False, cli_overrides={}, deploy_config_path=None, ) assert isinstance(omni_config, VllmOmniConfig) assert omni_config.stage_by_id(0).diffusion_config.model == "fake/model" def test_cosyvoice3_deploy_resolves_hash_snapshot_without_hf_metadata(self, tmp_path): snapshot = tmp_path / "models--FunAudioLLM--Fun-CosyVoice3-0.5B-2512" / "snapshots" / ("a" * 40) snapshot.mkdir(parents=True) (snapshot / "config.json").write_text("{}\n", encoding="utf-8") model = str(snapshot) deploy_path = get_deploy_config_path("cosyvoice3.yaml") assert StageConfigFactory.try_infer_model_type(model, trust_remote_code=False) is None assert StageConfigFactory.get_pipeline_config(model, trust_remote_code=False) is None omni_config = StageConfigFactory.create_from_model( model, trust_remote_code=False, cli_overrides={}, deploy_config_path=deploy_path, ) legacy_configs, _ = StageConfigFactory.create_legacy_stage_configs_from_model( model, trust_remote_code=False, cli_overrides={}, deploy_config_path=deploy_path, ) assert omni_config is not None assert omni_config.pipeline_config.model_type == "cosyvoice3" assert [ (stage.stage_id, stage.stage_type, stage.worker_type, stage.model_stage) for stage in omni_config.stage_configs ] == [ (0, StageType.LLM, "ar", "cosyvoice3_talker"), (1, StageType.LLM, "generation", "cosyvoice3_code2wav"), ] assert all(stage.model_config.model == model for stage in omni_config.stage_configs) assert all(stage.stage_type is not StageType.DIFFUSION for stage in omni_config.stage_configs) assert legacy_configs is not None assert [stage.stage_type for stage in legacy_configs] == [StageType.LLM, StageType.LLM] def test_pipeline_registration(self, clean_pipeline_registry): """Ensure that we can register and create a custom pipeline config.""" new_model_type = "new_model_type" pipe_cfg = PipelineConfig( model_type=new_model_type, stages=SINGLE_STAGE_PIPE_CFG, ) # Register the new PipelineConfig assert new_model_type not in OMNI_PIPELINES register_pipeline(pipe_cfg) assert new_model_type in OMNI_PIPELINES assert OMNI_PIPELINES[new_model_type] is pipe_cfg class FakeConfig(PretrainedConfig): model_type = new_model_type # Create the model fake_config = FakeConfig() with ( patch("vllm_omni.config.config_factory.get_config", return_value=fake_config), patch.object(VllmOmniConfig, "from_pipeline_config") as mock_create, ): StageConfigFactory.create_from_model( "fake/model", trust_remote_code=False, cli_overrides={}, deploy_config_path=None, ) mock_create.assert_called_once() assert mock_create.call_args.args == (pipe_cfg,) assert mock_create.call_args.kwargs["cli_overrides"] == { "trust_remote_code": False, "model": "fake/model", } assert pipe_cfg.model_type == new_model_type def test_resolver_registration(self, clean_pipeline_registry): """Ensure that we can register and create a custom resolver.""" new_model_type = "new_model_type" resolved_type = "resolved_type" class FakeConfig(PretrainedConfig): model_type = new_model_type @pipeline_cfg_resolver(config_type=FakeConfig) def custom_resolver( hf_config: FakeConfig, ) -> PipelineConfig: return PipelineConfig( model_type=resolved_type, stages=SINGLE_STAGE_PIPE_CFG, ) # Register the new PipelineConfig assert new_model_type not in OMNI_PIPELINES register_pipeline(custom_resolver, model_type=new_model_type) assert new_model_type in OMNI_PIPELINES assert OMNI_PIPELINES[new_model_type] is custom_resolver # Create the model fake_config = FakeConfig() with ( patch("vllm_omni.config.config_factory.get_config", return_value=fake_config) as mock_get_config, patch.object(VllmOmniConfig, "from_pipeline_config") as mock_create, ): StageConfigFactory.create_from_model( "fake/model", trust_remote_code=True, cli_overrides={}, deploy_config_path=None, ) mock_create.assert_called_once() assert mock_create.call_args.args[0].model_type == resolved_type assert mock_create.call_args.kwargs["cli_overrides"] == { "trust_remote_code": True, "model": "fake/model", } mock_get_config.assert_called_once_with("fake/model", trust_remote_code=True) with patch.object(StageConfigFactory, "_create_legacy_from_registry", return_value=([], None)) as mock_legacy: StageConfigFactory.create_legacy_stage_configs_from_model( "fake/model", trust_remote_code=True, cli_overrides={}, deploy_config_path=None, ) assert mock_legacy.call_args.args[0].model_type == resolved_type assert mock_legacy.call_args.args[1]["trust_remote_code"] is True def test_resolve_when_autodetect_resolves_none(self): """Regression test for: https://github.com/vllm-project/vllm-omni/issues/4726""" deploy_path = get_deploy_config_path("ming_tts.yaml") resolved_config = StageConfigFactory.create_from_model( model="inclusionAI/Ming-omni-tts-0.5B", trust_remote_code=False, cli_overrides={}, deploy_config_path=deploy_path, ) assert resolved_config is not None assert len(resolved_config.stage_configs) > 0 def test_legacy_and_structured_paths_share_deploy_pipeline_override(self, clean_pipeline_registry, tmp_path): pipeline_key = "deploy_only_pipeline" pipe_cfg = PipelineConfig( model_type=pipeline_key, model_arch="DeployOnlyArch", stages=SINGLE_STAGE_PIPE_CFG, ) register_pipeline(pipe_cfg) deploy_path = tmp_path / "deploy_only.yaml" deploy_path.write_text(f"pipeline: {pipeline_key}\n", encoding="utf-8") class FakeConfig(PretrainedConfig): model_type = "dense" architectures = ("UnregisteredArch",) with patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()): omni_config = StageConfigFactory.create_from_model( "fake/model", trust_remote_code=False, cli_overrides={}, deploy_config_path=str(deploy_path), ) legacy_configs, _ = StageConfigFactory.create_legacy_stage_configs_from_model( "fake/model", trust_remote_code=False, cli_overrides={}, deploy_config_path=str(deploy_path), ) assert omni_config is not None assert omni_config.pipeline_config is pipe_cfg assert legacy_configs is not None assert len(legacy_configs) == 1 assert legacy_configs[0].yaml_engine_args["model_arch"] == "DeployOnlyArch" def test_structured_path_loads_explicit_deploy_config_once(self, clean_pipeline_registry, tmp_path): pipeline_key = "single_load_pipeline" pipeline_cfg = PipelineConfig( model_type=pipeline_key, stages=SINGLE_STAGE_PIPE_CFG, ) register_pipeline(pipeline_cfg) deploy_path = tmp_path / "single_load.yaml" deploy_path.write_text(f"pipeline: {pipeline_key}\n", encoding="utf-8") class FakeConfig(PretrainedConfig): model_type = "unregistered" with ( patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()), patch( "vllm_omni.config.config_factory.load_deploy_config", wraps=load_deploy_config, ) as mock_load, ): StageConfigFactory.create_from_model( "fake/model", trust_remote_code=False, cli_overrides={}, deploy_config_path=str(deploy_path), ) mock_load.assert_called_once_with(deploy_path) def test_legacy_default_deploy_name_ignores_cwd(self, monkeypatch, tmp_path): default_name = "pipeline_default.yaml" (tmp_path / default_name).write_text("stages: []\n", encoding="utf-8") monkeypatch.chdir(tmp_path) pipeline = PipelineConfig( model_type="pipeline_with_default", default_deploy_config_name=default_name, stages=SINGLE_STAGE_PIPE_CFG, ) with patch( "vllm_omni.config.config_factory.load_deploy_config", return_value=DeployConfig(), ) as mock_load: StageConfigFactory._create_legacy_from_registry(pipeline, {}) mock_load.assert_called_once_with(_DEPLOY_DIR / default_name) def test_deploy_override_uses_correct_endpoint_restrictions(self, clean_pipeline_registry, tmp_path): """Ensure endpoint restrictions must come from the final pipeline after deploy config overrides, not the auto-detected pipeline. """ # Register two pipeline configs, where one has an endpoint restriction, and one doesn't restriction = EndpointRestriction( OmniServingCapability.COMPLETIONS, "pipeline_a blocks completions", ) pipe_a = PipelineConfig( model_type="detect_type", endpoint_restrictions=(restriction,), stages=SINGLE_STAGE_PIPE_CFG, ) pipe_b = PipelineConfig( model_type="override_type", endpoint_restrictions=(), stages=SINGLE_STAGE_PIPE_CFG, ) register_pipeline(pipe_a) register_pipeline(pipe_b) # Create a config with the autodetected type, and write the # deploy config specifying the override type to a temp path class FakeConfig(PretrainedConfig): model_type = "detect_type" deploy_yaml = tmp_path / "override.yaml" deploy_yaml.write_text("pipeline: override_type\n") # Get the endpoint restrictions, passing the deploy config with the override # type + patching the config for the detected type. Ensure that the endpoint # restrictions correspond to the type in the deploy config. with patch( "vllm_omni.config.config_factory.get_config", return_value=FakeConfig(), ): restrictions = StageConfigFactory.get_pipeline_endpoint_restrictions( model="fake/model", trust_remote_code=False, deploy_config_path=str(deploy_yaml), ) assert restrictions == () 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 TestDeployConfigLoading: def test_rejects_legacy_stage_args_schema(self, tmp_path): deploy_path = tmp_path / "legacy.yaml" deploy_path.write_text("stage_args:\n - stage_id: 0\n", encoding="utf-8") with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"): load_deploy_config(deploy_path) @pytest.mark.parametrize( ("filename", "max_sessions"), [ ("minicpmo_4_5.yaml", 4), ("minicpmo_4_5_2gpu.yaml", 4), ("minicpmo_4_5_3gpu.yaml", 4), ("minicpmo_4_5_8x4090.yaml", 1), ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4), ("minicpmo_4_5_4gpu_stage1_replicas.yaml", 4), ("minicpmo_4_5_8x4090_stage1_replicas.yaml", 1), ], ) def test_minicpmo_deploy_configs_enable_duplex(self, filename: str, max_sessions: int): """Every MiniCPM-o profile serves /v1/realtime, so the retired minicpmo_4_5_duplex.yaml overlay has no replacement to fall back on. """ deploy = load_deploy_config(Path(get_deploy_config_path(filename))) pipeline = resolve_pipeline_config("minicpmo_4_5") assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) assert deploy.session_mode == "duplex" # ``0`` disables the limiter: the configured session/stage capacity is # what bounds the pipeline. Any positive window serializes concurrent # audio first-packet generation (measured: TTFP -23% at 0 vs 4 on # L20X, 128 requests / concurrency 8). assert deploy.active_stream_window == 0 assert deploy.duplex_session.max_sessions == max_sessions assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"] assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"] assert stages[0].yaml_extras["default_sampling_params"]["stop_token_ids"] == [ 151704, 151645, ] def test_load_minicpmo_default_deploy_config(self): deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml")) deploy = load_deploy_config(deploy_path) pipeline = resolve_pipeline_config("minicpmo_4_5") assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) assert deploy.async_chunk is True assert stages[0].yaml_engine_args["async_scheduling"] is True assert stages[1].yaml_engine_args["async_scheduling"] is True assert "Async" in (stages[0].scheduler_cls or "") assert "Async" in (stages[1].scheduler_cls or "") assert [stage.devices for stage in deploy.stages] == ["0", "0", "0"] assert deploy.stages[1].enforce_eager is False assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 4096 assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 50 assert stages[1].yaml_extras["default_sampling_params"]["temperature"] == 0.8 assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561] assert "codec_sampling_params" not in stages[1].yaml_engine_args @pytest.mark.parametrize( ("filename", "stage0_devices", "stage1_devices", "stage2_devices", "stage1_replicas"), [ ("minicpmo_4_5_3gpu_stage1_replicas.yaml", "0", "1,2", "0", 2), ("minicpmo_4_5_4gpu_stage1_replicas.yaml", "0", "1,2,3", "0", 3), ("minicpmo_4_5_8x4090_stage1_replicas.yaml", "0,1,2,3", "4,5,6,7", "5", 4), ], ) def test_load_minicpmo_replica_deploy_configs( self, filename: str, stage0_devices: str, stage1_devices: str, stage2_devices: str, stage1_replicas: int, ): deploy_path = Path(get_deploy_config_path(filename)) deploy = load_deploy_config(deploy_path) pipeline = resolve_pipeline_config("minicpmo_4_5") assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) assert deploy.session_mode == "duplex" assert [stage.yaml_runtime["devices"] for stage in stages] == [ stage0_devices, stage1_devices, stage2_devices, ] assert stages[0].yaml_runtime["num_replicas"] == 1 assert stages[1].yaml_runtime["num_replicas"] == stage1_replicas assert all(stage.yaml_engine_args["async_scheduling"] is False for stage in stages[:2]) assert all("Async" not in (stage.scheduler_cls or "") for stage in stages[:2]) 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 = resolve_pipeline_config("qwen3_tts") stages = merge_pipeline_deploy(pipeline, deploy) assert isinstance(pipeline, PipelineConfig) 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(get_deploy_config_path("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 connector = deploy.connectors["connector_of_shared_memory"] assert "async_chunk_batch_min_size" not in connector["extra"] assert "async_chunk_batch_max_wait_ms" not in connector["extra"] @pytest.mark.parametrize( ("hf_config", "model"), [ (Q3_OMNI_THINKER_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Thinking"), (Q3_OMNI_ALL_STAGES_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Instruct"), ], ids=["thinker_checkpoint", "instruct_forced_thinker_only"], ) def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self, hf_config, model): deploy_path = Path(get_deploy_config_path("qwen3_omni_moe_thinking.yaml")) with patch( "vllm_omni.config.config_factory.get_config", return_value=hf_config, ): pipeline = StageConfigFactory.get_pipeline_config( model, trust_remote_code=False, deploy_config_path=str(deploy_path), ) assert pipeline is not None assert pipeline.model_type == "qwen3_omni_moe_thinker_only" deploy = load_deploy_config(deploy_path) assert deploy.pipeline == "qwen3_omni_moe_thinker_only" stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 1 thinker = stages[0] assert thinker.model_stage == "thinker" assert thinker.final_output is True assert thinker.final_output_type == "text" assert thinker.worker_type == "ar" assert thinker.scheduler_cls == "vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler" assert thinker.hf_config_name == "thinker_config" assert thinker.is_comprehension is True assert thinker.yaml_runtime["devices"] == "0,1" assert thinker.yaml_engine_args["engine_output_type"] == "text" assert thinker.yaml_engine_args["tensor_parallel_size"] == 2 assert thinker.yaml_engine_args["max_num_seqs"] == 1 assert thinker.yaml_engine_args["gpu_memory_utilization"] == 0.9 assert thinker.yaml_engine_args["enforce_eager"] is True assert thinker.yaml_engine_args["async_scheduling"] is False assert thinker.yaml_engine_args["distributed_executor_backend"] == "mp" assert thinker.yaml_engine_args["enable_prefix_caching"] is True assert thinker.yaml_engine_args["async_chunk"] is False assert thinker.yaml_extras["default_sampling_params"]["detokenize"] is True def test_load_voxtral_tts_deploy_config_schema_fields(self): deploy_path = Path(get_deploy_config_path("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(get_deploy_config_path("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_load_voxcpm2_deploy_config_preserves_engine_extras(self): deploy_path = get_deploy_config_path("voxcpm2.yaml") raw_stage = get_deploy_config_stage("voxcpm2.yaml", 0) expected_runtime_config = raw_stage["engine_extras"]["hf_overrides"]["voxcpm2_runtime_config"] deploy = load_deploy_config(deploy_path) runtime_config = deploy.stages[0].engine_extras["hf_overrides"]["voxcpm2_runtime_config"] assert runtime_config == expected_runtime_config @pytest.mark.parametrize( ("deploy_name", "pipeline_name", "stage_count", "final_output_type", "declared_pipeline"), [ ("mammoth_moda2.yaml", "mammoth_moda2", 2, "image", "mammoth_moda2"), ("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text", "mammoth_moda2_ar"), ("omnivoice.yaml", "omnivoice", 1, "audio", "omnivoice"), ("mimo_audio.yaml", "mimo_audio", 2, "audio", None), ("step_audio_2.yaml", "step_audio_2", 2, "audio", None), ("step_audio_2_asr.yaml", "step_audio_2_asr", 1, "text", "step_audio_2_asr"), ("step_audio_2_async_chunk.yaml", "step_audio_2", 2, "audio", None), ("step_audio2_ci.yaml", "step_audio_2_asr", 1, "text", "step_audio_2_asr"), ("hunyuan_video_15.yaml", "hunyuan_video_15", 1, "video", None), ("wan2_2_ti2v.yaml", "wan2_2_ti2v", 1, "video", None), ( "minimax_h3_disaggregated.yaml", "minimax_h3_disaggregated", 2, "video", "minimax_h3_disaggregated", ), ], ) def test_load_new_registry_backed_deploy_configs( self, deploy_name: str, pipeline_name: str, stage_count: int, final_output_type: str, declared_pipeline: str | None, ): deploy_path = Path(get_deploy_config_path(deploy_name)) deploy = load_deploy_config(deploy_path) assert deploy.pipeline == declared_pipeline with patch("vllm_omni.platforms.current_omni_platform") as platform: platform.device_name = "cuda" stages = merge_pipeline_deploy(OMNI_PIPELINES[pipeline_name], deploy) assert len(stages) == stage_count assert stages[-1].final_output is True assert stages[-1].final_output_type == final_output_type if pipeline_name == "minimax_h3_disaggregated": assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root") assert ( stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path") ) assert "model_path_resolver" not in stages[0].yaml_extras assert "model_path_resolver" not in stages[1].yaml_extras if deploy.trust_remote_code is not None: assert {s.yaml_engine_args.get("trust_remote_code") for s in stages} == {deploy.trust_remote_code} def test_minimax_h3_resolvers_reach_production_engine_args(self, tmp_path: Path, monkeypatch): from vllm_omni.engine import stage_init_utils model_root = tmp_path / "MiniMax-H3" (model_root / "FL2VA" / "text_encoder").mkdir(parents=True) (model_root / "Ref2VA" / "text_encoder").mkdir(parents=True) monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None) deploy = load_deploy_config(get_deploy_config_path("minimax_h3_disaggregated.yaml")) assert deploy.stages[1].engine_extras["model_loaded"] == { "text_encoder": False, "vae_encoder": False, } stages = merge_pipeline_deploy(OMNI_PIPELINES["minimax_h3_disaggregated"], deploy) resolved = [stage_init_utils.build_engine_args_dict(stage.to_omegaconf(), str(model_root)) for stage in stages] assert resolved[0]["model"] == str(model_root / "FL2VA" / "text_encoder") assert resolved[1]["model"] == str(model_root / "FL2VA") assert all("model_path_resolver" not in engine_args for engine_args in resolved) def test_minimax_h3_text_encoder_tp_reaches_only_stage_zero(self): pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"] stages, _ = StageConfigFactory._create_legacy_from_registry( pipeline, {"text_encoder_tp_size": 4}, ) assert stages[0].runtime_overrides["tensor_parallel_size"] == 4 assert "text_encoder_tp_size" not in stages[0].runtime_overrides assert "text_encoder_tp_size" not in stages[1].runtime_overrides assert stages[1].yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1 def test_minimax_h3_disaggregated_defaults_match_validated_topology(self): pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"] deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml"))) stages = merge_pipeline_deploy(pipeline, deploy) assert stages[0].yaml_engine_args["model_arch"] == "MiniMaxH3Encoder" assert stages[1].yaml_engine_args["model_arch"] == "MiniMaxH3Pipeline" assert stages[0].yaml_runtime["num_replicas"] == 1 assert stages[1].yaml_runtime["num_replicas"] == 1 assert stages[1].yaml_engine_args["model_loaded"] == { "text_encoder": False, "vae_encoder": False, } assert stages[0].yaml_engine_args["max_num_seqs"] == 1 assert stages[0].yaml_engine_args["hf_overrides"]["minimax_h3_encoder_components"] == { "text_encoder": {"parallel_mode": "tp"}, "video_vae": {"parallel_mode": "patch"}, "audio_vae": {"parallel_mode": "leader"}, } assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root") assert stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path") parallel = stages[1].yaml_engine_args["parallel_config"] assert parallel["tensor_parallel_size"] == 1 assert parallel["ulysses_degree"] == 4 assert parallel["vae_patch_parallel_size"] == 4 assert stages[1].yaml_engine_args["inline_diffusion"] is True assert stages[1].yaml_extras["default_sampling_params"]["num_inference_steps"] == 50 turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml"))) turbo_stages = merge_pipeline_deploy(pipeline, turbo) assert turbo_stages[0].yaml_engine_args["model_arch"] == stages[0].yaml_engine_args["model_arch"] assert turbo_stages[1].yaml_engine_args["model_loaded"] == stages[1].yaml_engine_args["model_loaded"] assert turbo_stages[0].yaml_engine_args["hf_overrides"] == stages[0].yaml_engine_args["hf_overrides"] turbo_sampling = turbo_stages[1].yaml_extras["default_sampling_params"] assert turbo_sampling["num_inference_steps"] == 5 assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0} def test_minimax_h3_text_encoder_tp_alias_targets_stage_zero(self): pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"] assert normalize_pipeline_cli_overrides(pipeline, {}) == {} assert normalize_pipeline_cli_overrides(pipeline, {"text_encoder_tp_size": 4}) == { "stage_0_tensor_parallel_size": 4 } def test_minimax_h3_stage_zero_tp_override_wins_alias_conflict(self): pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"] with pytest.warns(UserWarning, match="stage_0_tensor_parallel_size=2 takes precedence"): normalized = normalize_pipeline_cli_overrides( pipeline, {"text_encoder_tp_size": 4, "stage_0_tensor_parallel_size": 2}, ) assert normalized == {"stage_0_tensor_parallel_size": 2} def test_minimax_h3_rejects_text_encoder_tp_on_diffusion_stage(self): pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"] with pytest.raises(ValueError, match="stage_1_text_encoder_tp_size cannot be set"): normalize_pipeline_cli_overrides(pipeline, {"stage_1_text_encoder_tp_size": 4}) @pytest.mark.parametrize( ("config_json", "model_index", "expected_pipeline"), [ ({"model_type": "step_audio_2"}, None, "step_audio_2"), (None, {"_class_name": "HunyuanVideo15Pipeline"}, "hunyuan_video_15"), (None, {"_class_name": "WanPipeline"}, "wan2_2_ti2v"), (None, {"_class_name": "WanDMDPipeline"}, "wan2_2_ti2v"), ], ) def test_migrated_models_are_discovered_without_explicit_deploy( self, config_json, model_index, expected_pipeline, ): def get_model_file(filename, _model, revision=None): del revision if filename == "config.json": return config_json if filename == "model_index.json": return model_index return None with ( patch.object(StageConfigFactory, "get_hf_config", return_value=None), patch("vllm_omni.config.config_factory.get_hf_file_to_dict", side_effect=get_model_file), ): pipeline = StageConfigFactory.get_pipeline_config( model="/models/unrelated-checkpoint-name", trust_remote_code=False, ) assert pipeline is not None assert pipeline.model_type == expected_pipeline def test_minimax_h3_disaggregation_is_not_auto_discovered(self): def get_model_file(filename, _model, revision=None): del revision if filename == "config.json": return None if filename == "model_index.json": return {"_class_name": "MiniMaxH3Pipeline"} return None with ( patch.object(StageConfigFactory, "get_hf_config", return_value=None), patch( "vllm_omni.config.config_factory.get_hf_file_to_dict", side_effect=get_model_file, ), ): pipeline = StageConfigFactory.get_pipeline_config( model="/models/unrelated-checkpoint-name", trust_remote_code=False, ) assert pipeline is None @pytest.mark.parametrize("deploy_name", ["step_audio_2.yaml", "step_audio_2_async_chunk.yaml"]) def test_step_audio2_deploy_configs_fit_two_gpus(self, deploy_name): deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name))) assert deploy.stages[0].devices == "0,1" assert deploy.stages[0].tensor_parallel_size == 2 assert deploy.stages[1].devices == "1" def test_step_audio2_online_ci_deploy_is_single_thinker(self): deploy = load_deploy_config(Path(get_deploy_config_path("step_audio2_ci.yaml"))) assert deploy.pipeline == "step_audio_2_asr" assert deploy.async_chunk is False assert deploy.trust_remote_code is True assert deploy.enable_prefix_caching is False assert len(deploy.stages) == 1 stage = deploy.stages[0] assert stage.devices == "0" assert stage.tensor_parallel_size == 1 assert stage.max_model_len == 1024 assert stage.max_num_batched_tokens == 1024 assert stage.max_num_seqs == 1 assert stage.gpu_memory_utilization == 0.8 assert stage.skip_mm_profiling is True assert stage.enforce_eager is True assert stage.async_scheduling is False assert stage.engine_extras["hf_overrides"]["architectures"] == ["StepAudio2ThinkerForConditionalGeneration"] stages = merge_pipeline_deploy(OMNI_PIPELINES[deploy.pipeline], deploy) assert len(stages) == 1 assert stages[0].final_output is True assert stages[0].final_output_type == "text" assert stages[0].yaml_engine_args["engine_output_type"] == "text" def test_step_audio2_dispatches_sync_and_async_chunk_processors(self): pipeline = resolve_pipeline_config("step_audio_2") assert isinstance(pipeline, PipelineConfig) sync_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=False)) assert "custom_process_next_stage_input_func" not in sync_stages[0].yaml_engine_args assert sync_stages[1].custom_process_input_func.endswith("thinker2token2wav") async_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=True)) assert ( async_stages[0] .yaml_engine_args["custom_process_next_stage_input_func"] .endswith("thinker2token2wav_async_chunk") ) assert async_stages[1].custom_process_input_func is None def test_no_bundled_legacy_stage_config_yamls(self): repo_root = Path(__file__).resolve().parents[2] stage_config_dir = repo_root / "vllm_omni" / "model_executor" / "stage_configs" assert not list(stage_config_dir.glob("*.yaml")) def test_merge_pipeline_deploy(self): deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): pytest.skip("Deploy config not found") pipeline = resolve_pipeline_config( "qwen3_omni_moe", Q3_OMNI_ALL_STAGES_HF_CONFIG, ) assert isinstance(pipeline, PipelineConfig) 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 = resolve_pipeline_config( "qwen3_omni_moe", Q3_OMNI_ALL_STAGES_HF_CONFIG, ) assert isinstance(pipeline, PipelineConfig) base = Path(get_deploy_config_path("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, final_output=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_merge_pipeline_deploy_preserves_pipeline_scheduler_cls(self): scheduler_cls = "tests.fake.CustomScheduler" pipeline = PipelineConfig( model_type="test_scheduler", model_arch="TestModel", stages=( StagePipelineConfig( stage_id=0, model_stage="ar", execution_type=StageExecutionType.LLM_AR, scheduler_cls=scheduler_cls, final_output=True, ), ), ) deploy = DeployConfig(async_chunk=False, stages=[StageDeployConfig(stage_id=0)]) stages = merge_pipeline_deploy(pipeline, deploy) assert stages[0].scheduler_cls == scheduler_cls 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_explicit_engine_extras_merge_with_flat_and_engine_args(self): """Explicit engine_extras should be preserved with existing pass-through styles.""" fake_config = { "stages": [ { "stage_id": 0, "engine_extras": { "hf_overrides": { "runtime_config": { "explicit_only": True, "shared": "explicit", "nested": {"a": 1}, } } }, "hf_overrides": { "runtime_config": { "flat_only": True, "shared": "flat", "nested": {"b": 2}, } }, "engine_args": { "hf_overrides": { "runtime_config": { "engine_args_only": True, "shared": "engine_args", "nested": {"c": 3}, } } }, }, ] } with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): deploy = load_deploy_config("dummy.yaml") assert deploy.stages[0].engine_extras["hf_overrides"] == { "runtime_config": { "engine_args_only": True, "explicit_only": True, "flat_only": True, "nested": {"a": 1, "b": 2, "c": 3}, "shared": "engine_args", } } 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"] assert isinstance(original_foo, dict) 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 = resolve_pipeline_config( "qwen3_omni_moe", Q3_OMNI_ALL_STAGES_HF_CONFIG, ) assert isinstance(p, PipelineConfig) assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration" assert len(p.stages) == 3 def test_thinker(self): p = resolve_pipeline_config( "qwen3_omni_moe", Q3_OMNI_ALL_STAGES_HF_CONFIG, ) assert isinstance(p, PipelineConfig) s = p.get_stage(0) assert isinstance(s, StagePipelineConfig) 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): p = resolve_pipeline_config( "qwen3_omni_moe", Q3_OMNI_ALL_STAGES_HF_CONFIG, ) assert isinstance(p, PipelineConfig) s = p.get_stage(1) assert isinstance(s, StagePipelineConfig) assert s.input_sources == (0,) assert s.sampling_constraints["stop_token_ids"] == [2150] # thinker2talker was removed: sync_process_input_func always wins. assert s.custom_process_input_func is None assert s.sync_process_input_func is not None assert s.sync_process_input_func.endswith("thinker2talker_token_only") assert s.custom_process_next_stage_input_func is not None def test_code2wav(self): p = resolve_pipeline_config( "qwen3_omni_moe", Q3_OMNI_ALL_STAGES_HF_CONFIG, ) assert isinstance(p, PipelineConfig) s = p.get_stage(2) assert isinstance(s, StagePipelineConfig) assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.final_output_type == "audio" assert s.engine_output_type == "audio" class TestQwen2_5OmniPipeline: def test_registered(self): p = resolve_pipeline_config("qwen2_5_omni") assert isinstance(p, PipelineConfig) assert p.model_arch == "Qwen2_5OmniForConditionalGeneration" assert len(p.stages) == 3 def test_thinker(self): p = resolve_pipeline_config("qwen2_5_omni") assert isinstance(p, PipelineConfig) s = p.get_stage(0) assert isinstance(s, StagePipelineConfig) 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 assert s.hf_config_name == "thinker_config" def test_talker(self): p = resolve_pipeline_config("qwen2_5_omni") assert isinstance(p, PipelineConfig) s = p.get_stage(1) assert isinstance(s, StagePipelineConfig) assert s.input_sources == (0,) assert s.hf_config_name == "talker_config" assert s.sampling_constraints["stop_token_ids"] == [8294] # thinker2talker was removed: qwen2_5_omni has no async_chunk support, # so sync_process_input_func always wins and custom_process_input_func # was dead code. assert s.custom_process_input_func is None assert s.sync_process_input_func is not None def test_code2wav(self): p = resolve_pipeline_config("qwen2_5_omni") assert isinstance(p, PipelineConfig) s = p.get_stage(2) assert isinstance(s, StagePipelineConfig) assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.final_output_type == "audio" assert s.engine_output_type == "audio" assert s.hf_config_name == "thinker_config" class TestQwen3TTSPipeline: def test_registered(self): p = resolve_pipeline_config("qwen3_tts") assert isinstance(p, PipelineConfig) assert p is not None assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration" assert len(p.stages) == 2 def test_talker_stage(self): p = resolve_pipeline_config("qwen3_tts") assert isinstance(p, PipelineConfig) s = p.get_stage(0) assert isinstance(s, StagePipelineConfig) 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): p = resolve_pipeline_config("qwen3_tts") assert isinstance(p, PipelineConfig) s = p.get_stage(1) assert isinstance(s, StagePipelineConfig) 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(get_deploy_config_path("qwen3_tts.yaml")) if not deploy_path.exists(): pytest.skip("qwen3_tts deploy yaml not found") deploy = load_deploy_config(deploy_path) pipeline = OMNI_PIPELINES["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 = resolve_pipeline_config("ming_flash_omni") assert isinstance(p, PipelineConfig) assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 2 def test_thinker_stage(self): p = resolve_pipeline_config("ming_flash_omni") assert isinstance(p, PipelineConfig) s = p.get_stage(0) assert isinstance(s, StagePipelineConfig) 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): p = resolve_pipeline_config("ming_flash_omni") assert isinstance(p, PipelineConfig) s = p.get_stage(1) assert isinstance(s, StagePipelineConfig) 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" # thinker2talker was removed: ming_flash_omni has no async_chunk support # and both thinker2talker / thinker2talker_token_only called _build_talker_inputs # identically, so custom_process_input_func was dead code. assert s.custom_process_input_func is None assert s.sync_process_input_func is not None def test_talker_stage_processor_wiring_resolves(self): """The sync_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. """ p = resolve_pipeline_config("ming_flash_omni") assert isinstance(p, PipelineConfig) s = p.get_stage(1) assert isinstance(s, StagePipelineConfig) module_path, _, attr = s.sync_process_input_func.rpartition(".") module = importlib.import_module(module_path) assert callable(getattr(module, attr)) def test_tts_pipeline_registered(self): p = resolve_pipeline_config("ming_flash_omni_tts") assert isinstance(p, PipelineConfig) assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration" assert len(p.stages) == 1 def test_tts_stage(self): p = resolve_pipeline_config("ming_flash_omni_tts") assert isinstance(p, PipelineConfig) s = p.get_stage(0) assert isinstance(s, StagePipelineConfig) 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(get_deploy_config_path("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 = resolve_pipeline_config("ming_flash_omni") assert isinstance(pipeline, PipelineConfig) 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(get_deploy_config_path("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 = OMNI_PIPELINES["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 = resolve_pipeline_config("ming_flash_omni_thinker_only") assert isinstance(p, PipelineConfig) assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 1 def test_thinker_only_stage(self): p = resolve_pipeline_config("ming_flash_omni_thinker_only") assert isinstance(p, PipelineConfig) s = p.get_stage(0) assert isinstance(s, StagePipelineConfig) 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(get_deploy_config_path("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 = resolve_pipeline_config("ming_flash_omni_thinker_only") assert isinstance(pipeline, PipelineConfig) 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 = OMNI_PIPELINES.get("ming_flash_omni_image") assert p is not None assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 2 def test_image_thinker_stage(self): s = resolve_pipeline_config("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 = resolve_pipeline_config("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.""" pipeline = resolve_pipeline_config("ming_flash_omni_image") assert isinstance(pipeline, PipelineConfig) thinker = pipeline.get_stage(0) dit = pipeline.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(get_deploy_config_path("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 = resolve_pipeline_config("ming_flash_omni_image") assert isinstance(pipeline, PipelineConfig) 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_minicpmo_overlays_inherit_talker_sampling_params(self): """Overlays must keep the base Talker codec Sampler knobs.""" for filename in ( "minicpmo_4_5_2gpu.yaml", "minicpmo_4_5_3gpu.yaml", "minicpmo_4_5_3gpu_stage1_replicas.yaml", "minicpmo_4_5_4gpu_stage1_replicas.yaml", "minicpmo_4_5_8x4090_stage1_replicas.yaml", ): path = Path(get_deploy_config_path(filename)) if not path.exists(): pytest.skip(f"{filename} not found") deploy = load_deploy_config(path) sampling = deploy.stages[1].default_sampling_params assert sampling is not None, f"{filename} stage 1 lost default_sampling_params" assert sampling["temperature"] == 0.8, filename assert sampling["top_k"] == 25, filename assert sampling["top_p"] == 0.85, filename assert sampling["repetition_penalty"] == 1.05, filename 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_qwen3_omni_colocate_async_bounds_only_rocm_kv_cache(self): ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe_colocate_async.yaml")) pipeline = resolve_pipeline_config("qwen3_omni_moe_thinker_only") assert isinstance(pipeline, PipelineConfig) cuda = _apply_platform_overrides(load_deploy_config(ci_path), platform="cuda") cuda_stage = merge_pipeline_deploy(pipeline, cuda)[0] assert cuda_stage.yaml_engine_args["gpu_memory_utilization"] == 0.9 assert "kv_cache_memory_bytes" not in cuda_stage.yaml_engine_args rocm = _apply_platform_overrides(load_deploy_config(ci_path), platform="rocm") rocm_stage = merge_pipeline_deploy(pipeline, rocm)[0] assert "gpu_memory_utilization" not in rocm_stage.yaml_engine_args assert rocm_stage.yaml_engine_args["kv_cache_memory_bytes"] == 2 * 1024**3 def test_pure_inheritance_overlay(self, tmp_path): """An overlay with only ``base_config`` inherits everything.""" base = Path(get_deploy_config_path("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(get_deploy_config_path("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_qwen3_omni_musa_clears_parent_quantization_for_audio_stages(self): pipeline = resolve_pipeline_config( "qwen3_omni_moe", Q3_OMNI_ALL_STAGES_HF_CONFIG, ) assert isinstance(pipeline, PipelineConfig) deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) base_stages = merge_pipeline_deploy(pipeline, load_deploy_config(deploy_path)) assert "hf_overrides" not in base_stages[1].yaml_engine_args assert "hf_overrides" not in base_stages[2].yaml_engine_args musa = _apply_platform_overrides(load_deploy_config(deploy_path), platform="musa") musa_stages = merge_pipeline_deploy(pipeline, musa) expected = {"quantization_config": None} assert musa_stages[1].yaml_engine_args["hf_overrides"] == expected assert musa_stages[2].yaml_engine_args["hf_overrides"] == expected def test_qwen3_tts_rocm_disables_code2wav_outer_cudagraph(self): deploy_path = Path(get_deploy_config_path("qwen3_tts.yaml")) base = load_deploy_config(deploy_path) assert base.stages[0].enforce_eager is None assert base.stages[1].enforce_eager is False rocm = _apply_platform_overrides(base, platform="rocm") assert rocm.stages[0].enforce_eager is None assert rocm.stages[1].enforce_eager is True @pytest.mark.parametrize("deploy_name", ["qwen3_tts.yaml", "qwen3_tts_high_concurrency.yaml"]) @pytest.mark.parametrize("platform", ["cuda", "npu", "rocm"]) def test_qwen3_tts_default_code2wav_dtype_is_bf16(self, deploy_name, platform): deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name))) deploy = _apply_platform_overrides(deploy, platform=platform) stages = merge_pipeline_deploy(resolve_pipeline_config("qwen3_tts"), deploy) assert stages[1].yaml_engine_args["dtype"] == "bfloat16" @pytest.mark.parametrize("deploy_name", ["qwen3_tts.yaml", "qwen3_tts_high_concurrency.yaml"]) @pytest.mark.parametrize("platform", ["cuda", "npu"]) @pytest.mark.parametrize("dtype", ["float32", "bfloat16", "float16"]) def test_qwen3_tts_propagates_explicit_code2wav_dtype(self, deploy_name, platform, dtype): deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name))) deploy.stages[1].engine_extras["dtype"] = dtype deploy = _apply_platform_overrides(deploy, platform=platform) stages = merge_pipeline_deploy(resolve_pipeline_config("qwen3_tts"), deploy) assert stages[1].yaml_engine_args["dtype"] == dtype def test_moss_tts_rocm_disables_codec_cudagraph(self): deploy_path = Path(get_deploy_config_path("moss_tts.yaml")) base = load_deploy_config(deploy_path) assert base.stages[1].enforce_eager is False rocm = _apply_platform_overrides(base, platform="rocm") assert rocm.stages[1].enforce_eager is True def test_higgs_audio_v3_rocm_uses_triton_attention(self): deploy_path = Path(get_deploy_config_path("higgs_multimodal_qwen3.yaml")) base = load_deploy_config(deploy_path) assert base.stages[0].engine_extras["attention_backend"] == "FLASHINFER" rocm = _apply_platform_overrides(base, platform="rocm") assert rocm.stages[0].engine_extras["attention_backend"] == "TRITON_ATTN" pipeline = resolve_pipeline_config("higgs_multimodal_qwen3") assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, rocm) assert stages[0].yaml_engine_args["attention_backend"] == "TRITON_ATTN" def test_qwen3_omni_cuda_uses_thinker_rotary_custom_op(self): deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) pipeline = resolve_pipeline_config( "qwen3_omni_moe", Q3_OMNI_ALL_STAGES_HF_CONFIG, ) assert isinstance(pipeline, PipelineConfig) cuda = _apply_platform_overrides(load_deploy_config(deploy_path), platform="cuda") cuda_stages = merge_pipeline_deploy(pipeline, cuda) expected = {"custom_ops": ["+rotary_embedding"]} assert cuda_stages[0].yaml_engine_args["compilation_config"] == expected assert all("compilation_config" not in stage.yaml_engine_args for stage in cuda_stages[1:]) for platform in ("cpu", "musa", "npu", "rocm", "xpu"): deploy = _apply_platform_overrides( load_deploy_config(deploy_path), platform=platform, ) config = deploy.stages[0].compilation_config or {} assert "+rotary_embedding" not in config.get("custom_ops", []) def test_minicpmo_4_5_cuda_caps_talker_kv_cache(self): pipeline = resolve_pipeline_config("minicpmo_4_5") assert isinstance(pipeline, PipelineConfig) deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml")) cuda = _apply_platform_overrides(load_deploy_config(deploy_path), platform="cuda") cuda_stages = merge_pipeline_deploy(pipeline, cuda) assert cuda_stages[1].yaml_engine_args["kv_cache_memory_bytes"] == 2 * 1024**3 # The CUDA memory budget must not leak into the existing NPU profile. npu = _apply_platform_overrides(load_deploy_config(deploy_path), platform="npu") npu_stages = merge_pipeline_deploy(pipeline, npu) assert "kv_cache_memory_bytes" not in npu_stages[1].yaml_engine_args # Multi-GPU replica profiles place Talker on dedicated devices and # preserve their previous automatic KV-cache sizing. for filename in ( "minicpmo_4_5_3gpu_stage1_replicas.yaml", "minicpmo_4_5_4gpu_stage1_replicas.yaml", ): replica = _apply_platform_overrides( load_deploy_config(Path(get_deploy_config_path(filename))), platform="cuda" ) replica_stages = merge_pipeline_deploy(pipeline, replica) # Explicit null clears the inherited single-GPU 2 GiB CUDA cap. assert replica_stages[1].yaml_engine_args.get("kv_cache_memory_bytes") is None def test_fish_speech_npu_uses_ascend_kv_block_size(self): deploy_path = Path(get_deploy_config_path("fish_qwen3_omni.yaml")) deploy = _apply_platform_overrides(load_deploy_config(deploy_path), platform="npu") assert deploy.stages[0].engine_extras["block_size"] == 128 @pytest.mark.parametrize("platform", ["cuda", "npu"]) def test_recommended_native_runner_platform_defaults(self, platform): filename, pipeline_key = "qwen3_tts_mrv2.yaml", "qwen3_tts" deploy = load_deploy_config(Path(get_deploy_config_path(filename))) deploy = _apply_platform_overrides(deploy, platform=platform) expect_v2 = platform == "cuda" assert deploy.model_runner == ("v2" if expect_v2 else "v1") # The selection must reach the live worker-dispatch consumer, not just # the transport-level DeployConfig field. pipeline = resolve_pipeline_config(pipeline_key) stages = merge_pipeline_deploy(pipeline, deploy) assert [stage.yaml_engine_args["use_v2_model_runner"] for stage in stages] == [expect_v2] * len(stages) if expect_v2 and pipeline.stages and deploy.async_chunk: # v2 only engages the native plane on stages declaring support. assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages) @pytest.mark.parametrize("runner,native", [("v1", False), ("v2", False), ("v2", True)]) def test_mrv2_undeclared_transport_warns(self, monkeypatch, runner, native): from unittest.mock import Mock warning = Mock() monkeypatch.setattr("vllm_omni.config.stage_config.logger.warning", warning) pipeline = PipelineConfig( model_type="custom_pipeline", stages=[ StagePipelineConfig( stage_id=0, model_stage="custom_ar", final_output=True, supports_native_mrv2_data_plane=native ) ], ) merge_pipeline_deploy(pipeline, DeployConfig(model_runner=runner)) assert warning.call_count == int(runner == "v2" and not native) if warning.called: assert "legacy transport path" in warning.call_args.args[0] def test_runner_selection_rejects_engine_extras_override(self): pipeline = PipelineConfig( model_type="runner_selection_reserved", stages=SINGLE_STAGE_PIPE_CFG, ) for reserved in ("use_v2_model_runner", "supports_native_mrv2_data_plane"): deploy = DeployConfig( model_runner="v2", stages=[StageDeployConfig(stage_id=0, engine_extras={reserved: False})], ) with pytest.raises(ValueError, match=f"{reserved!r} must not be set"): merge_pipeline_deploy(pipeline, deploy) def test_npu_overrides(self): deploy_path = Path(get_deploy_config_path("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_qwen2_5_omni_xpu_uses_eager_ar_stages(self): deploy_path = Path(get_deploy_config_path("qwen2_5_omni.yaml")) deploy = load_deploy_config(deploy_path) deploy = _apply_platform_overrides(deploy, platform="xpu") assert deploy.stages[0].enforce_eager is True assert deploy.stages[1].enforce_eager is True assert deploy.stages[2].enforce_eager is True def test_xpu_overrides(self): deploy_path = Path(get_deploy_config_path("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(get_deploy_config_path("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): deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) pipeline = resolve_pipeline_config( "qwen3_omni_moe", Q3_OMNI_ALL_STAGES_HF_CONFIG, ) assert isinstance(pipeline, PipelineConfig) 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): deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) pipeline = resolve_pipeline_config( "qwen3_omni_moe", Q3_OMNI_ALL_STAGES_HF_CONFIG, ) assert isinstance(pipeline, PipelineConfig) 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 TestAuraOmniDeploy: def test_aura_omni_deploy_forces_pipeline_override(self): deploy_path = Path(get_deploy_config_path("aura_omni.yaml")) deploy = load_deploy_config(deploy_path) assert deploy.pipeline == "aura_omni" def test_aura_omni_deploy_resolves_four_native_stages(self): pipeline_cfg = resolve_pipeline_config("aura_omni") stages, _ = StageConfigFactory._create_legacy_from_registry( pipeline_cfg, cli_overrides={}, deploy_config_path=get_deploy_config_path("aura_omni.yaml"), ) assert [stage.model_stage for stage in stages] == [ "asr", "aura", "qwen3_tts", "code2wav", ] assert [stage.final_output for stage in stages] == [False, True, False, True] assert [stage.final_output_type for stage in stages] == [None, "text", None, "audio"] assert stages[0].yaml_engine_args["model_arch"] == "Qwen3ASRForConditionalGeneration" assert stages[1].yaml_engine_args["model_arch"] == "AuraQwen3VLForConditionalGeneration" assert stages[2].yaml_engine_args["model_arch"] == "Qwen3TTSTalkerForConditionalGeneration" assert stages[3].yaml_engine_args["model_arch"] == "Qwen3TTSCode2Wav" class TestDeployCliOverrideFlow: """Test deploy-YAML baselines overridden by CLI runtime overrides.""" def test_diffusion_deploy_fields_can_be_overridden_by_cli(self, tmp_path): deploy_path = tmp_path / "diffusion_stage.yaml" deploy_path.write_text( """ async_chunk: false stages: - stage_id: 0 devices: "4,5" parallel_config: pipeline_parallel_size: 1 data_parallel_size: 1 tensor_parallel_size: 1 enable_expert_parallel: false sequence_parallel_size: 1 ulysses_degree: 1 ring_degree: 1 cfg_parallel_size: 1 vae_patch_parallel_size: 1 use_hsdp: false hsdp_shard_size: -1 hsdp_replicate_size: 1 engine_args: cache_backend: cache_dit diffusion_attention_backend: FLASH_ATTN diffusion_kv_cache_dtype: auto step_execution: false vae_use_tiling: false enable_cpu_offload: false max_generated_image_size: 1048576 tts_max_instructions_length: 1000 """, encoding="utf-8", ) pipeline = PipelineConfig( model_type="test_diffusion", stages=( StagePipelineConfig( stage_id=0, model_stage="dit", execution_type=StageExecutionType.DIFFUSION, final_output=True, final_output_type="image", ), ), ) deploy = load_deploy_config(deploy_path) stages = merge_pipeline_deploy(pipeline, deploy) stage = stages[0] assert stage.yaml_engine_args["parallel_config"]["pipeline_parallel_size"] == 1 assert stage.yaml_engine_args["parallel_config"]["data_parallel_size"] == 1 assert stage.yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1 assert stage.yaml_engine_args["parallel_config"]["enable_expert_parallel"] is False assert stage.yaml_engine_args["parallel_config"]["sequence_parallel_size"] == 1 assert stage.yaml_engine_args["parallel_config"]["ulysses_degree"] == 1 assert stage.yaml_engine_args["parallel_config"]["ring_degree"] == 1 assert stage.yaml_engine_args["parallel_config"]["cfg_parallel_size"] == 1 assert stage.yaml_engine_args["parallel_config"]["vae_patch_parallel_size"] == 1 assert stage.yaml_engine_args["cache_backend"] == "cache_dit" assert stage.yaml_engine_args["parallel_config"]["use_hsdp"] is False assert stage.yaml_engine_args["parallel_config"]["hsdp_shard_size"] == -1 assert stage.yaml_engine_args["parallel_config"]["hsdp_replicate_size"] == 1 assert stage.yaml_engine_args["diffusion_attention_backend"] == "FLASH_ATTN" assert stage.yaml_engine_args["diffusion_kv_cache_dtype"] == "auto" assert stage.yaml_engine_args["step_execution"] is False assert stage.yaml_engine_args["vae_use_tiling"] is False assert stage.yaml_engine_args["enable_cpu_offload"] is False assert stage.yaml_engine_args["max_generated_image_size"] == 1048576 assert stage.yaml_engine_args["tts_max_instructions_length"] == 1000 stage.runtime_overrides = StageConfigFactory._merge_cli_overrides( stage, { "pipeline_parallel_size": 2, "data_parallel_size": 3, "tensor_parallel_size": 4, "enable_expert_parallel": True, "sequence_parallel_size": 24, "ulysses_degree": 2, "ring_degree": 4, "cfg_parallel_size": 2, "vae_patch_parallel_size": 2, "cache_backend": "tea_cache", "use_hsdp": True, "hsdp_shard_size": 8, "hsdp_replicate_size": 2, "diffusion_attention_backend": "SAGE_ATTN", "diffusion_kv_cache_dtype": "fp8", "step_execution": True, "vae_use_tiling": True, "enable_cpu_offload": True, "max_generated_image_size": 2097152, "tts_max_instructions_length": 2000, }, ) omega_config = stage.to_omegaconf() assert omega_config.engine_args.cache_backend == "tea_cache" assert omega_config.engine_args.diffusion_attention_backend == "SAGE_ATTN" assert omega_config.engine_args.diffusion_kv_cache_dtype == "fp8" assert omega_config.engine_args.step_execution is True assert omega_config.engine_args.vae_use_tiling is True assert omega_config.engine_args.enable_cpu_offload is True assert omega_config.engine_args.max_generated_image_size == 2097152 assert omega_config.engine_args.tts_max_instructions_length == 2000 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 == 4 assert omega_config.engine_args.parallel_config.enable_expert_parallel is True assert omega_config.engine_args.parallel_config.sequence_parallel_size == 24 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.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 def test_llm_deploy_fields_can_be_overridden_by_cli(self, tmp_path): deploy_path = tmp_path / "llm_stage.yaml" deploy_path.write_text( """ async_chunk: false stages: - stage_id: 0 devices: "0,1" engine_args: tensor_parallel_size: 1 enable_expert_parallel: false gpu_memory_utilization: 0.5 max_num_seqs: 16 max_num_batched_tokens: 1024 max_model_len: 4096 enforce_eager: false """, encoding="utf-8", ) pipeline = PipelineConfig( model_type="test_llm", stages=( StagePipelineConfig( stage_id=0, model_stage="thinker", execution_type=StageExecutionType.LLM_AR, final_output=True, final_output_type="text", ), ), ) deploy = load_deploy_config(deploy_path) stages = merge_pipeline_deploy(pipeline, deploy) stage = stages[0] assert stage.yaml_engine_args["tensor_parallel_size"] == 1 assert stage.yaml_engine_args["enable_expert_parallel"] is False assert stage.yaml_engine_args["gpu_memory_utilization"] == 0.5 assert stage.yaml_engine_args["max_num_seqs"] == 16 assert stage.yaml_engine_args["max_num_batched_tokens"] == 1024 assert stage.yaml_engine_args["max_model_len"] == 4096 assert stage.yaml_engine_args["enforce_eager"] is False stage.runtime_overrides = StageConfigFactory._merge_cli_overrides( stage, { "tensor_parallel_size": 2, "enable_expert_parallel": True, "gpu_memory_utilization": 0.9, "max_num_seqs": 32, "max_num_batched_tokens": 2048, "max_model_len": 8192, "enforce_eager": True, }, ) omega_config = stage.to_omegaconf() assert omega_config.engine_args.tensor_parallel_size == 2 assert omega_config.engine_args.enable_expert_parallel is True assert omega_config.engine_args.gpu_memory_utilization == 0.9 assert omega_config.engine_args.max_num_seqs == 32 assert omega_config.engine_args.max_num_batched_tokens == 2048 assert omega_config.engine_args.max_model_len == 8192 assert omega_config.engine_args.enforce_eager is True def test_cli_attention_shorthand_replaces_yaml_structured_default(self): stage = StageConfig( stage_id=0, model_stage="dit", stage_type=StageType.DIFFUSION, yaml_engine_args={ "diffusion_attention_config": { "default": {"backend": "FLASH_ATTN"}, "per_role": {"cross": {"backend": "TORCH_SDPA"}}, }, }, runtime_overrides={"diffusion_attention_backend": "SAGE_ATTN"}, ) engine_args = stage.to_omegaconf().engine_args assert engine_args.diffusion_attention_backend == "SAGE_ATTN" assert "default" not in engine_args.diffusion_attention_config assert engine_args.diffusion_attention_config.per_role.cross.backend == "TORCH_SDPA" def test_cli_attention_shorthand_keeps_yaml_per_role_only_config(self): stage = StageConfig( stage_id=0, model_stage="dit", stage_type=StageType.DIFFUSION, yaml_engine_args={"diffusion_attention_config": {"per_role": {"cross": {"backend": "TORCH_SDPA"}}}}, runtime_overrides={"diffusion_attention_backend": "SAGE_ATTN"}, ) engine_args = stage.to_omegaconf().engine_args assert engine_args.diffusion_attention_backend == "SAGE_ATTN" assert engine_args.diffusion_attention_config == {"per_role": {"cross": {"backend": "TORCH_SDPA"}}} def test_cli_attention_config_replaces_yaml_shorthand(self): stage = StageConfig( stage_id=0, model_stage="dit", stage_type=StageType.DIFFUSION, yaml_engine_args={"diffusion_attention_backend": "TORCH_SDPA"}, runtime_overrides={"diffusion_attention_config": {"default": {"backend": "SAGE_ATTN"}}}, ) engine_args = stage.to_omegaconf().engine_args assert "diffusion_attention_backend" not in engine_args assert engine_args.diffusion_attention_config.default.backend == "SAGE_ATTN" class TestSentinelDefaultPrecedence: """Caller-typed (non-None) values win over YAML; None values fall through to YAML / dataclass defaults (#3035).""" def _stages(self, cli_overrides): model_type = "qwen3_omni_moe" pipeline_cfg = resolve_pipeline_config( model_type, Q3_OMNI_ALL_STAGES_HF_CONFIG, ) stages, _ = StageConfigFactory._create_legacy_from_registry( pipeline_cfg, cli_overrides=cli_overrides, ) return stages 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_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 = resolve_pipeline_config("qwen3_tts") assert isinstance(pipeline, PipelineConfig) # 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.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 SIP cleanup: stage 0 has no ``custom_process_next_stage_input_func``, and stage 1 does not set ``requires_full_payload_input=True``. Stage 1 uses only ``thinker2talker_token_only`` (sync_process_input_func). The dead ``thinker2talker`` (custom_process_input_func) was removed because ming_flash_omni has no async_chunk support and both functions called ``_build_talker_inputs`` identically. Merge under either async_chunk mode must not re-introduce either full-payload transport gate.""" pipeline = resolve_pipeline_config("ming_flash_omni") assert isinstance(pipeline, PipelineConfig) 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 hook." ) assert stage1.requires_full_payload_input is False assert stage1.custom_process_input_func is None 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 scalar constraint precedence and additive required stop tokens.""" def test_constraints_win(self): deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) pipeline = resolve_pipeline_config( "qwen3_omni_moe", Q3_OMNI_ALL_STAGES_HF_CONFIG, ) assert isinstance(pipeline, PipelineConfig) 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] def test_required_stop_tokens_extend_yaml_defaults(self, tmp_path): pipeline = PipelineConfig( model_type="required_stop_tokens", stages=( StagePipelineConfig( stage_id=0, model_stage="thinker", final_output=True, sampling_constraints={ "detokenize": True, "stop_token_ids": [151704, 151645], }, ), ), ) deploy_path = tmp_path / "required-stop-tokens.yaml" deploy_path.write_text( """ stages: - stage_id: 0 default_sampling_params: max_tokens: 7 stop_token_ids: [100, 151704] """, encoding="utf-8", ) legacy_stage = merge_pipeline_deploy(pipeline, load_deploy_config(deploy_path))[0] structured_stage = VllmOmniConfig.from_pipeline_config( pipeline, deploy_config_path=str(deploy_path), ).stage_by_id(0) expected = { "max_tokens": 7, "detokenize": True, "stop_token_ids": [100, 151704, 151645], } assert legacy_stage.yaml_extras["default_sampling_params"] == expected assert structured_stage.model_config.default_sampling_params == expected class TestPipelineConfigResolvers: @pytest.mark.parametrize("resolver", [obj for obj in OMNI_PIPELINES.values() if callable(obj)]) def test_all_resolvers_reject_bad_types(self, resolver): """Ensure that all resolvers registered reject incorrect config types.""" class NotTheRightHfConfig(PretrainedConfig): pass assert resolver(NotTheRightHfConfig()) is None class TestObjectStorageConfigResolution: """Object-storage URIs must be materialized locally before any HF-style reads. Regression coverage for review on the Run:AI PR: ``get_config`` and the config.json/model_index.json fallbacks crashed with ``HFValidationError`` for ``s3://`` URIs, and the name-match fallback scanned the whole URI, so a bucket named after another pipeline could hijack pipeline selection. """ @pytest.fixture def fake_object_storage(self, monkeypatch, tmp_path): """Replace ObjectStorageModel with a fake that "pulls" into a tmp dir. Returns a recorder: calls to the returned callable write files into the materialized directory, and ``recorder.pulls`` records every ``pull_files`` invocation. """ pulls: list[tuple[str, list[str] | None, list[str] | None]] = [] materialized_files: list[tuple[str, str]] = [] class FakeObjectStorageModel: def __init__(self, url: str): self.dir = str(tmp_path) def pull_files(self, model_path, allow_pattern=None, ignore_pattern=None): pulls.append((model_path, allow_pattern, ignore_pattern)) for name, content in materialized_files: (tmp_path / name).write_text(content) monkeypatch.setattr(config_factory_module, "ObjectStorageModel", FakeObjectStorageModel) class Recorder: def __init__(self) -> None: self.pulls = pulls def set_files(self, files: list[tuple[str, str]]) -> None: materialized_files.clear() materialized_files.extend(files) return Recorder() def test_passthrough_for_non_uri(self, fake_object_storage): assert _materialize_object_storage_configs("org/model") == "org/model" assert _materialize_object_storage_configs("/local/model") == "/local/model" assert fake_object_storage.pulls == [] def test_materialize_pulls_configs_once_per_uri(self, fake_object_storage): uri = "s3://bucket/model" first = _materialize_object_storage_configs(uri) second = _materialize_object_storage_configs(uri) assert first == second assert len(fake_object_storage.pulls) == 1 pulled_uri, allow_pattern, ignore_pattern = fake_object_storage.pulls[0] assert pulled_uri == uri assert "*.json" in allow_pattern assert ignore_pattern is None def test_try_infer_model_type_reads_materialized_config(self, fake_object_storage): """A misleading bucket name must not beat the real config.json content.""" fake_object_storage.set_files( [ ( "config.json", '{"model_type": "qwen3_omni_moe", "architectures": ["Qwen3OmniMoeForConditionalGeneration"]}', ) ] ) uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct" model = StageConfigFactory.try_infer_model_type(model=uri, trust_remote_code=False) assert model == "qwen3_omni_moe" assert fake_object_storage.pulls # materialization happened def test_name_match_fallback_scans_basename_only(self, fake_object_storage): # Empty config.json (CosyVoice3 style) forces the name-match fallback. fake_object_storage.set_files([("config.json", "{}")]) deceptive = "s3://qwen3-tts-models/plain-checkpoint" assert StageConfigFactory.try_infer_model_type(model=deceptive, trust_remote_code=False) is None matching = "s3://any-bucket/my-cosyvoice3-model" assert StageConfigFactory.try_infer_model_type(model=matching, trust_remote_code=False) == "cosyvoice3"