Prepare for v0.0.8 release

- Remove TensorFlow from dependencies.
- Drop Python 3.6 support (JAX has dropped Python 3.6 support in July 2021).
- Re-enable github test workflow.

PiperOrigin-RevId: 408471735

.github/workflows/build_and_minimal_test.yml

@@ -1,47 +1,42 @@
 name: Build and minimal test
 
-on: [push, pull_request]
+on:
+  push:
+    branches:
+    - main
+  pull_request:
+    branches:
+    - main
 
 jobs:
   build:
 
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.6, 3.7, 3.8, 3.9]
+        python-version: [3.7, 3.8, 3.9]
 
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
-    - name: Cache pip
-      # Please read the documentation on
-      # https://docs.github.com/en/actions/guides/caching-dependencies-to-speed-up-workflows
-      # before making changes to this step.
-      uses: actions/cache@v2
-      with:
-        # Cache installed Python packages instead of the pip download cache.
-        path: ${{ env.pythonLocation }}/lib/python${{ matrix.python-version }}/site-packages
-        # Look to see if there is a cache hit for the corresponding setup.py.
-        key: ${{ runner.os }}-py${{ matrix.python-version }}-pip-${{ hashFiles('setup.py') }}
-        # Try the cache with a different setup.py.
-        restore-keys: ${{ runner.os }}-py${{ matrix.python-version }}-pip-
     - name: Install dependencies
       # Use `install -e` for local fedjax to prevent actual fedjax files from
       # being packaged into cache.
       run: |
         python -m pip install --upgrade pip setuptools wheel
         pip install -e .
+    - name: Basic test
+      # We run tests using `python -I` to prevent the script directory to be
+      # added to sys.path.
+      run: |
+        python -I fedjax/fedjax_test.py
+    - name: Full test
+      run: |
+        ./run_tests.sh
+    - name: Build documentation
+      run: |
         pip install -r docs/requirements.txt
-        pip install -r requirements-test.txt
-    # # TODO(b/191778948): Reselect once we understand absl flags and pytest.
-    # - name: Test
-    #   run: |
-    #     pytest -n auto -q \
-    #       -k "not SubsetFederatedDataTest and not SQLiteFederatedDataTest and not ForEachClientPmapTest and not DownloadsTest and not CheckpointTest and not LoggingTest" \
-    #       fedjax --ignore=fedjax/legacy/
-    # - name: Build documentation
-    #   run: |
-    #     sphinx-build -M html docs docs/_build
+        sphinx-build -M html docs docs/_build

docs/developer.md

@@ -25,55 +25,12 @@ symbolic links from site-packages into the repository.
 
 # Running the tests
 
