From 14d99dcac0ce025692d3a05686573f7de8078d0b Mon Sep 17 00:00:00 2001 From: Faruk Alpay <32020561+farukalpay@users.noreply.github.com> Date: Mon, 25 Aug 2025 15:14:11 +0200 Subject: [PATCH] Add deterministic seeding utility --- seeding.py | 122 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 122 insertions(+) create mode 100644 seeding.py diff --git a/seeding.py b/seeding.py new file mode 100644 index 0000000..1d83fd7 --- /dev/null +++ b/seeding.py @@ -0,0 +1,122 @@ +"""Deterministic seeding utilities. + +This module hashes a master seed together with component, run and stream +identifiers to derive 64-bit sub-seeds. The message is encoded as +``"{master_seed}|{component_id}|{run_id}|{stream_id}"`` and hashed with +SHA-256, taking the first eight bytes as a big-endian integer. The resulting +sub-seed can be fed to independent random number generators. + +Philox, a counter-based RNG available in NumPy, is chosen because it allows +reproducible, stateless streams that can be advanced independently across +parallel processes. +""" + +from __future__ import annotations + +import argparse +import hashlib +import random + +import numpy as np + +try: # Optional PyTorch integration + import torch +except Exception: # pragma: no cover - environment without torch + torch = None # type: ignore + + +def make_subseed( + master_seed: int | str, component_id: str, run_id: str, stream_id: int | str = 0 +) -> int: + """Derive a deterministic 64-bit sub-seed. + + Args: + master_seed: Global seed as an integer or string. + component_id: Identifier for the component (e.g., "dataloader"). + run_id: Identifier for the current experiment/run. + stream_id: Optional sub-stream identifier. + + Returns: + The first eight bytes of the SHA-256 digest interpreted as a + big-endian integer. + """ + + message = f"{master_seed}|{component_id}|{run_id}|{stream_id}" + digest = hashlib.sha256(message.encode("utf-8")).digest() + return int.from_bytes(digest[:8], "big") + + +def philox_rng(subseed: int) -> np.random.Generator: + """Create a NumPy Philox generator seeded with ``subseed``.""" + + return np.random.Generator(np.random.Philox(subseed)) + + +def python_rng(subseed: int) -> random.Random: + """Return a ``random.Random`` instance seeded with ``subseed``.""" + + rng = random.Random() + rng.seed(subseed) + return rng + + +def torch_rng(subseed: int, device: str | torch.device = "cpu"): + """Return a torch ``Generator`` seeded with ``subseed``. + + Args: + subseed: Seed value for the generator. + device: Torch device string or ``torch.device``. Defaults to ``"cpu"``. + + Raises: + ImportError: If PyTorch is not installed. + """ + + if torch is None: # pragma: no cover - only hit when torch missing + raise ImportError("PyTorch is not installed") + + gen = torch.Generator(device) + gen.manual_seed(int(subseed)) + return gen + + +def set_torch_deterministic(enabled: bool = True) -> None: + """Toggle deterministic algorithms and cuDNN flags in PyTorch. + + Args: + enabled: Whether to enable deterministic behaviour. + + Raises: + ImportError: If PyTorch is not installed. + """ + + if torch is None: # pragma: no cover - only hit when torch missing + raise ImportError("PyTorch is not installed") + + torch.backends.cudnn.deterministic = enabled + torch.backends.cudnn.benchmark = not enabled + torch.use_deterministic_algorithms(enabled) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Deterministic seeding utility") + parser.add_argument("--master-seed", required=True, help="Master seed (int or str)") + parser.add_argument("--component", required=True, help="Component identifier") + parser.add_argument("--run-id", required=True, help="Run identifier") + parser.add_argument("--stream-id", type=int, default=0, help="Stream identifier") + parser.add_argument("--n", type=int, default=5, help="How many numbers to draw") + args = parser.parse_args() + + subseed = make_subseed( + args.master_seed, args.component, args.run_id, args.stream_id + ) + print(f"Subseed: {subseed}") + + np_rng = philox_rng(subseed) + print("NumPy Philox:", np_rng.random(args.n)) + + py_rng = python_rng(subseed) + print("Python random:", [py_rng.random() for _ in range(args.n)]) + + if torch is not None: + tgen = torch_rng(subseed) + print("PyTorch:", torch.rand(args.n, generator=tgen).tolist())