Implement LRU cache eviction for persistent compilation cache

[ayaka14732]

This PR is part of the implementation of LRU cache eviction using the mtime attribute provided by the filesystem. The current PR does not support GCS, but this problem will be solved in a subsequent PR.

More details in the design doc: https://docs.google.com/document/d/111YibwGXOFb_hMm-lua1u63QooAzIBEH-xfRPGmibis/edit?usp=sharing

[nouiz]

Can you link to the design doc?
Also, would be good to have it documented somewhere?
Like in the https://jax.readthedocs.io/en/latest/persistent_compilation_cache.html file?

Note, I saw some NFS server being configured to not update mtime to speed up the server. Maybe document that this can happen and in that case, this will revert to creation time?
The first time I saw the behavior without knowing the reason, it took times to understand what was going on.

[hawkinsp]

Note, I saw some NFS server being configured to not update mtime to speed up the server. Maybe document that this can happen and in that case, this will revert to creation time? The first time I saw the behavior without knowing the reason, it took times to understand what was going on.

Yes, I suspect there's a chance you might see stale mtime values if you stick the cache on NFS and you're accessing it concurrently from multiple clients (see lookupcache in the NFS docs). I'm not sure there's a lot we can do about that, though.

[skye]

Overall nice work!

[skye]

This is a first cut at the LRU eviction implementation, so it isn't expected to work well with network file systems yet (notably GCS, which many Cloud TPU users use for their cache storage). We'll iterate from here. I don't think we should publicly document this until it works well across filesystems, but absolutely agree this should eventually be in https://jax.readthedocs.io/en/latest/persistent_compilation_cache.html.

[superbobry]

LGTM overall, but please address the comment in tests.

[ayaka14732]

Test fails because the test utilises filelock, which is not installed.

=================================== FAILURES ===================================
_______________________ LRUCacheTest.test_cache_eviction _______________________
[gw5] linux -- Python 3.12.3 /opt/hostedtoolcache/Python/3.12.3/x64/bin/python
tests/lru_cache_test.py:39: in test_cache_eviction
    cache = Impl(path, max_size=884700)
jax/_src/lru_cache.py:57: in __init__
    raise RuntimeError("Please install filelock package to set `jax_compilation_cache_max_size`")
E   RuntimeError: Please install filelock package to set `jax_compilation_cache_max_size`

[superbobry]

Add filelock to build/test-requirements.txt and to the deps in tests/BUILD.

Or skip the test for now if filelock is not importable.

[ayaka14732]

Just realised that JAX had a FileSystemCache that supports LRU cache eviction introduced in #6869, but was subsequently removed in #10771 to support GCS. This is exactly one of the challenges that I faced in this PR. Fortunately, I've devised potential solutions to simultaneously support LRU cache eviction and GCS compatibility. This is going to be completely solved in a subsequent PR.

[ayaka14732]

All comments resolved

[superbobry]

Add that to lru_cache deps?

You will want to use the py_deps macros for this: py_deps("filelock"). We have a list of deps in jax.bzl.

[ayaka14732]

Where can I find the lru_cache deps?

[gnecula]

Thank you for preparing this. Please squash the long chain of commits, or at least most of them.

[superbobry]

Please squash the commits.

[ayaka14732]

Can you link to the design doc?

@nouiz I've just added the link to the first comment.