[Bugfix] Provide Omni lifecycle state for auxiliary v2 stages

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

tests/engine/test_omni_request.py

@@ -13,15 +13,13 @@ substitutable with the base class — including base-style calls that pass
 
 import inspect
 
-import msgspec
 import numpy as np
 import pytest
 import torch
 from vllm.sampling_params import SamplingParams
-from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.request import Request
 
-from vllm_omni.engine import OmniEngineCoreRequest, PromptEmbedsPayload
+from vllm_omni.engine import PromptEmbedsPayload
 from vllm_omni.request import OmniRequest, OmniStreamingUpdate
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -113,24 +111,3 @@ def test_model_intermediate_buffer_round_trips_to_streaming_update():
     assert update is not None
     assert request.model_intermediate_buffer is buffer
     assert update.model_intermediate_buffer is buffer
-
-
-@pytest.mark.parametrize("session_id", [None, "session"])
-def test_session_survives_omni_upgrade_wire_transport_and_scheduler_conversion(session_id):
-    core = EngineCoreRequest(
-        request_id="session-request",
-        prompt_token_ids=[1, 2],
-        mm_features=None,
-        sampling_params=SamplingParams(max_tokens=2),
-        pooling_params=None,
-        arrival_time=0.0,
-        lora_request=None,
-        cache_salt=None,
-        data_parallel_rank=None,
-        session_id=session_id,
-    )
-    upgraded = OmniEngineCoreRequest.from_request(core)
-    decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(upgraded), type=OmniEngineCoreRequest)
-    scheduled = OmniRequest.from_engine_core_request(decoded, block_hasher=None)
-    assert decoded.session_id == session_id
-    assert scheduled.session_id == session_id

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/engine/__init__.py

@@ -123,7 +123,6 @@ class OmniEngineCoreRequest(EngineCoreRequest):
             trace_headers=request.trace_headers,
             resumable=request.resumable,
             external_req_id=request.external_req_id,
-            session_id=request.session_id,
             reasoning_ended=request.reasoning_ended,
             reasoning_parser_kwargs=request.reasoning_parser_kwargs,
             abort_immediately=request.abort_immediately,

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