From 5c665712867de37a6c75941a1b86ef058daa2a28 Mon Sep 17 00:00:00 2001
From: Blair Bonnett <blair.bonnett@gmail.com>
Date: Tue, 20 Aug 2024 11:34:05 +0200
Subject: [PATCH] Add QUATERNIONIC_DISABLE_CACHE environment variable.

It may not always be desirable to cache the Numba-generated code, for
example in a HPC environment where having lots of workers each compiling
the code may be preferable to them all trying to load it from the cache.

This adds a QUATERNIONIC_DISABLE_CACHE environment variable. If set to
1, the `cache` parameter of the `jit` and `guvectorize` decorators is
set to False. Default behaviour remains the same as before, i.e., to
use the cache.
---
 README.md                 |  9 +++++++++
 quaternionic/utilities.py |  6 ++++--
 tests/test_utilities.py   | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index ff17873..fbbeea1 100644
--- a/README.md
+++ b/README.md
@@ -273,6 +273,15 @@ Finally, there are also capabilities related to interpolation, for example as fu
   * squad (spherical quadratic interpolation)
 
 
+## Caching
+
+By default, the compiled code generated by numba is cached so that the compilation only needs to
+take place on the first import. If you want to disable this caching, for example in a
+high-performance computing environment where it may be preferable to compile the code than try to
+load a cache from disk, set the environment variable `QUATERNIONIC_DISABLE_CACHE` to `1` before
+importing this package.
+
+
 # Related packages
 
 Other python packages with some quaternion features include
diff --git a/quaternionic/utilities.py b/quaternionic/utilities.py
index 56108ac..458d991 100644
--- a/quaternionic/utilities.py
+++ b/quaternionic/utilities.py
@@ -2,6 +2,7 @@
 # See LICENSE file for details:
 # <https://github.com/moble/quaternionic/blob/master/LICENSE>
 
+import os
 import sys
 import functools
 import numpy as np
@@ -213,5 +214,6 @@ def __getitem__(self, *args, **kwargs):
 else:
     import numba
     from numba import float64, boolean
-    jit = functools.partial(numba.njit, cache=True)
-    guvectorize = functools.partial(numba.guvectorize, nopython=True, cache=True)
+    cache = os.environ.get("QUATERNIONIC_DISABLE_CACHE", "0") != "1"
+    jit = functools.partial(numba.njit, cache=cache)
+    guvectorize = functools.partial(numba.guvectorize, nopython=True, cache=cache)
diff --git a/tests/test_utilities.py b/tests/test_utilities.py
index 515c6fa..72e6f89 100644
--- a/tests/test_utilities.py
+++ b/tests/test_utilities.py
@@ -1,3 +1,4 @@
+import subprocess
 import sys
 import numpy as np
 import quaternionic
@@ -118,3 +119,37 @@ def test_pyguvectorize():
                     atol=0.0,
                     rtol=_quaternion_resolution
                 )
+
+
+@pytest.mark.skipif(sys.implementation.name.lower() == "pypy", reason="No numba on pypy")
+def test_cache_disable(tmp_path):
+    # First check caching works by default.
+    cache_dir = tmp_path / "enabled"
+    subprocess.run(
+        [sys.executable, "-c", "import quaternionic"],
+        env={
+            "NUMBA_CACHE_DIR": str(cache_dir),
+        },
+    )
+
+    # Numba uses a subdirectory named `quaternionic_<hashstr>`, with the hashstr computed
+    # from the full path to the source. We should have 1 directory in our cache, and that
+    # should contain many files.
+    contents = list(cache_dir.iterdir())
+    assert len(contents) == 1
+    subdir = contents[0]
+    assert subdir.name.startswith("quaternionic_")
+    assert len(list(subdir.iterdir())) > 10
+
+    # Change the cache location and check with the environment variable disabling caching.
+    # This seems to not create the cache directory, but include an alternative check that
+    # it is empty if it exists in case the directory does get created sometimes.
+    cache_dir = tmp_path / "disabled"
+    subprocess.run(
+        [sys.executable, "-c", "import quaternionic"],
+        env={
+            "QUATERNIONIC_DISABLE_CACHE": "1",
+            "NUMBA_CACHE_DIR": str(cache_dir),
+        },
+    )
+    assert (not cache_dir.exists()) or len(list(cache_dir.iterdir())) == 0