test_gpu.yml

name: GPU tests
on:
  push:
    branches: [main]
  pull_request:
    types: [labeled, opened, synchronize]

concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
  cancel-in-progress: true

defaults:
  run:
    shell: bash -euo pipefail {0}

jobs:
  check:
    name: "Triage: Check if GPU tests are allowed to run"
    if: ${{ !contains(github.event.pull_request.labels.*.name, 'skip-gpu-ci') }}
    runs-on: ubuntu-latest
    steps:
      - uses: flying-sheep/check@v1
        with:
          success: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'run-gpu-ci') }}

  get-environments:
    needs: check
    runs-on: ubuntu-latest
    outputs:
      envs: ${{ steps.get-envs.outputs.envs }}
    steps:
      - uses: actions/checkout@v5
        with:
          filter: blob:none
          fetch-depth: 0
      - name: Install uv
        uses: astral-sh/setup-uv@v7
      - name: Get GPU test environments
        id: get-envs
        run: |
          ENVS_JSON=$(uvx hatch env show --json | jq -c 'to_entries
            | map(
                select(.key | startswith("hatch-test-gpu"))
                | {
                    name: .key,
                    label: .key,
                    python: .value.python
                  }
              )')
          echo "envs=${ENVS_JSON}" | tee $GITHUB_OUTPUT

  test:
    name: ${{ matrix.env.label }}
    needs: get-environments
    runs-on: "cirun-aws-gpu--${{ github.run_id }}"
    timeout-minutes: 30
    strategy:
      fail-fast: false
      matrix:
        env: ${{ fromJSON(needs.get-environments.outputs.envs) }}
    steps:
      - uses: actions/checkout@v5
        with:
          filter: blob:none
          fetch-depth: 0
          submodules: true

      - name: Nvidia SMI sanity check
        run: nvidia-smi

      - name: Install uv
        uses: astral-sh/setup-uv@v7
        with:
          python-version: ${{ matrix.env.python }}
          cache-dependency-glob: pyproject.toml

      - name: Create hatch environment
        run: uvx hatch env create ${{ matrix.env.name }}
      - name: Show installed packages
        run: uvx hatch run ${{ matrix.env.name }}:uv pip freeze
      - name: Run GPU tests
        env:
          MPLBACKEND: agg
        run: uvx hatch run ${{ matrix.env.name }}:run -- -v --color=yes
      - name: Remove 'run-gpu-ci' label
        if: always()
        uses: actions-ecosystem/action-remove-labels@v1
        with:
          labels: "run-gpu-ci"
          github_token: ${{ secrets.GITHUB_TOKEN }}

  check-pass:
    name: GPU tests pass
    if: always()
    needs: [check, get-environments, test]
    runs-on: ubuntu-latest
    steps:
      - uses: re-actors/alls-green@release/v1
        with:
          allowed-skips: check, get-environments, test
          jobs: ${{ toJSON(needs) }}