[Bugfix] Provide Omni lifecycle state for auxiliary v2 stages

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py

@@ -216,40 +216,44 @@ def test_video_chat_stream_invalid_requests(
 # WS /v1/realtime
 # ─────────────────────────────────────────────────────────────────────────────
 
+# Resolved in ``test_realtime_invalid_requests`` (skip ``session.created`` handshake).
+_REALTIME_WS_MODEL_MISMATCH = object()
+_REALTIME_WS_INVALID_AUDIO_APPEND = object()
+
 
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize(
     "send_frames_spec, ws_error_code, err_message",
     [
-        pytest.param("{not-json", "invalid_event", "Invalid or unrecognized client event", id="invalid_json"),
+        pytest.param("{not-json", "invalid_json", "Invalid JSON", id="invalid_json"),
         pytest.param(
             json.dumps({"type": "session.update"}),
             "invalid_event",
-            "Invalid or unrecognized client event",
-            id="session_update_missing_session",
+            ("model", "Missing required field"),
+            id="session_update_missing_model",
         ),
         pytest.param(
-            json.dumps({"type": "session.update", "session": []}),
-            "invalid_event",
-            "Invalid or unrecognized client event",
-            id="session_update_invalid_session",
+            _REALTIME_WS_MODEL_MISMATCH,
+            "model_not_found",
+            ("model", "does not exist"),
+            id="session_update_model_mismatch",
         ),
         pytest.param(
-            json.dumps({"type": "input_audio_buffer.commit"}),
-            "invalid_request_error",
-            "Input audio buffer is empty",
-            id="commit_without_audio",
+            json.dumps({"type": "input_audio_buffer.commit", "final": False}),
+            "model_not_validated",
+            ("session.update", "validate the model"),
+            id="commit_without_model_validation",
         ),
         pytest.param(
             json.dumps({"type": "unknown.event"}),
-            "invalid_event",
-            "Invalid or unrecognized client event",
+            "unknown_event",
+            "Unknown event type",
             id="unknown_event_type",
         ),
         pytest.param(
-            json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
-            "invalid_request_error",
-            "Invalid base64 audio data",
+            _REALTIME_WS_INVALID_AUDIO_APPEND,
+            "invalid_audio",
+            "Invalid audio data",
             id="invalid_audio_append",
         ),
     ],
@@ -262,12 +266,27 @@ def test_realtime_invalid_requests(
     ws_error_code: str,
     err_message: str | tuple[str, ...],
 ) -> None:
+    send_frames: Any
+    if send_frames_spec is _REALTIME_WS_MODEL_MISMATCH:
+        send_frames = json.dumps(
+            {
+                "type": "session.update",
+                "model": "this-model-is-not-served-on-this-instance",
+            }
+        )
+    elif send_frames_spec is _REALTIME_WS_INVALID_AUDIO_APPEND:
+        send_frames = [
+            json.dumps({"type": "session.update", "model": omni_server.model}),
+            json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
+        ]
+    else:
+        send_frames = send_frames_spec
     online_client.send_realtime_ws_request(
         {
-            "send_frames": send_frames_spec,
+            "send_frames": send_frames,
             "timeout": 120,
             "ws_max_size": None,
-            "ws_skip_types": ["session.created", "conversation.created"],
+            "ws_skip_types": ["session.created"],
             "ws_json_type": "error",
             "ws_error_code": ws_error_code,
             "err_message": err_message,

tests/entrypoints/test_realtime_connection_helpers.py

@@ -4,10 +4,7 @@
 
 from __future__ import annotations
 
-import asyncio
 import base64
-import json
-from types import SimpleNamespace
 
 import numpy as np
 import pytest
@@ -87,56 +84,3 @@ class TestAsyncOmniStreamingParamsValidation:
         p = SamplingParams(n=1, stop=["\n"], output_kind=RequestOutputKind.DELTA)
         with pytest.raises(ValueError, match="Input streaming"):
             AsyncOmni._validate_streaming_input_sampling_params(p)
-
-
-@pytest.mark.parametrize(
-    "frame,code,message",
-    [
-        ("{not-json", "invalid_event", "Invalid or unrecognized client event"),
-        (json.dumps({"type": "session.update"}), "invalid_event", "Invalid or unrecognized client event"),
-        (
-            json.dumps({"type": "session.update", "session": []}),
-            "invalid_event",
-            "Invalid or unrecognized client event",
-        ),
-        (json.dumps({"type": "input_audio_buffer.commit"}), "invalid_request_error", "Input audio buffer is empty"),
-        (json.dumps({"type": "unknown.event"}), "invalid_event", "Invalid or unrecognized client event"),
-        (
-            json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
-            "invalid_request_error",
-            "Invalid base64 audio data",
-        ),
-    ],
-)
-def test_full_duplex_invalid_request_contract(frame, code, message):
-    from starlette.websockets import WebSocketDisconnect
-
-    from vllm_omni.entrypoints.openai.realtime.connection import OpenAIFullDuplexConnection
-
-    class Socket:
-        def __init__(self):
-            self.sent = []
-            self.frames = iter([frame])
-
-        async def accept(self):
-            pass
-
-        async def receive_text(self):
-            value = next(self.frames, None)
-            if value is None:
-                raise WebSocketDisconnect()
-            return value
-
-        async def send_text(self, payload):
-            self.sent.append(json.loads(payload))
-
-    async def run():
-        socket = Socket()
-        connection = OpenAIFullDuplexConnection(socket, SimpleNamespace(), "test-model", None)
-        await connection.handle_connection()
-        assert [event["type"] for event in socket.sent] == ["session.created", "conversation.created", "error"]
-        error = socket.sent[-1]["error"]
-        assert error["code"] == code
-        assert error["message"] == message
-
-    asyncio.run(run())

tests/worker_v2/test_omni_gpu_model_runner.py

@@ -13,6 +13,7 @@ from vllm import SamplingParams
 from vllm.config.compilation import CUDAGraphMode
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+from vllm.v1.worker.gpu.model_states.default import DefaultModelState
 from vllm.v1.worker.gpu.sample.sampler import Sampler
 
 from vllm_omni.config.model import OmniModelConfig
@@ -163,6 +164,28 @@ def test_init_model_state_factory_dispatches_omni_only(monkeypatch, flag, expect
         assert state is upstream.return_value
 
 
+@pytest.mark.parametrize("omni_stage", [False, True])
+@pytest.mark.parametrize("specialized", [False, True])
+def test_auxiliary_default_state_has_omni_lifecycle_but_preserves_upstream_states(monkeypatch, omni_stage, specialized):
+    class SpecializedState(DefaultModelState):
+        pass
+
+    state_cls = SpecializedState if specialized else DefaultModelState
+    upstream_state = object.__new__(state_cls)
+    monkeypatch.setattr("vllm_omni.worker_v2.model_states._upstream_init_model_state", lambda *_args: upstream_state)
+    monkeypatch.setattr(OmniModelState, "__init__", lambda *_args: None)
+    model_config = object.__new__(OmniModelConfig) if omni_stage else SimpleNamespace()
+    state = init_omni_model_state(
+        SimpleNamespace(model_config=model_config), SimpleNamespace(), None, torch.device("cpu")
+    )
+    if omni_stage and not specialized:
+        assert isinstance(state, OmniModelState)
+        assert callable(state.run_preprocess)
+        assert callable(state.postprocess_model_output)
+    else:
+        assert state is upstream_state
+
+
 def test_finish_requests_notifies_model_and_cleans_only_known_slots(monkeypatch):
     runner = _make_runner()
     calls = []

vllm_omni/worker_v2/model_states/__init__.py

@@ -4,7 +4,8 @@
 """Omni-aware ``init_model_state`` factory.
 
 Extends the upstream v2 factory with Omni architecture dispatch.
-Non-Omni architectures fall through to the upstream ``init_model_state``.
+Non-Omni architectures use upstream dispatch; its default state acquires
+Omni lifecycle support when owned by an Omni stage.
 """
 
 from __future__ import annotations
@@ -16,8 +17,11 @@ from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.model_states import (
     init_model_state as _upstream_init_model_state,
 )
+from vllm.v1.worker.gpu.model_states.default import DefaultModelState
 from vllm.v1.worker.gpu.model_states.interface import ModelState
 
+from vllm_omni.config.model import OmniModelConfig
+
 
 def init_omni_model_state(
     vllm_config: VllmConfig,
@@ -28,8 +32,8 @@ def init_omni_model_state(
     """Create the appropriate ``ModelState`` for *model*.
 
     Returns an ``OmniModelState`` when the model declares Omni lifecycle
-    capabilities; otherwise
-    delegates to the upstream v2 factory.
+    capabilities. Otherwise delegates to the upstream v2 factory, wrapping
+    its default state for auxiliary Omni stages.
     """
     factory = getattr(model, "create_omni_model_state", None)
     if factory is not None:
@@ -49,4 +53,12 @@ def init_omni_model_state(
         )
 
         return OmniModelState(vllm_config, model, encoder_cache, device)
-    return _upstream_init_model_state(vllm_config, model, encoder_cache, device)
+    state = _upstream_init_model_state(vllm_config, model, encoder_cache, device)
+    if type(state) is DefaultModelState and isinstance(getattr(vllm_config, "model_config", None), OmniModelConfig):
+        # Auxiliary stages such as the forced aligner have no model-owned
+        # hooks, but Omni's runner still needs its request/output lifecycle.
+        # Keep specialized upstream states and plain vLLM configs intact.
+        from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+
+        return OmniModelState(vllm_config, model, encoder_cache, device)
+    return state