chore: import upstream snapshot with attribution

python/ray/experimental/rdt/__init__.py

@@ -0,0 +1,29 @@
+from ray.experimental.rdt.rdt_manager import (
+    RDTManager,
+    set_target_for_ref,
+    wait_tensor_freed,
+)
+from ray.experimental.rdt.tensor_transport_manager import (
+    CommunicatorMetadata,
+    TensorTransportManager,
+    TensorTransportMetadata,
+)
+from ray.experimental.rdt.util import (
+    deregister_nixl_memory,
+    register_nixl_memory,
+    register_nixl_memory_pool,
+    register_tensor_transport,
+)
+
+__all__ = [
+    "RDTManager",
+    "wait_tensor_freed",
+    "register_tensor_transport",
+    "register_nixl_memory",
+    "deregister_nixl_memory",
+    "register_nixl_memory_pool",
+    "TensorTransportManager",
+    "TensorTransportMetadata",
+    "CommunicatorMetadata",
+    "set_target_for_ref",
+]

python/ray/experimental/rdt/collective_tensor_transport.py

@@ -0,0 +1,203 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List, Optional
+
+import ray
+from ray.experimental.rdt.tensor_transport_manager import (
+    CommunicatorMetadata,
+    TensorTransportManager,
+    TensorTransportMetadata,
+)
+
+if TYPE_CHECKING:
+    import torch
+
+
+@dataclass
+class CollectiveTransportMetadata(TensorTransportMetadata):
+    """Metadata for tensors stored in the GPU object store for collective transport."""
+
+
+@dataclass
+class CollectiveCommunicatorMetadata(CommunicatorMetadata):
+    """Metadata for the collective communicator (e.g. NCCL, GLOO).
+
+    Args:
+        src_rank: The rank of the source actor.
+        dst_rank: The rank of the destination actor.
+    """
+
+    communicator_name: str = ""
+    src_rank: Optional[int] = None
+    dst_rank: Optional[int] = None
+
+
+class CollectiveTensorTransport(TensorTransportManager):
+    def tensor_transport_backend(self) -> str:
+        raise NotImplementedError(
+            "NCCLTensorTransport or GLOOTensorTransport should be used instead of this base class."
+        )
+
+    @staticmethod
+    def is_one_sided() -> bool:
+        return False
+
+    @staticmethod
+    def can_abort_transport() -> bool:
+        return False
+
+    def actor_has_tensor_transport(self, actor: "ray.actor.ActorHandle") -> bool:
+        from ray.experimental.collective import get_collective_groups
+
+        communicators = get_collective_groups(
+            [actor], backend=self.tensor_transport_backend()
+        )
+        return len(communicators) > 0
+
+    def extract_tensor_transport_metadata(
+        self,
+        obj_id: str,
+        rdt_object: List["torch.Tensor"],
+    ) -> CollectiveTransportMetadata:
+        tensor_meta = []
+        device = None
+        if rdt_object:
+            device = rdt_object[0].device
+            for t in rdt_object:
+                if t.device.type != device.type:
+                    raise ValueError(
+                        "All tensors in an RDT object must have the same device type."
+                    )
+                tensor_meta.append((t.shape, t.dtype))
+        return CollectiveTransportMetadata(
+            tensor_meta=tensor_meta,
+            tensor_device=device.type if device else None,
+        )
+
+    def get_communicator_metadata(
+        self,
+        src_actor: "ray.actor.ActorHandle",
+        dst_actor: "ray.actor.ActorHandle",
+        backend: Optional[str] = None,
+    ) -> CollectiveCommunicatorMetadata:
+
+        from ray.experimental.collective import get_collective_groups
+
+        communicators = get_collective_groups(
+            [src_actor, dst_actor],
+            backend=backend,
+        )
+        # TODO(kevin85421): Support multiple communicators.
+        if len(communicators) == 0:
+            raise ValueError(
+                f"No communicators found for actors {src_actor} and {dst_actor}. "
+                "Create a communicator with "
+                "`ray.experimental.collective.create_collective_group` "
+                "before calling actor tasks. with non-default tensor_transport."
+            )
+        elif len(communicators) > 1:
+            raise ValueError(
+                f"There are {len(communicators)} possible communicators that contain actors {src_actor} and {dst_actor}. "
+                "Currently, RDT objects only support one communicator. Please make sure only "
+                "one communicator exists."
+            )
+        communicator = communicators[0]
+        src_rank = communicator.get_rank(src_actor)
+        if src_rank == -1:
+            raise ValueError(
+                f"Sender actor {src_actor} not found in communicator. "
+                "Please make sure the sender and receiver are in the same communicator."
+            )
+        dst_rank = communicator.get_rank(dst_actor)
+        if dst_rank == -1:
+            raise ValueError(
+                f"Receiver actor {dst_actor} not found in communicator. "
+                "Please make sure the sender and receiver are in the same communicator."
+            )
+
+        communicator_metadata = CollectiveCommunicatorMetadata(
+            communicator_name=communicator.name,
+            src_rank=src_rank,
+            dst_rank=dst_rank,
+        )
+        return communicator_metadata
+
+    def recv_multiple_tensors(
+        self,
+        obj_id: str,
+        tensor_transport_metadata: TensorTransportMetadata,
+        communicator_metadata: CommunicatorMetadata,
+        target_buffers: Optional[List["torch.Tensor"]] = None,
+    ):
+        from ray.experimental.rdt.util import (
+            create_empty_tensors_from_metadata,
+        )
+        from ray.util.collective.collective import recv
+
+        assert isinstance(tensor_transport_metadata, CollectiveTransportMetadata)
+        assert isinstance(communicator_metadata, CollectiveCommunicatorMetadata)
+
+        tensors = target_buffers or create_empty_tensors_from_metadata(
+            tensor_transport_metadata
+        )
+        for tensor in tensors:
+            recv(
+                tensor,
+                communicator_metadata.src_rank,
+                communicator_metadata.communicator_name,
+            )
+        return tensors
+
+    def send_multiple_tensors(
+        self,
+        tensors: List["torch.Tensor"],
+        tensor_transport_metadata: TensorTransportMetadata,
+        communicator_metadata: CommunicatorMetadata,
+    ):
+        import ray.util.collective as collective
+
+        assert isinstance(
+            tensor_transport_metadata, CollectiveTransportMetadata
+        ), "metadata must be a CollectiveTransportMetadata object for non-NIXL transport"
+        assert isinstance(
+            communicator_metadata, CollectiveCommunicatorMetadata
+        ), "metadata must be a CollectiveCommunicatorMetadata object for non-NIXL transport"
+
+        device = tensors[0].device if tensors else None
+
+        for tensor in tensors:
+            if tensor.device.type != device.type:
+                raise ValueError(
+                    f"tensor device {tensor.device} does not match device {device}"
+                )
+            collective.send(
+                tensor,
+                communicator_metadata.dst_rank,
+                communicator_metadata.communicator_name,
+            )
+
+    def garbage_collect(
+        self,
+        obj_id: str,
+        tensor_transport_meta: TensorTransportMetadata,
+        tensors: List["torch.Tensor"],
+    ):
+        pass
+
+    def abort_transport(
+        self,
+        obj_id: str,
+        communicator_metadata: CommunicatorMetadata,
+    ):
+        raise NotImplementedError(
+            "Collective transport does not support abort_transport for now."
+        )
+
+
+class NCCLTensorTransport(CollectiveTensorTransport):
+    def tensor_transport_backend(self) -> str:
+        return "NCCL"
+
+
+class GLOOTensorTransport(CollectiveTensorTransport):
+    def tensor_transport_backend(self) -> str:
+        return "GLOO"

python/ray/experimental/rdt/cuda_ipc_transport.py

@@ -0,0 +1,214 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, List, Optional
+
+import ray
+from ray.experimental.rdt.tensor_transport_manager import (
+    CommunicatorMetadata,
+    TensorTransportManager,
+    TensorTransportMetadata,
+)
+
+if TYPE_CHECKING:
+    import torch
+
+
+@dataclass
+class CudaIpcCommunicatorMetadata(CommunicatorMetadata):
+    """Metadata for the CUDA IPC communicator."""
+
+
+@dataclass
+class CudaIpcTransportMetadata(TensorTransportMetadata):
+    """Metadata for tensors stored in the GPU object store for CUDA IPC transport."""
+
+    # List of tuples, each containing the function and metadata to reconstruct the tensor.
+    cuda_ipc_handles: Optional[List[Any]] = None
+    # The IPC handle of the event that is used to synchronize the sender and receiver.
+    cuda_ipc_event_ipc_handle: Optional[bytes] = None
+    # The index of the GPU that the tensors are on. This requires that the GPU is
+    # assigned by Ray, e.g., using @ray.remote(num_gpus=1).
+    ray_gpu_idx: Optional[int] = None
+    # The node that the GPU that the tensors are on is on.
+    ray_node_id: Optional[str] = None
+
+
+class CudaIpcTransport(TensorTransportManager):
+    def __init__(self):
+        pass
+
+    @property
+    def tensor_transport_backend(self) -> str:
+        return "CUDA_IPC"
+
+    @staticmethod
+    def is_one_sided() -> bool:
+        return True
+
+    @staticmethod
+    def can_abort_transport() -> bool:
+        return False
+
+    def actor_has_tensor_transport(self, actor: "ray.actor.ActorHandle") -> bool:
+        # TODO: Ideally we would check if torch.cuda.is_available() on the actor
+        # and if so, return True. But we want to avoid blocking in ray.get() in
+        # this method since it gets called before submitting an actor task.
+        return True
+
+    def extract_tensor_transport_metadata(
+        self,
+        obj_id: str,
+        rdt_object: List["torch.Tensor"],
+    ) -> CudaIpcTransportMetadata:
+
+        tensor_meta = []
+        device = None
+        cuda_ipc_handles = []
+        event_ipc_handle = None
+        ray_gpu_idx = None
+        ray_node_id = None
+        if rdt_object:
+            import torch
+            from torch.multiprocessing.reductions import reduce_tensor
+
+            device = rdt_object[0].device
+            ray_gpu_idx = ray.get_gpu_ids()[device.index]
+            ray_node_id = ray.get_runtime_context().get_node_id()
+
+            # Create an interprocess-shareable CUDA event so that the receiver
+            # can wait for the sender's computations to complete.
+            event = torch.cuda.Event(interprocess=True)
+            torch.cuda.current_stream(device).record_event(event)
+
+            for t in rdt_object:
+                if t.device.type != device.type:
+                    raise ValueError(
+                        "All tensors in an RDT object must have the same device type."
+                    )
+                if t.device.index != device.index:
+                    raise ValueError(
+                        "All tensors in an RDT object must be on the same GPU."
+                    )
+                tensor_meta.append((t.shape, t.dtype))
+                ipc_handle = reduce_tensor(t)
+                cuda_ipc_handles.append(ipc_handle)
+
+            event_ipc_handle = event.ipc_handle()
+
+        return CudaIpcTransportMetadata(
+            tensor_meta=tensor_meta,
+            tensor_device=device.type if device else None,
+            cuda_ipc_handles=cuda_ipc_handles,
+            cuda_ipc_event_ipc_handle=event_ipc_handle,
+            ray_gpu_idx=ray_gpu_idx,
+            ray_node_id=ray_node_id,
+        )
+
+    def get_communicator_metadata(
+        self,
+        src_actor: "ray.actor.ActorHandle",
+        dst_actor: "ray.actor.ActorHandle",
+        backend: Optional[str] = None,
+    ) -> CudaIpcCommunicatorMetadata:
+
+        communicator_metadata = CudaIpcCommunicatorMetadata()
+        return communicator_metadata
+
+    def recv_multiple_tensors(
+        self,
+        obj_id: str,
+        tensor_transport_metadata: TensorTransportMetadata,
+        communicator_metadata: CommunicatorMetadata,
+        target_buffers: Optional[List["torch.Tensor"]] = None,
+    ) -> List["torch.Tensor"]:
+
+        assert isinstance(
+            tensor_transport_metadata, CudaIpcTransportMetadata
+        ), "metadata must be a CudaIpcTransportMetadata object for CUDA IPC transport"
+        assert isinstance(
+            communicator_metadata, CudaIpcCommunicatorMetadata
+        ), "metadata must be a CudaIpcCommunicatorMetadata object for CUDA IPC transport"
+
+        if target_buffers:
+            raise ValueError(
+                "The CUDA IPC transport does not support receiving into buffers."
+            )
+
+        tensors = []
+        if tensor_transport_metadata.tensor_meta:
+            import torch
+
+            cur_node_id = ray.get_runtime_context().get_node_id()
+            if cur_node_id != tensor_transport_metadata.ray_node_id:
+                raise ValueError(
+                    f"CUDA IPC transport only supports tensors on the same node, but the current node ID: {cur_node_id} and the sender node ID: {tensor_transport_metadata.ray_node_id} are different."
+                )
+
+            try:
+                device_idx = ray.get_gpu_ids().index(
+                    tensor_transport_metadata.ray_gpu_idx
+                )
+            except ValueError:
+                raise ValueError(
+                    f"CUDA IPC transport only supports tensors on the same GPU, but the receiver was not allocated the same GPUs by Ray as the sender (GPU: {tensor_transport_metadata.ray_gpu_idx}). To use the CUDA IPC RDT transport, ensure that the receiver is allocated the same GPU by Ray as the sender, and that CUDA_VISIBLE_DEVICES is set to `ray.get_gpu_ids()`, the GPUs assigned by Ray (this is the default behavior)."
+                )
+            device = torch.device(f"cuda:{device_idx}")
+
+            event_ipc_handle = tensor_transport_metadata.cuda_ipc_event_ipc_handle
+            if event_ipc_handle is not None:
+                # Reconstruct the event from IPC handle
+                event_remote = torch.cuda.Event.from_ipc_handle(
+                    device=device, handle=event_ipc_handle
+                )
+
+                # Make current stream wait for the sender's event
+                # This ensures sender's computation is complete before we use the tensor
+                # This is asynchronous - doesn't block CPU, only GPU stream
+                torch.cuda.current_stream(device).wait_event(event_remote)
+
+            for i, ipc_handle in enumerate(tensor_transport_metadata.cuda_ipc_handles):
+                # Reconstruct the tensor
+                func, args = ipc_handle
+                list_args = list(args)
+                # Fields specified in https://github.com/pytorch/pytorch/blob/1495b35d29512f303ab37780760c5e692158514b/torch/multiprocessing/reductions.py#L155
+                # Update device ID to match current process's device mapping
+                if not isinstance(list_args[6], int):
+                    raise RuntimeError(
+                        f"Expected CUDA IPC tensor reconstruction list_args[6] to be device ID, but got {list_args[6]}. Please file an issue at https://github.com/ray-project/ray/issues/new/choose."
+                    )
+                list_args[6] = device.index
+                try:
+                    tensor = func(*list_args)
+                except Exception as e:
+                    raise RuntimeError(
+                        "Error reconstructing CUDA IPC tensor. Source actor may have failed."
+                    ) from e
+                tensors.append(tensor)
+        return tensors
+
+    def send_multiple_tensors(
+        self,
+        tensors: List["torch.Tensor"],
+        tensor_transport_metadata: CudaIpcTransportMetadata,
+        communicator_metadata: CudaIpcCommunicatorMetadata,
+    ):
+        raise NotImplementedError(
+            "CUDA IPC transport does not support send_multiple_tensors, since it is a one-sided transport."
+        )
+
+    def garbage_collect(
+        self,
+        obj_id: str,
+        tensor_transport_meta: CudaIpcTransportMetadata,
+        tensors: List["torch.Tensor"],
+    ):
+        pass
+
+    def abort_transport(
+        self,
+        obj_id: str,
+        communicator_metadata: CudaIpcCommunicatorMetadata,
+    ):
+        # TODO: Implement CUDA IPC abort transport.
+        raise NotImplementedError(
+            "CUDA IPC transport does not support abort_transport for now."
+        )

python/ray/experimental/rdt/nixl_memory_pool.py

@@ -0,0 +1,289 @@
+"""Memory pool management for NIXL RDT optimization."""
+
+import logging
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+if TYPE_CHECKING:
+    import torch
+
+logger = logging.getLogger(__name__)
+
+
+class NixlOutOfMemoryError(RuntimeError):
+    """Raised when the NIXL memory pool runs out of space.
+
+    The pre-allocated memory pool does not have enough free space for the
+    requested allocation. Increase the pool size passed to
+    ``register_nixl_memory_pool`` to avoid this error.
+    """
+
+
+class MemoryBlock:
+    """Represents a memory block in the pool."""
+
+    def __init__(self, offset: int, size: int):
+        self.offset = offset
+        self.size = size
+
+    def __repr__(self):
+        return f"MemoryBlock(offset={self.offset}, size={self.size})"
+
+
+class MemoryPoolManager:
+    """Manages a pre-allocated memory pool for NIXL RDT transfers.
+
+    This class provides a memory allocator interface over a pre-allocated memory pool,
+    allowing reuse of registered memory descriptors across multiple transfers.
+
+    It also tracks which storage data pointers have allocated blocks, enabling
+    cross-call reuse (the same storage can reuse its pool slot across multiple
+    ray.put calls) and pool-level block management.
+    """
+
+    def __init__(self, pool_size: int, device: "torch.device"):
+        """Initialize the memory pool manager.
+
+        Args:
+            pool_size: Size of the memory pool in bytes.
+            device: Device to allocate the pool on.
+        """
+        import torch
+
+        self.pool_size = pool_size
+        self.device = device
+
+        # Allocate the memory pool as a single tensor
+        # We use a 1D tensor of uint8 to represent raw memory
+        self._pool_tensor = torch.zeros(
+            pool_size, dtype=torch.uint8, device=self.device
+        )
+
+        # Track free blocks using a largest-request-first, first-fit allocator.
+        # List of MemoryBlock for free blocks, sorted by offset.
+        self._free_blocks: List[MemoryBlock] = [MemoryBlock(offset=0, size=pool_size)]
+
+        # Track allocated blocks by storage data pointer.
+        # Maps storage_data_ptr -> MemoryBlock in the pool.
+        self._allocated_blocks: Dict[int, MemoryBlock] = {}
+
+    def get_pool_tensor(self) -> "torch.Tensor":
+        """Get the underlying pool tensor.
+
+        Returns:
+            The pre-allocated tensor representing the memory pool.
+        """
+        return self._pool_tensor
+
+    def has_block(self, tensor: "torch.Tensor") -> bool:
+        """Check if a tensor has an allocated block in the pool.
+
+        Args:
+            tensor: The tensor to check.
+
+        Returns:
+            True if the tensor's storage has an allocated block.
+        """
+        return tensor.untyped_storage().data_ptr() in self._allocated_blocks
+
+    def free_tensors(self, tensors: List["torch.Tensor"]) -> None:
+        """Return pool blocks for the given tensors back to the pool.
+
+        The caller is responsible for calling this method on the same tensors that were previously allocated in the pool before those tensors go out of scope.
+
+        Args:
+            tensors: Tensors whose pool blocks should be freed.
+        """
+        blocks = []
+        for tensor in tensors:
+            ptr = tensor.untyped_storage().data_ptr()
+            if ptr in self._allocated_blocks:
+                blocks.append(self._allocated_blocks.pop(ptr))
+        if blocks:
+            self._free_multiple(blocks)
+
+    def allocate_for_tensors(
+        self, tensors: List["torch.Tensor"]
+    ) -> List["torch.Tensor"]:
+        """Allocate pool blocks for unique storages, copy data in,
+        and return pool-backed tensor views for each input tensor. The caller is responsible for calling free on the original tensors to return the allocated tensor views back to the pool before the original tensors go out of scope.
+
+        Handles storage-level deduplication: views of the same storage share
+        one pool block within a single call, and the same storage reuses its
+        existing pool slot across calls.
+
+        Args:
+            tensors: Source tensors to allocate pool memory for.
+
+        Returns:
+            List of pool-backed tensor views, one per input tensor,
+            in the same order.
+
+        Raises:
+            NixlOutOfMemoryError: If the pool has insufficient space.
+        """
+        new_allocations = None
+        newly_tracked_ptrs: List[int] = []
+        try:
+            import torch
+
+            # Deduplicate storages: group tensors by storage data_ptr so
+            # views of the same storage share one pool allocation.
+            # Maps storage data_ptr -> index in alloc_sizes/new_allocations,
+            # or -1 for storages that already have a pool block (cache hit).
+            storage_idx: Dict[int, int] = {}
+            # Maps storage data_ptr -> a representative tensor (for copy).
+            ptr_to_tensor: Dict[int, "torch.Tensor"] = {}
+            alloc_sizes: List[int] = []
+
+            for tensor in tensors:
+                ptr = tensor.untyped_storage().data_ptr()
+                if ptr in storage_idx:
+                    continue
+                ptr_to_tensor[ptr] = tensor
+                if self.has_block(tensor):
+                    storage_idx[ptr] = -1
+                else:
+                    storage_idx[ptr] = len(alloc_sizes)
+                    alloc_sizes.append(tensor.untyped_storage().nbytes())
+
+            # Allocate new (non-cached) storages atomically.
+            if alloc_sizes:
+                new_allocations = self._allocate_multiple(alloc_sizes)
+                if new_allocations is None:
+                    raise NixlOutOfMemoryError(
+                        f"NIXL memory pool out of memory: cannot allocate "
+                        f"{len(alloc_sizes)} block(s) totaling "
+                        f"{sum(alloc_sizes)} bytes. Consider increasing "
+                        f"the pool size when calling "
+                        f"register_nixl_memory_pool."
+                    )
+
+            # Track and copy newly allocated blocks. Cache hits keep the
+            # originally copied data -- any mutations to the source storage
+            # since the first ray.put are not reflected in outstanding refs.
+            for ptr, idx in storage_idx.items():
+                if idx < 0:
+                    continue
+                blk = new_allocations[idx]
+                self._allocated_blocks[ptr] = blk
+                newly_tracked_ptrs.append(ptr)
+                # Copy the tensor's full underlying storage into the pool block.
+                src = ptr_to_tensor[ptr]
+                storage_size = src.untyped_storage().nbytes()
+                storage_bytes = torch.tensor(
+                    [], dtype=torch.uint8, device=src.device
+                ).set_(src.untyped_storage())
+                self._pool_tensor[blk.offset : blk.offset + storage_size].copy_(
+                    storage_bytes
+                )
+
+            # Build pool-backed tensor views for each input tensor.
+            pool_views: List["torch.Tensor"] = []
+            for tensor in tensors:
+                ptr = tensor.untyped_storage().data_ptr()
+                blk = self._allocated_blocks[ptr]
+                pool_offset = blk.offset + (
+                    tensor.storage_offset() * tensor.element_size()
+                )
+                view_byte_size = tensor.numel() * tensor.element_size()
+                pool_bytes = self._pool_tensor[
+                    pool_offset : pool_offset + view_byte_size
+                ]
+                pool_views.append(pool_bytes.view(tensor.dtype).reshape(tensor.shape))
+
+            return pool_views
+
+        except Exception:
+            # Roll back any pool mutations made in this call, then re-raise.
+            try:
+                if new_allocations is not None:
+                    self._free_multiple(new_allocations)
+                for ptr in newly_tracked_ptrs:
+                    self._allocated_blocks.pop(ptr, None)
+            except Exception as cleanup_err:
+                logger.error(f"Memory pool cleanup failed: {cleanup_err}.")
+            raise
+
+    def _allocate_multiple(self, sizes: List[int]) -> Optional[List[MemoryBlock]]:
+        """Allocate multiple memory blocks from the pool atomically.
+
+        Either all allocations succeed, or none of them do.
+
+        Args:
+            sizes: List of sizes to allocate in bytes.
+
+        Returns:
+            List of MemoryBlock if all allocations succeed, None otherwise.
+        """
+        if not sizes or any(s <= 0 for s in sizes):
+            raise ValueError("Invalid allocation request")
+
+        # If total free space is less than total requested, fail fast.
+        total_requested = sum(sizes)
+        total_free = sum(b.size for b in self._free_blocks)
+        if total_free < total_requested:
+            return None
+
+        # Allocate largest first to reduce fragmentation; then return in original order.
+        order = sorted(range(len(sizes)), key=lambda i: -sizes[i])
+        sorted_sizes = [sizes[i] for i in order]
+
+        # Try to allocate all blocks atomically.
+        allocations: List[MemoryBlock] = []
+        temp_free_blocks = [MemoryBlock(b.offset, b.size) for b in self._free_blocks]
+
+        for size in sorted_sizes:
+            allocated = False
+            for i, block in enumerate(temp_free_blocks):
+                if block.size >= size:
+                    # Allocate at the start of the current free block
+                    offset = block.offset
+                    remaining_after = block.size - size
+
+                    if remaining_after == 0:
+                        temp_free_blocks.pop(i)
+                    else:
+                        block.offset = offset + size
+                        block.size = remaining_after
+
+                    allocations.append(MemoryBlock(offset, size))
+                    allocated = True
+                    break
+
+            if not allocated:
+                # If any size cannot be allocated, the entire batch fails,
+                # do not modify the real state.
+                return None
+
+        # Reorder allocations back to original request order
+        result: List[MemoryBlock] = [MemoryBlock(0, 0)] * len(sizes)
+        for k, alloc in enumerate(allocations):
+            result[order[k]] = alloc
+
+        # All successful, submit modifications
+        temp_free_blocks.sort(key=lambda b: b.offset)
+        self._free_blocks = temp_free_blocks
+
+        return result
+
+    def _free_multiple(self, blocks: List[MemoryBlock]) -> None:
+        """Free multiple memory blocks back to the pool.
+
+        Args:
+            blocks: Memory blocks to free.
+        """
+        if not blocks:
+            raise ValueError("Invalid free request")
+        self._free_blocks.extend(blocks)
+
+        # Single pass: merge all adjacent free blocks
+        self._free_blocks.sort(key=lambda b: b.offset)
+        i = 0
+        while i < len(self._free_blocks) - 1:
+            curr = self._free_blocks[i]
+            next_block = self._free_blocks[i + 1]
+            if curr.offset + curr.size == next_block.offset:
+                curr.size += next_block.size
+                self._free_blocks.pop(i + 1)
+            else:
+                i += 1

python/ray/experimental/rdt/nixl_tensor_transport.py

@@ -0,0 +1,738 @@
+import functools
+import glob
+import logging
+import os
+import threading
+import time
+import traceback
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+import ray
+from ray._private.ray_constants import (
+    NIXL_REMOTE_AGENT_CACHE_MAXSIZE,
+)
+from ray.experimental.rdt.nixl_memory_pool import MemoryPoolManager
+from ray.experimental.rdt.tensor_transport_manager import (
+    CommunicatorMetadata,
+    FetchRequest,
+    TensorTransportManager,
+    TensorTransportMetadata,
+)
+
+if TYPE_CHECKING:
+    import torch
+
+logger = logging.getLogger(__name__)
+
+
+@functools.lru_cache(maxsize=1)
+def _is_efa_available() -> bool:
+    """Detect whether AWS EFA (Elastic Fabric Adapter) devices are present.
+
+    A bare host exposes ``efa*`` netdevs, but inside a container/Kubernetes pod
+    netdevs are network-namespaced away and only the rdma-verbs devices under
+    ``/sys/class/infiniband`` are mounted in. Those verbs devices are not
+    EFA-specific -- ordinary InfiniBand/RoCE NICs appear there too -- so we
+    confirm each one is bound to the kernel ``efa`` driver before treating it as
+    EFA. Without that check, non-AWS RDMA nodes would wrongly auto-select the
+    LIBFABRIC backend instead of UCX.
+    """
+    if glob.glob("/sys/class/net/efa*"):
+        return True
+    for ib_dev in glob.glob("/sys/class/infiniband/*"):
+        # A stale or broken sysfs entry shouldn't abort the scan; skip it and
+        # keep looking (defaulting to UCX if nothing resolves to the efa driver).
+        try:
+            driver = os.path.realpath(os.path.join(ib_dev, "device", "driver"))
+        except OSError:
+            continue
+        if os.path.basename(driver) == "efa":
+            return True
+    return False
+
+
+def _nixl_transport_available_in_process() -> bool:
+    """Returns whether the NIXL tensor transport can be initialized in this process.
+
+    Returns:
+        True if the NIXL agent initializes successfully, False on any failure
+        (e.g. nixl not installed, LIBFABRIC/EFA probe failure, or other backend
+        init errors).
+    """
+    try:
+        from ray.experimental.rdt.util import get_tensor_transport_manager
+
+        get_tensor_transport_manager("NIXL").get_nixl_agent()
+        return True
+    except Exception:
+        logger.debug("NIXL tensor transport unavailable on actor.", exc_info=True)
+        return False
+
+
+@dataclass
+class NixlCommunicatorMetadata(CommunicatorMetadata):
+    """Metadata for the NIXL communicator."""
+
+
+@dataclass
+class NixlTransportMetadata(TensorTransportMetadata):
+    """Metadata for tensors stored in the GPU object store for NIXL transport.
+
+    Args:
+        nixl_serialized_descs: Serialized tensor descriptors for NIXL transport.
+        nixl_agent_meta: The additional metadata of the remote NIXL agent.
+        nixl_agent_name: The name of the NIXL agent.
+        nixl_agent_meta_version: The version of the NIXL agent metadata.
+    """
+
+    nixl_serialized_descs: Optional[bytes] = None
+    nixl_agent_meta: Optional[bytes] = None
+    nixl_agent_name: Optional[str] = None
+    nixl_agent_meta_version: Optional[int] = 0
+
+    __eq__ = object.__eq__
+    __hash__ = object.__hash__
+
+
+@dataclass
+class TensorDesc:
+    # nixlRegDList handle, or None for pool-managed tensors (pool memory is
+    # registered once at pool creation, so individual tensors don't need their
+    # own NIXL registration).
+    reg_desc: Any
+    # tracks the number of NIXL metadata containing the tensor.
+    metadata_count: int
+
+
+@dataclass
+class NixlFetchRequest(FetchRequest):
+    """NIXL-specific FetchRequest carrying the async transfer state.
+
+    Returned by fetch_multiple_tensors and consumed by wait_fetch_complete.
+
+    Args:
+        obj_id: Inherited. The object ID for the transfer, used for abort checks and cleanup.
+        tensors: Inherited. Pre-allocated output tensors (populated before the transfer starts).
+        xfer_handle: NIXL transfer request handle.
+        nixl_agent: Reference to the NIXL agent.
+        remote_name: Name of the remote NIXL agent.
+        remove_tensor_descs: Whether to remove tensor descriptors from the cache during cleanup.
+    """
+
+    xfer_handle: Any = None
+    nixl_agent: Any = None
+    remote_name: Optional[str] = None
+    remove_tensor_descs: bool = False
+    transport: Any = None
+
+    def __del__(self):
+        if self.transport is not None:
+            self.transport._cleanup_transfer(
+                self.obj_id,
+                self.tensors,
+                self.xfer_handle,
+                self.remote_name,
+                self.remove_tensor_descs,
+            )
+
+
+class NixlTensorTransport(TensorTransportManager):
+    def __init__(self):
+        # This is lazily initialized because it requires NIXL to actually be installed and we want to allow an owner that is just coordinating to not need to have NIXL installed.
+        self._nixl_agent = None
+        self._aborted_transfer_obj_ids = set()
+        self._aborted_transfer_obj_ids_lock = threading.Lock()
+        # Mapping from tensor storage data pointer to the NIXL descriptor and reference count.
+        # Unlike _managed_meta_nixl, we only deregister tensors when ALL metadata containing the tensor is freed.
+        # For pool-managed tensors, reg_desc is None and the pool block is returned instead of deregistering.
+        self._tensor_desc_cache: Dict[int, TensorDesc] = {}
+        # Mapping from object ID to the NIXL managed meta.
+        # The lifetime of _managed_meta_nixl is tied to the object ref and freed when the ref goes out of scope.
+        self._managed_meta_nixl: Dict[str, Any] = {}
+        # Lock protecting _tensor_desc_cache and _managed_meta_nixl since they can be
+        # accessed from the main task execution thread or the _ray_system thread.
+        self._cache_lock = threading.RLock()
+        # LRU cache of remote agent names. When full, the least
+        # recently used remote agent is evicted and remove_remote_agent is called.
+        self._remote_agents: OrderedDict = OrderedDict()
+        # Increment the version whenever memory is deregistered.
+        self._nixl_agent_meta_version = 0
+        self._memory_pool: Optional[MemoryPoolManager] = None
+        # The NIXL backend the agent was actually created with ("UCX" or "LIBFABRIC").
+        self._backend: Optional[str] = None
+
+    def tensor_transport_backend(self) -> str:
+        return "NIXL"
+
+    @staticmethod
+    def is_one_sided() -> bool:
+        return True
+
+    @staticmethod
+    def can_abort_transport() -> bool:
+        return True
+
+    def register_nixl_memory(self, tensor: "torch.Tensor") -> None:
+        """Registers the tensor's memory with NIXL and bumps the reference count so the memory region is never deregistered."""
+        self._add_tensor_descs([tensor])
+
+    def register_nixl_memory_pool(self, size: int, device: "torch.device") -> None:
+        """Pre-allocates a memory pool and registers it with NIXL.
+
+        Args:
+            size: Size of the memory pool in bytes.
+            device: Device to allocate the pool on (cpu or cuda).
+
+        Raises:
+            ValueError: If a memory pool is already registered.
+        """
+        if self._memory_pool is not None:
+            raise ValueError(
+                "A memory pool is already registered. "
+                "Only one memory pool is supported."
+            )
+        nixl_agent = self.get_nixl_agent()
+        pool = MemoryPoolManager(pool_size=size, device=device)
+        nixl_agent.register_memory(pool.get_pool_tensor())
+        self._memory_pool = pool
+
+    def deregister_nixl_memory(self, tensor: "torch.Tensor") -> None:
+        """Decrements the reference count for the tensor's NIXL memory registration.
+        If the count reaches 0, the memory is deregistered from NIXL.
+        """
+        self._remove_tensor_descs([tensor])
+
+    def select_backend(self) -> str:
+        """Returns the NIXL backend to attempt.
+
+        Prefers LIBFABRIC when EFA devices are present and UCX everywhere else.
+        LIBFABRIC requires GPUDirect (GDR) for CUDA registration; if it isn't
+        available, ``_add_tensor_descs`` surfaces a clear error at registration
+        time with backend-specific troubleshooting guidance.
+        """
+        return "LIBFABRIC" if _is_efa_available() else "UCX"
+
+    def _make_nixl_agent(self, backend: str):
+        """Creates a NIXL agent configured with the given backend."""
+        from nixl._api import nixl_agent, nixl_agent_config
+
+        agent_config = nixl_agent_config(backends=[backend])
+        ctx = ray.get_runtime_context()
+        actor_id = ctx.get_actor_id()
+        if actor_id is None:
+            # If the actor id is None, it means the current process is a driver.
+            import uuid
+
+            actor_id = f"RAY-DRIVER-{uuid.uuid4()}"
+        return nixl_agent(actor_id, agent_config)
+
+    def get_nixl_agent(self):
+        """Returns the NIXL agent, building it once on first use."""
+        if self._nixl_agent is None:
+            self._nixl_agent = self._init_nixl_agent()
+        return self._nixl_agent
+
+    def _init_nixl_agent(self):
+        """Builds the NIXL agent for the selected backend."""
+        backend = self.select_backend()
+        agent = self._make_nixl_agent(backend)
+        self._backend = backend
+        logger.info("Using NIXL backend: %s", backend)
+        return agent
+
+    def actor_has_tensor_transport(self, actor: "ray.actor.ActorHandle") -> bool:
+        # TODO(dayshah): This is called on a .remote RDT call, so it's quite expensive.
+        def __ray_actor_has_tensor_transport__(
+            self: "ray.actor.ActorHandle",
+        ) -> bool:
+            return _nixl_transport_available_in_process()
+
+        return ray.get(
+            actor.__ray_call__.options(concurrency_group="_ray_system").remote(
+                __ray_actor_has_tensor_transport__
+            )
+        )
+
+    def extract_tensor_transport_metadata(
+        self,
+        obj_id: str,
+        rdt_object: List["torch.Tensor"],
+    ) -> NixlTransportMetadata:
+        import torch
+
+        with self._cache_lock:
+            device = None
+            tensor_meta = []
+
+            if rdt_object:
+                # We assume all tensors in one RDT object have the same device type,
+                # but we don't assume they're all on the same device.
+                devices = set()
+                device = rdt_object[0].device
+                for t in rdt_object:
+                    if t.device.type != device.type:
+                        raise ValueError(
+                            "All tensors in an RDT object must have the same device type."
+                        )
+                    if not t.is_contiguous():
+                        raise ValueError(
+                            "All tensors in an RDT object must be contiguous."
+                        )
+                    tensor_meta.append((t.shape, t.dtype))
+                    devices.add(t.device)
+                if device.type == "cuda":
+                    # We have to synchronize before memory registration to assure the
+                    # object has been created because nixl doesn't guarantee it will.
+                    for dev in devices:
+                        torch.cuda.synchronize(dev)
+
+                nixl_agent = self.get_nixl_agent()
+                # Use the pool only when every tensor lives on the exact same
+                # device as the pool, AND no tensor already has an existing
+                # NIXL registration (via register_nixl_memory).
+                pool_eligible = (
+                    self._memory_pool is not None
+                    and all(
+                        t.device == self._memory_pool.get_pool_tensor().device
+                        for t in rdt_object
+                    )
+                    and not any(self._tensor_memory_registered(t) for t in rdt_object)
+                )
+                if pool_eligible:
+                    xfer_descs = self._allocate_pool_xfer_descs(rdt_object)
+                else:
+                    self._add_tensor_descs(rdt_object)
+                    xfer_descs = nixl_agent.get_xfer_descs(rdt_object)
+
+                serialized_descs = nixl_agent.get_serialized_descs(xfer_descs)
+                agent_meta = nixl_agent.get_agent_metadata()
+                agent_name = nixl_agent.name
+                agent_meta_version = self._nixl_agent_meta_version
+            else:
+                serialized_descs, agent_meta = None, None
+                agent_name, agent_meta_version = None, None
+
+            ret = NixlTransportMetadata(
+                tensor_meta=tensor_meta,
+                tensor_device=device.type if device else None,
+                nixl_serialized_descs=serialized_descs,
+                nixl_agent_meta=agent_meta,
+                nixl_agent_name=agent_name,
+                nixl_agent_meta_version=agent_meta_version,
+            )
+            self._put_meta(obj_id, ret)
+            return ret
+
+    def get_communicator_metadata(
+        self,
+        src_actor: "ray.actor.ActorHandle",
+        dst_actor: "ray.actor.ActorHandle",
+        backend: Optional[str] = None,
+    ) -> NixlCommunicatorMetadata:
+        return NixlCommunicatorMetadata()
+
+    def fetch_multiple_tensors(
+        self,
+        obj_id: str,
+        tensor_transport_metadata: TensorTransportMetadata,
+        communicator_metadata: CommunicatorMetadata,
+        target_buffers: Optional[List["torch.Tensor"]] = None,
+    ) -> NixlFetchRequest:
+        """Initiates an async transfer for multiple tensors.
+
+        This triggers the transfer but does not wait for completion.
+        Call wait_fetch_complete(fetch_request) to wait for the transfer to
+        finish and retrieve the tensors.
+
+        Args:
+            obj_id: The object ID for the transfer.
+            tensor_transport_metadata: Metadata for the tensor transport.
+            communicator_metadata: Metadata for the communicator.
+            target_buffers: Optional pre-allocated buffers to receive tensors into.
+
+        Returns:
+            A NixlFetchRequest carrying the async transfer state.
+        """
+        from ray.experimental.rdt.util import (
+            create_empty_tensors_from_metadata,
+        )
+
+        tensors = target_buffers or create_empty_tensors_from_metadata(
+            tensor_transport_metadata
+        )
+
+        assert isinstance(tensor_transport_metadata, NixlTransportMetadata)
+        assert isinstance(communicator_metadata, NixlCommunicatorMetadata)
+
+        nixl_serialized_descs = tensor_transport_metadata.nixl_serialized_descs
+        remote_nixl_agent_meta = tensor_transport_metadata.nixl_agent_meta
+
+        with self._aborted_transfer_obj_ids_lock:
+            if obj_id in self._aborted_transfer_obj_ids:
+                self._aborted_transfer_obj_ids.remove(obj_id)
+                raise RuntimeError(f"NIXL transfer aborted for object id: {obj_id}")
+
+        remote_name = None
+        xfer_handle = None
+        added_tensor_descs = False
+
+        assert tensors
+
+        try:
+            nixl_agent = self.get_nixl_agent()
+            remote_xfer_descs = nixl_agent.deserialize_descs(nixl_serialized_descs)
+            # This creates a placeholder for the tensor in the tensor_desc_cache even though it doesn't have an object ref for caching purposes.
+            self._add_tensor_descs(tensors)
+            added_tensor_descs = True
+            local_xfer_descs = nixl_agent.get_xfer_descs(tensors)
+
+            remote_name = tensor_transport_metadata.nixl_agent_name
+            remote_agent_meta_version = (
+                tensor_transport_metadata.nixl_agent_meta_version
+            )
+
+            # Nixl agent reuse is enabled.
+            if NIXL_REMOTE_AGENT_CACHE_MAXSIZE > 0:
+                if remote_name in self._remote_agents:
+                    # If the remote agent metadata version is different from the cached one,
+                    # it means there was memory deregistered. We need to remove the remote agent
+                    # before adding it, because `nixlRemoteSection` currently does not support
+                    # updating descriptor list in such a case (there is potential memory overlap).
+                    if remote_agent_meta_version != self._remote_agents[remote_name]:
+                        nixl_agent.remove_remote_agent(remote_name)
+                    self._remote_agents.move_to_end(remote_name)
+                elif len(self._remote_agents) >= NIXL_REMOTE_AGENT_CACHE_MAXSIZE:
+                    evicted_agent_name, _ = self._remote_agents.popitem(last=False)
+                    nixl_agent.remove_remote_agent(evicted_agent_name)
+
+                self._remote_agents[remote_name] = remote_agent_meta_version
+
+            nixl_agent.add_remote_agent(remote_nixl_agent_meta)
+
+            xfer_handle = nixl_agent.initialize_xfer(
+                "READ",
+                local_xfer_descs,
+                remote_xfer_descs,
+                remote_name,
+                b"UUID",
+            )
+
+            state = nixl_agent.transfer(xfer_handle)
+            if state == "ERR":
+                raise RuntimeError("NIXL transfer got to Error state.")
+
+            return NixlFetchRequest(
+                tensors=tensors,
+                obj_id=obj_id,
+                xfer_handle=xfer_handle,
+                nixl_agent=nixl_agent,
+                remote_name=remote_name,
+                remove_tensor_descs=added_tensor_descs,
+                transport=self,
+            )
+        except Exception:
+            self._cleanup_transfer(
+                obj_id, tensors, xfer_handle, remote_name, added_tensor_descs
+            )
+            # TODO(swang): There is a circular import error because ray.util
+            # currently depends on ray.experimental.internal_kv.
+            from ray.exceptions import RayDirectTransportError
+
+            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()}"
+            ) from None
+
+    def wait_fetch_complete(
+        self, fetch_request: FetchRequest, timeout: float = -1
+    ) -> List["torch.Tensor"]:
+        """Waits for a previously initiated fetch to complete and returns the tensors.
+
+        Args:
+            fetch_request: The NixlFetchRequest returned by fetch_multiple_tensors.
+            timeout: Maximum time in seconds to wait. -1 means wait indefinitely.
+                0 means return immediately if not ready.
+
+        Returns:
+            List of tensors that were transferred.
+
+        Raises:
+            RayDirectTransportError: If the transfer failed.
+            TimeoutError: If the timeout is exceeded.
+        """
+        assert isinstance(fetch_request, NixlFetchRequest)
+        obj_id = fetch_request.obj_id
+
+        if not fetch_request.tensors:
+            return fetch_request.tensors
+
+        try:
+            # Check the state of the transfer continuously.
+            deadline = None if timeout < 0 else time.monotonic() + timeout
+            while True:
+                state = self.get_nixl_agent().check_xfer_state(
+                    fetch_request.xfer_handle
+                )
+                if state == "ERR":
+                    raise RuntimeError("NIXL transfer got to Error state.")
+                if state == "PROC":
+                    if deadline is not None and time.monotonic() >= deadline:
+                        raise TimeoutError(
+                            f"NIXL transfer timed out after {timeout}s for object id: {obj_id}"
+                        )
+                    with self._aborted_transfer_obj_ids_lock:
+                        if obj_id in self._aborted_transfer_obj_ids:
+                            self._aborted_transfer_obj_ids.remove(obj_id)
+                            raise RuntimeError(
+                                f"NIXL transfer aborted for object id: {obj_id}"
+                            )
+                    time.sleep(0.001)  # Avoid busy waiting
+                elif state == "DONE":
+                    break
+
+            return fetch_request.tensors
+        except TimeoutError:
+            raise
+        except Exception:
+            from ray.exceptions import RayDirectTransportError
+
+            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()}"
+            ) from None
+
+    def _cleanup_transfer(
+        self,
+        obj_id: str,
+        tensors: List["torch.Tensor"],
+        xfer_handle: Any,
+        remote_name: Optional[str],
+        remove_tensor_descs: bool,
+    ) -> 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)
+
+    def recv_multiple_tensors(
+        self,
+        obj_id: str,
+        tensor_transport_metadata: TensorTransportMetadata,
+        communicator_metadata: CommunicatorMetadata,
+        target_buffers: Optional[List["torch.Tensor"]] = None,
+    ) -> List["torch.Tensor"]:
+        """Receives multiple tensors synchronously."""
+        fetch_request = self.fetch_multiple_tensors(
+            obj_id, tensor_transport_metadata, communicator_metadata, target_buffers
+        )
+        return self.wait_fetch_complete(fetch_request)
+
+    def send_multiple_tensors(
+        self,
+        tensors: List["torch.Tensor"],
+        tensor_transport_metadata: TensorTransportMetadata,
+        communicator_metadata: CommunicatorMetadata,
+    ):
+        raise NotImplementedError(
+            "NIXL transport does not support send_multiple_tensors, since it is a one-sided transport."
+        )
+
+    def garbage_collect(
+        self,
+        obj_id: str,
+        tensor_transport_meta: TensorTransportMetadata,
+        tensors: List["torch.Tensor"],
+    ):
+        with self._cache_lock:
+            assert isinstance(tensor_transport_meta, NixlTransportMetadata)
+            if obj_id not in self._managed_meta_nixl:
+                return
+            self._managed_meta_nixl.pop(obj_id, None)
+            self._remove_tensor_descs(tensors)
+
+    def abort_transport(
+        self,
+        obj_id: str,
+        communicator_metadata: CommunicatorMetadata,
+    ):
+        with self._aborted_transfer_obj_ids_lock:
+            self._aborted_transfer_obj_ids.add(obj_id)
+
+    def _get_num_managed_meta_nixl(self) -> int:
+        with self._cache_lock:
+            return len(self._managed_meta_nixl)
+
+    def _get_meta(self, object_id: str) -> Optional[NixlTransportMetadata]:
+        """
+        Get the NIXL transport metadata for the given object ID if it exists
+        """
+        with self._cache_lock:
+            if object_id in self._managed_meta_nixl:
+                return self._managed_meta_nixl[object_id]
+            return None
+
+    def _put_meta(self, object_id: str, meta: NixlTransportMetadata):
+        """
+        Store the NIXL transport metadata for the given object ID
+        """
+        with self._cache_lock:
+            self._managed_meta_nixl[object_id] = meta
+
+    def _remove_tensor_descs(self, tensors: List["torch.Tensor"]):
+        """
+        Decrements the reference count for each tensor. If the count reaches 0,
+        traditionally-registered memory is deregistered from NIXL, while
+        pool-managed blocks (reg_desc is None) are returned to the pool.
+        """
+        with self._cache_lock:
+            pool_return_tensors: List["torch.Tensor"] = []
+            for tensor in tensors:
+                key = tensor.untyped_storage().data_ptr()
+                if key not in self._tensor_desc_cache:
+                    continue
+                tensor_desc = self._tensor_desc_cache[key]
+                tensor_desc.metadata_count -= 1
+                if tensor_desc.metadata_count == 0:
+                    self._tensor_desc_cache.pop(key)
+                    if tensor_desc.reg_desc is not None:
+                        # Traditional path: deregister NIXL memory.
+                        self.get_nixl_agent().deregister_memory(tensor_desc.reg_desc)
+                        self._nixl_agent_meta_version += 1
+                    else:
+                        # Pool path: return block to pool.
+                        pool_return_tensors.append(tensor)
+            if pool_return_tensors and self._memory_pool is not None:
+                self._memory_pool.free_tensors(pool_return_tensors)
+
+    def _add_tensor_descs(self, tensors: List["torch.Tensor"]):
+        """
+        If this is the first time the tensor is being registered, we register the
+        full underlying pytorch storage object with NIXL. Otherwise, we increment the reference count.
+        """
+        with self._cache_lock:
+            for tensor in tensors:
+                key = tensor.untyped_storage().data_ptr()
+                if key in self._tensor_desc_cache:
+                    self._tensor_desc_cache[key].metadata_count += 1
+                    continue
+                mem_type = "cuda" if tensor.is_cuda else "cpu"
+                # the GPU ID of the device the tensor is on.
+                # NOTE: we clip this to 0 since the GPU ID is not used for
+                # CPU tensors, and get_device returns -1 for CPU tensors.
+                # This triggers an error in nixl since it expects an unsigned.
+                gpu_id = max(tensor.get_device(), 0)
+                # Registering the full underlying pytorch storage object by
+                # constructing a memory region with the data pointer, size,
+                # GPU ID, and meta info. Doing the equivalent of what nixl
+                # does for pytorch tensors internally:
+                # https://github.com/ai-dynamo/nixl/blob/dd23ef01bd366aef89fa552f2b042f89a0b45fcb/src/api/python/_api.py#L1034
+                try:
+                    reg_desc = self.get_nixl_agent().register_memory(
+                        [
+                            (
+                                tensor.untyped_storage().data_ptr(),
+                                tensor.untyped_storage().nbytes(),
+                                gpu_id,
+                                "",
+                            )
+                        ],
+                        mem_type=mem_type,
+                    )
+                except Exception as e:
+                    # TODO(xyuzh): Remove the warning after nixl surfaces the error message
+                    if self._backend == "LIBFABRIC":
+                        troubleshooting = (
+                            "See https://github.com/ai-dynamo/nixl/blob/main/src/plugins/libfabric/README.md "
+                            "for LIBFABRIC troubleshooting. "
+                            "Set FI_LOG_LEVEL=Debug for libfabric diagnostics."
+                        )
+                    else:
+                        troubleshooting = (
+                            "See https://docs.ray.io/en/latest/ray-core/direct-transport/direct-transport.html "
+                            "for NIXL/UCX configuration. "
+                            "Set UCX_LOG_LEVEL=debug for UCX diagnostics."
+                        )
+                    vmm_hint = ""
+                    alloc_conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "").lower()
+                    if mem_type == "cuda" and "expandable_segments:true" in alloc_conf:
+                        vmm_hint = (
+                            " PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True is set; "
+                            "CUDA VMM memory can't be RDMA-registered — allocate "
+                            "transferred tensors without expandable_segments."
+                        )
+                    raise RuntimeError(
+                        f"Failed to register {mem_type} memory with NIXL "
+                        f"(backend={self._backend}, "
+                        f"size={tensor.untyped_storage().nbytes()} bytes, "
+                        f"gpu_id={gpu_id}).{vmm_hint} {troubleshooting}"
+                    ) from e
+                self._tensor_desc_cache[key] = TensorDesc(reg_desc, 1)
+
+    def _tensor_memory_registered(self, t: "torch.Tensor") -> bool:
+        """Check if the tensor's memory has been registered with NIXL."""
+        entry = self._tensor_desc_cache.get(t.untyped_storage().data_ptr())
+        return entry is not None and entry.reg_desc is not None
+
+    def _add_pool_tensor_descs(self, tensors: List["torch.Tensor"]):
+        """Add pool-managed tensor entries to the unified _tensor_desc_cache.
+
+        Pool-managed tensors use reg_desc=None since pool memory is registered
+        once at pool creation. The metadata_count tracks reference counting
+        just like traditional tensors.
+
+        Note: Entries are keyed by the source tensor's storage ``data_ptr()``.
+        If PyTorch frees and reallocates that storage address before GC runs,
+        a stale cache entry could map to an unrelated tensor. This is the same
+        constraint as the traditional (non-pool) path and is mitigated by the
+        fact that pool blocks hold a reference to pool memory, not the source
+        storage.
+        """
+        with self._cache_lock:
+            for tensor in tensors:
+                key = tensor.untyped_storage().data_ptr()
+                if key in self._tensor_desc_cache:
+                    self._tensor_desc_cache[key].metadata_count += 1
+                else:
+                    self._tensor_desc_cache[key] = TensorDesc(
+                        reg_desc=None, metadata_count=1
+                    )
+
+    def _allocate_pool_xfer_descs(self, tensors: List["torch.Tensor"]) -> Any:
+        """Allocate pool memory for tensors and return NIXL transfer descriptors.
+
+        Handles rollback of newly allocated pool blocks if get_xfer_descs
+        fails, without disturbing cached blocks from prior calls.
+        """
+        pool = self._memory_pool
+        # Remember which storages already have a pool block (cache hits)
+        # so we don't free them on rollback.
+        pre_existing = {
+            t.untyped_storage().data_ptr() for t in tensors if pool.has_block(t)
+        }
+        pool_tensor_views = pool.allocate_for_tensors(tensors)
+        try:
+            xfer_descs = self._nixl_agent.get_xfer_descs(pool_tensor_views)
+        except Exception:
+            # Only free newly allocated blocks, not cache hits.
+            new_tensors = [
+                t for t in tensors if t.untyped_storage().data_ptr() not in pre_existing
+            ]
+            if new_tensors:
+                pool.free_tensors(new_tensors)
+            raise
+        self._add_pool_tensor_descs(tensors)
+        return xfer_descs

