minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] [Docker] Upgrade Python 3.12 + torch 2.11 + triton 3.7 in ROCm 7.2.4 (#30984)

Browse Source 
Co-authored-by: Chen <bingxche@amd.com>
This commit is contained in:
chuyeh
2026-08-19 18:18:31 -07:00
committed by GitHub
co-authored by Chen
parent ab203663c4
commit c7478228dd
9 changed files with 428 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/ci/amd/amd_ci_install_dependency.sh
+88 -38
View File
@@ -30,7 +30,6 @@ EXTRAS="dev_hip,tracing"
if [ -n "$OPTIONAL_DEPS" ]; then
EXTRAS="dev_hip,tracing,${OPTIONAL_DEPS}"
fi
echo "Installing python extras: [${EXTRAS}]"
# Host names look like: linux-mi35x-gpu-1-xxxxx-runner-zzzzz
if [[ "${HOSTNAME_VALUE}" =~ ^linux-(mi[0-9]+[a-z]*)-gpu-[0-9]+ ]]; then
@@ -40,7 +39,56 @@ else
echo "Warning: could not parse GPU architecture from '${HOSTNAME_VALUE}', defaulting to ${GPU_ARCH}"
fi
# Identify the Dockerfile stage that built this image. Both the python extras
# below and the AITER pin lookup further down need the flavor, and they have to
# agree, so detect it once here.
#
# Prefer GPU_ARCH stamped into the image (gfx950-rocm724, gfx942, ...).
# Images built before that ENV existed: 724 stages already set
# PIP_CONSTRAINT and HSA_ENABLE_IPC_MODE_LEGACY; remaining HIP 7.2* is
# 720; else 7.0. Do not key off torch 2.11 — 720 may ship that later.
IMAGE_TORCH_VERSION=$(docker exec ci_sglang python3 -c 'import torch; print(torch.__version__)')
IMAGE_HIP_VERSION=$(docker exec ci_sglang python3 -c 'import torch; print(torch.version.hip or "")')
IMAGE_GPU_ARCH=$(docker exec ci_sglang printenv GPU_ARCH 2>/dev/null || true)
if [[ "${IMAGE_GPU_ARCH}" =~ ^(gfx942|gfx950)(-rocm720|-rocm724)?$ ]]; then
echo "[CI-IMAGE] Image GPU_ARCH=${IMAGE_GPU_ARCH}"
case "${IMAGE_GPU_ARCH}" in
*-rocm724) IMAGE_BASE_ARG_SUFFIX="_ROCM724"; IMAGE_STAGE_SUFFIX="-rocm724" ;;
*-rocm720) IMAGE_BASE_ARG_SUFFIX="_ROCM720"; IMAGE_STAGE_SUFFIX="-rocm720" ;;
*) IMAGE_BASE_ARG_SUFFIX=""; IMAGE_STAGE_SUFFIX="" ;;
esac
IMAGE_GFX="${IMAGE_GPU_ARCH%-*}"
else
IMAGE_PIP_CONSTRAINT=$(docker exec ci_sglang printenv PIP_CONSTRAINT 2>/dev/null || true)
IMAGE_HSA_LEGACY=$(docker exec ci_sglang printenv HSA_ENABLE_IPC_MODE_LEGACY 2>/dev/null || true)
if [[ -n "${IMAGE_PIP_CONSTRAINT}" || "${IMAGE_HSA_LEGACY}" == "1" ]]; then
IMAGE_BASE_ARG_SUFFIX="_ROCM724"
IMAGE_STAGE_SUFFIX="-rocm724"
elif [[ "${IMAGE_HIP_VERSION}" == 7.2* ]]; then
IMAGE_BASE_ARG_SUFFIX="_ROCM720"
IMAGE_STAGE_SUFFIX="-rocm720"
else
IMAGE_BASE_ARG_SUFFIX=""
IMAGE_STAGE_SUFFIX=""
fi
if [[ "${GPU_ARCH}" == "mi35x" ]]; then
IMAGE_GFX="gfx950"
else
IMAGE_GFX="gfx942"
fi
echo "[CI-IMAGE] Image has no GPU_ARCH stamp; inferred ${IMAGE_GFX}${IMAGE_STAGE_SUFFIX} (PIP_CONSTRAINT='${IMAGE_PIP_CONSTRAINT}', HSA_ENABLE_IPC_MODE_LEGACY='${IMAGE_HSA_LEGACY}', HIP=${IMAGE_HIP_VERSION})"
unset IMAGE_PIP_CONSTRAINT IMAGE_HSA_LEGACY
fi
unset IMAGE_GPU_ARCH
# Install the required dependencies in CI.
# ROCm 7.2.4 images ship torch 2.11, which srt_hip cannot satisfy (it pins
# compressed-tensors 0.15.0, requiring torch<2.11). Select the rocm724 extras.
if [[ "${IMAGE_STAGE_SUFFIX}" == "-rocm724" ]]; then
EXTRAS="${EXTRAS/dev_hip/dev_hip_rocm724}"
fi
echo "Image torch ${IMAGE_TORCH_VERSION}, HIP ${IMAGE_HIP_VERSION}; installing python extras: [${EXTRAS}]"
# Fix permissions on pip cache, ignore errors from concurrent access or missing temp files
docker exec ci_sglang chown -R root:root /sgl-data/pip-cache 2>/dev/null || true
docker exec ci_sglang pip install --cache-dir=/sgl-data/pip-cache --upgrade pip
@@ -221,6 +269,8 @@ if docker exec ci_sglang test -d /sgl-workspace/mori; then
cd /sgl-workspace/mori
git checkout '${MORI_COMMIT}'
git submodule update --init --recursive
apt-get update
apt-get install -y --no-install-recommends libgrpc++-dev 2>/dev/null || true
python3 setup.py develop
python3 -c 'import os, torch; print(os.path.join(os.path.dirname(torch.__file__), \"lib\"))' > /etc/ld.so.conf.d/torch.conf
ldconfig
@@ -248,35 +298,23 @@ DOCKERFILE="docker/rocm.Dockerfile"
GPU_ARCH="${GPU_ARCH:-mi30x}"
echo "[CI-AITER-CHECK] Runner GPU_ARCH=${GPU_ARCH}"
# ROCm 7.0 keeps the Triton its base image ships; later ROCm images run on the
# Triton AITER pins, so a rebuilt AITER has to bring its own along.
IMAGE_HIP_VERSION=$(docker exec ci_sglang python3 -c 'import torch; print(torch.version.hip or "")')
case "${IMAGE_HIP_VERSION}" in
7.0*) INSTALL_AITER_TRITON="false" ;;
7.*) INSTALL_AITER_TRITON="true" ;;
*)
echo "[CI-AITER-CHECK] ERROR: Unsupported or empty HIP version: '${IMAGE_HIP_VERSION}'"
exit 1
;;
esac
echo "[CI-AITER-CHECK] Container HIP=${IMAGE_HIP_VERSION}, install AITER's Triton on rebuild=${INSTALL_AITER_TRITON}"
# Image owns Triton (pinned in docker/rocm.Dockerfile). Rebuild AITER against it.
#############################################
# 1. Extract AITER_COMMIT from correct Dockerfile block
# 1. Extract AITER_COMMIT from the Dockerfile stage that built this image, as
# identified near the top of this script.
#############################################
if [[ "${GPU_ARCH}" == "mi35x" ]]; then
echo "[CI-AITER-CHECK] Using gfx950 block from Dockerfile..."
REPO_AITER_COMMIT=$(grep -F -A20 'FROM $BASE_IMAGE_950 AS gfx950' docker/rocm.Dockerfile \
| grep 'AITER_COMMIT_DEFAULT=' \
| head -n1 \
| sed 's/.*AITER_COMMIT_DEFAULT="\([^"]*\)".*/\1/')
if [[ "${IMAGE_GFX}" == "gfx950" ]]; then
_from_line="FROM \$BASE_IMAGE_950${IMAGE_BASE_ARG_SUFFIX} AS gfx950${IMAGE_STAGE_SUFFIX}"
else
echo "[CI-AITER-CHECK] Using gfx942 block from Dockerfile..."
REPO_AITER_COMMIT=$(grep -F -A20 'FROM $BASE_IMAGE_942 AS gfx942' docker/rocm.Dockerfile \
| grep 'AITER_COMMIT_DEFAULT=' \
| head -n1 \
| sed 's/.*AITER_COMMIT_DEFAULT="\([^"]*\)".*/\1/')
_from_line="FROM \$BASE_IMAGE_942${IMAGE_BASE_ARG_SUFFIX} AS gfx942${IMAGE_STAGE_SUFFIX}"
fi
echo "[CI-AITER-CHECK] Using ${_from_line} from Dockerfile..."
REPO_AITER_COMMIT=$(grep -F -A20 "${_from_line}" docker/rocm.Dockerfile \
| grep 'AITER_COMMIT_DEFAULT=' \
| head -n1 \
| sed 's/.*AITER_COMMIT_DEFAULT="\([^"]*\)".*/\1/')
unset _from_line
if [[ -z "${REPO_AITER_COMMIT}" ]]; then
@@ -344,6 +382,30 @@ if [[ "${NEED_REBUILD}" == "true" ]]; then
pip install -r requirements.txt
"
# Re-apply the Dockerfile torch.Stream patch after re-clone (ROCm/aiter#4817).
if [[ "${IMAGE_STAGE_SUFFIX}" == "-rocm724" ]]; then
docker exec -i ci_sglang python3 - <<'PY'
from pathlib import Path
p = Path("/sgl-workspace/aiter/csrc/cpp_itfs/torch_utils.py")
s = p.read_text()
old = """ elif isinstance(arg, torch.cuda.Stream):
c_args.append(ctypes.cast(arg.cuda_stream, ctypes.c_void_p))
"""
new = """ elif isinstance(arg, torch.Stream):
handle = getattr(arg, "cuda_stream", None)
if handle is None:
handle = torch.cuda.Stream(
stream_id=arg.stream_id,
device_index=arg.device_index,
device_type=arg.device_type,
).cuda_stream
c_args.append(ctypes.cast(handle, ctypes.c_void_p))
"""
if old in s:
p.write_text(s.replace(old, new))
PY
fi
if [[ "${GPU_ARCH}" == "mi35x" ]]; then
GPU_ARCH_LIST="gfx950"
else
@@ -351,19 +413,7 @@ if [[ "${NEED_REBUILD}" == "true" ]]; then
fi
echo "[CI-AITER-CHECK] GPU_ARCH_LIST=${GPU_ARCH_LIST}"
# Run the installer here rather than letting setup.py do it: setup.py
# swallows its errors, and the AITER_USE_SYSTEM_TRITON=1 below then keeps
# whatever Triton is already installed. Doing it up front fails closed.
if [[ "${INSTALL_AITER_TRITON}" == "true" ]]; then
docker exec ci_sglang bash -c "
set -euo pipefail
cd /sgl-workspace/aiter
test -f .github/scripts/install_triton.sh
bash .github/scripts/install_triton.sh
"
fi
# build AITER
# The image already has the Dockerfile-pinned Triton; compile against it.
docker exec ci_sglang bash -c "
cd /sgl-workspace/aiter && \
AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${GPU_ARCH_LIST} python3 setup.py develop
scripts/ci/amd/amd_ci_start_container.sh
+2 -2
View File
@@ -28,7 +28,7 @@ LOCAL_DOCKER_REGISTRY="10.44.14.109:5000"
# Parse command line arguments
MI30X_BASE_TAG="${DEFAULT_MI30X_BASE_TAG}"
MI35X_BASE_TAG="${DEFAULT_MI35X_BASE_TAG}"
CUSTOM_IMAGE=""
CUSTOM_IMAGE="${AMD_CI_IMAGE:-}"
BUILD_FROM_DOCKERFILE=""
GPU_ARCH_BUILD=""
@@ -50,7 +50,7 @@ while [[ $# -gt 0 ]]; do
echo "Options:"
echo " --mi30x-base-tag TAG Override MI30x base image tag"
echo " --mi35x-base-tag TAG Override MI35x base image tag"
echo " --custom-image IMAGE Use a specific Docker image directly"
echo " --custom-image IMAGE Use a specific Docker image directly (or set AMD_CI_IMAGE)"
echo " --build-from-dockerfile Build image from docker/rocm.Dockerfile"
echo " --gpu-arch ARCH GPU architecture for Dockerfile build (e.g., gfx950-rocm720)"
echo " --rocm-version VERSION Override ROCm version for image lookup (e.g., rocm720)"
scripts/ci/amd/amd_ci_start_container_disagg.sh
+26 -14
View File
@@ -28,11 +28,13 @@ LOCAL_DOCKER_REGISTRY="10.44.14.109:5000"
# Parse command line arguments
MI30X_BASE_TAG="${DEFAULT_MI30X_BASE_TAG}"
MI35X_BASE_TAG="${DEFAULT_MI35X_BASE_TAG}"
CUSTOM_IMAGE="${AMD_CI_IMAGE:-}"
while [[ $# -gt 0 ]]; do
case $1 in
--mi30x-base-tag) MI30X_BASE_TAG="$2"; shift 2;;
--mi35x-base-tag) MI35X_BASE_TAG="$2"; shift 2;;
--custom-image) CUSTOM_IMAGE="$2"; shift 2;;
--rocm-version)
ROCM_VERSION="$2"
MI30X_BASE_TAG="${SGLANG_VERSION}-${ROCM_VERSION}-mi30x"
@@ -40,7 +42,7 @@ while [[ $# -gt 0 ]]; do
echo "Using ROCm version override: ${ROCM_VERSION}"
shift 2;;
-h|--help)
echo "Usage: $0 [--mi30x-base-tag TAG] [--mi35x-base-tag TAG] [--rocm-version VERSION]"
echo "Usage: $0 [--mi30x-base-tag TAG] [--mi35x-base-tag TAG] [--custom-image IMAGE] [--rocm-version VERSION]"
exit 0
;;
*) echo "Unknown option $1"; exit 1;;
@@ -212,20 +214,30 @@ find_latest_image() {
esac
}
# Pull and run the latest image
IMAGE=$(find_latest_image "${GPU_ARCH}")
# Try the local docker registry first (avoids Docker Hub rate limits and is
# faster on the LAN); if that fails for any reason, fall back to the
# public registry with exponential-backoff retries. Capture stderr so the
# real failure reason (TLS handshake, 404, connection refused, etc.) is
# visible in the job log instead of being silently swallowed.
if local_pull_output=$(docker pull "${LOCAL_DOCKER_REGISTRY}/${IMAGE}" 2>&1); then
echo "Pulled from local docker registry: ${LOCAL_DOCKER_REGISTRY}/${IMAGE}"
docker tag "${LOCAL_DOCKER_REGISTRY}/${IMAGE}" "${IMAGE}"
# Determine which image to use
if [[ -n "${CUSTOM_IMAGE}" ]]; then
IMAGE="${CUSTOM_IMAGE}"
echo "Using custom image: ${IMAGE}"
if [[ "${IMAGE}" == "${LOCAL_DOCKER_REGISTRY}/"* ]]; then
docker pull "${IMAGE}"
else
retry_with_backoff 6 docker pull "${IMAGE}"
fi
else
echo "Local docker registry pull failed; falling back to public registry: ${IMAGE}" >&2
printf '%s\n' "${local_pull_output}" | sed 's/^/ [local-pull] /' >&2
retry_with_backoff 6 docker pull "${IMAGE}"
IMAGE=$(find_latest_image "${GPU_ARCH}")
# Try the local docker registry first (avoids Docker Hub rate limits and is
# faster on the LAN); if that fails for any reason, fall back to the
# public registry with exponential-backoff retries. Capture stderr so the
# real failure reason (TLS handshake, 404, connection refused, etc.) is
# visible in the job log instead of being silently swallowed.
if local_pull_output=$(docker pull "${LOCAL_DOCKER_REGISTRY}/${IMAGE}" 2>&1); then
echo "Pulled from local docker registry: ${LOCAL_DOCKER_REGISTRY}/${IMAGE}"
docker tag "${LOCAL_DOCKER_REGISTRY}/${IMAGE}" "${IMAGE}"
else
echo "Local docker registry pull failed; falling back to public registry: ${IMAGE}" >&2
printf '%s\n' "${local_pull_output}" | sed 's/^/ [local-pull] /' >&2
retry_with_backoff 6 docker pull "${IMAGE}"
fi
fi
CACHE_HOST=/home/runner/sglang-data