From 2d8c841ea1d1cdf35b5957f73b17d0fe05b00275 Mon Sep 17 00:00:00 2001
From: Ivan Zaitsev <ivanzaitsev@fb.com>
Date: Tue, 13 Dec 2022 17:29:22 -0800
Subject: [PATCH] Validate binary size linux_job

---
 .../validate-repackaged-binary-sizes.yml      | 114 ++++++++++--------
 1 file changed, 63 insertions(+), 51 deletions(-)

diff --git a/.github/workflows/validate-repackaged-binary-sizes.yml b/.github/workflows/validate-repackaged-binary-sizes.yml
index 722669df8..29d1bdc53 100644
--- a/.github/workflows/validate-repackaged-binary-sizes.yml
+++ b/.github/workflows/validate-repackaged-binary-sizes.yml
@@ -1,65 +1,77 @@
 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
+
 on:
   pull_request:
 
 jobs:
   validate-binary-size:
-    runs-on: linux.4xlarge.nvidia.gpu
     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: "3.10"
 #          - 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
-#          - 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
-#          - 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
-#          - 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
-
-    env:
-      GPU_ARCH_TYPE: cuda
-      GPU_ARCH_VERSION: "11.7"
-
-    steps:
-      - name: Checkout PyTorch builder
-        uses: actions/checkout@v3
-
-      - name: Install patchelf
-        run: |
-          chmod a+x common/install_patchelf.sh
-          sudo common/install_patchelf.sh
-
-      - name: Download torch whl
-        run: |
-          wget ${{ matrix.whl.url }}
-          FILENAME=$(ls -1 *.whl | head -n 1)
-          echo "::notice::Before repackaging: $(du -h $FILENAME | cut -f1)"
-          echo "FILENAME=$FILENAME" >> $GITHUB_ENV
-
-      - name: Repackage into manywheel
-        continue-on-error: true
-        run: |
-          release/pypi/prep_binary_for_pypi.sh $FILENAME
-          NEW_FILENAME=$(ls -1 *.whl | head -n 1)
-          echo "::notice::After repackaging: $(du -h $NEW_FILENAME | cut -f1)"
-          echo "NEW_FILENAME=$NEW_FILENAME" >> $GITHUB_ENV
-
-      - name: Run smoke test
-        continue-on-error: true
-        run: |
-          set -ex
-          # run smoke test to make sure the binary is not broken
-          conda create -n smoke-test python=${{ matrix.whl.python }} -y
-          conda activate smoke-test
-          pip install $NEW_FILENAME
-          python  ./test/smoke_test/smoke_test.py -- --package=torchonly
+#            python: "3.11"
+          - 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"
+          - 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"
+          - 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"
 
-      - name: Hold runner for 60 minutes or until ssh sessions have drained
-        timeout-minutes: 60
-        run: |
-          sleep infinity
\ No newline at end of file
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      runner: linux.4xlarge.nvidia.gpu
+      repository: "pytorch/builder"
+      ref: ${{ github.ref }}
+      job-name: "Validate binary size"
+      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)"
+        
+        # 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