scripts: Change promote pypi to be more flexible (#53774)

Summary:
Promotion to PyPI should be more flexible to allow any package to be
promoted to PyPI.

After we re-added a version suffix to cuda 10.2 it means that this
script needs to have the flexibility to designate which platform and
which version suffix will actually be uploaded to PyPI

Should coincide with https://github.com/pytorch/builder/pull/678

Signed-off-by: Eli Uriegas <eliuriegas@fb.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/53774

Reviewed By: jbschlosser

Differential Revision: D27052347

Pulled By: seemethere

fbshipit-source-id: 71129cc5afbd7de448c970ef721bc979c3420586

scripts/release/promote/prep_binary_for_pypi.sh

@@ -12,6 +12,8 @@
 set -eou pipefail
 shopt -s globstar
 
+OUTPUT_DIR=${OUTPUT_DIR:-$(pwd)}
+
 tmp_dir="$(mktemp -d)"
 trap 'rm -rf ${tmp_dir}' EXIT
 
@@ -27,7 +29,7 @@ for whl_file in "$@"; do
     version_with_suffix_escaped=${version_with_suffix/+/%2B}
     # Remove all suffixed +bleh versions
     version_no_suffix=${version_with_suffix/+*/}
-    new_whl_file=${whl_file/${version_with_suffix_escaped}/${version_no_suffix}}
+    new_whl_file=${OUTPUT_DIR}/$(basename "${whl_file/${version_with_suffix_escaped}/${version_no_suffix}}")
     dist_info_folder=$(find "${whl_dir}" -type d -name '*.dist-info' | head -1)
     basename_dist_info_folder=$(basename "${dist_info_folder}")
     dirname_dist_info_folder=$(dirname "${dist_info_folder}")

scripts/release/promote/wheel_to_pypi.sh

@@ -10,19 +10,26 @@ source "${DIR}/common_utils.sh"
 PACKAGE_NAME=${PACKAGE_NAME:-torch}
 
 pytorch_version="$(get_pytorch_version)"
+# Refers to the specific package we'd like to promote
+# i.e. VERSION_SUFFIX='%2Bcu102'
+#      torch-1.8.0+cu102 -> torch-1.8.0
+VERSION_SUFFIX=${VERSION_SUFFIX:-}
+# Refers to the specific platofmr we'd like to promote
+# i.e. PLATFORM=linux_x86_64
+# For domains like torchaudio / torchtext this is to be left blank
+PLATFORM=${PLATFORM:-}
 
-# This assumes you have already promoted the wheels to stable S3
 pkgs_to_promote=$(\
     curl -fsSL https://download.pytorch.org/whl/torch_stable.html \
-        | grep "${PACKAGE_NAME}-${pytorch_version}" \
-        | grep -v "%2B" \
-        | grep -v "win_amd64" \
+        | grep "${PACKAGE_NAME}-${pytorch_version}${VERSION_SUFFIX}-" \
+        | grep "${PLATFORM}" \
         | cut -d '"' -f2
 )
 
 tmp_dir="$(mktemp -d)"
-trap 'rm -rf ${tmp_dir}' EXIT
-pushd "${tmp_dir}"
+output_tmp_dir="$(mktemp -d)"
+trap 'rm -rf ${tmp_dir} ${output_tmp_dir}' EXIT
+pushd "${output_tmp_dir}"
 
 # Dry run by default
 DRY_RUN=${DRY_RUN:-enabled}
@@ -34,13 +41,25 @@ fi
 
 for pkg in ${pkgs_to_promote}; do
     pkg_basename="$(basename "${pkg//linux/manylinux1}")"
+    orig_pkg="${tmp_dir}/${pkg_basename}"
     (
         set -x
         # Download package, sub out linux for manylinux1
-        curl -fsSL -o "${pkg_basename}" "https://download.pytorch.org/whl/${pkg}"
+        curl -fsSL -o "${orig_pkg}" "https://download.pytorch.org/whl/${pkg}"
+    )
+
+    if [[ -n "${VERSION_SUFFIX}" ]]; then
+        OUTPUT_DIR="${output_tmp_dir}" ${DIR}/prep_binary_for_pypi.sh "${orig_pkg}"
+    else
+        mv "${orig_pkg}" "${output_tmp_dir}/"
+    fi
+
+    (
+        set -x
         ${TWINE_UPLOAD} \
             --disable-progress-bar \
             --non-interactive \
-            "${pkg_basename}"
+            ./*.whl
+        rm -rf ./*.whl
     )
 done