[Reusable workflows] Linux - Move codegened binary builds + upload to…

… a reusable workflow (take 2) (pytorch#81564)

Refactor the codgen used to generate linux binary builds into reusable workflows

This PR extracts out what's needed for binary builds & tests for linux, and the upload step that's shared by all OSes.

Other OS workflows are extracted in these PRs:
- Windows: pytorch#81572
- MacOS: pytorch#81571

The one exception here is I left the ROCm workflow steps as codegen since bugs with those workflows would bring down the entire ROCm runner for all other users as well.

Partially fixes pytorch#74480

### Testing
Ran CI with `ciflow/trunk`, `ciflow/binaries`,  `ciflow/binaries_conda`, and `ciflow/binaries_wheel` labels
Pull Request resolved: pytorch#81564
Approved by: https://github.com/huydhn

.github/actions/test-pytorch-binary/action.yml

@@ -0,0 +1,41 @@
+name: Test pytorch binary
+
+description: Pulls the docker image and tests the pytorch binary using it. All env variable referenced in the "Test PyTorch binary" step must be set in the GITHUB_ENV file
+
+runs:
+  using: composite
+  steps:
+    - name: Test PyTorch binary
+      shell: bash
+      run: |
+        set -x
+        # shellcheck disable=SC2086,SC2090
+        container_name=$(docker run \
+          ${GPU_FLAG:-} \
+          -e BINARY_ENV_FILE \
+          -e BUILDER_ROOT \
+          -e BUILD_ENVIRONMENT \
+          -e BUILD_SPLIT_CUDA \
+          -e DESIRED_CUDA \
+          -e DESIRED_DEVTOOLSET \
+          -e DESIRED_PYTHON \
+          -e GITHUB_ACTIONS \
+          -e GPU_ARCH_TYPE \
+          -e GPU_ARCH_VERSION \
+          -e LIBTORCH_VARIANT \
+          -e PACKAGE_TYPE \
+          -e PYTORCH_FINAL_PACKAGE_DIR \
+          -e PYTORCH_ROOT \
+          -e SKIP_ALL_TESTS \
+          --tty \
+          --detach \
+          -v "${GITHUB_WORKSPACE}/pytorch:/pytorch" \
+          -v "${GITHUB_WORKSPACE}/builder:/builder" \
+          -v "${RUNNER_TEMP}/artifacts:/final_pkgs" \
+          -w / \
+          "${DOCKER_IMAGE}"
+        )
+        docker exec -t -w "${PYTORCH_ROOT}" "${container_name}" bash -c "bash .circleci/scripts/binary_populate_env.sh"
+        # Generate test script
+        docker exec -t -w "${PYTORCH_ROOT}" -e OUTPUT_SCRIPT="/run.sh" "${container_name}" bash -c "bash .circleci/scripts/binary_linux_test.sh"
+        docker exec -t "${container_name}" bash -c "source ${BINARY_ENV_FILE} && bash -x /run.sh"

.github/templates/linux_binary_build_workflow.yml.j2

@@ -48,146 +48,56 @@ jobs:
 {%- for config in build_configs %}
   !{{ config["build_name"] }}-build:
     if: ${{ github.repository_owner == 'pytorch' }}
-    runs-on: linux.4xlarge
-    timeout-minutes: !{{ common.timeout_minutes }}
-    !{{ upload.binary_env(config) }}
-    steps:
-      !{{ common.setup_ec2_linux() }}
-      !{{ common.checkout(deep_clone=False, directory="pytorch") }}
-      !{{ common.checkout(deep_clone=False, directory="builder", repository="pytorch/builder", branch=common.builder_branch) }}
-{%- if config["gpu_arch_type"] == 'cuda' and config["gpu_arch_version"].startswith('11') %}
-      - name: Set BUILD_SPLIT_CUDA
-        run: |
-          echo "BUILD_SPLIT_CUDA='ON'" >> "$GITHUB_ENV"
-{%- endif %}
-      - name: Pull Docker image
-        run: |
-          !{{ common.add_retry_to_env() }}
-          retry docker pull "${DOCKER_IMAGE}"
-      - name: Build PyTorch binary
-        run: |
-          set -x
-          mkdir -p artifacts/
-          container_name=$(docker run \
-            -e BINARY_ENV_FILE \
-            -e BUILDER_ROOT \
-            -e BUILD_ENVIRONMENT \
-            -e BUILD_SPLIT_CUDA \
-            -e DESIRED_CUDA \
-            -e DESIRED_DEVTOOLSET \
-            -e DESIRED_PYTHON \
-            -e GITHUB_ACTIONS \
-            -e GPU_ARCH_TYPE \
-            -e GPU_ARCH_VERSION \
-            -e LIBTORCH_VARIANT \
-            -e PACKAGE_TYPE \
-            -e PYTORCH_FINAL_PACKAGE_DIR \
-            -e PYTORCH_ROOT \
-            -e SKIP_ALL_TESTS \
-            --tty \
-            --detach \
-            -v "${GITHUB_WORKSPACE}/pytorch:/pytorch" \
-            -v "${GITHUB_WORKSPACE}/builder:/builder" \
-            -v "${RUNNER_TEMP}/artifacts:/artifacts" \
-            -w / \
-            "${DOCKER_IMAGE}"
-          )
-          docker exec -t -w "${PYTORCH_ROOT}" "${container_name}" bash -c "bash .circleci/scripts/binary_populate_env.sh"
-          docker exec -t "${container_name}" bash -c "source ${BINARY_ENV_FILE} && bash /builder/!{{ config["package_type"] }}/build.sh"
-      !{{ common.chown_dir("${RUNNER_TEMP}/artifacts") }}
-      - uses: !{{ common.upload_artifact_s3_action }}
-        with:
-          name: !{{ config["build_name"] }}
-          retention-days: 14
-          if-no-files-found: error
-          path:
-            ${{ runner.temp }}/artifacts/*
-      !{{ common.teardown_ec2_linux("pytorch/") }}
+    uses: ./.github/workflows/_binary-build-linux.yml
+    with:!{{ upload.binary_env_as_input(config) }}
+      build_name: !{{ config["build_name"] }}
+      build_environment: !{{ build_environment }}
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+
   !{{ config["build_name"] }}-test:  # Testing
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: !{{ config["build_name"] }}-build
-{%- if config["gpu_arch_type"] == "rocm" %}
-    runs-on: linux.rocm.gpu
-{%- elif config["gpu_arch_type"] == "cuda" %}
-    runs-on: linux.4xlarge.nvidia.gpu
+{%- if config["gpu_arch_type"] != "rocm" %}
+    uses: ./.github/workflows/_binary-test-linux.yml
+    with:!{{ upload.binary_env_as_input(config) }}
+      build_name: !{{ config["build_name"] }}
+      build_environment: !{{ build_environment }}
+      {%- if config["gpu_arch_type"] == "rocm" %}
+      runs_on: linux.rocm.gpu
+      {%- elif config["gpu_arch_type"] == "cuda" %}
+      runs_on: linux.4xlarge.nvidia.gpu
+      {%- else %}
+      runs_on: linux.4xlarge
+      {%- endif %}
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
 {%- else %}
-    runs-on: linux.4xlarge
-{%- endif %}
+    runs-on: linux.rocm.gpu
     timeout-minutes: !{{ common.timeout_minutes }}
     !{{ upload.binary_env(config) }}
     steps:
-{%- if config["gpu_arch_type"] == "rocm" %}
       !{{ common.setup_rocm_linux() }}
-{%- else %}
-      !{{ common.setup_ec2_linux() }}
-{%- endif %}
       - uses: !{{ common.download_artifact_s3_action }}
         name: Download Build Artifacts
         with:
           name: !{{ config["build_name"] }}
           path: "${{ runner.temp }}/artifacts/"
       !{{ common.checkout(deep_clone=False, directory="pytorch") }}
       !{{ common.checkout(deep_clone=False, directory="builder", repository="pytorch/builder", branch=common.builder_branch) }}
-{%- if config["gpu_arch_type"] == "rocm" %}
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"
-{%- elif config["gpu_arch_type"] == "cuda" %}
-      - uses: nick-fields/retry@71062288b76e2b6214ebde0e673ce0de1755740a
-        name: Install nvidia driver, nvidia-docker runtime, set GPU_FLAG
-        with:
-          timeout_minutes: 10
-          max_attempts: 3
-          command: |
-            set -ex
-            pushd pytorch
-            bash .github/scripts/install_nvidia_utils_linux.sh
-            echo "GPU_FLAG=--gpus all" >> "${GITHUB_ENV}"
-            popd
-{%- endif %}
       - name: Pull Docker image
-        run: |
-          !{{ common.add_retry_to_env() }}
-          retry docker pull "${DOCKER_IMAGE}"
-      - name: Test PyTorch binary
-        run: |
-          set -x
-          # shellcheck disable=SC2086,SC2090
-          container_name=$(docker run \
-            ${GPU_FLAG:-} \
-            -e BINARY_ENV_FILE \
-            -e BUILDER_ROOT \
-            -e BUILD_ENVIRONMENT \
-            -e BUILD_SPLIT_CUDA \
-            -e DESIRED_CUDA \
-            -e DESIRED_DEVTOOLSET \
-            -e DESIRED_PYTHON \
-            -e GITHUB_ACTIONS \
-            -e GPU_ARCH_TYPE \
-            -e GPU_ARCH_VERSION \
-            -e LIBTORCH_VARIANT \
-            -e PACKAGE_TYPE \
-            -e PYTORCH_FINAL_PACKAGE_DIR \
-            -e PYTORCH_ROOT \
-            -e SKIP_ALL_TESTS \
-            --tty \
-            --detach \
-            -v "${GITHUB_WORKSPACE}/pytorch:/pytorch" \
-            -v "${GITHUB_WORKSPACE}/builder:/builder" \
-            -v "${RUNNER_TEMP}/artifacts:/final_pkgs" \
-            -w / \
-            "${DOCKER_IMAGE}"
-          )
-          docker exec -t -w "${PYTORCH_ROOT}" "${container_name}" bash -c "bash .circleci/scripts/binary_populate_env.sh"
-          # Generate test script
-          docker exec -t -w "${PYTORCH_ROOT}" -e OUTPUT_SCRIPT="/run.sh" "${container_name}" bash -c "bash .circleci/scripts/binary_linux_test.sh"
-          docker exec -t "${container_name}" bash -c "source ${BINARY_ENV_FILE} && bash -x /run.sh"
-{%- if config["gpu_arch_type"] == "rocm" %}
+        uses: ./pytorch/.github/actions/pull-docker-image
+        with:
+          docker-image: !{{ config["container_image"] }}
+      - name: Test Pytorch binary
+        uses: ./pytorch/.github/actions/test-pytorch-binary
       !{{ common.teardown_rocm_linux() }}
-{%- else %}
-      !{{ common.teardown_ec2_linux("pytorch/") }}
 {%- endif %}
-  {%- if branches == "nightly" %}
+
+{%- if branches == "nightly" %}
   !{{ upload.upload_binaries(config) }}
-  {%- endif %}
+{%- endif %}
 {%- endfor %}

.github/templates/upload.yml.j2

@@ -1,10 +1,16 @@
 {% import 'common.yml.j2' as common %}
 
 {%- macro binary_env(config, is_windows=False) -%}
-    env:
+    env:!{{ binary_env_as_input(config, is_windows, True) }}
+{%- endmacro %}
+
+{%- macro binary_env_as_input(config, is_windows=False, include_skip_tests=False) -%}
 {%- if is_windows %}
       PYTORCH_ROOT: ${{ github.workspace }}/pytorch
       BUILDER_ROOT: ${{ github.workspace }}/builder
+{%- else %}
+      PYTORCH_ROOT: /pytorch
+      BUILDER_ROOT: /builder
 {%- endif %}
       PACKAGE_TYPE: !{{ config["package_type"] }}
       # TODO: This is a legacy variable that we eventually want to get rid of in
@@ -14,23 +20,25 @@
       GPU_ARCH_VERSION: !{{ config["gpu_arch_version"] }}
 {%- endif %}
       GPU_ARCH_TYPE: !{{ config["gpu_arch_type"] }}
+{%- if include_skip_tests %}
+      SKIP_ALL_TESTS: 1
+{%- endif %}
 {%- if not is_windows %}
       DOCKER_IMAGE: !{{ config["container_image"] }}
 {%- endif %}
-      SKIP_ALL_TESTS: 1
 {%- if config["package_type"] == "libtorch" %}
-{%- if config["libtorch_config"] %}
+  {%- if config["libtorch_config"] %}
       LIBTORCH_CONFIG: !{{ config["libtorch_config"] }}
-{%- endif %}
+  {%- endif %}
       LIBTORCH_VARIANT: !{{ config["libtorch_variant"] }}
-{%- if config["devtoolset"] %}
+  {%- if config["devtoolset"] %}
       DESIRED_DEVTOOLSET: !{{ config["devtoolset"] }}
-{%- endif %}
-{%- if is_windows %}
+  {%- endif %}
+  {%- if is_windows %}
       # This is a dummy value for libtorch to work correctly with our batch scripts
       # without this value pip does not get installed for some reason
       DESIRED_PYTHON: "3.7"
-{%- endif %}
+  {%- endif %}
 {%- else %}
       DESIRED_PYTHON: "!{{ config["python_version"] }}"
 {%- endif %}
@@ -39,60 +47,21 @@
 
 {%- macro upload_binaries(config, is_windows=False, has_test=True, use_s3=True) -%}
 !{{ config["build_name"] }}-upload:  # Uploading
-    runs-on: linux.2xlarge  # self hosted runner to download ec2 artifacts
     if: ${{ github.repository_owner == 'pytorch' }}
 {%- if has_test %}
     needs: !{{ config["build_name"] }}-test
 {%- else %}
     needs: !{{ config["build_name"] }}-build
 {%- endif %}
-    !{{ binary_env(config, is_windows) }}
-    steps:
-      !{{ common.setup_ec2_linux() }}
-      - name: Clone pytorch/pytorch
-        uses: actions/checkout@v2
-{%- if use_s3 %}
-      - uses: !{{ common.download_artifact_s3_action }}
-{%- else %}
-      - uses: actions/download-artifact@v2
-{%- endif %}
-        name: Download Build Artifacts
-        with:
-          name: !{{ config["build_name"] }}
-          path: "${{ runner.temp }}/artifacts/"
-      - name: Set DRY_RUN (only for tagged pushes)
-        if: ${{ github.event_name == 'push' && (github.event.ref == 'refs/heads/nightly' || (startsWith(github.event.ref, 'refs/tags/') && !startsWith(github.event.ref, 'refs/tags/ciflow/'))) }}
-        run: |
-          echo "DRY_RUN=disabled" >> "$GITHUB_ENV"
-      - name: Set UPLOAD_CHANNEL (only for tagged pushes)
-        if: ${{ github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags/') && !startsWith(github.event.ref, 'refs/tags/ciflow/') }}
-        run: |
-          # reference ends with an RC suffix
-          if [[ ${GITHUB_REF_NAME} = *-rc[0-9]* ]]; then
-            echo "UPLOAD_CHANNEL=test" >> "$GITHUB_ENV"
-          fi
-      - name: Upload binaries
-        env:
-          PKG_DIR: "${{ runner.temp }}/artifacts"
-          UPLOAD_SUBFOLDER: "${{ env.DESIRED_CUDA }}"
-          # When running these on pull_request events these should be blank
-          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_PYTORCH_UPLOADER_ACCESS_KEY_ID }}
-          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_PYTORCH_UPLOADER_SECRET_ACCESS_KEY }}
-          ANACONDA_API_TOKEN: ${{ secrets.CONDA_PYTORCHBOT_TOKEN }}
-        run: |
-          docker run --rm -i \
-            -e ANACONDA_API_TOKEN \
-            -e AWS_ACCESS_KEY_ID \
-            -e AWS_SECRET_ACCESS_KEY \
-            -e DRY_RUN \
-            -e PACKAGE_TYPE \
-            -e PKG_DIR=/artifacts \
-            -e UPLOAD_CHANNEL \
-            -e UPLOAD_SUBFOLDER \
-            -v "${RUNNER_TEMP}/artifacts:/artifacts" \
-            -v "${GITHUB_WORKSPACE}:/v" \
-            -w /v \
-            308535385114.dkr.ecr.us-east-1.amazonaws.com/tool/miniconda3:4.10.3 \
-            bash -c '.circleci/scripts/binary_upload.sh'
-      !{{ common.teardown_ec2_linux() }}
-{%- endmacro -%}
+    with:!{{ binary_env_as_input(config, is_windows) }}
+      build_name: !{{ config["build_name"] }}
+      {%- if not use_s3 %}
+      use_s3: False
+      {%- endif %}
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+      aws-access-key-id: ${{ secrets.AWS_PYTORCH_UPLOADER_ACCESS_KEY_ID }}
+      aws-pytorch-uploader-secret-access-key: ${{ secrets.AWS_PYTORCH_UPLOADER_SECRET_ACCESS_KEY }}
+      conda-pytorchbot-token: ${{ secrets.CONDA_PYTORCHBOT_TOKEN }}
+    uses: ./.github/workflows/_binary-upload.yml
+{%- endmacro %}