chore: import upstream snapshot with attribution

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()
+            )
+        )