From e897640b439d3aa3e88c105684e357a77d6bacc0 Mon Sep 17 00:00:00 2001
From: tzhouam <tzhouam@connect.ust.hk>
Date: Wed, 11 Feb 2026 04:11:11 +0000
Subject: [PATCH] Update Dockerfile.ci to enhance flashinfer version retrieval
 logic and ensure package installation consistency

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
---
 docker/Dockerfile.ci | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci
index 4c49b9304..b6df5c223 100644
--- a/docker/Dockerfile.ci
+++ b/docker/Dockerfile.ci
@@ -23,13 +23,22 @@ RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm
 # Keep flashinfer runtime packages in version sync.
 # We derive the version from installed flashinfer-python (pulled by vllm)
 # and enforce the same version for flashinfer-cubin / flashinfer-jit-cache.
-# RUN FLASHINFER_VERSION="$(python3 -c 'import importlib.metadata as m; print(m.version("flashinfer-python"))')" && \
-#     FLASHINFER_CUDA_TAG="$(python3 -c 'import torch; print((torch.version.cuda or "12.4").replace(".", ""))')" && \
-#     uv pip install --system --upgrade --force-reinstall \
-#       "flashinfer-cubin==${FLASHINFER_VERSION}" \
-#       "flashinfer-jit-cache==${FLASHINFER_VERSION}" \
-#       --extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}" && \
-#     flashinfer show-config
+RUN FLASHINFER_VERSION="$(python3 -c 'import importlib.metadata as m; \
+for n in ("flashinfer", "flashinfer-python"): \
+    \
+    try: \
+        print(m.version(n)); \
+        break; \
+    except m.PackageNotFoundError: \
+        pass \
+else: \
+    raise SystemExit("flashinfer package not found")')" && \
+    FLASHINFER_CUDA_TAG="$(python3 -c 'import torch; print((torch.version.cuda or "12.4").replace(".", ""))')" && \
+    uv pip install --system --upgrade --force-reinstall \
+      "flashinfer-cubin==${FLASHINFER_VERSION}" \
+      "flashinfer-jit-cache==${FLASHINFER_VERSION}" \
+      --extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}" && \
+    flashinfer show-config
 
 # Install vllm-omni into the same uv-managed Python environment used by the base image.
 RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"