consolidate Linux workflows on CPU and GPU (#7189)

.github/unittest.sh

+#!/usr/bin/env bash
+
+set -euo pipefail
+
+# Prepare conda
+CONDA_PATH=$(which conda)
+eval "$(${CONDA_PATH} shell.bash hook)"
+conda config --set channel_priority strict
+
+echo '::group::Set PyTorch conda channel and wheel index'
+# TODO: Can we maybe have this as environment variable in the job template? For example, `IS_RELEASE`.
+if [[ (${GITHUB_EVENT_NAME} = 'pull_request' && (${GITHUB_BASE_REF} = 'release'*)) || (${GITHUB_REF} = 'refs/heads/release'*) ]]; then
+  CHANNEL_ID=test
+else
+  CHANNEL_ID=nightly
+fi
+PYTORCH_CONDA_CHANNEL=pytorch-"${CHANNEL_ID}"
+echo "PYTORCH_CONDA_CHANNEL=${PYTORCH_CONDA_CHANNEL}"
+
+case $GPU_ARCH_TYPE in
+  cpu)
+    GPU_ARCH_ID="cpu"
+    ;;
+  cuda)
+    VERSION_WITHOUT_DOT=$(echo "${GPU_ARCH_VERSION}" | sed 's/\.//')
+    GPU_ARCH_ID="cu${VERSION_WITHOUT_DOT}"
+    ;;
+  *)
+    echo "Unknown GPU_ARCH_TYPE=${GPU_ARCH_TYPE}"
+    exit 1
+    ;;
+esac
+PYTORCH_WHEEL_INDEX="https://download.pytorch.org/whl/${CHANNEL_ID}/${GPU_ARCH_ID}"
+echo "PYTORCH_WHEEL_INDEX=${PYTORCH_WHEEL_INDEX}"
+echo '::endgroup::'
+
+echo '::group::Create build environment'
+# See https://github.com/pytorch/vision/issues/7296 for ffmpeg
+conda create \
+  --name ci \
+  --quiet --yes \
+  python="${PYTHON_VERSION}" pip \
+  ninja libpng jpeg \
+  'ffmpeg<4.3' \
+  -c "${PYTORCH_CONDA_CHANNEL}" \
+  -c defaults
+conda activate ci
+pip install --progress-bar=off --upgrade setuptools
+
+# See https://github.com/pytorch/vision/issues/6790
+if [[ "${PYTHON_VERSION}" != "3.11" ]]; then
+  pip install --progress-bar=off av!=10.0.0
+fi
+
+echo '::endgroup::'
+
+echo '::group::Install PyTorch'
+pip install --progress-bar=off --pre torch --index-url="${PYTORCH_WHEEL_INDEX}"
+
+if [[ $GPU_ARCH_TYPE == 'cuda' ]]; then
+  python3 -c "import torch; exit(not torch.cuda.is_available())"
+fi
+echo '::endgroup::'
+
+echo '::group::Install TorchVision'
+python setup.py develop
+echo '::endgroup::'
+
+echo '::group::Collect PyTorch environment information'
+python -m torch.utils.collect_env
+echo '::endgroup::'
+
+echo '::group::Install testing utilities'
+pip install --progress-bar=off pytest pytest-mock pytest-cov
+echo '::endgroup::'
+
+echo '::group::Run tests'
+pytest --durations=25
+echo '::endgroup::'

.github/workflows/test-linux-cpu.yml

-name: Unit-tests on Linux CPU
-
-on:
-  pull_request:
-  push:
-    branches:
-      - nightly
-      - main
-      - release/*
-  workflow_dispatch:
-
-env:
-  CHANNEL: "nightly"
-
-jobs:
-  tests:
-    strategy:
-      matrix:
-        python_version: ["3.8", "3.9", "3.10"]
-      fail-fast: false
-    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
-    with:
-      runner: linux.12xlarge
-      repository: pytorch/vision
-      timeout: 120
-      script: |
-        # Mark Build Directory Safe
-        git config --global --add safe.directory /__w/vision/vision
-
-        # Set up Environment Variables
-        export PYTHON_VERSION="${{ matrix.python_version }}"
-        export VERSION="cpu"
-        export CUDATOOLKIT="cpuonly"
-
-        # Set CHANNEL
-        if [[ (${GITHUB_EVENT_NAME} = 'pull_request' && (${GITHUB_BASE_REF} = 'release'*)) || (${GITHUB_REF} = 'refs/heads/release'*) ]]; then
-          export CHANNEL=test
-        else
-          export CHANNEL=nightly
-        fi
-
-        # Create Conda Env
-        conda create -yp ci_env --quiet python="${PYTHON_VERSION}" numpy libpng jpeg scipy 'ffmpeg<4.3'
-        conda activate /work/ci_env
-
-        # Install PyTorch, Torchvision, and testing libraries
-        set -ex
-        conda install \
-          --yes \
-          -c "pytorch-${CHANNEL}" \
-          -c nvidia "pytorch-${CHANNEL}"::pytorch[build="*${VERSION}*"] \
-          "${CUDATOOLKIT}"
-        python3 setup.py develop
-        python3 -m pip install --progress-bar=off pytest pytest-mock 'av<10'
-
-        # Run Tests
-        python3 -m torch.utils.collect_env
-        python3 -m pytest --junitxml=test-results/junit.xml --durations 20

.github/workflows/test-linux-gpu.yml

-name: Unit-tests on Linux GPU
-
-on:
-  pull_request:
-  push:
-    branches:
-      - nightly
-      - main
-      - release/*
-  workflow_dispatch:
-
-env:
-  CHANNEL: "nightly"
-
-jobs:
-  tests:
-    strategy:
-      matrix:
-        python_version: ["3.8"]
-        cuda_arch_version: ["11.7"]
-      fail-fast: false
-    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
-    with:
-      runner: linux.g5.4xlarge.nvidia.gpu
-      repository: pytorch/vision
-      gpu-arch-type: cuda
-      gpu-arch-version: ${{ matrix.cuda_arch_version }}
-      timeout: 120
-      script: |
-        # Mark Build Directory Safe
-        git config --global --add safe.directory /__w/vision/vision
-
-        # Set up Environment Variables
-        export PYTHON_VERSION="${{ matrix.python_version }}"
-        export VERSION="${{ matrix.cuda_arch_version }}"
-        export CUDATOOLKIT="pytorch-cuda=${VERSION}"
-
-        # Set CHANNEL
-        if [[ (${GITHUB_EVENT_NAME} = 'pull_request' && (${GITHUB_BASE_REF} = 'release'*)) || (${GITHUB_REF} = 'refs/heads/release'*) ]]; then
-          export CHANNEL=test
-        else
-          export CHANNEL=nightly
-        fi
-
-        # Create Conda Env
-        conda create -yp ci_env --quiet python="${PYTHON_VERSION}" numpy libpng jpeg scipy 'ffmpeg<4.3'
-        conda activate /work/ci_env
-
-        # Install PyTorch, Torchvision, and testing libraries
-        set -ex
-        conda install \
-          --yes \
-          -c "pytorch-${CHANNEL}" \
-          -c nvidia "pytorch-${CHANNEL}"::pytorch[build="*${VERSION}*"] \
-          "${CUDATOOLKIT}"
-        python3 setup.py develop
-        python3 -m pip install --progress-bar=off pytest pytest-mock 'av<10'
-
-        # Run Tests
-        python3 -m torch.utils.collect_env
-        python3 -m pytest --junitxml=test-results/junit.xml --durations 20

.github/workflows/test-linux.yml

+name: Unit-tests on Linux
+
+on:
+  pull_request:
+  push:
+    branches:
+      - nightly
+      - main
+      - release/*
+  workflow_dispatch:
+
+jobs:
+  tests:
+    strategy:
+      matrix:
+        python-version:
+          - "3.8"
+          - "3.9"
+          - "3.10"
+          - "3.11"
+        runner: ["linux.12xlarge"]
+        gpu-arch-type: ["cpu"]
+        include:
+          - python-version: 3.8
+            runner: linux.g5.4xlarge.nvidia.gpu
+            gpu-arch-type: cuda
+            gpu-arch-version: "11.7"
+      fail-fast: false
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      repository: pytorch/vision
+      runner: ${{ matrix.runner }}
+      gpu-arch-type: ${{ matrix.gpu-arch-type }}
+      gpu-arch-version: ${{ matrix.gpu-arch-version }}
+      timeout: 120
+      script: |
+        export PYTHON_VERSION=${{ matrix.python-version }}
+        export GPU_ARCH_TYPE=${{ matrix.gpu-arch-type }}
+        export GPU_ARCH_VERSION=${{ matrix.gpu-arch-version }}
+
+        ./.github/unittest.sh