python/ray/experimental/rdt/rdt_manager.py

@@ -0,0 +1,963 @@
+import logging
+import threading
+import time
+import warnings
+import weakref
+from collections import defaultdict
+from dataclasses import dataclass
+from queue import Queue
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    NamedTuple,
+    Optional,
+    Set,
+    Tuple,
+    Union,
+)
+
+import ray
+from ray._private import ray_constants
+from ray._raylet import ObjectRef
+from ray.experimental.rdt.tensor_transport_manager import FetchRequest
+from ray.util.annotations import PublicAPI
+
+
+@dataclass
+class ObjectStoreFetchRequest(FetchRequest):
+    """Pending fetch via the object store. Holds the remote ObjectRef to ray.get on.
+
+    Args:
+        obj_id: The RDT object ID being fetched.
+        object_ref: The ObjectRef returned by the __ray_fetch_rdt_object__ remote call.
+        tensors: Unused. Tensors are returned directly by ray.get.
+    """
+
+    object_ref: Optional[ObjectRef] = None
+    tensors: Optional[List[Any]] = None
+
+
+if TYPE_CHECKING:
+
+    from ray.experimental.rdt.rdt_store import (
+        RDTStore,
+    )
+    from ray.experimental.rdt.tensor_transport_manager import (
+        CommunicatorMetadata,
+        FetchRequest,
+        TensorTransportMetadata,
+    )
+
+logger = logging.getLogger(__name__)
+
+# RDTMeta is a named tuple containing the source actor, tensor transport
+# backend, tensor metadata, and other information that needs to be recorded.
+# - The tensor transport backend is the backend used to transport the tensors.
+# - The tensor metadata is a list of tuples, each containing the shape and dtype
+#   of a tensor in the RDT store.
+class RDTMeta(NamedTuple):
+    src_actor: "ray.actor.ActorHandle"
+    tensor_transport_backend: str
+    # This is set when the actual object is created and the metadata makes it back to the owner.
+    # For ray.put the owner is the creator so it's immediately set.
+    tensor_transport_meta: Optional["TensorTransportMetadata"]
+    # sent_dest_actors tracks the set of actor IDs that this object has been sent to.
+    # Note that since the set is mutable, it shouldn't be accessed without a lock.
+    sent_dest_actors: Set[str]
+    # sent_to_src_actor_and_others_warned indicates whether the object has already triggered a warning about being sent back to the source actor and other actors simultaneously.
+    sent_to_src_actor_and_others_warned: bool
+    # If the user set buffers for the object, the object will be fetched directly into the buffers on a ray.get
+    target_buffers: Optional[List[weakref.ReferenceType[Any]]]
+
+
+# This is used to periodically check in on the RDT transfer through the refs from
+# __ray_send__ and __ray_recv__ and abort operations in case of failures / timeouts.
+class TransferMetadata(NamedTuple):
+    src_actor: "ray.actor.ActorHandle"
+    dst_actor: "ray.actor.ActorHandle"
+    send_ref: Optional[ObjectRef]
+    recv_ref: ObjectRef
+    communicator_meta: "CommunicatorMetadata"
+    backend: str
+    obj_id: str
+    timeout: float
+
+
+@PublicAPI(stability="alpha")
+def wait_tensor_freed(tensor: Any, timeout: Optional[float] = None):
+    """
+    Wait for the tensor to be freed.
+
+    This function is useful for cases where an actor keeps a reference to a
+    tensor after returning the tensor from a task annotated with
+    `@ray.method(tensor_transport=...)`. In this case, Ray will store a
+    *reference* to the tensor, so any in-place modifications made by the actor
+    that returned the tensor could be seen by other actors. See
+    :ref:`Ray Direct Transport (RDT) <direct-transport>` for more details.
+
+    Call this function for RDT objects to ensure that all corresponding
+    `ray.ObjectRefs` have gone out of scope and therefore the tensor is safe to
+    write to again.
+
+    Args:
+        tensor: The tensor to wait to be freed. This should be a tensor that was
+            previously returned by a task annotated with
+            `@ray.method(tensor_transport=...)` or stored via
+            `ray.put(_tensor_transport="...")`.
+        timeout: The timeout in seconds to wait for all references to the tensor
+            to go out of scope. Set to None to wait indefinitely. Note that if
+            None is used, this function could hang if the `ray.ObjectRefs` that
+            refer to this tensor never go out of scope.
+    """
+    rdt_manager = ray.worker.global_worker.rdt_manager
+    rdt_manager.rdt_store.wait_tensor_freed(tensor, timeout)
+
+
+@PublicAPI(stability="alpha")
+def set_target_for_ref(ref: ObjectRef, target: List[Any]):
+    """
+    Set target buffers for an RDT ObjectRef to fetch tensors into when `ray.get` is called.
+
+    This is only supported by some transports (e.g., NIXL). If the transport
+    does not support this feature, an exception will be raised during ray.get.
+
+    Before receiving, Ray validates that the provided target buffers match the metadata
+    of the tensors in the object (e.g., shape, dtype, device). If validation fails,
+    a `ValueError` is raised. We recommend sending over lists of tensors and passing a list
+    of the same length here because the serialization order from the sender-side must match
+    the order of the target tensors here.
+
+    Args:
+        ref: The ObjectRef to set the target buffers for. The ref must be for an RDT object.
+        target: A list of tensors to be used as the target buffers to receive into.
+    """
+    rdt_manager = ray.worker.global_worker.rdt_manager
+    rdt_manager.set_target_buffers_for_ref(ref, target)
+
+
+class RDTManager:
+    def __init__(self):
+        # This lock protects _managed_rdt_metadata, _queued_transfers, and _queued_frees since
+        # they can be accessed from the user's python thread or the CoreWorker's main io service thread.
+        self._lock = threading.Lock()
+
+        # A dictionary that maps from owned object's ID to RDTMeta.
+        # This dictionary is hosted on the "driver" process of the actors that
+        # store and send/receive RDT objects.
+        self._managed_rdt_metadata: Dict[str, RDTMeta] = {}
+        # Condition variable to wait for the tensor transport meta to be set.
+        self._tensor_transport_meta_cv = threading.Condition(self._lock)
+
+        # A dictionary that maps from an object id to a list of actors
+        # that are queued to receive the object.
+        self._queued_transfers: Dict[str, List["ray.actor.ActorHandle"]] = defaultdict(
+            list
+        )
+        # A set of object ids that are queued to be freed. This is used when the object is freed
+        # before the owner knows it's created (the tensor transport metadata is not available yet).
+        self._queued_frees: Set[str] = set()
+
+        # This lock makes sure the _rdt_store and _monitor_failures_thread are only created once.
+        self._init_lock = threading.Lock()
+
+        # Per-actor local storage for RDT objects. We create the RDT store
+        # lazily, if a user specifies a non-default tensor_transport, to avoid
+        # circular import and because it imports third-party dependencies like
+        # PyTorch.
+        self._rdt_store: Optional["RDTStore"] = None
+
+        # Thread safe queue of transport refs that the monitor thread needs to start monitoring
+        self._unmonitored_transfers: Queue[TransferMetadata] = Queue()
+        # Background thread to poll on the transfer operation.
+        self._monitor_failures_thread = None
+        # Event to signal the monitor_failures thread to shutdown
+        self._monitor_failures_shutdown_event = threading.Event()
+
+        # If the actor isn't in the dict, the task to launch the custom transport registration task hasn't been submitted yet.
+        # If the value is an object ref, we have to wait for the registration task to complete.
+        # If the value is True, the actor has registered any custom transports.
+        # The value should never be False.
+        # TODO: This is a short-term solution. In the future, we'll do registration with actor initialization
+        # to make actor restarts and submitting from another worker work.
+        self.actor_id_to_transports_registered: Dict[str, Union[ObjectRef, bool]] = {}
+
+    def register_custom_transports_on_actor(self, actor: "ray.actor.ActorHandle"):
+        from ray.experimental.rdt.util import (
+            register_custom_tensor_transports_on_actor,
+        )
+
+        ref = register_custom_tensor_transports_on_actor(actor)
+        # ref is None if there are no custom transports registered.
+        self.actor_id_to_transports_registered[actor._actor_id] = (
+            True if ref is None else ref
+        )
+
+    def wait_until_custom_transports_registered(self, actor: "ray.actor.ActorHandle"):
+        actor_id = actor._actor_id
+        if actor_id not in self.actor_id_to_transports_registered:
+            self.register_custom_transports_on_actor(actor)
+
+        if self.actor_id_to_transports_registered[actor_id] is not True:
+            ray.get(self.actor_id_to_transports_registered[actor_id])
+            self.actor_id_to_transports_registered[actor_id] = True
+
+    @property
+    def rdt_store(self) -> "ray.experimental.RDTStore":
+        with self._init_lock:
+            if self._rdt_store is None:
+                from ray.experimental.rdt.rdt_store import (
+                    RDTStore,
+                )
+
+                self._rdt_store = RDTStore()
+            return self._rdt_store
+
+    def shutdown(self):
+        """
+        Interrupt and join the _monitor_failures_thread
+        """
+        with self._init_lock:
+            if self._monitor_failures_thread:
+                self._monitor_failures_shutdown_event.set()
+                self._monitor_failures_thread.join()
+                self._monitor_failures_shutdown_event.clear()
+                self._monitor_failures_thread = None
+
+    def start_monitor_thread_if_needed(self):
+        with self._init_lock:
+            # To make sure _monitor_failures_thread is started only once
+            if self._monitor_failures_thread is None:
+                self._monitor_failures_thread = threading.Thread(
+                    target=self._monitor_failures, daemon=True
+                )
+                self._monitor_failures_thread.start()
+
+    def is_managed_object(self, obj_id: str) -> bool:
+        """
+        Check if the RDT object is owned or borrowed by this process.
+        """
+        with self._lock:
+            return obj_id in self._managed_rdt_metadata
+
+    def set_rdt_metadata(self, obj_id: str, rdt_meta: RDTMeta):
+        with self._lock:
+            self._managed_rdt_metadata[obj_id] = rdt_meta
+
+    def get_rdt_metadata(self, obj_id: str) -> Optional[RDTMeta]:
+        with self._lock:
+            return self._managed_rdt_metadata.get(obj_id, None)
+
+    def wait_for_tensor_transport_metadata(
+        self, obj_id: str, timeout: float
+    ) -> Optional["TensorTransportMetadata"]:
+        with self._tensor_transport_meta_cv:
+            if self._tensor_transport_meta_cv.wait_for(
+                lambda: self._managed_rdt_metadata[obj_id].tensor_transport_meta
+                is not None,
+                timeout=timeout,
+            ):
+                return self._managed_rdt_metadata[obj_id].tensor_transport_meta
+            else:
+                return None
+
+    def _monitor_failures(self):
+        """
+        Monitor the refs from send and recv tasks and abort the transfers
+        if they error out or timeout to prevent hanging.
+        """
+        not_done = []
+        done = []
+        ref_info_map = {}
+        while not self._monitor_failures_shutdown_event.is_set():
+            while not self._unmonitored_transfers.empty():
+                ref_info = self._unmonitored_transfers.get()
+                if ref_info.send_ref:
+                    not_done.append(ref_info.send_ref)
+                    ref_info_map[ref_info.send_ref.hex()] = ref_info
+                not_done.append(ref_info.recv_ref)
+                ref_info_map[ref_info.recv_ref.hex()] = ref_info
+            if len(not_done) > 0:
+                done, not_done = ray.wait(not_done, num_returns=1, timeout=1)
+            if len(done) > 0:
+                try:
+                    ray.get(done[0])
+                    ref_info_map.pop(done[0].hex(), None)
+                except Exception as e:
+                    self._abort_transport(done[0], ref_info_map, e)
+
+            while len(not_done) > 0:
+                if not_done[0].hex() not in ref_info_map:
+                    # The associated transfer was already aborted.
+                    not_done.pop(0)
+                elif ref_info_map[not_done[0].hex()].timeout < time.time():
+                    self._abort_transport(
+                        not_done[0],
+                        ref_info_map,
+                        TimeoutError(
+                            f"RDT transfer failed after {ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS}s. "
+                            "You can increase the timeout by setting RAY_rdt_fetch_fail_timeout_milliseconds"
+                        ),
+                    )
+                else:
+                    # wait returns lists in the same order they were passed in, so if
+                    # the timeout of first hasn't been reached, neither have the others.
+                    break
+            if len(not_done) == 0:
+                # If we emptied out _unmonitored_transfers on this iteration, wait for a bit.
+                self._monitor_failures_shutdown_event.wait(1)
+
+    def _abort_transport(
+        self,
+        failed_ref: ObjectRef,
+        ref_info_map: Dict[str, TransferMetadata],
+        exception: Exception,
+    ):
+        """
+        Cleans up the ref_info_map, kill the src and dst actors, and destroy the
+        collective group if necessary.
+        """
+        from ray.experimental.collective import destroy_collective_group
+        from ray.experimental.rdt.collective_tensor_transport import (
+            CollectiveCommunicatorMetadata,
+        )
+        from ray.experimental.rdt.rdt_store import (
+            __ray_abort_transport__,
+        )
+        from ray.experimental.rdt.util import (
+            get_tensor_transport_manager,
+        )
+
+        ref_info = ref_info_map.pop(failed_ref.hex(), None)
+        if ref_info is None:
+            return
+
+        if ref_info.send_ref:
+            ref_info_map.pop(ref_info.send_ref.hex(), None)
+        ref_info_map.pop(ref_info.recv_ref.hex(), None)
+
+        tensor_transport_manager = get_tensor_transport_manager(ref_info.backend)
+        if tensor_transport_manager.can_abort_transport():
+            if not tensor_transport_manager.__class__.is_one_sided():
+                # This is dead code until we implement a NCCL abort since NIXL
+                # is the only abortable transport for now and is one-sided.
+                ref_info.src_actor.__ray_call__.options(
+                    concurrency_group="_ray_system_error"
+                ).remote(
+                    __ray_abort_transport__,
+                    ref_info.obj_id,
+                    ref_info.communicator_meta,
+                    ref_info.backend,
+                )
+            ref_info.dst_actor.__ray_call__.options(
+                concurrency_group="_ray_system_error"
+            ).remote(
+                __ray_abort_transport__,
+                ref_info.obj_id,
+                ref_info.communicator_meta,
+                ref_info.backend,
+            )
+            logger.info(
+                "RDT transfer with src actor %s and dst actor %s failed due to %s.",
+                ref_info.src_actor,
+                ref_info.dst_actor,
+                exception,
+            )
+        else:
+            # TODO(#51276): Kill all actors in the collective group when we support more collective operations
+            ray.kill(ref_info.src_actor)
+            ray.kill(ref_info.dst_actor)
+            logger.error(
+                "RDT transfer with src actor %s and dst actor %s failed. Killing the actors. "
+                "Transfer failed with exception: %s",
+                ref_info.src_actor,
+                ref_info.dst_actor,
+                exception,
+            )
+
+        # isinstance does an implicit cast and makes communicator_name inaccessible
+        # so we have to get communicator_name before the cast.
+        if isinstance(ref_info.communicator_meta, CollectiveCommunicatorMetadata):
+            try:
+                collective_group_name = ref_info.communicator_meta.communicator_name
+                destroy_collective_group(collective_group_name)
+                logger.error(
+                    "Destroyed collective group %s due to a hanging/failed RDT transfer",
+                    collective_group_name,
+                )
+            except ValueError:
+                # Collective group was already destroyed
+                pass
+
+    def add_rdt_ref(
+        self,
+        obj_ref: ObjectRef,
+        src_actor: "ray.actor.ActorHandle",
+        tensor_transport: str,
+        tensor_transport_meta: Optional["TensorTransportMetadata"] = None,
+    ):
+        """Add an RDT object reference to the RDT manager. This should be
+        called whenever the current process calls a task that is annotated with
+        `@ray.method(tensor_transport=...)`.
+
+        Args:
+            obj_ref: The ObjectRef of the task output.
+            src_actor: The actor that executes the task and that creates the RDT object.
+            tensor_transport: The tensor transport protocol to use for the RDT object.
+            tensor_transport_meta: The tensor transport metadata that is pre-computed.
+                This is known at ref creation time if the object is created through ray.put.
+        """
+        self.set_rdt_metadata(
+            obj_ref.hex(),
+            RDTMeta(
+                src_actor=src_actor,
+                tensor_transport_backend=tensor_transport,
+                tensor_transport_meta=tensor_transport_meta,  # None if not from ray.put
+                sent_dest_actors=set(),
+                sent_to_src_actor_and_others_warned=False,
+                target_buffers=None,
+            ),
+        )
+
+    def set_tensor_transport_metadata_and_trigger_queued_operations(
+        self, obj_id: str, tensor_transport_meta: "TensorTransportMetadata"
+    ):
+        """
+        Sets the tensor transport metadata for an object and triggers any queued
+        up transfers or frees for that object.
+        """
+        dst_actors = None
+        free_object = False
+        with self._tensor_transport_meta_cv:
+            self._managed_rdt_metadata[obj_id] = self._managed_rdt_metadata[
+                obj_id
+            ]._replace(tensor_transport_meta=tensor_transport_meta)
+            dst_actors = self._queued_transfers.pop(obj_id, None)
+            free_object = obj_id in self._queued_frees
+            if free_object:
+                self._queued_frees.remove(obj_id)
+                # There shouldn't be any transfers queued if the free was queued,
+                # since we clear the queued transfers when queueing the free.
+                assert dst_actors is None
+            self._tensor_transport_meta_cv.notify_all()
+
+        if free_object:
+            self.free_object_primary_copy(obj_id)
+        if dst_actors:
+            for dst_actor in dst_actors:
+                # Trigger the transfer now that the metadata is available.
+                self.trigger_out_of_band_tensor_transfer(dst_actor, obj_id)
+
+    def set_target_buffers_for_ref(self, ref: ObjectRef, target_buffers: List[Any]):
+        with self._lock:
+            if ref.hex() not in self._managed_rdt_metadata:
+                raise ValueError(f"Ref {ref} is not an RDT object.")
+
+            self._managed_rdt_metadata[ref.hex()] = self._managed_rdt_metadata[
+                ref.hex()
+            ]._replace(
+                target_buffers=[
+                    weakref.ref(target_buffer) for target_buffer in target_buffers
+                ]
+            )
+
+    def _trigger_fetch(
+        self,
+        obj_id: str,
+        use_object_store: bool,
+    ) -> FetchRequest:
+        """
+        Start fetching an RDT object.
+
+        If the specified transport supports async fetches, this will trigger the
+        fetch without blocking. Note that this always triggers a fetch, even if
+        the object is already in the store.
+
+        Args:
+            obj_id: The object ID of the RDT object.
+            use_object_store: Whether to fetch through the object store or through
+                the designated one-sided tensor transport.
+
+        Returns:
+            A FetchRequest. Wait on the FetchRequest to get the tensors.
+        """
+        from ray.experimental.rdt.rdt_store import (
+            __ray_fetch_rdt_object__,
+        )
+        from ray.experimental.rdt.util import (
+            get_tensor_transport_manager,
+            is_one_sided_transport,
+        )
+
+        rdt_meta = self.get_rdt_metadata(obj_id)
+        assert rdt_meta is not None
+
+        if use_object_store:
+            if rdt_meta.target_buffers:
+                logger.warning(
+                    "Target buffers are not supported for use_object_store=True. Ignoring the target buffers."
+                )
+
+            src_actor = rdt_meta.src_actor
+            object_ref = src_actor.__ray_call__.options(
+                concurrency_group="_ray_system"
+            ).remote(__ray_fetch_rdt_object__, obj_id)
+            return ObjectStoreFetchRequest(
+                obj_id=obj_id, object_ref=object_ref, tensors=[]
+            )
+        else:
+            tensor_transport = rdt_meta.tensor_transport_backend
+            if not is_one_sided_transport(tensor_transport):
+                raise ValueError(
+                    f"ray.get is not allowed on RDT objects using the two-sided transport {tensor_transport}. "
+                    "Either use a one-sided RDT transport or pass _use_object_store=True to ray.get to fetch the object through the object store instead."
+                )
+            tensor_transport_manager = get_tensor_transport_manager(tensor_transport)
+            communicator_meta = tensor_transport_manager.get_communicator_metadata(
+                None, None, tensor_transport
+            )
+
+            tensor_transport_meta = rdt_meta.tensor_transport_meta
+            if tensor_transport_meta is None:
+                # We can't fetch the object until we know the creator has actually created the object.
+                timeout = ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS
+                tensor_transport_meta = self.wait_for_tensor_transport_metadata(
+                    obj_id, timeout
+                )
+                if tensor_transport_meta is None:
+                    raise TimeoutError(
+                        f"Timed out after {timeout}s waiting for object {obj_id} to be created while trying to get the object. "
+                        "You can increase the timeout by setting RAY_rdt_fetch_fail_timeout_milliseconds."
+                    )
+
+            target_buffers = None
+            if rdt_meta.target_buffers:
+                # Try to get the target buffers from the weak references. If any of the
+                # target buffers are not alive, we just won't use the target buffers.
+                target_buffers = []
+                for target_buffer in rdt_meta.target_buffers:
+                    buffer = target_buffer()
+                    if buffer is None:
+                        target_buffers = None
+                        break
+                    else:
+                        target_buffers.append(buffer)
+
+            if target_buffers is not None:
+                from ray.experimental.rdt.rdt_store import validate_tensor_buffers
+
+                device = tensor_transport_meta.tensor_device
+                tensor_meta = tensor_transport_meta.tensor_meta
+                validate_tensor_buffers(target_buffers, tensor_meta, device)
+
+            return tensor_transport_manager.fetch_multiple_tensors(
+                obj_id,
+                tensor_transport_meta,
+                communicator_meta,
+                target_buffers,
+            )
+
+    def _wait_fetch(
+        self, obj_id: str, fetch_request: FetchRequest, timeout: float = -1
+    ) -> List[Any]:
+        """
+        Waits for a previously triggered fetch to complete and returns the tensors.
+
+        Args:
+            obj_id: The object ID of the RDT object.
+            fetch_request: An ObjectStoreFetchRequest representing an object
+                transferred via Ray's object store or a FetchRequest
+                representing an object transferred via a tensor transport.
+            timeout: Maximum time in seconds to wait. -1 means wait indefinitely.
+                0 means return immediately if not ready.
+
+        Returns:
+            The list of tensors fetched.
+        """
+        if isinstance(fetch_request, ObjectStoreFetchRequest):
+            return ray.get(fetch_request.object_ref, timeout=timeout)
+        else:
+            from ray.experimental.rdt.util import get_tensor_transport_manager
+
+            rdt_meta = self.get_rdt_metadata(obj_id)
+            tensor_transport_manager = get_tensor_transport_manager(
+                rdt_meta.tensor_transport_backend
+            )
+            return tensor_transport_manager.wait_fetch_complete(
+                fetch_request, timeout=timeout
+            )
+
+    def queue_or_trigger_out_of_band_tensor_transfer(
+        self, dst_actor: "ray.actor.ActorHandle", task_args: Tuple[Any, ...]
+    ):
+        """
+        Triggers the transfer if the tensor metadata is available for the object. If it's
+        not available, the transfer is queued up until the metadata is available.
+        """
+        rdt_object_ids: Set[str] = set()
+        for arg in task_args:
+            # If an ObjectRef is managed, it means the actual value is a list of tensors stored
+            # on a remote actor. Therefore, this function will trigger a tensor communication
+            # operation between the sender and receiver actors.
+            if not isinstance(arg, ObjectRef):
+                continue
+            obj_id = arg.hex()
+            if self.is_managed_object(obj_id):
+                rdt_object_ids.add(obj_id)
+        if rdt_object_ids:
+            self.wait_until_custom_transports_registered(dst_actor)
+            for obj_id in rdt_object_ids:
+                # Atomically gets the tensor transport metadata for an object and queues up a transfer
+                # if the tensor transport metadata is not available.
+                with self._lock:
+                    tensor_transport_meta = self._managed_rdt_metadata[
+                        obj_id
+                    ].tensor_transport_meta
+                    if tensor_transport_meta is None:
+                        self._queued_transfers[obj_id].append(dst_actor)
+                if tensor_transport_meta is not None:
+                    self.trigger_out_of_band_tensor_transfer(dst_actor, obj_id)
+
+    def trigger_out_of_band_tensor_transfer(
+        self, dst_actor: "ray.actor.ActorHandle", obj_id: str
+    ):
+        """
+        Triggers tensor communication operations between actors. When a managed ObjectRef is passed
+        to another actor task, CPU data will still be passed through the object store, but the in-actor
+        tensors will be passed out-of-band.
+
+        This function triggers the out-of-band tensor transfer by submitting Ray actor
+        tasks `__ray_send__` to the sender actor and `__ray_recv__` to the receiver actor to initiate
+        tensor communication using protocols like NCCL or GLOO.
+
+        Before the receiver actor executes the actor task, the deserializer combines the
+        CPU data with the tensors from the sender actor to reconstruct the original task output
+        generated by the sender actor.
+
+        Args:
+            dst_actor: The target actor to receive tensors
+            obj_id: ID of the object to send to the dst_actor.
+
+        Returns:
+            None
+        """
+        from ray.experimental.rdt.rdt_store import (
+            __ray_recv__,
+            __ray_send__,
+        )
+        from ray.experimental.rdt.util import (
+            get_tensor_transport_manager,
+        )
+
+        with self._lock:
+            # Since sent_dest_actors is mutable, this whole block needs to be protected.
+            rdt_meta = self._managed_rdt_metadata[obj_id]
+            src_actor = rdt_meta.src_actor
+            tensor_transport_meta = rdt_meta.tensor_transport_meta
+
+            # Update the set of destination actors for this object
+            # The set inside NamedTuple is mutable, so we can modify it directly
+            rdt_meta.sent_dest_actors.add(dst_actor._actor_id)
+            # Check if a warning should be triggered for this object:
+            # 1. object has not triggered a warning yet.
+            # 2. object is sent back to its source actor.
+            # 3. object is also sent to at least one other actor
+            if (
+                not rdt_meta.sent_to_src_actor_and_others_warned
+                and src_actor._actor_id in rdt_meta.sent_dest_actors
+                and len(rdt_meta.sent_dest_actors) > 1
+            ):
+                warnings.warn(
+                    f"RDT ObjectRef({obj_id}) is being passed back to the actor that created it {src_actor}. "
+                    "Note that RDT objects are mutable. If the tensor is modified, Ray's internal copy will "
+                    "also be updated, and subsequent passes to other actors will receive the updated version "
+                    "instead of the original.",
+                    UserWarning,
+                )
+                # Mark the object as warned so that we don't warn again for this object.
+                self._managed_rdt_metadata[obj_id] = rdt_meta._replace(
+                    sent_to_src_actor_and_others_warned=True
+                )
+
+            if src_actor._actor_id == dst_actor._actor_id:
+                # If the source and destination actors are the same, the tensors can
+                # be transferred intra-process, so we skip the out-of-band tensor
+                # transfer.
+                return
+
+            tensor_transport_manager = get_tensor_transport_manager(
+                rdt_meta.tensor_transport_backend
+            )
+            communicator_meta = tensor_transport_manager.get_communicator_metadata(
+                src_actor,
+                dst_actor,
+                rdt_meta.tensor_transport_backend,
+            )
+
+            send_ref = None
+            if not tensor_transport_manager.__class__.is_one_sided():
+                # Send tensors stored in the `src_actor`'s GPU object store to the
+                # destination rank `dst_rank`.
+                # NOTE: We put this task on the background thread to avoid tasks
+                # executing on the main thread blocking the data transfer.
+                send_ref = src_actor.__ray_call__.options(
+                    concurrency_group="_ray_system"
+                ).remote(
+                    __ray_send__,
+                    obj_id,
+                    tensor_transport_meta,
+                    communicator_meta,
+                    rdt_meta.tensor_transport_backend,
+                )
+
+            # Receive tensors from the source rank and store them in the
+            # `dst_actor`'s GPU object store.
+            # NOTE: Putting this task on the background thread is technically only
+            # needed for the sender task, but we put the receiver task on the same
+            # background thread to ensure that all communication operations are
+            # executed in a global order.
+            recv_ref = dst_actor.__ray_call__.options(
+                concurrency_group="_ray_system"
+            ).remote(
+                __ray_recv__,
+                obj_id,
+                tensor_transport_meta,
+                communicator_meta,
+                rdt_meta.tensor_transport_backend,
+            )
+
+        self._unmonitored_transfers.put(
+            TransferMetadata(
+                src_actor=src_actor,
+                dst_actor=dst_actor,
+                send_ref=send_ref,
+                recv_ref=recv_ref,
+                communicator_meta=communicator_meta,
+                backend=rdt_meta.tensor_transport_backend,
+                obj_id=obj_id,
+                timeout=time.time() + ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS,
+            )
+        )
+        self.start_monitor_thread_if_needed()
+
+    def get_rdt_objects(
+        self,
+        object_ids: List[str],
+    ) -> Dict[str, List[Any]]:
+        """
+        Get RDT objects that have already been transferred (e.g. via __ray_recv__).
+
+        This is used in the task argument deserialization path where the
+        out-of-band tensor transfer has already been triggered by the caller.
+        It only waits on the local RDT store for the tensors to arrive.
+
+        Args:
+            object_ids: The object IDs of the RDT objects.
+
+        Returns:
+            A dict mapping object ID to the RDT object (list of tensors).
+        """
+        rdt_store = self.rdt_store
+        result: Dict[str, List[Any]] = {}
+        for object_id in object_ids:
+            pop_object = not rdt_store.is_primary_copy(object_id)
+            if pop_object:
+                result[object_id] = rdt_store.wait_and_pop_object(
+                    object_id, timeout=ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS
+                )
+            else:
+                result[object_id] = rdt_store.wait_and_get_object(
+                    object_id, timeout=ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS
+                )
+        return result
+
+    def fetch_and_get_rdt_objects(
+        self,
+        object_ids: List[str],
+        timeout: Optional[float] = None,
+        use_object_store: bool = False,
+    ) -> Dict[str, List[Any]]:
+        """
+        Fetch and get RDT objects for a list of object IDs, pipelining async fetches.
+
+        This is used in the ray.get codepath where the caller initiates the
+        tensor fetch. For one-sided transports (e.g. NIXL), all transfers are
+        triggered first before waiting, eliminating serial transfer latency.
+
+        Args:
+            object_ids: The object IDs of the RDT objects.
+            timeout: The user-specified timeout from ray.get, or None for no
+                user timeout. The actual deadline is the minimum of this and
+                RDT_FETCH_FAIL_TIMEOUT_SECONDS.
+            use_object_store: Whether to fetch through the object store or through
+                the designated tensor transport.
+
+        Returns:
+            A dict mapping object ID to the RDT object (list of tensors).
+
+        Raises:
+            GetTimeoutError: If the user-specified timeout is exceeded.
+            ObjectFetchTimedOutError: If RDT_FETCH_FAIL_TIMEOUT_SECONDS is exceeded.
+        """
+        from ray.exceptions import GetTimeoutError, ObjectFetchTimedOutError
+
+        rdt_timeout = ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS
+        now = time.time()
+        if timeout is not None and timeout >= 0:
+            rdt_deadline = now + rdt_timeout
+            user_deadline = now + timeout
+            if user_deadline < rdt_deadline:
+                deadline = user_deadline
+                user_timeout_is_smaller = True
+            else:
+                deadline = rdt_deadline
+                user_timeout_is_smaller = False
+        else:
+            deadline = now + rdt_timeout
+            user_timeout_is_smaller = False
+
+        rdt_store = self.rdt_store
+        result: Dict[str, List[Any]] = {}
+
+        # First, try to get objects that are already available in the store
+        # These are primary copies, or secondary copies created via
+        # __ray_recv__ that haven't been consumed yet.
+        if not use_object_store:
+            for object_id in object_ids:
+                try:
+                    result[object_id] = rdt_store.wait_and_get_object(
+                        object_id, timeout=0
+                    )
+                except TimeoutError:
+                    pass
+
+        # For remaining objects, trigger fetches.
+        fetch_requests: Dict[str, "FetchRequest"] = {}
+        for object_id in object_ids:
+            if object_id in result:
+                continue
+            assert self.is_managed_object(
+                object_id
+            ), f"No metadata found for {object_id}"
+
+            fetch_requests[object_id] = self._trigger_fetch(object_id, use_object_store)
+
+        # Wait for all in-flight fetches to complete.
+        while fetch_requests:
+            object_id, fetch_request = fetch_requests.popitem()
+            remaining = deadline - time.time()
+            if remaining < 0:
+                if user_timeout_is_smaller:
+                    # User passed a timeout to ray.get that expired.
+                    raise GetTimeoutError(f"ray.get timed out after {timeout}s.")
+                else:
+                    # Object fetch timeout expired. Throw an error in case we
+                    # hung.
+                    raise ObjectFetchTimedOutError(
+                        object_ref_hex=object_id,
+                        owner_address="",
+                        call_site="",
+                    )
+            try:
+                result[object_id] = self._wait_fetch(
+                    object_id, fetch_request, timeout=remaining
+                )
+            except (TimeoutError, GetTimeoutError):
+                if user_timeout_is_smaller:
+                    raise GetTimeoutError(f"ray.get timed out after {timeout}s.")
+                else:
+                    raise ObjectFetchTimedOutError(
+                        object_ref_hex=object_id,
+                        owner_address="",
+                        call_site="",
+                    )
+        return result
+
+    def queue_or_free_object_primary_copy(self, object_id: str):
+        """
+        Free the RDT object on the primary copy holder and free metadata
+        if the tensor metadata is available (the object has been created).
+        Otherwise, queue up the free operation until the tensor metadata is available.
+        """
+        # NOTE: This may have to change if we support lineage reconstruction for RDT
+        # TODO(#57962): Metadata is currently not removed on borrowers that borrow through
+        # the NIXL ray.put / ray.get
+        with self._lock:
+            self._queued_transfers.pop(object_id, None)
+            rdt_meta = self._managed_rdt_metadata[object_id]
+            tensor_transport_meta = rdt_meta.tensor_transport_meta
+            if tensor_transport_meta is None:
+                # The object hasn't been created at the time of the free.
+                self._queued_frees.add(object_id)
+
+        if tensor_transport_meta is not None:
+            self.free_object_primary_copy(object_id)
+
+    def free_object_primary_copy(self, object_id: str):
+        from ray.experimental.rdt.rdt_store import (
+            __ray_free__,
+        )
+
+        with self._lock:
+            rdt_meta = self._managed_rdt_metadata.pop(object_id)
+        # TODO(#60434): Once the driver has some notion about where rdt args are stored, we can
+        # integrate __ray_free__ with the current free objects RPC and avoid having to call
+        # an actor task here.
+        if not ray.is_initialized():
+            return
+        src_actor = rdt_meta.src_actor
+        tensor_transport_backend = rdt_meta.tensor_transport_backend
+        tensor_transport_meta = rdt_meta.tensor_transport_meta
+        src_actor.__ray_call__.options(concurrency_group="_ray_system").remote(
+            __ray_free__,
+            object_id,
+            tensor_transport_backend,
+            tensor_transport_meta,
+        )
+
+    @staticmethod
+    def actor_has_tensor_transport(
+        actor: "ray.actor.ActorHandle", tensor_transport: str
+    ):
+        """
+        Check if the actor has a communicator for the given tensor transport backend.
+
+        Args:
+            actor: The actor to check.
+            tensor_transport: The tensor transport backend to check.
+
+        Returns:
+            True if the actor has a communicator for the given tensor transport backend, False otherwise.
+        """
+        from ray.experimental.rdt.util import (
+            get_tensor_transport_manager,
+        )
+
+        tensor_transport_manager = get_tensor_transport_manager(tensor_transport)
+        return tensor_transport_manager.actor_has_tensor_transport(actor)
+
+    def put_object(
+        self,
+        obj_ref: ObjectRef,
+        tensor_transport: str,
+        tensors: List[Any],
+    ):
+        """
+        Put the RDT object into the RDT manager.
+
+        Args:
+            obj_ref: The object ref of the RDT object.
+            tensor_transport: The tensor transport backend to use.
+            tensors: The tensors to put into the RDT manager.
+        """
+        src_actor = ray.get_runtime_context().current_actor
+        tensor_transport_meta = self.rdt_store.add_object_primary(
+            obj_ref.hex(), tensors, tensor_transport
+        )
+        self.add_rdt_ref(
+            obj_ref,
+            src_actor,
+            tensor_transport,
+            tensor_transport_meta=tensor_transport_meta,
+        )

