test-linux-gpu.yml

name: Unit-tests on Linux GPU

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

env:
  CHANNEL: "nightly"

concurrency:
  # Documentation suggests ${{ github.head_ref }}, but that's only available on pull_request/pull_request_target triggers, so using ${{ github.ref }}.
  # On master, we want all builds to complete even if merging happens faster to make it easier to discover at which point something broke.
  group: ${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && format('ci-master-{0}', github.sha) || format('ci-{0}', github.ref) }}
  cancel-in-progress: true

jobs:
  tests:
    strategy:
      matrix:
        python_version: ["3.8"]
        cuda_arch_version: ["12.1"]
      fail-fast: false
    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
    with:
      runner: linux.g5.4xlarge.nvidia.gpu
      repository: pytorch/rl
      docker-image: "nvidia/cuda:12.1.0-devel-ubuntu22.04"
      gpu-arch-type: cuda
      gpu-arch-version: ${{ matrix.cuda_arch_version }}
      timeout: 90
      script: |
        # Set env vars from matrix
        export PYTHON_VERSION=${{ matrix.python_version }}
        # Commenting these out for now because the GPU test are not working inside docker
        export CUDA_ARCH_VERSION=${{ matrix.cuda_arch_version }}
        export CU_VERSION="cu${CUDA_ARCH_VERSION:0:2}${CUDA_ARCH_VERSION:3:1}"
        export TORCH_VERSION=nightly
        # Remove the following line when the GPU tests are working inside docker, and uncomment the above lines
        #export CU_VERSION="cpu"

        echo "PYTHON_VERSION: $PYTHON_VERSION"
        echo "CU_VERSION: $CU_VERSION"

        ## setup_env.sh
        bash .github/unittest/linux/scripts/run_all.sh