Using GFile for persistent compilation caching in JAX

[sshahrokhi]

The filesystem method that was used for persistent compilation caching, was chosen because it was able to implement Least Recently Used, but it didn't support gcs. GFile supports gcs which can be used between machines but unfortunately it is not following the LRU policy. Therefore, there is no maximum memory we allow for the cache as well.

[TESTING] I tested this method by using this test, tested for both local file and a gcs bucket to make sure both works fine.

[skye]

A few high-level comments:

• Please update the PR title and description to be more descriptive (see https://cbea.ms/git-commit/ for a long description of how to write a good commit message). This can also be the commit message once you squash everything before commit.
• Also describe how you tested this change in the PR description, since we don't have automatic Cloud TPU presubmits :(
• Delete the old implementation
• Add this as a breaking change to the changelog: https://github.com/google/jax/blob/main/CHANGELOG.md

[skye]

Also look at the lint failures please

[sshahrokhi]

Thanks for the reviews! I am working on them! Meanwhile:

• Add this as a breaking change to the changelog: https://github.com/google/jax/blob/main/CHANGELOG.md

I didn't remove the max_cache_size_bytes so that this change would not be breaking(only the logic behind the file system is changing)! So do you suggest removing it and also add this to changelog?

[skye]

I didn't remove the max_cache_size_bytes so that this change would not be breaking(only the logic behind the file system is changing)! So do you suggest removing it and also add this to changelog?

Yes. This is a breaking change, but I think it's better to raise an error than to silently stop working.

[sshahrokhi]

Thanks for all the comments, could you please take a look now? @skye

[skye]

Almost there!!

[skye]

Looks great! Can you squash all your changes to a single commit and then I'll get this in?

[skye]

Rename to GFileCacheTest

[skye]

Please squash your changes again too, and then we can get this in 🤞

[sshahrokhi]

done! thanks!