1#!/usr/bin/env bash 2 3set -eou pipefail 4 5DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" 6source "${DIR}/promote/common_utils.sh" 7 8if [[ -z "${RESTORE_FROM:-}" ]]; then 9 echo "ERROR: RESTORE_FROM environment variable must be specified" 10 echo " example: RESTORE_FROM=v1.6.0-rc3 ${0}" 11 exit 1 12fi 13 14DRY_RUN=${DRY_RUN:-enabled} 15 16PYTORCH_S3_BACKUP_BUCKET=${PYTORCH_S3_BACKUP_BUCKET:-s3://pytorch-backup/${RESTORE_FROM}} 17PYTORCH_S3_TEST_BUCKET=${PYTORCH_S3_TEST_BUCKET:-s3://pytorch/} 18PYTORCH_S3_FROM=${PYTORCH_S3_FROM:-${PYTORCH_S3_BACKUP_BUCKET}} 19PYTORCH_S3_TO=${PYTORCH_S3_TO:-s3://pytorch/} 20 21restore_wheels() { 22 aws_promote torch whl 23} 24 25restore_libtorch() { 26 aws_promote libtorch-* libtorch 27} 28 29ANACONDA="true anaconda" 30if [[ ${DRY_RUN} = "disabled" ]]; then 31 ANACONDA="anaconda" 32fi 33PYTORCH_CONDA_TO=${PYTORCH_CONDA_TO:-pytorch-test} 34 35upload_conda() { 36 local pkg 37 pkg=${1} 38 ( 39 set -x 40 ${ANACONDA} upload --skip -u "${PYTORCH_CONDA_TO}" "${pkg}" 41 ) 42} 43 44export -f upload_conda 45 46restore_conda() { 47 TMP_DIR="$(mktemp -d)" 48 trap 'rm -rf ${TMP_DIR}' EXIT 49 ( 50 set -x 51 aws s3 cp --recursive "${PYTORCH_S3_BACKUP_BUCKET}/conda" "${TMP_DIR}/" 52 ) 53 export ANACONDA 54 export PYTORCH_CONDA_TO 55 # Should upload all bz2 packages in parallel for quick restoration 56 find "${TMP_DIR}" -name '*.bz2' -type f \ 57 | xargs -P 10 -I % bash -c "(declare -t upload_conda); upload_conda %" 58} 59 60 61restore_wheels 62restore_libtorch 63restore_conda 64