Add custom policy for publishing nightly wheels

.github/workflows/_binary_upload.yml

@@ -41,6 +41,10 @@ on:
         required: false
         type: string
         default: ''
+      wheel-nightly-policy:
+        description: Custom wheel upload policy for nightly
+        type: string
+        default: ''
     secrets:
       PYPI_API_TOKEN:
         description: An optional token to upload to pypi
@@ -84,7 +88,7 @@ jobs:
         if: ${{ inputs.trigger-event == 'schedule' || (inputs.trigger-event == 'push' && startsWith(github.event.ref, 'refs/heads/nightly')) }}
         uses: aws-actions/configure-aws-credentials@50ac8dd1e1b10d09dac7b8727528b91bed831ac0 # v3.0.2
         with:
-          role-to-assume: arn:aws:iam::749337293305:role/gha_workflow_nightly_build_wheels
+          role-to-assume: arn:aws:iam::749337293305:role/${{ (inputs.wheel-nightly-policy != '' && inputs.wheel-nightly-policy) || 'gha_workflow_nightly_build_wheels'}}
           aws-region: us-east-1
 
       - name: Configure aws credentials (pytorch account)

.github/workflows/build_wheels_linux.yml

@@ -96,6 +96,15 @@ on:
         description: Custom wheel upload path
         required: false
         type: string
+      wheel-nightly-policy:
+        description: Custom wheel upload policy for nightly
+        type: string
+        default: ''
+      run-smoke-test:
+        description: Set to true if smoke test is needed
+        required: false
+        type: boolean
+        default: true
       pip-install-torch-extra-args:
         # NOTE: Why does this exist?
         # Well setuptools / python packaging doesn't actually allow you to specify dependencies
@@ -278,6 +287,7 @@ jobs:
           repository: ${{ inputs.repository  }}
           script: ${{ inputs.post-script }}
       - name: Smoke Test
+        if: ${{ inputs.run-smoke-test }}
         shell: bash -l {0}
         env:
           PACKAGE_NAME: ${{ inputs.package-name }}
@@ -334,6 +344,7 @@ jobs:
       trigger-event: ${{ inputs.trigger-event }}
       upload-to-pypi: ${{ inputs.upload-to-pypi }}
       wheel-upload-path: ${{ inputs.wheel-upload-path }}
+      wheel-nightly-policy: ${{ inputs.wheel-nightly-policy }}
     secrets:
       PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}