From a0107a6b15308fd2dc7dc2ac82fdb781ee7f7907 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Thu, 22 Oct 2026 22:03:16 +0000 Subject: [PATCH 1/3] fix: reseed shared sampler RNGs in forked children MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Samplers draw from numpy's global RandomState and Faker's shared Random. CPython reseeds the stdlib random module after fork but not these, so every child forked after a parent run generated the same sampler data. An after-in-child fork hook now reseeds both from OS entropy, matching CPython's behavior for random. Seeding in the process is unchanged, and the parent's RNG is not touched. 🤖 Generated with Claude Code Co-Authored-By: Claude Signed-off-by: Aaron Gonzales --- architecture/sampling.md | 1 + .../engine/sampling_gen/utils.py | 20 ++++++ .../tests/engine/sampling_gen/test_utils.py | 42 ++++++++++++ .../tests/interface/test_fork_safety.py | 64 +++++++++++++++++++ 4 files changed, 127 insertions(+) diff --git a/architecture/sampling.md b/architecture/sampling.md index 4fabbb566..4312ebe68 100644 --- a/architecture/sampling.md +++ b/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. ## Cross-References diff --git a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/utils.py b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/utils.py index 958b70e5d..ad423391a 100644 --- a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/utils.py +++ b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/utils.py @@ -4,6 +4,8 @@ from __future__ import annotations import numbers +import os +import sys import data_designer.lazy_heavy_imports as lazy @@ -40,3 +42,21 @@ 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) + + +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. + numpy = sys.modules.get("numpy") + if numpy 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(after_in_child=_reseed_shared_rngs_after_fork) diff --git a/packages/data-designer-engine/tests/engine/sampling_gen/test_utils.py b/packages/data-designer-engine/tests/engine/sampling_gen/test_utils.py index f59546746..a97d46403 100644 --- a/packages/data-designer-engine/tests/engine/sampling_gen/test_utils.py +++ b/packages/data-designer-engine/tests/engine/sampling_gen/test_utils.py @@ -7,6 +7,7 @@ import data_designer.lazy_heavy_imports as lazy from data_designer.engine.sampling_gen.utils import check_random_state +from tests.engine.fork_probe import requires_fork, run_fork_probe @pytest.mark.parametrize( @@ -36,3 +37,44 @@ 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" + """ + ) diff --git a/packages/data-designer/tests/interface/test_fork_safety.py b/packages/data-designer/tests/interface/test_fork_safety.py index 87dde0102..05536f3fc 100644 --- a/packages/data-designer/tests/interface/test_fork_safety.py +++ b/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:]}" From 56b49bebf398748f8aeb5f909346bc00c06ada1d Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 26 Oct 2026 19:02:12 +0000 Subject: [PATCH 2/3] fix: hold numpy's RNG lock across fork before reseeding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The after-fork reseed called numpy.random.seed(), which takes the global RandomState's lock. A child forked while another thread was drawing inherited that lock held by a thread it does not have, so the hook blocked forever before the child ran. Acquire the lock before fork and release it in the parent and in the child before reseeding, as logging does for its locks. 🤖 Generated with Claude Code Co-Authored-By: Claude Signed-off-by: Aaron Gonzales --- architecture/sampling.md | 2 +- .../engine/sampling_gen/utils.py | 45 +++++++++++++-- .../tests/engine/sampling_gen/test_utils.py | 55 ++++++++++++++++++- 3 files changed, 96 insertions(+), 6 deletions(-) diff --git a/architecture/sampling.md b/architecture/sampling.md index 4312ebe68..f7f4fcd47 100644 --- a/architecture/sampling.md +++ b/architecture/sampling.md @@ -69,7 +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. +- **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 diff --git a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/utils.py b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/utils.py index ad423391a..b4d0e5915 100644 --- a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/utils.py +++ b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/utils.py @@ -6,6 +6,7 @@ import numbers import os import sys +import threading import data_designer.lazy_heavy_imports as lazy @@ -44,19 +45,55 @@ def check_random_state(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. - numpy = sys.modules.get("numpy") - if numpy is not None: - numpy.random.seed() + _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(after_in_child=_reseed_shared_rngs_after_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, + ) diff --git a/packages/data-designer-engine/tests/engine/sampling_gen/test_utils.py b/packages/data-designer-engine/tests/engine/sampling_gen/test_utils.py index a97d46403..c4c1dfd89 100644 --- a/packages/data-designer-engine/tests/engine/sampling_gen/test_utils.py +++ b/packages/data-designer-engine/tests/engine/sampling_gen/test_utils.py @@ -3,10 +3,12 @@ from __future__ import annotations +import threading + 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 @@ -78,3 +80,54 @@ def test_forked_child_reseed_does_not_change_parent_draws() -> None: assert rng.random_sample(3).tolist() == expected, "forking reseeded the parent's RNG" """ ) + + +def test_numpy_global_rng_lock_is_the_lock_draws_hold() -> None: + # The fork hooks reach into numpy internals for this lock; fail here if numpy moves it. + rng = check_random_state(None) # imports numpy.random + lock = _numpy_global_rng_lock() + assert lock is not None + finished = threading.Event() + with lock: + drawer = threading.Thread(target=lambda: (rng.random_sample(), finished.set())) + drawer.start() + assert not finished.wait(0.2), "a global RandomState draw did not wait for the lock" + drawer.join(timeout=5) + assert finished.is_set() + + +@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" + """ + ) From 3412cdb3912d849aae86afe3b426515677f2728a Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 26 Oct 2026 19:05:54 +0000 Subject: [PATCH 3/3] test: check the fork-hook RNG lock against numpy's own draws and reseed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fork test holds the lock the hook helper returns, so a helper that returned a stable lock numpy does not use would still pass it. Check the helper against numpy itself: a sampler draw and np.random.seed() (what the child hook calls) must wait while the lock is held, and a separate RandomState, as a control, must not. 🤖 Generated with Claude Code Co-Authored-By: Claude Signed-off-by: Aaron Gonzales --- .../tests/engine/sampling_gen/test_utils.py | 51 +++++++++++++++---- 1 file changed, 41 insertions(+), 10 deletions(-) diff --git a/packages/data-designer-engine/tests/engine/sampling_gen/test_utils.py b/packages/data-designer-engine/tests/engine/sampling_gen/test_utils.py index c4c1dfd89..8afbc6fb4 100644 --- a/packages/data-designer-engine/tests/engine/sampling_gen/test_utils.py +++ b/packages/data-designer-engine/tests/engine/sampling_gen/test_utils.py @@ -4,6 +4,7 @@ from __future__ import annotations import threading +from collections.abc import Callable import pytest @@ -82,18 +83,48 @@ def test_forked_child_reseed_does_not_change_parent_draws() -> None: ) -def test_numpy_global_rng_lock_is_the_lock_draws_hold() -> None: - # The fork hooks reach into numpy internals for this lock; fail here if numpy moves it. - rng = check_random_state(None) # imports numpy.random +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 - finished = threading.Event() - with lock: - drawer = threading.Thread(target=lambda: (rng.random_sample(), finished.set())) - drawer.start() - assert not finished.wait(0.2), "a global RandomState draw did not wait for the lock" - drawer.join(timeout=5) - assert finished.is_set() + 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