[Core] Retain structured configs through runtime startup

docs/design/architecture_overview.md

@@ -323,7 +323,7 @@ inputs, fields, and ownership rules are described below.
 flowchart TB
     layer1["Layer 1 · Authoring inputs<br/>PipelineConfig + DeployConfig"]
     layer2["Layer 2 · Production resolution boundary<br/>resolve_omni_config()<br/>(StageConfigFactory + VllmOmniConfig internally)"]
-    layer3["Layer 3 · Startup hand-off<br/>OmniConfigResolution<br/>PipelineConfig + temporary OmegaConf stage bridge"]
+    layer3["Layer 3 · Startup hand-off<br/>OmniConfigResolution<br/>PipelineConfig + typed stage configs"]
     layer4["Layer 4 · Runtime launch planning<br/>StageRuntime"]
     layer5["Layer 5 · Engine materialization<br/>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.

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 @@
     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 emit(self, record: logging.LogRecord) -> None:
     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

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

tests/config/test_omni_config.py

@@ -11,6 +11,7 @@
 
 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
 

tests/e2e/features/comfyui/test_comfyui_integration.py

@@ -534,7 +534,7 @@ async def _mock_preprocess_chat(self, *args, **kwargs):
     # 1. The API layer uses its stage_list and stage_configs attributes
     # 2. Its __init__ method has slow side effects (model & config loading).
     # These cases use the non-duplex mock engine without loading model configs.
-    mocker.patch("vllm_omni.entrypoints.openai.api_server._is_duplex_model", return_value=False)
+    mocker.patch("vllm_omni.entrypoints.openai.api_server._should_serve_duplex", return_value=False)
     mock_async_omni_cls = mocker.patch("vllm_omni.entrypoints.openai.api_server.AsyncOmni")
     monkeypatch.setattr(
         "vllm_omni.entrypoints.openai.serving_chat.OmniOpenAIServingChat._preprocess_chat",

tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py

@@ -21,6 +21,7 @@
 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]