Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 3 additions & 15 deletions tests/config/test_pipeline_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,7 @@

import pytest

from vllm_omni.config.pipeline_registry import (
_DIFFUSION_PIPELINES,
_OMNI_PIPELINES,
_VLLM_OMNI_PIPELINES,
)
from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES
from vllm_omni.config.stage_config import (
_PIPELINE_REGISTRY,
PipelineConfig,
Expand All @@ -23,17 +19,9 @@
class TestCentralRegistryDeclarations:
"""Every in-tree pipeline must be declared exactly once in the central registry."""

def test_union_contains_all_omni(self):
def test_omni_entries_visible_in_registry(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}"
assert key in _PIPELINE_REGISTRY

def test_expected_omni_pipelines_present(self):
# Guard against accidental removal during future refactors.
Expand Down
2 changes: 1 addition & 1 deletion tests/test_config_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -675,7 +675,7 @@ def test_parse_missing_async_chunk_defaults_false(self, tmp_path):


class TestPipelineDiscovery:
"""Tests for the central pipeline registry (``pipeline_registry._VLLM_OMNI_PIPELINES``)."""
"""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
Expand Down
17 changes: 7 additions & 10 deletions vllm_omni/config/pipeline_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,13 @@
Adding a new pipeline:
1. Define the ``PipelineConfig`` instance as a module-level variable in
``vllm_omni/.../pipeline.py``.
2. Add one line to ``_OMNI_PIPELINES`` or ``_DIFFUSION_PIPELINES`` below.
2. Add one line to ``_OMNI_PIPELINES`` below.

Single-stage diffusion models continue to use the
``_create_default_diffusion_stage_cfg`` fallback in
``async_omni_engine.py`` — they don't need a registry entry. The empty
``_DIFFUSION_PIPELINES`` placeholder previously here (#2915) was removed
once #2987 (which would have populated it) was deferred.

``register_pipeline(config)`` in ``stage_config`` is still supported for
out-of-tree plugins and tests that create pipelines at runtime; those override
Expand Down Expand Up @@ -68,12 +74,3 @@
"FISH_SPEECH_PIPELINE",
),
}

# --- Single-stage diffusion pipelines (populated in PR 3/N) ---
_DIFFUSION_PIPELINES: dict[str, tuple[str, str]] = {}

# Union view used by ``_LazyPipelineRegistry``; don't mutate at runtime.
_VLLM_OMNI_PIPELINES: dict[str, tuple[str, str]] = {
**_OMNI_PIPELINES,
**_DIFFUSION_PIPELINES,
}
6 changes: 3 additions & 3 deletions vllm_omni/config/stage_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -265,9 +265,9 @@ def __init__(self) -> None:

def _get_lazy_map(self) -> dict[str, tuple[str, str]]:
if self._lazy_map is None:
from vllm_omni.config.pipeline_registry import _VLLM_OMNI_PIPELINES
from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES

self._lazy_map = _VLLM_OMNI_PIPELINES
self._lazy_map = _OMNI_PIPELINES
return self._lazy_map

def _load_lazy(self, model_type: str) -> PipelineConfig | None:
Expand Down Expand Up @@ -365,7 +365,7 @@ def __iter__(self):
def register_pipeline(pipeline: PipelineConfig) -> None:
"""Register a pipeline config dynamically.

In-tree pipelines are declared in ``pipeline_registry._VLLM_OMNI_PIPELINES``
In-tree pipelines are declared in ``pipeline_registry._OMNI_PIPELINES``
and loaded lazily; calling ``register_pipeline`` is only needed for
out-of-tree plugins or tests that build a ``PipelineConfig`` at runtime.
A dynamic registration overrides the central-registry entry with the same
Expand Down
Loading