[MUSA] Harden CI dependencies and diffusion warmup (#35610)
This commit is contained in:
R0CKSTAR
@@ -1,4 +1,5 @@
|
||||
#!/bin/bash
|
||||
# shellcheck disable=SC2034 # OPTIONAL_DEPS is retained for CLI compatibility.
|
||||
set -euo pipefail
|
||||
|
||||
# Parse command line arguments
|
||||
@@ -21,8 +22,40 @@ while [[ $# -gt 0 ]]; do
|
||||
esac
|
||||
done
|
||||
|
||||
PIP_INSTALL="python3 -m pip install --no-cache-dir"
|
||||
${PIP_INSTALL} --upgrade pip setuptools torchada --user
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
STACK_HELPER="${SCRIPT_DIR}/musa_python_stack.py"
|
||||
# Keep installed packages isolated in PYTHONUSERBASE, but let pip reuse its
|
||||
# content-addressed wheel cache across jobs. Disabling the cache made every
|
||||
# MUSA lane cold-download hundreds of megabytes and exceed the install timeout.
|
||||
PIP_INSTALL=(python3 -m pip install)
|
||||
readonly MUSA_TRITON_VERSION="3.2.0"
|
||||
# Current MUSA CI uses the CPython 3.10 x86_64 wheel. A Python upgrade must
|
||||
# update this digest together with the pinned artifact.
|
||||
readonly MUSA_TRITON_SHA256="65b15d42fac24a2eca4c0c9f0ac68c8bd7cbe6bcc9f619c3483fb4f323391303"
|
||||
readonly MUSA_TRITON_INDEX_URL="https://dl.mthreads.com/repo/api/pypi/pypi/simple"
|
||||
readonly MUSA_TORCHADA_VERSION="0.1.82"
|
||||
readonly MUSA_TORCHADA_SHA256="472663da083ef23502f08429618a36e5e6e9b2447cf72fff2568787c815b5903"
|
||||
readonly MUSA_TORCHADA_INDEX_URL="https://pypi.org/simple"
|
||||
readonly MUSA_SETUPTOOLS_SPEC="setuptools<82"
|
||||
MUSA_CI_SCRATCH="$(mktemp -d "${RUNNER_TEMP:-${TMPDIR:-/tmp}}/sglang-musa-ci.XXXXXX")"
|
||||
MUSA_CI_ISOLATED_USERBASE=""
|
||||
|
||||
if [ "${GITHUB_ACTIONS:-}" = "true" ]; then
|
||||
MUSA_CI_ISOLATED_USERBASE="1"
|
||||
PYTHONUSERBASE="$(mktemp -d "${RUNNER_TEMP:-${TMPDIR:-/tmp}}/sglang-musa-python.XXXXXX")"
|
||||
export PYTHONUSERBASE
|
||||
if [ -n "${GITHUB_ENV:-}" ]; then
|
||||
echo "PYTHONUSERBASE=${PYTHONUSERBASE}" >> "$GITHUB_ENV"
|
||||
fi
|
||||
if [ -n "${GITHUB_PATH:-}" ]; then
|
||||
echo "${PYTHONUSERBASE}/bin" >> "$GITHUB_PATH"
|
||||
fi
|
||||
echo "Using task-local Python user base: ${PYTHONUSERBASE}"
|
||||
fi
|
||||
|
||||
# torchada has an unpinned Torch dependency. Installing it from the public
|
||||
# index before the MUSA wheel bundle can pull CUDA Torch and CUDA Triton.
|
||||
"${PIP_INSTALL[@]}" --upgrade pip "$MUSA_SETUPTOOLS_SPEC" --user
|
||||
|
||||
echo "Checking stale torchada extension locks..."
|
||||
active_torchada_builds="$(
|
||||
@@ -43,11 +76,135 @@ if [ -d "$torch_extensions_dir" ]; then
|
||||
-delete
|
||||
fi
|
||||
|
||||
WHL_DIR="/sglang-checkout/whl"
|
||||
WHL_DIR="${WHL_DIR:-/sglang-checkout/whl}"
|
||||
MUSA_TRITON_WHEEL=""
|
||||
MUSA_TRITON_TASK_LOCAL=""
|
||||
MUSA_TORCHADA_WHEEL=""
|
||||
|
||||
wheel_sha256() {
|
||||
sha256sum "$1" | awk '{print $1}'
|
||||
}
|
||||
|
||||
find_bundled_musa_triton() {
|
||||
local candidate
|
||||
local -a candidates
|
||||
if [ ! -d "$WHL_DIR" ]; then
|
||||
return
|
||||
fi
|
||||
mapfile -t candidates < <(
|
||||
compgen -G "${WHL_DIR}/triton-${MUSA_TRITON_VERSION}-*.whl" || true
|
||||
)
|
||||
for candidate in "${candidates[@]}"; do
|
||||
if [ "$(wheel_sha256 "$candidate")" = "$MUSA_TRITON_SHA256" ]; then
|
||||
MUSA_TRITON_WHEEL="$candidate"
|
||||
return
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
download_musa_triton() {
|
||||
local download_dir="${MUSA_CI_SCRATCH}/musa-triton"
|
||||
local -a candidates
|
||||
mkdir -p "$download_dir"
|
||||
python3 -m pip --isolated download \
|
||||
--dest "$download_dir" \
|
||||
--index-url "$MUSA_TRITON_INDEX_URL" \
|
||||
--no-deps \
|
||||
--only-binary=:all: \
|
||||
"triton==${MUSA_TRITON_VERSION}"
|
||||
mapfile -t candidates < <(find "$download_dir" -maxdepth 1 -type f -name '*.whl')
|
||||
if [ "${#candidates[@]}" -ne 1 ]; then
|
||||
echo "::error::Expected one MUSA Triton wheel, found ${#candidates[@]}"
|
||||
exit 1
|
||||
fi
|
||||
MUSA_TRITON_WHEEL="${candidates[0]}"
|
||||
}
|
||||
|
||||
find_bundled_torchada() {
|
||||
local candidate
|
||||
local -a candidates
|
||||
if [ ! -d "$WHL_DIR" ]; then
|
||||
return
|
||||
fi
|
||||
mapfile -t candidates < <(
|
||||
compgen -G "${WHL_DIR}/torchada-${MUSA_TORCHADA_VERSION}-*.whl" || true
|
||||
)
|
||||
for candidate in "${candidates[@]}"; do
|
||||
if [ "$(wheel_sha256 "$candidate")" = "$MUSA_TORCHADA_SHA256" ]; then
|
||||
MUSA_TORCHADA_WHEEL="$candidate"
|
||||
return
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
download_torchada() {
|
||||
local download_dir="${MUSA_CI_SCRATCH}/torchada"
|
||||
local -a candidates
|
||||
mkdir -p "$download_dir"
|
||||
python3 -m pip --isolated download \
|
||||
--dest "$download_dir" \
|
||||
--index-url "$MUSA_TORCHADA_INDEX_URL" \
|
||||
--no-deps \
|
||||
--only-binary=:all: \
|
||||
"torchada==${MUSA_TORCHADA_VERSION}"
|
||||
mapfile -t candidates < <(find "$download_dir" -maxdepth 1 -type f -name '*.whl')
|
||||
if [ "${#candidates[@]}" -ne 1 ]; then
|
||||
echo "::error::Expected one torchada wheel, found ${#candidates[@]}"
|
||||
exit 1
|
||||
fi
|
||||
MUSA_TORCHADA_WHEEL="${candidates[0]}"
|
||||
}
|
||||
|
||||
find_bundled_musa_triton
|
||||
if [ -z "$MUSA_TRITON_WHEEL" ]; then
|
||||
if python3 "$STACK_HELPER" verify \
|
||||
--expected-triton-version "$MUSA_TRITON_VERSION" \
|
||||
--triton-only; then
|
||||
echo "Reusing the installed MUSA Triton"
|
||||
else
|
||||
# Existing runner bundles do not contain Triton yet. Keep a hash-pinned
|
||||
# vendor-index fallback until the wheel is shipped in /sglang-checkout/whl.
|
||||
download_musa_triton
|
||||
fi
|
||||
fi
|
||||
if [ -n "$MUSA_TRITON_WHEEL" ]; then
|
||||
if [ "$(wheel_sha256 "$MUSA_TRITON_WHEEL")" != "$MUSA_TRITON_SHA256" ]; then
|
||||
echo "::error::MUSA Triton wheel SHA256 does not match the pinned artifact"
|
||||
exit 1
|
||||
fi
|
||||
MUSA_TRITON_TASK_LOCAL="1"
|
||||
echo "Using MUSA Triton wheel: ${MUSA_TRITON_WHEEL}"
|
||||
fi
|
||||
|
||||
# torchada declares an unpinned dependency on Torch. Install the exact wheel
|
||||
# with --no-deps so a fresh user site cannot resolve public CUDA Torch/Triton.
|
||||
find_bundled_torchada
|
||||
if [ -z "$MUSA_TORCHADA_WHEEL" ]; then
|
||||
download_torchada
|
||||
fi
|
||||
if [ "$(wheel_sha256 "$MUSA_TORCHADA_WHEEL")" != "$MUSA_TORCHADA_SHA256" ]; then
|
||||
echo "::error::torchada wheel SHA256 does not match the pinned artifact"
|
||||
exit 1
|
||||
fi
|
||||
echo "Using torchada wheel: ${MUSA_TORCHADA_WHEEL}"
|
||||
|
||||
VENDOR_WHEELS=("$MUSA_TORCHADA_WHEEL")
|
||||
if [ -n "$MUSA_TRITON_WHEEL" ]; then
|
||||
VENDOR_WHEELS+=("$MUSA_TRITON_WHEEL")
|
||||
fi
|
||||
if [ -d "$WHL_DIR" ] && compgen -G "${WHL_DIR}"/*.whl > /dev/null; then
|
||||
for whl in "${WHL_DIR}"/*.whl; do
|
||||
case "$(basename "$whl")" in
|
||||
triton-*.whl|torchada-*.whl) continue;;
|
||||
esac
|
||||
VENDOR_WHEELS+=("$whl")
|
||||
done
|
||||
fi
|
||||
|
||||
if [ -z "$MUSA_CI_ISOLATED_USERBASE" ]; then
|
||||
echo "Uninstall old packages based on wheel METADATA..."
|
||||
PKGS=$(
|
||||
for whl in "${WHL_DIR}"/*.whl; do
|
||||
for whl in "${VENDOR_WHEELS[@]}"; do
|
||||
meta_file=$(zipinfo -1 "$whl" | awk '/\.dist-info\/METADATA$/ {print; exit}')
|
||||
[ -n "$meta_file" ] || continue
|
||||
unzip -p "$whl" "$meta_file" 2>/dev/null | sed -n 's/^Name: //p' | head -n1
|
||||
@@ -55,18 +212,36 @@ if [ -d "$WHL_DIR" ] && compgen -G "${WHL_DIR}"/*.whl > /dev/null; then
|
||||
)
|
||||
for pkg in $PKGS; do
|
||||
echo "Uninstalling $pkg"
|
||||
pip uninstall -y "$pkg" || true
|
||||
python3 -m pip uninstall -y "$pkg" || true
|
||||
done
|
||||
echo "Installing wheel files without dependency resolution..."
|
||||
${PIP_INSTALL} "${WHL_DIR}"/*.whl --user
|
||||
fi
|
||||
|
||||
if [ "${#VENDOR_WHEELS[@]}" -gt 0 ]; then
|
||||
echo "Installing vendor wheels without dependency resolution..."
|
||||
if [ -n "$MUSA_CI_ISOLATED_USERBASE" ]; then
|
||||
"${PIP_INSTALL[@]}" --ignore-installed --no-deps --user "${VENDOR_WHEELS[@]}"
|
||||
else
|
||||
"${PIP_INSTALL[@]}" --no-deps --user "${VENDOR_WHEELS[@]}"
|
||||
fi
|
||||
fi
|
||||
|
||||
MUSA_CONSTRAINTS="${MUSA_CI_SCRATCH}/constraints.txt"
|
||||
python3 "$STACK_HELPER" constraints --output "$MUSA_CONSTRAINTS"
|
||||
if [ -n "${GITHUB_ENV:-}" ]; then
|
||||
echo "MUSA_CONSTRAINTS=${MUSA_CONSTRAINTS}" >> "$GITHUB_ENV"
|
||||
fi
|
||||
python3 "$STACK_HELPER" verify \
|
||||
--expected-triton-version "$MUSA_TRITON_VERSION" \
|
||||
${MUSA_TRITON_TASK_LOCAL:+--require-user-site}
|
||||
|
||||
if [ -n "$SKIP_SGLANG_BUILD" ]; then
|
||||
echo "Didn't build checkout SGLang"
|
||||
exit 0
|
||||
else
|
||||
pip uninstall sgl-kernel -y || true
|
||||
pip uninstall sglang -y || true
|
||||
if [ -z "$MUSA_CI_ISOLATED_USERBASE" ]; then
|
||||
python3 -m pip uninstall sgl-kernel -y || true
|
||||
python3 -m pip uninstall sglang -y || true
|
||||
fi
|
||||
# Clear Python cache to ensure latest code is used (works for any env: venv, system, conda)
|
||||
REPO_ROOT="${GITHUB_WORKSPACE:-$(pwd)}"
|
||||
find "$REPO_ROOT" -name "*.pyc" -delete 2>/dev/null || true
|
||||
@@ -83,9 +258,17 @@ else
|
||||
bash "${REPO_ROOT}/scripts/ci/utils/install_rustup.sh"
|
||||
export PATH="${CARGO_HOME:-$HOME/.cargo}/bin:${PATH}"
|
||||
|
||||
cd "${REPO_ROOT}" && ${PIP_INSTALL} -v -e "python[dev_musa]" --user
|
||||
cd "${REPO_ROOT}" && "${PIP_INSTALL[@]}" \
|
||||
--constraint "$MUSA_CONSTRAINTS" \
|
||||
-v \
|
||||
-e "python[dev_musa]" \
|
||||
--user
|
||||
|
||||
cd "${REPO_ROOT}/python/sglang/kernels/aot"
|
||||
rm -f pyproject.toml && mv pyproject_musa.toml pyproject.toml && MTGPU_TARGET=mp_31 python3 setup_musa.py install --user
|
||||
echo "$HOME/.local/bin" >> "$GITHUB_PATH"
|
||||
python3 "$STACK_HELPER" verify \
|
||||
--expected-triton-version "$MUSA_TRITON_VERSION" \
|
||||
--require-driver \
|
||||
--require-resolved-dependencies \
|
||||
${MUSA_TRITON_TASK_LOCAL:+--require-user-site}
|
||||
fi
|
||||
|
||||
Reference in New Issue
Block a user