fix(processing): shuffle with a local RNG instead of seeding global generators

src/litdata/processing/data_processor.py

@@ -19,7 +19,6 @@
 import multiprocessing.queues
 import os
 import pickle
-import random
 import shutil
 import signal
 import sys
@@ -40,7 +39,6 @@
 from urllib import parse
 
 import numpy as np
-import torch
 
 from litdata.constants import (
     _DEFAULT_FAST_DEV_RUN_ITEMS,
@@ -711,9 +709,13 @@ def _map_items_to_workers_weighted(
     user_items: list[Any],
     weights: list[int] | None = None,
     file_size: bool = True,
+    rng: np.random.RandomState | None = None,
 ) -> list[list[Any]]:
     """Map the items to the workers based on the weights.
 
+    ``rng`` shuffles each worker's items. Pass a local ``np.random.RandomState`` to keep the global NumPy
+    generator untouched; ``None`` falls back to ``np.random``.
+
     >>> workers_user_items = _map_items_to_workers_weighted(2, list(range(5)), weights=[1, 2, 3, 4, 5])
     >>> assert workers_user_items == [[1, 4, 0], [3, 2]]
     """
@@ -735,7 +737,8 @@ def _map_items_to_workers_weighted(
         else:
             print(f"{label} {worker_id} gets ({len(worker_items[worker_id])}) items for a total weight of {size}.")
 
-    return [np.random.permutation(worker_items[worker_id]).tolist() for worker_id in worker_ids_this_node]
+    shuffler = np.random if rng is None else rng
+    return [shuffler.permutation(worker_items[worker_id]).tolist() for worker_id in worker_ids_this_node]
 
 
 def _get_num_bytes(item: Any, base_path: str) -> int:
@@ -1827,7 +1830,9 @@ def __init__(
             num_uploaders: The number of file uploaders to use.
             delete_cached_files: Whether to delete the cached files.
             fast_dev_run: Whether to run a quick dev run.
-            random_seed: The random seed to be set before shuffling the data.
+            random_seed: Seed for the local generator that shuffles the items assigned to each worker. The global
+                ``random`` / NumPy / torch generators are not touched, so seed those yourself if the recipe's
+                ``prepare_structure`` or ``prepare_item`` needs reproducible randomness.
             reorder_files: By default, reorders the files by file size to distribute work equally among all workers.
                 Set this to ``False`` if the order in which samples are processed should be preserved.
             weights: Provide a list of weights associated to the inputs.
@@ -1951,11 +1956,6 @@ def run(self, data_recipe: DataRecipe) -> None:
         if self.verbose:
             print(f"Setup started with fast_dev_run={self.fast_dev_run}.")
 
-        # Force random seed to be fixed
-        random.seed(self.random_seed)
-        np.random.seed(self.random_seed)
-        torch.manual_seed(self.random_seed)
-
         # Call the setup method of the user
         user_items: list[Any] | StreamingDataLoader | Queue = data_recipe.prepare_structure(
             self.input_dir.path if self.input_dir else None
@@ -1976,22 +1976,24 @@ def run(self, data_recipe: DataRecipe) -> None:
         # Unordered: one bin per node; workers steal from a shared node queue.
         map_workers = 1 if not self.keep_data_ordered else self.num_workers
 
+        # Local generator: shuffle per ``random_seed`` without touching the caller's global RNG state
+        rng = np.random.RandomState(self.random_seed)
         if isinstance(user_items, list):
             assert isinstance(user_items, list)
 
             if self.weights is not None:
                 if len(self.weights) != len(user_items):
                     raise ValueError("The provided weights length should match the inputs' length.")
                 workers_user_items = _map_items_to_workers_weighted(
-                    num_workers=map_workers, user_items=user_items, weights=self.weights, file_size=False
+                    num_workers=map_workers, user_items=user_items, weights=self.weights, file_size=False, rng=rng
                 )
 
             elif self.reorder_files and self.input_dir.path and not _is_studio_fuse_path(self.input_dir.path):
                 # TODO: Only do this on node 0, and broadcast the item sizes to the other nodes.
                 # Skip FUSE mounts: ``stat`` / ``getsize`` on those paths is extremely slow.
                 item_sizes = _get_item_filesizes(user_items, base_path=self.input_dir.path)
                 workers_user_items = _map_items_to_workers_weighted(
-                    num_workers=map_workers, user_items=user_items, weights=item_sizes
+                    num_workers=map_workers, user_items=user_items, weights=item_sizes, rng=rng
                 )
             else:
                 if self.align_chunking and data_recipe.chunk_size is None:

src/litdata/processing/functions.py

@@ -319,6 +319,7 @@ def map(
     storage_options: dict[str, Any] = {},
     keep_data_ordered: bool | None = None,
     broadcast_paths: bool = False,
+    random_seed: int | None = 42,
 ) -> None:
     """Maps a callable over a collection of inputs, possibly in a distributed way.
 
@@ -349,6 +350,8 @@ def map(
             preserved. If True, each worker processes a static slice in order. ``None`` uses False.
         broadcast_paths: Broadcast resolved input/output dirs across multi-node ranks. Defaults to ``False``.
             Auto-enabled when ``input_dir`` or ``output_dir`` contains a ``{%strftime}`` time template.
+        random_seed: Seed for the local generator that shuffles the items assigned to each worker. It does not
+            seed the global ``random`` / NumPy / torch generators, so ``fn`` should seed its own randomness if needed.
     """
     _check_version_and_prompt_upgrade(__version__)
 
@@ -420,6 +423,7 @@ def map(
             storage_options=storage_options,
             keep_data_ordered=keep_data_ordered,
             broadcast_paths=should_broadcast_paths,
+            random_seed=random_seed,
         )
 
         with optimize_dns_context(optimize_dns if optimize_dns is not None else False):
@@ -482,6 +486,7 @@ def optimize(
     verbose: bool = True,
     broadcast_paths: bool = False,
     key_fn: Callable[[Any], Any] | None = None,
+    random_seed: int | None = 42,
 ) -> None:
     """This function converts a dataset into chunks, possibly in a distributed way.
 
@@ -544,6 +549,8 @@ def optimize(
         key_fn: Optional callable ``sample -> key`` (str or int). When set, writes a ``keys/``
             sidecar mapping each key to its sample ``index`` (and chunk location) for
             ``dataset[key]`` / ``dataset_update``.
+        random_seed: Seed for the local generator that shuffles the items assigned to each worker. It does not
+            seed the global ``random`` / NumPy / torch generators, so ``fn`` should seed its own randomness if needed.
     """
     _check_version_and_prompt_upgrade(__version__)
 
@@ -667,6 +674,7 @@ def optimize(
             keep_data_ordered=keep_data_ordered,
             verbose=verbose,
             broadcast_paths=should_broadcast_paths,
+            random_seed=random_seed,
         )
 
         with optimize_dns_context(optimize_dns if optimize_dns is not None else False):

tests/processing/test_data_processor.py

@@ -2696,3 +2696,94 @@ def test_studio_lightning_storage_shared_node_queue(tmpdir):
         assert names == {f"{i:02d}.bin" for i in range(len(sizes))}
     finally:
         shutil.rmtree(root, ignore_errors=True)
+
+
+@pytest.mark.parametrize("seed", [0, 1])
+def test_optimize_preserves_the_callers_random_state(tmpdir, seed):
+    """`optimize` runs in the caller's process, so it must not reseed the caller's generators (#944)."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    expected = (random.random(), np.random.rand(), torch.rand(1).item())  # noqa: S311
+
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+    functions.optimize(
+        fn=_identity_optimize_fn,
+        inputs=list(range(4)),
+        output_dir=os.path.join(tmpdir, f"optimized-{seed}"),
+        chunk_bytes="1MB",
+        num_workers=0,
+    )
+
+    # Without the fix the draw after `optimize` is the one for `random_seed`, not for `seed`.
+    assert (random.random(), np.random.rand(), torch.rand(1).item()) == pytest.approx(expected)  # noqa: S311
+
+
+def test_optimize_does_not_seed_the_global_generators(tmpdir, monkeypatch):
+    """Only the worker shuffle needs randomness, and it uses a local generator (#944).
+
+    Seeding torch globally would also queue a seed for CUDA / MPS / XPU, which torch applies lazily on first use, so
+    ``optimize`` must not call any of the global seeders.
+    """
+    for target in ("random.seed", "numpy.random.seed", "torch.manual_seed"):
+        monkeypatch.setattr(target, mock.MagicMock(side_effect=AssertionError(f"{target} was called")))
+
+    functions.optimize(
+        fn=_identity_optimize_fn,
+        inputs=list(range(4)),
+        output_dir=os.path.join(tmpdir, "optimized"),
+        chunk_bytes="1MB",
+        num_workers=0,
+        random_seed=7,
+    )
+
+
+@pytest.mark.parametrize("use_weights", [True, False], ids=["weights", "file_sizes"])
+def test_weighted_shuffle_uses_local_rng(tmpdir, monkeypatch, use_weights):
+    """Both weighted branches shuffle with a local generator seeded by ``random_seed`` (#944)."""
+    input_dir = os.path.join(tmpdir, "inputs")
+    os.makedirs(input_dir)
+    inputs = []
+    for i in range(8):
+        path = os.path.join(input_dir, f"{i}.bin")
+        with open(path, "wb") as f:
+            f.write(b"x" * (i + 1))
+        inputs.append(path)
+
+    class _StopAfterMapping(Exception):
+        pass
+
+    calls = []
+
+    def spy(**kwargs):
+        calls.append((kwargs.get("rng"), _map_items_to_workers_weighted(**kwargs)))
+        raise _StopAfterMapping
+
+    monkeypatch.setattr(data_processor_module, "_map_items_to_workers_weighted", spy)
+
+    for seed in (0, 1):
+        np.random.seed(seed)
+        expected = np.random.rand()
+        np.random.seed(seed)
+
+        data_processor = DataProcessor(
+            input_dir=input_dir,
+            num_workers=2,
+            keep_data_ordered=True,
+            random_seed=7,
+            weights=list(range(1, 9)) if use_weights else None,
+            verbose=False,
+        )
+        with pytest.raises(_StopAfterMapping):
+            data_processor.run(LambdaMapRecipe(_map_copy_fn, inputs))
+
+        assert np.random.rand() == pytest.approx(expected)
+
+    (rng_0, items_0), (rng_1, items_1) = calls
+    assert isinstance(rng_0, np.random.RandomState)
+    assert isinstance(rng_1, np.random.RandomState)
+    # Same ``random_seed`` -> same shuffle, whatever the caller seeded.
+    assert items_0 == items_1