From 87d913968239178f6eeec25636f3c7786c77984d Mon Sep 17 00:00:00 2001 From: Torax team Date: Fri, 22 Mar 2024 14:13:06 -0700 Subject: [PATCH] Add test sharding and enable more tests `sim_custom_sources.py` is disabled until it runs under pytest. PiperOrigin-RevId: 618284210 --- .github/workflows/pytest.yml | 103 ++++++++++++++++++++++++++++++----- 1 file changed, 89 insertions(+), 14 deletions(-) diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index c7f527f0..9f346897 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -3,18 +3,48 @@ 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: 4 # Controls tests sharding enabled by `pytest-shard` jobs: - pytest-job: + shards-job: + name: Generate shards runs-on: ubuntu-latest - timeout-minutes: 80 - concurrency: - group: ${{ github.workflow }}-${{ github.ref }} - cancel-in-progress: true + 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-shards: + 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 @@ -22,7 +52,7 @@ jobs: # Install deps - uses: actions/setup-python@v5 with: - python-version: "3.10" + 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" @@ -30,19 +60,64 @@ jobs: - run: pip --version - run: pip install -e .[dev] # TODO(b/323504363): [dev] should install these - - run: pip install pytest pytest-xdist + - run: pip install pytest pytest-xdist pytest-shard - run: pip freeze # Run tests (in parallel) # TODO(b/323504363): tests should be discovered automatically - name: Run core tests - run: pytest -vv -n auto torax/tests/{boundary_conditions,config,config_slice,geometry,interpolated_param,jax_utils,math_utils,opt,sim_time_dependence}.py + run: | + pytest \ + 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/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_config.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/config_slice.py \ + torax/tests/config.py \ + torax/tests/geometry.py \ + torax/tests/interpolated_param.py \ + torax/tests/jax_utils.py \ + torax/tests/math_utils.py \ + torax/tests/opt.py \ + torax/tests/physics.py \ + torax/tests/sim_time_dependence.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 }} + + # # TODO(b/323504363): fix UnparsedFlagAccessError and run under pytest + # - name: "Run torax/tests/sim_custom_sources.py" + # run: python torax/tests/sim_custom_sources.py - # TODO(b/323504363): these tests should also run under pytest - - name: "Run sim_custom_sources.py" - run: python torax/tests/sim_custom_sources.py - # TODO(b/323504363): these tests should be parallelized with multiple workers - # - name: "Run sim_no_compile.py" + # # TODO(b/323504363): fix UnparsedFlagAccessError and run under pytest + # - name: "Run torax/sources/tests/formulas.py" + # run: python torax/sources/tests/formulas.py + + # # TODO(b/323504363): fix UnparsedFlagAccessError, pass env variable, and run under pytest + # - name: "Run torax/tests/sim_no_compile.py" # run: TORAX_COMPILATION_ENABLED=False python torax/tests/sim_no_compile.py - # # - name: "Run sim.py" + + # # TODO(b/323504363): fix UnparsedFlagAccessError, parallelize, and run under pytest + # - name: "Run torax/tests/sim.py" # run: python torax/tests/sim.py + + pytest-job: + # Dummy job to enable smooth transition of Copybara configs. + # TODO(b/323504363): remove once Copybara configs are updated. + runs-on: ubuntu-latest + timeout-minutes: 1 + steps: + - run: echo OK