""" A pytorch memory cache that can be allocated by ConnectionHandler (on cpu) and used over multiple calls to Runtime. For now, the only purpose of this code is to ensure that allocated memory will be deleted properly. """ import asyncio import contextlib import ctypes import multiprocessing as mp import os import time from typing import AsyncContextManager, Dict, Optional, Sequence import async_timeout import torch from hivemind.utils import TensorDescriptor, enter_asynchronously, get_logger from petals.data_structures import Handle from petals.utils.asyncio import shield_and_wait from petals.utils.misc import get_size_in_bytes logger = get_logger(__name__) class MemoryCache: """A shared cache for storing tensors that persist across calls. Main use case: storing past attention KVs""" def __init__(self, max_size_bytes: Optional[int], max_alloc_timeout: Optional[float] = None): self.max_size_bytes = max_size_bytes if max_size_bytes is not None else (2**64 - 1) self.max_alloc_timeout = max_alloc_timeout self._lock_metadata = mp.Lock() self._current_size = mp.Value(ctypes.c_int64, 0, lock=False) self._enqueued_size = mp.Value(ctypes.c_int64, 0, lock=True) self._handle_counter = mp.Value(ctypes.c_int64, 0, lock=False) self._allocated_tensors: Dict[Handle, torch.Tensor] = {} self.runtime_pid = os.getpid() self._pipe_recv, self._pipe_send = mp.Pipe(duplex=False) # any ConnectionHandler -> runtime self._lock_acquire_memory = mp.Lock() self._memory_freed_event = mp.Event() @property def current_size_bytes(self) -> int: return self._current_size.value @current_size_bytes.setter def current_size_bytes(self, value: int): self._current_size.value = value @property def enqueued_size_bytes(self) -> int: return self._enqueued_size.value @enqueued_size_bytes.setter def enqueued_size_bytes(self, value: int): self._enqueued_size.value = value @property def bytes_left(self) -> int: return self.max_size_bytes - self.current_size_bytes @property def handle_counter(self) -> int: return self._handle_counter.value @handle_counter.setter def handle_counter(self, value: int): self._handle_counter.value = value @contextlib.asynccontextmanager async def allocate_cache( self, *descriptors: TensorDescriptor, timeout: float ) -> AsyncContextManager[Sequence[Handle]]: """ Create a handle that is associated with buffers on unique device. If cache full, raises AllocationFailed. :param descriptors: one or more tensors tensor of this size, dtype, etc :param timeout: optional maximum time to wait for cache allocation; None (default) means no time limit :note: if descriptors reside on different devices, it is expected that they are approximately balanced across devices; if not, it will count maximum tensor allocation across devices for the purposes of size limit :note: This function should be called by connection handlers, it can be called concurrently from multiple processes. Furthermore, it can be called concurrently with at most one use_cache call in runtime. """ assert os.getpid() != self.runtime_pid, "must be called by a ConnectionHandler, not runtime" assert all(descr.device is not None for descr in descriptors), "please specify allocated devices" if self.max_alloc_timeout is not None: timeout = min(timeout, self.max_alloc_timeout) max_alloc_size = self.get_allocation_size(*descriptors) gib = 1024**3 cur_size, max_size = self.current_size_bytes, self.max_size_bytes friendly_max_size = f"{max_size / gib:.2f}" if max_size != 2**64 - 1 else "inf" logger.info( f"rpc_inference.wait_for_alloc(size={max_alloc_size / gib:.2f} GiB), " f"already used {cur_size / gib:.2f}/{friendly_max_size} GiB ({cur_size / max_size * 100:.1f}%)" ) alloc_task = asyncio.create_task(self._schedule_alloc(max_alloc_size, *descriptors, timeout=timeout)) try: handles = await shield_and_wait(alloc_task) logger.info(f"rpc_inference.alloc_done(size={max_alloc_size / gib:.2f} GiB)") yield handles finally: self._free(max_alloc_size, alloc_task) @staticmethod def get_allocation_size(*descriptors: TensorDescriptor) -> int: """Return the memory size (bytes) to be allocated on a device. If there are many devices, return maximum""" alloc_size_by_device = {} for descr in descriptors: tensor_size = descr.numel() * get_size_in_bytes(descr.dtype) alloc_size_by_device[descr.device] = alloc_size_by_device.get(descr.device, 0) + tensor_size return max(alloc_size_by_device.values()) async def _schedule_alloc( self, alloc_size: int, *descriptors: TensorDescriptor, timeout: Optional[float] ) -> Sequence[Handle]: """ This method should be called inside asyncio.shield() because: - hivemind.utils.enter_asynchronously() does not always release the lock on cancellation """ try: async with self._wait_for_free_memory(alloc_size, timeout): with self._lock_metadata: handles = tuple(int(self.handle_counter) + i for i in range(len(descriptors))) self.current_size_bytes += alloc_size self.handle_counter += len(handles) # note: this will eventually overflow and it is okay self._pipe_send.send((handles, descriptors)) return handles except TimeoutError: raise AllocationFailed(f"Could not allocate {alloc_size} (timeout={timeout})") @contextlib.asynccontextmanager async def _wait_for_free_memory(self, alloc_size: int, timeout: Optional[float]): start_time = time.perf_counter() loop = asyncio.get_event_loop() with self._enqueued_size.get_lock(): self._enqueued_size.value += alloc_size allocated = False try: context_manager = async_timeout.timeout(timeout) if timeout != 0 else contextlib.AsyncExitStack() # contextlib.AsyncExitStack() is used as a null context here async with context_manager: if timeout == 0 and self.current_size_bytes + self.enqueued_size_bytes > self.max_size_bytes: raise AllocationFailed(f"Could not allocate {alloc_size} bytes immediately: out of memory") async with enter_asynchronously(self._lock_acquire_memory): if self.current_size_bytes + alloc_size > self.max_size_bytes: if timeout == 0: raise AllocationFailed(f"Could not allocate {alloc_size} bytes immediately: out of memory") elapsed_time = time.perf_counter() - start_time remaining_timeout = max(0.0, timeout - elapsed_time) if timeout is not None else None await loop.run_in_executor(None, self._wait_until_available, alloc_size, remaining_timeout) allocated = True with self._enqueued_size.get_lock(): self._enqueued_size.value -= alloc_size yield except asyncio.TimeoutError: raise AllocationFailed(f"Could not allocate {alloc_size} within {timeout} seconds") finally: if not allocated: with self._enqueued_size.get_lock(): self._enqueued_size.value -= alloc_size def _free(self, alloc_size: int, alloc_task: asyncio.Task): if alloc_task.exception() is not None: return handles = alloc_task.result() with self._lock_metadata: self._pipe_send.send((handles, None)) # signal runtime to free these handles self.current_size_bytes -= alloc_size self._memory_freed_event.set() def _wait_until_available(self, allocated_size: int, timeout: Optional[float] = None): # note: this function should only be called inside _lock_acquire_memory! if allocated_size > self.max_size_bytes: raise AllocationFailed( f"Could not allocate {allocated_size} bytes, max cache size = {self.max_size_bytes} bytes" ) timeout = timeout if timeout != float("inf") else None deadline = None if timeout is None else time.perf_counter() + timeout while self.current_size_bytes + allocated_size > self.max_size_bytes: remaining_time = None if timeout is None else deadline - time.perf_counter() if not self._memory_freed_event.wait(remaining_time): raise AllocationFailed( f"Server's attention cache is full, failed to allocate {allocated_size} bytes in {timeout} seconds" ) self._memory_freed_event.clear() @contextlib.contextmanager def use_cache(self, *handles: Handle) -> Sequence[torch.Tensor]: """ Return one or more tensors previously allocated with allocate_cache, :note: This method is called by ModuleBackend in runtime: a single process with NO process parallelism. However, runtime may call use_cache concurrently with one or more connection handlers calling allocate_cache """ assert os.getpid() == self.runtime_pid # note: this specific function is not concurrent, so you can safely allocate/offload/defragment data here # read creation/deletion requests from connection handlers while self._pipe_recv.poll(): recv_handles, recv_data = self._pipe_recv.recv() if recv_data is not None: # create new tensors assert len(recv_handles) == len(recv_data) for handle, descr in zip(recv_handles, recv_data): self._allocated_tensors[handle] = descr.make_zeros() assert handle in self._allocated_tensors, f"Sanity check failed: no such handle ({handle})" else: # delete tensors by handle for handle in recv_handles: if handle not in self._allocated_tensors: logger.warning( f"Sanity check failed: asked to delete handle {handle}, but there is no such handle" ) self._allocated_tensors.pop(handle, None) yield tuple(self._allocated_tensors[handle] for handle in handles) class AllocationFailed(Exception): pass