diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md index 862516831..8b7168824 100644 --- a/docs/design/architecture_overview.md +++ b/docs/design/architecture_overview.md @@ -323,7 +323,7 @@ inputs, fields, and ownership rules are described below. flowchart TB layer1["Layer 1 · Authoring inputs
PipelineConfig + DeployConfig"] layer2["Layer 2 · Production resolution boundary
resolve_omni_config()
(StageConfigFactory + VllmOmniConfig internally)"] - layer3["Layer 3 · Startup hand-off
OmniConfigResolution
PipelineConfig + temporary OmegaConf stage bridge"] + layer3["Layer 3 · Startup hand-off
OmniConfigResolution
PipelineConfig + typed stage configs"] layer4["Layer 4 · Runtime launch planning
StageRuntime"] layer5["Layer 5 · Engine materialization
VllmConfig / OmniDiffusionConfig"] @@ -335,10 +335,10 @@ The production resolution boundary is `resolve_omni_config()` in 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. +This envelope carries the effective `PipelineConfig` alongside typed +`stage_configs`, which `StageRuntime` consumes directly. Both views describe the +same resolved topology, including injected stages. Backend arguments are +projected from these typed stages when each engine is initialized. The legacy `stage_args` YAML path has been removed. Model topology now resolves through `PipelineConfig`, with runtime overrides supplied by `DeployConfig`. @@ -360,7 +360,7 @@ The important ownership rules are: 3. CLI and Python overrides are applied at the resolution boundary, with per-stage overrides taking precedence over global values where supported. 4. `OmniConfigResolution` is the sole startup hand-off. Its `pipeline_config` - and temporary `stage_configs` compatibility view must describe the same + and typed `stage_configs` must describe the same topology. 5. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan` is runtime-private state, not a user configuration object. diff --git a/tests/config/composable_parallel/test_factory_integration.py b/tests/config/composable_parallel/test_factory_integration.py index 468b3db2c..d49d32042 100644 --- a/tests/config/composable_parallel/test_factory_integration.py +++ b/tests/config/composable_parallel/test_factory_integration.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Integration tests: strategy overlay on the real merge_pipeline_deploy seam.""" from __future__ import annotations @@ -20,6 +20,7 @@ from vllm_omni.config.composable_parallel import ( apply_strategy_specs, ) from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config.omni_config import VllmOmniConfig from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy @@ -137,3 +138,31 @@ def test_cli_overrides_strategy_with_warning(): assert _resolved(stages, "talker").runtime.num_replicas == 3 # ...and the override was warned about, naming the conflicting field. assert any("num_replicas" in m and "overrides the strategy-derived" in m for m in messages) + + +def test_typed_cli_overrides_strategy(): + pipeline = OMNI_PIPELINES["qwen2_5_omni"] + config = VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=load_deploy_config(_DEPLOY), + cli_overrides={"stage_1_num_replicas": 3}, + strategy_specs={"talker": [_stage_replica(2, "round_robin")]}, + ) + + talker = next(stage for stage in config.stage_configs if stage.model_stage == "talker") + assert talker.runtime_config.num_replicas == 3 + assert config.orchestrator_config.omni_lb_policy == "round-robin" + assert config.strategy_omni_lb_policy == "round-robin" + + +def test_typed_tp_only_strategy_has_no_derived_lb_policy(): + pipeline = OMNI_PIPELINES["qwen2_5_omni"] + config = VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=load_deploy_config(_DEPLOY), + cli_overrides={"omni_lb_policy": "round-robin"}, + strategy_specs={"thinker": [_tp(1)]}, + ) + + assert config.orchestrator_config.omni_lb_policy == "round-robin" + assert config.strategy_omni_lb_policy is None diff --git a/tests/config/test_forced_aligner_injection.py b/tests/config/test_forced_aligner_injection.py index 65c315229..2037b22f1 100644 --- a/tests/config/test_forced_aligner_injection.py +++ b/tests/config/test_forced_aligner_injection.py @@ -2,13 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest +from vllm.pooling_params import PoolingParams import vllm_omni.config.pipeline_registry # noqa: F401 (populate registry) +from vllm_omni.config.omni_config import VllmOmniConfig from vllm_omni.config.pipeline_registry import OMNI_PIPELINES as _PIPELINE_REGISTRY from vllm_omni.config.stage_config import ( DeployConfig, StageExecutionType, ) +from vllm_omni.engine.stage_init_utils import ( + build_engine_args_dict_from_omni_stage_config, + extract_stage_metadata_from_omni_stage_config, +) from vllm_omni.model_executor.stage_input_processors.forced_aligner import ( POOLING_OUTPUT_DECODER_PATH, TIMESTAMPS_MODALITY, @@ -94,3 +100,33 @@ def test_aligner_uses_completed_audio_without_disabling_upstream_chunks(async_ch assert stages[2].engine_args.runner == "pooling" assert len(pipeline.stages) == 2 assert deploy.stages == [] + + +def test_injected_aligner_survives_typed_config_and_runtime_projections(): + pipeline, deploy = inject_forced_aligner_stage( + _PIPELINE_REGISTRY["qwen3_tts"], + DeployConfig(), + {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"}, + ) + config = VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=deploy, + cli_overrides={"model": "/models/Qwen3-TTS"}, + ) + stage = config.stage_configs[-1] + + assert stage.pooling_config.runner == "pooling" + assert stage.pooling_config.pooling_output_decoder == POOLING_OUTPUT_DECODER_PATH + assert isinstance(stage.pooling_config.default_pooling_params, PoolingParams) + assert stage.pooling_config.default_pooling_params.task == deploy.stages[-1].default_pooling_params["task"] + + metadata = extract_stage_metadata_from_omni_stage_config(stage) + assert isinstance(metadata.default_sampling_params, PoolingParams) + assert metadata.default_sampling_params.task == deploy.stages[-1].default_pooling_params["task"] + + engine_args = build_engine_args_dict_from_omni_stage_config( + stage, + stage.model_config.model or "/models/Qwen3-ForcedAligner-0.6B", + ) + assert engine_args["runner"] == "pooling" + assert engine_args["pooling_output_decoder"] == POOLING_OUTPUT_DECODER_PATH diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 1ed62f7f1..97f9a5370 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -11,6 +11,7 @@ from pathlib import Path import msgspec import pytest +import torch from pydantic import ValidationError from pydantic.fields import FieldInfo from transformers import Qwen3OmniMoeConfig @@ -123,6 +124,18 @@ def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) - assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["turn", "turn"] +def test_nested_stage_override_deep_merges_structured_model_config() -> None: + config = _from_pipeline_key( + "cosmos3_policy", + deploy_config_path=get_deploy_config_path("cosmos3_policy_droid.yaml"), + cli_overrides={"stage_0_model_config": {"guardrails": False}}, + ) + + model_config = config.stage_by_id(0).diffusion_config.model_config + assert model_config["guardrails"] is False + assert model_config["policy_server_config"]["action_space"] == "joint_position" + + @pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES)) def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str): pipeline = _resolve_pipeline_or_skip(model_type) @@ -654,6 +667,7 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations(): "load_config", "cache_config", "scheduler_config", + "pooling_config", "connector_config", "runtime_config", "parallel_config", @@ -889,6 +903,17 @@ def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_saf msgspec.msgpack.encode(payload) +def test_structured_diffusion_torch_dtype_is_msgpack_transport_safe(): + stage_config = _from_pipeline_key( + "hunyuan_image3_dit", + cli_overrides={"dtype": torch.bfloat16}, + ).stage_by_id(0) + + serialized = _serialize_stage_config(stage_config) + assert serialized["diffusion_config"]["dtype"] == "bfloat16" + msgspec.msgpack.encode(serialized) + + def test_diffusion_parallel_config_fields_cover_legacy_surface(): from vllm_omni.diffusion.data import DiffusionParallelConfig diff --git a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py index 56a960121..4bea6563b 100644 --- a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py +++ b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py @@ -21,6 +21,7 @@ import torch.distributed as dist from transformers import AutoTokenizer from tests.helpers.mark import hardware_test +from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient from vllm_omni.entrypoints.async_omni import AsyncOmni @@ -136,6 +137,15 @@ async def _generate_once( def _assert_live_uniproc_worker(engine: AsyncOmni) -> None: """Default single-GPU path must use the in-process executor, not mp IPC.""" + stage_configs = engine.engine.stage_configs + assert len(stage_configs) == 1 + stage_config = stage_configs[0] + assert isinstance(stage_config, VllmOmniDiffusionStageConfig) + assert stage_config.model_config.model == MODEL + assert stage_config.model_config.enforce_eager is True + assert stage_config.scheduler_config.max_num_seqs == 1 + assert stage_config.diffusion_config.custom_pipeline_args == {"pipeline_class": CUSTOM_PIPELINE_CLASS} + assert stage_config.diffusion_config.worker_extension_cls == WORKER_EXTENSION_CLASS clients = engine.engine.stage_clients assert len(clients) == 1 client = clients[0] diff --git a/tests/e2e/offline_inference/test_hunyuanimage3.py b/tests/e2e/offline_inference/test_hunyuanimage3.py index 8655bd63f..21bfe9b39 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3.py @@ -1,5 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # ruff: noqa: E501 +import copy import importlib.util +import os import sys from collections.abc import Generator from pathlib import Path @@ -10,20 +15,37 @@ import torch.nn.functional as F from PIL import Image from transformers import CLIPModel, CLIPProcessor +from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner -from tests.helpers.stage_config import get_deploy_config_path +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni import Omni +from vllm_omni.config.omni_config import ( + VllmOmniARStageConfig, + VllmOmniDiffusionStageConfig, +) +from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_extras import get_ar_input_builder, get_ar_tokenizer_validator, get_model_class_name from vllm_omni.platforms import current_omni_platform PROMPT = "A brown and white dog is running on the grass" -MODEL_NAME = "tencent/HunyuanImage-3.0" +MODEL_NAME = os.environ.get("HUNYUAN_IMAGE3_MODEL", "tencent/HunyuanImage-3.0-Instruct") LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32" -DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image_3_moe.yaml") +_PRODUCTION_DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image_3_moe.yaml") +# Keep the production Mooncake connector, but use its TCP transport for this +# single-host test. The shared server's RDMA device is not usable reliably; +# TCP still exercises the Mooncake AR→DiT data plane without changing deploy +# semantics in production. +DEPLOY_CONFIG_PATH = modify_stage_config( + _PRODUCTION_DEPLOY_CONFIG_PATH, + updates={"connectors.rdma_connector.extra.protocol": "tcp"}, +) -pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] +pytestmark = [ + pytest.mark.advanced_model, + pytest.mark.diffusion, +] # System prompt type. Options: None, dynamic, en_vanilla, en_recaption, en_think_recaption, en_unified # Below are the CLIP embedding tensors from the official HunyuanImage model (seed=1234, prompt: "A brown and white dog is running on the grass"). @@ -285,25 +307,97 @@ def omni() -> Generator[Omni, None, None]: yield runner.omni +@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.") +@hardware_test(res={"cuda": "H100"}, num_cards=4) +def test_structured_mixed_pipeline_config_reaches_runtime(omni: Omni) -> None: + """Mixed-pipeline deploy settings reach both terminal engine configs.""" + engine = omni.engine + + # Topology, connector selection, and placement remain Omni-owned. + resolved_stages = engine.stage_configs + assert len(resolved_stages) == 2 + ar_stage, diffusion_stage = resolved_stages + assert isinstance(ar_stage, VllmOmniARStageConfig) + assert ar_stage.stage_id == 0 + assert ar_stage.model_stage == "AR" + assert ar_stage.final_output_type == "text" + assert ar_stage.runtime_config.devices == "0,1" + assert ar_stage.connector_config.output_connectors == {"to_stage_1": "rdma_connector"} + assert isinstance(diffusion_stage, VllmOmniDiffusionStageConfig) + assert diffusion_stage.stage_id == 1 + assert diffusion_stage.model_stage == "dit" + assert diffusion_stage.final_output_type == "image" + assert diffusion_stage.input_sources == [0] + assert diffusion_stage.runtime_config.devices == "2,3" + assert diffusion_stage.connector_config.input_connectors == {"from_stage_0": "rdma_connector"} + + # Stage 0 is materialized as the VllmConfig actually passed to its engine. + ar_vllm_config, diffusion_vllm_config = engine.stage_vllm_configs + assert ar_vllm_config is not None + assert diffusion_vllm_config is None + assert ar_vllm_config.model_config.enforce_eager is True + assert ar_vllm_config.model_config.trust_remote_code is True + assert ar_vllm_config.parallel_config.tensor_parallel_size == 2 + assert ar_vllm_config.scheduler_config.max_num_seqs == 1 + assert ar_vllm_config.scheduler_config.max_num_batched_tokens == 32768 + assert ar_vllm_config.cache_config.gpu_memory_utilization == 0.95 + ar_omni_config = ar_vllm_config.model_config.omni_config + assert ar_omni_config.omni_kv_config["need_send_cache"] is True + + expected_ar_sampling = { + "temperature": 0.0, + "top_p": 1, + "top_k": -1, + "max_tokens": 8192, + "detokenize": True, + "skip_special_tokens": False, + "include_stop_str_in_output": True, + } + ar_sampling, diffusion_sampling = engine.default_sampling_params_list + assert all(getattr(ar_sampling, name) == value for name, value in expected_ar_sampling.items()) + assert diffusion_sampling.num_inference_steps == 50 + assert diffusion_sampling.guidance_scale == 0 + + # Stage 1 is a diffusion engine, so its terminal owner is the launched + # OmniDiffusionConfig rather than a VllmConfig. + od_config = omni.get_diffusion_od_config() + assert od_config is not None + assert od_config.model == MODEL_NAME + assert od_config.enforce_eager is True + assert od_config.distributed_executor_backend == "mp" + assert od_config.parallel_config.tensor_parallel_size == 2 + assert od_config.parallel_config.enable_expert_parallel is True + assert od_config.max_num_seqs == 1 + diffusion_kv = od_config.omni_kv_config + assert diffusion_kv["need_recv_cache"] is True + assert diffusion_kv["enable_kv_async_prefetch"] is True + assert diffusion_kv["kv_prefetch_min_free_mem_ratio"] == 0.2 + + def _extract_generated_image(outputs: list[object]) -> Image.Image: if not outputs: raise AssertionError("No outputs were returned from Omni.generate()") - first_output = outputs[0] - if images := getattr(first_output, "images", None): - return images[0] - - request_output = first_output - if request_output is not None and (images := getattr(request_output, "images", None)): - return images[0] + for output in outputs: + if images := getattr(output, "images", None): + return images[0] + multimodal_output = getattr(output, "multimodal_output", None) + if isinstance(multimodal_output, dict) and (images := multimodal_output.get("images")): + return images[0] raise AssertionError("No generated image found in Omni output") def extract_embedding(image: Image.Image, clip_model: CLIPModel, clip_processor: CLIPProcessor) -> torch.Tensor: - inputs = clip_processor(images=image.convert("RGB"), return_tensors="pt") + # The full CLIP forward API is stable across supported Transformers + # versions, but it expects text inputs as well as pixels. An empty text + # prompt is sufficient when only the projected image embedding is needed. + inputs = clip_processor(text=[""], images=[image.convert("RGB")], return_tensors="pt", padding=True) with torch.inference_mode(): - features = clip_model.get_image_features(**inputs) + # Use the public CLIP forward output, matching the shared CLIPScorer + # helper. ``get_image_features`` changed from Tensor to + # BaseModelOutputWithPooling in Transformers 5.x. + features = clip_model(**inputs).image_embeds features = F.normalize(features, p=2, dim=-1) return features.squeeze(0) @@ -320,16 +414,36 @@ def compare_semantic( def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image: - generator_device = current_omni_platform.device_type or "cuda" - sampling_params = OmniDiffusionSamplingParams( - seed=1234, - generator=torch.Generator(device=generator_device).manual_seed(1234), - num_outputs_per_prompt=1, - ) - if use_system_prompt is not None: - sampling_params.extra_args = {"use_system_prompt": use_system_prompt} + from transformers import AutoTokenizer - outputs = omni.generate({"prompt": PROMPT}, sampling_params) + tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) + prompt_tokens = build_prompt_tokens( + PROMPT, + tokenizer, + task="t2i", + bot_task="think", + sys_type=use_system_prompt or "None", + ) + stop_token_ids = resolve_stop_token_ids(task="t2i", bot_task="think", tokenizer=tokenizer, image_size="1024x1024") + params_list = copy.deepcopy(omni.default_sampling_params_list) + generator_device = current_omni_platform.device_type or "cuda" + for params in params_list: + if isinstance(params, OmniDiffusionSamplingParams): + params.seed = 1234 + params.generator = torch.Generator(device=generator_device).manual_seed(1234) + params.num_outputs_per_prompt = 1 + elif hasattr(params, "stop_token_ids"): + params.stop_token_ids = stop_token_ids + + prompt = { + "prompt": PROMPT, + "prompt_token_ids": prompt_tokens.token_ids, + "use_system_prompt": prompt_tokens.system_prompt_type, + "modalities": ["image"], + "height": 1024, + "width": 1024, + } + outputs = omni.generate(prompt, params_list) return _extract_generated_image(outputs) @@ -342,7 +456,7 @@ def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image: MIN_SEMANTIC_SIMILARITY = 0.5 -@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.") +@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.") @pytest.mark.parametrize("system_prompt_name,use_system_prompt,expected_embedding", SYSTEM_PROMPT_CASES) def test_system_prompt_scores( omni: Omni, @@ -378,7 +492,7 @@ def _load_text_to_image_module(): return module -@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.") +@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs for this test.") def test_shared_script_ar_path_reaches_generation( omni: Omni, clip_bundle: tuple[CLIPModel, CLIPProcessor], diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py index 32c345194..f5937265c 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni.py +++ b/tests/e2e/offline_inference/test_qwen3_omni.py @@ -17,6 +17,10 @@ import pytest from tests.helpers.mark import hardware_test from tests.helpers.media import generate_synthetic_video from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config +from vllm_omni.config.omni_config import ( + VllmOmniARStageConfig, + VllmOmniGenerationStageConfig, +) from vllm_omni.platforms import current_omni_platform models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] @@ -25,6 +29,7 @@ thinker_only_models = ["Qwen/Qwen3-Omni-30B-A3B-Captioner"] # Single CI deploy YAML; rocm/xpu deltas are picked automatically via the # platforms: section. Only CUDA needs an extra enforce_eager tweak. _CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml") +_PRODUCTION_DEPLOY = get_deploy_config_path("qwen3_omni_moe.yaml") def get_cuda_graph_config(): @@ -46,6 +51,7 @@ else: # Create parameter combinations for model and stage config test_params = [(model, stage_config) for model in models for stage_config in stage_configs] +production_test_params = [(model, _PRODUCTION_DEPLOY) for model in models] # we can use the same config for a model that only has thinker (i.e., does not # enable audio output) because the resolver should figure out that it doesn't # need the full pipeline based on the HF config. @@ -59,6 +65,85 @@ def get_question(prompt_type="video"): return prompts.get(prompt_type, prompts["video"]) +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", production_test_params, indirect=True) +def test_structured_multistage_config_reaches_runtime(omni_runner, offline_client) -> None: + """Deploy settings reach the materialized configs used by live stages.""" + engine = omni_runner.omni.engine + + # The resolved typed stages remain the Omni-owned source for topology and + # placement, which are not fields of the terminal vLLM config. + resolved_stages = engine.stage_configs + assert len(resolved_stages) == 3 + thinker, talker, code2wav = resolved_stages + assert isinstance(thinker, VllmOmniARStageConfig) + assert isinstance(talker, VllmOmniARStageConfig) + assert isinstance(code2wav, VllmOmniGenerationStageConfig) + assert [stage.stage_id for stage in resolved_stages] == [0, 1, 2] + assert [stage.model_stage for stage in resolved_stages] == ["thinker", "talker", "code2wav"] + assert code2wav.final_output_type == "audio" + assert thinker.runtime_config.devices == "0" + assert talker.runtime_config.devices == code2wav.runtime_config.devices == "1" + assert talker.connector_config.input_connectors == {"from_stage_0": "connector_of_shared_memory"} + assert code2wav.connector_config.input_connectors == {"from_stage_1": "connector_of_shared_memory"} + + # Engine-owned deploy settings must survive the startup projection into + # the actual VllmConfig objects retained by the stage pools. + assert len(engine.stage_vllm_configs) == 3 + thinker_vllm, talker_vllm, code2wav_vllm = engine.stage_vllm_configs + assert thinker_vllm is not None + assert talker_vllm is not None + assert code2wav_vllm is not None + assert thinker_vllm.scheduler_config.max_num_seqs == 64 + assert talker_vllm.scheduler_config.max_num_seqs == 64 + assert code2wav_vllm.scheduler_config.max_num_seqs == 64 + assert thinker_vllm.scheduler_config.max_num_batched_tokens == 32768 + assert talker_vllm.scheduler_config.max_num_batched_tokens == 32768 + assert code2wav_vllm.scheduler_config.max_num_batched_tokens == 65536 + assert thinker_vllm.cache_config.gpu_memory_utilization == 0.9 + assert talker_vllm.cache_config.gpu_memory_utilization == 0.6 + assert code2wav_vllm.cache_config.gpu_memory_utilization == 0.1 + assert all(config.cache_config.enable_prefix_caching is False for config in engine.stage_vllm_configs) + assert all(config.model_config.trust_remote_code is True for config in engine.stage_vllm_configs) + assert code2wav_vllm.scheduler_config.enable_chunked_prefill is False + assert code2wav_vllm.scheduler_config.async_scheduling is False + + # Sampling defaults are consumed from StageClient metadata, rather than + # from either the resolver output or VllmConfig. + expected_sampling = ( + {"temperature": 0.0, "max_tokens": 2048}, + {"temperature": 0.9, "top_k": 50, "max_tokens": 4096, "repetition_penalty": 1.05}, + { + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 65536, + "repetition_penalty": 1.1, + }, + ) + assert len(engine.default_sampling_params_list) == len(expected_sampling) + for runtime_params, expected in zip(engine.default_sampling_params_list, expected_sampling, strict=True): + # SamplingParams normalization may add backend defaults (for example + # ``detokenize=True``). Verify deploy-owned values without rejecting + # those non-lossy normalized fields. + assert all(getattr(runtime_params, name) == value for name, value in expected.items()) + + if current_omni_platform.is_cuda(): + assert thinker_vllm.model_config.enforce_eager is False + assert talker_vllm.model_config.enforce_eager is False + + # Exercise the materialized stages and their connector path with a real + # request; the assertions above are not merely resolver-shape checks. + offline_client.send_omni_request( + { + "prompts": "Answer with one short sentence: what is the capital of China?", + "modalities": ["text"], + } + ) + + @pytest.mark.advanced_model @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) diff --git a/tests/e2e/offline_inference/test_wan22_t2v.py b/tests/e2e/offline_inference/test_wan22_t2v.py index 805fbd1fe..07ac5c8a6 100644 --- a/tests/e2e/offline_inference/test_wan22_t2v.py +++ b/tests/e2e/offline_inference/test_wan22_t2v.py @@ -5,6 +5,7 @@ import pytest from tests.helpers.mark import hardware_test from tests.helpers.runtime import OfflineOmniClient +from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" @@ -12,6 +13,31 @@ PROMPT = "Two anthropomorphic cats in boxing gear on a spotlighted stage." NEGATIVE_PROMPT = "low quality, blurry, watermark, text" +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 1}) +@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True) +def test_structured_diffusion_config_reaches_runtime(omni_runner) -> None: + """The default WAN resolver launches the typed diffusion config itself.""" + omni = omni_runner.omni + resolved_stages = omni.engine.stage_configs + assert len(resolved_stages) == 1 + stage_config = resolved_stages[0] + + assert isinstance(stage_config, VllmOmniDiffusionStageConfig) + assert stage_config.stage_id == 0 + assert stage_config.model_stage == "dit" + assert stage_config.final_output_type == "video" + + assert omni.engine.stage_vllm_configs == [None] + od_config = omni.get_diffusion_od_config() + assert od_config is not None + assert od_config.stage_id == 0 + assert od_config.model == MODEL + assert od_config.model_class_name == "WanPipeline" + assert od_config.output_type == "pil" + + @pytest.mark.advanced_model @pytest.mark.diffusion @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 1}) diff --git a/tests/engine/test_async_omni_engine_abort_ack.py b/tests/engine/test_async_omni_engine_abort_ack.py index 43a29bf3a..366e640f6 100644 --- a/tests/engine/test_async_omni_engine_abort_ack.py +++ b/tests/engine/test_async_omni_engine_abort_ack.py @@ -66,9 +66,10 @@ def test_abort_async_waits_for_ack(mocker: MockerFixture): assert msg.rpc_id == "abort-rpc-1" assert not task.done() - rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True)) + partial_output = SimpleNamespace(request_id="req-a", finished=True) + rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True, abort_outputs=[partial_output])) result = await task - assert result == [] + assert result == [partial_output] try: asyncio.run(_run()) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 15d31cbbc..637835122 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -12,6 +12,7 @@ import types import pytest from omegaconf import OmegaConf +from vllm_omni.config.omni_config import OmniStageRuntimeConfig from vllm_omni.diffusion.data import AttentionConfig from vllm_omni.engine import omni_engine_base as async_omni_engine_module from vllm_omni.engine.async_omni_engine import AsyncOmniEngine @@ -21,12 +22,23 @@ from vllm_omni.engine.stage_init_utils import ( build_stage0_input_processor, compute_replica_layout, split_devices_for_replicas, + stage_runtime_env, ) from vllm_omni.engine.stage_runtime import StageRuntime pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_stage_runtime_env_accepts_typed_runtime_config(monkeypatch): + env_key = "VLLM_OMNI_TEST_TYPED_STAGE_ENV" + monkeypatch.delenv(env_key, raising=False) + + with stage_runtime_env(0, OmniStageRuntimeConfig(env={env_key: "typed-value"})): + assert os.environ[env_key] == "typed-value" + + assert env_key not in os.environ + + def test_orchestrator_startup_timeout_warns_how_to_raise_limits(monkeypatch): engine = object.__new__(AsyncOmniEngine) engine.orchestrator_thread = types.SimpleNamespace(is_alive=lambda: True) @@ -223,6 +235,50 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp assert captured["log_stats"] is True +def test_async_omni_engine_initialize_stages_retains_stage0_prompt_transform(monkeypatch): + import vllm_omni.engine.omni_engine_base as engine_mod + + engine = object.__new__(AsyncOmniEngine) + engine.stage_configs = [types.SimpleNamespace()] + engine.model = "dummy-model" + engine.config_path = "dummy-config" + engine.single_stage_mode = False + engine.async_chunk = False + engine.tokenizer = None + engine._single_stage_id_filter = None + engine._omni_master_address = None + engine._omni_master_port = None + engine._omni_dp_size_local = 1 + engine._omni_heartbeat_timeout = 30.0 + engine._omni_lb_policy = "random" + engine.request_queue = types.SimpleNamespace() + engine._log_stats = False + engine._parallel_stage_init = False + + prompt_transform = object() + client = types.SimpleNamespace( + prompt_transform_func=prompt_transform, + prompt_expand_func=None, + default_sampling_params=types.SimpleNamespace(), + final_output=True, + final_output_type="text", + stage_type="llm", + model_stage="text_encoder", + is_comprehension=True, + ) + pool = types.SimpleNamespace( + stage_client=client, + stage_vllm_config=None, + output_processor=None, + ) + runtime = types.SimpleNamespace(stage_pools=[pool], initialize=lambda: None) + monkeypatch.setattr(engine_mod, "create_stage_runtime", lambda **_kwargs: runtime) + + engine._initialize_stages(stage_init_timeout=7) + + assert engine.prompt_transform_func is prompt_transform + + def test_compute_replica_layout_splits_diffusion_devices_by_world_size(): stage_cfg = types.SimpleNamespace( stage_id=0, @@ -1715,6 +1771,51 @@ 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_inject_kv_stage_info_updates_typed_connector_config(): + from vllm_omni.config.omni_config import ( + OmniStageConnectorConfig, + OmniStageDiffusionParallelConfig, + VllmOmniDiffusionStageConfig, + ) + from vllm_omni.config.stage_config import StageExecutionType, StagePipelineConfig + from vllm_omni.engine.stage_init_utils import inject_kv_stage_info + from vllm_omni.entrypoints.utils import inject_omni_kv_config + + stage0 = VllmOmniDiffusionStageConfig( + stage_pipeline_config=StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + ), + connector_config=OmniStageConnectorConfig( + omni_kv_config={ + "need_send_cache": True, + "omni_from_stage": "0", + "omni_to_stage": "1", + } + ), + parallel_config=OmniStageDiffusionParallelConfig(tensor_parallel_size=4), + ) + stage1 = VllmOmniDiffusionStageConfig( + stage_pipeline_config=StagePipelineConfig( + stage_id=1, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(0,), + ), + connector_config=OmniStageConnectorConfig(omni_kv_config={"need_recv_cache": True}), + parallel_config=OmniStageDiffusionParallelConfig(tensor_parallel_size=2), + ) + + inject_omni_kv_config(stage0, {"kv_connector": "P2pNcclConnector"}, "0", "1") + inject_kv_stage_info(stage0, 0, [stage0, stage1]) + + assert stage0.connector_config.omni_kv_config["stage_id"] == 0 + assert stage0.connector_config.omni_kv_config["connector_config"] == {"kv_connector": "P2pNcclConnector"} + assert stage0.connector_config.omni_kv_config["engine_input_source"] == [] + assert stage0.connector_config.omni_kv_config["rank_mapping"] == {"from_tp": 4, "to_tp": 2} + + 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 @@ -1852,3 +1953,80 @@ def test_engine_async_chunk_includes_downstream_stages(monkeypatch, stage_modes, with pytest.raises(ConfigResolvedError): engine.__init__("test-model") assert engine.async_chunk is expected + + +@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"]) +def test_dist_stage_runtime_applies_local_dp_to_stage_config(typed): + from vllm_omni.config.omni_config import VllmOmniARStageConfig + from vllm_omni.config.stage_config import StagePipelineConfig + from vllm_omni.engine.stage_runtime import DistStageRuntime + + if typed: + stage = VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="ar")) + runtime_cfg = stage.runtime_config + else: + runtime_cfg = types.SimpleNamespace(num_replicas=1) + stage = types.SimpleNamespace(stage_id=0, runtime=runtime_cfg) + runtime = DistStageRuntime( + stage_configs=[stage], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + async_chunk=False, + single_stage_id_filter=0, + omni_master_address="127.0.0.1", + omni_master_port=12345, + omni_dp_size_local=2, + ) + + runtime._validate_single_stage_mode_replica_constraints() + + assert runtime_cfg.num_replicas == 2 + + +@pytest.mark.parametrize("num_replicas", [1, 2, 3]) +def test_typed_diffusion_replicas_share_one_config_between_planning_and_launch(mocker, num_replicas): + from vllm_omni.config.config_factory import StageConfigFactory + from vllm_omni.engine import stage_runtime as runtime_module + + stage = StageConfigFactory.create_typed_default_diffusion( + "generic-diffusion", {"model_class_name": "QwenImagePipeline"} + ).stage_configs[0] + devices = ",".join(str(i) for i in range(num_replicas)) + stage.runtime_config.devices = devices + stage.runtime_config.num_replicas = num_replicas + runtime = StageRuntime( + stage_configs=[stage], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + async_chunk=False, + ) + mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={}) + mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None)) + mocker.patch.object(runtime, "_stage_device_scope", side_effect=lambda *_: contextlib.nullcontext()) + client = mocker.Mock() + launch = mocker.patch.object(runtime_module, "launch_diffusion_stage_replica", return_value=(client, None)) + + # Replanning must not inherit the first replica's narrowed device slice. + for _ in range(2): + counts, device_map = compute_replica_layout([stage]) + plans = runtime._build_logical_stage_init_plans( + omni_transfer_config=None, + replicas_per_stage=counts, + replica_devices_map=device_map, + ) + replicas = plans[0].replicas + assert stage.runtime_config.devices == devices + assert len({id(plan.stage_cfg) for plan in replicas}) == num_replicas + for i, plan in enumerate(replicas): + assert plan.metadata.runtime_cfg is plan.stage_cfg.runtime_config + assert plan.stage_cfg.runtime_config.devices == str(i) + if num_replicas > 1: + assert plan.stage_cfg is not stage + assert runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1) is client + assert launch.call_args.kwargs["stage_config"] is plan.stage_cfg + assert launch.call_args.kwargs["metadata"] is plan.metadata + assert launch.call_args.kwargs["use_inline"] is (num_replicas == 1) + + assert launch.call_count == 2 * num_replicas diff --git a/tests/engine/test_parallel_stage_init.py b/tests/engine/test_parallel_stage_init.py index d88d49825..9f89aabca 100644 --- a/tests/engine/test_parallel_stage_init.py +++ b/tests/engine/test_parallel_stage_init.py @@ -306,6 +306,23 @@ def test_check_admission_admits_diffusion_with_explicit_util(): assert ledgers[0].kv_budget_bytes == int(cap * 0.3) +def test_check_admission_reads_typed_diffusion_cache_config(): + cap = 80 * 1024**3 + diff = _local_diffusion_replica(0, 0, "0", util=None) + diff.stage_cfg.cache_config = types.SimpleNamespace(gpu_memory_utilization=0.35) + plan = LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[diff]) + + ledgers = check_admission( + [plan], + resolve_physical_devices=lambda _r: [0], + device_total_memory=lambda _d: cap, + external_reserve_bytes=0, + safety_margin_bytes=0, + ) + + assert ledgers[0].kv_budget_bytes == int(cap * 0.35) + + def test_check_admission_unresolved_local_raises(): """A local replica invisible to the ledger must fail admission, not skip it: it would otherwise initialize concurrently with unbounded demand.""" diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 41d63af9e..9ce707dbc 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -16,7 +16,8 @@ from pytest_mock import MockerFixture from vllm.v1.engine.utils import EngineZmqAddresses from vllm_omni.config.config_factory import StageConfigFactory -from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.config.omni_config import VllmOmniARStageConfig, VllmOmniDiffusionStageConfig +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, StageExecutionType, StagePipelineConfig from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase from vllm_omni.engine.stage_engine_startup import ( @@ -57,6 +58,10 @@ def _make_llm_plan( vllm_config: Any | None = None, ) -> LogicalStageInitPlan: stage_cfg = _make_stage_cfg(stage_id) + if launch_mode == "remote": + stage_cfg = VllmOmniARStageConfig( + stage_pipeline_config=StagePipelineConfig(stage_id=stage_id, model_stage="thinker"), + ) metadata = SimpleNamespace( stage_id=stage_id, stage_type="llm", @@ -98,6 +103,16 @@ def _make_diffusion_plan( launch_mode: str, ) -> LogicalStageInitPlan: stage_cfg = _make_stage_cfg(stage_id, stage_type="diffusion") + if launch_mode == "remote": + stage_cfg = VllmOmniDiffusionStageConfig( + stage_pipeline_config=StagePipelineConfig( + stage_id=stage_id, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + final_output=True, + final_output_type="image", + ), + ) metadata = SimpleNamespace( stage_id=stage_id, stage_type="diffusion", @@ -960,6 +975,7 @@ class TestSingleStageReplicaInitialization: assert mock_connect.call_args.kwargs["stage_id"] == 1 assert mock_connect.call_args.kwargs["replica_id"] == 0 assert client_kwargs["log_stats"] is True + assert client_kwargs["metadata"].model_stage == "thinker" assert events == ["enter", "exit", "attach"] def test_initialize_llm_replica_remote_missing_registered_stage_config_raises(self, mocker: MockerFixture): @@ -1125,11 +1141,9 @@ class TestSingleStageReplicaInitialization: return _ctx() - remote_metadata = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0].metadata plan = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0] sentinel_client = SimpleNamespace() - mocker.patch.object(runtime_mod, "extract_legacy_stage_metadata", return_value=remote_metadata) mock_connect = mocker.patch.object(runtime_mod, "connect_remote_diffusion_proc", side_effect=_fake_connect) mock_from_addresses = mocker.patch( "vllm_omni.diffusion.stage_diffusion_client.StageDiffusionClient.from_addresses", @@ -1146,6 +1160,7 @@ class TestSingleStageReplicaInitialization: replica_id=0, ) mock_from_addresses.assert_called_once() + assert mock_from_addresses.call_args.args[0].final_output_type == "image" def test_initialize_local_diffusion_replica_registers_with_master(self, mocker: MockerFixture): import vllm_omni.engine.stage_runtime as runtime_mod diff --git a/tests/engine/test_stage_device_layout.py b/tests/engine/test_stage_device_layout.py index 0473dc341..beb66dfba 100644 --- a/tests/engine/test_stage_device_layout.py +++ b/tests/engine/test_stage_device_layout.py @@ -156,6 +156,43 @@ def test_replica_guard_and_splitter_share_local_world_size(devices, expected): assert replica_devices_map == {0: expected} +@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"]) +@pytest.mark.parametrize("local_dp", [None, 0, 1, 2]) +@pytest.mark.parametrize("pool_devices", [False, True], ids=["template", "pool"]) +def test_replica_layout_preserves_local_dp_semantics(typed, local_dp, pool_devices): + from vllm_omni.config.omni_config import VllmOmniConfig + from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig + from vllm_omni.engine.stage_engine_startup import get_headless_replica_devices + + engine_args = { + "tensor_parallel_size": 2, + "data_parallel_size": 4, + "data_parallel_size_local": local_dp, + "pipeline_parallel_size": 2, + } + width = 2 * max(1, local_dp if local_dp is not None else 4) * 2 + devices = ",".join(str(i) for i in range(width * (2 if pool_devices else 1))) + if typed: + pipeline = PipelineConfig( + model_type="test", + model_arch="TestModel", + stages=(StagePipelineConfig(stage_id=0, model_stage="ar", final_output=True),), + ) + stage = VllmOmniConfig.from_pipeline_config( + pipeline, + cli_overrides={ + f"stage_0_{key}": value for key, value in {**engine_args, "devices": devices, "num_replicas": 2}.items() + }, + ).stage_configs[0] + else: + stage = _stage(0, devices=devices, num_replicas=2, engine_args=engine_args) + expected = [",".join(str(i) for i in range(r * width, (r + 1) * width)) for r in range(2)] + + assert get_stage_devices_per_replica(stage) == width + assert compute_replica_layout([stage]) == ([2], {0: expected}) + assert get_headless_replica_devices(stage, stage_id=0, omni_dp_size_local=2) == expected + + def test_non_tp_layout_error_uses_generic_guidance(): """A PP mismatch must not be reported as a top-level TP broadcast.""" stage = _stage(0, devices="0") diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 7ac5b6b69..6650c8e45 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -95,6 +95,7 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset( "num_weight_load_threads", "omni_kv_config", "parallel_config", + "pooling_output_decoder", "silence_ban_frames", "subtalker_sampling_params", "task_type", @@ -366,6 +367,21 @@ def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults(): assert inherited_defaults.isdisjoint(engine_args) +def test_typed_llm_projection_omits_diffusion_only_and_process_only_defaults(): + stage_config = VllmOmniARStageConfig( + stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"), + ) + + engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config) + + assert { + "enable_multithread_weight_load", + "num_weight_load_threads", + "disable_autocast", + "log_level", + }.isdisjoint(engine_args) + + def test_mammoth_fp8_kv_deploy_projects_only_ar_stage(monkeypatch): monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None) diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index 37e2dc430..6e28bb5b1 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -64,6 +64,13 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] # TODO: extend this set and the expected census if @router.head / @router.options appear. _HTTP_METHODS = {"GET", "POST", "DELETE", "PUT", "PATCH"} + +def test_profiler_endpoints_detect_typed_stage_config(): + stage = SimpleNamespace(profiler_config=SimpleNamespace(profiler="torch")) + + assert api_server._should_enable_profiler_endpoints([stage]) + + # Exact Omni router census on main. Update intentionally when routes change. _EXPECTED_ROUTER_ROUTES = { ("POST", "/v1/chat/completions"), diff --git a/tests/entrypoints/openai_api/test_audex_serving_guards.py b/tests/entrypoints/openai_api/test_audex_serving_guards.py index 87524825e..141dd8cac 100644 --- a/tests/entrypoints/openai_api/test_audex_serving_guards.py +++ b/tests/entrypoints/openai_api/test_audex_serving_guards.py @@ -15,6 +15,7 @@ from types import SimpleNamespace import pytest import torch +from vllm_omni.config.stage_config import StagePipelineConfig from vllm_omni.entrypoints.openai.serving_speech import ( _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES, OmniOpenAIServingSpeech, @@ -251,9 +252,28 @@ class TestAudexStageDetection: stage = serving._find_tts_stage() assert stage is not None and stage.engine_args.model_stage == "audex_thinker" + def test_typed_pipeline_detected_without_legacy_engine_args(self): + serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) + stages = [ + SimpleNamespace( + stage_pipeline_config=StagePipelineConfig(stage_id=index, model_stage=model_stage), + model_config=SimpleNamespace(model_arch=None), + worker_type="generation", + ) + for index, model_stage in enumerate(("audex_thinker", "audex_code2wav")) + ] + serving.engine_client = SimpleNamespace(stage_configs=stages) + + stage = serving._find_tts_stage() + serving._tts_stage = stage + + assert stage is stages[0] + assert serving._detect_tts_model_type() == "audex" + def test_guided_injection_sets_guided_temperature(): """Guided requests decode at the measured CFG temperature (0.05).""" + serving = TestAudexCfgInjection()._serving_with_fake_tokenizer() cond = TestAudexCfgInjection()._cond_prompt() params = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1) @@ -265,6 +285,31 @@ def test_guided_injection_sets_guided_temperature(): assert out.temperature == 0.1 +def test_typed_full_pipeline_uses_full_checkpoint_tokenizer(mocker): + from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext + + serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) + serving._tts_stage = SimpleNamespace( + stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="audex_omni"), + model_config=SimpleNamespace(model_arch=None), + worker_type="ar", + ) + serving.engine_client = SimpleNamespace( + model_config=SimpleNamespace(model="/models/audex"), + ) + ensure_snapshot = mocker.patch( + "vllm_omni.model_executor.models.audex.checkpoint.ensure_audex_snapshot", + return_value="/snapshots/audex", + ) + from_pretrained = mocker.patch("transformers.AutoTokenizer.from_pretrained", return_value=object()) + + adapter = AudexAdapter(SpeechServingContext(server=serving, engine_client=serving.engine_client)) + adapter._get_tokenizer() + + ensure_snapshot.assert_called_once_with("/models/audex", profile="full") + from_pretrained.assert_called_once_with("/snapshots/audex/checkpoint_folder_full") + + # ---------------------------------------------------------------- TTA serving surface from vllm_omni.entrypoints.openai.tts_adapters import resolve_adapter # noqa: E402 diff --git a/tests/entrypoints/openai_api/test_openpi_serving.py b/tests/entrypoints/openai_api/test_openpi_serving.py index 46ff0f822..d274eadda 100644 --- a/tests/entrypoints/openai_api/test_openpi_serving.py +++ b/tests/entrypoints/openai_api/test_openpi_serving.py @@ -138,6 +138,36 @@ def test_policy_server_config_reads_stage_config_model_config(): assert serving.policy_server_config.to_dict() == policy_config +@pytest.mark.parametrize("inline", [False, True], ids=["out_of_process", "inline"]) +def test_policy_server_config_reads_typed_runtime(inline): + from tests.helpers.stage_config import get_deploy_config_path + from vllm_omni.config.omni_config import VllmOmniConfig + from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE + from vllm_omni.entrypoints.async_omni import AsyncOmni + + typed_stage = VllmOmniConfig.from_pipeline_config( + PI0_PIPELINE, deploy_config_path=get_deploy_config_path("pi0.yaml") + ).stage_configs[0] + policy_config = typed_stage.diffusion_config.model_config["policy_server_config"] + stage_client = SimpleNamespace(stage_type="diffusion") + if inline: + # Worker-enriched inline metadata remains authoritative over the + # startup stage config; only the process client needs the fallback. + policy_config = {**policy_config, "worker_metadata": "loaded"} + stage_client.od_config = SimpleNamespace(model_config={"policy_server_config": policy_config}) + engine_client = AsyncOmni.__new__(AsyncOmni) + engine_client.engine = SimpleNamespace( + stage_clients=[stage_client], + stage_configs=[typed_stage], + stage_vllm_configs=[None], + get_diffusion_od_config=lambda: SimpleNamespace(model_class_name="Pi0Pipeline"), + ) + + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client) + + assert serving.policy_server_config.to_dict() == policy_config + + def test_policy_server_config_reads_omegaconf_stage_config(): engine_client = SimpleNamespace( get_diffusion_od_config=lambda: None, diff --git a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py index 59d070a39..417f457d3 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py +++ b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py @@ -27,7 +27,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @dataclass class _EngineClientStub: - stage_configs: list[dict[str, object]] + stage_configs: list[object] @dataclass @@ -149,6 +149,37 @@ async def test_chat_reference_audio_reaches_code2wav_transport( assert code2wav_payload.meta.ref_audio_sr == 16000 +@pytest.mark.asyncio +async def test_typed_stage_enables_reference_audio( + serving_chat: OmniOpenAIServingChat, + mocker: MockerFixture, +) -> None: + reference_audio_url = "data:audio/wav;base64,AAAA" + reference_waveform = np.array([0.25, -0.5], dtype=np.float32) + serving_chat.engine_client = _EngineClientStub( + [ + SimpleNamespace( + model_config=SimpleNamespace(model_arch="MiniCPMO45OmniForConditionalGeneration"), + model_stage="llm", + ) + ] + ) + serving_chat.model_config = _ModelConfigStub() + engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]} + media_connector = mocker.patch( + "vllm_omni.entrypoints.openai.serving_chat.MediaConnector", + ).return_value + media_connector.fetch_audio_async = mocker.AsyncMock(return_value=(reference_waveform, 16000)) + + await serving_chat._attach_minicpmo45_reference_audio( + engine_prompt, + _chat_request(reference_audio_url), + ) + + media_connector.fetch_audio_async.assert_awaited_once_with(reference_audio_url) + assert MINICPMO45_REFERENCE_AUDIO_KEY in cast(dict[str, object], engine_prompt) + + @pytest.mark.parametrize("request_shape", ["extra_body", "model_extra"]) @pytest.mark.asyncio async def test_reference_audio_accepts_request_extra_mappings( diff --git a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py index 1a3a4f22e..5ce129d15 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py +++ b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_stream_output_kinds.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for MiniCPM-o 4.5 audio-streaming output-kind coercion. Chat ``stream=true`` coerces every AR stage to DELTA. For MiniCPM-o 4.5 that @@ -33,6 +34,13 @@ def _stage(model_arch: str, model_stage: str) -> SimpleNamespace: return SimpleNamespace(engine_args=SimpleNamespace(model_arch=model_arch, model_stage=model_stage)) +def _typed_stage(model_arch: str, model_stage: str) -> SimpleNamespace: + return SimpleNamespace( + model_config=SimpleNamespace(model_arch=model_arch), + model_stage=model_stage, + ) + + def _minicpmo45_stages(): return [_stage(MINICPMO45_ARCH, "llm"), _stage(MINICPMO45_ARCH, "tts")] @@ -50,6 +58,20 @@ def test_thinker_forced_final_only_talker_stays_delta(serving_chat): assert out[1].output_kind is RequestOutputKind.DELTA +def test_typed_stages_force_thinker_final_only(serving_chat): + serving_chat.engine_client = SimpleNamespace( + stage_configs=[ + _typed_stage(MINICPMO45_ARCH, "llm"), + _typed_stage(MINICPMO45_ARCH, "tts"), + ] + ) + + out = serving_chat._fix_minicpmo45_audio_stream_output_kinds(_delta_params(), ["text", "audio"]) + + assert out[0].output_kind is RequestOutputKind.FINAL_ONLY + assert out[1].output_kind is RequestOutputKind.DELTA + + @pytest.mark.parametrize("modalities", [None, [], ["text"]]) def test_non_audio_request_is_untouched(serving_chat, modalities): serving_chat.engine_client = SimpleNamespace(stage_configs=_minicpmo45_stages()) diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index f8dfaafb6..5ac895d86 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -27,6 +27,7 @@ from pydantic import ValidationError from pytest_mock import MockerFixture from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse +from vllm_omni.config.stage_config import StagePipelineConfig from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler from vllm_omni.diffusion.sched.step_scheduler import StepScheduler @@ -1039,6 +1040,38 @@ class TestTTSMethods: asyncio.run(server._prepare_speech_generation(request)) server.shutdown() + def test_prepare_speech_rejects_typed_non_tts_omni_model(self, mocker: MockerFixture): + """Typed Qwen3-Omni stages must hit the same talker guard as legacy stages.""" + mock_engine_client = mocker.MagicMock() + mock_engine_client.errored = False + mock_engine_client.tts_max_instructions_length = None + mock_engine_client.stage_configs = [ + SimpleNamespace( + stage_pipeline_config=StagePipelineConfig(stage_id=index, model_stage=model_stage), + model_config=SimpleNamespace(model_arch=None), + worker_type=worker_type, + ) + for index, (model_stage, worker_type) in enumerate( + ( + ("thinker", "ar"), + ("talker", "ar"), + ("code2wav", "generation"), + ) + ) + ] + + mock_models = mocker.MagicMock() + mock_models.is_base_model.return_value = True + server = OmniOpenAIServingSpeech( + engine_client=mock_engine_client, + models=mock_models, + request_logger=mocker.MagicMock(), + ) + request = OpenAICreateSpeechRequest(input="Hello world") + with pytest.raises(ValueError, match="only supported for dedicated TTS models"): + asyncio.run(server._prepare_speech_generation(request)) + server.shutdown() + def test_estimate_prompt_len_fallback(self, speech_server): """Test prompt length estimation falls back to 2048 when model is unavailable.""" tts_params = {"text": ["Hello"], "task_type": ["CustomVoice"]} diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index e0a165a55..500feab35 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -22,6 +22,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( detect_tts_model_type, resolve_adapter, ) +from vllm_omni.entrypoints.openai.tts_adapters.base import resolve_stage_model_path from vllm_omni.entrypoints.openai.tts_adapters.covo_audio import CovoAudioAdapter from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v2 import HiggsAudioV2Adapter from vllm_omni.entrypoints.openai.tts_adapters.indextts2 import ( @@ -99,6 +100,31 @@ def test_resolve_unknown_returns_none(): assert resolve_adapter(None) is None +def test_stage_model_path_prefers_typed_override(): + engine_client = SimpleNamespace( + stage_configs=[ + SimpleNamespace(engine_args=SimpleNamespace(model="legacy-stage-model")), + SimpleNamespace( + model_config=SimpleNamespace(model="typed-stage-model"), + ), + ], + model="served-model", + ) + + assert resolve_stage_model_path(engine_client) == "typed-stage-model" + + +def test_stage_model_path_falls_back_to_legacy_then_served_model(): + legacy_client = SimpleNamespace( + stage_configs=[SimpleNamespace(engine_args=SimpleNamespace(model="legacy-stage-model"))], + model="served-model", + ) + served_client = SimpleNamespace(stage_configs=[SimpleNamespace()], model="served-model") + + assert resolve_stage_model_path(legacy_client) == "legacy-stage-model" + assert resolve_stage_model_path(served_client) == "served-model" + + def test_voxcpm2_resolves(): """VoxCPM2 (the served ``latent_generator`` model) resolves cleanly. diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 2de5c4056..0b44cabdf 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -1134,6 +1134,32 @@ def test_mixed_reference_capability_uses_model_metadata_when_config_defaults_fal assert handler.supports_mixed_reference_inputs +def test_typed_stage_drives_video_capability_checks(): + from vllm_omni.config.config_factory import StageConfigFactory + + minimax_stage = StageConfigFactory.create_typed_default_diffusion( + "minimax-h3", + {"model_class_name": "MiniMaxH3Pipeline"}, + ).stage_configs[0] + minimax_handler = OmniOpenAIServingVideo.for_diffusion( + SimpleNamespace(od_config=SimpleNamespace(model_class_name=None)), + model_name="minimax-h3", + stage_configs=[minimax_stage], + ) + assert minimax_handler.supports_mixed_reference_inputs + + cosmos_stage = StageConfigFactory.create_typed_default_diffusion( + "cosmos3", + {"model_class_name": "Cosmos3OmniDiffusersPipeline"}, + ).stage_configs[0] + cosmos_handler = OmniOpenAIServingVideo.for_diffusion( + SimpleNamespace(od_config=SimpleNamespace(model_class_name=None)), + model_name="cosmos3", + stage_configs=[cosmos_stage], + ) + assert cosmos_handler.supported_control_upload_types == frozenset({"edge", "blur", "depth", "seg", "wsm"}) + + @pytest.mark.parametrize("model_class_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"]) def test_control_upload_capability_is_declared_only_by_cosmos3(test_client, model_class_name): handler = test_client.app.state.openai_serving_video @@ -1191,6 +1217,35 @@ def test_cosmos3_reference_video_limit_uses_v2v_condition_frames(): assert spec.keep == "first" +@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"]) +@pytest.mark.parametrize( + ("num_frames", "extra_params", "expected"), + [ + (189, {"condition_frame_indexes_vision": [0, 2], "condition_video_keep": "last"}, (9, "last")), + (189, {"condition_frame_indexes_vision": [0, 2]}, (9, "first")), + (5, {"condition_frame_indexes_vision": [0, 20]}, (5, "first")), + (None, {"action_mode": "inverse_dynamics", "action_chunk_size": 16}, (17, "first")), + ], +) +def test_cosmos3_reference_video_decode_policy_with_runtime_configs(typed, num_frames, extra_params, expected): + from vllm_omni.config.config_factory import StageConfigFactory + + if typed: + stages = list( + StageConfigFactory.create_typed_default_diffusion( + "cosmos3", + {"model_class_name": "Cosmos3OmniDiffusersPipeline"}, + ).stage_configs + ) + else: + stages = _cosmos3_stage_configs() + request = VideoGenerationRequest(prompt="Continue this motion.", num_frames=num_frames, extra_params=extra_params) + + spec = _reference_video_decode_spec(request, stages) + + assert (spec.max_frames, spec.keep) == expected + + def test_cosmos3_reference_video_limit_preserves_action_frames(): request = VideoGenerationRequest( prompt="Predict the action.", diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 3b378930a..9b99af712 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -2,13 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace +from typing import Any import pytest import torch from pydantic import ValidationError from vllm_omni.config.config_factory import StageConfigFactory -from vllm_omni.config.resolver import OmniConfigResolution +from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig +from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config 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 @@ -117,7 +119,7 @@ def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocke trust_remote_code=False, ) - assert stage_configs[0]["engine_args"]["extras"]["ltx2_use_conv_vae"] is True + assert stage_configs[0].diffusion_config.extras["ltx2_use_conv_vae"] is True def test_default_stage_rejects_unknown_nested_parallel_config_key(): @@ -239,6 +241,33 @@ def test_default_stage_config_includes_default_sampling_params(): } +@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"]) +@pytest.mark.parametrize("sampling_defaults", [{"0": {"guidance_scale": 7.5}}, '{"0":{"guidance_scale":7.5}}']) +def test_generic_diffusion_sampling_defaults_remain_overridable(typed, sampling_defaults): + from vllm_omni.config.yaml_util import create_config + from vllm_omni.entrypoints.omni_base import OmniBase + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + kwargs = {"model_class_name": "QwenImagePipeline", "default_sampling_params": sampling_defaults} + if typed: + stage = StageConfigFactory.create_typed_default_diffusion("generic-diffusion", kwargs).stage_configs[0] + metadata = stage_init_utils.extract_stage_metadata_from_omni_stage_config(stage) + else: + stage = create_config(StageConfigFactory.create_default_diffusion(kwargs))[0] + metadata = stage_init_utils.extract_legacy_stage_metadata(stage) + base = OmniBase.__new__(OmniBase) + base.engine = SimpleNamespace(num_stages=1, stage_configs=[stage]) + base.default_sampling_params_list = [metadata.default_sampling_params] + base.sampling_constraints_list = base._get_sampling_constraints_list([stage]) + + assert base.resolve_sampling_params_list(None)[0].guidance_scale == 7.5 + requested = OmniDiffusionSamplingParams(guidance_scale=2.0) + assert base.resolve_sampling_params_list(requested)[0].guidance_scale == 2.0 + assert requested.guidance_scale == 2.0 + assert base.default_sampling_params_list[0].guidance_scale == 7.5 + assert base.sampling_constraints_list == [{}] + + def test_default_stage_config_includes_diffusion_attention_backend(): """Ensure diffusion attention shorthand lands in engine_args.diffusion_attention_config.""" stage_cfg = StageConfigFactory.create_default_diffusion( @@ -726,3 +755,106 @@ def test_default_stage_config_includes_quantization_config(): stage_cfg = StageConfigFactory.create_default_diffusion({"quantization_config": quantization_config})[0] assert stage_cfg["engine_args"]["quantization_config"] == quantization_config + + +@pytest.mark.parametrize("model_class_name", ["HeliosPipeline", "HunyuanVideo15Pipeline"]) +def test_generic_diffusion_uses_canonical_video_output_type(model_class_name): + config = StageConfigFactory.create_typed_default_diffusion( + "generic-video", + {"model_class_name": model_class_name}, + ) + + assert config.stage_configs[0].final_output_type == "video" + + +def test_generic_diffusion_resolves_structured_stage_without_legacy_conversion(mocker): + """Generic diffusion reaches runtime as the structured stage itself.""" + mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None) + mocker.patch( + "vllm_omni.config.resolver._resolve_generic_diffusion_model_class", + return_value=(True, "FakeDiffusionPipeline"), + ) + + resolved = resolve_omni_config( + "generic-diffusion", + trust_remote_code=False, + deploy_config_path=None, + cli_overrides={ + "num_gpus": 4, + "tensor_parallel_size": 2, + "default_sampling_params": '{"0": {"guidance_scale": 7.5}}', + }, + stage_overrides=None, + strategy_config_path=None, + ) + + assert resolved.pipeline_config is not None + assert resolved.pipeline_config.model_type == "generic_diffusion" + assert len(resolved.stage_configs) == 1 + stage = resolved.stage_configs[0] + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.model_config.model == "generic-diffusion" + assert stage.model_config.default_sampling_params == {"guidance_scale": 7.5} + assert stage.parallel_config.tensor_parallel_size == 2 + assert stage.parallel_config.data_parallel_size == 2 + assert stage.parallel_config.world_size == 4 + assert stage.runtime_config.devices == "0,1,2,3" + + +def test_generic_diffusion_structured_stage_reaches_standard_startup(mocker): + """Standard runtime resolves and starts the typed stage through the real launcher.""" + from vllm_omni.engine import stage_engine_startup as startup_module + from vllm_omni.engine import stage_runtime as runtime_module + from vllm_omni.engine.stage_runtime import StageRuntime + + mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None) + mocker.patch( + "vllm_omni.config.resolver._resolve_generic_diffusion_model_class", + return_value=(True, "FakeDiffusionPipeline"), + ) + resolved = resolve_omni_config( + "generic-diffusion", + trust_remote_code=False, + deploy_config_path=None, + cli_overrides={"num_gpus": 1}, + stage_overrides=None, + strategy_config_path=None, + ) + stage = resolved.stage_configs[0] + launched: dict[str, Any] = {} + client = SimpleNamespace(input_address=None, shutdown=mocker.Mock()) + + mocker.patch.object(runtime_module, "prepare_engine_environment") + mocker.patch.object(runtime_module, "load_omni_transfer_config_for_model", return_value=None) + mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={}) + mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None)) + + def _initialize_typed_stage(stage_id, model, stage_config, metadata, **kwargs): + launched.update( + stage_id=stage_id, + model=model, + stage_config=stage_config, + metadata=metadata, + **kwargs, + ) + return client + + mocker.patch.object(startup_module, "initialize_diffusion_stage", side_effect=_initialize_typed_stage) + + runtime = StageRuntime( + stage_configs=list(resolved.stage_configs), + model="generic-diffusion", + config_path="", + stage_init_timeout=10, + async_chunk=False, + ) + runtime.initialize() + + assert launched["stage_config"] is stage + assert isinstance(launched["stage_config"], VllmOmniDiffusionStageConfig) + assert launched["metadata"].stage_type == "diffusion" + assert launched["metadata"].model_stage == "diffusion" + assert launched["use_inline"] is True + assert launched["model"] == "generic-diffusion" + assert launched["stage_id"] == 0 + assert runtime.stage_pools[0].clients == [client] diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index ea57d07ab..969bbaf91 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -17,6 +17,8 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError +from vllm_omni.config.omni_config import VllmOmniConfig +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.stage_config import StageConfig from vllm_omni.engine.async_omni_engine import StageRuntimeInfo from vllm_omni.engine.messages import ErrorMessage, OutputMessage @@ -202,6 +204,49 @@ def _make_base(): return obj +@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"]) +def test_resolve_sampling_params_list_preserves_stage_constraints(typed): + from vllm_omni.config.omni_config import VllmOmniARStageConfig + from vllm_omni.config.stage_config import StagePipelineConfig + + base = _make_base() + base.engine.num_stages = 1 + base.default_sampling_params_list = [SamplingParams(max_tokens=1000, detokenize=False, stop_token_ids=[42])] + if typed: + stage = VllmOmniARStageConfig( + stage_pipeline_config=StagePipelineConfig( + stage_id=0, + model_stage="dummy-model", + sampling_constraints={"detokenize": False, "stop_token_ids": [42]}, + ) + ) + else: + stage = StageConfig( + stage_id=0, + model_stage="dummy-model", + sampling_constraints={"detokenize": False, "stop_token_ids": [42]}, + ).to_omegaconf() + base.engine.stage_configs = [stage] + base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs) + assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": [42]}] + + resolved_defaults = base.resolve_sampling_params_list(None) + assert resolved_defaults[0].stop_token_ids == [42] + + caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[7]) + resolved = base.resolve_sampling_params_list(caller_params) + + assert resolved[0] is not caller_params + assert resolved[0].seed == 1234 + assert resolved[0].max_tokens == 7 + assert resolved[0].detokenize is False + assert resolved[0].stop_token_ids == [7, 42] + assert 7 in resolved[0]._all_stop_token_ids + assert 42 in resolved[0]._all_stop_token_ids + assert caller_params.detokenize is True + assert caller_params.stop_token_ids == [7] + + def test_resolve_sampling_params_list_merges_required_stop_tokens(): base = _make_base() base.engine.num_stages = 1 @@ -219,22 +264,13 @@ def test_resolve_sampling_params_list_merges_required_stop_tokens(): base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs) assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": required_stop_ids}] - resolved_defaults = base.resolve_sampling_params_list(None) - assert resolved_defaults[0].stop_token_ids == [99, 151704, 151645] +def test_sampling_constraints_are_forwarded_by_typed_stage_configs(): + config = VllmOmniConfig.from_pipeline_config(OMNI_PIPELINES["qwen3_tts"]) - caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[100, 151704]) + constraints = OmniBase._get_sampling_constraints_list(config.stage_configs) - resolved = base.resolve_sampling_params_list(caller_params) - - assert resolved[0] is not caller_params - assert resolved[0].seed == 1234 - assert resolved[0].max_tokens == 7 - assert resolved[0].detokenize is False - assert resolved[0].stop_token_ids == [100, 151704, 151645] - assert {100, 151704, 151645}.issubset(resolved[0]._all_stop_token_ids) - assert caller_params.detokenize is True - assert caller_params.stop_token_ids == [100, 151704] + assert constraints == [dict(stage.stage_pipeline_config.sampling_constraints) for stage in config.stage_configs] @pytest.mark.parametrize("use_defaults", [False, True]) diff --git a/tests/entrypoints/test_resolve_dreamzero_config.py b/tests/entrypoints/test_resolve_dreamzero_config.py index 450e1be39..3cfa0f0e5 100644 --- a/tests/entrypoints/test_resolve_dreamzero_config.py +++ b/tests/entrypoints/test_resolve_dreamzero_config.py @@ -32,15 +32,18 @@ def test_dreamzero_resolves_through_registry_with_model_defaults(monkeypatch): stage_overrides=None, strategy_config_path=None, ) - engine_args = resolved.stage_configs[0].engine_args + stage = resolved.stage_configs[0] 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" + assert stage.diffusion_config.model_class_name == "DreamZeroPipeline" + assert stage.diffusion_config.model_config["policy_server_config"]["action_space"] == "joint_position" def test_dreamzero_enrich_config_preserves_explicit_model_class_name(monkeypatch): + def _get_hf_config(path, _model): + return None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]} + monkeypatch.setattr( "vllm.transformers_utils.config.get_hf_file_to_dict", lambda path, _model, **_kwargs: ( diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 1f3b58170..8d3766df0 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -7,10 +7,12 @@ from __future__ import annotations import argparse from types import SimpleNamespace +from typing import Any import pytest from pytest_mock import MockerFixture +from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig from vllm_omni.config.resolver import OmniConfigResolution from vllm_omni.entrypoints.cli.serve import ( OmniServeCommand, @@ -334,6 +336,26 @@ def test_parse_stage_overrides_invalid_json_raises() -> None: assert f"Got: {bad!r}" in message +@pytest.mark.parametrize( + "payload", + ['{"abc": {}}', '{"-1": {}}', '{"1.5": {}}', '{"\uff10": {}}'], +) +def test_parse_stage_overrides_rejects_invalid_stage_ids(payload: str) -> None: + with pytest.raises(argparse.ArgumentTypeError, match="non-negative integer stage ids"): + _parse_stage_overrides(payload) + + +def test_parse_stage_overrides_preserves_arbitrary_override_fields() -> None: + parsed = _parse_stage_overrides( + '{"0": {"extras": {"ltx2_use_conv_vae": true}, "kv_cache_dtype": "fp8", "typo_field_xyz": 1}}' + ) + assert parsed["0"] == { + "extras": {"ltx2_use_conv_vae": True}, + "kv_cache_dtype": "fp8", + "typo_field_xyz": 1, + } + + def test_parse_stage_overrides_rejects_non_dict_top_level() -> None: """Top-level must be a JSON object (dict). A list, scalar, or non-dict mapping is rejected with a ValueError naming ``--stage-overrides`` and @@ -706,6 +728,77 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) assert manager_kwargs["omni_replica_id"] == 0 +def test_run_headless_generic_diffusion_launches_structured_stage(mocker: MockerFixture) -> None: + """Headless resolution starts the typed stage through the real group launcher.""" + from vllm_omni.engine import stage_engine_startup as startup_module + + mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None) + mocker.patch( + "vllm_omni.config.resolver._resolve_generic_diffusion_model_class", + return_value=(True, "FakeDiffusionPipeline"), + ) + mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") + mocker.patch.object(startup_module.stage_init_utils, "load_omni_transfer_config_for_model", return_value=None) + mocker.patch.object( + startup_module.initialization, + "resolve_omni_kv_config_for_stage", + return_value=(None, None, None), + ) + captured: dict[str, Any] = {} + od_config = SimpleNamespace() + + def _build_diffusion_config(model, stage_config, metadata): + captured.update(model=model, stage_config=stage_config, metadata=metadata) + return od_config + + def _launch_replica_group(**kwargs): + captured.update(group_kwargs=kwargs) + captured["manager"] = kwargs["launch_one"](0) + + def _launch_replica(**kwargs): + captured.update(replica_kwargs=kwargs) + return SimpleNamespace(exitcode=None) + + mocker.patch.object(startup_module.stage_init_utils, "build_diffusion_config", side_effect=_build_diffusion_config) + mocker.patch.object(startup_module, "launch_headless_replica_group", side_effect=_launch_replica_group) + mocker.patch.object(startup_module, "launch_headless_diffusion_replica", side_effect=_launch_replica) + + explicit_keys = frozenset( + { + "model", + "stage_id", + "omni_master_address", + "omni_master_port", + "worker_backend", + "model_class_name", + "num_gpus", + } + ) + args = _make_headless_args( + explicit_keys=explicit_keys, + model="generic-diffusion", + model_class_name="FakeDiffusionPipeline", + num_gpus=1, + ) + + run_headless(args) + + stage = captured["stage_config"] + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert captured["metadata"].stage_type == "diffusion" + assert captured["metadata"].model_stage == "diffusion" + assert captured["model"] == "generic-diffusion" + assert captured["group_kwargs"]["stage_id"] == 0 + assert captured["group_kwargs"]["omni_dp_size_local"] == 1 + assert captured["group_kwargs"]["per_replica_devices"] == ["0"] + assert captured["replica_kwargs"]["stage_config"] is stage + assert captured["replica_kwargs"]["stage_id"] == 0 + assert captured["replica_kwargs"]["omni_master_address"] == "127.0.0.1" + assert captured["replica_kwargs"]["omni_master_port"] == 26000 + assert captured["replica_kwargs"]["od_config"] is od_config + assert stage.runtime_config.devices == "0" + + def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixture) -> None: """A diffusion replica that exits with a non-zero code must surface as a RuntimeError from ``run_headless`` (the head needs the signal to roll diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 680bb110c..e7f94ffef 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -13,6 +13,14 @@ import torch from pytest_mock import MockerFixture from vllm.sampling_params import RequestOutputKind, SamplingParams +from vllm_omni.config.composable_parallel import ( + Broadcast, + FanInByStage, + MeshAxisSpec, + RouteByStage, + StrategySpec, + TakeRank, +) from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.resolver import ( OmniConfigResolution, @@ -21,7 +29,6 @@ from vllm_omni.config.resolver import ( 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.entrypoints.utils import ( @@ -356,16 +363,12 @@ class TestResolveOmniConfig: ) 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" + assert resolved.stage_configs[0].diffusion_config.dtype is torch.float32 + assert resolved.stage_configs[0].diffusion_config.engine_backend == engine_backend + assert resolved.stage_configs[0].diffusion_config.revision == "pinned-revision" def test_generic_diffusion_uses_registered_model_metadata(self, mocker: MockerFixture): - mocker.patch( - "vllm_omni.config.resolver.StageConfigFactory.create_from_model", - return_value=None, - ) + mocker.patch("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", @@ -374,59 +377,98 @@ class TestResolveOmniConfig: "vllm_omni.config.resolver.DiffusionModelRegistry.get_supported_archs", return_value={"WanImageToVideoPipeline"}, ) + load_model_class = mocker.patch("vllm_omni.config.resolver.DiffusionModelRegistry._try_load_model_cls") 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", - }, + cli_overrides={"diffusion_load_format": "diffusers", "revision": "pinned-revision"}, stage_overrides=None, strategy_config_path=None, ) + resolve_model_class.assert_called_once_with("/models/Wan2.2-I2V", "diffusers", "pinned-revision") + load_model_class.assert_not_called() 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.diffusion_config.model_class_name == "WanImageToVideoPipeline" + assert stage.diffusion_config.revision == "pinned-revision" assert stage.final_output_type == "video" + def test_generic_diffusion_stage_overrides_reach_typed_backend(self, mocker: MockerFixture): + from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config + + mocker.patch("vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=None) + mocker.patch( + "vllm_omni.config.resolver._resolve_generic_diffusion_model_class", + return_value=(True, "LTX2Pipeline"), + ) + resolved = resolve_omni_config( + "/models/LTX-2.5-Diffusers", + trust_remote_code=False, + deploy_config_path=None, + cli_overrides={ + "extras": {"ltx2_use_conv_vae": False, "keep": "global"}, + "tensor_parallel_size": 1, + "diffusion_streaming_output": True, + }, + stage_overrides={ + "0": { + "extras": {"ltx2_use_conv_vae": True}, + "tensor_parallel_size": 2, + "devices": "2,3", + "engine_backend": "custom.backend", + }, + }, + strategy_config_path=None, + ) + stage = resolved.stage_by_id(0) + engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="/models/LTX-2.5-Diffusers") + + assert stage.runtime_config.devices == "2,3" + assert stage.parallel_config.world_size == 2 + assert engine_args["engine_backend"] == "custom.backend" + assert engine_args["extras"]["ltx2_use_conv_vae"] is True + assert engine_args["extras"]["keep"] == "global" + assert engine_args["streaming_output"] is True + def test_registered_pipeline_uses_structured_metadata_and_preserves_override_trust(self, mocker: MockerFixture): endpoint_restriction = SimpleNamespace(name="chat") + typed_stage = SimpleNamespace(stage_id=1) structured_config = SimpleNamespace( - orchestrator_config=SimpleNamespace(deploy_config_path="/resolved/deploy.yaml"), + orchestrator_config=SimpleNamespace( + deploy_config_path="/resolved/deploy.yaml", + omni_lb_policy="round_robin", + ), pipeline_config=SimpleNamespace(endpoint_restrictions=(endpoint_restriction,)), + stage_configs=(typed_stage,), + strategy_omni_lb_policy="round-robin", ) - 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", - ), + "vllm_omni.config.resolver.StageConfigFactory.create_legacy_stage_configs_from_model", ) - strategy_specs = {"stage_1": {"dp": 3}} + strategy_specs = { + "stage_1": [ + StrategySpec( + "stage_replica", + MeshAxisSpec("stage_replica", 2), + RouteByStage("round_robin"), + FanInByStage(), + ) + ] + } load_strategy = mocker.patch( "vllm_omni.config.resolver._load_strategy_specs", return_value=strategy_specs, ) + apply_strategy = mocker.patch( + "vllm_omni.config.composable_parallel.apply_strategy_specs", + return_value=SimpleNamespace(omni_lb_policy="round_robin"), + ) resolved = resolve_omni_config( "dummy-model", @@ -447,31 +489,63 @@ class TestResolveOmniConfig: 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, ) + create_legacy.assert_not_called() load_strategy.assert_called_once_with("strategy.yaml") + apply_strategy.assert_not_called() 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.omni_lb_policy == "round-robin" assert resolved.endpoint_restrictions == (endpoint_restriction,) - assert resolved.stage_configs == (runtime_stage,) + assert resolved.stage_configs == (typed_stage,) - def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture): - pipeline = OMNI_PIPELINES["qwen3_tts"] - assert isinstance(pipeline, PipelineConfig) + def test_tp_only_strategy_does_not_report_default_lb_policy_as_derived(self, mocker: MockerFixture): structured_config = SimpleNamespace( - orchestrator_config=SimpleNamespace(deploy_config_path=None), - pipeline_config=pipeline, + orchestrator_config=SimpleNamespace( + deploy_config_path="/resolved/deploy.yaml", + omni_lb_policy="round-robin", + ), + pipeline_config=SimpleNamespace(endpoint_restrictions=()), + stage_configs=(SimpleNamespace(stage_id=0),), + strategy_omni_lb_policy=None, ) mocker.patch( "vllm_omni.config.resolver.StageConfigFactory.create_from_model", return_value=structured_config, ) + mocker.patch( + "vllm_omni.config.resolver._load_strategy_specs", + return_value={ + "thinker": [ + StrategySpec( + "tp", + MeshAxisSpec("tp", 2), + Broadcast(), + TakeRank(), + ) + ] + }, + ) + + resolved = resolve_omni_config( + "dummy-model", + trust_remote_code=True, + deploy_config_path="deploy.yaml", + cli_overrides={"omni_lb_policy": "round-robin"}, + stage_overrides=None, + strategy_config_path="strategy.yaml", + ) + + assert resolved.omni_lb_policy is None + + def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture): + pipeline = OMNI_PIPELINES["qwen3_tts"] + assert isinstance(pipeline, PipelineConfig) + mocker.patch( + "vllm_omni.config.config_factory.StageConfigFactory.get_pipeline_config", + return_value=pipeline, + ) resolved = resolve_omni_config( "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", @@ -483,9 +557,9 @@ class TestResolveOmniConfig: ) 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 [stage.stage_id for stage in resolved.pipeline_config.stages] == [ + stage.stage_id for stage in resolved.stage_configs + ] assert resolved.pipeline_config.stages[-1].model_stage == "forced_aligner" assert len(resolved.pipeline_config.stages) == len(pipeline.stages) + 1 diff --git a/vllm_omni/config/composable_parallel/apply.py b/vllm_omni/config/composable_parallel/apply.py index 879c7aec3..acfb9e229 100644 --- a/vllm_omni/config/composable_parallel/apply.py +++ b/vllm_omni/config/composable_parallel/apply.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Apply declarative parallel strategies onto merged stage configs. This is the *override-after-merge* seam. ``merge_pipeline_deploy`` first fuses @@ -198,10 +198,41 @@ def _check_devices(runtime: dict[str, Any], cfg: OmniParallelConfig, *, role: Ro def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> None: + declared = set(cfg.l1_owners.keys()) + if hasattr(stage, "parallel_config") and hasattr(stage, "runtime_config"): + parallel = stage.parallel_config + for kind, field_name in _ENGINE_FIELD_BY_KIND.items(): + if kind in declared: + existing = getattr(parallel, field_name, None) + value = getattr(cfg, field_name) + if existing not in (None, 1, value): + _fail( + f"role {role!r}: strategy derives {field_name}={value} " + f"but structured config already set {existing}" + ) + setattr(parallel, field_name, value) + if "ep" in declared and cfg.enable_expert_parallel: + setattr(parallel, "enable_expert_parallel", True) + if "stage_replica" in declared: + existing = stage.runtime_config.num_replicas + if existing not in (None, 1, cfg.stage_replica_size): + _fail( + f"role {role!r}: strategy derives " + f"num_replicas={cfg.stage_replica_size} but structured config already set {existing}" + ) + stage.runtime_config.num_replicas = cfg.stage_replica_size + check_device_layout( + stage.runtime_config.devices, + tensor_parallel_size=cfg.tensor_parallel_size, + data_parallel_size=cfg.data_parallel_size, + pipeline_parallel_size=cfg.pipeline_parallel_size, + num_replicas=int(stage.runtime_config.num_replicas or 1), + role=role, + ) + return + engine_args = stage.yaml_engine_args runtime = stage.yaml_runtime - - declared = set(cfg.l1_owners.keys()) for kind, field_name in _ENGINE_FIELD_BY_KIND.items(): if kind in declared: _set_explicit(engine_args, field_name, getattr(cfg, field_name), role=role) @@ -209,7 +240,6 @@ def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> No _set_explicit(engine_args, "enable_expert_parallel", True, role=role) if "stage_replica" in declared: _set_num_replicas(runtime, cfg.stage_replica_size, role=role) - _check_devices(runtime, cfg, role=role) diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py index 90108e7df..94decce9a 100644 --- a/vllm_omni/config/config_factory.py +++ b/vllm_omni/config/config_factory.py @@ -25,6 +25,9 @@ from vllm_omni.config.stage_config import ( DeployConfig, PipelineConfig, StageConfig, + StageDeployConfig, + StageExecutionType, + StagePipelineConfig, StageType, build_stage_runtime_overrides, load_deploy_config, @@ -373,6 +376,7 @@ class StageConfigFactory: trust_remote_code: bool | None, cli_overrides: dict[str, Any], deploy_config_path: str | None, + strategy_specs: Mapping[Any, Any] | None = None, ) -> VllmOmniConfig | None: """Build the structured Omni config for a model/deploy pair.""" user_deploy_config = cls._load_user_deploy_config(deploy_config_path) @@ -387,6 +391,20 @@ class StageConfigFactory: if pipeline_cfg is None: return None + # Keep the typed representation's topology in lockstep with the + # compatibility launch representation. The aligner is a real extra + # stage, not an engine-only setting, so it must be injected before + # either representation is built. + if user_deploy_config is None: + if pipeline_cfg.default_deploy_config_name is not None: + default_deploy_path = _DEPLOY_DIR / pipeline_cfg.default_deploy_config_name + user_deploy_config = load_deploy_config(default_deploy_path) + deploy_config_path = str(default_deploy_path) + else: + user_deploy_config = DeployConfig() + from vllm_omni.utils.forced_aligner import inject_forced_aligner_stage + + pipeline_cfg, user_deploy_config = inject_forced_aligner_stage(pipeline_cfg, user_deploy_config, cli_overrides) registry_cli_overrides = with_trust_remote_code_override( {**cli_overrides, "model": model}, trust_remote_code, @@ -396,6 +414,7 @@ class StageConfigFactory: user_deploy_config=user_deploy_config, deploy_config_path=deploy_config_path, cli_overrides=registry_cli_overrides, + strategy_specs=strategy_specs, ) @classmethod @@ -621,13 +640,11 @@ class StageConfigFactory: ) @classmethod - def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]: - """Build the temporary runtime ABI for a generic diffusion stage. - - The terminal diffusion config owns engine defaults and normalization; - this compatibility builder only adds Omni stage topology, request - defaults, and device placement. - """ + def _normalize_default_diffusion( + cls, + kwargs: dict[str, Any], + ) -> tuple[dict[str, Any], DiffusionParallelConfig, dict[str, Any], str]: + """Normalize inputs shared by typed and compatibility diffusion builders.""" raw_sampling_params = kwargs.get("default_sampling_params") if isinstance(raw_sampling_params, str): try: @@ -643,8 +660,6 @@ class StageConfigFactory: 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" extras = dict(engine_args.get("extras") or {}) for key, default in ( @@ -665,6 +680,21 @@ class StageConfigFactory: final_output_type, model_class_name, ) + return engine_args, parallel_config, default_sampling_params, final_output_type + + @classmethod + def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]: + """Build the temporary runtime ABI for a generic diffusion stage. + + The terminal diffusion config owns engine defaults and normalization; + this compatibility builder only adds Omni stage topology, request + defaults, and device placement. + """ + engine_args, parallel_config, default_sampling_params, final_output_type = cls._normalize_default_diffusion( + kwargs + ) + engine_args["parallel_config"] = asdict(parallel_config) + engine_args["model_stage"] = "diffusion" return [ { @@ -681,6 +711,50 @@ class StageConfigFactory: } ] + @classmethod + def create_typed_default_diffusion( + cls, + model: str, + kwargs: dict[str, Any], + ) -> VllmOmniConfig: + """Build generic diffusion directly into the structured runtime config.""" + engine_overrides, parallel_config, default_sampling_params, final_output_type = ( + cls._normalize_default_diffusion(kwargs) + ) + engine_overrides.update(asdict(parallel_config)) + model_class_name = engine_overrides.get("model_class_name") + pipeline = PipelineConfig( + model_type="generic_diffusion", + model_arch=str(model_class_name or ""), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + final_output=True, + final_output_type=final_output_type, + ), + ), + ) + + # These values have already been normalized for this diffusion stage. + # Scope them explicitly so diffusion-only fields such as engine_backend + # and extras are not filtered by the global LLM CLI argument surface. + stage_overrides = { + f"stage_0_{key}": value for key, value in engine_overrides.items() if key not in {"model", "stage_id"} + } + stage_overrides["stage_0_devices"] = ( + kwargs.get("stage_0_devices") + or kwargs.get("devices") + or ",".join(str(i) for i in range(parallel_config.world_size)) + ) + # Caller defaults belong to deployment, not immutable pipeline constraints: + # explicit request sampling parameters must still be able to override them. + deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0, default_sampling_params=default_sampling_params)]) + return VllmOmniConfig.from_pipeline_config( + pipeline, user_deploy_config=deploy, cli_overrides={"model": model, **stage_overrides} + ) + @classmethod def _merge_cli_overrides( cls, diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index c13935489..1407615ab 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -29,7 +29,9 @@ from vllm.config import ProfilerConfig as VllmProfilerConfig from vllm.config import SchedulerConfig as VllmSchedulerConfig from vllm.config.utils import config from vllm.engine.arg_utils import EngineArgs as VllmEngineArgs +from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.pooling_params import PoolingParams from vllm_omni.config.stage_config import ( _DEPLOY_DIR, @@ -42,11 +44,13 @@ from vllm_omni.config.stage_config import ( StagePipelineConfig, StageType, _apply_platform_overrides, + _get_recursively_merged_dict, _resolve_scheduler, _scheduler_path, _select_processor_funcs, build_stage_runtime_overrides, load_deploy_config, + merge_pipeline_deploy, merge_sampling_constraints, normalize_pipeline_cli_overrides, reconcile_diffusion_attention_overrides, @@ -55,6 +59,8 @@ from vllm_omni.config.stage_config import ( ) from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +logger = init_logger(__name__) + _EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str | None]] = { StageExecutionType.LLM_AR: (StageType.LLM, "ar"), StageExecutionType.LLM_GENERATION: (StageType.LLM, "generation"), @@ -186,6 +192,11 @@ class _ModelEngineOverrides(TypedDict, total=False): enable_prompt_embeds: bool +class _PoolingEngineOverrides(TypedDict, total=False): + runner: str + pooling_output_decoder: str + + class _LoadEngineOverrides(TypedDict, total=False): tokenizer: str download_dir: str @@ -261,6 +272,7 @@ class _StageEngineValues: quantization: _QuantizationEngineOverrides model: _ModelEngineOverrides load: _LoadEngineOverrides + pooling: _PoolingEngineOverrides cache: _CacheEngineOverrides scheduler: _SchedulerEngineOverrides connector: _ConnectorEngineOverrides @@ -476,6 +488,15 @@ class OmniStageModelConfig(_TrackExplicitConfigFields): enable_prompt_embeds: bool | None = None +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class OmniStagePoolingConfig: + """Typed inputs owned by vLLM pooling stages.""" + + runner: str | None = None + pooling_output_decoder: str | None = None + default_pooling_params: PoolingParams | None = None + + @_enforce_keyword_only_init @config(kw_only=True) class OmniStageLoadConfig(_TrackExplicitConfigFields, VllmLoadConfig): @@ -833,12 +854,12 @@ class _DiffusionConfigProjection: from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs from vllm_omni.diffusion.offloader.config import parse_diffusion_offload_config - normalized_kwargs = normalize_omni_diffusion_kwargs(kwargs) + config_kwargs = normalize_omni_diffusion_kwargs(kwargs) # Validate before stage construction while retaining the raw mapping # needed by dataclass/config serialization across process boundaries. - parse_diffusion_offload_config(normalized_kwargs.get("diffusion_offload_config")) + parse_diffusion_offload_config(config_kwargs.get("diffusion_offload_config")) valid_fields = {f.name for f in fields(cast(Any, cls))} - return cls(**{k: v for k, v in normalized_kwargs.items() if k in valid_fields}) + return cls(**{k: v for k, v in config_kwargs.items() if k in valid_fields}) def __post_init__(self) -> None: # Keep diffusion imports lazy so importing vllm_omni.config does not @@ -1145,6 +1166,7 @@ _MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__) _LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__) _CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__) _SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__) +_POOLING_ENGINE_FIELDS = frozenset(_PoolingEngineOverrides.__annotations__) _CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__) _RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__) _DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset({"compilation_config", "profiler_config"}) @@ -1175,6 +1197,7 @@ _LLM_STAGE_ENGINE_FIELDS = ( | _LLM_CACHE_ENGINE_FIELDS | _LLM_SCHEDULER_ENGINE_FIELDS | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS + | _POOLING_ENGINE_FIELDS | {"parallel_config"} ) _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS = ( @@ -1308,7 +1331,12 @@ def _stage_engine_values( if topology.execution_type == StageExecutionType.DIFFUSION: # Mirror StageConfig.to_omegaconf so both projections resolve alike. reconcile_diffusion_attention_overrides(engine, stage_cli_overrides) - engine.update(_copy_value(stage_cli_overrides)) + for key, value in stage_cli_overrides.items(): + existing = engine.get(key) + if key != "omni_kv_config" and isinstance(existing, dict) and isinstance(value, Mapping): + engine[key] = _get_recursively_merged_dict(existing, dict(value)) + else: + engine[key] = _copy_value(value) _validate_stage_engine_override_ownership( topology.stage_id, topology.execution_type, @@ -1334,6 +1362,10 @@ def _stage_engine_values( ), model=cast(_ModelEngineOverrides, _select_engine_overrides(engine, _MODEL_ENGINE_FIELDS)), load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, load_engine_fields)), + pooling=cast( + _PoolingEngineOverrides, + _select_engine_overrides(engine, _POOLING_ENGINE_FIELDS), + ), cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, cache_engine_fields)), scheduler=cast( _SchedulerEngineOverrides, @@ -1406,6 +1438,7 @@ class BaseVllmOmniStageConfig: cache_config: OmniStageCacheConfig = field(default_factory=OmniStageCacheConfig) scheduler_config: OmniStageSchedulerConfig = field(default_factory=OmniStageSchedulerConfig) connector_config: OmniStageConnectorConfig = field(default_factory=OmniStageConnectorConfig) + pooling_config: OmniStagePoolingConfig = field(default_factory=OmniStagePoolingConfig) runtime_config: OmniStageRuntimeConfig = field(default_factory=OmniStageRuntimeConfig) parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig) compilation_config: VllmCompilationConfig | None = None @@ -1490,6 +1523,10 @@ class BaseVllmOmniStageConfig: def prompt_transform_func(self) -> str | None: return self.stage_pipeline_config.prompt_transform_func + @property + def sampling_constraints(self) -> dict[str, Any]: + return dict(self.stage_pipeline_config.sampling_constraints) + @property def cfg_kv_collect_func(self) -> str | None: return self.stage_pipeline_config.cfg_kv_collect_func @@ -1543,6 +1580,7 @@ def _build_common_stage_config_kwargs( model=model, ), "load_config": _build_load_config(topology, engine.load), + "pooling_config": _build_pooling_config(stage_deploy, engine.pooling), "cache_config": _build_cache_config( deploy, engine.cache, @@ -1756,6 +1794,20 @@ def _build_model_config( ) +def _build_pooling_config( + stage_deploy: StageDeployConfig | None, + engine: _PoolingEngineOverrides, +) -> OmniStagePoolingConfig: + default_pooling_params = None + if stage_deploy is not None and stage_deploy.default_pooling_params: + default_pooling_params = PoolingParams(**dict(stage_deploy.default_pooling_params)) + return cast(Any, OmniStagePoolingConfig)( + runner=_copy_value(engine.get("runner")), + pooling_output_decoder=_copy_value(engine.get("pooling_output_decoder")), + default_pooling_params=default_pooling_params, + ) + + def _build_load_config( topology: StagePipelineConfig, engine: _LoadEngineOverrides, @@ -1934,6 +1986,10 @@ class VllmOmniConfig: pipeline_config: PipelineConfig stage_configs: tuple[StageConfigType, ...] orchestrator_config: VllmOmniOrchestratorConfig = field(default_factory=VllmOmniOrchestratorConfig) + # Keep strategy provenance separate from the effective orchestrator value: + # an explicit CLI policy may be present even when the strategy declares no + # stage-replica axis. + strategy_omni_lb_policy: str | None = None def stage_by_id(self, stage_id: int) -> StageConfigType: for stage in self.stage_configs: @@ -1949,6 +2005,7 @@ class VllmOmniConfig: user_deploy_config: DeployConfig | None = None, deploy_config_path: str | None = None, cli_overrides: dict[str, Any] | None = None, + strategy_specs: Mapping[Any, Any] | None = None, ) -> VllmOmniConfig: """Create a structured config from a resolved pipeline and deploy YAML.""" if cli_overrides is None: @@ -1973,6 +2030,62 @@ class VllmOmniConfig: deploy.async_chunk = False _validate_async_chunk_support(pipeline_cfg, deploy) validate_stage_async_chunk_edges(pipeline_cfg, deploy) + + strategy_result = None + if strategy_specs: + from vllm_omni.config.composable_parallel import apply_strategy_specs + + strategy_stages = merge_pipeline_deploy(pipeline_cfg, copy.deepcopy(deploy), {}) + strategy_result = apply_strategy_specs(strategy_stages, strategy_specs) + strategy_overrides: dict[str, Any] = {} + axis_fields = { + "tp": "tensor_parallel_size", + "dp": "data_parallel_size", + "pp": "pipeline_parallel_size", + } + for stage_id, parallel in strategy_result.per_stage_config.items(): + declared = set(parallel.l1_owners) + explicit = build_stage_runtime_overrides(stage_id, cli_overrides) + for axis, field_name in axis_fields.items(): + if axis not in declared: + continue + derived = getattr(parallel, field_name) + value = explicit.get(field_name) + if value is None: + value = derived + elif value != derived: + logger.warning( + "[composable_parallel] stage %s: CLI %s=%s overrides the " + "strategy-derived %s=%s. The CLI value wins; remove one to avoid ambiguity.", + stage_id, + field_name, + value, + field_name, + derived, + ) + strategy_overrides[f"stage_{stage_id}_{field_name}"] = value + if "ep" in declared and parallel.enable_expert_parallel: + strategy_overrides[f"stage_{stage_id}_enable_expert_parallel"] = bool( + explicit.get("enable_expert_parallel", True) + ) + if "stage_replica" in declared: + derived = parallel.stage_replica_size + value = explicit.get("num_replicas") + if value is None: + value = derived + elif value != derived: + logger.warning( + "[composable_parallel] stage %s: CLI num_replicas=%s overrides the " + "strategy-derived num_replicas=%s. The CLI value wins; remove one to avoid ambiguity.", + stage_id, + value, + derived, + ) + strategy_overrides[f"stage_{stage_id}_num_replicas"] = value + cli_overrides = {**strategy_overrides, **cli_overrides} + if strategy_result.omni_lb_policy is not None and cli_overrides.get("omni_lb_policy") is None: + cli_overrides["omni_lb_policy"] = strategy_result.omni_lb_policy + deploy_by_id = {stage.stage_id: stage for stage in deploy.stages} model = cli_overrides.get("model") @@ -1995,6 +2108,20 @@ class VllmOmniConfig: ) for topology in pipeline_cfg.stages ) + if strategy_result is not None: + from vllm_omni.config.composable_parallel import check_device_layout + + by_id = {stage.stage_id: stage for stage in stage_configs} + for stage_id in strategy_result.per_stage_config: + stage = by_id[stage_id] + check_device_layout( + stage.runtime_config.devices, + tensor_parallel_size=stage.parallel_config.tensor_parallel_size, + data_parallel_size=stage.parallel_config.data_parallel_size, + pipeline_parallel_size=stage.parallel_config.pipeline_parallel_size, + num_replicas=stage.runtime_config.num_replicas, + role=stage.model_stage, + ) orchestrator_config = cast(Any, VllmOmniOrchestratorConfig)( deploy_config_path=loaded_deploy_config_path, @@ -2004,6 +2131,7 @@ class VllmOmniConfig: pipeline_config=pipeline_cfg, stage_configs=stage_configs, orchestrator_config=orchestrator_config, + strategy_omni_lb_policy=(strategy_result.omni_lb_policy if strategy_result is not None else None), ) diff --git a/vllm_omni/config/resolver.py b/vllm_omni/config/resolver.py index 4f5ee8a6c..0db436768 100644 --- a/vllm_omni/config/resolver.py +++ b/vllm_omni/config/resolver.py @@ -14,7 +14,6 @@ from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remot 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 @@ -24,19 +23,10 @@ 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. - """ + """Canonical typed result returned by the production config resolver.""" config_path: str | None - stage_configs: tuple[Any, ...] # Temporary StageConfig/OmegaConf bridge. + stage_configs: tuple[Any, ...] pipeline_config: PipelineConfig | None = None omni_lb_policy: str | None = None @@ -176,27 +166,13 @@ def _load_strategy_specs(strategy_config_path: str | None) -> Mapping[Any, Any] 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 the canonical structured runtime view for a resolved pipeline.""" 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, + config_path=structured_config.orchestrator_config.deploy_config_path, + stage_configs=tuple(structured_config.stage_configs), + pipeline_config=structured_config.pipeline_config, + omni_lb_policy=getattr(structured_config, "strategy_omni_lb_policy", None), ) @@ -231,18 +207,16 @@ def resolve_omni_config( registry_overrides = dict(normalized_overrides) _flatten_stage_overrides(registry_overrides, stage_overrides) + strategy_specs = _load_strategy_specs(strategy_config_path) structured_config = StageConfigFactory.create_from_model( model, trust_remote_code=trust_remote_code, cli_overrides=registry_overrides, deploy_config_path=deploy_config_path, + strategy_specs=strategy_specs, ) if structured_config is not None: - return _build_registered_resolution( - structured_config, - cli_overrides=registry_overrides, - strategy_config_path=strategy_config_path, - ) + return _build_registered_resolution(structured_config) _apply_generic_stage_overrides(normalized_overrides, stage_overrides) supported, model_class_name = _resolve_generic_diffusion_model_class(model, normalized_overrides) @@ -252,11 +226,12 @@ def resolve_omni_config( ) if model_class_name is not None: normalized_overrides.setdefault("model_class_name", model_class_name) - default_stages = StageConfigFactory.create_default_diffusion(normalized_overrides) + default_config = StageConfigFactory.create_typed_default_diffusion(model, normalized_overrides) return OmniConfigResolution( config_path=deploy_config_path, - stage_configs=tuple(create_config(_convert_dataclasses_to_dict(default_stages))), + stage_configs=tuple(default_config.stage_configs), + pipeline_config=default_config.pipeline_config, ) diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py index 85310e9dd..d615ebca0 100644 --- a/vllm_omni/engine/omni_engine_base.py +++ b/vllm_omni/engine/omni_engine_base.py @@ -272,7 +272,14 @@ class OmniEngineBase: self.num_stages = len(self.stage_configs) self.async_chunk = any( - bool(getattr(getattr(stage, "engine_args", None), "async_chunk", False)) for stage in self.stage_configs + bool( + getattr( + getattr(stage, "connector_config", None), + "async_chunk", + getattr(getattr(stage, "engine_args", None), "async_chunk", False), + ) + ) + for stage in self.stage_configs ) self.stage_pools: list[StagePool] = [] self.stage_clients: list[StageClient] = [] # logical-stage view for external readers diff --git a/vllm_omni/engine/stage_admission.py b/vllm_omni/engine/stage_admission.py index 469787bca..bd57f805d 100644 --- a/vllm_omni/engine/stage_admission.py +++ b/vllm_omni/engine/stage_admission.py @@ -287,8 +287,13 @@ def check_admission( def _diffusion_utilization(replica: Any) -> float | None: - """Best-effort gpu_memory_utilization for a diffusion replica from raw args.""" + """Resolve a diffusion replica's memory budget from typed or legacy config.""" stage_cfg = getattr(replica, "stage_cfg", None) + cache_config = getattr(stage_cfg, "cache_config", None) + util = getattr(cache_config, "gpu_memory_utilization", None) + if util is not None: + return float(util) + engine_args = getattr(stage_cfg, "engine_args", None) if isinstance(engine_args, dict): util = engine_args.get("gpu_memory_utilization") diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 71ba17c7c..612c7481f 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -17,6 +17,7 @@ from types import SimpleNamespace from typing import Any import msgspec +import torch import zmq from omegaconf import OmegaConf from vllm.config import VllmConfig @@ -33,6 +34,7 @@ from vllm.v1.engine.utils import ( ) from vllm.v1.executor import Executor +from vllm_omni.config.omni_config import BaseVllmOmniStageConfig from vllm_omni.distributed.omni_connectors.utils import initialization from vllm_omni.engine import stage_init_utils from vllm_omni.engine.stage_init_utils import ( @@ -67,11 +69,13 @@ def _serialize_stage_config(stage_config: Any) -> Any: """Convert a stage config to msgpack-friendly builtins.""" if stage_config is None or isinstance(stage_config, (str, bytes, int, float, bool)): return stage_config + if isinstance(stage_config, torch.dtype): + return str(stage_config).removeprefix("torch.") if OmegaConf.is_config(stage_config): return _serialize_stage_config(OmegaConf.to_container(stage_config, resolve=True)) - if dataclasses.is_dataclass(stage_config): + if dataclasses.is_dataclass(stage_config) and not isinstance(stage_config, type): return _serialize_stage_config(dataclasses.asdict(stage_config)) if isinstance(stage_config, dict): @@ -1403,7 +1407,7 @@ def get_headless_replica_devices( omni_dp_size_local: int, ) -> list[str | None]: """Return per-replica device slices for a headless stage.""" - runtime_cfg = getattr(stage_cfg, "runtime", None) + runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", None)) devices_str: str | None = None if runtime_cfg is not None: devices_str = ( @@ -1509,10 +1513,11 @@ def launch_headless_diffusion_replicas( stage_id, ) - # Headless diffusion startup and its downstream helpers still consume the - # legacy StageConfig shape; switch this with the coordinated RFC #4021 - # stage-init cutover. - metadata = stage_init_utils.extract_legacy_stage_metadata(stage_cfg) + metadata = ( + stage_init_utils.extract_stage_metadata_from_omni_stage_config(stage_cfg) + if isinstance(stage_cfg, BaseVllmOmniStageConfig) + else stage_init_utils.extract_legacy_stage_metadata(stage_cfg) + ) if omni_conn_cfg: inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) # Headless single-stage launch must still infer cross-stage TP topology diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 9668dfc99..06f916394 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -402,6 +402,11 @@ def _tp_size_for_stage(stage_configs: Sequence[Any], stage_id: Any) -> int | Non for stage_cfg in stage_configs: if str(getattr(stage_cfg, "stage_id", None)) not in id_strs: continue + if isinstance(stage_cfg, BaseVllmOmniStageConfig): + try: + return max(1, int(stage_cfg.parallel_config.tensor_parallel_size)) + except (TypeError, ValueError): + return 1 engine_args = getattr(stage_cfg, "engine_args", None) if engine_args is None: return 1 @@ -489,11 +494,15 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[ rank mappings automatically. """ try: - engine_args = stage_cfg.engine_args - if hasattr(engine_args, "get"): - omni_kv = engine_args.get("omni_kv_config", None) + connector_config = getattr(stage_cfg, "connector_config", None) + if connector_config is not None: + omni_kv = connector_config.omni_kv_config else: - omni_kv = getattr(engine_args, "omni_kv_config", None) + engine_args = stage_cfg.engine_args + if hasattr(engine_args, "get"): + omni_kv = engine_args.get("omni_kv_config", None) + else: + omni_kv = getattr(engine_args, "omni_kv_config", None) if omni_kv is None: return @@ -504,7 +513,11 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[ if "stage_id" not in omni_kv: omni_kv["stage_id"] = stage_id - engine_input_source = getattr(stage_cfg, "engine_input_source", None) + engine_input_source = ( + getattr(stage_cfg, "input_sources", None) + if connector_config is not None + else getattr(stage_cfg, "engine_input_source", None) + ) if engine_input_source is not None: if hasattr(omni_kv, "setdefault"): omni_kv.setdefault("engine_input_source", list(engine_input_source)) @@ -703,17 +716,16 @@ def _resolve_omni_metadata_hook(path: str | None) -> Callable | None: def extract_stage_metadata_from_omni_stage_config( stage_config: BaseVllmOmniStageConfig, ) -> StageMetadata: - """Project one typed stage config into metadata for a future cutover. - - This projection is not used by production startup yet. Current replica - layout, engine-argument, remote-diffusion, and platform setup paths still - require the legacy StageConfig/OmegaConf shape. - """ + """Project one typed stage config into production runtime metadata.""" stage_type: Literal["llm", "diffusion"] = "diffusion" if stage_config.stage_type == StageType.DIFFUSION else "llm" - sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams - sampling_params: OmniSamplingParams = sampling_params_cls( - **(stage_config.model_config.default_sampling_params or {}) - ) + pooling_config = stage_config.pooling_config + if stage_type == "llm" and (pooling_config.runner or "").lower() == "pooling": + sampling_params: OmniSamplingParams | PoolingParams = ( + copy.deepcopy(pooling_config.default_pooling_params) or PoolingParams() + ) + else: + sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams + sampling_params = sampling_params_cls(**(stage_config.model_config.default_sampling_params or {})) custom_process_input_func = _resolve_omni_metadata_hook(stage_config.custom_process_input_func) if stage_type == "diffusion": @@ -878,7 +890,9 @@ def _get_local_llm_parallel_sizes( unset. """ if engine_args is None: - engine_args = getattr(stage_cfg, "engine_args", {}) + engine_args = getattr(stage_cfg, "parallel_config", None) + if engine_args is None: + engine_args = getattr(stage_cfg, "engine_args", {}) tp_size = int(_get_attr_or_item(engine_args, "tensor_parallel_size", 1) or 1) pp_size = int(_get_attr_or_item(engine_args, "pipeline_parallel_size", 1) or 1) local_dp_size = _get_attr_or_item(engine_args, "data_parallel_size_local", None) @@ -889,9 +903,12 @@ def _get_local_llm_parallel_sizes( def get_stage_devices_per_replica(stage_cfg: Any, engine_args: Any | None = None) -> int: """Return the number of devices consumed by one replica of *stage_cfg*.""" - if engine_args is None: - engine_args = getattr(stage_cfg, "engine_args", {}) if getattr(stage_cfg, "stage_type", "llm") == "diffusion": + typed_parallel_config = getattr(stage_cfg, "parallel_config", None) + if typed_parallel_config is not None: + return max(1, int(typed_parallel_config.world_size)) + if engine_args is None: + engine_args = getattr(stage_cfg, "engine_args", {}) parallel_config = _get_attr_or_item(engine_args, "parallel_config") if parallel_config is None: return 1 @@ -933,7 +950,7 @@ def compute_replica_layout( """ replicas_per_stage: list[int] = [] for stage_cfg in stage_configs: - runtime_cfg = getattr(stage_cfg, "runtime", {}) + runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", {})) num_replicas = int( runtime_cfg.get("num_replicas", 1) if hasattr(runtime_cfg, "get") @@ -948,7 +965,7 @@ def compute_replica_layout( num_replicas = replicas_per_stage[stage_id] if num_replicas <= 1: continue - runtime_cfg = getattr(stage_cfg, "runtime", {}) + runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", {})) devices_str = ( runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) ) @@ -1086,15 +1103,31 @@ def _project_omni_stage_engine_args( diffusion_stage = cast(VllmOmniDiffusionStageConfig, stage_config) engine_args.update(_project_omni_config_fields(diffusion_stage.diffusion_config)) + model_excluded_fields = { + "default_sampling_params", + "has_sampling_extra_args", + } + runtime_excluded_fields = {"devices", "num_replicas", "env", "num_gpus"} + if not is_diffusion: + # These values configure OmniDiffusionConfig or its worker process; + # OmniEngineArgs has no matching fields for LLM stages. + model_excluded_fields.update( + { + "disable_autocast", + "enable_multithread_weight_load", + "num_weight_load_threads", + } + ) + runtime_excluded_fields.add("log_level") + for config, excluded_fields in ( ( stage_config.model_config, - frozenset({"default_sampling_params", "has_sampling_extra_args"}), + frozenset(model_excluded_fields), ), ( stage_config.runtime_config, - frozenset({"devices", "num_replicas", "env", "num_gpus"}) - | (frozenset({"additional_config"}) if is_diffusion else frozenset()), + frozenset(runtime_excluded_fields) | (frozenset({"additional_config"}) if is_diffusion else frozenset()), ), ): engine_args.update( @@ -1103,6 +1136,11 @@ def _project_omni_stage_engine_args( exclude=excluded_fields, ) ) + pooling_config = stage_config.pooling_config + if pooling_config.runner is not None: + engine_args["runner"] = pooling_config.runner + if pooling_config.pooling_output_decoder is not None: + engine_args["pooling_output_decoder"] = pooling_config.pooling_output_decoder if is_diffusion: for config, field_map in ( @@ -1330,11 +1368,9 @@ def build_engine_args_dict( stage_connector_spec: dict[str, Any] | None = None, cli_tokenizer: str | None = None, ) -> dict[str, Any]: - """Build engine arguments through the stable production entry point. + """Preserve the legacy engine-argument API for compatibility callers. - Production inputs still use the legacy stage representation. Keep that - compatibility choice behind this function so callers do not bind directly - to a representation-specific implementation. + Typed runtime stages use ``build_engine_args_dict_from_omni_stage_config``. """ return build_legacy_engine_args_dict( stage_config, @@ -1350,13 +1386,7 @@ def build_engine_args_dict_from_omni_stage_config( stage_connector_spec: dict[str, Any] | None = None, cli_tokenizer: str | None = None, ) -> dict[str, Any]: - """Project one typed stage config into backend engine arguments. - - This projection is prepared for the RFC #4021 stage-init cutover. Current - production startup reaches the legacy implementation through - ``build_engine_args_dict`` while strategy and startup-plan inputs still - use the legacy representation. - """ + """Project one typed production stage config into backend engine arguments.""" engine_args_dict = _project_omni_stage_engine_args(stage_config) _apply_rocm_attention_backend(engine_args_dict, stage_config.stage_type) return _finalize_engine_args_dict( @@ -1391,7 +1421,7 @@ def _check_stage_device_layout(stage_config: Any, engine_args_dict: dict[str, An """ from vllm_omni.config.composable_parallel import StrategyApplyError, check_device_layout - runtime = getattr(stage_config, "runtime", None) + runtime = getattr(stage_config, "runtime_config", getattr(stage_config, "runtime", None)) devices = _get_attr_or_item(runtime, "devices", None) if runtime is not None else None if devices is None: # No explicit placement -> vLLM assigns devices itself; nothing to check. @@ -1452,10 +1482,12 @@ def build_vllm_config( (vllm_config, executor_class) """ if engine_args_dict is None: - engine_args_dict = build_engine_args_dict( - stage_config, - model, - stage_connector_spec=stage_connector_spec, + engine_args_dict = ( + build_engine_args_dict_from_omni_stage_config( + stage_config, model, stage_connector_spec=stage_connector_spec + ) + if isinstance(stage_config, BaseVllmOmniStageConfig) + else build_engine_args_dict(stage_config, model, stage_connector_spec=stage_connector_spec) ) filtered_engine_args_dict = filter_dataclass_kwargs(OmniEngineArgs, engine_args_dict) @@ -1927,7 +1959,11 @@ def build_diffusion_config( ) -> Any: """Build diffusion config for a stage.""" - engine_args_dict = build_engine_args_dict(stage_cfg, model) + engine_args_dict = ( + build_engine_args_dict_from_omni_stage_config(stage_cfg, model) + if isinstance(stage_cfg, BaseVllmOmniStageConfig) + else build_engine_args_dict(stage_cfg, model) + ) od_config = OmniDiffusionConfig.from_kwargs(**engine_args_dict) num_devices_per_stage = od_config.parallel_config.world_size diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 53400df5e..d17e3b08f 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -16,9 +16,9 @@ from dataclasses import dataclass from typing import Any, cast import janus -from omegaconf import OmegaConf from vllm.logger import init_logger +from vllm_omni.config.omni_config import BaseVllmOmniStageConfig from vllm_omni.distributed.omni_connectors.utils.initialization import ( resolve_omni_kv_config_for_stage, ) @@ -49,11 +49,13 @@ from vllm_omni.engine.stage_init_utils import ( _inject_inferred_kv_tp_topology, acquire_device_locks, build_engine_args_dict, + build_engine_args_dict_from_omni_stage_config, build_llm_stage_output_processor, build_vllm_config, compute_replica_layout, device_overlap_group_keys, extract_legacy_stage_metadata, + extract_stage_metadata_from_omni_stage_config, get_stage_connector_spec, inject_kv_stage_info, inject_omni_kv_connector_config, @@ -345,12 +347,12 @@ class StageRuntime: """Build startup plans for every logical stage and replica.""" stage_plans: list[LogicalStageInitPlan] = [] - # RFC #4021 transition boundary: stage planning still relies on legacy - # StageConfig.runtime and StageConfig.engine_args for replica and engine - # setup. Keep metadata extraction on the legacy path until the - # coordinated stage-init cutover. for stage_idx, stage_cfg in enumerate(self._stage_configs): - base_metadata = extract_legacy_stage_metadata(stage_cfg) + base_metadata = ( + extract_stage_metadata_from_omni_stage_config(stage_cfg) + if isinstance(stage_cfg, BaseVllmOmniStageConfig) + else extract_legacy_stage_metadata(stage_cfg) + ) stage_id = int(base_metadata.stage_id) if stage_id != stage_idx: raise ValueError( @@ -374,14 +376,20 @@ class StageRuntime: executor_class = None engine_args_dict = None if base_metadata.stage_type != "diffusion": - # The stable adapter entry point still receives the same - # legacy stage object as replica planning. Its implementation - # switches only at the coordinated RFC #4021 cutover. - engine_args_dict = build_engine_args_dict( - stage_cfg, - self._model, - stage_connector_spec=stage_connector_spec, - cli_tokenizer=self._tokenizer, + engine_args_dict = ( + build_engine_args_dict_from_omni_stage_config( + stage_cfg, + self._model, + stage_connector_spec=stage_connector_spec, + cli_tokenizer=self._tokenizer, + ) + if isinstance(stage_cfg, BaseVllmOmniStageConfig) + else build_engine_args_dict( + stage_cfg, + self._model, + stage_connector_spec=stage_connector_spec, + cli_tokenizer=self._tokenizer, + ) ) inject_omni_kv_connector_config( engine_args_dict, @@ -401,11 +409,20 @@ class StageRuntime: ) for replica_id in range(num_replicas): - replica_cfg = copy.deepcopy(stage_cfg) if replica_id > 0 else stage_cfg + # Keep the logical stage's device pool intact; each replica owns + # the same config throughout planning, metadata and launch. + replica_cfg = copy.deepcopy(stage_cfg) if num_replicas > 1 else stage_cfg if stage_idx in replica_devices_map: - replica_cfg.runtime.devices = replica_devices_map[stage_idx][replica_id] + devices = replica_devices_map[stage_idx][replica_id] + runtime_cfg = getattr(replica_cfg, "runtime_config", getattr(replica_cfg, "runtime", None)) + if runtime_cfg is not None: + runtime_cfg.devices = devices - replica_metadata = extract_legacy_stage_metadata(replica_cfg) + replica_metadata = ( + extract_stage_metadata_from_omni_stage_config(replica_cfg) + if isinstance(replica_cfg, BaseVllmOmniStageConfig) + else extract_legacy_stage_metadata(replica_cfg) + ) replica_metadata.replica_id = replica_id if launch_mode == "remote" and replica_metadata.stage_type != "diffusion": replica_metadata.runtime_cfg = None @@ -859,17 +876,25 @@ class StageRuntime: raise RuntimeError("Omni KV connector requires source and destination stages") inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to) inject_kv_stage_info(plan.stage_cfg, plan.metadata.stage_id, self._stage_configs) - engine_args = getattr(plan.stage_cfg, "engine_args", {}) - inline_diffusion = ( - engine_args.get("inline_diffusion", False) - if hasattr(engine_args, "get") - else getattr(engine_args, "inline_diffusion", False) - ) - custom_pipeline_args = ( - engine_args.get("custom_pipeline_args") - if hasattr(engine_args, "get") - else getattr(engine_args, "custom_pipeline_args", None) - ) + if isinstance(plan.stage_cfg, BaseVllmOmniStageConfig): + inline_diffusion = plan.stage_cfg.stage_pipeline_config.inline_diffusion + custom_pipeline_args = getattr( + getattr(plan.stage_cfg, "diffusion_config", None), + "custom_pipeline_args", + None, + ) + else: + engine_args = getattr(plan.stage_cfg, "engine_args", {}) + inline_diffusion = ( + engine_args.get("inline_diffusion", False) + if hasattr(engine_args, "get") + else getattr(engine_args, "inline_diffusion", False) + ) + custom_pipeline_args = ( + engine_args.get("custom_pipeline_args") + if hasattr(engine_args, "get") + else getattr(engine_args, "custom_pipeline_args", None) + ) client, resources = launch_diffusion_stage_replica( model=self._model, stage_config=plan.stage_cfg, @@ -1028,7 +1053,7 @@ class DistStageRuntime(StageRuntime): for idx, stage_cfg in enumerate(self._stage_configs): stage_id = int(getattr(stage_cfg, "stage_id", idx)) - runtime_cfg = getattr(stage_cfg, "runtime", None) + runtime_cfg = getattr(stage_cfg, "runtime_config", getattr(stage_cfg, "runtime", None)) if runtime_cfg is None: continue if stage_id == target_stage_id: @@ -1100,13 +1125,9 @@ class DistStageRuntime(StageRuntime): if registered_stage_cfg is None: raise ValueError(f"Remote stage {plan.metadata.stage_id} registered without stage config") - # Remote diffusion registration still transports the legacy mapping - # shape. Reconstruct and project that shape until its RFC #4021 cutover. - metadata = ( - extract_legacy_stage_metadata(OmegaConf.create(registered_stage_cfg)) - if plan.metadata.stage_type == "diffusion" - else copy.deepcopy(plan.metadata) - ) + # Registration is transport-only: the head-side typed plan remains + # authoritative for metadata and topology. + metadata = extract_stage_metadata_from_omni_stage_config(plan.stage_cfg) metadata.replica_id = plan.replica_id ctx = StageRemoteFactoryContext( stage_id=plan.metadata.stage_id, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index bc7f81132..a35f06202 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -1048,6 +1048,7 @@ def run_headless(args: TrackingNamespace) -> None: ) from vllm_omni.engine.stage_init_utils import ( build_engine_args_dict, + build_engine_args_dict_from_omni_stage_config, build_vllm_config, get_stage_connector_spec, inject_omni_kv_connector_config, @@ -1141,21 +1142,21 @@ def run_headless(args: TrackingNamespace) -> None: stage_connector_spec = get_stage_connector_spec( omni_transfer_config=omni_transfer_config, stage_id=stage_id, - async_chunk=bool(stage_cfg.engine_args.get("async_chunk", False)), + async_chunk=bool( + getattr(getattr(stage_cfg, "connector_config", None), "async_chunk", None) + if hasattr(stage_cfg, "connector_config") + else stage_cfg.engine_args.get("async_chunk", False) + ), ) - # ``runtime_cfg`` is mostly inherited from the parent's - # CUDA_VISIBLE_DEVICES; when ``--omni-dp-size-local > 1`` we additionally - # bracket each replica's spawn below with setup_stage_devices so they - # don't all stack on cuda:0 (see ``per_replica_devices`` above). - # Headless startup still supplies the legacy OmegaConf stage shape through - # the stable adapter entry point. The implementation switches only when - # RFC #4021 threads structured stage configs through the launch plan. - engine_args_dict = build_engine_args_dict( - stage_cfg, - model, - stage_connector_spec=stage_connector_spec, - cli_tokenizer=getattr(args, "tokenizer", None), + engine_args_dict = ( + build_engine_args_dict_from_omni_stage_config( + stage_cfg, model, stage_connector_spec=stage_connector_spec, cli_tokenizer=getattr(args, "tokenizer", None) + ) + if hasattr(stage_cfg, "connector_config") + else build_engine_args_dict( + stage_cfg, model, stage_connector_spec=stage_connector_spec, cli_tokenizer=getattr(args, "tokenizer", None) + ) ) inject_omni_kv_connector_config(engine_args_dict, omni_kv_connector, stage_id) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 484d5d170..5b3f2d54b 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -248,10 +248,27 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): return obj.get(key, default) return getattr(obj, key, default) + @classmethod + def _stage_model_metadata(cls, stage: Any) -> tuple[str | None, str | None]: + """Return ``(model_arch, model_stage)`` for typed or legacy stages.""" + model_config = cls._stage_get(stage, "model_config") + model_arch = cls._stage_get(model_config, "model_arch") + model_stage = cls._stage_get(stage, "model_stage") + + topology = cls._stage_get(stage, "stage_pipeline_config") + if model_stage is None: + model_stage = cls._stage_get(topology, "model_stage") + + engine_args = cls._stage_get(stage, "engine_args") + if model_arch is None: + model_arch = cls._stage_get(engine_args, "model_arch") + if model_stage is None: + model_stage = cls._stage_get(engine_args, "model_stage") + return model_arch, model_stage + def _has_minicpmo45_stage(self) -> bool: for stage in getattr(self.engine_client, "stage_configs", []) or []: - engine_args = self._stage_get(stage, "engine_args") - model_arch = self._stage_get(engine_args, "model_arch") + model_arch, _ = self._stage_model_metadata(stage) if model_arch == "MiniCPMO45OmniForConditionalGeneration": return True return False @@ -324,11 +341,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if not hasattr(sp, "output_kind"): continue - engine_args = self._stage_get(stage, "engine_args") - if self._stage_get(engine_args, "model_arch") != "MiniCPMO45OmniForConditionalGeneration": + model_arch, model_stage = self._stage_model_metadata(stage) + if model_arch != "MiniCPMO45OmniForConditionalGeneration": continue - model_stage = self._stage_get(engine_args, "model_stage") if model_stage == "llm": sp.output_kind = RequestOutputKind.FINAL_ONLY elif model_stage == "tts": diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 80f357651..8aa32899a 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -36,6 +36,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.utils import random_uuid from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError +from vllm_omni.config.stage_config import StagePipelineConfig from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin, StreamingAudioResampler from vllm_omni.entrypoints.openai.protocol.audio import ( AudioResponse, @@ -76,6 +77,24 @@ _SPEECH_USAGE_TOTAL_TOKENS_HEADER = "X-VLLM-OMNI-TOTAL-TOKENS" _SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TEXT-TOKENS" _SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-AUDIO-TOKENS" + +def _stage_speech_metadata(stage: Any) -> tuple[str | None, str | None, str | None]: + """Return speech routing metadata from typed or legacy stage configs.""" + topology = getattr(stage, "stage_pipeline_config", None) + if isinstance(topology, StagePipelineConfig): + return ( + getattr(topology, "model_stage", None), + getattr(getattr(stage, "model_config", None), "model_arch", None), + getattr(stage, "worker_type", None), + ) + engine_args = getattr(stage, "engine_args", None) + return ( + getattr(engine_args, "model_stage", None), + getattr(engine_args, "model_arch", None), + getattr(engine_args, "worker_type", None), + ) + + # TTS Configuration # # The stage-key -> model-type mapping is NOT declared here: it is derived from @@ -472,12 +491,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): """Find and return the TTS stage config, or None if not found.""" tts_stage_keys = all_tts_stage_keys() entry_stage_archs = tts_entry_stage_archs() - all_stages = frozenset(stage.engine_args.model_stage for stage in self.engine_client.stage_configs) + all_stages = frozenset( + model_stage + for stage in self.engine_client.stage_configs + if (model_stage := _stage_speech_metadata(stage)[0]) is not None + ) for stage in self.engine_client.stage_configs: - engine_args = stage.engine_args - model_stage = engine_args.model_stage - model_arch = getattr(engine_args, "model_arch", None) - worker_type = getattr(engine_args, "worker_type", None) + model_stage, model_arch, worker_type = _stage_speech_metadata(stage) if model_stage in tts_stage_keys: # Owning the stage key is not always enough: a model may be # speech-capable only in some deployment topologies. Ask the @@ -501,10 +521,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): """ if self._tts_stage is None: return None - return detect_tts_model_type( - getattr(self._tts_stage.engine_args, "model_stage", None), - getattr(self._tts_stage.engine_args, "model_arch", None), - ) + return detect_tts_model_type(*_stage_speech_metadata(self._tts_stage)[:2]) def _compute_max_instructions_length(self) -> int: """Compute max instructions length with precedence: CLI > stage config > default. @@ -518,7 +535,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # 2. Try to get from TTS stage config if self._tts_stage is not None: - tts_args = getattr(self._tts_stage, "tts_args", {}) + topology = getattr(self._tts_stage, "stage_pipeline_config", None) + tts_args = ( + getattr(topology, "extras", {}).get("tts_args", {}) + if isinstance(topology, StagePipelineConfig) + else getattr(self._tts_stage, "tts_args", {}) + ) if "max_instructions_length" in tts_args: return tts_args["max_instructions_length"] @@ -1896,9 +1918,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # chat-template markers (im_start_token_id 151644). A raw-text # prompt produces a 1-token placeholder that crashes the talker's # prefill/decode handoff. Reject early with an actionable message. - stage_names = { - getattr(getattr(s, "engine_args", None), "model_stage", None) for s in self.engine_client.stage_configs - } + stage_names = {_stage_speech_metadata(s)[0] for s in self.engine_client.stage_configs} if "talker" in stage_names: raise ValueError( "The /v1/audio/speech endpoint is only supported for " diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 9a24a1ef9..717496362 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -53,6 +53,30 @@ logger = init_logger(__name__) _VIDEO_RESPONSE_FRAME_CONVERSION_WORKERS = 8 + +def _config_value(config: Any, key: str, default: Any = None) -> Any: + if isinstance(config, Mapping): + return config.get(key, default) + return getattr(config, key, default) + + +def _stage_diffusion_model_class_name(stage_config: Any) -> str | None: + """Resolve a stage's diffusion class across typed and legacy shapes.""" + diffusion_config = _config_value(stage_config, "diffusion_config") + model_class_name = _config_value(diffusion_config, "model_class_name") + if model_class_name: + return str(model_class_name) + + model_config = _config_value(stage_config, "model_config") + model_arch = _config_value(model_config, "model_arch") or _config_value(stage_config, "model_arch") + if model_arch: + return str(model_arch) + + engine_args = _config_value(stage_config, "engine_args", {}) + model_class_name = _config_value(engine_args, "model_class_name") + return str(model_class_name) if model_class_name else None + + if TYPE_CHECKING: from vllm_omni.diffusion.data import OmniDiffusionConfig @@ -195,18 +219,7 @@ class OmniOpenAIServingVideo: model_class_name = getattr(od_config, "model_class_name", None) model_archs = [model_class_name] for stage_config in self.stage_configs or (): - stage_get = ( - stage_config.get if isinstance(stage_config, Mapping) else lambda key: getattr(stage_config, key, None) - ) - engine_args = stage_get("engine_args") or {} - model_archs.extend( - ( - stage_get("model_arch"), - engine_args.get("model_class_name") - if isinstance(engine_args, Mapping) - else getattr(engine_args, "model_class_name", None), - ) - ) + model_archs.append(_stage_diffusion_model_class_name(stage_config)) metadata_capability = any( get_diffusion_model_metadata(model_arch).supports_mixed_reference_inputs for model_arch in model_archs ) @@ -223,18 +236,7 @@ class OmniOpenAIServingVideo: od_config = self._resolve_diffusion_od_config() model_archs = [None if od_config is None else getattr(od_config, "model_class_name", None)] for stage_config in self.stage_configs or (): - stage_get = ( - stage_config.get if isinstance(stage_config, Mapping) else lambda key: getattr(stage_config, key, None) - ) - engine_args = stage_get("engine_args") or {} - model_archs.extend( - ( - stage_get("model_arch"), - engine_args.get("model_class_name") - if isinstance(engine_args, Mapping) - else getattr(engine_args, "model_class_name", None), - ) - ) + model_archs.append(_stage_diffusion_model_class_name(stage_config)) supported: set[str] = set() for model_arch in model_archs: diff --git a/vllm_omni/entrypoints/openai/tts_adapters/audex.py b/vllm_omni/entrypoints/openai/tts_adapters/audex.py index 0e31b633f..3b9e4a9a5 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/audex.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/audex.py @@ -42,7 +42,11 @@ class AudexAdapter(ARTTSAdapter): from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot - stage = getattr(getattr(self.ctx.server._tts_stage, "engine_args", None), "model_stage", None) + stage_config = self.ctx.server._tts_stage + topology = getattr(stage_config, "stage_pipeline_config", None) + stage = getattr(topology, "model_stage", None) + if stage is None: + stage = getattr(getattr(stage_config, "engine_args", None), "model_stage", None) profile, folder = ( ("full", "checkpoint_folder_full") if stage == "audex_omni" else ("tts", "checkpoint_folder_audiogen") ) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index ff43b27a5..cb760e865 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -39,6 +39,21 @@ DEFAULT_TTS_LANGUAGES = frozenset( ) +def resolve_stage_model_path(engine_client: Any) -> str | None: + """Prefer typed stage model overrides, then legacy and served-model values.""" + stages = getattr(engine_client, "stage_configs", ()) or () + for stage in stages: + model_path = getattr(getattr(stage, "model_config", None), "model", None) + if model_path: + return str(model_path) + for stage in stages: + model_path = getattr(getattr(stage, "engine_args", None), "model", None) + if model_path: + return str(model_path) + model_path = getattr(engine_client, "model", None) + return str(model_path) if model_path else None + + def conditioning_cache_salt(request: "OpenAICreateSpeechRequest", tts_params: dict | None = None) -> str: """Stable hash of the real Stage 0 conditioning for the prefix cache. diff --git a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py index 207d8b47b..33f16cf0a 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py @@ -9,7 +9,12 @@ import numpy as np from vllm.inputs import tokens_input from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens +from vllm_omni.entrypoints.openai.tts_adapters.base import ( + ARTTSAdapter, + PreparedRequest, + apply_max_new_tokens, + resolve_stage_model_path, +) if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest @@ -81,17 +86,7 @@ class HiggsAudioV2Adapter(ARTTSAdapter): from transformers import AutoProcessor - model_path = None - for stage in self.engine_client.stage_configs: - model_path = getattr(getattr(stage, "engine_args", None), "model", None) - if model_path: - break - if model_path is None: - # Fallback: the orchestrator stores the served model id on the engine - # itself (set by AsyncOmniEngine.__init__). Stage-level engine_args - # may not surface ``model`` when the deploy yaml doesn't set it per - # stage (the CLI-passed model id is the single source of truth). - model_path = getattr(self.engine_client, "model", None) + model_path = resolve_stage_model_path(self.engine_client) if model_path is None: raise RuntimeError("higgs_audio_v2 serving could not resolve the model path from the engine stage configs") processor = AutoProcessor.from_pretrained(model_path) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py index 1e64f7a63..d993a7395 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py @@ -16,6 +16,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ( PreparedRequest, apply_max_new_tokens, conditioning_cache_salt, + resolve_stage_model_path, ) _REF_CODE_CACHE_MAX_ENTRIES = 256 @@ -170,13 +171,7 @@ class HiggsAudioV3Adapter(ARTTSAdapter): HiggsAudioV3TokenizerAdapter, ) - model_path = None - for stage in self.engine_client.stage_configs: - model_path = getattr(getattr(stage, "engine_args", None), "model", None) - if model_path: - break - if model_path is None: - model_path = getattr(self.engine_client, "model", None) + model_path = resolve_stage_model_path(self.engine_client) if model_path is None: raise RuntimeError("higgs_audio_v3 serving could not resolve model path") tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) diff --git a/vllm_omni/entrypoints/openai/video/generation/helpers.py b/vllm_omni/entrypoints/openai/video/generation/helpers.py index 4473fce07..0ffd1ee83 100644 --- a/vllm_omni/entrypoints/openai/video/generation/helpers.py +++ b/vllm_omni/entrypoints/openai/video/generation/helpers.py @@ -58,6 +58,7 @@ from vllm_omni.entrypoints.openai.serving_video import ( ReferenceAudio, ReferenceImage, ReferenceVideo, + _stage_diffusion_model_class_name, ) from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER from vllm_omni.entrypoints.openai.stores import VIDEO_STORE @@ -136,10 +137,6 @@ def _config_get(config: Any, key: str, default: Any = None) -> Any: return getattr(config, key, default) -def _stage_engine_args(stage_cfg: Any) -> Any: - return _config_get(stage_cfg, "engine_args", {}) or {} - - def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]: if not stage_configs: return [] @@ -150,7 +147,7 @@ def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]: for stage_cfg in stage_configs: if get_stage_type(stage_cfg) != "diffusion": continue - model_class_name = _config_get(_stage_engine_args(stage_cfg), "model_class_name") + model_class_name = _stage_diffusion_model_class_name(stage_cfg) if not model_class_name: continue model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py index 562ed70d4..11fde2bf8 100644 --- a/vllm_omni/entrypoints/openpi/serving.py +++ b/vllm_omni/entrypoints/openpi/serving.py @@ -103,6 +103,14 @@ class ServingRealtimeRobotOpenPI: for stage_config in getattr(engine_client, "stage_configs", []) or []: if getattr(stage_config, "stage_type", None) != "diffusion": continue + # Typed diffusion stages keep model-owned OpenPI handshake + # metadata in diffusion_config.model_config. The out-of-process + # head only has this stage view because full od_config lives in + # the worker. + diffusion_config = getattr(stage_config, "diffusion_config", None) + model_config = getattr(diffusion_config, "model_config", None) + if model_config is not None: + break engine_args = getattr(stage_config, "engine_args", None) model_config = getattr(engine_args, "model_config", None) if model_config is not None: diff --git a/vllm_omni/entrypoints/serve/profile/utils.py b/vllm_omni/entrypoints/serve/profile/utils.py index f444a3534..d91ef2cad 100644 --- a/vllm_omni/entrypoints/serve/profile/utils.py +++ b/vllm_omni/entrypoints/serve/profile/utils.py @@ -4,18 +4,20 @@ def _should_enable_profiler_endpoints(stage_configs: list | None) -> bool: - """Check if any stage has profiler_config set in its engine_args.""" + """Check whether any typed or legacy stage enables a profiler.""" if not stage_configs: return False for stage in stage_configs: - engine_args = stage.get("engine_args") if isinstance(stage, dict) else getattr(stage, "engine_args", None) - if engine_args is None: - continue profiler_config = ( - engine_args.get("profiler_config") - if isinstance(engine_args, dict) - else getattr(engine_args, "profiler_config", None) + stage.get("profiler_config") if isinstance(stage, dict) else getattr(stage, "profiler_config", None) ) + if profiler_config is None: + engine_args = stage.get("engine_args") if isinstance(stage, dict) else getattr(stage, "engine_args", None) + profiler_config = ( + engine_args.get("profiler_config") + if isinstance(engine_args, dict) + else getattr(engine_args, "profiler_config", None) + ) if profiler_config is not None: profiler = ( profiler_config.get("profiler") diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 281cbc55c..8c2ba0747 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -1,7 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import logging import os +from dataclasses import asdict, is_dataclass from multiprocessing import shared_memory as _shm from typing import Any @@ -274,6 +278,8 @@ def _to_dict(x: Any) -> dict[str, Any]: try: if isinstance(x, dict): return dict(x) + if is_dataclass(x) and not isinstance(x, type): + return asdict(x) return _omega_to_dict(x) except Exception: try: diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 39131fd58..74234d8d6 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -18,6 +18,13 @@ logger = init_logger(__name__) 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.""" + typed_connector_config = getattr(stage, "connector_config", None) + if typed_connector_config is not None: + omni_conf_dict = dict(typed_connector_config.omni_kv_config or {}) + omni_conf_dict.update(connector_config=omni_conn_cfg, omni_from_stage=omni_from, omni_to_stage=omni_to) + typed_connector_config.omni_kv_config = omni_conf_dict + return + # Prepare omni_kv_config dict omni_conf_dict = {} try: