rebase: align vllm-omni with vLLM ac9126e58aa7

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

.buildkite/common/scripts/upload_pipeline.py

@@ -102,6 +102,7 @@ CI_SOURCE_FILE_DEPENDENCIES_PATH = ROOT / ".buildkite/common/ci_source_file_depe
 SOURCE_FILTER_FALLBACK_KEY = "source_filter_fallback"
 CUDA_HF_TOKEN_ENV = "VLLM_CI_HF_TOKEN"
 CUDA_HF_TOKEN_EXPORT = f'if [ -n "$${{{CUDA_HF_TOKEN_ENV}:-}}" ]; then export HF_TOKEN="$${{{CUDA_HF_TOKEN_ENV}}}"; fi'
+CUDA_RUNTIME_CHECK = "python3 .buildkite/cuda/scripts/check_vllm_runtime.py"
 
 # Bootstrap Buildkite ``if`` expressions.
 # ``*_MAIN_IF``: main + env schedule. ``*_LABEL_IF``: PR label (and/or composed with MAIN).
@@ -620,11 +621,11 @@ def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None:
     if any(preset_name == chip or preset_name.startswith(f"{chip}_") for chip in _cuda_mirror_chips()):
         commands = merged.get("commands")
         if isinstance(commands, list):
-            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, *commands]
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK, *commands]
         elif commands is None:
-            merged["commands"] = [CUDA_HF_TOKEN_EXPORT]
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK]
         else:
-            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, commands]
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK, commands]
     # Preset retry (K8S_RETRY on l4_*) must not clobber a step that opted out.
     if "retry" in step:
         merged["retry"] = step["retry"]

.buildkite/cuda/scripts/check_vllm_runtime.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Verify the pinned vLLM and compiled CUDA imports on a GPU CI worker."""
+
+from __future__ import annotations
+
+import importlib
+import re
+from pathlib import Path
+
+
+def main() -> None:
+    dockerfile = Path(__file__).resolve().parents[3] / "docker" / "Dockerfile.ci"
+    match = re.search(
+        r"^ARG VLLM_PRECOMPILED_WHEEL_COMMIT=([0-9a-f]{40})$",
+        dockerfile.read_text(),
+        re.MULTILINE,
+    )
+    if match is None:
+        raise RuntimeError(f"Missing vLLM commit pin in {dockerfile}")
+    expected_commit = match.group(1)
+
+    import torch
+    import vllm
+
+    assert expected_commit[:8] in vllm.__version__, vllm.__version__
+    assert torch.cuda.is_available(), "vLLM runtime check requires a GPU worker"
+    for module in ("flashinfer", "vllm._C_stable_libtorch"):
+        importlib.import_module(module)
+    print(
+        f"Verified vLLM CUDA runtime: vllm={vllm.__version__} "
+        f"commit={expected_commit} torch={torch.__version__} "
+        f"cuda={torch.version.cuda} device={torch.cuda.get_device_name()}"
+    )
+
+
+if __name__ == "__main__":
+    main()

docker/Dockerfile.ci

@@ -44,9 +44,10 @@ ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
 # Do NOT use a non-editable install here — that copies an empty stub into
 # site-packages and breaks ``import vllm_omni.entrypoints``.
 RUN uv pip install --system -e ".[dev]"
-# Fail the image build if dependency resolution replaced the selected vLLM
-# or its compiled CUDA runtime no longer imports.
-RUN python3 -c "import vllm, torch, flashinfer; import vllm._C_stable_libtorch; assert '${VLLM_PRECOMPILED_WHEEL_COMMIT}'[:8] in vllm.__version__, vllm.__version__"
+# Fail the image build if dependency resolution replaced the selected vLLM.
+# The CPU builder has no NVIDIA driver; CUDA imports are checked at the start
+# of each GPU test job by .buildkite/cuda/scripts/check_vllm_runtime.py.
+RUN python3 -c "import vllm; assert '${VLLM_PRECOMPILED_WHEEL_COMMIT}'[:8] in vllm.__version__, vllm.__version__"
 
 # The final stage just copies the working dir on top of the
 # dependencies image. This uses --link so that we can skip

tests/buildkite/test_upload_pipeline.py

@@ -172,7 +172,11 @@ def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins(monkeypatch: pytest
     env_names = {item["name"] for item in container["env"]}
     assert "VLLM_CI_HF_TOKEN" in env_names
     assert "HF_TOKEN" not in env_names
-    assert step["commands"] == [CUDA_HF_TOKEN_EXPORT, "pytest -sv tests/example"]
+    assert step["commands"] == [
+        CUDA_HF_TOKEN_EXPORT,
+        "python3 .buildkite/cuda/scripts/check_vllm_runtime.py",
+        "pytest -sv tests/example",
+    ]
 
 
 def test_mirror_hardwares_l4_preserves_explicit_retry() -> None:
@@ -239,6 +243,10 @@ def test_all_cuda_mirror_hardwares_restore_hf_token_at_runtime() -> None:
         assert "VLLM_CI_HF_TOKEN" in env_names
         assert "HF_TOKEN" not in env_names
         assert step["commands"][0] == CUDA_HF_TOKEN_EXPORT
+        assert step["commands"][1:] == [
+            "python3 .buildkite/cuda/scripts/check_vllm_runtime.py",
+            "pytest -sv tests/example",
+        ]
 
 
 def _gpu_limit(step: dict) -> int: