| 1 | #!/bin/bash |
| 2 | set -e |
| 3 | |
| 4 | if [[ $# -lt 3 ]]; then |
| 5 | echo "Usage: $0 <registry> <repo> <commit>" |
| 6 | exit 1 |
| 7 | fi |
| 8 | |
| 9 | REGISTRY=$1 |
| 10 | REPO=$2 |
| 11 | BUILDKITE_COMMIT=$3 |
| 12 | |
| 13 | # When TORCH_NIGHTLY=1, build the CPU image against torch nightly and tag it |
| 14 | # with the -torch-nightly-cpu suffix the test steps pull on the nightly lane. |
| 15 | PYTORCH_NIGHTLY_ARGS=() |
| 16 | if [[ "${TORCH_NIGHTLY:-0}" == "1" ]]; then |
| 17 | IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-torch-nightly-cpu" |
| 18 | PYTORCH_NIGHTLY_ARGS=(--build-arg PYTORCH_NIGHTLY=1) |
| 19 | else |
| 20 | IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu" |
| 21 | fi |
| 22 | |
| 23 | # replace invalid characters in Docker image tags and truncate to 128 chars |
| 24 | clean_docker_tag() { |
| 25 | local input="$1" |
| 26 | echo "$input" | sed 's/[^a-zA-Z0-9._-]/_/g' | cut -c1-128 |
| 27 | } |
| 28 | |
| 29 | # resolve and set: CACHE_TO, CACHE_FROM, CACHE_FROM_BASE_BRANCH, CACHE_FROM_MAIN |
| 30 | # Reuses the same ECR cache repos as the CUDA image build, with an |
| 31 | # "x86_cpu" suffix on every tag so CPU and CUDA cache blobs stay isolated. |
| 32 | # The nightly lane appends "-nightly": it resolves a different torch, so |
| 33 | # sharing a cache ref with the regular lane would have the two evict each |
| 34 | # other from the same ref on every run. |
| 35 | prepare_cache_tags() { |
| 36 | TEST_CACHE_ECR="936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache" |
| 37 | MAIN_CACHE_ECR="936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache" |
| 38 | CACHE_SUFFIX="x86_cpu" |
| 39 | if [[ "${TORCH_NIGHTLY:-0}" == "1" ]]; then |
| 40 | CACHE_SUFFIX="x86_cpu-nightly" |
| 41 | fi |
| 42 | |
| 43 | if [[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" ]]; then |
| 44 | if [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then |
| 45 | cache="${MAIN_CACHE_ECR}:latest-${CACHE_SUFFIX}" |
| 46 | else |
| 47 | clean_branch=$(clean_docker_tag "${BUILDKITE_BRANCH:-unknown}") |
| 48 | cache="${TEST_CACHE_ECR}:${clean_branch}-${CACHE_SUFFIX}" |
| 49 | fi |
| 50 | CACHE_TO="$cache" |
| 51 | CACHE_FROM="$cache" |
| 52 | CACHE_FROM_BASE_BRANCH="$cache" |
| 53 | else |
| 54 | CACHE_TO="${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}-${CACHE_SUFFIX}" |
| 55 | CACHE_FROM="${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}-${CACHE_SUFFIX}" |
| 56 | if [[ "${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" == "main" ]]; then |
| 57 | CACHE_FROM_BASE_BRANCH="${MAIN_CACHE_ECR}:latest-${CACHE_SUFFIX}" |
| 58 | else |
| 59 | clean_base=$(clean_docker_tag "${BUILDKITE_PULL_REQUEST_BASE_BRANCH}") |
| 60 | CACHE_FROM_BASE_BRANCH="${TEST_CACHE_ECR}:${clean_base}-${CACHE_SUFFIX}" |
| 61 | fi |
| 62 | fi |
| 63 | |
| 64 | CACHE_FROM_MAIN="${MAIN_CACHE_ECR}:latest-${CACHE_SUFFIX}" |
| 65 | } |
| 66 | |
| 67 | # authenticate with AWS ECR (public, for the image; private, for the cache repos) |
| 68 | aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true |
| 69 | aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true |
| 70 | |
| 71 | # skip build if image already exists |
| 72 | if docker manifest inspect "$IMAGE" >/dev/null 2>&1; then |
| 73 | echo "Image found" |
| 74 | else |
| 75 | echo "Image not found, proceeding with build..." |
| 76 | |
| 77 | prepare_cache_tags |
| 78 | echo "--- :mag: Cache tags" |
| 79 | echo "CACHE_TO: ${CACHE_TO}" |
| 80 | echo "CACHE_FROM: ${CACHE_FROM}" |
| 81 | echo "CACHE_FROM_BASE_BRANCH: ${CACHE_FROM_BASE_BRANCH}" |
| 82 | echo "CACHE_FROM_MAIN: ${CACHE_FROM_MAIN}" |
| 83 | |
| 84 | # dedupe cache-from refs so identical branch/main fallbacks aren't repeated |
| 85 | CACHE_FROM_ARGS=() |
| 86 | SEEN_REFS="" |
| 87 | for ref in "$CACHE_FROM" "$CACHE_FROM_BASE_BRANCH" "$CACHE_FROM_MAIN"; do |
| 88 | if [[ "$SEEN_REFS" != *"|${ref}|"* ]]; then |
| 89 | CACHE_FROM_ARGS+=(--cache-from "type=registry,ref=${ref}") |
| 90 | SEEN_REFS="${SEEN_REFS}|${ref}|" |
| 91 | fi |
| 92 | done |
| 93 | |
| 94 | echo "--- :docker: Setting up buildx" |
| 95 | # network=host so the isolated docker-container builder can still reach the |
| 96 | # EC2 instance metadata service for sccache's S3 credentials; recreate on |
| 97 | # every run in case a stale builder without this driver-opt already exists |
| 98 | # on a warm CI host. |
| 99 | docker buildx rm vllm-cpu-builder >/dev/null 2>&1 || true |
| 100 | docker buildx create --name vllm-cpu-builder --driver docker-container --driver-opt network=host --use |
| 101 | docker buildx inspect --bootstrap |
| 102 | |
| 103 | # build and push |
| 104 | docker buildx build --file docker/Dockerfile.cpu \ |
| 105 | --build-arg max_jobs=16 \ |
| 106 | --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ |
| 107 | --build-arg VLLM_CPU_X86=true \ |
| 108 | --build-arg USE_SCCACHE=1 \ |
| 109 | "${PYTORCH_NIGHTLY_ARGS[@]}" \ |
| 110 | --tag "$IMAGE" \ |
| 111 | --target vllm-test \ |
| 112 | "${CACHE_FROM_ARGS[@]}" \ |
| 113 | --cache-to "type=registry,ref=${CACHE_TO},mode=max" \ |
| 114 | --push \ |
| 115 | --progress plain . |
| 116 | fi |
| 117 | |
| 118 | .buildkite/scripts/annotate-image-build.sh "$IMAGE" |