| """Deterministic random-state management. |
| |
| The generator must produce byte-identical output for a given ``(master_seed, |
| sample_index)`` pair regardless of process/thread scheduling, batch size, or |
| worker count. To guarantee this we never share a single global RNG across |
| samples or pipeline stages. Instead, every stage of every sample draws from |
| an independently *derived* RNG whose seed is a pure function of |
| ``(master_seed, sample_index, stage_name)``. |
| |
| This makes the pipeline safe for multiprocessing: worker N generating sample |
| 57 produces exactly the same random draws as a single-process run would have, |
| because nothing depends on execution order. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import random |
| from typing import Any |
|
|
| import numpy as np |
|
|
| _MOD32 = 2**32 |
|
|
|
|
| def derive_seed(master_seed: int, *parts: Any) -> int: |
| """Derive a deterministic 32-bit seed from a master seed and arbitrary parts. |
| |
| Uses SHA-256 over the concatenated string representation of all parts so |
| that the mapping is stable across Python versions and processes (unlike |
| Python's salted ``hash()``). |
| |
| Args: |
| master_seed: The top-level seed configured by the user. |
| *parts: Any number of hashable/stringifiable disambiguators, e.g. |
| sample index and pipeline stage name. |
| |
| Returns: |
| A non-negative integer in ``[0, 2**32)`` suitable for seeding |
| ``numpy.random.Generator`` or ``random.Random``. |
| """ |
| key = f"{master_seed}|" + "|".join(str(p) for p in parts) |
| digest = hashlib.sha256(key.encode("utf-8")).digest() |
| return int.from_bytes(digest[:4], byteorder="big") % _MOD32 |
|
|
|
|
| class SeedManager: |
| """Provides independent, deterministic RNGs for each pipeline stage of one sample. |
| |
| Example: |
| >>> sm = SeedManager(master_seed=42, sample_index=7) |
| >>> rng = sm.numpy_rng("placement") |
| >>> rng.uniform(0, 1) |
| """ |
|
|
| def __init__(self, master_seed: int, sample_index: int) -> None: |
| self.master_seed = master_seed |
| self.sample_index = sample_index |
|
|
| def seed_for(self, stage: str) -> int: |
| """Return the deterministic integer seed for a named pipeline stage.""" |
| return derive_seed(self.master_seed, self.sample_index, stage) |
|
|
| def numpy_rng(self, stage: str) -> np.random.Generator: |
| """Return a fresh ``numpy.random.Generator`` (PCG64) seeded for ``stage``.""" |
| return np.random.default_rng(self.seed_for(stage)) |
|
|
| def python_rng(self, stage: str) -> random.Random: |
| """Return a fresh ``random.Random`` seeded for ``stage``.""" |
| return random.Random(self.seed_for(stage)) |
|
|
|
|
| def resolve_scalar_or_range(value: Any, rng: np.random.Generator) -> float: |
| """Resolve a config parameter that is either a fixed scalar or a ``{min, max}`` range. |
| |
| Args: |
| value: Either a plain number, or an object/mapping exposing ``min`` |
| and ``max`` attributes/keys (a :class:`~src.config.schema.Range` |
| or an equivalent dict). |
| rng: The seeded generator to draw from when ``value`` is a range. |
| |
| Returns: |
| A resolved float value. |
| """ |
| if isinstance(value, (int, float)): |
| return float(value) |
| if isinstance(value, dict) and "min" in value and "max" in value: |
| lo, hi = float(value["min"]), float(value["max"]) |
| return float(rng.uniform(lo, hi)) if hi > lo else lo |
| |
| lo = getattr(value, "min", None) |
| hi = getattr(value, "max", None) |
| if lo is not None and hi is not None: |
| lo, hi = float(lo), float(hi) |
| return float(rng.uniform(lo, hi)) if hi > lo else lo |
| raise TypeError(f"Cannot resolve scalar-or-range value: {value!r}") |
|
|
