2026-09-01 08:41:23 +08:00
|
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
|
|
2026-02-24 22:20:01 +08:00
|
|
|
"""Unit tests for vllm_omni.entrypoints.utils module."""
|
|
|
|
|
|
2026-09-03 14:19:08 +08:00
|
|
|
import logging
|
2026-02-24 22:20:01 +08:00
|
|
|
from collections import Counter
|
|
|
|
|
from dataclasses import dataclass
|
2026-07-31 17:38:26 +08:00
|
|
|
from types import SimpleNamespace
|
2026-02-24 22:20:01 +08:00
|
|
|
|
|
|
|
|
import pytest
|
2026-04-06 21:24:19 -06:00
|
|
|
import torch
|
2026-02-26 17:08:39 +08:00
|
|
|
from pytest_mock import MockerFixture
|
2026-04-25 21:11:53 -06:00
|
|
|
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
2026-02-24 22:20:01 +08:00
|
|
|
|
2026-09-17 12:47:38 +09:00
|
|
|
from vllm_omni.config.composable_parallel import (
|
|
|
|
|
Broadcast,
|
|
|
|
|
FanInByStage,
|
|
|
|
|
MeshAxisSpec,
|
|
|
|
|
RouteByStage,
|
|
|
|
|
StrategySpec,
|
|
|
|
|
TakeRank,
|
|
|
|
|
)
|
2026-09-08 10:30:26 +08:00
|
|
|
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
|
|
|
|
from vllm_omni.config.resolver import (
|
|
|
|
|
OmniConfigResolution,
|
|
|
|
|
_convert_dataclasses_to_dict,
|
|
|
|
|
_filter_dict_like_object,
|
|
|
|
|
resolve_omni_config,
|
|
|
|
|
)
|
|
|
|
|
from vllm_omni.config.stage_config import PipelineConfig
|
2026-02-24 22:20:01 +08:00
|
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
|
from vllm_omni.engine.arg_utils import OmniEngineArgs
|
|
|
|
|
from vllm_omni.entrypoints.utils import (
|
2026-04-25 21:11:53 -06:00
|
|
|
coerce_param_message_types,
|
2026-02-24 22:20:01 +08:00
|
|
|
filter_dataclass_kwargs,
|
|
|
|
|
)
|
2026-09-08 18:17:00 +08:00
|
|
|
from vllm_omni.model_executor.models.qwen3_omni.pipeline import QWEN3_OMNI_PIPELINE
|
2026-02-24 22:20:01 +08:00
|
|
|
|
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class TestFilterDictLikeObject:
|
|
|
|
|
"""Test suite for _filter_dict_like_object function."""
|
|
|
|
|
|
|
|
|
|
def test_simple_dict(self):
|
|
|
|
|
"""Test filtering a simple dictionary with no callables."""
|
|
|
|
|
input_dict = {"key1": "value1", "key2": 42, "key3": [1, 2, 3]}
|
|
|
|
|
result = _filter_dict_like_object(input_dict)
|
|
|
|
|
|
|
|
|
|
assert result == input_dict
|
|
|
|
|
assert isinstance(result, dict)
|
|
|
|
|
|
|
|
|
|
def test_dict_with_nested_values(self):
|
|
|
|
|
"""Test filtering dict with nested dict and list values."""
|
|
|
|
|
input_dict = {
|
|
|
|
|
"level1": {
|
|
|
|
|
"level2": {"key": "value"},
|
|
|
|
|
"list": [1, 2, 3],
|
|
|
|
|
},
|
|
|
|
|
"simple": "string",
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
result = _filter_dict_like_object(input_dict)
|
|
|
|
|
|
|
|
|
|
# Nested dicts and lists should be recursively processed
|
|
|
|
|
assert result["simple"] == "string"
|
|
|
|
|
assert isinstance(result["level1"], dict)
|
|
|
|
|
|
|
|
|
|
def test_dict_with_dataclass_values(self):
|
|
|
|
|
"""Test filtering dict containing dataclass values."""
|
|
|
|
|
|
|
|
|
|
@dataclass
|
|
|
|
|
class TestDataclass:
|
|
|
|
|
field1: str
|
|
|
|
|
field2: int
|
|
|
|
|
|
|
|
|
|
obj = TestDataclass(field1="test", field2=42)
|
|
|
|
|
input_dict = {"data": obj, "normal": "value"}
|
|
|
|
|
|
|
|
|
|
result = _filter_dict_like_object(input_dict)
|
|
|
|
|
|
|
|
|
|
# Dataclass should be converted to dict by recursive _convert_dataclasses_to_dict
|
|
|
|
|
assert "data" in result
|
|
|
|
|
assert "normal" in result
|
|
|
|
|
assert result["normal"] == "value"
|
|
|
|
|
|
|
|
|
|
def test_dict_with_counter_values(self):
|
|
|
|
|
"""Test filtering dict containing Counter objects."""
|
|
|
|
|
counter_obj = Counter({"a": 1, "b": 2})
|
|
|
|
|
input_dict = {"counter": counter_obj, "normal": "value"}
|
|
|
|
|
|
|
|
|
|
result = _filter_dict_like_object(input_dict)
|
|
|
|
|
|
|
|
|
|
# Counter should be converted to regular dict
|
|
|
|
|
assert "counter" in result
|
|
|
|
|
assert "normal" in result
|
|
|
|
|
assert result["normal"] == "value"
|
|
|
|
|
|
|
|
|
|
def test_empty_dict(self):
|
|
|
|
|
"""Test filtering an empty dictionary."""
|
|
|
|
|
result = _filter_dict_like_object({})
|
|
|
|
|
assert result == {}
|
|
|
|
|
assert isinstance(result, dict)
|
|
|
|
|
|
|
|
|
|
def test_dict_with_set_values(self):
|
|
|
|
|
"""Test filtering dict with set values."""
|
|
|
|
|
input_dict = {"set_key": {1, 2, 3}, "normal": "value"}
|
|
|
|
|
|
|
|
|
|
result = _filter_dict_like_object(input_dict)
|
|
|
|
|
|
|
|
|
|
assert "set_key" in result
|
|
|
|
|
assert "normal" in result
|
|
|
|
|
# Set should be converted to list by _convert_dataclasses_to_dict
|
|
|
|
|
assert result["normal"] == "value"
|
|
|
|
|
|
|
|
|
|
def test_dict_with_none_values(self):
|
|
|
|
|
"""Test filtering dict with None values."""
|
|
|
|
|
input_dict = {"key1": None, "key2": "value", "key3": 0}
|
|
|
|
|
|
|
|
|
|
result = _filter_dict_like_object(input_dict)
|
|
|
|
|
|
|
|
|
|
assert result == input_dict
|
|
|
|
|
|
|
|
|
|
def test_dict_with_mixed_types(self):
|
|
|
|
|
"""Test filtering dict with mixed value types."""
|
|
|
|
|
input_dict = {
|
|
|
|
|
"string": "hello",
|
|
|
|
|
"int": 42,
|
|
|
|
|
"float": 3.14,
|
|
|
|
|
"bool": True,
|
|
|
|
|
"none": None,
|
|
|
|
|
"list": [1, 2, 3],
|
|
|
|
|
"tuple": (1, 2, 3),
|
|
|
|
|
"set": {1, 2, 3},
|
|
|
|
|
"dict": {"nested": "value"},
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
result = _filter_dict_like_object(input_dict)
|
|
|
|
|
|
|
|
|
|
assert "string" in result
|
|
|
|
|
assert "int" in result
|
|
|
|
|
assert "float" in result
|
|
|
|
|
assert "bool" in result
|
|
|
|
|
assert "none" in result
|
|
|
|
|
assert "list" in result
|
|
|
|
|
assert "tuple" in result
|
|
|
|
|
assert "set" in result
|
|
|
|
|
assert "dict" in result
|
|
|
|
|
|
|
|
|
|
def test_dict_preserves_key_types(self):
|
|
|
|
|
"""Test that dict key types are preserved."""
|
|
|
|
|
input_dict = {
|
|
|
|
|
"string_key": "value1",
|
|
|
|
|
42: "value2",
|
|
|
|
|
(1, 2): "value3", # tuple as key
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
result = _filter_dict_like_object(input_dict)
|
|
|
|
|
|
|
|
|
|
# Keys should remain the same
|
|
|
|
|
assert "string_key" in result
|
|
|
|
|
assert 42 in result
|
|
|
|
|
assert (1, 2) in result
|
|
|
|
|
|
2026-02-26 17:08:39 +08:00
|
|
|
def test_dict_with_recursive_structure(self, mocker: MockerFixture):
|
2026-02-24 22:20:01 +08:00
|
|
|
"""Test filtering dict with recursive/complex nested structure."""
|
|
|
|
|
input_dict = {
|
|
|
|
|
"level1": {
|
|
|
|
|
"level2": {
|
|
|
|
|
"level3": {"key": "value"},
|
|
|
|
|
"callable": lambda x: x,
|
|
|
|
|
}
|
|
|
|
|
},
|
|
|
|
|
"normal": "value",
|
|
|
|
|
}
|
|
|
|
|
|
2026-09-08 10:30:26 +08:00
|
|
|
mocker.patch("vllm_omni.config.resolver.logger")
|
2026-02-26 17:08:39 +08:00
|
|
|
result = _filter_dict_like_object(input_dict)
|
2026-02-24 22:20:01 +08:00
|
|
|
|
2026-02-26 17:08:39 +08:00
|
|
|
# Normal key should exist
|
|
|
|
|
assert "normal" in result
|
|
|
|
|
# Level1 should exist
|
|
|
|
|
assert "level1" in result
|
2026-02-24 22:20:01 +08:00
|
|
|
|
2026-02-26 17:08:39 +08:00
|
|
|
def test_integration_with_convert_dataclasses(self, mocker: MockerFixture):
|
2026-02-24 22:20:01 +08:00
|
|
|
"""Test that _filter_dict_like_object integrates properly with _convert_dataclasses_to_dict."""
|
|
|
|
|
|
|
|
|
|
@dataclass
|
|
|
|
|
class Config:
|
|
|
|
|
name: str
|
|
|
|
|
count: int
|
|
|
|
|
|
|
|
|
|
input_dict = {
|
|
|
|
|
"config": Config(name="test", count=5),
|
|
|
|
|
"func": lambda x: x,
|
|
|
|
|
"normal": "value",
|
|
|
|
|
}
|
|
|
|
|
|
2026-09-08 10:30:26 +08:00
|
|
|
mocker.patch("vllm_omni.config.resolver.logger")
|
2026-02-26 17:08:39 +08:00
|
|
|
result = _filter_dict_like_object(input_dict)
|
2026-02-24 22:20:01 +08:00
|
|
|
|
2026-02-26 17:08:39 +08:00
|
|
|
# Callable should be filtered
|
|
|
|
|
assert "func" not in result
|
|
|
|
|
# Config should be converted to dict
|
|
|
|
|
assert "config" in result
|
|
|
|
|
assert "normal" in result
|
2026-02-24 22:20:01 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class TestConvertDataclassesToDict:
|
|
|
|
|
"""Test suite for _convert_dataclasses_to_dict function."""
|
|
|
|
|
|
2026-02-26 17:08:39 +08:00
|
|
|
def test_uses_filter_dict_like_object(self, mocker: MockerFixture):
|
2026-02-24 22:20:01 +08:00
|
|
|
"""Test that _convert_dataclasses_to_dict uses _filter_dict_like_object for dicts."""
|
|
|
|
|
input_dict = {
|
|
|
|
|
"normal": "value",
|
|
|
|
|
"callable": lambda x: x,
|
|
|
|
|
}
|
|
|
|
|
|
2026-09-08 10:30:26 +08:00
|
|
|
mocker.patch("vllm_omni.config.resolver.logger")
|
2026-02-26 17:08:39 +08:00
|
|
|
result = _convert_dataclasses_to_dict(input_dict)
|
2026-02-24 22:20:01 +08:00
|
|
|
|
2026-02-26 17:08:39 +08:00
|
|
|
# Callable should be filtered out by _filter_dict_like_object
|
|
|
|
|
assert "normal" in result
|
|
|
|
|
assert "callable" not in result
|
2026-02-24 22:20:01 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class TestFilterDataclassKwargs:
|
|
|
|
|
"""Test basic functionality of filter_dataclass_kwargs."""
|
|
|
|
|
|
|
|
|
|
def test_simple_filtering(self):
|
|
|
|
|
"""Test basic dataclass kwargs filtering."""
|
|
|
|
|
|
|
|
|
|
@dataclass
|
|
|
|
|
class SimpleConfig:
|
|
|
|
|
name: str
|
|
|
|
|
count: int
|
|
|
|
|
|
|
|
|
|
kwargs = {"name": "test", "count": 42, "invalid": "should_be_removed"}
|
|
|
|
|
result = filter_dataclass_kwargs(SimpleConfig, kwargs)
|
|
|
|
|
|
|
|
|
|
assert "name" in result
|
|
|
|
|
assert "count" in result
|
|
|
|
|
assert "invalid" not in result
|
|
|
|
|
|
|
|
|
|
def test_invalid_dataclass_raises_error(self):
|
|
|
|
|
"""Test that non-dataclass raises ValueError."""
|
|
|
|
|
with pytest.raises(ValueError, match="is not a dataclass"):
|
|
|
|
|
filter_dataclass_kwargs(dict, {})
|
|
|
|
|
|
|
|
|
|
def test_invalid_kwargs_type_raises_error(self):
|
|
|
|
|
"""Test that non-dict kwargs raises ValueError."""
|
|
|
|
|
|
|
|
|
|
@dataclass
|
|
|
|
|
class SimpleConfig:
|
|
|
|
|
name: str
|
|
|
|
|
|
|
|
|
|
with pytest.raises(ValueError, match="kwargs must be a dictionary"):
|
|
|
|
|
filter_dataclass_kwargs(SimpleConfig, "invalid")
|
|
|
|
|
|
2026-09-03 14:19:08 +08:00
|
|
|
def test_filters_omni_engine_args_unknown_fields(self, caplog):
|
|
|
|
|
"""Test that OmniEngineArgs kwargs are filtered to valid fields only,
|
|
|
|
|
and that the WARNING contract fires for every drop — the affordance
|
|
|
|
|
that lets a ``--stage-overrides`` typo (``{"0":{"kv_cache_dtpye":...}}``)
|
|
|
|
|
surface in the log rather than being silently misapplied downstream.
|
|
|
|
|
"""
|
2026-02-24 22:20:01 +08:00
|
|
|
kwargs = {
|
|
|
|
|
"model": "dummy",
|
|
|
|
|
"stage_id": 1,
|
|
|
|
|
"engine_output_type": "image",
|
|
|
|
|
"unknown_field": "drop_me",
|
|
|
|
|
}
|
|
|
|
|
|
2026-09-03 14:19:08 +08:00
|
|
|
with caplog.at_level(logging.WARNING, logger="vllm_omni.entrypoints.utils"):
|
|
|
|
|
result = filter_dataclass_kwargs(OmniEngineArgs, kwargs)
|
2026-02-24 22:20:01 +08:00
|
|
|
|
|
|
|
|
assert "model" in result
|
|
|
|
|
assert "stage_id" in result
|
|
|
|
|
assert "engine_output_type" in result
|
|
|
|
|
assert "unknown_field" not in result
|
2026-09-03 14:19:08 +08:00
|
|
|
assert any(rec.levelno == logging.WARNING and "unknown_field" in rec.message for rec in caplog.records), (
|
|
|
|
|
f"expected WARNING naming 'unknown_field'; got {[r.message for r in caplog.records]}"
|
|
|
|
|
)
|
2026-02-24 22:20:01 +08:00
|
|
|
|
|
|
|
|
def test_filters_omni_diffusion_config_union_dataclass(self):
|
|
|
|
|
"""Test that OmniDiffusionConfig filters nested dataclass in Union fields."""
|
|
|
|
|
kwargs = {
|
|
|
|
|
"model": "dummy",
|
|
|
|
|
"cache_config": {
|
|
|
|
|
"rel_l1_thresh": 0.3,
|
|
|
|
|
"extra_param": "should_drop",
|
|
|
|
|
},
|
|
|
|
|
"unknown_top": "drop_me",
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
result = filter_dataclass_kwargs(OmniDiffusionConfig, kwargs)
|
|
|
|
|
|
|
|
|
|
assert "model" in result
|
|
|
|
|
assert "cache_config" in result
|
|
|
|
|
assert "unknown_top" not in result
|
|
|
|
|
assert result["cache_config"]["rel_l1_thresh"] == 0.3
|
|
|
|
|
assert "extra_param" not in result["cache_config"]
|
2026-03-26 09:47:01 +08:00
|
|
|
|
|
|
|
|
|
2026-09-08 10:30:26 +08:00
|
|
|
class TestResolveOmniConfig:
|
2026-09-19 21:46:38 +05:30
|
|
|
@pytest.mark.parametrize("model_class_name", ["AnimaPipeline", "AnimaModularPipeline"])
|
|
|
|
|
def test_native_anima_checkpoint_uses_default_diffusion_stage_without_model_config(
|
|
|
|
|
self, tmp_path, mocker: MockerFixture, model_class_name: str
|
|
|
|
|
):
|
|
|
|
|
checkpoint = tmp_path / "anima.safetensors"
|
|
|
|
|
checkpoint.write_text("dummy")
|
|
|
|
|
custom_args = {"components_path": "/tmp/anima-components"}
|
|
|
|
|
mocker.patch(
|
|
|
|
|
"vllm_omni.config.resolver.StageConfigFactory.create_from_model",
|
|
|
|
|
side_effect=AssertionError("native Anima checkpoints should not require model config discovery"),
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
resolved = resolve_omni_config(
|
|
|
|
|
str(checkpoint),
|
|
|
|
|
trust_remote_code=False,
|
|
|
|
|
deploy_config_path=None,
|
|
|
|
|
cli_overrides={"model_class_name": model_class_name, "custom_pipeline_args": custom_args},
|
|
|
|
|
stage_overrides=None,
|
|
|
|
|
strategy_config_path=None,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert resolved.config_path is None
|
|
|
|
|
assert len(resolved.stage_configs) == 1
|
|
|
|
|
stage_config = resolved.stage_configs[0]
|
|
|
|
|
assert stage_config.model_stage == "diffusion"
|
|
|
|
|
assert stage_config.diffusion_config.model_class_name == "AnimaPipeline"
|
|
|
|
|
assert stage_config.diffusion_config.custom_pipeline_args == custom_args
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
|
|
|
("model_class_name", "checkpoint_exists", "deploy_config_path"),
|
|
|
|
|
[
|
|
|
|
|
("AnimaPipeline", True, "deploy.yaml"),
|
|
|
|
|
("AnimaPipeline", False, None),
|
|
|
|
|
("FluxPipeline", True, None),
|
|
|
|
|
],
|
|
|
|
|
)
|
|
|
|
|
def test_native_single_file_fallback_preserves_config_discovery(
|
|
|
|
|
self, tmp_path, mocker: MockerFixture, model_class_name, checkpoint_exists, deploy_config_path
|
|
|
|
|
):
|
|
|
|
|
checkpoint = tmp_path / "model.safetensors"
|
|
|
|
|
if checkpoint_exists:
|
|
|
|
|
checkpoint.write_text("dummy")
|
|
|
|
|
create_from_model = mocker.patch(
|
|
|
|
|
"vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
resolved = resolve_omni_config(
|
|
|
|
|
str(checkpoint),
|
|
|
|
|
trust_remote_code=False,
|
|
|
|
|
deploy_config_path=deploy_config_path,
|
|
|
|
|
cli_overrides={"model_class_name": model_class_name},
|
|
|
|
|
stage_overrides=None,
|
|
|
|
|
strategy_config_path=None,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
create_from_model.assert_called_once_with(
|
|
|
|
|
str(checkpoint),
|
|
|
|
|
trust_remote_code=False,
|
|
|
|
|
cli_overrides={"model_class_name": model_class_name, "trust_remote_code": False},
|
|
|
|
|
deploy_config_path=deploy_config_path,
|
|
|
|
|
strategy_specs=None,
|
|
|
|
|
)
|
|
|
|
|
assert resolved.config_path == deploy_config_path
|
|
|
|
|
|
2026-09-08 18:17:00 +08:00
|
|
|
def test_bare_deploy_name_returns_packaged_resolved_path(self, tmp_path, mocker: MockerFixture):
|
|
|
|
|
deploy_name = "qwen3_omni_moe.yaml"
|
|
|
|
|
deploy_path = tmp_path / deploy_name
|
|
|
|
|
deploy_path.write_text(
|
|
|
|
|
"duplex_session:\n server_vad_model_path: /models/silero_vad.onnx\nstages: []\n",
|
|
|
|
|
encoding="utf-8",
|
|
|
|
|
)
|
|
|
|
|
mocker.patch("vllm_omni.config.config_factory._DEPLOY_DIR", tmp_path)
|
|
|
|
|
mocker.patch("vllm_omni.config.omni_config._DEPLOY_DIR", tmp_path)
|
|
|
|
|
mocker.patch(
|
|
|
|
|
"vllm_omni.config.config_factory.StageConfigFactory.get_pipeline_config",
|
|
|
|
|
return_value=QWEN3_OMNI_PIPELINE,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
resolved = resolve_omni_config(
|
|
|
|
|
"dummy-model",
|
|
|
|
|
trust_remote_code=False,
|
|
|
|
|
deploy_config_path=deploy_name,
|
|
|
|
|
cli_overrides={},
|
|
|
|
|
stage_overrides=None,
|
|
|
|
|
strategy_config_path=None,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert resolved.config_path == str(deploy_path)
|
|
|
|
|
|
2026-09-08 10:30:26 +08:00
|
|
|
def test_stage_lookup_error_lists_resolved_ids(self):
|
|
|
|
|
resolved = OmniConfigResolution(
|
|
|
|
|
config_path=None,
|
|
|
|
|
stage_configs=(SimpleNamespace(stage_id=1), SimpleNamespace(stage_id=3)),
|
2026-03-26 09:47:01 +08:00
|
|
|
)
|
|
|
|
|
|
2026-09-08 10:30:26 +08:00
|
|
|
with pytest.raises(KeyError, match=r"no stage 2; resolved stages: \[1, 3\]"):
|
|
|
|
|
resolved.stage_by_id(2)
|
2026-03-26 09:47:01 +08:00
|
|
|
|
2026-07-18 17:21:07 +08:00
|
|
|
def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture):
|
2026-09-08 10:30:26 +08:00
|
|
|
"""Ensure that generic diffusion overrides survive resolution."""
|
2026-08-16 09:11:20 +08:00
|
|
|
engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
|
2026-09-08 10:30:26 +08:00
|
|
|
mocker.patch(
|
|
|
|
|
"vllm_omni.config.resolver.StageConfigFactory.create_from_model",
|
|
|
|
|
return_value=None,
|
2026-04-06 21:24:19 -06:00
|
|
|
)
|
2026-09-08 10:30:26 +08:00
|
|
|
mocker.patch(
|
|
|
|
|
"vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
|
|
|
|
|
return_value=(True, "FluxPipeline"),
|
2026-05-06 19:33:18 +08:00
|
|
|
)
|
2026-09-08 10:30:26 +08:00
|
|
|
kwargs = {
|
|
|
|
|
"dtype": torch.float32,
|
|
|
|
|
"engine_backend": engine_backend,
|
|
|
|
|
"revision": "pinned-revision",
|
2026-08-13 21:29:03 +08:00
|
|
|
}
|
2026-09-08 10:30:26 +08:00
|
|
|
resolved = resolve_omni_config(
|
|
|
|
|
"black-forest-labs/FLUX.2-klein-4B",
|
|
|
|
|
trust_remote_code=False,
|
|
|
|
|
deploy_config_path=None,
|
|
|
|
|
cli_overrides=kwargs,
|
2026-05-06 19:33:18 +08:00
|
|
|
stage_overrides=None,
|
2026-07-10 22:29:24 -05:00
|
|
|
strategy_config_path=None,
|
2026-05-06 19:33:18 +08:00
|
|
|
)
|
2026-09-08 10:30:26 +08:00
|
|
|
assert resolved.config_path is None
|
|
|
|
|
assert len(resolved.stage_configs) == 1
|
2026-09-17 12:47:38 +09:00
|
|
|
assert resolved.stage_configs[0].diffusion_config.dtype is torch.float32
|
|
|
|
|
assert resolved.stage_configs[0].diffusion_config.engine_backend == engine_backend
|
|
|
|
|
assert resolved.stage_configs[0].diffusion_config.revision == "pinned-revision"
|
2026-05-06 19:33:18 +08:00
|
|
|
|
2026-09-08 10:30:26 +08:00
|
|
|
def test_generic_diffusion_uses_registered_model_metadata(self, mocker: MockerFixture):
|
2026-09-17 12:47:38 +09:00
|
|
|
mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
|
2026-09-08 10:30:26 +08:00
|
|
|
resolve_model_class = mocker.patch(
|
|
|
|
|
"vllm_omni.config.resolver.resolve_model_class_name",
|
|
|
|
|
return_value="WanImageToVideoPipeline",
|
|
|
|
|
)
|
|
|
|
|
mocker.patch(
|
|
|
|
|
"vllm_omni.config.resolver.DiffusionModelRegistry.get_supported_archs",
|
|
|
|
|
return_value={"WanImageToVideoPipeline"},
|
2026-04-10 14:40:37 +08:00
|
|
|
)
|
2026-09-17 12:47:38 +09:00
|
|
|
load_model_class = mocker.patch("vllm_omni.config.resolver.DiffusionModelRegistry._try_load_model_cls")
|
2026-04-10 14:40:37 +08:00
|
|
|
|
2026-09-08 10:30:26 +08:00
|
|
|
resolved = resolve_omni_config(
|
|
|
|
|
"/models/Wan2.2-I2V",
|
|
|
|
|
trust_remote_code=False,
|
|
|
|
|
deploy_config_path=None,
|
2026-09-17 12:47:38 +09:00
|
|
|
cli_overrides={"diffusion_load_format": "diffusers", "revision": "pinned-revision"},
|
2026-09-08 10:30:26 +08:00
|
|
|
stage_overrides=None,
|
|
|
|
|
strategy_config_path=None,
|
|
|
|
|
)
|
2026-04-10 14:40:37 +08:00
|
|
|
|
2026-09-17 12:47:38 +09:00
|
|
|
resolve_model_class.assert_called_once_with("/models/Wan2.2-I2V", "diffusers", "pinned-revision")
|
|
|
|
|
load_model_class.assert_not_called()
|
2026-09-08 10:30:26 +08:00
|
|
|
stage = resolved.stage_configs[0]
|
2026-09-17 12:47:38 +09:00
|
|
|
assert stage.diffusion_config.model_class_name == "WanImageToVideoPipeline"
|
|
|
|
|
assert stage.diffusion_config.revision == "pinned-revision"
|
2026-09-08 10:30:26 +08:00
|
|
|
assert stage.final_output_type == "video"
|
|
|
|
|
|
2026-09-17 12:47:38 +09:00
|
|
|
def test_generic_diffusion_stage_overrides_reach_typed_backend(self, mocker: MockerFixture):
|
|
|
|
|
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
|
|
|
|
|
|
|
|
|
mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None)
|
|
|
|
|
mocker.patch(
|
|
|
|
|
"vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
|
|
|
|
|
return_value=(True, "LTX2Pipeline"),
|
|
|
|
|
)
|
|
|
|
|
resolved = resolve_omni_config(
|
|
|
|
|
"/models/LTX-2.5-Diffusers",
|
|
|
|
|
trust_remote_code=False,
|
|
|
|
|
deploy_config_path=None,
|
|
|
|
|
cli_overrides={
|
|
|
|
|
"extras": {"ltx2_use_conv_vae": False, "keep": "global"},
|
|
|
|
|
"tensor_parallel_size": 1,
|
|
|
|
|
"diffusion_streaming_output": True,
|
|
|
|
|
},
|
|
|
|
|
stage_overrides={
|
|
|
|
|
"0": {
|
|
|
|
|
"extras": {"ltx2_use_conv_vae": True},
|
|
|
|
|
"tensor_parallel_size": 2,
|
|
|
|
|
"devices": "2,3",
|
|
|
|
|
"engine_backend": "custom.backend",
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
strategy_config_path=None,
|
|
|
|
|
)
|
|
|
|
|
stage = resolved.stage_by_id(0)
|
|
|
|
|
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="/models/LTX-2.5-Diffusers")
|
|
|
|
|
|
|
|
|
|
assert stage.runtime_config.devices == "2,3"
|
|
|
|
|
assert stage.parallel_config.world_size == 2
|
|
|
|
|
assert engine_args["engine_backend"] == "custom.backend"
|
|
|
|
|
assert engine_args["extras"]["ltx2_use_conv_vae"] is True
|
|
|
|
|
assert engine_args["extras"]["keep"] == "global"
|
|
|
|
|
assert engine_args["streaming_output"] is True
|
|
|
|
|
|
2026-09-08 10:30:26 +08:00
|
|
|
def test_registered_pipeline_uses_structured_metadata_and_preserves_override_trust(self, mocker: MockerFixture):
|
|
|
|
|
endpoint_restriction = SimpleNamespace(name="chat")
|
2026-09-17 12:47:38 +09:00
|
|
|
typed_stage = SimpleNamespace(stage_id=1)
|
2026-09-08 10:30:26 +08:00
|
|
|
structured_config = SimpleNamespace(
|
2026-09-17 12:47:38 +09:00
|
|
|
orchestrator_config=SimpleNamespace(
|
|
|
|
|
deploy_config_path="/resolved/deploy.yaml",
|
|
|
|
|
omni_lb_policy="round_robin",
|
|
|
|
|
),
|
2026-09-08 10:30:26 +08:00
|
|
|
pipeline_config=SimpleNamespace(endpoint_restrictions=(endpoint_restriction,)),
|
2026-09-17 12:47:38 +09:00
|
|
|
stage_configs=(typed_stage,),
|
|
|
|
|
strategy_omni_lb_policy="round-robin",
|
2026-09-08 10:30:26 +08:00
|
|
|
)
|
|
|
|
|
create_structured = mocker.patch(
|
|
|
|
|
"vllm_omni.config.resolver.StageConfigFactory.create_from_model",
|
|
|
|
|
return_value=structured_config,
|
|
|
|
|
)
|
|
|
|
|
create_legacy = mocker.patch(
|
2026-09-17 12:47:38 +09:00
|
|
|
"vllm_omni.config.resolver.StageConfigFactory.create_legacy_stage_configs_from_model",
|
2026-09-08 10:30:26 +08:00
|
|
|
)
|
2026-09-17 12:47:38 +09:00
|
|
|
strategy_specs = {
|
|
|
|
|
"stage_1": [
|
|
|
|
|
StrategySpec(
|
|
|
|
|
"stage_replica",
|
|
|
|
|
MeshAxisSpec("stage_replica", 2),
|
|
|
|
|
RouteByStage("round_robin"),
|
|
|
|
|
FanInByStage(),
|
|
|
|
|
)
|
|
|
|
|
]
|
|
|
|
|
}
|
2026-09-08 10:30:26 +08:00
|
|
|
load_strategy = mocker.patch(
|
|
|
|
|
"vllm_omni.config.resolver._load_strategy_specs",
|
|
|
|
|
return_value=strategy_specs,
|
|
|
|
|
)
|
2026-09-17 12:47:38 +09:00
|
|
|
apply_strategy = mocker.patch(
|
|
|
|
|
"vllm_omni.config.composable_parallel.apply_strategy_specs",
|
|
|
|
|
return_value=SimpleNamespace(omni_lb_policy="round_robin"),
|
|
|
|
|
)
|
2026-09-08 10:30:26 +08:00
|
|
|
|
|
|
|
|
resolved = resolve_omni_config(
|
|
|
|
|
"dummy-model",
|
|
|
|
|
trust_remote_code=None,
|
|
|
|
|
deploy_config_path="deploy.yaml",
|
|
|
|
|
cli_overrides={"dtype": "bfloat16", "trust_remote_code": True},
|
|
|
|
|
stage_overrides={"1": {"tensor_parallel_size": 2}},
|
|
|
|
|
strategy_config_path="strategy.yaml",
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
expected_overrides = {
|
|
|
|
|
"dtype": "bfloat16",
|
|
|
|
|
"trust_remote_code": True,
|
|
|
|
|
"stage_1_tensor_parallel_size": 2,
|
|
|
|
|
}
|
|
|
|
|
create_structured.assert_called_once_with(
|
|
|
|
|
"dummy-model",
|
|
|
|
|
trust_remote_code=None,
|
|
|
|
|
cli_overrides=expected_overrides,
|
|
|
|
|
deploy_config_path="deploy.yaml",
|
|
|
|
|
strategy_specs=strategy_specs,
|
|
|
|
|
)
|
2026-09-17 12:47:38 +09:00
|
|
|
create_legacy.assert_not_called()
|
2026-09-08 10:30:26 +08:00
|
|
|
load_strategy.assert_called_once_with("strategy.yaml")
|
2026-09-17 12:47:38 +09:00
|
|
|
apply_strategy.assert_not_called()
|
2026-09-08 10:30:26 +08:00
|
|
|
assert resolved.config_path == "/resolved/deploy.yaml"
|
|
|
|
|
assert resolved.pipeline_config is structured_config.pipeline_config
|
2026-09-17 12:47:38 +09:00
|
|
|
assert resolved.omni_lb_policy == "round-robin"
|
2026-09-08 10:30:26 +08:00
|
|
|
assert resolved.endpoint_restrictions == (endpoint_restriction,)
|
2026-09-17 12:47:38 +09:00
|
|
|
assert resolved.stage_configs == (typed_stage,)
|
2026-09-08 10:30:26 +08:00
|
|
|
|
2026-09-17 12:47:38 +09:00
|
|
|
def test_tp_only_strategy_does_not_report_default_lb_policy_as_derived(self, mocker: MockerFixture):
|
2026-09-08 10:30:26 +08:00
|
|
|
structured_config = SimpleNamespace(
|
2026-09-17 12:47:38 +09:00
|
|
|
orchestrator_config=SimpleNamespace(
|
|
|
|
|
deploy_config_path="/resolved/deploy.yaml",
|
|
|
|
|
omni_lb_policy="round-robin",
|
|
|
|
|
),
|
|
|
|
|
pipeline_config=SimpleNamespace(endpoint_restrictions=()),
|
|
|
|
|
stage_configs=(SimpleNamespace(stage_id=0),),
|
|
|
|
|
strategy_omni_lb_policy=None,
|
2026-09-08 10:30:26 +08:00
|
|
|
)
|
|
|
|
|
mocker.patch(
|
|
|
|
|
"vllm_omni.config.resolver.StageConfigFactory.create_from_model",
|
|
|
|
|
return_value=structured_config,
|
|
|
|
|
)
|
2026-09-17 12:47:38 +09:00
|
|
|
mocker.patch(
|
|
|
|
|
"vllm_omni.config.resolver._load_strategy_specs",
|
|
|
|
|
return_value={
|
|
|
|
|
"thinker": [
|
|
|
|
|
StrategySpec(
|
|
|
|
|
"tp",
|
|
|
|
|
MeshAxisSpec("tp", 2),
|
|
|
|
|
Broadcast(),
|
|
|
|
|
TakeRank(),
|
|
|
|
|
)
|
|
|
|
|
]
|
|
|
|
|
},
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
resolved = resolve_omni_config(
|
|
|
|
|
"dummy-model",
|
|
|
|
|
trust_remote_code=True,
|
|
|
|
|
deploy_config_path="deploy.yaml",
|
|
|
|
|
cli_overrides={"omni_lb_policy": "round-robin"},
|
|
|
|
|
stage_overrides=None,
|
|
|
|
|
strategy_config_path="strategy.yaml",
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert resolved.omni_lb_policy is None
|
|
|
|
|
|
|
|
|
|
def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture):
|
|
|
|
|
pipeline = OMNI_PIPELINES["qwen3_tts"]
|
|
|
|
|
assert isinstance(pipeline, PipelineConfig)
|
|
|
|
|
mocker.patch(
|
|
|
|
|
"vllm_omni.config.config_factory.StageConfigFactory.get_pipeline_config",
|
|
|
|
|
return_value=pipeline,
|
|
|
|
|
)
|
2026-09-08 10:30:26 +08:00
|
|
|
|
|
|
|
|
resolved = resolve_omni_config(
|
|
|
|
|
"Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
|
|
|
|
|
trust_remote_code=False,
|
|
|
|
|
deploy_config_path=None,
|
|
|
|
|
cli_overrides={"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"},
|
|
|
|
|
stage_overrides=None,
|
|
|
|
|
strategy_config_path=None,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert resolved.pipeline_config is not None
|
2026-09-17 12:47:38 +09:00
|
|
|
assert [stage.stage_id for stage in resolved.pipeline_config.stages] == [
|
|
|
|
|
stage.stage_id for stage in resolved.stage_configs
|
|
|
|
|
]
|
2026-09-08 10:30:26 +08:00
|
|
|
assert resolved.pipeline_config.stages[-1].model_stage == "forced_aligner"
|
|
|
|
|
assert len(resolved.pipeline_config.stages) == len(pipeline.stages) + 1
|
2026-04-25 21:11:53 -06:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class TestCumulativeStreamingCoercion:
|
|
|
|
|
@pytest.mark.parametrize("skip_clone", [True, False])
|
|
|
|
|
def test_cumulative_default_becomes_delta_if_stream(self, skip_clone):
|
|
|
|
|
"""Ensure cumulative messages are coercible to delta if streaming."""
|
|
|
|
|
sp = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
|
|
|
|
|
sp.skip_clone = skip_clone
|
|
|
|
|
result = coerce_param_message_types([sp], is_streaming=True)[0]
|
|
|
|
|
assert isinstance(result, SamplingParams)
|
|
|
|
|
assert result.output_kind == RequestOutputKind.DELTA
|
|
|
|
|
assert (skip_clone and sp is result) or (not skip_clone and sp is not result)
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("skip_clone", [True, False])
|
|
|
|
|
def test_cumulative_default_becomes_final_only_if_not_stream(self, skip_clone):
|
|
|
|
|
"""Ensure cumulative messages are coercible to final only if not streaming."""
|
|
|
|
|
sp = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
|
|
|
|
|
sp.skip_clone = skip_clone
|
|
|
|
|
result = coerce_param_message_types([sp], is_streaming=False)[0]
|
|
|
|
|
assert isinstance(result, SamplingParams)
|
|
|
|
|
assert result.output_kind == RequestOutputKind.FINAL_ONLY
|
|
|
|
|
assert (skip_clone and sp is result) or (not skip_clone and sp is not result)
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("is_streaming", [True, False])
|
|
|
|
|
@pytest.mark.parametrize("output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY])
|
|
|
|
|
def test_non_cumulative_are_coerced(self, output_kind, is_streaming):
|
|
|
|
|
"""Ensure non-cumulative params are coerced to the target type."""
|
|
|
|
|
sp = SamplingParams(output_kind=output_kind)
|
|
|
|
|
expected = RequestOutputKind.DELTA if is_streaming else RequestOutputKind.FINAL_ONLY
|
|
|
|
|
result = coerce_param_message_types([sp], is_streaming=is_streaming)[0]
|
|
|
|
|
assert isinstance(result, SamplingParams)
|
|
|
|
|
assert result.output_kind == expected
|
|
|
|
|
|
|
|
|
|
def test_coercion_applies_to_all_stages(self):
|
|
|
|
|
"""Ensure all stages are coerced to DELTA for streaming."""
|
|
|
|
|
sp0 = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
|
|
|
|
|
sp1 = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
|
|
|
|
|
result = coerce_param_message_types([sp0, sp1], is_streaming=True)
|
|
|
|
|
assert all([isinstance(r, SamplingParams) for r in result])
|
|
|
|
|
assert result[0].output_kind == RequestOutputKind.DELTA
|
|
|
|
|
assert result[1].output_kind == RequestOutputKind.DELTA
|
2026-09-16 22:29:31 -07:00
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("api_server_count", [1, 2])
|
|
|
|
|
def test_prepare_stage_config_inputs_filters_frontend_options(api_server_count):
|
|
|
|
|
from vllm_omni.entrypoints.utils import prepare_stage_config_inputs
|
|
|
|
|
|
|
|
|
|
kwargs = {
|
|
|
|
|
"api_server_count": api_server_count,
|
|
|
|
|
"disable_log_stats": True,
|
|
|
|
|
"model_tag": "example",
|
|
|
|
|
"max_num_seqs": 8,
|
|
|
|
|
"stage_overrides": '{"0": {"max_num_seqs": 4}}',
|
|
|
|
|
}
|
|
|
|
|
inputs = prepare_stage_config_inputs("example", kwargs, trust_remote_code=None)
|
|
|
|
|
|
|
|
|
|
assert inputs.kwargs == {"max_num_seqs": 8}
|
|
|
|
|
assert inputs.stage_overrides == {"0": {"max_num_seqs": 4}}
|
|
|
|
|
assert kwargs["api_server_count"] == api_server_count
|