[CI] Skip flashinfer-cubin reinstall when version matches (#19470)
Co-authored-by: Alison Shao <alisonshao@MacBook-Pro-D2W773R9CD.local> Co-authored-by: Alison Shao <alisonshao@Mac.attlocal.net> Co-authored-by: Alison Shao <alisonshao@mac.lan>
This commit is contained in:
Alison Shao
co-authored by
Alison Shao
Alison Shao
Alison Shao
Alison Shao
Alison Shao
Alison Shao
parent
b3cfad0a80
commit
1c1712d8e5
@@ -5,9 +5,35 @@
|
||||
# (e.g. sm_100, sm_120) due to PyPI size limits. This script checks the local
|
||||
# cubin status against the flashinfer artifact repository and downloads any
|
||||
# missing files.
|
||||
set -euxo pipefail
|
||||
#
|
||||
# This script is best-effort: if the status check or download times out (e.g.
|
||||
# due to a GPU in error state blocking CUDA init), we warn and continue.
|
||||
# The pip package already includes cubins for common architectures (sm_80, sm_90).
|
||||
set -uxo pipefail
|
||||
|
||||
CUBIN_STATUS=$(FLASHINFER_LOGGING_LEVEL=warning python3 -c "
|
||||
# Early exit: the pip package already includes cubins for sm_80 and sm_90.
|
||||
# Only sm_100+ (Blackwell) needs extra cubins downloaded. Skip the expensive
|
||||
# Python status check entirely if no such GPU is present.
|
||||
if COMPUTE_CAPS=$(timeout 10 nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null); then
|
||||
NEEDS_EXTRA_CUBINS=false
|
||||
while IFS= read -r cap; do
|
||||
major="${cap%%.*}"
|
||||
if [ "$major" -ge 10 ] 2>/dev/null; then
|
||||
NEEDS_EXTRA_CUBINS=true
|
||||
break
|
||||
fi
|
||||
done <<< "$COMPUTE_CAPS"
|
||||
if [ "$NEEDS_EXTRA_CUBINS" = false ]; then
|
||||
echo "All GPUs are sm_9x or older (compute caps: $(echo $COMPUTE_CAPS | tr '\n' ' ')), pip cubins sufficient — skipping download"
|
||||
exit 0
|
||||
fi
|
||||
fi
|
||||
|
||||
# Use timeout to prevent hangs when GPUs are in error state (the flashinfer
|
||||
# import can trigger CUDA init which blocks on bad GPUs).
|
||||
CUBIN_STATUS=$(timeout 60 python3 -c "
|
||||
import os
|
||||
os.environ.setdefault('CUDA_VISIBLE_DEVICES', '')
|
||||
from flashinfer.artifacts import get_artifacts_status
|
||||
status = get_artifacts_status()
|
||||
total = len(status)
|
||||
@@ -24,9 +50,13 @@ if echo "$CUBIN_STATUS" | grep -qE '^[0-9]+/[0-9]+$'; then
|
||||
echo "All flashinfer cubins already present (${CUBIN_STATUS}), skipping download"
|
||||
else
|
||||
echo "Cubins incomplete (${CUBIN_STATUS}), downloading..."
|
||||
FLASHINFER_LOGGING_LEVEL=warning python3 -m flashinfer --download-cubin
|
||||
if ! timeout 300 env FLASHINFER_LOGGING_LEVEL=warning python3 -m flashinfer --download-cubin; then
|
||||
echo "WARNING: flashinfer cubin download failed or timed out, continuing with existing cubins"
|
||||
fi
|
||||
fi
|
||||
else
|
||||
echo "Could not determine cubin status, downloading as fallback..."
|
||||
FLASHINFER_LOGGING_LEVEL=warning python3 -m flashinfer --download-cubin
|
||||
echo "Could not determine cubin status (status check timed out or failed), attempting download..."
|
||||
if ! timeout 300 env FLASHINFER_LOGGING_LEVEL=warning python3 -m flashinfer --download-cubin; then
|
||||
echo "WARNING: flashinfer cubin download failed or timed out, continuing with existing cubins"
|
||||
fi
|
||||
fi
|
||||
|
||||
@@ -118,7 +118,35 @@ fi
|
||||
|
||||
# Clean up existing installations
|
||||
$PIP_UNINSTALL_CMD sgl-kernel sglang $PIP_UNINSTALL_SUFFIX || true
|
||||
$PIP_UNINSTALL_CMD flashinfer-python flashinfer-cubin flashinfer-jit-cache $PIP_UNINSTALL_SUFFIX || true
|
||||
# Keep flashinfer packages installed if version matches to avoid re-downloading:
|
||||
# - flashinfer-cubin: 150+ MB, plus extra cubins from ci_download_flashinfer_cubin.sh
|
||||
# - flashinfer-jit-cache: 1.2+ GB, by far the largest download in CI
|
||||
FLASHINFER_CUBIN_REQUIRED=$(grep -Po -m1 '(?<=flashinfer_cubin==)[0-9A-Za-z\.\-]+' python/pyproject.toml || echo "")
|
||||
FLASHINFER_CUBIN_INSTALLED=$(pip show flashinfer-cubin 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
|
||||
FLASHINFER_JIT_INSTALLED=$(pip show flashinfer-jit-cache 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//' || echo "")
|
||||
|
||||
UNINSTALL_CUBIN=true
|
||||
UNINSTALL_JIT_CACHE=true
|
||||
|
||||
if [ "$FLASHINFER_CUBIN_INSTALLED" = "$FLASHINFER_CUBIN_REQUIRED" ] && [ -n "$FLASHINFER_CUBIN_REQUIRED" ]; then
|
||||
echo "flashinfer-cubin==${FLASHINFER_CUBIN_REQUIRED} already installed, keeping it"
|
||||
UNINSTALL_CUBIN=false
|
||||
else
|
||||
echo "flashinfer-cubin version mismatch (installed: ${FLASHINFER_CUBIN_INSTALLED:-none}, required: ${FLASHINFER_CUBIN_REQUIRED}), reinstalling"
|
||||
fi
|
||||
|
||||
if [ "$FLASHINFER_JIT_INSTALLED" = "$FLASHINFER_VERSION" ] && [ -n "$FLASHINFER_VERSION" ]; then
|
||||
echo "flashinfer-jit-cache==${FLASHINFER_VERSION} already installed, keeping it"
|
||||
UNINSTALL_JIT_CACHE=false
|
||||
else
|
||||
echo "flashinfer-jit-cache version mismatch (installed: ${FLASHINFER_JIT_INSTALLED:-none}, required: ${FLASHINFER_VERSION}), will reinstall"
|
||||
fi
|
||||
|
||||
# Build uninstall list based on what needs updating
|
||||
FLASHINFER_UNINSTALL="flashinfer-python"
|
||||
[ "$UNINSTALL_CUBIN" = true ] && FLASHINFER_UNINSTALL="$FLASHINFER_UNINSTALL flashinfer-cubin"
|
||||
[ "$UNINSTALL_JIT_CACHE" = true ] && FLASHINFER_UNINSTALL="$FLASHINFER_UNINSTALL flashinfer-jit-cache"
|
||||
$PIP_UNINSTALL_CMD $FLASHINFER_UNINSTALL $PIP_UNINSTALL_SUFFIX || true
|
||||
$PIP_UNINSTALL_CMD opencv-python opencv-python-headless $PIP_UNINSTALL_SUFFIX || true
|
||||
|
||||
# Install the main package
|
||||
@@ -236,52 +264,60 @@ fi
|
||||
$PIP_CMD uninstall xformers || true
|
||||
|
||||
# Install flashinfer-jit-cache with caching and retry logic (flashinfer.ai can have transient DNS issues)
|
||||
# Cache directory for flashinfer wheels (persists across CI runs on self-hosted runners)
|
||||
FLASHINFER_CACHE_DIR="${HOME}/.cache/flashinfer-wheels"
|
||||
mkdir -p "${FLASHINFER_CACHE_DIR}"
|
||||
|
||||
# Clean up old versions to avoid cache bloat
|
||||
find "${FLASHINFER_CACHE_DIR}" -name "flashinfer_jit_cache-*.whl" ! -name "flashinfer_jit_cache-${FLASHINFER_VERSION}*" -type f -delete 2>/dev/null || true
|
||||
|
||||
FLASHINFER_WHEEL_PATTERN="flashinfer_jit_cache-${FLASHINFER_VERSION}*.whl"
|
||||
CACHED_WHEEL=$(find "${FLASHINFER_CACHE_DIR}" -name "${FLASHINFER_WHEEL_PATTERN}" -type f 2>/dev/null | head -n 1)
|
||||
|
||||
# The jit-cache wheel is 1.2+ GB, so we skip the download entirely if already installed.
|
||||
FLASHINFER_INSTALLED=false
|
||||
|
||||
# Try to install from cache first
|
||||
if [ -n "$CACHED_WHEEL" ] && [ -f "$CACHED_WHEEL" ]; then
|
||||
echo "Found cached flashinfer wheel: $CACHED_WHEEL"
|
||||
if $PIP_CMD install "$CACHED_WHEEL" $PIP_INSTALL_SUFFIX; then
|
||||
FLASHINFER_INSTALLED=true
|
||||
echo "Successfully installed flashinfer-jit-cache from cache"
|
||||
else
|
||||
echo "Failed to install from cache, will try downloading..."
|
||||
rm -f "$CACHED_WHEEL"
|
||||
fi
|
||||
if [ "$UNINSTALL_JIT_CACHE" = false ]; then
|
||||
FLASHINFER_INSTALLED=true
|
||||
echo "flashinfer-jit-cache already at correct version, skipping download"
|
||||
fi
|
||||
|
||||
# If not installed from cache, download with retry logic
|
||||
if [ "$FLASHINFER_INSTALLED" = false ]; then
|
||||
for i in {1..5}; do
|
||||
# Download wheel to cache directory (use pip directly as uv pip doesn't support download)
|
||||
if pip download flashinfer-jit-cache==${FLASHINFER_VERSION} \
|
||||
--index-url https://flashinfer.ai/whl/${CU_VERSION} \
|
||||
-d "${FLASHINFER_CACHE_DIR}"; then
|
||||
# Cache directory for flashinfer wheels (persists across CI runs on self-hosted runners)
|
||||
FLASHINFER_CACHE_DIR="${HOME}/.cache/flashinfer-wheels"
|
||||
mkdir -p "${FLASHINFER_CACHE_DIR}"
|
||||
|
||||
CACHED_WHEEL=$(find "${FLASHINFER_CACHE_DIR}" -name "${FLASHINFER_WHEEL_PATTERN}" -type f 2>/dev/null | head -n 1)
|
||||
if [ -n "$CACHED_WHEEL" ] && [ -f "$CACHED_WHEEL" ]; then
|
||||
if $PIP_CMD install "$CACHED_WHEEL" $PIP_INSTALL_SUFFIX; then
|
||||
FLASHINFER_INSTALLED=true
|
||||
echo "Successfully downloaded and installed flashinfer-jit-cache"
|
||||
break
|
||||
fi
|
||||
else
|
||||
echo "Warning: Download succeeded but wheel file not found"
|
||||
fi
|
||||
# Clean up old versions to avoid cache bloat
|
||||
find "${FLASHINFER_CACHE_DIR}" -name "flashinfer_jit_cache-*.whl" ! -name "flashinfer_jit_cache-${FLASHINFER_VERSION}*" -type f -delete 2>/dev/null || true
|
||||
|
||||
FLASHINFER_WHEEL_PATTERN="flashinfer_jit_cache-${FLASHINFER_VERSION}*.whl"
|
||||
CACHED_WHEEL=$(find "${FLASHINFER_CACHE_DIR}" -name "${FLASHINFER_WHEEL_PATTERN}" -type f 2>/dev/null | head -n 1)
|
||||
|
||||
# Try to install from cache first
|
||||
if [ -n "$CACHED_WHEEL" ] && [ -f "$CACHED_WHEEL" ]; then
|
||||
echo "Found cached flashinfer wheel: $CACHED_WHEEL"
|
||||
if $PIP_CMD install "$CACHED_WHEEL" $PIP_INSTALL_SUFFIX; then
|
||||
FLASHINFER_INSTALLED=true
|
||||
echo "Successfully installed flashinfer-jit-cache from cache"
|
||||
else
|
||||
echo "Failed to install from cache, will try downloading..."
|
||||
rm -f "$CACHED_WHEEL"
|
||||
fi
|
||||
echo "Attempt $i to download flashinfer-jit-cache failed, retrying in 10 seconds..."
|
||||
sleep 10
|
||||
done
|
||||
fi
|
||||
|
||||
# If not installed from cache, download with retry logic
|
||||
if [ "$FLASHINFER_INSTALLED" = false ]; then
|
||||
for i in {1..5}; do
|
||||
# Download wheel to cache directory (use pip directly as uv pip doesn't support download)
|
||||
# Timeout after 10 minutes — the wheel is ~1.2 GB
|
||||
if timeout 600 pip download flashinfer-jit-cache==${FLASHINFER_VERSION} \
|
||||
--index-url https://flashinfer.ai/whl/${CU_VERSION} \
|
||||
-d "${FLASHINFER_CACHE_DIR}"; then
|
||||
|
||||
CACHED_WHEEL=$(find "${FLASHINFER_CACHE_DIR}" -name "${FLASHINFER_WHEEL_PATTERN}" -type f 2>/dev/null | head -n 1)
|
||||
if [ -n "$CACHED_WHEEL" ] && [ -f "$CACHED_WHEEL" ]; then
|
||||
if $PIP_CMD install "$CACHED_WHEEL" $PIP_INSTALL_SUFFIX; then
|
||||
FLASHINFER_INSTALLED=true
|
||||
echo "Successfully downloaded and installed flashinfer-jit-cache"
|
||||
break
|
||||
fi
|
||||
else
|
||||
echo "Warning: Download succeeded but wheel file not found"
|
||||
fi
|
||||
fi
|
||||
echo "Attempt $i to download flashinfer-jit-cache failed, retrying in 10 seconds..."
|
||||
sleep 10
|
||||
done
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ "$FLASHINFER_INSTALLED" = false ]; then
|
||||
|
||||
Reference in New Issue
Block a user