1#!/usr/bin/env bash 2# Script used only in CD pipeline 3 4set -eou pipefail 5 6TOPDIR=$(git rev-parse --show-toplevel) 7 8image="$1" 9shift 10 11if [ -z "${image}" ]; then 12 echo "Usage: $0 IMAGE" 13 exit 1 14fi 15 16DOCKER_IMAGE="pytorch/${image}" 17 18DOCKER_REGISTRY="${DOCKER_REGISTRY:-docker.io}" 19 20GPU_ARCH_TYPE=${GPU_ARCH_TYPE:-cpu} 21GPU_ARCH_VERSION=${GPU_ARCH_VERSION:-} 22MANY_LINUX_VERSION=${MANY_LINUX_VERSION:-} 23DOCKERFILE_SUFFIX=${DOCKERFILE_SUFFIX:-} 24WITH_PUSH=${WITH_PUSH:-} 25 26case ${GPU_ARCH_TYPE} in 27 cpu) 28 TARGET=cpu_final 29 DOCKER_TAG=cpu 30 GPU_IMAGE=centos:7 31 DOCKER_GPU_BUILD_ARG=" --build-arg DEVTOOLSET_VERSION=9" 32 ;; 33 cpu-manylinux_2_28) 34 TARGET=cpu_final 35 DOCKER_TAG=cpu 36 GPU_IMAGE=amd64/almalinux:8 37 DOCKER_GPU_BUILD_ARG=" --build-arg DEVTOOLSET_VERSION=11" 38 MANY_LINUX_VERSION="2_28" 39 ;; 40 cpu-aarch64) 41 TARGET=final 42 DOCKER_TAG=cpu-aarch64 43 GPU_IMAGE=arm64v8/centos:7 44 DOCKER_GPU_BUILD_ARG=" --build-arg DEVTOOLSET_VERSION=10" 45 MANY_LINUX_VERSION="aarch64" 46 ;; 47 cpu-aarch64-2_28) 48 TARGET=final 49 DOCKER_TAG=cpu-aarch64 50 GPU_IMAGE=arm64v8/almalinux:8 51 DOCKER_GPU_BUILD_ARG=" --build-arg DEVTOOLSET_VERSION=11" 52 MANY_LINUX_VERSION="2_28_aarch64" 53 ;; 54 cpu-cxx11-abi) 55 TARGET=final 56 DOCKER_TAG=cpu-cxx11-abi 57 GPU_IMAGE="" 58 DOCKER_GPU_BUILD_ARG=" --build-arg DEVTOOLSET_VERSION=9" 59 MANY_LINUX_VERSION="cxx11-abi" 60 ;; 61 cpu-s390x) 62 TARGET=final 63 DOCKER_TAG=cpu-s390x 64 GPU_IMAGE=redhat/ubi9 65 DOCKER_GPU_BUILD_ARG="" 66 MANY_LINUX_VERSION="s390x" 67 ;; 68 cuda) 69 TARGET=cuda_final 70 DOCKER_TAG=cuda${GPU_ARCH_VERSION} 71 # Keep this up to date with the minimum version of CUDA we currently support 72 GPU_IMAGE=centos:7 73 DOCKER_GPU_BUILD_ARG="--build-arg BASE_CUDA_VERSION=${GPU_ARCH_VERSION} --build-arg DEVTOOLSET_VERSION=9" 74 ;; 75 cuda-manylinux_2_28) 76 TARGET=cuda_final 77 DOCKER_TAG=cuda${GPU_ARCH_VERSION} 78 GPU_IMAGE=amd64/almalinux:8 79 DOCKER_GPU_BUILD_ARG="--build-arg BASE_CUDA_VERSION=${GPU_ARCH_VERSION} --build-arg DEVTOOLSET_VERSION=11" 80 MANY_LINUX_VERSION="2_28" 81 ;; 82 cuda-aarch64) 83 TARGET=cuda_final 84 DOCKER_TAG=cuda${GPU_ARCH_VERSION} 85 GPU_IMAGE=arm64v8/centos:7 86 DOCKER_GPU_BUILD_ARG="--build-arg BASE_CUDA_VERSION=${GPU_ARCH_VERSION} --build-arg DEVTOOLSET_VERSION=11" 87 MANY_LINUX_VERSION="aarch64" 88 DOCKERFILE_SUFFIX="_cuda_aarch64" 89 ;; 90 rocm) 91 TARGET=rocm_final 92 DOCKER_TAG=rocm${GPU_ARCH_VERSION} 93 GPU_IMAGE=rocm/dev-centos-7:${GPU_ARCH_VERSION}-complete 94 PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx1030;gfx1100" 95 ROCM_REGEX="([0-9]+)\.([0-9]+)[\.]?([0-9]*)" 96 if [[ $GPU_ARCH_VERSION =~ $ROCM_REGEX ]]; then 97 ROCM_VERSION_INT=$((${BASH_REMATCH[1]}*10000 + ${BASH_REMATCH[2]}*100 + ${BASH_REMATCH[3]:-0})) 98 else 99 echo "ERROR: rocm regex failed" 100 exit 1 101 fi 102 if [[ $ROCM_VERSION_INT -ge 60000 ]]; then 103 PYTORCH_ROCM_ARCH+=";gfx942" 104 fi 105 DOCKER_GPU_BUILD_ARG="--build-arg ROCM_VERSION=${GPU_ARCH_VERSION} --build-arg PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} --build-arg DEVTOOLSET_VERSION=9" 106 ;; 107 xpu) 108 TARGET=xpu_final 109 DOCKER_TAG=xpu 110 GPU_IMAGE=amd64/almalinux:8 111 DOCKER_GPU_BUILD_ARG=" --build-arg DEVTOOLSET_VERSION=11" 112 MANY_LINUX_VERSION="2_28" 113 ;; 114 *) 115 echo "ERROR: Unrecognized GPU_ARCH_TYPE: ${GPU_ARCH_TYPE}" 116 exit 1 117 ;; 118esac 119 120IMAGES='' 121 122if [[ -n ${MANY_LINUX_VERSION} && -z ${DOCKERFILE_SUFFIX} ]]; then 123 DOCKERFILE_SUFFIX=_${MANY_LINUX_VERSION} 124fi 125( 126 set -x 127 DOCKER_BUILDKIT=1 docker build \ 128 ${DOCKER_GPU_BUILD_ARG} \ 129 --build-arg "GPU_IMAGE=${GPU_IMAGE}" \ 130 --target "${TARGET}" \ 131 -t "${DOCKER_IMAGE}" \ 132 $@ \ 133 -f "${TOPDIR}/.ci/docker/manywheel/Dockerfile${DOCKERFILE_SUFFIX}" \ 134 "${TOPDIR}/.ci/docker/" 135) 136 137GITHUB_REF=${GITHUB_REF:-$(git symbolic-ref -q HEAD || git describe --tags --exact-match)} 138GIT_BRANCH_NAME="2.5" 139GIT_COMMIT_SHA=${GITHUB_SHA:-$(git rev-parse HEAD)} 140DOCKER_IMAGE_BRANCH_TAG=${DOCKER_IMAGE}-${GIT_BRANCH_NAME} 141DOCKER_IMAGE_SHA_TAG=${DOCKER_IMAGE}-${GIT_COMMIT_SHA} 142 143if [[ "${WITH_PUSH}" == true ]]; then 144 ( 145 set -x 146 docker push "${DOCKER_IMAGE}" 147 if [[ -n ${GITHUB_REF} ]]; then 148 docker tag ${DOCKER_IMAGE} ${DOCKER_IMAGE_BRANCH_TAG} 149 docker tag ${DOCKER_IMAGE} ${DOCKER_IMAGE_SHA_TAG} 150 docker push "${DOCKER_IMAGE_BRANCH_TAG}" 151 docker push "${DOCKER_IMAGE_SHA_TAG}" 152 fi 153 ) 154fi 155