ci: refactor CUDA dependency install script (#21017)
This commit is contained in:
Lianmin Zheng
8
.github/workflows/pr-test.yml
vendored
8
.github/workflows/pr-test.yml
vendored
@@ -915,10 +915,16 @@ jobs:
|
||||
with:
|
||||
python-version: '3.10'
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
|
||||
# uv pip targets a venv by default; setup-python has no venv — install into that interpreter (see UV_SYSTEM_PYTHON in https://docs.astral.sh/uv/guides/integration/github/)
|
||||
- name: Install dependencies
|
||||
timeout-minutes: 20
|
||||
env:
|
||||
UV_SYSTEM_PYTHON: "1"
|
||||
run: |
|
||||
pip install -e "python/[dev]"
|
||||
uv pip install -e "python[dev]" --index-strategy unsafe-best-match --prerelease allow
|
||||
|
||||
- name: Run test
|
||||
timeout-minutes: 10
|
||||
|
||||
71
scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh
Executable file
71
scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh
Executable file
@@ -0,0 +1,71 @@
|
||||
#!/bin/bash
|
||||
# Install flashinfer-jit-cache with caching and retry logic (flashinfer.ai can have transient DNS issues).
|
||||
# The jit-cache wheel is 1.2+ GB, so we skip the download entirely if already installed.
|
||||
#
|
||||
# Required environment (caller must export or set):
|
||||
# UNINSTALL_JIT_CACHE — literal true/false (skip download when false)
|
||||
# FLASHINFER_PYTHON_REQUIRED — e.g. from python/pyproject.toml (flashinfer_python)
|
||||
# CU_VERSION — e.g. cu129
|
||||
# PIP_CMD — e.g. "pip" or "uv pip"
|
||||
# PIP_INSTALL_SUFFIX — extra pip args for this runner
|
||||
set -euxo pipefail
|
||||
|
||||
: "${UNINSTALL_JIT_CACHE:?must be set}"
|
||||
: "${FLASHINFER_PYTHON_REQUIRED:?must be set}"
|
||||
: "${CU_VERSION:?must be set}"
|
||||
: "${PIP_CMD:?must be set}"
|
||||
|
||||
FLASHINFER_JIT_CACHE_INSTALLED=false
|
||||
if [ "$UNINSTALL_JIT_CACHE" = false ]; then
|
||||
FLASHINFER_JIT_CACHE_INSTALLED=true
|
||||
echo "flashinfer-jit-cache already at correct version, skipping download"
|
||||
fi
|
||||
|
||||
if [ "$FLASHINFER_JIT_CACHE_INSTALLED" = false ]; then
|
||||
FLASHINFER_CACHE_DIR="${HOME}/.cache/flashinfer-wheels"
|
||||
mkdir -p "${FLASHINFER_CACHE_DIR}"
|
||||
|
||||
find "${FLASHINFER_CACHE_DIR}" -name "flashinfer_jit_cache-*.whl" ! -name "flashinfer_jit_cache-${FLASHINFER_PYTHON_REQUIRED}*" -type f -delete 2>/dev/null || true
|
||||
|
||||
FLASHINFER_WHEEL_PATTERN="flashinfer_jit_cache-${FLASHINFER_PYTHON_REQUIRED}*.whl"
|
||||
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
|
||||
echo "Found cached flashinfer wheel: $CACHED_WHEEL"
|
||||
if $PIP_CMD install "$CACHED_WHEEL" $PIP_INSTALL_SUFFIX; then
|
||||
FLASHINFER_JIT_CACHE_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
|
||||
|
||||
if [ "$FLASHINFER_JIT_CACHE_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 timeout 600 pip download "flashinfer-jit-cache==${FLASHINFER_PYTHON_REQUIRED}" \
|
||||
--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_JIT_CACHE_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_JIT_CACHE_INSTALLED" = false ]; then
|
||||
echo "ERROR: Failed to install flashinfer-jit-cache after 5 attempts"
|
||||
exit 1
|
||||
fi
|
||||
@@ -1,37 +1,110 @@
|
||||
#!/bin/bash
|
||||
# Install the dependency in CI.
|
||||
#
|
||||
# Structure (see section banners below):
|
||||
# - Configuration & timing
|
||||
# - Host / runner detection (arch, Blackwell, pip vs uv)
|
||||
# - Kill existing processes
|
||||
# - Install apt packages
|
||||
# - Python package site hygiene & install protoc
|
||||
# - Pip / uv toolchain & stale package cleanup
|
||||
# - Uninstall Flashinfer
|
||||
# - Install main package
|
||||
# - Install sglang-kernel
|
||||
# - Install sglang-router
|
||||
# - Download flashinfer artifacts
|
||||
# - Install extra dependency
|
||||
# - Fix other dependencies
|
||||
# - Prepare runner
|
||||
# - Verify imports
|
||||
set -euxo pipefail
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Configuration & timing
|
||||
# ------------------------------------------------------------------------------
|
||||
# Set up environment variables
|
||||
IS_BLACKWELL=${IS_BLACKWELL:-0}
|
||||
CU_VERSION="cu129"
|
||||
FLASHINFER_VERSION=0.6.6
|
||||
OPTIONAL_DEPS="${1:-}"
|
||||
|
||||
# Detect system architecture
|
||||
SECONDS=0
|
||||
_CI_MARK_PREV=${SECONDS}
|
||||
|
||||
mark_step_done() {
|
||||
local label=$1
|
||||
local now=${SECONDS}
|
||||
local step=$((now - _CI_MARK_PREV))
|
||||
printf '\n[STEP DONE] %s, step: %ss, total: %ss, date: %s\n' \
|
||||
"${label}" "${step}" "${now}" "$(date -u '+%Y-%m-%dT%H:%M:%SZ')"
|
||||
_CI_MARK_PREV=${now}
|
||||
}
|
||||
|
||||
mark_step_done "Configuration"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Host / runner detection (CPU arch, Blackwell, USE_UV)
|
||||
# ------------------------------------------------------------------------------
|
||||
# Detect CPU architecture (x86_64 or aarch64)
|
||||
ARCH=$(uname -m)
|
||||
echo "Detected architecture: ${ARCH}"
|
||||
|
||||
if [ "$CU_VERSION" = "cu130" ]; then
|
||||
NVRTC_SPEC="nvidia-cuda-nvrtc"
|
||||
# Detect GPU architecture (blackwell or not)
|
||||
if [ "${IS_BLACKWELL+set}" = set ]; then
|
||||
case "$IS_BLACKWELL" in 1 | true | yes) IS_BLACKWELL=1 ;; *) IS_BLACKWELL=0 ;; esac
|
||||
echo "IS_BLACKWELL=${IS_BLACKWELL} (manually set via environment)"
|
||||
else
|
||||
NVRTC_SPEC="nvidia-cuda-nvrtc-cu12"
|
||||
IS_BLACKWELL=0
|
||||
if command -v nvidia-smi >/dev/null 2>&1; then
|
||||
while IFS= read -r cap; do
|
||||
major="${cap%%.*}"
|
||||
if [ "${major:-0}" -ge 10 ] 2>/dev/null; then
|
||||
IS_BLACKWELL=1
|
||||
break
|
||||
fi
|
||||
done <<< "$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)"
|
||||
fi
|
||||
echo "IS_BLACKWELL=${IS_BLACKWELL} (auto-detected via nvidia-smi)"
|
||||
fi
|
||||
|
||||
# Whether to use pip or uv to install dependencies
|
||||
if [ "${USE_UV+set}" != set ]; then
|
||||
if [ "$IS_BLACKWELL" = "1" ]; then
|
||||
# Our current b200 runners have some issues with uv, so we default to pip
|
||||
# It is a runner specific issue, not a GPU architecture issue.
|
||||
USE_UV=false
|
||||
else
|
||||
USE_UV=true
|
||||
fi
|
||||
fi
|
||||
case "$(printf '%s' "$USE_UV" | tr '[:upper:]' '[:lower:]')" in 1 | true | yes) USE_UV=1 ;; *) USE_UV=0 ;; esac
|
||||
echo "USE_UV=${USE_UV}"
|
||||
|
||||
mark_step_done "Host / runner detection"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Kill existing processes
|
||||
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
|
||||
# ------------------------------------------------------------------------------
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
bash "${SCRIPT_DIR}/../../killall_sglang.sh"
|
||||
echo "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-}"
|
||||
|
||||
mark_step_done "Kill existing processes"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install apt packages
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install apt packages (including python3/pip which may be missing on some runners)
|
||||
# Use --no-install-recommends and ignore errors from unrelated broken packages on the runner
|
||||
# The NVIDIA driver packages may have broken dependencies that are unrelated to these packages
|
||||
# Run apt-get update first to refresh package index (stale index causes 404 on security.ubuntu.com)
|
||||
apt-get update || true
|
||||
apt-get install -y --no-install-recommends python3 python3-pip python3-venv python3-dev git libnuma-dev libssl-dev pkg-config libibverbs-dev libibverbs1 ibverbs-providers ibverbs-utils ffmpeg libavcodec-dev libavformat-dev libavutil-dev libswscale-dev || {
|
||||
CI_APT_PACKAGES=(
|
||||
python3 python3-pip python3-venv python3-dev git libnuma-dev libssl-dev pkg-config
|
||||
libibverbs-dev libibverbs1 ibverbs-providers ibverbs-utils
|
||||
ffmpeg libavcodec-dev libavformat-dev libavutil-dev libswscale-dev
|
||||
)
|
||||
apt-get install -y --no-install-recommends "${CI_APT_PACKAGES[@]}" || {
|
||||
echo "Warning: apt-get install failed, checking if required packages are available..."
|
||||
# Verify the packages we need are actually installed
|
||||
for pkg in python3 python3-pip python3-venv python3-dev git libnuma-dev libssl-dev pkg-config libibverbs-dev libibverbs1 ibverbs-providers ibverbs-utils ffmpeg; do
|
||||
for pkg in "${CI_APT_PACKAGES[@]}"; do
|
||||
if ! dpkg -l "$pkg" 2>/dev/null | grep -q "^ii"; then
|
||||
echo "ERROR: Required package $pkg is not installed and apt-get failed"
|
||||
exit 1
|
||||
@@ -40,66 +113,19 @@ apt-get install -y --no-install-recommends python3 python3-pip python3-venv pyth
|
||||
echo "All required packages are already installed, continuing..."
|
||||
}
|
||||
|
||||
mark_step_done "Install apt packages"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Python package site hygiene & install protoc
|
||||
# ------------------------------------------------------------------------------
|
||||
# Clear torch compilation cache
|
||||
python3 -c 'import os, shutil, tempfile, getpass; cache_dir = os.environ.get("TORCHINDUCTOR_CACHE_DIR") or os.path.join(tempfile.gettempdir(), "torchinductor_" + getpass.getuser()); shutil.rmtree(cache_dir, ignore_errors=True)'
|
||||
|
||||
# Check if protoc of correct architecture is already installed
|
||||
if command -v protoc >/dev/null 2>&1; then
|
||||
if protoc --version >/dev/null 2>&1; then
|
||||
echo "protoc already installed: $(protoc --version)"
|
||||
else
|
||||
echo "protoc found but not runnable, reinstalling..."
|
||||
INSTALL_PROTOC=1
|
||||
fi
|
||||
else
|
||||
INSTALL_PROTOC=1
|
||||
fi
|
||||
|
||||
# Install protoc for router build (gRPC protobuf compilation)
|
||||
if [ "${INSTALL_PROTOC:-0}" = "1" ]; then
|
||||
# TODO: move this to a separate script
|
||||
echo "Installing protoc..."
|
||||
if command -v apt-get &> /dev/null; then
|
||||
# Ubuntu/Debian
|
||||
apt-get update || true # May fail due to unrelated broken packages
|
||||
apt-get install -y --no-install-recommends wget unzip gcc g++ perl make || {
|
||||
echo "Warning: apt-get install failed, checking if required packages are available..."
|
||||
for pkg in wget unzip gcc g++ perl make; do
|
||||
if ! dpkg -l "$pkg" 2>/dev/null | grep -q "^ii"; then
|
||||
echo "ERROR: Required package $pkg is not installed and apt-get failed"
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
echo "All required packages are already installed, continuing..."
|
||||
}
|
||||
elif command -v yum &> /dev/null; then
|
||||
# RHEL/CentOS
|
||||
yum update -y
|
||||
yum install -y wget unzip gcc gcc-c++ perl-core make
|
||||
fi
|
||||
|
||||
cd /tmp
|
||||
# Determine protoc architecture
|
||||
if [ "$ARCH" = "aarch64" ] || [ "$ARCH" = "arm64" ]; then
|
||||
PROTOC_ARCH="aarch_64"
|
||||
else
|
||||
PROTOC_ARCH="x86_64"
|
||||
fi
|
||||
PROTOC_ZIP="protoc-32.0-linux-${PROTOC_ARCH}.zip"
|
||||
wget https://github.com/protocolbuffers/protobuf/releases/download/v32.0/${PROTOC_ZIP}
|
||||
unzip -o ${PROTOC_ZIP} -d /usr/local
|
||||
rm ${PROTOC_ZIP}
|
||||
protoc --version
|
||||
cd -
|
||||
else
|
||||
echo "protoc already installed: $(protoc --version)"
|
||||
fi
|
||||
|
||||
# Remove broken dist-info directories (missing METADATA per PEP 376)
|
||||
SITE_PACKAGES=$(python3 -c "import site; print(site.getsitepackages()[0])")
|
||||
if [ -d "$SITE_PACKAGES" ]; then
|
||||
{ set +x; } 2>/dev/null
|
||||
find "$SITE_PACKAGES" -maxdepth 1 -name "*.dist-info" -type d | while read d; do
|
||||
find "$SITE_PACKAGES" -maxdepth 1 -name "*.dist-info" -type d | while read -r d; do
|
||||
if [ ! -f "$d/METADATA" ]; then
|
||||
echo "Removing broken dist-info: $d"
|
||||
rm -rf "$d"
|
||||
@@ -108,18 +134,23 @@ if [ -d "$SITE_PACKAGES" ]; then
|
||||
set -x
|
||||
fi
|
||||
|
||||
# Install uv (use python3 -m pip for robustness since some runners only have pip3)
|
||||
# Install protoc
|
||||
bash "${SCRIPT_DIR}/../utils/install_protoc.sh"
|
||||
|
||||
mark_step_done "Python package site hygiene & install protoc"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Pip / uv toolchain & stale package cleanup
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install pip and uv (use python3 -m pip for robustness since some runners only have pip3)
|
||||
python3 -m pip install --upgrade pip
|
||||
|
||||
if [ "$IS_BLACKWELL" = "1" ]; then
|
||||
# The blackwell CI runner has some issues with pip and uv,
|
||||
# so we can only use pip with `--break-system-packages`
|
||||
if [ "$USE_UV" = "0" ]; then
|
||||
PIP_CMD="pip"
|
||||
PIP_INSTALL_SUFFIX="--break-system-packages"
|
||||
PIP_UNINSTALL_CMD="pip uninstall -y"
|
||||
PIP_UNINSTALL_SUFFIX="--break-system-packages"
|
||||
else
|
||||
# In normal cases, we use uv, which is much faster than pip.
|
||||
pip install uv
|
||||
export UV_SYSTEM_PYTHON=true
|
||||
|
||||
@@ -132,9 +163,15 @@ fi
|
||||
# Clean up existing installations
|
||||
$PIP_UNINSTALL_CMD sgl-kernel sglang-kernel sglang sgl-fa4 flash-attn-4 $PIP_UNINSTALL_SUFFIX || true
|
||||
|
||||
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-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 "")
|
||||
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 "")
|
||||
@@ -149,11 +186,11 @@ 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"
|
||||
if [ "$FLASHINFER_JIT_INSTALLED" = "$FLASHINFER_PYTHON_REQUIRED" ] && [ -n "$FLASHINFER_PYTHON_REQUIRED" ]; then
|
||||
echo "flashinfer-jit-cache==${FLASHINFER_PYTHON_REQUIRED} 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"
|
||||
echo "flashinfer-jit-cache version mismatch (installed: ${FLASHINFER_JIT_INSTALLED:-none}, required: ${FLASHINFER_PYTHON_REQUIRED}), will reinstall"
|
||||
fi
|
||||
|
||||
# Build uninstall list based on what needs updating
|
||||
@@ -163,34 +200,24 @@ FLASHINFER_UNINSTALL="flashinfer-python"
|
||||
$PIP_UNINSTALL_CMD $FLASHINFER_UNINSTALL $PIP_UNINSTALL_SUFFIX || true
|
||||
$PIP_UNINSTALL_CMD opencv-python opencv-python-headless $PIP_UNINSTALL_SUFFIX || true
|
||||
|
||||
mark_step_done "Uninstall Flashinfer"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install main package
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install the main package
|
||||
EXTRAS="dev"
|
||||
if [ -n "$OPTIONAL_DEPS" ]; then
|
||||
EXTRAS="dev,${OPTIONAL_DEPS}"
|
||||
fi
|
||||
echo "Installing python extras: [${EXTRAS}]"
|
||||
|
||||
$PIP_CMD install -e "python[${EXTRAS}]" --extra-index-url https://download.pytorch.org/whl/${CU_VERSION} $PIP_INSTALL_SUFFIX
|
||||
|
||||
# Fix CUDA version mismatch between torch and torchaudio.
|
||||
# PyPI's torch 2.9.1 bundles cu128 but torchaudio from pytorch.org/cu129 uses cu129.
|
||||
# This mismatch causes torchaudio's C extension to fail loading, producing:
|
||||
# "partially initialized module 'torchaudio' has no attribute 'lib'"
|
||||
# We cannot replace torch with cu129 (breaks sgl_kernel ABI), so instead we reinstall
|
||||
# torchaudio/torchvision from an index matching torch's CUDA version.
|
||||
TORCH_CUDA_VER=$(python3 -c "import torch; v=torch.version.cuda; parts=v.split('.'); print(f'cu{parts[0]}{parts[1]}')")
|
||||
echo "Detected torch CUDA version: ${TORCH_CUDA_VER}"
|
||||
if [ "${TORCH_CUDA_VER}" != "${CU_VERSION}" ]; then
|
||||
# Pin versions to match what was installed by pyproject.toml (strip +cuXYZ suffix)
|
||||
TORCHAUDIO_VER=$(pip show torchaudio 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')
|
||||
TORCHVISION_VER=$(pip show torchvision 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')
|
||||
echo "Reinstalling torchaudio==${TORCHAUDIO_VER} torchvision==${TORCHVISION_VER} from ${TORCH_CUDA_VER} index to match torch..."
|
||||
$PIP_CMD install "torchaudio==${TORCHAUDIO_VER}" "torchvision==${TORCHVISION_VER}" --index-url "https://download.pytorch.org/whl/${TORCH_CUDA_VER}" --force-reinstall --no-deps $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
|
||||
# Install router for pd-disagg test
|
||||
$PIP_CMD install sglang-router $PIP_INSTALL_SUFFIX
|
||||
mark_step_done "Install main package"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install sglang-kernel
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install sgl-kernel
|
||||
SGL_KERNEL_VERSION_FROM_KERNEL=$(grep -Po '(?<=^version = ")[^"]*' sgl-kernel/pyproject.toml)
|
||||
SGL_KERNEL_VERSION_FROM_SRT=$(grep -Po -m1 '(?<=sglang-kernel==)[0-9A-Za-z\.\-]+' python/pyproject.toml)
|
||||
@@ -227,117 +254,106 @@ else
|
||||
fi
|
||||
fi
|
||||
|
||||
mark_step_done "Install sglang-kernel"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install sglang-router
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install router for pd-disagg test
|
||||
$PIP_CMD install sglang-router $PIP_INSTALL_SUFFIX
|
||||
|
||||
# Show current packages
|
||||
$PIP_CMD list
|
||||
|
||||
mark_step_done "Install sglang-router"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Download flashinfer artifacts
|
||||
# ------------------------------------------------------------------------------
|
||||
# Download flashinfer jit cache
|
||||
UNINSTALL_JIT_CACHE="$UNINSTALL_JIT_CACHE" \
|
||||
FLASHINFER_PYTHON_REQUIRED="$FLASHINFER_PYTHON_REQUIRED" \
|
||||
CU_VERSION="$CU_VERSION" \
|
||||
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"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install extra dependency
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install other python dependencies
|
||||
if [ "$CU_VERSION" = "cu130" ]; then
|
||||
NVRTC_SPEC="nvidia-cuda-nvrtc"
|
||||
else
|
||||
NVRTC_SPEC="nvidia-cuda-nvrtc-cu12"
|
||||
fi
|
||||
$PIP_CMD install mooncake-transfer-engine==0.3.9 "${NVRTC_SPEC}" py-spy scipy huggingface_hub[hf_xet] pytest $PIP_INSTALL_SUFFIX
|
||||
|
||||
# Install other test dependencies
|
||||
if [ "$IS_BLACKWELL" != "1" ]; then
|
||||
# For lmms_evals evaluating MMMU
|
||||
git clone --branch v0.5 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git
|
||||
$PIP_CMD install -e lmms-eval/ $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
|
||||
# DeepEP depends on nvshmem 3.4.5
|
||||
# On Blackwell machines, skip reinstall if correct version already installed to avoid race conditions
|
||||
if [ "$IS_BLACKWELL" = "1" ]; then
|
||||
INSTALLED_NVSHMEM=$(pip show nvidia-nvshmem-cu12 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
|
||||
if [ "$INSTALLED_NVSHMEM" = "3.4.5" ]; then
|
||||
echo "nvidia-nvshmem-cu12==3.4.5 already installed, skipping reinstall"
|
||||
else
|
||||
$PIP_CMD install nvidia-nvshmem-cu12==3.4.5 $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
else
|
||||
$PIP_CMD install nvidia-nvshmem-cu12==3.4.5 --force-reinstall $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
|
||||
# Cudnn with version less than 9.16.0.29 will cause performance regression on Conv3D kernel
|
||||
# On Blackwell machines, skip reinstall if correct version already installed to avoid race conditions
|
||||
if [ "$IS_BLACKWELL" = "1" ]; then
|
||||
INSTALLED_CUDNN=$(pip show nvidia-cudnn-cu12 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
|
||||
if [ "$INSTALLED_CUDNN" = "9.16.0.29" ]; then
|
||||
echo "nvidia-cudnn-cu12==9.16.0.29 already installed, skipping reinstall"
|
||||
else
|
||||
$PIP_CMD install nvidia-cudnn-cu12==9.16.0.29 $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
else
|
||||
$PIP_CMD install nvidia-cudnn-cu12==9.16.0.29 --force-reinstall $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
$PIP_CMD uninstall xformers || true
|
||||
|
||||
# Install flashinfer-jit-cache with caching and retry logic (flashinfer.ai can have transient DNS issues)
|
||||
# The jit-cache wheel is 1.2+ GB, so we skip the download entirely if already installed.
|
||||
FLASHINFER_INSTALLED=false
|
||||
if [ "$UNINSTALL_JIT_CACHE" = false ]; then
|
||||
FLASHINFER_INSTALLED=true
|
||||
echo "flashinfer-jit-cache already at correct version, skipping download"
|
||||
mark_step_done "Install extra dependency"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Fix other dependencies
|
||||
# ------------------------------------------------------------------------------
|
||||
# Fix CUDA version mismatch between torch and torchaudio.
|
||||
# PyPI's torch 2.9.1 bundles cu128 but torchaudio from pytorch.org/cu129 uses cu129.
|
||||
# This mismatch causes torchaudio's C extension to fail loading, producing:
|
||||
# "partially initialized module 'torchaudio' has no attribute 'lib'"
|
||||
# We cannot replace torch with cu129 (breaks sgl_kernel ABI), so instead we reinstall
|
||||
# torchaudio/torchvision from an index matching torch's CUDA version.
|
||||
TORCH_CUDA_VER=$(python3 -c "import torch; v=torch.version.cuda; parts=v.split('.'); print(f'cu{parts[0]}{parts[1]}')")
|
||||
echo "Detected torch CUDA version: ${TORCH_CUDA_VER}"
|
||||
if [ "${TORCH_CUDA_VER}" != "${CU_VERSION}" ]; then
|
||||
# Pin versions to match what was installed by pyproject.toml (strip +cuXYZ suffix)
|
||||
TORCHAUDIO_VER=$(pip show torchaudio 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')
|
||||
TORCHVISION_VER=$(pip show torchvision 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')
|
||||
echo "Reinstalling torchaudio==${TORCHAUDIO_VER} torchvision==${TORCHVISION_VER} from ${TORCH_CUDA_VER} index to match torch..."
|
||||
$PIP_CMD install "torchaudio==${TORCHAUDIO_VER}" "torchvision==${TORCHVISION_VER}" --index-url "https://download.pytorch.org/whl/${TORCH_CUDA_VER}" --force-reinstall --no-deps $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
|
||||
if [ "$FLASHINFER_INSTALLED" = false ]; 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}"
|
||||
|
||||
# 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
|
||||
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
|
||||
# Fix dependencies: DeepEP depends on nvshmem 3.4.5 — skip reinstall when already correct (avoids pip races / wasted work)
|
||||
INSTALLED_NVSHMEM=$(pip show nvidia-nvshmem-cu12 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
|
||||
if [ "$INSTALLED_NVSHMEM" = "3.4.5" ]; then
|
||||
echo "nvidia-nvshmem-cu12==3.4.5 already installed, skipping reinstall"
|
||||
else
|
||||
$PIP_CMD install nvidia-nvshmem-cu12==3.4.5 $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
|
||||
if [ "$FLASHINFER_INSTALLED" = false ]; then
|
||||
echo "ERROR: Failed to install flashinfer-jit-cache after 5 attempts"
|
||||
exit 1
|
||||
# Fix dependencies: Cudnn with version less than 9.16.0.29 will cause performance regression on Conv3D kernel
|
||||
INSTALLED_CUDNN=$(pip show nvidia-cudnn-cu12 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
|
||||
if [ "$INSTALLED_CUDNN" = "9.16.0.29" ]; then
|
||||
echo "nvidia-cudnn-cu12==9.16.0.29 already installed, skipping reinstall"
|
||||
else
|
||||
$PIP_CMD install nvidia-cudnn-cu12==9.16.0.29 $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
|
||||
# Download flashinfer cubins if the local set is incomplete
|
||||
bash "${SCRIPT_DIR}/ci_download_flashinfer_cubin.sh"
|
||||
mark_step_done "Fix other dependencies"
|
||||
|
||||
# Force reinstall nvidia-cutlass-dsl to avoid potential version mismatch issues
|
||||
$PIP_CMD install "nvidia-cutlass-dsl>=4.4.1" "nvidia-cutlass-dsl-libs-base>=4.4.1" --no-deps --force-reinstall $PIP_INSTALL_SUFFIX || true
|
||||
# ------------------------------------------------------------------------------
|
||||
# Prepare runner
|
||||
# ------------------------------------------------------------------------------
|
||||
# Prepare the CI runner (cleanup HuggingFace cache, etc.)
|
||||
bash "${SCRIPT_DIR}/prepare_runner.sh"
|
||||
|
||||
mark_step_done "Prepare runner"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Verify imports
|
||||
# ------------------------------------------------------------------------------
|
||||
# Show current packages
|
||||
$PIP_CMD list
|
||||
python3 -c "import torch; print(torch.version.cuda)"
|
||||
python3 -c "import cutlass; import cutlass.cute;"
|
||||
|
||||
# Prepare the CI runner (cleanup HuggingFace cache, etc.)
|
||||
bash "${SCRIPT_DIR}/prepare_runner.sh"
|
||||
mark_step_done "Verify imports"
|
||||
|
||||
53
scripts/ci/utils/install_protoc.sh
Executable file
53
scripts/ci/utils/install_protoc.sh
Executable file
@@ -0,0 +1,53 @@
|
||||
#!/bin/bash
|
||||
# Ensure protoc is installed for router build (gRPC protobuf compilation).
|
||||
set -euxo pipefail
|
||||
|
||||
if command -v protoc >/dev/null 2>&1 && protoc --version >/dev/null 2>&1; then
|
||||
echo "protoc already installed: $(protoc --version)"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
if command -v protoc >/dev/null 2>&1; then
|
||||
echo "protoc found but not runnable, reinstalling..."
|
||||
else
|
||||
echo "protoc not found, installing..."
|
||||
fi
|
||||
|
||||
ARCH=$(uname -m)
|
||||
|
||||
if command -v apt-get &> /dev/null; then
|
||||
# Ubuntu/Debian
|
||||
apt-get update || true # May fail due to unrelated broken packages
|
||||
PROTOC_APT_PACKAGES=(wget unzip gcc g++ perl make)
|
||||
apt-get install -y --no-install-recommends "${PROTOC_APT_PACKAGES[@]}" || {
|
||||
echo "Warning: apt-get install failed, checking if required packages are available..."
|
||||
for pkg in "${PROTOC_APT_PACKAGES[@]}"; do
|
||||
if ! dpkg -l "$pkg" 2>/dev/null | grep -q "^ii"; then
|
||||
echo "ERROR: Required package $pkg is not installed and apt-get failed"
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
echo "All required packages are already installed, continuing..."
|
||||
}
|
||||
elif command -v yum &> /dev/null; then
|
||||
# RHEL/CentOS
|
||||
yum update -y
|
||||
yum install -y wget unzip gcc gcc-c++ perl-core make
|
||||
else
|
||||
echo "ERROR: Neither apt-get nor yum found; cannot install protoc build deps"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [ "$ARCH" = "aarch64" ] || [ "$ARCH" = "arm64" ]; then
|
||||
PROTOC_ARCH="aarch_64"
|
||||
else
|
||||
PROTOC_ARCH="x86_64"
|
||||
fi
|
||||
PROTOC_ZIP="protoc-32.0-linux-${PROTOC_ARCH}.zip"
|
||||
(
|
||||
cd /tmp
|
||||
wget "https://github.com/protocolbuffers/protobuf/releases/download/v32.0/${PROTOC_ZIP}"
|
||||
unzip -o "${PROTOC_ZIP}" -d /usr/local
|
||||
rm -f "${PROTOC_ZIP}"
|
||||
)
|
||||
protoc --version
|
||||
@@ -10,7 +10,6 @@ from utils import compare_versions, get_repo_root, normalize_version, validate_v
|
||||
FILES_TO_UPDATE = [
|
||||
Path("python/pyproject.toml"),
|
||||
Path("docker/Dockerfile"),
|
||||
Path("scripts/ci/cuda/ci_install_dependency.sh"),
|
||||
Path("python/sglang/srt/entrypoints/engine.py"),
|
||||
Path("python/sglang/srt/utils/common.py"),
|
||||
]
|
||||
@@ -52,12 +51,6 @@ def replace_flashinfer_version(
|
||||
rf"\g<1>{new_version}",
|
||||
new_content,
|
||||
)
|
||||
elif name == "ci_install_dependency.sh":
|
||||
new_content = re.sub(
|
||||
rf"(FLASHINFER_VERSION=){re.escape(old_version)}",
|
||||
rf"\g<1>{new_version}",
|
||||
new_content,
|
||||
)
|
||||
elif name == "engine.py":
|
||||
new_content = re.sub(
|
||||
r'(assert_pkg_version\(\s*"flashinfer_python",\s*)"'
|
||||
|
||||
Reference in New Issue
Block a user