perf(streaming): async S3 chunk prefetch, cache eviction, obstore downloads by tchaton · Pull Request #852 · Lightning-AI/litData

tchaton · 2026-07-24T16:02:39Z

Summary

Default remote path: process workers + async chunk prefetch (obstore). Opt out with LITDATA_ASYNC_CHUNK_PREFETCH=0.
Async prepare thread: stream to disk (adownload_file), reuse event loop, non-blocking batch drain, floor max_pre_download to 4 for gather width.
Sync S3 downloads also prefer obstore (boto3 fallback). Stream chunk size override: LITDATA_OBSTORE_STREAM_MIN_CHUNK_MIB (default 8).
Cache eviction: size-first deletes, shared .litdata_cache_slots for multi-worker budgets, force-redownload keeps slots; drain all pending deletes on END_TOKEN so tiny caches fully clear.
JPEG deserialize uses ImageReadMode.RGB.
Removed experimental use_threading (real S3 ~2.3× slower than process workers).
Benches under scripts/bench/ (including bench_obstore_*.py, bench_s3_full_epochs.py).

Real-S3 ImageNet (this Studio)

Dataset: lightning_data_search, 48 workers, batch 256, wipe /cache/chunks, PYTHON_GIL=0, transforms RandomResizedCrop(224)+flip+fp16.

Mode	ep0	ep1	mean
process + sync (boto3)	10525	12161	11284
process + async (`max_pre=4`)	10836	12154	11457
process + async (`max_pre=8`)	10554	12266	11345

Warm epoch ~12.2k is decode/transform-bound. Prefer max_pre_download=4.

Cache budget

Warn below 25GB. Slot gate only for realistic budgets (≥10MB + delete-when-processed). Real S3 check at 10GB / 48w: peak ~1.07× limit.

Obstore benches

PYTHON_GIL=0 python -Xgil=0 scripts/bench/bench_obstore_vs_boto3.py
PYTHON_GIL=0 python -Xgil=0 scripts/bench/bench_obstore_imagenet_grid.py --workers 48 --chunk-sizes 2,4,8,16
PYTHON_GIL=0 python -Xgil=0 scripts/bench/bench_s3_full_epochs.py --label run --epochs 2 --workers 48 --batch-size 256 --max-pre-download 4

Test plan

tests/streaming/test_cache_eviction.py
tests/streaming/test_async_prefetch.py
tests/streaming/test_next_streaming_optimizations.py
lock cleanup / download-reader overlap / JPEG writer updates
CI

…loader spike - Use deque for upcoming_indexes (avoid O(n) pop(0) per sample). - Add LITDATA_TIMING counters + prefetch/thread microbenches. - Advisory cache-byte tracking for eviction; unify Parquet/Tokens/PyTree readiness waits; remove dead skip-deletion machinery (eager .cnt wins). - Experimental StreamingDataLoader(use_threading=True) spike inspired by FFCV's thread-based loader (no process IPC); not the default path. Co-authored-by: Cursor <cursoragent@cursor.com>

Fix prefetch matrix for num_workers=0, report item counts in the threaded vs process bench, and add a deque popleft microbench. Co-authored-by: Cursor <cursoragent@cursor.com>

Threaded StreamingDataLoader only helps when the GIL is disabled; reject use_threading=True on regular CPython so the experimental path is not enabled under GIL contention. Co-authored-by: Cursor <cursoragent@cursor.com>

tchaton · 2026-07-24T16:04:26Z

Follow-up: use_threading=True now requires a free-threaded (no-GIL) Python runtime (not sys._is_gil_enabled()). Regular GIL builds raise a clear RuntimeError.

These were incorrectly removed as "dead code" in the next-opt spike. Eager .cnt refcounting does not replace the skip-deletion helpers / ChunksConfig.skip_chunk_indexes_deletion API; put them back. Co-authored-by: Cursor <cursoragent@cursor.com>

Keep existing tests untouched. Add a dedicated threaded-loader suite for gates, coverage, determinism, resume counts, failures, and lifecycle. Fix per-worker prefetch so in-order round-robin cannot deadlock. Co-authored-by: Cursor <cursoragent@cursor.com>