python/ray/experimental/rdt/rdt_store.py

@@ -0,0 +1,370 @@
+import threading
+from collections import defaultdict, deque
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
+
+from ray.experimental.rdt.tensor_transport_manager import (
+    CommunicatorMetadata,
+    TensorTransportMetadata,
+)
+from ray.experimental.rdt.util import (
+    device_match_transport,
+    get_tensor_transport_manager,
+)
+
+if TYPE_CHECKING:
+    import torch
+
+
+def __ray_send__(
+    self,
+    obj_id: str,
+    tensor_transport_meta: TensorTransportMetadata,
+    communicator_meta: CommunicatorMetadata,
+    backend: str,
+):
+    """Helper function that runs on the src actor to send tensors to the dst actor."""
+    from ray._private.worker import global_worker
+
+    rdt_store = global_worker.rdt_manager._rdt_store
+    assert rdt_store.has_object(obj_id), f"obj_id={obj_id} not found in RDT store"
+
+    tensors = rdt_store.get_object(obj_id)
+
+    tensor_transport_manager = get_tensor_transport_manager(backend)
+    tensor_transport_manager.send_multiple_tensors(
+        tensors,
+        tensor_transport_meta,
+        communicator_meta,
+    )
+
+
+def validate_tensor_buffers(
+    tensor_buffers: List["torch.Tensor"],
+    tensor_meta: List[Tuple["torch.Size", "torch.dtype"]],
+    device: str,
+):
+    if len(tensor_buffers) != len(tensor_meta):
+        raise ValueError(
+            f"Length of tensor_buffers ({len(tensor_buffers)}) does not match length from object metadata ({len(tensor_meta)})."
+        )
+
+    def tensor_buffer_mismatch_msg(prop, idx, actual, expected):
+        return f"{prop} of tensor_buffer at index {idx} ({actual}) does not match {prop.lower()} from object metadata ({expected})."
+
+    for idx, single_buffer in enumerate(tensor_buffers):
+        shape, dtype = tensor_meta[idx]
+        if single_buffer.shape != shape:
+            raise ValueError(
+                tensor_buffer_mismatch_msg("Shape", idx, single_buffer.shape, shape)
+            )
+        if single_buffer.dtype != dtype:
+            raise ValueError(
+                tensor_buffer_mismatch_msg("Dtype", idx, single_buffer.dtype, dtype)
+            )
+        if single_buffer.device.type != device:
+            raise ValueError(
+                tensor_buffer_mismatch_msg(
+                    "Device", idx, single_buffer.device.type, device
+                )
+            )
+        if not single_buffer.is_contiguous():
+            raise ValueError(f"Tensor buffer at index {idx} is not contiguous.")
+
+
+def __ray_recv__(
+    self,
+    obj_id: str,
+    tensor_transport_meta: TensorTransportMetadata,
+    communicator_meta: CommunicatorMetadata,
+    backend: str,
+    target_buffers: Optional[List[Any]] = None,
+):
+    """Helper function that runs on the dst actor to receive tensors from the src actor."""
+    from ray._private.worker import global_worker
+
+    rdt_store = global_worker.rdt_manager.rdt_store
+    try:
+        tensor_transport_manager = get_tensor_transport_manager(backend)
+        if target_buffers:
+            # Currently only torch tensors are supported as target buffers. We could make this
+            # more generic in the future by adding a pluggable buffer validation function.
+            validate_tensor_buffers(
+                target_buffers,
+                tensor_transport_meta.tensor_meta,
+                tensor_transport_meta.tensor_device,
+            )
+        tensors = tensor_transport_manager.recv_multiple_tensors(
+            obj_id,
+            tensor_transport_meta,
+            communicator_meta,
+            target_buffers,
+        )
+        assert len(tensors) == len(tensor_transport_meta.tensor_meta)
+        rdt_store.add_object(obj_id, tensors)
+    except Exception as e:
+        # Store the error as an RDT object if the recv fails, so waiters will raise the error.
+        rdt_store.add_object(obj_id, e)
+
+
+def __ray_abort_transport__(
+    self, obj_id: str, communicator_meta: CommunicatorMetadata, backend: str
+):
+    """Helper function that can run on an actor doing a send or recv to abort the transport."""
+    tensor_transport_manager = get_tensor_transport_manager(backend)
+    tensor_transport_manager.abort_transport(obj_id, communicator_meta)
+
+
+def __ray_free__(
+    self,
+    obj_id: str,
+    tensor_transport_backend: str,
+    tensor_transport_meta: TensorTransportMetadata,
+):
+    try:
+        from ray._private.worker import global_worker
+
+        tensor_transport_manager = get_tensor_transport_manager(
+            tensor_transport_backend
+        )
+        rdt_manager = global_worker.rdt_manager
+        rdt_store = rdt_manager.rdt_store
+
+        if not rdt_store.has_object(obj_id):
+            return
+        tensors = rdt_store.get_object(obj_id)
+        tensor_transport_manager.garbage_collect(obj_id, tensor_transport_meta, tensors)
+
+        rdt_store.pop_object(obj_id)
+    except AssertionError:
+        # This could fail if this is a retry and it's already been freed.
+        pass
+
+
+def __ray_fetch_rdt_object__(self, obj_id: str):
+    """Helper function that runs on the src actor to fetch tensors from the RDT store via the object store."""
+    from ray._private.worker import global_worker
+
+    rdt_store = global_worker.rdt_manager.rdt_store
+    rdt_object = rdt_store.wait_and_get_object(obj_id)
+    return rdt_object
+
+
+@dataclass
+class _RDTObject:
+    # A list of tensors representing the RDT object.
+    data: List[Any]
+    # Whether the RDT object is the primary copy.
+    is_primary: bool
+    # If a recv failed, we store the error here.
+    error: Optional[Exception] = None
+
+
+class RDTStore:
+    """
+    This class is thread-safe. The GPU object store is meant to be read and
+    written by the following threads:
+    1. The main thread, which is executing user code. This thread may get, put,
+    and pop objects.
+    2. The background _ray_system thread, which executes data transfers. This
+    thread may get and put objects.
+    3. The background CoreWorker server thread, which executes garbage
+    collection callbacks that pop objects that are no longer in use.
+    """
+
+    def __init__(self):
+        # A dictionary that maps from an object ID to a queue of tensor lists.
+        #
+        # Note: Currently, `_rdt_store` is only supported for Ray Actors.
+        self._rdt_store: Dict[str, deque[_RDTObject]] = defaultdict(deque)
+        # Mapping from tensor data pointer to the IDs of objects that contain it.
+        self._tensor_to_object_ids: Dict[int, Set[str]] = defaultdict[int, Set[str]](
+            set
+        )
+        # Synchronization for the RDT store.
+        self._lock = threading.RLock()
+        # Signal when an object becomes present in the object store.
+        self._object_present_cv = threading.Condition(self._lock)
+        # Signal when an object is freed from the object store.
+        self._object_freed_cv = threading.Condition(self._lock)
+
+    def has_object(self, obj_id: str) -> bool:
+        with self._lock:
+            existed = obj_id in self._rdt_store
+            if existed:
+                return len(self._rdt_store[obj_id]) > 0
+            return existed
+
+    def has_tensor(self, tensor: Any) -> bool:
+        # Method only used for testing.
+        with self._lock:
+            return id(tensor) in self._tensor_to_object_ids
+
+    def get_object(self, obj_id: str) -> Optional[List[Any]]:
+        with self._lock:
+            if self._rdt_store[obj_id][0].error:
+                raise self._rdt_store[obj_id][0].error
+            return self._rdt_store[obj_id][0].data
+
+    def add_object(
+        self,
+        obj_id: str,
+        rdt_object: Union[List[Any], Exception],
+        is_primary: bool = False,
+    ):
+        """
+        Add an RDT object to the RDT store.
+
+        Args:
+            obj_id: The object ID of the RDT object.
+            rdt_object: A list of tensors representing the RDT object.
+            is_primary: Whether the RDT object is the primary copy.
+        """
+        with self._object_present_cv:
+            if isinstance(rdt_object, Exception):
+                self._rdt_store[obj_id].append(
+                    _RDTObject([], is_primary, error=rdt_object)
+                )
+            else:
+                for tensor in rdt_object:
+                    self._tensor_to_object_ids[id(tensor)].add(obj_id)
+                # Append to the queue instead of overwriting
+                self._rdt_store[obj_id].append(
+                    _RDTObject(
+                        rdt_object,
+                        is_primary,
+                    )
+                )
+            self._object_present_cv.notify_all()
+
+    def add_object_primary(
+        self, obj_id: str, tensors: List[Any], tensor_transport: str
+    ) -> TensorTransportMetadata:
+        with self._object_present_cv:
+            # A primary entry may already exist from a prior attempt of the
+            # same task (e.g., a task that succeeded and populated the RDT
+            # store but whose reply was lost, then got retried). Keep the
+            # existing primary — do not re-store — and return metadata
+            # derived from it so the metadata matches what `__ray_send__`
+            # will actually transmit.
+            queue = self._rdt_store.get(obj_id)
+            if queue:
+                tensors_to_describe = queue[0].data
+            else:
+                self.add_object(obj_id, tensors, is_primary=True)
+                tensors_to_describe = tensors
+
+        tensor_transport_manager = get_tensor_transport_manager(tensor_transport)
+        tensor_transport_meta = (
+            tensor_transport_manager.extract_tensor_transport_metadata(
+                obj_id, tensors_to_describe
+            )
+        )
+
+        if tensor_transport_meta.tensor_meta and not device_match_transport(
+            tensor_transport_meta.tensor_device, tensor_transport
+        ):
+            raise ValueError(
+                f"Tensor transport backend {tensor_transport} does not support "
+                f"tensor transfer on device {tensor_transport_meta.tensor_device}."
+            )
+
+        return tensor_transport_meta
+
+    def is_primary_copy(self, obj_id: str) -> bool:
+        with self._lock:
+            return self.has_object(obj_id) and self._rdt_store[obj_id][0].is_primary
+
+    def wait_and_get_object(
+        self, obj_id: str, timeout: Optional[float] = None
+    ) -> List[Any]:
+        """Atomically waits for the RDT object to be present in the RDT
+        store, then gets it. If the object is not present after the optional
+        timeout, raise a TimeoutError.
+
+        Args:
+            obj_id: The object ID to wait for.
+            timeout: The maximum time in seconds to wait for the object to be
+                present in the RDT store. If not specified, wait indefinitely.
+
+        Returns:
+            The tensors in the RDT object.
+        """
+        with self._lock:
+            self._wait_object(obj_id, timeout)
+            return self.get_object(obj_id)
+
+    def wait_and_pop_object(
+        self, obj_id: str, timeout: Optional[float] = None
+    ) -> List[Any]:
+        """Atomically waits for the RDT object to be present in the RDT
+        store, then pops it.  If the object is not present after the optional
+        timeout, raise a TimeoutError.
+
+        Args:
+            obj_id: The object ID to wait for.
+            timeout: The maximum time in seconds to wait for the object to be
+                present in the RDT store. If not specified, wait indefinitely.
+
+        Returns:
+            The RDT object.
+        """
+        with self._lock:
+            self._wait_object(obj_id, timeout)
+            return self.pop_object(obj_id)
+
+    def _wait_object(self, obj_id: str, timeout: Optional[float] = None) -> None:
+        """Helper method to wait for the RDT object to be present in the RDT store.
+        If the object is not present after the optional timeout, raise a
+        TimeoutError.
+
+        Args:
+            obj_id: The object ID to wait for.
+            timeout: The maximum time in seconds to wait for the object to be
+                present in the RDT store. If not specified, wait indefinitely.
+        """
+        with self._object_present_cv:
+            if not self._object_present_cv.wait_for(
+                lambda: self.has_object(obj_id),
+                timeout=timeout,
+            ):
+                raise TimeoutError(
+                    f"ObjectRef({obj_id}) not found in RDT object store after {timeout}s, transfer may have failed. Please report this issue on GitHub: https://github.com/ray-project/ray/issues/new/choose"
+                )
+
+    def pop_object(self, obj_id: str) -> List[Any]:
+        with self._lock:
+            queue = self._rdt_store.get(obj_id)
+            assert queue is not None, f"obj_id={obj_id} not found in RDT store"
+            rdt_object = queue.popleft()
+            if len(queue) == 0:
+                del self._rdt_store[obj_id]
+            if rdt_object.error:
+                raise rdt_object.error
+            for tensor in rdt_object.data:
+                self._tensor_to_object_ids[id(tensor)].remove(obj_id)
+                if len(self._tensor_to_object_ids[id(tensor)]) == 0:
+                    self._tensor_to_object_ids.pop(id(tensor))
+            self._object_freed_cv.notify_all()
+            return rdt_object.data
+
+    def wait_tensor_freed(self, tensor: Any, timeout: Optional[float] = None) -> None:
+        """
+        Wait for the object to be freed from the RDT store.
+        """
+        with self._object_freed_cv:
+            if not self._object_freed_cv.wait_for(
+                lambda: id(tensor) not in self._tensor_to_object_ids,
+                timeout=timeout,
+            ):
+                raise TimeoutError(
+                    f"Tensor {tensor} not freed from RDT object store after {timeout}s. The tensor will not be freed until all ObjectRefs containing the tensor have gone out of scope."
+                )
+
+    def get_num_objects(self) -> int:
+        """
+        Return the number of objects in the RDT store.
+        """
+        with self._lock:
+            # Count total objects across all queues
+            return sum(len(queue) for queue in self._rdt_store.values())

