Refactor flashinfer version retrieval in Dockerfile.ci for improved accuracy and consistency in package installations

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

docker/Dockerfile.ci

@@ -21,18 +21,9 @@ RUN apt-get update && \
 RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/4a1550d22d7058e129d0e1257e726b3bf4a77025
 
 # Keep flashinfer runtime packages in version sync.
-# We derive the version from installed flashinfer-python (pulled by vllm)
+# We derive the version from installed flashinfer (pulled by vllm)
 # and enforce the same version for flashinfer-cubin / flashinfer-jit-cache.
-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")')" && \
+RUN FLASHINFER_VERSION="$(python3 -c 'import importlib.metadata as m; print(next(m.version(n) for n in ("flashinfer","flashinfer-python") if m.packages_distributions().get(n)))')" && \
     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}" \