diff --git a/src/litdata/processing/data_processor.py b/src/litdata/processing/data_processor.py index 9a11881b8..be037bdca 100644 --- a/src/litdata/processing/data_processor.py +++ b/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,6 +1976,8 @@ 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) @@ -1983,7 +1985,7 @@ def run(self, data_recipe: DataRecipe) -> 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): @@ -1991,7 +1993,7 @@ def run(self, data_recipe: DataRecipe) -> None: # 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: diff --git a/src/litdata/processing/functions.py b/src/litdata/processing/functions.py index 7e0658e86..c7d592e84 100644 --- a/src/litdata/processing/functions.py +++ b/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): diff --git a/tests/processing/test_data_processor.py b/tests/processing/test_data_processor.py index 61b46c1a4..335ee9f27 100644 --- a/tests/processing/test_data_processor.py +++ b/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.""" + 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. + + 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``.""" + 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