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

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

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"),

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

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,

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(

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())

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"]}

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.
 

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.",