87 lines
3.3 KiB
Python
87 lines
3.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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, StagePipelineConfig, pipeline_cfg_resolver
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def build_fake_pipeline_config(model_type: str) -> PipelineConfig:
|
|
return PipelineConfig(
|
|
model_type=model_type, stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
|
|
)
|
|
|
|
|
|
@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 build_fake_pipeline_config("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 = build_fake_pipeline_config(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 = build_fake_pipeline_config(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_omni_pipelines_sorted_alphabetically():
|
|
"""OMNI_PIPELINES keys must stay sorted (case-insensitive) by model_type."""
|
|
assert list(OMNI_PIPELINES) == sorted(OMNI_PIPELINES, key=str.lower)
|