Refactor conda upload job to a separate job running on GH ephemeral runner (#4886)

huydhn · huydhn · commit 6128c1ab7043 · 2024-02-12T12:38:17.000-08:00
Similar to #4877, this moves conda upload into a separate job on GH ephemeral runner: * I need a new `_binary_conda_upload` reusable workflow because conda upload uses anaconda client to upload to conda, not awscli to upload to S3. * The build job doesn't have access to `pytorchbot-env` anymore, thus it has no access to `CONDA_PYTORCHBOT_TOKEN` and `CONDA_PYTORCHBOT_TOKEN_TEST` secrets. Only the upload job has this access.
diff --git a/.github/workflows/_binary_conda_upload.yml b/.github/workflows/_binary_conda_upload.yml
@@ -0,0 +1,111 @@
+name: upload conda
+
+on:
+  workflow_call:
+    inputs:
+      repository:
+        description: 'Repository to checkout, defaults to ""'
+        default: ''
+        type: string
+      ref:
+        description: 'Reference to checkout, defaults to "nightly"'
+        default: 'nightly'
+        type: string
+      test-infra-repository:
+        description: 'Test infra repository to use'
+        default: "pytorch/test-infra"
+        type: string
+      test-infra-ref:
+        description: 'Test infra reference to use'
+        default: ""
+        type: string
+      build-matrix:
+        description: 'Build matrix to utilize'
+        default: ''
+        type: string
+      trigger-event:
+        description: 'Trigger Event in caller that determines whether or not to upload'
+        type: string
+        default: ''
+    secrets:
+      CONDA_PYTORCHBOT_TOKEN:
+        description: 'Access Token needed to upload binaries to anaconda nightly channel'
+        required: false
+      CONDA_PYTORCHBOT_TOKEN_TEST:
+        description: 'Access Token needed to upload binaries to anaconda test channel'
+        required: false
+
+jobs:
+  upload:
+    runs-on: ubuntu-22.04
+    environment: ${{(inputs.trigger-event == 'push' && (startsWith(github.event.ref, 'refs/heads/nightly') || startsWith(github.event.ref, 'refs/tags/v'))) && 'pytorchbot-env' || ''}}
+    strategy:
+      fail-fast: false
+      matrix: ${{ fromJSON(inputs.build-matrix) }}
+    timeout-minutes: 30
+    name: ${{ matrix.build_name }}
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          repository: ${{ inputs.test-infra-repository }}
+          ref: ${{ inputs.test-infra-ref }}
+          path: test-infra
+
+      # For pytorch_pkg_helpers which we need to run to generate the artifact name and target S3 buckets
+      - uses: ./test-infra/.github/actions/setup-binary-upload
+        with:
+          repository: ${{ inputs.repository }}
+          ref: ${{ inputs.ref }}
+          python-version: ${{ matrix.python_version }}
+          cuda-version: ${{ matrix.desired_cuda }}
+          upload-to-base-bucket: ${{ matrix.upload_to_base_bucket }}
+
+      - uses: ./test-infra/.github/actions/set-channel
+
+      - name: Download the artifact
+        uses: actions/download-artifact@v3
+        with:
+          name: ${{ env.ARTIFACT_NAME }}
+          path: ${{ inputs.repository }}/distr
+
+      - name: Nightly or release RC
+        if: ${{ (inputs.trigger-event == 'push' && startsWith(github.event.ref, 'refs/heads/nightly')) || (env.CHANNEL == 'test' && startsWith(github.event.ref, 'refs/tags/')) }}
+        shell: bash
+        env:
+          CONDA_PYTORCHBOT_TOKEN: ${{ secrets.CONDA_PYTORCHBOT_TOKEN }}
+          CONDA_PYTORCHBOT_TOKEN_TEST: ${{ secrets.CONDA_PYTORCHBOT_TOKEN_TEST }}
+        run: |
+          set -ex
+          echo "NIGHTLY_OR_TEST=1" >> "${GITHUB_ENV}"
+
+          if [[ "${CHANNEL}" = "nightly" ]]; then
+            echo "CONDA_TOKEN=${CONDA_PYTORCHBOT_TOKEN}" >> "${GITHUB_ENV}"
+          else
+            echo "CONDA_TOKEN=${CONDA_PYTORCHBOT_TOKEN_TEST}" >> "${GITHUB_ENV}"
+          fi
+
+      - name: Upload package to conda
+        working-directory: ${{ inputs.repository }}
+        run: |
+          set -ex
+
+          # shellcheck disable=SC1090
+          source "${BUILD_ENV_FILE}"
+          conda install --yes --quiet anaconda-client
+
+          if [[ "${NIGHTLY_OR_TEST:-0}" == "1" ]]; then
+            for pkg in distr/**/*.tar.bz2; do
+              anaconda \
+                -t "${CONDA_TOKEN}" \
+                upload "${pkg}" \
+                -u "pytorch-${CHANNEL}" \
+                --label main \
+                --no-progress \
+                --force
+            done
+          else
+            echo "Testing the upload of the following files to pytorch-${CHANNEL} conda channel:"
+            for pkg in distr/**/*.tar.bz2; do
+              ls -lah "${pkg}"
+            done
+          fi
diff --git a/.github/workflows/build_conda_linux.yml b/.github/workflows/build_conda_linux.yml
@@ -81,7 +81,6 @@ jobs:
       CU_VERSION: ${{ matrix.desired_cuda }}
     name: ${{ matrix.build_name }}
     runs-on: ${{ matrix.validation_runner }}
-    environment: ${{(inputs.trigger-event == 'push' || startsWith(github.event.ref, 'refs/tags/')) && 'pytorchbot-env' || ''}}
     container:
       image: ${{ matrix.container_image }}
       options: ${{ matrix.gpu_arch_type == 'cuda' && '--gpus all' || ' ' }}
@@ -102,6 +101,7 @@ jobs:
           ref: ${{ inputs.ref }}
           setup-miniconda: true
           python-version: ${{ env.PYTHON_VERSION }}
+          cuda-version: ${{ env.CU_VERSION }}
       - name: Combine Env Var and Build Env Files
         if: ${{ inputs.env-var-script != '' }}
         working-directory: ${{ inputs.repository }}
@@ -144,7 +144,6 @@ jobs:
             --python "${PYTHON_VERSION}" \
             --output-folder distr/ \
             "${CONDA_PACKAGE_DIRECTORY}"
-
       - name: Upload artifact to GitHub
         continue-on-error: true
         uses: actions/upload-artifact@v3
@@ -207,30 +206,20 @@ jobs:
             ${CONDA_RUN_SMOKE} python "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT}"
           fi
           conda env remove -p "${CONDA_ENV_SMOKE}"
-      - name: Upload package to conda
-        if: ${{ (inputs.trigger-event == 'push' && env.CHANNEL != 'test') || (env.CHANNEL == 'test' && startsWith(github.event.ref, 'refs/tags/')) }}
-        working-directory: ${{ inputs.repository }}
-        env:
-          CONDA_PYTORCHBOT_TOKEN: ${{ secrets.CONDA_PYTORCHBOT_TOKEN }}
-          CONDA_PYTORCHBOT_TOKEN_TEST: ${{ secrets.CONDA_PYTORCHBOT_TOKEN_TEST }}
-        run: |
-          if [[ "${CHANNEL}" = "nightly" ]]; then
-            export CONDA_TOKEN="${CONDA_PYTORCHBOT_TOKEN}"
-          else
-            export CONDA_TOKEN="${CONDA_PYTORCHBOT_TOKEN_TEST}"
-          fi
 
-          set -euxo pipefail
-          # shellcheck disable=SC1090
-          source "${BUILD_ENV_FILE}"
-          ${CONDA_RUN} conda install --yes --quiet anaconda-client
-          ${CONDA_RUN} anaconda \
-            -t "${CONDA_TOKEN}" \
-            upload distr/linux-64/*.tar.bz2 \
-            -u "pytorch-${CHANNEL}" \
-            --label main \
-            --no-progress \
-            --force
+  upload:
+    needs: build
+    uses: ./.github/workflows/_binary_conda_upload.yml
+    with:
+      repository: ${{ inputs.repository }}
+      ref: ${{ inputs.ref }}
+      test-infra-repository: ${{ inputs.test-infra-repository }}
+      test-infra-ref: ${{ inputs.test-infra-ref }}
+      build-matrix: ${{ inputs.build-matrix }}
+      trigger-event: ${{ inputs.trigger-event }}
+    secrets:
+      CONDA_PYTORCHBOT_TOKEN: ${{ secrets.CONDA_PYTORCHBOT_TOKEN }}
+      CONDA_PYTORCHBOT_TOKEN_TEST: ${{ secrets.CONDA_PYTORCHBOT_TOKEN_TEST }}
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}
diff --git a/.github/workflows/build_conda_macos.yml b/.github/workflows/build_conda_macos.yml
@@ -85,7 +85,6 @@ jobs:
       CU_VERSION: cpu
     name: ${{ matrix.build_name }}
     runs-on: ${{ inputs.runner-type }}
-    environment: ${{(inputs.trigger-event == 'push' || startsWith(github.event.ref, 'refs/tags/')) && 'pytorchbot-env' || ''}}
     # If a build is taking longer than 60 minutes on these runners we need
     # to have a conversation
     timeout-minutes: 60
@@ -112,6 +111,7 @@ jobs:
           ref: ${{ inputs.ref }}
           setup-miniconda: false
           python-version: ${{ env.PYTHON_VERSION }}
+          cuda-version: ${{ env.CU_VERSION }}
       - name: Combine Env Var and Build Env Files
         if: ${{ inputs.env-var-script != '' }}
         working-directory: ${{ inputs.repository }}
@@ -148,7 +148,6 @@ jobs:
             --python "${PYTHON_VERSION}" \
             --output-folder distr/ \
             "${CONDA_PACKAGE_DIRECTORY}"
-
       - name: Upload artifact to GitHub
         continue-on-error: true
         uses: actions/upload-artifact@v3
@@ -211,42 +210,25 @@ jobs:
 
           export PATH=${OLD_PATH}
           conda env remove -p "${CONDA_ENV_SMOKE}"
-      - name: Upload package to conda
-        if: ${{ (inputs.trigger-event == 'push' && env.CHANNEL != 'test') || (env.CHANNEL == 'test' && startsWith(github.event.ref, 'refs/tags/')) }}
-        working-directory: ${{ inputs.repository }}
-        env:
-          CONDA_PYTORCHBOT_TOKEN: ${{ secrets.CONDA_PYTORCHBOT_TOKEN }}
-          CONDA_PYTORCHBOT_TOKEN_TEST: ${{ secrets.CONDA_PYTORCHBOT_TOKEN_TEST }}
-        run: |
-          if [[ "${CHANNEL}" = "nightly" ]]; then
-            export CONDA_TOKEN="${CONDA_PYTORCHBOT_TOKEN}"
-          else
-            export CONDA_TOKEN="${CONDA_PYTORCHBOT_TOKEN_TEST}"
-          fi
-
-          set -euxo pipefail
-          # shellcheck disable=SC1090
-          source "${BUILD_ENV_FILE}"
-
-          ${CONDA_RUN} conda install --yes --quiet anaconda-client
-          arch_name="$(uname -m)"
-          if [ "${arch_name}" = "arm64" ]; then
-            export ARCH_NAME="osx-arm64"
-          else
-            export ARCH_NAME="osx-64"
-          fi
-          ${CONDA_RUN} anaconda \
-            -t "${CONDA_TOKEN}" \
-            upload "distr/${ARCH_NAME}"/*.tar.bz2 \
-            -u "pytorch-${CHANNEL}" \
-            --label main \
-            --no-progress \
-            --force
       - name: Clean up disk space
         if: always()
         continue-on-error: true
         uses: ./test-infra/.github/actions/check-disk-space
 
+  upload:
+    needs: build
+    uses: ./.github/workflows/_binary_conda_upload.yml
+    with:
+      repository: ${{ inputs.repository }}
+      ref: ${{ inputs.ref }}
+      test-infra-repository: ${{ inputs.test-infra-repository }}
+      test-infra-ref: ${{ inputs.test-infra-ref }}
+      build-matrix: ${{ inputs.build-matrix }}
+      trigger-event: ${{ inputs.trigger-event }}
+    secrets:
+      CONDA_PYTORCHBOT_TOKEN: ${{ secrets.CONDA_PYTORCHBOT_TOKEN }}
+      CONDA_PYTORCHBOT_TOKEN_TEST: ${{ secrets.CONDA_PYTORCHBOT_TOKEN_TEST }}
+
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}
   cancel-in-progress: true
diff --git a/.github/workflows/build_conda_windows.yml b/.github/workflows/build_conda_windows.yml
@@ -81,7 +81,6 @@ jobs:
       CU_VERSION: ${{ matrix.desired_cuda }}
     name: ${{ matrix.build_name }}
     runs-on: ${{ matrix.validation_runner }}
-    environment: ${{(inputs.trigger-event == 'push' || startsWith(github.event.ref, 'refs/tags/')) && 'pytorchbot-env' || ''}}
     defaults:
       run:
         shell: bash -l {0}
@@ -111,6 +110,7 @@ jobs:
           ref: ${{ inputs.ref }}
           setup-miniconda: false
           python-version: ${{ env.PYTHON_VERSION }}
+          cuda-version: ${{ env.CU_VERSION }}
       - name: Run Pre-Script with Caching
         if: ${{ inputs.pre-script != '' }}
         uses: ./test-infra/.github/actions/run-script-with-cache
@@ -228,30 +228,20 @@ jobs:
             ${CONDA_RUN_SMOKE} python "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT}"
           fi
           conda env remove -p "${CONDA_ENV_SMOKE}"
-      - name: Upload package to conda
-        if: ${{ (inputs.trigger-event == 'push' && env.CHANNEL != 'test') || (env.CHANNEL == 'test' && startsWith(github.event.ref, 'refs/tags/')) }}
-        working-directory: ${{ inputs.repository }}
-        env:
-          PACKAGE_NAME: ${{ inputs.package-name }}
-          CONDA_PYTORCHBOT_TOKEN: ${{ secrets.CONDA_PYTORCHBOT_TOKEN }}
-          CONDA_PYTORCHBOT_TOKEN_TEST: ${{ secrets.CONDA_PYTORCHBOT_TOKEN_TEST }}
-        run: |
-          if [[ "${CHANNEL}" = "nightly" ]]; then
-            export CONDA_TOKEN="${CONDA_PYTORCHBOT_TOKEN}"
-          else
-            export CONDA_TOKEN="${CONDA_PYTORCHBOT_TOKEN_TEST}"
-          fi
 
-          set -euxo pipefail
-          source "${BUILD_ENV_FILE}"
-          ${CONDA_RUN} conda install --yes --quiet anaconda-client
-          ${CONDA_RUN} anaconda \
-            -t "${CONDA_TOKEN}" \
-            upload "distr/win-64/${PACKAGE_NAME}*.tar.bz2" \
-            -u "pytorch-${CHANNEL}" \
-            --label main \
-            --no-progress \
-            --force
+  upload:
+    needs: build
+    uses: ./.github/workflows/_binary_conda_upload.yml
+    with:
+      repository: ${{ inputs.repository }}
+      ref: ${{ inputs.ref }}
+      test-infra-repository: ${{ inputs.test-infra-repository }}
+      test-infra-ref: ${{ inputs.test-infra-ref }}
+      build-matrix: ${{ inputs.build-matrix }}
+      trigger-event: ${{ inputs.trigger-event }}
+    secrets:
+      CONDA_PYTORCHBOT_TOKEN: ${{ secrets.CONDA_PYTORCHBOT_TOKEN }}
+      CONDA_PYTORCHBOT_TOKEN_TEST: ${{ secrets.CONDA_PYTORCHBOT_TOKEN_TEST }}
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}
