[Refactor 1/N] Centralize Omni config resolution (#5140)
Signed-off-by: Chenguang ZHENG <645327136@qq.com>
This commit is contained in:
Chenguang Zheng authored and GitHub committed
1 parent
1413eaab7c
commit
087ac46318
23 files changed
+1035
-1812
No files matched your search
+163
-249
@@ -4,10 +4,8 @@
|
||||
"""Unit tests for vllm_omni.entrypoints.utils module."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from collections import Counter
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
@@ -15,24 +13,20 @@ import torch
|
||||
from pytest_mock import MockerFixture
|
||||
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||||
|
||||
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
|
||||
from vllm_omni.config.yaml_util import create_config
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.engine.arg_utils import OmniEngineArgs
|
||||
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
|
||||
from vllm_omni.engine.stage_init_utils import (
|
||||
get_stage_connector_spec,
|
||||
load_omni_transfer_config_for_model,
|
||||
)
|
||||
from vllm_omni.entrypoints.utils import (
|
||||
_convert_dataclasses_to_dict,
|
||||
_filter_dict_like_object,
|
||||
_try_resolve_omni_model_type,
|
||||
coerce_param_message_types,
|
||||
filter_dataclass_kwargs,
|
||||
filter_stages,
|
||||
load_and_resolve_stage_configs,
|
||||
load_stage_configs_from_model,
|
||||
resolve_model_config_path,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -173,7 +167,7 @@ class TestFilterDictLikeObject:
|
||||
"normal": "value",
|
||||
}
|
||||
|
||||
mocker.patch("vllm_omni.entrypoints.utils.logger")
|
||||
mocker.patch("vllm_omni.config.resolver.logger")
|
||||
result = _filter_dict_like_object(input_dict)
|
||||
|
||||
# Normal key should exist
|
||||
@@ -195,7 +189,7 @@ class TestFilterDictLikeObject:
|
||||
"normal": "value",
|
||||
}
|
||||
|
||||
mocker.patch("vllm_omni.entrypoints.utils.logger")
|
||||
mocker.patch("vllm_omni.config.resolver.logger")
|
||||
result = _filter_dict_like_object(input_dict)
|
||||
|
||||
# Callable should be filtered
|
||||
@@ -215,7 +209,7 @@ class TestConvertDataclassesToDict:
|
||||
"callable": lambda x: x,
|
||||
}
|
||||
|
||||
mocker.patch("vllm_omni.entrypoints.utils.logger")
|
||||
mocker.patch("vllm_omni.config.resolver.logger")
|
||||
result = _convert_dataclasses_to_dict(input_dict)
|
||||
|
||||
# Callable should be filtered out by _filter_dict_like_object
|
||||
@@ -300,254 +294,174 @@ class TestFilterDataclassKwargs:
|
||||
assert "extra_param" not in result["cache_config"]
|
||||
|
||||
|
||||
class TestResolveModelConfigPath:
|
||||
"""Test suite for resolve_model_config_path function with diffusers format models."""
|
||||
|
||||
def test_glm_image_diffusers_format_resolution(self, mocker: MockerFixture):
|
||||
"""Test GlmImagePipeline diffusers class resolves to glm_image config."""
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.utils.file_or_path_exists",
|
||||
return_value=True,
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.utils._try_get_class_name_from_diffusers_config",
|
||||
return_value="GlmImagePipeline",
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.utils.current_omni_platform.get_default_stage_config_path",
|
||||
return_value="vllm_omni/model_executor/stage_configs",
|
||||
class TestResolveOmniConfig:
|
||||
def test_stage_lookup_error_lists_resolved_ids(self):
|
||||
resolved = OmniConfigResolution(
|
||||
config_path=None,
|
||||
stage_configs=(SimpleNamespace(stage_id=1), SimpleNamespace(stage_id=3)),
|
||||
)
|
||||
|
||||
original_exists = os.path.exists
|
||||
with pytest.raises(KeyError, match=r"no stage 2; resolved stages: \[1, 3\]"):
|
||||
resolved.stage_by_id(2)
|
||||
|
||||
def mock_exists(path):
|
||||
if "glm_image.yaml" in str(path):
|
||||
return True
|
||||
return original_exists(path)
|
||||
|
||||
mocker.patch("os.path.exists", side_effect=mock_exists)
|
||||
|
||||
result = resolve_model_config_path("zai-org/GLM-Image")
|
||||
|
||||
assert result is not None
|
||||
assert "glm_image.yaml" in result
|
||||
|
||||
def test_uses_registry_default_deploy_when_hf_model_type_differs(
|
||||
self,
|
||||
mocker: MockerFixture,
|
||||
tmp_path,
|
||||
):
|
||||
from vllm_omni.config.stage_config import _DEPLOY_DIR as _REAL_DEPLOY_DIR
|
||||
from vllm_omni.entrypoints import utils as utils_mod
|
||||
|
||||
hf_config = SimpleNamespace(
|
||||
model_type="minicpmo",
|
||||
architectures=["MiniCPMO"],
|
||||
version="4.5",
|
||||
)
|
||||
model = "local/minicpmo-registry-default-test"
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.utils.get_config",
|
||||
return_value=hf_config,
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.config.config_factory.get_config",
|
||||
return_value=hf_config,
|
||||
)
|
||||
# A bare deploy/<hf_model_type>.yaml must not win over the registered
|
||||
# pipeline default when HF model_type is not an OMNI_PIPELINES key.
|
||||
(tmp_path / "minicpmo.yaml").write_text("stages: []\n", encoding="utf-8")
|
||||
(tmp_path / "minicpmo_4_5.yaml").write_bytes((_REAL_DEPLOY_DIR / "minicpmo_4_5.yaml").read_bytes())
|
||||
mocker.patch.object(utils_mod, "_DEPLOY_DIR", tmp_path)
|
||||
|
||||
result = resolve_model_config_path(model)
|
||||
|
||||
assert result is not None
|
||||
assert Path(result).as_posix().endswith("minicpmo_4_5.yaml")
|
||||
assert Path(result).name == "minicpmo_4_5.yaml"
|
||||
transfer_config = load_omni_transfer_config_for_model(model, result)
|
||||
assert transfer_config is not None
|
||||
|
||||
sender = get_stage_connector_spec(
|
||||
omni_transfer_config=transfer_config,
|
||||
stage_id=1,
|
||||
async_chunk=True,
|
||||
)
|
||||
receiver = get_stage_connector_spec(
|
||||
omni_transfer_config=transfer_config,
|
||||
stage_id=2,
|
||||
async_chunk=True,
|
||||
)
|
||||
assert sender["extra"]["role"] == "sender"
|
||||
assert receiver["extra"]["role"] == "receiver"
|
||||
assert sender["extra"]["connector_get_sleep_s"] == 0.01
|
||||
assert sender["extra"]["connector_get_max_wait_first_chunk"] == 3000
|
||||
assert sender["extra"]["connector_get_max_wait"] == 300
|
||||
|
||||
def test_object_storage_uri_resolves_via_materialized_configs(
|
||||
self,
|
||||
mocker: MockerFixture,
|
||||
tmp_path,
|
||||
):
|
||||
"""An object-storage URI must not reach HF helpers; resolution reads the
|
||||
locally materialized config files instead."""
|
||||
hf_config = SimpleNamespace(
|
||||
model_type="qwen3_omni_moe",
|
||||
architectures=["Qwen3OmniMoeForConditionalGeneration"],
|
||||
)
|
||||
uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct"
|
||||
materialize = mocker.patch(
|
||||
"vllm_omni.entrypoints.utils._materialize_object_storage_configs",
|
||||
return_value=str(tmp_path),
|
||||
)
|
||||
get_config = mocker.patch(
|
||||
"vllm_omni.entrypoints.utils.get_config",
|
||||
return_value=hf_config,
|
||||
)
|
||||
|
||||
result = resolve_model_config_path(uri)
|
||||
|
||||
materialize.assert_called_once_with(uri)
|
||||
get_config.assert_called_once()
|
||||
assert get_config.call_args[0][0] == str(tmp_path)
|
||||
assert result is not None
|
||||
assert Path(result).as_posix().endswith("qwen3_omni_moe.yaml")
|
||||
|
||||
|
||||
class TestTryResolveOmniModelType:
|
||||
"""Name matching must only scan the last path component."""
|
||||
|
||||
def test_uri_bucket_name_cannot_hijack_pipeline(self):
|
||||
# "qwen3-tts-models" normalizes to "qwen3ttsmodels", which contains the
|
||||
# registered key "qwen3tts" — only the basename may participate, so a
|
||||
# misleadingly named bucket selects nothing.
|
||||
assert _try_resolve_omni_model_type("s3://qwen3-tts-models/plain-checkpoint") is None
|
||||
|
||||
def test_basename_still_matches_pipeline(self):
|
||||
assert _try_resolve_omni_model_type("gs://any-bucket/Fun-CosyVoice3-0.5B-2512") == "cosyvoice3"
|
||||
|
||||
|
||||
class TestLoadAndResolveStageConfigs:
|
||||
def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture):
|
||||
"""Ensure that kwargs survive default stage creation."""
|
||||
"""Ensure that generic diffusion overrides survive resolution."""
|
||||
engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
|
||||
kwargs = {"dtype": torch.float32, "engine_backend": engine_backend}
|
||||
kwargs["revision"] = "pinned-revision"
|
||||
mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path", return_value=None)
|
||||
mocker.patch("vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=([], None))
|
||||
|
||||
config_path, stage_configs, _ = load_and_resolve_stage_configs(
|
||||
model="black-forest-labs/FLUX.2-klein-4B",
|
||||
kwargs=kwargs,
|
||||
trust_remote_code=False,
|
||||
default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs),
|
||||
mocker.patch(
|
||||
"vllm_omni.config.resolver.StageConfigFactory.create_from_model",
|
||||
return_value=None,
|
||||
)
|
||||
assert config_path is None
|
||||
assert len(stage_configs) == 1
|
||||
assert "dtype" in stage_configs[0]["engine_args"]
|
||||
assert stage_configs[0]["engine_args"]["engine_backend"] == engine_backend
|
||||
assert stage_configs[0]["engine_args"]["revision"] == "pinned-revision"
|
||||
|
||||
def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
|
||||
deploy_path = tmp_path / "qwen3_multi.yaml"
|
||||
deploy_path.write_text(
|
||||
'stages:\n - stage_id: 0\n devices: "0"\n - stage_id: 1\n devices: "1,2,3"\n num_replicas: 3\n',
|
||||
encoding="utf-8",
|
||||
mocker.patch(
|
||||
"vllm_omni.config.resolver._resolve_generic_diffusion_model_class",
|
||||
return_value=(True, "FluxPipeline"),
|
||||
)
|
||||
|
||||
returned_stage_configs = [
|
||||
create_config({"stage_id": 0, "runtime": {"devices": "0"}, "engine_args": {"model": "dummy"}}),
|
||||
create_config(
|
||||
{
|
||||
"stage_id": 1,
|
||||
"runtime": {"devices": "1,2,3", "num_replicas": 3},
|
||||
"engine_args": {"model": "dummy"},
|
||||
}
|
||||
),
|
||||
]
|
||||
load_stage_configs = mocker.patch(
|
||||
"vllm_omni.entrypoints.utils.load_stage_configs_from_model",
|
||||
return_value=(returned_stage_configs, None),
|
||||
)
|
||||
cli_overrides = {
|
||||
"interleave_mm_strings": True,
|
||||
"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
|
||||
"gpu_memory_utilization": 0.55,
|
||||
kwargs = {
|
||||
"dtype": torch.float32,
|
||||
"engine_backend": engine_backend,
|
||||
"revision": "pinned-revision",
|
||||
}
|
||||
|
||||
config_path, stage_configs, _ = load_and_resolve_stage_configs(
|
||||
model="dummy-model",
|
||||
kwargs=cli_overrides,
|
||||
trust_remote_code=True,
|
||||
deploy_config_path=str(deploy_path),
|
||||
)
|
||||
|
||||
load_stage_configs.assert_called_once_with(
|
||||
"dummy-model",
|
||||
trust_remote_code=True,
|
||||
base_engine_args=cli_overrides,
|
||||
deploy_config_path=str(deploy_path),
|
||||
resolved = resolve_omni_config(
|
||||
"black-forest-labs/FLUX.2-klein-4B",
|
||||
trust_remote_code=False,
|
||||
deploy_config_path=None,
|
||||
cli_overrides=kwargs,
|
||||
stage_overrides=None,
|
||||
strategy_config_path=None,
|
||||
)
|
||||
assert config_path == str(deploy_path)
|
||||
assert len(stage_configs) == 2
|
||||
assert stage_configs[1].runtime.num_replicas == 3
|
||||
assert stage_configs[1].runtime.devices == "1,2,3"
|
||||
assert resolved.config_path is None
|
||||
assert len(resolved.stage_configs) == 1
|
||||
engine_args = resolved.stage_configs[0]["engine_args"]
|
||||
assert "dtype" in engine_args
|
||||
assert engine_args["engine_backend"] == engine_backend
|
||||
assert engine_args["revision"] == "pinned-revision"
|
||||
|
||||
def test_filter_stages_selects_mode_stages_without_mutating_stage_config(self, tmp_path):
|
||||
config_path = tmp_path / "deploy.yaml"
|
||||
config_path.write_text(
|
||||
"""modes:
|
||||
- mode: text-to-text
|
||||
stages: [0]
|
||||
- mode: text-to-image
|
||||
stages: [0, 1]
|
||||
""",
|
||||
encoding="utf-8",
|
||||
)
|
||||
stages = [
|
||||
create_config(
|
||||
{
|
||||
"stage_id": 0,
|
||||
"runtime": {"requires_multimodal_data": True},
|
||||
"final_output": False,
|
||||
"final_output_type": None,
|
||||
}
|
||||
),
|
||||
create_config(
|
||||
{
|
||||
"stage_id": 1,
|
||||
"runtime": {"requires_multimodal_data": True},
|
||||
"final_output": True,
|
||||
"final_output_type": "image",
|
||||
}
|
||||
),
|
||||
]
|
||||
|
||||
filtered = filter_stages(str(config_path), stages, {"mode": "text-to-text"})
|
||||
|
||||
assert len(filtered) == 1
|
||||
assert filtered[0].stage_id == 0
|
||||
assert filtered[0].runtime.requires_multimodal_data is True
|
||||
assert filtered[0].final_output is False
|
||||
assert filtered[0].final_output_type is None
|
||||
|
||||
|
||||
class TestLoadStageConfigsFromModel:
|
||||
def test_unresolved_model_does_not_fall_back_to_yaml(self, mocker: MockerFixture, caplog):
|
||||
def test_generic_diffusion_uses_registered_model_metadata(self, mocker: MockerFixture):
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.utils.StageConfigFactory.create_legacy_stage_configs_from_model",
|
||||
return_value=(None, None),
|
||||
"vllm_omni.config.resolver.StageConfigFactory.create_from_model",
|
||||
return_value=None,
|
||||
)
|
||||
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"},
|
||||
)
|
||||
resolve_path = mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path")
|
||||
|
||||
result = load_stage_configs_from_model("unregistered-model", trust_remote_code=False)
|
||||
resolved = resolve_omni_config(
|
||||
"/models/Wan2.2-I2V",
|
||||
trust_remote_code=False,
|
||||
deploy_config_path=None,
|
||||
cli_overrides={
|
||||
"diffusion_load_format": "diffusers",
|
||||
"revision": "pinned-revision",
|
||||
},
|
||||
stage_overrides=None,
|
||||
strategy_config_path=None,
|
||||
)
|
||||
|
||||
assert result == ([], None)
|
||||
resolve_path.assert_not_called()
|
||||
assert "No registered PipelineConfig resolved" in caplog.text
|
||||
assert "deploy_config" in caplog.text
|
||||
stage = resolved.stage_configs[0]
|
||||
resolve_model_class.assert_called_once_with(
|
||||
"/models/Wan2.2-I2V",
|
||||
"diffusers",
|
||||
"pinned-revision",
|
||||
)
|
||||
assert stage.engine_args.model_class_name == "WanImageToVideoPipeline"
|
||||
assert stage.final_output_type == "video"
|
||||
|
||||
def test_registered_pipeline_uses_structured_metadata_and_preserves_override_trust(self, mocker: MockerFixture):
|
||||
endpoint_restriction = SimpleNamespace(name="chat")
|
||||
structured_config = SimpleNamespace(
|
||||
orchestrator_config=SimpleNamespace(deploy_config_path="/resolved/deploy.yaml"),
|
||||
pipeline_config=SimpleNamespace(endpoint_restrictions=(endpoint_restriction,)),
|
||||
)
|
||||
runtime_stage = create_config(
|
||||
{
|
||||
"stage_id": 1,
|
||||
"runtime": {"devices": "1,2,3", "num_replicas": 3},
|
||||
"engine_args": {"model": "dummy-model"},
|
||||
}
|
||||
)
|
||||
legacy_stage = SimpleNamespace(to_omegaconf=lambda: runtime_stage)
|
||||
create_structured = mocker.patch(
|
||||
"vllm_omni.config.resolver.StageConfigFactory.create_from_model",
|
||||
return_value=structured_config,
|
||||
)
|
||||
create_legacy = mocker.patch(
|
||||
"vllm_omni.config.resolver.StageConfigFactory._resolve_legacy_from_registry",
|
||||
return_value=SimpleNamespace(
|
||||
stage_configs=[legacy_stage],
|
||||
pipeline_config=structured_config.pipeline_config,
|
||||
omni_lb_policy="round_robin",
|
||||
),
|
||||
)
|
||||
strategy_specs = {"stage_1": {"dp": 3}}
|
||||
load_strategy = mocker.patch(
|
||||
"vllm_omni.config.resolver._load_strategy_specs",
|
||||
return_value=strategy_specs,
|
||||
)
|
||||
|
||||
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",
|
||||
)
|
||||
create_legacy.assert_called_once_with(
|
||||
structured_config.pipeline_config,
|
||||
expected_overrides,
|
||||
"/resolved/deploy.yaml",
|
||||
strategy_specs=strategy_specs,
|
||||
)
|
||||
load_strategy.assert_called_once_with("strategy.yaml")
|
||||
assert resolved.config_path == "/resolved/deploy.yaml"
|
||||
assert resolved.pipeline_config is structured_config.pipeline_config
|
||||
assert resolved.omni_lb_policy == "round_robin"
|
||||
assert resolved.endpoint_restrictions == (endpoint_restriction,)
|
||||
assert resolved.stage_configs == (runtime_stage,)
|
||||
|
||||
def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture):
|
||||
pipeline = OMNI_PIPELINES["qwen3_tts"]
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
structured_config = SimpleNamespace(
|
||||
orchestrator_config=SimpleNamespace(deploy_config_path=None),
|
||||
pipeline_config=pipeline,
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.config.resolver.StageConfigFactory.create_from_model",
|
||||
return_value=structured_config,
|
||||
)
|
||||
|
||||
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
|
||||
pipeline_stage_ids = [stage.stage_id for stage in resolved.pipeline_config.stages]
|
||||
runtime_stage_ids = [stage.stage_id for stage in resolved.stage_configs]
|
||||
assert pipeline_stage_ids == runtime_stage_ids
|
||||
assert resolved.pipeline_config.stages[-1].model_stage == "forced_aligner"
|
||||
assert len(resolved.pipeline_config.stages) == len(pipeline.stages) + 1
|
||||
|
||||
|
||||
class TestCumulativeStreamingCoercion:
|
||||
|
||||
Reference in new issue
Block a user