Fixed tests OOM'ing/crashing/etc.

Author: patrick-kidger

.github/workflows/release.yml

@@ -14,9 +14,9 @@ jobs:
         with:
             python-version: "3.11"
             test-script: |
-                python -m pip install pytest psutil jax jaxlib equinox scipy optax
+                python -m pip install pytest jax jaxlib equinox scipy optax
                 cp -r ${{ github.workspace }}/test ./test
-                pytest
+                python -m test
             pypi-token: ${{ secrets.pypi_token }}
             github-user: patrick-kidger
             github-token: ${{ github.token }}

.github/workflows/run_tests.yml

@@ -23,12 +23,12 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pytest psutil wheel scipy numpy optax jaxlib
+          python -m pip install pytest wheel scipy numpy optax jaxlib
 
       - name: Checks with pre-commit
         uses: pre-commit/action@v2.0.3
 
       - name: Test with pytest
         run: |
           python -m pip install .
-          python -m pytest --durations=0
+          python -m test

test/__main__.py

@@ -0,0 +1,15 @@
+import pathlib
+import subprocess
+import sys
+
+
+here = pathlib.Path(__file__).resolve().parent
+
+
+# Each file is ran separately to avoid out-of-memorying.
+running_out = 0
+for file in here.iterdir():
+    if file.is_file() and file.name.startswith("test"):
+        out = subprocess.run(f"pytest {file}", shell=True).returncode
+        running_out = max(running_out, out)
+sys.exit(running_out)

test/conftest.py

@@ -1,11 +1,8 @@
-import gc
 import random
-import sys
 
 import jax
 import jax.config
-import jax.random as jrandom
-import psutil
+import jax.random as jr
 import pytest
 
 
@@ -16,31 +13,6 @@
 def getkey():
     def _getkey():
         # Not sure what the maximum actually is but this will do
-        return jrandom.PRNGKey(random.randint(0, 2**31 - 1))
+        return jr.PRNGKey(random.randint(0, 2**31 - 1))
 
     return _getkey
-
-
-# Hugely hacky way of reducing memory usage in tests.
-# JAX can be a little over-happy with its caching; this is especially noticable when
-# performing tests and therefore doing an unusual amount of compilation etc.
-# This can be enough to exceed the 8GB RAM available to Ubuntu instances on GitHub
-# Actions.
-@pytest.fixture(autouse=True)
-def clear_caches():
-    process = psutil.Process()
-    if process.memory_info().vms > 4 * 2**30:  # >4GB memory usage
-        jax.clear_backends()
-        for module_name, module in sys.modules.copy().items():
-            if module_name.startswith("jax"):
-                if module_name not in ["jax.interpreters.partial_eval"]:
-                    for obj_name in dir(module):
-                        obj = getattr(module, obj_name)
-                        if hasattr(obj, "cache_clear"):
-                            try:
-                                print(f"Clearing {obj}")
-                                if "Weakref" not in type(obj).__name__:
-                                    obj.cache_clear()
-                            except Exception:
-                                pass
-        gc.collect()