Files
vllm-omni/tests/test_config_factory.py
T
2026-05-11 22:49:19 +08:00

1625 lines
62 KiB
Python

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Unit tests for StageConfigFactory and related classes.
"""
import warnings
from dataclasses import dataclass
from pathlib import Path
import pytest
from vllm_omni.config.stage_config import (
_PIPELINE_REGISTRY,
ModelPipeline,
PipelineConfig,
StageConfig,
StageConfigFactory,
StageExecutionType,
StagePipelineConfig,
StageType,
_resolve_scheduler,
build_stage_runtime_overrides,
register_pipeline,
strip_parent_engine_args,
)
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
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."""
import warnings
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."""
import warnings
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_omits_none_deploy_overrides_for_engine_args(self):
"""None deploy overrides must fall through to EngineArgs defaults."""
from vllm_omni.config.stage_config import deploy_override_field_names
config = StageConfig(
stage_id=0,
model_stage="thinker",
runtime_overrides={name: None for name in deploy_override_field_names()},
)
omega_config = config.to_omegaconf()
engine_args = dict(omega_config.engine_args)
assert "devices" not in engine_args
assert "max_batch_size" not in engine_args
for name in deploy_override_field_names() - {"devices"}:
assert name not in engine_args
class TestModelPipeline:
"""Tests for ModelPipeline class."""
def test_valid_linear_dag(self):
"""Test validation of a valid linear DAG."""
stages = [
StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
StageConfig(stage_id=2, model_stage="code2wav", input_sources=[1]),
]
pipeline = ModelPipeline(model_type="test", stages=stages)
errors = pipeline.validate_pipeline()
assert errors == [], f"Unexpected errors: {errors}"
def test_valid_branching_dag(self):
"""Test validation of a valid branching DAG."""
stages = [
StageConfig(stage_id=0, model_stage="input", input_sources=[]),
StageConfig(stage_id=1, model_stage="branch_a", input_sources=[0]),
StageConfig(stage_id=2, model_stage="branch_b", input_sources=[0]),
]
pipeline = ModelPipeline(model_type="test", stages=stages)
errors = pipeline.validate_pipeline()
assert errors == [], f"Unexpected errors: {errors}"
def test_missing_entry_point(self):
"""Test that missing entry point is detected."""
stages = [
StageConfig(stage_id=0, model_stage="stage_a", input_sources=[1]),
StageConfig(stage_id=1, model_stage="stage_b", input_sources=[0]),
]
pipeline = ModelPipeline(model_type="test", stages=stages)
errors = pipeline.validate_pipeline()
assert any("entry point" in e.lower() for e in errors)
def test_missing_dependency(self):
"""Test that missing stage reference is detected."""
stages = [
StageConfig(stage_id=0, model_stage="input", input_sources=[]),
StageConfig(stage_id=1, model_stage="output", input_sources=[99]), # Invalid
]
pipeline = ModelPipeline(model_type="test", stages=stages)
errors = pipeline.validate_pipeline()
assert any("non-existent" in e.lower() for e in errors)
def test_duplicate_stage_ids(self):
"""Test that duplicate stage IDs are detected."""
stages = [
StageConfig(stage_id=0, model_stage="stage_a", input_sources=[]),
StageConfig(stage_id=0, model_stage="stage_b", input_sources=[]), # Duplicate
]
pipeline = ModelPipeline(model_type="test", stages=stages)
errors = pipeline.validate_pipeline()
assert any("duplicate" in e.lower() for e in errors)
def test_self_reference(self):
"""Test that self-references are detected."""
stages = [
StageConfig(stage_id=0, model_stage="entry", input_sources=[]),
StageConfig(stage_id=1, model_stage="self_ref", input_sources=[1]), # Self
]
pipeline = ModelPipeline(model_type="test", stages=stages)
errors = pipeline.validate_pipeline()
assert any("itself" in e.lower() for e in errors)
def test_get_stage_by_id(self):
"""Test getting stage by ID."""
stages = [
StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
]
pipeline = ModelPipeline(model_type="test", stages=stages)
stage = pipeline.get_stage(1)
assert stage is not None
assert stage.model_stage == "talker"
missing = pipeline.get_stage(99)
assert missing is None
def test_empty_pipeline(self):
"""Test validation of empty pipeline."""
pipeline = ModelPipeline(model_type="test", stages=[])
errors = pipeline.validate_pipeline()
assert any("no stages" in e.lower() for e in errors)
class TestStageConfigFactory:
"""Tests for StageConfigFactory class."""
def test_default_diffusion_no_yaml(self):
"""Test single-stage diffusion works without YAML config (@ZJY0516)."""
kwargs = {
"cache_backend": "none",
"cache_config": None,
"dtype": "bfloat16",
}
configs = StageConfigFactory.create_default_diffusion(kwargs)
assert len(configs) == 1
cfg = configs[0]
assert cfg["stage_id"] == 0
assert cfg["stage_type"] == "diffusion"
assert cfg["final_output"] is True
assert cfg["final_output_type"] == "image"
def test_default_diffusion_with_parallel_config(self):
"""Test diffusion config calculates devices from parallel_config."""
@dataclass
class MockParallelConfig:
world_size: int = 4
kwargs = {
"parallel_config": MockParallelConfig(),
"cache_backend": "tea_cache",
}
configs = StageConfigFactory.create_default_diffusion(kwargs)
assert configs[0]["runtime"]["devices"] == "0,1,2,3"
def test_per_stage_override_precedence(self):
"""Test that --stage-0-gpu-memory-utilization overrides global."""
stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[])
cli_overrides = {
"gpu_memory_utilization": 0.5, # Global
"stage_0_gpu_memory_utilization": 0.9, # Per-stage override
}
overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides)
# Per-stage should override global
assert overrides["gpu_memory_utilization"] == 0.9
def test_cli_override_forwards_engine_registered_args(self):
"""Test that any engine-registered CLI arg is forwarded (@wuhang2014)."""
stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[])
cli_overrides = {
"gpu_memory_utilization": 0.9, # Well-known param
"custom_engine_flag": True, # Not orchestrator-owned, so forwarded
}
overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides)
assert overrides["gpu_memory_utilization"] == 0.9
assert overrides["custom_engine_flag"] is True
def test_cli_override_excludes_internal_keys(self):
"""Test that internal/orchestrator keys are not forwarded."""
stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[])
cli_overrides = {
"gpu_memory_utilization": 0.9,
"model": "some_model", # Internal
"stage_configs_path": "/path", # Internal
"batch_timeout": 10, # Internal
}
overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides)
assert overrides["gpu_memory_utilization"] == 0.9
assert "model" not in overrides
assert "stage_configs_path" not in overrides
assert "batch_timeout" not in overrides
def test_per_stage_override_excludes_internal_keys(self):
"""Test that per-stage overrides also skip internal keys."""
stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[])
cli_overrides = {
"stage_0_gpu_memory_utilization": 0.9,
"stage_0_model": "override_model", # Internal, should be skipped
"stage_0_batch_timeout": 5, # Internal, should be skipped
}
overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides)
assert overrides["gpu_memory_utilization"] == 0.9
assert "model" not in overrides
assert "batch_timeout" not in overrides
class TestStageResolutionHelpers:
"""Tests for shared stage override / filtering helpers."""
def test_build_stage_runtime_overrides_ignores_other_stage_and_internal_keys(self):
# Pass the same filter set the function uses by default
# (orchestrator-only fields plus SHARED_FIELDS so ``model`` is
# treated as not-per-stage-overridable).
overrides = build_stage_runtime_overrides(
0,
{
"gpu_memory_utilization": 0.5,
"stage_0_gpu_memory_utilization": 0.9,
"stage_1_gpu_memory_utilization": 0.1,
"stage_0_model": "should_be_ignored",
"parallel_config": {"world_size": 2},
},
internal_keys=internal_blacklist_keys() | SHARED_FIELDS,
)
assert overrides["gpu_memory_utilization"] == 0.9
assert "model" not in overrides
assert "parallel_config" not in overrides
def test_strip_parent_engine_args_reports_only_surprising_parent_overrides(self):
from dataclasses import fields as dc_fields
from vllm.engine.arg_utils import EngineArgs
parent_fields = {f.name: f for f in dc_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):
from dataclasses import fields as dc_fields
from vllm.engine.arg_utils import EngineArgs
parent_fields = {f.name: f for f in dc_fields(EngineArgs)}
filtered, overridden = strip_parent_engine_args(
{
"model": "some/model",
"stage_configs_path": "/tmp/stages.yaml",
"allowed_local_media_path": "/data/qwentts",
"allowed_media_domains": ["example.com"],
},
parent_fields=parent_fields,
keep_keys={"allowed_local_media_path", "allowed_media_domains"},
strip_keys={"stage_configs_path"},
no_warn_keys={"model"},
)
assert filtered == {
"allowed_local_media_path": "/data/qwentts",
"allowed_media_domains": ["example.com"],
}
assert overridden == []
class TestPipelineYamlParsing:
"""Tests for pipeline YAML file parsing (@ZJY0516)."""
def test_parse_qwen3_omni_moe_yaml(self, tmp_path):
"""Test parsing the qwen3_omni_moe pipeline YAML."""
yaml_content = """\
model_type: qwen3_omni_moe
async_chunk: false
stages:
- stage_id: 0
model_stage: thinker
stage_type: llm
input_sources: []
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
hf_config_name: thinker_config
final_output: true
final_output_type: text
is_comprehension: true
- stage_id: 1
model_stage: talker
stage_type: llm
input_sources: [0]
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
hf_config_name: talker_config
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
- stage_id: 2
model_stage: code2wav
stage_type: llm
input_sources: [1]
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
hf_config_name: thinker_config
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
final_output: true
final_output_type: audio
"""
yaml_file = tmp_path / "qwen3_omni_moe.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "qwen3_omni_moe")
assert pipeline.model_type == "qwen3_omni_moe"
assert len(pipeline.stages) == 3
# Stage 0: thinker
s0 = pipeline.stages[0]
assert s0.stage_id == 0
assert s0.model_stage == "thinker"
assert s0.stage_type == StageType.LLM
assert s0.input_sources == []
assert s0.worker_type == "ar"
assert s0.final_output is True
assert s0.final_output_type == "text"
assert s0.is_comprehension is True
# Stage 1: talker
s1 = pipeline.stages[1]
assert s1.stage_id == 1
assert s1.input_sources == [0]
assert s1.custom_process_input_func == (
"vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker"
)
assert s1.final_output is False
# Stage 2: code2wav
s2 = pipeline.stages[2]
assert s2.stage_id == 2
assert s2.input_sources == [1]
assert s2.worker_type == "generation"
assert s2.final_output_type == "audio"
def test_parse_yaml_with_legacy_engine_input_source(self, tmp_path):
"""Test backward compatibility with engine_input_source field."""
yaml_content = """\
model_type: legacy_model
stages:
- stage_id: 0
model_stage: entry
stage_type: llm
- stage_id: 1
model_stage: downstream
stage_type: llm
engine_input_source: [0]
"""
yaml_file = tmp_path / "legacy.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_model")
assert pipeline.stages[1].input_sources == [0]
def test_parse_yaml_with_connectors_and_edges(self, tmp_path):
"""Test parsing pipeline with optional connectors and edges."""
yaml_content = """\
model_type: test_model
stages:
- stage_id: 0
model_stage: entry
stage_type: llm
input_sources: []
connectors:
type: ray
edges:
- from: 0
to: 1
"""
yaml_file = tmp_path / "with_connectors.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_model")
assert pipeline.connectors == {"type": "ray"}
assert pipeline.edges == [{"from": 0, "to": 1}]
def test_parsed_pipeline_passes_validation(self, tmp_path):
"""Test that a well-formed YAML produces a valid pipeline."""
yaml_content = """\
model_type: valid_model
stages:
- stage_id: 0
model_stage: entry
stage_type: llm
input_sources: []
final_output: true
final_output_type: text
- stage_id: 1
model_stage: next
stage_type: llm
input_sources: [0]
"""
yaml_file = tmp_path / "valid.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "valid_model")
errors = pipeline.validate_pipeline()
assert errors == [], f"Unexpected validation errors: {errors}"
def test_parse_yaml_migrates_runtime_max_batch_size(self, tmp_path):
"""Test that runtime.max_batch_size is migrated to engine_args.max_num_seqs."""
yaml_content = """\
model_type: legacy_batch
stages:
- stage_id: 0
model_stage: entry
stage_type: llm
input_sources: []
runtime:
devices: "0"
max_batch_size: 16
engine_args:
model_arch: SomeModel
"""
yaml_file = tmp_path / "legacy_batch.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_batch")
s0 = pipeline.stages[0]
assert "max_batch_size" not in s0.yaml_runtime
assert s0.yaml_engine_args.get("max_num_seqs") == 16
def test_parse_diffusion_stage_type(self, tmp_path):
"""Test parsing a diffusion stage type from YAML."""
yaml_content = """\
model_type: diff_model
stages:
- stage_id: 0
model_stage: dit
stage_type: diffusion
input_sources: []
final_output: true
final_output_type: image
"""
yaml_file = tmp_path / "diffusion.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "diff_model")
assert pipeline.stages[0].stage_type == StageType.DIFFUSION
def test_parse_mimo_audio_yaml(self, tmp_path):
"""Test parsing the MiMo Audio pipeline YAML."""
yaml_content = """\
model_type: mimo_audio
stages:
- stage_id: 0
model_stage: fused_thinker_talker
stage_type: llm
input_sources: []
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
is_comprehension: true
final_output: true
final_output_type: text
- stage_id: 1
model_stage: code2wav
stage_type: llm
input_sources: [0]
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav
final_output: true
final_output_type: audio
"""
yaml_file = tmp_path / "mimo_audio.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "mimo_audio")
assert pipeline.model_type == "mimo_audio"
assert len(pipeline.stages) == 2
# Stage 0: fused_thinker_talker
s0 = pipeline.stages[0]
assert s0.stage_id == 0
assert s0.model_stage == "fused_thinker_talker"
assert s0.input_sources == []
assert s0.worker_type == "ar"
assert s0.is_comprehension is True
assert s0.final_output is True
assert s0.final_output_type == "text"
# Stage 1: code2wav
s1 = pipeline.stages[1]
assert s1.stage_id == 1
assert s1.input_sources == [0]
assert s1.worker_type == "generation"
assert s1.custom_process_input_func == (
"vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav"
)
assert s1.final_output is True
assert s1.final_output_type == "audio"
# Pipeline validation
errors = pipeline.validate_pipeline()
assert errors == [], f"Unexpected errors: {errors}"
class TestAsyncChunk:
"""Tests for async_chunk pipeline flag."""
def test_model_pipeline_async_chunk_default(self):
"""Test ModelPipeline defaults async_chunk to False."""
stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
pipeline = ModelPipeline(model_type="test", stages=stages)
assert pipeline.async_chunk is False
def test_model_pipeline_async_chunk_set(self):
"""Test ModelPipeline with async_chunk=True."""
stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
pipeline = ModelPipeline(model_type="test", stages=stages, async_chunk=True)
assert pipeline.async_chunk is True
def test_parse_async_chunk_from_yaml(self, tmp_path):
"""Test that async_chunk is parsed from pipeline YAML."""
yaml_content = """\
model_type: test_async
async_chunk: true
stages:
- stage_id: 0
model_stage: entry
stage_type: llm
input_sources: []
"""
yaml_file = tmp_path / "async.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_async")
assert pipeline.async_chunk is True
def test_parse_missing_async_chunk_defaults_false(self, tmp_path):
"""Test that missing async_chunk defaults to False."""
yaml_content = """\
model_type: test_no_async
stages:
- stage_id: 0
model_stage: entry
stage_type: llm
input_sources: []
"""
yaml_file = tmp_path / "no_async.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_no_async")
assert pipeline.async_chunk is False
class TestPipelineDiscovery:
"""Tests for the central pipeline registry (``pipeline_registry._OMNI_PIPELINES``)."""
def test_registry_has_known_models(self):
"""Built-in pipelines are lazy-loaded from the central declaration
on first access; no eager import or discovery walk needed."""
# ``in`` triggers the lazy-map lookup without forcing a load.
assert "qwen2_5_omni" in _PIPELINE_REGISTRY
assert "qwen3_omni_moe" in _PIPELINE_REGISTRY
assert "qwen3_tts" in _PIPELINE_REGISTRY
def test_registry_loads_pipeline_on_getitem(self):
"""Looking up a registered model_type returns the matching PipelineConfig."""
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
assert pipeline.model_type == "qwen3_omni_moe"
assert len(pipeline.stages) == 3 # thinker + talker + code2wav
def test_registry_returns_none_for_unknown(self):
"""Unknown model_types aren't found; ``get()`` returns None."""
assert "definitely_not_a_real_model" not in _PIPELINE_REGISTRY
assert _PIPELINE_REGISTRY.get("definitely_not_a_real_model") is None
def test_pipeline_config_supports_hf_architectures(self):
"""PipelineConfig accepts hf_architectures for HF-arch fallback
(replaces the old _ARCHITECTURE_MODELS dict)."""
p = PipelineConfig(
model_type="custom_collide",
hf_architectures=("SomeCollidingArch",),
)
assert p.hf_architectures == ("SomeCollidingArch",)
class TestStagePipelineConfig:
def test_frozen(self):
s = StagePipelineConfig(stage_id=0, model_stage="a")
with pytest.raises(AttributeError):
s.model_stage = "changed"
def test_defaults(self):
s = StagePipelineConfig(stage_id=0, model_stage="a")
assert s.execution_type == StageExecutionType.LLM_AR
assert s.input_sources == ()
assert s.final_output is False
assert s.sampling_constraints == {}
assert s.engine_output_type is None
class TestPipelineConfigNew:
def test_frozen(self):
p = PipelineConfig(model_type="t", model_arch="A")
with pytest.raises(AttributeError):
p.model_type = "changed"
def test_validate_valid(self):
p = PipelineConfig(
model_type="t",
model_arch="A",
stages=(
StagePipelineConfig(stage_id=0, model_stage="a"),
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
),
)
assert p.validate() == []
def test_validate_no_stages(self):
p = PipelineConfig(model_type="t", model_arch="A")
assert any("no stages" in e.lower() for e in p.validate())
class TestResolveScheduler:
def test_all_execution_types_handled(self):
for et in StageExecutionType:
_resolve_scheduler(et)
def test_ar_sync_when_false(self):
cls = _resolve_scheduler(StageExecutionType.LLM_AR, async_scheduling=False)
assert cls is not None
assert "Async" not in cls.__name__
def test_ar_async_when_true(self):
cls = _resolve_scheduler(StageExecutionType.LLM_AR, async_scheduling=True)
assert cls is not None
assert "Async" in cls.__name__
def test_generation(self):
cls = _resolve_scheduler(StageExecutionType.LLM_GENERATION)
assert cls is not None
assert "Generation" in cls.__name__
def test_diffusion_returns_none(self):
assert _resolve_scheduler(StageExecutionType.DIFFUSION) is None
class TestPipelineRegistry:
def test_register_and_lookup(self):
p = PipelineConfig(
model_type="__test_only__",
model_arch="A",
stages=(StagePipelineConfig(stage_id=0, model_stage="a"),),
)
register_pipeline(p)
assert _PIPELINE_REGISTRY["__test_only__"] is p
del _PIPELINE_REGISTRY["__test_only__"]
class TestDeployConfigLoading:
def test_deploy_override_fields_include_deploy_schema_fields(self):
from vllm_omni.config.stage_config import deploy_override_field_names
expected_fields = {
"async_chunk",
"async_scheduling",
"compilation_config",
"config_format",
"data_parallel_size",
"devices",
"disable_hybrid_kv_cache_manager",
"distributed_executor_backend",
"dtype",
"enable_chunked_prefill",
"enable_flashinfer_autotune",
"enable_prefix_caching",
"enforce_eager",
"gpu_memory_utilization",
"load_format",
"max_model_len",
"max_num_batched_tokens",
"max_num_seqs",
"mm_processor_cache_gb",
"pipeline_parallel_size",
"profiler_config",
"quantization",
"skip_mm_profiling",
"subtalker_sampling_params",
"tensor_parallel_size",
"tokenizer_mode",
"trust_remote_code",
}
actual_fields = deploy_override_field_names()
assert expected_fields == actual_fields, (
f"added={actual_fields - expected_fields}, removed={expected_fields - actual_fields}"
)
def test_load_qwen3_omni_moe_deploy_config(self):
from pathlib import Path
from vllm_omni.config.stage_config import load_deploy_config
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):
from pathlib import Path
from vllm_omni.config.stage_config import load_deploy_config
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):
from pathlib import Path
from vllm_omni.config.stage_config import load_deploy_config
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):
from pathlib import Path
import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
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):
from pathlib import Path
import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
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):
from vllm_omni.config.stage_config import (
DeployConfig,
PipelineConfig,
StageDeployConfig,
StageExecutionType,
StagePipelineConfig,
merge_pipeline_deploy,
)
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
class TestQwen3OmniPipeline:
def test_registered(self):
import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.qwen2_5_omni.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.qwen2_5_omni.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.qwen2_5_omni.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.qwen2_5_omni.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.qwen3_tts.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.qwen3_tts.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.qwen3_tts.pipeline # noqa: F401
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."""
import vllm_omni.model_executor.models.qwen3_tts.pipeline # noqa: F401
from vllm_omni.config.stage_config import load_deploy_config, 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."""
from vllm_omni.config.stage_config import _deep_merge_stage
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):
import vllm_omni.model_executor.models.ming_flash_omni.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.ming_flash_omni.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.ming_flash_omni.pipeline # noqa: F401
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.
"""
import importlib
import vllm_omni.model_executor.models.ming_flash_omni.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.ming_flash_omni.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.ming_flash_omni.pipeline # noqa: F401
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."""
import vllm_omni.model_executor.models.ming_flash_omni.pipeline # noqa: F401
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
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."""
import vllm_omni.model_executor.models.ming_flash_omni.pipeline # noqa: F401
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
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):
import vllm_omni.model_executor.models.ming_flash_omni.pipeline # noqa: F401
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):
import vllm_omni.model_executor.models.ming_flash_omni.pipeline # noqa: F401
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."""
import vllm_omni.model_executor.models.ming_flash_omni.pipeline # noqa: F401
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
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"
class TestBaseConfigInheritance:
"""Test deploy YAML base_config inheritance."""
def test_ci_inherits_from_main(self):
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.config.stage_config import load_deploy_config
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 == "dummy"
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):
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.config.stage_config import load_deploy_config
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
# Inherited from base
assert s0["temperature"] == 0.4
assert s0["seed"] == 42
def test_pure_inheritance_overlay(self, tmp_path):
"""An overlay with only ``base_config`` inherits everything."""
from vllm_omni.config.stage_config import load_deploy_config
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."""
from vllm_omni.config.stage_config import load_deploy_config
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):
from pathlib import Path
from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config
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 True
def test_xpu_overrides(self):
from pathlib import Path
from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config
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):
from pathlib import Path
from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config
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."""
from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config
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):
from pathlib import Path
import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
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):
from pathlib import Path
import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
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):
import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401
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 "max_num_seqs" not in stages[2].yaml_engine_args
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):
import inspect
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_cli_explicit_keys_kwarg_emits_deprecation(self):
import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
StageConfigFactory._create_from_registry(
"qwen3_omni_moe",
cli_overrides={},
cli_explicit_keys={"max_num_seqs"},
)
assert any(issubclass(x.category, DeprecationWarning) for x in w)
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."""
import vllm_omni.model_executor.models.qwen3_tts.pipeline # noqa: F401
from vllm_omni.config.stage_config import (
_PIPELINE_REGISTRY,
DeployConfig,
merge_pipeline_deploy,
)
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 has no streaming processor, stage 1's
# batch-end processor wires up.
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 is not None
assert sync_stages[1].custom_process_input_func.endswith("talker2code2wav")
class TestSamplingConstraintsPrecedence:
"""Test that pipeline sampling_constraints override deploy defaults."""
def test_constraints_win(self):
from pathlib import Path
import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
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]
# Deploy temperature still flows through
assert stages[0].yaml_extras["default_sampling_params"]["temperature"] == 0.4