[Bugfix] Preserve session IDs when upgrading Omni engine requests

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

tests/e2e/accuracy/sana_video/conftest.py

@@ -1,26 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-import pytest
-
-
-@pytest.fixture
-def sana_model_parallel(tmp_path):
-    """Construct native parallel layers with real single-rank GPU groups."""
-    from vllm_omni.diffusion.distributed.parallel_state import (
-        destroy_distributed_env,
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-
-    init_distributed_environment(
-        world_size=1,
-        rank=0,
-        local_rank=0,
-        distributed_init_method=f"file://{tmp_path / 'sana-distributed'}",
-    )
-    try:
-        initialize_model_parallel()
-        yield
-    finally:
-        destroy_distributed_env()

tests/e2e/accuracy/sana_video/test_sana_video_golden.py

@@ -202,7 +202,6 @@ def test_sana_video_transformer_matches_frozen_golden(
     variant: str,
     _hardware,
     accuracy_artifact_root: Path,
-    sana_model_parallel,
 ) -> None:
     del _hardware
     from vllm.config import LoadConfig

tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py

@@ -189,7 +189,6 @@ def _cfg_prediction(
 def test_sana_video_quick_pipeline_stages_match_diffusers(
     variant: str,
     _hardware,
-    sana_model_parallel,
 ) -> None:
     """Compare the first denoising step without replacing native attention.
 

tests/engine/test_omni_request.py

@@ -13,13 +13,15 @@ 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 PromptEmbedsPayload
+from vllm_omni.engine import OmniEngineCoreRequest, PromptEmbedsPayload
 from vllm_omni.request import OmniRequest, OmniStreamingUpdate
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -111,3 +113,24 @@ 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

vllm_omni/engine/__init__.py

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