Skip to content

Add a persistent cache#5800

Merged
jonb377 merged 4 commits intomasterfrom
jonbolin/persistent-cache
Dec 5, 2023
Merged

Add a persistent cache#5800
jonb377 merged 4 commits intomasterfrom
jonbolin/persistent-cache

Conversation

@jonb377
Copy link
Copy Markdown
Collaborator

@jonb377 jonb377 commented Nov 15, 2023

This changes adds a new PersistentCache instance which will read and write to disk to enable persistent compilation caching. The PersistentCache is currently not used anywhere, but in a later change, the PersistentCache will replace the Cache in the XlaGraphExecutor when persistent caching is enabled.

@jonb377 jonb377 self-assigned this Nov 15, 2023
// will go through the wrapped Cache.
template <typename K, typename T, typename H = std::hash<K>,
typename E = std::equal_to<K>>
class PersistentCache : public AbstractCache<K, T, H, E> {
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we specialize the PersistentCache a bit more? Do we really need to carry that many template parameters?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't see a use case for the H and E templates, so those can be dropped and the default value directly applied the parent class. Keeping the same interface as Cache lets this be a drop-in replacement though, maybe we can follow up to remove the unnecessary templates from both?

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let me take a closer look (probably after my break...). I need to play around with the patch to awake my C++ muscles...

Copy link
Copy Markdown
Collaborator

@will-cromar will-cromar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

Comment thread torch_xla/csrc/runtime/cache.h Outdated
Comment thread torch_xla/csrc/runtime/cache.h Outdated
std::lock_guard<std::mutex> slock(lock_);
memory_cache_.Clear();
// Delete and recreate the cache directory on disk.
if (!readonly_) {
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe instead of conditioning on readonly_ here, introduce an optional arg for Clear to clear the persistent cache, if cache_dir is provided.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should respect the readonly_ setting here - the reason for it is in shared mounts, we only want the writer to modify the objects on disk.

Copy link
Copy Markdown
Contributor

@yeounoh yeounoh left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks @jonb377 left minor comments

Copy link
Copy Markdown
Collaborator

@alanwaketan alanwaketan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A few nit, but LGTM in general.

Comment thread torch_xla/csrc/runtime/cache.h Outdated
Comment thread torch_xla/csrc/runtime/cache.h Outdated
@jonb377 jonb377 merged commit f06f6c0 into master Dec 5, 2023
@jonb377 jonb377 deleted the jonbolin/persistent-cache branch December 5, 2023 22:42
chunnienc pushed a commit to chunnienc/xla that referenced this pull request Dec 14, 2023
* Add a persistent cache

* Add locking

* Make serialize functions accept const parameters

* Clarify readonly parameter
golechwierowicz pushed a commit that referenced this pull request Jan 12, 2024
* Add a persistent cache

* Add locking

* Make serialize functions accept const parameters

* Clarify readonly parameter
bhavya01 pushed a commit that referenced this pull request Apr 22, 2024
* Add a persistent cache

* Add locking

* Make serialize functions accept const parameters

* Clarify readonly parameter
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants