[Test] Restore stage-client class identities after the module reload test

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

@@ -268,11 +268,21 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
     """Regression test for forward references in make_async_mp_client."""
     import vllm_omni.engine.stage_engine_core_client as client_mod
 
-    importlib.reload(client_mod)
+    # Snapshot the module namespace so a reload cannot strand other test
+    # modules that captured the pre-reload class objects (isinstance checks
+    # against the stale class would then fail). ``importlib.reload`` mutates
+    # the module dict in place; restoring it leaves the runtime exactly as it
+    # was before this test ran.
+    saved = dict(client_mod.__dict__)
+    try:
+        importlib.reload(client_mod)
 
-    assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
-        "StageEngineCoreClient | DPLBStageEngineCoreClient"
-    )
+        assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
+            "StageEngineCoreClient | DPLBStageEngineCoreClient"
+        )
+    finally:
+        client_mod.__dict__.clear()
+        client_mod.__dict__.update(saved)
 
 
 def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch):