[Config Refactor 2.5/N] Centralize pipeline registry (#2915)

Signed-off-by: lishunyang <lishunyang12@163.com>
This commit is contained in:
SYLAR authored and GitHub committed 2026-04-20 01:18:03 +08:00
1 parent 68f28f9631
commit cd384d95f3
8 files changed
+337 -63

No files matched your search

View File
Whitespace-only changes.
+111
View File
@@ -0,0 +1,111 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the central pipeline registry (2.5/N)."""
from __future__ import annotations
import pytest
from vllm_omni.config.pipeline_registry import (
_DIFFUSION_PIPELINES,
_OMNI_PIPELINES,
_VLLM_OMNI_PIPELINES,
)
from vllm_omni.config.stage_config import (
_PIPELINE_REGISTRY,
PipelineConfig,
StageExecutionType,
StagePipelineConfig,
register_pipeline,
)
class TestCentralRegistryDeclarations:
"""Every in-tree pipeline must be declared exactly once in the central registry."""
def test_union_contains_all_omni(self):
for key in _OMNI_PIPELINES:
assert key in _VLLM_OMNI_PIPELINES
def test_union_contains_all_diffusion(self):
for key in _DIFFUSION_PIPELINES:
assert key in _VLLM_OMNI_PIPELINES
def test_no_duplicate_model_type_between_omni_and_diffusion(self):
overlap = set(_OMNI_PIPELINES) & set(_DIFFUSION_PIPELINES)
assert not overlap, f"Duplicate model_types across omni/diffusion: {overlap}"
def test_expected_omni_pipelines_present(self):
# Guard against accidental removal during future refactors.
assert "qwen2_5_omni" in _OMNI_PIPELINES
assert "qwen2_5_omni_thinker_only" in _OMNI_PIPELINES
assert "qwen3_omni_moe" in _OMNI_PIPELINES
assert "qwen3_tts" in _OMNI_PIPELINES
class TestLazyLoading:
"""Pipelines are imported only on first access."""
def test_contains_without_import(self):
# ``in`` hits the lazy map, not the loaded cache.
assert "qwen3_omni_moe" in _PIPELINE_REGISTRY
def test_getitem_loads_correct_pipeline(self):
pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
assert pipeline.model_type == "qwen3_omni_moe"
assert pipeline.model_arch == "Qwen3OmniMoeForConditionalGeneration"
def test_unknown_model_type_returns_none_via_get(self):
assert _PIPELINE_REGISTRY.get("not_a_real_pipeline") is None
def test_unknown_model_type_raises_keyerror_via_getitem(self):
with pytest.raises(KeyError):
_PIPELINE_REGISTRY["not_a_real_pipeline"]
def test_iteration_yields_registered_pipelines(self):
keys = set(_PIPELINE_REGISTRY)
assert "qwen2_5_omni" in keys
assert "qwen3_omni_moe" in keys
class TestDynamicRegistration:
"""``register_pipeline()`` still works for plugins and tests."""
def test_register_adds_to_registry(self):
custom = PipelineConfig(
model_type="_test_dynamic_registration",
model_arch="DynamicTestModel",
stages=(
StagePipelineConfig(
stage_id=0,
model_stage="test",
execution_type=StageExecutionType.LLM_AR,
input_sources=(),
final_output=True,
),
),
)
register_pipeline(custom)
try:
assert "_test_dynamic_registration" in _PIPELINE_REGISTRY
assert _PIPELINE_REGISTRY["_test_dynamic_registration"] is custom
finally:
# Don't leak the test registration into other tests.
if "_test_dynamic_registration" in _PIPELINE_REGISTRY:
del _PIPELINE_REGISTRY["_test_dynamic_registration"]
def test_dynamic_registration_overrides_lazy_entry(self):
# Build a substitute for qwen3_omni_moe that we can distinguish.
original = _PIPELINE_REGISTRY["qwen3_omni_moe"]
override = PipelineConfig(
model_type="qwen3_omni_moe",
model_arch="OverriddenArch",
stages=original.stages,
)
register_pipeline(override)
try:
assert _PIPELINE_REGISTRY["qwen3_omni_moe"].model_arch == "OverriddenArch"
finally:
# Remove the dynamic override so later tests see the original.
if "qwen3_omni_moe" in _PIPELINE_REGISTRY._loaded:
del _PIPELINE_REGISTRY["qwen3_omni_moe"]
+25 -11
View File
@@ -23,7 +23,7 @@ from vllm_omni.config.stage_config import (
register_pipeline,
strip_parent_engine_args,
)
from vllm_omni.engine.arg_utils import internal_blacklist_keys
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
class TestStageType:
@@ -330,6 +330,9 @@ 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,
{
@@ -339,7 +342,7 @@ class TestStageResolutionHelpers:
"stage_0_model": "should_be_ignored",
"parallel_config": {"world_size": 2},
},
internal_keys=internal_blacklist_keys(),
internal_keys=internal_blacklist_keys() | SHARED_FIELDS,
)
assert overrides["gpu_memory_utilization"] == 0.9
@@ -672,19 +675,27 @@ stages:
class TestPipelineDiscovery:
"""Tests for auto-discovery of pipelines from models/*/pipeline.py."""
"""Tests for the central pipeline registry (``pipeline_registry._VLLM_OMNI_PIPELINES``)."""
def test_discover_populates_registry_with_known_models(self):
"""``_discover_all_pipelines`` imports every pipeline.py so the
registry is populated with the built-in models after one call."""
from vllm_omni.config.stage_config import _discover_all_pipelines
_discover_all_pipelines()
# These models have a pipeline.py in-tree and must be registered.
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)."""
@@ -950,7 +961,10 @@ class TestBaseConfigInheritance:
assert deploy.stages[0].gpu_memory_utilization == 0.9
assert deploy.connectors is not None
assert "connector_of_shared_memory" in deploy.connectors
assert deploy.async_chunk is True
# CI overlay explicitly sets async_chunk: False (see
# tests/utils.py::_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.utils import get_deploy_config_path