[core][rdt] Implement receive side memory pool

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