File size: 3,723 Bytes
3f820d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """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
# Pydantic Range model (duck-typed access)
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}")
|