Add a CIFAR-shaped three-way microbench so PR numbers include a non-LitData torch Dataset+DataLoader baseline alongside StreamingDataLoader modes. Co-authored-by: Cursor <cursoragent@cursor.com>

Keep use_threading as the no-GIL in-process loader (no use_asyncio). Add an opt-in LITDATA_ASYNC_CHUNK_PREFETCH path that overlaps remote chunk downloads via asyncio.gather inside PrepareChunksThread only. Co-authored-by: Cursor <cursoragent@cursor.com>

codecov-commenter · 2026-07-24T16:50:13Z

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 75.09804% with 127 lines in your changes missing coverage. Please review.
✅ Project coverage is 81%. Comparing base (ed7409d) to head (c95cc24).
⚠️ Report is 1 commits behind head on main.
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@         Coverage Diff          @@
##           main   #852    +/-   ##
====================================
- Coverage    81%    81%    -0%     
====================================
  Files        54     56     +2     
  Lines      7923   8389   +466     
====================================
+ Hits       6450   6794   +344     
- Misses     1473   1595   +122

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

…tion Make process+async the default remote path (opt out with LITDATA_ASYNC_CHUNK_PREFETCH=0), stream S3 chunks to disk, reuse the prepare-thread event loop, and floor max_pre_download to 4 for gather width. Enforce max_cache_size with shared chunk slots and size-first eviction so multi-worker prefetch cannot ignore the disk budget. Co-authored-by: Cursor <cursoragent@cursor.com>

for more information, see https://pre-commit.ci

Keep all streaming/S3 bench entrypoints in one folder and fix REPO_ROOT path depth after the move. Co-authored-by: Cursor <cursoragent@cursor.com>

Real-S3 ImageNet benches showed threaded workers ~2.3× slower than process workers under decode+transform load. Drop the API, threaded_loader module, and related tests/benches. Co-authored-by: Cursor <cursoragent@cursor.com>

for more information, see https://pre-commit.ci

Use obstore for S3 sync downloads with tunable stream chunk size, drain pending chunk deletes at END_TOKEN so tiny max_cache_size caches clear, and add obstore ImageNet/chunk-size bench helpers. Co-authored-by: Cursor <cursoragent@cursor.com>

for more information, see https://pre-commit.ci

Co-authored-by: Cursor <cursoragent@cursor.com>

for more information, see https://pre-commit.ci

Co-authored-by: Cursor <cursoragent@cursor.com>

…exit Prevents leaked asyncio_N daemon threads from to_thread under async chunk prefetch, which tripped Windows session thread-police teardown. Co-authored-by: Cursor <cursoragent@cursor.com>

…lots Async floor raised max_pre_download to 4 while sub-10MB budgets skipped the cap, so reader eviction tests could briefly hold too many chunks. Co-authored-by: Cursor <cursoragent@cursor.com>

Floor the delete-when-processed prefetch cap at 2 and tear down the asyncio default executor without waiting so prepare-thread exit cannot hang. Co-authored-by: Cursor <cursoragent@cursor.com>

Decompression os.replace can briefly deny open() on Windows even after the file is size-ready; retry instead of failing the read. Co-authored-by: Cursor <cursoragent@cursor.com>

for more information, see https://pre-commit.ci

tchaton and others added 2 commits July 24, 2026 17:01

chore(streaming): harden next-opt microbenches

1ca5835

Fix prefetch matrix for num_workers=0, report item counts in the threaded vs process bench, and add a deque popleft microbench. Co-authored-by: Cursor <cursoragent@cursor.com>

tchaton requested a review from justusschock as a code owner July 24, 2026 16:02

fix(streaming): require no-GIL Python for use_threading

a2b83ac

Threaded StreamingDataLoader only helps when the GIL is disabled; reject use_threading=True on regular CPython so the experimental path is not enabled under GIL contention. Co-authored-by: Cursor <cursoragent@cursor.com>

tchaton and others added 4 commits July 24, 2026 17:10

chore(bench): compare LitData process/threaded vs raw DataLoader

