fix(processing): leave the caller's global RNG state alone after optimize

[Yi-111-a]

Summary

optimize() runs in the caller's process, but DataProcessor.run seeded random, numpy and torch at the top of run and never restored them. The seed the caller set before calling optimize was silently replaced, so everything random afterwards (model initialization, shuffling, sampling) was identical whatever seed the caller chose — and differed depending on whether an optimize call had happened earlier in the run, for example only when a cache did not exist yet.

The states are now captured before seeding and restored in a finally, so the fixed seed covers the processing work only and the caller's own generators are handed back untouched.

Changes

• src/litdata/processing/data_processor.py: run saves random.getstate(), np.random.get_state(), torch.random.get_rng_state() (plus the CUDA states when CUDA is already initialized) before seeding, then restores them in a finally. The existing body moves verbatim into a new _run, which keeps the restore on every exit path, including an exception raised out of the recipe.

The diff is deliberately indentation-only for the moved body.

Test

tests/processing/test_data_processor.py::test_optimize_preserves_the_callers_random_state seeds the caller with 0 and with 1, runs optimize, and asserts the next draw from each of the three generators matches the draw the caller would have made without the call.

It reproduces the reported behaviour on the unpatched tree:

FAILED seed=0 - assert (0.6394267984...2692632675171) == approx((0.844...37 ± 5.0e-07))
FAILED seed=1 - assert (0.6394267984...2692632675171) == approx((0.134...67 ± 7.6e-07))

which is exactly the random=0.639427 reported in the issue for both seeds. With the fix both parameters pass.

Note on approach

The issue suggests private generators (random.Random, np.random.default_rng) or save/restore. Save/restore is the smaller change here: the existing seeding deliberately covers the whole of run including prepare_structure and the worker setup, and several call sites rely on optimize being reproducible. Private generators would have to be threaded through all of that. Restoring keeps optimize exactly as deterministic as before and only narrows its blast radius.

Fixes #944

[codecov-commenter]

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

Codecov Report

❌ Patch coverage is 93.75000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 82%. Comparing base (c6f075c) to head (761c028).
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@          Coverage Diff          @@
##            main    #945   +/-   ##
=====================================
  Coverage     82%     82%           
=====================================
  Files         72      72           
  Lines      16445   16457   +12     
=====================================
+ Hits       13505   13517   +12     
  Misses      2940    2940           

🚀 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.