56 lines
2.4 KiB
Python
56 lines
2.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Pipelines for text models vLLM implements and Omni does not own.
|
|
|
|
A model being in vLLM's registry does not make it reachable through
|
|
``Omni(model=...)``: the orchestrator resolves a pipeline by ``model_type``
|
|
before any weight is read, so a ``model_type`` with no entry fails early and
|
|
opaquely. Spark-X2.5 hit exactly that; these tests pin the general case.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
|
from vllm_omni.config.stage_config import StageExecutionType
|
|
from vllm_omni.config.vllm_native_pipelines import LLAMA_PIPELINE, QWEN3_5_PIPELINE
|
|
|
|
|
|
def test_llama_resolves_from_the_registry():
|
|
"""MiniCPM5-2B declares model_type "llama"; without this entry the
|
|
orchestrator cannot resolve it at all."""
|
|
assert OMNI_PIPELINES["llama"] is LLAMA_PIPELINE
|
|
|
|
|
|
def test_it_is_one_autoregressive_text_stage():
|
|
assert LLAMA_PIPELINE.get_validation_errors() == []
|
|
assert len(LLAMA_PIPELINE.stages) == 1
|
|
stage = LLAMA_PIPELINE.stages[0]
|
|
assert stage.execution_type is StageExecutionType.LLM_AR
|
|
assert stage.input_sources == ()
|
|
assert stage.final_output and stage.final_output_type == "text"
|
|
assert stage.owns_tokenizer
|
|
assert stage.model_arch == "LlamaForCausalLM"
|
|
|
|
|
|
def test_the_stage_declares_no_cross_chunk_state():
|
|
"""A causal LM's state is its KV cache, which the backend owns. Claiming
|
|
otherwise would make the scheduler park requests to hold capacity."""
|
|
assert not LLAMA_PIPELINE.stages[0].retains_state_across_chunks
|
|
|
|
|
|
def test_registering_llama_is_keyed_on_architecture_not_checkpoint():
|
|
"""Worth pinning because it is the surprising part: this entry covers every
|
|
checkpoint whose config says "llama", not only the one it was added for.
|
|
That is the right granularity — pipeline shape is a property of the
|
|
architecture — but it is a claim about shape, not a verification of any
|
|
particular weights."""
|
|
assert LLAMA_PIPELINE.model_type == "llama"
|
|
assert LLAMA_PIPELINE.hf_architectures == ("LlamaForCausalLM",)
|
|
|
|
|
|
def test_qwen3_5_conditional_generation_resolves():
|
|
assert OMNI_PIPELINES["qwen3_5"] is QWEN3_5_PIPELINE
|
|
assert QWEN3_5_PIPELINE.get_validation_errors() == []
|
|
assert QWEN3_5_PIPELINE.hf_architectures == ("Qwen3_5ForConditionalGeneration",)
|
|
assert QWEN3_5_PIPELINE.stages[0].execution_type is StageExecutionType.LLM_AR
|