test-gpu.yml

name: AWS GPU

on:
  push:
    branches: [main, "[0-9]+.[0-9]+.x"]
  pull_request:
    types:
      - labeled
      - opened
      - synchronize

env:
  PYTEST_ADDOPTS: "-v --color=yes"
  FORCE_COLOR: "1"

permissions: {}

# Cancel the job if new commits are pushed
# https://stackoverflow.com/questions/66335225/how-to-cancel-previous-runs-in-the-pr-when-you-push-new-commitsupdate-the-curre
concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
  cancel-in-progress: true

# There are two jobs:
# 1. `check` determines if the second job (`test`) will be run (through a job dependency).
# 2. `test` runs on an AWS runner and executes the GPU tests.
jobs:
  # If the `skip-gpu-ci` label is set, this job is skipped, and consequently the `test` job too.
  # If the `run-gpu-ci` label is set or we reacted to a `push` event, this job succeeds (and `test` is run).
  # If neither is set, this job fails, `test` is skipped, and the whole workflow fails.
  check:
    name: "Triage: Check if GPU tests are allowed to run"
    if: ${{ !(github.event.pull_request.user.login == 'pre-commit-ci[bot]' || contains(github.event.pull_request.labels.*.name, 'skip-gpu-ci')) }}
    runs-on: ubuntu-latest
    permissions: {}
    steps:
      - uses: flying-sheep/check@6162e3da6b50abdfe00c3233e6f0ce2d1bdcf69c # v1.0.4
        with:
          success: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'run-gpu-ci') }}
  # If `check` wasn’t skipped or failed, start an AWS runner and run the GPU tests on it.
  test:
    name: GPU Tests
    needs: check
    runs-on: "cirun-aws-gpu--${{ github.run_id }}"
    # Setting a timeout of 30 minutes, as the AWS costs money
    # At time of writing, a typical run takes about 5 minutes
    timeout-minutes: 30
    permissions:
      pull-requests: write

    steps:
      # TODO: upgrade once cirun image supports node 24
      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
        with:
          fetch-depth: 0
          filter: blob:none
          persist-credentials: false

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

      - name: Install UV
        # TODO: upgrade once cirun image supports node 24
        uses: astral-sh/setup-uv@d0cc045d04ccac9d8b7881df0226f9e82c39688e # v6.8.0
        with:
          enable-cache: true

      - name: Extract max Python version from classifiers
        run: |
          import os
          import check_python_versions.sources.pyproject as cpv
          *_, max_ver = cpv.get_supported_python_versions()
          with open(os.environ["GITHUB_ENV"], "a") as env:
            print(f"UV_PYTHON={max_ver}", file=env)
        shell: uvx --with=check-python-versions==0.24 python {0}

      - name: Install AnnData
        run: |
          uv venv
          uv pip install -e ".[cu12]" --group=test --group=test-jax -c ci/constraints.txt

      - name: Env list
        run: uv pip list

      - name: Run tests
        run: |
          uv run coverage run -m pytest -m "array_api or gpu" -n auto
          uv run coverage combine
          uv run coverage xml

      - uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
        with:
          token: ${{ secrets.CODECOV_TOKEN }}
          fail_ci_if_error: true
          files: test-data/coverage.xml

      - name: Remove 'run-gpu-ci' Label
        if: always()
        uses: actions-ecosystem/action-remove-labels@2ce5d41b4b6aa8503e285553f75ed56e0a40bae0 # v1.3.0
        with:
          labels: "run-gpu-ci"
          github_token: ${{ secrets.GITHUB_TOKEN }}