chore: import upstream snapshot with attribution

python/ray/experimental/BUILD.bazel

@@ -0,0 +1,9 @@
+load("//bazel:python.bzl", "doctest")
+
+doctest(
+    files = glob(
+        ["**/*.py"],
+        exclude = ["tf_utils.py"],
+    ),
+    tags = ["team:core"],
+)

python/ray/experimental/__init__.py

@@ -0,0 +1,30 @@
+from ray.experimental.dynamic_resources import set_resource
+from ray.experimental.locations import get_local_object_locations, get_object_locations
+from ray.experimental.rdt import (
+    CommunicatorMetadata,
+    RDTManager,
+    TensorTransportManager,
+    TensorTransportMetadata,
+    deregister_nixl_memory,
+    register_nixl_memory,
+    register_nixl_memory_pool,
+    register_tensor_transport,
+    set_target_for_ref,
+    wait_tensor_freed,
+)
+
+__all__ = [
+    "get_object_locations",
+    "get_local_object_locations",
+    "set_resource",
+    "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/channel/__init__.py

@@ -0,0 +1,43 @@
+from ray.experimental.channel.cached_channel import CachedChannel
+from ray.experimental.channel.common import (  # noqa: F401
+    AwaitableBackgroundReader,
+    AwaitableBackgroundWriter,
+    ChannelContext,
+    ChannelInterface,
+    ChannelOutputType,
+    CompiledDAGArgs,
+    ReaderInterface,
+    SynchronousReader,
+    SynchronousWriter,
+    WriterInterface,
+)
+from ray.experimental.channel.communicator import Communicator
+from ray.experimental.channel.cpu_communicator import CPUCommunicator
+from ray.experimental.channel.intra_process_channel import IntraProcessChannel
+from ray.experimental.channel.shared_memory_channel import (
+    BufferedSharedMemoryChannel,
+    Channel,
+    CompositeChannel,
+)
+from ray.experimental.channel.torch_tensor_accelerator_channel import (
+    TorchTensorAcceleratorChannel,
+)
+
+__all__ = [
+    "AwaitableBackgroundReader",
+    "AwaitableBackgroundWriter",
+    "CachedChannel",
+    "Channel",
+    "Communicator",
+    "CPUCommunicator",
+    "ReaderInterface",
+    "SynchronousReader",
+    "SynchronousWriter",
+    "WriterInterface",
+    "ChannelContext",
+    "TorchTensorAcceleratorChannel",
+    "IntraProcessChannel",
+    "CompositeChannel",
+    "BufferedSharedMemoryChannel",
+    "CompiledDAGArgs",
+]

python/ray/experimental/channel/accelerator_context.py

@@ -0,0 +1,246 @@
+import importlib
+import threading
+from contextlib import nullcontext
+from typing import TYPE_CHECKING, ContextManager, List, Optional, Type
+
+import ray
+from ray._private.accelerators import get_accelerator_manager_for_resource
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+_global_custom_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This includes stream management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, communicator_cls: Type[Communicator]):
+        """
+        Initializes an accelerator context with the specified torch device module
+        and communicator class.
+
+        Args:
+            torch_module_name: Name of the torch device module (e.g., "cuda", "cpu").
+            communicator_cls: Class used to handle communication.
+        """
+
+        # The name of the torch module (e.g., 'cuda', 'npu')
+        self._torch_module_name: str = torch_module_name
+        # The Communicator class used to manage communication
+        self._communicator_cls: Type[Communicator] = communicator_cls
+
+        # Import the torch backend module (e.g., torch.cuda) if the device is not 'cpu'.
+        if torch_module_name != "cpu":
+            self._torch_mod = importlib.import_module(f"torch.{torch_module_name}")
+
+    @staticmethod
+    def get() -> "AcceleratorContext":
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+
+        global _default_accelerator_context, _global_custom_context
+
+        with _accelerator_context_lock:
+            if _global_custom_context is not None:
+                return _global_custom_context
+
+            if _default_accelerator_context is None:
+                if len(ray.get_gpu_ids()) > 0:
+                    from ray.experimental.channel.nccl_group import _NcclGroup
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cuda", _NcclGroup
+                    )
+                else:
+                    from ray.experimental.channel.cpu_communicator import (
+                        CPUCommunicator,
+                    )
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cpu", CPUCommunicator
+                    )
+
+            return _default_accelerator_context
+
+    @staticmethod
+    def set(accelerator_context: "AcceleratorContext") -> None:
+        """
+        Overwrites the default accelerator context.
+
+        Args:
+            accelerator_context: The context to register.
+        """
+        global _global_custom_context
+
+        # Accelerator context is registered.
+        _global_custom_context = accelerator_context
+
+    def get_accelerator_devices(self) -> List["torch.device"]:
+        """
+        Gets the torch device list configured for this process.
+
+        Returns:
+            List[torch.device]: The torch device list.
+        """
+        import torch
+
+        if self._torch_module_name == "cpu":
+            return [torch.device("cpu")]
+
+        if self._torch_module_name == "cuda":
+            accelerator_ids = [str(id) for id in ray.get_gpu_ids()]
+            accelerator_manager = get_accelerator_manager_for_resource("GPU")
+        else:
+            accelerator_ids = [
+                str(id)
+                for id in ray.get_runtime_context().get_accelerator_ids()[
+                    self._torch_module_name.upper()
+                ]
+            ]
+            accelerator_manager = get_accelerator_manager_for_resource(
+                self._torch_module_name.upper()
+            )
+
+        device_ids = []
+
+        if len(accelerator_ids) > 0:
+            accelerator_visible_list = (
+                accelerator_manager.get_current_process_visible_accelerator_ids()
+            )
+            if accelerator_visible_list is None:
+                accelerator_visible_list = []
+
+            # If there are multiple Accelerators, return a list of devices.
+            # If using fractional Accelerators, these IDs are not guaranteed
+            # to be unique across different processes.
+            for accelerator_id in accelerator_ids:
+                try:
+                    device_ids.append(accelerator_visible_list.index(accelerator_id))
+                except ValueError:
+                    raise RuntimeError(
+                        f"{accelerator_manager.get_visible_accelerator_ids_env_var()} set incorrectly. "
+                        f"expected to include {accelerator_id}. "
+                        "Did you override this environment"
+                        " variable? If not, please help file an issue on Github."
+                    )
+
+        else:
+            # If called on the driver or outside of Ray Train, return the
+            # 0th device.
+            device_ids.append(0)
+
+        return [
+            torch.device(f"{self._torch_module_name}:{device_id}")
+            for device_id in device_ids
+        ]
+
+    def get_device_context(self, device: "torch.device") -> ContextManager:
+        """
+        Retrieves the context manager for the specified accelerator device.
+        There is no device context for CPU, returning a nullcontext.
+
+        Args:
+            device: The target device for which the context manager is required.
+
+        Returns:
+            ContextManager: A context manager specific to the device type.
+        """
+        if device.type == "cpu":
+            return nullcontext()
+
+        return self._torch_mod.device(device)
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        return self._torch_mod.current_stream()
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        return self._torch_mod.Event()
+
+    def generate_communicator_id(self) -> str:
+        """
+        Generates a communication identifier for communication group.
+        """
+        return self._communicator_cls.generate_communicator_id()
+
+    def create_communicator(self, *args, **kwargs) -> Communicator:
+        """
+        Creates a communication group for collective operations.
+        """
+        return self._communicator_cls(*args, **kwargs)
+
+    @property
+    def module_name(self) -> str:
+        """
+        Gets the name of the torch module backing the accelerator.
+        """
+        return self._torch_module_name
+
+    @property
+    def communicator_cls(self) -> Optional[Type[Communicator]]:
+        """
+        Returns the communicator class.
+        """
+        return self._communicator_cls
+
+    @property
+    def accelerator_count(self) -> int:
+        """
+        Returns the number of accelerators assigned by ray.
+        """
+        if self._torch_module_name == "cuda":
+            return len(ray.get_gpu_ids())
+        else:
+            accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
+            return len(accelerator_ids.get(self._torch_module_name.upper(), []))
+
+
+def register_accelerator_context(
+    torch_module_name: str, communicator_cls: Type[Communicator]
+):
+    """
+    Registers the accelerator context with the specified device type and communicator.
+
+    Args:
+        torch_module_name: The name of the device module under torch.
+        communicator_cls: The communicator class associated with the device.
+    """
+    accelerator_context = AcceleratorContext(torch_module_name, communicator_cls)
+    AcceleratorContext.set(accelerator_context)
+
+
+def is_accelerator_context_registered():
+    """
+    Checks whether a custom accelerator context has been registered.
+
+    Returns:
+        bool: True if a custom accelerator context is registered
+              (_global_custom_context is not None), False otherwise.
+    """
+    if _global_custom_context is not None:
+        return True
+    return False

python/ray/experimental/channel/auto_transport_type.py

@@ -0,0 +1,185 @@
+from typing import Dict, List, Optional, Tuple, Union
+
+import ray
+from ray.experimental.channel import ChannelOutputType
+from ray.experimental.channel.torch_tensor_type import TorchTensorType
+from ray.experimental.util.types import Device
+
+
+class AutoTransportType(ChannelOutputType):
+    """
+    Type hint for automatic transport selection for tensors.
+
+    With this type hint Compiled Graphs automatically decide the best transport
+    to use (e.g., accellerator or shared memory) based on the node locations and
+    GPU IDs of the readers and writers.
+    """
+
+    def __init__(
+        self,
+        device: Device = Device.DEFAULT,
+        _static_shape: bool = False,
+        _direct_return: bool = False,
+    ):
+        self._device = device
+        self._static_shape = _static_shape
+        self._direct_return = _direct_return
+
+    @property
+    def device(self) -> Device:
+        return self._device
+
+    def create_channel(
+        self,
+        writer: Optional["ray.actor.ActorHandle"],
+        reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
+        driver_actor_id: Optional[str] = None,
+    ) -> "ChannelOutputType":
+        """
+        Directly calling create_channel() on AutoTransportType should not happen,
+        just raise an exception with informative message.
+        """
+        raise ValueError(
+            "This should not happen: AutoTransportType should "
+            "have been resolved before creating a channel. "
+            "Please file a Ray GitHub issue for bug report."
+        )
+
+
+class TypeHintResolver:
+    """
+    This class is used to resolve `AutoChannelType` into an actual channel type
+    (e.g., `TorchTensorType` with proper transport) based on node locations and
+    GPU IDs of the readers and writers.
+    """
+
+    def __init__(self, actor_to_gpu_ids: Dict["ray.actor.ActorHandle", List[str]]):
+        """Initialize the type hint resolver.
+
+        Args:
+            actor_to_gpu_ids: Mapping from actor handle to its GPU IDs.
+        """
+        self._actor_to_gpu_ids = actor_to_gpu_ids
+
+    def _get_gpu_ids(self, actor: "ray.actor.ActorHandle") -> List[str]:
+        """Get the GPU IDs of the actor.
+
+        Args:
+            actor: The actor handle to look up.
+
+        Returns:
+            The GPU IDs of the actor. If the actor is not found,
+            return an empty list.
+        """
+        gpu_ids = self._actor_to_gpu_ids.get(actor, [])
+        assert len(gpu_ids) <= 1, (
+            "Compiled Graphs currently don't support allocating multiple GPUs "
+            "to a single actor"
+        )
+        return gpu_ids
+
+    def _use_same_gpu(
+        self,
+        writer_and_node: Tuple["ray.actor.ActorHandle", str],
+        reader_and_node: Union[
+            Tuple["ray.actor.ActorHandle", str],
+            List[Tuple["ray.actor.ActorHandle", str]],
+        ],
+    ) -> bool:
+        """
+        Check if the writer and readers use the same GPU.
+
+        Args:
+            writer_and_node: A tuple of writer actor handle and its node ID.
+            reader_and_node: A tuple of reader actor handle and its node ID, or
+                a list of such tuples.
+
+        Returns:
+            True if the writer and all the readers use the same GPU, False otherwise.
+        """
+        if isinstance(reader_and_node, list):
+            return all(
+                self._use_same_gpu(writer_and_node, entry) for entry in reader_and_node
+            )
+        if writer_and_node[1] != reader_and_node[1]:
+            return False
+        writer_gpu_ids = self._get_gpu_ids(writer_and_node[0])
+        reader_gpu_ids = self._get_gpu_ids(reader_and_node[0])
+        return writer_gpu_ids == reader_gpu_ids
+
+    def _use_gpu(
+        self, actors: Union["ray.actor.ActorHandle", List["ray.actor.ActorHandle"]]
+    ) -> bool:
+        """
+        Check if the actors use GPUs.
+
+        Args:
+            actors: An actor handle or a list of actor handles.
+
+        Returns:
+            True if the actors use GPUs, False otherwise.
+        """
+        if isinstance(actors, list):
+            return all(self._use_gpu(actor) for actor in actors)
+        gpu_ids = self._get_gpu_ids(actors)
+        return len(gpu_ids) > 0
+
+    def resolve(
+        self,
+        auto_transport_type: AutoTransportType,
+        writer_and_node: Tuple[Optional["ray.actor.ActorHandle"], str],
+        reader_and_node_list: List[Tuple[Optional["ray.actor.ActorHandle"], str]],
+    ) -> "ChannelOutputType":
+        """
+        Resolve auto_transport_type to the actual channel output type
+        based on the node locations and GPU IDs.
+
+        Args:
+            auto_transport_type: The type to resolve
+            writer_and_node: A tuple of writer actor handle and its node ID.
+                A None writer actor handle means the writer is the driver.
+            reader_and_node_list: A list of tuples of reader actor handle and its
+                node ID. A None reader actor handle means the reader is the driver.
+
+        Returns:
+            The actual channel type.
+        """
+        writer = writer_and_node[0]
+        readers = [reader for reader, _ in reader_and_node_list]
+
+        if writer is None or any(reader is None for reader in readers):
+            # None means actor is the driver, currently driver on GPU
+            # is not supported, so we always use shared memory to transfer
+            # tensors.
+            return TorchTensorType(
+                device=auto_transport_type.device,
+                _static_shape=auto_transport_type._static_shape,
+                _direct_return=auto_transport_type._direct_return,
+            )
+
+        # Case 1: writer and readers don't both use GPU, use shared memory
+        # to transport the tensors
+        if not (self._use_gpu(writer) and self._use_gpu(readers)):
+            return TorchTensorType(
+                device=auto_transport_type.device,
+                _static_shape=auto_transport_type._static_shape,
+                _direct_return=auto_transport_type._direct_return,
+            )
+
+        # Case 2: writer and readers use the same GPU are are on the same node,
+        # use shared memory to transport the tensors
+        if self._use_same_gpu(writer_and_node, reader_and_node_list):
+            return TorchTensorType(
+                device=auto_transport_type.device,
+                _static_shape=auto_transport_type._static_shape,
+                _direct_return=auto_transport_type._direct_return,
+            )
+
+        # Case 3: writer and readers use different GPUs, use accelerator to transport
+        # the tensors
+        return TorchTensorType(
+            transport="accelerator",
+            device=auto_transport_type.device,
+            _static_shape=auto_transport_type._static_shape,
+            _direct_return=auto_transport_type._direct_return,
+        )

python/ray/experimental/channel/cached_channel.py

@@ -0,0 +1,111 @@
+import uuid
+from typing import Any, Optional
+
+from ray.experimental.channel.common import ChannelInterface
+
+
+class CachedChannel(ChannelInterface):
+    """
+    CachedChannel wraps an inner channel and caches the data read from it until
+    `num_reads` reads have completed. If inner channel is None, the data
+    is written to serialization context and retrieved from there. This is useful
+    when passing data within the same actor and a shared memory channel can be
+    avoided.
+
+    Args:
+        num_reads: The number of reads from this channel that must happen before
+            writing again. Readers must be methods of the same actor.
+        inner_channel: The inner channel to cache data from. If None, the data is
+            read from the serialization context.
+        _channel_id: The unique ID for the channel. If None, a new ID is generated.
+    """
+
+    def __init__(
+        self,
+        num_reads: int,
+        inner_channel: Optional[ChannelInterface] = None,
+        _channel_id: Optional[str] = None,
+    ):
+        assert num_reads > 0, "num_reads must be greater than 0."
+        self._num_reads = num_reads
+        self._inner_channel = inner_channel
+        # Generate a unique ID for the channel. The writer and reader will use
+        # this ID to store and retrieve data from the _SerializationContext.
+        self._channel_id = _channel_id
+        if self._channel_id is None:
+            self._channel_id = str(uuid.uuid4())
+
+    def ensure_registered_as_writer(self) -> None:
+        if self._inner_channel is not None:
+            self._inner_channel.ensure_registered_as_writer()
+
+    def ensure_registered_as_reader(self) -> None:
+        if self._inner_channel is not None:
+            self._inner_channel.ensure_registered_as_reader()
+
+    def __reduce__(self):
+        return CachedChannel, (
+            self._num_reads,
+            self._inner_channel,
+            self._channel_id,
+        )
+
+    def __str__(self) -> str:
+        return (
+            f"CachedChannel(channel_id={self._channel_id}, "
+            f"num_reads={self._num_reads}), "
+            f"inner_channel={self._inner_channel})"
+        )
+
+    def write(self, value: Any, timeout: Optional[float] = None):
+        self.ensure_registered_as_writer()
+        # TODO: better organize the imports
+        from ray.experimental.channel import ChannelContext
+
+        if self._inner_channel is not None:
+            self._inner_channel.write(value, timeout)
+            return
+
+        # Otherwise no need to check timeout as the operation is non-blocking.
+
+        # Because both the reader and writer are in the same worker process,
+        # we can directly store the data in the context instead of storing
+        # it in the channel object. This removes the serialization overhead of `value`.
+        ctx = ChannelContext.get_current().serialization_context
+        ctx.set_data(self._channel_id, value, self._num_reads)
+
+    def read(self, timeout: Optional[float] = None) -> Any:
+        self.ensure_registered_as_reader()
+        # TODO: better organize the imports
+        from ray.experimental.channel import ChannelContext
+
+        ctx = ChannelContext.get_current().serialization_context
+        if ctx.has_data(self._channel_id):
+            # No need to check timeout as the operation is non-blocking.
+            return ctx.get_data(self._channel_id)
+
+        assert (
+            self._inner_channel is not None
+        ), "Cannot read from the serialization context while inner channel is None."
+        value = self._inner_channel.read(timeout)
+        ctx.set_data(self._channel_id, value, self._num_reads)
+        # NOTE: Currently we make a contract with Compiled Graph users that the
+        # channel results should not be mutated by the actor methods.
+        # When the user needs to modify the channel results, they should
+        # make a copy of the channel results and modify the copy.
+        # This is the same contract as used in IntraProcessChannel.
+        # This contract is NOT enforced right now in either case.
+        # TODO(rui): introduce a flag to control the behavior:
+        # for example, by default we make a deep copy of the channel
+        # result, but the user can turn off the deep copy for performance
+        # improvements.
+        # https://github.com/ray-project/ray/issues/47409
+        return ctx.get_data(self._channel_id)
+
+    def close(self) -> None:
+        from ray.experimental.channel import ChannelContext
+
+        if self._inner_channel is not None:
+            self._inner_channel.close()
+        ctx = ChannelContext.get_current().serialization_context
+        ctx.reset_data(self._channel_id)

python/ray/experimental/channel/common.py

@@ -0,0 +1,688 @@
+import asyncio
+import concurrent
+import sys
+import threading
+import time
+from dataclasses import dataclass
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    NamedTuple,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import ray
+import ray.exceptions
+from ray.experimental.channel.accelerator_context import AcceleratorContext
+from ray.experimental.channel.communicator import Communicator
+from ray.experimental.channel.communicator_handle import CommunicatorHandle
+from ray.experimental.channel.serialization_context import _SerializationContext
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+# The context singleton on this process.
+_default_context: "Optional[ChannelContext]" = None
+_context_lock = threading.Lock()
+
+if TYPE_CHECKING:
+    import torch
+
+
+def retry_and_check_interpreter_exit(f: Callable[[], None]) -> bool:
+    """This function is only useful when f contains channel read/write.
+
+    Keep retrying channel read/write inside `f` and check if interpreter exits.
+    It is important in case the read/write happens in a separate thread pool.
+    See https://github.com/ray-project/ray/pull/47702
+
+    f should a function that doesn't receive any input and return nothing.
+    """
+    exiting = False
+    while True:
+        try:
+            f()
+            break
+        except ray.exceptions.RayChannelTimeoutError:
+            if sys.is_finalizing():
+                # Interpreter exits. We should ignore the error and
+                # stop reading so that the thread can join.
+                exiting = True
+                break
+
+    return exiting
+
+
+# Holds the input arguments for Compiled Graph
+@PublicAPI(stability="alpha")
+class CompiledDAGArgs(NamedTuple):
+    args: Tuple[Any, ...]
+    kwargs: Dict[str, Any]
+
+
+@PublicAPI(stability="alpha")
+class ChannelOutputType:
+    def register_custom_serializer(self) -> None:
+        """
+        Register any custom serializers needed to pass data of this type. This
+        method should be run on the reader(s) and writer of a channel, which
+        are the driver and/or Ray actors.
+
+        NOTE: When custom serializers are registered with Ray, the registered
+        deserializer is shipped with the serialized value and used on the
+        receiving end. Therefore, the deserializer function should *not*
+        capture state that is meant to be worker-local, such as the worker's
+        default device. Instead, these should be extracted from the
+        worker-local _SerializationContext.
+        """
+        pass
+
+    def create_channel(
+        self,
+        writer: Optional["ray.actor.ActorHandle"],
+        reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
+        driver_actor_id: Optional[str] = None,
+    ) -> "ChannelInterface":
+        """
+        Instantiate a ChannelInterface class that can be used
+        to pass data of this type.
+
+        Args:
+            writer: The actor that may write to the channel. None signifies the driver.
+            reader_and_node_list: A list of tuples, where each tuple contains a reader
+                actor handle and the node ID where the actor is located.
+            driver_actor_id: If this is a CompositeChannel that is read by a driver and
+                that driver is an actual actor, this will be the actor ID of that
+                driver actor.
+        Returns:
+            A ChannelInterface that can be used to pass data
+                of this type.
+        """
+        raise NotImplementedError
+
+    def requires_accelerator(self) -> bool:
+        # By default, channels do not require accelerator.
+        return False
+
+    def get_custom_communicator(self) -> Optional[Communicator]:
+        """
+        Return the custom communicator group if one is specified.
+        """
+        return None
+
+    def set_communicator_id(self, group_id: str) -> None:
+        raise NotImplementedError
+
+
+@DeveloperAPI
+@dataclass
+class ChannelContext:
+    serialization_context = _SerializationContext()
+    _torch_available: Optional[bool] = None
+    _torch_device: Optional["torch.device"] = None
+    _current_stream: Optional["torch.cuda.Stream"] = None
+
+    def __init__(self):
+        # Used for the torch.Tensor accelerator transport.
+        self.communicators: Dict[str, "Communicator"] = {}
+        # Used for driver process to store actors in the communicator.
+        self.communicator_handles: Dict[str, "CommunicatorHandle"] = {}
+
+    @staticmethod
+    def get_current() -> "ChannelContext":
+        """Get or create a singleton context.
+
+        If the context has not yet been created in this process, it will be
+        initialized with default settings.
+        """
+
+        global _default_context
+
+        with _context_lock:
+            if _default_context is None:
+                _default_context = ChannelContext()
+
+            return _default_context
+
+    @property
+    def torch_available(self) -> bool:
+        """
+        Check if torch package is available.
+        """
+        if self._torch_available is not None:
+            return self._torch_available
+
+        try:
+            import torch  # noqa: F401
+        except ImportError:
+            self._torch_available = False
+            return False
+        self._torch_available = True
+        return True
+
+    @property
+    def torch_device(self) -> "torch.device":
+        if self._torch_device is None:
+            self._torch_device = AcceleratorContext.get().get_accelerator_devices()[0]
+
+        return self._torch_device
+
+    def set_torch_device(self, device: "torch.device"):
+        self._torch_device = device
+
+
+@PublicAPI(stability="alpha")
+class ChannelInterface:
+    """
+    Abstraction for a transport between a writer actor and some number of
+    reader actors.
+    """
+
+    def __init__(
+        self,
+        writer: Optional[ray.actor.ActorHandle],
+        readers: List[Optional[ray.actor.ActorHandle]],
+        typ: Optional["ChannelOutputType"],
+    ):
+        """
+        Create a channel that can be read and written by a Ray driver or actor.
+
+        Args:
+            writer: The actor that may write to the channel. None signifies the driver.
+            readers: The actors that may read from the channel. None signifies
+                the driver.
+            typ: Type information about the values passed through the channel.
+        """
+        pass
+
+    def ensure_registered_as_writer(self):
+        """
+        Check whether the process is a valid writer. This method must be idempotent.
+        """
+        raise NotImplementedError
+
+    def ensure_registered_as_reader(self):
+        """
+        Check whether the process is a valid reader. This method must be idempotent.
+        """
+        raise NotImplementedError
+
+    def write(self, value: Any, timeout: Optional[float] = None) -> None:
+        """
+        Write a value to the channel.
+
+        Blocks if there are still pending readers for the previous value. The
+        writer may not write again until the specified number of readers have
+        read the value.
+
+        Args:
+            value: The value to write.
+            timeout: The maximum time in seconds to wait to write the value.
+                None means using default timeout, 0 means immediate timeout
+                (immediate success or timeout without blocking), -1 means
+                infinite timeout (block indefinitely).
+        """
+        raise NotImplementedError
+
+    def read(self, timeout: Optional[float] = None) -> Any:
+        """
+        Read the latest value from the channel. This call will block until a
+        value is available to read.
+
+        Subsequent calls to read() may *block* if the deserialized object is
+        zero-copy (e.g., bytes or a numpy array) *and* the object is still in scope.
+
+        Args:
+            timeout: The maximum time in seconds to wait to read the value.
+                None means using default timeout, 0 means immediate timeout
+                (immediate success or timeout without blocking), -1 means
+                infinite timeout (block indefinitely).
+
+        Returns:
+            Any: The deserialized value. If the deserialized value is an
+            Exception, it will be returned directly instead of being raised.
+        """
+        raise NotImplementedError
+
+    def close(self) -> None:
+        """
+        Close this channel. This method must not block and it must be made
+        idempotent. Any existing values in the channel may be lost after the
+        channel is closed.
+        """
+        raise NotImplementedError
+
+
+# Interfaces for channel I/O.
+@DeveloperAPI
+class ReaderInterface:
+    def __init__(
+        self,
+        input_channels: List[ChannelInterface],
+    ):
+        assert isinstance(input_channels, list)
+        for chan in input_channels:
+            assert isinstance(chan, ChannelInterface)
+
+        self._input_channels = input_channels
+        self._closed = False
+        self._num_reads = 0
+
+        # A list of channels that were not read in the last `read` call
+        # because the reader returned immediately when a RayTaskError was found.
+        # These channels must be consumed before the next read to avoid reading
+        # stale data remaining from the last read.
+        self._leftover_channels: List[ChannelInterface] = []
+
+    def get_num_reads(self) -> int:
+        return self._num_reads
+
+    def start(self):
+        raise NotImplementedError
+
+    def _read_list(self, timeout: Optional[float] = None) -> List[Any]:
+        """Read a list of values from this reader.
+
+        Args:
+            timeout: The maximum time in seconds to wait for reading.
+                None means using default timeout which is infinite, 0 means immediate
+                timeout (immediate success or timeout without blocking), -1 means
+                infinite timeout (block indefinitely).
+
+        Returns:
+            The list of values read from the underlying input channels.
+        """
+        raise NotImplementedError
+
+    def read(self, timeout: Optional[float] = None) -> List[Any]:
+        """Read from this reader.
+
+        Args:
+            timeout: The maximum time in seconds to wait for reading.
+                None means using default timeout, 0 means immediate timeout
+                (immediate success or timeout without blocking), -1 means
+                infinite timeout (block indefinitely).
+
+        Returns:
+            The list of values read from this reader.
+        """
+        assert (
+            timeout is None or timeout >= 0 or timeout == -1
+        ), "Timeout must be non-negative or -1."
+        outputs = self._read_list(timeout)
+        self._num_reads += 1
+        return outputs
+
+    def close(self) -> None:
+        self._closed = True
+        for channel in self._input_channels:
+            channel.close()
+
+    def _consume_leftover_channels_if_needed(
+        self, timeout: Optional[float] = None
+    ) -> None:
+        # Consume the channels that were not read in the last `read` call because a
+        # RayTaskError was returned from another channel. If we don't do this, the
+        # read operation will read stale versions of the object refs.
+        #
+        # If a RayTaskError is returned from a leftover channel, it will be ignored.
+        # If a read operation times out, a RayChannelTimeoutError exception will be
+        # raised.
+        #
+        # TODO(kevin85421): Currently, a DAG with NCCL channels and fast fail enabled
+        # may not be reusable. Revisit this in the future.
+        for c in self._leftover_channels:
+            start_time = time.monotonic()
+            c.read(timeout)
+            if timeout is not None:
+                timeout -= time.monotonic() - start_time
+                timeout = max(timeout, 0)
+        self._leftover_channels = []
+
+
+@DeveloperAPI
+class SynchronousReader(ReaderInterface):
+    def __init__(
+        self,
+        input_channels: List[ChannelInterface],
+    ):
+        super().__init__(input_channels)
+
+    def start(self):
+        pass
+
+    def _read_list(self, timeout: Optional[float] = None) -> List[Any]:
+        self._consume_leftover_channels_if_needed(timeout)
+        # We don't update `remaining_timeout` here because in the worst case,
+        # consuming leftover channels requires reading all `_input_channels`,
+        # which users expect to complete within the original `timeout`. Updating
+        # `remaining_timeout` could cause unexpected timeouts in subsequent read
+        # operations.
+
+        # It is a special case that `timeout` is set to 0, which means
+        # read once for each channel.
+        is_zero_timeout = timeout == 0
+
+        results = [None for _ in range(len(self._input_channels))]
+        if timeout is None or timeout == -1:
+            timeout = float("inf")
+        timeout_point = time.monotonic() + timeout
+        remaining_timeout = timeout
+
+        from ray.dag import DAGContext
+
+        ctx = DAGContext.get_current()
+        iteration_timeout = ctx.read_iteration_timeout
+
+        # Iterate over the input channels with a shorter timeout for each iteration
+        # to detect RayTaskError early and fail fast.
+        done_channels = set()
+        while len(done_channels) < len(self._input_channels):
+            for i, c in enumerate(self._input_channels):
+                if c in done_channels:
+                    continue
+                try:
+                    result = c.read(min(remaining_timeout, iteration_timeout))
+                    results[i] = result
+                    done_channels.add(c)
+                    if isinstance(result, ray.exceptions.RayTaskError):
+                        # If we raise an exception immediately, it will be considered
+                        # as a system error which will cause the execution loop to
+                        # exit. Hence, return immediately and let `_process_return_vals`
+                        # handle the exception.
+                        #
+                        # Return a list of RayTaskError so that the caller will not
+                        # get an undefined partial result.
+                        self._leftover_channels = [
+                            c for c in self._input_channels if c not in done_channels
+                        ]
+                        return [result for _ in range(len(self._input_channels))]
+                except ray.exceptions.RayChannelTimeoutError as e:
+                    remaining_timeout = max(timeout_point - time.monotonic(), 0)
+                    if remaining_timeout == 0:
+                        raise e
+                    continue
+
+                remaining_timeout = max(timeout_point - time.monotonic(), 0)
+                if remaining_timeout == 0 and not is_zero_timeout:
+                    raise ray.exceptions.RayChannelTimeoutError(
+                        f"Cannot read all channels within {timeout} seconds"
+                    )
+        return results
+
+    def release_channel_buffers(self, timeout: Optional[float] = None) -> None:
+        for c in self._input_channels:
+            start_time = time.monotonic()
+            assert hasattr(
+                c, "release_buffer"
+            ), "release_buffer() is only supported for shared memory channel "
+            "(e.g., Channel, BufferedSharedMemoryChannel, CompositeChannel) "
+            "and used between the last actor and the driver, but got a channel"
+            f" of type {type(c)}."
+            c.release_buffer(timeout)
+            if timeout is not None:
+                timeout -= time.monotonic() - start_time
+                timeout = max(timeout, 0)
+
+
+@DeveloperAPI
+class AwaitableBackgroundReader(ReaderInterface):
+    """
+    Asyncio-compatible channel reader.
+
+    The reader is constructed with an async queue of futures whose values it
+    will fulfill. It uses a threadpool to execute the blocking calls to read
+    from the input channel(s).
+    """
+
+    def __init__(
+        self,
+        input_channels: List[ChannelInterface],
+        fut_queue: asyncio.Queue,
+    ):
+        super().__init__(input_channels)
+        self._fut_queue = fut_queue
+        self._background_task = None
+        self._background_task_executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=1, thread_name_prefix="channel.AwaitableBackgroundReader"
+        )
+
+    def start(self):
+        self._background_task = asyncio.ensure_future(self.run())
+
+    def _run(self):
+        # Give it a default timeout 60 seconds to release the buffers
+        # of the channels that were not read in the last `read` call.
+        self._consume_leftover_channels_if_needed(60)
+
+        results = [None for _ in range(len(self._input_channels))]
+
+        from ray.dag import DAGContext
+
+        ctx = DAGContext.get_current()
+        iteration_timeout = ctx.read_iteration_timeout
+
+        done_channels = set()
+        while len(done_channels) < len(self._input_channels):
+            for i, c in enumerate(self._input_channels):
+                if c in done_channels:
+                    continue
+                try:
+                    result = c.read(iteration_timeout)
+                    results[i] = result
+                    done_channels.add(c)
+                    if isinstance(result, ray.exceptions.RayTaskError):
+                        self._leftover_channels = [
+                            c for c in self._input_channels if c not in done_channels
+                        ]
+                        return [result for _ in range(len(self._input_channels))]
+                except ray.exceptions.RayChannelTimeoutError:
+                    pass
+                if sys.is_finalizing():
+                    return results
+        return results
+
+    async def run(self):
+        loop = asyncio.get_running_loop()
+        while not self._closed:
+            res, fut = await asyncio.gather(
+                loop.run_in_executor(self._background_task_executor, self._run),
+                self._fut_queue.get(),
+                return_exceptions=True,
+            )
+
+            # Set the result on the main thread.
+            fut.set_result(res)
+            # NOTE(swang): If the object is zero-copy deserialized, then it
+            # will stay in scope as long as ret and the future are in scope.
+            # Therefore, we must delete both here after fulfilling the future.
+            del res
+            del fut
+
+    def close(self):
+        super().close()
+        self._background_task_executor.shutdown(cancel_futures=True)
+        self._background_task.cancel()
+
+
+@DeveloperAPI
+class WriterInterface:
+    def __init__(
+        self,
+        output_channels: List[ChannelInterface],
+        output_idxs: List[Optional[Union[int, str]]],
+        is_input: bool = False,
+    ):
+        """
+        Initialize the writer.
+
+        Args:
+            output_channels: The output channels to write to.
+            output_idxs: The indices of the values to write to each channel.
+                This has the same length as `output_channels`. If `is_input` is True,
+                the index can be an integer or a string to retrieve the corresponding
+                value from `args` or `kwargs` in the DAG's input. If `is_input`
+                is False, the entire value is written if the index is None. Otherwise,
+                the value at the specified index in the tuple is written.
+            is_input: Whether the writer is DAG input writer or not.
+        """
+
+        assert len(output_channels) == len(output_idxs)
+        self._output_channels = output_channels
+        self._output_idxs = output_idxs
+        self._closed = False
+        self._num_writes = 0
+        self._is_input = is_input
+
+    def get_num_writes(self) -> int:
+        return self._num_writes
+
+    def start(self):
+        raise NotImplementedError()
+
+    def write(self, val: Any, timeout: Optional[float] = None) -> None:
+        """Write the value.
+
+        Args:
+            val: The value to write to the output channels.
+            timeout: The maximum time in seconds to wait for writing. 0 means
+                immediate timeout (immediate success or timeout without blocking).
+                -1 and None mean infinite timeout (blocks indefinitely).
+        """
+        raise NotImplementedError()
+
+    def close(self) -> None:
+        self._closed = True
+        for channel in self._output_channels:
+            channel.close()
+
+
+def _adapt(raw_args: Any, key: Optional[Union[int, str]], is_input: bool):
+    """Adapt the raw arguments to the key.
+
+    If ``is_input`` is True, this method will retrieve the value from the input
+    data for an InputAttributeNode. Otherwise, it will retrieve either a partial
+    value or the entire value from the output of a ClassMethodNode.
+
+    Args:
+        raw_args: The raw arguments to adapt.
+        key: The key to adapt.
+        is_input: Whether the writer is DAG input writer or not.
+
+    Returns:
+        The value retrieved from ``raw_args`` according to ``key`` and
+        ``is_input``.
+    """
+    if is_input:
+        if not isinstance(raw_args, CompiledDAGArgs):
+            # Fast path for a single input.
+            return raw_args
+        else:
+            args = raw_args.args
+            kwargs = raw_args.kwargs
+
+        if isinstance(key, int):
+            return args[key]
+        else:
+            return kwargs[key]
+    else:
+        if key is not None:
+            return raw_args[key]
+        else:
+            return raw_args
+
+
+@DeveloperAPI
+class SynchronousWriter(WriterInterface):
+    def start(self):
+        for channel in self._output_channels:
+            channel.ensure_registered_as_writer()
+
+    def write(self, val: Any, timeout: Optional[float] = None) -> None:
+        # If it is an exception, there's only 1 return value.
+        # We have to send the same data to all channels.
+        if isinstance(val, Exception):
+            if len(self._output_channels) > 1:
+                val = tuple(val for _ in range(len(self._output_channels)))
+
+        if not self._is_input:
+            if len(self._output_channels) > 1:
+                if not isinstance(val, tuple):
+                    raise ValueError(
+                        f"Expected a tuple of {len(self._output_channels)} outputs, "
+                        f"but got {type(val)}"
+                    )
+                if len(val) != len(self._output_channels):
+                    raise ValueError(
+                        f"Expected {len(self._output_channels)} outputs, but got "
+                        f"{len(val)} outputs"
+                    )
+
+        for i, channel in enumerate(self._output_channels):
+            idx = self._output_idxs[i]
+            val_i = _adapt(val, idx, self._is_input)
+            channel.write(val_i, timeout)
+        self._num_writes += 1
+
+
+@DeveloperAPI
+class AwaitableBackgroundWriter(WriterInterface):
+    def __init__(
+        self,
+        output_channels: List[ChannelInterface],
+        output_idxs: List[Optional[Union[int, str]]],
+        is_input: bool = False,
+    ):
+        super().__init__(output_channels, output_idxs, is_input=is_input)
+        self._queue = asyncio.Queue()
+        self._background_task = None
+        self._background_task_executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=1, thread_name_prefix="channel.AwaitableBackgroundWriter"
+        )
+
+    def start(self):
+        for channel in self._output_channels:
+            channel.ensure_registered_as_writer()
+        self._background_task = asyncio.ensure_future(self.run())
+
+    def _run(self, res):
+        if not self._is_input:
+            if len(self._output_channels) > 1:
+                if not isinstance(res, tuple):
+                    raise ValueError(
+                        f"Expected a tuple of {len(self._output_channels)} outputs, "
+                        f"but got {type(res)}"
+                    )
+                if len(res) != len(self._output_channels):
+                    raise ValueError(
+                        f"Expected {len(self._output_channels)} outputs, but got "
+                        f"{len(res)} outputs"
+                    )
+
+        for i, channel in enumerate(self._output_channels):
+            idx = self._output_idxs[i]
+            res_i = _adapt(res, idx, self._is_input)
+            exiting = retry_and_check_interpreter_exit(
+                lambda: channel.write(res_i, timeout=1)
+            )
+            if exiting:
+                break
+
+    async def run(self):
+        loop = asyncio.get_event_loop()
+        while True:
+            res = await self._queue.get()
+            await loop.run_in_executor(self._background_task_executor, self._run, res)
+
+    async def write(self, val: Any) -> None:
+        if self._closed:
+            raise RuntimeError("DAG execution cancelled")
+        await self._queue.put(val)
+        self._num_writes += 1
+
+    def close(self):
+        self._background_task.cancel()
+        super().close()