python/ray/experimental/rdt/tensor_transport_manager.py

@@ -0,0 +1,295 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
+
+if TYPE_CHECKING:
+    import numpy as np
+    import torch
+
+    import ray
+
+
+# NOTE: This is a public facing abstract interface for custom tensor transports.
+# Be sure to update the direct-transport docs when making changes to this interface, especially if changing the path to the file.
+
+
+@dataclass
+class CommunicatorMetadata:
+    """Metadata for the communicator."""
+
+
+@dataclass
+class TensorTransportMetadata:
+    """Metadata for tensors stored in the GPU object store.
+
+    Args:
+        tensor_meta: A list of tuples, each containing the shape and dtype of a tensor.
+        tensor_device: The device of the tensor. Currently, we require all tensors in the
+        list have the same device type.
+    """
+
+    tensor_meta: List[
+        Union[Tuple["torch.Size", "torch.dtype"], Tuple[Tuple[int, ...], "np.dtype"]]
+    ]
+    tensor_device: Optional[str] = None
+
+
+@dataclass
+class FetchRequest:
+    """Represents a pending or completed tensor fetch operation.
+
+    The default fetch/wait implementation stores the tensors here directly
+    after a synchronous recv. Transports with true async capability may
+    subclass this to carry additional state needed by wait_fetch_complete.
+
+    Subclasses should handle all resource cleanup in __del__ rather than
+    in wait_fetch_complete, so that resources are released even if the
+    caller never waits on the request.
+
+    Args:
+        obj_id: The object ID for the fetch operation.
+        tensors: The fetched tensors.
+    """
+
+    obj_id: str
+    tensors: List[Any]
+
+
+class TensorTransportManager(ABC):
+    """
+    Interface with which to implement custom tensor transports.
+    """
+
+    @abstractmethod
+    def tensor_transport_backend(self) -> str:
+        """
+        Returns the name of your tensor transport backend.
+        Ray uses this name to match your transport with the ``tensor_transport`` argument
+        on the method.
+
+        Returns:
+            str: The backend of the tensor transport.
+        """
+
+    @staticmethod
+    @abstractmethod
+    def is_one_sided() -> bool:
+        """
+        Indicates whether your transport uses one-sided communication where only the receiver
+        initiates the transfer.
+
+        One-sided transports: The receiver can directly read the sender's memory without the sender
+        actively participating. NIXL and CUDA-IPC are examples.
+
+        Two-sided transports: Both sender and receiver must actively participate in the transfer.
+        Collective communication libraries like NCCL and GLOO are examples.
+
+        This affects how Ray orchestrates the transfer and handles failures. Two-sided transports
+        have extra limitations described in :ref:`limitations <limitations>`. Ray will not call
+        `send_multiple_tensors` for one-sided transports; the transfer is expected to happen through
+        just `recv_multiple_tensors`.
+
+        Returns:
+            bool: True if the backend is one-sided, False otherwise.
+        """
+
+    @staticmethod
+    @abstractmethod
+    def can_abort_transport() -> bool:
+        """
+        Indicates whether your transport can safely abort an in-progress transfer.
+
+        If ``True``, Ray calls `abort_transport` on both the source and destination actors when a
+        send / recv error, allowing your transport to clean up gracefully.
+
+        If ``False``, Ray kills the involved actors to prevent deadlocks when errors occur during
+        transfer.
+
+        Return ``True`` only if your transport can reliably interrupt an in-progress send or receive
+        operation without leaving either party in a blocked state.
+
+        Returns:
+            bool: True if the backend can abort the transport.
+        """
+
+    @abstractmethod
+    def actor_has_tensor_transport(self, actor: "ray.actor.ActorHandle") -> bool:
+        """Whether the actor has the tensor transport available.
+
+        Args:
+            actor: The actor to check.
+
+        Returns:
+            bool: True if the actor has the tensor transport available, False otherwise.
+        """
+
+    @abstractmethod
+    def extract_tensor_transport_metadata(
+        self,
+        obj_id: str,
+        rdt_object: List[Any],
+    ) -> TensorTransportMetadata:
+        """
+        Implement this method to create the TensorTransportMetadata you defined previously.
+        Ray calls this on the source actor immediately after the actor task creates the result tensors.
+        Implement this to:
+
+        1. Record tensor shapes, dtypes, and devices.
+        2. Perform any transport-specific tensor registration such as registering memory for RDMA.
+        3. Store any handles or identifiers needed for the transfer.
+
+        Args:
+            obj_id: The ID of the RDT object to extract the tensor transport metadata from.
+            rdt_object: The RDT object to extract the tensor transport metadata from.
+
+        Returns:
+            TensorTransportMetadata: The tensor transport metadata.
+        """
+
+    @abstractmethod
+    def get_communicator_metadata(
+        self,
+        src_actor: "ray.actor.ActorHandle",
+        dst_actor: "ray.actor.ActorHandle",
+        backend: Optional[str] = None,
+    ) -> CommunicatorMetadata:
+        """
+        Gets the CommunicatorMetadata for a send/recv. Ray calls this on the owner/driver process before
+        orchestrating the transfer. You can typically implement this to return information both actors
+        need to identify each other such as ranks in a collective group. Many forms of transports such
+        as one-sided RDMA reads may be ok just returning empty CommunicatorMetadata here.
+
+        Args:
+            src_actor: The actor that runs this function.
+            dst_actor: The actor that runs this function.
+            backend: The backend to use for the collective operation.
+
+        Returns:
+            CommunicatorMetadata: The communicator metadata.
+        """
+
+    @abstractmethod
+    def recv_multiple_tensors(
+        self,
+        obj_id: str,
+        tensor_transport_metadata: TensorTransportMetadata,
+        communicator_metadata: CommunicatorMetadata,
+        target_buffers: Optional[List[Any]] = None,
+    ) -> List[Any]:
+        """
+        Receives tensors on the destination actor. Ray calls this on the destination
+        actor during the transfer.
+
+        Args:
+            obj_id: The object ID for related GPU object.
+            tensor_transport_metadata: The tensor transport metadata for the GPU object.
+            communicator_metadata: The communicator metadata for the send/recv operation.
+            target_buffers: Pre-allocated buffers to receive the tensors into if possible.
+        Returns:
+            List[Any]: The received tensors.
+        """
+
+    def fetch_multiple_tensors(
+        self,
+        obj_id: str,
+        tensor_transport_metadata: TensorTransportMetadata,
+        communicator_metadata: CommunicatorMetadata,
+        target_buffers: Optional[List[Any]] = None,
+    ) -> FetchRequest:
+        """Initiate a fetch for multiple tensors without waiting for completion.
+
+        The default implementation calls recv_multiple_tensors synchronously and
+        stores the result in a FetchRequest. Transports with true async capability
+        should override both this method and wait_fetch_complete.
+
+        Call wait_fetch_complete(fetch_request) afterward to retrieve the tensors.
+
+        Args:
+            obj_id: The object ID for the related GPU object.
+            tensor_transport_metadata: The tensor transport metadata for the GPU object.
+            communicator_metadata: The communicator metadata for the send/recv operation.
+            target_buffers: Pre-allocated buffers to receive the tensors into if possible.
+
+        Returns:
+            A FetchRequest whose tensors field is already populated.
+        """
+        tensors = self.recv_multiple_tensors(
+            obj_id, tensor_transport_metadata, communicator_metadata, target_buffers
+        )
+        return FetchRequest(obj_id=obj_id, tensors=tensors)
+
+    def wait_fetch_complete(
+        self, fetch_request: FetchRequest, timeout: float = -1
+    ) -> List[Any]:
+        """Wait for a previously initiated fetch to complete and return the tensors.
+
+        The default implementation returns the tensors stored in the FetchRequest
+        directly, since the default fetch_multiple_tensors is synchronous.
+
+        Args:
+            fetch_request: The FetchRequest returned by fetch_multiple_tensors.
+            timeout: Maximum time in seconds to wait. -1 means wait indefinitely.
+                0 means return immediately if not ready.
+
+        Returns:
+            The received tensors.
+
+        Raises:
+            TimeoutError: If timeout is exceeded.
+        """
+        return fetch_request.tensors
+
+    @abstractmethod
+    def send_multiple_tensors(
+        self,
+        tensors: List[Any],
+        tensor_transport_metadata: TensorTransportMetadata,
+        communicator_metadata: CommunicatorMetadata,
+    ):
+        """
+        Sends tensors from the source actor to the destination actor. Ray calls this on the source actor
+        during the transfer. Implement this to perform the actual data transfer using your transport's
+        send mechanism. For one-sided transports, you can simply avoid implementing this method or even
+        raise a NotImplementedError to ensure it's not being called.
+
+        Args:
+            tensors: The tensors or jax arrays to send.
+            tensor_transport_metadata: The tensor transport metadata for the RDT object.
+            communicator_metadata: The communicator metadata for the send/recv operation.
+        """
+
+    @abstractmethod
+    def garbage_collect(
+        self,
+        obj_id: str,
+        tensor_transport_meta: TensorTransportMetadata,
+        tensors: List[Any],
+    ):
+        """
+        Clean up resources for an RDT object. Ray calls this on the source actor
+        after Ray's distributed reference counting protocol determines the object is out of scope.
+
+        Use this to release any resources your transport allocated, such as deregistering memory buffers.
+        On the receiver side, no cleanup is needed — Ray does not hold onto the tensor after
+        returning it to the user, so it is garbage collected normally when the user releases it.
+
+        Args:
+            obj_id: The ID of the GPU object to garbage collect.
+            tensor_transport_meta: The tensor transport metadata.
+            tensors: The tensors that are contained in the ObjectRef that is being freed.
+        """
+
+    @abstractmethod
+    def abort_transport(
+        self,
+        obj_id: str,
+        communicator_metadata: CommunicatorMetadata,
+    ):
+        """
+        Aborts an in-progress transfer. Ray calls this on both the source and destination actors
+        when a system error occurs if `can_abort_transport` returns ``True``.
+
+        Args:
+            obj_id: The object ID for related GPU object.
+            communicator_metadata: The communicator metadata for the send/recv operation.
+        """

