perf(streaming): Speed up PyTreeLoader reads with mmap and compiled unflatten#849
Merged
Conversation
The PyTreeLoader opened each chunk unbuffered and did two `seek`+`read` syscalls per item (offset pair, then the item). Memory-mapping the chunk and caching its offset table turns per-item reads into memory slices (~+8% single-process read throughput on a local dataset). mmap is applied ONLY to chunks that are safe to map: not encrypted, and NOT shared with another worker. A shared chunk can be deleted/replaced by a co-worker while it is still mapped, which crashes with SIGSEGV rather than a recoverable error (see issues #459 and #756). The set of non-shared chunks is exactly the per-worker shared-chunk analysis already computed for the deletion-race fix; shared and encrypted chunks keep the previous unbuffered read path. The mmap allow-set defaults to empty, so any path that does not opt in (e.g. random __getitem__ access) uses normal reads. The map is opened/closed per chunk transition and released on teardown and when pickled (spawn workers), so no map outlives the chunk file. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
tchaton
force-pushed
the
perf/mmap-nonshared
branch
from
July 24, 2026 12:41
3e188dc to
e74cc74
Compare
justusschock
approved these changes
Jul 24, 2026
|
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #849 +/- ##
====================================
Coverage 81% 81%
====================================
Files 54 54
Lines 7736 7842 +106
====================================
+ Hits 6268 6372 +104
- Misses 1468 1470 +2 🚀 New features to boost your workflow:
|
Cut PyTreeLoader hot-path overhead by compiling a specialized unflatten once per dataset, tightening non-shared mmap item reads, and covering the new paths with targeted tests. Co-authored-by: Cursor <cursoragent@cursor.com>
tchaton
changed the title
Replace nested closures with module-level callables so item loaders pickle cleanly under DataLoader spawn workers. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
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
3 participants
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.
Before submitting
What does this PR do?
Speeds up the
StreamingDataset/PyTreeLoaderread hot path for local cached chunks.structunpacking for size headers and int/float leaves.Local microbench on 80k dict-of-leaves items (warm sequential reads): about 2.9× vs the previous stock file path (~224k → ~644k items/s).
PR review
Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in GitHub issues there's a high chance it will not be merged.
Did you have fun?
Make sure you had fun coding 🙃