name: Tests on Linux

on:
  pull_request:
  push:
    branches:
      - nightly
      - main
      - release/*
  workflow_dispatch:

jobs:
  unittests:
    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: |
        set -euo pipefail

        export PYTHON_VERSION=${{ matrix.python-version }}
        export GPU_ARCH_TYPE=${{ matrix.gpu-arch-type }}
        export GPU_ARCH_VERSION=${{ matrix.gpu-arch-version }}

        ./.github/scripts/unittest.sh

  onnx:
    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
    with:
      repository: pytorch/vision
      script: |
        set -euo pipefail

        export PYTHON_VERSION=3.8
        export GPU_ARCH_TYPE=cpu

        ./.github/scripts/setup-env.sh
        
        # Prepare conda
        CONDA_PATH=$(which conda)
        eval "$(${CONDA_PATH} shell.bash hook)"
        conda activate ci
        
        echo '::group::Install ONNX'
        pip install --progress-bar=off onnx onnxruntime
        echo '::endgroup::'
        
        echo '::group::Install testing utilities'
        pip install --progress-bar=off pytest
        echo '::endgroup::'
        
        echo '::group::Run ONNX tests'
        pytest --durations=25 -v test/test_onnx.py
        echo '::endgroup::'

  unittests-extended:
    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
    with:
      repository: pytorch/vision
      script: |
        set -euo pipefail

        export PYTHON_VERSION=3.8
        export GPU_ARCH_TYPE=cpu

        ./.github/scripts/setup-env.sh
        
        # Prepare conda
        CONDA_PATH=$(which conda)
        eval "$(${CONDA_PATH} shell.bash hook)"
        conda activate ci
        
        echo '::group::Pre-download model weights'
        pip install --progress-bar=off aiohttp aiofiles tqdm
        python scripts/download_model_urls.py
        echo '::endgroup::'
        
        echo '::group::Install testing utilities'
        pip install --progress-bar=off pytest
        echo '::endgroup::'
        
        echo '::group::Run extended unittests'
        export PYTORCH_TEST_WITH_EXTENDED=1
        pytest --durations=25 -v test/test_extended_*.py
        echo '::endgroup::'