From 3c73b14a4a6a3c5ea2e6de0366bb58dba4a29a4c Mon Sep 17 00:00:00 2001
From: Andrey Talman <atalman@fb.com>
Date: Fri, 17 Jan 2025 19:28:58 +0000
Subject: [PATCH] Update wheel binary size validations (#6185)

Enable to run these workflows on workflow_dispatch
Pass channel to Validate PyPI Wheel Binary Size to be able to validate
nightly or test binaries
---
 ...ml => validate-pypi-wheel-binary-size.yml} | 20 +++++++++++++++----
 .../validate-repackaged-binary-sizes.yml      |  1 +
 2 files changed, 17 insertions(+), 4 deletions(-)
 rename .github/workflows/{validate-nightly-pypi-wheel-binary-size.yml => validate-pypi-wheel-binary-size.yml} (53%)

diff --git a/.github/workflows/validate-nightly-pypi-wheel-binary-size.yml b/.github/workflows/validate-pypi-wheel-binary-size.yml
similarity index 53%
rename from .github/workflows/validate-nightly-pypi-wheel-binary-size.yml
rename to .github/workflows/validate-pypi-wheel-binary-size.yml
index 24fffc16eb..c92961d030 100644
--- a/.github/workflows/validate-nightly-pypi-wheel-binary-size.yml
+++ b/.github/workflows/validate-pypi-wheel-binary-size.yml
@@ -1,16 +1,27 @@
-name: Validate Nightly PyPI Wheel Binary Size
+name: Validate PyPI Wheel Binary Size
 on:
   pull_request:
     paths:
-      - .github/workflows/validate-nightly-pypi-wheel-binary-size.yml
+      - .github/workflows/validate-pypi-wheel-binary-size.yml
   workflow_dispatch:
+    inputs:
+      channel:
+          description: "Channel to use (nightly, test)"
+          required: true
+          type: choice
+          default: test
+          options:
+            - nightly
+            - test
   schedule:
     # At 2:30 pm UTC (7:30 am PDT)
     - cron: "30 14 * * *"
 
 jobs:
-  nightly-pypi-binary-size-validation:
+  pypi-binary-size-validation:
     runs-on: ubuntu-latest
+    env:
+      CHANNEL: ${{ inputs.channel || 'nightly' }}
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -21,6 +32,7 @@ jobs:
           pip3 install -r tools/binary_size_validation/requirements.txt
       - name: Run validation
         run: |
+          # shellcheck disable=SC2086
           python tools/binary_size_validation/binary_size_validation.py \
-              --url https://download.pytorch.org/whl/nightly/cu121/torch/ \
+              --url https://download.pytorch.org/whl/${CHANNEL}/cu124/torch/ \
               --include "linux" --only-latest-version --threshold 750
diff --git a/.github/workflows/validate-repackaged-binary-sizes.yml b/.github/workflows/validate-repackaged-binary-sizes.yml
index e819465a26..6325aff29f 100644
--- a/.github/workflows/validate-repackaged-binary-sizes.yml
+++ b/.github/workflows/validate-repackaged-binary-sizes.yml
@@ -12,6 +12,7 @@ name: Validate manywheel binaries
 #  * optionally upload the repackaged binaries as artifacts (for debug or promotion)
 
 on:
+  workflow_dispatch:
   pull_request:
     paths:
       - .github/workflows/validate-repackaged-binary-sizes.yml