93 lines
3.6 KiB
Python
93 lines
3.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for out of tree registration to OMNI_PIPELINES."""
|
|
|
|
import pytest
|
|
from transformers import PretrainedConfig
|
|
|
|
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
|
|
from vllm_omni.config.stage_config import PipelineConfig, pipeline_cfg_resolver
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.fixture
|
|
def custom_resolver():
|
|
"""Build a reusable custom resolver for PipelineConfigs."""
|
|
|
|
class CustomConfigType(PretrainedConfig):
|
|
pass
|
|
|
|
@pipeline_cfg_resolver(config_type=CustomConfigType)
|
|
def custom_resolver(
|
|
hf_config: CustomConfigType,
|
|
) -> PipelineConfig:
|
|
return PipelineConfig(model_type="resolved_type")
|
|
|
|
return custom_resolver
|
|
|
|
|
|
def test_register_pipeline_config(clean_pipeline_registry):
|
|
"""Ensure that we can register a custom pipeline config to OMNI_PIPELINES."""
|
|
new_model_type = "new_model_type"
|
|
pipe_cfg = PipelineConfig(model_type=new_model_type)
|
|
assert new_model_type not in OMNI_PIPELINES
|
|
register_pipeline(pipe_cfg)
|
|
assert new_model_type in OMNI_PIPELINES
|
|
assert OMNI_PIPELINES[new_model_type] is pipe_cfg
|
|
|
|
|
|
def test_register_pipeline_config_with_model_type(clean_pipeline_registry):
|
|
"""Ensure that we can register a custom pipeline config with an explicit model_type to OMNI_PIPELINES."""
|
|
new_model_type = "new_model_type"
|
|
unused_model_type = "foo"
|
|
pipe_cfg = PipelineConfig(model_type=unused_model_type)
|
|
assert new_model_type not in OMNI_PIPELINES
|
|
assert unused_model_type not in OMNI_PIPELINES
|
|
|
|
# Registering with an explicitly provided model_type uses
|
|
# the passed value instead of the pipeline_cfg.model_type
|
|
register_pipeline(pipe_cfg, new_model_type)
|
|
assert new_model_type in OMNI_PIPELINES
|
|
assert unused_model_type not in OMNI_PIPELINES
|
|
assert OMNI_PIPELINES[new_model_type] is pipe_cfg
|
|
|
|
|
|
def test_register_resolver(custom_resolver, clean_pipeline_registry):
|
|
"""Ensure that we can register a custom resolver to OMNI_PIPELINES."""
|
|
new_model_type = "new_model_type"
|
|
assert new_model_type not in OMNI_PIPELINES
|
|
register_pipeline(custom_resolver, new_model_type)
|
|
assert new_model_type in OMNI_PIPELINES
|
|
assert OMNI_PIPELINES[new_model_type] is custom_resolver
|
|
|
|
|
|
def test_register_resolver_requires_model_type(custom_resolver, clean_pipeline_registry):
|
|
"""Ensure that registering a custom resolver to OMNI_PIPELINES requires an explicit model_type."""
|
|
with pytest.raises(ValueError):
|
|
register_pipeline(custom_resolver)
|
|
|
|
|
|
def test_minimax_h3_disaggregation_is_explicit_opt_in():
|
|
assert "minimax_h3" not in OMNI_PIPELINES
|
|
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
|
|
assert isinstance(pipeline, PipelineConfig)
|
|
assert pipeline.model_type == "minimax_h3_disaggregated"
|
|
|
|
|
|
def test_in_tree_pipelines_pass_topology_validation():
|
|
"""Every statically registered in-tree pipeline must have a valid topology.
|
|
|
|
``PipelineConfig.validate()`` only runs on ``register_pipeline``, which
|
|
in-tree pipelines never go through, so a shipped topology error would
|
|
otherwise stay invisible until serving hangs or misroutes.
|
|
Resolver entries are skipped: they need an ``hf_config`` to produce a
|
|
``PipelineConfig``.
|
|
"""
|
|
failures = {
|
|
model_type: errors
|
|
for model_type, pipeline in OMNI_PIPELINES.items()
|
|
if isinstance(pipeline, PipelineConfig) and (errors := pipeline.validate())
|
|
}
|
|
assert not failures, f"in-tree pipelines failed topology validation: {failures}"
|