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_NAME="pytorch/${image}" 15 16 17export DOCKER_BUILDKIT=1 18TOPDIR=$(git rev-parse --show-toplevel) 19 20CUDA_VERSION=${CUDA_VERSION:-12.1} 21 22case ${CUDA_VERSION} in 23 cpu) 24 BASE_TARGET=base 25 DOCKER_TAG=cpu 26 ;; 27 all) 28 BASE_TARGET=all_cuda 29 DOCKER_TAG=latest 30 ;; 31 *) 32 BASE_TARGET=cuda${CUDA_VERSION} 33 DOCKER_TAG=cuda${CUDA_VERSION} 34 ;; 35esac 36 37 38( 39 set -x 40 docker build \ 41 --target final \ 42 --progress plain \ 43 --build-arg "BASE_TARGET=${BASE_TARGET}" \ 44 --build-arg "CUDA_VERSION=${CUDA_VERSION}" \ 45 --build-arg "DEVTOOLSET_VERSION=9" \ 46 -t ${DOCKER_IMAGE_NAME} \ 47 $@ \ 48 -f "${TOPDIR}/.ci/docker/conda/Dockerfile" \ 49 ${TOPDIR}/.ci/docker/ 50) 51 52if [[ "${DOCKER_TAG}" =~ ^cuda* ]]; then 53 # Test that we're using the right CUDA compiler 54 ( 55 set -x 56 docker run --rm "${DOCKER_IMAGE_NAME}" nvcc --version | grep "cuda_${CUDA_VERSION}" 57 ) 58fi 59 60GITHUB_REF=${GITHUB_REF:-$(git symbolic-ref -q HEAD || git describe --tags --exact-match)} 61GIT_BRANCH_NAME=${GITHUB_REF##*/} 62GIT_COMMIT_SHA=${GITHUB_SHA:-$(git rev-parse HEAD)} 63DOCKER_IMAGE_BRANCH_TAG=${DOCKER_IMAGE_NAME}-${GIT_BRANCH_NAME} 64DOCKER_IMAGE_SHA_TAG=${DOCKER_IMAGE_NAME}-${GIT_COMMIT_SHA} 65if [[ "${WITH_PUSH:-}" == true ]]; then 66 ( 67 set -x 68 docker push "${DOCKER_IMAGE_NAME}" 69 if [[ -n ${GITHUB_REF} ]]; then 70 docker tag ${DOCKER_IMAGE_NAME} ${DOCKER_IMAGE_BRANCH_TAG} 71 docker tag ${DOCKER_IMAGE_NAME} ${DOCKER_IMAGE_SHA_TAG} 72 docker push "${DOCKER_IMAGE_BRANCH_TAG}" 73 docker push "${DOCKER_IMAGE_SHA_TAG}" 74 fi 75 ) 76fi 77