File size: 1,743 Bytes
6cc35b0 | 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 | """Dependency-light RNG scoping for ESMFold2 workflows."""
from __future__ import annotations
import random
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any
import numpy as np
import torch
from torch import Tensor
@dataclass(frozen=True)
class _RandomState:
python: object
numpy: tuple[Any, ...]
torch_cpu: Tensor
torch_cuda: list[Tensor] | None
def _capture_random_state() -> _RandomState:
cuda_state = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
return _RandomState(
python=random.getstate(),
numpy=np.random.get_state(),
torch_cpu=torch.random.get_rng_state(),
torch_cuda=cuda_state,
)
def _restore_random_state(state: _RandomState) -> None:
random.setstate(state.python)
np.random.set_state(state.numpy)
torch.random.set_rng_state(state.torch_cpu)
if state.torch_cuda is not None:
torch.cuda.set_rng_state_all(state.torch_cuda)
@contextmanager
def seed_context(seed: int | None) -> Iterator[None]:
"""Seed Python, NumPy, and Torch temporarily, then restore every stream."""
if seed is None:
yield
return
if isinstance(seed, bool) or not isinstance(seed, int):
raise TypeError("seed must be None or an integer (excluding bool).")
state = _capture_random_state()
normalized_seed = seed % (2**32)
random.seed(normalized_seed)
np.random.seed(normalized_seed)
torch.manual_seed(normalized_seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(normalized_seed)
try:
yield
finally:
_restore_random_state(state)
__all__ = ["seed_context"]
|