76ab644

Add a CIFAR-shaped three-way microbench so PR numbers include a non-LitData torch Dataset+DataLoader baseline alongside StreamingDataLoader modes. Co-authored-by: Cursor <cursoragent@cursor.com>

tchaton and others added 8 commits July 24, 2026 21:06

[pre-commit.ci] auto fixes from pre-commit.com hooks

62ce3b2

for more information, see https://pre-commit.ci

chore(bench): move microbenches under scripts/bench/

ebd3c4f

Keep all streaming/S3 bench entrypoints in one folder and fix REPO_ROOT path depth after the move. Co-authored-by: Cursor <cursoragent@cursor.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

9fd92d5

for more information, see https://pre-commit.ci

[pre-commit.ci] auto fixes from pre-commit.com hooks

62b321d

for more information, see https://pre-commit.ci

fix(typing): satisfy mypy for obstore ClientConfig and cache cap

b1b2fb5

Co-authored-by: Cursor <cursoragent@cursor.com>

tchaton changed the title ~~perf(streaming): opt-in threaded StreamingDataLoader + next read opts~~ perf(streaming): async S3 chunk prefetch, cache eviction, obstore downloads Jul 24, 2026

tchaton and others added 10 commits July 24, 2026 22:23

chore(bench): ignore ruff docstring/security noise under scripts/bench

cef7d12

Co-authored-by: Cursor <cursoragent@cursor.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

5d485a5

for more information, see https://pre-commit.ci

Merge branch 'main' into perf/next-streaming-optimizations

a00453c

fix(streaming): remove duplicate chunk-ready helpers in item_loader

70d84f2

Co-authored-by: Cursor <cursoragent@cursor.com>

test(streaming): force boto3 path in atomic S3 download unit test

49afe63

Co-authored-by: Cursor <cursoragent@cursor.com>

fix(streaming): shut down asyncio default executor on prepare-thread …

267376b

…exit Prevents leaked asyncio_N daemon threads from to_thread under async chunk prefetch, which tripped Windows session thread-police teardown. Co-authored-by: Cursor <cursoragent@cursor.com>

fix(streaming): cap prefetch under tiny max_cache_size even without s…

0d42d34

…lots Async floor raised max_pre_download to 4 while sub-10MB budgets skipped the cap, so reader eviction tests could briefly hold too many chunks. Co-authored-by: Cursor <cursoragent@cursor.com>

fix(streaming): avoid max_pre=1 deadlock under tiny cache budgets

2df8363

Floor the delete-when-processed prefetch cap at 2 and tear down the asyncio default executor without waiting so prepare-thread exit cannot hang. Co-authored-by: Cursor <cursoragent@cursor.com>

fix(streaming): retry chunk open on Windows PermissionError

6fc241c

Decompression os.replace can briefly deny open() on Windows even after the file is size-ready; retry instead of failing the read. Co-authored-by: Cursor <cursoragent@cursor.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

c95cc24

for more information, see https://pre-commit.ci

tchaton merged commit 3a07199 into main Jul 25, 2026
33 of 35 checks passed

tchaton deleted the perf/next-streaming-optimizations branch July 25, 2026 00:18

tchaton mentioned this pull request Jul 25, 2026

fix(streaming): unblock macOS 3.14 DataLoader restore CI hangs #855

Merged

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

perf(streaming): async S3 chunk prefetch, cache eviction, obstore downloads#852

perf(streaming): async S3 chunk prefetch, cache eviction, obstore downloads#852
tchaton merged 25 commits into
mainfrom
perf/next-streaming-optimizations

tchaton commented Jul 24, 2026 •

edited

Loading

Uh oh!

tchaton commented Jul 24, 2026

Uh oh!

codecov-commenter commented Jul 24, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

tchaton commented Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Real-S3 ImageNet (this Studio)

Cache budget

Obstore benches

Test plan

Uh oh!

tchaton commented Jul 24, 2026

Uh oh!

codecov-commenter commented Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

tchaton commented Jul 24, 2026 •

edited

Loading

codecov-commenter commented Jul 24, 2026 •

edited

Loading