-## Running specific tests
+We created a simple [`run_tests.sh`](https://www.google.com/url?sa=D&q=https%3A%2F%2Fgithub.com%2Fgoogle%2Ffedjax%2Fblob%2Fmain%2Frun_tests.sh)
+script for running tests. See its comments for examples. Before creating a pull
+request, we recommend running all the FedJAX tests (i.e. running `run_tests.sh`
+with no arguments) to verify the correctness of a change.
 
-You can run a specific set of tests using
-[pytest](https://docs.pytest.org/en/latest/usage.html#specifying-tests-selecting-tests)'s
-built-in selection mechanisms, or alternatively, you can run a specific test
-file directly to see more detailed information about the cases being run.
-
-We recommend using `pytest-xdist`, which can run tests in parallel. First,
-install `pytest-xdist` and other test-only dependencies:
-
-```bash
-pip install -r requirements-test.txt
-```
-
-Then, from the repository root directory, run:
-
-```bash
-# Run all tests in algorithms
-pytest -n auto fedjax/algorithms
-
-# Run only fedjax/core/metrics_test.py
-pytest -n auto fedjax/core/metrics_test.py
-```
-
-The `-n auto` tells pytest to use as many processes as your computer has CPU
-cores. For more details, see the
-[pytest-xdist docs](https://github.com/pytest-dev/pytest-xdist#speed-up-test-runs-by-sending-tests-to-multiple-cpus).
-
-We also use a handful of custom configurations in the pytest.ini file. For more
-details, see the
-[pytest docs](https://docs.pytest.org/en/6.2.x/reference.html#ini-options-ref).
-
-## Running all CI tests
-
-Before creating a pull request, we recommend running all the FedJAX tests that
-are used in the continuous integration:
-
-```bash
-pytest -n auto -q \
-  -k "not SubsetFederatedDataTest and not SQLiteFederatedDataTest and not ForEachClientPmapTest and not DownloadsTest and not CheckpointTest and not LoggingTest" \
-  fedjax --ignore=fedjax/legacy/
-```
-
-`-q` will reduce verbosity levels, `-k` selects/deselects specific tests, and
-`--ignore=fedjax/legacy/` is used to skip the entire fedjax.legacy module. If
-there are errors or failures, you can run those specific tests using the
-commands in the previous section to see more focused details.
-
-# Update the docs
+# Updating the docs
 
 Install the requirements
 
@@ -141,12 +98,12 @@ You can use the second method if the `path_to_obj` is very long.
 # :func:`my_func`
 
 # Create a reference "Module.apply()" to method fedjax.experimental.model.Model.apply_for_train.
-# :meth:`Model.apply_for_train <fedjax.experimental.model.Model.apply_for_train>` 
-``` 
+# :meth:`Model.apply_for_train <fedjax.experimental.model.Model.apply_for_train>`
+```
 
 To create a hyperlink, use the following syntax:
 ```bash
-# Note the double underscore at the end: 
+# Note the double underscore at the end:
 # `Link to Google <http://www.google.com>`__
 ```
 
@@ -185,7 +142,7 @@ Update it as needed, `Run all cells` then `Download ipynb` to your local repo.
 You may want to test that it executes properly, using `sphinx-build` as explained above.
 We recommend making changes this way to avoid introducing format errors into the `.ipynb` files.
 
-In the future, we may build and re-execute the notebooks as part of the 
+In the future, we may build and re-execute the notebooks as part of the
 [Read the docs](https://fedjax.readthedocs.io/en/latest) build.
 However, for now, we exclude all notebooks from the build due to long durations (downloading dataset files, expensive model training, etc.).
 See `exclude_patterns` in [conf.py](https://github.com/google/fedjax/blob/main/docs/conf.py)

fedjax/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 """FedJAX version."""
 
-__version__ = '0.0.7'
+__version__ = '0.0.8'

requirements-test.txt

@@ -1,2 +1,2 @@
-pytest-xdist
 flax
+tensorflow-cpu

run_tests.sh

@@ -0,0 +1,45 @@
+#!/bin/bash
+#
+# Script for running tests in OSS development.
+#
+# Examples:
+#
+# * Running all tests under fedjax/ (default):
+#
+#     ./run_tests.sh
+#
+# * Running tests in directory fedjax/core:
+#
+#     ./run_tests.sh fedjax/core
+#
+# * Running tests in fedjax/fedjax_test.py and directory fedjax/core:
+#
+#     ./run_tests.sh fedjax/fedjax_test.py fedjax/core
+
+set -e
+set -o pipefail
+
+if [[ "$#" -eq 0 ]]; then
+  TO_RUN=(fedjax)
+else
+  TO_RUN="$@"
+fi
+
+# Install build/test dependencies.
+pip install -e .
+pip install -r requirements-test.txt
+
+# Run tests in serial.
+#
+# We use 'python -I' to prevent the script directory from being included in
+# sys.path.
+#
+# TODO(wuke): Improve readability of test failures.
+#
+# TODO(wuke): Run tests in parallel. We need to deal with race conditions caused
+# by multiple tests accessing temporary files with the same name (e.g.
+# "test_sqlite_federated_data.sqlite" in both
+# fedjax/core/sqlite_federated_data_test.py and
+# fedjax/core/federated_data_test.py
+find "${TO_RUN[@]}" -name '*_test.py' -print0 | \
+  xargs -0 -n1 -- sh -c 'python -I $0 2>&1 || exit 255'

setup.py

@@ -42,13 +42,12 @@
         'optax',
         'requests',
     ],
-    python_requires='>=3.6',
+    python_requires='>=3.7',
     classifiers=[
         'Development Status :: 3 - Alpha',
         'Intended Audience :: Developers',
         'Intended Audience :: Science/Research',
         'License :: OSI Approved :: Apache Software License',
-        'Programming Language :: Python :: 3.6',
         'Programming Language :: Python :: 3.7',
         'Programming Language :: Python :: 3.8',
         'Programming Language :: Python :: 3.9',