[Bugfix] Provide Omni lifecycle state for auxiliary v2 stages

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'
-  key: image-build
-  commands:
-  - bash .buildkite/cuda/scripts/build-ci-image.sh
-  agents:
-    queue: cpu_queue_premerge
-- label: Upload selected weekly baseline
-  depends_on: image-build
-  key: upload-weekly-pipeline
-  commands:
-  - |
-    python3 - <<'PYBASELINE'
-    import tempfile
-    from pathlib import Path
-    import yaml
-    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())
-    selected = set()
-    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__)\''
-    def select(steps):
-        kept = []
-        for step in steps:
-            if not isinstance(step, dict):
-                continue
-            if 'steps' in step:
-                children = select(step['steps'])
-                if children:
-                    kept.append(dict(step, steps=children))
-            elif step.get('label') in labels:
-                item = dict(step)
-                commands = item['commands']
-                item['commands'] = [guard] + ([commands] if isinstance(commands, str) else commands)
-                kept.append(item)
-                selected.add(item['label'])
-        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
+  - label: ":docker: Build image"
+    key: image-build
+    commands:
+      - "bash .buildkite/cuda/scripts/build-ci-image.sh"
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L2 Test
+  - label: "Upload Ready Pipeline"
+    depends_on: image-build
+    key: upload-ready-pipeline
+    commands:
+      - python3 .buildkite/common/scripts/upload_pipeline.py --upload --all .buildkite/cuda/test-ready.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L3 Test
+  - label: "Upload Merge Pipeline"
+    depends_on: image-build
+    key: upload-merge-pipeline
+    commands:
+      - python3 .buildkite/common/scripts/upload_pipeline.py --upload --all .buildkite/cuda/test-merge.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L4 Test
+  - label: "Upload Nightly Pipeline"
+    depends_on: image-build
+    key: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - python3 .buildkite/common/scripts/upload_pipeline.py --upload --all .buildkite/cuda/test-nightly.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 VLLM_OMNI_VERSION_OVERRIDE=0.0.0
+ENV SETUPTOOLS_SCM_PRETEND_VERSION=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/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