[Core] Retain structured configs through runtime startup (#6849)

Signed-off-by: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com>

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

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 @@ 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