fix: reseed shared sampler RNGs in forked children

architecture/sampling.md

@@ -69,6 +69,7 @@ The engine-side generator for sampler columns. Extends `FromScratchColumnGenerat
 - **Managed datasets for person data** provide realistic, locale-consistent person records that Faker alone cannot guarantee (e.g., matching name ethnicity to locale, consistent address formatting).
 - **Separate DAG from the main execution DAG** — sampler columns have their own dependency graph within the `DatasetGenerator`, independent of the broader column execution DAG in `DatasetBuilder`. This is because sampler columns are generated as a batch before LLM columns.
 - **Discriminated union for sampler params** mirrors the column config pattern — each sampler type has its own params class with a `Literal` discriminator, enabling type-safe deserialization and validation.
+- **Shared random state** — samplers draw from numpy's global `RandomState` and Faker's shared `Random` unless given a `random_state`, so seeding them in the process (`np.random.seed`, `Faker.seed`) makes sampler output repeatable. An `os.register_at_fork` hook in `sampling_gen/utils.py` reseeds both from OS entropy in a forked child, as CPython does for the stdlib `random` module, so children forked after a run do not generate identical data; a child that needs repeatable output seeds itself. The hook also holds numpy's global RNG lock across the fork, so a child forked while another thread is drawing does not inherit that lock held by a thread it does not have.
 
 ## Cross-References
 

packages/data-designer-engine/src/data_designer/engine/sampling_gen/utils.py

@@ -4,6 +4,9 @@
 from __future__ import annotations
 
 import numbers
+import os
+import sys
+import threading
 
 import data_designer.lazy_heavy_imports as lazy
 
@@ -40,3 +43,57 @@ def check_random_state(seed):
     if isinstance(seed, lazy.np.random.RandomState):
         return seed
     raise ValueError("%r cannot be used to seed a numpy.random.RandomState instance" % seed)
+
+
+_held_numpy_rng_lock: threading.RLock | None = None
+
+
+def _numpy_global_rng_lock() -> threading.RLock | None:
+    # Every draw from numpy's global RandomState, and seed(), holds its bit generator's RLock.
+    # Only look it up when numpy.random is already imported, to keep the fork hooks from
+    # importing it.
+    numpy_random = sys.modules.get("numpy.random")
+    if numpy_random is None:
+        return None
+    return numpy_random.mtrand._rand._bit_generator.lock
+
+
+def _acquire_numpy_rng_lock_before_fork() -> None:
+    # Hold the lock across fork so a child never inherits it locked by a thread that does not
+    # exist there, which would block the child's reseed (and any later draw) forever. This waits
+    # for a draw in progress on another thread to finish.
+    global _held_numpy_rng_lock
+    lock = _numpy_global_rng_lock()
+    if lock is not None:
+        lock.acquire()
+    _held_numpy_rng_lock = lock
+
+
+def _release_numpy_rng_lock_after_fork() -> None:
+    global _held_numpy_rng_lock
+    lock, _held_numpy_rng_lock = _held_numpy_rng_lock, None
+    if lock is not None:
+        lock.release()
+
+
+def _reseed_shared_rngs_after_fork() -> None:
+    # Samplers draw from numpy's global RandomState (check_random_state(None)) and Faker's
+    # shared Random. CPython reseeds the stdlib random module in a forked child but not these,
+    # so children forked after a run would all draw the same samples. Reseed them from OS
+    # entropy, as CPython does for random; a child that needs reproducible draws seeds itself.
+    # Only reseed modules already imported, to keep the hook from importing them.
+    _release_numpy_rng_lock_after_fork()
+    numpy_random = sys.modules.get("numpy.random")
+    if numpy_random is not None:
+        numpy_random.seed()
+    faker_generator = sys.modules.get("faker.generator")
+    if faker_generator is not None:
+        faker_generator.random.seed()
+
+
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(
+        before=_acquire_numpy_rng_lock_before_fork,
+        after_in_parent=_release_numpy_rng_lock_after_fork,
+        after_in_child=_reseed_shared_rngs_after_fork,
+    )

packages/data-designer-engine/tests/engine/sampling_gen/test_utils.py

@@ -3,10 +3,14 @@
 
 from __future__ import annotations
 
+import threading
+from collections.abc import Callable
+
 import pytest
 
 import data_designer.lazy_heavy_imports as lazy
-from data_designer.engine.sampling_gen.utils import check_random_state
+from data_designer.engine.sampling_gen.utils import _numpy_global_rng_lock, check_random_state
+from tests.engine.fork_probe import requires_fork, run_fork_probe
 
 
 @pytest.mark.parametrize(
@@ -36,3 +40,125 @@ def test_check_random_state_scenarios(test_case, input_value, expected_type, exp
 def test_check_random_state_invalid():
     with pytest.raises(ValueError, match="'invalid' cannot be used to seed a numpy.random.RandomState instance"):
         check_random_state("invalid")
+
+
+@requires_fork
+def test_forked_children_draw_different_samples() -> None:
+    run_fork_probe(
+        """
+        import faker
+
+        from data_designer.engine.sampling_gen.utils import check_random_state
+
+        rng = check_random_state(None)
+        fake = faker.Faker("en_US")
+        rng.random_sample()  # a parent run has drawn from both, as sampler runs do
+        fake.name()
+
+
+        def draw(queue):
+            queue.put((float(rng.random_sample()), [fake.name() for _ in range(5)]))
+
+
+        first, second = run_forked_child(draw), run_forked_child(draw)
+        assert first[0] != second[0], "forked children drew the same numpy samples"
+        assert first[1] != second[1], "forked children drew the same Faker samples"
+        """
+    )
+
+
+@requires_fork
+def test_forked_child_reseed_does_not_change_parent_draws() -> None:
+    run_fork_probe(
+        """
+        from data_designer.engine.sampling_gen.utils import check_random_state
+
+        rng = check_random_state(None)
+        rng.seed(7)
+        expected = rng.random_sample(3).tolist()
+        rng.seed(7)
+        run_forked_child(lambda queue: queue.put(None))
+        assert rng.random_sample(3).tolist() == expected, "forking reseeded the parent's RNG"
+        """
+    )
+
+
+def _call_waits_for_lock(lock: threading.RLock, call: Callable[[], object]) -> bool:
+    """Return whether ``call`` on another thread blocks while this thread holds ``lock``."""
+    started, finished = threading.Event(), threading.Event()
+
+    def run() -> None:
+        started.set()
+        call()
+        finished.set()
+
+    with lock:
+        thread = threading.Thread(target=run)
+        thread.start()
+        started.wait()
+        waited = not finished.wait(0.5)
+    thread.join(timeout=5)
+    assert finished.is_set(), "call did not finish after the lock was released"
+    return waited
+
+
+@pytest.mark.parametrize(
+    "call",
+    [
+        pytest.param(lambda: check_random_state(None).random_sample(), id="sampler-draw"),
+        pytest.param(lambda: lazy.np.random.seed(), id="child-reseed"),
+    ],
+)
+def test_fork_hook_lock_is_the_lock_global_random_state_takes(call: Callable[[], object]) -> None:
+    # The fork hooks take numpy's lock through a private path (mtrand._rand._bit_generator.lock),
+    # and the fork test holds the lock this helper returns, so it cannot tell whether numpy uses
+    # it. Check against numpy itself: sampler draws and the child's reseed must wait for it.
+    check_random_state(None)  # imports numpy.random
+    lock = _numpy_global_rng_lock()
+    assert lock is not None
+    assert _call_waits_for_lock(lock, call), "numpy's global RandomState does not take the lock the fork hooks hold"
+
+
+def test_call_waits_for_lock_detects_a_call_that_does_not_take_the_lock() -> None:
+    # Control for the test above: a separate RandomState has its own lock and must not wait.
+    check_random_state(None)
+    lock = _numpy_global_rng_lock()
+    assert lock is not None
+    assert not _call_waits_for_lock(lock, lambda: check_random_state(1).random_sample())
+
+
+@requires_fork
+def test_fork_while_another_thread_draws_does_not_hang_child() -> None:
+    run_fork_probe(
+        """
+        import threading
+        import time
+
+        from data_designer.engine.sampling_gen.utils import _numpy_global_rng_lock, check_random_state
+
+        rng = check_random_state(None)
+        rng.random_sample()
+        lock = _numpy_global_rng_lock()
+        held = threading.Event()
+
+
+        def draw_in_progress():
+            # Holds the lock as a draw on another thread does when fork is called.
+            with lock:
+                held.set()
+                time.sleep(1.0)
+
+
+        threading.Thread(target=draw_in_progress).start()
+        held.wait()
+        started = time.monotonic()
+        value = run_forked_child(lambda queue: queue.put(float(rng.random_sample())))
+        assert isinstance(value, float)
+        assert time.monotonic() - started >= 0.9, "fork did not wait for the draw in progress"
+
+        drawer = threading.Thread(target=rng.random_sample)
+        drawer.start()
+        drawer.join(timeout=5)
+        assert not drawer.is_alive(), "the parent's RNG lock stayed held after fork"
+        """
+    )

packages/data-designer/tests/interface/test_fork_safety.py

@@ -54,3 +54,67 @@ def test_preview_in_child_forked_after_parent_preview() -> None:
         env={**os.environ, "NEMO_TELEMETRY_ENABLED": "false"},
     )
     assert result.returncode == 0, f"preview() in a forked child failed:\n{result.stderr[-4000:]}"
+
+
+_PREVIEW_IN_SIBLING_FORKED_CHILDREN = """
+import json
+import multiprocessing
+import sys
+
+import data_designer.config as dd
+from data_designer.interface import DataDesigner
+
+builder = dd.DataDesignerConfigBuilder()
+builder.add_column(
+    dd.SamplerColumnConfig(
+        name="category",
+        sampler_type=dd.SamplerType.CATEGORY,
+        params=dd.CategorySamplerParams(values=list("abcdefgh")),
+    )
+)
+builder.add_column(
+    dd.SamplerColumnConfig(
+        name="uniform",
+        sampler_type=dd.SamplerType.UNIFORM,
+        params=dd.UniformSamplerParams(low=0, high=1000),
+    )
+)
+
+
+def preview(queue=None):
+    dataset = DataDesigner().preview(builder, num_records=6).dataset
+    if queue is not None:
+        queue.put(json.dumps(dataset.to_dict(orient="list")))
+
+
+def preview_in_child():
+    ctx = multiprocessing.get_context("fork")
+    queue = ctx.Queue()
+    child = ctx.Process(target=preview, args=(queue,))
+    child.start()
+    child.join(timeout=30)
+    if child.is_alive():
+        child.kill()
+        child.join()
+        sys.exit("forked child hung in preview()")
+    if child.exitcode != 0:
+        sys.exit(f"forked child exited with code {child.exitcode}")
+    return queue.get(timeout=5)
+
+
+preview()  # the parent's run advances the shared sampler RNG
+if preview_in_child() == preview_in_child():
+    sys.exit("sibling forked children generated identical sampler data")
+"""
+
+
+@pytest.mark.skipif("fork" not in multiprocessing.get_all_start_methods(), reason="requires the fork start method")
+def test_sibling_forked_children_generate_different_sampler_data() -> None:
+    result = subprocess.run(
+        [sys.executable, "-c", _PREVIEW_IN_SIBLING_FORKED_CHILDREN],
+        capture_output=True,
+        text=True,
+        timeout=120,
+        env={**os.environ, "NEMO_TELEMETRY_ENABLED": "false"},
+    )
+    assert result.returncode == 0, f"sibling forked children failed:\n{result.stderr[-4000:]}"