[wolf-06]

Before submitting

• Was this discussed/agreed via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

Fixes #790.

This PR adds an optional global cache budget to CombinedStreamingDataset so that a single total disk budget can be enforced across all child StreamingDatasets, instead of each dataset independently growing its cache (which easily leads to runaway disk usage when composing large mixtures, e.g. 50+ datasets).

New API (purely additive, fully backward compatible)

CombinedStreamingDataset(
    datasets=[...],
    max_cache_size="500GB",            # total budget: int bytes / "500GB" / 0.90 of free disk (default None = unchanged behavior)
    cache_allocation="proportional",   # "proportional" (default) or "equal"
)

cache_allocation	Behavior
"proportional" (default)	Each child gets a share of the budget matching its sampling weight (more-sampled datasets get more cache headroom).
"equal"	The budget is split evenly across children.

Behavior details

• max_cache_size=None (default) is a strict no-op: children keep their own max_cache_size (existing behavior untouched).
• When a combined budget is set, it overrides per-child max_cache_size values.
• Budgets are resolved once via the existing _resolve_max_cache_size machinery (int bytes, human-readable sizes like "100GB", or free-disk fractions like 0.90) and applied to each child's max_cache_size in __iter__ before its cache/eviction thread is created — so the proven per-directory eviction in PrepareChunksThread enforces the budget unchanged (no changes to cache.py / reader.py / dataset.py).
• Because each DataLoader worker re-runs __iter__ on its dataset copy, budgets are correctly applied inside every worker process (covered by a dedicated test that records the budget each worker copy sees).
• Safety guard: zero-weight datasets and non-StreamingDataset children are skipped, so a child's budget is never set to a falsy 0 (which would silently disable eviction in PrepareChunksThread).
• Warnings: a per-dataset budget below 25GB emits the same guidance as StreamingDataset; if the MAX_CACHE_SIZE environment variable is set alongside a combined budget, a warning documents that the env var takes precedence.

Example

combined = CombinedStreamingDataset(
    datasets=train_datasets,      # e.g. SlimPajama + StarCoder mixture
    seed=42,
    weights=weights,
    iterate_over_all=False,
    max_cache_size="500GB",       # total across ALL child datasets
)

With datasets of 8 and 24 items (no explicit weights), a "10GB" budget distributes as 2.5GB / 7.5GB.

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.

---

Testing

Followed the TDD workflow from CONTRIBUTING.md (failing tests first, then implementation).

• New tests (16) in tests/streaming/test_combined.py: no-op default, equal/proportional allocation (explicit weights + length-based default weights), child-budget override, invalid cache_allocation raises ValueError, end-to-end enforcement (cache._reader._max_cache_size), string/fraction parsing, per-DataLoader-worker enforcement (records budgets from worker copies, spawn-safe), MAX_CACHE_SIZE env-var warning, epoch re-allocation stability, small-budget warning, and helper-level stub/None tests.
• Results:
  • pytest tests/streaming/test_combined.py tests/streaming/test_cache_eviction.py tests/streaming/test_dataloader.py → 88 passed, 2 skipped
  • Full CI-equivalent run (pytest tests --ignore=tests/processing --ignore=tests/raw -n 2 --dist=loadgroup --timeout=120) → 741 passed, 4 skipped, 3 failed; all 3 failures are pre-existing parallel-load flakes in unrelated distributed tests (test_cache_for_image_dataset_distributed[2], test_dataset_reshuffling_every_epoch, test_parallel_dataset_partial_iteration_resume[False-False-None]) — all pass when run serially.
  • ruff check / ruff format --check / mypy → clean.
  • Verified worker budgets under the spawn start method (macOS/Windows CI path).

Documentation

• README.md: new "Global cache budget" subsection in the Combined streaming section (example + allocation-mode table + caveats).
• src/litdata/CHANGELOG.md: [unreleased] ### Added entry.

[codecov-commenter]

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

Codecov Report

❌ Patch coverage is 97.29730% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82%. Comparing base (c6f075c) to head (ae0b715).
⚠️ Report is 2 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    #940   +/-   ##
=====================================
  Coverage     82%     82%           
=====================================
  Files         72      72           
  Lines      16445   16519   +74     
=====================================
+ Hits       13505   13580   +75     
+ Misses      2940    2939    -1     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[tchaton]

Found two P2 correctness issues at 40c2b03:

• combined.py:62–63: Fractional budgets probe the working directory when cache_dir is omitted, although the actual cache can reside elsewhere. Mocking 100GB free on the working filesystem and 10GB on the cache filesystem allocated 90GB instead of 9GB. Resolve the actual cache location before probing free space.

• combined.py:77–83: Positive weights that round to zero retain the child’s previous budget, bypassing the combined limit. Reproduced with a 1,000-byte total and weights [0.0001, 0.9999]: reader budgets became 45GB and 999 bytes. Distinguish zero weights from rounded-zero allocations; reject insufficient budgets or allocate a minimum within the total.

All 16 focused allocation/budget tests passed; neither case is covered.

_{— GPT-6, on direct instruction; reviewed before posting}

[wolf-06]

Hi @bhimrazy @tchaton @dhedey @pwgardipee @deependujha

I have made the changes as per the suggestions.
Could PTAL ?
thank you