From 087ac46318c42d1fa1e4de13a8053f94683898bd Mon Sep 17 00:00:00 2001
From: Chenguang Zheng <645327136@qq.com>
Date: Tue, 8 Sep 2026 10:30:26 +0800
Subject: [PATCH] [Refactor 1/N] Centralize Omni config resolution (#5140)
Signed-off-by: Chenguang ZHENG <645327136@qq.com>
---
docs/design/architecture_overview.md | 23 +-
docs/mkdocs/hooks/generate_argparse.py | 1 +
tests/config/test_config_factory.py | 9 +-
tests/config/test_omni_config.py | 64 ++
.../minimax_h3/test_minimax_h3_contract.py | 13 +-
.../test_diffusion_config_propagation.py | 2 +-
.../diffusion/test_diffusion_step_pipeline.py | 4 +-
.../test_async_omni_engine_stage_init.py | 129 +---
.../test_async_omni_diffusion_config.py | 273 +++-----
tests/entrypoints/test_pd_disaggregation.py | 7 +-
.../test_resolve_dreamzero_config.py | 32 +-
tests/entrypoints/test_serve.py | 112 ++--
tests/entrypoints/test_utils.py | 412 +++++-------
vllm_omni/config/__init__.py | 4 +
vllm_omni/config/config_factory.py | 162 +++--
vllm_omni/config/pipeline_registry.py | 5 +-
vllm_omni/config/resolver.py | 266 ++++++++
vllm_omni/diffusion/data.py | 121 ++--
vllm_omni/engine/arg_utils.py | 2 +-
vllm_omni/engine/async_omni_engine.py | 537 +++-------------
vllm_omni/engine/stage_init_utils.py | 9 +-
vllm_omni/entrypoints/cli/serve.py | 54 +-
vllm_omni/entrypoints/utils.py | 606 +-----------------
23 files changed, 1035 insertions(+), 1812 deletions(-)
create mode 100644 vllm_omni/config/resolver.py
diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md
index 5b14dd8ec..58722ffa0 100644
--- a/docs/design/architecture_overview.md
+++ b/docs/design/architecture_overview.md
@@ -323,18 +323,24 @@ inputs, fields, and ownership rules are described below.
```mermaid
flowchart TB
layer1["Layer 1 · Authoring inputs
PipelineConfig + DeployConfig"]
- layer2["Layer 2 · Resolve once
StageConfigFactory.create_from_model()
VllmOmniConfig.from_pipeline_config()"]
- layer3["Layer 3 · Transport-safe control plane
VllmOmniConfig"]
+ layer2["Layer 2 · Production resolution boundary
resolve_omni_config()
(StageConfigFactory + VllmOmniConfig internally)"]
+ layer3["Layer 3 · Startup hand-off
OmniConfigResolution
PipelineConfig + temporary OmegaConf stage bridge"]
layer4["Layer 4 · Runtime launch planning
StageRuntime"]
layer5["Layer 5 · Engine materialization
VllmConfig / OmniDiffusionConfig"]
layer1 --> layer2 --> layer3 --> layer4 --> layer5
```
-The single resolution boundary is implemented by
-`StageConfigFactory.create_from_model()` and
-`VllmOmniConfig.from_pipeline_config()` in
+The production resolution boundary is `resolve_omni_config()` in
[`vllm_omni/config`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/config).
+It delegates typed construction to `StageConfigFactory.create_from_model()` and
+`VllmOmniConfig.from_pipeline_config()`, then returns an
+`OmniConfigResolution` consumed by both `AsyncOmniEngine` and headless startup.
+Until `StageRuntime` consumes typed stage configs directly, this envelope carries
+the effective `PipelineConfig` alongside OmegaConf-compatible `stage_configs` as
+a temporary runtime bridge; both views describe the same resolved topology,
+including injected stages.
+
The legacy `stage_args` YAML path has been removed. Model topology now resolves
through `PipelineConfig`, with runtime overrides supplied by `DeployConfig`.
@@ -354,9 +360,12 @@ The important ownership rules are:
deploy-time defaults; it does not redefine the model graph.
3. CLI and Python overrides are applied at the resolution boundary, with
per-stage overrides taking precedence over global values where supported.
-4. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
+4. `OmniConfigResolution` is the sole startup hand-off. Its `pipeline_config`
+ and temporary `stage_configs` compatibility view must describe the same
+ topology.
+5. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
is runtime-private state, not a user configuration object.
-5. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the
+6. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the
process that owns the corresponding engine.
## Main features
diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/hooks/generate_argparse.py
index b695185d1..4c3ad3373 100644
--- a/docs/mkdocs/hooks/generate_argparse.py
+++ b/docs/mkdocs/hooks/generate_argparse.py
@@ -124,6 +124,7 @@ def extract_omni_serve_subparser_init():
"make_arg_parser": lambda parser: parser, # no-op for doc
"_ensure_vllm_platform": lambda: None, # no-op for doc
"_nonneg_finite_float": float,
+ "_parse_stage_overrides": __import__("json").loads,
"VLLM_SUBCMD_PARSER_EPILOG": "",
"logger": logger,
"DummySubparsers": DummySubparsers,
diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py
index 3723e5523..54d04cb82 100644
--- a/tests/config/test_config_factory.py
+++ b/tests/config/test_config_factory.py
@@ -6,7 +6,6 @@ Unit tests for StageConfigFactory and related classes.
import importlib
import warnings
-from dataclasses import dataclass
from pathlib import Path
from unittest.mock import patch
@@ -37,6 +36,7 @@ from vllm_omni.config.stage_config import (
normalize_pipeline_cli_overrides,
pipeline_cfg_resolver,
)
+from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -451,13 +451,8 @@ class TestStageConfigFactory:
def test_default_diffusion_with_parallel_config(self):
"""Test diffusion config calculates devices from parallel_config."""
-
- @dataclass
- class MockParallelConfig:
- world_size: int = 4
-
kwargs = {
- "parallel_config": MockParallelConfig(),
+ "parallel_config": DiffusionParallelConfig(tensor_parallel_size=4),
"cache_backend": "tea_cache",
}
configs = StageConfigFactory.create_default_diffusion(kwargs)
diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py
index f91e6af42..aed9c970d 100644
--- a/tests/config/test_omni_config.py
+++ b/tests/config/test_omni_config.py
@@ -1253,6 +1253,70 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
assert cfg.diffusers_call_kwargs == {}
+def test_diffusion_config_none_values_preserve_dataclass_defaults():
+ from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+ normalized = OmniDiffusionConfig.normalize_init_kwargs(
+ {
+ "lora_scale": None,
+ "enable_sleep_mode": None,
+ "diffusers_load_kwargs": None,
+ }
+ )
+
+ assert "lora_scale" not in normalized
+ assert "enable_sleep_mode" not in normalized
+ assert normalized["diffusers_load_kwargs"] == {}
+
+ config = OmniDiffusionConfig.from_kwargs(
+ lora_scale=None,
+ enable_sleep_mode=None,
+ )
+ assert config.lora_scale == 1.0
+ assert config.enable_sleep_mode is False
+
+
+@pytest.mark.parametrize(
+ ("canonical_key", "alias_key", "canonical_value", "alias_value"),
+ [
+ ("lora_scale", "static_lora_scale", 0.75, 0.25),
+ (
+ "quantization_config",
+ "diffusion_quantization_config",
+ {"method": "canonical"},
+ {"method": "diffusion-alias"},
+ ),
+ (
+ "quantization_config",
+ "quantization",
+ {"method": "canonical"},
+ "legacy-alias",
+ ),
+ ("diffusion_kv_cache_dtype", "kv_cache_dtype", "fp8", "fp16"),
+ ("diffusion_kv_cache_skip_steps", "kv_cache_skip_steps", "0-1", "2-3"),
+ ("diffusion_kv_cache_skip_layers", "kv_cache_skip_layers", "1-2", "3-4"),
+ ("streaming_output", "diffusion_streaming_output", False, True),
+ ],
+)
+def test_diffusion_alias_conflicts_prefer_canonical_key(
+ canonical_key,
+ alias_key,
+ canonical_value,
+ alias_value,
+):
+ from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs
+
+ normalized = normalize_omni_diffusion_kwargs(
+ {
+ canonical_key: canonical_value,
+ alias_key: alias_value,
+ }
+ )
+
+ assert normalized[canonical_key] == canonical_value
+ assert alias_key not in normalized
+
+
def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch):
from vllm_omni.platforms import current_omni_platform
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
index 114f5eab3..a94f842c5 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
@@ -177,9 +177,9 @@ def _write_partition_index(path, *, partition, tasks):
def test_modular_diffusers_index_is_resolved_generically(tmp_path):
+ from vllm_omni.config.resolver import resolve_omni_config
from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
- from vllm_omni.entrypoints.utils import resolve_model_config_path
(tmp_path / "modular_model_index.json").write_text(
json.dumps(
@@ -193,7 +193,16 @@ def test_modular_diffusers_index_is_resolved_generically(tmp_path):
assert is_diffusion_model(str(tmp_path))
assert resolve_model_class_name(str(tmp_path)) == "MiniMaxH3ModularPipeline"
- assert resolve_model_config_path(str(tmp_path)) is None
+ resolved = resolve_omni_config(
+ str(tmp_path),
+ trust_remote_code=False,
+ deploy_config_path=None,
+ cli_overrides=None,
+ stage_overrides=None,
+ strategy_config_path=None,
+ )
+ assert resolved.config_path is None
+ assert resolved.stage_by_id(0).engine_args.model_class_name == "MiniMaxH3ModularPipeline"
config = OmniDiffusionConfig(model=str(tmp_path))
config.enrich_config()
diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py
index 164fcfbeb..394f7c04f 100644
--- a/tests/diffusion/test_diffusion_config_propagation.py
+++ b/tests/diffusion/test_diffusion_config_propagation.py
@@ -131,7 +131,7 @@ class TestCreateDefaultDiffusion:
def test_dtype_serialized_as_string(self):
stages = StageConfigFactory.create_default_diffusion({"dtype": torch.float16, "model": "x"})
- assert stages[0]["engine_args"]["dtype"] == "torch.float16"
+ assert stages[0]["engine_args"]["dtype"] == "float16"
def test_cache_backend_defaults_to_none(self):
stages = StageConfigFactory.create_default_diffusion({"model": "x"})
diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py
index 94c3eb48f..28bba6e4d 100644
--- a/tests/diffusion/test_diffusion_step_pipeline.py
+++ b/tests/diffusion/test_diffusion_step_pipeline.py
@@ -15,6 +15,7 @@ from pytest_mock import MockerFixture
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
from tests.helpers.mark import hardware_test
+from vllm_omni.config.config_factory import StageConfigFactory
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
@@ -46,7 +47,6 @@ from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunn
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.utils import RunnerOutput, StepRequestState
-from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.errors import OmniClientError
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
@@ -1182,7 +1182,7 @@ class TestSupportedPipelines:
"""Step-execution protocol checks for supported pipelines."""
def test_default_stage_config_includes_step_execution(self):
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"step_execution": True,
}
diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py
index 528e46f3a..b1771da28 100644
--- a/tests/engine/test_async_omni_engine_stage_init.py
+++ b/tests/engine/test_async_omni_engine_stage_init.py
@@ -11,7 +11,7 @@ import types
import pytest
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, normalize_omni_diffusion_kwargs
+from vllm_omni.diffusion.data import AttentionConfig
from vllm_omni.engine import async_omni_engine as async_omni_engine_module
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.stage_init_utils import (
@@ -1732,133 +1732,6 @@ def test_inject_kv_stage_info_infers_receiver_tp_topology():
assert stage1.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2}
-def test_resolve_stage_configs_injects_global_diffusion_attention_when_missing(monkeypatch):
- import vllm_omni.engine.async_omni_engine as engine_mod
-
- engine = object.__new__(AsyncOmniEngine)
- stage_cfg = types.SimpleNamespace(
- stage_type="diffusion",
- engine_args=types.SimpleNamespace(
- diffusion_attention_config=None,
- lora_path=None,
- lora_scale=None,
- enable_sleep_mode=None,
- quantization_config=None,
- ),
- )
-
- monkeypatch.setattr(
- engine_mod,
- "load_and_resolve_stage_configs",
- lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
- )
-
- _config_path, stage_configs = engine._resolve_stage_configs(
- model="dummy-model",
- kwargs={"diffusion_attention_backend": "FLASH_ATTN"},
- trust_remote_code=False,
- )
-
- diffusion_attention_config = stage_configs[0].engine_args.diffusion_attention_config
- assert isinstance(diffusion_attention_config, AttentionConfig)
- assert diffusion_attention_config.default is not None
- assert diffusion_attention_config.default.backend == "FLASH_ATTN"
-
-
-def test_resolve_stage_configs_preserves_stage_diffusion_attention(monkeypatch):
- import vllm_omni.engine.async_omni_engine as engine_mod
-
- engine = object.__new__(AsyncOmniEngine)
- existing_attention = AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA"))
- stage_cfg = types.SimpleNamespace(
- stage_type="diffusion",
- engine_args=types.SimpleNamespace(
- diffusion_attention_config=existing_attention,
- lora_path=None,
- lora_scale=None,
- enable_sleep_mode=None,
- quantization_config=None,
- ),
- )
-
- monkeypatch.setattr(
- engine_mod,
- "load_and_resolve_stage_configs",
- lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
- )
-
- _config_path, stage_configs = engine._resolve_stage_configs(
- model="dummy-model",
- kwargs={"diffusion_attention_backend": "FLASH_ATTN"},
- trust_remote_code=False,
- )
-
- assert stage_configs[0].engine_args.diffusion_attention_config is existing_attention
-
-
-def test_resolve_stage_configs_does_not_inject_over_stage_diffusion_attention_backend(monkeypatch):
- import vllm_omni.engine.async_omni_engine as engine_mod
-
- engine = object.__new__(AsyncOmniEngine)
- stage_cfg = types.SimpleNamespace(
- stage_type="diffusion",
- engine_args=types.SimpleNamespace(
- diffusion_attention_backend="TORCH_SDPA",
- diffusion_attention_config=None,
- lora_path=None,
- lora_scale=None,
- enable_sleep_mode=None,
- quantization_config=None,
- ),
- )
-
- monkeypatch.setattr(
- engine_mod,
- "load_and_resolve_stage_configs",
- lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
- )
-
- _config_path, stage_configs = engine._resolve_stage_configs(
- model="dummy-model",
- kwargs={"diffusion_attention_backend": "FLASH_ATTN"},
- trust_remote_code=False,
- )
-
- engine_args = stage_configs[0].engine_args
- assert engine_args.diffusion_attention_backend == "TORCH_SDPA"
- assert engine_args.diffusion_attention_config is None
- normalized = normalize_omni_diffusion_kwargs(vars(engine_args))
- assert normalized["diffusion_attention_config"].default.backend == "TORCH_SDPA"
-
-
-def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stage(monkeypatch):
- import vllm_omni.engine.async_omni_engine as engine_mod
-
- engine = object.__new__(AsyncOmniEngine)
- stage_cfg = types.SimpleNamespace(
- stage_type="llm",
- engine_args=types.SimpleNamespace(
- attention_config={"backend": "FLASH_ATTN"},
- enable_sleep_mode=None,
- ),
- )
-
- monkeypatch.setattr(
- engine_mod,
- "load_and_resolve_stage_configs",
- lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
- )
-
- _config_path, stage_configs = engine._resolve_stage_configs(
- model="dummy-model",
- kwargs={"diffusion_attention_backend": "TORCH_SDPA"},
- trust_remote_code=False,
- )
-
- assert stage_configs[0].engine_args.attention_config == {"backend": "FLASH_ATTN"}
- assert not hasattr(stage_configs[0].engine_args, "diffusion_attention_config")
-
-
def test_extract_legacy_stage_metadata_rocm_does_not_inject_diffusion_attention(monkeypatch):
"""ROCm default attention logic only applies to LLM stages, not diffusion."""
from vllm_omni.engine.stage_init_utils import extract_legacy_stage_metadata
diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py
index 616d5486e..dda9ab38c 100644
--- a/tests/entrypoints/test_async_omni_diffusion_config.py
+++ b/tests/entrypoints/test_async_omni_diffusion_config.py
@@ -4,8 +4,12 @@
from types import SimpleNamespace
import pytest
+import torch
+from pydantic import ValidationError
-from vllm_omni.diffusion.data import AttentionConfig
+from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config.resolver import OmniConfigResolution
+from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
from vllm_omni.engine import stage_init_utils
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.entrypoints.cli.serve import OmniServeCommand
@@ -14,9 +18,13 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+def _terminal_config(stage_cfg: dict) -> OmniDiffusionConfig:
+ return OmniDiffusionConfig.from_kwargs(**stage_cfg["engine_args"])
+
+
def test_default_stage_config_includes_cache_backend():
"""Ensure cache knobs survive the default diffusion-stage builder."""
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"cache_backend": "cache_dit",
"cache_config": '{"Fn_compute_blocks": 2}',
@@ -30,12 +38,12 @@ def test_default_stage_config_includes_cache_backend():
assert engine_args["cache_backend"] == "cache_dit"
assert engine_args["cache_config"]["Fn_compute_blocks"] == 2
assert engine_args["vae_use_slicing"] is True
- assert engine_args["parallel_config"].ulysses_degree == 2
+ assert engine_args["parallel_config"]["ulysses_degree"] == 2
assert engine_args["model_stage"] == "diffusion"
def test_default_stage_config_preserves_ulysses_a2a_permute() -> None:
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"ulysses_degree": 4,
"ulysses_a2a_permute": True,
@@ -43,18 +51,18 @@ def test_default_stage_config_preserves_ulysses_a2a_permute() -> None:
)[0]
parallel_config = stage_cfg["engine_args"]["parallel_config"]
- assert parallel_config.ulysses_degree == 4
- assert parallel_config.ulysses_a2a_permute is True
+ assert parallel_config["ulysses_degree"] == 4
+ assert parallel_config["ulysses_a2a_permute"] is True
def test_default_stage_config_preserves_model_extras():
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"extras": {"ltx2_use_conv_vae": True}})[0]
+ stage_cfg = StageConfigFactory.create_default_diffusion({"extras": {"ltx2_use_conv_vae": True}})[0]
assert stage_cfg["engine_args"]["extras"]["ltx2_use_conv_vae"] is True
def test_default_stage_config_preserves_and_overrides_promoted_extras():
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"extras": {
"auxiliary_text_encoder": "/models/extras-llama",
@@ -68,7 +76,7 @@ def test_default_stage_config_preserves_and_overrides_promoted_extras():
assert extras["auxiliary_text_encoder"] == "/models/extras-llama"
assert extras["default_llama_model_id"] == "extras/default-llama"
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"extras": {
"auxiliary_text_encoder": "/models/extras-llama",
@@ -84,43 +92,57 @@ def test_default_stage_config_preserves_and_overrides_promoted_extras():
assert extras["default_llama_model_id"] == "top-level/default-llama"
-def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocker):
+@pytest.mark.parametrize(
+ "stage_overrides",
+ [
+ {"0": {"extras": {"ltx2_use_conv_vae": True}}},
+ '{"0":{"extras":{"ltx2_use_conv_vae":true}}}',
+ ],
+)
+def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocker, stage_overrides):
"""Local/unregistered Diffusers checkpoints still honor stage-0 extras."""
-
- def resolve_with_default(*_args, default_stage_cfg_factory, **_kwargs):
- return None, default_stage_cfg_factory(), None
-
mocker.patch(
- "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
- side_effect=resolve_with_default,
+ "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"),
)
engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
_, stage_configs = engine._resolve_stage_configs(
"/models/LTX-2.5-Diffusers",
- {"stage_overrides": '{"0":{"extras":{"ltx2_use_conv_vae":true}}}'},
+ {"stage_overrides": stage_overrides},
trust_remote_code=False,
)
assert stage_configs[0]["engine_args"]["extras"]["ltx2_use_conv_vae"] is True
+def test_default_stage_rejects_unknown_nested_parallel_config_key():
+ unknown_key = "unknown_parallel_field"
+ with pytest.raises(ValidationError, match=unknown_key):
+ StageConfigFactory.create_default_diffusion(
+ {"parallel_config": {unknown_key: 2}},
+ )
+
+
def test_default_cache_config_used_when_missing():
"""Ensure default cache_config is synthesized when only backend is given."""
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"cache_backend": "cache_dit",
}
)[0]
- cache_config = stage_cfg["engine_args"]["cache_config"]
- assert cache_config is not None
- assert cache_config["Fn_compute_blocks"] == 1
+ cache_config = _terminal_config(stage_cfg).cache_config
+ assert cache_config.Fn_compute_blocks == 1
def test_default_stage_devices_from_sequence_parallel():
"""Ensure runtime devices reflect computed diffusion world size."""
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"ulysses_degree": 2,
"ring_degree": 2,
@@ -132,7 +154,7 @@ def test_default_stage_devices_from_sequence_parallel():
def test_default_stage_devices_and_dp_from_num_gpus():
"""Resolve omitted DP before deriving devices from an explicit GPU count."""
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"num_gpus": 8,
"tensor_parallel_size": 2,
@@ -141,15 +163,14 @@ def test_default_stage_devices_and_dp_from_num_gpus():
)[0]
parallel_config = stage_cfg["engine_args"]["parallel_config"]
- assert parallel_config.data_parallel_size == 2
- assert parallel_config.world_size == 8
+ assert parallel_config["data_parallel_size"] == 2
assert stage_cfg["engine_args"]["num_gpus"] == 8
assert stage_cfg["runtime"]["devices"] == "0,1,2,3,4,5,6,7"
def test_default_stage_config_uses_parallel_size_kwargs():
"""Ensure default diffusion parallel_config uses CLI/API parallel sizes."""
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"pipeline_parallel_size": 2,
"data_parallel_size": 3,
@@ -159,15 +180,15 @@ def test_default_stage_config_uses_parallel_size_kwargs():
)[0]
parallel_config = stage_cfg["engine_args"]["parallel_config"]
- assert parallel_config.pipeline_parallel_size == 2
- assert parallel_config.data_parallel_size == 3
- assert parallel_config.tensor_parallel_size == 4
- assert parallel_config.enable_expert_parallel is True
+ assert parallel_config["pipeline_parallel_size"] == 2
+ assert parallel_config["data_parallel_size"] == 3
+ assert parallel_config["tensor_parallel_size"] == 4
+ assert parallel_config["enable_expert_parallel"] is True
def test_default_stage_config_preserves_omitted_dp_for_runtime_inference():
"""Keep omitted DP unresolved until the runtime WORLD size is known."""
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"pipeline_parallel_size": None,
"data_parallel_size": None,
@@ -180,18 +201,19 @@ def test_default_stage_config_preserves_omitted_dp_for_runtime_inference():
)[0]
parallel_config = stage_cfg["engine_args"]["parallel_config"]
- assert parallel_config.pipeline_parallel_size == 1
- assert parallel_config.data_parallel_size is None
- assert parallel_config.tensor_parallel_size == 1
- assert parallel_config.enable_expert_parallel is False
- assert stage_cfg["engine_args"]["enforce_eager"] is False
- assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "regional"
- assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is True
+ assert parallel_config["pipeline_parallel_size"] == 1
+ assert parallel_config["data_parallel_size"] is None
+ assert parallel_config["tensor_parallel_size"] == 1
+ assert parallel_config["enable_expert_parallel"] is False
+ terminal_config = _terminal_config(stage_cfg)
+ assert terminal_config.enforce_eager is False
+ assert terminal_config.diffusion_compile_granularity == "regional"
+ assert terminal_config.diffusion_compile_dynamic is True
def test_default_stage_config_propagates_ulysses_mode():
"""Ensure UAA mode survives default diffusion-stage creation."""
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"ulysses_degree": 4,
"ulysses_mode": "advanced_uaa",
@@ -199,13 +221,13 @@ def test_default_stage_config_propagates_ulysses_mode():
)[0]
parallel_config = stage_cfg["engine_args"]["parallel_config"]
- assert parallel_config.ulysses_degree == 4
- assert parallel_config.ulysses_mode == "advanced_uaa"
+ assert parallel_config["ulysses_degree"] == 4
+ assert parallel_config["ulysses_mode"] == "advanced_uaa"
def test_default_stage_config_includes_default_sampling_params():
"""Ensure default sampling params survive the default diffusion-stage builder."""
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"default_sampling_params": '{"0": {"generator_device":"cpu", "guidance_scale":7.5}}',
}
@@ -219,13 +241,13 @@ def test_default_stage_config_includes_default_sampling_params():
def test_default_stage_config_includes_diffusion_attention_backend():
"""Ensure diffusion attention shorthand lands in engine_args.diffusion_attention_config."""
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"diffusion_attention_backend": "FLASH_ATTN",
}
)[0]
- diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
+ diffusion_attention_config = _terminal_config(stage_cfg).diffusion_attention_config
assert isinstance(diffusion_attention_config, AttentionConfig)
assert diffusion_attention_config.default is not None
assert diffusion_attention_config.default.backend == "FLASH_ATTN"
@@ -233,7 +255,7 @@ def test_default_stage_config_includes_diffusion_attention_backend():
def test_default_stage_config_includes_diffusion_attention_config():
"""Ensure structured diffusion attention config survives default stage creation."""
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"diffusion_attention_config": {
"default": {"backend": "FLASH_ATTN"},
@@ -242,7 +264,7 @@ def test_default_stage_config_includes_diffusion_attention_config():
}
)[0]
- diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
+ diffusion_attention_config = _terminal_config(stage_cfg).diffusion_attention_config
assert isinstance(diffusion_attention_config, AttentionConfig)
assert diffusion_attention_config.default is not None
assert diffusion_attention_config.default.backend == "FLASH_ATTN"
@@ -252,7 +274,7 @@ def test_default_stage_config_includes_diffusion_attention_config():
def test_default_stage_config_rejects_conflicting_diffusion_attention_inputs():
"""Ensure shorthand and default.backend stay mutually exclusive."""
with pytest.raises(ValueError, match="mutually exclusive"):
- AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ StageConfigFactory.create_default_diffusion(
{
"diffusion_attention_backend": "FLASH_ATTN",
"diffusion_attention_config": {
@@ -264,7 +286,7 @@ def test_default_stage_config_rejects_conflicting_diffusion_attention_inputs():
def test_default_stage_config_engine_args():
"""Ensure default diffusion-stage builder sets and propagates engine_args."""
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
"distributed_executor_backend": "ray",
"boundary_ratio": 0.875,
@@ -283,7 +305,7 @@ def test_default_stage_config_engine_args():
def test_default_stage_config_whitelist_none_fallback():
"""DeployConfig / StageDeployConfig whitelist fields with value None
fall back to OmniDiffusionConfig dataclass defaults."""
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ stage_cfg = StageConfigFactory.create_default_diffusion(
{
# DeployConfig pipeline-wide
"trust_remote_code": None,
@@ -294,12 +316,11 @@ def test_default_stage_config_whitelist_none_fallback():
}
)[0]
- engine_args = stage_cfg["engine_args"]
-
- assert engine_args["trust_remote_code"] is False
- assert engine_args["distributed_executor_backend"] is None
- assert engine_args["dtype"] == "auto"
- assert engine_args["enforce_eager"] is False
+ terminal = _terminal_config(stage_cfg)
+ assert terminal.trust_remote_code is False
+ assert terminal.distributed_executor_backend is None
+ assert terminal.dtype == torch.bfloat16
+ assert terminal.enforce_eager is False
def test_serve_cli_accepts_ulysses_mode():
@@ -321,12 +342,12 @@ def test_serve_cli_accepts_ulysses_mode():
)
explicit_kwargs = args.get_explicit_kwargs_dict()
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+ stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
parallel_config = stage_cfg["engine_args"]["parallel_config"]
assert args.ulysses_mode == "advanced_uaa"
- assert parallel_config.ulysses_degree == 4
- assert parallel_config.ulysses_mode == "advanced_uaa"
+ assert parallel_config["ulysses_degree"] == 4
+ assert parallel_config["ulysses_mode"] == "advanced_uaa"
def test_serve_cli_accepts_text_encoder_tp_size():
@@ -345,11 +366,11 @@ def test_serve_cli_accepts_text_encoder_tp_size():
)
explicit_kwargs = args.get_explicit_kwargs_dict()
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+ stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
parallel_config = stage_cfg["engine_args"]["parallel_config"]
assert args.text_encoder_tp_size == 4
- assert parallel_config.text_encoder_tp_size == 4
+ assert parallel_config["text_encoder_tp_size"] == 4
def test_serve_cli_forwards_model_defined_task_type_to_diffusion_stage():
@@ -368,7 +389,7 @@ def test_serve_cli_forwards_model_defined_task_type_to_diffusion_stage():
)
explicit_kwargs = args.get_explicit_kwargs_dict()
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+ stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
assert args.task_type == "fl2va"
assert stage_cfg["engine_args"]["task_type"] == "fl2va"
@@ -390,7 +411,7 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
)
explicit_kwargs = args.get_explicit_kwargs_dict()
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+ stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
assert args.enable_diffusion_pipeline_profiler is True
assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True
@@ -416,7 +437,7 @@ def test_serve_cli_forwards_distilled_lora_to_diffusion_stage():
)
explicit_kwargs = args.get_explicit_kwargs_dict()
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+ stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
engine_args = stage_cfg["engine_args"]
assert explicit_kwargs["lora_backend"] == "distill"
@@ -447,7 +468,7 @@ def test_serve_cli_forwards_compact_diffusion_offload_config():
)
explicit_kwargs = args.get_explicit_kwargs_dict()
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+ stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
engine_args = stage_cfg["engine_args"]
expected = {
@@ -513,7 +534,7 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
)
explicit_kwargs = args.get_explicit_kwargs_dict()
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+ stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
engine_args = stage_cfg["engine_args"]
assert explicit_kwargs["host_weight_runtime_mode"] == "preferred"
@@ -542,7 +563,7 @@ def test_serve_cli_accepts_diffusion_compile_controls():
)
explicit_kwargs = args.get_explicit_kwargs_dict()
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+ stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
assert args.diffusion_compile_granularity == "full"
assert args.diffusion_compile_dynamic is False
@@ -569,7 +590,7 @@ def test_serve_cli_accepts_diffusion_attention_backend():
)
explicit_kwargs = args.get_explicit_kwargs_dict()
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+ stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
assert args.diffusion_attention_backend == "FASTVIDEO_VSA"
@@ -597,7 +618,7 @@ def test_serve_cli_accepts_request_batch_max_wait_ms():
)
explicit_kwargs = args.get_explicit_kwargs_dict()
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+ stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0]
assert args.request_batch_max_wait_ms == 250.0
assert stage_cfg["engine_args"]["request_batch_max_wait_ms"] == 250.0
@@ -637,7 +658,7 @@ def test_serve_cli_accepts_additional_config():
]
)
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(vars(args))[0]
+ stage_cfg = StageConfigFactory.create_default_diffusion(vars(args))[0]
engine_args = stage_cfg["engine_args"]
@@ -645,79 +666,19 @@ def test_serve_cli_accepts_additional_config():
assert engine_args["additional_config"] == {"torchair_graph_config": {"enabled": True}}
-def test_default_stage_resolves_video_output_from_checkpoint(mocker):
- captured = {}
-
- def resolve_with_default(*args, default_stage_cfg_factory, **kwargs):
- del args, kwargs
- captured.update(default_stage_cfg_factory()[0])
- stage = SimpleNamespace(stage_type="diffusion", engine_args=SimpleNamespace())
- return ("", [stage], None)
-
- resolver = mocker.patch(
- "vllm_omni.engine.async_omni_engine.resolve_model_class_name",
- return_value="MiniMaxH3Pipeline",
- )
- mocker.patch(
- "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
- side_effect=resolve_with_default,
- )
- engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
- engine._strip_single_engine_args = lambda kwargs: kwargs
-
- engine._resolve_stage_configs(
- "/models/MiniMax-H3/FL2VA",
- {},
- trust_remote_code=False,
- )
-
- resolver.assert_called_once_with("/models/MiniMax-H3/FL2VA", "default")
- assert captured["final_output_type"] == "video"
-
-
-def test_default_diffusers_stage_preserves_video_model_identity(mocker):
- captured = {}
-
- def resolve_with_default(*args, default_stage_cfg_factory, **kwargs):
- del args, kwargs
- captured.update(default_stage_cfg_factory()[0])
- stage = SimpleNamespace(stage_type="diffusion", engine_args=SimpleNamespace())
- return ("", [stage], None)
-
- mocker.patch(
- "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index",
- return_value={"_class_name": "WanImageToVideoPipeline"},
- )
- mocker.patch(
- "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
- side_effect=resolve_with_default,
- )
- engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
- engine._strip_single_engine_args = lambda kwargs: kwargs
-
- engine._resolve_stage_configs(
- "/models/Wan2.2-I2V",
- {"diffusion_load_format": "diffusers"},
- trust_remote_code=False,
- )
-
- assert captured["engine_args"]["model_class_name"] == "WanImageToVideoPipeline"
- assert captured["final_output_type"] == "video"
-
-
-def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mocker):
- """Ensure YAML/deploy stage resolution forwards top-level additional_config."""
+def test_resolve_stage_configs_delegates_overrides_to_resolver(mocker):
+ """The engine consumes resolver output without a second merge pass."""
+ additional_config = {"torchair_graph_config": {"enabled": True}}
fake_diffusion_stage = SimpleNamespace(
stage_type="diffusion",
- engine_args=SimpleNamespace(),
+ engine_args=SimpleNamespace(additional_config=additional_config),
)
- fake_llm_stage = SimpleNamespace(
- stage_type="llm",
- engine_args=SimpleNamespace(),
- )
- mocker.patch(
- "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
- return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage], None),
+ resolve_config = mocker.patch(
+ "vllm_omni.engine.async_omni_engine.resolve_omni_config",
+ return_value=OmniConfigResolution(
+ config_path="dummy.yaml",
+ stage_configs=(fake_diffusion_stage,),
+ ),
)
engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
@@ -726,13 +687,15 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
"dummy-model",
{
"deploy_config": "dummy.yaml",
- "additional_config": {"torchair_graph_config": {"enabled": True}},
+ "additional_config": additional_config,
},
trust_remote_code=False,
)
- assert not hasattr(stage_configs[0].engine_args, "additional_config")
- assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}}
+ assert stage_configs == [fake_diffusion_stage]
+ assert resolve_config.call_args.args == ("dummy-model",)
+ assert resolve_config.call_args.kwargs["deploy_config_path"] == "dummy.yaml"
+ assert resolve_config.call_args.kwargs["cli_overrides"]["additional_config"] is additional_config
@pytest.mark.parametrize(
@@ -760,30 +723,6 @@ def test_default_stage_config_includes_quantization_config():
"weights": "weights.bin",
}
- stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"quantization_config": quantization_config})[0]
+ stage_cfg = StageConfigFactory.create_default_diffusion({"quantization_config": quantization_config})[0]
assert stage_cfg["engine_args"]["quantization_config"] == quantization_config
-
-
-def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage(mocker):
- fake_diffusion_stage = SimpleNamespace(
- stage_type="diffusion",
- engine_args=SimpleNamespace(quantization_config=None),
- )
- mocker.patch(
- "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
- return_value=("dummy.yaml", [fake_diffusion_stage], None),
- )
-
- engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
-
- _, stage_configs = engine._resolve_stage_configs(
- "dummy-model",
- {
- "deploy_config": "dummy.yaml",
- "quantization_config": {"method": "bitsandbytes"},
- },
- trust_remote_code=False,
- )
-
- assert stage_configs[0].engine_args.quantization_config == {"method": "bitsandbytes"}
diff --git a/tests/entrypoints/test_pd_disaggregation.py b/tests/entrypoints/test_pd_disaggregation.py
index c896ccb88..e9a3ba54c 100644
--- a/tests/entrypoints/test_pd_disaggregation.py
+++ b/tests/entrypoints/test_pd_disaggregation.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Unit tests for PD (Prefill-Decode) disaggregation in the Omni orchestrator.
Tests the PD detection, validation, config parsing, sampling param
@@ -74,7 +77,7 @@ class _FakeStageConfig:
class _FakeQueue:
def __init__(self, maxsize=0):
- self._queue = Queue(maxsize=maxsize)
+ self._queue: Queue[Any] = Queue(maxsize=maxsize)
def put(self, item):
self._queue.put(item)
@@ -343,7 +346,7 @@ def mock_get_config(monkeypatch):
monkeypatch.setattr(
"vllm.transformers_utils.config.get_config", lambda model, **kwargs: fake_hf_config, raising=False
)
- monkeypatch.setattr("vllm_omni.entrypoints.utils.get_config", lambda model, **kwargs: fake_hf_config, raising=False)
+ monkeypatch.setattr("vllm_omni.config.config_factory.get_config", lambda model, **kwargs: fake_hf_config)
def _mock_cached_file(path_or_repo_id, *args, **kwargs):
import os
diff --git a/tests/entrypoints/test_resolve_dreamzero_config.py b/tests/entrypoints/test_resolve_dreamzero_config.py
index 660203710..450e1be39 100644
--- a/tests/entrypoints/test_resolve_dreamzero_config.py
+++ b/tests/entrypoints/test_resolve_dreamzero_config.py
@@ -1,28 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import pytest
+from vllm_omni.config.resolver import resolve_omni_config
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
-from vllm_omni.entrypoints.utils import load_stage_configs_from_model, resolve_model_config_path
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch):
- monkeypatch.setattr(
- "vllm_omni.entrypoints.utils.get_config",
- lambda _model, trust_remote_code=True: type("Cfg", (), {"model_type": "vla"})(),
- )
- monkeypatch.setattr(
- "vllm_omni.entrypoints.utils._looks_like_dreamzero",
- lambda _model: True,
- )
- result = resolve_model_config_path("GEAR-Dreams/DreamZero-DROID")
-
- assert result is not None
- assert result.endswith("vllm_omni/deploy/dreamzero.yaml")
-
-
-def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
+def test_dreamzero_resolves_through_registry_with_model_defaults(monkeypatch):
monkeypatch.setattr(
"vllm_omni.config.config_factory.StageConfigFactory._try_infer_model_type",
classmethod(lambda _cls, model, trust_remote_code=True: "vla"),
@@ -36,12 +24,18 @@ def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
lambda _model: True,
)
- stage_configs, _ = load_stage_configs_from_model(
+ resolved = resolve_omni_config(
"GEAR-Dreams/DreamZero-DROID",
trust_remote_code=False,
+ deploy_config_path=None,
+ cli_overrides=None,
+ stage_overrides=None,
+ strategy_config_path=None,
)
- engine_args = stage_configs[0].engine_args
+ engine_args = resolved.stage_configs[0].engine_args
+ assert resolved.config_path is not None
+ assert resolved.config_path.endswith("vllm_omni/deploy/dreamzero.yaml")
assert engine_args.model_class_name == "DreamZeroPipeline"
assert engine_args.model_config.policy_server_config.action_space == "joint_position"
diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py
index 29090c425..b64758e5d 100644
--- a/tests/entrypoints/test_serve.py
+++ b/tests/entrypoints/test_serve.py
@@ -11,13 +11,22 @@ from types import SimpleNamespace
import pytest
from pytest_mock import MockerFixture
-from vllm_omni.entrypoints.cli.serve import OmniServeCommand, run_headless
+from vllm_omni.config.resolver import OmniConfigResolution
+from vllm_omni.entrypoints.cli.serve import (
+ OmniServeCommand,
+ _parse_stage_overrides,
+ run_headless,
+)
from vllm_omni.entrypoints.utils import parse_stage_overrides
from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+def _resolved(*stages: SimpleNamespace) -> OmniConfigResolution:
+ return OmniConfigResolution(config_path="/fake/stages.yaml", stage_configs=tuple(stages))
+
+
def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None:
"""``--no-async-chunk`` should parse to ``async_chunk=False`` and mark the
shared deploy-level dest as explicitly provided by the user."""
@@ -110,7 +119,7 @@ def test_serve_parser_accepts_diffusion_quantization_config() -> None:
assert args.get_explicit_kwargs_dict()["diffusion_quantization_config"] == expected
-def _make_headless_args(**kwargs) -> TrackingNamespace:
+def _make_headless_args(*, explicit_keys: frozenset[str] | None = None, **kwargs) -> TrackingNamespace:
defaults = {
"model": "fake-model",
"stage_id": 0,
@@ -130,7 +139,7 @@ def _make_headless_args(**kwargs) -> TrackingNamespace:
ns = argparse.Namespace(**ns_kwargs)
return TrackingNamespace(
unfiltered_ns=ns,
- explicit_keys=frozenset(ns.__dict__.keys()),
+ explicit_keys=frozenset(ns.__dict__.keys()) if explicit_keys is None else explicit_keys,
)
@@ -159,41 +168,49 @@ def test_run_headless_rejects_non_multiprocess_worker_backend() -> None:
# ---------------------------------------------------------------------------
-# --stage-overrides parsing parity (headless vs standard path)
+# --stage-overrides parsing at the serving boundary
# ---------------------------------------------------------------------------
def test_parse_stage_overrides_valid_json() -> None:
"""A valid JSON string is parsed into the nested per-stage dict."""
- parsed = parse_stage_overrides('{"0": {"devices": "0,1"}, "1": {"devices": "2"}}')
+ parsed = _parse_stage_overrides('{"0": {"devices": "0,1"}, "1": {"devices": "2"}}')
assert parsed == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
-def test_parse_stage_overrides_none_and_empty_return_none() -> None:
- """No overrides (None / empty string) resolve to ``None``."""
- assert parse_stage_overrides(None) is None
- assert parse_stage_overrides("") is None
+def test_serve_parser_parses_stage_overrides_before_resolution() -> None:
+ parser = TrackingArgumentParser()
+ subparsers = parser.add_subparsers(dest="subcommand")
+ OmniServeCommand().subparser_init(subparsers)
+
+ args = parser.parse_args(
+ [
+ "serve",
+ "fake-model",
+ "--omni",
+ "--stage-overrides",
+ '{"0": {"devices": "0,1"}}',
+ ]
+ )
+
+ assert args.stage_overrides == {"0": {"devices": "0,1"}}
-def test_parse_stage_overrides_empty_dict_returns_none() -> None:
- """An empty dict is falsy and must resolve to ``None``, locking in parity
- with the original standard-path ``if stage_overrides_json:`` falsy check."""
- assert parse_stage_overrides({}) is None
+def test_serve_parser_accepts_empty_stage_overrides_as_noop() -> None:
+ parser = TrackingArgumentParser()
+ subparsers = parser.add_subparsers(dest="subcommand")
+ OmniServeCommand().subparser_init(subparsers)
+ args = parser.parse_args(["serve", "fake-model", "--omni", "--stage-overrides", "{}"])
-def test_parse_stage_overrides_passes_through_non_str() -> None:
- """An already-parsed mapping is returned unchanged (identity)."""
- overrides = {"0": {"devices": "0"}}
- assert parse_stage_overrides(overrides) is overrides
+ assert args.stage_overrides == {}
def test_parse_stage_overrides_invalid_json_raises() -> None:
- """Invalid JSON raises ValueError whose message matches the standard path
- verbatim: the ``--stage-overrides is not valid JSON:`` prefix AND the
- ``Got: `` suffix echoing the raw input."""
+ """Invalid JSON fails at the serving boundary with the raw input."""
bad = "{not valid json}"
- with pytest.raises(ValueError) as excinfo:
- parse_stage_overrides(bad)
+ with pytest.raises(argparse.ArgumentTypeError) as excinfo:
+ _parse_stage_overrides(bad)
message = str(excinfo.value)
assert message.startswith("--stage-overrides is not valid JSON:")
assert f"Got: {bad!r}" in message
@@ -271,11 +288,8 @@ def test_parse_stage_overrides_accepts_empty_inner_dict() -> None:
assert parsed == {"0": {}, "1": {}}
-def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture) -> None:
- """Regression: the headless path must parse ``--stage-overrides`` (a JSON
- string) and forward the parsed dict to ``load_and_resolve_stage_configs``,
- mirroring the standard engine path. Previously it was dropped entirely,
- silently producing a different per-stage device layout."""
+def test_run_headless_forwards_parsed_stage_overrides(mocker: MockerFixture) -> None:
+ """The headless resolver receives the mapping parsed by argparse."""
captured: dict = {}
def _fake_resolve(*args, **kwargs):
@@ -283,10 +297,10 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture)
captured.update(kwargs)
# Return a stage that does NOT match stage_id=0 so run_headless stops
# right after the resolver call (we only care about how it was called).
- return ("/fake/stages.yaml", [SimpleNamespace(stage_id=99)], None)
+ return _resolved(SimpleNamespace(stage_id=99))
mocker.patch(
- "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+ "vllm_omni.config.resolver.resolve_omni_config",
side_effect=_fake_resolve,
)
@@ -294,28 +308,22 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture)
stage_id=0,
deploy_config="/tmp/deploy.yaml",
strategy_config="/tmp/strategy.yaml",
- stage_overrides='{"0": {"devices": "0,1"}, "1": {"devices": "2"}}',
+ stage_overrides={"0": {"devices": "0,1"}, "1": {"devices": "2"}},
)
with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
run_headless(args)
- assert len(captured["args"]) == 2
+ assert captured["args"] == ("fake-model",)
assert captured["deploy_config_path"] == "/tmp/deploy.yaml"
assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
assert captured["strategy_config_path"] == "/tmp/strategy.yaml"
-def test_run_headless_invalid_stage_overrides_raises(mocker: MockerFixture) -> None:
- """Invalid ``--stage-overrides`` JSON in headless mode fails fast with the
- shared ValueError instead of being silently ignored."""
- mocker.patch(
- "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
- return_value=("/fake/stages.yaml", [SimpleNamespace(stage_id=0)], None),
- )
-
- args = _make_headless_args(stage_id=0, stage_overrides="{not valid json}")
- with pytest.raises(ValueError, match="--stage-overrides is not valid JSON"):
- run_headless(args)
+def test_parse_stage_overrides_rejects_non_mapping_values() -> None:
+ with pytest.raises(argparse.ArgumentTypeError, match="JSON object"):
+ _parse_stage_overrides('["not", "a", "mapping"]')
+ with pytest.raises(argparse.ArgumentTypeError, match="must be an object"):
+ _parse_stage_overrides('{"0": "not a mapping"}')
def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture) -> None:
@@ -323,8 +331,8 @@ def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture)
fails fast when the launcher's --stage-id doesn't match any entry."""
other_stage = SimpleNamespace(stage_id=99)
mocker.patch(
- "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
- return_value=("/fake/stages.yaml", [other_stage], None),
+ "vllm_omni.config.resolver.resolve_omni_config",
+ return_value=_resolved(other_stage),
)
args = _make_headless_args(stage_id=0)
@@ -372,8 +380,8 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker
engine_manager = mocker.Mock()
mocker.patch(
- "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
- return_value=("/fake/stages.yaml", [stage_cfg], None),
+ "vllm_omni.config.resolver.resolve_omni_config",
+ return_value=_resolved(stage_cfg),
)
mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -452,8 +460,8 @@ def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: Mo
manager_b = mocker.Mock()
mocker.patch(
- "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
- return_value=("/fake/stages.yaml", [stage_cfg], None),
+ "vllm_omni.config.resolver.resolve_omni_config",
+ return_value=_resolved(stage_cfg),
)
mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -511,8 +519,8 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture)
proc.is_alive.return_value = False
mocker.patch(
- "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
- return_value=("/fake/stages.yaml", [stage_cfg], None),
+ "vllm_omni.config.resolver.resolve_omni_config",
+ return_value=_resolved(stage_cfg),
)
mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -591,8 +599,8 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur
proc.is_alive.return_value = False
mocker.patch(
- "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
- return_value=("/fake/stages.yaml", [stage_cfg], None),
+ "vllm_omni.config.resolver.resolve_omni_config",
+ return_value=_resolved(stage_cfg),
)
mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py
index 6bc6cdc2a..fa6bf8dbe 100644
--- a/tests/entrypoints/test_utils.py
+++ b/tests/entrypoints/test_utils.py
@@ -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/.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:
diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py
index f53a1ef74..073cf4f59 100644
--- a/vllm_omni/config/__init__.py
+++ b/vllm_omni/config/__init__.py
@@ -45,8 +45,10 @@ from vllm_omni.config.yaml_util import (
# `from vllm_omni.config.lora import LoRAConfig` (used while data.py is still
# loading) cannot close a circular import through DiffusionOutput.
_LAZY_ATTRS = {
+ "OmniConfigResolution": ("vllm_omni.config.resolver", "OmniConfigResolution"),
"StageConfigFactory": ("vllm_omni.config.config_factory", "StageConfigFactory"),
"register_pipeline": ("vllm_omni.config.pipeline_registry", "register_pipeline"),
+ "resolve_omni_config": ("vllm_omni.config.resolver", "resolve_omni_config"),
}
__all__ = [
@@ -60,6 +62,8 @@ __all__ = [
"VllmOmniGenerationStageConfig",
"VllmOmniDiffusionStageConfig",
"StageConfigType",
+ "OmniConfigResolution",
+ "resolve_omni_config",
# Structured Omni sub-configs.
"OmniStageCacheConfig",
"OmniStageConnectorConfig",
diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py
index b2eb64bca..90108e7df 100644
--- a/vllm_omni/config/config_factory.py
+++ b/vllm_omni/config/config_factory.py
@@ -4,10 +4,10 @@
from __future__ import annotations
-import dataclasses
import functools
+import json
from collections.abc import Mapping
-from dataclasses import asdict
+from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any
@@ -31,7 +31,7 @@ from vllm_omni.config.stage_config import (
merge_pipeline_deploy,
normalize_pipeline_cli_overrides,
)
-from vllm_omni.config.yaml_util import create_config
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
from vllm_omni.diffusion.io_support import get_diffusion_output_type
from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
@@ -48,6 +48,15 @@ logger = init_logger(__name__)
_DEFAULT_PARALLEL_DEGREE = 1
+@dataclass(frozen=True)
+class _LegacyConfigResolution:
+ """Factory-owned result for the temporary legacy runtime bridge."""
+
+ stage_configs: list[StageConfig]
+ pipeline_config: PipelineConfig
+ omni_lb_policy: str | None
+
+
@functools.cache
def _materialize_object_storage_configs(model: str) -> str:
"""Materialize an object-storage model URI's config files locally.
@@ -104,6 +113,11 @@ def with_trust_remote_code_override(
class StageConfigFactory:
"""Factory that loads pipeline YAML and merges CLI overrides.
+ Production startup source selection is owned by
+ :func:`vllm_omni.config.resolver.resolve_omni_config`. Factory methods are
+ lower-level construction primitives for that resolver, config internals,
+ and focused tests; entrypoints and engines must not call them directly.
+
Handles both single-stage and multi-stage models.
Pipelines are declared in ``vllm_omni/config/pipeline_registry.py`` and
@@ -394,11 +408,13 @@ class StageConfigFactory:
deploy_config_path: str | None,
strategy_specs: Mapping[Any, Any] | None = None,
) -> tuple[list[StageConfig] | None, str | None]:
- """Build current runtime stage configs from the shared resolution.
+ """Build the migration-only runtime ABI from the shared resolution.
- The engine still consumes the legacy StageConfig/OmegaConf shape.
- RFC #4021 will replace this transitional path as runtime consumers move
- to VllmOmniConfig.
+ The engine still consumes the legacy ``StageConfig``/OmegaConf shape.
+ This method is the resolver's temporary compatibility bridge, not an
+ alternative production source-selection entrypoint and not a stable
+ public contract. RFC #4021 will remove it as runtime consumers move to
+ ``VllmOmniConfig``.
"""
user_deploy_config = cls._load_user_deploy_config(deploy_config_path)
pipeline_cfg = cls.get_pipeline_config(
@@ -429,14 +445,29 @@ class StageConfigFactory:
user_deploy_config: DeployConfig | None = None,
strategy_specs: Mapping[Any, Any] | None = None,
) -> tuple[list[StageConfig], str | None]:
- """Create current runtime StageConfigs from registry + deploy YAML.
+ """Return the existing two-value legacy runtime ABI."""
+ resolved = cls._resolve_legacy_from_registry(
+ pipeline_cfg,
+ cli_overrides,
+ deploy_config_path,
+ user_deploy_config,
+ strategy_specs,
+ )
+ return resolved.stage_configs, resolved.omni_lb_policy
+
+ @classmethod
+ def _resolve_legacy_from_registry(
+ cls,
+ pipeline_cfg: PipelineConfig,
+ cli_overrides: dict[str, Any],
+ deploy_config_path: str | None = None,
+ user_deploy_config: DeployConfig | None = None,
+ strategy_specs: Mapping[Any, Any] | None = None,
+ ) -> _LegacyConfigResolution:
+ """Create runtime stages and retain their effective topology.
Precedence: caller-typed (non-None) value > deploy YAML >
StageDeployConfig dataclass default.
-
- Returns ``(stages, omni_lb_policy)`` — the strategy-derived pipeline-wide
- load-balancer policy (``None`` when no strategy set one) travels with the
- stages instead of through a mutable out-param.
"""
cli_overrides = normalize_pipeline_cli_overrides(pipeline_cfg, cli_overrides)
deploy_cfg: DeployConfig | None
@@ -475,7 +506,11 @@ class StageConfigFactory:
cls._reconcile_strategy_with_cli(stages, applied)
omni_lb_policy = applied.omni_lb_policy if applied is not None else None
- return stages, omni_lb_policy
+ return _LegacyConfigResolution(
+ stage_configs=stages,
+ pipeline_config=pipeline_cfg,
+ omni_lb_policy=omni_lb_policy,
+ )
@staticmethod
def _apply_strategy_specs(
@@ -587,65 +622,64 @@ class StageConfigFactory:
@classmethod
def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
- """Single-stage diffusion - no YAML needed.
+ """Build the temporary runtime ABI for a generic diffusion stage.
- Creates a default diffusion stage configuration for single-stage
- diffusion models. Returns a legacy OmegaConf-compatible dict for
- backward compatibility with OmniStage.
-
- Args:
- kwargs: Engine arguments from CLI/API.
-
- Returns:
- List containing a single config dict for the diffusion stage.
+ The terminal diffusion config owns engine defaults and normalization;
+ this compatibility builder only adds Omni stage topology, request
+ defaults, and device placement.
"""
- # Calculate devices based on parallel config
- devices = "0"
- if "parallel_config" in kwargs:
- num_devices = kwargs["parallel_config"].world_size
- for i in range(1, num_devices):
- devices += f",{i}"
+ raw_sampling_params = kwargs.get("default_sampling_params")
+ if isinstance(raw_sampling_params, str):
+ try:
+ raw_sampling_params = json.loads(raw_sampling_params)
+ except json.JSONDecodeError:
+ logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.")
+ raw_sampling_params = None
+ if not isinstance(raw_sampling_params, Mapping):
+ raw_sampling_params = None
+ default_sampling_params = dict(raw_sampling_params.get("0", {})) if raw_sampling_params else {}
- engine_args: dict[str, Any] = {}
- for key, value in kwargs.items():
- if key in ("parallel_config",):
- continue
- engine_args[key] = value
-
- # Serialize parallel_config as dict for OmegaConf. Test helpers
- # sometimes pass SimpleNamespace rather than a dataclass instance.
- if "parallel_config" in kwargs:
- parallel_config = kwargs["parallel_config"]
- if dataclasses.is_dataclass(parallel_config) and not isinstance(parallel_config, type):
- engine_args["parallel_config"] = asdict(parallel_config)
- elif hasattr(parallel_config, "__dict__"):
- engine_args["parallel_config"] = dict(vars(parallel_config))
- else:
- engine_args["parallel_config"] = parallel_config
-
- engine_args.setdefault("cache_backend", "none")
+ parallel_config = DiffusionParallelConfig.from_stage_overrides(kwargs)
+ if kwargs.get("num_gpus") is not None:
+ parallel_config.resolve_data_parallel_size(int(kwargs["num_gpus"]))
+ engine_args = OmniDiffusionConfig.normalize_init_kwargs(kwargs)
+ engine_args["parallel_config"] = asdict(parallel_config)
engine_args["model_stage"] = "diffusion"
- # Convert dtype to string for OmegaConf
- if "dtype" in engine_args:
- engine_args["dtype"] = str(engine_args["dtype"])
+ extras = dict(engine_args.get("extras") or {})
+ for key, default in (
+ ("auxiliary_text_encoder", None),
+ ("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
+ ):
+ value = kwargs.get(key)
+ if value is not None:
+ extras[key] = value
+ else:
+ extras.setdefault(key, default)
+ engine_args["extras"] = extras
- engine_args.setdefault("max_num_seqs", 1)
model_class_name = engine_args.get("model_class_name")
+ final_output_type = get_diffusion_output_type(model_class_name)
+ logger.info(
+ "Resolved generic diffusion final_output_type=%r for model_class_name=%r.",
+ final_output_type,
+ model_class_name,
+ )
- config_dict: dict[str, Any] = {
- "stage_id": 0,
- "stage_type": StageType.DIFFUSION.value,
- "runtime": {
- "process": True,
- "devices": devices,
- },
- "engine_args": create_config(engine_args),
- "final_output": True,
- "final_output_type": get_diffusion_output_type(model_class_name),
- }
-
- return [config_dict]
+ return [
+ {
+ "stage_id": 0,
+ "stage_type": StageType.DIFFUSION.value,
+ "runtime": {
+ "process": True,
+ "devices": ",".join(str(i) for i in range(parallel_config.world_size)),
+ },
+ "engine_args": engine_args,
+ "default_sampling_params": default_sampling_params,
+ "final_output": True,
+ "final_output_type": final_output_type,
+ }
+ ]
@classmethod
def _merge_cli_overrides(
diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py
index 5ac47dddc..7b14f49b1 100644
--- a/vllm_omni/config/pipeline_registry.py
+++ b/vllm_omni/config/pipeline_registry.py
@@ -17,9 +17,8 @@ To add a new pipeline:
Out of tree pipeline configs or resolvers can also be registered with register_pipeline.
-NOTE: Single-stage diffusion models continue to use the
-``_create_default_diffusion_stage_cfg`` fallback in
-``async_omni_engine.py``; for now we do not add them to registry.
+NOTE: Generic single-stage diffusion is selected by ``config.resolver`` when
+no registered Omni pipeline matches.
"""
from __future__ import annotations
diff --git a/vllm_omni/config/resolver.py b/vllm_omni/config/resolver.py
new file mode 100644
index 000000000..4f5ee8a6c
--- /dev/null
+++ b/vllm_omni/config/resolver.py
@@ -0,0 +1,266 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, fields, is_dataclass
+from typing import Any
+
+from vllm.logger import init_logger
+
+from vllm_omni.config.composable_parallel.strategy_loader import load_strategy_specs
+from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override
+from vllm_omni.config.endpoint_policy import EndpointRestriction
+from vllm_omni.config.omni_config import VllmOmniConfig
+from vllm_omni.config.stage_config import PipelineConfig
+from vllm_omni.config.yaml_util import create_config
+from vllm_omni.diffusion.data import resolve_model_class_name
+from vllm_omni.diffusion.registry import DiffusionModelRegistry
+from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class OmniConfigResolution:
+ """Migration envelope returned by the production config resolver.
+
+ ``stage_configs`` intentionally carries the current OmegaConf-compatible
+ runtime ABI only until stage startup consumes ``VllmOmniConfig`` directly.
+ ``pipeline_config`` is the effective topology after runtime stage injection,
+ so both fields always describe the same set of stages.
+ It is not a stable authoring or extension API; new production callers
+ should resolve through :func:`resolve_omni_config` and must not construct or
+ merge this compatibility shape themselves.
+ """
+
+ config_path: str | None
+ stage_configs: tuple[Any, ...] # Temporary StageConfig/OmegaConf bridge.
+ pipeline_config: PipelineConfig | None = None
+ omni_lb_policy: str | None = None
+
+ @property
+ def endpoint_restrictions(self) -> tuple[EndpointRestriction, ...]:
+ if self.pipeline_config is None:
+ return ()
+ return tuple(self.pipeline_config.endpoint_restrictions)
+
+ def stage_by_id(self, stage_id: int) -> Any:
+ for stage in self.stage_configs:
+ if stage.stage_id == stage_id:
+ return stage
+ available = [stage.stage_id for stage in self.stage_configs]
+ raise KeyError(f"no stage {stage_id}; resolved stages: {available}")
+
+
+def _filter_dict_like_object(obj: dict | Any) -> dict:
+ """Convert a dict-like object while dropping OmegaConf-incompatible callables."""
+ result = {}
+ filtered_keys = []
+ for key, value in obj.items():
+ # Preserve class objects as import paths for consumers such as
+ # custom_pipeline_args.pipeline_class.
+ if isinstance(value, type):
+ module = getattr(value, "__module__", None)
+ qualname = getattr(value, "__qualname__", getattr(value, "__name__", None))
+ result[key] = f"{module}.{qualname}" if module and qualname and module != "builtins" else qualname
+ elif callable(value):
+ filtered_keys.append(str(key))
+ else:
+ result[key] = _convert_dataclasses_to_dict(value)
+ if filtered_keys:
+ logger.warning(
+ "Filtered out %d callable object(s) from base_engine_args that are not compatible with OmegaConf: %s.",
+ len(filtered_keys),
+ filtered_keys,
+ )
+ return result
+
+
+def _convert_dataclasses_to_dict(obj: Any) -> Any:
+ """Recursively convert caller values to OmegaConf-compatible types."""
+ # Check by class name before dict to cover both collections.Counter and
+ # vllm.utils.Counter without importing either implementation.
+ if hasattr(obj, "__class__") and obj.__class__.__name__ == "Counter":
+ try:
+ return dict(obj)
+ except (TypeError, ValueError):
+ return {}
+ if isinstance(obj, set):
+ return list(obj)
+ if is_dataclass(obj) and not isinstance(obj, type):
+ result = {}
+ for config_field in fields(obj):
+ if not config_field.init:
+ continue
+ value = getattr(obj, config_field.name)
+ # At the CLI/deploy boundary, None means "unset" rather than
+ # "clear an inherited value". Preserve the dataclass default by
+ # omitting fields whose declared default is already None.
+ if value is None and config_field.default is None:
+ continue
+ result[config_field.name] = _convert_dataclasses_to_dict(value)
+ return result
+ if isinstance(obj, dict):
+ return _filter_dict_like_object(obj)
+ if isinstance(obj, type):
+ module = getattr(obj, "__module__", None)
+ qualname = getattr(obj, "__qualname__", getattr(obj, "__name__", None))
+ return f"{module}.{qualname}" if module and qualname and module != "builtins" else qualname
+ if callable(obj):
+ logger.warning(
+ "Cannot convert callable %r to an OmegaConf-compatible value.",
+ obj,
+ )
+ raise TypeError(f"callable {obj!r} is not an OmegaConf-compatible value")
+ if isinstance(obj, (list, tuple)):
+ converted = []
+ for item in obj:
+ if callable(item):
+ logger.warning(
+ "Filtered callable %r from an OmegaConf-compatible sequence.",
+ item,
+ )
+ continue
+ converted.append(_convert_dataclasses_to_dict(item))
+ return type(obj)(converted)
+ if hasattr(obj, "keys") and hasattr(obj, "values") and not isinstance(obj, (str, bytes)):
+ try:
+ return _filter_dict_like_object(obj)
+ except (TypeError, ValueError, AttributeError) as exc:
+ logger.warning(
+ "Failed to convert dict-like %s to an OmegaConf-compatible mapping.",
+ type(obj).__name__,
+ )
+ raise TypeError(f"cannot convert dict-like {type(obj).__name__}") from exc
+ return obj
+
+
+def _flatten_stage_overrides(
+ cli_overrides: dict[str, Any],
+ stage_overrides: Mapping[str, Mapping[str, Any]] | None,
+) -> None:
+ """Translate per-stage mappings to the factory's flat CLI override ABI."""
+ if not stage_overrides:
+ return
+ for stage_id, overrides in stage_overrides.items():
+ if not isinstance(overrides, Mapping):
+ raise TypeError(f"stage override {stage_id!r} must be a mapping, got {type(overrides).__name__}")
+ for key, value in overrides.items():
+ cli_overrides[f"stage_{stage_id}_{key}"] = value
+
+
+def _apply_generic_stage_overrides(
+ cli_overrides: dict[str, Any],
+ stage_overrides: Mapping[str, Mapping[str, Any]] | None,
+) -> None:
+ """Apply stage-zero overrides to the generic single-stage fallback."""
+ if not stage_overrides or not (stage_zero := stage_overrides.get("0")):
+ return
+ for key, value in stage_zero.items():
+ if key == "extras":
+ if not isinstance(value, Mapping):
+ raise TypeError(f"stage override '0'.extras must be a mapping, got {type(value).__name__}")
+ cli_overrides[key] = {**dict(cli_overrides.get(key) or {}), **value}
+ else:
+ cli_overrides[key] = value
+
+
+def _load_strategy_specs(strategy_config_path: str | None) -> Mapping[Any, Any] | None:
+ if strategy_config_path is None:
+ return None
+
+ return load_strategy_specs(strategy_config_path)
+
+
+def _build_registered_resolution(
+ structured_config: VllmOmniConfig,
+ *,
+ cli_overrides: dict[str, Any],
+ strategy_config_path: str | None,
+) -> OmniConfigResolution:
+ """Build the temporary runtime view for an already-resolved pipeline."""
+ # Runtime consumers have not yet moved to typed per-stage configs. Use the
+ # factory-owned compatibility bridge instead of reimplementing legacy YAML
+ # discovery and merging in this resolver.
+ effective_deploy_path = structured_config.orchestrator_config.deploy_config_path
+ legacy_resolution = StageConfigFactory._resolve_legacy_from_registry(
+ structured_config.pipeline_config,
+ cli_overrides,
+ effective_deploy_path,
+ strategy_specs=_load_strategy_specs(strategy_config_path),
+ )
+
+ return OmniConfigResolution(
+ config_path=effective_deploy_path,
+ stage_configs=tuple(stage.to_omegaconf() for stage in legacy_resolution.stage_configs),
+ pipeline_config=legacy_resolution.pipeline_config,
+ omni_lb_policy=legacy_resolution.omni_lb_policy,
+ )
+
+
+def _resolve_generic_diffusion_model_class(
+ model: str,
+ cli_overrides: Mapping[str, Any],
+) -> tuple[bool, str | None]:
+ """Detect generic diffusion support and its serving pipeline class."""
+ model_class_name = cli_overrides.get("model_class_name") or resolve_model_class_name(
+ model,
+ str(cli_overrides.get("diffusion_load_format") or "default"),
+ cli_overrides.get("revision"),
+ )
+ supported = bool(model_class_name and model_class_name in DiffusionModelRegistry.get_supported_archs())
+ if not supported:
+ supported = is_diffusion_model(model)
+ return supported, str(model_class_name) if model_class_name else None
+
+
+def resolve_omni_config(
+ model: str,
+ *,
+ trust_remote_code: bool | None,
+ deploy_config_path: str | None,
+ cli_overrides: Mapping[str, Any] | None,
+ stage_overrides: Mapping[str, Mapping[str, Any]] | None,
+ strategy_config_path: str | None,
+) -> OmniConfigResolution:
+ """Resolve registry/deploy inputs through the single public entrypoint."""
+ normalized_overrides = _convert_dataclasses_to_dict(dict(cli_overrides or {}))
+ normalized_overrides = with_trust_remote_code_override(normalized_overrides, trust_remote_code)
+ registry_overrides = dict(normalized_overrides)
+ _flatten_stage_overrides(registry_overrides, stage_overrides)
+
+ structured_config = StageConfigFactory.create_from_model(
+ model,
+ trust_remote_code=trust_remote_code,
+ cli_overrides=registry_overrides,
+ deploy_config_path=deploy_config_path,
+ )
+ if structured_config is not None:
+ return _build_registered_resolution(
+ structured_config,
+ cli_overrides=registry_overrides,
+ strategy_config_path=strategy_config_path,
+ )
+
+ _apply_generic_stage_overrides(normalized_overrides, stage_overrides)
+ supported, model_class_name = _resolve_generic_diffusion_model_class(model, normalized_overrides)
+ if not supported:
+ raise ValueError(
+ f"Model {model!r} did not resolve to a registered Omni pipeline or a supported diffusion model."
+ )
+ if model_class_name is not None:
+ normalized_overrides.setdefault("model_class_name", model_class_name)
+ default_stages = StageConfigFactory.create_default_diffusion(normalized_overrides)
+
+ return OmniConfigResolution(
+ config_path=deploy_config_path,
+ stage_configs=tuple(create_config(_convert_dataclasses_to_dict(default_stages))),
+ )
+
+
+__all__ = [
+ "OmniConfigResolution",
+ "resolve_omni_config",
+]
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index fd6f0237d..c30529520 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -2,6 +2,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import copy
+import json
import math
import os
import random
@@ -42,45 +43,57 @@ if TYPE_CHECKING:
logger = init_logger(__name__)
-def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]:
+def normalize_omni_diffusion_kwargs(raw_kwargs: Mapping[str, Any]) -> dict[str, Any]:
"""Normalize legacy diffusion kwargs before config construction."""
- normalized = dict(kwargs)
+ config_kwargs = dict(raw_kwargs)
+
+ dtype = config_kwargs.get("dtype")
+ if dtype is None:
+ config_kwargs["dtype"] = "auto"
+ elif isinstance(dtype, torch.dtype):
+ config_kwargs["dtype"] = str(dtype).removeprefix("torch.")
+ elif not isinstance(dtype, str):
+ raise TypeError(f"Provided dtype must be a string or torch.dtype, got {type(dtype).__name__}")
# Backwards-compatibility: older callers may use a diffusion-specific
# "static_lora_scale" kwarg. Normalize it to the canonical "lora_scale".
- if "static_lora_scale" in normalized:
- if "lora_scale" not in normalized:
- normalized["lora_scale"] = normalized["static_lora_scale"]
- normalized.pop("static_lora_scale", None)
+ if "static_lora_scale" in config_kwargs:
+ if "lora_scale" not in config_kwargs:
+ config_kwargs["lora_scale"] = config_kwargs["static_lora_scale"]
+ config_kwargs.pop("static_lora_scale", None)
+
+ diffusion_quantization = config_kwargs.pop("diffusion_quantization_config", None)
+ if config_kwargs.get("quantization_config") is None and diffusion_quantization is not None:
+ config_kwargs["quantization_config"] = diffusion_quantization
# Backwards-compatibility: map "quantization" to "quantization_config"
# so callers using the old field name still work.
- if "quantization" in normalized and normalized.get("quantization_config", None) is None:
- normalized["quantization_config"] = normalized.pop("quantization")
+ if "quantization" in config_kwargs and config_kwargs.get("quantization_config", None) is None:
+ config_kwargs["quantization_config"] = config_kwargs.pop("quantization")
else:
- normalized.pop("quantization", None)
+ config_kwargs.pop("quantization", None)
# Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype.
- if normalized.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in normalized:
- normalized["diffusion_kv_cache_dtype"] = normalized.pop("kv_cache_dtype")
+ if config_kwargs.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in config_kwargs:
+ config_kwargs["diffusion_kv_cache_dtype"] = config_kwargs.pop("kv_cache_dtype")
else:
- normalized.pop("kv_cache_dtype", None)
- if normalized.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in normalized:
- normalized["diffusion_kv_cache_skip_steps"] = normalized.pop("kv_cache_skip_steps")
+ config_kwargs.pop("kv_cache_dtype", None)
+ if config_kwargs.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in config_kwargs:
+ config_kwargs["diffusion_kv_cache_skip_steps"] = config_kwargs.pop("kv_cache_skip_steps")
else:
- normalized.pop("kv_cache_skip_steps", None)
- if normalized.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in normalized:
- normalized["diffusion_kv_cache_skip_layers"] = normalized.pop("kv_cache_skip_layers")
+ config_kwargs.pop("kv_cache_skip_steps", None)
+ if config_kwargs.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in config_kwargs:
+ config_kwargs["diffusion_kv_cache_skip_layers"] = config_kwargs.pop("kv_cache_skip_layers")
else:
- normalized.pop("kv_cache_skip_layers", None)
+ config_kwargs.pop("kv_cache_skip_layers", None)
# Handle "diffusion_attention_backend" shorthand: merge into
# diffusion_attention_config before field filtering.
- diffusion_attn_backend = normalized.pop("diffusion_attention_backend", None)
- fastvideo_vsa_topk = normalized.pop("fastvideo_vsa_topk", None)
+ diffusion_attn_backend = config_kwargs.pop("diffusion_attention_backend", None)
+ fastvideo_vsa_topk = config_kwargs.pop("fastvideo_vsa_topk", None)
if diffusion_attn_backend is not None or fastvideo_vsa_topk is not None:
- existing = normalized.get("diffusion_attention_config")
- normalized["diffusion_attention_config"] = parse_attention_config(
+ existing = config_kwargs.get("diffusion_attention_config")
+ config_kwargs["diffusion_attention_config"] = parse_attention_config(
existing,
attention_backend=diffusion_attn_backend,
fastvideo_vsa_topk=fastvideo_vsa_topk,
@@ -88,21 +101,33 @@ def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]
# Check environment variable as fallback for cache_backend.
# Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND.
- if "cache_backend" not in normalized:
+ if "cache_backend" not in config_kwargs:
cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
- normalized["cache_backend"] = cache_backend.lower() if cache_backend else "none"
- elif normalized["cache_backend"] is None:
+ config_kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none"
+ elif config_kwargs["cache_backend"] is None:
# Callers (e.g. example CLIs with `default=None`) pass an explicit
# None for "no cache"; canonicalize it so every consumer sees the
# declared `str` value instead of relying on per-model None handling.
- normalized["cache_backend"] = "none"
+ config_kwargs["cache_backend"] = "none"
+
+ cache_config = config_kwargs.get("cache_config")
+ if isinstance(cache_config, str):
+ try:
+ config_kwargs["cache_config"] = json.loads(cache_config)
+ except json.JSONDecodeError:
+ logger.warning("Invalid cache_config JSON, using backend defaults.")
+ config_kwargs.pop("cache_config", None)
+
+ if config_kwargs.get("streaming_output") is None and config_kwargs.get("diffusion_streaming_output") is not None:
+ config_kwargs["streaming_output"] = config_kwargs["diffusion_streaming_output"]
+ config_kwargs.pop("diffusion_streaming_output", None)
# Convert optional YAML null values to empty containers.
for key in ("diffusers_load_kwargs", "diffusers_call_kwargs"):
- if key in normalized and normalized[key] is None:
- normalized[key] = {}
+ if key in config_kwargs and config_kwargs[key] is None:
+ config_kwargs[key] = {}
- return normalized
+ return config_kwargs
def validate_host_weight_runtime_options(*, mode: object, root: object) -> None:
@@ -439,6 +464,27 @@ class DiffusionParallelConfig:
raise TypeError(f"Expected parallel config dict, got {type(data)!r}")
return cls(**data)
+ @classmethod
+ def from_stage_overrides(cls, kwargs: Mapping[str, Any]) -> "DiffusionParallelConfig":
+ """Resolve nested or flat stage inputs through the config's own defaults."""
+ parallel_config = kwargs.get("parallel_config")
+ if isinstance(parallel_config, cls):
+ return parallel_config
+
+ valid_fields = {item.name for item in fields(cls) if item.init}
+ if isinstance(parallel_config, Mapping):
+ return cls.from_dict(dict(parallel_config))
+ elif parallel_config is not None and hasattr(parallel_config, "__dict__"):
+ return cls.from_dict(dict(vars(parallel_config)))
+ elif parallel_config is not None:
+ raise TypeError(
+ f"parallel_config must be a mapping or DiffusionParallelConfig, got {type(parallel_config).__name__}"
+ )
+ else:
+ values = {key: value for key, value in kwargs.items() if key in valid_fields}
+
+ return cls(**{key: value for key, value in values.items() if value is not None})
+
@dataclass
class TransformerConfig:
@@ -1550,15 +1596,18 @@ class OmniDiffusionConfig:
raise
@classmethod
- def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
- kwargs = normalize_omni_diffusion_kwargs(kwargs)
-
- # Filter kwargs to only include valid fields
+ def normalize_init_kwargs(cls, raw_kwargs: Mapping[str, Any]) -> dict[str, Any]:
+ config_kwargs = normalize_omni_diffusion_kwargs(raw_kwargs)
valid_fields = {f.name for f in fields(cls)}
- filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}
+ # Remaining ``None`` values mean "unset" at the CLI/deploy boundary.
+ # Drop them so non-optional dataclass defaults are not overwritten.
+ # Fields where ``None`` has normalization semantics (for example dtype
+ # and nullable container inputs) are handled above before this filter.
+ return {key: value for key, value in config_kwargs.items() if key in valid_fields and value is not None}
- instance = cls(**filtered_kwargs)
- return instance
+ @classmethod
+ def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
+ return cls(**cls.normalize_init_kwargs(kwargs))
@dataclass
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index 812f4c3e7..f9ee40cbb 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -518,7 +518,7 @@ class OrchestratorArgs:
# === Config Files ===
deploy_config: str | None = None
- stage_overrides: str | None = None # raw JSON string; parsed downstream
+ stage_overrides: dict[str, dict[str, Any]] | None = None
# Optional composable-parallel strategy.yaml; orchestrator reads it, overlays
# derived sizing onto merged stages, then drops it before per-stage engine args.
strategy_config: str | None = None
diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py
index 4cec16382..4f1a5beb5 100644
--- a/vllm_omni/engine/async_omni_engine.py
+++ b/vllm_omni/engine/async_omni_engine.py
@@ -13,7 +13,6 @@ from __future__ import annotations
import asyncio
import concurrent.futures
import copy
-import json
import queue
import shutil
import threading
@@ -21,30 +20,23 @@ import time
import uuid
import weakref
from collections.abc import Mapping, Sequence
-from dataclasses import asdict
from typing import Any, Literal, cast
import janus
-import torch
-from omegaconf import OmegaConf
from vllm import envs as vllm_envs
from vllm.inputs import PromptType
from vllm.logger import init_logger
from vllm.v1.engine import EngineCoreRequest
from vllm.v1.engine.input_processor import InputProcessor
-from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override
+from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config
from vllm_omni.config.stage_config import (
DuplexSessionRuntimeConfig,
+ PipelineConfig,
load_deploy_config,
)
from vllm_omni.data_entry_keys import REQUEST_ARTIFACT_DIRS_KEY, TRANSFORM_OWNED_META_KEYS
-from vllm_omni.diffusion.data import (
- DiffusionParallelConfig,
- parse_attention_config,
- resolve_model_class_name,
-)
-from vllm_omni.diffusion.io_support import get_diffusion_output_type
from vllm_omni.engine import OmniEngineCoreRequest
from vllm_omni.engine.async_engine_utils import (
SHUTDOWN_ENQUEUE_TIMEOUT_S,
@@ -88,10 +80,7 @@ from vllm_omni.engine.stage_runtime import (
create_stage_runtime,
)
from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
-from vllm_omni.entrypoints.utils import (
- load_and_resolve_stage_configs,
- parse_stage_overrides,
-)
+from vllm_omni.entrypoints.utils import parse_stage_overrides
from vllm_omni.inputs.data import OmniInteractionPrompt, OmniSamplingParams
from vllm_omni.metrics.prometheus import OmniRequestCounter
@@ -99,6 +88,27 @@ logger = init_logger(__name__)
_STARTUP_POLL_INTERVAL_S = 1.0
_REQUEST_QUEUE_MAXSIZE = 256
+_ConfigResolutionResult = OmniConfigResolution | tuple[str | None, list[Any], str | None]
+
+
+def load_and_resolve_stage_configs(
+ model: str,
+ kwargs: dict[str, Any],
+ *,
+ trust_remote_code: bool | None,
+ deploy_config_path: str | None,
+ stage_overrides: Mapping[str, Mapping[str, Any]] | None,
+ strategy_config_path: str | None,
+) -> OmniConfigResolution:
+ """Compatibility seam delegating to the single config resolver."""
+ return resolve_omni_config(
+ model,
+ trust_remote_code=trust_remote_code,
+ cli_overrides=kwargs,
+ deploy_config_path=deploy_config_path,
+ stage_overrides=stage_overrides,
+ strategy_config_path=strategy_config_path,
+ )
class AsyncOmniEngine:
@@ -198,41 +208,28 @@ class AsyncOmniEngine:
self._omni_master_port,
)
- # Stage resolution pops deploy_config, so get pipeline-wide settings beforehand.
+ # Keep the historical tuple return from _resolve_stage_configs while
+ # retaining the richer resolver result for pipeline-wide settings.
+ # Overrides of that private seam fall back to the factory below.
deploy_config_path = kwargs.get("deploy_config")
- # ``trust_remote_code`` is tri-state (bool | None): ``None`` means "not
- # specified" so stage-config resolution can defer to the deploy yaml's
- # per-stage value (see ``with_trust_remote_code_override``). The
- # restriction path below loads the top-level HF config via vLLM's
- # ``get_config``, which needs a real bool, so collapse ``None`` to the
- # default ``False`` here (#5495).
- pipeline_config = StageConfigFactory.get_pipeline_config(
- model=model,
- trust_remote_code=bool(trust_remote_code),
- deploy_config_path=deploy_config_path,
- )
- self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else ()
- self._duplex_runtime_extension_path = (
- pipeline_config.duplex_runtime_extension if pipeline_config is not None else None
- )
- self.duplex_serving_adapter_path = (
- pipeline_config.duplex_serving_adapter if pipeline_config is not None else None
- )
- self._duplex_control_enabled = bool(pipeline_config and pipeline_config.duplex_control_enabled)
- self.duplex_session_config = DuplexSessionRuntimeConfig()
- if deploy_config_path is not None:
- self.duplex_session_config = load_deploy_config(deploy_config_path).duplex_session
-
- # Tri-state: None means "not specified" — the deploy yaml's per-stage
- # trust_remote_code stays in effect. An explicit True/False here is a
- # global override (precedence: caller > deploy yaml > default False);
- # the merge rule lives in with_trust_remote_code_override.
- kwargs = with_trust_remote_code_override(kwargs, trust_remote_code)
+ self._config_resolution: OmniConfigResolution | None = None
self.config_path, self.stage_configs = self._resolve_stage_configs(
model,
kwargs,
trust_remote_code=trust_remote_code,
)
+ if self._config_resolution is None:
+ pipeline_config = StageConfigFactory.get_pipeline_config(
+ model=model,
+ trust_remote_code=bool(trust_remote_code),
+ deploy_config_path=deploy_config_path,
+ )
+ self._set_pipeline_runtime_config(pipeline_config, deploy_config_path)
+ else:
+ self._set_pipeline_runtime_config(
+ self._config_resolution.pipeline_config,
+ self._config_resolution.config_path,
+ )
self.num_stages = len(self.stage_configs)
stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None
@@ -921,52 +918,6 @@ class AsyncOmniEngine:
len(companions),
)
- @staticmethod
- def _get_default_cache_config(cache_backend: str | None) -> dict[str, Any] | None:
- if cache_backend == "cache_dit":
- return {
- "Fn_compute_blocks": 1,
- "Bn_compute_blocks": 0,
- "max_warmup_steps": 4,
- "residual_diff_threshold": 0.24,
- "max_continuous_cached_steps": 3,
- "enable_taylorseer": False,
- "taylorseer_order": 1,
- "scm_steps_mask_policy": None,
- "scm_steps_policy": "dynamic",
- }
- if cache_backend == "tea_cache":
- return {
- "rel_l1_thresh": 0.2,
- }
- if cache_backend == "mag_cache":
- return {
- "mag_threshold": 0.24,
- "mag_max_skip_steps": 5,
- "mag_retention_ratio": 0.1,
- }
- if cache_backend == "step_cache":
- return {
- "step_cache_dit_enabled": True,
- "velocity_sim_thresholds": [0.95, 0.93],
- "velocity_skip_countdowns": [4, 2],
- "step_cache_dit_min_history": 2,
- "step_cache_dit_max_history": 2,
- }
- return None
-
- @staticmethod
- def _normalize_cache_config(cache_backend: str | None, cache_config: Any | None) -> Any | None:
- if isinstance(cache_config, str):
- try:
- cache_config = json.loads(cache_config)
- except json.JSONDecodeError:
- logger.warning("Invalid cache_config JSON, using defaults.")
- cache_config = None
- if cache_config is None and cache_backend not in (None, "", "none"):
- cache_config = AsyncOmniEngine._get_default_cache_config(cache_backend)
- return cache_config
-
def _detect_pd_config(self) -> dict[str, Any] | None:
"""Detect PD (Prefill-Decode) disaggregation config from stage_configs.
Returns a dict with 'pd_pair' and 'bootstrap_addr', or None.
@@ -1009,219 +960,6 @@ class AsyncOmniEngine:
"prefill_engine_id": prefill_engine_id,
}
- @staticmethod
- def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list:
- """Create a default single-stage diffusion config from kwargs."""
- # We temporally create a default config for diffusion stage.
- # In the future, we should merge the default config with the user-provided config.
- normalized_kwargs = dict(kwargs)
- default_sampling_params = normalized_kwargs.get("default_sampling_params")
- if isinstance(default_sampling_params, str):
- try:
- default_sampling_params = json.loads(default_sampling_params)
- except json.JSONDecodeError:
- logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.")
- default_sampling_params = None
- if not isinstance(default_sampling_params, dict):
- default_sampling_params = None
- stage_default_sampling_params = default_sampling_params.get("0", {}) if default_sampling_params else {}
- if normalized_kwargs.get("dtype") is None:
- normalized_kwargs["dtype"] = "auto"
-
- # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error.
- if "dtype" in normalized_kwargs and not isinstance(normalized_kwargs["dtype"], str):
- if not isinstance(normalized_kwargs["dtype"], torch.dtype):
- raise TypeError(
- f"Provided dtype must be a string or torch.dtype, got {type(normalized_kwargs['dtype']).__name__}"
- )
- normalized_kwargs["dtype"] = str(normalized_kwargs["dtype"]).removeprefix("torch.")
-
- cache_backend = normalized_kwargs.get("cache_backend", "none")
- cache_config = AsyncOmniEngine._normalize_cache_config(
- cache_backend,
- normalized_kwargs.get("cache_config", None),
- )
-
- parallel_config = normalized_kwargs.get("parallel_config")
- if isinstance(parallel_config, dict):
- parallel_config = DiffusionParallelConfig.from_dict(parallel_config)
- if parallel_config is None:
- ulysses_degree = normalized_kwargs.get("ulysses_degree") or 1
- ring_degree = normalized_kwargs.get("ring_degree") or 1
- allgather_degree = normalized_kwargs.get("allgather_degree") or 1
- ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict"
- ulysses_a2a_permute = bool(normalized_kwargs.get("ulysses_a2a_permute", False))
- sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size")
- pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
- data_parallel_size = normalized_kwargs.get("data_parallel_size")
- tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1
- cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1
- pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
- vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1
- vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile"
- text_encoder_tp_size = normalized_kwargs.get("text_encoder_tp_size") or 1
- enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False
- use_hsdp = normalized_kwargs.get("use_hsdp", False)
- hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1)
- hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size", 1)
- if sequence_parallel_size is None:
- sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree
-
- parallel_config = DiffusionParallelConfig(
- pipeline_parallel_size=pipeline_parallel_size,
- data_parallel_size=data_parallel_size,
- tensor_parallel_size=tensor_parallel_size,
- enable_expert_parallel=enable_expert_parallel,
- sequence_parallel_size=sequence_parallel_size,
- ulysses_degree=ulysses_degree,
- ring_degree=ring_degree,
- allgather_degree=allgather_degree,
- ulysses_mode=ulysses_mode,
- ulysses_a2a_permute=ulysses_a2a_permute,
- cfg_parallel_size=cfg_parallel_size,
- vae_patch_parallel_size=vae_patch_parallel_size,
- vae_parallel_mode=vae_parallel_mode,
- text_encoder_tp_size=text_encoder_tp_size,
- use_hsdp=use_hsdp,
- hsdp_shard_size=hsdp_shard_size,
- hsdp_replicate_size=hsdp_replicate_size,
- )
-
- num_gpus = normalized_kwargs.get("num_gpus")
- if num_gpus is not None:
- num_gpus = int(num_gpus)
- parallel_config.resolve_data_parallel_size(num_gpus)
-
- num_devices = max(1, int(parallel_config.world_size))
- devices = ",".join(str(i) for i in range(num_devices))
- model_class_name = kwargs.get("model_class_name", None)
- final_output_type = get_diffusion_output_type(model_class_name)
-
- attention_config = None
- if (
- kwargs.get("diffusion_attention_config") is not None
- or kwargs.get("diffusion_attention_backend") is not None
- or kwargs.get("fastvideo_vsa_topk") is not None
- ):
- attention_config = parse_attention_config(
- kwargs.get("diffusion_attention_config"),
- attention_backend=kwargs.get("diffusion_attention_backend"),
- fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"),
- )
-
- extras = dict(kwargs.get("extras") or {})
- for key, default in (
- ("auxiliary_text_encoder", None),
- ("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
- ):
- top_level_value = kwargs.get(key)
- if top_level_value is not None:
- extras[key] = top_level_value
- else:
- extras.setdefault(key, default)
-
- stage_engine_args = {
- "max_num_seqs": kwargs.get("max_num_seqs") or 1,
- "parallel_config": parallel_config,
- # Default-stage construction bypasses the structured projection.
- # Runner selection remains owned by the selected engine/platform.
- "engine_backend": kwargs.get("engine_backend", "default"),
- "model_class_name": kwargs.get("model_class_name", None),
- "task_type": kwargs.get("task_type", None),
- "model_config": kwargs.get("model_config", None),
- "additional_config": kwargs.get("additional_config", None),
- "step_execution": kwargs.get("step_execution", False),
- "request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0),
- "vae_use_slicing": kwargs.get("vae_use_slicing", False),
- "vae_use_tiling": kwargs.get("vae_use_tiling", False),
- "cache_backend": cache_backend,
- "cache_config": cache_config,
- "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
- "diffusion_offload_config": kwargs.get("diffusion_offload_config", None),
- "enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
- "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
- "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
- "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
- "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
- "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"),
- "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"),
- "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0),
- "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
- "diffusion_compile_granularity": (
- "regional"
- if kwargs.get("diffusion_compile_granularity") is None
- else kwargs["diffusion_compile_granularity"]
- ),
- "diffusion_compile_dynamic": (
- True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"]
- ),
- "fa_deterministic": bool(kwargs.get("fa_deterministic", False)),
- "boundary_ratio": kwargs.get("boundary_ratio", None),
- "flow_shift": kwargs.get("flow_shift", None),
- "diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
- "lora_path": kwargs.get("lora_path", None),
- "lora_scale": kwargs.get("lora_scale", 1.0),
- "lora_backend": kwargs.get("lora_backend", "peft"),
- "custom_pipeline_args": kwargs.get("custom_pipeline_args", None),
- "worker_extension_cls": kwargs.get("worker_extension_cls", None),
- "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]),
- "distributed_executor_backend": kwargs.get("distributed_executor_backend"),
- "enable_sleep_mode": kwargs.get("enable_sleep_mode", False),
- "enable_prompt_embed_cache": kwargs.get("enable_prompt_embed_cache", False),
- "prompt_embed_cache_size": kwargs.get("prompt_embed_cache_size", 32),
- "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
- "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
- "quantization": kwargs.get("quantization", None),
- "quantization_config": kwargs.get("quantization_config", None),
- "diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None),
- "diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None),
- "diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None),
- **({"diffusion_attention_config": attention_config} if attention_config is not None else {}),
- "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)),
- "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
- "streaming_output": kwargs.get("diffusion_streaming_output", False),
- "enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
- "extras": extras,
- **(
- {
- "profiler_config": asdict(kwargs["profiler_config"])
- if hasattr(kwargs["profiler_config"], "__dataclass_fields__")
- else kwargs["profiler_config"]
- }
- if kwargs.get("profiler_config") is not None
- else {}
- ),
- }
- if num_gpus is not None:
- stage_engine_args["num_gpus"] = num_gpus
- # Only set dtype if it was already explicitly passed and normalized
- if "dtype" in normalized_kwargs:
- stage_engine_args["dtype"] = normalized_kwargs["dtype"]
-
- # New split fields for diffusers adapter kwargs.
- if kwargs.get("diffusers_load_kwargs") is not None:
- stage_engine_args["diffusers_load_kwargs"] = kwargs["diffusers_load_kwargs"]
- if kwargs.get("diffusers_call_kwargs") is not None:
- stage_engine_args["diffusers_call_kwargs"] = kwargs["diffusers_call_kwargs"]
-
- default_stage_cfg = [
- {
- "stage_id": 0,
- "stage_type": "diffusion",
- "runtime": {
- "process": True,
- "devices": devices,
- },
- "engine_args": stage_engine_args,
- "engine_input_source": [],
- "default_sampling_params": stage_default_sampling_params,
- "final_output": True,
- "final_output_type": final_output_type,
- }
- ]
- default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion"
- return default_stage_cfg
-
def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None:
"""Apply a strategy-derived ``omni_lb_policy`` to the engine.
@@ -1251,6 +989,29 @@ class AsyncOmniEngine:
)
self._omni_lb_policy = str(derived)
+ @staticmethod
+ def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list[dict[str, Any]]:
+ """Compatibility seam for the factory-owned diffusion fallback."""
+ return StageConfigFactory.create_default_diffusion(kwargs)
+
+ def _set_pipeline_runtime_config(
+ self,
+ pipeline_config: PipelineConfig | None,
+ config_path: str | None,
+ ) -> None:
+ """Initialize engine-wide settings resolved from pipeline metadata."""
+ self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else ()
+ self._duplex_runtime_extension_path = (
+ pipeline_config.duplex_runtime_extension if pipeline_config is not None else None
+ )
+ self.duplex_serving_adapter_path = (
+ pipeline_config.duplex_serving_adapter if pipeline_config is not None else None
+ )
+ self._duplex_control_enabled = bool(pipeline_config and pipeline_config.duplex_control_enabled)
+ self.duplex_session_config = DuplexSessionRuntimeConfig()
+ if config_path is not None:
+ self.duplex_session_config = load_deploy_config(config_path).duplex_session
+
def _resolve_stage_configs(
self,
model: str,
@@ -1266,159 +1027,43 @@ class AsyncOmniEngine:
deploy_config_path = kwargs.pop("deploy_config", None)
strategy_config_path = kwargs.pop("strategy_config", None)
- stage_overrides_json = kwargs.pop("stage_overrides", None)
+ # CLI callers arrive pre-parsed; offline Python callers may use the
+ # JSON-string form documented in recipes.
+ stage_overrides = parse_stage_overrides(kwargs.pop("stage_overrides", None))
# ``diffusion_streaming_output`` is the public AsyncOmni/serve kwarg;
- # stage configs know the field as ``streaming_output``. The unregistered
- # single-stage fallback translates it in
- # ``_create_default_diffusion_stage_cfg``, but a registered pipeline
- # resolves through StageConfigFactory, which only passes through keys the
- # stage schema recognizes — so mirror the translation here. Only a truthy
- # value is mirrored: the serve CLI always carries the flag's ``False``
- # default, which must not override a deploy YAML's ``streaming_output``.
+ # stage configs know the field as ``streaming_output``. Mirror only a
+ # truthy value so the CLI's False default does not override deploy YAML.
if kwargs.get("diffusion_streaming_output") and kwargs.get("streaming_output") is None:
kwargs["streaming_output"] = True
- # Parse --stage-overrides JSON string if provided
- stage_overrides = parse_stage_overrides(stage_overrides_json)
-
- # Unregistered diffusion checkpoints use the single-stage fallback
- # below instead of StageConfigFactory, so fold stage-0 model extras
- # into the fallback's input as well. Registered pipelines still get
- # the complete override mapping through load_and_resolve_stage_configs.
- default_stage_kwargs = kwargs
- stage_zero_overrides = (stage_overrides or {}).get("0", {})
- if "extras" in stage_zero_overrides:
- override_extras = stage_zero_overrides["extras"]
- if not isinstance(override_extras, Mapping):
- raise TypeError("stage 0 extras must be a mapping")
- default_stage_kwargs = {
- **kwargs,
- "extras": {
- **(kwargs.get("extras") or {}),
- **override_extras,
- },
- }
-
- def create_default_stage_config() -> list:
- fallback_kwargs = dict(default_stage_kwargs)
- if not fallback_kwargs.get("model_class_name"):
- model_class_name = resolve_model_class_name(
- model,
- fallback_kwargs.get("diffusion_load_format", "default"),
- )
- if model_class_name is not None:
- fallback_kwargs["model_class_name"] = model_class_name
- return self._create_default_diffusion_stage_cfg(fallback_kwargs)
-
- config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs(
- model,
- kwargs,
- trust_remote_code=trust_remote_code,
- default_stage_cfg_factory=create_default_stage_config,
- deploy_config_path=deploy_config_path,
- stage_overrides=stage_overrides,
- strategy_config_path=strategy_config_path,
+ resolution = cast(
+ _ConfigResolutionResult,
+ load_and_resolve_stage_configs(
+ model,
+ kwargs,
+ trust_remote_code=trust_remote_code,
+ deploy_config_path=deploy_config_path,
+ stage_overrides=stage_overrides,
+ strategy_config_path=strategy_config_path,
+ ),
)
+ if isinstance(resolution, OmniConfigResolution):
+ self._config_resolution = resolution
+ config_path = resolution.config_path
+ stage_configs = list(resolution.stage_configs)
+ strategy_lb_policy = resolution.omni_lb_policy
+ else:
+ # Compatibility for overrides of the historical tuple-returning
+ # seam. Production always receives OmniConfigResolution above.
+ config_path, stage_configs, strategy_lb_policy = resolution
# A strategy.yaml may derive a pipeline-wide load-balancer policy. It is
# an orchestrator-level knob (read once at construction), so apply it here
# rather than as a per-stage config field.
self._apply_strategy_lb_policy(strategy_lb_policy, kwargs)
- # Inject diffusion LoRA-related knobs from kwargs if not present in the stage config.
- for cfg in stage_configs:
- try:
- if not hasattr(cfg, "engine_args") or cfg.engine_args is None:
- cfg.engine_args = OmegaConf.create({})
- global_sleep_mode = kwargs.get("enable_sleep_mode")
- if global_sleep_mode is not None:
- if not hasattr(cfg.engine_args, "enable_sleep_mode") or cfg.engine_args.enable_sleep_mode is None:
- cfg.engine_args.enable_sleep_mode = global_sleep_mode
- if getattr(cfg, "stage_type", None) != "diffusion":
- continue
- if not hasattr(cfg, "engine_args") or cfg.engine_args is None:
- cfg.engine_args = OmegaConf.create({})
- additional_config = kwargs.get("additional_config")
- if additional_config is not None:
- current_additional_config = getattr(cfg.engine_args, "additional_config", None)
- if current_additional_config in (None, {}):
- cfg.engine_args.additional_config = additional_config
- if kwargs.get("lora_path") is not None:
- if not hasattr(cfg.engine_args, "lora_path") or cfg.engine_args.lora_path is None:
- cfg.engine_args.lora_path = kwargs["lora_path"]
- lora_scale = kwargs.get("lora_scale")
- if lora_scale is None:
- # Backwards compatibility for older callers.
- lora_scale = kwargs.get("static_lora_scale")
- if lora_scale is not None:
- if not hasattr(cfg.engine_args, "lora_scale") or cfg.engine_args.lora_scale is None:
- cfg.engine_args.lora_scale = lora_scale
- if kwargs.get("lora_backend") is not None:
- if not hasattr(cfg.engine_args, "lora_backend") or cfg.engine_args.lora_backend is None:
- cfg.engine_args.lora_backend = kwargs["lora_backend"]
- if (
- kwargs.get("diffusion_attention_config") is not None
- or kwargs.get("diffusion_attention_backend") is not None
- or kwargs.get("fastvideo_vsa_topk") is not None
- ):
- has_stage_attention = (
- getattr(cfg.engine_args, "diffusion_attention_config", None) is not None
- or getattr(cfg.engine_args, "diffusion_attention_backend", None) is not None
- )
- if not has_stage_attention:
- cfg.engine_args.diffusion_attention_config = parse_attention_config(
- kwargs.get("diffusion_attention_config"),
- attention_backend=kwargs.get("diffusion_attention_backend"),
- fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"),
- )
- quantization_config = kwargs.get("diffusion_quantization_config") or kwargs.get("quantization_config")
- if quantization_config is not None:
- if (
- not hasattr(cfg.engine_args, "quantization_config")
- or cfg.engine_args.quantization_config is None
- ):
- cfg.engine_args.quantization_config = quantization_config
- # Inject profiler flags for diffusion stages
- for profiler_key in (
- "enable_diffusion_pipeline_profiler",
- "enable_ar_profiler",
- ):
- val = kwargs.get(profiler_key)
- if val:
- if not hasattr(cfg.engine_args, profiler_key) or not getattr(
- cfg.engine_args, profiler_key, False
- ):
- setattr(cfg.engine_args, profiler_key, val)
- quantization = kwargs.get("quantization")
- if quantization is not None:
- if not hasattr(cfg.engine_args, "quantization") or cfg.engine_args.quantization is None:
- cfg.engine_args.quantization = quantization
- diffusion_kv_cache_dtype = kwargs.get("diffusion_kv_cache_dtype")
- if diffusion_kv_cache_dtype is not None:
- if (
- not hasattr(cfg.engine_args, "diffusion_kv_cache_dtype")
- or cfg.engine_args.diffusion_kv_cache_dtype is None
- ):
- cfg.engine_args.diffusion_kv_cache_dtype = diffusion_kv_cache_dtype
- diffusion_kv_cache_skip_steps = kwargs.get("diffusion_kv_cache_skip_steps")
- if diffusion_kv_cache_skip_steps is not None:
- if (
- not hasattr(cfg.engine_args, "diffusion_kv_cache_skip_steps")
- or cfg.engine_args.diffusion_kv_cache_skip_steps is None
- ):
- cfg.engine_args.diffusion_kv_cache_skip_steps = diffusion_kv_cache_skip_steps
- diffusion_kv_cache_skip_layers = kwargs.get("diffusion_kv_cache_skip_layers")
- if diffusion_kv_cache_skip_layers is not None:
- if (
- not hasattr(cfg.engine_args, "diffusion_kv_cache_skip_layers")
- or cfg.engine_args.diffusion_kv_cache_skip_layers is None
- ):
- cfg.engine_args.diffusion_kv_cache_skip_layers = diffusion_kv_cache_skip_layers
- except Exception as e:
- logger.warning("Failed to inject LoRA config for stage: %s", e)
-
- return config_path, stage_configs
+ return cast(str, config_path), stage_configs
# ==================== Public API ====================
@@ -1944,7 +1589,7 @@ class AsyncOmniEngine:
tokens generated before abort (empty for diffusion / no OP state).
"""
if not request_ids or getattr(self, "_shutdown_called", False):
- return
+ return []
if self.request_queue is None:
raise RuntimeError("request_queue is not initialized")
transport = self._correlated_rpc_client
@@ -1971,7 +1616,7 @@ class AsyncOmniEngine:
result_msg = await loop.run_in_executor(None, _wait)
except Exception as exc:
if getattr(self, "_shutdown_called", False) and is_abort_transport_shutdown(exc):
- return
+ return []
raise
if not result_msg.success:
raise RuntimeError(result_msg.error or "abort failed")
diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py
index f2bc42ab1..eed686837 100644
--- a/vllm_omni/engine/stage_init_utils.py
+++ b/vllm_omni/engine/stage_init_utils.py
@@ -53,7 +53,7 @@ from vllm_omni.config.stage_config import StageType
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.engine.arg_utils import OmniEngineArgs
from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices
-from vllm_omni.entrypoints.utils import filter_dataclass_kwargs, resolve_model_config_path
+from vllm_omni.entrypoints.utils import filter_dataclass_kwargs
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams
from vllm_omni.inputs.preprocess import OmniInputPreprocessor
from vllm_omni.outputs.output_processor import MultimodalOutputProcessor
@@ -1811,7 +1811,7 @@ def release_device_locks(lock_fds: list[int]) -> None:
def load_omni_transfer_config_for_model(model: str, config_path: str | None) -> Any:
- """Load omni transfer config from an explicit path or resolved model config.
+ """Load omni transfer config from the resolver-selected deploy config.
Resolves ``base_config`` inheritance (CI overlay → base deploy YAML) so
that connectors defined in the base config are visible to the transfer
@@ -1820,12 +1820,11 @@ def load_omni_transfer_config_for_model(model: str, config_path: str | None) ->
from vllm_omni.distributed.omni_connectors import load_omni_transfer_config
try:
- resolved_config_path = config_path or resolve_model_config_path(model)
- if resolved_config_path is None:
+ if config_path is None:
return None
from vllm_omni.config.stage_config import resolve_deploy_yaml
- resolved_dict = resolve_deploy_yaml(resolved_config_path)
+ resolved_dict = resolve_deploy_yaml(config_path)
return load_omni_transfer_config(config_dict=resolved_dict)
except Exception as e:
logger.warning("[stage_init] Failed to load transfer config: %s", e)
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index ea56acb1a..25f25d23c 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -14,6 +14,7 @@ import math
import os
import signal
from types import FrameType
+from typing import Any, cast
import uvloop
from vllm.entrypoints.cli.types import CLISubcommand
@@ -23,6 +24,7 @@ from vllm.logger import init_logger
from vllm_omni.entrypoints.cli.logo import log_logo
from vllm_omni.entrypoints.openai.api_server import omni_run_server
+from vllm_omni.entrypoints.utils import parse_stage_overrides
from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
logger = init_logger(__name__)
@@ -47,6 +49,17 @@ Search by using: `--help=` to explore options by section (e.g.,
"""
+def _parse_stage_overrides(value: str) -> dict[str, dict[str, Any]]:
+ """Adapt shared stage-override validation to argparse's error type."""
+ try:
+ parsed = parse_stage_overrides(value)
+ except ValueError as exc:
+ raise argparse.ArgumentTypeError(str(exc)) from exc
+ if parsed is None:
+ raise argparse.ArgumentTypeError("--stage-overrides requires a JSON object")
+ return parsed
+
+
def _nonneg_finite_float(value: str) -> float:
"""Argparse type for finite, non-negative floats (rejects nan/inf)."""
try:
@@ -286,7 +299,7 @@ class OmniServeCommand(CLISubcommand):
)
omni_config_group.add_argument(
"--stage-overrides",
- type=str,
+ type=_parse_stage_overrides,
default=None,
help="Per-stage JSON overrides. Example: "
'\'{"0": {"gpu_memory_utilization": 0.8}, "2": {"enforce_eager": true}}\'',
@@ -945,6 +958,7 @@ def run_headless(args: TrackingNamespace) -> None:
from vllm.v1.executor.multiproc_executor import MultiprocExecutor
from vllm.version import __version__ as VLLM_VERSION
+ from vllm_omni.config.resolver import resolve_omni_config
from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_omni_kv_config_for_stage
from vllm_omni.engine.stage_engine_startup import (
get_headless_replica_devices,
@@ -959,10 +973,6 @@ def run_headless(args: TrackingNamespace) -> None:
load_omni_transfer_config_for_model,
prepare_engine_environment,
)
- from vllm_omni.entrypoints.utils import (
- load_and_resolve_stage_configs,
- parse_stage_overrides,
- )
model = args.model
stage_id: int | None = args.stage_id
@@ -984,6 +994,10 @@ def run_headless(args: TrackingNamespace) -> None:
# Filter down to a dict of things explicitly requested by the user
args_dict = args.get_explicit_kwargs_dict()
+ deploy_config_path = args_dict.pop("deploy_config", None)
+ strategy_config_path = args_dict.pop("strategy_config", None)
+ stage_overrides = args_dict.pop("stage_overrides", None)
+
# ``--replica-id`` is deprecated and ignored — replica ids are
# auto-assigned by ``OmniMasterServer`` so headless processes carry
# no knowledge of their per-replica id at launch time. Warn (don't
@@ -996,33 +1010,27 @@ def run_headless(args: TrackingNamespace) -> None:
"master server.",
args.replica_id,
)
+ args_dict.pop("replica_id")
- # Parse --stage-overrides (raw JSON string) exactly like the standard
- # engine path (AsyncOmniEngine._resolve_stage_configs) so headless and
- # standard launches resolve to the same per-stage device layout.
- stage_overrides = parse_stage_overrides(args_dict.get("stage_overrides"))
-
- config_path, stage_configs, _ = load_and_resolve_stage_configs(
+ resolved = resolve_omni_config(
model,
- args_dict,
# store_true cannot express an explicit False: absent maps to None
# ("not specified") so the deploy yaml's per-stage value applies.
trust_remote_code=getattr(args, "trust_remote_code", None) or None,
- deploy_config_path=args_dict.get("deploy_config"),
+ cli_overrides=args_dict,
+ deploy_config_path=deploy_config_path,
stage_overrides=stage_overrides,
- strategy_config_path=args_dict.get("strategy_config"),
+ strategy_config_path=strategy_config_path,
)
+ config_path = resolved.config_path
+ stage_configs = list(resolved.stage_configs)
- # Locate the stage config that matches stage_id.
- stage_cfg = None
- for cfg in stage_configs:
- if cfg.stage_id == stage_id:
- stage_cfg = cfg
- break
- if stage_cfg is None:
+ try:
+ stage_cfg = resolved.stage_by_id(stage_id)
+ except KeyError:
raise ValueError(
f"No stage config found for stage_id={stage_id}. Available stage ids: {[c.stage_id for c in stage_configs]}"
- )
+ ) from None
prepare_engine_environment()
per_replica_devices = get_headless_replica_devices(stage_cfg, stage_id, omni_dp_size_local)
@@ -1037,7 +1045,7 @@ def run_headless(args: TrackingNamespace) -> None:
omni_master_port=omni_master_port,
omni_dp_size_local=omni_dp_size_local,
per_replica_devices=per_replica_devices,
- config_path=config_path,
+ config_path=cast(str, config_path),
replica_bind_address=omni_replica_address,
)
return
diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py
index 4b91f752a..e1efd7ac4 100644
--- a/vllm_omni/entrypoints/utils.py
+++ b/vllm_omni/entrypoints/utils.py
@@ -4,45 +4,18 @@
import json
import os
import types
-from collections import Counter
from dataclasses import fields, is_dataclass
-from pathlib import Path
from typing import Any, get_args, get_origin
from vllm.logger import init_logger
from vllm.sampling_params import RequestOutputKind, SamplingParams
-from vllm.transformers_utils.config import get_config, get_hf_file_to_dict
-from vllm.transformers_utils.repo_utils import file_or_path_exists
-from vllm_omni.config.config_factory import (
- StageConfigFactory,
- _materialize_object_storage_configs,
- _name_match_candidate,
- with_trust_remote_code_override,
-)
-from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
-from vllm_omni.config.stage_config import _DEPLOY_DIR
-from vllm_omni.config.yaml_util import create_config, load_yaml_config
-from vllm_omni.diffusion.utils.hf_utils import (
- _looks_like_dreamzero,
- get_diffusion_model_index,
- resolve_native_diffusion_model_class,
-)
from vllm_omni.entrypoints.stage_utils import _to_dict
from vllm_omni.inputs.data import OmniSamplingParams
-from vllm_omni.platforms import current_omni_platform
-
-# Get the project root directory (2 levels up from this file)
-PROJECT_ROOT = Path(__file__).parent.parent.parent
logger = init_logger(__name__)
-_DIFFUSERS_CLASS_TO_CONFIG: dict[str, str] = {
- "GlmImagePipeline": "glm_image",
-}
-
-
def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from: str, omni_to: str) -> None:
"""Inject connector configuration into stage engine arguments."""
# Prepare omni_kv_config dict
@@ -76,501 +49,9 @@ def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from:
logger.error(f"Failed to inject omni connector config into stage: {e}")
-def _try_get_class_name_from_diffusers_config(model: str) -> str | None:
- """Try to get class name from diffusers model configuration files.
-
- Args:
- model: Model name or path
-
- Returns:
- Model type string if found, None otherwise
- """
- model_index = get_diffusion_model_index(model)
- if model_index and isinstance(model_index, dict) and "_class_name" in model_index:
- logger.debug(f"Found Diffusers model type '{model_index['_class_name']}'")
- return model_index["_class_name"]
-
- return None
-
-
-def _filter_dict_like_object(obj: dict | Any) -> dict:
- """Filter dict-like object by removing callables and recursively converting values.
-
- Converts dict-like objects to regular dicts while filtering out callable values
- that are incompatible with OmegaConf. Recursively processes values through
- _convert_dataclasses_to_dict for nested object conversion.
-
- Args:
- obj: Dict or dict-like object to filter
-
- Returns:
- Regular dict with callables filtered out and values recursively converted
-
- Raises:
- TypeError: If obj doesn't support .items() method
- ValueError: If dict conversion fails unexpectedly
- """
-
- def _is_callable_value(value: Any) -> bool:
- if callable(value):
- return True
- return isinstance(
- value,
- types.FunctionType | types.MethodType | types.BuiltinFunctionType | types.BuiltinMethodType,
- )
-
- result = {}
- filtered_keys = []
- for k, v in obj.items():
- # Preserve class objects by converting to a fully qualified name string
- # so callers that resolve via import path (e.g. custom_pipeline_args.pipeline_class)
- # still work after OmegaConf round-trip.
- if isinstance(v, type):
- module = getattr(v, "__module__", None)
- qualname = getattr(v, "__qualname__", getattr(v, "__name__", None))
- if module and qualname and module != "builtins":
- result[k] = f"{module}.{qualname}"
- else:
- result[k] = qualname
- elif _is_callable_value(v):
- filtered_keys.append(str(k))
- else:
- result[k] = _convert_dataclasses_to_dict(v)
- if filtered_keys:
- logger.warning(
- f"Filtered out {len(filtered_keys)} callable object(s) from base_engine_args "
- f"that are not compatible with OmegaConf: {filtered_keys}. "
- )
- return result
-
-
-def _convert_dataclasses_to_dict(obj: Any) -> Any:
- """Recursively convert non-serializable objects to OmegaConf-compatible types.
-
- This is needed because OmegaConf cannot handle:
- - Dataclass objects with Literal type annotations (e.g., StructuredOutputsConfig)
- - Counter objects (from collections or vllm.utils)
- - Set objects
- - Callable objects (functions, methods, etc.)
- - Other non-primitive types
- """
- # IMPORTANT: Check Counter BEFORE dict, since Counter is a subclass of dict
- # Handle Counter objects (convert to dict)
- # Check by class name first to catch both collections.Counter and vllm.utils.Counter
- if hasattr(obj, "__class__") and obj.__class__.__name__ == "Counter":
- try:
- return dict(obj)
- except (TypeError, ValueError):
- # If Counter can't be converted to dict, return empty dict
- return {}
- # Also check isinstance for collections.Counter (must be before dict check)
- if isinstance(obj, Counter):
- return dict(obj)
- # Handle set objects (convert to list)
- if isinstance(obj, set):
- return list(obj)
- # Handle dataclass objects
- # Use field iteration instead of asdict() to:
- # 1. Only include init fields (non-init fields cause "unexpected kwarg" errors)
- # 2. Skip None values matching field defaults (avoids Pydantic validation
- # failures when None is explicitly passed for non-Optional typed fields,
- # e.g. CompilationConfig.cudagraph_capture_sizes: list[int] = None)
- if is_dataclass(obj) and not isinstance(obj, type):
- result = {}
- for f in fields(obj):
- if not f.init:
- continue
- value = getattr(obj, f.name)
- if value is None and f.default is None:
- continue
- result[f.name] = _convert_dataclasses_to_dict(value)
- return result
- # Handle dictionaries (recurse into values) and filter out callables(cause error in OmegaConf.create)
- # Note: This must come AFTER Counter check since Counter is a dict subclass
- if isinstance(obj, dict):
- return _filter_dict_like_object(obj)
- # Preserve class objects by converting to a fully qualified name string.
- if isinstance(obj, type):
- module = getattr(obj, "__module__", None)
- qualname = getattr(obj, "__qualname__", getattr(obj, "__name__", None))
- if module and qualname and module != "builtins":
- return f"{module}.{qualname}"
- return qualname
- # Handle callable objects (functions, methods, etc.) - skip them
- # Note: This comes after dict/list checks to avoid misclassifying dict-like objects
- if callable(obj):
- return None
- # Handle lists and tuples (recurse into items)
- if isinstance(obj, list | tuple):
- return type(obj)(_convert_dataclasses_to_dict(item) for item in obj if not callable(item))
- # Try to convert any dict-like object (has keys/values methods) to dict
- if hasattr(obj, "keys") and hasattr(obj, "values") and not isinstance(obj, str | bytes):
- try:
- return _filter_dict_like_object(obj)
- except (TypeError, ValueError, AttributeError):
- # If conversion fails, return as-is
- return obj
- # Primitive types and other objects that OmegaConf can handle
- return obj
-
-
-def _try_resolve_omni_model_type(model: str) -> str | None:
- """Try to resolve model_type for omni models with empty config.json.
-
- Searches both the legacy ``stage_configs/*.yaml`` directory and the
- migrated ``deploy/*.yaml`` directory for a stem that substring-matches
- the model path basename (e.g. ``cosyvoice3`` in
- ``FunAudioLLM/Fun-CosyVoice3-0.5B-2512``). Only the basename is scanned
- so URI segments such as the bucket name cannot select an unrelated
- pipeline. The longest match wins so ``cosyvoice3`` beats ``cosyvoice``
- and ``bagel_single_stage`` beats ``bagel``.
- """
- model_lower = _name_match_candidate(model).lower().replace("-", "").replace("_", "")
- best_match: str | None = None
- best_len = 0
- for subdir in ("model_executor/stage_configs", "deploy"):
- config_dir = PROJECT_ROOT / "vllm_omni" / subdir
- if not config_dir.exists():
- continue
- for config_file in sorted(config_dir.glob("*.yaml")):
- candidate = config_file.stem.replace("-", "").replace("_", "")
- if candidate and candidate in model_lower and len(candidate) > best_len:
- best_match = config_file.stem
- best_len = len(candidate)
- return best_match
-
-
-def _registry_default_deploy_path(model: str) -> str | None:
- """Return the registered pipeline's default deploy YAML path, if any.
-
- A checkpoint's HF ``model_type`` is not always the registered Omni pipeline
- key. MiniCPM-o 4.5, for example, reports ``minicpmo`` while its
- predicate-selected pipeline is ``minicpmo_4_5``. Stage construction already
- loads that pipeline's default deploy config, so connector discovery must
- resolve the same path.
- """
- pipeline_config = StageConfigFactory.get_pipeline_config(
- model=model,
- trust_remote_code=True,
- )
- if pipeline_config is None or pipeline_config.default_deploy_config_name is None:
- return None
- default_deploy_path = _DEPLOY_DIR / pipeline_config.default_deploy_config_name
- if default_deploy_path.is_file():
- return str(default_deploy_path)
- return None
-
-
-def resolve_model_config_path(model: str) -> str | None:
- """Resolve the stage/deploy config file path from the model name.
-
- Resolves configuration path based on the model type and device type.
- Order:
- 1. Device-specific stage config under ``stage_configs/{device_type}/``
- 2. If HF ``model_type`` is not an ``OMNI_PIPELINES`` key, the registered
- pipeline's ``default_deploy_config_name`` (keeps connectors aligned with
- stage construction when HF type and pipeline key differ)
- 3. ``deploy/{model_type}.yaml``
- 4. Legacy ``stage_configs/{model_type}.yaml``
- 5. Registered pipeline default deploy config (final fallthrough)
-
- Args:
- model: Model name or path (used to determine model_type)
-
- Returns:
- String path to the stage/deploy configuration file, or ``None`` if no
- matching config file exists.
-
- Raises:
- ValueError: If model_type cannot be determined
- """
- # Object-storage URIs are streamed by vLLM's Run:AI streamer only after
- # each stage builds its ModelConfig, so config resolution here reads the
- # local copy that vLLM-Omni materializes for such URIs. Name-based
- # fallbacks keep the original string (the materialized path is a hash).
- config_source = _materialize_object_storage_configs(model)
- # Try to get config from standard transformers format first
- try:
- hf_config = get_config(config_source, trust_remote_code=True)
- model_type = hf_config.model_type
- except (ValueError, Exception):
- # If standard transformers format fails, try diffusers format
- native_model_class = resolve_native_diffusion_model_class(config_source)
- if native_model_class is not None:
- model_type = native_model_class
- elif get_diffusion_model_index(config_source) is not None:
- model_type = _try_get_class_name_from_diffusers_config(config_source)
- if model_type is None:
- raise ValueError(
- f"Could not determine model_type for diffusers model: {model}. "
- "Please ensure its Diffusers pipeline index contains '_class_name'"
- )
- elif file_or_path_exists(config_source, "config.json", revision=None):
- # Try to read config.json manually for custom models like Bagel that fail get_config
- # but have a valid config.json with model_type
- try:
- config_dict = get_hf_file_to_dict("config.json", config_source, revision=None)
- if config_dict and "model_type" in config_dict:
- model_type = config_dict["model_type"]
- else:
- # For models with empty config.json (e.g. CosyVoice3),
- # try matching against registered omni stage configs.
- model_type = _try_resolve_omni_model_type(model)
- if model_type is None:
- raise ValueError(f"config.json found but missing 'model_type' for model: {model}")
- except Exception as e:
- raise ValueError(f"Failed to read config.json for model: {model}. Error: {e}") from e
- else:
- # No config.json at repo root (e.g. GLM-TTS stores configs in
- # subdirectories only). Try matching against registered deploy
- # YAML filenames before giving up.
- model_type = _try_resolve_omni_model_type(model)
- if model_type is None:
- raise ValueError(
- f"Could not determine model_type for model: {model}. "
- "Model is not in standard transformers or Diffusers format. "
- f"Please ensure the model has proper configuration files with 'model_type' field"
- )
-
- default_config_path = current_omni_platform.get_default_stage_config_path()
- if model_type == "vla" and _looks_like_dreamzero(config_source):
- model_type = "dreamzero"
-
- if model_type in _DIFFUSERS_CLASS_TO_CONFIG:
- normalized_model_type = _DIFFUSERS_CLASS_TO_CONFIG[model_type]
- else:
- normalized_model_type = model_type.replace("-", "_")
- model_type_str = f"{normalized_model_type}.yaml"
- complete_config_path = PROJECT_ROOT / default_config_path / model_type_str
- if os.path.exists(complete_config_path):
- return str(complete_config_path)
-
- # Prefer the registry default before deploy/.yaml when the
- # HF type is not itself a pipeline key. Otherwise a future
- # deploy/minicpmo.yaml would desync connectors from stages for MiniCPM-o 4.5.
- if normalized_model_type not in OMNI_PIPELINES:
- registry_path = _registry_default_deploy_path(model)
- if registry_path is not None:
- return registry_path
-
- deploy_config_path = _DEPLOY_DIR / model_type_str
- if os.path.exists(deploy_config_path):
- return str(deploy_config_path)
-
- stage_config_file = f"vllm_omni/model_executor/stage_configs/{normalized_model_type}.yaml"
- stage_config_path = PROJECT_ROOT / stage_config_file
- if os.path.exists(stage_config_path):
- return str(stage_config_path)
-
- return _registry_default_deploy_path(model)
-
-
-def load_stage_configs_from_model(
- model: str,
- *,
- trust_remote_code: bool | None,
- base_engine_args: dict | None = None,
- deploy_config_path: str | None = None,
- stage_overrides: dict[str, dict[str, Any]] | None = None,
- strategy_config_path: str | None = None,
-) -> tuple[list, str | None]:
- """Load stage configurations from model's default config file.
-
- For models registered in the pipeline registry, uses
- ``StageConfigFactory.create_legacy_stage_configs_from_model()`` which merges
- PipelineConfig + DeployConfig + CLI overrides.
-
- Models that cannot be resolved through the registry return no stages so the
- caller can apply its default stage configuration, when available.
-
- Args:
- model: Model name or path (used to determine model_type)
- trust_remote_code: Whether to trust remote code while resolving the model config.
- base_engine_args: Base engine args to merge as CLI overrides.
- deploy_config_path: Optional explicit deploy config path.
- stage_overrides: Per-stage overrides from --stage-overrides.
- strategy_config_path: Optional path to a composable-parallel
- ``strategy.yaml`` whose derived sizing is overlaid onto the
- registry-merged stages.
-
- Returns:
- ``(stage_configs, omni_lb_policy)``: the list of stage configuration
- dictionaries plus the strategy-derived pipeline-wide ``omni_lb_policy``
- (``None`` when no strategy set one). The policy is returned rather than
- written into a caller-provided mutable dict.
- """
- if base_engine_args is None:
- base_engine_args = {}
-
- cli_overrides = _convert_dataclasses_to_dict(dict(base_engine_args))
- # A False inherited from the engine-args dump is the store_true flag's
- # default, not an explicit choice — drop it so only the tri-state
- # parameter below decides (see with_trust_remote_code_override).
- if not cli_overrides.get("trust_remote_code"):
- cli_overrides.pop("trust_remote_code", None)
- cli_overrides = with_trust_remote_code_override(cli_overrides, trust_remote_code)
- if stage_overrides:
- for stage_id_str, overrides in stage_overrides.items():
- for key, val in overrides.items():
- cli_overrides[f"stage_{stage_id_str}_{key}"] = val
-
- # Current runtime initialization still consumes legacy OmegaConf stage
- # configs. ``StageConfigFactory.create_from_model`` now produces the
- # structured ``VllmOmniConfig`` object; future RFC #4021 changes will
- # migrate the engine/runtime consumers and replace this legacy resolver.
- strategy_specs = None
- if strategy_config_path is not None:
- from vllm_omni.config.composable_parallel.strategy_loader import load_strategy_specs
-
- strategy_specs = load_strategy_specs(strategy_config_path)
-
- stages, omni_lb_policy = StageConfigFactory.create_legacy_stage_configs_from_model(
- model,
- trust_remote_code=trust_remote_code,
- cli_overrides=cli_overrides,
- deploy_config_path=deploy_config_path,
- strategy_specs=strategy_specs,
- )
- if stages is not None:
- # Convert StageConfig objects to OmegaConf for backward compat
- return [stage.to_omegaconf() for stage in stages], omni_lb_policy
-
- strategy_note = ""
- if strategy_config_path is not None:
- strategy_note = f" Strategy config {strategy_config_path!r} was not applied."
- logger.warning(
- "No registered PipelineConfig resolved for model %r. Legacy `stage_args` "
- "YAMLs are no longer supported; register the pipeline and provide deployment "
- "overrides through `deploy_config`.%s",
- model,
- strategy_note,
- )
- return [], None
-
-
-def filter_stages(
- config_path: str | None,
- stage_configs: list,
- kwargs: dict | None,
-) -> list:
- """Filter stage configs by mode when YAML defines a `modes` section.
-
- The YAML can define, e.g.:
-
- modes:
- - mode: text-to-image
- stages: [1]
- - mode: image-to-text
- stages: [0]
-
- When users pass `mode="image-to-text"` into Omni(**kwargs), only the stages
- listed for that mode are returned. If no mode is provided, defaults to
- "text-to-image". If no modes are defined or filtering fails, returns the
- original stage_configs unchanged.
-
- Args:
- config_path: Path to the YAML config (used to read `modes`).
- stage_configs: Loaded list of stage configs.
- kwargs: Engine/caller kwargs; may contain "mode".
-
- Returns:
- Filtered list of stage configs (or original list if filtering not applied).
- """
- if not stage_configs or config_path is None:
- return stage_configs
-
- try:
- cfg = load_yaml_config(config_path)
- yaml_modes = getattr(cfg, "modes", None)
- if yaml_modes is None:
- return stage_configs
-
- mode_to_stage_ids: dict[str, list[int]] = {}
- if yaml_modes is not None:
- for entry in yaml_modes:
- mode_name = None
- stages = None
- if hasattr(entry, "mode") or hasattr(entry, "stages"):
- mode_name = getattr(entry, "mode", None)
- stages = getattr(entry, "stages", None)
- elif isinstance(entry, dict):
- mode_name = entry.get("mode")
- stages = entry.get("stages")
-
- if mode_name is None or stages is None:
- continue
-
- if isinstance(stages, int):
- stage_list = [stages]
- else:
- stage_list = list(stages)
-
- mode_to_stage_ids[str(mode_name)] = [int(sid) for sid in stage_list]
-
- # No modes section or empty mapping: use all stages and return early.
- active_mode: str | None = None
- if isinstance(kwargs, dict):
- active_mode = kwargs.get("mode")
-
- if active_mode is None:
- active_mode = "text-to-image"
-
- if active_mode not in mode_to_stage_ids:
- logger.warning(
- "Requested mode '%s' not found in config '%s'; available modes: %s. Using all stages.",
- active_mode,
- config_path,
- sorted(mode_to_stage_ids.keys()),
- )
- return stage_configs
-
- allowed_ids = set(mode_to_stage_ids[active_mode])
- filtered_stage_configs = [sc for sc in stage_configs if getattr(sc, "stage_id", None) in allowed_ids]
- if not filtered_stage_configs:
- logger.warning(
- "Mode '%s' in config '%s' resolved to stage ids %s, but none matched loaded stage_args. "
- "Falling back to all stages.",
- active_mode,
- config_path,
- sorted(allowed_ids),
- )
- return stage_configs
-
- return filtered_stage_configs
- except Exception as e:
- logger.warning("Failed to apply mode-based stage filtering: %s", e)
- return stage_configs
-
-
def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
- """Parse the ``--stage-overrides`` value into a per-stage override dict.
-
- ``value`` may be a raw JSON string (as supplied on the CLI) or an
- already-parsed mapping. Returns ``None`` when no overrides are given.
-
- Only the **shape** of the override mapping is validated here:
-
- - top-level must be a dict
- - keys must be non-negative ASCII integer strings (stage ids)
- - values must be dicts
-
- Field-name / type / range validation is intentionally not done at this
- layer. The downstream resolver forwards every surviving key as
- ``stage__`` into ``cli_overrides`` (see
- ``load_stage_configs_from_model``); field semantics live in
- ``StageConfigFactory`` for registered pipelines, and the
- default-diffusion fallback in ``async_omni_engine.py`` reads
- ``stage_zero_overrides["extras"]`` for unregistered ones. Unknown
- engine-arg keys that don't match ``OmniEngineArgs`` are dropped with
- a warning at ``filter_dataclass_kwargs`` rather than failing here,
- so a typo surfaces in the log instead of being silently lost.
-
- Raises:
- ValueError: when ``value`` is not a valid per-stage override
- mapping (invalid JSON, non-dict top level, non-ASCII-digit
- stage id, non-dict value).
- """
- if not value:
+ """Parse and validate the shape of per-stage JSON overrides."""
+ if value is None:
return None
if isinstance(value, str):
try:
@@ -582,93 +63,22 @@ def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
if not isinstance(parsed, dict):
raise ValueError(
- f"--stage-overrides must be a JSON object mapping stage_id -> overrides, "
+ "--stage-overrides must be a JSON object mapping stage_id -> overrides, "
f"got {type(parsed).__name__}: {parsed!r}"
)
- if not parsed:
- return None
-
- for stage_id_str, overrides in parsed.items():
- # ``str.isdigit()`` accepts Unicode digit classes (e.g. fullwidth
- # ``"０"``); downstream's ``_STAGE_OVERRIDE_PATTERN`` is ASCII-only
- # (``^stage_(\d+)_(.+)$``), so a non-ASCII-digit key would parse
- # here and silently misroute as a global key. ``isascii()`` closes
- # that hole.
- if not isinstance(stage_id_str, str) or not stage_id_str.isdigit() or not stage_id_str.isascii():
+ for stage_id, overrides in parsed.items():
+ if not isinstance(stage_id, str) or not stage_id.isascii() or not stage_id.isdigit():
raise ValueError(
- f"--stage-overrides keys must be non-negative integer stage ids (as strings), got {stage_id_str!r}"
+ f"--stage-overrides keys must be non-negative integer stage ids (as strings), got {stage_id!r}"
)
if not isinstance(overrides, dict):
raise ValueError(
- f"--stage-overrides[{stage_id_str!r}] must be an object, got {type(overrides).__name__}: {overrides!r}"
+ f"--stage-overrides[{stage_id!r}] must be an object, got {type(overrides).__name__}: {overrides!r}"
)
return parsed
-def load_and_resolve_stage_configs(
- model: str,
- kwargs: dict | None,
- *,
- trust_remote_code: bool | None,
- default_stage_cfg_factory: Any = None,
- deploy_config_path: str | None = None,
- stage_overrides: dict[str, dict[str, Any]] | None = None,
- strategy_config_path: str | None = None,
-) -> tuple[str, list, str | None]:
- """Load stage configurations from a deploy YAML or model defaults.
-
- Args:
- model: Model name or path
- kwargs: Engine arguments to merge with stage configs
- trust_remote_code: Whether to trust remote code while resolving the model config.
- default_stage_cfg_factory: Optional callable that takes no args and returns
- default stage config list when no configs are found
- deploy_config_path: Optional path to a deploy YAML.
- stage_overrides: Per-stage overrides from ``--stage-overrides`` JSON.
- Keys are stage_id strings, values are dicts of overrides.
- strategy_config_path: Optional path to a composable-parallel
- ``strategy.yaml`` overlaid onto the registry-merged stages.
-
- Returns:
- Tuple of ``(config_path, stage_configs, omni_lb_policy)`` — the last is
- the strategy-derived pipeline-wide load-balancer policy (``None`` when no
- strategy set one), returned for the engine to apply.
- """
- config_path = deploy_config_path if deploy_config_path is not None else resolve_model_config_path(model)
- stage_configs, omni_lb_policy = load_stage_configs_from_model(
- model,
- trust_remote_code=trust_remote_code,
- base_engine_args=kwargs,
- deploy_config_path=deploy_config_path,
- stage_overrides=stage_overrides,
- strategy_config_path=strategy_config_path,
- )
- if not stage_configs:
- if default_stage_cfg_factory is not None:
- default_stage_cfg = default_stage_cfg_factory()
- stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg))
- else:
- stage_configs = []
-
- stage_configs = filter_stages(config_path, stage_configs, kwargs)
-
- # ``revision`` selects the model source for every stage. Keep the global
- # CLI pin on each stage's engine args so diffusion config construction can
- # pass it through to all component loaders. The registry-to-legacy-stage
- # conversion currently drops this inherited vLLM EngineArgs field.
- revision = (kwargs or {}).get("revision")
- if revision is not None:
- for stage_config in stage_configs:
- engine_args = stage_config.get("engine_args") if hasattr(stage_config, "get") else None
- if engine_args is not None and engine_args.get("revision") is None:
- engine_args["revision"] = revision
-
- logger.debug(f"stage_configs: {stage_configs}")
-
- return config_path, stage_configs, omni_lb_policy
-
-
def get_final_stage_id_for_e2e(
output_modalities: list[str] | None, default_modalities: list[str], stage_list: list
) -> int:
@@ -839,7 +249,7 @@ def detect_pid_host() -> bool:
if not ic:
return True
- return has_pid_host()
+ return has_pid_host() is True
### Helpers for handling delta messages