| 1 | #!/bin/bash |
| 2 | set -euo pipefail |
| 3 | |
| 4 | # Build a vLLM test image with PyTorch nightly installed. |
| 5 | # Called by the pipeline generator's "vLLM Against PyTorch Nightly" group. |
| 6 | |
| 7 | if [[ $# -lt 5 ]]; then |
| 8 | echo "Usage: $0 <registry> <repo> <commit> <branch> <image_tag>" |
| 9 | exit 1 |
| 10 | fi |
| 11 | |
| 12 | REGISTRY=$1 |
| 13 | REPO=$2 |
| 14 | BUILDKITE_COMMIT=$3 |
| 15 | BRANCH=$4 |
| 16 | IMAGE_TAG=$5 |
| 17 | |
| 18 | # --- Arguments --- |
| 19 | echo "--- :mag: Arguments" |
| 20 | echo "REGISTRY: ${REGISTRY}" |
| 21 | echo "REPO: ${REPO}" |
| 22 | echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}" |
| 23 | echo "BRANCH: ${BRANCH}" |
| 24 | echo "IMAGE_TAG: ${IMAGE_TAG}" |
| 25 | |
| 26 | # --- ECR login --- |
| 27 | echo "--- :key: ECR login" |
| 28 | aws ecr-public get-login-password --region us-east-1 \ |
| 29 | | docker login --username AWS --password-stdin "$REGISTRY" |
| 30 | aws ecr get-login-password --region us-east-1 \ |
| 31 | | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com |
| 32 | |
| 33 | # --- Set up buildx --- |
| 34 | echo "--- :docker: Setting up buildx" |
| 35 | docker buildx create --name vllm-builder --driver docker-container --use || true |
| 36 | docker buildx inspect --bootstrap |
| 37 | docker buildx ls |
| 38 | |
| 39 | # --- Skip if image already exists --- |
| 40 | echo "--- :mag: Checking if image already exists" |
| 41 | if docker manifest inspect "$IMAGE_TAG" >/dev/null 2>&1; then |
| 42 | echo "Image found: $IMAGE_TAG — skipping build" |
| 43 | .buildkite/scripts/annotate-image-build.sh "$IMAGE_TAG" |
| 44 | exit 0 |
| 45 | fi |
| 46 | echo "Image not found, proceeding with build..." |
| 47 | |
| 48 | # --- CUDA 13.0 for nightly builds --- |
| 49 | # Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9 |
| 50 | NIGHTLY_CUDA_VERSION="13.0.2" |
| 51 | NIGHTLY_BUILD_BASE_IMAGE="pytorch/manylinux2_28-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359" |
| 52 | # Must track ARG UBUNTU_VERSION in docker/Dockerfile. The nightly lane exists to |
| 53 | # run the normal pipeline against nightly torch, so it has to differ from the |
| 54 | # regular image in torch only -- not in the OS underneath it. This was pinned to |
| 55 | # 22.04 while everything else moved to 24.04, and the runtime dependency list in |
| 56 | # tools/install_snapshot_runtime.sh is 24.04-only (libgnutls30t64 is the 64-bit |
| 57 | # time_t rename and does not exist on jammy), so the nightly build failed with |
| 58 | # "E: Unable to locate package libgnutls30t64" while regular CI stayed green on |
| 59 | # a cached layer. |
| 60 | NIGHTLY_UBUNTU_VERSION="24.04" |
| 61 | NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu${NIGHTLY_UBUNTU_VERSION}" |
| 62 | |
| 63 | echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})" |
| 64 | docker buildx build --file docker/Dockerfile \ |
| 65 | --build-arg max_jobs=16 \ |
| 66 | --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ |
| 67 | --build-arg USE_SCCACHE=1 \ |
| 68 | --build-arg PYTORCH_NIGHTLY=1 \ |
| 69 | --build-arg CUDA_VERSION="${NIGHTLY_CUDA_VERSION}" \ |
| 70 | --build-arg BUILD_BASE_IMAGE="${NIGHTLY_BUILD_BASE_IMAGE}" \ |
| 71 | --build-arg FINAL_BASE_IMAGE="${NIGHTLY_FINAL_BASE_IMAGE}" \ |
| 72 | --build-arg torch_cuda_arch_list="8.0 8.9 9.0 10.0 12.0" \ |
| 73 | --tag "$IMAGE_TAG" \ |
| 74 | --push \ |
| 75 | --target test \ |
| 76 | --progress plain . |
| 77 | |
| 78 | echo "--- :white_check_mark: Torch nightly image build complete: $IMAGE_TAG" |
| 79 | |
| 80 | .buildkite/scripts/annotate-image-build.sh "$IMAGE_TAG" |