tests.yml

name: Tests

on:
  schedule:
  - cron: 00 00 * * 1
  push:
    branches: [main]
  pull_request:
    branches: [main]

concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}
  cancel-in-progress: true

jobs:
  fast-tests:
    name: Fast tests Python ${{ matrix.python-version }} ${{ matrix.jax-version }}
    runs-on: ubuntu-latest
    # allow tests using the latest JAX to fail
    continue-on-error: ${{ matrix.jax-version == 'jax-latest' }}
    strategy:
      fail-fast: false
      matrix:
        python-version: ['3.10']
        jax-version: [jax-default, jax-latest]

    steps:
    - uses: actions/checkout@v4
    - name: Set up Python 3.9
      uses: actions/setup-python@v5
      with:
        python-version: ${{ matrix.python-version }}

    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
        python -m pip install tox

    - name: Setup environment
      run: |
        tox -e py${{ matrix.python-version }}-${{ matrix.jax-version }} --notest -v

    - name: Run tests
      run: |
        tox -e py${{ matrix.python-version }}-${{ matrix.jax-version }} --skip-pkg-install -- -m fast --memray -n auto -vv

  gpu-tests:
    name: Fast GPU tests Python 3.10 on ubuntu-22.04
    runs-on: [self-hosted, ott-gpu]
    container:
      image: docker://michalk8/cuda:12.2.2-cudnn8-devel-ubuntu22.04
      options: --gpus="device=2"
    steps:
    - uses: actions/checkout@v4

    - name: Install dependencies
      run: |
        python3 -m pip install --upgrade pip
        python3 -m pip install -e".[test]"
        python3 -m pip install "jax[cuda12]"

    - name: Run nvidia-smi
      run: |
        nvidia-smi

    - name: Run tests
      run: |
        python3 -m pytest -m "fast and not cpu" --memray --durations 10 -vv
      env:
        XLA_PYTHON_CLIENT_PREALLOCATE: 'false'

  tests:
    name: Python ${{ matrix.python-version }} on ${{ matrix.os }}
    runs-on: ${{ matrix.os }}
    strategy:
      fail-fast: false
      matrix:
        python-version: ['3.10', '3.11', '3.12', '3.13']
        os: [ubuntu-latest]
        include:
        - python-version: '3.9'
          os: macos-13
        - python-version: '3.10'
          os: macos-14

    steps:
    - uses: actions/checkout@v4
    - name: Set up Python ${{ matrix.python-version }}
      uses: actions/setup-python@v5
      with:
        python-version: ${{ matrix.python-version }}

    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
        python -m pip install tox

    - name: Setup environment
      run: |
        tox -e py${{ matrix.python-version }} --notest -v

    - name: Run tests
      run: |
        tox -e py${{ matrix.python-version }} --skip-pkg-install
      env:
        PYTEST_ADDOPTS: --memray -vv

    - name: Upload coverage
      uses: codecov/codecov-action@v4
      with:
        files: ./coverage.xml
        flags: tests-${{ matrix.os }}-${{ matrix.python-version }}
        name: unittests
        token: ${{ secrets.CODECOV_TOKEN }}
        env_vars: OS,PYTHON
        fail_ci_if_error: false
        verbose: true