minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CI] Skip flashinfer-cubin reinstall when version matches (#19470)

Browse Source 
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
2026-03-05 13:30:44 -08:00
committed by GitHub
co-authored by Alison Shao Alison Shao Alison Shao
parent b3cfad0a80
commit 1c1712d8e5
2 changed files with 112 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/ci/cuda/ci_download_flashinfer_cubin.sh
+35 -5
View File
@@ -5,9 +5,35 @@
# (e.g. sm_100, sm_120) due to PyPI size limits. This script checks the local # (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 # cubin status against the flashinfer artifact repository and downloads any
# missing files. # 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 from flashinfer.artifacts import get_artifacts_status
status = get_artifacts_status() status = get_artifacts_status()
total = len(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" echo "All flashinfer cubins already present (${CUBIN_STATUS}), skipping download"
else else
echo "Cubins incomplete (${CUBIN_STATUS}), downloading..." 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 fi
else else
echo "Could not determine cubin status, downloading as fallback..." echo "Could not determine cubin status (status check timed out or failed), attempting download..."
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 fi
scripts/ci/cuda/ci_install_dependency.sh
+77 -41
View File
@@ -118,7 +118,35 @@ fi
# Clean up existing installations # Clean up existing installations
$PIP_UNINSTALL_CMD sgl-kernel sglang $PIP_UNINSTALL_SUFFIX || true $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 $PIP_UNINSTALL_CMD opencv-python opencv-python-headless $PIP_UNINSTALL_SUFFIX || true
# Install the main package # Install the main package
@@ -236,52 +264,60 @@ fi
$PIP_CMD uninstall xformers || true $PIP_CMD uninstall xformers || true
# Install flashinfer-jit-cache with caching and retry logic (flashinfer.ai can have transient DNS issues) # 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) # The jit-cache wheel is 1.2+ GB, so we skip the download entirely if already installed.
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)
FLASHINFER_INSTALLED=false FLASHINFER_INSTALLED=false
if [ "$UNINSTALL_JIT_CACHE" = false ]; then
# Try to install from cache first FLASHINFER_INSTALLED=true
if [ -n "$CACHED_WHEEL" ] && [ -f "$CACHED_WHEEL" ]; then echo "flashinfer-jit-cache already at correct version, skipping download"
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
fi fi
# If not installed from cache, download with retry logic
if [ "$FLASHINFER_INSTALLED" = false ]; then if [ "$FLASHINFER_INSTALLED" = false ]; then
for i in {1..5}; do # Cache directory for flashinfer wheels (persists across CI runs on self-hosted runners)
# Download wheel to cache directory (use pip directly as uv pip doesn't support download) FLASHINFER_CACHE_DIR="${HOME}/.cache/flashinfer-wheels"
if pip download flashinfer-jit-cache==${FLASHINFER_VERSION} \ mkdir -p "${FLASHINFER_CACHE_DIR}"
--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) # Clean up old versions to avoid cache bloat
if [ -n "$CACHED_WHEEL" ] && [ -f "$CACHED_WHEEL" ]; then find "${FLASHINFER_CACHE_DIR}" -name "flashinfer_jit_cache-*.whl" ! -name "flashinfer_jit_cache-${FLASHINFER_VERSION}*" -type f -delete 2>/dev/null || true
if $PIP_CMD install "$CACHED_WHEEL" $PIP_INSTALL_SUFFIX; then
FLASHINFER_INSTALLED=true FLASHINFER_WHEEL_PATTERN="flashinfer_jit_cache-${FLASHINFER_VERSION}*.whl"
echo "Successfully downloaded and installed flashinfer-jit-cache" CACHED_WHEEL=$(find "${FLASHINFER_CACHE_DIR}" -name "${FLASHINFER_WHEEL_PATTERN}" -type f 2>/dev/null | head -n 1)
break
fi # Try to install from cache first
else if [ -n "$CACHED_WHEEL" ] && [ -f "$CACHED_WHEEL" ]; then
echo "Warning: Download succeeded but wheel file not found" echo "Found cached flashinfer wheel: $CACHED_WHEEL"
fi 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 fi
echo "Attempt $i to download flashinfer-jit-cache failed, retrying in 10 seconds..." fi
sleep 10
done # 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 fi
if [ "$FLASHINFER_INSTALLED" = false ]; then if [ "$FLASHINFER_INSTALLED" = false ]; then