diff --git a/python/ray/experimental/rdt/nixl_memory_pool.py b/python/ray/experimental/rdt/nixl_memory_pool.py index 7b68b32a1a3..d526e56f1d0 100644 --- a/python/ray/experimental/rdt/nixl_memory_pool.py +++ b/python/ray/experimental/rdt/nixl_memory_pool.py @@ -1,6 +1,16 @@ """Memory pool management for NIXL RDT optimization.""" -from typing import TYPE_CHECKING, Dict, List, NamedTuple, Sequence, Tuple +import threading +from typing import ( + TYPE_CHECKING, + Dict, + List, + NamedTuple, + Optional, + Sequence, + Tuple, + Union, +) if TYPE_CHECKING: import torch @@ -193,6 +203,10 @@ def __init__(self, pool_size: int, device: "torch.device"): pool_size, dtype=torch.uint8, device=self.device ) + # Guards _free_blocks and _allocated_by_obj. This is the innermost lock: + # the pool never calls back into the transport, so it is always safe to + # take while holding the transport's cache lock. + self._lock = threading.Lock() # List of MemoryBlock for free blocks, sorted by offset. self._free_blocks: List[MemoryBlock] = [MemoryBlock(offset=0, size=pool_size)] # Blocks allocated per object ID. @@ -235,91 +249,94 @@ def allocate_group( ] sizes = [layout.nbytes for layout in tensor_layouts] - # Snapshot the free list so the whole group is atomic. If this obj_id - # already owns blocks (re-extract), treat them as free for packing so - # the new allocation can reuse that space; on failure the real state - # is left untouched. - temp_free = [MemoryBlock(b.offset, b.size) for b in self._free_blocks] - prior = self._allocated_by_obj.get(obj_id) - if prior: - temp_free.extend(MemoryBlock(b.offset, b.size) for b in prior) - _merge_free_blocks(temp_free) - - if sum(b.size for b in temp_free) < sum(sizes): - raise NixlOutOfMemoryError( - f"NIXL memory pool out of memory: cannot allocate {len(sizes)} " - f"tensor(s) totaling {sum(sizes)} bytes. Consider increasing the " - f"pool size when calling register_nixl_memory_pool." - ) + with self._lock: + # Snapshot the free list so the whole group is atomic. If this obj_id + # already owns blocks (re-extract), treat them as free for packing so + # the new allocation can reuse that space; on failure the real state + # is left untouched. + temp_free = [MemoryBlock(b.offset, b.size) for b in self._free_blocks] + prior = self._allocated_by_obj.get(obj_id) + if prior: + temp_free.extend(MemoryBlock(b.offset, b.size) for b in prior) + _merge_free_blocks(temp_free) + + if sum(b.size for b in temp_free) < sum(sizes): + raise NixlOutOfMemoryError( + f"NIXL memory pool out of memory: cannot allocate {len(sizes)} " + f"tensor(s) totaling {sum(sizes)} bytes. Consider increasing the " + f"pool size when calling register_nixl_memory_pool." + ) - blocks: List[MemoryBlock] = [] - # Bytes actually packed into each block, and the absolute pool offset of - # each tensor. Tensors are taken in order, so pool_starts ends up in - # tensor order. - block_nbytes: List[int] = [] - pool_starts: List[int] = [] - remaining = list(range(len(tensors))) - - while remaining: - rem_layouts = [tensor_layouts[i] for i in remaining] - offsets, total_nbytes = packed_offsets(rem_layouts) - - # Prefer the smallest free block that fits everything remaining. - free_idx = min( - (i for i, b in enumerate(temp_free) if b.size >= total_nbytes), - key=lambda i: temp_free[i].size, - default=None, - ) - if free_idx is not None: - take_count = len(remaining) - placed_nbytes = total_nbytes - else: - # Take the largest free block and pack as many as fit in order. - free_idx = max( - range(len(temp_free)), + blocks: List[MemoryBlock] = [] + # Bytes actually packed into each block, and the absolute pool offset + # of each tensor. Tensors are taken in order, so pool_starts ends up + # in tensor order. + block_nbytes: List[int] = [] + pool_starts: List[int] = [] + remaining = list(range(len(tensors))) + + while remaining: + rem_layouts = [tensor_layouts[i] for i in remaining] + offsets, total_nbytes = packed_offsets(rem_layouts) + + # Prefer the smallest free block that fits everything remaining. + free_idx = min( + (i for i, b in enumerate(temp_free) if b.size >= total_nbytes), key=lambda i: temp_free[i].size, default=None, ) - hole_size = 0 if free_idx is None else temp_free[free_idx].size - take_count = 0 - placed_nbytes = 0 - for n, layout in enumerate(rem_layouts): - if offsets[n] + layout.nbytes > hole_size: - break - take_count = n + 1 - placed_nbytes = offsets[n] + layout.nbytes - if take_count == 0: - raise NixlOutOfMemoryError( - f"NIXL memory pool out of memory: cannot allocate next " - f"tensor of {rem_layouts[0].nbytes} bytes (largest free " - f"block is {hole_size} bytes). Consider increasing the " - f"pool size when calling register_nixl_memory_pool." + if free_idx is not None: + take_count = len(remaining) + placed_nbytes = total_nbytes + else: + # Take the largest free block and pack as many as fit in order. + free_idx = max( + range(len(temp_free)), + key=lambda i: temp_free[i].size, + default=None, ) - - # Round the carved block up so subsequent offsets stay aligned. - free_block = temp_free[free_idx] - block_offset = free_block.offset - carved = min(_align_up(placed_nbytes, _MAX_ALIGNMENT), free_block.size) - if carved == free_block.size: - temp_free.pop(free_idx) - else: - free_block.offset += carved - free_block.size -= carved - - blocks.append(MemoryBlock(block_offset, carved)) - block_nbytes.append(placed_nbytes) - pool_starts.extend(block_offset + off for off in offsets[:take_count]) - remaining = remaining[take_count:] - - # Commit only after the full group packs successfully. - temp_free.sort(key=lambda b: b.offset) - self._free_blocks = temp_free - self._allocated_by_obj[obj_id] = blocks + hole_size = 0 if free_idx is None else temp_free[free_idx].size + take_count = 0 + placed_nbytes = 0 + for n, layout in enumerate(rem_layouts): + if offsets[n] + layout.nbytes > hole_size: + break + take_count = n + 1 + placed_nbytes = offsets[n] + layout.nbytes + if take_count == 0: + raise NixlOutOfMemoryError( + f"NIXL memory pool out of memory: cannot allocate next " + f"tensor of {rem_layouts[0].nbytes} bytes (largest free " + f"block is {hole_size} bytes). Consider increasing the " + f"pool size when calling register_nixl_memory_pool." + ) + + # Round the carved block up so subsequent offsets stay aligned. + free_block = temp_free[free_idx] + block_offset = free_block.offset + carved = min(_align_up(placed_nbytes, _MAX_ALIGNMENT), free_block.size) + if carved == free_block.size: + temp_free.pop(free_idx) + else: + free_block.offset += carved + free_block.size -= carved + + blocks.append(MemoryBlock(block_offset, carved)) + block_nbytes.append(placed_nbytes) + pool_starts.extend(block_offset + off for off in offsets[:take_count]) + remaining = remaining[take_count:] + + # Commit only after the full group packs successfully. + temp_free.sort(key=lambda b: b.offset) + self._free_blocks = temp_free + self._allocated_by_obj[obj_id] = blocks regions = [ self._pool_tensor[b.offset : b.offset + nbytes] for b, nbytes in zip(blocks, block_nbytes) ] + # Safe outside the lock: obj_id owns these blocks now, so no other + # thread can hand the same bytes out while they are being written. self._copy_into_pool(tensors, sizes, pool_starts) return regions @@ -344,6 +361,142 @@ def _copy_into_pool( src_bytes = tensor.flatten().view(torch.uint8) self._pool_tensor[pool_start : pool_start + nbytes].copy_(src_bytes) + def allocate_regions( + self, + region_nbytes: Sequence[int], + ) -> Tuple[List["torch.Tensor"], List[MemoryBlock]]: + """Carve one contiguous pool region per requested byte count. + + This is the receive-side counterpart to ``allocate_group``: the sender + packs tensors it already has, while the receiver only knows how many + bytes each incoming NIXL descriptor carries. One region per descriptor + keeps the read contiguous, and the caller lays the individual tensors + out inside a region with ``packed_offsets``, which works because blocks + start on ``_MAX_ALIGNMENT``. + + Unlike ``allocate_group`` this copies nothing and takes no ``obj_id``. + Receive buffers live for exactly one transfer, so the caller owns the + returned blocks and returns them with ``copy_out_and_free`` once the + transfer lands, or ``free_blocks`` if it fails. + + Args: + region_nbytes: Byte count of each region, in descriptor order. + + Returns: + (regions, blocks) in the requested order. Each region is a ``uint8`` + view sized exactly as requested; its block may be slightly larger + because of alignment. + + Raises: + NixlOutOfMemoryError: If the pool has insufficient space. + """ + with self._lock: + # Snapshot the free list so the whole group is atomic: a later region + # that does not fit leaves the real free list untouched. + temp_free = [MemoryBlock(b.offset, b.size) for b in self._free_blocks] + blocks: List[MemoryBlock] = [] + + for nbytes in region_nbytes: + # Prefer the smallest block that fits so large holes stay intact + # for whichever region needs them. + free_idx = min( + (i for i, b in enumerate(temp_free) if b.size >= nbytes), + key=lambda i: temp_free[i].size, + default=None, + ) + if free_idx is None: + largest = max((b.size for b in temp_free), default=0) + raise NixlOutOfMemoryError( + f"NIXL memory pool out of memory: cannot allocate a " + f"contiguous receive buffer of {nbytes} bytes (largest free " + f"block is {largest} bytes). Consider increasing the pool " + f"size when calling register_nixl_memory_pool." + ) + + free_block = temp_free[free_idx] + block_offset = free_block.offset + # Round up so the next block still starts aligned, but never past + # the end of the hole we are carving from. + carved = min(_align_up(nbytes, _MAX_ALIGNMENT), free_block.size) + if carved == free_block.size: + temp_free.pop(free_idx) + else: + free_block.offset += carved + free_block.size -= carved + blocks.append(MemoryBlock(block_offset, carved)) + + # Commit only after every region has been placed. + temp_free.sort(key=lambda b: b.offset) + self._free_blocks = temp_free + + regions = [ + self._pool_tensor[b.offset : b.offset + nbytes] + for b, nbytes in zip(blocks, region_nbytes) + ] + return regions, blocks + + def copy_out_and_free( + self, + tensors: List["torch.Tensor"], + blocks: List[MemoryBlock], + target_device: Optional[Union[str, "torch.device"]] = None, + ) -> List["torch.Tensor"]: + """Copy pool-backed tensors into independent tensors and free their blocks. + + This decouples the returned tensors from the pool, so their lifetime is + no longer tied to the pool's free list. Callers can hand the copies to + user code and the blocks are immediately reusable by the next transfer. + + Args: + tensors: Views into pool regions from ``allocate_regions``. + blocks: The blocks backing those views. + target_device: Device the copies should land on. Defaults to the + pool's own device. Staging through a pool on a different device + costs nothing extra, since the copy out happens either way. + + Returns: + One independently allocated tensor per input, in the same order. + """ + import torch + + device = self.device if target_device is None else target_device + # The copies run without the pool lock. The blocks stay allocated until + # the free below, so no other thread can hand these bytes out while they + # are being read, and the device sync does not block every other + # allocation behind it. + try: + # copy=True because .to() is a no-op when the device already + # matches, which would keep the result aliasing the pool block we + # are about to hand back. + # TODO(#65828): Allow a user to specify a stream for the copies. + copies = [tensor.to(device, copy=True) for tensor in tensors] + finally: + try: + if self.device.type == "cuda": + # TODO(#65829): Synchronize lazily. The copy only has to + # finish before the next NIXL transfer writes into the + # block, not before this returns. + torch.cuda.synchronize(self.device) + finally: + self.free_blocks(blocks) + return copies + + def free_blocks(self, blocks: List[MemoryBlock]) -> None: + """Return blocks from ``allocate_regions`` to the free list. + + Args: + blocks: Memory blocks to free. An empty list is a no-op. + """ + if not blocks: + return + with self._lock: + self._free_blocks_locked(blocks) + + def _free_blocks_locked(self, blocks: List[MemoryBlock]) -> None: + """Return blocks to the free list. Caller must hold ``_lock``.""" + self._free_blocks.extend(blocks) + _merge_free_blocks(self._free_blocks) + def free_object(self, obj_id: str) -> bool: """Return pool blocks for ``obj_id`` if any. @@ -353,9 +506,9 @@ def free_object(self, obj_id: str) -> bool: Returns: True if blocks were freed, False if ``obj_id`` had no allocation. """ - blocks = self._allocated_by_obj.pop(obj_id, None) - if blocks is None: - return False - self._free_blocks.extend(blocks) - _merge_free_blocks(self._free_blocks) - return True + with self._lock: + blocks = self._allocated_by_obj.pop(obj_id, None) + if blocks is None: + return False + self._free_blocks_locked(blocks) + return True diff --git a/python/ray/experimental/rdt/nixl_tensor_transport.py b/python/ray/experimental/rdt/nixl_tensor_transport.py index 29dd5962084..1579f5288eb 100644 --- a/python/ray/experimental/rdt/nixl_tensor_transport.py +++ b/python/ray/experimental/rdt/nixl_tensor_transport.py @@ -15,6 +15,7 @@ NIXL_REMOTE_AGENT_CACHE_MAXSIZE, ) from ray.experimental.rdt.nixl_memory_pool import ( + MemoryBlock, MemoryPoolManager, TensorLayout, group_tensors_by_desc, @@ -125,7 +126,12 @@ class NixlFetchRequest(FetchRequest): remove_tensor_descs: Whether to remove tensor descriptors from the cache during cleanup. registered_tensors: Tensors that were registered with NIXL, to deregister on cleanup. These are the buffers behind ``tensors``, which for a - group transfer are views into fewer, larger buffers. + group transfer are views into fewer, larger buffers. Empty when the + receive-side pool is used, since the pool is registered as a whole. + pool_blocks: Pool blocks backing the receive buffers, or None if the + receive-side pool was not used. Cleared once returned to the pool. + pool_target_device: Device the pool-backed buffers are copied out to, + which may differ from the pool's own device. """ xfer_handle: Any = None @@ -134,16 +140,32 @@ class NixlFetchRequest(FetchRequest): remove_tensor_descs: bool = False transport: Any = None registered_tensors: List["torch.Tensor"] = field(default_factory=list) + pool_blocks: Optional[List[MemoryBlock]] = None + pool_target_device: Optional[str] = None + + def release(self) -> None: + """Releases the transfer's resources, and is safe to call repeatedly. + + Callers that fail before handing the tensors back should call this + instead of waiting for ``__del__``: a raised exception keeps this + request alive through its traceback, which would hold pool blocks that + later transfers need. + """ + transport, self.transport = self.transport, None + if transport is None: + return + pool_blocks, self.pool_blocks = self.pool_blocks, None + transport._cleanup_transfer( + self.obj_id, + self.registered_tensors, + self.xfer_handle, + self.remote_name, + self.remove_tensor_descs, + pool_blocks, + ) def __del__(self): - if self.transport is not None: - self.transport._cleanup_transfer( - self.obj_id, - self.registered_tensors, - self.xfer_handle, - self.remote_name, - self.remove_tensor_descs, - ) + self.release() class NixlTensorTransport(TensorTransportManager): @@ -413,6 +435,8 @@ def fetch_multiple_tensors( added_tensor_descs = False registered_tensors: List["torch.Tensor"] = [] tensors: List["torch.Tensor"] = [] + pool_blocks: Optional[List[MemoryBlock]] = None + pool_target_device: Optional[str] = None try: nixl_agent = self.get_nixl_agent() @@ -456,10 +480,24 @@ def fetch_multiple_tensors( ) else: # One buffer per remote descriptor; views at recovered offsets. - group_buffers = [ - torch.empty(nbytes, dtype=torch.uint8, device=device) - for nbytes in packed_group_nbytes - ] + if self._memory_pool is not None: + # Carve the group buffers out of the pool, which is already + # registered with NIXL. wait_fetch_complete copies them out + # into independent tensors before returning, so the pool's + # own device doesn't have to be the tensors' device. + pool_target_device = device + group_buffers, pool_blocks = self._memory_pool.allocate_regions( + packed_group_nbytes + ) + else: + group_buffers = [ + torch.empty(nbytes, dtype=torch.uint8, device=device) + for nbytes in packed_group_nbytes + ] + self._add_tensor_descs(group_buffers) + added_tensor_descs = True + registered_tensors = group_buffers + tensors = [None] * len(tensor_meta) # type: ignore[list-item] for desc_idx, desc_group in enumerate(desc_groups): offsets, _ = packed_offsets([tensor_layouts[j] for j in desc_group]) @@ -472,9 +510,6 @@ def fetch_multiple_tensors( .reshape(shape) ) - self._add_tensor_descs(group_buffers) - added_tensor_descs = True - registered_tensors = group_buffers local_xfer_descs = nixl_agent.get_xfer_descs(group_buffers) remote_name = tensor_transport_metadata.nixl_agent_name @@ -521,6 +556,8 @@ def fetch_multiple_tensors( remove_tensor_descs=added_tensor_descs, transport=self, registered_tensors=registered_tensors, + pool_blocks=pool_blocks, + pool_target_device=pool_target_device, ) except Exception: self._cleanup_transfer( @@ -529,6 +566,7 @@ def fetch_multiple_tensors( xfer_handle, remote_name, added_tensor_descs, + pool_blocks, ) # TODO(swang): There is a circular import error because ray.util # currently depends on ray.experimental.internal_kv. @@ -586,12 +624,28 @@ def wait_fetch_complete( elif state == "DONE": break + if fetch_request.pool_blocks is not None: + blocks, fetch_request.pool_blocks = fetch_request.pool_blocks, None + # The tensors are views into pool blocks that the next transfer + # can overwrite, so copy them out before the caller sees them. + fetch_request.tensors = self._memory_pool.copy_out_and_free( + fetch_request.tensors, blocks, fetch_request.pool_target_device + ) return fetch_request.tensors except TimeoutError: + # The transfer is still in flight and may keep writing into the + # receive buffers, so leave them allocated. raise except Exception: from ray.exceptions import RayDirectTransportError + try: + fetch_request.release() + except Exception: + logger.exception( + f"Failed to release NIXL transfer resources for object id: {obj_id}." + ) + raise RayDirectTransportError( f"The NIXL transfer failed for object id: {obj_id}. The source actor may have died during the transfer. " f"The exception thrown from nixl transfer was:\n {traceback.format_exc()}" @@ -604,23 +658,29 @@ def _cleanup_transfer( xfer_handle: Any, remote_name: Optional[str], remove_tensor_descs: bool, + pool_blocks: Optional[List[MemoryBlock]] = None, ) -> None: """Cleans up resources after a transfer completes or fails.""" # We could raise errors or NIXL could raise errors like NIXL_ERR_REMOTE_DISCONNECT, # so doing best effort cleanup. nixl_agent = self._nixl_agent - if nixl_agent is None: - return - # We could raise errors or NIXL could raise errors like NIXL_ERR_REMOTE_DISCONNECT, - # so doing best effort cleanup. - with self._aborted_transfer_obj_ids_lock: - self._aborted_transfer_obj_ids.discard(obj_id) - if xfer_handle: - nixl_agent.release_xfer_handle(xfer_handle) - if NIXL_REMOTE_AGENT_CACHE_MAXSIZE == 0 and remote_name: - nixl_agent.remove_remote_agent(remote_name) - if remove_tensor_descs: - self._remove_tensor_descs(tensors) + try: + if nixl_agent is None: + return + with self._aborted_transfer_obj_ids_lock: + self._aborted_transfer_obj_ids.discard(obj_id) + if xfer_handle: + nixl_agent.release_xfer_handle(xfer_handle) + if NIXL_REMOTE_AGENT_CACHE_MAXSIZE == 0 and remote_name: + nixl_agent.remove_remote_agent(remote_name) + if remove_tensor_descs: + self._remove_tensor_descs(tensors) + finally: + # Reclaim the receive buffers only after the transfer handle is + # released, so that a block cannot be reused while NIXL still + # references it. + if pool_blocks and self._memory_pool is not None: + self._memory_pool.free_blocks(pool_blocks) def recv_multiple_tensors( self, diff --git a/python/ray/experimental/rdt/util.py b/python/ray/experimental/rdt/util.py index 2499a15c250..161cf9cc9b6 100644 --- a/python/ray/experimental/rdt/util.py +++ b/python/ray/experimental/rdt/util.py @@ -296,16 +296,24 @@ def register_nixl_memory_pool(size: int, device: "torch.device") -> None: This enables pool-based memory management for NIXL transfers, which can improve performance by avoiding repeated memory registration/deregistration. The pool is - registered once with NIXL and individual tensors are copied into it on ``ray.put``. + registered once with NIXL. - Only the tensors passed to ``ray.put`` are copied (by their own byte size), - not their full underlying storage. Contiguous tensors from a single - ``ray.put`` are packed into as few pool blocks as the free list allows. + On the sender side, only the tensors passed to ``ray.put`` are copied (by their + own byte size), not their full underlying storage. Contiguous tensors from a + single ``ray.put`` are packed into as few pool blocks as the free list allows. Pool blocks are freed when the ``ObjectRef`` goes out of scope. Each tensor is placed on a multiple of its own element size, so tensors sharing a dtype pack with no padding between them. + On the receiver side, incoming tensors are read into the pool when ``ray.get`` + fetches via NIXL, then copied out into ordinary tensors and the pool blocks + are returned as soon as the transfer completes. The tensor you get back is + independent of the pool, so the pool size only bounds how much data can be + in flight at once, not how much received data you can hold. If user-supplied + target buffers are used instead, the pool is bypassed and the traditional + register/deregister path is used. + If the pool has insufficient space for an allocation, :class:`NixlOutOfMemoryError` is raised. diff --git a/python/ray/tests/rdt/test_nixl_memory_pool.py b/python/ray/tests/rdt/test_nixl_memory_pool.py index de56fa1dec5..fe5de106762 100644 --- a/python/ray/tests/rdt/test_nixl_memory_pool.py +++ b/python/ray/tests/rdt/test_nixl_memory_pool.py @@ -8,6 +8,7 @@ import torch from ray.experimental.rdt.nixl_memory_pool import ( + _MAX_ALIGNMENT, MemoryPoolManager, NixlOutOfMemoryError, TensorLayout, @@ -455,5 +456,158 @@ def test_block_merging(self): assert regions[0].numel() == _nbytes(t_big) +# --------------------------------------------------------------------------- +# Receive side: allocate_regions / copy_out_and_free / free_blocks +# --------------------------------------------------------------------------- + + +class TestAllocateRegions: + def test_regions_are_exact_size_and_backed_by_pool(self): + pool = MemoryPoolManager(pool_size=128, device=torch.device("cpu")) + regions, blocks = pool.allocate_regions([10, 20]) + + assert [r.numel() for r in regions] == [10, 20] + assert len(blocks) == 2 + for region in regions: + assert region.dtype == torch.uint8 + assert ( + region.untyped_storage().data_ptr() + == pool.get_pool_tensor().untyped_storage().data_ptr() + ) + + def test_blocks_are_aligned_for_any_dtype(self): + """Unaligned region sizes must still leave later blocks viewable.""" + pool = MemoryPoolManager(pool_size=256, device=torch.device("cpu")) + _, blocks = pool.allocate_regions([1, 3, 7, 9]) + for block in blocks: + assert block.offset % _MAX_ALIGNMENT == 0 + + def test_round_trips_a_sender_packed_group(self): + """A region sized like a sender descriptor decodes back to the tensors. + + This is the wire contract the receive path relies on: the pool only + needs to hand back a correctly aligned region of the right length. + """ + tensors = [ + _make_tensor([1], dtype=torch.int8), + _make_tensor([2.0, 3.0, 4.0]), + _make_tensor([5.0, 6.0], dtype=torch.float64), + ] + layouts = _layout(tensors) + _, packed_nbytes = packed_offsets(layouts) + + pool = MemoryPoolManager(pool_size=256, device=torch.device("cpu")) + # Occupy the front of the pool so the region is not trivially at + # offset 0, the way a receive into a used pool would land. + pool.allocate_regions([1]) + regions, _ = pool.allocate_regions([packed_nbytes]) + + views = _unpack(regions, tensors) + for view, tensor in zip(views, tensors): + view.copy_(tensor) + assert torch.equal(view, tensor) + + def test_atomic_allocation_failure(self): + pool = MemoryPoolManager(pool_size=16, device=torch.device("cpu")) + with pytest.raises(NixlOutOfMemoryError): + pool.allocate_regions([8, 12]) + # Pool state unchanged: the whole pool is still allocatable. + regions, blocks = pool.allocate_regions([16]) + assert regions[0].numel() == 16 + pool.free_blocks(blocks) + assert sum(b.size for b in pool._free_blocks) == 16 + + def test_oom_on_fragmentation(self): + """A region must be contiguous, so split free space cannot serve it.""" + pool = MemoryPoolManager(pool_size=48, device=torch.device("cpu")) + _, first = pool.allocate_regions([16]) + _, middle = pool.allocate_regions([16]) + _, last = pool.allocate_regions([16]) + pool.free_blocks(first + last) + + with pytest.raises(NixlOutOfMemoryError): + pool.allocate_regions([32]) + assert sum(b.size for b in pool._free_blocks) == 32 + + pool.free_blocks(middle) + regions, _ = pool.allocate_regions([32]) + assert regions[0].numel() == 32 + + +class TestCopyOutAndFree: + def test_copies_are_independent_of_pool(self): + pool = MemoryPoolManager(pool_size=64, device=torch.device("cpu")) + regions, blocks = pool.allocate_regions([12]) + view = regions[0].view(torch.float32) + view.copy_(_make_tensor([1.0, 2.0, 3.0])) + + copies = pool.copy_out_and_free([view], blocks) + + assert torch.equal(copies[0], _make_tensor([1.0, 2.0, 3.0])) + assert ( + copies[0].untyped_storage().data_ptr() + != pool.get_pool_tensor().untyped_storage().data_ptr() + ) + + # Reusing the block must not disturb the copy. + reused, _ = pool.allocate_regions([12]) + reused[0].fill_(99) + assert torch.equal(copies[0], _make_tensor([1.0, 2.0, 3.0])) + + def test_blocks_are_reusable_without_gc(self): + """Blocks come back immediately, so sequential receives can exceed the + pool size in aggregate.""" + pool = MemoryPoolManager(pool_size=12, device=torch.device("cpu")) + for _ in range(3): + regions, blocks = pool.allocate_regions([12]) + pool.copy_out_and_free([regions[0].view(torch.float32)], blocks) + + assert sum(b.size for b in pool._free_blocks) == 12 + + def test_frees_blocks_even_if_copy_fails(self): + """A failed copy out must not strand the blocks it was handed.""" + pool = MemoryPoolManager(pool_size=64, device=torch.device("cpu")) + regions, blocks = pool.allocate_regions([16]) + + with pytest.raises(RuntimeError): + pool.copy_out_and_free(regions, blocks, target_device="not_a_device") + + assert sum(b.size for b in pool._free_blocks) == 64 + + def test_syncs_before_freeing_when_a_copy_fails(self, monkeypatch): + """A failure partway through still has to wait for the queued copies. + + Copies for earlier tensors are only ordered on the CUDA stream, so + freeing the block first would let the next NIXL transfer overwrite + memory they are still reading. + """ + pool = MemoryPoolManager(pool_size=64, device=torch.device("cpu")) + regions, blocks = pool.allocate_regions([16]) + # Take the CUDA path without a GPU: the copy fails on the bad target + # device before any real device work is queued. + pool.device = torch.device("cuda") + + order = [] + free_blocks = pool.free_blocks + + def record_free(freed): + order.append("free") + free_blocks(freed) + + monkeypatch.setattr(torch.cuda, "synchronize", lambda *_: order.append("sync")) + monkeypatch.setattr(pool, "free_blocks", record_free) + + with pytest.raises(RuntimeError): + pool.copy_out_and_free(regions, blocks, target_device="not_a_device") + + assert order == ["sync", "free"] + assert sum(b.size for b in pool._free_blocks) == 64 + + def test_free_blocks_is_noop_for_empty_list(self): + pool = MemoryPoolManager(pool_size=64, device=torch.device("cpu")) + pool.free_blocks([]) + assert sum(b.size for b in pool._free_blocks) == 64 + + if __name__ == "__main__": sys.exit(pytest.main(["-sv", __file__])) diff --git a/python/ray/tests/rdt/test_rdt_nixl.py b/python/ray/tests/rdt/test_rdt_nixl.py index d25a56aee71..d0b54eadd34 100644 --- a/python/ray/tests/rdt/test_rdt_nixl.py +++ b/python/ray/tests/rdt/test_rdt_nixl.py @@ -1115,6 +1115,130 @@ def consume(self, refs): assert result == [[1.0, 2.0], [3.0, 4.0, 5.0]] +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 2}], indirect=True) +def test_nixl_recv_memory_pool(ray_start_regular, device): + """ + Test receiver-side NIXL memory pool: incoming tensors are read into the + pre-registered pool and then copied out, so the tensor handed to the user + does not alias the pool and the block is returned right after the transfer. + """ + + @ray.remote(num_gpus=1, num_cpus=0, enable_tensor_transport=True) + class RecvPoolActor: + def __init__(self, pool_device, pool_size): + from ray.experimental import register_nixl_memory_pool + + register_nixl_memory_pool(pool_size, torch.device(pool_device)) + self._held_tensors = [] + + def consume_one(self, refs, device): + tensor = ray.get(refs[0]) + assert tensor.device.type == device + transport = get_tensor_transport_manager("NIXL") + pool = transport._memory_pool + # The received tensor is a copy, so it does not alias the pool. + assert ( + tensor.untyped_storage().data_ptr() + != pool.get_pool_tensor().untyped_storage().data_ptr() + ) + # The block is back in the pool even though the tensor is alive. + assert self._pool_free_bytes() == pool.pool_size + value = tensor.sum().item() + self._held_tensors.append(tensor) + return value + + def _pool_free_bytes(self): + transport = get_tensor_transport_manager("NIXL") + pool = transport._memory_pool + return sum(block.size for block in pool._free_blocks) + + def get_pool_free_bytes(self): + return self._pool_free_bytes() + + src_actor = GPUTestActor.remote() + # Each int64 tensor [1,2,3] is 24 bytes, so the pool holds exactly one. + dst_actor = RecvPoolActor.remote(device, 24) + + # Receives totaling more than the pool succeed, because each block is + # freed once its data has been copied out, even though the receiver keeps + # every tensor alive. + refs = src_actor.produce.remote([torch.tensor([1, 2, 3]).to(device)]) + assert ray.get(dst_actor.consume_one.remote(refs, device)) == 6 + + refs = src_actor.produce.remote([torch.tensor([4, 5, 6]).to(device)]) + assert ray.get(dst_actor.consume_one.remote(refs, device)) == 15 + + refs = src_actor.produce.remote([torch.tensor([7, 8, 9]).to(device)]) + assert ray.get(dst_actor.consume_one.remote(refs, device)) == 24 + + # A single transfer that doesn't fit in the pool still fails. + refs = src_actor.produce.remote([torch.tensor([1, 2, 3, 4, 5, 6]).to(device)]) + with pytest.raises(ray.exceptions.RayTaskError) as excinfo: + ray.get(dst_actor.consume_one.remote(refs, device)) + assert "NixlOutOfMemoryError" in str(excinfo.value) and "out of memory" in str( + excinfo.value + ) + + # The failed transfer left the pool intact. + assert ray.get(dst_actor.get_pool_free_bytes.remote()) == 24 + + +@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 2}], indirect=True) +def test_nixl_recv_memory_pool_packed_group(ray_start_regular): + """A sender-packed group is received into one pool region and unpacked. + + The sender packs both tensors into a single descriptor, so the receiver + allocates one pool region for the whole group and recovers the individual + tensors from it. Mixed dtypes make the packed offsets non-trivial. + """ + + @ray.remote(num_gpus=1, num_cpus=0, enable_tensor_transport=True) + class PoolSrc: + def __init__(self): + from ray.experimental import register_nixl_memory_pool + + register_nixl_memory_pool(1024, torch.device("cuda")) + + def put_list(self): + tensors = [ + torch.tensor([1], dtype=torch.int8).to("cuda"), + torch.tensor([2.0, 3.0], dtype=torch.float32).to("cuda"), + torch.tensor([4.0, 5.0], dtype=torch.float64).to("cuda"), + ] + ref = ray.put(tensors, _tensor_transport="nixl") + meta = get_tensor_transport_manager("NIXL")._get_meta(ref.hex()) + descs = ( + get_tensor_transport_manager("NIXL") + .get_nixl_agent() + .deserialize_descs(meta.nixl_serialized_descs) + ) + return ref, descs.descCount() + + @ray.remote(num_gpus=1, num_cpus=0, enable_tensor_transport=True) + class PoolDst: + def __init__(self): + from ray.experimental import register_nixl_memory_pool + + register_nixl_memory_pool(1024, torch.device("cuda")) + + def consume(self, refs): + tensors = ray.get(refs[0]) + pool = get_tensor_transport_manager("NIXL")._memory_pool + pool_ptr = pool.get_pool_tensor().untyped_storage().data_ptr() + for tensor in tensors: + assert tensor.untyped_storage().data_ptr() != pool_ptr + # The region is back even though every tensor is still alive. + assert sum(b.size for b in pool._free_blocks) == pool.pool_size + return [t.cpu().tolist() for t in tensors] + + src = PoolSrc.remote() + dst = PoolDst.remote() + ref, desc_count = ray.get(src.put_list.remote()) + assert desc_count == 1 + assert ray.get(dst.consume.remote([ref])) == [[1], [2.0, 3.0], [4.0, 5.0]] + + @pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 2}], indirect=True) def test_set_nixl_cuda_stream(ray_start_regular): """set_nixl_cuda_stream restricts the pre-registration sync to the given