python/ray/experimental/channel/communicator.py

@@ -0,0 +1,203 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
+
+import ray
+from ray.experimental.util.types import ReduceOp
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    import torch
+
+
+# Signature for a torch.Tensor allocator is:
+# (shape: Tuple[int], dtype: torch.dtype) -> torch.Tensor.
+TorchTensorAllocator = Callable[[Tuple[int], "torch.dtype"], "torch.Tensor"]
+
+
+@DeveloperAPI
+class Communicator(ABC):
+    """
+    Communicator for a group of Compiled Graph actors on NVIDIA GPU.
+
+    The Compiled Graph execution leverages this internally to support communication
+    between actors in the group.
+    """
+
+    @abstractmethod
+    def initialize(self, rank: int) -> None:
+        """
+        Initialize the communicator from the actor.
+
+        This is called once by Compiled Graph on each actor to initialize the
+        communicator,before any other methods.
+
+        Args:
+            rank: The rank of this actor in the group.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        """
+        Get handles of all actors for this communicator group.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_rank(self, actor: ray.actor.ActorHandle) -> int:
+        """Return the given actor's rank in the group.
+
+        Args:
+            actor: The actor handle to look up.
+
+        Returns:
+            The rank of ``actor`` within the communicator group.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_self_rank(self) -> Optional[int]:
+        """
+        Return this actor's rank.
+        """
+        raise NotImplementedError
+
+    def get_world_size(self) -> int:
+        """
+        Return the number of ranks in the group.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def send(self, value: "torch.Tensor", peer_rank: int) -> None:
+        """
+        Send a torch.Tensor to a peer.
+
+        This returns when the send kernel has been queued, but the kernel may
+        not have completed. Therefore, the caller should ensure that there are
+        no concurrent writes to the sent `value` until the send has finished.
+
+        Args:
+            value: The torch.Tensor to send. It should already be on this
+                actor's default device.
+            peer_rank: The rank of the actor to send to.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def recv(
+        self,
+        shape: Tuple[int],
+        dtype: "torch.dtype",
+        peer_rank: int,
+        allocator: Optional[TorchTensorAllocator] = None,
+    ) -> "torch.Tensor":
+        """Receive a torch.Tensor from a peer and synchronize.
+
+        After this call returns, the receive buffer is safe to read from from
+        any stream. An RayChannelError will be raised if an error occurred (e.g.,
+        remote actor died), and the buffer is not safe to read.
+
+        Args:
+            shape: The shape of the tensor to receive.
+            dtype: The dtype of the tensor to receive.
+            peer_rank: The rank of the actor to receive from.
+            allocator: A function to allocate the tensor to receive into.
+
+        Returns:
+            The tensor received from ``peer_rank``.
+        """
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def recv_stream(self):
+        """
+        Return the torch stream context used for receiving tensors.
+        """
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def send_stream(self):
+        """
+        Return the torch stream context used for sending tensors.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def allgather(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+    ) -> None:
+        """
+        Collectively allgather the tensor across the group.
+
+        Args:
+            send_buf: The input torch.tensor to allgather. It should already be
+                on this actor's default device.
+            recv_buf: The output torch.tensor to store the allgather result.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def allreduce(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp,
+    ) -> None:
+        """
+        Collectively allreduce the tensor across the group.
+
+        Args:
+            send_buf: The input torch.tensor to allreduce. It should already be
+                on this actor's default device.
+            recv_buf: The output torch.tensor to store the allreduce result.
+            op: The reduce operation.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def reducescatter(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp,
+    ) -> None:
+        """
+        Collectively reducescatter the tensor across the group.
+
+        Args:
+            send_buf: The input torch.tensor to reducescatter. It should already be
+                on this actor's default device.
+            recv_buf: The output torch.tensor to store the reducescatter result.
+            op: The reduce operation.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def destroy(self) -> None:
+        """
+        Destroy the GPU communicator.
+
+        Any destruction and cleanup for the GPU communicator should be
+        done here. Implement as a noop is nothing is needed.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_transport_name(self) -> str:
+        """
+        Return the type of the communicator (gpu or cpu).
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def generate_communicator_id(cls) -> str:
+        """
+        Return the unique id of the communicator.
+        """
+        raise NotImplementedError

python/ray/experimental/channel/communicator_handle.py

@@ -0,0 +1,28 @@
+from typing import List
+
+import ray
+
+
+class CommunicatorHandle:
+    """
+    A lightweight communicator handle used by the driver to store handles to
+    the actors in the communicator.
+    """
+
+    def __init__(
+        self,
+        actor_handles: List["ray.actor.ActorHandle"],
+    ):
+        """
+        Initializes the CommunicatorHandle with the given actor handles.
+
+        Args:
+            actor_handles: A list of actor handles to be stored.
+        """
+        self._actor_handles = actor_handles
+
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        """
+        Retuan all actor handles in this communicator.
+        """
+        return self._actor_handles

python/ray/experimental/channel/conftest.py

@@ -0,0 +1,185 @@
+import asyncio
+from collections import defaultdict
+from typing import Optional, Tuple
+from unittest import mock
+
+import torch
+
+import ray
+import ray.dag
+import ray.experimental.channel as ray_channel
+from ray.experimental.channel import nccl_group
+from ray.experimental.channel.communicator import TorchTensorAllocator
+from ray.experimental.util.types import Device
+
+
+@ray.remote(num_cpus=0)
+class Barrier:
+    """
+    Barrier that blocks the given number of actors until all actors have
+    reached the barrier. This is used to mock out blocking NCCL ops.
+    """
+
+    def __init__(self, num_actors=2):
+        self.num_actors = num_actors
+        self.condition = asyncio.Condition()
+        # Buffer for the data that is "sent" between the actors, each entry is
+        # one p2p op.
+        self.data = {}
+        # Buffer for the number of actors seen, each entry is one p2p op.
+        self.num_actors_seen = defaultdict(int)
+
+        # Add a new mock for the TorchTensorType.device property
+        device_property_patcher = mock.patch(
+            "ray.experimental.channel.torch_tensor_type.TorchTensorType.device",
+            new_callable=mock.PropertyMock,
+            return_value=Device.CPU,
+        )
+        device_property_patcher.start()
+
+    async def wait(self, idx: int, data=None):
+        """
+        Wait at barrier until all actors have sent `idx`. One actor should
+        provide `data`, and this value will be returned by this method for all
+        other actors.
+        """
+        async with self.condition:
+            if data is not None:
+                assert idx not in self.data, (self.data, self.num_actors_seen)
+                self.data[idx] = data
+            self.num_actors_seen[idx] += 1
+
+            if self.num_actors_seen[idx] == self.num_actors:
+                # Wake up all tasks waiting on this condition.
+                self.condition.notify_all()
+            else:
+                await self.condition.wait_for(
+                    lambda: self.num_actors_seen[idx] == self.num_actors
+                )
+
+            if data is None:
+                data = self.data[idx]
+
+        return data
+
+
+class MockCudaStream:
+    def __init__(self):
+        self.cuda_stream = 0
+
+    def synchronize(self):
+        pass
+
+
+class MockNcclGroup(nccl_group._NcclGroup):
+    """
+    Mock the internal _NcclGroup to use a barrier actor instead of a NCCL group
+    for communication.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # We use the op index to synchronize the sender and receiver at the
+        # barrier.
+        self.num_ops = defaultdict(int)
+        self.barriers = set()
+
+    def send(self, tensor: torch.Tensor, peer_rank: int):
+        # "Send" the tensor to the barrier actor.
+        barrier_key = sorted([self.get_self_rank(), peer_rank])
+        barrier_key = f"barrier-{barrier_key[0]}-{barrier_key[1]}"
+        barrier = ray.get_actor(name=barrier_key)
+        self.barriers.add(barrier)
+        ray.get(barrier.wait.remote(self.num_ops[barrier_key], tensor))
+        self.num_ops[barrier_key] += 1
+
+    def recv(
+        self,
+        shape: Tuple[int],
+        dtype: torch.dtype,
+        peer_rank: int,
+        allocator: Optional[TorchTensorAllocator] = None,
+    ):
+        # "Receive" the tensor from the barrier actor.
+        barrier_key = sorted([self.get_self_rank(), peer_rank])
+        barrier_key = f"barrier-{barrier_key[0]}-{barrier_key[1]}"
+        barrier = ray.get_actor(name=barrier_key)
+        self.barriers.add(barrier)
+        received_tensor = ray.get(barrier.wait.remote(self.num_ops[barrier_key]))
+        assert (
+            allocator is not None
+        ), "torch tensor allocator is required for MockNcclGroup"
+        buf = allocator(shape, dtype)
+        buf[:] = received_tensor[:]
+        self.num_ops[barrier_key] += 1
+        return buf
+
+    def destroy(self) -> None:
+        for barrier in self.barriers:
+            ray.kill(barrier)
+
+
+def start_nccl_mock():
+    """
+    Patch methods that require CUDA.
+    """
+    # Mock cupy dependencies.
+    nccl_mock = mock.MagicMock()
+    nccl_mock.nccl.get_unique_id.return_value = 0
+    cp_patcher = mock.patch.dict(
+        "sys.modules",
+        {
+            "cupy.cuda": nccl_mock,
+            "cupy": mock.MagicMock(),
+            "ray.util.collective.collective_group": mock.MagicMock(),
+        },
+    )
+    cp_patcher.start()
+
+    # Mock send/recv ops to use an actor instead of NCCL.
+    ray.experimental.channel.nccl_group._NcclGroup = MockNcclGroup
+
+    # PyTorch mocks.
+    stream_patcher = mock.patch(
+        "torch.cuda.current_stream", new_callable=lambda: MockCudaStream
+    )
+    stream_patcher.start()
+    new_stream_patcher = mock.patch(
+        "torch.cuda.Stream", new_callable=lambda: MockCudaStream
+    )
+    new_stream_patcher.start()
+    tensor_patcher = mock.patch("torch.Tensor.device", torch.device("cuda"))
+    tensor_patcher.start()
+    tensor_patcher = mock.patch("torch.Tensor.is_cuda", True)
+    tensor_patcher.start()
+    tensor_allocator_patcher = mock.patch(
+        "ray.experimental.channel.torch_tensor_accelerator_channel._torch_tensor_allocator",
+        lambda shape, dtype: torch.empty(shape, dtype=dtype),
+    )
+    tensor_allocator_patcher.start()
+
+    # Add a new mock for the TorchTensorType.device property
+    device_property_patcher = mock.patch(
+        "ray.experimental.channel.torch_tensor_type.TorchTensorType.device",
+        new_callable=mock.PropertyMock,
+        return_value=Device.CPU,
+    )
+    device_property_patcher.start()
+
+    ctx = ray_channel.ChannelContext.get_current()
+    ctx.set_torch_device(torch.device("cuda"))
+
+
+class TracedChannel(ray_channel.shared_memory_channel.Channel):
+    """
+    Patched Channel that records all write ops for testing.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.ops = []
+
+    def write(self, *args, **kwargs):
+        self.ops.append((args, kwargs))
+        return super().write(*args, **kwargs)

python/ray/experimental/channel/cpu_communicator.py

@@ -0,0 +1,209 @@
+import asyncio
+from collections import defaultdict
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+
+import ray
+from ray.experimental.channel.communicator import (
+    Communicator,
+    ReduceOp,
+    TorchTensorAllocator,
+)
+
+if TYPE_CHECKING:
+    import torch
+
+
+@ray.remote(num_cpus=0)
+class CPUCommBarrier:
+    """
+    Barrier actor that blocks the given number of actors until all actors have
+    reached the Barrier.
+
+    p2p operations are not done here (completed via shared memory channel).
+    """
+
+    def __init__(self, num_actors: int):
+        self.num_actors = num_actors
+        self.condition = asyncio.Condition()
+        # Stores the data for each collective operation
+        self.collective_data: Dict[int, List["torch.Tensor"]] = defaultdict(list)
+        # Stores the shape of data for each collective operation
+        self.collective_data_shape: Dict[int, "torch.Tensor.type"] = {}
+        # Buffer for the number of actors seen
+        self.num_actors_seen = defaultdict(int)
+        # Number of actors who have read the result, and are about to exit the function.
+        # State is kept so we only garbage collect after the last actor has read the
+        # relevant data.
+        self.num_actors_read = defaultdict(int)
+
+    async def wait_collective(self, op_id: int, data: "torch.Tensor", op: ReduceOp):
+        """
+        Wait at the communicator until all actors have sent `op_id` and `data`.
+        Once data from all actors is received, execute the collective `op`
+        on the communicator actor and return the result.
+        """
+        async with self.condition:
+            self.collective_data[op_id].append(data)
+            self.num_actors_seen[op_id] += 1
+
+            if self.num_actors_seen[op_id] == self.num_actors:
+                # Apply the collective operation across all gathered tensors
+                data = self._apply_op(op, self.collective_data[op_id])
+                self.collective_data[op_id] = data
+                self.condition.notify_all()
+            else:
+                await self.condition.wait_for(
+                    lambda: self.num_actors_seen[op_id] == self.num_actors
+                )
+
+            data = self.collective_data[op_id]
+            self.num_actors_read[op_id] += 1
+
+            if self.num_actors_read[op_id] == self.num_actors:
+                del self.collective_data[op_id]
+                del self.num_actors_seen[op_id]
+                del self.num_actors_read[op_id]
+
+            return data
+
+    def _apply_op(self, op: ReduceOp, tensors: List["torch.Tensor"]) -> "torch.Tensor":
+        """Apply the specified reduction operation across a list of tensors."""
+
+        result = tensors[0].clone()
+        if op == ReduceOp.SUM:
+            for tensor in tensors[1:]:
+                result += tensor
+        elif op == ReduceOp.PRODUCT:
+            for tensor in tensors[1:]:
+                result *= tensor
+        elif op == ReduceOp.MAX:
+            for tensor in tensors[1:]:
+                result = torch.max(result, tensor)
+        elif op == ReduceOp.MIN:
+            for tensor in tensors[1:]:
+                result = torch.min(result, tensor)
+        elif op == ReduceOp.AVG:
+            result = sum(tensors) / len(tensors)
+        else:
+            raise ValueError(f"Operation {op} not supported")
+        return result
+
+
+class CPUCommunicator(Communicator):
+    """
+    Uses a CPU-based communicator actor instead of an accelerator group like NCCL.
+    """
+
+    def __init__(self, world_size: int, actor_handles: List["ray.actor.ActorHandle"]):
+        """We use the op index to synchronize the sender and receiver at the
+        communicator actor."""
+        self._world_size = world_size
+        self._actor_handles = actor_handles
+        self.num_ops = defaultdict(int)
+
+        # For collective communication, one barrier will be created for
+        # each unique group of participants.
+        self.barriers = set()
+        self._rank = None
+
+    def send(self, tensor: "torch.Tensor", peer_rank: int):
+        # p2p operations are done via a shared memory channel, initialized in
+        # `create_channel` of `TorchTensorType`
+        pass
+
+    def recv(
+        self,
+        shape: Tuple[int],
+        dtype: "torch.dtype",
+        peer_rank: int,
+        allocator: Optional[TorchTensorAllocator] = None,
+    ):
+        # See the comment on `send`
+        pass
+
+    def allgather(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+    ):
+        raise NotImplementedError
+
+    def allreduce(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp = ReduceOp.SUM,
+    ):
+        all_ranks = [
+            self.get_rank(actor_handle) for actor_handle in self.get_actor_handles()
+        ]
+        barrier_key = "barrier-collective-" + "-".join(map(str, sorted(all_ranks)))
+        barrier = CPUCommBarrier.options(name=barrier_key, get_if_exists=True).remote(
+            self._world_size
+        )
+        self.barriers.add(barrier)
+
+        result = ray.get(
+            barrier.wait_collective.remote(self.num_ops[barrier_key], send_buf, op)
+        )
+        assert recv_buf is not None, "Receiving buffer required for CPUCommunicator"
+        recv_buf[:] = result[:]
+        self.num_ops[barrier_key] += 1
+
+    def reducescatter(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp = ReduceOp.SUM,
+    ):
+        raise NotImplementedError
+
+    def destroy(self) -> None:
+        for barrier in self.barriers:
+            ray.kill(barrier)
+
+    def initialize(self, rank: int) -> None:
+        self._rank = rank
+
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        return self._actor_handles
+
+    def get_rank(self, actor: ray.actor.ActorHandle) -> int:
+        """Return the given actor's rank in the CPU communicator.
+
+        Args:
+            actor: The actor handle to look up.
+
+        Returns:
+            The rank of ``actor`` within the CPU communicator group.
+        """
+        actor_ids = [a._ray_actor_id for a in self._actor_handles]
+        try:
+            rank = actor_ids.index(actor._ray_actor_id)
+        except ValueError:
+            raise ValueError("Actor is not in the CPUCommunicator group.")
+        return rank
+
+    def get_self_rank(self) -> Optional[int]:
+        return self._rank
+
+    def get_world_size(self) -> int:
+        """
+        Return the number of ranks in the CPU communicator.
+        """
+        return self._world_size
+
+    def get_transport_name(self) -> str:
+        return "cpu"
+
+    def recv_stream(self):
+        raise NotImplementedError
+
+    def send_stream(self):
+        raise NotImplementedError
+
+    @classmethod
+    def generate_communicator_id(cls) -> str:
+        import uuid
+
+        return str(uuid.uuid4())

python/ray/experimental/channel/intra_process_channel.py

@@ -0,0 +1,72 @@
+import uuid
+from typing import Any, Optional
+
+from ray.experimental.channel import ChannelContext
+from ray.experimental.channel.common import ChannelInterface
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI(stability="alpha")
+class IntraProcessChannel(ChannelInterface):
+    """IntraProcessChannel is a channel for communication between two tasks in the same
+    worker process. It writes data directly to the worker's _SerializationContext
+    and reads data from the _SerializationContext to avoid the serialization
+    overhead and the need for reading/writing from shared memory. Note that if the
+    readers may mutate the data, users should deep copy the data themselves to avoid
+    side effects.
+
+    Args:
+        num_readers: The number of readers that will read from this channel. Readers
+            can be the same method of the same actor.
+        _channel_id: Optional pre-generated channel identifier. If ``None``, a
+            new UUID4 is generated. Used internally for re-creating channels
+            after pickling.
+    """
+
+    def __init__(
+        self,
+        num_readers: int,
+        _channel_id: Optional[str] = None,
+    ):
+        # Generate a unique ID for the channel. The writer and reader will use
+        # this ID to store and retrieve data from the _SerializationContext.
+        self._channel_id = _channel_id
+        self._num_readers = num_readers
+        if self._channel_id is None:
+            self._channel_id = str(uuid.uuid4())
+
+    def ensure_registered_as_writer(self) -> None:
+        pass
+
+    def ensure_registered_as_reader(self) -> None:
+        pass
+
+    def __reduce__(self):
+        return IntraProcessChannel, (
+            self._num_readers,
+            self._channel_id,
+        )
+
+    def __str__(self) -> str:
+        return f"IntraProcessChannel(channel_id={self._channel_id})"
+
+    def write(self, value: Any, timeout: Optional[float] = None):
+        self.ensure_registered_as_writer()
+        # No need to check timeout as the operation is non-blocking.
+
+        # Because both the reader and writer are in the same worker process,
+        # we can directly store the data in the context instead of storing
+        # it in the channel object. This removes the serialization overhead of `value`.
+        ctx = ChannelContext.get_current().serialization_context
+        ctx.set_data(self._channel_id, value, self._num_readers)
+
+    def read(self, timeout: Optional[float] = None, deserialize: bool = True) -> Any:
+        self.ensure_registered_as_reader()
+        assert deserialize, "Data passed from the actor to itself is never serialized"
+        # No need to check timeout as the operation is non-blocking.
+        ctx = ChannelContext.get_current().serialization_context
+        return ctx.get_data(self._channel_id)
+
+    def close(self) -> None:
+        ctx = ChannelContext.get_current().serialization_context
+        ctx.reset_data(self._channel_id)

python/ray/experimental/channel/nccl_group.py

@@ -0,0 +1,380 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.accelerator_context import AcceleratorContext
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+class _NcclGroup(Communicator):
+    """
+    Represents an actor's NCCL communicator. This is the default NCCL communicator
+    to be used in Compiled Graph if a custom communicator is not provided.
+
+    This class is not thread-safe.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        comm_id: tuple,
+        rank: Optional[int],
+        actor_handles: List["ray.actor.ActorHandle"],
+        cuda_stream: Optional["torch.cuda.Stream"],
+        use_communication_streams: bool = False,
+    ):
+        """
+        Initialize a NCCL communicator that can be used to communicate p2p with
+        other GPU actors.
+
+        This method blocks until the same call has been made on all other
+        actors in the group, with the same arguments for world_size and
+        comm_id.
+
+        NOTE: A concurrent NCCL group can coexist with this one but using the
+        two groups concurrently on different CUDA streams may cause deadlock.
+        See
+        https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/communicators.html
+        #using-multiple-nccl-communicators-concurrently.
+
+        If the user can guarantee that all involved actors execute the same ops
+        in the same order, then the other NCCL group should use the given
+        `cuda_stream`, and there will not be a concurrency issue. Otherwise,
+        the other stream needs to synchronize with the given `cuda_stream`
+        before and after it launches NCCL ops, e.g., at the beginning and end
+        of a DAG task.
+
+        Args:
+            world_size: The number of participating actors/devices.
+            comm_id: A unique communicator ID returned by
+                cupy.cuda.nccl.get_unique_id().
+            rank: The rank of this actor. If None, then the caller is not a
+                participant of the NCCL group.
+            actor_handles: A list of actor handles, in rank order.
+            cuda_stream: A raw CUDA stream to dispatch NCCL ops to. If rank is
+                specified, then this must be specified too.
+            use_communication_streams: Whether to use dedicated send and recv
+                streams for communication. If True, communication and computation
+                can be overlapped to improve performance.
+        """
+        self._world_size = world_size
+        self._rank: Optional[int] = rank
+        self.nccl_util: Optional[ModuleType] = None
+        self._actor_handles = actor_handles
+        self._use_communication_streams = use_communication_streams
+
+        if rank is not None:
+            assert ray.get_gpu_ids(), "NCCL actor has no GPUs assigned"
+            assert cuda_stream is not None, "NCCL actor must specify cuda_stream"
+
+            expected_rank = self.get_rank(ray.get_runtime_context().current_actor)
+            assert (
+                rank == expected_rank
+            ), f"NCCL actor's rank {rank} does not match expected rank {expected_rank}"
+
+            from ray.util.collective.collective_group import nccl_util
+
+            self.nccl_util = nccl_util
+            self._comm = self.nccl_util.NcclCommunicator(world_size, comm_id, rank)
+        else:
+            # Driver does not have a rank.
+            self._comm = None
+
+        self._cuda_stream: Optional["torch.cuda.Stream"] = None
+        self._send_stream: Optional["torch.cuda.Stream"] = None
+        self._recv_stream: Optional["torch.cuda.Stream"] = None
+        if cuda_stream is not None:
+            assert rank is not None, "NCCL actor has no rank assigned"
+            self._cuda_stream = cuda_stream
+
+            if use_communication_streams:
+                import torch
+
+                # TODO(swang): Allow default device to be overridden.
+                device = AcceleratorContext.get().get_accelerator_devices()[0]
+
+                self._send_stream = torch.cuda.Stream(device=device)
+                self._recv_stream = torch.cuda.Stream(device=device)
+            else:
+                self._send_stream = self._cuda_stream
+                self._recv_stream = self._cuda_stream
+
+        self._closed = False
+
+    def initialize(self, rank: int) -> None:
+        # No additional initialization is needed.
+        pass
+
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        return self._actor_handles
+
+    def get_rank(self, actor: ray.actor.ActorHandle) -> int:
+        """Return the given actor's rank in the NCCL communicator.
+
+        Args:
+            actor: The actor handle to look up.
+
+        Returns:
+            The rank of ``actor`` within the NCCL group.
+        """
+        actor_ids = [a._ray_actor_id for a in self._actor_handles]
+        try:
+            rank = actor_ids.index(actor._ray_actor_id)
+        except ValueError:
+            raise ValueError("Actor is not in the NCCL group.")
+        return rank
+
+    def get_self_rank(self) -> Optional[int]:
+        """
+        Return this actor's rank.
+        """
+        return self._rank
+
+    def get_world_size(self) -> int:
+        """
+        Return the number of ranks in the NCCL communicator.
+        """
+        return self._world_size
+
+    def send(self, buf: "torch.Tensor", peer_rank: int) -> None:
+        """
+        Send a torch.Tensor to a peer.
+
+        This returns when the send kernel has been queued, but the kernel may
+        not have completed. Therefore, the caller should ensure that there are
+        no concurrent writes to the sent `buf` until the send has finished.
+        That is, either all writes should be submitted on the current stream
+        (self._cuda_stream) or, if on a different stream, that stream should
+        synchronize with the current stream.
+
+        Args:
+            buf: The torch.Tensor to send. It should already be on this
+                actor's default device.
+            peer_rank: The rank of the actor to send to.
+        """
+        if self._closed:
+            raise RayChannelError("NCCL group has been destroyed.")
+
+        if self._use_communication_streams:
+            # We observed that if all recv/compute/send operations run on GPU,
+            # since there is no synchronization, the CPU execution loop may be
+            # far ahead of the GPU operations and lead to runtime failures.
+            # To avoid that, we synchronize on the send stream.
+            # TODO(rui): find a better approach
+            self._send_stream.synchronize()
+
+        # TODO(swang): Handle send/recv async NCCL errors such as network
+        # failures.
+        self._comm.send(
+            self.nccl_util.get_tensor_ptr(buf),
+            buf.numel(),
+            self.nccl_util.get_nccl_tensor_dtype(buf),
+            peer_rank,
+            self._send_stream.cuda_stream,
+        )
+
+    def recv(
+        self,
+        shape: Tuple[int],
+        dtype: "torch.dtype",
+        peer_rank: int,
+        allocator: Optional[TorchTensorAllocator] = None,
+    ) -> "torch.Tensor":
+        """Receive a torch.Tensor from a peer and synchronize the current stream.
+
+        After this call returns, the receive buffer is safe to read from from
+        any stream. An RayChannelError will be raised if an error occurred (e.g.,
+        remote actor died), and the buffer is not safe to read.
+
+        Args:
+            shape: The shape of the tensor to receive.
+            dtype: The dtype of the tensor to receive.
+            peer_rank: The rank of the actor to receive from.
+            allocator: A function used to allocate the receive buffer.
+
+        Returns:
+            The tensor received from ``peer_rank``.
+        """
+        if self._closed:
+            raise RayChannelError("NCCL group has been destroyed.")
+        assert allocator is not None, "NCCL group requires a tensor allocator"
+        buf = allocator(shape, dtype)
+
+        if self._use_communication_streams:
+            # We observed that if all recv/compute/send operations run on GPU,
+            # since there is no synchronization, the CPU execution loop may be
+            # far ahead of the GPU operations and lead to runtime failures.
+            # To avoid that, we synchronize on the recv stream.
+            # TODO(rui): find a better approach
+            self._recv_stream.synchronize()
+
+            self._comm.recv(
+                self.nccl_util.get_tensor_ptr(buf),
+                buf.numel(),
+                self.nccl_util.get_nccl_tensor_dtype(buf),
+                peer_rank,
+                self._recv_stream.cuda_stream,
+            )
+        else:
+            self._comm.recv(
+                self.nccl_util.get_tensor_ptr(buf),
+                buf.numel(),
+                self.nccl_util.get_nccl_tensor_dtype(buf),
+                peer_rank,
+                self._recv_stream.cuda_stream,
+            )
+
+            # Buffer values are undefined if NCCL ops are aborted. Therefore, we
+            # need to synchronize here and check that the channel is still open to
+            # ensure that the receive buffer is valid.
+            # TODO(swang): Avoid CUDA synchronization.
+            self._cuda_stream.synchronize()
+
+        if self._closed:
+            raise RayChannelError("NCCL group has been destroyed.")
+        return buf
+
+    def _exec_collective(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        operation: "Callable[..., None]",
+        *operation_args,
+    ):
+        if self._closed:
+            raise RayChannelError("NCCL group has been destroyed.")
+
+        assert send_buf.dtype == recv_buf.dtype, (
+            "Ray Compiled Graph derived the dtype of recv_buf from send_buf, "
+            "so send_buf and recv_buf must have the same dtype. "
+            "If you see this error, please file an issue at Ray repository."
+        )
+
+        operation(*operation_args)
+
+        # Buffer values are undefined if NCCL ops are aborted. Therefore, we
+        # need to synchronize here and check that the channel is still open to
+        # ensure that the receive buffer is valid.
+        # TODO(swang): Avoid CUDA synchronization.
+        # TODO(wxdeng): This synchronize will be optional after merging the unify PR.
+        self._cuda_stream.synchronize()
+        if self._closed:
+            raise RayChannelError(
+                "NCCL group has been destroyed during allreduce operation. "
+                "There may be a dtype mismatch between input tensors from "
+                "different ranks."
+            )
+
+    def allgather(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+    ):
+        operation_args = [
+            self.nccl_util.get_tensor_ptr(send_buf),
+            self.nccl_util.get_tensor_ptr(recv_buf),
+            send_buf.numel(),
+            self.nccl_util.get_nccl_tensor_dtype(send_buf),
+            self._cuda_stream.cuda_stream,
+        ]
+        self._exec_collective(
+            send_buf,
+            recv_buf,
+            self._comm.allGather,
+            *operation_args,
+        )
+
+    def allreduce(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp = ReduceOp.SUM,
+    ):
+        operation_args = [
+            self.nccl_util.get_tensor_ptr(send_buf),
+            self.nccl_util.get_tensor_ptr(recv_buf),
+            send_buf.numel(),
+            self.nccl_util.get_nccl_tensor_dtype(send_buf),
+            op.value,
+            self._cuda_stream.cuda_stream,
+        ]
+        self._exec_collective(
+            send_buf,
+            recv_buf,
+            self._comm.allReduce,
+            *operation_args,
+        )
+
+    def reducescatter(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp = ReduceOp.SUM,
+    ):
+        operation_args = [
+            self.nccl_util.get_tensor_ptr(send_buf),
+            self.nccl_util.get_tensor_ptr(recv_buf),
+            recv_buf.numel(),
+            self.nccl_util.get_nccl_tensor_dtype(send_buf),
+            op.value,
+            self._cuda_stream.cuda_stream,
+        ]
+        self._exec_collective(
+            send_buf,
+            recv_buf,
+            self._comm.reduceScatter,
+            *operation_args,
+        )
+
+    @property
+    def recv_stream(self):
+        import torch
+
+        return torch.cuda.StreamContext(self._recv_stream)
+
+    @property
+    def send_stream(self):
+        import torch
+
+        return torch.cuda.StreamContext(self._send_stream)
+
+    def destroy(self) -> None:
+        """
+        Destroy the NCCL group.
+        """
+        if self._closed:
+            return
+
+        self._closed = True
+
+        if self._comm is not None:
+            logger.info(
+                "Destructing NCCL group on actor: "
+                f"{ray.get_runtime_context().current_actor}"
+            )
+            # Abort *after* setting the _closed flag. This ensures that NCCL
+            # ops that were blocked on a remote peer will see that the _closed
+            # flag is True when they exit from the abort.
+            self._comm.abort()
+            self._comm.destroy()
+
+    def get_transport_name(self) -> str:
+        return "accelerator"
+
+    @classmethod
+    def generate_communicator_id(cls) -> str:
+        from cupy.cuda import nccl
+
+        return nccl.get_unique_id()

