Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,643 @@
|
||||
import collections
|
||||
import contextlib
|
||||
import ctypes
|
||||
import pickle
|
||||
import sys
|
||||
from typing import Any, Literal
|
||||
|
||||
import torch
|
||||
from torch._utils import _augment_memory_snapshot_stack_traces, _dummy_type
|
||||
from torch.types import Device
|
||||
|
||||
from . import _get_device_index, _is_compiled, _lazy_init, is_initialized
|
||||
|
||||
|
||||
if not _is_compiled():
|
||||
# Define dummy base classes
|
||||
torch._C.__dict__["_xpu_XPUAllocator"] = _dummy_type("_xpu_XPUAllocator")
|
||||
torch._C.__dict__["_XPUMemPool"] = _dummy_type("_XPUMemPool")
|
||||
torch._C.__dict__["_xpu_beginAllocateCurrentThreadToPool"] = _dummy_type(
|
||||
"_xpu_beginAllocateCurrentThreadToPool"
|
||||
)
|
||||
torch._C.__dict__["_xpu_endAllocateToPool"] = _dummy_type("_xpu_endAllocateToPool")
|
||||
torch._C.__dict__["_xpu_releasePool"] = _dummy_type("_xpu_releasePool")
|
||||
|
||||
|
||||
def empty_cache() -> None:
|
||||
r"""Release all unoccupied cached memory currently held by the caching
|
||||
allocator so that those can be used in other XPU application.
|
||||
|
||||
.. note::
|
||||
:func:`~torch.xpu.empty_cache` doesn't increase the amount of XPU
|
||||
memory available for PyTorch. However, it may help reduce fragmentation
|
||||
of XPU memory in certain cases.
|
||||
"""
|
||||
if is_initialized():
|
||||
torch._C._xpu_emptyCache()
|
||||
|
||||
|
||||
def reset_peak_memory_stats(device: Device = None) -> None:
|
||||
r"""Reset the "peak" stats tracked by the XPU memory allocator.
|
||||
|
||||
See :func:`~torch.xpu.memory_stats` for details. Peak stats correspond to the
|
||||
`"peak"` key in each individual stat dict.
|
||||
|
||||
Args:
|
||||
device (torch.device or int or str, optional): selected device. Returns
|
||||
statistic for the current device, given by :func:`~torch.xpu.current_device`,
|
||||
if :attr:`device` is ``None`` (default).
|
||||
"""
|
||||
device = _get_device_index(device, optional=True)
|
||||
return torch._C._xpu_resetPeakMemoryStats(device)
|
||||
|
||||
|
||||
def reset_accumulated_memory_stats(device: Device = None) -> None:
|
||||
r"""Reset the "accumulated" (historical) stats tracked by the XPU memory allocator.
|
||||
|
||||
See :func:`~torch.xpu.memory_stats` for details. Accumulated stats correspond to
|
||||
the `"allocated"` and `"freed"` keys in each individual stat dict.
|
||||
|
||||
Args:
|
||||
device (torch.device or int or str, optional): selected device. Returns
|
||||
statistic for the current device, given by :func:`~torch.xpu.current_device`,
|
||||
if :attr:`device` is ``None`` (default).
|
||||
"""
|
||||
device = _get_device_index(device, optional=True)
|
||||
return torch._C._xpu_resetAccumulatedMemoryStats(device)
|
||||
|
||||
|
||||
def memory_stats_as_nested_dict(device: Device = None) -> dict[str, Any]:
|
||||
r"""Return the result of :func:`~torch.xpu.memory_stats` as a nested dictionary."""
|
||||
if not is_initialized():
|
||||
return {}
|
||||
device = _get_device_index(device, optional=True)
|
||||
return torch._C._xpu_memoryStats(device)
|
||||
|
||||
|
||||
def memory_stats(device: Device = None) -> dict[str, Any]:
|
||||
r"""Return a dictionary of XPU memory allocator statistics for a given device.
|
||||
|
||||
The return value of this function is a dictionary of statistics, each of
|
||||
which is a non-negative integer.
|
||||
|
||||
Core statistics:
|
||||
|
||||
- ``"allocated_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
|
||||
amount of allocated memory.
|
||||
- ``"reserved_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
|
||||
amount of reserved memory.
|
||||
- ``"active_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
|
||||
amount of active memory.
|
||||
- ``"requested_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
|
||||
memory requested by client code, compare this with allocated_bytes to check if
|
||||
allocation rounding adds too much overhead.
|
||||
|
||||
For these core statistics, values are broken down as follows.
|
||||
|
||||
Pool type:
|
||||
|
||||
- ``all``: combined statistics across all memory pools.
|
||||
- ``large_pool``: statistics for the large allocation pool (for size >= 1MB allocations).
|
||||
- ``small_pool``: statistics for the small allocation pool (for size < 1MB allocations).
|
||||
|
||||
Metric type:
|
||||
|
||||
- ``current``: current value of this metric.
|
||||
- ``peak``: maximum value of this metric.
|
||||
- ``allocated``: historical total increase in this metric.
|
||||
- ``freed``: historical total decrease in this metric.
|
||||
|
||||
Args:
|
||||
device (torch.device or int or str, optional): selected device. Returns
|
||||
statistics for the current device, given by :func:`~torch.xpu.current_device`,
|
||||
if :attr:`device` is ``None`` (default).
|
||||
"""
|
||||
result = []
|
||||
|
||||
def _recurse_add_to_result(prefix: str, obj: Any) -> None:
|
||||
if isinstance(obj, dict):
|
||||
if len(prefix) > 0:
|
||||
prefix += "."
|
||||
for k, v in obj.items():
|
||||
_recurse_add_to_result(prefix + k, v)
|
||||
else:
|
||||
result.append((prefix, obj))
|
||||
|
||||
stats = memory_stats_as_nested_dict(device=device)
|
||||
_recurse_add_to_result("", stats)
|
||||
result.sort()
|
||||
|
||||
return collections.OrderedDict(result)
|
||||
|
||||
|
||||
def memory_allocated(device: Device = None) -> int:
|
||||
r"""Return the current GPU memory occupied by tensors in bytes for a given device.
|
||||
|
||||
Args:
|
||||
device (torch.device or int or str, optional): selected device. Returns
|
||||
statistic for the current device, given by :func:`~torch.xpu.current_device`,
|
||||
if :attr:`device` is ``None`` (default).
|
||||
|
||||
.. note::
|
||||
This is likely less than the amount shown in `xpu-smi` since some
|
||||
unused memory can be held by the caching allocator and some context
|
||||
needs to be created on GPU.
|
||||
"""
|
||||
return memory_stats(device=device).get("allocated_bytes.all.current", 0)
|
||||
|
||||
|
||||
def max_memory_allocated(device: Device = None) -> int:
|
||||
r"""Return the maximum GPU memory occupied by tensors in bytes for a given device.
|
||||
|
||||
By default, this returns the peak allocated memory since the beginning of
|
||||
this program. :func:`~torch.xpu.reset_peak_memory_stats` can be used to
|
||||
reset the starting point in tracking this metric. For example, these two
|
||||
functions can measure the peak allocated memory usage of each iteration in a
|
||||
training loop.
|
||||
|
||||
Args:
|
||||
device (torch.device or int or str, optional): selected device. Returns
|
||||
statistic for the current device, given by :func:`~torch.xpu.current_device`,
|
||||
if :attr:`device` is ``None`` (default).
|
||||
"""
|
||||
return memory_stats(device=device).get("allocated_bytes.all.peak", 0)
|
||||
|
||||
|
||||
def memory_reserved(device: Device = None) -> int:
|
||||
r"""Return the current GPU memory managed by the caching allocator in bytes for a given device.
|
||||
|
||||
Args:
|
||||
device (torch.device or int or str, optional): selected device. Returns
|
||||
statistic for the current device, given by :func:`~torch.xpu.current_device`,
|
||||
if :attr:`device` is ``None`` (default).
|
||||
"""
|
||||
return memory_stats(device=device).get("reserved_bytes.all.current", 0)
|
||||
|
||||
|
||||
def max_memory_reserved(device: Device = None) -> int:
|
||||
r"""Return the maximum GPU memory managed by the caching allocator in bytes for a given device.
|
||||
|
||||
By default, this returns the peak cached memory since the beginning of this
|
||||
program. :func:`~torch.xpu.reset_peak_memory_stats` can be used to reset
|
||||
the starting point in tracking this metric. For example, these two functions
|
||||
can measure the peak cached memory amount of each iteration in a training
|
||||
loop.
|
||||
|
||||
Args:
|
||||
device (torch.device or int or str, optional): selected device. Returns
|
||||
statistic for the current device, given by :func:`~torch.xpu.current_device`,
|
||||
if :attr:`device` is ``None`` (default).
|
||||
"""
|
||||
return memory_stats(device=device).get("reserved_bytes.all.peak", 0)
|
||||
|
||||
|
||||
def mem_get_info(device: Device = None) -> tuple[int, int]:
|
||||
r"""Return the global free and total GPU memory for a given device.
|
||||
|
||||
Args:
|
||||
device (torch.device or int or str, optional): selected device. Returns
|
||||
statistic for the current device, given by :func:`~torch.xpu.current_device`,
|
||||
if :attr:`device` is ``None`` (default).
|
||||
|
||||
Returns:
|
||||
tuple[int, int]: a tuple of two integers (free_memory, total_memory) in bytes.
|
||||
The first value is the free memory on the device (available across all processes and applications),
|
||||
The second value is the device's total hardware memory capacity.
|
||||
"""
|
||||
_lazy_init()
|
||||
device = _get_device_index(device, optional=True)
|
||||
return torch._C._xpu_getMemoryInfo(device)
|
||||
|
||||
|
||||
def get_per_process_memory_fraction(device: Device = None) -> float:
|
||||
r"""
|
||||
Retrieve the memory fraction currently set for a process on a given XPU device.
|
||||
This fraction represents the portion of the total device memory that
|
||||
the caching allocator is allowed to use. The allowed memory is calculated as:
|
||||
|
||||
.. math:: \text{allowed\_memory} = \text{total\_memory} \times \text{fraction}
|
||||
|
||||
Args:
|
||||
device (torch.device or int or str, optional): selected device. It uses the current device,
|
||||
given by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` (default).
|
||||
|
||||
Returns:
|
||||
float: The memory fraction in the range 0.0 to 1.0.
|
||||
"""
|
||||
_lazy_init()
|
||||
device = _get_device_index(device, optional=True)
|
||||
return torch._C._xpu_getMemoryFraction(device)
|
||||
|
||||
|
||||
def set_per_process_memory_fraction(fraction: float, device: Device = None) -> None:
|
||||
r"""
|
||||
Set the memory fraction for a single process on XPU device.
|
||||
This function limits the amount of memory that the caching allocator can allocate
|
||||
on the specified XPU device. The allowed memory is computed as:
|
||||
|
||||
.. math:: \text{allowed\_memory} = \text{total\_memory} \times \text{fraction}
|
||||
|
||||
If the process attempts to allocate more than this allowed memory,
|
||||
an out-of-memory error will be raised by the allocator.
|
||||
|
||||
Arguments:
|
||||
fraction (float): Range: 0~1. Allowed memory equals total_memory * fraction.
|
||||
device (torch.device or int or str, optional): selected device. It uses the current device,
|
||||
given by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` (default).
|
||||
|
||||
.. note:: In general, the total available free memory is less than the total capacity.
|
||||
"""
|
||||
_lazy_init()
|
||||
device = _get_device_index(device, optional=True)
|
||||
if not isinstance(fraction, float):
|
||||
raise TypeError("Invalid type for fraction argument, must be `float`")
|
||||
torch._C._xpu_setMemoryFraction(fraction, device)
|
||||
|
||||
|
||||
def memory_snapshot(
|
||||
mempool_id: tuple[int, int] | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
r"""
|
||||
Return a snapshot of the XPU memory allocator state across all devices.
|
||||
Provides detailed information for each memory segment managed by the allocator
|
||||
including its size, owning pool, associated stream, call stack traces, and other relevant attributes.
|
||||
|
||||
Arguments:
|
||||
mempool_id (tuple[int, int] or None, optional): The memory pool id. If None, the default memory pool is used.
|
||||
|
||||
Returns:
|
||||
list[dict[str, Any]]: List of memory segments and their attributes.
|
||||
"""
|
||||
if not is_initialized():
|
||||
return []
|
||||
return torch._C._xpu_memorySnapshot(mempool_id)["segments"]
|
||||
|
||||
|
||||
def _snapshot(device: Device = None, augment_with_fx_traces: bool = False):
|
||||
"""
|
||||
Capture a snapshot of the XPU memory state at the time this function is called.
|
||||
|
||||
The returned snapshot is a dictionary with the following structure.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
class Snapshot(TypedDict):
|
||||
segments: List[Segment]
|
||||
device_traces: List[List[TraceEntry]]
|
||||
|
||||
|
||||
class Segment(TypedDict):
|
||||
# A Segment represents a contiguous memory region returned by the SYCL runtime.
|
||||
#
|
||||
# All reserved memory is composed of these segments. Segments are
|
||||
# cached and reused by the allocator. When allocations are smaller
|
||||
# than the segment, the segment may be split into multiple Blocks.
|
||||
#
|
||||
# Calling :func:`~torch.xpu.memory.empty_cache` releases segments that are entirely inactive.
|
||||
address: int
|
||||
total_size: int # total size of segment
|
||||
stream: int
|
||||
segment_type: Literal["small", "large"] # 'large' (>1MB)
|
||||
allocated_size: int # size of memory in use
|
||||
active_size: int # size of memory in use or in active_awaiting_free state
|
||||
blocks: List[Block]
|
||||
|
||||
|
||||
class Block(TypedDict):
|
||||
# A sub-region of a Segment, either currently allocated or cached for reuse.
|
||||
size: int
|
||||
requested_size: int # Original requested size (may be smaller than `size`)
|
||||
address: int
|
||||
state: Literal[
|
||||
"active_allocated", # used by a tensor
|
||||
"active_awaiting_free", # waiting for another stream synchronization, then become free
|
||||
"inactive", # free for reuse
|
||||
]
|
||||
frames: List[Frame] # stack trace from where the allocation occurred
|
||||
|
||||
|
||||
class Frame(TypedDict):
|
||||
filename: str
|
||||
line: int
|
||||
name: str
|
||||
# Optional fields when `augment_with_fx_traces=True` and the frame
|
||||
# corresponds to FX-generated code.
|
||||
fx_node_op: str # FX node operation type (e.g., 'call_function', 'output')
|
||||
fx_node_name: str # FX node name (e.g., 'linear', 'relu_1')
|
||||
fx_original_trace: str # Original model source code stack trace
|
||||
|
||||
|
||||
class TraceEntry(TypedDict):
|
||||
# Trace entries are recorded only when :func:`~torch.xpu.memory._record_memory_history` is enabled.
|
||||
action: Literal[
|
||||
"alloc" # memory allocated
|
||||
"free_requested", # received a call to free memory
|
||||
"free_completed", # memory reclaimed and reusable
|
||||
"segment_alloc", # ask SYCL runtime for more memory
|
||||
"segment_free", # called SYCL runtime to return memory to XPU
|
||||
"segment_map", # ask SYCL runtime to map memory
|
||||
"segment_unmap", # called SYCL runtime to unmap memory
|
||||
"snapshot", # snapshot taken
|
||||
"oom", # threw an OOM exception
|
||||
]
|
||||
addr: int # not present for OOM
|
||||
frames: List[Frame]
|
||||
size: int
|
||||
stream: int
|
||||
device_free: int # only present for OOM, the amount of free memory reported by the device
|
||||
|
||||
Arguments:
|
||||
device (torch.device or int or str, optional): selected device. It uses the current device,
|
||||
given by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` (default).
|
||||
augment_with_fx_traces (bool, optional): If True, augment stack trace frames with FX debug information
|
||||
that maps generated FX code back to original model source code. This adds the FX-related
|
||||
fields (fx_node_op, fx_node_name, fx_original_trace) to Frame objects. Default is ``False``.
|
||||
|
||||
Returns:
|
||||
The Snapshot dictionary object
|
||||
"""
|
||||
s = torch._C._xpu_memorySnapshot(None)
|
||||
if augment_with_fx_traces:
|
||||
s = _augment_memory_snapshot_stack_traces(s) # type: ignore[assignment, arg-type]
|
||||
return s
|
||||
|
||||
|
||||
def _dump_snapshot(
|
||||
filename: str = "dump_snapshot.pickle", augment_with_fx_traces: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
Save a pickled version of the `torch.memory._snapshot()` dictionary to a file.
|
||||
|
||||
This file can be opened by the interactive snapshot viewer at pytorch.org/memory_viz
|
||||
|
||||
Snapshot file sizes scale with `max_entries` and stack trace depth per entry,
|
||||
with several KB per entry. These can easily be in the GB range for longer running
|
||||
workflows with large `max_entries`.
|
||||
|
||||
Arguments:
|
||||
filename (str, optional): Name of the file to create. Defaults to "dump_snapshot.pickle".
|
||||
augment_with_fx_traces (bool, optional): If True, augment the snapshot with FX debug information
|
||||
before dumping. This maps generated FX code stack traces back to original model
|
||||
source code. Defaults to ``False``.
|
||||
"""
|
||||
s = _snapshot(augment_with_fx_traces=augment_with_fx_traces)
|
||||
|
||||
with open(filename, "wb") as f:
|
||||
pickle.dump(s, f)
|
||||
|
||||
|
||||
def _record_memory_history(
|
||||
enabled: Literal["state", "all"] | None = "all",
|
||||
context: Literal["state", "alloc", "all"] | None = "all",
|
||||
stacks: Literal["python", "all"] = "all",
|
||||
max_entries: int = sys.maxsize,
|
||||
clear_history: bool = False,
|
||||
skip_actions: list[str] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Enable recording of stack traces associated with memory allocations, so you can
|
||||
tell what allocated any piece of memory in :func:`~torch.xpu.memory._snapshot()`.
|
||||
|
||||
In addition to keeping stack traces with each current allocation and free,
|
||||
this will also enable recording of a history of all alloc/free events.
|
||||
|
||||
Use :func:`~torch.xpu.memory._snapshot()` to retrieve this information,
|
||||
and the tools in `_memory_viz.py` to visualize snapshots.
|
||||
|
||||
Buffer behavior
|
||||
---------------
|
||||
|
||||
This will store up to `max_entries` instances of `TraceEntry` when enabled.
|
||||
Python trace collection defaults to `sys.maxsize`, meaning long-running
|
||||
or indefinitely running jobs should set a reasonable limit to avoid excessive
|
||||
memory use. Expect each entry to be several KB.
|
||||
|
||||
Longer running workflows or those with smaller `max_entries` values will only
|
||||
store the last accumulated `max_entries` entries, meaning new entries overwrite
|
||||
older entries, reference to ring buffer behavior.
|
||||
|
||||
Latency impact
|
||||
--------------
|
||||
|
||||
The Python trace collection is fast (2us per trace), so you may consider
|
||||
enabling this on production jobs if you anticipate ever having to debug
|
||||
memory issues.
|
||||
|
||||
C++ trace collection is also fast (~50ns/frame), which for many typical programs
|
||||
works out to ~2us per trace, but can vary depending on stack depth.
|
||||
|
||||
Arguments:
|
||||
enabled (Literal["state", "all"], optional):
|
||||
`None`, disable recording memory history.
|
||||
`"state"`, keep information for currently allocated memory.
|
||||
`"all"`, additionally keep a history of all alloc/free calls.
|
||||
Defaults to "all".
|
||||
context (Literal["state", "alloc", "all"], optional):
|
||||
`None`, Do not record any tracebacks.
|
||||
`"state"`, Record tracebacks for currently allocated memory.
|
||||
`"alloc"`, additionally keep tracebacks for alloc calls.
|
||||
`"all"`, additionally keep tracebacks for free calls.
|
||||
Defaults to "all".
|
||||
stacks (Literal["python", "all"], optional):
|
||||
`"python"`, include Python, TorchScript, and inductor frames in tracebacks.
|
||||
`"all"`, additionally include C++ frames.
|
||||
Defaults to "all".
|
||||
max_entries (int, optional): Keep a maximum of `max_entries`
|
||||
alloc/free events in the recorded history recorded.
|
||||
clear_history (bool, optional): Clear history when enabling, defaults to ``False``.
|
||||
skip_actions (list[str], optional): List of action types to skip when recording
|
||||
memory history. This can be used to reduce memory overhead by excluding
|
||||
certain types of events from being recorded. Valid action types are:
|
||||
|
||||
- `"alloc"`: Memory allocation events
|
||||
- `"free_requested"`: Free requests (memory marked for freeing)
|
||||
- `"free_completed"`: Completed free operations (memory actually freed)
|
||||
- `"segment_alloc"`: Segment allocation from SYCL runtime
|
||||
- `"segment_free"`: Segment freed back to XPU via SYCL runtime
|
||||
- `"segment_map"`: Segment map events
|
||||
- `"segment_unmap"`: Segment unmap events
|
||||
- `"snapshot"`: Memory snapshot generation events
|
||||
- `"oom"`: Out-of-memory exceptions
|
||||
|
||||
For example, to skip recording free_requested events:
|
||||
`skip_actions=["free_requested"]`
|
||||
|
||||
Defaults to ``None`` (record all actions).
|
||||
"""
|
||||
torch._C._xpu_recordMemoryHistory(
|
||||
enabled,
|
||||
context,
|
||||
stacks,
|
||||
max_entries,
|
||||
clear_history,
|
||||
skip_actions if skip_actions is not None else [],
|
||||
)
|
||||
|
||||
|
||||
class _XPUAllocator:
|
||||
r"""Wrapper over internal XPU memory allocators."""
|
||||
|
||||
def __init__(self, allocator: torch._C._xpu_XPUAllocator):
|
||||
self._allocator = allocator
|
||||
|
||||
def allocator(self):
|
||||
return self._allocator
|
||||
|
||||
|
||||
class XPUPluggableAllocator(_XPUAllocator):
|
||||
r"""
|
||||
XPU memory allocator loaded dynamically from a shared library.
|
||||
|
||||
This lets users provide custom allocation and free functions implemented
|
||||
in a separate shared library. The allocator is registered and could become
|
||||
available for use via :func:`~torch.xpu.memory.change_current_allocator`.
|
||||
|
||||
Arguments:
|
||||
path_to_lib_file (str):
|
||||
Filesystem path to the shared library file containing the allocation
|
||||
and free functions.
|
||||
alloc_fn_name (str):
|
||||
Name of the allocation function exported from the shared library.
|
||||
The function must have the signature:
|
||||
|
||||
``void* alloc_fn(size_t size, int device, sycl::queue* queue);``
|
||||
|
||||
free_fn_name (str):
|
||||
Name of the free function exported from the shared library.
|
||||
The function must have the signature:
|
||||
|
||||
``void free_fn(void* ptr, size_t size, int device, sycl::queue* queue);``
|
||||
"""
|
||||
|
||||
def __init__(self, path_to_lib_file: str, alloc_fn_name: str, free_fn_name: str):
|
||||
allocator_lib = ctypes.CDLL(path_to_lib_file)
|
||||
|
||||
alloc_fn_ptr = getattr(allocator_lib, alloc_fn_name)
|
||||
free_fn_ptr = getattr(allocator_lib, free_fn_name)
|
||||
|
||||
alloc_fn_addr = ctypes.cast(alloc_fn_ptr, ctypes.c_void_p).value
|
||||
free_fn_addr = ctypes.cast(free_fn_ptr, ctypes.c_void_p).value
|
||||
|
||||
if alloc_fn_addr is None or free_fn_addr is None:
|
||||
raise RuntimeError(
|
||||
"Failed to load allocator symbols from the shared library."
|
||||
)
|
||||
|
||||
self._allocator = torch._C._xpu_customAllocator(alloc_fn_addr, free_fn_addr)
|
||||
|
||||
|
||||
def change_current_allocator(allocator: _XPUAllocator) -> None:
|
||||
r"""Change the currently used memory allocator to be the one provided.
|
||||
|
||||
.. note::
|
||||
If the current allocator has already been used/initialized, this function will error.
|
||||
|
||||
Arguments:
|
||||
allocator (torch.xpu.memory._XPUAllocator): allocator to be set as the active one.
|
||||
"""
|
||||
torch._C._xpu_changeCurrentAllocator(allocator.allocator())
|
||||
|
||||
|
||||
def _get_current_allocator() -> _XPUAllocator:
|
||||
r"""Return the allocator being currently used.
|
||||
|
||||
Returns:
|
||||
_XPUAllocator: the allocator being currently used.
|
||||
"""
|
||||
return _XPUAllocator(torch._C._xpu_getAllocator())
|
||||
|
||||
|
||||
class MemPool(torch._C._XPUMemPool):
|
||||
r"""MemPool represents a pool of memory in a caching allocator. Currently,
|
||||
it's just the ID of the pool object maintained in the XPUCachingAllocator.
|
||||
|
||||
Args:
|
||||
allocator(torch._C._xpu_XPUAllocator, optional): a
|
||||
torch._C._xpu_XPUAllocator object that can be used to
|
||||
define how memory gets allocated in the pool. If :attr:`allocator`
|
||||
is ``None`` (default), memory allocation follows the default/
|
||||
current configuration of the XPUCachingAllocator.
|
||||
use_on_oom(bool): a bool that indicates if this pool can be used
|
||||
as a last resort if a memory allocation outside of the pool fails due
|
||||
to Out Of Memory. This is ``False`` by default.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
allocator: torch._C._xpu_XPUAllocator | None = None,
|
||||
use_on_oom: bool = False,
|
||||
):
|
||||
super().__init__(allocator, True, use_on_oom)
|
||||
|
||||
@property
|
||||
def id(self) -> tuple[int, int]:
|
||||
r"""Returns the ID of this pool as a tuple of two ints."""
|
||||
return super().id
|
||||
|
||||
@property
|
||||
def allocator(self) -> torch._C._xpu_XPUAllocator | None:
|
||||
r"""Returns the allocator this MemPool routes allocations to."""
|
||||
return super().allocator
|
||||
|
||||
def use_count(self) -> int:
|
||||
r"""Returns the reference count of this pool."""
|
||||
return super().use_count()
|
||||
|
||||
def snapshot(self):
|
||||
r"""Return a snapshot of the XPU memory allocator pool state across all
|
||||
devices.
|
||||
|
||||
Interpreting the output of this function requires familiarity with the
|
||||
memory allocator internals.
|
||||
"""
|
||||
snapshot = torch.xpu.memory_snapshot(self.id)
|
||||
return snapshot
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def use_mem_pool(pool: MemPool, device: "Device" = None):
|
||||
r"""A context manager that routes allocations to a given pool.
|
||||
|
||||
Args:
|
||||
pool(torch.xpu.MemPool): a :class:`MemPool` object to be made active so that
|
||||
allocations route to this pool.
|
||||
device (torch.device or int, optional): selected device. Uses :class:`MemPool on
|
||||
the current device, given by :func:`~torch.xpu.current_device`,
|
||||
if :attr:`device` is ``None`` (default).
|
||||
|
||||
.. note::
|
||||
This context manager makes only current thread's allocations route to
|
||||
the given pool. If a new thread is spawned inside the context manager
|
||||
(e.g. by calling backward) the allocations in that thread will not
|
||||
route to the given pool.
|
||||
"""
|
||||
device_index = (
|
||||
torch.xpu.current_device() if device is None else _get_device_index(device)
|
||||
)
|
||||
torch._C._xpu_beginAllocateCurrentThreadToPool(device_index, pool.id)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch._C._xpu_endAllocateToPool(device_index, pool.id)
|
||||
torch._C._xpu_releasePool(device_index, pool.id)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MemPool",
|
||||
"XPUPluggableAllocator",
|
||||
"change_current_allocator",
|
||||
"empty_cache",
|
||||
"get_per_process_memory_fraction",
|
||||
"max_memory_allocated",
|
||||
"max_memory_reserved",
|
||||
"mem_get_info",
|
||||
"memory_allocated",
|
||||
"memory_reserved",
|
||||
"memory_snapshot",
|
||||
"memory_stats",
|
||||
"memory_stats_as_nested_dict",
|
||||
"reset_accumulated_memory_stats",
|
||||
"reset_peak_memory_stats",
|
||||
"set_per_process_memory_fraction",
|
||||
"use_mem_pool",
|
||||
]
|
||||
Reference in New Issue
Block a user