Skip to content

Commit

Permalink
Merge pull request #59 from patrick-kidger/reduce-test-memory-usage
Browse files Browse the repository at this point in the history
Hacky way of reducing test memory usage.
  • Loading branch information
patrick-kidger authored Feb 10, 2022
2 parents 09d8684 + bed2b88 commit 2b4e4d8
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 2 deletions.
4 changes: 2 additions & 2 deletions .github/workflows/run_tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ jobs:
strategy:
matrix:
python-version: [ 3.7, 3.8, 3.9 ]
os: [ macOS-latest ]
os: [ ubuntu-latest ]
fail-fast: false
runs-on: ${{ matrix.os }}
steps:
Expand All @@ -25,7 +25,7 @@ jobs:
- name: Install dependencies
run: |
python -m pip install --upgrade pip
python -m pip install pytest wheel scipy numpy
python -m pip install pytest psutil wheel scipy numpy
- name: Checks with pre-commit
uses: pre-commit/[email protected]
Expand Down
21 changes: 21 additions & 0 deletions test/conftest.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,10 @@
import gc
import random
import sys

import jax.config
import jax.random as jrandom
import psutil
import pytest


Expand All @@ -15,3 +18,21 @@ def _getkey():
return jrandom.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
for module_name, module in sys.modules.items():
if module_name.startswith("jax"):
for obj_name in dir(module):
obj = getattr(module, obj_name)
if hasattr(obj, "cache_clear"):
obj.cache_clear()
gc.collect()

0 comments on commit 2b4e4d8

Please sign in to comment.