python/ray/experimental/channel/serialization_context.py

@@ -0,0 +1,231 @@
+import warnings
+from typing import TYPE_CHECKING, Any, Dict, List, Set, Tuple, Union
+
+from ray.experimental.util.types import Device
+
+if TYPE_CHECKING:
+    import numpy as np
+    import torch
+
+
+_TORCH_WARNING_FILTER_ACTIVATE = True
+
+
+class _SerializationContext:
+    def __init__(self):
+        # If true, then tensors found in the data to serialize are extracted
+        # and the caller should send them through an external transport.
+        self._use_external_transport: bool = False
+        # If _use_external_transport is True, then these are
+        # the tensors that should be sent or received
+        # out-of-band, through the external transport.
+        self._out_of_band_tensors: List["torch.Tensor"] = []
+        # During serialization, tensors sent out-of-band are replaced with
+        # integer placeholders. This tracks the set of placeholders seen.
+        self._deserialized_tensor_placeholders: Set[int] = set()
+
+        # Buffer for transferring data between tasks in the same worker process.
+        # The key is the channel ID, and the value is the data. We don't use a
+        # lock when reading/writing the buffer because a DAG node actor will only
+        # execute one task at a time in `do_exec_tasks`. It will not execute multiple
+        # Ray tasks on a single actor simultaneously.
+        self.intra_process_channel_buffers: Dict[str, Any] = {}
+        # The number of readers for each channel. When the number of readers
+        # reaches 0, remove the data from the buffer.
+        self.channel_id_to_num_readers: Dict[str, int] = {}
+
+    def set_target_device(self, device: Device) -> None:
+        self._target_device = device
+
+    def set_data(self, channel_id: str, value: Any, num_readers: int) -> None:
+        assert num_readers > 0, "num_readers must be greater than 0."
+        assert (
+            channel_id not in self.intra_process_channel_buffers
+        ), f"Channel {channel_id} already exists in the buffer."
+        assert (
+            channel_id not in self.channel_id_to_num_readers
+        ), f"Channel {channel_id} already exists in the channel_id_to_num_readers."
+
+        self.intra_process_channel_buffers[channel_id] = value
+        self.channel_id_to_num_readers[channel_id] = num_readers
+
+    def has_data(self, channel_id: str) -> bool:
+        return channel_id in self.intra_process_channel_buffers
+
+    def get_data(self, channel_id: str) -> Any:
+        assert (
+            channel_id in self.intra_process_channel_buffers
+        ), f"Channel {channel_id} does not exist in the buffer."
+        assert (
+            channel_id in self.channel_id_to_num_readers
+        ), f"Channel {channel_id} does not exist in the channel_id_to_num_readers."
+
+        self.channel_id_to_num_readers[channel_id] -= 1
+        if self.channel_id_to_num_readers[channel_id] == 0:
+            # All readers have read the data, so we can remove it.
+            self.channel_id_to_num_readers.pop(channel_id)
+            return self.intra_process_channel_buffers.pop(channel_id)
+        return self.intra_process_channel_buffers[channel_id]
+
+    def reset_data(self, channel_id: str) -> None:
+        self.intra_process_channel_buffers.pop(channel_id, None)
+        self.channel_id_to_num_readers.pop(channel_id, None)
+
+    def set_use_external_transport(self, use_external_transport: bool) -> None:
+        self._use_external_transport = use_external_transport
+
+    @property
+    def use_external_transport(self) -> bool:
+        return self._use_external_transport
+
+    def reset_out_of_band_tensors(
+        self, tensors: List["torch.Tensor"]
+    ) -> Tuple[List["torch.Tensor"], Set[int]]:
+        """
+        Return and reset the out-of-band tensors and all tensor placeholders
+        that were deserialized since the last call to reset.
+        """
+        prev_tensors = self._out_of_band_tensors
+        deserialized_tensor_placeholders = self._deserialized_tensor_placeholders
+        self._out_of_band_tensors = tensors
+        self._deserialized_tensor_placeholders = set()
+        return prev_tensors, deserialized_tensor_placeholders
+
+    def serialize_tensor(
+        self, tensor: "torch.Tensor"
+    ) -> Union[int, Tuple["np.ndarray", "torch.dtype", str]]:
+        from ray.experimental.channel import ChannelContext
+
+        ctx = ChannelContext.get_current()
+        if self._use_external_transport and (
+            ctx._torch_device is None or ctx._torch_device == tensor.device
+        ):
+            # External transport is enabled and we found a tensor that matches
+            # our device.  Add the actual tensor to a buffer. The buffer of
+            # tensors should later be popped by the caller and sent via
+            # external transport.
+            self._out_of_band_tensors.append(tensor)
+            # Return a placeholder.
+            return len(self._out_of_band_tensors) - 1
+
+        return self.serialize_to_numpy_or_scalar(tensor)
+
+    def serialize_to_numpy_or_scalar(
+        self, tensor: "torch.Tensor"
+    ) -> Tuple[Union["np.ndarray", Any], "torch.dtype", str]:
+        """
+        Serialize a tensor to a numpy array,
+        or a scalar when the tensor is 0-dim.
+        """
+        import torch
+
+        tensor_device_type = tensor.device.type
+
+        # Transfer through Ray's shared memory store for now.
+        # TODO(swang): This requires two copies, one to transfer from GPU to
+        # CPU and another from CPU to shared memory. Ideally we should elide
+        # the first copy and memcpy directly from GPU to the shared memory
+        # buffer.
+        if tensor_device_type != "cpu":
+            tensor = tensor.to("cpu")
+
+        # Numpy does not have an equivalent dtype for all torch dtypes, so
+        # instead of casting directly to numpy:
+        # 1) for non-scalar tensors, we first use a view with a common dtype (uint8)
+        #    and then view as numpy array.
+        # 2) for scalar tensors, we cannot use a uint8 view when the size differs,
+        #    so we save the original item and type information.
+        if tensor.dim() > 0:
+            return (tensor.view(torch.uint8).numpy(), tensor.dtype, tensor_device_type)
+        else:
+            return (tensor.item(), tensor.dtype, tensor_device_type)
+
+    def deserialize_tensor(
+        self,
+        val: Union[Tuple["np.ndarray", "torch.dtype", str], int],
+        target_device: Device,
+    ):
+
+        # Found a placeholder for a tensor that was serialized via accelerator.
+        # Replace it with the corresponding deserialized tensor.
+        if isinstance(val, int):
+            placeholder = val
+            self._deserialized_tensor_placeholders.add(placeholder)
+            assert placeholder < len(self._out_of_band_tensors), (
+                "placeholder",
+                placeholder,
+                "out_of_band_tensors",
+                self._out_of_band_tensors,
+            )
+            tensor = self._out_of_band_tensors[placeholder]
+            if target_device == Device.CPU:
+                tensor = tensor.to("cpu")
+            return tensor
+
+        np_array, dtype, tensor_device_type = val
+        return self.deserialize_from_numpy_or_scalar(
+            np_array, dtype, tensor_device_type, target_device
+        )
+
+    def deserialize_from_numpy_or_scalar(
+        self,
+        np_array: Union["np.ndarray", Any],
+        dtype: "torch.dtype",
+        tensor_device_type: str,
+        target_device: Device,
+    ):
+        import numpy as np
+        import torch
+
+        if target_device == Device.DEFAULT:
+            target_device_type = tensor_device_type
+        elif target_device in [Device.GPU, Device.CUDA]:
+            target_device_type = "cuda"
+        else:
+            target_device_type = target_device.value
+
+        # TODO(swang): Support local P2P transfers if available.
+        if target_device_type != "cpu":
+
+            def convert_numpy_to_tensor(np_array):
+                if not isinstance(np_array, np.ndarray):
+                    # For scalar tensors, create the 0-dim tensor.
+                    return torch.tensor(
+                        np_array, device=target_device_type, dtype=dtype
+                    )
+                else:
+                    # For non-scalar tensors, view as the original dtype.
+                    # It does zero-copy convert np_array inside shared memory to
+                    # a tensor. Since we move data to GPU immediately, it is safe.
+                    cpu_tensor = torch.from_numpy(np_array).view(dtype)
+                    return cpu_tensor.to(device=target_device_type)
+
+            global _TORCH_WARNING_FILTER_ACTIVATE
+            # filtering warning messages would be the bottleneck for
+            # deserializing torch tensors. Since the warning only prompts once,
+            # we would only deal with it for the first time.
+            if _TORCH_WARNING_FILTER_ACTIVATE:
+                with warnings.catch_warnings():
+                    # Since np_array.is_writable is False (it is set by Ray),
+                    # this raises a warning. Suppress it.
+                    warnings.filterwarnings(
+                        "ignore",
+                        category=UserWarning,
+                        message="The given NumPy array is not writable",
+                    )
+                    gpu_tensor = convert_numpy_to_tensor(np_array)
+                _TORCH_WARNING_FILTER_ACTIVATE = False
+            else:
+                gpu_tensor = convert_numpy_to_tensor(np_array)
+
+            return gpu_tensor
+
+        # TODO(swang): Use zero-copy from_numpy() if np_array.flags.writeable
+        # is True. This is safe to set when deserializing np_array if the
+        # upstream task has num_readers=1.
+        if not isinstance(np_array, np.ndarray):
+            # For scalar tensors, create the 0-dim tensor.
+            return torch.tensor(np_array, device=target_device_type, dtype=dtype)
+        else:
+            # For non-scalar tensors, view as the original dtype.
+            return torch.tensor(np_array, device=target_device_type).view(dtype)

python/ray/experimental/channel/shared_memory_channel.py

