diff --git a/autoconf/jax_wrapper.py b/autoconf/jax_wrapper.py index 2fc3b22..17ece83 100644 --- a/autoconf/jax_wrapper.py +++ b/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. diff --git a/test_autoconf/test_jax_wrapper.py b/test_autoconf/test_jax_wrapper.py new file mode 100644 index 0000000..74c8f9c --- /dev/null +++ b/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"