import hashlib
import importlib.abc
import os
import struct
import random
import sys
from contextlib import AbstractContextManager
from typing import Any, Union, Optional, Tuple, cast, List
import numpy as np
from numpy.random import RandomState
# Torch is optional and expensive: importing it costs several hundred MiB of
# RSS per process. Seeding must therefore never import Torch itself. Torch
# RNG state is managed only once Torch is loaded; if Torch is imported later
# (e.g., by a Brain module loaded inside an RNGScope, or by a Brain's setup()
# after seed_globals()), a one-shot import hook applies exactly the seeding
# the eager implementation would have applied, so all Torch random streams
# stay identical.
def _loaded_torch() -> Any:
"""Return the Torch module if it is already imported, else ``None``."""
return sys.modules.get("torch")
_active_scopes: List["RNGScope"] = [] # in entry order, outermost first
_pending_global_seed: Optional[int] = None
def _torch_hook_wanted() -> bool:
return bool(_active_scopes) or _pending_global_seed is not None
def _on_torch_imported(torch: Any) -> None:
"""Apply deferred seeding right after Torch finished importing.
Order mirrors the eager implementation: the process-global seed from
:func:`seed_globals` came first, then each active scope (outermost
first) captured the then-current state as its ambient and seeded its own
stream.
"""
global _pending_global_seed
if _pending_global_seed is not None:
torch.manual_seed(_pending_global_seed)
_pending_global_seed = None
for scope in list(_active_scopes):
scope._adopt_torch(torch)
class _NotifyingLoader(importlib.abc.Loader):
"""Wraps Torch's real loader to run :func:`_on_torch_imported`."""
def __init__(self, loader):
self._loader = loader
def create_module(self, spec):
return self._loader.create_module(spec)
def exec_module(self, module):
self._loader.exec_module(module)
# Hide the wrapper from introspection (inspect, importlib.resources)
module.__loader__ = self._loader
if getattr(module, "__spec__", None) is not None:
module.__spec__.loader = self._loader
_on_torch_imported(module)
def __getattr__(self, name):
return getattr(self._loader, name)
class _TorchImportHook(importlib.abc.MetaPathFinder):
"""Meta path finder that only acts on a top-level ``import torch``
while deferred seeding is pending; otherwise it is a no-op."""
def find_spec(self, fullname, path=None, target=None):
if fullname != "torch" or not _torch_hook_wanted():
return None
for finder in sys.meta_path:
if finder is self or not hasattr(finder, "find_spec"):
continue
spec = finder.find_spec(fullname, path, target)
if spec is not None:
if spec.loader is not None and hasattr(
spec.loader, "exec_module"
):
spec.loader = _NotifyingLoader(spec.loader)
return spec
return None
_TORCH_IMPORT_HOOK = _TorchImportHook()
def _install_torch_import_hook() -> None:
if not any(f is _TORCH_IMPORT_HOOK for f in sys.meta_path):
sys.meta_path.insert(0, _TORCH_IMPORT_HOOK)
[docs]
class RNGScope(AbstractContextManager):
"""Temporarily isolate Python/NumPy and optional Torch CPU RNG streams.
Scopes are synchronous and must not overlap across threads: these libraries
expose process-global RNGs. CUDA RNG state is intentionally unmanaged.
Torch is never imported by the scope. If it is loaded on entry, its CPU
RNG is isolated like Python's and NumPy's; if it is imported while the
scope is active, it is seeded right after its import, which yields the
same random streams as importing it on entry.
"""
def __init__(self, seed: Optional[int]):
# Seeds from ExperimentRun.create_subseed() are numpy integer scalars.
self.seed = None if seed is None else int(seed)
self._active = False
self._python_state = None
self._numpy_state = None
self._torch_state = None
self._ambient = None
self._ambient_torch = None
self._torch_managed = False
def _adopt_torch(self, torch) -> None:
"""Start managing Torch's CPU RNG for the current activation."""
self._ambient_torch = torch.random.get_rng_state()
if self._torch_state is None:
# Equivalent to the eager first-entry seeding: an unused stream
# seeded on an earlier entry is still in its freshly seeded state.
torch.random.default_generator.manual_seed(self.seed)
else:
torch.random.set_rng_state(self._torch_state)
self._torch_managed = True
def __enter__(self):
if self.seed is None:
return self
if self._active:
raise RuntimeError("RNGScope cannot be entered recursively")
ambient_python = random.getstate()
ambient_numpy = np.random.get_state()
self._ambient = (ambient_python, ambient_numpy)
if self._python_state is None:
random.seed(self.seed)
np.random.seed(self.seed % (2**32))
else:
random.setstate(self._python_state)
np.random.set_state(self._numpy_state)
self._torch_managed = False
torch = _loaded_torch()
if torch is not None:
self._adopt_torch(torch)
else:
_install_torch_import_hook()
self._active = True
_active_scopes.append(self)
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if self._active:
try:
_active_scopes.remove(self)
except ValueError:
pass
self._python_state = random.getstate()
self._numpy_state = np.random.get_state()
torch = _loaded_torch()
if self._torch_managed and torch is not None:
self._torch_state = torch.random.get_rng_state()
torch.random.set_rng_state(self._ambient_torch)
ambient_python, ambient_numpy = self._ambient
random.setstate(ambient_python)
np.random.set_state(ambient_numpy)
self._ambient = None
self._ambient_torch = None
self._torch_managed = False
self._active = False
return False
def __getstate__(self):
if self._active:
raise RuntimeError("Cannot pickle an active RNGScope")
return self.__dict__.copy()
[docs]
def seed_globals(seed: int) -> None:
"""Seed process-global Python, NumPy and optional Torch RNGs.
Torch is not imported for this. If it is not loaded yet, its seed is
applied as soon as it is imported (the latest call wins), which yields
the same Torch random stream as seeding it eagerly.
"""
global _pending_global_seed
seed = int(seed)
if seed < 0:
raise ValueError("seed must be non-negative")
random.seed(seed)
np.random.seed(seed % (2**32))
torch = _loaded_torch()
if torch is not None:
_pending_global_seed = None
torch.manual_seed(seed)
else:
_pending_global_seed = seed
_install_torch_import_hook()
[docs]
def np_random(seed: Optional[int] = None) -> Tuple[RandomState, int]:
if seed is not None and not (isinstance(seed, int) and 0 <= seed):
raise RuntimeError(
"Seed must be a non-negative integer "
"or omitted, not {}".format(seed)
)
seed = create_seed(seed)
rng: RandomState = RandomState()
rng.seed(_int_list_from_bigint(hash_seed(seed)))
return rng, seed
[docs]
def hash_seed(seed: Union[int, str, None] = None, max_bytes: int = 8) -> int:
"""Any given evaluation is likely to have many PRNG's active at
once. (Most commonly, because the environment is running in
multiple processes.) There's literature indicating that having
linear correlations between seeds of multiple PRNG's can correlate
the outputs:
http://blogs.unity3d.com/2015/01/07/a-primer-on-repeatable-random-numbers/
http://stackoverflow.com/questions/1554958/how-different-do-random-seeds-need-to-be
http://dl.acm.org/citation.cfm?id=1276928
Thus, for sanity we hash the seeds before using them. (This scheme
is likely not crypto-strength, but it should be good enough to get
rid of simple correlations.)
Args:
seed (Union[int, str, None]): Hashed via its string form; None
seeds from an operating system specific randomness source.
max_bytes: Maximum number of bytes to use in the hashed seed.
"""
if seed is None:
seed = create_seed(max_bytes=max_bytes)
hsh = hashlib.sha512(str(seed).encode("utf8")).digest()
return _bigint_from_bytes(hsh[:max_bytes])
[docs]
def create_seed(a: Union[int, str, None] = None, max_bytes: int = 8) -> int:
"""Create a strong random seed. Otherwise, Python 2 would seed using
the system time, which might be non-robust especially in the
presence of concurrency.
Args:
a (Union[int, str, None]): None seeds from an operating system specific randomness source.
max_bytes: Maximum number of bytes to use in the seed.
"""
# Adapted from https://svn.python.org/projects/python/tags/r32/Lib/random.py
if a is None:
a = _bigint_from_bytes(os.urandom(max_bytes))
elif isinstance(a, str):
try:
a = int_to_seed(int(a), max_bytes)
except ValueError:
a = cast(str, a)
a_bytes = a.encode("utf8")
a_bytes = a_bytes + hashlib.sha512(a_bytes).digest()
a = _bigint_from_bytes(a_bytes[:max_bytes])
elif isinstance(a, int):
a = int_to_seed(a, max_bytes)
else:
raise RuntimeError("Invalid type for seed: %s (%s)" % (type(a), a))
return a
[docs]
def int_to_seed(a: int, max_bytes: int) -> int:
"""Turns integer into a seed.
Parameters
----------
a Seed starting point
max_bytes Maximum number of bytes to use in the seed
Returns
-------
a Seed
"""
a = a % 2 ** (8 * max_bytes)
return a
def _bigint_from_bytes(bytes_: bytes) -> int:
sizeof_int = 4
padding = sizeof_int - len(bytes_) % sizeof_int
bytes_ += b"\0" * padding
int_count = int(len(bytes_) / sizeof_int)
unpacked = struct.unpack("{}I".format(int_count), bytes_)
accum = 0
for i, val in enumerate(unpacked):
accum += 2 ** (sizeof_int * 8 * i) * val
return accum
def _int_list_from_bigint(bigint: int) -> List[int]:
# Special case 0
if bigint < 0:
raise TypeError("Seed must be non-negative, not {}".format(bigint))
elif bigint == 0:
return [0]
ints = []
while bigint > 0:
bigint, mod = divmod(bigint, 2**32)
ints.append(mod)
return ints