Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 57 additions & 0 deletions tests/dreamzero/test_deploy_configs.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

"""Every shipped DreamZero deploy config must select the AR-Diffusion engine.

``DreamZeroPipeline.__init__`` rejects any other engine, so a config missing
``engine_backend`` fails at worker startup rather than at request time. That
only surfaces in slow multi-GPU e2e serving tests (see issue #5046), so assert
it here on CPU instead.
"""

from pathlib import Path
from types import SimpleNamespace

import pytest
import yaml

from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

# Discover siblings of a known config so a future dreamzero*.yaml is covered
# without editing this test.
DEPLOY_DIR = Path(get_deploy_config_path("dreamzero.yaml")).parent


def _dreamzero_stages() -> list[tuple[str, int, dict]]:
stages = []
for name in sorted(path.name for path in DEPLOY_DIR.glob("dreamzero*.yaml")):
config = yaml.safe_load(Path(get_deploy_config_path(name)).read_text())
for index, stage in enumerate(config.get("stages") or []):
stages.append((name, index, stage))
return stages


def test_dreamzero_deploy_configs_exist() -> None:
assert _dreamzero_stages(), f"no DreamZero deploy configs found under {DEPLOY_DIR}"


@pytest.mark.parametrize(
("config_name", "stage_index", "stage"),
_dreamzero_stages(),
ids=lambda value: value if isinstance(value, str) else None,
)
def test_dreamzero_stage_selects_ar_diffusion_engine(config_name: str, stage_index: int, stage: dict) -> None:
backend = stage.get("engine_backend")
assert backend, (
f"{config_name} stage {stage_index} does not set engine_backend; "
"DreamZeroPipeline requires the AR-Diffusion engine"
)
# Resolving also proves the import path is real, not just a lookalike string.
engine_class = DiffusionEngine.resolve_engine_class(SimpleNamespace(engine_backend=backend))
assert issubclass(engine_class, ARDiffusionEngine), (
f"{config_name} stage {stage_index} resolves to {engine_class.__name__}, not an AR-Diffusion engine"
)
2 changes: 2 additions & 0 deletions vllm_omni/deploy/dreamzero_tp1_cfg2.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ stages:
max_num_seqs: 1
enforce_eager: false
model_class_name: DreamZeroPipeline
# DreamZero runs on the AR-Diffusion engine (engine-level paged KV cache).
engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine
parallel_config:
tensor_parallel_size: 1
cfg_parallel_size: 2
Expand Down
Loading