From cc8d20c765f92ae75416d318170a07a8d05beea6 Mon Sep 17 00:00:00 2001
From: Ivan Zaitsev <108101595+izaitsevfb@users.noreply.github.com>
Date: Wed, 14 Dec 2022 12:32:30 -0800
Subject: [PATCH] PyPi binary validation and size check (#1230)

* Validate binary size

* Validate binary size linux_job

* evaluate the fix from #1231

* Add an optional artifact upload, consolidate fixes to `prep_binary_for_pypi.sh`
---
 .../validate-repackaged-binary-sizes.yml      | 88 +++++++++++++++++++
 release/pypi/prep_binary_for_pypi.sh          |  1 +
 2 files changed, 89 insertions(+)
 create mode 100644 .github/workflows/validate-repackaged-binary-sizes.yml

diff --git a/.github/workflows/validate-repackaged-binary-sizes.yml b/.github/workflows/validate-repackaged-binary-sizes.yml
new file mode 100644
index 000000000..695c68d3a
--- /dev/null
+++ b/.github/workflows/validate-repackaged-binary-sizes.yml
@@ -0,0 +1,88 @@
+name: Validate manywheel binaries
+
+# This workflow validates the size of the manywheel binaries after repackaging for PyPi
+# Specify the direct URLs to the binaries (from https://download.pytorch.org/whl/test/torch/) in the matrix
+# along with the python version.
+#
+# The workflow will:
+#  * download the binaries,
+#  * run release/pypi/prep_binary_for_pypi.sh
+#  * run smoke tests on the repackaged binaries
+#  * display the size before and after repackaging as the workflow annotation
+#  * optionally upload the repackaged binaries as artifacts (for debug or promotion)
+
+on:
+  pull_request:
+    paths:
+      - .github/workflows/validate-repackaged-binary-sizes.yml
+      - release/pypi/prep_binary_for_pypi.sh
+
+jobs:
+  validate-binary-size:
+    strategy:
+      fail-fast: false
+      matrix:
+        whl:
+          - url: https://download.pytorch.org/whl/test/cu117_pypi_cudnn/torch-1.13.1%2Bcu117.with.pypi.cudnn-cp310-cp310-linux_x86_64.whl
+            python: "3.10"  # python version to use for smoke tests
+            upload_artifact: false # upload the repackaged binary as an artifact
+          - url: https://download.pytorch.org/whl/test/cu117_pypi_cudnn/torch-1.13.1%2Bcu117.with.pypi.cudnn-cp37-cp37m-linux_x86_64.whl
+            python: "3.7"
+            artifact: false
+          - url: https://download.pytorch.org/whl/test/cu117_pypi_cudnn/torch-1.13.1%2Bcu117.with.pypi.cudnn-cp38-cp38-linux_x86_64.whl
+            python: "3.8"
+            artifact: false
+          - url: https://download.pytorch.org/whl/test/cu117_pypi_cudnn/torch-1.13.1%2Bcu117.with.pypi.cudnn-cp39-cp39-linux_x86_64.whl
+            python: "3.9"
+            artifact: false
+     #    - url: https://download.pytorch.org/whl/test/cu117_pypi_cudnn/torch-1.13.1%2Bcu117.with.pypi.cudnn-cp311-cp311-linux_x86_64.whl
+     #      python: "3.11"
+     #      artifact: false
+
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      runner: linux.4xlarge.nvidia.gpu
+      job-name: "Validate binary size"
+      upload-artifact: ${{ matrix.whl.upload_artifact == 'true' && 'repackaged-binary' || '' }}
+      script: |
+        set -ex
+        export ENV_NAME="conda-env-${{ github.run_id }}"
+        export GPU_ARCH_VER="11.7"
+        export GPU_ARCH_TYPE="cuda"
+        export CUDA_VER="11.7"
+        export DESIRED_PYTHON="${{ matrix.whl.python }}"
+        export DESIRED_CUDA="cu117"
+        export PACKAGE_TYPE="wheel"
+        export TARGET_OS="linux"
+        export INSTALLATION=""
+        
+        # install zip
+        sudo yum install zip -y
+        
+        # install patchelf
+        chmod a+x common/install_patchelf.sh
+        sudo common/install_patchelf.sh
+        
+        # download torch whl
+        wget ${{ matrix.whl.url }}
+        FILENAME=$(ls -1 *.whl | head -n 1)
+        SIZE_BEFORE=$(du -h $FILENAME | cut -f1)
+        
+        # repackage into manywheel
+        release/pypi/prep_binary_for_pypi.sh $FILENAME
+        
+        NEW_FILENAME=$(ls -1 *.whl | head -n 1)
+        echo "::notice:: $FILENAME before: $SIZE_BEFORE after: $(du -h $NEW_FILENAME | cut -f1)"
+        
+        # cp to ${RUNNER_ARTIFACT_DIR}
+        cp $NEW_FILENAME ${RUNNER_ARTIFACT_DIR}/
+        
+        # create conda env
+        conda create -y -n $ENV_NAME python=$DESIRED_PYTHON
+        conda activate $ENV_NAME
+        
+        # install torch
+        pip install numpy pillow $NEW_FILENAME
+        
+        # run smoke test
+        python ./test/smoke_test/smoke_test.py --package=torchonly
\ No newline at end of file
diff --git a/release/pypi/prep_binary_for_pypi.sh b/release/pypi/prep_binary_for_pypi.sh
index 639d834a6..f1d9edb6d 100755
--- a/release/pypi/prep_binary_for_pypi.sh
+++ b/release/pypi/prep_binary_for_pypi.sh
@@ -95,6 +95,7 @@ for whl_file in "$@"; do
             fi
         )
 
+        rm -rf "${new_whl_file}"
         zip -qr9 "${new_whl_file}" .
     )
 done