[CI] Add per-job uv venv isolation and upgrade CI version to Cuda 13 (#23119)
Co-authored-by: Kangyan Zhou <zky314343421@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Alison Shao <a.shao@wustl.edu> Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Baizhou Zhang
co-authored by
Kangyan Zhou
Claude Opus 4.7
Alison Shao
Mick
Kangyan Zhou
Claude Opus 4.7
Alison Shao
Mick
parent
03828f4205
commit
6ecd6f84db
@@ -21,14 +21,20 @@ NVIDIA_PIP_WHEELS="/root/.cache/nvidia-pip-wheels"
|
||||
mkdir -p "$NVIDIA_WHEEL_CACHE"
|
||||
|
||||
for url in \
|
||||
"https://pypi.nvidia.com/nvidia-cudnn-cu12/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl" \
|
||||
"https://pypi.nvidia.com/nvidia-nvshmem-cu12/nvidia_nvshmem_cu12-3.3.20-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl"; do
|
||||
"https://pypi.nvidia.com/nvidia-cudnn-cu13/nvidia_cudnn_cu13-9.16.0.29-py3-none-manylinux_2_27_x86_64.whl" \
|
||||
"https://pypi.nvidia.com/nvidia-nvshmem-cu13/nvidia_nvshmem_cu13-3.3.20-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl"; do
|
||||
whl="$NVIDIA_WHEEL_CACHE/$(basename "$url")"
|
||||
[ -f "$whl" ] && unzip -tq "$whl" &>/dev/null || curl -fL -o "$whl" "$url"
|
||||
done
|
||||
|
||||
pip install --no-deps "$NVIDIA_WHEEL_CACHE"/nvidia_cudnn_cu12-*.whl \
|
||||
"$NVIDIA_WHEEL_CACHE"/nvidia_nvshmem_cu12-*.whl 2>/dev/null || true
|
||||
# Caller (ci_install_dependency.sh) sets $PIP_CMD/$PIP_INSTALL_SUFFIX to route
|
||||
# installs into the active environment (venv or system). The `:-pip` fallback
|
||||
# keeps the file runnable ad-hoc for debugging; in CI the caller always sets
|
||||
# these. Silent failure here is deliberate — the pinned cudnn/nvshmem installs
|
||||
# later in ci_install_dependency.sh are the source of truth; this is only a
|
||||
# download optimization.
|
||||
${PIP_CMD:-pip} install --no-deps "$NVIDIA_WHEEL_CACHE"/nvidia_cudnn_cu13-*.whl \
|
||||
"$NVIDIA_WHEEL_CACHE"/nvidia_nvshmem_cu13-*.whl ${PIP_INSTALL_SUFFIX:-} 2>/dev/null || true
|
||||
|
||||
# If pre-cached NVIDIA pip wheels exist, tell pip to check there first.
|
||||
# This avoids re-downloading ~2 GB of cublas/cufft/nvrtc/etc. every run
|
||||
|
||||
Executable
+58
@@ -0,0 +1,58 @@
|
||||
#!/bin/bash
|
||||
# Remove the per-job uv venv created by ci_install_dependency.sh.
|
||||
#
|
||||
# Meant to run in a post-job workflow step with `if: always()` so the venv is
|
||||
# destroyed even on job failure/cancel. Runner-level safety net: a cron or
|
||||
# startup task should also purge stale /tmp/sglang-ci-* directories to catch
|
||||
# cancelled or crashed jobs that never reached this cleanup.
|
||||
|
||||
# Best-effort cleanup: never fail the job.
|
||||
set +e
|
||||
set -u
|
||||
|
||||
# Skip entirely when venv mode is disabled — no /tmp/sglang-ci-* dir exists
|
||||
# and there's nothing to sweep. Matches the USE_VENV parsing in
|
||||
# ci_install_dependency.sh (accepts 1/true/yes, case-insensitive).
|
||||
USE_VENV_RAW="${USE_VENV:-true}"
|
||||
case "$(printf '%s' "$USE_VENV_RAW" | tr '[:upper:]' '[:lower:]')" in
|
||||
1 | true | yes) ;;
|
||||
*)
|
||||
echo "USE_VENV=${USE_VENV_RAW}: skipping venv cleanup"
|
||||
exit 0
|
||||
;;
|
||||
esac
|
||||
|
||||
# Prefer the path propagated via GITHUB_ENV. Fallback: glob for any venv from
|
||||
# this run+job (covers the case where install crashed before exporting the path).
|
||||
if [ -n "${SGLANG_CI_VENV_PATH:-}" ] && [ -d "$SGLANG_CI_VENV_PATH" ]; then
|
||||
if rm -rf "$SGLANG_CI_VENV_PATH"; then
|
||||
echo "Cleaned up venv: $SGLANG_CI_VENV_PATH"
|
||||
else
|
||||
echo "::warning::Failed to remove $SGLANG_CI_VENV_PATH — runner cron should sweep /tmp/sglang-ci-*"
|
||||
fi
|
||||
else
|
||||
matched=0
|
||||
for venv in /tmp/sglang-ci-${GITHUB_RUN_ID:-unknownrun}-${GITHUB_JOB:-unknownjob}-*; do
|
||||
[ -d "$venv" ] || continue
|
||||
matched=1
|
||||
if rm -rf "$venv"; then
|
||||
echo "Cleaned up venv (via glob): $venv"
|
||||
else
|
||||
echo "::warning::Failed to remove $venv — runner cron should sweep /tmp/sglang-ci-*"
|
||||
fi
|
||||
done
|
||||
[ "$matched" -eq 0 ] && echo "No venv to clean for run=${GITHUB_RUN_ID:-?} job=${GITHUB_JOB:-?}"
|
||||
fi
|
||||
|
||||
# Sweep stale venvs from cancelled/crashed jobs that never reached cleanup.
|
||||
# Any /tmp/sglang-ci-* dir older than 4 hours is considered orphaned.
|
||||
stale_count=0
|
||||
for venv in /tmp/sglang-ci-*; do
|
||||
[ -d "$venv" ] || continue
|
||||
if find "$venv" -maxdepth 0 -mmin +240 -print -quit | grep -q .; then
|
||||
rm -rf "$venv" && stale_count=$((stale_count + 1))
|
||||
fi
|
||||
done
|
||||
[ "$stale_count" -gt 0 ] && echo "Swept $stale_count stale venv(s) older than 4h"
|
||||
|
||||
exit 0
|
||||
@@ -5,7 +5,7 @@
|
||||
# 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
|
||||
# CU_VERSION — e.g. cu130
|
||||
# PIP_CMD — e.g. "pip" or "uv pip"
|
||||
# PIP_INSTALL_SUFFIX — extra pip args for this runner
|
||||
set -euxo pipefail
|
||||
|
||||
@@ -2,7 +2,23 @@
|
||||
# Install the dependency in CI.
|
||||
set -euxo pipefail
|
||||
|
||||
bash scripts/ci/cuda/ci_install_dependency.sh
|
||||
# Source (not bash) so that venv activation, $PIP_CMD, $CU_VERSION, $NVCC_VER, and
|
||||
# $PIP_INSTALL_SUFFIX all propagate into this shell. Without sourcing, the subshell
|
||||
# exits and this script would fall back to system Python.
|
||||
#
|
||||
# Note: any `exit N` or `set -e` trip inside the sourced script terminates *this*
|
||||
# script too (bash runs sourced commands in the current shell, so `exit` is not
|
||||
# caught by `if`/`||`). The real error message appears upstream in the log.
|
||||
# shellcheck disable=SC1091
|
||||
source scripts/ci/cuda/ci_install_dependency.sh
|
||||
|
||||
# In venv mode, PIP_CMD must be set by the sourced script. If it isn't, the
|
||||
# source chain is broken and we'd silently fall back to system `pip` below —
|
||||
# exactly the split-install bug the migration is meant to prevent.
|
||||
if [ -z "${PIP_CMD:-}" ]; then
|
||||
echo "FATAL:PIP_CMD is unset after sourcing ci_install_dependency.sh"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
export GDRCOPY_HOME=/usr/src/gdrdrv-2.5.1/
|
||||
export CUDA_HOME=/usr/local/cuda
|
||||
@@ -96,24 +112,41 @@ fi
|
||||
|
||||
cd ${DEEPEP_DIR}
|
||||
if [ "$GRACE_BLACKWELL" = "1" ]; then
|
||||
CUDA_VERSION=$(nvidia-smi | grep "CUDA Version" | head -n1 | awk '{print $9}')
|
||||
# Resolve the toolkit CUDA version. Preference order:
|
||||
# 1. $NVCC_VER inherited from the sourced ci_install_dependency.sh
|
||||
# (both scripts agree on the detected value, no re-detection cost).
|
||||
# 2. Local `nvcc --version` (authoritative — container toolkit).
|
||||
# 3. `nvidia-smi` (host driver; last resort).
|
||||
if [ -n "${NVCC_VER:-}" ]; then
|
||||
CUDA_VERSION="$NVCC_VER"
|
||||
elif command -v nvcc >/dev/null 2>&1; then
|
||||
CUDA_VERSION=$(nvcc --version | grep -oP 'release \K[0-9]+\.[0-9]+')
|
||||
else
|
||||
CUDA_VERSION=$(nvidia-smi | grep "CUDA Version" | head -n1 | awk '{print $9}' || true)
|
||||
fi
|
||||
if [ -z "${CUDA_VERSION:-}" ]; then
|
||||
echo "FATAL: could not determine CUDA toolkit version (NVCC_VER unset, nvcc missing, nvidia-smi empty)"
|
||||
exit 1
|
||||
fi
|
||||
if [ "$CUDA_VERSION" = "12.8" ]; then
|
||||
CHOSEN_TORCH_CUDA_ARCH_LIST='10.0'
|
||||
elif awk -v ver="$CUDA_VERSION" 'BEGIN {exit !(ver > 12.8)}'; then
|
||||
# With cuda > 12.8, the compiler supports 10.3, so we should use
|
||||
# CHOSEN_TORCH_CUDA_ARCH_LIST='10.0;10.3'
|
||||
#
|
||||
# However, our CI machine has a weird setup and nvidia-smi reports wrong CUDA version in the container.
|
||||
# The container is actually cuda 12.8, but nvidia-smi reports 13.0, leading to compilation errors. so we
|
||||
# drop 10.3.
|
||||
CHOSEN_TORCH_CUDA_ARCH_LIST='10.0'
|
||||
# CUDA > 12.8 supports sm_103 (Blackwell)
|
||||
CHOSEN_TORCH_CUDA_ARCH_LIST='10.0;10.3'
|
||||
else
|
||||
echo "Unsupported CUDA version for Grace Blackwell: $CUDA_VERSION" && exit 1
|
||||
fi && \
|
||||
if [ "${CUDA_VERSION%%.*}" = "13" ]; then \
|
||||
sed -i "/^ include_dirs = \['csrc\/'\]/a\ include_dirs.append('${CUDA_HOME}/include/cccl')" setup.py; \
|
||||
fi
|
||||
TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" pip install --no-build-isolation .
|
||||
TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" ${PIP_CMD:-pip} install --no-build-isolation . ${PIP_INSTALL_SUFFIX:-}
|
||||
else
|
||||
# CUDA 13.0 puts CCCL headers in /usr/local/cuda/include/cccl/ but nvshmem
|
||||
# includes them as <cuda/__cccl_config> expecting /usr/local/cuda/include/cuda/.
|
||||
# Add the cccl path to setup.py include_dirs so the compiler finds them.
|
||||
NVCC_MAJOR=$(nvcc --version 2>/dev/null | grep -oP 'release \K[0-9]+' || echo "0")
|
||||
if [ "$NVCC_MAJOR" = "13" ]; then
|
||||
sed -i "/^ include_dirs = \['csrc\/'\]/a\ include_dirs.append('${CUDA_HOME:-/usr/local/cuda}/include/cccl')" setup.py
|
||||
fi
|
||||
python3 setup.py install
|
||||
fi
|
||||
|
||||
@@ -23,7 +23,15 @@ set -euxo pipefail
|
||||
# Configuration & timing
|
||||
# ------------------------------------------------------------------------------
|
||||
# Set up environment variables
|
||||
CU_VERSION="cu129"
|
||||
#
|
||||
# CU_VERSION controls:
|
||||
# - PyTorch index URL (pytorch.org/whl/${CU_VERSION})
|
||||
# - FlashInfer JIT cache index (flashinfer.ai/whl/${CU_VERSION})
|
||||
# - nvrtc variant selection (cu12 vs cu13)
|
||||
|
||||
CU_VERSION="${CU_VERSION:-cu130}"
|
||||
CU_STRIP="${CU_VERSION#cu}"
|
||||
CU_MAJOR="${CU_STRIP:0:2}"
|
||||
|
||||
# Nvidia package versions we override (torch pins older versions).
|
||||
# Used both as pip constraints during install and for post-install verification.
|
||||
@@ -31,6 +39,55 @@ NVIDIA_CUDNN_VERSION="9.16.0.29"
|
||||
NVIDIA_NVSHMEM_VERSION="3.4.5"
|
||||
OPTIONAL_DEPS="${1:-}"
|
||||
|
||||
# Whether to create a uv venv. Default false; set USE_VENV=false to install
|
||||
# directly into system Python (useful for runners where uv venv misbehaves).
|
||||
USE_VENV="${USE_VENV:-0}"
|
||||
echo "USE_VENV=${USE_VENV}"
|
||||
|
||||
# uv must be available on system Python (to create the venv, or to run
|
||||
# `uv pip install --system` when venv mode is disabled). Install if missing.
|
||||
python3 -m pip install --upgrade pip
|
||||
if ! command -v uv >/dev/null 2>&1; then
|
||||
pip install uv
|
||||
fi
|
||||
|
||||
SYS_PYTHON_VER=$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
|
||||
|
||||
if [ "$USE_VENV" = "1" ]; then
|
||||
# Per-job unique path. Include $$ (shell PID) so concurrent/back-to-back jobs
|
||||
# on the same runner never target the same directory even if GITHUB_JOB
|
||||
# doesn't differentiate matrix partitions.
|
||||
UV_VENV="/tmp/sglang-ci-${GITHUB_RUN_ID:-norun}-${GITHUB_JOB:-nojob}-$$"
|
||||
# --seed installs pip/setuptools into the venv so bare `pip` calls in
|
||||
# cache_nvidia_wheels.sh and the human-eval setup resolve to the venv's
|
||||
# pip (rather than silently falling back to system Python).
|
||||
uv venv "$UV_VENV" --python "python${SYS_PYTHON_VER}" --seed
|
||||
# shellcheck disable=SC1091
|
||||
source "$UV_VENV/bin/activate"
|
||||
# Assert activation actually took effect. A misconfigured activate script
|
||||
# would otherwise leave us silently running against system Python.
|
||||
[ "${VIRTUAL_ENV:-}" = "$UV_VENV" ] || { echo "FATAL: venv activation did not set VIRTUAL_ENV correctly"; exit 1; }
|
||||
[ "$(command -v python3)" = "$UV_VENV/bin/python3" ] || { echo "FATAL: python3 still resolves outside venv (got $(command -v python3))"; exit 1; }
|
||||
|
||||
# Propagate to subsequent workflow steps. GITHUB_ENV/GITHUB_PATH only
|
||||
# affect *later* steps, never the current one.
|
||||
if [ -n "${GITHUB_ENV:-}" ]; then
|
||||
echo "VIRTUAL_ENV=$UV_VENV" >> "$GITHUB_ENV"
|
||||
echo "SGLANG_CI_VENV_PATH=$UV_VENV" >> "$GITHUB_ENV"
|
||||
# Set BASH_ENV early so subsequent steps auto-source the venv's env script.
|
||||
# LD_LIBRARY_PATH is written to this file later (after packages are installed)
|
||||
# and gets picked up even if GITHUB_ENV becomes unavailable at that point.
|
||||
echo "BASH_ENV=$UV_VENV/env.sh" >> "$GITHUB_ENV"
|
||||
touch "$UV_VENV/env.sh"
|
||||
fi
|
||||
if [ -n "${GITHUB_PATH:-}" ]; then
|
||||
echo "$UV_VENV/bin" >> "$GITHUB_PATH"
|
||||
fi
|
||||
else
|
||||
echo "USE_VENV=0: skipping uv venv creation, installing into system Python"
|
||||
UV_VENV=""
|
||||
fi
|
||||
|
||||
SECONDS=0
|
||||
_CI_MARK_PREV=${SECONDS}
|
||||
|
||||
@@ -159,24 +216,26 @@ mark_step_done "Python package site hygiene & install protoc + rust"
|
||||
# ------------------------------------------------------------------------------
|
||||
# Pip / uv toolchain & stale package cleanup
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install pip and uv (use python3 -m pip for robustness since some runners only have pip3)
|
||||
# Install pip and uv (use python3 -m pip for robustness since some runners only have pip3).
|
||||
# In venv mode this upgrades the venv's pip (the bootstrap block near the top
|
||||
# already upgraded system pip before `uv venv`).
|
||||
python3 -m pip install --upgrade pip
|
||||
|
||||
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
|
||||
pip install uv
|
||||
export UV_SYSTEM_PYTHON=true
|
||||
|
||||
PIP_CMD="uv pip"
|
||||
PIP_INSTALL_SUFFIX="--index-strategy unsafe-best-match --prerelease allow"
|
||||
PIP_UNINSTALL_CMD="uv pip uninstall"
|
||||
PIP_UNINSTALL_SUFFIX=""
|
||||
# uv is already installed on system Python (above).
|
||||
# - Venv mode: the venv is active and `uv pip` targets it automatically.
|
||||
# - Non-venv mode: UV_SYSTEM_PYTHON=1 makes `uv pip` operate on system Python
|
||||
# (otherwise uv refuses to run outside a venv).
|
||||
if [ "$USE_VENV" != "1" ]; then
|
||||
export UV_SYSTEM_PYTHON=1
|
||||
fi
|
||||
|
||||
export UV_LINK_MODE=copy
|
||||
PIP_CMD="uv pip"
|
||||
PIP_INSTALL_SUFFIX="--index-strategy unsafe-best-match --prerelease allow"
|
||||
PIP_UNINSTALL_CMD="uv pip uninstall"
|
||||
PIP_UNINSTALL_SUFFIX=""
|
||||
|
||||
|
||||
# Clean up existing installations
|
||||
$PIP_UNINSTALL_CMD sgl-kernel sglang-kernel sglang sgl-fa4 flash-attn-4 $PIP_UNINSTALL_SUFFIX || true
|
||||
|
||||
@@ -234,19 +293,19 @@ if [ -n "$OPTIONAL_DEPS" ]; then
|
||||
EXTRAS="dev,runai,tracing,${OPTIONAL_DEPS}"
|
||||
fi
|
||||
echo "Installing python extras: [${EXTRAS}]"
|
||||
source "$(dirname "$0")/cache_nvidia_wheels.sh"
|
||||
$PIP_CMD install -e "python[${EXTRAS}]" --extra-index-url https://download.pytorch.org/whl/${CU_VERSION} $PIP_INSTALL_SUFFIX
|
||||
# source "${SCRIPT_DIR}/cache_nvidia_wheels.sh"
|
||||
$PIP_CMD install -e "python[${EXTRAS}]" $PIP_INSTALL_SUFFIX
|
||||
|
||||
mark_step_done "Install main package"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install sglang-kernel
|
||||
# Install torch/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)
|
||||
echo "SGL_KERNEL_VERSION_FROM_KERNEL=${SGL_KERNEL_VERSION_FROM_KERNEL} SGL_KERNEL_VERSION_FROM_SRT=${SGL_KERNEL_VERSION_FROM_SRT}"
|
||||
|
||||
|
||||
if [ "${CUSTOM_BUILD_SGL_KERNEL:-}" = "true" ] && [ -d "sgl-kernel/dist" ]; then
|
||||
ls -alh sgl-kernel/dist
|
||||
# Determine wheel architecture
|
||||
@@ -255,26 +314,50 @@ if [ "${CUSTOM_BUILD_SGL_KERNEL:-}" = "true" ] && [ -d "sgl-kernel/dist" ]; then
|
||||
else
|
||||
WHEEL_ARCH="x86_64"
|
||||
fi
|
||||
$PIP_CMD install sgl-kernel/dist/sglang_kernel-${SGL_KERNEL_VERSION_FROM_KERNEL}-cp310-abi3-manylinux2014_${WHEEL_ARCH}.whl --force-reinstall $PIP_INSTALL_SUFFIX
|
||||
elif [ "${CUSTOM_BUILD_SGL_KERNEL:-}" = "true" ] && [ ! -d "sgl-kernel/dist" ]; then
|
||||
# CUSTOM_BUILD_SGL_KERNEL was set but artifacts not available (e.g., stage rerun without wheel build)
|
||||
# Fail instead of falling back to PyPI - we need to test the built kernel, not PyPI version
|
||||
echo "ERROR: CUSTOM_BUILD_SGL_KERNEL=true but sgl-kernel/dist not found."
|
||||
echo "This usually happens when rerunning a stage without the sgl-kernel-build-wheels job."
|
||||
echo "Please re-run the full workflow using /tag-and-rerun-ci to rebuild the kernel."
|
||||
exit 1
|
||||
# Wheel may have +cuXYZ suffix (e.g. sglang_kernel-0.4.0+cu130-...) depending on CUDA version
|
||||
KERNEL_WHL=$(ls sgl-kernel/dist/sglang_kernel-${SGL_KERNEL_VERSION_FROM_KERNEL}*-cp310-abi3-manylinux2014_${WHEEL_ARCH}.whl 2>/dev/null | head -1)
|
||||
if [ -z "$KERNEL_WHL" ]; then
|
||||
echo "ERROR: No matching sgl-kernel wheel found in sgl-kernel/dist/ for version ${SGL_KERNEL_VERSION_FROM_KERNEL} arch ${WHEEL_ARCH}"
|
||||
ls -alh sgl-kernel/dist/
|
||||
exit 1
|
||||
fi
|
||||
echo "Installing sgl-kernel wheel: $KERNEL_WHL"
|
||||
$PIP_CMD install "$KERNEL_WHL" --force-reinstall $PIP_INSTALL_SUFFIX
|
||||
else
|
||||
# On Blackwell machines, skip reinstall if correct version already installed to avoid race conditions
|
||||
if [ "$IS_BLACKWELL" = "1" ]; then
|
||||
INSTALLED_SGL_KERNEL=$(pip show sglang-kernel 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
|
||||
if [ "$INSTALLED_SGL_KERNEL" = "$SGL_KERNEL_VERSION_FROM_SRT" ]; then
|
||||
echo "sglang-kernel==${SGL_KERNEL_VERSION_FROM_SRT} already installed, skipping reinstall"
|
||||
else
|
||||
echo "Installing sglang-kernel==${SGL_KERNEL_VERSION_FROM_SRT} (current: ${INSTALLED_SGL_KERNEL:-none})"
|
||||
$PIP_CMD install sglang-kernel==${SGL_KERNEL_VERSION_FROM_SRT} $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
if [ "${CUSTOM_BUILD_SGL_KERNEL:-}" = "true" ] && [ ! -d "sgl-kernel/dist" ]; then
|
||||
# CUSTOM_BUILD_SGL_KERNEL was set but artifacts not available (e.g., stage rerun without wheel build)
|
||||
# Fail instead of falling back to PyPI - we need to test the built kernel, not PyPI version
|
||||
echo "ERROR: CUSTOM_BUILD_SGL_KERNEL=true but sgl-kernel/dist not found."
|
||||
echo "This usually happens when rerunning a stage without the sgl-kernel-build-wheels job."
|
||||
echo "Please re-run the full workflow using /tag-and-rerun-ci to rebuild the kernel."
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
# Now we are running torch with cuda13 in CI environment, so the torch packages will be reinstalled if they are still at CU129 version
|
||||
# TODO: Remove this part after torch has been upgraded to 2.11, where cu13 is enabled by default
|
||||
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
|
||||
TORCH_VER=$(pip show torch 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')
|
||||
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 torch==${TORCH_VER} torchaudio==${TORCHAUDIO_VER} torchvision==${TORCHVISION_VER} from ${CU_VERSION} index to match torch..."
|
||||
$PIP_CMD install "torch==${TORCH_VER}" "torchaudio==${TORCHAUDIO_VER}" "torchvision==${TORCHVISION_VER}" --index-url "https://download.pytorch.org/whl/${CU_VERSION}" --force-reinstall --no-deps $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
|
||||
# sglang-kernel wheels carry a +cuXYZ local version tag (e.g. 0.4.1+cu130).
|
||||
# If it doesn't match CU_VERSION, reinstall from the matching index.
|
||||
SGL_KERNEL_FULL_VER=$(pip show sglang-kernel 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
|
||||
SGL_KERNEL_CUDA_VER=$(printf '%s' "$SGL_KERNEL_FULL_VER" | sed -n 's/.*+//p')
|
||||
echo "Detected sglang-kernel version: ${SGL_KERNEL_FULL_VER} (CUDA tag: ${SGL_KERNEL_CUDA_VER:-none})"
|
||||
if [ -n "$SGL_KERNEL_CUDA_VER" ] && [ "$SGL_KERNEL_CUDA_VER" != "$CU_VERSION" ]; then
|
||||
SGL_KERNEL_VER="${SGL_KERNEL_FULL_VER%+*}"
|
||||
echo "Reinstalling sglang-kernel==${SGL_KERNEL_VER} from ${CU_VERSION} index to match torch..."
|
||||
if [ "$CU_MAJOR" = "13" ]; then
|
||||
$PIP_CMD install "sglang-kernel==${SGL_KERNEL_VER}" --index-url "https://docs.sglang.ai/whl/${CU_VERSION}/" --force-reinstall --no-deps $PIP_INSTALL_SUFFIX
|
||||
else
|
||||
$PIP_CMD install sglang-kernel==${SGL_KERNEL_VERSION_FROM_SRT} --force-reinstall $PIP_INSTALL_SUFFIX
|
||||
$PIP_CMD install "sglang-kernel==${SGL_KERNEL_VER}" --force-reinstall --no-deps $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
fi
|
||||
|
||||
@@ -304,16 +387,83 @@ UNINSTALL_JIT_CACHE="$UNINSTALL_JIT_CACHE" \
|
||||
|
||||
mark_step_done "Download flashinfer artifacts"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Stabilize FlashInfer JIT cache paths
|
||||
# ------------------------------------------------------------------------------
|
||||
# FlashInfer JIT writes build.ninja with hardcoded -isystem paths pointing to the
|
||||
# venv's flashinfer/data/ and tvm_ffi/include/. With per-job venvs each job gets
|
||||
# a unique /tmp/sglang-ci-<run>-<job>-<pid>/ path, but the JIT cache is shared
|
||||
# on the host mount. When the next job's venv has a different path and the old one
|
||||
# is cleaned up, ninja fails because source files no longer exist at the cached path.
|
||||
#
|
||||
# Fix (two parts):
|
||||
# 1. Clear only STALE cached_ops (build.ninja referencing non-existent venv paths).
|
||||
# Do NOT clear all cached_ops — they contain compiled .so files that take 10-20 min
|
||||
# to recompile. Only remove entries where the source paths no longer exist.
|
||||
# 2. Copy source files to a stable host-mounted path and symlink each venv's
|
||||
# copy there. build.ninja then references the stable path across all jobs.
|
||||
#
|
||||
# Part 1: Clear stale cached_ops (keep valid compiled kernels)
|
||||
if [ "$USE_VENV" = "1" ]; then
|
||||
STABLE_FI_DIR="${HOME}/.cache/flashinfer/_stable_src"
|
||||
if [ -d "${HOME}/.cache/flashinfer" ]; then
|
||||
STALE_COUNT=0
|
||||
while IFS= read -r ninja_file; do
|
||||
# Check for stale venv paths (/tmp/sglang-ci-*) or old stable path (flashinfer-src)
|
||||
STALE_PATH=$(grep -o '/tmp/sglang-ci-[^ ]*\|flashinfer-src' "$ninja_file" 2>/dev/null | head -1 || true)
|
||||
if [ -n "$STALE_PATH" ]; then
|
||||
if echo "$STALE_PATH" | grep -q "flashinfer-src" || [ ! -d "$STALE_PATH" ]; then
|
||||
rm -rf "$(dirname "$ninja_file")"
|
||||
STALE_COUNT=$((STALE_COUNT + 1))
|
||||
fi
|
||||
fi
|
||||
done < <(find "${HOME}/.cache/flashinfer" -name "build.ninja" -type f 2>/dev/null)
|
||||
echo "Cleaned $STALE_COUNT stale FlashInfer cached_ops (kept valid ones)"
|
||||
fi
|
||||
|
||||
# Part 2: Stabilize paths (STABLE_FI_DIR set above in Part 1)
|
||||
FI_DATA=$(python3 -c "import flashinfer, os; print(os.path.join(os.path.dirname(flashinfer.__file__), 'data'))")
|
||||
TVM_INC=$(python3 -c "import tvm_ffi, os; print(os.path.join(os.path.dirname(tvm_ffi.__file__), 'include'))")
|
||||
|
||||
FI_VERSION="${FLASHINFER_PYTHON_REQUIRED}"
|
||||
if [ ! -d "$STABLE_FI_DIR/flashinfer-data" ] || [ "$(cat "$STABLE_FI_DIR/.version" 2>/dev/null)" != "$FI_VERSION" ]; then
|
||||
rm -rf "$STABLE_FI_DIR"
|
||||
mkdir -p "$STABLE_FI_DIR"
|
||||
cp -a "$FI_DATA" "$STABLE_FI_DIR/flashinfer-data"
|
||||
cp -a "$TVM_INC" "$STABLE_FI_DIR/tvm-ffi-include"
|
||||
echo "$FI_VERSION" > "$STABLE_FI_DIR/.version"
|
||||
echo "Copied flashinfer source files to stable path: $STABLE_FI_DIR (version=$FI_VERSION)"
|
||||
else
|
||||
echo "Stable flashinfer source path up to date (version=$FI_VERSION)"
|
||||
fi
|
||||
|
||||
rm -rf "$FI_DATA"
|
||||
ln -s "$STABLE_FI_DIR/flashinfer-data" "$FI_DATA"
|
||||
TVM_INC_PARENT=$(dirname "$TVM_INC")
|
||||
rm -rf "$TVM_INC_PARENT/include"
|
||||
ln -s "$STABLE_FI_DIR/tvm-ffi-include" "$TVM_INC_PARENT/include"
|
||||
echo "Symlinked venv flashinfer/tvm_ffi -> $STABLE_FI_DIR"
|
||||
|
||||
mark_step_done "Stabilize FlashInfer JIT cache paths"
|
||||
fi
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install extra dependency
|
||||
# ------------------------------------------------------------------------------
|
||||
# Install other python dependencies
|
||||
if [ "$CU_VERSION" = "cu130" ]; then
|
||||
NVRTC_SPEC="nvidia-cuda-nvrtc"
|
||||
# Install other python dependencies.
|
||||
# Match on CUDA major version so future minor bumps (cu131, etc.) don't fall
|
||||
# through to the wrong branch. Prefer NVCC_VER (set in the venv path); otherwise
|
||||
# parse the first two digits of CU_VERSION (pytorch convention is cu{major}{minor}
|
||||
# with a single-digit minor, e.g. cu126, cu129, cu130).
|
||||
if [ "$CU_MAJOR" = "13" ]; then
|
||||
MOONCAKE_PKG="mooncake-transfer-engine-cuda13==0.3.10.post1"
|
||||
EXTRA_NVIDIA_SPECS="nvidia-cuda-nvrtc"
|
||||
else
|
||||
NVRTC_SPEC="nvidia-cuda-nvrtc-cu12"
|
||||
MOONCAKE_PKG="mooncake-transfer-engine==0.3.10.post1"
|
||||
EXTRA_NVIDIA_SPECS="nvidia-cuda-nvrtc-cu12"
|
||||
fi
|
||||
$PIP_CMD install mooncake-transfer-engine==0.3.10.post1 "${NVRTC_SPEC}" py-spy scipy huggingface_hub[hf_xet] pytest $PIP_INSTALL_SUFFIX
|
||||
$PIP_CMD install ${MOONCAKE_PKG} ${EXTRA_NVIDIA_SPECS} py-spy scipy huggingface_hub[hf_xet] pytest $PIP_INSTALL_SUFFIX
|
||||
|
||||
# Install other test dependencies
|
||||
if [ "$IS_BLACKWELL" != "1" ]; then
|
||||
@@ -328,55 +478,51 @@ 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
|
||||
|
||||
# Pick cu12 vs cu13 variants of nvshmem / cudnn based on CU_VERSION
|
||||
if [ "$CU_MAJOR" = "13" ]; then
|
||||
NVSHMEM_PKG="nvidia-nvshmem-cu13"
|
||||
CUDNN_PKG="nvidia-cudnn-cu13"
|
||||
else
|
||||
NVSHMEM_PKG="nvidia-nvshmem-cu12"
|
||||
CUDNN_PKG="nvidia-cudnn-cu12"
|
||||
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 "")
|
||||
INSTALLED_NVSHMEM=$(pip show ${NVSHMEM_PKG} 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
|
||||
if [ "$INSTALLED_NVSHMEM" = "$NVIDIA_NVSHMEM_VERSION" ]; then
|
||||
echo "nvidia-nvshmem-cu12==${NVIDIA_NVSHMEM_VERSION} already installed, skipping reinstall"
|
||||
echo "${NVSHMEM_PKG}==${NVIDIA_NVSHMEM_VERSION} already installed, skipping reinstall"
|
||||
else
|
||||
$PIP_CMD install nvidia-nvshmem-cu12==${NVIDIA_NVSHMEM_VERSION} $PIP_INSTALL_SUFFIX
|
||||
$PIP_CMD install ${NVSHMEM_PKG}==${NVIDIA_NVSHMEM_VERSION} $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
|
||||
# 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 "")
|
||||
INSTALLED_CUDNN=$(pip show ${CUDNN_PKG} 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
|
||||
if [ "$INSTALLED_CUDNN" = "$NVIDIA_CUDNN_VERSION" ]; then
|
||||
echo "nvidia-cudnn-cu12==${NVIDIA_CUDNN_VERSION} already installed, skipping reinstall"
|
||||
echo "${CUDNN_PKG}==${NVIDIA_CUDNN_VERSION} already installed, skipping reinstall"
|
||||
else
|
||||
$PIP_CMD install nvidia-cudnn-cu12==${NVIDIA_CUDNN_VERSION} $PIP_INSTALL_SUFFIX
|
||||
$PIP_CMD install ${CUDNN_PKG}==${NVIDIA_CUDNN_VERSION} $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
|
||||
mark_step_done "Fix other dependencies"
|
||||
|
||||
# Force reinstall nvidia-cutlass-dsl to ensure the .pth file exists.
|
||||
# The Docker image ships nvidia-cutlass-dsl-libs-base 4.3.5; upgrading to 4.4.2
|
||||
# can delete the .pth file without reliably recreating it (pip race condition).
|
||||
$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
|
||||
|
||||
# Download kernels from kernels community
|
||||
kernels download python || true
|
||||
kernels lock python || true
|
||||
mv python/kernels.lock ${HOME}/.cache/sglang || true
|
||||
# Ensure target is a directory — on fresh containers or after a previous buggy
|
||||
# `mv` that created a FILE at this path, mkdir -p would fail silently.
|
||||
[ -e "${HOME}/.cache/sglang" ] && [ ! -d "${HOME}/.cache/sglang" ] && rm -f "${HOME}/.cache/sglang"
|
||||
mkdir -p "${HOME}/.cache/sglang/"
|
||||
mv python/kernels.lock "${HOME}/.cache/sglang/" || true
|
||||
|
||||
# Install human-eval
|
||||
pip install "setuptools==70.0.0"
|
||||
git clone https://github.com/merrymercy/human-eval.git
|
||||
cd human-eval
|
||||
pip install -e . --no-build-isolation
|
||||
# Install human-eval. This script is sourced from ci_install_deepep.sh, so a
|
||||
# bare `cd human-eval` would leave the caller stuck in that directory for the
|
||||
# rest of its execution. The subshell keeps the cd local to the pip install.
|
||||
$PIP_CMD install "setuptools==70.0.0" $PIP_INSTALL_SUFFIX
|
||||
[ -d human-eval ] || git clone https://github.com/merrymercy/human-eval.git
|
||||
(
|
||||
cd human-eval
|
||||
$PIP_CMD install -e . --no-build-isolation $PIP_INSTALL_SUFFIX)
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Prepare runner
|
||||
@@ -386,6 +532,35 @@ bash "${SCRIPT_DIR}/prepare_runner.sh"
|
||||
|
||||
mark_step_done "Prepare runner"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# LD_LIBRARY_PATH discovery
|
||||
# ------------------------------------------------------------------------------
|
||||
# NVIDIA pip packages (cublas, cudnn, nccl, nvrtc, ...) and torch ship .so files
|
||||
# under site-packages. In venv mode these are NOT on the default LD_LIBRARY_PATH,
|
||||
# so dlopen('libcublas.so.12') from torch would fail. Prepend them here.
|
||||
# In non-venv mode, system site-packages may also need this if the runner's
|
||||
# default ld config doesn't cover the NVIDIA pip layout.
|
||||
SITE_PACKAGES=$(python3 -c "import site, sys; print(site.getsitepackages()[0])")
|
||||
# Glob matches NVIDIA pip-package layout:
|
||||
# site-packages/nvidia/<component>/lib/lib*.so. If NVIDIA restructures
|
||||
# packaging, this may need updating.
|
||||
NVIDIA_LIBS=$(find "$SITE_PACKAGES" -path "*/nvidia/*/lib" -type d 2>/dev/null | tr '\n' ':')
|
||||
TORCH_LIB="$SITE_PACKAGES/torch/lib"
|
||||
VENV_LD="${NVIDIA_LIBS}${TORCH_LIB}"
|
||||
export LD_LIBRARY_PATH="${VENV_LD}${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
|
||||
# Write LD_LIBRARY_PATH to the venv's env.sh (always succeeds — local file)
|
||||
# so subsequent steps auto-source it via BASH_ENV. In non-venv mode, skip the
|
||||
# env.sh write and rely on GITHUB_ENV propagation.
|
||||
if [ "$USE_VENV" = "1" ] && [ -n "$UV_VENV" ]; then
|
||||
echo "export LD_LIBRARY_PATH=\"$LD_LIBRARY_PATH\"" >> "$UV_VENV/env.sh"
|
||||
fi
|
||||
# Also try GITHUB_ENV (may fail if runner temp file was cleaned up during long installs).
|
||||
if [ -n "${GITHUB_ENV:-}" ]; then
|
||||
echo "LD_LIBRARY_PATH=$LD_LIBRARY_PATH" >> "$GITHUB_ENV" || echo "WARNING: GITHUB_ENV write failed; LD_LIBRARY_PATH will be set via BASH_ENV instead"
|
||||
fi
|
||||
echo "LD_LIBRARY_PATH=$LD_LIBRARY_PATH"
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Verify imports
|
||||
# ------------------------------------------------------------------------------
|
||||
@@ -393,5 +568,3 @@ mark_step_done "Prepare runner"
|
||||
$PIP_CMD list
|
||||
python3 -c "import torch; print(torch.version.cuda)"
|
||||
python3 -c "import cutlass; import cutlass.cute;"
|
||||
|
||||
mark_step_done "Verify imports"
|
||||
|
||||
Reference in New Issue
Block a user