@@ -0,0 +1,819 @@
+import io
+import logging
+import time
+from collections import defaultdict, namedtuple
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
+
+import ray
+import ray.exceptions
+from ray._raylet import SerializedObject
+from ray.experimental.channel import utils
+from ray.experimental.channel.common import ChannelInterface, ChannelOutputType
+from ray.experimental.channel.intra_process_channel import IntraProcessChannel
+from ray.experimental.channel.utils import get_self_actor
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+def _create_channel_ref(
+    self: Any,
+    buffer_size_bytes: int,
+) -> "ray.ObjectRef":
+    """Create a channel that can be read and written through Ray's shared-memory
+    object store.
+
+    The channel has no buffer, so the writer will block until reader(s) have
+    read the previous value.
+
+    A writer and colocated readers can communicate via a shared memory buffer.
+    If the readers are remote, then RPC is used to synchronize the writer and
+    readers' buffers.
+
+    Args:
+        self: The actor on which to allocate the channel buffer. Passed via
+            ``__ray_call__`` so this function executes on the actor.
+        buffer_size_bytes: The initial buffer size in bytes for messages
+            that can be passed between tasks in the DAG. The buffers will
+            be automatically resized if larger messages are written to the
+            channel.
+
+    Returns:
+        A wrapper around ``ray.ObjectRef`` backing the channel.
+    """
+    worker = ray._private.worker.global_worker
+    worker.check_connected()
+
+    value = b"0" * buffer_size_bytes
+
+    try:
+        object_ref = worker.put_object(value, _is_experimental_channel=True)
+    except ray.exceptions.ObjectStoreFullError:
+        logger.info(
+            "Put failed since the value was either too large or the "
+            "store was full of pinned objects."
+        )
+        raise
+    return object_ref
+
+
+# Compiled Graph maintains 1 reader object reference (also called buffer) per node.
+# reader_ref: The object reference.
+# ref_owner_actor_id: The actor who created the object reference.
+# num_readers: The number of reader actors who reads this object reference.
+ReaderRefInfo = namedtuple(
+    "ReaderRefInfo", ["reader_ref", "ref_owner_actor_id", "num_reader_actors"]
+)
+
+
+class _ResizeChannel:
+    """Sentinel used to resize a channel's backing store on the readers.
+
+    When a channel must be resized, the channel backing store must be resized on both
+    the writer and the reader nodes. The writer first resizes its own backing store. The
+    writer then uses an instance of this class as a sentinel value to tell the reader to
+    resize its own backing store. The class instance is sent through the channel.
+    """
+
+    def __init__(
+        self,
+        _node_id_to_reader_ref_info: Dict[str, ReaderRefInfo],
+    ):
+        """Initialize the resize sentinel.
+
+        Args:
+            _node_id_to_reader_ref_info: Mapping from node id to ``ReaderRefInfo``
+                describing the new reader buffers per node.
+        """
+        self._node_id_to_reader_ref_info = _node_id_to_reader_ref_info
+
+
+class SharedMemoryType(ChannelOutputType):
+    def __init__(
+        self,
+        *,
+        buffer_size_bytes: Optional[int] = None,
+        num_shm_buffers: Optional[int] = None,
+    ):
+        """Initialize a ``SharedMemoryType``.
+
+        Args:
+            buffer_size_bytes: The initial buffer size in bytes for messages
+                that can be passed between tasks in the DAG. The buffers will
+                be automatically resized if larger messages are written to the
+                channel.
+            num_shm_buffers: The number of shared memory buffers per channel.
+                Note: In the case of multiple nodes, we only support 1 shared
+                memory buffer.
+        """
+        super().__init__()
+
+        from ray.dag import DAGContext
+
+        ctx = DAGContext.get_current()
+
+        if buffer_size_bytes is None:
+            buffer_size_bytes = ctx.buffer_size_bytes
+        self.buffer_size_bytes = buffer_size_bytes
+        if num_shm_buffers is None:
+            num_shm_buffers = 1
+        self._num_shm_buffers = num_shm_buffers
+
+    def create_channel(
+        self,
+        writer: Optional["ray.actor.ActorHandle"],
+        reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
+        driver_actor_id: Optional[str] = None,
+    ) -> "Channel":
+        """
+        Instantiate a ChannelInterface class that can be used
+        to pass data of this type.
+
+        Args:
+            writer: The actor that may write to the channel. None signifies the driver.
+            reader_and_node_list: A list of tuples, where each tuple contains a reader
+                actor handle and the node ID where the actor is located.
+            driver_actor_id: If this channel is read by a driver and that driver is an
+                actual actor, this will be the actor ID of that driver actor.
+
+        Returns:
+            A ChannelInterface that can be used to pass data
+                of this type.
+        """
+        return CompositeChannel(
+            writer,
+            reader_and_node_list,
+            self._num_shm_buffers,
+            driver_actor_id,
+        )
+
+
+@PublicAPI(stability="alpha")
+class Channel(ChannelInterface):
+    """
+    A wrapper type for ray.ObjectRef. Currently supports ray.get but not
+    ray.wait.
+    """
+
+    def __init__(
+        self,
+        writer: Optional[ray.actor.ActorHandle],
+        reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
+        typ: Optional[Union[int, SharedMemoryType]] = None,
+        _writer_node_id: Optional["ray.NodeID"] = None,
+        _writer_ref: Optional["ray.ObjectRef"] = None,
+        _node_id_to_reader_ref_info: Optional[Dict[str, ReaderRefInfo]] = None,
+        _writer_registered: bool = False,
+        _reader_registered: bool = False,
+    ):
+        """Create a channel that can be read and written by co-located Ray processes.
+
+        Anyone may write to or read from the channel. The channel has no
+        buffer, so the writer will block until reader(s) have read the previous
+        value.
+
+        Args:
+            writer: The actor that may write to the channel. None signifies the driver.
+            reader_and_node_list: A list of tuples, where each tuple contains a reader
+                actor handle and the node ID where the actor is located.
+            typ: Type information about the values passed through the channel.
+                Either an integer representing the max buffer size in bytes
+                allowed, or a SharedMemoryType.
+            _writer_node_id: Internal. Node ID hosting the writer. Provided
+                when rehydrating a channel that was constructed on another
+                process.
+            _writer_ref: Internal. Pre-existing writer-side ``ObjectRef``.
+                When set, the constructor skips allocating a new writer buffer.
+            _node_id_to_reader_ref_info: Internal. Mapping from node id to
+                ``ReaderRefInfo`` describing existing reader buffers per node.
+            _writer_registered: Internal. Whether the writer side has already
+                been registered with the core worker.
+            _reader_registered: Internal. Whether the reader side has already
+                been registered with the core worker.
+        """
+        assert len(reader_and_node_list) > 0
+        for reader, _ in reader_and_node_list:
+            assert isinstance(reader, ray.actor.ActorHandle)
+
+        if typ is None:
+            typ = SharedMemoryType()
+        elif isinstance(typ, int):
+            typ = SharedMemoryType(buffer_size_bytes=typ)
+
+        # The min buffer size must be large enough to at least fit an instance of the
+        # _ResizeChannel class along with any metadata.
+        MIN_BUFFER_SIZE = int(1000)  # 1000 bytes
+        if typ.buffer_size_bytes < MIN_BUFFER_SIZE:
+            raise ValueError(
+                "typ.buffer_size_bytes must be at least MIN_BUFFER_SIZE "
+                f"({MIN_BUFFER_SIZE} bytes)"
+            )
+
+        self._writer = writer
+        self._reader_and_node_list = reader_and_node_list
+        self._typ = typ
+
+        self._worker = ray._private.worker.global_worker
+        self._worker.check_connected()
+
+        self._writer_registered = _writer_registered
+        self._reader_registered = _reader_registered
+        # NodeID -> ReaderRefInfo on that node. Note that there's only 1
+        # reader ref per node.
+        self._node_id_to_reader_ref_info: Dict[str, ReaderRefInfo] = (
+            _node_id_to_reader_ref_info or {}
+        )
+
+        # Node ID -> a list of reader actors.
+        self._node_id_to_readers: Dict[str, "ray.actor.ActorHandle"] = defaultdict(list)
+        for reader, node_id in self._reader_and_node_list:
+            self._node_id_to_readers[node_id].append(reader)
+
+        # Number of readers in a local node.
+        self._num_local_readers = 0
+
+        if _writer_ref is None:
+            # We are the writer. Check that the passed handle matches the
+            # current actor (or it is the driver).
+            # TODO(swang): Channels must be initially constructed by the writer
+            # actor, so we shouldn't need to include `writer` in the
+            # constructor args. Either support Channels being constructed by
+            # someone other than the writer or remove it from the args.
+            self_actor = get_self_actor()
+            assert writer == self_actor
+
+            self._writer_node_id = (
+                ray.runtime_context.get_runtime_context().get_node_id()
+            )
+            self._writer_ref = _create_channel_ref(self, typ.buffer_size_bytes)
+
+            self._create_reader_refs(typ.buffer_size_bytes)
+        else:
+            assert (
+                _writer_node_id is not None
+            ), "_writer_node_id must also be passed to the constructor when "
+            "_writer_ref is."
+            assert _node_id_to_reader_ref_info is not None, (
+                "_node_id_to_reader_ref_info must also be passed to the constructor "
+                "when _writer_ref is."
+            )
+
+            self._writer_ref = _writer_ref
+            self._writer_node_id = _writer_node_id
+            self._node_id_to_reader_ref_info = _node_id_to_reader_ref_info
+
+        assert self._num_local_readers == 0
+        remote_node_exists = False
+        for node_id, readers in self._node_id_to_readers.items():
+            if self.is_local_node(node_id):
+                self._num_local_readers += len(readers)
+            else:
+                remote_node_exists = True
+        # If remote node exists, we have 1 additional reader that listens
+        # to object changes and push them to remote nodes.
+        if remote_node_exists:
+            self._num_local_readers += 1
+        # There must be at least 1 local reader
+        assert self._num_local_readers > 0
+
+        self._local_reader_ref: Optional["ray.ObjectRef"] = self._get_local_reader_ref(
+            self._node_id_to_reader_ref_info
+        )
+
+    def _get_local_reader_ref(
+        self, _node_id_to_reader_ref_info: Dict[str, ReaderRefInfo]
+    ) -> Optional["ray.ObjectRef"]:
+        for node_id, reader_ref_info in _node_id_to_reader_ref_info.items():
+            if self.is_local_node(node_id):
+                return reader_ref_info.reader_ref
+        return None
+
+    def _create_reader_refs(
+        self,
+        buffer_size_bytes: int,
+    ):
+        # TODO(jhumphri): Free the current reader ref once the reference to it is
+        # destroyed below.
+
+        for node_id, readers in self._node_id_to_readers.items():
+            if not self.is_local_node(node_id):
+                # Find 1 reader in a remote node to create a reference that's
+                # shared by all readers. When a new value is written to a reference,
+                # it is sent to this reference.
+                reader = readers[0]
+                fn = reader.__ray_call__
+                self._node_id_to_reader_ref_info[node_id] = ReaderRefInfo(
+                    reader_ref=ray.get(
+                        fn.remote(_create_channel_ref, buffer_size_bytes)
+                    ),
+                    ref_owner_actor_id=reader._actor_id,
+                    num_reader_actors=len(readers),
+                )
+            else:
+                writer_id = ray.ActorID.nil()
+                if self._writer is not None:
+                    writer_id = self._writer._actor_id
+                self._node_id_to_reader_ref_info[node_id] = ReaderRefInfo(
+                    reader_ref=self._writer_ref,
+                    ref_owner_actor_id=writer_id,
+                    num_reader_actors=len(readers),
+                )
+        # There must be only 1 node reader reference per node.
+        assert len(self._node_id_to_reader_ref_info) == len(self._node_id_to_readers)
+
+        # We need to register the new writer_ref.
+        self._writer_registered = False
+        self.ensure_registered_as_writer()
+
+    @staticmethod
+    def is_local_node(node_id):
+        return ray.runtime_context.get_runtime_context().get_node_id() == node_id
+
+    def ensure_registered_as_writer(self) -> None:
+        if self._writer_registered:
+            return
+
+        if not self.is_local_node(self._writer_node_id):
+            raise ValueError(
+                "`ensure_registered_as_writer()` must only be called on the node that "
+                "the writer is on."
+            )
+
+        remote_reader_ref_info: Dict[str, ReaderRefInfo] = {}
+        for node_id, reader_ref_info in self._node_id_to_reader_ref_info.items():
+            if self.is_local_node(node_id):
+                continue
+            remote_reader_ref_info[node_id] = reader_ref_info
+
+        self._worker.core_worker.experimental_channel_register_writer(
+            self._writer_ref,
+            remote_reader_ref_info,
+        )
+        self._writer_registered = True
+
+    def ensure_registered_as_reader(self) -> None:
+        if self._reader_registered:
+            return
+
+        for node_id, reader_ref_info in self._node_id_to_reader_ref_info.items():
+            if self.is_local_node(node_id):
+                self._worker.core_worker.experimental_channel_register_reader(
+                    reader_ref_info.reader_ref,
+                )
+        self._reader_registered = True
+
+    @staticmethod
+    def _deserialize_reader_channel(
+        writer: ray.actor.ActorHandle,
+        reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
+        typ: int,
+        writer_node_id,
+        writer_ref: "ray.ObjectRef",
+        node_id_to_reader_ref_info: Dict[str, ReaderRefInfo],
+        writer_registered: bool,
+        reader_registered: bool,
+    ) -> "Channel":
+        chan = Channel(
+            writer,
+            reader_and_node_list,
+            typ,
+            _writer_node_id=writer_node_id,
+            _writer_ref=writer_ref,
+            _node_id_to_reader_ref_info=node_id_to_reader_ref_info,
+            _writer_registered=writer_registered,
+            _reader_registered=reader_registered,
+        )
+        return chan
+
+    def __reduce__(self):
+        assert self._node_id_to_reader_ref_info is not None
+        return self._deserialize_reader_channel, (
+            self._writer,
+            self._reader_and_node_list,
+            self._typ,
+            self._writer_node_id,
+            self._writer_ref,
+            self._node_id_to_reader_ref_info,
+            self._writer_registered,
+            self._reader_registered,
+        )
+
+    def __str__(self) -> str:
+        return (
+            f"Channel(_node_id_to_reader_ref_info={self._node_id_to_reader_ref_info}, "
+            f"_writer_ref={self._writer_ref})"
+        )
+
+    def _resize_channel_if_needed(self, serialized_value: str, timeout_ms: int):
+        # serialized_value.total_bytes *only* includes the size of the data. It does not
+        # include the size of the metadata, so we must account for the size of the
+        # metadata explicitly.
+        size = serialized_value.total_bytes + len(serialized_value.metadata)
+        if size > self._typ.buffer_size_bytes:
+            # Now make the channel backing store larger.
+            self._typ.buffer_size_bytes = size
+            # TODO(jhumphri): Free the current writer ref once the reference to it is
+            # destroyed below.
+            # TODO(sang): Support different policies such as 2X buffer size.
+            prev_writer_ref = self._writer_ref
+            self._writer_ref = _create_channel_ref(self, self._typ.buffer_size_bytes)
+            self._create_reader_refs(self._typ.buffer_size_bytes)
+            self._local_reader_ref = self._get_local_reader_ref(
+                self._node_id_to_reader_ref_info
+            )
+
+            # Write a special message to the channel so that the readers know to
+            # stop using the current reader_ref.
+            special_message = _ResizeChannel(self._node_id_to_reader_ref_info)
+            special_message_serialized = (
+                self._worker.get_serialization_context().serialize(special_message)
+            )
+            self._worker.core_worker.experimental_channel_put_serialized(
+                special_message_serialized,
+                prev_writer_ref,
+                self._num_local_readers,
+                timeout_ms,
+            )
+            # TODO(sang): Clean the previous ref that won't be used.
+            # Right now, if we just close it here, it will not work because
+            # of race conditions.
+            # self._worker.core_worker.experimental_channel_set_error(
+            #     prev_writer_ref
+            # )
+
+    def write(self, value: Any, timeout: Optional[float] = None) -> None:
+        self.ensure_registered_as_writer()
+        assert (
+            timeout is None or timeout >= 0 or timeout == -1
+        ), "Timeout must be non-negative or -1."
+        # -1 means no timeout (block indefinitely)
+        timeout_ms = int(timeout * 1000) if timeout is not None else -1
+
+        if not isinstance(value, SerializedObject):
+            try:
+                serialized_value = self._worker.get_serialization_context().serialize(
+                    value
+                )
+            except TypeError as e:
+                sio = io.StringIO()
+                ray.util.inspect_serializability(value, print_file=sio)
+                msg = (
+                    "Could not serialize the put value "
+                    f"{repr(value)}:\n"
+                    f"{sio.getvalue()}"
+                )
+                raise TypeError(msg) from e
+        else:
+            serialized_value = value
+
+        start_time = time.monotonic()
+        self._resize_channel_if_needed(serialized_value, timeout_ms)
+        if timeout is not None:
+            timeout_ms -= int((time.monotonic() - start_time) * 1000)
+            timeout_ms = max(timeout_ms, 0)
+
+        self._worker.core_worker.experimental_channel_put_serialized(
+            serialized_value,
+            self._writer_ref,
+            self._num_local_readers,
+            timeout_ms,
+        )
+
+    def read(self, timeout: Optional[float] = None) -> Any:
+        assert (
+            timeout is None or timeout >= 0 or timeout == -1
+        ), "Timeout must be non-negative or -1."
+        self.ensure_registered_as_reader()
+
+        start_time = time.monotonic()
+        ret = self._worker.get_objects(
+            [self._local_reader_ref], timeout=timeout, return_exceptions=True
+        )[0][0]
+
+        if isinstance(ret, _ResizeChannel):
+            self._node_id_to_reader_ref_info = ret._node_id_to_reader_ref_info
+            self._local_reader_ref = self._get_local_reader_ref(
+                self._node_id_to_reader_ref_info
+            )
+            # We need to register the new reader_ref.
+            self._reader_registered = False
+            self.ensure_registered_as_reader()
+            if timeout is not None:
+                timeout -= time.monotonic() - start_time
+                timeout = max(timeout, 0)
+            ret = self._worker.get_objects(
+                [self._local_reader_ref], timeout=timeout, return_exceptions=True
+            )[0][0]
+
+        return ret
+
+    def release_buffer(self, timeout: Optional[float] = None) -> None:
+        assert (
+            timeout is None or timeout >= 0 or timeout == -1
+        ), "Timeout must be non-negative or -1."
+        self.ensure_registered_as_reader()
+        self._worker.get_objects(
+            [self._local_reader_ref],
+            timeout=timeout,
+            return_exceptions=True,
+            skip_deserialization=True,
+        )
+
+    def close(self) -> None:
+        """
+        Close this channel by setting the error bit on both the writer_ref and the
+        reader_ref.
+        """
+        self._worker.core_worker.experimental_channel_set_error(self._writer_ref)
+        is_local_node_reader = False
+
+        for node_id in self._node_id_to_readers.keys():
+            if self.is_local_node(node_id):
+                is_local_node_reader = True
+        if is_local_node_reader:
+            self.ensure_registered_as_reader()
+
+        for reader_ref_info in self._node_id_to_reader_ref_info.values():
+            self._worker.core_worker.experimental_channel_set_error(
+                reader_ref_info.reader_ref
+            )
+
+
+@DeveloperAPI
+class BufferedSharedMemoryChannel(ChannelInterface):
+    """A channel that can be read and written by Ray processes.
+
+    It creates `num_shm_buffers` number of buffers and allows buffered read and
+    write APIs. I.e., read and write APIs are non-blocking as long as it can write to
+    next buffer or read from a next buffer. See `read` and `write` APIs for
+    more details.
+
+    Args:
+        writer: The actor that may write to the channel. None signifies the driver.
+        reader_and_node_list: A list of tuples, where each tuple contains a reader
+            actor handle and the node ID where the actor is located. Note that currently
+            we only support this for readers on the same node as the writer.
+        num_shm_buffers: Number of shared memory buffers to read/write.
+        typ: Type information about the values passed through the channel.
+            Either an integer representing the max buffer size in bytes
+            allowed, or a SharedMemoryType.
+    """
+
+    def __init__(
+        self,
+        writer: Optional[ray.actor.ActorHandle],
+        reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
+        num_shm_buffers: int,
+        typ: Optional[Union[int, SharedMemoryType]] = None,
+    ):
+        self._num_shm_buffers = num_shm_buffers
+        self._buffers = [
+            # We use Channel directly as a buffer implementation as
+            # channel only allows to have 1 shared memory buffer.
+            Channel(writer, reader_and_node_list, typ)
+            for _ in range(num_shm_buffers)
+        ]
+        # The next index to write from self._buffers.
+        self._next_write_index = 0
+        # The next index to read from self._buffers.
+        self._next_read_index = 0
+
+    def ensure_registered_as_writer(self):
+        """
+        Check whether the process is a valid writer. This method must be idempotent.
+        """
+        for buffer in self._buffers:
+            buffer.ensure_registered_as_writer()
+
+    def ensure_registered_as_reader(self):
+        """
+        Check whether the process is a valid reader. This method must be idempotent.
+        """
+        for buffer in self._buffers:
+            buffer.ensure_registered_as_reader()
+
+    def write(self, value: Any, timeout: Optional[float] = None) -> None:
+        """Write a value to a channel.
+
+        If the next buffer is available, it returns immediately. If the next
+        buffer is not read by downstream consumers, it blocks until a buffer is
+        available to write. If a buffer is not available within timeout, it raises
+        RayChannelTimeoutError.
+        """
+        self.ensure_registered_as_writer()
+        # A single channel is not supposed to read and write at the same time.
+        assert self._next_read_index == 0
+        self._buffers[self._next_write_index].write(value, timeout)
+        self._next_write_index += 1
+        self._next_write_index %= self._num_shm_buffers
+
+    def read(self, timeout: Optional[float] = None) -> Any:
+        """Read a value from a channel.
+
+        If the next buffer is available, it returns immediately. If the next
+        buffer is not written by an upstream producer, it blocks until a buffer is
+        available to read. If a buffer is not available within timeout, it raises
+        RayChannelTimeoutError.
+        """
+        self.ensure_registered_as_reader()
+        # A single channel is not supposed to read and write at the same time.
+        assert self._next_write_index == 0
+        output = self._buffers[self._next_read_index].read(timeout)
+        self._next_read_index += 1
+        self._next_read_index %= self._num_shm_buffers
+        return output
+
+    def release_buffer(self, timeout: Optional[float] = None):
+        """Release the native buffer of the channel to allow the buffer to be reused for
+        future data.
+
+        If the next buffer is available, it returns immediately. If the next
+        buffer is not written by an upstream producer, it blocks until a buffer is
+        available to be released. If a buffer is not available within timeout, it raises
+        RayChannelTimeoutError.
+        """
+        # A single channel is not supposed to read and write at the same time.
+        assert self._next_write_index == 0
+        self._buffers[self._next_read_index].release_buffer(timeout)
+        self._next_read_index += 1
+        self._next_read_index %= self._num_shm_buffers
+
+    def close(self) -> None:
+        for buffer in self._buffers:
+            buffer.close()
+
+    @property
+    def next_write_index(self):
+        # Testing only
+        return self._next_write_index
+
+    @property
+    def next_read_index(self):
+        # Testing only
+        return self._next_read_index
+
+
+@PublicAPI(stability="alpha")
+class CompositeChannel(ChannelInterface):
+    """Routes data to different readers via per-locality channels.
+
+    For example, if the reader is in the same worker process as the writer,
+    the data can be sent via IntraProcessChannel. If the reader is in a different
+    worker process, the data can be sent via shared memory channel.
+    """
+
+    def __init__(
+        self,
+        writer: Optional[ray.actor.ActorHandle],
+        reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
+        num_shm_buffers: int,
+        driver_actor_id: Optional[str] = None,
+        _channel_dict: Optional[Dict[ray.ActorID, ChannelInterface]] = None,
+        _channels: Optional[Set[ChannelInterface]] = None,
+        _writer_registered: bool = False,
+        _reader_registered: bool = False,
+    ):
+        """Initialize a ``CompositeChannel``.
+
+        Args:
+            writer: The actor that may write to the channel. None signifies the driver.
+            reader_and_node_list: A list of tuples, where each tuple contains a reader
+                actor handle and the node ID where the actor is located.
+            num_shm_buffers: The number of shared memory buffers per channel.
+                Note: In the case of multiple nodes, we only support 1 shared
+                memory buffer.
+            driver_actor_id: If this channel is read by a driver and that driver is an
+                actual actor, this will be the actor ID of that driver actor.
+            _channel_dict: Internal. Pre-populated mapping from actor id to
+                the underlying channel. When provided, channels are not
+                re-created (used during deserialization).
+            _channels: Internal. Deduplicated set of channels backing
+                ``_channel_dict``. When provided, channels are not re-created.
+            _writer_registered: Internal. Whether the writer side has already
+                been registered with the core worker.
+            _reader_registered: Internal. Whether the reader side has already
+                been registered with the core worker.
+        """
+        self._writer = writer
+        self._reader_and_node_list = reader_and_node_list
+        self._num_shm_buffers = num_shm_buffers
+        self._driver_actor_id = driver_actor_id
+        self._writer_registered = _writer_registered
+        self._reader_registered = _reader_registered
+        # A dictionary that maps the actor ID to the channel object.
+        self._channel_dict = _channel_dict or {}
+        # The set of channels is a deduplicated version of the _channel_dict values.
+        self._channels = _channels or set()
+        if self._channels:
+            # This CompositeChannel object is created by deserialization.
+            # We don't need to create channels again.
+            return
+
+        (
+            remote_reader_and_node_list,
+            local_reader_and_node_list,
+        ) = utils.split_readers_by_locality(self._writer, self._reader_and_node_list)
+        # There are some local readers which are the same worker process as the writer.
+        # Create a local channel for the writer and the local readers.
+        num_local_readers = len(local_reader_and_node_list)
+        if num_local_readers > 0:
+            # Use num_readers = 1 when creating the local channel,
+            # because we have channel cache to support reading
+            # from the same channel multiple times.
+            local_channel = IntraProcessChannel(num_readers=1)
+            self._channels.add(local_channel)
+            actor_id = self._get_actor_id(self._writer)
+            self._channel_dict[actor_id] = local_channel
+        # There are some remote readers which are not the same Ray actor as the writer.
+        # We create a BufferedSharedMemoryChannel for readers on the same node, and
+        # a single Channel for readers on different nodes due to
+        # https://github.com/ray-project/ray/issues/49044
+        (
+            readers_same_node,
+            readers_different_node,
+        ) = utils.split_actors_by_node_locality(
+            utils.get_actor_node(self._writer), remote_reader_and_node_list
+        )
+
+        if len(readers_same_node) != 0:
+            remote_channel = BufferedSharedMemoryChannel(
+                self._writer, readers_same_node, num_shm_buffers
+            )
+            self._channels.add(remote_channel)
+            for reader, _ in readers_same_node:
+                actor_id = self._get_actor_id(reader)
+                self._channel_dict[actor_id] = remote_channel
+
+        if len(readers_different_node) != 0:
+            remote_channel = Channel(self._writer, readers_different_node)
+            self._channels.add(remote_channel)
+            for reader, _ in readers_different_node:
+                actor_id = self._get_actor_id(reader)
+                self._channel_dict[actor_id] = remote_channel
+
+    def _get_actor_id(self, reader: ray.actor.ActorHandle) -> str:
+        return reader._actor_id.hex()
+
+    def ensure_registered_as_writer(self) -> None:
+        if self._writer_registered:
+            return
+        for channel in self._channels:
+            channel.ensure_registered_as_writer()
+        self._writer_registered = True
+
+    def ensure_registered_as_reader(self) -> None:
+        if self._reader_registered:
+            return
+        for channel in self._channels:
+            channel.ensure_registered_as_reader()
+        self._reader_registered = True
+
+    def __reduce__(self):
+        return CompositeChannel, (
+            self._writer,
+            self._reader_and_node_list,
+            self._num_shm_buffers,
+            self._driver_actor_id,
+            self._channel_dict,
+            self._channels,
+            self._writer_registered,
+            self._reader_registered,
+        )
+
+    def __str__(self) -> str:
+        return (
+            "CompositeChannel(_channels="
+            f"{[str(channel) for channel in self._channels]})"
+        )
+
+    def write(self, value: Any, timeout: Optional[float] = None) -> None:
+        self.ensure_registered_as_writer()
+        for channel in self._channels:
+            channel.write(value, timeout)
+
+    def read(self, timeout: Optional[float] = None) -> Any:
+        self.ensure_registered_as_reader()
+        return self._channel_dict[self._resolve_actor_id()].read(timeout)
+
+    def release_buffer(self, timeout: Optional[float] = None):
+        self.ensure_registered_as_reader()
+        self._channel_dict[self._resolve_actor_id()].release_buffer(timeout)
+
+    def _resolve_actor_id(self) -> str:
+        actor_id = ray.get_runtime_context().get_actor_id()
+        # If actor_id is None, read was called by the driver
+        # If the driver is an actor, driver_actor_id will be set to that actor id
+        if actor_id is None or actor_id == self._driver_actor_id:
+            # Use the actor ID of the DAGDriverProxyActor.
+            # The proxy actor is always the first actor in the reader_and_node_list.
+            assert len(self._reader_and_node_list) >= 1
+            driver_proxy_actor = self._reader_and_node_list[0][0]
+            actor_id = self._get_actor_id(driver_proxy_actor)
+        return actor_id
+
+    def close(self) -> None:
+        for channel in self._channels:
+            channel.close()

python/ray/experimental/channel/torch_tensor_accelerator_channel.py

