Merge pull request #128 from PyAutoLabs/feature/jax-cache-default

feat: JAX persistent compilation cache by default + XLA_FLAGS append fix (#127)

Author: Jammy2211

autoconf/jax_wrapper.py

@@ -16,21 +16,28 @@
 if not xla_env_set:
     logger.info(
         """
-        For fast JAX compile times, the envirment variable XLA_FLAGS must be set to "--xla_disable_hlo_passes=constant_folding",
+        For fast JAX compile times, the envirment variable XLA_FLAGS must include "--xla_disable_hlo_passes=constant_folding",
         which is currently not.
-        
-        In Python, to do this manually, use the code: 
-        
+
+        In Python, to do this manually, use the code:
+
         import os
         os.environ["XLA_FLAGS"] = "--xla_disable_hlo_passes=constant_folding"
-        
-        The environment variable has been set automatically for you now, however if JAX has already been imported, 
-        this change will not take effect and JAX function compiling times may be slow. 
-        
+
+        The environment variable has been set automatically for you now, however if JAX has already been imported,
+        this change will not take effect and JAX function compiling times may be slow.
+
         Therefore, it is recommended to set this environment variable before running your script, e.g. in your terminal.
         """)
 
-    os.environ['XLA_FLAGS'] = "--xla_disable_hlo_passes=constant_folding"
+    # Append rather than overwrite: replacing the value silently discarded any
+    # flags the user or a batch script had set (e.g. --xla_dump_to=...,
+    # --xla_gpu_autotune_level=0), which is indistinguishable from those flags
+    # having no effect.
+    if xla_env:
+        os.environ["XLA_FLAGS"] = f"{xla_env} --xla_disable_hlo_passes=constant_folding"
+    else:
+        os.environ["XLA_FLAGS"] = "--xla_disable_hlo_passes=constant_folding"
 
 jax_enable_x64 = os.environ.get("JAX_ENABLE_X64")
 
@@ -53,6 +60,33 @@
         """
     )
 
+if "JAX_COMPILATION_CACHE_DIR" not in os.environ:
+
+    _cache_root = os.environ.get("XDG_CACHE_HOME") or os.path.join(
+        os.path.expanduser("~"), ".cache"
+    )
+    _cache_dir = os.path.join(_cache_root, "pyauto_jax")
+
+    os.environ["JAX_COMPILATION_CACHE_DIR"] = _cache_dir
+
+    logger.info(
+        f"""
+        The JAX persistent compilation cache has been enabled at {_cache_dir}
+        (JAX_COMPILATION_CACHE_DIR). The first fit of a given model and data shape
+        on this machine compiles its JAX functions, which can take minutes; the
+        compiled code is cached on disk, so later runs (including after restarting
+        Python) skip this cost.
+
+        To use a different location, set JAX_COMPILATION_CACHE_DIR before running
+        your script. To disable the cache entirely, set it to an empty string.
+        """
+    )
+
+# An explicitly empty JAX_COMPILATION_CACHE_DIR means "cache disabled"; do not
+# force a compile-time threshold in that case.
+if os.environ.get("JAX_COMPILATION_CACHE_DIR"):
+    os.environ.setdefault("JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS", "1")
+
 
 def register_pytree_node_class(cls):
     """Opt-in JAX pytree class registration that defers the JAX import.

test_autoconf/test_jax_wrapper.py

@@ -0,0 +1,94 @@
+"""Tests for the env-var handling in autoconf.jax_wrapper.
+
+The wrapper's env logic runs at import time, so each test reloads the module
+under a controlled os.environ. No test imports jax — the wrapper only sets
+environment variables.
+"""
+
+import importlib
+import os
+
+import pytest
+
+import autoconf.jax_wrapper
+
+CONSTANT_FOLDING = "--xla_disable_hlo_passes=constant_folding"
+
+
+@pytest.fixture
+def clean_env(monkeypatch):
+    for key in (
+        "XLA_FLAGS",
+        "JAX_ENABLE_X64",
+        "JAX_COMPILATION_CACHE_DIR",
+        "JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS",
+        "XDG_CACHE_HOME",
+    ):
+        monkeypatch.delenv(key, raising=False)
+    yield monkeypatch
+    importlib.reload(autoconf.jax_wrapper)
+
+
+def reload_wrapper():
+    return importlib.reload(autoconf.jax_wrapper)
+
+
+def test_xla_flags_set_when_unset(clean_env):
+    reload_wrapper()
+    assert os.environ["XLA_FLAGS"] == CONSTANT_FOLDING
+
+
+def test_xla_flags_appended_not_clobbered(clean_env):
+    clean_env.setenv("XLA_FLAGS", "--xla_dump_to=/tmp/foo --xla_gpu_autotune_level=0")
+    reload_wrapper()
+    flags = os.environ["XLA_FLAGS"]
+    assert "--xla_dump_to=/tmp/foo" in flags
+    assert "--xla_gpu_autotune_level=0" in flags
+    assert CONSTANT_FOLDING in flags
+
+
+def test_xla_flags_unchanged_when_already_present(clean_env):
+    preset = f"--xla_dump_to=/tmp/foo {CONSTANT_FOLDING}"
+    clean_env.setenv("XLA_FLAGS", preset)
+    reload_wrapper()
+    assert os.environ["XLA_FLAGS"] == preset
+
+
+def test_cache_dir_defaulted_when_unset(clean_env):
+    reload_wrapper()
+    expected = os.path.join(os.path.expanduser("~"), ".cache", "pyauto_jax")
+    assert os.environ["JAX_COMPILATION_CACHE_DIR"] == expected
+    assert os.environ["JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS"] == "1"
+
+
+def test_cache_dir_respects_xdg_cache_home(clean_env):
+    clean_env.setenv("XDG_CACHE_HOME", "/custom/cache")
+    reload_wrapper()
+    assert os.environ["JAX_COMPILATION_CACHE_DIR"] == os.path.join(
+        "/custom/cache", "pyauto_jax"
+    )
+
+
+def test_cache_dir_respects_preset_value(clean_env):
+    clean_env.setenv("JAX_COMPILATION_CACHE_DIR", "/my/cache")
+    reload_wrapper()
+    assert os.environ["JAX_COMPILATION_CACHE_DIR"] == "/my/cache"
+    assert os.environ["JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS"] == "1"
+
+
+def test_cache_disabled_by_empty_string(clean_env):
+    clean_env.setenv("JAX_COMPILATION_CACHE_DIR", "")
+    reload_wrapper()
+    assert os.environ["JAX_COMPILATION_CACHE_DIR"] == ""
+    assert "JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS" not in os.environ
+
+
+def test_min_compile_time_respects_preset_value(clean_env):
+    clean_env.setenv("JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS", "10")
+    reload_wrapper()
+    assert os.environ["JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS"] == "10"
+
+
+def test_x64_enabled_by_default(clean_env):
+    reload_wrapper()
+    assert os.environ["JAX_ENABLE_X64"] == "True"