{{ message }}
Give each xdist worker its own torch.compile cache directory - #9134
Merged
Conversation
Inductor's on-disk caches default to one directory per USER, not per process, so the four workers under pytest -n 4 share /tmp/torchinductor_<user> and its fxgraph, aotautograd and Triton subtrees. The upstream recipe is explicit that a common TORCHINDUCTOR_CACHE_DIR is what makes processes SHARE compiled artifacts, so a different value per worker is how they are kept apart. Sharing is not obviously wrong, since the entries are content-addressed. It is still a write-write interaction between processes that nothing here controls, and a cache is exactly the sort of thing that turns a deterministic suite into an intermittent one. This removes the interaction for the price of some recompilation. What it buys, measured rather than assumed, so nobody has to guess later: a full unsloth_zoo run against an empty dedicated cache directory took 578s and left ZERO entries in it. That suite never populates the on-disk cache, so for that step this is neither a saving nor a cost. It is insurance for the suites that do compile, applied wherever tests run in parallel rather than only where a problem has already been seen. It has to run before torch is imported, so it is a module-level import in the two conftests rather than a fixture. TRITON_CACHE_DIR is set explicitly rather than left to follow, because it only derives from the inductor directory when unset and an environment exporting it would keep all four workers on one Triton cache. An explicit TORCHINDUCTOR_CACHE_DIR is split underneath rather than replaced, so a CI path chosen on purpose is respected. A single-process run is left alone entirely. Verified under real xdist, not only in unit tests: four workers report four distinct directories, and the parallel suites still pass.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.

Inductor's on-disk caches default to one directory per user, not per process, so the four workers under
pytest -n 4share/tmp/torchinductor_<user>and itsfxgraph,aotautogradand Triton subtrees. The upstream recipe is explicit that a commonTORCHINDUCTOR_CACHE_DIRis what makes processes share compiled artifacts, so a different value per worker is how they are kept apart.Sharing is not obviously wrong — the entries are content-addressed. It is still a write-write interaction between processes that nothing here controls, and a cache is exactly the sort of thing that turns a deterministic suite into an intermittent one. This removes the interaction for the price of some recompilation.
What it buys, measured rather than assumed
A full
unsloth_zoorun against an empty, dedicated cache directory:That suite never populates the on-disk cache at all, so for the largest step in the Core legs this is neither a saving nor a cost — and the 75s cold/warm gap is page cache, not compilation, since there is nothing in the directory either time. Worth recording, because it also rules out persisting
TORCHINDUCTOR_CACHE_DIRacross CI runs for that suite: there is nothing to persist.So this is insurance for the suites that do compile, applied wherever tests run in parallel rather than only where a problem has already been seen.
Details that matter
TRITON_CACHE_DIRis set explicitly rather than left to follow: it only derives from the inductor directory when unset, and an environment that exports it would keep all four workers on one Triton cache.TORCHINDUCTOR_CACHE_DIRis split underneath rather than replaced, so a CI path chosen on purpose is respected.Verification
Under real xdist, not only unit tests: four workers report four distinct directories, and the parallel suites still pass. The guard also asserts that both conftests still reach the helper, because a helper nobody imports is the silent failure here.