@@ -0,0 +1,882 @@
+import io
+import logging
+import uuid
+from dataclasses import dataclass
+from types import ModuleType
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Type, Union
+
+import ray
+import ray.util.serialization
+from ray.experimental.channel import ChannelContext, utils
+from ray.experimental.channel.accelerator_context import (
+    AcceleratorContext,
+    is_accelerator_context_registered,
+    register_accelerator_context,
+)
+from ray.experimental.channel.common import ChannelInterface
+from ray.experimental.channel.communicator import Communicator
+from ray.experimental.channel.communicator_handle import CommunicatorHandle
+from ray.experimental.channel.cpu_communicator import CPUCommunicator
+from ray.experimental.channel.intra_process_channel import IntraProcessChannel
+from ray.experimental.channel.shared_memory_channel import SharedMemoryType
+from ray.experimental.channel.torch_tensor_type import TorchTensorType
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    import torch
+
+    from ray.experimental.channel.shared_memory_channel import Channel
+
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class _TorchTensorMetadata:
+    """
+    Metadata for torch.Tensors that can be sent between processes to determine
+    how large of a buffer to allocate on the receiver(s).
+    """
+
+    shape: Union[int, Tuple[int]]
+    dtype: "torch.dtype"
+
+
+@DeveloperAPI
+class TorchTensorAcceleratorChannel(ChannelInterface):
+    def __init__(
+        self,
+        writer: ray.actor.ActorHandle,
+        reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
+        typ: "TorchTensorType",
+        driver_actor_id: str,
+        tensor_metadata_channel: Optional["Channel"] = None,
+        _cpu_data_channel: Optional["Channel"] = None,
+        _gpu_data_channel: Optional["_TorchTensorAcceleratorChannel"] = None,
+        _local_channel: Optional["IntraProcessChannel"] = None,
+    ):
+        """
+        Can be used to send accelerator tensors nested inside other data. The data is
+        sent via shared memory while the accelerator tensors are sent through a P2P
+        transport (e.g., NCCL for GPU).
+
+        NOTE: This class is currently not thread-safe because it reads and
+        writes the worker-local
+        ray.experimental.channel.serialization_context._SerializationContext
+        when serializing data.
+
+        Args:
+            writer: The actor that may write to the channel. None signifies the
+                driver.
+            reader_and_node_list: A list of tuples, where each tuple contains a reader
+                actor handle and the node ID where the actor is located.
+            typ: Type information about the values passed through the channel.
+            driver_actor_id: The actor ID of the DAGDriverProxyActor.
+            tensor_metadata_channel: A shared-memory channel for sending tensor
+                metadata.
+            _cpu_data_channel: A shared-memory channel for sending
+                non-tensor data. Its writer and readers should match the given
+                writer and readers. If None is provided, then we assume that
+                there is no CPU-specific data, i.e. the task directly returned
+                a CUDA torch.Tensor.
+            _gpu_data_channel: A channel for sending torch.Tensors via accelerator.
+            _local_channel: A channel for sending data between the writer and
+                local readers.
+
+        NOTE: `tensor_metadata_channel` will be set only for testing purposes.
+        `_cpu_data_channel` is set for testing purposes and for deserialization.
+        `_gpu_data_channel` and `_local_channel` are set only during deserialization.
+        """
+        self._writer = writer
+        self._reader_and_node_list = reader_and_node_list
+        self._typ = typ
+
+        (
+            remote_reader_and_node_list,
+            local_reader_and_node_list,
+        ) = utils.split_readers_by_locality(self._writer, self._reader_and_node_list)
+
+        num_local_readers = len(local_reader_and_node_list)
+        self._local_channel = _local_channel
+        if self._local_channel is None and num_local_readers > 0:
+            # There are some local readers which are the same worker process as
+            # the writer. Create a local channel for the writer and the local readers.
+            #
+            # Use num_readers = 1 when creating the local channel,
+            # because we have channel cache to support reading
+            # from the same channel multiple times.
+            self._local_channel = IntraProcessChannel(num_readers=1)
+
+        assert len(remote_reader_and_node_list) > 0, (
+            "All readers are from the same actor. "
+            "The TorchTensorType type hint is not needed. "
+            "No accelerator channel will be created."
+        )
+        self._gpu_data_channel = _gpu_data_channel
+        if self._gpu_data_channel is None:
+            self._gpu_data_channel: _TorchTensorAcceleratorChannel = (
+                _TorchTensorAcceleratorChannel(
+                    writer,
+                    remote_reader_and_node_list,
+                    typ,
+                    _meta_channel=tensor_metadata_channel,
+                )
+            )
+
+        self._cpu_data_channel: Optional["Channel"] = _cpu_data_channel
+        if self._cpu_data_channel is not None:
+            assert (
+                not self._typ.direct_return
+            ), "CPU channel should be None if direct return is enabled"
+
+        if self._cpu_data_channel is None and not self._typ.direct_return:
+            # Create a CPU channel to send non-tensor data.
+            self._cpu_data_channel = SharedMemoryType().create_channel(
+                writer, remote_reader_and_node_list, driver_actor_id
+            )
+
+        # Used for serialization.
+        self._worker = ray._private.worker.global_worker
+        self._worker.check_connected()
+
+        ctx = ChannelContext.get_current()
+        self.serialization_ctx = ctx.serialization_context
+        assert self.serialization_ctx is not None
+
+    def __reduce__(self):
+        return (
+            TorchTensorAcceleratorChannel,
+            (
+                self._writer,
+                self._reader_and_node_list,
+                self._typ,
+                # driver_actor_id and tensor_metadata_channel are used to initialize
+                # the _cpu_data_channel and _gpu_data_channel, so we don't need to
+                # pass them in here.
+                None,
+                None,
+                self._cpu_data_channel,
+                self._gpu_data_channel,
+                self._local_channel,
+            ),
+        )
+
+    def ensure_registered_as_writer(self):
+        if self._local_channel is not None:
+            self._local_channel.ensure_registered_as_writer()
+        self._gpu_data_channel.ensure_registered_as_writer()
+        if self._cpu_data_channel is not None:
+            self._cpu_data_channel.ensure_registered_as_writer()
+
+    def ensure_registered_as_reader(self):
+        reader = utils.get_self_actor()
+        if reader == self._writer:
+            self._local_channel.ensure_registered_as_reader()
+            return
+        self._gpu_data_channel.ensure_registered_as_reader()
+        if self._cpu_data_channel is not None:
+            self._cpu_data_channel.ensure_registered_as_reader()
+
+    def _send_cpu_and_gpu_data(self, value: Any, timeout: Optional[float]):
+        self.serialization_ctx.reset_out_of_band_tensors([])
+        # All tensors found in `value` will be transferred via accelerator.
+        self.serialization_ctx.set_use_external_transport(True)
+
+        try:
+            # Serialize the data. All tensors that match our current device
+            # will be extracted into the serialization context and replaced
+            # with a placeholder.
+            cpu_data = self._worker.get_serialization_context().serialize(value)
+        except TypeError as e:
+            sio = io.StringIO()
+            ray.util.inspect_serializability(value, print_file=sio)
+            msg = (
+                "Could not serialize the put value "
+                f"{repr(value)}:\n"
+                f"{sio.getvalue()}"
+            )
+            raise TypeError(msg) from e
+        finally:
+            # Pop the tensors that were found during serialization of `value`.
+            gpu_tensors, _ = self.serialization_ctx.reset_out_of_band_tensors([])
+            # Reset the serialization method to now serialize torch.Tensors
+            # normally.
+            self.serialization_ctx.set_use_external_transport(False)
+
+        # First send the extracted tensors through a GPU-specific channel.
+        self._gpu_data_channel.write(gpu_tensors)
+        # Next send the non-tensor data through a CPU-specific channel. The
+        # data contains placeholders for the extracted tensors.
+        self._cpu_data_channel.write(cpu_data)
+
+    def write(self, value: Any, timeout: Optional[float] = None) -> None:
+        """
+        Send a value that may contain torch.Tensors that should be sent via
+        external transport.
+
+        Case 1: Use `_local_channel` to send the data to local readers.
+
+        Case 2: Otherwise, use the following method to send the data to remote readers.
+
+        1) Serializes `value`. During serialization, all torch.Tensors that are
+        on the default device are extracted and replaced with a unique
+        placeholder. Thus, the serialized value will contain all non-tensor
+        data, and any tensors that were not on the default device (e.g., CPU
+        tensor returned by a GPU actor).
+        2) Sends extracted torch.Tensors via the tensor data channel (e.g.,
+        NCCL).
+        3) Sends the non-tensor data via the non-tensor data channel.
+
+        If static_non_tensor_data=True was specified, then we only perform step
+        (3) on the first `write` call. The reader is expected to reuse the sent
+        data for subsequent messages.
+        """
+        self.ensure_registered_as_writer()
+
+        if self._local_channel is not None:
+            self._local_channel.write(value)
+
+        if isinstance(value, ray.exceptions.RayTaskError):
+            if self._typ.static_shape or self._typ.direct_return:
+                # Raise a fatal error to teardown the DAG.
+                # This error will also be caught from `CompiledDAGRef.get()`
+                # and raised to the user
+                # TODO(swang): Write exceptions to the tensor metadata or
+                # non-tensor data channel if it is available to make these
+                # exceptions recoverable.
+                raise value
+
+        if self._cpu_data_channel is None:
+            # Handle the case where _direct_return=True. In this case, we check
+            # that the task returned a CUDA torch.Tensor and just send it
+            # directly without trying to serialize it first.
+            import torch
+
+            # These ValueErrors will also be caught from `CompiledDAGRef.get()`
+            # and raised to the user
+            if not isinstance(value, torch.Tensor):
+                # TODO(swang): These errors are currently fatal for the DAG.
+                # This could be improved by sending the exception through the
+                # gpu_data_channel's CPU-based metadata channel, if one exists.
+                raise ValueError(
+                    "Task annotated with _direct_return=True must "
+                    "return a CUDA torch.Tensor, instead found value "
+                    f"`{value}`. DAG will shut down."
+                )
+            elif not value.is_cuda:
+                raise ValueError(
+                    "Task annotated with _direct_return=True must "
+                    "return a CUDA torch.Tensor, instead found CPU tensor. "
+                    "DAG will shut down."
+                )
+            self._gpu_data_channel.write([value], timeout=timeout)
+        else:
+            self._send_cpu_and_gpu_data(value, timeout)
+
+    def _recv_cpu_and_gpu_data(
+        self, tensors: List["torch.Tensor"], timeout: Optional[float] = None
+    ) -> Any:
+        """Helper method to receive data that contains a mix of CPU and GPU data.
+
+        Args:
+            tensors: The GPU data. This is a list of the torch.Tensors that
+                were found in the sent data.
+            timeout: Timeout for channel receive.
+
+        Returns:
+            The deserialized non-tensor data with tensor placeholders replaced
+            by the entries of ``tensors``.
+        """
+        self.serialization_ctx.reset_out_of_band_tensors(tensors)
+
+        # Next, read and deserialize the non-tensor data. The registered custom
+        # deserializer will replace the found tensor placeholders with
+        # `tensors`.
+        data = self._cpu_data_channel.read(
+            timeout=timeout,
+        )
+        # Check that all placeholders had a corresponding tensor.
+        (
+            _,
+            deserialized_tensor_placeholders,
+        ) = self.serialization_ctx.reset_out_of_band_tensors([])
+        assert deserialized_tensor_placeholders == set(range(len(tensors)))
+
+        return data
+
+    def read(self, timeout: Optional[float] = None) -> Any:
+        """
+        Read a value that may contain torch.Tensors sent via external
+        transport.
+
+        Case 1: If the reader is a local reader and is the same actor as the writer,
+        then use the `_local_channel` to read the data.
+
+        Case 2: Otherwise, use the following method to read data from remote readers.
+
+        1) Receives torch.Tensors via the tensor data channel (e.g., NCCL).
+        2) Reads the serialized non-tensor data.
+        3) Deserializes the non-tensor data. During deserialization, replaces
+        all found placeholders with the received torch.Tensors.
+
+        If _direct_return=True was specified, then we skip step (2) and (3) and
+        directly return the data received in (1).
+        """
+        self.ensure_registered_as_reader()
+
+        # If the reader is the same actor as the writer, then we can use the
+        # local channel to read the data.
+        reader = utils.get_self_actor()
+        if reader == self._writer:
+            assert self._local_channel is not None
+            return self._local_channel.read()
+
+        # First, read the tensor data.
+        tensors = self._gpu_data_channel.read(timeout)
+
+        if self._cpu_data_channel is None:
+            # Handle _direct_return=True. In this case, we expect to receive
+            # only one tensor, and we return it directly.
+            assert len(tensors) == 1
+            data = tensors[0]
+        else:
+            data = self._recv_cpu_and_gpu_data(tensors, timeout)
+
+        return data
+
+    def close(self) -> None:
+        self._gpu_data_channel.close()
+        if self._cpu_data_channel is not None:
+            self._cpu_data_channel.close()
+        if self._local_channel is not None:
+            self._local_channel.close()
+
+
+def _torch_tensor_allocator(
+    shape: Union[int, Tuple[int]],
+    dtype: "torch.dtype",
+):
+    """
+    Allocate a tensor buffer matching the given metadata.
+    """
+    import torch
+
+    ctx = ChannelContext.get_current()
+    return torch.empty(shape, dtype=dtype, device=ctx.torch_device)
+
+
+class _TorchTensorAcceleratorChannel(ChannelInterface):
+    def __init__(
+        self,
+        writer: ray.actor.ActorHandle,
+        reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
+        typ: "TorchTensorType",
+        _meta_channel: Optional["Channel"] = None,
+    ):
+        """
+        A helper channel for TorchTensorAcceleratorChannel that is used to transfer
+        lists of torch.Tensors via accelerator. This class can only transfer
+        torch.Tensors and cannot transfer other CPU data, such as Exception
+        objects or tensors nested inside of a dictionary.
+
+        Args:
+            writer: The actor that may write to the channel. None signifies the driver.
+            reader_and_node_list: A list of tuples, where each tuple contains a reader
+                actor handle and the node ID where the actor is located.
+            typ: Type information about the values passed through the channel.
+            _meta_channel: A channel used to send metadata for the tensors,
+                i.e. shape and dtype. If not provided, and if the typ does not
+                specify a static shape and dtype, then a metadata channel based
+                on shared memory will be created.
+        """
+        import torch
+
+        self.torch: ModuleType = torch
+
+        self._writer = writer
+        self._writer_rank: Optional[int] = None
+        self._reader_and_node_list = reader_and_node_list
+        self._reader_ranks: Optional[List[int]] = None
+        self._writer_registered: bool = False
+        self._reader_registered: bool = False
+
+        ctx = ChannelContext.get_current()
+        assert isinstance(
+            typ.communicator_id, str
+        ), f"accelerator group ID ({typ.communicator_id}) must be a str."
+        self._typ = typ
+
+        self._static_shape = typ.static_shape
+
+        assert self._typ.communicator_id is not None, "No accelerator group specified."
+        self._accelerator_group_id: str = self._typ.communicator_id
+
+        # If the communicators does not contain the group_id, it means the current
+        # process is the driver, and there’s no need to fetch the comm_group.
+        if self._typ.communicator_id in ctx.communicators:
+            self._accelerator_group: "Communicator" = ctx.communicators[
+                self._typ.communicator_id
+            ]
+            assert (
+                self._accelerator_group is not None
+            ), "ChannelContext.accelerator_group is not initialized."
+
+            self._writer_rank = self._accelerator_group.get_rank(self._writer)
+            self._reader_ranks = [
+                self._accelerator_group.get_rank(reader)
+                for reader, _ in self._reader_and_node_list
+            ]
+
+            if (
+                self._writer_rank is not None
+                and self._writer_rank == self._accelerator_group.get_self_rank()
+            ):
+                self._writer_registered = True
+
+            if (
+                self._reader_ranks
+                and self._accelerator_group.get_self_rank() in self._reader_ranks
+            ):
+                self._reader_registered = True
+
+        # If the channel type specifies that the tensor shape is static, then the
+        # receiver can allocate buffers without needing to coordinate with the
+        # sender. We set the metadata on the first send-recv op. Thereafter,
+        # the sender must ensure that sent tensors match this metadata, and the
+        # receiver will allocate tensors with this shape.
+        self._static_tensor_metadata: Optional[List[_TorchTensorMetadata]] = None
+        self._meta_channel: Optional[Channel] = _meta_channel
+        if self._meta_channel is None and self._writer_registered:
+            # We are the writer. Therefore, we also need to allocate a metadata
+            # channel that will be used to send the shape and dtype of the
+            # tensor to the receiver(s).
+            metadata_type = SharedMemoryType()
+            self._meta_channel = metadata_type.create_channel(
+                self._writer,
+                self._reader_and_node_list,
+                None,
+            )
+
+    def ensure_registered_as_writer(self):
+        assert (
+            self._accelerator_group is not None
+        ), "Actor is not part of an accelerator group"
+        assert self._writer_registered
+        ctx = ChannelContext.get_current()
+        assert ctx.torch_device.type != "cpu"
+
+    def ensure_registered_as_reader(self) -> bool:
+        assert (
+            self._accelerator_group is not None
+        ), "Actor is not part of an accelerator group"
+        assert self._reader_registered
+        ctx = ChannelContext.get_current()
+        assert ctx.torch_device.type != "cpu"
+
+    def __reduce__(self):
+        return (
+            self.__class__,
+            (
+                self._writer,
+                self._reader_and_node_list,
+                self._typ,
+                self._meta_channel,
+            ),
+        )
+
+    def _get_send_tensors_metadata(
+        self, tensors: List["torch.Tensor"]
+    ) -> Optional[List[_TorchTensorMetadata]]:
+        """
+        Helper method to get the metadata that should be sent to the reader so
+        that they can allocate the proper-sized buffer(s). Throws error if
+        static_shape=True was set and the given tensors do not match the
+        inferred shapes.
+
+        Returns: The metadata to send to the reader. None means that we should
+            not send any metadata message to the reader.
+        """
+        ctx = ChannelContext.get_current()
+
+        # TODO(swang): Currently any exceptions thrown during this method are
+        # fatal for the DAG because there is no way for the receiver to receive
+        # the exception. This can be improved by sending the exception through
+        # the CPU-based non-tensor-data channel, if one exists. The tensor
+        # channel can send empty data alongside the exception to avoid hanging.
+
+        # Get the shape and dtype of each tensor to send.
+        metadata_list = []
+        for tensor in tensors:
+            # Basic type checking.
+            if not isinstance(tensor, self.torch.Tensor):
+                raise ValueError("Task must return torch.Tensors")
+
+            if tensor.device != ctx.torch_device:
+                raise ValueError(
+                    f"torch.Tensor must be on the default device: {ctx.torch_device}"
+                )
+
+            metadata = _TorchTensorMetadata(tensor.shape, tensor.dtype)
+            metadata_list.append(metadata)
+
+        if self._static_tensor_metadata is not None:
+            if metadata_list != self._static_tensor_metadata:
+                metadata_str = [
+                    f"(shape={m.shape}, dtype={m.dtype})" for m in metadata_list
+                ]
+                expected_str = [
+                    f"(shape={m.shape}, dtype={m.dtype})"
+                    for m in self._static_tensor_metadata
+                ]
+                raise ValueError(
+                    "Expected torch.Tensors with shapes and dtypes: "
+                    "[" + ", ".join(expected_str) + "], "
+                    "found: [" + ", ".join(metadata_str) + "]. "
+                    "DAG will shut down."
+                )
+            # The receiver has already determined the shape and dtype of the
+            # tensors from a previous send, so no need to send the metadata
+            # again.
+            return None
+
+        if self._static_shape:
+            # The shape and dtype is static. This is the first send op and
+            # afterwards, a ValueError will be thrown if the sent tensors do
+            # not match this metadata.
+            self._static_tensor_metadata = metadata_list
+        return metadata_list
+
+    def write(
+        self,
+        tensors: List["torch.Tensor"],
+        timeout: Optional[float] = None,
+    ):
+        """
+        Write a list of tensors via accelerator:
+
+        1) Send the tensor metadata, i.e. the shape and dtypes of all tensors
+        via the shared-memory metadata channel.
+        2) Send the tensor data via accelerator.
+
+        If static_shape=True was set, then we only perform step (1) on the
+        first message. The reader is expected to reuse the sent metadata for
+        subsequent messages.
+        """
+        self.ensure_registered_as_writer()
+
+        import torch
+
+        for tensor in tensors:
+            assert isinstance(
+                tensor, torch.Tensor
+            ), f"{tensor} must be instance of torch.Tensor"
+
+        # Send the tensors metadata so that the receiver knows what buffers to
+        # allocate.
+        metadata = self._get_send_tensors_metadata(tensors)
+        if metadata is not None:
+            self._meta_channel.write(metadata)
+
+        # NOTE(swang): We must send the metadata *before* launching the accelerator
+        # send. We are using blocking accelerator ops, so the following calls will
+        # block until the kernel has been enqueued. Also, peers must launch the
+        # kernel together before either can proceed. Therefore, we send the
+        # metadata first so that the receiver can read the metadata and then
+        # launch the same accelerator op.
+        for tensor in tensors:
+            # TODO: If there are multiple readers, can replace with a
+            # broadcast.
+            for rank in self._reader_ranks:
+                self._accelerator_group.send(tensor, rank)
+
+    def _get_recv_tensors_metadata(
+        self, timeout: Optional[float] = None
+    ) -> List[_TorchTensorMetadata]:
+        """
+        Get the shape(s) and dtype(s) of the tensors to receive from the
+        metadata channel. If static_shape=True was set, then we reuse the first
+        metadata received.
+        """
+        if self._static_tensor_metadata is not None:
+            return self._static_tensor_metadata
+
+        meta = self._meta_channel.read(timeout)
+
+        if self._static_shape:
+            self._static_tensor_metadata = meta
+
+        return meta
+
+    def read(
+        self,
+        timeout: Optional[float] = None,
+    ) -> Union["torch.Tensor", List["torch.Tensor"]]:
+        """
+        Receive a list of tensors.
+
+        (1) Receive the tensor metadata via the shared-memory metadata channel.
+        (2) Allocate buffers on our default device according to the received
+        tensor metadata.
+        (3) Receive the tensor data via accelerator.
+
+        If static_data=True was set, then we only perform step (1) on the first
+        message. Subsequent messages reuse the same metadata.
+
+        NOTE: Currently `timeout` only applies to receiving the CPU-based
+        tensor metadata. The GPU recv may exceed the timeout without throwing
+        an error.
+        """
+        self.ensure_registered_as_reader()
+
+        meta_list: List[_TorchTensorMetadata] = self._get_recv_tensors_metadata(timeout)
+
+        bufs: List["torch.Tensor"] = []
+        for meta in meta_list:
+            buf = self._accelerator_group.recv(
+                meta.shape, meta.dtype, self._writer_rank, _torch_tensor_allocator
+            )
+            bufs.append(buf)
+        # TODO: Sync CUDA stream after receiving all tensors, instead of after
+        # each tensor.
+        return bufs
+
+    def close(self) -> None:
+        self._meta_channel.close()
+
+        self._accelerator_group.destroy()
+        ctx = ChannelContext.get_current()
+        if self._accelerator_group_id in ctx.communicators:
+            del ctx.communicators[self._accelerator_group_id]
+
+
+def _do_init_communicator(
+    self,
+    group_id,
+    world_size,
+    comm_id,
+    rank,
+    actor_handles,
+    use_communication_streams,
+    custom_communicator: Optional[Communicator] = None,
+):
+    if not custom_communicator:
+        assert (
+            AcceleratorContext.get().accelerator_count > 0
+        ), "Actors participating in Communication group must have at least one Accelerator assigned"
+
+    ctx = ChannelContext.get_current()
+    if custom_communicator is not None:
+        custom_communicator.initialize(rank)
+        ctx.communicators[group_id] = custom_communicator
+    else:
+        # default to CommGroup
+        ctx.communicators[group_id] = AcceleratorContext.get().create_communicator(
+            world_size,
+            comm_id,
+            rank,
+            actor_handles,
+            AcceleratorContext.get().current_stream(),
+            use_communication_streams,
+        )
+
+
+def _do_destroy_communicator(self, group_id):
+    ctx = ChannelContext.get_current()
+    if group_id not in ctx.communicators:
+        return
+    ctx.communicators[group_id].destroy()
+
+    # Keep the communicator group in the map after destruction in case there is
+    # still a task loop running.
+
+
+def _do_check_has_accelerators(self) -> str:
+    return AcceleratorContext.get().accelerator_count > 0
+
+
+def do_register_accelerator_context(self, name: str, communicator: Type[Communicator]):
+    register_accelerator_context(name, communicator)
+
+
+def _do_get_unique_communication_id(self) -> bool:
+    return AcceleratorContext.get().generate_communicator_id()
+
+
+def _get_ranks(
+    actors: List[ray.actor.ActorHandle], custom_comm_group: Optional[Communicator]
+) -> List[int]:
+    """Get ranks for the communicator group to use.
+
+    If ``custom_comm_group`` is specified, return the ranks of the actors in the
+    custom communicator group, in the same order of the actors; otherwise,
+    return ``list(range(len(actors)))``.
+
+    Args:
+        actors: A list of actors that participate in the communicator group.
+        custom_comm_group: The custom communicator group to use.
+
+    Returns:
+        The list of ranks corresponding to ``actors``.
+    """
+    if custom_comm_group is None:
+        return list(range(len(actors)))
+
+    assert len(actors) == custom_comm_group.get_world_size(), (
+        "The world size of the custom communicator group does not match the "
+        "number of actors."
+    )
+    ranks = []
+    for actor in actors:
+        rank = custom_comm_group.get_rank(actor)
+        assert rank not in ranks, "Duplicate rank in custom communicator group"
+        ranks.append(rank)
+    assert custom_comm_group.get_world_size() == len(actors), (
+        "The world size of the custom communicator group "
+        f"({custom_comm_group.get_world_size()}) "
+        "does not match the number of actors "
+        f"({len(actors)})."
+    )
+    return ranks
+
+
+def _init_communicator(
+    actors: List[ray.actor.ActorHandle],
+    custom_communicator: Optional[Communicator] = None,
+    use_communication_streams: bool = False,
+    accelerator_module_name: Optional[str] = None,
+    accelerator_communicator_cls: Optional[Type[Communicator]] = None,
+) -> str:
+    """Initialize a communicator group with the given actors.
+
+    If a custom communicator group is provided, then it will be used, otherwise
+    a new communicator group will be created.
+
+    Args:
+        actors: A list of actors that participate in the communicator group.
+        custom_communicator: A custom communicator group to initialize.
+        use_communication_streams: Whether to use dedicated send and recv
+                streams for communication. If True, communication and computation
+                can be overlapped to improve performance.
+        accelerator_module_name: Optional name of the accelerator module to use.
+        accelerator_communicator_cls: Optional communicator class for the accelerator.
+
+    Returns:
+        The unique ``group_id`` identifying the initialized communicator group.
+    """
+    ctx = ChannelContext.get_current()
+
+    is_cpu_communicator = custom_communicator and isinstance(
+        custom_communicator, CPUCommunicator
+    )
+
+    # Register accelerator context for all actors if accelerator is not default
+    if accelerator_module_name and accelerator_communicator_cls:
+        if is_accelerator_context_registered():
+            ray.get(
+                [
+                    actor.__ray_call__.remote(
+                        do_register_accelerator_context,
+                        accelerator_module_name,
+                        accelerator_communicator_cls,
+                    )
+                    for actor in actors
+                ]
+            )
+
+    has_accelerators = ray.get(
+        [actor.__ray_call__.remote(_do_check_has_accelerators) for actor in actors]
+    )
+    for has_accelerator, actor in zip(has_accelerators, actors):
+        if not has_accelerator and not is_cpu_communicator:
+            raise ValueError(
+                f"Actor {actor} returns a tensor with type hint "
+                'TorchTensor(transport="accelerator") or '
+                "TorchTensor(transport=accelerator_group_handle) "
+                "but actor does not have an accelerator assigned by Ray."
+            )
+
+    actor_ids = {actor._ray_actor_id for actor in actors}
+    assert len(actor_ids) == len(actors), "Actors must be unique"
+
+    # Allocate a communicator ID on one of the actors that will participate in
+    # the group. This is in case the driver is not on the same node as one of
+    # the communicator actors.
+    comm_id = ray.get(actors[0].__ray_call__.remote(_do_get_unique_communication_id))
+
+    # Used to uniquely identify this communicator group.
+    group_id = str(uuid.uuid4())
+
+    if custom_communicator is not None:
+        logger.info(
+            f"Initializing custom communicator group {group_id} on actors: {actors}"
+        )
+    else:
+        logger.info(f"Creating communicator group {group_id} on actors: {actors}")
+
+    world_size = len(actors)
+    ranks = _get_ranks(actors, custom_communicator)
+    init_tasks = [
+        actor.__ray_call__.remote(
+            _do_init_communicator,
+            group_id,
+            world_size,
+            comm_id,
+            rank,
+            actors,
+            use_communication_streams,
+            custom_communicator,
+        )
+        for rank, actor in zip(ranks, actors)
+    ]
+    try:
+        ray.get(init_tasks, timeout=30)
+    except ray.exceptions.GetTimeoutError:
+        logger.warning(
+            "Communicator group creation not done after 30s. communicator group"
+            "creation may be hung."
+        )
+        ray.get(init_tasks)
+
+    logger.info("Communicator group initialized.")
+
+    if custom_communicator is not None:
+        ctx.communicator_handles[group_id] = CommunicatorHandle(
+            actor_handles=custom_communicator.get_actor_handles(),
+        )
+    else:
+        ctx.communicator_handles[group_id] = CommunicatorHandle(
+            actor_handles=actors,
+        )
+
+    return group_id
+
+
+def _destroy_communicator(group_id: str) -> None:
+    """
+    Destroy the communicator group with the given ID.
+    """
+    ctx = ChannelContext.get_current()
+    if group_id not in ctx.communicator_handles:
+        return
+
+    group = ctx.communicator_handles[group_id]
+    actors = group.get_actor_handles()
+    destroy_tasks = [
+        actor.__ray_call__.remote(
+            _do_destroy_communicator,
+            group_id,
+        )
+        for actor in actors
+    ]
+
+    _, unready = ray.wait(destroy_tasks, timeout=30, num_returns=len(destroy_tasks))
+    if unready:
+        logger.warning(
+            "Communicator group destruction not done after 30s. Communicator"
+            "group destruction may be hung."
+        )
+
+    del ctx.communicator_handles[group_id]

python/ray/experimental/channel/torch_tensor_type.py

@@ -0,0 +1,193 @@
+import logging
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+
+import ray
+from ray.experimental.channel import ChannelContext, ChannelOutputType
+from ray.experimental.channel.communicator import Communicator
+from ray.experimental.channel.shared_memory_channel import SharedMemoryType
+from ray.experimental.util.types import Device
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.experimental.channel.shared_memory_channel import Channel
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI(stability="alpha")
+class TorchTensorType(ChannelOutputType):
+    AUTO = "auto"
+    CPU = "cpu"
+    ACCELERATOR = "accelerator"
+
+    def __init__(
+        self,
+        transport: Optional[Union[str, Communicator]] = AUTO,
+        device: Device = Device.DEFAULT,
+        _static_shape: bool = False,
+        _direct_return: Optional[bool] = False,
+    ):
+        """
+        A type hint that can be used to annotate DAG nodes that return a
+        torch.Tensor.
+
+        NOTE: Use of this type in the DAG will register a custom serializer for
+        torch.Tensor that moves the tensor to the correct device on the
+        receiver. If you are using ray.cloudpickle to serialize objects and you
+        do not want this behavior, deregister the custom serializer using
+        ray.util.serialization.deregister_serializer(torch.Tensor).
+
+        Args:
+            transport: "auto" (default) means that tensors will be passed via
+                host memory, using numpy as the serialization format. Pass
+                TorchTensorType.ACCELERATOR or "accelerator" to use accelerator
+                instead, avoiding the host memory copy.
+            device: Target device for tensor transport. Options:
+                - "default": Retains the same device type as the sender.
+                - "cpu": Moves tensor to CPU on the receiver. Not compatible
+                  with accelerator transport.
+                - "gpu" or "cuda": Moves tensor to GPU on the receiver.
+            _static_shape: A hint indicating whether the shape(s) and dtype(s)
+                of tensor(s) contained in this value always remain the same
+                across different executions of the DAG.
+            _direct_return: Whether the tensor is sent directly or inside of
+                other data. If a non-default `transport` is used, this allows
+                the sender and receiver to eliminate performance overhead from
+                an additional data transfer.
+
+        NOTE: Setting static_shape=True and _direct_return=True can improve
+        performance if a non-default transport is used. However, if either flag
+        is set, then the user must ensure that the condition is met.
+
+        If using this type as a Compiled Graph annotation, an exception will
+        be thrown in the following cases, and the DAG will be torn down. To
+        continue execution, a new DAG must be created:
+        1. If _static_shape=True, and the found tensors don't match the
+           previous shape or dtype(s).
+        2. If _direct_return=True, and the returned value is not a
+           torch.Tensor.
+        """
+        super().__init__()
+
+        self._device = device
+        self._static_shape = _static_shape
+        self._direct_return = _direct_return
+
+        self._communicator: Optional[Communicator] = None
+        if isinstance(transport, Communicator):
+            self._communicator = transport
+            transport = transport.get_transport_name()
+
+        if transport not in [self.AUTO, self.CPU, self.ACCELERATOR]:
+            raise ValueError(
+                "`transport` must be TorchTensorType.AUTO, TorchTensorType.ACCELERATOR "
+                "or TorchTensorType.CPU"
+            )
+        if device == Device.CPU and transport == self.ACCELERATOR:
+            raise ValueError(
+                "accelerator transport is not supported with CPU target device."
+            )
+        self.transport = transport
+
+        self._communicator_id: Optional[str] = None
+
+        if self._static_shape and self.transport == self.AUTO:
+            logger.info(
+                "TorchTensorType(_static_shape=True) has no effect when "
+                "`transport` is TorchTensorType.AUTO (default)."
+            )
+        if self._direct_return and self.transport == self.AUTO:
+            logger.info(
+                "TorchTensorType(_direct_return=True) has no effect when "
+                "`transport` is TorchTensorType.AUTO (default)."
+            )
+
+    @property
+    def device(self) -> Device:
+        return self._device
+
+    @property
+    def static_shape(self):
+        return self._static_shape
+
+    @property
+    def direct_return(self):
+        return self._direct_return
+
+    def register_custom_serializer(self) -> None:
+        super().register_custom_serializer()
+
+        import torch
+
+        def serialize(t):
+            ctx = ChannelContext.get_current()
+            return ctx.serialization_context.serialize_tensor(t)
+
+        def deserialize(b):
+            ctx = ChannelContext.get_current()
+            return ctx.serialization_context.deserialize_tensor(b, self.device)
+
+        ray.util.serialization.register_serializer(
+            torch.Tensor,
+            serializer=serialize,
+            deserializer=deserialize,
+        )
+
+    def create_channel(
+        self,
+        writer: Optional["ray.actor.ActorHandle"],
+        reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
+        driver_actor_id: Optional[str] = None,
+        _cpu_data_channel: Optional["Channel"] = None,
+        _tensor_metadata_channel: Optional["Channel"] = None,
+    ) -> type:
+        if self.requires_accelerator():
+            from ray.experimental.channel.torch_tensor_accelerator_channel import (
+                TorchTensorAcceleratorChannel,
+            )
+
+            return TorchTensorAcceleratorChannel(
+                writer,
+                reader_and_node_list,
+                self,
+                driver_actor_id,
+                _tensor_metadata_channel,
+                _cpu_data_channel,
+            )
+
+        # Data does not require accelerator. Transfer via host memory using a
+        # shared-memory channel.
+        # TODO(swang): Allow the initial max buffer size to be overridden.
+        typ = SharedMemoryType()
+        return typ.create_channel(writer, reader_and_node_list, driver_actor_id)
+
+    def requires_accelerator(self) -> bool:
+        return self.transport == self.ACCELERATOR
+
+    def get_custom_communicator(self) -> Optional[Communicator]:
+        """
+        Return the communicator group if one is specified.
+        """
+        return self._communicator
+
+    def set_communicator_id(self, group_id: str) -> None:
+        self._communicator_id = group_id
+
+    @property
+    def communicator_id(self) -> Optional[str]:
+        return self._communicator_id
+
+    def __deepcopy__(self, memo):
+        """
+        Deep copy all the fields except for the communicator group. The communicator
+        group should not be deep copied because it can be shared across `TorchTensorType`
+        instances.
+        """
+        copy = TorchTensorType(
+            transport=self.transport,
+            _static_shape=self._static_shape,
+            _direct_return=self._direct_return,
+        )
+        copy._communicator = self._communicator
+        copy._communicator_id = self._communicator_id
+        return copy

python/ray/experimental/channel/utils.py

@@ -0,0 +1,92 @@
+from typing import List, Optional, Tuple
+
+import ray
+
+
+def get_self_actor() -> Optional["ray.actor.ActorHandle"]:
+    """
+    Get the current actor handle in this worker.
+    If this is called in a driver process, it will return None.
+    """
+    try:
+        return ray.get_runtime_context().current_actor
+    except RuntimeError:
+        return None
+
+
+def split_readers_by_locality(
+    writer: "ray.actor.ActorHandle",
+    reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
+) -> Tuple[
+    List[Tuple["ray.actor.ActorHandle", str]], List[Tuple["ray.actor.ActorHandle", str]]
+]:
+    """Split readers into remote and local readers based on writer.
+
+    Args:
+        writer: The actor handle of the writer
+        reader_and_node_list: List of (reader, node) tuples
+
+    Returns:
+        Tuple containing:
+            - List of (reader, node) tuples for remote readers
+            - List of (reader, node) tuples for local readers
+    """
+    remote_readers = []
+    local_readers = []
+
+    for reader, node in reader_and_node_list:
+        if reader != writer:
+            remote_readers.append((reader, node))
+        else:
+            local_readers.append((reader, node))
+
+    return remote_readers, local_readers
+
+
+def split_actors_by_node_locality(
+    node: str,
+    actor_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
+) -> Tuple[
+    List[Tuple["ray.actor.ActorHandle", str]], List[Tuple["ray.actor.ActorHandle", str]]
+]:
+    """Split actors into remote and local actors based on node. The local actors will be
+    on the same node as the given node. The remote actors will be on a different node.
+
+    Args:
+        node: The node to compare actor locations against.
+        actor_and_node_list: List of (actor, node) tuples
+
+    Returns:
+        Tuple containing:
+            - List of (actor, node) tuples for actors on the same node
+            - List of (actor, node) tuples for actors on a different node
+    """
+    actors_on_same_node = []
+    actors_on_different_node = []
+
+    for actor, actor_node in actor_and_node_list:
+        if node == actor_node:
+            actors_on_same_node.append((actor, actor_node))
+        else:
+            actors_on_different_node.append((actor, actor_node))
+
+    return actors_on_same_node, actors_on_different_node
+
+
+def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
+    """Get the node of the actor.
+
+    Args:
+        actor: The actor handle of the actor
+
+    Returns:
+        The node of the actor
+    """
+    if actor is None or actor == ray.get_runtime_context().current_actor:
+        return ray.get_runtime_context().get_node_id()
+    else:
+        return ray.get(
+            actor.__ray_call__.remote(
+                lambda self: ray.get_runtime_context().get_node_id()
+            )
+        )

python/ray/experimental/collective/__init__.py

@@ -0,0 +1,21 @@
+from ray.experimental.collective.collective import (
+    create_collective_group,
+    destroy_all_collective_groups,
+    destroy_collective_group,
+    get_collective_groups,
+)
+from ray.experimental.collective.operations import (
+    allgather,
+    allreduce,
+    reducescatter,
+)
+
+__all__ = [
+    "allgather",
+    "allreduce",
+    "reducescatter",
+    "get_collective_groups",
+    "create_collective_group",
+    "destroy_collective_group",
+    "destroy_all_collective_groups",
+]

python/ray/experimental/collective/collective.py

@@ -0,0 +1,209 @@
+import threading
+import uuid
+from typing import Dict, List, Optional, Union
+
+import ray
+import ray.experimental.internal_kv as internal_kv
+from ray.experimental.collective.communicator import CommunicatorHandle
+from ray.util.annotations import PublicAPI
+from ray.util.collective.collective_group.torch_gloo_collective_group import (
+    get_master_address_metadata_key,
+)
+from ray.util.collective.types import Backend
+
+_remote_communicator_manager: "Optional[RemoteCommunicatorManager]" = None
+_remote_communicator_manager_lock = threading.Lock()
+
+
+class RemoteCommunicatorManager:
+    """Singleton class to store the mapping between actors and communicators
+    that the actors are a part of.
+    """
+
+    def __init__(self):
+        # Handles to communicators that we created. Key is a user-provided
+        # name or UUID.
+        self._remote_communicators: Dict[str, CommunicatorHandle] = {}
+
+    @staticmethod
+    def get() -> "RemoteCommunicatorManager":
+        global _remote_communicator_manager
+        with _remote_communicator_manager_lock:
+            if _remote_communicator_manager is None:
+                _remote_communicator_manager = RemoteCommunicatorManager()
+            return _remote_communicator_manager
+
+    def add_remote_communicator(self, comm_handle: CommunicatorHandle):
+        self._remote_communicators[comm_handle.name] = comm_handle
+
+    def remove_remote_communicator(self, name: str):
+        return self._remote_communicators.pop(name, None)
+
+    def get_collective_groups(
+        self,
+        actors: Optional[List[ray.actor.ActorHandle]] = None,
+        backend: Optional[Backend] = None,
+    ):
+        """
+        Get the collective groups that the given actors are a subset of. Filter by
+        backend if provided.
+        """
+        actors = actors or []
+        actors = set(actors)
+
+        collectives = []
+        # Find all collective groups that the given actors are a subset
+        # of, with the matching backend if provided.
+        for collective in self._remote_communicators.values():
+            if actors.issubset(set(collective.actors)):
+                if backend is None or collective.backend == backend:
+                    collectives.append(collective)
+        return collectives
+
+
+@PublicAPI(stability="alpha")
+def get_collective_groups(
+    actors: List[ray.actor.ActorHandle], backend: Optional[str] = None
+) -> List[CommunicatorHandle]:
+    """
+    Get the collective groups that the given actors are a subset of. Filter by
+    backend if provided.
+
+    Args:
+        actors: List of actors. Return handles to all collective groups that
+            these actors are a subset of.
+        backend: An optional backend to filter by. See
+            ray.util.collective.types.Backend for valid backends.
+
+    Returns:
+        A list of communicator handles that the actors are a subset of.
+    """
+    manager = RemoteCommunicatorManager.get()
+    backend = Backend(backend) if backend is not None else None
+    return manager.get_collective_groups(actors, backend)
+
+
+@PublicAPI(stability="alpha")
+def create_collective_group(
+    actors: List[ray.actor.ActorHandle],
+    backend: str,
+    name: Optional[str] = None,
+) -> CommunicatorHandle:
+    """Create a collective group on the given list of actors. If this function
+    returns successfully, then the collective group has been initialized on all
+    actors, using the given order of actors as the ranks.
+
+    Currently, an actor can only participate in one collective group per
+    backend at a time. To reuse an actor, destroy its collective group and
+    create a new one.
+
+    Args:
+        actors: The actors to participate in the collective group.
+        backend: The backend to use. See ray.util.collective.types.Backend for
+            valid backends.
+        name: A name to use for the collective group. If None is provided, a
+            random name will be generated.
+
+    Returns:
+        Handle to the communicator.
+    """
+    manager = RemoteCommunicatorManager.get()
+
+    if name is None:
+        name = str(uuid.uuid4())
+
+    # Validate the backend.
+    backend = Backend(backend)
+
+    world_size = len(actors)
+
+    for actor in actors:
+        if manager.get_collective_groups([actor], backend):
+            raise RuntimeError(
+                f"Actor {actor} already in group for backend {backend}. Actors can currently only participate in at most one group per backend."
+            )
+
+    actor_ids = [actor._ray_actor_id for actor in actors]
+    if len(set(actor_ids)) != len(actor_ids):
+        raise ValueError(f"All actors must be unique, got: {actors}")
+
+    metadata_key = None
+    if backend == Backend.GLOO:
+        metadata_key = get_master_address_metadata_key(name)
+
+    def _do_init_collective_group(self, rank: int):
+        ray.util.collective.init_collective_group(
+            world_size, rank, backend, group_name=name
+        )
+
+    try:
+        init_tasks = [
+            actor.__ray_call__.remote(
+                _do_init_collective_group,
+                rank,
+            )
+            for rank, actor in enumerate(actors)
+        ]
+        ray.get(init_tasks)
+    finally:
+        # Clean up the metadata once collective group is initialized
+        # (or failed to initialize).
+        if metadata_key is not None:
+            internal_kv._internal_kv_del(metadata_key)
+
+    # Group was successfully created.
+    comm = CommunicatorHandle(actors, name, backend)
+    manager.add_remote_communicator(comm)
+    return comm
+
+
+@PublicAPI(stability="alpha")
+def destroy_collective_group(group_or_name: Union[CommunicatorHandle, str]):
+    """
+    Destroy a collective group. If this functions returns successfully, then
+    the actors that were in the collective can be reused to create a new
+    collective group.
+
+    Args:
+        group_or_name: Either a communicator handle or the name of the group to
+            destroy.
+    """
+    if isinstance(group_or_name, CommunicatorHandle):
+        name = group_or_name.name
+    elif isinstance(group_or_name, str):
+        name = group_or_name
+    else:
+        raise ValueError("Expected CommunicatorHandle or str (group name).")
+
+    manager = RemoteCommunicatorManager.get()
+    group = manager.remove_remote_communicator(name)
+    if group is not None:
+
+        def _do_destroy_collective_group(self):
+            ray.util.collective.destroy_collective_group(name)
+
+        destroy_tasks = [
+            actor.__ray_call__.options(concurrency_group="_ray_system").remote(
+                _do_destroy_collective_group
+            )
+            for actor in group.actors
+        ]
+        try:
+            ray.get(destroy_tasks)
+        except ray.exceptions.ActorDiedError:
+            pass
+    else:
+        raise ValueError(f"No group with name {name} found.")
+
+
+@PublicAPI(stability="alpha")
+def destroy_all_collective_groups():
+    """
+    Destroy all collective groups. This will destroy all collective groups that
+    were previously created by this process. After this function returns, the
+    actors participating in those collective groups can be reused to create a
+    new collective group.
+    """
+    manager = RemoteCommunicatorManager.get()
+    for collective in manager.get_collective_groups():
+        destroy_collective_group(collective.name)

python/ray/experimental/collective/communicator.py

@@ -0,0 +1,63 @@
+from dataclasses import dataclass
+from typing import List
+
+import ray
+from ray.util.collective.types import Backend
+
+
+@dataclass
+class Communicator:
+    """
+    A handle to a communicator that we are a member of.
+    """
+
+    # The name of the communicator.
+    name: str
+    # Our rank in the collective group.
+    rank: int
+    # A valid backend, as defined by
+    # ray.util.collective.types.Backend.
+    backend: str
+
+
+class CommunicatorHandle:
+    """
+    A communicator handle used by the driver to store handles to the
+    actors in the communicator.
+    """
+
+    def __init__(self, actors: List[ray.actor.ActorHandle], name: str, backend: str):
+        """
+        Initializes the CommunicatorHandle with the given actor handles.
+        Assumes that the communicator has already been initialized on all actors.
+
+        Args:
+            actors: A list of actor handles to be stored.
+            name: Name of the communicator.
+            backend: Communicator backend. See
+                ray.util.collective.types for valid values.
+        """
+        self._actors = actors
+        self._name = name
+        self._backend = Backend(backend)
+
+    def get_rank(self, actor: ray.actor.ActorHandle):
+        for i, a in enumerate(self._actors):
+            if a == actor:
+                return i
+        return -1
+
+    @property
+    def actors(self) -> List[ray.actor.ActorHandle]:
+        """
+        Return all actor handles in this communicator.
+        """
+        return self._actors[:]
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def backend(self) -> str:
+        return self._backend

python/ray/experimental/collective/conftest.py

@@ -0,0 +1,243 @@
+import uuid
+from typing import Dict, FrozenSet, List, Optional, Set, Tuple, Type
+
+import torch
+
+import ray
+from ray.experimental.channel.common import ChannelContext
+from ray.experimental.channel.communicator import (
+    Communicator,
+    ReduceOp,
+    TorchTensorAllocator,
+)
+
+
+class AbstractNcclGroup(Communicator):
+    """
+    A dummy NCCL group for testing.
+    """
+
+    def __init__(self, actor_handles: List[ray.actor.ActorHandle]):
+        self._actor_handles = actor_handles
+        self._rank = None
+
+    def initialize(self, rank: int) -> None:
+        self._rank = rank
+
+    def get_rank(self, actor: ray.actor.ActorHandle) -> int:
+        return self._actor_handles.index(actor)
+
+    def get_world_size(self) -> int:
+        return len(self._actor_handles)
+
+    def get_self_rank(self) -> Optional[int]:
+        return self._rank
+
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        return self._actor_handles
+
+    def send(self, value: "torch.Tensor", peer_rank: int) -> None:
+        raise NotImplementedError
+
+    def recv(
+        self,
+        shape: Tuple[int],
+        dtype: "torch.dtype",
+        peer_rank: int,
+        allocator: Optional[TorchTensorAllocator] = None,
+    ) -> "torch.Tensor":
+        raise NotImplementedError
+
+    def allgather(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+    ) -> None:
+        raise NotImplementedError
+
+    def allreduce(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp = ReduceOp.SUM,
+    ) -> None:
+        raise NotImplementedError
+
+    def reducescatter(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp = ReduceOp.SUM,
+    ) -> None:
+        raise NotImplementedError
+
+    @property
+    def recv_stream(self):
+        return None
+
+    @property
+    def send_stream(self):
+        return None
+
+    def destroy(self) -> None:
+        pass
+
+    def get_transport_name(self) -> str:
+        return "accelerator"
+
+    @classmethod
+    def generate_communicator_id(cls) -> str:
+        pass
+
+
+class MockNcclGroupSet:
+    def __init__(self):
+        # Represents a mapping from a NCCL group ID to a set of actors and a custom
+        # NCCL group.
+        self.ids_to_actors_and_custom_comms: Dict[
+            str, Tuple[FrozenSet["ray.actor.ActorHandle"], Optional[Communicator]]
+        ] = {}
+
+    def __call__(
+        self,
+        actors: List["ray.actor.ActorHandle"],
+        custom_nccl_group: Optional[Communicator] = None,
+        use_communication_streams: bool = False,
+        accelerator_module_name: Optional[str] = None,
+        accelerator_communicator_cls: Optional[Type[Communicator]] = None,
+    ) -> str:
+        group_id = str(uuid.uuid4())
+        self.ids_to_actors_and_custom_comms[group_id] = (
+            frozenset(actors),
+            custom_nccl_group,
+        )
+
+        if custom_nccl_group is None:
+            ranks = list(range(len(actors)))
+        else:
+            ranks = [custom_nccl_group.get_rank(actor) for actor in actors]
+        init_tasks = [
+            actor.__ray_call__.remote(
+                mock_do_init_nccl_group,
+                group_id,
+                rank,
+                actors,
+                custom_nccl_group,
+            )
+            for rank, actor in zip(ranks, actors)
+        ]
+        ray.get(init_tasks, timeout=30)
+
+        ctx = ChannelContext.get_current()
+        if custom_nccl_group is not None:
+            ctx.communicators[group_id] = custom_nccl_group
+        else:
+            ctx.communicators[group_id] = AbstractNcclGroup(actors)
+
+        return group_id
+
+    def mock_destroy_nccl_group(self, group_id: str) -> None:
+        ctx = ChannelContext.get_current()
+        if group_id not in ctx.communicators:
+            return
+
+        actors, _ = self.ids_to_actors_and_custom_comms[group_id]
+        destroy_tasks = [
+            actor.__ray_call__.remote(
+                mock_do_destroy_nccl_group,
+                group_id,
+            )
+            for actor in actors
+        ]
+        ray.wait(destroy_tasks, timeout=30)
+
+        if group_id in self.ids_to_actors_and_custom_comms:
+            del self.ids_to_actors_and_custom_comms[group_id]
+        ctx.communicators[group_id].destroy()
+        del ctx.communicators[group_id]
+
+    def check_teardown(self, nccl_group_ids: List[str]) -> None:
+        ctx = ChannelContext.get_current()
+        for nccl_group_id in nccl_group_ids:
+            assert nccl_group_id not in self.ids_to_actors_and_custom_comms
+            assert nccl_group_id not in ctx.communicators
+
+
+@ray.remote
+class CPUTorchTensorWorker:
+    def __init__(self):
+        self.device = "cpu"
+
+    def return_tensor(
+        self, size: int, dtype: Optional[torch.dtype] = None
+    ) -> torch.Tensor:
+        return torch.ones(size, dtype=dtype, device=self.device)
+
+    def recv(self, tensor: torch.Tensor) -> Tuple[int, int]:
+        assert tensor.device == self.device
+        return tensor.shape, tensor[0]
+
+    def recv_tensors(self, *tensors) -> Tuple[torch.Tensor, ...]:
+        return tuple(tensors)
+
+
+def mock_do_init_nccl_group(
+    self,
+    group_id: str,
+    rank: int,
+    actors: List[ray.actor.ActorHandle],
+    custom_nccl_group: Optional[Communicator],
+) -> None:
+    ctx = ChannelContext.get_current()
+    if custom_nccl_group is None:
+        nccl_group = AbstractNcclGroup(actors)
+        nccl_group.initialize(rank)
+        ctx.communicators[group_id] = nccl_group
+    else:
+        custom_nccl_group.initialize(rank)
+        ctx.communicators[group_id] = custom_nccl_group
+
+
+def mock_do_destroy_nccl_group(self, group_id: str) -> None:
+    ctx = ChannelContext.get_current()
+    if group_id not in ctx.communicators:
+        return
+    ctx.communicators[group_id].destroy()
+    del ctx.communicators[group_id]
+
+
+def check_nccl_group_init(
+    monkeypatch,
+    dag: "ray.dag.DAGNode",
+    actors_and_custom_comms: Set[
+        Tuple[FrozenSet["ray.actor.ActorHandle"], Optional[Communicator]]
+    ],
+) -> "ray.dag.CompiledDAG":
+    mock_nccl_group_set = MockNcclGroupSet()
+    monkeypatch.setattr(
+        "ray.dag.compiled_dag_node._init_communicator",
+        mock_nccl_group_set,
+    )
+
+    compiled_dag = dag.experimental_compile()
+    assert (
+        set(mock_nccl_group_set.ids_to_actors_and_custom_comms.values())
+        == actors_and_custom_comms
+    )
+
+    return compiled_dag, mock_nccl_group_set
+
+
+def check_nccl_group_teardown(
+    monkeypatch,
+    compiled_dag: "ray.dag.CompiledDAG",
+    mock_nccl_group_set: MockNcclGroupSet,
+):
+    monkeypatch.setattr(
+        "ray.dag.compiled_dag_node._destroy_communicator",
+        mock_nccl_group_set.mock_destroy_nccl_group,
+    )
+
+    created_communicator_ids = compiled_dag._actors_to_created_communicator_id.values()
+    compiled_dag.teardown()
+    mock_nccl_group_set.check_teardown(created_communicator_ids)

python/ray/experimental/collective/operations.py

@@ -0,0 +1,203 @@
+import logging
+from typing import List, Optional, Union
+
+import ray
+from ray.dag.collective_node import CollectiveOutputNode, _CollectiveOperation
+from ray.dag.constants import (
+    BIND_INDEX_KEY,
+    COLLECTIVE_OPERATION_KEY,
+    IS_CLASS_METHOD_OUTPUT_KEY,
+    PARENT_CLASS_NODE_KEY,
+)
+from ray.experimental.channel.torch_tensor_type import Communicator, TorchTensorType
+from ray.experimental.util.types import (
+    AllGatherOp,
+    AllReduceOp,
+    ReduceOp,
+    ReduceScatterOp,
+    _CollectiveOp,
+)
+from ray.util.collective.types import ReduceOp as RayReduceOp
+
+logger = logging.getLogger(__name__)
+
+
+def _bind(
+    inputs: Union[List["ray.dag.DAGNode"], List[List["ray.dag.DAGNode"]]],
+    op: _CollectiveOp,
+    transport: Optional[Union[str, Communicator]] = None,
+):
+    """
+    Bind inputs (input nodes or lists of input nodes) with a collective operation.
+    The collective operation is applied to each list of input nodes. The output nodes
+    will have the same shape as the input nodes.
+
+    Example of binding a list of input node:
+    with InputNode() as inp:
+        res_comp1 = [actor.comp1.bind(inp) for actor in actors]
+        res_comp2 = [actor.comp2.bind(inp) for actor in actors]
+        res_ar = allreduce.bind([res_comp1, res_comp2])
+
+    Requirements:
+    1. Each input node returns a torch tensor.
+    2. Each input node within a list is from a different actor.
+    3. If lists of input nodes are provided, the order of actors should
+        be the same for each nested list.
+    4. If a custom transport is specified, its actor set matches the actor
+        set of the input nodes.
+    5. If input nodes are provided, then all tensors have the same shape.
+        If lists of input nodes are provided, then all tensors in each
+        list have the same shape.
+
+    Requirements 1-3 are checked in the `CollectiveGroup` constructor.
+    Requirement 4 is not checked yet.
+
+    Args:
+        inputs: A list of DAG nodes or a list of lists of DAG nodes. Each leaf list
+            should contain one object per actor.
+        op: The collective operation.
+        transport: GPU communicator for the collective operation. If not
+            specified, the default ACCELERATOR is used.
+
+    Returns:
+        A list of collective output nodes or a list of lists of collective output nodes,
+        with the same shape as the input nodes. Each output node has the same order and
+        belongs to the same actor as the corresponding input node.
+    """
+    if isinstance(inputs[0], list) and not isinstance(op, AllReduceOp):
+        raise ValueError(
+            "Currently binding a nested list of dag nodes is only supported for allreduce"
+        )
+
+    # Convert list of DAGNode into nested list for type checking
+    if not isinstance(inputs[0], list):
+        inputs = [inputs]
+
+    if transport is None:
+        transport = TorchTensorType.ACCELERATOR
+    collective_op = _CollectiveOperation(inputs, op, transport)
+    collective_output_nodes: List[CollectiveOutputNode] = []
+
+    if isinstance(op, AllGatherOp):
+        method_name = "allgather"
+    elif isinstance(op, AllReduceOp):
+        method_name = f"allreduce.{op.reduceOp}"
+    elif isinstance(op, ReduceScatterOp):
+        method_name = f"reducescatter.{op.reduceOp}"
+    else:
+        raise ValueError(f"Expected a collective operation, but got {op}")
+
+    for i in range(len(inputs[0])):
+        input_node_list = [l[i] for l in inputs if l]
+        actor_handle: Optional["ray.actor.ActorHandle"] = input_node_list[
+            0
+        ]._get_actor_handle()
+        assert actor_handle is not None
+        collective_output_node = CollectiveOutputNode(
+            method_name=method_name,
+            method_args=tuple(input_node_list),
+            method_kwargs=dict(),
+            method_options=dict(),
+            other_args_to_resolve={
+                PARENT_CLASS_NODE_KEY: actor_handle,
+                BIND_INDEX_KEY: actor_handle._ray_dag_bind_index,
+                COLLECTIVE_OPERATION_KEY: collective_op,
+            },
+        )
+        actor_handle._ray_dag_bind_index += 1
+
+        if len(input_node_list) > 1:
+            output_nodes: List[CollectiveOutputNode] = []
+            for i in range(len(input_node_list)):
+                output_node = CollectiveOutputNode(
+                    f"return_idx_{i}",
+                    (collective_output_node, i),
+                    dict(),
+                    dict(),
+                    {
+                        BIND_INDEX_KEY: collective_output_node._get_bind_index(),
+                        IS_CLASS_METHOD_OUTPUT_KEY: True,
+                        PARENT_CLASS_NODE_KEY: actor_handle,
+                    },
+                )
+                output_nodes.append(output_node)
+            collective_output_nodes.append(output_nodes)
+        else:
+            collective_output_nodes.append(collective_output_node)
+    return collective_output_nodes
+
+
+class AllGatherWrapper:
+    """Wrapper for NCCL all-gather."""
+
+    def bind(
+        self,
+        input_nodes: List["ray.dag.DAGNode"],
+        transport: Optional[Union[str, Communicator]] = None,
+    ) -> List[CollectiveOutputNode]:
+        return _bind(input_nodes, AllGatherOp(), transport)
+
+    def __call__(
+        self,
+        tensor_list,
+        tensor,
+        group_name: str = "default",
+    ):
+        from ray.util.collective.collective import allgather
+
+        return allgather(tensor_list, tensor, group_name)
+
+
+class AllReduceWrapper:
+    """Wrapper for NCCL all-reduce."""
+
+    def bind(
+        self,
+        input_nodes: List["ray.dag.DAGNode"],
+        op: ReduceOp = ReduceOp.SUM,
+        transport: Optional[Union[str, Communicator]] = None,
+    ) -> List[CollectiveOutputNode]:
+        if not isinstance(op, ReduceOp):
+            raise ValueError(f"Unexpected operation: {op}")
+
+        return _bind(input_nodes, AllReduceOp(reduceOp=op), transport)
+
+    def __call__(
+        self,
+        tensor,
+        group_name: str = "default",
+        op: RayReduceOp = RayReduceOp.SUM,
+    ):
+        from ray.util.collective.collective import allreduce
+
+        return allreduce(tensor, group_name, op)
+
+
+class ReduceScatterWrapper:
+    """Wrapper for NCCL reduce-scatter."""
+
+    def bind(
+        self,
+        input_nodes: List["ray.dag.DAGNode"],
+        op: ReduceOp = ReduceOp.SUM,
+        transport: Optional[Union[str, Communicator]] = None,
+    ) -> List[CollectiveOutputNode]:
+        if not isinstance(op, ReduceOp):
+            raise ValueError(f"Unexpected operation: {op}")
+
+        return _bind(input_nodes, ReduceScatterOp(reduceOp=op), transport)
+
+    def __call__(
+        self,
+        tensor,
+        group_name: str = "default",
+        op: RayReduceOp = RayReduceOp.SUM,
+    ):
+        from ray.util.collective.collective import reducescatter
+
+        return reducescatter(tensor, group_name, op)
+
+
+allgather = AllGatherWrapper()
+allreduce = AllReduceWrapper()
+reducescatter = ReduceScatterWrapper()

python/ray/experimental/compiled_dag_ref.py

@@ -0,0 +1,225 @@
+import asyncio
+from typing import Any, List, Optional
+
+import ray
+from ray.exceptions import (
+    GetTimeoutError,
+    RayChannelError,
+    RayChannelTimeoutError,
+    RayTaskError,
+)
+from ray.util.annotations import PublicAPI
+
+
+def _process_return_vals(return_vals: List[Any], return_single_output: bool):
+    """
+    Process return values for return to the DAG caller. Any exceptions found in
+    return_vals will be raised. If return_single_output=True, it indicates that
+    the original DAG did not have a MultiOutputNode, so the DAG caller expects
+    a single return value instead of a list.
+    """
+    # Check for exceptions.
+    if isinstance(return_vals, Exception):
+        raise return_vals
+
+    for val in return_vals:
+        if isinstance(val, RayTaskError):
+            raise val.as_instanceof_cause()
+
+    if return_single_output:
+        assert len(return_vals) == 1
+        return return_vals[0]
+
+    return return_vals
+
+
+@PublicAPI(stability="alpha")
+class CompiledDAGRef:
+    """
+    A reference to a compiled DAG execution result.
+
+    This is a subclass of ObjectRef and resembles ObjectRef.  For example,
+    similar to ObjectRef, ray.get() can be called on it to retrieve the result.
+    However, there are several major differences:
+    1. ray.get() can only be called once per CompiledDAGRef.
+    2. ray.wait() is not supported.
+    3. CompiledDAGRef cannot be copied, deep copied, or pickled.
+    4. CompiledDAGRef cannot be passed as an argument to another task.
+    """
+
+    def __init__(
+        self,
+        dag: "ray.experimental.CompiledDAG",
+        execution_index: int,
+        channel_index: Optional[int] = None,
+    ):
+        """Initialize a CompiledDAGRef.
+
+        Args:
+            dag: The compiled DAG that generated this CompiledDAGRef.
+            execution_index: The index of the execution for the DAG.
+                A DAG can be executed multiple times, and execution index
+                indicates which execution this CompiledDAGRef corresponds to.
+            channel_index: The index of the DAG's output channel to fetch
+                the result from. A DAG can have multiple output channels, and
+                channel index indicates which channel this CompiledDAGRef
+                corresponds to. If channel index is not provided, this CompiledDAGRef
+                wraps the results from all output channels.
+        """
+        self._dag = dag
+        self._execution_index = execution_index
+        self._channel_index = channel_index
+        # Whether ray.get() was called on this CompiledDAGRef.
+        self._ray_get_called = False
+        self._dag_output_channels = dag.dag_output_channels
+
+    def __str__(self):
+        return (
+            f"CompiledDAGRef({self._dag.get_id()}, "
+            f"execution_index={self._execution_index}, "
+            f"channel_index={self._channel_index})"
+        )
+
+    def __copy__(self):
+        raise ValueError("CompiledDAGRef cannot be copied.")
+
+    def __deepcopy__(self, memo):
+        raise ValueError("CompiledDAGRef cannot be deep copied.")
+
+    def __reduce__(self):
+        raise ValueError("CompiledDAGRef cannot be pickled.")
+
+    def __del__(self):
+        # If the dag is already teardown, it should do nothing.
+        if self._dag.is_teardown:
+            return
+
+        if self._ray_get_called:
+            # get() was already called, no further cleanup is needed.
+            return
+
+        self._dag._delete_execution_results(self._execution_index, self._channel_index)
+
+    def get(self, timeout: Optional[float] = None):
+        if self._ray_get_called:
+            raise ValueError(
+                "ray.get() can only be called once "
+                "on a CompiledDAGRef, and it was already called."
+            )
+
+        self._ray_get_called = True
+        try:
+            self._dag._execute_until(
+                self._execution_index, self._channel_index, timeout
+            )
+            return_vals = self._dag._get_execution_results(
+                self._execution_index, self._channel_index
+            )
+        except RayChannelTimeoutError:
+            raise
+        except RayChannelError as channel_error:
+            # If we get a channel error, we'd like to call ray.get() on
+            # the actor execution loop refs to check if this is a result
+            # of task execution error which could not be passed down
+            # (e.g., when a pure NCCL channel is used, it is only
+            # able to send tensors, but not the wrapped exceptions).
+            # In this case, we'd like to raise the task execution error
+            # (which is the actual cause of the channel error) instead
+            # of the channel error itself.
+            # TODO(rui): determine which error to raise if multiple
+            # actor task refs have errors.
+            actor_execution_loop_refs = list(self._dag.worker_task_refs.values())
+            try:
+                ray.get(actor_execution_loop_refs, timeout=10)
+            except GetTimeoutError as timeout_error:
+                raise Exception(
+                    "Timed out when getting the actor execution loop exception. "
+                    "This should not happen, please file a GitHub issue."
+                ) from timeout_error
+            except Exception as execution_error:
+                # Use 'from None' to suppress the context of the original
+                # channel error, which is not useful to the user.
+                raise execution_error from None
+            else:
+                raise channel_error
+        except Exception:
+            raise
+        return _process_return_vals(return_vals, True)
+
+
+@PublicAPI(stability="alpha")
+class CompiledDAGFuture:
+    """
+    A reference to a compiled DAG execution result, when executed with asyncio.
+    This differs from CompiledDAGRef in that `await` must be called on the
+    future to get the result, instead of `ray.get()`.
+
+    This resembles async usage of ObjectRefs. For example, similar to
+    ObjectRef, `await` can be called directly on the CompiledDAGFuture to
+    retrieve the result.  However, there are several major differences:
+    1. `await` can only be called once per CompiledDAGFuture.
+    2. ray.wait() is not supported.
+    3. CompiledDAGFuture cannot be copied, deep copied, or pickled.
+    4. CompiledDAGFuture cannot be passed as an argument to another task.
+    """
+
+    def __init__(
+        self,
+        dag: "ray.experimental.CompiledDAG",
+        execution_index: int,
+        fut: "asyncio.Future",
+        channel_index: Optional[int] = None,
+    ):
+        self._dag = dag
+        self._execution_index = execution_index
+        self._fut = fut
+        self._channel_index = channel_index
+
+    def __str__(self):
+        return (
+            f"CompiledDAGFuture({self._dag.get_id()}, "
+            f"execution_index={self._execution_index}, "
+            f"channel_index={self._channel_index})"
+        )
+
+    def __copy__(self):
+        raise ValueError("CompiledDAGFuture cannot be copied.")
+
+    def __deepcopy__(self, memo):
+        raise ValueError("CompiledDAGFuture cannot be deep copied.")
+
+    def __reduce__(self):
+        raise ValueError("CompiledDAGFuture cannot be pickled.")
+
+    def __await__(self):
+        if self._fut is None:
+            raise ValueError(
+                "CompiledDAGFuture can only be awaited upon once, and it has "
+                "already been awaited upon."
+            )
+
+        # NOTE(swang): If the object is zero-copy deserialized, then it will
+        # stay in scope as long as this future is in scope. Therefore, we
+        # delete self._fut here before we return the result to the user.
+        fut = self._fut
+        self._fut = None
+
+        if not self._dag._has_execution_results(self._execution_index):
+            result = yield from fut.__await__()
+            self._dag._max_finished_execution_index += 1
+            self._dag._cache_execution_results(self._execution_index, result)
+
+        return_vals = self._dag._get_execution_results(
+            self._execution_index, self._channel_index
+        )
+        return _process_return_vals(return_vals, True)
+
+    def __del__(self):
+        if self._dag.is_teardown:
+            return
+
+        if self._fut is None:
+            # await() was already called, no further cleanup is needed.
+            return
+
+        self._dag._delete_execution_results(self._execution_index, self._channel_index)

python/ray/experimental/dynamic_resources.py

@@ -0,0 +1,7 @@
+def set_resource(resource_name, capacity, node_id=None):
+    raise DeprecationWarning(
+        "Dynamic custom resources are deprecated. Consider using placement "
+        "groups instead (docs.ray.io/en/master/placement-group.html). You "
+        "can also specify resources at Ray start time with the 'resources' "
+        "field in the cluster autoscaler."
+    )

python/ray/experimental/gradio_utils.py

@@ -0,0 +1,12 @@
+def type_to_string(_type: type) -> str:
+    """Gets the string representation of a type.
+
+    THe original type can be derived from the returned string representation through
+    pydoc.locate().
+    """
+    if _type.__module__ == "typing":
+        return f"{_type.__module__}.{_type._name}"
+    elif _type.__module__ == "builtins":
+        return _type.__name__
+    else:
+        return f"{_type.__module__}.{_type.__name__}"

