anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
bf03413c3651bc01dc1f690f951cf98eee61dd45
vllm-omni/tests/test_config_factory.py
T
Honghan ZhuClaude Opus 4.8zhuhonghan
 b8f68174ad [Refactor] Remove dead legacy pipeline.yaml loader and duplicate diff… (#4023) 
Signed-off-by: dph97 <nenbaying@163.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: zhuhonghan <z00806815@china.huawei.com>
2026-06-09 07:34:35 +08:00

1574 lines
65 KiB
Python
Raw Blame History

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