jax-ml/jax · error · RuntimeError

Please install the `filelock` package to set `jax_compilatio

Error message

Please install the `filelock` package to set `jax_compilation_cache_max_size`

What it means

Enabling cache eviction (any jax_compilation_cache_max_size other than -1) requires the filelock package for cross-process locking. If filelock is not installed, LRUCache.__init__ raises RuntimeError at setup time.

Source

Thrown at jax/_src/lru_cache.py:73

    """Args:

      path: The path to the cache directory.
      max_size: The maximum size of the cache in bytes. Caching will be
        disabled if this value is set to ``0``. A special value of ``-1``
        indicates no limit, allowing the cache size to grow indefinitely.
      lock_timeout_secs: (optional) The timeout for acquiring a file lock.
    """
    if not _is_local_filesystem(path) and not pathlib.epath_installed:
      raise RuntimeError("Please install the `etils[epath]` package to specify a cache directory on a non-local filesystem")

    self.path = self._path = pathlib.Path(path)
    self.path.mkdir(parents=True, exist_ok=True)

    self.eviction_enabled = max_size != -1  # no eviction if `max_size` is set to -1

    if self.eviction_enabled:
      if filelock is None:
        raise RuntimeError("Please install the `filelock` package to set `jax_compilation_cache_max_size`")

      self.max_size = max_size
      self.lock_timeout_secs = lock_timeout_secs

      self.lock_path = self.path / ".lockfile"
      if _is_local_filesystem(path):
        self.lock = filelock.FileLock(self.lock_path)
      else:
        self.lock = filelock.SoftFileLock(self.lock_path)

  def get(self, key: str) -> bytes | None:
    """Retrieves the cached value for the given key.

    Args:
      key: The key for which the cache value is retrieved.

    Returns:
      The cached data as bytes if available; ``None`` otherwise.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

pip install filelock
Or set jax_compilation_cache_max_size to -1 (no eviction, no lock needed)
Or disable the compilation cache entirely

Example fix

# before
jax.config.update('jax_compilation_cache_max_size', 10 * 2**30)

# after
# pip install filelock
jax.config.update('jax_compilation_cache_max_size', 10 * 2**30)

Defensive patterns

Strategy: validation

Validate before calling

import importlib.util
if max_size not in (-1, 0):
    assert importlib.util.find_spec('filelock'), 'pip install filelock'

Try / catch

try:
    jax.config.update('jax_compilation_cache_max_size', size)
except RuntimeError:
    jax.config.update('jax_compilation_cache_max_size', -1)

Prevention

Install filelock alongside any cache-size cap
Default max_size to -1 in configs where the extra may be missing

When it happens

Trigger: Setting jax_compilation_cache_max_size to a positive byte value without filelock installed.

Common situations: Adding a cache size cap to prevent unbounded disk growth in long-running training jobs; slim docker images where jax extras were not installed.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/25548a0d78699dc5. Report an issue: GitHub.