wehub-resource-sync
69 lines
2.5 KiB
Python
69 lines
2.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Ensure diffusion stage YAML configs only use valid OmniDiffusionConfig fields.
|
|
|
|
Regression test for https://github.com/vllm-project/vllm-omni/issues/2563
|
|
"""
|
|
|
|
from dataclasses import fields
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import yaml
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
try:
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
except Exception:
|
|
OmniDiffusionConfig = None
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
OmniDiffusionConfig is None,
|
|
reason="OmniDiffusionConfig could not be imported (missing torch?)",
|
|
)
|
|
def test_diffusion_stage_configs_only_contain_valid_fields():
|
|
"""Diffusion stage engine_args must only contain OmniDiffusionConfig fields.
|
|
|
|
Regression test for https://github.com/vllm-project/vllm-omni/issues/2563
|
|
"""
|
|
# Scan both main configs and test configs
|
|
repo_root = Path(__file__).parent.parent
|
|
config_dirs = [
|
|
repo_root / "vllm_omni" / "model_executor" / "stage_configs",
|
|
]
|
|
# Also scan test directories recursively
|
|
test_dir = repo_root / "tests"
|
|
|
|
yaml_paths: list[Path] = []
|
|
for config_dir in config_dirs:
|
|
yaml_paths.extend(sorted(config_dir.glob("*.yaml")))
|
|
yaml_paths.extend(sorted(test_dir.rglob("*.yaml")))
|
|
|
|
valid_fields = {f.name for f in fields(OmniDiffusionConfig)}
|
|
# model_stage is consumed by the stage init layer, not OmniDiffusionConfig
|
|
valid_fields.add("model_stage")
|
|
# model_arch is consumed by the stage init layer for diffusion model class resolution
|
|
valid_fields.add("model_arch")
|
|
# "quantization" is mapped to "quantization_config" by from_kwargs() backwards-compat
|
|
valid_fields.add("quantization")
|
|
|
|
invalid_entries: list[tuple[str, set[str]]] = []
|
|
for yaml_path in yaml_paths:
|
|
with open(yaml_path) as fh:
|
|
config = yaml.safe_load(fh)
|
|
|
|
stages = config.get("stage_args", config.get("stages", []))
|
|
for stage in stages:
|
|
if stage.get("stage_type") != "diffusion":
|
|
continue
|
|
engine_args = stage.get("engine_args", {})
|
|
invalid = set(engine_args.keys()) - valid_fields
|
|
if invalid:
|
|
invalid_entries.append((yaml_path.relative_to(repo_root), invalid))
|
|
|
|
assert not invalid_entries, "Diffusion stage configs contain fields not in OmniDiffusionConfig:\n" + "\n".join(
|
|
f" {name}: {sorted(bad)}" for name, bad in invalid_entries
|
|
)
|