rpx_benchmark.determinism

  1"""Deterministic seeding for reproducible benchmark runs.
  2
  3Seeds every stochastic source the toolkit is likely to touch:
  4``random``, ``numpy``, ``torch`` (CPU + CUDA + MPS), and
  5HuggingFace ``transformers`` / ``datasets`` where present. Safe to
  6call even when those libraries aren't installed — the function
  7touches only the backends that import successfully.
  8
  9Example
 10-------
 11
 12Imperative::
 13
 14    from rpx_benchmark.determinism import seed_all
 15    seed_all(42)
 16
 17Context manager (restores RNG state on exit)::
 18
 19    from rpx_benchmark.determinism import deterministic
 20
 21    with deterministic(42):
 22        result = runner.run()
 23
 24Design notes
 25------------
 26- Seeding is *best-effort across backends*. Calling code should treat
 27  this as "best chance at determinism", not a hard guarantee —
 28  non-deterministic CUDA kernels, cuDNN algorithm selection, and
 29  multi-threaded data loaders can still introduce run-to-run variance.
 30- The context manager snapshots and restores the standard-library and
 31  numpy RNG state so nested determinism blocks don't leak state.
 32  Torch / transformers / datasets expose no cheap state-snapshot
 33  primitive, so their post-exit state is the sequence advanced during
 34  the block (same as re-seeding without a context manager).
 35"""
 36
 37from __future__ import annotations
 38
 39import os
 40import random
 41from contextlib import contextmanager
 42from typing import Iterator
 43
 44__all__ = ["seed_all", "deterministic", "RPX_SEED"]
 45
 46
 47#: Canonical project-wide seed. MMDDYYYY-encoded 05/06/2026 — the date
 48#: the team committed to a single deterministic seed across the whole
 49#: benchmark (bootstrap CI, test data generation, sparse-depth /
 50#: keypoint-pair samplers, ORD pixel-pair sampler, etc.). Stored here so
 51#: there's exactly one place to change if the project policy shifts.
 52#: Python's int literal grammar disallows a leading zero, so the
 53#: storage form is ``5_062_026``; the underscores keep the
 54#: month-day-year segmentation visible at a glance.
 55RPX_SEED: int = 5_062_026
 56
 57
 58def _seed_python(seed: int) -> None:
 59    random.seed(seed)
 60    # PYTHONHASHSEED affects dict ordering in child subprocesses; the
 61    # running interpreter snapshots it at startup so setting it here
 62    # only helps processes spawned later (e.g. DataLoader workers).
 63    os.environ["PYTHONHASHSEED"] = str(seed)
 64
 65
 66def _seed_numpy(seed: int) -> None:
 67    try:
 68        import numpy as np  # noqa: PLC0415
 69    except ImportError:
 70        return
 71    np.random.seed(seed)
 72
 73
 74def _seed_torch(seed: int) -> None:
 75    try:
 76        import torch  # noqa: PLC0415
 77    except ImportError:
 78        return
 79    torch.manual_seed(seed)
 80    if hasattr(torch, "cuda") and torch.cuda.is_available():
 81        torch.cuda.manual_seed_all(seed)
 82    mps = getattr(torch, "mps", None)
 83    if mps is not None and hasattr(mps, "manual_seed"):
 84        try:
 85            mps.manual_seed(seed)
 86        except Exception:
 87            pass
 88
 89
 90def _seed_hf(seed: int) -> None:
 91    """Seed transformers / datasets when they're importable.
 92
 93    Both libraries expose a ``set_seed`` helper that internally seeds
 94    Python, NumPy, and PyTorch — we still call our own seeders first
 95    so the HF helpers only fill gaps we might have missed.
 96    """
 97    try:
 98        from transformers import set_seed as _tf_set_seed  # noqa: PLC0415
 99
100        _tf_set_seed(seed)
101    except ImportError:
102        pass
103
104
105def seed_all(seed: int) -> None:
106    """Seed every stochastic backend detectable in the current env.
107
108    Parameters
109    ----------
110    seed : int
111        Seed forwarded to Python ``random``, NumPy, PyTorch (CPU +
112        CUDA + MPS), and ``transformers.set_seed`` when available.
113    """
114    _seed_python(seed)
115    _seed_numpy(seed)
116    _seed_torch(seed)
117    _seed_hf(seed)
118
119
120@contextmanager
121def deterministic(seed: int) -> Iterator[None]:
122    """Run a block with deterministic RNG and restore on exit.
123
124    Python ``random`` and NumPy RNG state are snapshotted and
125    restored. Torch / transformers state is merely re-seeded; see the
126    module docstring for the full scope.
127
128    Parameters
129    ----------
130    seed : int
131
132    Yields
133    ------
134    None
135    """
136    py_state = random.getstate()
137    np_state = None
138    try:
139        import numpy as np  # noqa: PLC0415
140
141        np_state = np.random.get_state()
142    except ImportError:
143        np = None  # type: ignore[assignment]
144
145    seed_all(seed)
146    try:
147        yield
148    finally:
149        random.setstate(py_state)
150        if np_state is not None:
151            import numpy as _np  # noqa: PLC0415
152
153            _np.random.set_state(np_state)