Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,270 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import contextlib
|
||||
import warnings
|
||||
from collections.abc import Generator
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
__all__ = [
|
||||
"set_rng_state",
|
||||
"get_rng_state",
|
||||
"manual_seed",
|
||||
"seed",
|
||||
"initial_seed",
|
||||
"fork_rng",
|
||||
"thread_safe_generator",
|
||||
]
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.utils.data._utils.worker import WorkerInfo
|
||||
|
||||
from torch._C import default_generator
|
||||
|
||||
|
||||
def set_rng_state(new_state: torch.Tensor) -> None:
|
||||
r"""Sets the random number generator state.
|
||||
|
||||
.. note:: This function only works for CPU. For CUDA, please use
|
||||
:func:`torch.manual_seed`, which works for both CPU and CUDA.
|
||||
|
||||
Args:
|
||||
new_state (torch.ByteTensor): The desired state
|
||||
"""
|
||||
default_generator.set_state(new_state)
|
||||
|
||||
|
||||
def get_rng_state() -> torch.Tensor:
|
||||
r"""Returns the random number generator state as a `torch.ByteTensor`.
|
||||
|
||||
.. note:: The returned state is for the default generator on CPU only.
|
||||
|
||||
See also: :func:`torch.random.fork_rng`.
|
||||
"""
|
||||
return default_generator.get_state()
|
||||
|
||||
|
||||
def manual_seed(seed) -> torch._C.Generator:
|
||||
r"""Sets the seed for generating random numbers on all devices. Returns a
|
||||
`torch.Generator` object.
|
||||
|
||||
Args:
|
||||
seed (int): The desired seed. Value must be within the inclusive range
|
||||
`[-0x8000_0000_0000_0000, 0xffff_ffff_ffff_ffff]`. Otherwise, a RuntimeError
|
||||
is raised. Negative inputs are remapped to positive values with the formula
|
||||
`0xffff_ffff_ffff_ffff + seed`.
|
||||
"""
|
||||
return _manual_seed_impl(seed)
|
||||
|
||||
|
||||
def _manual_seed_impl(seed) -> torch._C.Generator:
|
||||
seed = int(seed)
|
||||
import torch.cuda
|
||||
|
||||
if not torch.cuda._is_in_bad_fork():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
import torch.mps
|
||||
|
||||
if not torch.mps._is_in_bad_fork():
|
||||
torch.mps.manual_seed(seed)
|
||||
|
||||
import torch.xpu
|
||||
|
||||
if not torch.xpu._is_in_bad_fork():
|
||||
torch.xpu.manual_seed_all(seed)
|
||||
|
||||
import torch.mtia
|
||||
|
||||
if not torch.mtia._is_in_bad_fork():
|
||||
torch.mtia.manual_seed_all(seed)
|
||||
|
||||
_seed_custom_device(seed)
|
||||
|
||||
return default_generator.manual_seed(seed)
|
||||
|
||||
|
||||
def seed() -> int:
|
||||
r"""Sets the seed for generating random numbers to a non-deterministic
|
||||
random number on all devices. Returns a 64 bit number used to seed the RNG.
|
||||
"""
|
||||
seed = default_generator.seed()
|
||||
import torch.cuda
|
||||
|
||||
if not torch.cuda._is_in_bad_fork():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
import torch.mps
|
||||
|
||||
if not torch.mps._is_in_bad_fork():
|
||||
torch.mps.manual_seed(seed)
|
||||
|
||||
import torch.xpu
|
||||
|
||||
if not torch.xpu._is_in_bad_fork():
|
||||
torch.xpu.manual_seed_all(seed)
|
||||
|
||||
import torch.mtia
|
||||
|
||||
if not torch.mtia._is_in_bad_fork():
|
||||
torch.mtia.manual_seed_all(seed)
|
||||
|
||||
_seed_custom_device(seed)
|
||||
|
||||
return seed
|
||||
|
||||
|
||||
def _seed_custom_device(seed) -> None:
|
||||
r"""Sets the seed to generate random numbers for custom device.
|
||||
|
||||
Args:
|
||||
seed (int): The desired seed.
|
||||
|
||||
See [Note: support the custom device with privateuse1]
|
||||
"""
|
||||
seed = int(seed)
|
||||
custom_backend_name = torch._C._get_privateuse1_backend_name()
|
||||
if hasattr(torch, custom_backend_name):
|
||||
custom_device_mod = getattr(torch, custom_backend_name)
|
||||
_bad_fork_name = "_is_in_bad_fork"
|
||||
_seed_all_name = "manual_seed_all"
|
||||
if hasattr(custom_device_mod, _bad_fork_name) and hasattr(
|
||||
custom_device_mod, _seed_all_name
|
||||
):
|
||||
if not getattr(custom_device_mod, _bad_fork_name)():
|
||||
getattr(custom_device_mod, _seed_all_name)(seed)
|
||||
else:
|
||||
message = f"Set seed for `{custom_backend_name}` device does not take effect, please add API's "
|
||||
message += f"`{_bad_fork_name}` and `{_seed_all_name}` to `{custom_backend_name}` device module."
|
||||
warnings.warn(message, UserWarning, stacklevel=3)
|
||||
|
||||
|
||||
def initial_seed() -> int:
|
||||
r"""Returns the initial seed for generating random numbers as a
|
||||
Python `long`.
|
||||
|
||||
.. note:: The returned seed is for the default generator on CPU only.
|
||||
"""
|
||||
return default_generator.initial_seed()
|
||||
|
||||
|
||||
_fork_rng_warned_already = False
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def fork_rng(
|
||||
devices=None,
|
||||
enabled=True,
|
||||
_caller="fork_rng",
|
||||
_devices_kw="devices",
|
||||
device_type="cuda",
|
||||
) -> Generator:
|
||||
"""
|
||||
Forks the RNG, so that when you return, the RNG is reset
|
||||
to the state that it was previously in.
|
||||
|
||||
Args:
|
||||
devices (iterable of Device IDs): devices for which to fork
|
||||
the RNG. CPU RNG state is always forked. By default, :meth:`fork_rng` operates
|
||||
on all devices, but will emit a warning if your machine has a lot
|
||||
of devices, since this function will run very slowly in that case.
|
||||
If you explicitly specify devices, this warning will be suppressed
|
||||
enabled (bool): if ``False``, the RNG is not forked. This is a convenience
|
||||
argument for easily disabling the context manager without having
|
||||
to delete it and unindent your Python code under it.
|
||||
device_type (str): device type str, default is `cuda`. As for supported device,
|
||||
see details in :ref:`accelerator<accelerators>`
|
||||
"""
|
||||
|
||||
if device_type == "meta":
|
||||
yield
|
||||
return
|
||||
|
||||
device_type = torch.device(device_type).type
|
||||
device_mod = getattr(torch, device_type, None)
|
||||
if device_mod is None:
|
||||
raise RuntimeError(
|
||||
f"torch has no module of `{device_type}`, you should register "
|
||||
+ "a module by `torch._register_device_module`."
|
||||
)
|
||||
global _fork_rng_warned_already
|
||||
|
||||
# Internal arguments:
|
||||
# _caller: the function which called fork_rng, which the user used
|
||||
# _devices_kw: the devices keyword of _caller
|
||||
|
||||
if not enabled:
|
||||
yield
|
||||
return
|
||||
|
||||
if devices is None:
|
||||
num_devices = device_mod.device_count()
|
||||
if num_devices > 1 and not _fork_rng_warned_already:
|
||||
message = (
|
||||
f"{device_type.upper()} reports that you have {num_devices} available devices, and "
|
||||
f"you have used {_caller} without explicitly specifying which devices are being used. "
|
||||
f"For safety, we initialize *every* {device_type.upper()} device by default, which can "
|
||||
f"be quite slow if you have a lot of {device_type.upper()}s. If you know that you are only"
|
||||
f" making use of a few {device_type.upper()} devices, set the environment variable "
|
||||
f"{device_type.upper()}_VISIBLE_DEVICES or the '{_devices_kw}' keyword argument of {_caller} "
|
||||
"with the set of devices you are actually using. For example, if you are using CPU only, "
|
||||
"set device.upper()_VISIBLE_DEVICES= or devices=[]; if you are using device 0 only, "
|
||||
f"set {device_type.upper()}_VISIBLE_DEVICES=0 or devices=[0]. To initialize all devices "
|
||||
f"and suppress this warning, set the '{_devices_kw}' keyword argument to "
|
||||
f"`range(torch.{device_type}.device_count())`."
|
||||
)
|
||||
warnings.warn(message, stacklevel=2)
|
||||
_fork_rng_warned_already = True
|
||||
devices = list(range(num_devices))
|
||||
else:
|
||||
# Protect against user passing us a generator; we need to traverse this
|
||||
# multiple times but a generator will be exhausted upon first traversal
|
||||
devices = list(devices)
|
||||
|
||||
cpu_rng_state = torch.get_rng_state()
|
||||
device_rng_states = [device_mod.get_rng_state(device) for device in devices]
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch.set_rng_state(cpu_rng_state)
|
||||
for device, device_rng_state in zip(devices, device_rng_states):
|
||||
device_mod.set_rng_state(device_rng_state, device)
|
||||
|
||||
|
||||
def thread_safe_generator() -> torch.Generator | None:
|
||||
"""Returns a thread-safe random number generator for use in DataLoader workers.
|
||||
This function provides a convenient way for transforms and user code to use
|
||||
thread-safe random number generation without manually checking worker context.
|
||||
When called in a DataLoader thread worker, returns the worker's thread-local
|
||||
:class:`torch.Generator`. When called in the main process or process workers,
|
||||
returns ``None`` (which causes PyTorch functions to use the default global RNG).
|
||||
Returns:
|
||||
Optional[torch.Generator]: Thread-local generator in thread workers, None otherwise.
|
||||
Example::
|
||||
>>> from torch.random import thread_safe_generator
|
||||
>>> generator = thread_safe_generator()
|
||||
>>> torch.randint(0, 10, (5,), generator=generator)
|
||||
Example with transforms::
|
||||
>>> from torch.random import thread_safe_generator
|
||||
>>> class MyRandomTransform:
|
||||
... def __call__(self, img):
|
||||
... generator = thread_safe_generator()
|
||||
... offset = torch.randint(0, 10, (2,), generator=generator)
|
||||
... return img[..., offset[0]:, offset[1]:]
|
||||
"""
|
||||
# Lazy import to avoid circular dependency during torch module initialization
|
||||
# torch.__init__ loads torch.random early, but torch.utils.data triggers
|
||||
# torch.distributed which needs torch to be fully initialized
|
||||
from torch.utils.data import get_worker_info
|
||||
|
||||
worker_info: WorkerInfo | None = get_worker_info()
|
||||
if (
|
||||
worker_info is not None
|
||||
and worker_info.worker_method == "thread"
|
||||
and worker_info.rng is not None
|
||||
):
|
||||
return worker_info.rng.torch_generator
|
||||
return None
|
||||
Reference in New Issue
Block a user