From e9f730a3eb30823bef22c3519e15c173f727d4a9 Mon Sep 17 00:00:00 2001
From: Sevag H <shanssian@nvidia.com>
Date: Thu, 23 Feb 2023 11:40:53 -0500
Subject: [PATCH] Modify rapids-twine to discover and upload all wheels (#45)

This changes rapids-twine to more closely resemble the anaconda upload;
all wheels are automatically discovered by looking for the
`wheel_python` string in the S3 path. Then, all the wheels are uploaded
with twine.

This way, the wheel publish workflows don't need to be aware of the
different wheel variants (python versions, architectures, CUDA toolkits,
etc.) that are being built.
---
 tools/rapids-twine | 65 +++++++++++++++++++++++++++++++---------------
 1 file changed, 44 insertions(+), 21 deletions(-)

diff --git a/tools/rapids-twine b/tools/rapids-twine
index 9b2656c..b5e4215 100755
--- a/tools/rapids-twine
+++ b/tools/rapids-twine
@@ -1,34 +1,57 @@
 #!/bin/bash
-# A utility script that wraps twine to support build tags (PEP427)
-#        
+# A utility script that wraps twine to upload all pip wheels of a workflow run
+# 
 # Positional Arguments:
-#   1) wheel dir
+#   1) wheel name
 set -exou pipefail
+source rapids-constants
 export RAPIDS_SCRIPT_NAME="rapids-twine"
 
 if [ -z "$1" ]; then
-  rapids-echo-stderr "Must specify input arguments: WHEEL_DIR"
+  rapids-echo-stderr "Must specify input arguments: WHEEL_NAME"
   exit 1
 fi
-wheeldir="$1"
+WHEEL_NAME="$1"
 
-build_tag="${RAPIDS_PY_WHEEL_BUILD_TAG:-""}"
+WHEEL_SEARCH_KEY="wheel_python_${WHEEL_NAME}"
 
-if [ "${build_tag}" != "" ]; then
-  echo "Checking if build tag ${build_tag} is legal..."
-  if [[ "${build_tag}" =~ [^[:digit:]] ]]; then
-    rapids-echo-stderr "Build tag can only be digits"
-    exit 1
-  fi
+WHEEL_DIR="./dist"
+mkdir -p "${WHEEL_DIR}"
 
-  echo "Need to apply build tag ${build_tag}..."
-  for wheelfile in "${wheeldir}"/*.whl; do
-    wheel_version=$(echo "$wheelfile" | cut -d'-' -f2)
-    replacement_wheel_version="${wheel_version}-${build_tag}"
-    wheelnewfile="${wheelfile/$wheel_version/$replacement_wheel_version}"
+S3_PATH=$(rapids-s3-path)
+BUCKET_PREFIX=${S3_PATH/s3:\/\/${RAPIDS_DOWNLOADS_BUCKET}\//} # removes s3://rapids-downloads/ from s3://rapids-downloads/ci/rmm/...
 
-    mv "${wheelfile}" "${wheelnewfile}"
-  done
-fi
+# shellcheck disable=SC2016
+WHEEL_TARBALLS=$(
+  set -eo pipefail;
+  aws \
+    --output json \
+    s3api list-objects \
+    --bucket "${RAPIDS_DOWNLOADS_BUCKET}" \
+    --prefix "${BUCKET_PREFIX}" \
+    --page-size 100 \
+    --query "Contents[?contains(Key, '${WHEEL_SEARCH_KEY}')].Key" \
+    | jq -c
+)
+export WHEEL_TARBALLS
+
+# first untar them all
+for OBJ in $(jq -nr 'env.WHEEL_TARBALLS | fromjson | .[]'); do
+  FILENAME=$(basename "${OBJ}")
+  S3_URI="${S3_PATH}${FILENAME}"
+
+  rapids-echo-stderr "Untarring ${S3_URI} into ${WHEEL_DIR}"
+  aws s3 cp --only-show-errors "${S3_URI}" - | tar xzf - -C "${WHEEL_DIR}"
+done
+
+# then run twine on all wheels
+export RAPIDS_RETRY_SLEEP=180
+# shellcheck disable=SC2086
+rapids-retry twine \
+  upload \
+  --disable-progress-bar \
+  --non-interactive \
+  --skip-existing \
+  "${WHEEL_DIR}"/*.whl
 
-twine upload --disable-progress-bar --non-interactive --skip-existing "${wheeldir}"/*.whl
+echo ""