From 701a0e0c2551bd56ae9d3b35540af4928e8661bf Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad Date: Mon, 13 Apr 2026 03:30:41 +0800 Subject: [PATCH] [CI/Docker] Clean up redundant flashinfer cubin downloads (#22491) --- docker/Dockerfile | 5 +- .../ci/cuda/ci_download_flashinfer_cubin.sh | 62 ------------------- scripts/ci/cuda/ci_install_dependency.sh | 4 +- 3 files changed, 2 insertions(+), 69 deletions(-) delete mode 100755 scripts/ci/cuda/ci_download_flashinfer_cubin.sh diff --git a/docker/Dockerfile b/docker/Dockerfile index 2287f9e15..487594aa2 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -302,9 +302,7 @@ FROM torch_deps AS flashinfer_cache ARG CUDA_VERSION ARG INSTALL_FLASHINFER_JIT_CACHE ARG FLASHINFER_VERSION -ARG BUILD_AND_DOWNLOAD_PARALLEL -# Download flashinfer cubins # Stage jit-cache artifacts into /flashinfer_jit_output for clean COPY later RUN --mount=type=cache,target=/root/.cache/pip \ case "$CUDA_VERSION" in \ @@ -319,8 +317,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \ python3 -m pip install flashinfer-jit-cache==${FLASHINFER_VERSION} --index-url https://flashinfer.ai/whl/cu${CUINDEX} \ && cp -r /usr/local/lib/python3.12/dist-packages/flashinfer_jit_cache /flashinfer_jit_output/ \ && cp -r /usr/local/lib/python3.12/dist-packages/flashinfer_jit_cache-*.dist-info /flashinfer_jit_output/ ; \ - fi \ - && FLASHINFER_CUBIN_DOWNLOAD_THREADS=${BUILD_AND_DOWNLOAD_PARALLEL} FLASHINFER_LOGGING_LEVEL=warning python3 -m flashinfer --download-cubin + fi ######################################################## # PARALLEL STAGE 4: Dev Tools Builder (starts from base) diff --git a/scripts/ci/cuda/ci_download_flashinfer_cubin.sh b/scripts/ci/cuda/ci_download_flashinfer_cubin.sh deleted file mode 100755 index 153f93d91..000000000 --- a/scripts/ci/cuda/ci_download_flashinfer_cubin.sh +++ /dev/null @@ -1,62 +0,0 @@ -#!/bin/bash -# Download flashinfer cubins if the local set is incomplete. -# -# The flashinfer-cubin pip package may not include cubins for newer architectures -# (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. -# -# 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 - -# 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) -downloaded = sum(1 for _, exists in status if exists) -print(f'{downloaded}/{total}') -" 2>/dev/null) || CUBIN_STATUS="unknown" - -echo "Flashinfer cubin status: ${CUBIN_STATUS}" - -if echo "$CUBIN_STATUS" | grep -qE '^[0-9]+/[0-9]+$'; then - CUBIN_DOWNLOADED="${CUBIN_STATUS%/*}" - CUBIN_TOTAL="${CUBIN_STATUS#*/}" - if [ "$CUBIN_DOWNLOADED" = "$CUBIN_TOTAL" ] && [ "$CUBIN_TOTAL" != "0" ]; then - echo "All flashinfer cubins already present (${CUBIN_STATUS}), skipping download" - else - echo "Cubins incomplete (${CUBIN_STATUS}), downloading..." - 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 (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 c10a79e62..7735977ef 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/scripts/ci/cuda/ci_install_dependency.sh @@ -181,7 +181,7 @@ mark_step_done "Pip / uv toolchain & stale package cleanup" # Uninstall Flashinfer # ------------------------------------------------------------------------------ # 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-cubin: 150+ MB # - flashinfer-jit-cache: 1.2+ GB, by far the largest download in CI FLASHINFER_PYTHON_REQUIRED=$(grep -Po -m1 '(?<=flashinfer_python==)[0-9A-Za-z\.\-]+' python/pyproject.toml || echo "") FLASHINFER_CUBIN_REQUIRED=$(grep -Po -m1 '(?<=flashinfer_cubin==)[0-9A-Za-z\.\-]+' python/pyproject.toml || echo "") @@ -290,8 +290,6 @@ UNINSTALL_JIT_CACHE="$UNINSTALL_JIT_CACHE" \ PIP_CMD="$PIP_CMD" \ PIP_INSTALL_SUFFIX="$PIP_INSTALL_SUFFIX" \ bash "${SCRIPT_DIR}/ci_download_flashinfer_jit_cache.sh" -# Download flashinfer cubins -bash "${SCRIPT_DIR}/ci_download_flashinfer_cubin.sh" mark_step_done "Download flashinfer artifacts"