1#!/usr/bin/env bash 2# Script used only in CD pipeline 3 4set -eou pipefail 5 6image="$1" 7shift 8 9if [ -z "${image}" ]; then 10 echo "Usage: $0 IMAGE" 11 exit 1 12fi 13 14DOCKER_IMAGE="pytorch/${image}" 15 16TOPDIR=$(git rev-parse --show-toplevel) 17 18GPU_ARCH_TYPE=${GPU_ARCH_TYPE:-cpu} 19GPU_ARCH_VERSION=${GPU_ARCH_VERSION:-} 20 21WITH_PUSH=${WITH_PUSH:-} 22 23DOCKER=${DOCKER:-docker} 24 25case ${GPU_ARCH_TYPE} in 26 cpu) 27 BASE_TARGET=cpu 28 DOCKER_TAG=cpu 29 GPU_IMAGE=ubuntu:20.04 30 DOCKER_GPU_BUILD_ARG="" 31 ;; 32 cuda) 33 BASE_TARGET=cuda${GPU_ARCH_VERSION} 34 DOCKER_TAG=cuda${GPU_ARCH_VERSION} 35 GPU_IMAGE=ubuntu:20.04 36 DOCKER_GPU_BUILD_ARG="" 37 ;; 38 rocm) 39 BASE_TARGET=rocm 40 DOCKER_TAG=rocm${GPU_ARCH_VERSION} 41 GPU_IMAGE=rocm/dev-ubuntu-20.04:${GPU_ARCH_VERSION}-complete 42 PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx1030;gfx1100" 43 ROCM_REGEX="([0-9]+)\.([0-9]+)[\.]?([0-9]*)" 44 if [[ $GPU_ARCH_VERSION =~ $ROCM_REGEX ]]; then 45 ROCM_VERSION_INT=$((${BASH_REMATCH[1]}*10000 + ${BASH_REMATCH[2]}*100 + ${BASH_REMATCH[3]:-0})) 46 else 47 echo "ERROR: rocm regex failed" 48 exit 1 49 fi 50 if [[ $ROCM_VERSION_INT -ge 60000 ]]; then 51 PYTORCH_ROCM_ARCH+=";gfx942" 52 fi 53 DOCKER_GPU_BUILD_ARG="--build-arg PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}" 54 ;; 55 *) 56 echo "ERROR: Unrecognized GPU_ARCH_TYPE: ${GPU_ARCH_TYPE}" 57 exit 1 58 ;; 59esac 60 61 62( 63 set -x 64 DOCKER_BUILDKIT=1 ${DOCKER} build \ 65 --target final \ 66 ${DOCKER_GPU_BUILD_ARG} \ 67 --build-arg "GPU_IMAGE=${GPU_IMAGE}" \ 68 --build-arg "BASE_TARGET=${BASE_TARGET}" \ 69 -t "${DOCKER_IMAGE}" \ 70 $@ \ 71 -f "${TOPDIR}/.ci/docker/libtorch/Dockerfile" \ 72 "${TOPDIR}/.ci/docker/" 73 74) 75 76GITHUB_REF=${GITHUB_REF:-$(git symbolic-ref -q HEAD || git describe --tags --exact-match)} 77GIT_BRANCH_NAME=${GITHUB_REF##*/} 78GIT_COMMIT_SHA=${GITHUB_SHA:-$(git rev-parse HEAD)} 79DOCKER_IMAGE_BRANCH_TAG=${DOCKER_IMAGE}-${GIT_BRANCH_NAME} 80DOCKER_IMAGE_SHA_TAG=${DOCKER_IMAGE}-${GIT_COMMIT_SHA} 81 82if [[ "${WITH_PUSH}" == true ]]; then 83 ( 84 set -x 85 ${DOCKER} push "${DOCKER_IMAGE}" 86 if [[ -n ${GITHUB_REF} ]]; then 87 ${DOCKER} tag ${DOCKER_IMAGE} ${DOCKER_IMAGE_BRANCH_TAG} 88 ${DOCKER} tag ${DOCKER_IMAGE} ${DOCKER_IMAGE_SHA_TAG} 89 ${DOCKER} push "${DOCKER_IMAGE_BRANCH_TAG}" 90 ${DOCKER} push "${DOCKER_IMAGE_SHA_TAG}" 91 fi 92 ) 93fi 94