From 1c1712d8e54124c1c0b6c6cd89165a77970523ab Mon Sep 17 00:00:00 2001 From: Alison Shao <54658187+alisonshao@users.noreply.github.com> Date: Thu, 5 Mar 2026 13:30:44 -0800 Subject: [PATCH] [CI] Skip flashinfer-cubin reinstall when version matches (#19470) Co-authored-by: Alison Shao Co-authored-by: Alison Shao Co-authored-by: Alison Shao --- .../ci/cuda/ci_download_flashinfer_cubin.sh | 40 +++++- scripts/ci/cuda/ci_install_dependency.sh | 118 ++++++++++++------ 2 files changed, 112 insertions(+), 46 deletions(-) diff --git a/scripts/ci/cuda/ci_download_flashinfer_cubin.sh b/scripts/ci/cuda/ci_download_flashinfer_cubin.sh index 6241cc78b..153f93d91 100755 --- a/scripts/ci/cuda/ci_download_flashinfer_cubin.sh +++ b/scripts/ci/cuda/ci_download_flashinfer_cubin.sh @@ -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 diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh index 3d47af3f7..5e6dc3a3a 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/scripts/ci/cuda/ci_install_dependency.sh @@ -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