Files
vLLM/.buildkite/scripts/rocm/stack.sh
T
b2456c622b [ROCm][CI] Make ROCm 10.0 the default image; keep ROCm 7.2 as -rocm72 (#58761)
Signed-off-by: Rohan Potdar <rohan.potdar@amd.com>
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Micah Williamson <micah.williamson@amd.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
2026-10-17 21:43:59 -05:00

79 lines
3.4 KiB
Bash

#!/usr/bin/env bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# Derive everything the release pipeline needs about a ROCm stack from its base
# Dockerfile, so the stack's image is the single input to every ROCm script.
#
# Usage: source .buildkite/scripts/rocm/stack.sh <base-dockerfile>
# e.g. docker/Dockerfile.rocm_base (ROCm 10.0) or docker/Dockerfile.rocm_72_base
#
# Exports:
# ROCM_STACK_BASE_DOCKERFILE the base Dockerfile
# ROCM_STACK_DOCKERFILE the matching vLLM Dockerfile (base name minus _base)
# ROCM_STACK_THEROCK 1 if the ROCm SDK comes from TheRock wheels
# ROCM_STACK_VERSION ROCm version, e.g. 10.0.0 or 7.2.3
# ROCM_STACK_VARIANT wheel variant and Docker Hub tag flavor from major.minor,
# same rule as setup.py (rocm100, rocm72)
# ROCM_STACK_IMAGE_KEY hash of the base Dockerfile (base image cache key)
# ROCM_STACK_WHEEL_KEY hash of the base Dockerfile plus files it COPYs from
# the repo (base wheel cache key)
# ROCM_STACK_ECR_BASE ECR tag of the cached base image
# ROCM_STACK_ECR_IMAGE ECR tag of this commit's release image
# ROCM_STACK_BASE_WHEELS_DIR local dir for base wheels (also the artifact path)
# ROCM_STACK_VLLM_WHEEL_DIR local dir for the vLLM wheel (also the artifact path)
rocm_stack_init() {
local base="$1"
if [[ ! -f "$base" ]]; then
echo "ROCm base Dockerfile not found: $base" >&2
return 2
fi
ROCM_STACK_BASE_DOCKERFILE="$base"
ROCM_STACK_DOCKERFILE="${base%_base}"
local sdk
sdk=$(sed -nE 's/^ARG ROCM_SDK_VERSION=([0-9.]+).*/\1/p' "$base" | head -1)
if [[ -n "$sdk" ]]; then
ROCM_STACK_THEROCK=1
ROCM_STACK_VERSION="$sdk"
else
ROCM_STACK_THEROCK=0
# BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.2.3-complete -> 7.2.3
ROCM_STACK_VERSION=$(sed -nE 's/^ARG BASE_IMAGE=.*:([0-9]+(\.[0-9]+)+).*/\1/p' "$base" | head -1)
fi
if [[ -z "$ROCM_STACK_VERSION" ]]; then
echo "Could not determine the ROCm version from $base" >&2
return 2
fi
local major minor
IFS=. read -r major minor _ <<< "$ROCM_STACK_VERSION"
ROCM_STACK_VARIANT="rocm${major}${minor}"
ROCM_STACK_IMAGE_KEY=$(sha256sum "$base" | cut -c1-16)
local inputs
inputs=$(grep -E '^COPY ' "$base" | grep -v -- '--from' | awk '{print $2}' || true)
if [[ -z "$inputs" ]]; then
ROCM_STACK_WHEEL_KEY="$ROCM_STACK_IMAGE_KEY"
else
# shellcheck disable=SC2086
ROCM_STACK_WHEEL_KEY=$(sha256sum "$base" $inputs | sha256sum | cut -c1-16)
fi
local ecr="public.ecr.aws/q9t5s3a7/vllm-release-repo"
ROCM_STACK_ECR_BASE="${ecr}:${ROCM_STACK_IMAGE_KEY}-${ROCM_STACK_VARIANT}-base"
ROCM_STACK_ECR_IMAGE="${ecr}:${BUILDKITE_COMMIT:-unknown}-${ROCM_STACK_VARIANT}"
ROCM_STACK_BASE_WHEELS_DIR="artifacts/${ROCM_STACK_VARIANT}-base-wheels"
ROCM_STACK_VLLM_WHEEL_DIR="artifacts/${ROCM_STACK_VARIANT}-vllm-wheel"
export ROCM_STACK_BASE_DOCKERFILE ROCM_STACK_DOCKERFILE ROCM_STACK_THEROCK \
ROCM_STACK_VERSION ROCM_STACK_VARIANT ROCM_STACK_IMAGE_KEY \
ROCM_STACK_WHEEL_KEY ROCM_STACK_ECR_BASE ROCM_STACK_ECR_IMAGE \
ROCM_STACK_BASE_WHEELS_DIR ROCM_STACK_VLLM_WHEEL_DIR
}
if [[ $# -gt 0 ]]; then
rocm_stack_init "$1"
fi