| 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 against torch nightly and tag it with the |
| 14 | # -torch-nightly-arm64 suffix the DGX Spark test steps pull on the nightly lane |
| 15 | # (ci-infra get_image(arm64=True) appends -torch-nightly before -arm64). |
| 16 | PYTORCH_NIGHTLY_ARGS=() |
| 17 | if [[ "${TORCH_NIGHTLY:-0}" == "1" ]]; then |
| 18 | IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-torch-nightly-arm64" |
| 19 | PYTORCH_NIGHTLY_ARGS=(--build-arg PYTORCH_NIGHTLY=1) |
| 20 | else |
| 21 | IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-arm64" |
| 22 | fi |
| 23 | |
| 24 | # authenticate with AWS ECR |
| 25 | aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true |
| 26 | |
| 27 | # skip build if image already exists |
| 28 | if docker manifest inspect "$IMAGE" >/dev/null 2>&1; then |
| 29 | echo "Image found" |
| 30 | else |
| 31 | echo "Image not found, proceeding with build..." |
| 32 | # build for arm64 GPU targets: Grace/GH200 (sm_90), |
| 33 | # Blackwell/Thor (sm_100/sm_103/sm_110), and DGX Spark/GB10 |
| 34 | # (sm_121, family-covered by 12.0 under CUDA 13) |
| 35 | docker build --file docker/Dockerfile \ |
| 36 | --platform linux/arm64 \ |
| 37 | --build-arg max_jobs=16 \ |
| 38 | --build-arg nvcc_threads=4 \ |
| 39 | --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0-b8b5f17a7d9ccfc25bbc5cf17b3fcea12964a042 \ |
| 40 | --build-arg torch_cuda_arch_list="9.0 10.0 11.0 12.0" \ |
| 41 | --build-arg USE_SCCACHE=1 \ |
| 42 | --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ |
| 43 | "${PYTORCH_NIGHTLY_ARGS[@]}" \ |
| 44 | --tag "$IMAGE" \ |
| 45 | --target test \ |
| 46 | --progress plain . |
| 47 | # push |
| 48 | docker push "$IMAGE" |
| 49 | fi |
| 50 | |
| 51 | .buildkite/scripts/annotate-image-build.sh "$IMAGE" |