[Bugfix] Preserve session IDs when upgrading Omni engine requests

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

.buildkite/cuda/rebase-pipeline.yml

@@ -1,46 +1,35 @@
-# Focused weekly comparison against unchanged Omni main and official vLLM 0.30.0.
 steps:
-- label: ':docker: Build image'
+  - label: ":docker: Build image"
-  key: image-build
+    key: image-build
-  commands:
+    commands:
-  - bash .buildkite/cuda/scripts/build-ci-image.sh
+      - "bash .buildkite/cuda/scripts/build-ci-image.sh"
-  agents:
+    agents:
-    queue: cpu_queue_premerge
+      queue: "cpu_queue_premerge"
-- label: Upload selected weekly baseline
+
-  depends_on: image-build
+  # L2 Test
-  key: upload-weekly-pipeline
+  - label: "Upload Ready Pipeline"
-  commands:
+    depends_on: image-build
-  - |
+    key: upload-ready-pipeline
-    python3 - <<'PYBASELINE'
+    commands:
-    import tempfile
+      - python3 .buildkite/common/scripts/upload_pipeline.py --upload --all .buildkite/cuda/test-ready.yml
-    from pathlib import Path
+    agents:
-    import yaml
+      queue: "cpu_queue_premerge"
-    labels = set(['TTS · Audex · CosyVoice · Qwen3-TTS · L4', 'TTS · H100', 'Diffusion · Bagel · H100 · Single-GPU', 'Diffusion · SenseNova-U1 · H100 · Single-GPU', 'Diffusion · DreamZero · H100 · 2-GPU', 'Diffusion · H100 · 4-GPU', 'Diffusion · MAGI-2 Preview Serving Smoke · H100 · 4-GPU'])
+
-    source = yaml.safe_load(Path('.buildkite/cuda/test-weekly.yml').read_text())
+  # L3 Test
-    selected = set()
+  - label: "Upload Merge Pipeline"
-    guard = 'python3 -c \'import vllm, torch, flashinfer, vllm._C_stable_libtorch; assert vllm.__version__ == \'"\'"\'0.30.0\'"\'"\'; assert torch.cuda.is_available(); print(\'"\'"\'Verified baseline runtime:\'"\'"\', vllm.__version__, torch.__version__)\''
+    depends_on: image-build
-    def select(steps):
+    key: upload-merge-pipeline
-        kept = []
+    commands:
-        for step in steps:
+      - python3 .buildkite/common/scripts/upload_pipeline.py --upload --all .buildkite/cuda/test-merge.yml
-            if not isinstance(step, dict):
+    agents:
-                continue
+      queue: "cpu_queue_premerge"
-            if 'steps' in step:
+
-                children = select(step['steps'])
+  # L4 Test
-                if children:
+  - label: "Upload Nightly Pipeline"
-                    kept.append(dict(step, steps=children))
+    depends_on: image-build
-            elif step.get('label') in labels:
+    key: upload-nightly-pipeline
-                item = dict(step)
+    if: build.env("NIGHTLY") == "1"
-                commands = item['commands']
+    commands:
-                item['commands'] = [guard] + ([commands] if isinstance(commands, str) else commands)
+      - python3 .buildkite/common/scripts/upload_pipeline.py --upload --all .buildkite/cuda/test-nightly.yml
-                kept.append(item)
+    agents:
-                selected.add(item['label'])
+      queue: "cpu_queue_premerge"
-        return kept
-    source['steps'] = select(source['steps'])
-    assert selected == labels, (selected, labels)
-    path = Path(tempfile.gettempdir()) / 'vllm-omni-selected-weekly-baseline.yml'
-    path.write_text(yaml.safe_dump(source, sort_keys=False, allow_unicode=True))
-    print('Selected baseline suites:', sorted(selected), flush=True)
-    PYBASELINE
-    python3 .buildkite/common/scripts/upload_pipeline.py --upload --all /tmp/vllm-omni-selected-weekly-baseline.yml
-  agents:
-    queue: cpu_queue_premerge

docker/Dockerfile.ci

@@ -36,8 +36,7 @@ RUN apt-get update && \
 # for CI
 COPY pyproject.toml setup.py ./
 COPY requirements/ requirements/
-# Keep source-built dependency versions independent of the Omni CI version.
+ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
-ENV VLLM_OMNI_VERSION_OVERRIDE=0.0.0
 # Editable install: deps + project metadata are cached in this layer, while
 # the package itself is linked to APP_DIR (populated later by COPY --link).
 # Do NOT use a non-editable install here — that copies an empty stub into

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,