Skip to content

Conversation

@jameszianxuTT
Copy link

@jameszianxuTT jameszianxuTT commented Nov 24, 2025

Ticket: #9712

There is currently no way to clear the XLAGraphExecutor's computation cache.

This change exposes a cache clearing function torch_xla.runtime.clear_computation_cache to clear the computation cache. It also adds some simple tests to check that the in-memory and persistent disk cache are correctly cleared by this function.

* Add torch_xla.runtime.clear_computation_cache() binding
* Run linters
* Add optional workflow dispatch to build wheel from branch instead of torch-xla main
* Change assertion to warning
* Return bool so caller can optionally check if cache was cleared
@jameszianxuTT
Copy link
Author

Hey @ysiraichi, apologies for the ping but would it be possible to either approve this PR to run CI, or otherwise comment on it?

We've been using it in our fork without issue so far for several weeks, and it's an opt-in feature.

@ysiraichi
Copy link
Collaborator

Could you add tests for memory and persistent cache?

@qihqi qihqi requested a review from zhanyong-wan January 12, 2026 04:56
@jameszianxuTT
Copy link
Author

Added some simple tests. Some small instrumentation shows these results, as expected.

 python test/test_compilation_cache_utils.py TestClearCompilationCache.test_clear_persistent_cache
[...]
[ RUN      ] TestClearCompilationCache.test_clear_persistent_cache
2026-01-19 20:51:51.124016: W torch_xla/csrc/runtime/profiler.cpp:88] Profiler API not found for PJRT plugin
graph_cnt: 1
cache_files_before: ['339896306817392351095887498798814495545']

[run xr.clear_computation_cache()]

new_graph_cnt: 0
cache_files_after: []
python test/test_compilation_cache_utils.py TestClearCompilationCache.test_clear_memory_cache
[...]
[ RUN      ] TestClearCompilationCache.test_clear_memory_cache
2026-01-19 20:48:47.183612: W torch_xla/csrc/runtime/profiler.cpp:88] Profiler API not found for PJRT plugin
graph_cnt: 1

[run xr.clear_computation_cache()]

new_graph_cnt: 0
[       OK ] TestClearCompilationCache.test_clear_memory_cache

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants