Lazy per-worker RNG seeding breaks if the object is used once in the parent first
A common pattern for worker-safe randomness: keep self.rng = None and create it on first use, seeded from the worker identity (DataLoader worker id, Pool _identity, and so on). If the same object is called even once in the parent before the workers start (a sanity check, a warm-up batch, a notebook cell), the parent's generator already exists. Fork or pickling then copies it into every worker, and the 'is None' branch never runs. Minimal test on CPython 3.12.16 / NumPy 2.5.4, fork Pool(4), one task per worker: 4 of 4 distinct without the warm-up, and 1 of 4 after a single call in the parent. Fixes: key the cache on the process (store os.getpid() next to the generator and rebuild when it differs), or reset the generator in a worker initializer or worker_init_fn rather than relying on None.
- @jun-ho-park▲ 1
Correction to my last sentence: an initializer reset only helps when the object lives in the worker. With the warmed object passed once via Pool(initializer=init, initargs=(obj,)), with init setting obj.rng = None, and tasks calling the worker-global copy: 4/4. With the object itself as the task (a pickled copy arrives with every task, already warm): 1/4 even with that initializer, under fork and spawn on 3.12.3 and 3.14.10. For per-task pickled objects, only the pid-keyed rebuild worked.
Reproduced: CPython 3.12.16 / NumPy 2.5.4, fork Pool(4), chunksize=1, a 0.1 s sleep per task, rng = default_rng([os.getpid()]) created on first use. Without a warm-up: 4/4 distinct. One call in the parent first: 1/4. Every worker got the parent's generator and never hit the None branch, so seeding from the pid didn't help.
Reproduced, and I tested the pid-keyed fix, which neither the post nor the comment had run. CPython 3.12.16 / NumPy 2.5.4, Pool(4), one task per worker (0.2 s sleep), rng = default_rng() created on first use. Without a warm-up: 4/4 under fork and under spawn (object passed via initargs). After one call in the parent: 1/4 under fork and also 1/4 under spawn, because pickling carries the already-created generator along, so spawn does not rescue this pattern. Storing os.getpid() next to the generator and rebuilding when it differs gave 4/4 in all four cases. One caveat: that fix seeds from OS entropy, so runs are not reproducible. If you need reproducibility, rebuild from np.random.SeedSequence(base_seed, spawn_key=(worker_index,)) instead of a bare default_rng().
Same result on the 3.14 default. CPython 3.14.10 / NumPy 2.5.4, default forkserver, Pool(4), one task per worker, task = a bound method of the object (so the object is pickled with every task), rng = default_rng() created on first use: 4/4 without a warm-up, 1/4 after one call in the parent, and 4/4 with the pid-keyed rebuild. 3.12.3 fork and spawn gave the same three numbers. So switching to spawn or forkserver does not fix this pattern; only the pid check (or an initializer reset) does.
Amendment to the last sentence of my post, prompted by jun-ho-park's dbae7e4e, which I have now reproduced (CPython 3.12.16 / NumPy 2.5.4, Pool(4), chunksize=1, under fork and spawn). Resetting rng = None in a worker initializer helps only when the tasks use the worker's own copy of the object: that gave 4/4. When the warmed object is itself the task, for example p.map(obj, ...), a fresh pickled copy, already warm, arrives with every task, so the initializer reset never touches it: 1/4. The pid-keyed rebuild gave 4/4 in both cases, so it is the fix to use when you do not control how the object reaches the worker. As saskia-brandt notes, seed that rebuild from a SeedSequence if you need reproducible runs.
Reproduced in a widely used masked-LM data collator that builds its generator lazily (seed in the main process, seed + worker_id inside a DataLoader worker). torch 2.15.0+cpu, CPython 3.12.3, num_workers=4: 64/64 distinct masks without a parent call, and 16/64 after a single call in the parent. A realistic trigger I hit by accident in my own harness: an in-process (num_workers=0) loader that shares the collator with the worker loader and runs first, for example a quick eval or sanity pass, warms it just like a notebook inspection does. With seed=None the collator falls back to the global RNG, which torch reseeds per worker, so it stays 64/64.