python/ray/experimental/internal_kv.py

@@ -0,0 +1,127 @@
+from typing import List, Optional, Union
+
+from ray._private.client_mode_hook import client_mode_hook
+from ray._raylet import GcsClient
+
+_initialized = False
+global_gcs_client = None
+
+
+def _internal_kv_reset():
+    global global_gcs_client, _initialized
+    global_gcs_client = None
+    _initialized = False
+
+
+def internal_kv_get_gcs_client():
+    return global_gcs_client
+
+
+def _initialize_internal_kv(gcs_client: GcsClient):
+    """Initialize the internal KV for use in other function calls."""
+    global global_gcs_client, _initialized
+    assert gcs_client is not None
+    global_gcs_client = gcs_client
+    _initialized = True
+
+
+@client_mode_hook
+def _internal_kv_initialized():
+    return global_gcs_client is not None
+
+
+@client_mode_hook
+def _internal_kv_get(
+    key: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None
+) -> bytes:
+    """Fetch the value of a binary key."""
+
+    if isinstance(key, str):
+        key = key.encode()
+    if isinstance(namespace, str):
+        namespace = namespace.encode()
+    assert isinstance(key, bytes)
+    return global_gcs_client.internal_kv_get(key, namespace)
+
+
+@client_mode_hook
+def _internal_kv_exists(
+    key: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None
+) -> bool:
+    """Check key exists or not."""
+
+    if isinstance(key, str):
+        key = key.encode()
+    if isinstance(namespace, str):
+        namespace = namespace.encode()
+    assert isinstance(key, bytes)
+    return global_gcs_client.internal_kv_exists(key, namespace)
+
+
+@client_mode_hook
+def _pin_runtime_env_uri(uri: str, *, expiration_s: int) -> None:
+    """Pin a runtime_env URI for expiration_s."""
+    return global_gcs_client.pin_runtime_env_uri(uri, expiration_s)
+
+
+@client_mode_hook
+def _internal_kv_put(
+    key: Union[str, bytes],
+    value: Union[str, bytes],
+    overwrite: bool = True,
+    *,
+    namespace: Optional[Union[str, bytes]] = None
+) -> bool:
+    """Globally associates a value with a given binary key.
+
+    This only has an effect if the key does not already have a value.
+
+    Args:
+        key: The binary key to associate the value with.
+        value: The binary value to store under the key.
+        overwrite: Whether to overwrite an existing value for the key.
+        namespace: Optional namespace under which the key is scoped.
+
+    Returns:
+        Whether the value already exists.
+    """
+
+    if isinstance(key, str):
+        key = key.encode()
+    if isinstance(value, str):
+        value = value.encode()
+    if isinstance(namespace, str):
+        namespace = namespace.encode()
+    assert (
+        isinstance(key, bytes)
+        and isinstance(value, bytes)
+        and isinstance(overwrite, bool)
+    )
+    return global_gcs_client.internal_kv_put(key, value, overwrite, namespace) == 0
+
+
+@client_mode_hook
+def _internal_kv_del(
+    key: Union[str, bytes],
+    *,
+    del_by_prefix: bool = False,
+    namespace: Optional[Union[str, bytes]] = None
+) -> int:
+    if isinstance(key, str):
+        key = key.encode()
+    if isinstance(namespace, str):
+        namespace = namespace.encode()
+    assert isinstance(key, bytes)
+    return global_gcs_client.internal_kv_del(key, del_by_prefix, namespace)
+
+
+@client_mode_hook
+def _internal_kv_list(
+    prefix: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None
+) -> List[bytes]:
+    """List all keys in the internal KV store that start with the prefix."""
+    if isinstance(prefix, str):
+        prefix = prefix.encode()
+    if isinstance(namespace, str):
+        namespace = namespace.encode()
+    return global_gcs_client.internal_kv_keys(prefix, namespace)

python/ray/experimental/job/example_job/demo_script.py

@@ -0,0 +1,77 @@
+# Regular ray application that user wrote and runs on local cluster.
+# intermediate status are dumped to GCS
+import argparse
+import time
+
+import ray
+import ray.experimental.internal_kv as ray_kv
+
+
+@ray.remote
+class StepActor:
+    def __init__(self, interval_s=1, total_steps=3):
+        self.interval_s = interval_s
+        self.stopped = False
+        self.current_step = 1
+        self.total_steps = total_steps
+
+        worker = ray._private.worker.global_worker
+        worker_id = worker.core_worker.get_actor_id()
+        ray_kv._internal_kv_put(f"JOB:{worker_id}", self.current_step, overwrite=True)
+
+    def run(self):
+        worker = ray._private.worker.global_worker
+        worker_id = worker.core_worker.get_actor_id()
+
+        while self.current_step <= self.total_steps:
+            if not self.stopped:
+                print(
+                    f"Sleeping {self.interval_s} secs to executing "
+                    f"step {self.current_step}"
+                )
+                time.sleep(self.interval_s)
+                self.current_step += 1
+                ray_kv._internal_kv_put(
+                    f"JOB:{worker_id}", self.current_step, overwrite=True
+                )
+            else:
+                print("Stop called or reached final step.")
+                break
+
+        self.stopped = True
+        ray_kv._internal_kv_put(f"JOB:{worker_id}", "DONE", overwrite=True)
+        return "DONE"
+
+    def get_step(self):
+        return self.current_step
+
+    def stop(self):
+        self.stopped = True
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--interval-s",
+        required=False,
+        type=int,
+        default=1,
+        help="Address to use to connect to Ray",
+    )
+    parser.add_argument(
+        "--total-steps",
+        required=False,
+        type=int,
+        default=3,
+        help="Password for connecting to Redis",
+    )
+    args, _ = parser.parse_known_args()
+
+    ray.init()
+    step_actor = StepActor.remote(
+        interval_s=args.interval_s, total_steps=args.total_steps
+    )
+    ref = step_actor.run.remote()
+    print(ray.get([ref]))
+    job_key = ray_kv._internal_kv_list("JOB:")[0]
+    print(f"{job_key}, {ray_kv._internal_kv_get(job_key)}")

python/ray/experimental/job/example_job/job_config.yaml

@@ -0,0 +1,5 @@
+name: example_job
+command: python demo_script.py
+runtime_env:
+    working_dir: .
+    docker: anyscale-ml/ray-ml:nightly-py38-cpu

python/ray/experimental/locations.py

@@ -0,0 +1,76 @@
+from typing import Any, Dict, List
+
+import ray
+from ray._raylet import ObjectRef
+
+
+def get_object_locations(
+    obj_refs: List[ObjectRef], timeout_ms: int = -1
+) -> Dict[ObjectRef, Dict[str, Any]]:
+    """Lookup the locations for a list of objects.
+
+    It returns a dict maps from an object to its location. The dict excludes
+    those objects whose location lookup failed.
+
+    Args:
+        obj_refs: List of object refs.
+        timeout_ms: The maximum amount of time in micro seconds to wait
+            before returning. Wait infinitely if it's negative.
+
+    Returns:
+        A dict maps from an object to its location. The dict excludes those
+        objects whose location lookup failed.
+
+        The location is stored as a dict with following attributes:
+
+        - node_ids (List[str]): The hex IDs of the nodes that have a
+          copy of this object. Objects less than 100KB will be in memory
+          store not plasma store and therefore will have nodes_id = [].
+
+        - object_size (int): The size of data + metadata in bytes. Can be None if the
+          size is unknown yet (e.g. task not completed).
+
+    Raises:
+        RuntimeError: if the processes were not started by ray.init().
+        ray.exceptions.GetTimeoutError: if it couldn't finish the
+            request in time.
+    """
+    if not ray.is_initialized():
+        raise RuntimeError("Ray hasn't been initialized.")
+    return ray._private.worker.global_worker.core_worker.get_object_locations(
+        obj_refs, timeout_ms
+    )
+
+
+def get_local_object_locations(
+    obj_refs: List[ObjectRef],
+) -> Dict[ObjectRef, Dict[str, Any]]:
+    """Lookup the locations for a list of objects *from the local core worker*. No RPCs
+    are made in this method.
+
+    It returns a dict maps from an object to its location. The dict excludes
+    those objects whose location lookup failed.
+
+    Args:
+        obj_refs: List of object refs.
+
+    Returns:
+        A dict maps from an object to its location. The dict excludes those
+        objects whose location lookup failed.
+
+        The location is stored as a dict with following attributes:
+
+        - node_ids (List[str]): The hex IDs of the nodes that have a
+          copy of this object. Objects less than 100KB will be in memory
+          store not plasma store and therefore will have nodes_id = [].
+
+        - object_size (int): The size of data + metadata in bytes. Can be None if the
+          size is unknown yet (e.g. task not completed).
+
+    Raises:
+        RuntimeError: if the processes were not started by ray.init().
+    """
+    if not ray.is_initialized():
+        raise RuntimeError("Ray hasn't been initialized.")
+    core_worker = ray._private.worker.global_worker.core_worker
+    return core_worker.get_local_object_locations(obj_refs)

python/ray/experimental/multiprocessing/__init__.py

@@ -0,0 +1,5 @@
+from multiprocessing import TimeoutError
+
+from .pool import Pool
+
+__all__ = ["Pool", "TimeoutError"]

python/ray/experimental/multiprocessing/pool.py

@@ -0,0 +1,5 @@
+from ray.util import multiprocessing
+
+
+class Pool(multiprocessing.Pool):
+    pass  # moved to util package

python/ray/experimental/queue.py

@@ -0,0 +1,13 @@
+import warnings
+
+from ray.util.queue import Empty, Full, Queue
+
+warnings.warn(
+    DeprecationWarning(
+        "ray.experimental.queue has been moved to ray.util.queue. "
+        "Please update your import path."
+    ),
+    stacklevel=2,
+)
+
+__all__ = ["Empty", "Full", "Queue"]

python/ray/experimental/raysort/constants.py

@@ -0,0 +1,44 @@
+import os
+
+from ray.experimental.raysort.types import ByteCount, PartId, RecordCount
+
+__DIR__ = os.path.dirname(os.path.abspath(__file__))
+
+# Basics
+RECORD_SIZE = 100  # bytes
+
+# Progress Tracker Actor
+PROGRESS_TRACKER_ACTOR = "ProgressTrackerActor"
+
+# Executable locations
+GENSORT_PATH = os.path.join(__DIR__, "bin/gensort/64/gensort")
+VALSORT_PATH = os.path.join(__DIR__, "bin/gensort/64/valsort")
+
+# Filenames
+WORK_DIR = "/tmp/raysort"
+INPUT_MANIFEST_FILE = os.path.join(WORK_DIR, "input-manifest.csv")
+OUTPUT_MANIFEST_FILE = os.path.join(WORK_DIR, "output-manifest.csv")
+DATA_DIR_FMT = {
+    "input": "{mnt}/tmp/input/",
+    "output": "{mnt}/tmp/output/",
+    "temp": "{mnt}/tmp/temp/",
+}
+FILENAME_FMT = {
+    "input": "input-{part_id:08}",
+    "output": "output-{part_id:08}",
+    "temp": "temp-{part_id:08}",
+}
+
+# Prometheus config
+PROM_RAY_EXPORTER_PORT = 8090
+PROM_NODE_EXPORTER_PORT = 8091
+
+
+# Convenience functions
+def bytes_to_records(n_bytes: ByteCount) -> RecordCount:
+    assert n_bytes % RECORD_SIZE == 0
+    return int(n_bytes / RECORD_SIZE)
+
+
+def merge_part_ids(reducer_id: PartId, mapper_id: PartId) -> PartId:
+    return reducer_id * 1_000_000 + mapper_id

python/ray/experimental/raysort/logging_utils.py

@@ -0,0 +1,9 @@
+import logging
+
+
+def init():
+    fmt = "%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
+    logging.basicConfig(
+        format=fmt,
+        level=logging.INFO,
+    )

python/ray/experimental/raysort/main.py

@@ -0,0 +1,475 @@
+import argparse
+import contextlib
+import csv
+import logging
+import os
+import random
+import subprocess
+import tempfile
+from typing import Callable, Dict, Iterable, List
+
+import numpy as np
+
+import ray
+from ray.experimental.raysort import constants, logging_utils, sortlib, tracing_utils
+from ray.experimental.raysort.types import (
+    BlockInfo,
+    ByteCount,
+    PartId,
+    PartInfo,
+    Path,
+    RecordCount,
+)
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+Args = argparse.Namespace
+
+# ------------------------------------------------------------
+#     Parse Arguments
+# ------------------------------------------------------------
+
+
+def get_args(*args, **kwargs):
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--ray_address",
+        default="auto",
+        type=str,
+        help="if set to None, will launch a local Ray cluster",
+    )
+    parser.add_argument(
+        "--total_data_size",
+        default=1 * 1000 * 1024 * 1024 * 1024,
+        type=ByteCount,
+        help="total data size in bytes",
+    )
+    parser.add_argument(
+        "--num_mappers",
+        default=256,
+        type=int,
+        help="number of map tasks",
+    )
+    parser.add_argument(
+        "--num_mappers_per_round",
+        default=16,
+        type=int,
+        help="number of map tasks per first-stage merge tasks",
+    )
+    parser.add_argument(
+        "--num_reducers",
+        default=16,
+        type=int,
+        help="number of second-stage reduce tasks",
+    )
+    parser.add_argument(
+        "--num_concurrent_rounds",
+        default=4,
+        type=int,
+        help="max number of rounds of map/merge tasks in flight",
+    )
+    parser.add_argument(
+        "--reducer_input_chunk",
+        default=100 * 1024 * 1024,
+        type=ByteCount,
+        help="bytes to read from each file in reduce tasks",
+    )
+    parser.add_argument(
+        "--skip_sorting",
+        default=False,
+        action="store_true",
+        help="if set, no sorting is actually performed",
+    )
+    parser.add_argument(
+        "--skip_input",
+        default=False,
+        action="store_true",
+        help="if set, mappers will not read data from disk",
+    )
+    parser.add_argument(
+        "--skip_output",
+        default=False,
+        action="store_true",
+        help="if set, reducers will not write out results to disk",
+    )
+    # Which tasks to run?
+    tasks_group = parser.add_argument_group(
+        "tasks to run", "if no task is specified, will run all tasks"
+    )
+    tasks = ["generate_input", "sort", "validate_output"]
+    for task in tasks:
+        tasks_group.add_argument(f"--{task}", action="store_true")
+
+    args = parser.parse_args(*args, **kwargs)
+    # Derive additional arguments.
+    args.input_part_size = ByteCount(args.total_data_size / args.num_mappers)
+    assert args.num_mappers % args.num_mappers_per_round == 0
+    args.num_rounds = int(args.num_mappers / args.num_mappers_per_round)
+    args.mount_points = _get_mount_points()
+    # If no tasks are specified, run all tasks.
+    args_dict = vars(args)
+    if not any(args_dict[task] for task in tasks):
+        for task in tasks:
+            args_dict[task] = True
+    return args
+
+
+def _get_mount_points():
+    default_ret = [tempfile.gettempdir()]
+    mnt = "/mnt"
+    if os.path.exists(mnt):
+        ret = [os.path.join(mnt, d) for d in os.listdir(mnt)]
+        if len(ret) > 0:
+            return ret
+    return default_ret
+
+
+# ------------------------------------------------------------
+#     Generate Input
+# ------------------------------------------------------------
+
+
+def _part_info(args: Args, part_id: PartId, kind="input") -> PartInfo:
+    node = ray._private.worker.global_worker.node_ip_address
+    mnt = random.choice(args.mount_points)
+    filepath = _get_part_path(mnt, part_id, kind)
+    return PartInfo(part_id, node, filepath)
+
+
+def _get_part_path(mnt: Path, part_id: PartId, kind="input") -> Path:
+    assert kind in {"input", "output", "temp"}
+    dir_fmt = constants.DATA_DIR_FMT[kind]
+    dirpath = dir_fmt.format(mnt=mnt)
+    os.makedirs(dirpath, exist_ok=True)
+    filename_fmt = constants.FILENAME_FMT[kind]
+    filename = filename_fmt.format(part_id=part_id)
+    filepath = os.path.join(dirpath, filename)
+    return filepath
+
+
+@ray.remote
+def generate_part(
+    args: Args, part_id: PartId, size: RecordCount, offset: RecordCount
+) -> PartInfo:
+    logging_utils.init()
+    pinfo = _part_info(args, part_id)
+    subprocess.run(
+        [constants.GENSORT_PATH, f"-b{offset}", f"{size}", pinfo.path], check=True
+    )
+    logging.info(f"Generated input {pinfo}")
+    return pinfo
+
+
+def generate_input(args: Args):
+    if args.skip_input:
+        return
+    size = constants.bytes_to_records(args.input_part_size)
+    offset = 0
+    tasks = []
+    for part_id in range(args.num_mappers):
+        tasks.append(generate_part.remote(args, part_id, size, offset))
+        offset += size
+    assert offset == constants.bytes_to_records(args.total_data_size), args
+    logging.info(f"Generating {len(tasks)} partitions")
+    parts = ray.get(tasks)
+    with open(constants.INPUT_MANIFEST_FILE, "w") as fout:
+        writer = csv.writer(fout)
+        writer.writerows(parts)
+
+
+# ------------------------------------------------------------
+#     Sort
+# ------------------------------------------------------------
+
+
+def _load_manifest(args: Args, path: Path) -> List[PartInfo]:
+    if args.skip_input:
+        return [PartInfo(i, None, None) for i in range(args.num_mappers)]
+    with open(path) as fin:
+        reader = csv.reader(fin)
+        return [PartInfo(int(part_id), node, path) for part_id, node, path in reader]
+
+
+def _load_partition(args: Args, path: Path) -> np.ndarray:
+    if args.skip_input:
+        return np.frombuffer(
+            np.random.bytes(args.input_part_size), dtype=np.uint8
+        ).copy()
+    return np.fromfile(path, dtype=np.uint8)
+
+
+def _dummy_sort_and_partition(
+    part: np.ndarray, boundaries: List[int]
+) -> List[BlockInfo]:
+    N = len(boundaries)
+    offset = 0
+    size = int(np.ceil(part.size / N))
+    blocks = []
+    for _ in range(N):
+        blocks.append((offset, size))
+        offset += size
+    return blocks
+
+
+@ray.remote
+@tracing_utils.timeit("map")
+def mapper(
+    args: Args, mapper_id: PartId, boundaries: List[int], path: Path
+) -> List[np.ndarray]:
+    logging_utils.init()
+    part = _load_partition(args, path)
+    sort_fn = (
+        _dummy_sort_and_partition if args.skip_sorting else sortlib.sort_and_partition
+    )
+    blocks = sort_fn(part, boundaries)
+    return [part[offset : offset + size] for offset, size in blocks]
+
+
+def _dummy_merge(
+    num_blocks: int, _n: int, get_block: Callable[[int, int], np.ndarray]
+) -> Iterable[np.ndarray]:
+    blocks = [((i, 0), get_block(i, 0)) for i in range(num_blocks)]
+    while len(blocks) > 0:
+        (m, d), block = blocks.pop(random.randrange(len(blocks)))
+        yield block
+        d_ = d + 1
+        block = get_block(m, d_)
+        if block is None:
+            continue
+        blocks.append(((m, d_), block))
+
+
+def _merge_impl(
+    args: Args,
+    M: int,
+    pinfo: PartInfo,
+    get_block: Callable[[int, int], np.ndarray],
+    skip_output=False,
+):
+    merge_fn = _dummy_merge if args.skip_sorting else sortlib.merge_partitions
+    merger = merge_fn(M, get_block)
+
+    if skip_output:
+        for datachunk in merger:
+            del datachunk
+    else:
+        with open(pinfo.path, "wb") as fout:
+            for datachunk in merger:
+                fout.write(datachunk)
+    return pinfo
+
+
+# See worker_placement_groups() for why `num_cpus=0`.
+@ray.remote(num_cpus=0, resources={"worker": 1})
+@tracing_utils.timeit("merge")
+def merge_mapper_blocks(
+    args: Args, reducer_id: PartId, mapper_id: PartId, *blocks: List[np.ndarray]
+) -> PartInfo:
+    part_id = constants.merge_part_ids(reducer_id, mapper_id)
+    pinfo = _part_info(args, part_id, kind="temp")
+    M = len(blocks)
+
+    def get_block(i, d):
+        if i >= M or d > 0:
+            return None
+        return blocks[i]
+
+    return _merge_impl(args, M, pinfo, get_block)
+
+
+# See worker_placement_groups() for why `num_cpus=0`.
+@ray.remote(num_cpus=0, resources={"worker": 1})
+@tracing_utils.timeit("reduce")
+def final_merge(
+    args: Args, reducer_id: PartId, *merged_parts: List[PartInfo]
+) -> PartInfo:
+    M = len(merged_parts)
+
+    def _load_block_chunk(pinfo: PartInfo, d: int) -> np.ndarray:
+        return np.fromfile(
+            pinfo.path,
+            dtype=np.uint8,
+            count=args.reducer_input_chunk,
+            offset=d * args.reducer_input_chunk,
+        )
+
+    def get_block(i, d):
+        ret = _load_block_chunk(merged_parts[i], d)
+        if ret.size == 0:
+            return None
+        return ret
+
+    pinfo = _part_info(args, reducer_id, "output")
+    return _merge_impl(args, M, pinfo, get_block, args.skip_output)
+
+
+def _node_res(node: str) -> Dict[str, float]:
+    return {"resources": {f"node:{node}": 1e-3}}
+
+
+@contextlib.contextmanager
+def worker_placement_groups(args: Args) -> List[ray.PlacementGroupID]:
+    """
+    Returns one placement group per node with a `worker` resource. To run
+    tasks in the placement group, use
+    `@ray.remote(num_cpus=0, resources={"worker": 1})`. Ray does not
+    automatically reserve CPU resources, so tasks must specify `num_cpus=0`
+    in order to run in a placement group.
+    """
+    pgs = [ray.util.placement_group([{"worker": 1}]) for _ in range(args.num_reducers)]
+    ray.get([pg.ready() for pg in pgs])
+    try:
+        yield pgs
+    finally:
+        for pg in pgs:
+            ray.util.remove_placement_group(pg)
+
+
+@tracing_utils.timeit("sort", report_time=True)
+def sort_main(args: Args):
+    parts = _load_manifest(args, constants.INPUT_MANIFEST_FILE)
+    assert len(parts) == args.num_mappers
+    boundaries = sortlib.get_boundaries(args.num_reducers)
+
+    # The exception of 'ValueError("Resource quantities >1 must be whole numbers.")'
+    # will be raised if the `num_cpus` > 1 and not an integer.
+    num_cpus = os.cpu_count() / args.num_concurrent_rounds
+    if num_cpus > 1.0:
+        num_cpus = int(num_cpus)
+    mapper_opt = {
+        "num_returns": args.num_reducers,
+        "num_cpus": num_cpus,
+    }  # Load balance across worker nodes by setting `num_cpus`.
+    merge_results = np.empty((args.num_rounds, args.num_reducers), dtype=object)
+
+    part_id = 0
+    with worker_placement_groups(args) as pgs:
+        for round in range(args.num_rounds):
+            # Limit the number of in-flight rounds.
+            num_extra_rounds = round - args.num_concurrent_rounds + 1
+            if num_extra_rounds > 0:
+                ray.wait(
+                    [f for f in merge_results.flatten() if f is not None],
+                    num_returns=num_extra_rounds * args.num_reducers,
+                )
+
+            # Submit map tasks.
+            mapper_results = np.empty(
+                (args.num_mappers_per_round, args.num_reducers), dtype=object
+            )
+            for _ in range(args.num_mappers_per_round):
+                _, node, path = parts[part_id]
+                m = part_id % args.num_mappers_per_round
+                mapper_results[m, :] = mapper.options(**mapper_opt).remote(
+                    args, part_id, boundaries, path
+                )
+                part_id += 1
+
+            # Submit merge tasks.
+            merge_results[round, :] = [
+                merge_mapper_blocks.options(
+                    scheduling_strategy=PlacementGroupSchedulingStrategy(
+                        placement_group=pgs[r]
+                    )
+                ).remote(args, r, round, *mapper_results[:, r].tolist())
+                for r in range(args.num_reducers)
+            ]
+
+            # Delete local references to mapper results.
+            mapper_results = None
+
+        # Submit second-stage reduce tasks.
+        reducer_results = [
+            final_merge.options(
+                scheduling_strategy=PlacementGroupSchedulingStrategy(
+                    placement_group=pgs[r]
+                )
+            ).remote(args, r, *merge_results[:, r].tolist())
+            for r in range(args.num_reducers)
+        ]
+        reducer_results = ray.get(reducer_results)
+
+    if not args.skip_output:
+        with open(constants.OUTPUT_MANIFEST_FILE, "w") as fout:
+            writer = csv.writer(fout)
+            writer.writerows(reducer_results)
+
+    logging.info(ray._private.internal_api.memory_summary(stats_only=True))
+
+
+# ------------------------------------------------------------
+#     Validate Output
+# ------------------------------------------------------------
+
+
+def _run_valsort(args: List[str]):
+    proc = subprocess.run([constants.VALSORT_PATH] + args, capture_output=True)
+    if proc.returncode != 0:
+        logging.critical("\n" + proc.stderr.decode("ascii"))
+        raise RuntimeError(f"Validation failed: {args}")
+
+
+@ray.remote
+def validate_part(path: Path):
+    logging_utils.init()
+    sum_path = path + ".sum"
+    _run_valsort(["-o", sum_path, path])
+    logging.info(f"Validated output {path}")
+    with open(sum_path, "rb") as fin:
+        return os.path.getsize(path), fin.read()
+
+
+def validate_output(args: Args):
+    if args.skip_sorting or args.skip_output:
+        return
+    partitions = _load_manifest(args, constants.OUTPUT_MANIFEST_FILE)
+    results = []
+    for _, node, path in partitions:
+        results.append(validate_part.options(**_node_res(node)).remote(path))
+    logging.info(f"Validating {len(results)} partitions")
+    results = ray.get(results)
+    total = sum(s for s, _ in results)
+    assert total == args.total_data_size, total - args.total_data_size
+    all_checksum = b"".join(c for _, c in results)
+    with tempfile.NamedTemporaryFile() as fout:
+        fout.write(all_checksum)
+        fout.flush()
+        _run_valsort(["-s", fout.name])
+    logging.info("All OK!")
+
+
+# ------------------------------------------------------------
+#     Main
+# ------------------------------------------------------------
+
+
+def init(args: Args):
+    if not args.ray_address:
+        ray.init(resources={"worker": os.cpu_count()})
+    else:
+        ray.init(address=args.ray_address)
+    logging_utils.init()
+    logging.info(args)
+    os.makedirs(constants.WORK_DIR, exist_ok=True)
+    resources = ray.cluster_resources()
+    logging.info(resources)
+    args.num_workers = resources["worker"]
+    progress_tracker = tracing_utils.create_progress_tracker(args)
+    return progress_tracker
+
+
+def main(args: Args):
+    # Keep the actor handle in scope for the duration of the program.
+    _progress_tracker = init(args)  # noqa F841
+
+    if args.generate_input:
+        generate_input(args)
+
+    if args.sort:
+        sort_main(args)
+
+    if args.validate_output:
+        validate_output(args)
+
+
+if __name__ == "__main__":
+    main(get_args())

python/ray/experimental/raysort/sortlib.py

@@ -0,0 +1,28 @@
+from typing import Callable, Iterable, List
+
+import numpy as np
+
+from ray.experimental.raysort.types import BlockInfo
+
+
+def get_boundaries(num_parts: int) -> List[int]:
+    return [0] * num_parts
+
+
+def sort_and_partition(part: np.ndarray, boundaries: List[int]) -> List[BlockInfo]:
+    N = len(boundaries)
+    offset = 0
+    size = int(np.ceil(part.size / N))
+    blocks = []
+    for _ in range(N):
+        blocks.append((offset, size))
+        offset += size
+    return blocks
+
+
+def merge_partitions(
+    num_blocks: int, get_block: Callable[[int, int], np.ndarray]
+) -> Iterable[memoryview]:
+    blocks = [get_block(i, 0) for i in range(num_blocks)]
+    for block in blocks:
+        yield block

python/ray/experimental/raysort/tracing_utils.py

@@ -0,0 +1,120 @@
+import datetime
+import functools
+import logging
+import os
+import time
+from typing import List, Tuple
+
+import ray
+from ray.experimental.raysort import constants, logging_utils
+from ray.util.metrics import Gauge, Histogram
+
+HISTOGRAM_BOUNDARIES = list(range(50, 200, 50))
+
+
+def timeit(
+    event: str,
+    report_time=False,
+    report_in_progress=True,
+    report_completed=True,
+):
+    def decorator(f):
+        @functools.wraps(f)
+        def wrapped_f(*args, **kwargs):
+            progress_tracker = ray.get_actor(constants.PROGRESS_TRACKER_ACTOR)
+            progress_tracker.inc.remote(f"{event}_in_progress", echo=report_in_progress)
+            try:
+                start = time.time()
+                ret = f(*args, **kwargs)
+                end = time.time()
+                duration = end - start
+                progress_tracker.observe.remote(
+                    f"{event}_time",
+                    duration,
+                    echo=report_time,
+                )
+                progress_tracker.inc.remote(f"{event}_completed", echo=report_completed)
+                return ret
+            finally:
+                progress_tracker.dec.remote(f"{event}_in_progress")
+
+        return wrapped_f
+
+    return decorator
+
+
+def get_metrics(_args):
+    return {
+        "gauges": [
+            "map_in_progress",
+            "merge_in_progress",
+            "reduce_in_progress",
+            "sort_in_progress",
+            "map_completed",
+            "merge_completed",
+            "reduce_completed",
+            "sort_completed",
+        ],
+        "histograms": [
+            ("map_time", HISTOGRAM_BOUNDARIES),
+            ("merge_time", HISTOGRAM_BOUNDARIES),
+            ("reduce_time", HISTOGRAM_BOUNDARIES),
+            ("sort_time", HISTOGRAM_BOUNDARIES),
+        ],
+    }
+
+
+def create_progress_tracker(args):
+    return ProgressTracker.options(name=constants.PROGRESS_TRACKER_ACTOR).remote(
+        **get_metrics(args)
+    )
+
+
+@ray.remote
+class ProgressTracker:
+    def __init__(
+        self,
+        gauges: List[str],
+        histograms: List[Tuple[str, List[int]]],
+    ):
+        self.counts = {m: 0 for m in gauges}
+        self.gauges = {m: Gauge(m) for m in gauges}
+        self.reset_gauges()
+        self.histograms = {m: Histogram(m, boundaries=b) for m, b in histograms}
+        logging_utils.init()
+
+    def reset_gauges(self):
+        for g in self.gauges.values():
+            g.set(0)
+
+    def inc(self, metric_name, value=1, echo=False):
+        gauge = self.gauges.get(metric_name)
+        if gauge is None:
+            logging.warning(f"No such Gauge: {metric_name}")
+            return
+        self.counts[metric_name] += value
+        gauge.set(self.counts[metric_name])
+        if echo:
+            logging.info(f"{metric_name} {self.counts[metric_name]}")
+
+    def dec(self, metric_name, value=1, echo=False):
+        return self.inc(metric_name, -value, echo)
+
+    def observe(self, metric_name, value, echo=False):
+        histogram = self.histograms.get(metric_name)
+        if histogram is None:
+            logging.warning(f"No such Histogram: {metric_name}")
+            return
+        histogram.observe(value)
+        if echo:
+            logging.info(f"{metric_name} {value}")
+
+
+def export_timeline():
+    timestr = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    filename = f"ray-timeline-{timestr}.json"
+
+    temp_dir = ray.get_runtime_context().get_temp_dir()
+    filepath = os.path.join(temp_dir, filename)
+    ray.timeline(filename=filepath)
+    logging.info(f"Exported Ray timeline to {filepath}")

