Unified cache under sys.prefix

AGENTS.md

@@ -79,10 +79,15 @@ Firedrake's toolchain, in order:
 * **`petsc4py`/PETSc version skew:** Rebuild `petsc4py` (`pip install --no-build-isolation -e .`) after
   switching the PETSc branch/commit under an existing venv. A stale `petsc4py` fails `import firedrake`
   with an `undefined symbol: ...` error, not a Firedrake traceback.
-* **Caching:** TSFC kernels and PyOP2 code are cached under
-  `FIREDRAKE_TSFC_KERNEL_CACHE_DIR`/`PYOP2_CACHE_DIR` (default `$VIRTUAL_ENV/.cache/{tsfc,pyop2}`), set
-  by `firedrake.configuration.setup_cache_dirs()` on `import firedrake`. Run `firedrake-clean` if a
-  change to the code generator does not take effect.
+* **Caching:** `import firedrake` sets `FIREDRAKE_TSFC_KERNEL_CACHE_DIR`, `PYOP2_CACHE_DIR`, and
+  `XDG_CACHE_HOME`, all derived from `FIREDRAKE_CACHE_DIR` (default `sys.prefix/.cache`, or `~/.cache`
+  if that isn't writable). Run `firedrake-clean` if a change to the code generator does not take
+  effect.
+* **Isolated cache:** Set `FIREDRAKE_CACHE_DIR` to a scratch directory before running Firedrake to keep
+  one variant's cached kernels from leaking into another's, e.g. for an A/B comparison.
+* **Scratch caches are yours to clean up:** `firedrake-clean` only clears the caches implied by the
+  *current* environment, so it can't find a `FIREDRAKE_CACHE_DIR` from an earlier, different run.
+  Delete a scratch cache directory yourself once you're done with it.
 * **Smoke test:** `firedrake-check` runs a process-count-grouped subset of the regression suite; use it
   before a full run.
 

firedrake/configuration.py

@@ -1,12 +1,23 @@
 """Replaces functionality from the removed `firedrake_configuration` module."""
 
 import os
+import sys
 from pathlib import Path
 
 
 def setup_cache_dirs():
-    root = Path(os.environ.get("VIRTUAL_ENV", Path.home())).joinpath(".cache")
+    root = os.environ.get("FIREDRAKE_CACHE_DIR")
+    if root is None:
+        prefix = Path(sys.prefix)
+        base = prefix if os.access(prefix, os.W_OK) else Path.home()
+        root = str(base.joinpath(".cache"))
+
+    root = Path(root)
     if "PYOP2_CACHE_DIR" not in os.environ:
         os.environ["PYOP2_CACHE_DIR"] = str(root.joinpath("pyop2"))
-    if 'FIREDRAKE_TSFC_KERNEL_CACHE_DIR' not in os.environ:
+    if "FIREDRAKE_TSFC_KERNEL_CACHE_DIR" not in os.environ:
         os.environ["FIREDRAKE_TSFC_KERNEL_CACHE_DIR"] = str(root.joinpath("tsfc"))
+    # loopy's persistent caches go through pytools, which only listens to
+    # XDG_CACHE_HOME (or platformdirs' default) for its cache location.
+    if "XDG_CACHE_HOME" not in os.environ:
+        os.environ["XDG_CACHE_HOME"] = str(root)

tests/firedrake/test_0init.py

@@ -1,6 +1,7 @@
 import os
+import pytest
 from firedrake import *
-from pathlib import Path
+from firedrake.configuration import setup_cache_dirs
 
 
 def test_pyop2_custom_init():
@@ -15,6 +16,75 @@ def test_pyop2_custom_init():
 
 
 def test_pyop2_cache_dir_set_correctly():
-    root = Path(os.environ.get("VIRTUAL_ENV", "~")).joinpath(".cache")
-    cache_dir = os.environ.get("PYOP2_CACHE_DIR", str(root.joinpath("pyop2")))
-    assert op2.configuration["cache_dir"] == cache_dir
+    assert "PYOP2_CACHE_DIR" in os.environ
+    assert op2.configuration["cache_dir"] == os.environ["PYOP2_CACHE_DIR"]
+
+
+CACHE_ENV_VARS = (
+    "FIREDRAKE_CACHE_DIR",
+    "PYOP2_CACHE_DIR",
+    "FIREDRAKE_TSFC_KERNEL_CACHE_DIR",
+    "XDG_CACHE_HOME",
+)
+
+
+@pytest.fixture
+def clean_cache_env(monkeypatch):
+    # monkeypatch.delenv() snapshots each variable's prior value (set or
+    # unset) and restores it automatically at teardown, even though
+    # setup_cache_dirs() itself writes to os.environ directly.
+    for var in CACHE_ENV_VARS:
+        monkeypatch.delenv(var, raising=False)
+
+
+def test_setup_cache_dirs_uses_writable_sys_prefix(clean_cache_env, monkeypatch, tmp_path):
+    monkeypatch.setattr("sys.prefix", str(tmp_path))
+
+    setup_cache_dirs()
+
+    root = tmp_path.joinpath(".cache")
+    assert os.environ["PYOP2_CACHE_DIR"] == str(root.joinpath("pyop2"))
+    assert os.environ["FIREDRAKE_TSFC_KERNEL_CACHE_DIR"] == str(root.joinpath("tsfc"))
+    assert os.environ["XDG_CACHE_HOME"] == str(root)
+
+
+def test_setup_cache_dirs_falls_back_when_sys_prefix_is_not_writable(clean_cache_env, monkeypatch, tmp_path):
+    # os.access(..., os.W_OK) always reports True for root regardless of the
+    # file mode, so a real chmod can't simulate "not writable" in CI; fake
+    # os.access itself instead.
+    unwritable = tmp_path.joinpath("unwritable")
+    unwritable.mkdir()
+    monkeypatch.setattr("sys.prefix", str(unwritable))
+    monkeypatch.setattr("pathlib.Path.home", lambda: tmp_path.joinpath("home"))
+    real_access = os.access
+    monkeypatch.setattr(
+        "os.access",
+        lambda path, mode, *a, **kw: False if str(path) == str(unwritable) else real_access(path, mode, *a, **kw),
+    )
+
+    setup_cache_dirs()
+
+    root = tmp_path.joinpath("home", ".cache")
+    assert os.environ["PYOP2_CACHE_DIR"] == str(root.joinpath("pyop2"))
+
+
+def test_setup_cache_dirs_honours_firedrake_cache_dir(clean_cache_env, monkeypatch, tmp_path):
+    monkeypatch.setenv("FIREDRAKE_CACHE_DIR", str(tmp_path))
+
+    setup_cache_dirs()
+
+    assert os.environ["PYOP2_CACHE_DIR"] == str(tmp_path.joinpath("pyop2"))
+    assert os.environ["FIREDRAKE_TSFC_KERNEL_CACHE_DIR"] == str(tmp_path.joinpath("tsfc"))
+    assert os.environ["XDG_CACHE_HOME"] == str(tmp_path)
+
+
+def test_setup_cache_dirs_does_not_override_explicit_settings(clean_cache_env, monkeypatch, tmp_path):
+    monkeypatch.setenv("FIREDRAKE_CACHE_DIR", str(tmp_path))
+    monkeypatch.setenv("PYOP2_CACHE_DIR", str(tmp_path.joinpath("custom-pyop2")))
+    monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path.joinpath("custom-xdg")))
+
+    setup_cache_dirs()
+
+    assert os.environ["PYOP2_CACHE_DIR"] == str(tmp_path.joinpath("custom-pyop2"))
+    assert os.environ["XDG_CACHE_HOME"] == str(tmp_path.joinpath("custom-xdg"))
+    assert os.environ["FIREDRAKE_TSFC_KERNEL_CACHE_DIR"] == str(tmp_path.joinpath("tsfc"))