Create test matrix to speedup tests

.github/workflows/build.yml

@@ -59,44 +59,13 @@ jobs:
           echo "See $url for help on how to resolve this."
           exit 1
         fi
-  test-docs:
-    name: Test documentation
-    runs-on: ubuntu-latest
-    strategy:
-      matrix:
-        python-version: [3.7]
-    steps:
-    - name: Cancel previous
-      uses: styfle/[email protected]
-      with:
-        access_token: ${{ github.token }}          
-    - name: Install native dependencies
-      run: |
-        sudo apt-get -y install pandoc
-    - uses: actions/checkout@v2
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v1
-      with:
-        python-version: ${{ matrix.python-version }}
-    - name: Install dependencies
-      run: |
-        pip install .
-        pip install .[testing]
-        pip install -r docs/requirements.txt
-    - name: Test testcode blocks in documentation
-      run: |
-        sphinx-build -M doctest docs docs/_build -T
-        # test docstrings
-        pytest -n auto flax --doctest-modules --suppress-no-test-exit-code
-    - name: Build documentation
-      run: |
-        sphinx-build -M html docs docs/_build -T
-  run-tests:
-    name: Run tests
+  tests:
+    name: Run Tests
     runs-on: ubuntu-latest
     strategy:
       matrix:
         python-version: [3.7]
+        test-type: [doctest, pytest, pytype]
     steps:
     - name: Cancel previous
       uses: styfle/[email protected]
@@ -112,10 +81,20 @@ jobs:
         pip install .
         pip install .[testing]
         pip install -r docs/requirements.txt
-    - name: Test with pytest and generate coverage report
+    - name: Test with ${{ matrix.test-type }}
       run: |
-        tests/run_all_tests.sh --with-cov --no-doctest
+        if [[ "${{ matrix.test-type }}" == "doctest" ]]; then
+          tests/run_all_tests.sh --no-pytest --no-pytype
+        elif [[ "${{ matrix.test-type }}" == "pytest" ]]; then
+          tests/run_all_tests.sh --no-doctest --no-pytype --with-cov
+        elif [[ "${{ matrix.test-type }}" == "pytype" ]]; then
+          tests/run_all_tests.sh --no-doctest --no-pytest
+        else
+          echo "Unknown test type: ${{ matrix.test-type }}"
+          exit 1
+        fi
     - name: Upload coverage to Codecov
+      if: matrix.test-type == 'pytest'
       uses: codecov/codecov-action@v1
       with:
         file: ./coverage.xml

tests/run_all_tests.sh

@@ -1,11 +1,10 @@
 #!/bin/bash
 
-export JAX_NUMPY_RANK_PROMOTION=raise
-export FLAX_PROFILE=1
-export FLAX_LAZY_RNG=1
-
 PYTEST_OPTS=
 RUN_DOCTEST=true
+RUN_PYTEST=true
+RUN_PYTYPE=true
+
 for flag in "$@"; do
 case $flag in
   --with-cov)
@@ -19,13 +18,29 @@ case $flag in
   --no-doctest)
   RUN_DOCTEST=false
   ;;
+  --no-pytest)
+  RUN_PYTEST=false
+  ;;
+  --no-pytype)
+  RUN_PYTYPE=false
+  ;;
   *)
   echo "Unknown flag: $flag"
   exit 1
   ;;
 esac
 done
 
+echo "====== test config ======="
+echo "PYTEST_OPTS: $PYTEST_OPTS"
+echo "RUN_DOCTEST: $RUN_DOCTEST"
+echo "RUN_PYTEST: $RUN_PYTEST"
+echo "RUN_PYTYPE: $RUN_PYTYPE"
+echo "jax: $(python -c 'import jax; print(jax.__version__)')"
+echo "flax: $(python -c 'import flax; print(flax.__version__)')"
+echo "=========================="
+echo ""
+
 sh $(dirname "$0")/download_dataset_metadata.sh || exit
 
 # Instead of using set -e, we have a manual error trap that
@@ -45,40 +60,55 @@ handle_errors () {
 
 # Run embedded tests inside docs
 if $RUN_DOCTEST; then
+  echo "=== RUNNING DOCTESTS ==="
+  # test doctest
   sphinx-build -M doctest docs docs/_build -T
+  # test build html
+  sphinx-build -M html docs docs/_build -T
   # test docstrings
   pytest -n auto flax --doctest-modules --suppress-no-test-exit-code
 fi
 
-# Run some test on separate process, avoiding device configs poluting each other
-PYTEST_IGNORE=
-for file in "tests/jax_utils_test.py"; do
-    echo "pytest -n auto $file $PYTEST_OPTS"
-    pytest -n auto $file $PYTEST_OPTS
-    PYTEST_IGNORE+=" --ignore=$file"
-done
-# Run battery of core FLAX API tests.
-pytest -n auto tests $PYTEST_OPTS $PYTEST_IGNORE
+# env vars must be set after doctest
+export JAX_NUMPY_RANK_PROMOTION=raise
+export FLAX_PROFILE=1
+export FLAX_LAZY_RNG=1
 
-# Per-example tests.
-#
-# we apply pytest within each example to avoid pytest's annoying test-filename collision.
-# In pytest foo/bar/baz_test.py and baz/bleep/baz_test.py will collide and error out when
-# /foo/bar and /baz/bleep aren't set up as packages.
-for egd in $(find examples -maxdepth 1 -mindepth 1 -type d); do
-    pytest $egd
-done
+if $RUN_PYTEST; then
+  echo "=== RUNNING PYTESTS ==="
+  # Run some test on separate process, avoiding device configs poluting each other
+  PYTEST_IGNORE=
+  for file in "tests/jax_utils_test.py"; do
+      echo "pytest -n auto $file $PYTEST_OPTS"
+      pytest -n auto $file $PYTEST_OPTS
+      PYTEST_IGNORE+=" --ignore=$file"
+  done
+  # Run battery of core FLAX API tests.
+  pytest -n auto tests $PYTEST_OPTS $PYTEST_IGNORE
 
-# Validate types in library code.
-pytype --config pytype.cfg flax/
+  # Per-example tests.
+  #
+  # we apply pytest within each example to avoid pytest's annoying test-filename collision.
+  # In pytest foo/bar/baz_test.py and baz/bleep/baz_test.py will collide and error out when
+  # /foo/bar and /baz/bleep aren't set up as packages.
+  for egd in $(find examples -maxdepth 1 -mindepth 1 -type d); do
+      pytest $egd
+  done
+fi
 
-# Validate types in examples.
-for egd in $(find examples -maxdepth 1 -mindepth 1 -type d); do
-    # use cd to make sure pytpe cache lives in example dir and doesn't name clash
-    # use *.py to avoid importing configs as a top-level import which leads to import errors
-    # because config files use relative imports (e.g. from config import ...).
-    (cd $egd ; pytype --config ../../pytype.cfg "*.py")
-done
+if $RUN_PYTYPE; then
+  echo "=== RUNNING PYTYPE ==="
+  # Validate types in library code.
+  pytype --config pytype.cfg flax/
+
+  # Validate types in examples.
+  for egd in $(find examples -maxdepth 1 -mindepth 1 -type d); do
+      # use cd to make sure pytpe cache lives in example dir and doesn't name clash
+      # use *.py to avoid importing configs as a top-level import which leads to import errors
+      # because config files use relative imports (e.g. from config import ...).
+      (cd $egd ; pytype --config ../../pytype.cfg "*.py")
+  done
+fi
 
 # Return error code 0 if no real failures happened.
 echo "finished all tests."