python/ray/experimental/rdt/util.py

@@ -0,0 +1,350 @@
+import threading
+from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional
+
+import ray
+from ray._raylet import ObjectRef
+from ray.experimental.rdt.collective_tensor_transport import (
+    GLOOTensorTransport,
+    NCCLTensorTransport,
+)
+from ray.experimental.rdt.cuda_ipc_transport import CudaIpcTransport
+from ray.experimental.rdt.nixl_tensor_transport import (
+    NixlTensorTransport,
+)
+from ray.experimental.rdt.tensor_transport_manager import (
+    TensorTransportManager,
+    TensorTransportMetadata,
+)
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    import torch
+
+
+class TransportManagerInfo(NamedTuple):
+    # Class that implements TensorTransportManager
+    transport_manager_class: type[TensorTransportManager]
+    # List of supported device types for the transport
+    devices: List[str]
+    # Data type for this transport (e.g. torch.Tensor or jax.Array)
+    # If not provided, defaults to torch.Tensor
+    data_type: type
+
+
+transport_manager_info: Dict[str, TransportManagerInfo] = {}
+
+# Singleton instances of transport managers
+transport_managers: Dict[str, TensorTransportManager] = {}
+
+# To protect the singleton instances of transport managers
+transport_managers_lock = threading.Lock()
+
+# Flipped to True when the first custom transport is registered.
+has_custom_transports = False
+
+
+@PublicAPI(stability="alpha")
+def register_tensor_transport(
+    transport_name: str,
+    devices: List[str],
+    transport_manager_class: type[TensorTransportManager],
+    data_type: type,
+):
+    """
+    Register a new tensor transport for use in Ray. Note that this needs to be called
+    before you create the actors that will use the transport. The actors also
+    need to be created in the same process from which you call this function.
+
+    Args:
+        transport_name: The name of the transport protocol.
+        devices: List of PyTorch device types supported by this transport (e.g., ["cuda", "cpu"]).
+        transport_manager_class: A class that implements TensorTransportManager.
+        data_type: The data type for this transport (e.g. torch.Tensor or jax.Array).
+    Raises:
+        ValueError: If transport_manager_class is not a subclass of TensorTransportManager.
+    """
+    global transport_manager_info
+    global has_custom_transports
+
+    transport_name = transport_name.upper()
+
+    if transport_name in transport_manager_info:
+        raise ValueError(f"Transport {transport_name} already registered.")
+
+    if not issubclass(transport_manager_class, TensorTransportManager):
+        raise ValueError(
+            f"transport_manager_class {transport_manager_class.__name__} must be a subclass of TensorTransportManager."
+        )
+
+    transport_manager_info[transport_name] = TransportManagerInfo(
+        transport_manager_class, devices, data_type
+    )
+
+    if transport_name not in DEFAULT_TRANSPORTS:
+        has_custom_transports = True
+
+
+DEFAULT_TRANSPORTS = ["NIXL", "GLOO", "NCCL", "CUDA_IPC"]
+
+_default_transports_registered = False
+
+
+def _ensure_default_transports_registered():
+    global _default_transports_registered
+    with transport_managers_lock:
+        if _default_transports_registered:
+            return
+        _default_transports_registered = True
+        try:
+            import torch
+
+            register_tensor_transport(
+                "NIXL", ["cuda", "cpu"], NixlTensorTransport, torch.Tensor
+            )
+            register_tensor_transport(
+                "GLOO", ["cpu"], GLOOTensorTransport, torch.Tensor
+            )
+            register_tensor_transport(
+                "NCCL", ["cuda"], NCCLTensorTransport, torch.Tensor
+            )
+            register_tensor_transport(
+                "CUDA_IPC", ["cuda"], CudaIpcTransport, torch.Tensor
+            )
+        except ImportError:
+            pass
+
+
+def get_transport_data_type(tensor_transport: str) -> type:
+    _ensure_default_transports_registered()
+    if tensor_transport not in transport_manager_info:
+        raise ValueError(f"Unsupported tensor transport protocol: {tensor_transport}")
+
+    return transport_manager_info[tensor_transport].data_type
+
+
+def get_tensor_transport_manager(
+    transport_name: str,
+) -> "TensorTransportManager":
+    """Get the tensor transport manager for the given tensor transport protocol.
+
+    Args:
+        transport_name: The tensor transport protocol to use for the GPU object.
+
+    Returns:
+        TensorTransportManager: The tensor transport manager for the given tensor transport protocol.
+    """
+    global transport_manager_info
+    global transport_managers
+    global transport_managers_lock
+
+    _ensure_default_transports_registered()
+    with transport_managers_lock:
+        if transport_name in transport_managers:
+            return transport_managers[transport_name]
+
+        if transport_name not in transport_manager_info:
+            raise ValueError(f"Unsupported tensor transport protocol: {transport_name}")
+
+        transport_managers[transport_name] = transport_manager_info[
+            transport_name
+        ].transport_manager_class()
+        return transport_managers[transport_name]
+
+
+def register_custom_tensor_transports_on_actor(
+    actor: "ray.actor.ActorHandle",
+) -> Optional[ObjectRef]:
+    """
+    If there's no custom transports to register, returns None.
+    Otherwise returns an ObjectRef for a task on the actor that will register the custom transports.
+    """
+    global transport_manager_info
+    global has_custom_transports
+
+    _ensure_default_transports_registered()
+    if not has_custom_transports:
+        return None
+
+    def register_transport_on_actor(
+        self, owner_transport_manager_info: Dict[str, TransportManagerInfo]
+    ):
+        from ray.experimental.rdt.util import (
+            _ensure_default_transports_registered,
+            register_tensor_transport,
+            transport_manager_info,
+        )
+
+        _ensure_default_transports_registered()
+        for transport_name, transport_info in owner_transport_manager_info.items():
+            if transport_name not in transport_manager_info:
+                register_tensor_transport(
+                    transport_name,
+                    transport_info.devices,
+                    transport_info.transport_manager_class,
+                    transport_info.data_type,
+                )
+
+    return actor.__ray_call__.options(concurrency_group="_ray_system").remote(
+        register_transport_on_actor, transport_manager_info
+    )
+
+
+def device_match_transport(device: str, tensor_transport: str) -> bool:
+    """Check if the device matches the transport."""
+    _ensure_default_transports_registered()
+    if tensor_transport not in transport_manager_info:
+        raise ValueError(f"Unsupported tensor transport protocol: {tensor_transport}")
+
+    return device in transport_manager_info[tensor_transport].devices
+
+
+def normalize_and_validate_tensor_transport(tensor_transport: str) -> str:
+    _ensure_default_transports_registered()
+    tensor_transport = tensor_transport.upper()
+
+    if tensor_transport not in transport_manager_info:
+        raise ValueError(f"Invalid tensor transport: {tensor_transport}")
+
+    return tensor_transport
+
+
+def is_one_sided_transport(tensor_transport: str) -> bool:
+    _ensure_default_transports_registered()
+    return transport_manager_info[
+        tensor_transport
+    ].transport_manager_class.is_one_sided()
+
+
+@PublicAPI(stability="alpha")
+def register_nixl_memory(tensor: "torch.Tensor") -> None:
+    """Registers the tensor's memory with NIXL and bumps the reference count so the memory region is never deregistered.
+
+    By default, the lifetime of the NIXL memory registration is tied to the ObjectRef. This means that only when the ObjectRef is created
+    do we register the memory with NIXL and deregister it when the ObjectRef goes out of scope. However, this function can be used
+    to pre-register a tensor's memory with NIXL and keep it registered for the lifetime of the process which can improve performance
+    if the same tensor is re-used in multiple RDT objects.
+
+    If called on a tensor that is already registered with NIXL, we still prevent the tensor's memory from being deregistered.
+
+    Args:
+        tensor: A PyTorch tensor whose memory should be registered with NIXL.
+
+    Example:
+
+        .. code-block:: python
+
+            import torch
+            import ray
+            from ray.experimental import register_nixl_memory
+
+            @ray.remote(num_gpus=1, enable_tensor_transport=True)
+            class Trainer:
+                def __init__(self):
+                    self.weight = torch.randn(1000, 1000, device="cuda")
+                    # Pre-register the memory with NIXL for the lifetime of the process
+                    register_nixl_memory(self.weight)
+
+                # Both of the below methods will use the cached NIXL memory registration on multiple calls. You can also mix them,
+                # i.e. call get_weight_ref_by_rows then get_weight_ref and get_weight_ref will not trigger a new NIXL memory registration.
+
+                # You can ray.put views to each row of the weight matrix if you want to use them separately in your code
+                def get_weight_ref_by_rows(self):
+                    views = [self.weight[i] for i in range(1000)]
+                    # Each put call does not trigger a new NIXL memory registration
+                    return ray.put(views, _tensor_transport="nixl")
+
+                # You can also ray.put the entire weight matrix at once
+                def get_weight_ref(self):
+                    return ray.put(self.weight, _tensor_transport="nixl")
+    """
+    nixl_transport = get_tensor_transport_manager("NIXL")
+    nixl_transport.register_nixl_memory(tensor)
+
+
+@PublicAPI(stability="alpha")
+def deregister_nixl_memory(tensor: "torch.Tensor") -> None:
+    """Decrements the reference count for the tensor's NIXL memory registration added by :func:`ray.experimental.register_nixl_memory`.
+
+    If the reference count reaches 0, the memory is deregistered from NIXL.
+    This should only be called after :func:`ray.experimental.register_nixl_memory` has been called for this tensor.
+    Any existing ``ray.ObjectRef`` instances that reference this tensor's memory will keep the
+    NIXL memory registration alive independently until they go out of scope.
+
+    Args:
+        tensor: A PyTorch tensor whose NIXL memory registration reference count should be decremented.
+
+    Example:
+
+        .. code-block:: python
+
+            # Extending the example from register_nixl_memory:
+            @ray.remote(num_gpus=1, enable_tensor_transport=True)
+            class Trainer:
+                def deregister_weight(self):
+                    # Remove the NIXL memory registration added by register_nixl_memory.
+                    # The memory may still be registered if there are live ObjectRefs
+                    # that reference it.
+                    deregister_nixl_memory(self.weight)
+    """
+    nixl_transport = get_tensor_transport_manager("NIXL")
+    nixl_transport.deregister_nixl_memory(tensor)
+
+
+@PublicAPI(stability="alpha")
+def register_nixl_memory_pool(size: int, device: "torch.device") -> None:
+    """Pre-allocates a memory pool and registers it with NIXL.
+
+    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``.
+
+    Within a single ``ray.put`` call, tensors sharing the same underlying storage
+    (including views) are automatically deduplicated — only one copy of each unique
+    storage is allocated. Across multiple ``ray.put`` calls, if the same storage
+    appears again, the existing pool slot is reused without re-copying the data.
+    As a result, data can be potentially stale once you ``ray.put`` the storage
+    tensor — subsequent mutations to that storage may not be reflected in outstanding refs.
+    Clone the tensor before ``ray.put`` if snapshot semantics are required.
+
+    If the pool has insufficient space for an allocation,
+    :class:`NixlOutOfMemoryError` is raised.
+
+    Args:
+        size: Size of the memory pool in bytes.
+        device: Device to allocate the pool on (e.g., ``torch.device("cpu")``
+            or ``torch.device("cuda")``).
+
+    Example:
+
+        .. code-block:: python
+
+            import torch
+            import ray
+            from ray.experimental import register_nixl_memory_pool
+
+            @ray.remote(num_gpus=1, enable_tensor_transport=True)
+            class Trainer:
+                def __init__(self):
+                    # Pre-allocate a 1GB GPU memory pool for NIXL transfers
+                    register_nixl_memory_pool(1024 * 1024 * 1024, torch.device("cuda"))
+
+                def get_weight_ref(self):
+                    weight = torch.randn(1000, 1000, device="cuda")
+                    return ray.put(weight, _tensor_transport="nixl")
+    """
+    nixl_transport = get_tensor_transport_manager("NIXL")
+    nixl_transport.register_nixl_memory_pool(size, device)
+
+
+def create_empty_tensors_from_metadata(
+    tensor_transport_meta: TensorTransportMetadata,
+) -> List["torch.Tensor"]:
+    import torch
+
+    tensors = []
+    device = tensor_transport_meta.tensor_device
+    for meta in tensor_transport_meta.tensor_meta:
+        shape, dtype = meta
+        tensor = torch.empty(shape, dtype=dtype, device=device)
+        tensors.append(tensor)
+    return tensors