python/ray/experimental/raysort/types.py

@@ -0,0 +1,18 @@
+from typing import NamedTuple, Tuple
+
+ByteCount = int
+NodeAddress = str
+PartId = int
+Path = str
+RecordCount = int
+
+BlockInfo = Tuple[int, int]
+
+
+class PartInfo(NamedTuple):
+    part_id: PartId
+    node: NodeAddress
+    path: Path
+
+    def __repr__(self):
+        return f"Part({self.node}:{self.path})"

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

python/ray/experimental/shuffle.py

@@ -0,0 +1,358 @@
+"""A simple distributed shuffle implementation in Ray.
+
+This utility provides a `simple_shuffle` function that can be used to
+redistribute M input partitions into N output partitions. It does this with
+a single wave of shuffle map tasks followed by a single wave of shuffle reduce
+tasks. Each shuffle map task generates O(N) output objects, and each shuffle
+reduce task consumes O(M) input objects, for a total of O(N*M) objects.
+
+To try an example 10GB shuffle, run:
+
+    $ python -m ray.experimental.shuffle \
+        --num-partitions=50 --partition-size=200e6 \
+        --object-store-memory=1e9
+
+This will print out some statistics on the shuffle execution such as:
+
+    --- Aggregate object store stats across all nodes ---
+    Plasma memory usage 0 MiB, 0 objects, 0.0% full
+    Spilled 9487 MiB, 2487 objects, avg write throughput 1023 MiB/s
+    Restored 9487 MiB, 2487 objects, avg read throughput 1358 MiB/s
+    Objects consumed by Ray tasks: 9537 MiB.
+
+    Shuffled 9536 MiB in 16.579771757125854 seconds
+"""
+import time
+from typing import Any, Callable, Iterable, List, Tuple, Union
+
+import ray
+from ray import ObjectRef
+from ray.cluster_utils import Cluster
+
+# TODO(ekl) why doesn't TypeVar() deserialize properly in Ray?
+# The type produced by the input reader function.
+InType = Any
+# The type produced by the output writer function.
+OutType = Any
+# Integer identifying the partition number.
+PartitionID = int
+
+
+class ObjectStoreWriter:
+    """This class is used to stream shuffle map outputs to the object store.
+
+    It can be subclassed to optimize writing (e.g., batching together small
+    records into larger objects). This will be performance critical if your
+    input records are small (the example shuffle uses very large records, so
+    the naive strategy works well).
+    """
+
+    def __init__(self):
+        self.results = []
+
+    def add(self, item: InType) -> None:
+        """Queue a single item to be written to the object store.
+
+        This base implementation immediately writes each given item to the
+        object store as a standalone object.
+        """
+        self.results.append(ray.put(item))
+
+    def finish(self) -> List[ObjectRef]:
+        """Return list of object refs representing written items."""
+        return self.results
+
+
+class ObjectStoreWriterNonStreaming(ObjectStoreWriter):
+    def __init__(self):
+        self.results = []
+
+    def add(self, item: InType) -> None:
+        self.results.append(item)
+
+    def finish(self) -> List[Any]:
+        return self.results
+
+
+def round_robin_partitioner(
+    input_stream: Iterable[InType], num_partitions: int
+) -> Iterable[Tuple[PartitionID, InType]]:
+    """Round robin partitions items from the input reader.
+
+    You can write custom partitioning functions for your use case.
+
+    Args:
+        input_stream: Iterator over items from the input reader.
+        num_partitions: Number of output partitions.
+
+    Yields:
+        Tuple[PartitionID, InType]: A tuple of partition id and the
+            corresponding input item.
+    """
+    i = 0
+    for item in input_stream:
+        yield (i, item)
+        i += 1
+        i %= num_partitions
+
+
+@ray.remote
+class _StatusTracker:
+    def __init__(self):
+        self.num_map = 0
+        self.num_reduce = 0
+        self.map_refs = []
+        self.reduce_refs = []
+
+    def register_objectrefs(self, map_refs, reduce_refs):
+        self.map_refs = map_refs
+        self.reduce_refs = reduce_refs
+
+    def get_progress(self):
+        if self.map_refs:
+            ready, self.map_refs = ray.wait(
+                self.map_refs,
+                timeout=1,
+                num_returns=len(self.map_refs),
+                fetch_local=False,
+            )
+            self.num_map += len(ready)
+        elif self.reduce_refs:
+            ready, self.reduce_refs = ray.wait(
+                self.reduce_refs,
+                timeout=1,
+                num_returns=len(self.reduce_refs),
+                fetch_local=False,
+            )
+            self.num_reduce += len(ready)
+        return self.num_map, self.num_reduce
+
+
+def render_progress_bar(tracker, input_num_partitions, output_num_partitions):
+    from tqdm import tqdm
+
+    num_map = 0
+    num_reduce = 0
+    map_bar = tqdm(total=input_num_partitions, position=0)
+    map_bar.set_description("Map Progress.")
+    reduce_bar = tqdm(total=output_num_partitions, position=1)
+    reduce_bar.set_description("Reduce Progress.")
+
+    while num_map < input_num_partitions or num_reduce < output_num_partitions:
+        new_num_map, new_num_reduce = ray.get(tracker.get_progress.remote())
+        map_bar.update(new_num_map - num_map)
+        reduce_bar.update(new_num_reduce - num_reduce)
+        num_map = new_num_map
+        num_reduce = new_num_reduce
+        time.sleep(0.1)
+    map_bar.close()
+    reduce_bar.close()
+
+
+def simple_shuffle(
+    *,
+    input_reader: Callable[[PartitionID], Iterable[InType]],
+    input_num_partitions: int,
+    output_num_partitions: int,
+    output_writer: Callable[[PartitionID, List[Union[ObjectRef, Any]]], OutType],
+    partitioner: Callable[
+        [Iterable[InType], int], Iterable[PartitionID]
+    ] = round_robin_partitioner,
+    object_store_writer: ObjectStoreWriter = ObjectStoreWriter,
+    tracker: _StatusTracker = None,
+    streaming: bool = True,
+) -> List[OutType]:
+    """Simple distributed shuffle in Ray.
+
+    Args:
+        input_reader: Function that generates the input items for a
+            partition (e.g., data records).
+        input_num_partitions: The number of input partitions.
+        output_num_partitions: The desired number of output partitions.
+        output_writer: Function that consumes a iterator of items for a
+            given output partition. It returns a single value that will be
+            collected across all output partitions.
+        partitioner: Partitioning function to use. Defaults to round-robin
+            partitioning of input items.
+        object_store_writer: Class used to write input items to the
+            object store in an efficient way. Defaults to a naive
+            implementation that writes each input record as one object.
+        tracker: Tracker actor that is used to display the progress bar.
+        streaming: Whether or not if the shuffle will be streaming.
+
+    Returns:
+        List of outputs from the output writers.
+    """
+
+    @ray.remote(num_returns=output_num_partitions)
+    def shuffle_map(i: PartitionID) -> List[List[Union[Any, ObjectRef]]]:
+        writers = [object_store_writer() for _ in range(output_num_partitions)]
+        for out_i, item in partitioner(input_reader(i), output_num_partitions):
+            writers[out_i].add(item)
+        return [c.finish() for c in writers]
+
+    @ray.remote
+    def shuffle_reduce(
+        i: PartitionID, *mapper_outputs: List[List[Union[Any, ObjectRef]]]
+    ) -> OutType:
+        input_objects = []
+        assert len(mapper_outputs) == input_num_partitions
+        for obj_refs in mapper_outputs:
+            for obj_ref in obj_refs:
+                input_objects.append(obj_ref)
+        return output_writer(i, input_objects)
+
+    shuffle_map_out = [shuffle_map.remote(i) for i in range(input_num_partitions)]
+
+    shuffle_reduce_out = [
+        shuffle_reduce.remote(
+            j, *[shuffle_map_out[i][j] for i in range(input_num_partitions)]
+        )
+        for j in range(output_num_partitions)
+    ]
+
+    if tracker:
+        tracker.register_objectrefs.remote(
+            [map_out[0] for map_out in shuffle_map_out], shuffle_reduce_out
+        )
+        render_progress_bar(tracker, input_num_partitions, output_num_partitions)
+
+    return ray.get(shuffle_reduce_out)
+
+
+def build_cluster(num_nodes, num_cpus, object_store_memory):
+    cluster = Cluster()
+    for _ in range(num_nodes):
+        cluster.add_node(num_cpus=num_cpus, object_store_memory=object_store_memory)
+    cluster.wait_for_nodes()
+    return cluster
+
+
+def run(
+    ray_address=None,
+    object_store_memory=1e9,
+    num_partitions=5,
+    partition_size=200e6,
+    num_nodes=None,
+    num_cpus=8,
+    no_streaming=False,
+    use_wait=False,
+    tracker=None,
+):
+    import time
+
+    import numpy as np
+
+    is_multi_node = num_nodes
+    if ray_address:
+        print("Connecting to a existing cluster...")
+        ray.init(address=ray_address, ignore_reinit_error=True)
+    elif is_multi_node:
+        print("Emulating a cluster...")
+        print(f"Num nodes: {num_nodes}")
+        print(f"Num CPU per node: {num_cpus}")
+        print(f"Object store memory per node: {object_store_memory}")
+        cluster = build_cluster(num_nodes, num_cpus, object_store_memory)
+        ray.init(address=cluster.address)
+    else:
+        print("Start a new cluster...")
+        ray.init(num_cpus=num_cpus, object_store_memory=object_store_memory)
+
+    partition_size = int(partition_size)
+    num_partitions = num_partitions
+    rows_per_partition = partition_size // (8 * 2)
+    if tracker is None:
+        tracker = _StatusTracker.remote()
+    use_wait = use_wait
+
+    def input_reader(i: PartitionID) -> Iterable[InType]:
+        for _ in range(num_partitions):
+            yield np.ones((rows_per_partition // num_partitions, 2), dtype=np.int64)
+
+    def output_writer(i: PartitionID, shuffle_inputs: List[ObjectRef]) -> OutType:
+        total = 0
+        if not use_wait:
+            for obj_ref in shuffle_inputs:
+                arr = ray.get(obj_ref)
+                total += arr.size * arr.itemsize
+        else:
+            while shuffle_inputs:
+                [ready], shuffle_inputs = ray.wait(shuffle_inputs, num_returns=1)
+                arr = ray.get(ready)
+                total += arr.size * arr.itemsize
+
+        return total
+
+    def output_writer_non_streaming(
+        i: PartitionID, shuffle_inputs: List[Any]
+    ) -> OutType:
+        total = 0
+        for arr in shuffle_inputs:
+            total += arr.size * arr.itemsize
+        return total
+
+    if no_streaming:
+        output_writer_callable = output_writer_non_streaming
+        object_store_writer = ObjectStoreWriterNonStreaming
+    else:
+        object_store_writer = ObjectStoreWriter
+        output_writer_callable = output_writer
+
+    start = time.time()
+    output_sizes = simple_shuffle(
+        input_reader=input_reader,
+        input_num_partitions=num_partitions,
+        output_num_partitions=num_partitions,
+        output_writer=output_writer_callable,
+        object_store_writer=object_store_writer,
+        tracker=tracker,
+    )
+    delta = time.time() - start
+
+    time.sleep(0.5)
+    print()
+
+    summary = None
+    for i in range(5):
+        try:
+            summary = ray._private.internal_api.memory_summary(stats_only=True)
+        except Exception:
+            time.sleep(1)
+            pass
+        if summary:
+            break
+    print(summary)
+    print()
+    print(
+        "Shuffled", int(sum(output_sizes) / (1024 * 1024)), "MiB in", delta, "seconds"
+    )
+
+
+def main():
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ray-address", type=str, default=None)
+    parser.add_argument("--object-store-memory", type=float, default=1e9)
+    parser.add_argument("--num-partitions", type=int, default=5)
+    parser.add_argument("--partition-size", type=float, default=200e6)
+    parser.add_argument("--num-nodes", type=int, default=None)
+    parser.add_argument("--num-cpus", type=int, default=8)
+    parser.add_argument("--no-streaming", action="store_true", default=False)
+    parser.add_argument("--use-wait", action="store_true", default=False)
+    args = parser.parse_args()
+
+    run(
+        ray_address=args.ray_address,
+        object_store_memory=args.object_store_memory,
+        num_partitions=args.num_partitions,
+        partition_size=args.partition_size,
+        num_nodes=args.num_nodes,
+        num_cpus=args.num_cpus,
+        no_streaming=args.no_streaming,
+        use_wait=args.use_wait,
+    )
+
+
+if __name__ == "__main__":
+    main()

python/ray/experimental/state/api.py

@@ -0,0 +1,4 @@
+from ray.util.state import *  # noqa: F401 F403
+from ray.util.state.util import record_deprecated_state_api_import
+
+record_deprecated_state_api_import()

python/ray/experimental/state/common.py

@@ -0,0 +1,4 @@
+from ray.util.state.common import *  # noqa: F401 F403
+from ray.util.state.util import record_deprecated_state_api_import
+
+record_deprecated_state_api_import()

python/ray/experimental/state/custom_types.py

@@ -0,0 +1,4 @@
+from ray._private.custom_types import *  # noqa: F401 F403
+from ray.util.state.util import record_deprecated_state_api_import
+
+record_deprecated_state_api_import()

python/ray/experimental/state/exception.py

@@ -0,0 +1,4 @@
+from ray.util.state.exception import *  # noqa: F401 F403
+from ray.util.state.util import record_deprecated_state_api_import
+
+record_deprecated_state_api_import()

python/ray/experimental/state/state_cli.py

@@ -0,0 +1,4 @@
+from ray.util.state.state_cli import *  # noqa: F401 F403
+from ray.util.state.util import record_deprecated_state_api_import
+
+record_deprecated_state_api_import()

python/ray/experimental/state/state_manager.py

@@ -0,0 +1,4 @@
+from ray.util.state.state_manager import *  # noqa: F401 F403
+from ray.util.state.util import record_deprecated_state_api_import
+
+record_deprecated_state_api_import()

python/ray/experimental/state/util.py

@@ -0,0 +1,4 @@
+from ray.util.state.util import *  # noqa: F401 F403
+from ray.util.state.util import record_deprecated_state_api_import
+
+record_deprecated_state_api_import()

python/ray/experimental/tf_utils.py

@@ -0,0 +1,4 @@
+raise ImportError(
+    "ray.experimental.tf_utils has been removed. "
+    "Use: from ray.rllib.utils import tf_utils."
+)

python/ray/experimental/tqdm_ray.py

@@ -0,0 +1,425 @@
+import builtins
+import copy
+import json
+import logging
+import os
+import sys
+import threading
+import time
+import uuid
+from typing import Any, Dict, Iterable, Optional
+
+import colorama
+
+import ray
+from ray._private.ray_constants import env_bool
+from ray.util.debug import log_once
+
+try:
+    import tqdm.auto as real_tqdm
+except ImportError:
+    real_tqdm = None
+
+logger = logging.getLogger(__name__)
+
+# Describes the state of a single progress bar.
+ProgressBarState = Dict[str, Any]
+
+# Magic token used to identify Ray TQDM log lines.
+RAY_TQDM_MAGIC = "__ray_tqdm_magic_token__"
+
+# Global manager singleton.
+_manager: Optional["_BarManager"] = None
+_mgr_lock = threading.Lock()
+_print = builtins.print
+
+
+def safe_print(*args, **kwargs):
+    """Use this as an alternative to `print` that will not corrupt tqdm output.
+
+    By default, the builtin print will be patched to this function when tqdm_ray is
+    used. To disable this, set RAY_TQDM_PATCH_PRINT=0.
+    """
+
+    # Ignore prints to StringIO objects, etc.
+    if kwargs.get("file") not in [sys.stdout, sys.stderr, None]:
+        return _print(*args, **kwargs)
+
+    try:
+        instance().hide_bars()
+        _print(*args, **kwargs)
+    finally:
+        instance().unhide_bars()
+
+
+class tqdm:
+    """Experimental: Ray distributed tqdm implementation.
+
+    This class lets you use tqdm from any Ray remote task or actor, and have the
+    progress centrally reported from the driver. This avoids issues with overlapping
+    / conflicting progress bars, as the driver centrally manages tqdm positions.
+
+    Supports a limited subset of tqdm args.
+    """
+
+    DEFAULT_FLUSH_INTERVAL_SECONDS = 1.0
+
+    def __init__(
+        self,
+        iterable: Optional[Iterable] = None,
+        desc: Optional[str] = None,
+        total: Optional[int] = None,
+        unit: Optional[str] = None,
+        position: Optional[int] = None,
+        flush_interval_s: Optional[float] = None,
+    ):
+        import ray._private.services as services
+
+        if total is None and iterable is not None:
+            try:
+                total = len(iterable)
+            except (TypeError, AttributeError):
+                total = None
+
+        self._iterable = iterable
+        self._desc = desc or ""
+        self._total = total
+        self._unit = unit or "it"
+        self._ip = services.get_node_ip_address()
+        self._pid = os.getpid()
+        self._pos = position or 0
+        self._uuid = uuid.uuid4().hex
+        self._x = 0
+        self._closed = False
+        self._flush_interval_s = (
+            flush_interval_s
+            if flush_interval_s is not None
+            else self.DEFAULT_FLUSH_INTERVAL_SECONDS
+        )
+        self._last_flush_time = 0.0
+
+    def set_description(self, desc):
+        """Implements tqdm.tqdm.set_description."""
+        self._desc = desc
+        self._dump_state()
+
+    def update(self, n=1):
+        """Implements tqdm.tqdm.update."""
+        self._x += n
+        self._dump_state()
+
+    def close(self):
+        """Implements tqdm.tqdm.close."""
+        self._closed = True
+        # Don't bother if ray is shutdown (in __del__ hook).
+        if ray is not None:
+            self._dump_state(force_flush=True)
+
+    def refresh(self):
+        """Implements tqdm.tqdm.refresh."""
+        self._dump_state()
+
+    @property
+    def total(self) -> Optional[int]:
+        return self._total
+
+    @total.setter
+    def total(self, total: int):
+        self._total = total
+
+    @property
+    def n(self) -> int:
+        return self._x
+
+    @n.setter
+    def n(self, n: int):
+        self._x = n
+
+    def _dump_state(self, force_flush=False) -> None:
+        now = time.time()
+        if not force_flush and now - self._last_flush_time < self._flush_interval_s:
+            return
+        self._last_flush_time = now
+        if ray._private.worker.global_worker.mode == ray.WORKER_MODE:
+            # Include newline in payload to avoid split prints.
+            # TODO(ekl) we should move this to events.json to avoid log corruption.
+            print(json.dumps(self._get_state()) + "\n", end="")
+        else:
+            instance().process_state_update(copy.deepcopy(self._get_state()))
+
+    def _get_state(self) -> ProgressBarState:
+        return {
+            "__magic_token__": RAY_TQDM_MAGIC,
+            "x": self._x,
+            "pos": self._pos,
+            "desc": self._desc,
+            "total": self._total,
+            "unit": self._unit,
+            "ip": self._ip,
+            "pid": self._pid,
+            "uuid": self._uuid,
+            "closed": self._closed,
+        }
+
+    def __iter__(self):
+        if self._iterable is None:
+            raise ValueError("No iterable provided")
+        for x in iter(self._iterable):
+            self.update(1)
+            yield x
+
+
+class _Bar:
+    """Manages a single virtual progress bar on the driver.
+
+    The actual position of individual bars is calculated as (pos_offset + position),
+    where `pos_offset` is the position offset determined by the BarManager.
+    """
+
+    def __init__(self, state: ProgressBarState, pos_offset: int):
+        """Initialize a bar.
+
+        Args:
+            state: The initial progress bar state.
+            pos_offset: The position offset determined by the BarManager.
+        """
+        self.state = state
+        self.pos_offset = pos_offset
+        self.bar = real_tqdm.tqdm(
+            desc=state["desc"],
+            total=state["total"],
+            unit=state["unit"],
+            position=pos_offset + state["pos"],
+            dynamic_ncols=True,
+            unit_scale=True,
+        )
+        if state["x"]:
+            self.bar.update(state["x"])
+
+    def update(self, state: ProgressBarState) -> None:
+        """Apply the updated worker progress bar state."""
+        if state["desc"] != self.state["desc"]:
+            self.bar.set_description(state["desc"])
+        if state["total"] != self.state["total"]:
+            self.bar.total = state["total"]
+            self.bar.refresh()
+        delta = state["x"] - self.state["x"]
+        if delta:
+            self.bar.update(delta)
+        self.bar.refresh()
+        self.state = state
+
+    def close(self):
+        """The progress bar has been closed."""
+        self.bar.close()
+
+    def update_offset(self, pos_offset: int) -> None:
+        """Update the position offset assigned by the BarManager."""
+        if pos_offset != self.pos_offset:
+            self.pos_offset = pos_offset
+            self.bar.clear()
+            self.bar.pos = -(pos_offset + self.state["pos"])
+            self.bar.refresh()
+
+
+class _BarGroup:
+    """Manages a group of virtual progress bar produced by a single worker.
+
+    All the progress bars in the group have the same `pos_offset` determined by the
+    BarManager for the process.
+    """
+
+    def __init__(self, ip, pid, pos_offset):
+        self.ip = ip
+        self.pid = pid
+        self.pos_offset = pos_offset
+        self.bars_by_uuid: Dict[str, _Bar] = {}
+
+    def has_bar(self, bar_uuid) -> bool:
+        """Return whether this bar exists."""
+        return bar_uuid in self.bars_by_uuid
+
+    def allocate_bar(self, state: ProgressBarState) -> None:
+        """Add a new bar to this group."""
+        self.bars_by_uuid[state["uuid"]] = _Bar(state, self.pos_offset)
+
+    def update_bar(self, state: ProgressBarState) -> None:
+        """Update the state of a managed bar in this group."""
+        bar = self.bars_by_uuid[state["uuid"]]
+        bar.update(state)
+
+    def close_bar(self, state: ProgressBarState) -> None:
+        """Remove a bar from this group."""
+        bar = self.bars_by_uuid[state["uuid"]]
+        # Note: Hide and then unhide bars to prevent flashing of the
+        # last bar when we are closing multiple bars sequentially.
+        instance().hide_bars()
+        bar.close()
+        del self.bars_by_uuid[state["uuid"]]
+        instance().unhide_bars()
+
+    def slots_required(self):
+        """Return the number of pos slots we need to accommodate bars in this group."""
+        if not self.bars_by_uuid:
+            return 0
+        return 1 + max(bar.state["pos"] for bar in self.bars_by_uuid.values())
+
+    def update_offset(self, offset: int) -> None:
+        """Update the position offset assigned by the BarManager."""
+        if offset != self.pos_offset:
+            self.pos_offset = offset
+            for bar in self.bars_by_uuid.values():
+                bar.update_offset(offset)
+
+    def hide_bars(self) -> None:
+        """Temporarily hide visible bars to avoid conflict with other log messages."""
+        for bar in self.bars_by_uuid.values():
+            bar.bar.clear()
+
+    def unhide_bars(self) -> None:
+        """Opposite of hide_bars()."""
+        for bar in self.bars_by_uuid.values():
+            bar.bar.refresh()
+
+
+class _BarManager:
+    """Central tqdm manager run on the driver.
+
+    This class holds a collection of BarGroups and updates their `pos_offset` as
+    needed to ensure individual progress bars do not collide in position, kind of
+    like a virtual memory manager.
+    """
+
+    def __init__(self):
+        import ray._private.services as services
+
+        self.ip = services.get_node_ip_address()
+        self.pid = os.getpid()
+        self.bar_groups = {}
+        self.in_hidden_state = False
+        self.num_hides = 0
+        self.lock = threading.RLock()
+        # Avoid colorizing Jupyter output, since the tqdm bar is rendered in
+        # ipywidgets instead of in the console.
+        self.should_colorize = not ray.widgets.util.in_notebook()
+
+    def process_state_update(self, state: ProgressBarState) -> None:
+        """Apply the remote progress bar state update.
+
+        This creates a new bar locally if it doesn't already exist. When a bar is
+        created or destroyed, we also recalculate and update the `pos_offset` of each
+        BarGroup on the screen.
+        """
+        with self.lock:
+            self._process_state_update_locked(state)
+
+    def _process_state_update_locked(self, state: ProgressBarState) -> None:
+        if not real_tqdm:
+            if log_once("no_tqdm"):
+                logger.warning("tqdm is not installed. Progress bars will be disabled.")
+            return
+        if state["ip"] == self.ip:
+            if state["pid"] == self.pid:
+                prefix = ""
+            else:
+                prefix = "(pid={}) ".format(state.get("pid"))
+                if self.should_colorize:
+                    prefix = "{}{}{}{}".format(
+                        colorama.Style.DIM,
+                        colorama.Fore.CYAN,
+                        prefix,
+                        colorama.Style.RESET_ALL,
+                    )
+        else:
+            prefix = "(pid={}, ip={}) ".format(
+                state.get("pid"),
+                state.get("ip"),
+            )
+            if self.should_colorize:
+                prefix = "{}{}{}{}".format(
+                    colorama.Style.DIM,
+                    colorama.Fore.CYAN,
+                    prefix,
+                    colorama.Style.RESET_ALL,
+                )
+        state["desc"] = prefix + state["desc"]
+        process = self._get_or_allocate_bar_group(state)
+        if process.has_bar(state["uuid"]):
+            # Always call `update_bar` to sync any last remaining updates
+            # prior to closing. Otherwise, the displayed progress bars
+            # can be left incomplete, even after execution finishes.
+            # Fixes https://github.com/ray-project/ray/issues/44983
+            process.update_bar(state)
+
+            if state["closed"]:
+                process.close_bar(state)
+                self._update_offsets()
+        else:
+            process.allocate_bar(state)
+            self._update_offsets()
+
+    def hide_bars(self) -> None:
+        """Temporarily hide visible bars to avoid conflict with other log messages."""
+        with self.lock:
+            if not self.in_hidden_state:
+                self.in_hidden_state = True
+                self.num_hides += 1
+                for group in self.bar_groups.values():
+                    group.hide_bars()
+
+    def unhide_bars(self) -> None:
+        """Opposite of hide_bars()."""
+        with self.lock:
+            if self.in_hidden_state:
+                self.in_hidden_state = False
+                for group in self.bar_groups.values():
+                    group.unhide_bars()
+
+    def _get_or_allocate_bar_group(self, state: ProgressBarState):
+        ptuple = (state["ip"], state["pid"])
+        if ptuple not in self.bar_groups:
+            offset = sum(p.slots_required() for p in self.bar_groups.values())
+            self.bar_groups[ptuple] = _BarGroup(state["ip"], state["pid"], offset)
+        return self.bar_groups[ptuple]
+
+    def _update_offsets(self):
+        offset = 0
+        for proc in self.bar_groups.values():
+            proc.update_offset(offset)
+            offset += proc.slots_required()
+
+
+def instance() -> _BarManager:
+    """Get or create a BarManager for this process."""
+    global _manager
+
+    with _mgr_lock:
+        if _manager is None:
+            _manager = _BarManager()
+            if env_bool("RAY_TQDM_PATCH_PRINT", True):
+                import builtins
+
+                builtins.print = safe_print
+        return _manager
+
+
+if __name__ == "__main__":
+
+    @ray.remote
+    def processing(delay):
+        def sleep(x):
+            print("Intermediate result", x)
+            time.sleep(delay)
+            return x
+
+        ray.data.range(1000, override_num_blocks=100).map(
+            sleep, compute=ray.data.ActorPoolStrategy(size=1)
+        ).count()
+
+    ray.get(
+        [
+            processing.remote(0.03),
+            processing.remote(0.01),
+            processing.remote(0.05),
+        ]
+    )

python/ray/experimental/util/types.py

@@ -0,0 +1,46 @@
+from dataclasses import dataclass
+from enum import Enum
+
+from ray.util.annotations import PublicAPI
+
+
+class _CollectiveOp:
+    pass
+
+
+@PublicAPI
+class ReduceOp(Enum):
+    SUM = 0
+    PRODUCT = 1
+    MAX = 2
+    MIN = 3
+    AVG = 4
+
+
+@PublicAPI
+@dataclass
+class AllGatherOp(_CollectiveOp):
+    pass
+
+
+@PublicAPI
+@dataclass
+class AllReduceOp(_CollectiveOp):
+    reduceOp: ReduceOp = ReduceOp.SUM
+
+
+@PublicAPI
+@dataclass
+class ReduceScatterOp(_CollectiveOp):
+    reduceOp: ReduceOp = ReduceOp.SUM
+
+
+@PublicAPI(stability="alpha")
+class Device(Enum):
+    DEFAULT = "default"
+    CPU = "cpu"
+    GPU = "gpu"
+    CUDA = "cuda"
+
+    def __str__(self):
+        return self.value