pytest.yml

name: Unittests

# Allow to trigger the workflow manually (e.g. when deps changes)
on: [push, workflow_dispatch]

# Concurrency config borrowed from tensorflow_datasets.
concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.ref != 'refs/heads/master' || github.run_number }}
  # Cancel only PR intermediate builds
  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}

# Set the correct env variables for all the unit tests
env:
  TORAX_ERRORS_ENABLED: 1
  PYTEST_NUM_SHARDS: 8  # Controls tests sharding enabled by `pytest-shard`

jobs:
  shards-job:
    name: Generate shards
    runs-on: ubuntu-latest

    steps:
      - name: Create variables
        id: create-vars
        run: |
          echo "num-shards=$(jq -n -c '[${{ env.PYTEST_NUM_SHARDS }}]')" >> $GITHUB_OUTPUT
          echo "shard-ids=$(jq -n -c '[range(1;${{ env.PYTEST_NUM_SHARDS }}+1)]')" >> $GITHUB_OUTPUT

    outputs:
      num-shards: ${{ steps.create-vars.outputs.num-shards }}
      shard-ids: ${{ steps.create-vars.outputs.shard-ids }}

  pytest-job:
    needs: shards-job

    name: '[${{ matrix.os-version }}][Python ${{ matrix.python-version }}][${{ matrix.shard-id }}/${{ matrix.num-shards }}] Core TORAX tests'
    runs-on: ${{ matrix.os-version }}
    timeout-minutes: 30
    strategy:
      # Do not cancel in-progress jobs if any matrix job fails.
      fail-fast: false
      matrix:
        # Can't reference env variables in matrix
        num-shards: ${{ fromJson(needs.shards-job.outputs.num-shards) }}
        shard-id: ${{ fromJson(needs.shards-job.outputs.shard-ids) }}
        python-version: ['3.10']
        os-version: [ubuntu-latest]

    steps:
    - uses: actions/checkout@v4

    # Install deps
    - uses: actions/setup-python@v5
      with:
        python-version: ${{ matrix.python-version }}

    - run: git clone https://gitlab.com/qualikiz-group/qlknn-hyper.git
    - run: echo "TORAX_QLKNN_MODEL_PATH=$PWD/qlknn-hyper" >> "$GITHUB_ENV"

    - run: pip --version
    - run: pip install -e .[dev]
    - run: pip freeze

    # Run tests (in parallel)
    # TODO(b/326579500): tests should be discovered automatically
    - name: Run core tests
      run: |
        pytest \
          torax/config/tests/runtime_params_slice.py \
          torax/config/tests/runtime_params.py \
          torax/fvm/tests/fvm.py \
          torax/sources/tests/bootstrap_current_source.py \
          torax/sources/tests/current_density_sources.py \
          torax/sources/tests/electron_density_sources.py \
          torax/sources/tests/external_current_source.py \
          torax/sources/tests/formulas.py \
          torax/sources/tests/fusion_heat_source.py \
          torax/sources/tests/generic_ion_el_heat_source.py \
          torax/sources/tests/ion_el_heat_sources.py \
          torax/sources/tests/qei_source.py \
          torax/sources/tests/source_models.py \
          torax/sources/tests/source.py \
          torax/spectators/tests/plotting.py \
          torax/spectators/tests/spectator.py \
          torax/tests/boundary_conditions.py \
          torax/tests/geometry.py \
          torax/tests/interpolated_param.py \
          torax/tests/jax_utils.py \
          torax/tests/math_utils.py \
          torax/tests/persistent_cache.py \
          torax/tests/physics.py \
          torax/tests/test_run_simulation_main.py \
          torax/tests/sim_custom_sources.py \
          torax/tests/sim_time_dependence.py \
          torax/tests/sim.py \
          torax/tests/state.py \
          torax/transport_model/tests/qlknn_wrapper.py \
          torax/transport_model/tests/transport_model.py \
          -vv -n auto \
          --shard-id=$((${{ matrix.shard-id }} - 1)) --num-shards=${{ env.PYTEST_NUM_SHARDS }}

    # One test requires an extra environment variable, so we run it separately.
    - name: Run sim_no_compile
      if: ${{ matrix.shard-id == 1 }}
      run: |
        TORAX_COMPILATION_ENABLED=False pytest -vv -n auto torax/tests/sim_no_compile.py