Allow custom s3 upload path within pytorch (#7280)

To be used by forge and torchcomms

Author: atalman

.github/workflows/_binary_upload.yml

@@ -36,6 +36,11 @@ on:
         description: The comma-separated list of CUDA arch to be uploaded to pypi
         default: ''
         type: string
+      wheel-upload-path:
+        description: Custom wheel upload path
+        required: false
+        type: string
+        default: ''
     secrets:
       PYPI_API_TOKEN:
         description: An optional token to upload to pypi
@@ -99,6 +104,8 @@ jobs:
       - name: Upload package to pytorch.org
         shell: bash
         working-directory: ${{ inputs.repository }}
+        env:
+          WHEEL_UPLOAD_PATH: ${{ inputs.wheel-upload-path }}
         run: |
           set -ex
 
@@ -112,6 +119,11 @@ jobs:
             AWS_CMD="aws s3 cp"
           fi
 
+           # set path to specific wheel-upload-path
+          if [[ -n "$WHEEL_UPLOAD_PATH" ]]; then
+            PYTORCH_S3_BUCKET_PATH="s3://pytorch/${WHEEL_UPLOAD_PATH}"
+          fi
+
           for pkg in dist/*; do
             shm_id=$(sha256sum "${pkg}" | awk '{print $1}')
             ${AWS_CMD} "$pkg" "${PYTORCH_S3_BUCKET_PATH}" --acl public-read \

.github/workflows/build_wheels_linux.yml

@@ -88,6 +88,10 @@ on:
         required: false
         default: "python -m build --wheel"
         type: string
+      wheel-upload-path:
+        description: Custom wheel upload path
+        required: false
+        type: string
       pip-install-torch-extra-args:
         # NOTE: Why does this exist?
         # Well setuptools / python packaging doesn't actually allow you to specify dependencies
@@ -325,6 +329,7 @@ jobs:
       architecture: ${{ inputs.architecture }}
       trigger-event: ${{ inputs.trigger-event }}
       upload-to-pypi: ${{ inputs.upload-to-pypi }}
+      wheel-upload-path: ${{ inputs.wheel-upload-path }}
     secrets:
       PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
 

.github/workflows/test_build_wheels_linux_without_cuda.yml

@@ -45,6 +45,7 @@ jobs:
       package-name: ${{ matrix.package-name }}
       trigger-event: "${{ github.event_name }}"
       build-platform: python-build-package
+      wheel-upload-path: "test/custom/path/"
       pip-install-torch-extra-args:
         torchvision
         torchao