xref: /aosp_15_r20/external/pytorch/scripts/release/promote/conda_to_conda.sh (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1#!/usr/bin/env bash
2
3set -eou pipefail
4
5DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
6source "${DIR}/common_utils.sh"
7
8
9# Allow for users to pass PACKAGE_NAME
10# For use with other packages, i.e. torchvision, etc.
11PACKAGE_NAME=${PACKAGE_NAME:-pytorch}
12PYTORCH_CONDA_FROM=${PYTORCH_CONDA_FROM:-pytorch-test}
13PYTORCH_CONDA_TO=${PYTORCH_CONDA_TO:-pytorch}
14CONDA_PLATFORMS="linux-64 osx-64 win-64 noarch"
15
16pytorch_version="$(get_pytorch_version)"
17
18tmp_dir="$(mktemp -d)"
19pushd "${tmp_dir}"
20trap 'rm -rf ${tmp_dir}' EXIT
21
22conda_search() {
23    conda search -q "${PYTORCH_CONDA_FROM}::${PACKAGE_NAME}==${pytorch_version}" -c "${PYTORCH_CONDA_FROM}" --platform "${platform}" \
24        | grep -e "^${PACKAGE_NAME}" \
25        | awk -F ' *' '{print $3}' \
26        | xargs -I % echo "https://anaconda.org/${PYTORCH_CONDA_FROM}/${PACKAGE_NAME}/${pytorch_version}/download/${platform}/${PACKAGE_NAME}-${pytorch_version}-%.tar.bz2"
27}
28
29pkgs_to_download=()
30for platform in ${CONDA_PLATFORMS}; do
31    pkgs_to_download+=($(\
32        conda_search 2>/dev/null || true
33    ))
34    # Create directory where packages will eventually be downloaded
35    mkdir -p "${platform}"
36done
37
38my_curl() {
39    local dl_url=$1
40    local start=$(date +%s)
41    # downloads should be distinguished by platform which should be the second
42    # to last field in the url, this is to avoid clobbering same named files
43    # for different platforms
44    dl_dir=$(echo "${dl_url}" | rev | cut -d'/' -f 2 | rev)
45    dl_name=$(echo "${dl_url}" | rev | cut -d'/' -f 1 | rev)
46    curl -fsSL -o "${dl_dir}/${dl_name}" "${dl_url}"
47    local end=$(date +%s)
48    local diff=$(( end - start ))
49    echo "+ ${dl_url} took ${diff}s"
50}
51export -f my_curl
52
53# Download all packages in parallel
54printf '%s\n' "${pkgs_to_download[@]}" \
55    | xargs -P 10 -I % bash -c '(declare -t my_curl); my_curl %'
56
57# dry run by default
58DRY_RUN=${DRY_RUN:-enabled}
59ANACONDA="true anaconda"
60if [[ $DRY_RUN = "disabled" ]]; then
61    ANACONDA="anaconda"
62fi
63(
64    # We use --skip here to avoid re-uploading files we've already uploaded
65    set -x
66    ${ANACONDA} upload --skip -u ${PYTORCH_CONDA_TO} $(find . -name '*.bz2')
67)
68
69popd
70