Add files using upload-large-folder tool

lib/python3.12/site-packages/ray/dag/__init__.py

@@ -0,0 +1,46 @@
+from ray.dag.dag_node import DAGNode
+from ray.dag.function_node import FunctionNode
+from ray.dag.class_node import (
+    ClassNode,
+    ClassMethodNode,
+)
+from ray.dag.collective_node import CollectiveOutputNode
+from ray.dag.input_node import (
+    InputNode,
+    InputAttributeNode,
+    DAGInputData,
+)
+from ray.dag.output_node import MultiOutputNode
+from ray.dag.dag_operation_future import DAGOperationFuture, GPUFuture
+from ray.dag.constants import (
+    PARENT_CLASS_NODE_KEY,
+    PREV_CLASS_METHOD_CALL_KEY,
+    BIND_INDEX_KEY,
+    IS_CLASS_METHOD_OUTPUT_KEY,
+    COLLECTIVE_OPERATION_KEY,
+    DAGNODE_TYPE_KEY,
+)
+from ray.dag.vis_utils import plot
+from ray.dag.context import DAGContext
+
+__all__ = [
+    "ClassNode",
+    "ClassMethodNode",
+    "CollectiveOutputNode",
+    "DAGNode",
+    "DAGOperationFuture",
+    "FunctionNode",
+    "GPUFuture",
+    "InputNode",
+    "InputAttributeNode",
+    "DAGInputData",
+    "PARENT_CLASS_NODE_KEY",
+    "PREV_CLASS_METHOD_CALL_KEY",
+    "BIND_INDEX_KEY",
+    "IS_CLASS_METHOD_OUTPUT_KEY",
+    "COLLECTIVE_OPERATION_KEY",
+    "DAGNODE_TYPE_KEY",
+    "plot",
+    "MultiOutputNode",
+    "DAGContext",
+]

lib/python3.12/site-packages/ray/dag/base.py

@@ -0,0 +1,8 @@
+"""This module defines the base class for object scanning and gets rid of
+reference cycles."""
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class DAGNodeBase:
+    """Common base class for a node in a Ray task graph."""

lib/python3.12/site-packages/ray/dag/class_node.py

@@ -0,0 +1,320 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+from weakref import ReferenceType
+
+import ray
+from ray.dag.constants import (
+    BIND_INDEX_KEY,
+    IS_CLASS_METHOD_OUTPUT_KEY,
+    PARENT_CLASS_NODE_KEY,
+    PREV_CLASS_METHOD_CALL_KEY,
+)
+from ray.dag.dag_node import DAGNode
+from ray.dag.format_utils import get_dag_node_str
+from ray.dag.input_node import InputNode
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class ClassNode(DAGNode):
+    """Represents an actor creation in a Ray task DAG."""
+
+    def __init__(
+        self,
+        cls,
+        cls_args,
+        cls_kwargs,
+        cls_options,
+        other_args_to_resolve=None,
+    ):
+        self._body = cls
+        self._last_call: Optional["ClassMethodNode"] = None
+        super().__init__(
+            cls_args,
+            cls_kwargs,
+            cls_options,
+            other_args_to_resolve=other_args_to_resolve,
+        )
+
+        if self._contains_input_node():
+            raise ValueError(
+                "InputNode handles user dynamic input the DAG, and "
+                "cannot be used as args, kwargs, or other_args_to_resolve "
+                "in ClassNode constructor because it is not available at "
+                "class construction or binding time."
+            )
+
+    def _copy_impl(
+        self,
+        new_args: List[Any],
+        new_kwargs: Dict[str, Any],
+        new_options: Dict[str, Any],
+        new_other_args_to_resolve: Dict[str, Any],
+    ):
+        return ClassNode(
+            self._body,
+            new_args,
+            new_kwargs,
+            new_options,
+            other_args_to_resolve=new_other_args_to_resolve,
+        )
+
+    def _execute_impl(self, *args, **kwargs):
+        """Executor of ClassNode by ray.remote()
+
+        Args and kwargs are to match base class signature, but not in the
+        implementation. All args and kwargs should be resolved and replaced
+        with value in bound_args and bound_kwargs via bottom-up recursion when
+        current node is executed.
+        """
+        return (
+            ray.remote(self._body)
+            .options(**self._bound_options)
+            .remote(*self._bound_args, **self._bound_kwargs)
+        )
+
+    def _contains_input_node(self) -> bool:
+        """Check if InputNode is used in children DAGNodes with current node
+        as the root.
+        """
+        children_dag_nodes = self._get_all_child_nodes()
+        for child in children_dag_nodes:
+            if isinstance(child, InputNode):
+                return True
+        return False
+
+    def __getattr__(self, method_name: str):
+        # User trying to call .bind() without a bind class method
+        if method_name == "bind" and "bind" not in dir(self._body):
+            raise AttributeError(f".bind() cannot be used again on {type(self)} ")
+        # Raise an error if the method is invalid.
+        getattr(self._body, method_name)
+        call_node = _UnboundClassMethodNode(self, method_name, {})
+        return call_node
+
+    def __str__(self) -> str:
+        return get_dag_node_str(self, str(self._body))
+
+
+class _UnboundClassMethodNode(object):
+    def __init__(self, actor: ClassNode, method_name: str, options: dict):
+        # TODO(sang): Theoretically, We should use weakref cuz it is
+        # a circular dependency but when I used weakref, it fails
+        # because we cannot serialize the weakref.
+        self._actor = actor
+        self._method_name = method_name
+        self._options = options
+
+    def bind(self, *args, **kwargs):
+        other_args_to_resolve = {
+            PARENT_CLASS_NODE_KEY: self._actor,
+            PREV_CLASS_METHOD_CALL_KEY: self._actor._last_call,
+        }
+
+        node = ClassMethodNode(
+            self._method_name,
+            args,
+            kwargs,
+            self._options,
+            other_args_to_resolve=other_args_to_resolve,
+        )
+        self._actor._last_call = node
+        return node
+
+    def __getattr__(self, attr: str):
+        if attr == "remote":
+            raise AttributeError(
+                ".remote() cannot be used on ClassMethodNodes. Use .bind() instead "
+                "to express an symbolic actor call."
+            )
+        else:
+            return self.__getattribute__(attr)
+
+    def options(self, **options):
+        self._options = options
+        return self
+
+
+class _ClassMethodOutput:
+    """Represents a class method output in a Ray function DAG."""
+
+    def __init__(self, class_method_call: "ClassMethodNode", output_idx: int):
+        # The upstream class method call that returns multiple values.
+        self._class_method_call = class_method_call
+        # The output index of the return value from the upstream class method call.
+        self._output_idx = output_idx
+
+    @property
+    def class_method_call(self) -> "ClassMethodNode":
+        return self._class_method_call
+
+    @property
+    def output_idx(self) -> int:
+        return self._output_idx
+
+
+@DeveloperAPI
+class ClassMethodNode(DAGNode):
+    """Represents an actor method invocation in a Ray function DAG."""
+
+    def __init__(
+        self,
+        method_name: str,
+        method_args: Tuple[Any],
+        method_kwargs: Dict[str, Any],
+        method_options: Dict[str, Any],
+        other_args_to_resolve: Dict[str, Any],
+    ):
+        self._bound_args = method_args or []
+        self._bound_kwargs = method_kwargs or {}
+        self._bound_options = method_options or {}
+        self._method_name: str = method_name
+        # Parse other_args_to_resolve and assign to variables
+        self._parent_class_node: Union[
+            ClassNode, ReferenceType["ray._private.actor.ActorHandle"]
+        ] = other_args_to_resolve.get(PARENT_CLASS_NODE_KEY)
+        # Used to track lineage of ClassMethodCall to preserve deterministic
+        # submission and execution order.
+        self._prev_class_method_call: Optional[
+            ClassMethodNode
+        ] = other_args_to_resolve.get(PREV_CLASS_METHOD_CALL_KEY, None)
+        # The index/order when bind() is called on this class method
+        self._bind_index: Optional[int] = other_args_to_resolve.get(
+            BIND_INDEX_KEY, None
+        )
+        # Represent if the ClassMethodNode is a class method output. If True,
+        # the node is a placeholder for a return value from the ClassMethodNode
+        # that returns multiple values. If False, the node is a class method call.
+        self._is_class_method_output: bool = other_args_to_resolve.get(
+            IS_CLASS_METHOD_OUTPUT_KEY, False
+        )
+        # Represents the return value from the upstream ClassMethodNode that
+        # returns multiple values. If the node is a class method call, this is None.
+        self._class_method_output: Optional[_ClassMethodOutput] = None
+        if self._is_class_method_output:
+            # Set the upstream ClassMethodNode and the output index of the return
+            # value from `method_args`.
+            self._class_method_output = _ClassMethodOutput(
+                method_args[0], method_args[1]
+            )
+
+        # The actor creation task dependency is encoded as the first argument,
+        # and the ordering dependency as the second, which ensures they are
+        # executed prior to this node.
+        super().__init__(
+            method_args,
+            method_kwargs,
+            method_options,
+            other_args_to_resolve=other_args_to_resolve,
+        )
+
+    def _copy_impl(
+        self,
+        new_args: List[Any],
+        new_kwargs: Dict[str, Any],
+        new_options: Dict[str, Any],
+        new_other_args_to_resolve: Dict[str, Any],
+    ):
+        return ClassMethodNode(
+            self._method_name,
+            new_args,
+            new_kwargs,
+            new_options,
+            other_args_to_resolve=new_other_args_to_resolve,
+        )
+
+    def _execute_impl(self, *args, **kwargs):
+        """Executor of ClassMethodNode by ray.remote()
+
+        Args and kwargs are to match base class signature, but not in the
+        implementation. All args and kwargs should be resolved and replaced
+        with value in bound_args and bound_kwargs via bottom-up recursion when
+        current node is executed.
+        """
+        if self.is_class_method_call:
+            method_body = getattr(self._parent_class_node, self._method_name)
+            # Execute with bound args.
+            return method_body.options(**self._bound_options).remote(
+                *self._bound_args,
+                **self._bound_kwargs,
+            )
+        else:
+            assert self._class_method_output is not None
+            return self._bound_args[0][self._class_method_output.output_idx]
+
+    def __str__(self) -> str:
+        return get_dag_node_str(self, f"{self._method_name}()")
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+    def get_method_name(self) -> str:
+        return self._method_name
+
+    def _get_bind_index(self) -> int:
+        return self._bind_index
+
+    def _get_remote_method(self, method_name):
+        method_body = getattr(self._parent_class_node, method_name)
+        return method_body
+
+    def _get_actor_handle(self) -> Optional["ray.actor.ActorHandle"]:
+        if not isinstance(self._parent_class_node, ray.actor.ActorHandle):
+            return None
+        return self._parent_class_node
+
+    @property
+    def num_returns(self) -> int:
+        """
+        Return the number of return values from the class method call. If the
+        node is a class method output, return the number of return values from
+        the upstream class method call.
+        """
+
+        if self.is_class_method_call:
+            num_returns = self._bound_options.get("num_returns", None)
+            if num_returns is None:
+                method = self._get_remote_method(self._method_name)
+                num_returns = method.__getstate__()["num_returns"]
+            return num_returns
+        else:
+            assert self._class_method_output is not None
+            return self._class_method_output.class_method_call.num_returns
+
+    @property
+    def is_class_method_call(self) -> bool:
+        """
+        Return True if the node is a class method call, False if the node is a
+        class method output.
+        """
+        return not self._is_class_method_output
+
+    @property
+    def is_class_method_output(self) -> bool:
+        """
+        Return True if the node is a class method output, False if the node is a
+        class method call.
+        """
+        return self._is_class_method_output
+
+    @property
+    def class_method_call(self) -> Optional["ClassMethodNode"]:
+        """
+        Return the upstream class method call that returns multiple values. If
+        the node is a class method output, return None.
+        """
+
+        if self._class_method_output is None:
+            return None
+        return self._class_method_output.class_method_call
+
+    @property
+    def output_idx(self) -> Optional[int]:
+        """
+        Return the output index of the return value from the upstream class
+        method call that returns multiple values. If the node is a class method
+        call, return None.
+        """
+
+        if self._class_method_output is None:
+            return None
+        return self._class_method_output.output_idx

lib/python3.12/site-packages/ray/dag/collective_node.py

@@ -0,0 +1,307 @@
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+if TYPE_CHECKING:
+    import torch
+
+import ray
+from ray.dag import (
+    ClassMethodNode,
+    DAGNode,
+)
+from ray.dag.constants import COLLECTIVE_OPERATION_KEY, IS_CLASS_METHOD_OUTPUT_KEY
+from ray.experimental.channel import ChannelContext
+from ray.experimental.channel.torch_tensor_type import Communicator, TorchTensorType
+from ray.experimental.util.types import (
+    AllGatherOp,
+    AllReduceOp,
+    ReduceScatterOp,
+    _CollectiveOp,
+)
+from ray.util.annotations import DeveloperAPI
+
+
+class _CollectiveOperation:
+    """
+    Represent metadata for a collective communicator collective operation.
+
+    Args:
+        inputs: A list of lists of DAGNode. Each nested list inside
+            of inputs should contain exactly one object per actor.
+            If multiple nested lists are provided, then the order of
+            actors should be the same for each nested list.
+        op: The collective operation to perform.
+        transport: The transport to use for the collective operation.
+
+    Requirements:
+    1. Input nodes are unique.
+    2. Actor handles are unique.
+    3. Actor handles match the custom communicator group if specified.
+    """
+
+    def __init__(
+        self,
+        inputs: List[List[DAGNode]],
+        op: _CollectiveOp,
+        transport: Optional[Union[str, Communicator]] = None,
+    ):
+        self._actor_handles: List["ray.actor.ActorHandle"] = []
+        for i, input_nodes in enumerate(inputs):
+            # Check non-empty input list
+            if len(input_nodes) == 0:
+                nested_list_error_msg = f" at index {i}" if len(inputs) > 1 else ""
+                raise ValueError(
+                    f"Expected non-empty input list{nested_list_error_msg}."
+                )
+
+            # Check input nodes are DAGNode
+            if not all(isinstance(node, DAGNode) for node in input_nodes):
+                nested_list_error_msg = (
+                    f" at list at index {i}" if len(inputs) > 1 else ""
+                )
+                raise ValueError(
+                    f"Expected all input nodes to be DAGNode{nested_list_error_msg}, "
+                    f"but got {input_nodes}."
+                )
+
+            # Check unique input nodes
+            if len(set(input_nodes)) != len(input_nodes):
+                duplicates = [
+                    input_node
+                    for input_node in input_nodes
+                    if input_nodes.count(input_node) > 1
+                ]
+                nested_list_error_msg = (
+                    f" at list at index {i}" if len(inputs) > 1 else ""
+                )
+                raise ValueError(
+                    f"Expected unique input nodes{nested_list_error_msg}, but found duplicates: "
+                    f"{duplicates}"
+                )
+
+            current_actor_handles = []
+            for input_node in input_nodes:
+                actor_handle = input_node._get_actor_handle()
+                if actor_handle is None:
+                    nested_list_error_msg = (
+                        f" at list at index {i}" if len(inputs) > 1 else ""
+                    )
+                    raise ValueError(
+                        f"Expected an actor handle from the input node{nested_list_error_msg}"
+                    )
+                current_actor_handles.append(actor_handle)
+
+            # Check unique actor handles
+            if len(set(current_actor_handles)) != len(current_actor_handles):
+                invalid_input_nodes = [
+                    input_node
+                    for input_node in input_nodes
+                    if current_actor_handles.count(input_node._get_actor_handle()) > 1
+                ]
+                nested_list_error_msg = (
+                    f" at list at index {i}" if len(inputs) > 1 else ""
+                )
+                raise ValueError(
+                    f"Expected unique actor handles{nested_list_error_msg}, "
+                    "but found duplicate actor handles from input nodes: "
+                    f"{invalid_input_nodes}"
+                )
+
+            if i == 0:
+                first_actor_handles = current_actor_handles
+
+            # Check all lists of DAGNode have the same number of nodes
+            if len(inputs[0]) != len(inputs[i]):
+                raise ValueError(
+                    f"Expected all input lists to have the same number of nodes. "
+                    f"List at index 0 has length {len(inputs[0])}, but list at "
+                    f"index {i} has length {len(inputs[i])}."
+                )
+
+            # Check all lists of DAGNode have same set of actor handles
+            if set(first_actor_handles) != set(current_actor_handles):
+                raise ValueError(
+                    f"Expected all input lists to have the same set of actor handles. "
+                    f"List at index 0 has actors {set(first_actor_handles)}, but list at "
+                    f"index {i} has actors {set(current_actor_handles)}."
+                )
+
+            # Check all lists of DAGNode have same order of actor handles
+            for j, (first, current) in enumerate(
+                zip(first_actor_handles, current_actor_handles)
+            ):
+                if first != current:
+                    raise ValueError(
+                        f"Expected all input lists to have the same order of actor handles. "
+                        f"List at index 0 has actor {first} at position {j}, but list at "
+                        f"index {i} has actor {current} at position {j}."
+                    )
+        self._actor_handles = current_actor_handles
+
+        self._op = op
+        if transport is None:
+            transport = TorchTensorType.ACCELERATOR
+        self._type_hint = TorchTensorType(transport=transport, _direct_return=True)
+        if isinstance(transport, Communicator):
+            if set(transport.get_actor_handles()) != set(self._actor_handles):
+                raise ValueError(
+                    "Expected actor handles to match the custom communicator group"
+                )
+
+    def __str__(self) -> str:
+        return (
+            f"CollectiveOperation("
+            f"_actor_handles={self._actor_handles}, "
+            f"_op={self._op}, "
+            f"_type_hint={self._type_hint})"
+        )
+
+    @property
+    def actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        return self._actor_handles
+
+    @property
+    def type_hint(self) -> TorchTensorType:
+        return self._type_hint
+
+    def get_communicator(self) -> Communicator:
+        if self._type_hint.communicator_id is not None:
+            ctx = ChannelContext.get_current()
+            communicator = ctx.communicators[self._type_hint.communicator_id]
+        elif self._type_hint.get_custom_communicator() is not None:
+            communicator = self._type_hint.get_custom_communicator()
+        else:
+            raise ValueError("Expected a communicator group")
+        return communicator
+
+    def execute(
+        self, *send_buf: "torch.Tensor"
+    ) -> Union["torch.Tensor", Tuple["torch.Tensor", ...]]:
+        """
+        Call the collective operation on the input tensor(s). Output tensor(s) are
+        allocated and returned.
+
+        Args:
+            *send_buf: A variable number of torch tensors to send to the collective
+                operation. The tensors have the same order as the input nodes.
+
+        Returns:
+            A torch tensor or a tuple of torch tensors containing the results of the
+            collective operation. The output tensors have the same length and order
+            as the input node list of the actor of this operation.
+        """
+        import torch
+
+        if not all(isinstance(t, torch.Tensor) for t in send_buf):
+            raise ValueError("Expected a torch tensor for each input node")
+
+        communicator = self.get_communicator()
+        if isinstance(self._op, AllGatherOp):
+            assert len(send_buf) == 1
+            t = send_buf[0]
+            world_size = len(self._actor_handles)
+            recv_buf = torch.empty(
+                (t.shape[0] * world_size, *t.shape[1:]),
+                dtype=t.dtype,
+                device=t.device,
+            )
+            communicator.allgather(t, recv_buf)
+        elif isinstance(self._op, AllReduceOp):
+            if len(send_buf) == 1:
+                t = send_buf[0]
+                recv_buf = torch.empty_like(t)
+                communicator.allreduce(t, recv_buf, self._op.reduceOp)
+            else:
+                if not all(t.dtype == send_buf[0].dtype for t in send_buf):
+                    raise ValueError(
+                        "Expected all input tensors to have the same dtype, "
+                        f"but got {[t.dtype for t in send_buf]}"
+                    )
+
+                def unflatten_from(flat_buf, bufs):
+                    views = []
+                    offset = 0
+                    for t in bufs:
+                        numel = t.numel()
+                        t = flat_buf[offset : offset + numel].view(t.shape)
+                        views.append(t)
+                        offset += numel
+                    return tuple(views)
+
+                flat_buf = torch.nn.utils.parameters_to_vector(send_buf)
+                communicator.allreduce(flat_buf, flat_buf, self._op.reduceOp)
+                recv_buf = unflatten_from(flat_buf, send_buf)
+        elif isinstance(self._op, ReduceScatterOp):
+            assert len(send_buf) == 1
+            t = send_buf[0]
+            world_size = len(self._actor_handles)
+            if t.shape[0] % world_size != 0:
+                raise ValueError(
+                    "Expected the first dimension of the input tensor to be divisible "
+                    f"by the world size {world_size}"
+                )
+            recv_buf = torch.empty(
+                (t.shape[0] // world_size, *t.shape[1:]),
+                dtype=t.dtype,
+                device=t.device,
+            )
+            communicator.reducescatter(t, recv_buf, self._op.reduceOp)
+        return recv_buf
+
+
+@DeveloperAPI
+class CollectiveOutputNode(ClassMethodNode):
+    """Represent an output node from a communicator collective operation in a Ray DAG."""
+
+    def __init__(
+        self,
+        method_name: str,
+        method_args: Tuple[
+            DAGNode,
+        ],
+        method_kwargs: Dict[str, Any],
+        method_options: Dict[str, Any],
+        other_args_to_resolve: Dict[str, Any],
+    ):
+        # Parse the input node(s).
+        self._inputs = method_args
+        # Parse the collective operation.
+        self._collective_op: _CollectiveOperation = other_args_to_resolve.get(
+            COLLECTIVE_OPERATION_KEY, None
+        )
+        self._is_class_method_output: bool = other_args_to_resolve.get(
+            IS_CLASS_METHOD_OUTPUT_KEY, False
+        )
+        if self._collective_op is None and not self._is_class_method_output:
+            raise ValueError("Expected a collective operation")
+
+        super().__init__(
+            method_name,
+            method_args,
+            method_kwargs,
+            method_options,
+            other_args_to_resolve,
+        )
+
+    def _copy_impl(
+        self,
+        new_args: List[Any],
+        new_kwargs: Dict[str, Any],
+        new_options: Dict[str, Any],
+        new_other_args_to_resolve: Dict[str, Any],
+    ):
+        return CollectiveOutputNode(
+            self._method_name,
+            new_args,
+            new_kwargs,
+            new_options,
+            other_args_to_resolve=new_other_args_to_resolve,
+        )
+
+    def _execute_impl(self, *args, **kwargs):
+        raise NotImplementedError(
+            "CollectiveOutputNode is only supported with dag.experimental_compile()"
+        )
+
+    @property
+    def collective_op(self) -> _CollectiveOperation:
+        return self._collective_op

lib/python3.12/site-packages/ray/dag/conftest.py

@@ -0,0 +1,17 @@
+import os
+
+import pytest
+
+import ray
+
+TEST_NAMESPACE = "ray_dag_test_namespace"
+
+
+@pytest.fixture(scope="session")
+def shared_ray_instance():
+    # Remove ray address for test ray cluster in case we have
+    # lingering RAY_ADDRESS="http://127.0.0.1:8265" from previous local job
+    # submissions.
+    if "RAY_ADDRESS" in os.environ:
+        del os.environ["RAY_ADDRESS"]
+    yield ray.init(num_cpus=16, namespace=TEST_NAMESPACE, log_to_driver=True)

lib/python3.12/site-packages/ray/dag/constants.py

@@ -0,0 +1,40 @@
+import os
+
+# Reserved keys used to handle ClassMethodNode in Ray DAG building.
+PARENT_CLASS_NODE_KEY = "parent_class_node"
+PREV_CLASS_METHOD_CALL_KEY = "prev_class_method_call"
+BIND_INDEX_KEY = "bind_index"
+IS_CLASS_METHOD_OUTPUT_KEY = "is_class_method_output"
+
+# Reserved keys used to handle CollectiveOutputNode in Ray DAG building.
+COLLECTIVE_OPERATION_KEY = "collective_operation"
+
+# Reserved key to distinguish DAGNode type and avoid collision with user dict.
+DAGNODE_TYPE_KEY = "__dag_node_type__"
+
+# Feature flag to turn off the deadlock detection.
+RAY_CGRAPH_ENABLE_DETECT_DEADLOCK = (
+    os.environ.get("RAY_CGRAPH_ENABLE_DETECT_DEADLOCK", "1") == "1"
+)
+
+# Feature flag to turn on profiling.
+RAY_CGRAPH_ENABLE_PROFILING = os.environ.get("RAY_CGRAPH_ENABLE_PROFILING", "0") == "1"
+
+# Feature flag to turn on NVTX (NVIDIA Tools Extension Library) profiling.
+# With this flag, Compiled Graph uses nvtx to automatically annotate and profile
+# function calls during each actor's execution loop.
+# This cannot be used together with RAY_CGRAPH_ENABLE_TORCH_PROFILING.
+RAY_CGRAPH_ENABLE_NVTX_PROFILING = (
+    os.environ.get("RAY_CGRAPH_ENABLE_NVTX_PROFILING", "0") == "1"
+)
+
+# Feature flag to turn on torch profiling.
+# This cannot be used together with RAY_CGRAPH_ENABLE_NVTX_PROFILING.
+RAY_CGRAPH_ENABLE_TORCH_PROFILING = (
+    os.environ.get("RAY_CGRAPH_ENABLE_TORCH_PROFILING", "0") == "1"
+)
+
+# Feature flag to turn on visualization of the execution schedule.
+RAY_CGRAPH_VISUALIZE_SCHEDULE = (
+    os.environ.get("RAY_CGRAPH_VISUALIZE_SCHEDULE", "0") == "1"
+)

lib/python3.12/site-packages/ray/dag/context.py

@@ -0,0 +1,108 @@
+import os
+import threading
+from dataclasses import dataclass
+from typing import Optional
+
+from ray.util.annotations import DeveloperAPI
+
+# The context singleton on this process.
+_default_context: "Optional[DAGContext]" = None
+_context_lock = threading.Lock()
+
+DEFAULT_SUBMIT_TIMEOUT_S = int(os.environ.get("RAY_CGRAPH_submit_timeout", 10))
+DEFAULT_GET_TIMEOUT_S = int(os.environ.get("RAY_CGRAPH_get_timeout", 10))
+DEFAULT_TEARDOWN_TIMEOUT_S = int(os.environ.get("RAY_CGRAPH_teardown_timeout", 30))
+DEFAULT_READ_ITERATION_TIMEOUT_S = float(
+    os.environ.get("RAY_CGRAPH_read_iteration_timeout_s", 0.1)
+)
+# Default buffer size is 1MB.
+DEFAULT_BUFFER_SIZE_BYTES = int(os.environ.get("RAY_CGRAPH_buffer_size_bytes", 1e6))
+# The default number of in-flight executions that can be submitted before consuming the
+# output.
+DEFAULT_MAX_INFLIGHT_EXECUTIONS = int(
+    os.environ.get("RAY_CGRAPH_max_inflight_executions", 10)
+)
+
+# The default number of results that can be buffered at the driver.
+DEFAULT_MAX_BUFFERED_RESULTS = int(
+    os.environ.get("RAY_CGRAPH_max_buffered_results", 1000)
+)
+
+DEFAULT_OVERLAP_GPU_COMMUNICATION = bool(
+    os.environ.get("RAY_CGRAPH_overlap_gpu_communication", 0)
+)
+
+
+@DeveloperAPI
+@dataclass
+class DAGContext:
+    """Global settings for Ray DAG.
+
+    You can configure parameters in the DAGContext by setting the environment
+    variables, `RAY_CGRAPH_<param>` (e.g., `RAY_CGRAPH_buffer_size_bytes`) or Python.
+
+    Examples:
+        >>> from ray.dag import DAGContext
+        >>> DAGContext.get_current().buffer_size_bytes
+        1000000
+        >>> DAGContext.get_current().buffer_size_bytes = 500
+        >>> DAGContext.get_current().buffer_size_bytes
+        500
+
+    Args:
+        submit_timeout: The maximum time in seconds to wait for execute()
+            calls.
+        get_timeout: The maximum time in seconds to wait when retrieving
+            a result from the DAG during `ray.get`. This should be set to a
+            value higher than the expected time to execute the entire DAG.
+        teardown_timeout: The maximum time in seconds to wait for the DAG to
+            cleanly shut down.
+        read_iteration_timeout: The timeout in seconds for each read iteration
+            that reads one of the input channels. If the timeout is reached, the
+            read operation will be interrupted and will try to read the next
+            input channel. It must be less than or equal to `get_timeout`.
+        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.
+        max_inflight_executions: The maximum number of in-flight executions that
+            can be submitted via `execute` or `execute_async` before consuming
+            the output using `ray.get()`. If the caller submits more executions,
+            `RayCgraphCapacityExceeded` is raised.
+        overlap_gpu_communication: (experimental) Whether to overlap GPU
+            communication with computation during DAG execution. If True, the
+            communication and computation can be overlapped, which can improve
+            the performance of the DAG execution.
+    """
+
+    submit_timeout: int = DEFAULT_SUBMIT_TIMEOUT_S
+    get_timeout: int = DEFAULT_GET_TIMEOUT_S
+    teardown_timeout: int = DEFAULT_TEARDOWN_TIMEOUT_S
+    read_iteration_timeout: float = DEFAULT_READ_ITERATION_TIMEOUT_S
+    buffer_size_bytes: int = DEFAULT_BUFFER_SIZE_BYTES
+    max_inflight_executions: int = DEFAULT_MAX_INFLIGHT_EXECUTIONS
+    max_buffered_results: int = DEFAULT_MAX_BUFFERED_RESULTS
+    overlap_gpu_communication: bool = DEFAULT_OVERLAP_GPU_COMMUNICATION
+
+    def __post_init__(self):
+        if self.read_iteration_timeout > self.get_timeout:
+            raise ValueError(
+                "RAY_CGRAPH_read_iteration_timeout_s "
+                f"({self.read_iteration_timeout}) must be less than or equal to "
+                f"RAY_CGRAPH_get_timeout ({self.get_timeout})"
+            )
+
+    @staticmethod
+    def get_current() -> "DAGContext":
+        """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 = DAGContext()
+
+            return _default_context

lib/python3.12/site-packages/ray/dag/dag_node.py

@@ -0,0 +1,705 @@
+import asyncio
+import copy
+import uuid
+from itertools import chain
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
+import ray
+from ray.dag.base import DAGNodeBase
+from ray.dag.compiled_dag_node import build_compiled_dag_from_ray_dag
+from ray.dag.py_obj_scanner import _PyObjScanner
+from ray.experimental.channel import ChannelOutputType
+from ray.experimental.channel.auto_transport_type import AutoTransportType
+from ray.experimental.channel.communicator import Communicator
+from ray.experimental.channel.torch_tensor_type import TorchTensorType
+from ray.experimental.util.types import Device
+from ray.util.annotations import DeveloperAPI
+
+T = TypeVar("T")
+
+
+@DeveloperAPI
+class DAGNode(DAGNodeBase):
+    """Abstract class for a node in a Ray task graph.
+
+    A node has a type (e.g., FunctionNode), data (e.g., function options and
+    body), arguments (Python values, DAGNodes, and DAGNodes nested within Python
+    argument values) and options (Ray API .options() used for function, class
+    or class method)
+    """
+
+    def __init__(
+        self,
+        args: Tuple[Any],
+        kwargs: Dict[str, Any],
+        options: Dict[str, Any],
+        other_args_to_resolve: Dict[str, Any],
+    ):
+        """
+        args:
+            args (Tuple[Any]): Bound node arguments.
+                ex: func_or_class.bind(1)
+            kwargs (Dict[str, Any]): Bound node keyword arguments.
+                ex: func_or_class.bind(a=1)
+            options (Dict[str, Any]): Bound node options arguments.
+                ex: func_or_class.options(num_cpus=2)
+            other_args_to_resolve (Dict[str, Any]): Bound kwargs to resolve
+                that's specific to subclass implementation without exposing
+                as args in base class, example: ClassMethodNode
+        """
+        self._bound_args: Tuple[Any] = args or []
+        self._bound_kwargs: Dict[str, Any] = kwargs or {}
+        self._bound_options: Dict[str, Any] = options or {}
+        self._bound_other_args_to_resolve: Optional[Dict[str, Any]] = (
+            other_args_to_resolve or {}
+        )
+
+        # The list of nodes that use this DAG node as an argument.
+        self._downstream_nodes: List["DAGNode"] = []
+
+        # UUID that is not changed over copies of this node.
+        self._stable_uuid = uuid.uuid4().hex
+
+        # Indicates whether this DAG node contains nested DAG nodes.
+        # Nested DAG nodes are allowed in traditional DAGs but not
+        # in Ray Compiled Graphs, except for MultiOutputNode.
+        self._args_contain_nested_dag_node = False
+
+        # The list of nodes that this DAG node uses as an argument.
+        self._upstream_nodes: List["DAGNode"] = self._collect_upstream_nodes()
+
+        # Cached values from last call to execute()
+        self.cache_from_last_execute = {}
+
+        self._type_hint: ChannelOutputType = ChannelOutputType()
+
+        # If the original type hint is an AutoTransportType, we make a copy
+        # here when it is resolved to the actual type, as additional debugging
+        # information. Otherwise, it is None.
+        self._original_type_hint: Optional[ChannelOutputType] = None
+
+        # Whether this node calls `experimental_compile`.
+        self.is_cgraph_output_node = False
+
+    def _collect_upstream_nodes(self) -> List["DAGNode"]:
+        """
+        Retrieve upstream nodes and update their downstream dependencies.
+
+        Currently, the DAG assumes that all DAGNodes in `args`, `kwargs`, and
+        `other_args_to_resolve` are upstream nodes. However, Ray Compiled Graphs
+        builds the upstream/downstream relationship based only on args. Be cautious
+        when persisting DAGNodes in `other_args_to_resolve` and kwargs in the future.
+
+        TODO (kevin85421): Currently, the upstream nodes and downstream nodes have
+        circular references. Therefore, it relies on the garbage collector to clean
+        them up instead of reference counting. We should consider using weak references
+        to avoid circular references.
+        """
+        upstream_nodes: List["DAGNode"] = []
+
+        # Ray Compiled Graphs do not allow nested DAG nodes in arguments.
+        # Specifically, a DAGNode should not be placed inside any type of
+        # container. However, we only know if this is a compiled graph
+        # when calling `experimental_compile`. Therefore, we need to check
+        # in advance if the arguments contain nested DAG nodes and raise
+        # an error after compilation.
+        assert hasattr(self._bound_args, "__iter__")
+        for arg in self._bound_args:
+            if isinstance(arg, DAGNode):
+                upstream_nodes.append(arg)
+            else:
+                scanner = _PyObjScanner()
+                dag_nodes = scanner.find_nodes(arg)
+                upstream_nodes.extend(dag_nodes)
+                scanner.clear()
+                self._args_contain_nested_dag_node = len(dag_nodes) > 0
+
+        scanner = _PyObjScanner()
+        other_upstream_nodes: List["DAGNode"] = scanner.find_nodes(
+            [
+                self._bound_kwargs,
+                self._bound_other_args_to_resolve,
+            ]
+        )
+        upstream_nodes.extend(other_upstream_nodes)
+        scanner.clear()
+        # Update dependencies.
+        for upstream_node in upstream_nodes:
+            upstream_node._downstream_nodes.append(self)
+        return upstream_nodes
+
+    def with_tensor_transport(
+        self,
+        transport: Optional[Union[str, Communicator]] = "auto",
+        device: Literal["default", "cpu", "gpu", "cuda"] = "default",
+        _static_shape: bool = False,
+        _direct_return: bool = False,
+    ):
+        """
+        Configure the torch tensor transport for this node.
+
+        Args:
+            transport: Specifies the tensor transport mechanism.
+                - "accelerator": Tensors are communicated using accelerator-specific backends
+                (e.g., NCCL, XLA, or vendor-provided transport). This is the recommended option
+                for most use cases, as it supports extensibility and future hardware backends.
+                - "nccl": Tensors are passed explicitly via NCCL. This option is kept for
+                backwards compatibility and may be removed in the future. Use "accelerator"
+                instead unless you have legacy requirements.
+                - "shm": Tensors are passed via host shared memory and gRPC. Typically used
+                when accelerator-based transport is unavailable or not suitable.
+                - "auto" (default): The system automatically selects the appropriate transport
+                mechanism based on the sender and receiver, usually preferring accelerator-based
+                transport when available.
+            device: The target device to use for the tensor transport.
+                "default": The tensor will maintain its original device placement from the sender
+                "cpu": The tensor will be explicitly moved to CPU device in the receiver
+                "gpu" or "cuda": The tensor will be explicitly moved to GPU device in 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. If this is True, the
+                transport will be more efficient.
+            _direct_return: Whether the tensor is sent directly or inside of
+                other data. If a "nccl" transport is used, this allows the
+                sender and receiver to eliminate performance overhead from
+                an additional data transfer.
+        """
+        try:
+            device = Device(device)
+        except ValueError:
+            valid_devices = ", ".join(f"'{d.value}'" for d in Device)
+            raise ValueError(
+                f"Invalid device '{device}'. Valid options are: {valid_devices}."
+            )
+        if transport == "auto":
+            self._type_hint = AutoTransportType(
+                device=device,
+                _static_shape=_static_shape,
+                _direct_return=_direct_return,
+            )
+        elif transport == "nccl":
+            self._type_hint = TorchTensorType(
+                transport="accelerator",
+                device=device,
+                _static_shape=_static_shape,
+                _direct_return=_direct_return,
+            )
+        elif transport == "accelerator":
+            self._type_hint = TorchTensorType(
+                transport="accelerator",
+                device=device,
+                _static_shape=_static_shape,
+                _direct_return=_direct_return,
+            )
+        elif transport == "shm":
+            self._type_hint = TorchTensorType(
+                device=device,
+                _static_shape=_static_shape,
+                _direct_return=_direct_return,
+            )
+        else:
+            if not isinstance(transport, Communicator):
+                raise ValueError(
+                    f"Invalid transport type: {transport}. "
+                    "Transport must be one of 'auto', 'nccl', 'shm', 'accelerator' or "
+                    "an instance of Communicator type."
+                )
+            self._type_hint = TorchTensorType(
+                transport=transport,
+                device=device,
+                _static_shape=_static_shape,
+                _direct_return=_direct_return,
+            )
+        return self
+
+    @property
+    def type_hint(self) -> ChannelOutputType:
+        return self._type_hint
+
+    @type_hint.setter
+    def type_hint(self, type_hint: ChannelOutputType) -> None:
+        if isinstance(self._type_hint, AutoTransportType):
+            self._original_type_hint = self._type_hint
+        self._type_hint = type_hint
+
+    def get_args(self) -> Tuple[Any]:
+        """Return the tuple of arguments for this node."""
+
+        return self._bound_args
+
+    def get_kwargs(self) -> Dict[str, Any]:
+        """Return the dict of keyword arguments for this node."""
+
+        return self._bound_kwargs.copy()
+
+    def get_options(self) -> Dict[str, Any]:
+        """Return the dict of options arguments for this node."""
+
+        return self._bound_options.copy()
+
+    def get_other_args_to_resolve(self) -> Dict[str, Any]:
+        """Return the dict of other args to resolve arguments for this node."""
+        return self._bound_other_args_to_resolve.copy()
+
+    def get_stable_uuid(self) -> str:
+        """Return stable uuid for this node.
+        1) Generated only once at first instance creation
+        2) Stable across pickling, replacement and JSON serialization.
+        """
+        return self._stable_uuid
+
+    async def get_object_refs_from_last_execute(self) -> Dict[str, Any]:
+        """Gets cached object refs from the last call to execute().
+
+        After this DAG is executed through execute(), retrieves a map between node
+        UUID to a reference to the return value of the default executor on that node.
+        """
+        cache = {}
+        for node_uuid, value in self.cache_from_last_execute.items():
+            if isinstance(value, asyncio.Task):
+                cache[node_uuid] = await value
+            else:
+                cache[node_uuid] = value
+
+        return cache
+
+    def clear_cache(self):
+        self.cache_from_last_execute = {}
+
+    def experimental_compile(
+        self,
+        _submit_timeout: Optional[float] = None,
+        _buffer_size_bytes: Optional[int] = None,
+        enable_asyncio: bool = False,
+        _max_inflight_executions: Optional[int] = None,
+        _max_buffered_results: Optional[int] = None,
+        _overlap_gpu_communication: Optional[bool] = None,
+        _default_communicator: Optional[Union[Communicator, str]] = "create",
+    ) -> "ray.dag.CompiledDAG":
+        """Compile an accelerated execution path for this DAG.
+
+        Args:
+            _submit_timeout: The maximum time in seconds to wait for execute() calls.
+                None means using default timeout, 0 means immediate timeout
+                (immediate success or timeout without blocking), -1 means
+                infinite timeout (block indefinitely).
+            _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.
+            enable_asyncio: Whether to enable asyncio for this DAG.
+            _max_inflight_executions: The maximum number of in-flight executions that
+                can be submitted via `execute` or `execute_async` before consuming
+                the output using `ray.get()`. If the caller submits more executions,
+                `RayCgraphCapacityExceeded` is raised.
+            _max_buffered_results: The maximum number of results that can be
+                buffered at the driver. If more than this number of results
+                are buffered, `RayCgraphCapacityExceeded` is raised. Note that
+                when result corresponding to an execution is retrieved
+                (by calling `ray.get()` on a `CompiledDAGRef` or
+                `CompiledDAGRef` or await on a `CompiledDAGFuture`), results
+                corresponding to earlier executions that have not been retrieved
+                yet are buffered.
+            _overlap_gpu_communication: (experimental) Whether to overlap GPU
+                communication with computation during DAG execution. If True, the
+                communication and computation can be overlapped, which can improve
+                the performance of the DAG execution. If None, the default value
+                will be used.
+            _default_communicator: The default communicator to use to transfer
+                tensors. Three types of values are valid. (1) Communicator:
+                For p2p operations, this is the default communicator
+                to use for nodes annotated with `with_tensor_transport()` and when
+                shared memory is not the desired option (e.g., when transport="nccl",
+                or when transport="auto" for communication between two different GPUs).
+                For collective operations, this is the default communicator to use
+                when a custom communicator is not specified.
+                (2) "create": for each collective operation without a custom communicator
+                specified, a communicator is created and initialized on its involved actors,
+                or an already created communicator is reused if the set of actors is the same.
+                For all p2p operations without a custom communicator specified, it reuses
+                an already created collective communicator if the p2p actors are a subset.
+                Otherwise, a new communicator is created.
+                (3) None: a ValueError will be thrown if a custom communicator is not specified.
+
+        Returns:
+            A compiled DAG.
+        """
+        from ray.dag import DAGContext
+
+        ctx = DAGContext.get_current()
+        if _buffer_size_bytes is None:
+            _buffer_size_bytes = ctx.buffer_size_bytes
+
+        # Validate whether this DAG node has already been compiled.
+        if self.is_cgraph_output_node:
+            raise ValueError(
+                "It is not allowed to call `experimental_compile` on the same DAG "
+                "object multiple times no matter whether `teardown` is called or not. "
+                "Please reuse the existing compiled DAG or create a new one."
+            )
+        # Whether this node is an output node in the DAG. We cannot determine
+        # this in the constructor because the output node is determined when
+        # `experimental_compile` is called.
+        self.is_cgraph_output_node = True
+        return build_compiled_dag_from_ray_dag(
+            self,
+            _submit_timeout,
+            _buffer_size_bytes,
+            enable_asyncio,
+            _max_inflight_executions,
+            _max_buffered_results,
+            _overlap_gpu_communication,
+            _default_communicator,
+        )
+
+    def execute(
+        self, *args, _ray_cache_refs: bool = False, **kwargs
+    ) -> Union[ray.ObjectRef, "ray.actor.ActorHandle"]:
+        """Execute this DAG using the Ray default executor _execute_impl().
+
+        Args:
+            _ray_cache_refs: If true, stores the default executor's return values
+                on each node in this DAG in a cache. These should be a mix of:
+                - ray.ObjectRefs pointing to the outputs of method and function nodes
+                - Serve handles for class nodes
+                - resolved values representing user input at runtime
+        """
+
+        def executor(node):
+            return node._execute_impl(*args, **kwargs)
+
+        result = self.apply_recursive(executor)
+        if _ray_cache_refs:
+            self.cache_from_last_execute = executor.cache
+        return result
+
+    def _get_toplevel_child_nodes(self) -> List["DAGNode"]:
+        """Return the list of nodes specified as top-level args.
+
+        For example, in `f.remote(a, [b])`, only `a` is a top-level arg.
+
+        This list of nodes are those that are typically resolved prior to
+        task execution in Ray. This does not include nodes nested within args.
+        For that, use ``_get_all_child_nodes()``.
+        """
+
+        # we use List instead of Set here because the hash key of the node
+        # object changes each time we create it. So if using Set here, the
+        # order of returned children can be different if we create the same
+        # nodes and dag one more time.
+        children = []
+        for a in self.get_args():
+            if isinstance(a, DAGNode):
+                if a not in children:
+                    children.append(a)
+        for a in self.get_kwargs().values():
+            if isinstance(a, DAGNode):
+                if a not in children:
+                    children.append(a)
+        for a in self.get_other_args_to_resolve().values():
+            if isinstance(a, DAGNode):
+                if a not in children:
+                    children.append(a)
+        return children
+
+    def _get_all_child_nodes(self) -> List["DAGNode"]:
+        """Return the list of nodes referenced by the args, kwargs, and
+        args_to_resolve in current node, even they're deeply nested.
+
+        Examples:
+            f.remote(a, [b]) -> [a, b]
+            f.remote(a, [b], key={"nested": [c]}) -> [a, b, c]
+        """
+
+        scanner = _PyObjScanner()
+        # we use List instead of Set here, reason explained
+        # in `_get_toplevel_child_nodes`.
+        children = []
+        for n in scanner.find_nodes(
+            [
+                self._bound_args,
+                self._bound_kwargs,
+                self._bound_other_args_to_resolve,
+            ]
+        ):
+            if n not in children:
+                children.append(n)
+        scanner.clear()
+        return children
+
+    def _apply_and_replace_all_child_nodes(
+        self, fn: "Callable[[DAGNode], T]"
+    ) -> "DAGNode":
+        """Apply and replace all immediate child nodes using a given function.
+
+        This is a shallow replacement only. To recursively transform nodes in
+        the DAG, use ``apply_recursive()``.
+
+        Args:
+            fn: Callable that will be applied once to each child of this node.
+
+        Returns:
+            New DAGNode after replacing all child nodes.
+        """
+
+        replace_table = {}
+        # CloudPickler scanner object for current layer of DAGNode. Same
+        # scanner should be use for a full find & replace cycle.
+        scanner = _PyObjScanner()
+        # Find all first-level nested DAGNode children in args.
+        # Update replacement table and execute the replace.
+        for node in scanner.find_nodes(
+            [
+                self._bound_args,
+                self._bound_kwargs,
+                self._bound_other_args_to_resolve,
+            ]
+        ):
+            if node not in replace_table:
+                replace_table[node] = fn(node)
+        new_args, new_kwargs, new_other_args_to_resolve = scanner.replace_nodes(
+            replace_table
+        )
+        scanner.clear()
+
+        # Return updated copy of self.
+        return self._copy(
+            new_args, new_kwargs, self.get_options(), new_other_args_to_resolve
+        )
+
+    def apply_recursive(self, fn: "Callable[[DAGNode], T]") -> T:
+        """Apply callable on each node in this DAG in a bottom-up tree walk.
+
+        Args:
+            fn: Callable that will be applied once to each node in the
+                DAG. It will be applied recursively bottom-up, so nodes can
+                assume the fn has been applied to their args already.
+
+        Returns:
+            Return type of the fn after application to the tree.
+        """
+
+        if not type(fn).__name__ == "_CachingFn":
+
+            class _CachingFn:
+                def __init__(self, fn):
+                    self.cache = {}
+                    self.fn = fn
+                    self.fn.cache = self.cache
+                    self.input_node_uuid = None
+
+                def __call__(self, node: "DAGNode"):
+                    from ray.dag.input_node import InputNode
+
+                    if node._stable_uuid not in self.cache:
+                        self.cache[node._stable_uuid] = self.fn(node)
+                    if isinstance(node, InputNode):
+                        if not self.input_node_uuid:
+                            self.input_node_uuid = node._stable_uuid
+                        elif self.input_node_uuid != node._stable_uuid:
+                            raise AssertionError(
+                                "Each DAG should only have one unique InputNode."
+                            )
+                    return self.cache[node._stable_uuid]
+
+            fn = _CachingFn(fn)
+        else:
+            if self._stable_uuid in fn.cache:
+                return fn.cache[self._stable_uuid]
+
+        return fn(
+            self._apply_and_replace_all_child_nodes(
+                lambda node: node.apply_recursive(fn)
+            )
+        )
+
+    def traverse_and_apply(self, fn: "Callable[[DAGNode], T]"):
+        """
+        Traverse all nodes in the connected component of the DAG that contains
+        the `self` node, and apply the given function to each node.
+        """
+        visited = set()
+        queue = [self]
+        cgraph_output_node: Optional[DAGNode] = None
+
+        while queue:
+            node = queue.pop(0)
+            if node._args_contain_nested_dag_node:
+                self._raise_nested_dag_node_error(node._bound_args)
+
+            if node not in visited:
+                if node.is_cgraph_output_node:
+                    # Validate whether there are multiple nodes that call
+                    # `experimental_compile`.
+                    if cgraph_output_node is not None:
+                        raise ValueError(
+                            "The DAG was compiled more than once. The following two "
+                            "nodes call `experimental_compile`: "
+                            f"(1) {cgraph_output_node}, (2) {node}"
+                        )
+                    cgraph_output_node = node
+                fn(node)
+                visited.add(node)
+                """
+                Add all unseen downstream and upstream nodes to the queue.
+                This function should be called by the root of the DAG. However,
+                in some invalid cases, some nodes may not be descendants of the
+                root. Therefore, we also add upstream nodes to the queue so that
+                a meaningful error message can be raised when the DAG is compiled.
+
+                ```
+                with InputNode() as inp:
+                    dag = MultiOutputNode([a1.inc.bind(inp), a2.inc.bind(1)])
+                ```
+
+                In the above example, `a2.inc` is not a descendant of inp. If we only
+                add downstream nodes to the queue, the `a2.inc` node will not be visited
+                , and the error message will be hard to understand, such as a key error
+                in the compiled DAG.
+                """
+                for neighbor in chain.from_iterable(
+                    [node._downstream_nodes, node._upstream_nodes]
+                ):
+                    if neighbor not in visited:
+                        queue.append(neighbor)
+
+    def _raise_nested_dag_node_error(self, args):
+        """
+        Raise an error for nested DAGNodes in Ray Compiled Graphs.
+
+        Args:
+            args: The arguments of the DAGNode.
+        """
+        for arg in args:
+            if isinstance(arg, DAGNode):
+                continue
+            else:
+                scanner = _PyObjScanner()
+                dag_nodes = scanner.find_nodes([arg])
+                scanner.clear()
+                if len(dag_nodes) > 0:
+                    raise ValueError(
+                        f"Found {len(dag_nodes)} DAGNodes from the arg {arg} "
+                        f"in {self}. Please ensure that the argument is a "
+                        "single DAGNode and that a DAGNode is not allowed to "
+                        "be placed inside any type of container."
+                    )
+        raise AssertionError(
+            "A DAGNode's args should contain nested DAGNodes as args, "
+            "but none were found during the compilation process. This is a "
+            "Ray internal error. Please report this issue to the Ray team."
+        )
+
+    def _find_root(self) -> "DAGNode":
+        """
+        Return the root node of the DAG. The root node must be an InputNode.
+        """
+        from ray.dag.input_node import InputNode
+
+        node = self
+        while not isinstance(node, InputNode):
+            if len(node._upstream_nodes) == 0:
+                raise ValueError(
+                    "No InputNode found in the DAG: when traversing upwards, "
+                    f"no upstream node was found for {node}."
+                )
+            node = node._upstream_nodes[0]
+        return node
+
+    def apply_functional(
+        self,
+        source_input_list: Any,
+        predicate_fn: Callable,
+        apply_fn: Callable,
+    ):
+        """
+        Apply a given function to DAGNodes in source_input_list, and return
+        the replaced inputs without mutating or coping any DAGNode.
+
+        Args:
+            source_input_list: Source inputs to extract and apply function on
+                all children DAGNode instances.
+            predicate_fn: Applied on each DAGNode instance found and determine
+                if we should apply function to it. Can be used to filter node
+                types.
+            apply_fn: Function to apply on the node on bound attributes. Example::
+
+                apply_fn = lambda node: node._get_serve_deployment_handle(
+                    node._deployment, node._bound_other_args_to_resolve
+                )
+
+        Returns:
+            replaced_inputs: Outputs of apply_fn on DAGNodes in
+                source_input_list that passes predicate_fn.
+        """
+        replace_table = {}
+        scanner = _PyObjScanner()
+        for node in scanner.find_nodes(source_input_list):
+            if predicate_fn(node) and node not in replace_table:
+                replace_table[node] = apply_fn(node)
+
+        replaced_inputs = scanner.replace_nodes(replace_table)
+        scanner.clear()
+
+        return replaced_inputs
+
+    def _execute_impl(
+        self, *args, **kwargs
+    ) -> Union[ray.ObjectRef, "ray.actor.ActorHandle"]:
+        """Execute this node, assuming args have been transformed already."""
+        raise NotImplementedError
+
+    def _copy_impl(
+        self,
+        new_args: List[Any],
+        new_kwargs: Dict[str, Any],
+        new_options: Dict[str, Any],
+        new_other_args_to_resolve: Dict[str, Any],
+    ) -> "DAGNode":
+        """Return a copy of this node with the given new args."""
+        raise NotImplementedError
+
+    def _copy(
+        self,
+        new_args: List[Any],
+        new_kwargs: Dict[str, Any],
+        new_options: Dict[str, Any],
+        new_other_args_to_resolve: Dict[str, Any],
+    ) -> "DAGNode":
+        """Return a copy of this node with the given new args."""
+        instance = self._copy_impl(
+            new_args, new_kwargs, new_options, new_other_args_to_resolve
+        )
+        instance._stable_uuid = self._stable_uuid
+        instance._type_hint = copy.deepcopy(self._type_hint)
+        instance._original_type_hint = copy.deepcopy(self._original_type_hint)
+        return instance
+
+    def __getstate__(self):
+        """Required due to overriding `__getattr__` else pickling fails."""
+        return self.__dict__
+
+    def __setstate__(self, d: Dict[str, Any]):
+        """Required due to overriding `__getattr__` else pickling fails."""
+        self.__dict__.update(d)
+
+    def __getattr__(self, attr: str):
+        if attr == "bind":
+            raise AttributeError(f".bind() cannot be used again on {type(self)} ")
+        elif attr == "remote":
+            raise AttributeError(
+                f".remote() cannot be used on {type(self)}. To execute the task "
+                "graph for this node, use .execute()."
+            )
+        else:
+            return self.__getattribute__(attr)

lib/python3.12/site-packages/ray/dag/dag_node_operation.py

@@ -0,0 +1,853 @@
+import copy
+import heapq
+import logging
+from collections import defaultdict
+from enum import Enum
+from functools import total_ordering
+from typing import Dict, List, Optional, Set, Tuple
+
+import ray
+
+logger = logging.getLogger(__name__)
+
+
+class _DAGNodeOperationType(Enum):
+    """
+    There are three types of operations that a DAG node can perform:
+    1. READ: Read from an input channel.
+    2. COMPUTE: Execute the method corresponding to the node.
+    3. WRITE: Write to an output channel.
+    """
+
+    READ = "READ"
+    COMPUTE = "COMPUTE"
+    WRITE = "WRITE"
+
+    def viz_str(self):
+        """
+        A string representation of the operation type to be used in visualization.
+
+        The result string is a single character because conciseness is preferred.
+        """
+        if self == _DAGNodeOperationType.READ:
+            return "R"
+        elif self == _DAGNodeOperationType.COMPUTE:
+            return "C"
+        elif self == _DAGNodeOperationType.WRITE:
+            return "W"
+        assert False, f"Unknown operation type: {self}"
+
+
+class _DAGNodeOperation:
+    def __init__(
+        self,
+        exec_task_idx: int,
+        operation_type: _DAGNodeOperationType,
+        method_name: Optional[str] = None,
+    ):
+        """
+        Args:
+            exec_task_idx: The index of the task that this operation belongs to
+                in the actor's ExecutableTask list. The index is not the same
+                as bind_index because there may be more tasks bound to an actor
+                than tasks that appear in the current compiled DAG.
+            operation_type: The type of operation to perform.
+            method_name: The name of the method that this operation originates
+                from. This is only for visualization and debugging purposes.
+        """
+        self.exec_task_idx = exec_task_idx
+        self.type = operation_type
+        self.method_name = method_name
+
+    def __repr__(self):
+        return (
+            f"_DAGNodeOperation("
+            f"exec_task_idx: {self.exec_task_idx}, "
+            f"type: {self.type}, "
+            f"method_name: {self.method_name})"
+        )
+
+    def viz_str(self):
+        """
+        A string representation of the node to be used in visualization.
+        """
+        return f"[{self.exec_task_idx}] {self.method_name} {self.type.viz_str()}"
+
+    def __hash__(self):
+        return hash((self.exec_task_idx, self.type))
+
+    def __eq__(self, other):
+        # An operation is uniquely identified by its `exec_task_idx` and type.
+        # `method_name` is only for debugging purposes.
+        return self.exec_task_idx == other.exec_task_idx and self.type == other.type
+
+
+@total_ordering
+class _DAGOperationGraphNode:
+    def __init__(
+        self,
+        operation: _DAGNodeOperation,
+        task_idx: int,
+        actor_handle: "ray.actor.ActorHandle",
+        requires_accelerator: bool,
+    ):
+        """
+        _DAGOperationGraphNode represents a node in the DAG operation graph.
+        It contains information about the node's in-degree, out-degree, edges,
+        and the operation it performs.
+
+        Args:
+            operation: The operation that this node performs. The operation
+                can be a READ, COMPUTE, or WRITE operation.
+            task_idx: A unique index which can be used to index into
+                `CompiledDAG.idx_to_task` to get the corresponding task.
+            actor_handle: The actor handle to which this operation belongs.
+            requires_accelerator: Whether this operation requires accelerator.
+        """
+        self.operation = operation
+        self.task_idx = task_idx
+        self.actor_handle = actor_handle
+        self.requires_accelerator = requires_accelerator
+        # The in_edges and out_edges are dicts of tuples to strings.
+        # Each tuple (the key) contains an integer `task_idx`, which can be
+        # used to index into `idx_to_task` to get the corresponding task,
+        # and a `_DAGNodeOperationType`, which can be READ, COMPUTE, or WRITE.
+        # The string (the value) is the visualization information of the edge,
+        # it is a tuple of a label of the edge and a boolean indicating whether
+        # the edge is a control dependency.
+        self.in_edges: Dict[Tuple[int, _DAGNodeOperationType], Tuple[str, bool]] = {}
+        self.out_edges: Dict[Tuple[int, _DAGNodeOperationType], Tuple[str, bool]] = {}
+        # The synchronous nodes are all the nodes that belong to the same accelerator
+        # operation. Each node is represented by a tuple of its task idx and type.
+        self.sync_idxs: Set[Tuple[int, _DAGNodeOperationType]] = set()
+        # The pending synchronous nodes are the nodes that are pending to be executed,
+        # i.e., their in-degrees are zero. When a synchronous node is pending, it
+        # will be added to the pending synchronous nodes of all the nodes in the
+        # accelerator operation.
+        self.pending_sync_idxs: Set[Tuple[int, _DAGNodeOperationType]] = set()
+
+    def __repr__(self):
+        return (
+            f"_DAGOperationGraphNode("
+            f"operation: {self.operation}, "
+            f"task_idx: {self.task_idx}, "
+            f"actor_id: {self.actor_handle._ray_actor_id}, "
+            f"requires_accelerator: {self.requires_accelerator})"
+        )
+
+    def __lt__(self, other: "_DAGOperationGraphNode"):
+        """
+        This function defines the order of the nodes in the priority queue used in
+        `_select_next_nodes`. The priority queue is a min-heap, so the node with
+        higher priority is considered "less than" the other node.
+        """
+        if self.is_accelerator_op != other.is_accelerator_op:
+            # When one node is an accelerator operation and the other is not,
+            # prioritize the accelerator operation.
+            return self.is_accelerator_op
+        else:
+            # When either both nodes are accelerator operations or both nodes
+            # are not accelerator operations, prioritize the earlier task within
+            # the same actor and load balance tasks across actors. The tie is
+            # broken by the `task_idx`.
+            return (self.operation.exec_task_idx, self.task_idx) < (
+                other.operation.exec_task_idx,
+                other.task_idx,
+            )
+
+    def __eq__(self, other: "_DAGOperationGraphNode"):
+        """
+        Two operations are equal only when they have the same `exec_task_idx` and `type`
+        and belong to the same actor.
+        """
+        return (
+            self.actor_handle == other.actor_handle
+            and self.operation.exec_task_idx == other.operation.exec_task_idx
+            and self.operation.type == other.operation.type
+        )
+
+    def __hash__(self):
+        """
+        An operation is uniquely identified by its `task_idx` and type.
+        """
+        return hash((self.operation, self.task_idx))
+
+    @property
+    def in_degree(self) -> int:
+        return len(self.in_edges)
+
+    @property
+    def is_ready(self) -> bool:
+        """
+        If a node is not an accelerator operation, it is ready when it has a zero
+        in-degree.
+        If it is an accelerator operation, it is ready when all the nodes in the
+        operation have zero in-degrees.
+        """
+        return self.in_degree == 0 and (
+            len(self.pending_sync_idxs) == len(self.sync_idxs)
+        )
+
+    @property
+    def is_read(self) -> bool:
+        return self.operation.type == _DAGNodeOperationType.READ
+
+    @property
+    def is_accelerator_read(self) -> bool:
+        """
+        A node is an accelerator read if it is a read node and requires accelerator.
+        """
+        return (
+            self.operation.type == _DAGNodeOperationType.READ
+            and self.requires_accelerator
+        )
+
+    @property
+    def is_accelerator_compute(self) -> bool:
+        """
+        A node is an accelerator compute if it is a compute node and requires accelerator.
+        """
+        return (
+            self.operation.type == _DAGNodeOperationType.COMPUTE
+            and self.requires_accelerator
+        )
+
+    @property
+    def is_accelerator_write(self) -> bool:
+        """
+        A node is an accelerator write if it is a write node and requires accelerator.
+        """
+        return (
+            self.operation.type == _DAGNodeOperationType.WRITE
+            and self.requires_accelerator
+        )
+
+    @property
+    def is_accelerator_op(self) -> bool:
+        return (
+            self.is_accelerator_read
+            or self.is_accelerator_compute
+            or self.is_accelerator_write
+        )
+
+    def viz_str(self):
+        """
+        A string representation of the node to be used in visualization.
+        """
+        return self.operation.viz_str()
+
+    @property
+    def _actor_id(self):
+        return self.actor_handle._ray_actor_id.hex()
+
+
+def _add_edge(
+    from_node: _DAGOperationGraphNode,
+    to_node: _DAGOperationGraphNode,
+    label: str = "",
+    control_dependency: bool = False,
+):
+    """
+    Add an edge from `from_node` to `to_node`.
+
+    Args:
+        from_node: The node from which the edge originates.
+        to_node: The node to which the edge points.
+        label: The label of the edge. This will be used to annotate the edge
+            in the visualization of the execution schedule.
+    """
+    from_node.out_edges[(to_node.task_idx, to_node.operation.type)] = (
+        label,
+        control_dependency,
+    )
+    to_node.in_edges[(from_node.task_idx, from_node.operation.type)] = (
+        label,
+        control_dependency,
+    )
+
+
+def _update_pending_sync_idxs(
+    graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]],
+    node: _DAGOperationGraphNode,
+) -> None:
+    """
+    Update the node as pending for its synchronous nodes.
+    """
+    idx = (node.task_idx, node.operation.type)
+    for task_idx, op_type in node.sync_idxs:
+        sync_node = graph[task_idx][op_type]
+        sync_node.pending_sync_idxs.add(idx)
+
+
+def _push_candidate_node_if_ready(
+    actor_to_candidates: Dict["ray._raylet.ActorID", List[_DAGOperationGraphNode]],
+    graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]],
+    node: _DAGOperationGraphNode,
+) -> None:
+    """
+    Push the node with a zero in-degree to the candidates if its operation is ready.
+    If it has synchronous nodes, its accelerator operation is not ready until all
+    the nodes are pending, then all the nodes will be pushed to the candidates.
+    """
+    assert node.in_degree == 0, "Expected to have a zero in-degree"
+    # For the accelerator write node, update the in-degrees of the downstream
+    # accelerator read nodes and update them as pending. This is necessary because
+    # the data dependency edges between accelerator write and read nodes are only
+    # updated here. The accelerator P2P operation becomes ready after both the write
+    # and read nodes are marked as pending.
+    if node.is_accelerator_write:
+        for task_idx, op_type in node.out_edges:
+            read_node = graph[task_idx][op_type]
+            read_node.in_edges.pop((node.task_idx, node.operation.type))
+            assert read_node.is_accelerator_read and len(read_node.in_edges) == 0
+            _update_pending_sync_idxs(graph, read_node)
+    # For the accelerator operation node, update it as pending.
+    if len(node.sync_idxs) != 0:
+        _update_pending_sync_idxs(graph, node)
+    # The accelerator operation is ready when all the nodes have zero in-degrees.
+    # When the last node in the operation is updated as pending, push all the nodes
+    # to the candidates.
+    if node.is_ready:
+        if len(node.sync_idxs) == 0:
+            heapq.heappush(
+                actor_to_candidates[node.actor_handle._actor_id],
+                node,
+            )
+        else:
+            for task_idx, op_type in node.sync_idxs:
+                sync_node = graph[task_idx][op_type]
+                heapq.heappush(
+                    actor_to_candidates[sync_node.actor_handle._actor_id],
+                    sync_node,
+                )
+
+
+def _select_next_nodes(
+    actor_to_candidates: Dict["ray._raylet.ActorID", List[_DAGOperationGraphNode]],
+    graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]],
+) -> Optional[List[_DAGOperationGraphNode]]:
+    """
+    This function selects the next nodes for the topological sort to generate
+    execution schedule. If there are multiple candidate _DAGOperationGraphNodes,
+    select the node with the top priority. The priority is defined in
+    `_DAGOperationGraphNode.__lt__`.
+
+    For the implementation details, we maintain a priority queue for each actor,
+    where the head of the priority queue is the node with the smallest `exec_task_idx`.
+    When a node has a zero in-degree, it is added to the corresponding actor's
+    priority queue. For a node other than an accelerator collective node, it is ready to be
+    executed if it has a zero in-degree. For an accelerator collective node, it is ready to
+    be executed when all the nodes in its collective operation have zero in-degrees.
+
+    If a node is an accelerator collective node, it updates the `ready_collective_nodes` of
+    all the nodes in its collective operation. Unless all the nodes in its collective
+    group have zero in-degrees, this node is removed from the candidate list.
+    Eventually, exactly one accelerator collective node from its collective operation is
+    selected from the candidate list.
+
+    If the selected node is an accelerator write node, select all the downstream accelerator
+    read nodes. If the selected node is an accelerator collective node, select all the accelerator
+    compute nodes in its collective operation.
+
+    Args:
+        actor_to_candidates: A dictionary mapping an actor id to a list of
+            candidate nodes. The list is maintained as a priority queue, so
+            the head of the queue, i.e., `candidates[0]`, is the node with
+            the smallest `bind_index`.
+        graph: A dictionary mapping the index of a task to a dictionary of its
+            _DAGOperationGraphNodes for different operations.
+
+    Returns:
+        A list of _DAGOperationGraphNodes to be placed into the corresponding
+        execution schedules.
+    """
+    top_priority_node = None
+    for candidates in actor_to_candidates.values():
+        if len(candidates) == 0:
+            continue
+        if top_priority_node is None or candidates[0] < top_priority_node:
+            top_priority_node = candidates[0]
+
+    if top_priority_node is None:
+        return None
+    next_nodes = [top_priority_node]
+
+    # Select all the synchronous nodes in the accelerator operation.
+    if len(top_priority_node.sync_idxs) != 0:
+        for task_idx, op_type in top_priority_node.sync_idxs:
+            node = graph[task_idx][op_type]
+            if node != top_priority_node:
+                next_nodes.append(node)
+
+    # Remove the selected nodes from the candidates.
+    for node in next_nodes:
+        candidates = actor_to_candidates[node.actor_handle._actor_id]
+        candidates.remove(node)
+        heapq.heapify(candidates)
+
+    # Remove the selected nodes from the candidates.
+    for node in next_nodes:
+        candidates = actor_to_candidates[node.actor_handle._actor_id]
+        #  The accelerator read nodes are not added to the candidates.
+        if node in candidates:
+            candidates.remove(node)
+            heapq.heapify(candidates)
+
+    return next_nodes
+
+
+def _build_dag_node_operation_graph(
+    idx_to_task: Dict[int, "ray.dag.compiled_dag_node.CompiledTask"],
+    actor_to_operation_nodes: Dict[
+        "ray.actor.ActorHandle", List[List[_DAGOperationGraphNode]]
+    ],
+) -> Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]]:
+    """
+    Generate a DAG node operation graph by adding edges based on the
+    following rules:
+
+    #1  Add edges from READ to COMPUTE, and from COMPUTE to WRITE, which
+        belong to the same task.
+    #2  Add an edge from COMPUTE with bind_index i to COMPUTE with bind_index
+        i+1 if they belong to the same actor.
+    #3  Add an edge from WRITE of the writer task to READ of the reader task.
+
+    This is the step one of building an execution schedule for each actor.
+
+    Args:
+        idx_to_task: A dictionary that maps the `task_idx` to the `CompiledTask`.
+            `CompiledTask` contains information about a DAGNode and its downstream
+            nodes.
+
+        actor_to_operation_nodes: A dictionary that maps an actor handle to
+            a list of lists of _DAGOperationGraphNode. For the same actor, the
+            index of the outer list corresponds to the index of the ExecutableTask
+            in the list of `executable_tasks` in `actor_to_executable_tasks`. In
+            the inner list, the order of operations is READ, COMPUTE, and WRITE.
+
+    Returns:
+        A graph where each node is a _DAGOperationGraphNode. The key is `task_idx`,
+        the index to retrieve its task from `idx_to_task`, and the value is a
+        dictionary that maps the _DAGNodeOperationType (READ, COMPUTE, or WRITE)
+        to the corresponding _DAGOperationGraphNode
+    """
+    assert idx_to_task
+    graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]] = {}
+
+    for _, operation_nodes_list in actor_to_operation_nodes.items():
+        prev_compute_node = None
+        for operation_nodes in operation_nodes_list:
+            task_idx = operation_nodes[0].task_idx
+            read_node, compute_node, write_node = (
+                operation_nodes[0],
+                operation_nodes[1],
+                operation_nodes[2],
+            )
+            # Add edges from READ to COMPUTE, and from COMPUTE to WRITE, which
+            # belong to the same task.
+            _add_edge(read_node, compute_node)
+            _add_edge(compute_node, write_node)
+            # Add an edge from COMPUTE with `bind_index` i to COMPUTE with
+            # `bind_index` i+1 if they belong to the same actor.
+            if prev_compute_node is not None:
+                _add_edge(prev_compute_node, compute_node, "", True)
+            prev_compute_node = compute_node
+            assert task_idx not in graph
+            graph[task_idx] = {
+                _DAGNodeOperationType.READ: read_node,
+                _DAGNodeOperationType.COMPUTE: compute_node,
+                _DAGNodeOperationType.WRITE: write_node,
+            }
+
+    # Import `ray.dag` here to avoid circular import.
+    from ray.dag import ClassMethodNode, CollectiveOutputNode, MultiOutputNode
+    from ray.dag.collective_node import _CollectiveOperation
+
+    # Add an edge from WRITE of the writer task to READ of the reader task.
+    # Set synchronous nodes for accelerator P2P operations.
+    for task_idx, task in idx_to_task.items():
+        if not (
+            isinstance(task.dag_node, ClassMethodNode)
+            or isinstance(task.dag_node, CollectiveOutputNode)
+        ):
+            # The graph is used to generate an execution schedule for each actor.
+            # The edge from the InputNode has no impact on the final execution
+            # schedule.
+            continue
+        if (
+            isinstance(task.dag_node, ClassMethodNode)
+            and task.dag_node.is_class_method_output
+        ):
+            # Class method output node dependencies are handled at its upstream:
+            # i.e., class method node
+            continue
+        for downstream_task_idx in task.downstream_task_idxs:
+            downstream_dag_node = idx_to_task[downstream_task_idx].dag_node
+            if isinstance(downstream_dag_node, MultiOutputNode):
+                continue
+            write_node = graph[task_idx][_DAGNodeOperationType.WRITE]
+            if (
+                isinstance(downstream_dag_node, ClassMethodNode)
+                and downstream_dag_node.is_class_method_output
+            ):
+                consumer_idxs = idx_to_task[downstream_task_idx].downstream_task_idxs
+                for consumer_idx in consumer_idxs:
+                    if consumer_idx in graph:
+                        read_node = graph[consumer_idx][_DAGNodeOperationType.READ]
+                        _add_edge(
+                            write_node,
+                            read_node,
+                            "accelerator" if write_node.requires_accelerator else "shm",
+                        )
+                        if write_node.requires_accelerator:
+                            idxs = {
+                                (task_idx, _DAGNodeOperationType.WRITE),
+                                (consumer_idx, _DAGNodeOperationType.READ),
+                            }
+                            for node in [write_node, read_node]:
+                                node.sync_idxs.update(idxs)
+                continue
+            read_node = graph[downstream_task_idx][_DAGNodeOperationType.READ]
+            _add_edge(
+                write_node,
+                read_node,
+                "accelerator" if write_node.requires_accelerator else "shm",
+            )
+            if write_node.requires_accelerator:
+                idxs = {
+                    (task_idx, _DAGNodeOperationType.WRITE),
+                    (downstream_task_idx, _DAGNodeOperationType.READ),
+                }
+                for node in [write_node, read_node]:
+                    node.sync_idxs.update(idxs)
+
+    # Set synchronous nodes for accelerator collective operations.
+    collective_op_to_idxs: Dict[
+        _CollectiveOperation, Set[Tuple[int, _DAGNodeOperationType]]
+    ] = defaultdict(set)
+    for task_idx, task in idx_to_task.items():
+        if (
+            isinstance(task.dag_node, CollectiveOutputNode)
+            and not task.dag_node.is_class_method_output
+        ):
+            collective_op_to_idxs[task.dag_node.collective_op].add(
+                (task_idx, _DAGNodeOperationType.COMPUTE)
+            )
+    for idxs in collective_op_to_idxs.values():
+        for task_idx, op_type in idxs:
+            graph[task_idx][op_type].sync_idxs = idxs
+
+    return graph
+
+
+def _actor_viz_label(actor: "ray.actor.ActorHandle"):
+    """
+    Returns the label of an actor in the visualization of the execution schedule.
+
+    Args:
+        actor: The actor to be represented.
+    """
+    class_name = actor._ray_actor_creation_function_descriptor.class_name
+    actor_id = actor._ray_actor_id.hex()
+    return f"Actor class name: {class_name}\nActor ID: {actor_id}"
+
+
+def _node_viz_id_and_label(
+    node: _DAGOperationGraphNode, idx: int, optimized_index: int
+):
+    """
+    Returns the visualization id and label of a node. The visualization id is unique
+    across all nodes.
+
+    Args:
+        node: The node to be represented.
+        idx: The index of the node in the execution schedule.
+        optimized_index: The index of the node in the optimized execution schedule.
+    """
+    node_viz_label = node.viz_str() + f" {idx},{optimized_index}"
+    node_viz_id = f"{node._actor_id}_{node_viz_label}"
+    return node_viz_id, node_viz_label
+
+
+def _visualize_execution_schedule(
+    actor_to_execution_schedule: Dict[
+        "ray.actor.ActorHandle", List[_DAGOperationGraphNode]
+    ],
+    actor_to_overlapped_schedule: Optional[
+        Dict["ray.actor.ActorHandle", List[_DAGOperationGraphNode]]
+    ],
+    graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]],
+):
+    """
+    Visualize the execution schedule for each actor.
+
+    The visualization will be saved as a PNG file named `compiled_graph_schedule.png`.
+    Details of the visualization: # noqa
+
+        Node description format:
+            [<task_index>] <method_name> <operation> <orig_index>, <overlap_index>
+
+        Node description fields:
+            operation: is R(READ), C(COMPUTE), or W(WRITE)
+            orig_index: the index in the original execution schedule
+            overlap_index: the index in the overlap-communication optimized execution schedule
+            If this is different from orig_index, the node is highlighted in red color
+
+        Node grouping:
+            The nodes belonging to the same actor are grouped in the same rectangle
+            The actor class name and the actor id are shown in the rectangle
+
+        Edges:
+            black color (without label): data dependency
+            black color (annotated with "shm"): shared memory channel
+            blue color (annotated with "accelerator): accelerator channel
+            dashed edge: control dependency between compute operations
+
+    Args:
+        actor_to_execution_schedule: A dictionary that maps an actor handle to
+            the execution schedule which is a list of operation nodes.
+        actor_to_overlapped_schedule: A dictionary that maps an actor handle to the
+            optimized execution schedule which is a list of operation nodes.
+        graph: A graph where each node is a _DAGOperationGraphNode. The key is
+            `task_idx`, the index to retrieve its task from `idx_to_task`, and
+            the value is a dictionary that maps the _DAGNodeOperationType (READ,
+            COMPUTE, or WRITE) to the corresponding _DAGOperationGraphNode. It is
+            generated by `_build_dag_node_operation_graph`.
+    """
+    try:
+        import graphviz
+    except ImportError:
+        raise ImportError(
+            "Please install graphviz to visualize the execution schedule. "
+            "You can install it by running `pip install graphviz`."
+        )
+
+    dot = graphviz.Digraph(comment="DAG")
+    # A dictionary that maps a node to its visualization id
+    node_to_viz_id: Dict[_DAGOperationGraphNode, str] = {}
+
+    if actor_to_overlapped_schedule is None:
+        # TODO(rui): make the visualization more concise by only displaying
+        # the original schedule
+        actor_to_overlapped_schedule = actor_to_execution_schedule
+    for actor, execution_nodes in actor_to_execution_schedule.items():
+        overlapped_schedule = actor_to_overlapped_schedule[actor]
+        node_to_optimized_index = {
+            node: i for i, node in enumerate(overlapped_schedule)
+        }
+
+        actor_id = actor._ray_actor_id.hex()
+        with dot.subgraph(name=f"cluster_{actor_id}") as subgraph:
+            subgraph.attr(rank=actor_id, label=_actor_viz_label(actor))
+            for i, node in enumerate(execution_nodes):
+                optimized_index = node_to_optimized_index.get(node)
+                node_viz_id, node_viz_label = _node_viz_id_and_label(
+                    node, i, optimized_index
+                )
+                color = "red" if optimized_index != i else "black"
+                subgraph.node(node_viz_id, node_viz_label, color=color)
+                node_to_viz_id[node] = node_viz_id
+
+    for actor, execution_nodes in actor_to_execution_schedule.items():
+        for i, node in enumerate(execution_nodes):
+            node_viz_id = node_to_viz_id[node]
+            for out_edge, viz_info in node.out_edges.items():
+                label, control_dependency = viz_info
+                out_task_idx, out_op_type = out_edge
+                out_node = graph[out_task_idx][out_op_type]
+                out_node_viz_id = node_to_viz_id[out_node]
+                color = "blue" if label == "accelerator" else "black"
+                style = "dashed" if control_dependency else "solid"
+                dot.edge(
+                    node_viz_id, out_node_viz_id, label=label, color=color, style=style
+                )
+
+    # Add legend
+    with dot.subgraph(name="cluster_legend") as legend:
+        legend.attr(label="Legend", labelloc="t", fontsize="20", bgcolor="lightgrey")
+
+        # Single node and its explanation
+        legend.node("example_node", "[0] bwd C 10,10\n")
+        explanation = (
+            '<<TABLE BORDER="0" CELLBORDER="0" CELLSPACING="0">'  # noqa
+            '<TR><TD ALIGN="LEFT"><B>Node description format:</B></TD></TR>'
+            '<TR><TD ALIGN="LEFT">[&lt;task_index&gt;] &lt;method_name&gt; &lt;operation&gt; &lt;orig_index&gt;, &lt;overlap_index&gt;</TD></TR>'  # noqa
+            "<TR><TD></TD></TR>"
+            '<TR><TD ALIGN="LEFT"><B>Node description fields:</B></TD></TR>'
+            '<TR><TD ALIGN="LEFT">operation: is R(READ), C(COMPUTE), or W(WRITE)</TD></TR>'  # noqa
+            '<TR><TD ALIGN="LEFT">orig_index: the index in the original execution schedule</TD></TR>'  # noqa
+            '<TR><TD ALIGN="LEFT">overlap_index: the index in the overlap-communication optimized execution schedule</TD></TR>'  # noqa
+            '<TR><TD ALIGN="LEFT">If this is different from orig_index, the node is highlighted in <FONT COLOR="red">red color</FONT></TD></TR>'  # noqa
+            "<TR><TD></TD></TR>"
+            '<TR><TD ALIGN="LEFT"><B>Node grouping:</B></TD></TR>'
+            '<TR><TD ALIGN="LEFT">The nodes belonging to the same actor are grouped in the same rectangle</TD></TR>'  # noqa
+            '<TR><TD ALIGN="LEFT">The actor class name and the actor id are shown in the rectangle</TD></TR>'  # noqa
+            "<TR><TD></TD></TR>"
+            '<TR><TD ALIGN="LEFT"><B>Edges:</B></TD></TR>'
+            '<TR><TD ALIGN="LEFT">black color (without label): data dependency</TD></TR>'  # noqa
+            '<TR><TD ALIGN="LEFT">black color (annotated with "shm"): shared memory channel</TD></TR>'  # noqa
+            '<TR><TD ALIGN="LEFT"><FONT COLOR="blue">blue color</FONT> (annotated with "accelerator): accelerator channel</TD></TR>'  # noqa
+            '<TR><TD ALIGN="LEFT">dashed edge: control dependency between compute operations</TD></TR>'  # noqa
+            "</TABLE>>"
+        )
+
+        legend.node("example_explanation", explanation, shape="plaintext")
+        legend.edge("example_node", "example_explanation", style="invis")
+
+    logger.info(
+        "Writing compiled graph schedule visualization "
+        "to compiled_graph_schedule.png"
+    )
+    dot.render("compiled_graph_schedule", format="png", view=False)
+
+
+def _generate_actor_to_execution_schedule(
+    graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]],
+) -> Dict["ray.actor.ActorHandle", List[_DAGOperationGraphNode]]:
+    """
+    Generate an execution schedule for each actor. The schedule is a list of
+    operation nodes to be executed. The function uses a topological sort
+    algorithm to generate the schedule.
+
+    Args:
+        graph: A graph where each node is a _DAGOperationGraphNode. The key is
+            `task_idx`, the index to retrieve its task from `idx_to_task`, and
+            the value is a dictionary that maps the _DAGNodeOperationType (READ,
+            COMPUTE, or WRITE) to the corresponding _DAGOperationGraphNode. It is
+            generated by `_build_dag_node_operation_graph`.
+
+    Returns:
+        actor_to_execution_schedule: A dictionary that maps an actor handle to
+            the execution schedule which is a list of operation nodes to be
+            executed.
+    """
+
+    # Mapping from the actor handle to the execution schedule which is a list
+    # of operations to be executed.
+    actor_to_execution_schedule: Dict[
+        "ray.actor.ActorHandle", List[_DAGOperationGraphNode]
+    ] = defaultdict(list)
+
+    # A dictionary mapping an actor id to a list of candidate nodes. The list
+    # is maintained as a priority queue, so the head of the queue, i.e.,
+    # `candidates[0]`, is the node with the smallest `bind_index`.
+    actor_to_candidates: Dict[
+        "ray._raylet.ActorID", List[_DAGOperationGraphNode]
+    ] = defaultdict(list)
+    for _, node_dict in graph.items():
+        for _, node in node_dict.items():
+            # A node with a zero in-degree edge means all of its dependencies
+            # have been satisfied, including both data and control dependencies.
+            # Therefore, it is a candidate for execution.
+            if node.in_degree == 0:
+                _push_candidate_node_if_ready(actor_to_candidates, graph, node)
+
+    visited_nodes = set()
+
+    # Use topological sort algorithm to generate the execution schedule.
+    while True:
+        # Select a list of nodes to be executed. There are three cases:
+        # 1. If a selected node is not an accelerator operation, only itself is returned.
+        # 2. If a selected node is an accelerator write operation, the corresponding accelerator
+        #    read operations are also returned.
+        # 3. If a selected node is an accelerator collective operation, all the nodes in
+        #    its collective operation are returned.
+        nodes = _select_next_nodes(actor_to_candidates, graph)
+        if nodes is None:
+            break
+        # Add the selected nodes to the execution schedule.
+        for node in nodes:
+            assert node not in visited_nodes
+            visited_nodes.add(node)
+            actor_to_execution_schedule[node.actor_handle].append(node)
+        # Update the in-degree of the downstream nodes.
+        for node in nodes:
+            for out_node_task_idx, out_node_type in node.out_edges:
+                out_node = graph[out_node_task_idx][out_node_type]
+                if out_node in visited_nodes:
+                    # If the downstream node is already visited, it has been added
+                    # to the execution schedule. They are the accelerator read nodes in
+                    # case 2.
+                    continue
+                out_node.in_edges.pop((node.task_idx, node.operation.type))
+                if out_node.in_degree == 0:
+                    _push_candidate_node_if_ready(actor_to_candidates, graph, out_node)
+    assert len(visited_nodes) == len(graph) * 3, "Expected all nodes to be visited"
+
+    return actor_to_execution_schedule
+
+
+def _generate_overlapped_execution_schedule(
+    actor_to_execution_schedule: Dict[
+        "ray.actor.ActorHandle", List[_DAGOperationGraphNode]
+    ],
+) -> Dict["ray.actor.ActorHandle", List[_DAGOperationGraphNode]]:
+    """
+    From an existing execution schedule, generate a new schedule by overlapping
+    computation and communication.
+
+    Currently, the algorithm generates a new schedule for each actor as follows:
+    For each accelerator read operation (i.e., recv), scan backwards to find the nearest
+    compute node to swap with so that the accelerator read operation can be overlapped
+    with computation.
+
+    Collective operations are not yet supported.
+
+    Args:
+        actor_to_execution_schedule: A dictionary that maps an actor handle to
+            the existing execution schedule for the actor. The schedule is a list
+            is a list of operations to be executed.
+
+    Returns:
+        A dictionary that maps an actor handle to the overlapped execution schedule
+        for the actor.
+    """
+
+    actor_to_overlapped_schedule: Dict[
+        "ray.actor.ActorHandle", List[_DAGOperationGraphNode]
+    ] = copy.deepcopy(actor_to_execution_schedule)
+    for overlapped_schedule in actor_to_overlapped_schedule.values():
+        for i in range(len(overlapped_schedule)):
+            if (
+                overlapped_schedule[i].operation.type == _DAGNodeOperationType.READ
+                and overlapped_schedule[i].requires_accelerator
+            ):
+                # For each accelerator read operation (i.e., recv), scan backwards
+                # to find the nearest compute node to swap with so that
+                # the accelerator read operation can be overlapped with computation.
+                for j in range(i - 1, -1, -1):
+                    if (
+                        overlapped_schedule[j].operation.type
+                        == _DAGNodeOperationType.COMPUTE
+                    ):
+                        # Found a desired compute operation, make the swap
+                        accelerator_read_op = overlapped_schedule[i]
+                        prev_ops = overlapped_schedule[j:i]
+                        overlapped_schedule[j + 1 : i + 1] = prev_ops
+                        overlapped_schedule[j] = accelerator_read_op
+                        break
+                    if (
+                        overlapped_schedule[j].operation.type
+                        == _DAGNodeOperationType.READ
+                        or overlapped_schedule[j].operation.type
+                        == _DAGNodeOperationType.WRITE
+                    ) and overlapped_schedule[j].requires_accelerator:
+                        # Found an accelerator read/write operation, skip the overlap
+                        # optimization to keep relative order of accelerator operations
+                        break
+    return actor_to_overlapped_schedule
+
+
+def _extract_execution_schedule(
+    actor_to_execution_schedule: Dict[
+        "ray.actor.ActorHandle", List[_DAGOperationGraphNode]
+    ],
+) -> Dict["ray.actor.ActorHandle", List[_DAGNodeOperation]]:
+    """
+    Extract _DAGNodeOperation from _DAGOperationGraphNode in the schedule
+    and discard unnecessary information.
+    """
+    return {
+        actor: [node.operation for node in nodes]
+        for actor, nodes in actor_to_execution_schedule.items()
+    }

lib/python3.12/site-packages/ray/dag/dag_operation_future.py

@@ -0,0 +1,144 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Generic, TypeVar
+
+from ray.experimental.channel.accelerator_context import AcceleratorContext
+from ray.util.annotations import DeveloperAPI
+
+T = TypeVar("T")
+
+
+@DeveloperAPI
+class DAGOperationFuture(ABC, Generic[T]):
+    """
+    A future representing the result of a DAG operation.
+
+    This is an abstraction that is internal to each actor,
+    and is not exposed to the DAG caller.
+    """
+
+    @abstractmethod
+    def wait(self):
+        """
+        Wait for the future and return the result of the operation.
+        """
+        raise NotImplementedError
+
+
+@DeveloperAPI
+class ResolvedFuture(DAGOperationFuture):
+    """
+    A future that is already resolved. Calling `wait()` on this will
+    immediately return the result without blocking.
+    """
+
+    def __init__(self, result):
+        """
+        Initialize a resolved future.
+
+        Args:
+            result: The result of the future.
+        """
+        self._result = result
+
+    def wait(self):
+        """
+        Wait and immediately return the result. This operation will not block.
+        """
+        return self._result
+
+
+@DeveloperAPI
+class GPUFuture(DAGOperationFuture[Any]):
+    """
+    A future for a GPU event on a CUDA stream.
+
+    This future wraps a buffer, and records an event on the given stream
+    when it is created. When the future is waited on, it makes the current
+    CUDA stream wait on the event, then returns the buffer.
+
+    The buffer must be a GPU tensor produced by an earlier operation launched
+    on the given stream, or it could be CPU data. Then the future guarantees
+    that when the wait() returns, the buffer is ready on the current stream.
+
+    The `wait()` does not block CPU.
+    """
+
+    # Caching GPU futures ensures CUDA events associated with futures are properly
+    # destroyed instead of relying on garbage collection. The CUDA event contained
+    # in a GPU future is destroyed right before removing the future from the cache.
+    # The dictionary key is the future ID, which is the task idx of the dag operation
+    # that produced the future. When a future is created, it is immediately added to
+    # the cache. When a future has been waited on, it is removed from the cache.
+    # When adding a future, if its ID is already a key in the cache, the old future
+    # is removed. This can happen when an exception is thrown in a previous execution
+    # of the dag, in which case the old future is never waited on.
+    # Upon dag teardown, all pending futures produced by the dag are removed.
+    gpu_futures: Dict[int, "GPUFuture"] = {}
+
+    @staticmethod
+    def add_gpu_future(fut_id: int, fut: "GPUFuture") -> None:
+        """
+        Cache the GPU future.
+        Args:
+            fut_id: GPU future ID.
+            fut: GPU future to be cached.
+        """
+        if fut_id in GPUFuture.gpu_futures:
+            # The old future was not waited on because of an execution exception.
+            GPUFuture.gpu_futures.pop(fut_id).destroy_event()
+        GPUFuture.gpu_futures[fut_id] = fut
+
+    @staticmethod
+    def remove_gpu_future(fut_id: int) -> None:
+        """
+        Remove the cached GPU future and destroy its CUDA event.
+        Args:
+            fut_id: GPU future ID.
+        """
+        if fut_id in GPUFuture.gpu_futures:
+            GPUFuture.gpu_futures.pop(fut_id).destroy_event()
+
+    def __init__(self, buf: Any, fut_id: int, stream: Any = None):
+        """
+        Initialize a GPU future on the given stream.
+
+        Args:
+            buf: The buffer to return when the future is resolved.
+            fut_id: The future ID to cache the future.
+            stream: The torch stream to record the event on, this event is waited
+                on when the future is resolved. If None, the current stream is used.
+        """
+        if stream is None:
+            stream = AcceleratorContext.get().current_stream()
+
+        self._buf = buf
+        self._event = AcceleratorContext.get().create_event()
+        self._event.record(stream)
+        self._fut_id = fut_id
+        self._waited: bool = False
+
+        # Cache the GPU future such that its CUDA event is properly destroyed.
+        GPUFuture.add_gpu_future(fut_id, self)
+
+    def wait(self) -> Any:
+        """
+        Wait for the future on the current CUDA stream and return the result from
+        the GPU operation. This operation does not block CPU.
+        """
+        current_stream = AcceleratorContext.get().current_stream()
+        if not self._waited:
+            self._waited = True
+            current_stream.wait_event(self._event)
+            # Destroy the CUDA event after it is waited on.
+            GPUFuture.remove_gpu_future(self._fut_id)
+
+        return self._buf
+
+    def destroy_event(self) -> None:
+        """
+        Destroy the CUDA event associated with this future.
+        """
+        if self._event is None:
+            return
+
+        self._event = None

lib/python3.12/site-packages/ray/dag/format_utils.py

@@ -0,0 +1,155 @@
+from ray.dag import DAGNode
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+def get_dag_node_str(
+    dag_node: DAGNode,
+    body_line,
+):
+    indent = _get_indentation()
+    other_args_to_resolve_lines = _get_other_args_to_resolve_lines(
+        dag_node._bound_other_args_to_resolve
+    )
+    return (
+        f"({dag_node.__class__.__name__}, {dag_node._stable_uuid})(\n"
+        f"{indent}body={body_line}\n"
+        f"{indent}args={_get_args_lines(dag_node._bound_args)}\n"
+        f"{indent}kwargs={_get_kwargs_lines(dag_node._bound_kwargs)}\n"
+        f"{indent}options={_get_options_lines(dag_node._bound_options)}\n"
+        f"{indent}other_args_to_resolve={other_args_to_resolve_lines}\n"
+        f")"
+    )
+
+
+def _get_indentation(num_spaces=4):
+    return " " * num_spaces
+
+
+def _get_args_lines(bound_args):
+    """Pretty prints bounded args of a DAGNode, and recursively handle
+    DAGNode in list / dict containers.
+    """
+    indent = _get_indentation()
+    lines = []
+    for arg in bound_args:
+        if isinstance(arg, DAGNode):
+            node_repr_lines = str(arg).split("\n")
+            for node_repr_line in node_repr_lines:
+                lines.append(f"{indent}" + node_repr_line)
+        elif isinstance(arg, list):
+            for ele in arg:
+                node_repr_lines = str(ele).split("\n")
+                for node_repr_line in node_repr_lines:
+                    lines.append(f"{indent}" + node_repr_line)
+        elif isinstance(arg, dict):
+            for _, val in arg.items():
+                node_repr_lines = str(val).split("\n")
+                for node_repr_line in node_repr_lines:
+                    lines.append(f"{indent}" + node_repr_line)
+        # TODO: (jiaodong) Handle nested containers and other obj types
+        else:
+            lines.append(f"{indent}" + str(arg) + ", ")
+
+    if len(lines) == 0:
+        args_line = "[]"
+    else:
+        args_line = "["
+        for args in lines:
+            args_line += f"\n{indent}{args}"
+        args_line += f"\n{indent}]"
+
+    return args_line
+
+
+def _get_kwargs_lines(bound_kwargs):
+    """Pretty prints bounded kwargs of a DAGNode, and recursively handle
+    DAGNode in list / dict containers.
+    """
+    # TODO: (jiaodong) Nits, we're missing keys and indentation was a bit off.
+    if not bound_kwargs:
+        return "{}"
+    indent = _get_indentation()
+    kwargs_lines = []
+    for key, val in bound_kwargs.items():
+        if isinstance(val, DAGNode):
+            node_repr_lines = str(val).split("\n")
+            for index, node_repr_line in enumerate(node_repr_lines):
+                if index == 0:
+                    kwargs_lines.append(
+                        f"{indent}{key}:" + f"{indent}" + node_repr_line
+                    )
+                else:
+                    kwargs_lines.append(f"{indent}{indent}" + node_repr_line)
+
+        elif isinstance(val, list):
+            for ele in val:
+                node_repr_lines = str(ele).split("\n")
+                for node_repr_line in node_repr_lines:
+                    kwargs_lines.append(f"{indent}" + node_repr_line)
+        elif isinstance(val, dict):
+            for _, inner_val in val.items():
+                node_repr_lines = str(inner_val).split("\n")
+                for node_repr_line in node_repr_lines:
+                    kwargs_lines.append(f"{indent}" + node_repr_line)
+        # TODO: (jiaodong) Handle nested containers and other obj types
+        else:
+            kwargs_lines.append(val)
+
+    if len(kwargs_lines) > 0:
+        kwargs_line = "{"
+        for line in kwargs_lines:
+            kwargs_line += f"\n{indent}{line}"
+        kwargs_line += f"\n{indent}}}"
+    else:
+        kwargs_line = "{}"
+
+    return kwargs_line
+
+
+def _get_options_lines(bound_options):
+    """Pretty prints .options() in DAGNode. Only prints non-empty values."""
+    if not bound_options:
+        return "{}"
+    indent = _get_indentation()
+    options_lines = []
+    for key, val in bound_options.items():
+        if val:
+            options_lines.append(f"{indent}{key}: " + str(val))
+
+    options_line = "{"
+    for line in options_lines:
+        options_line += f"\n{indent}{line}"
+    options_line += f"\n{indent}}}"
+    return options_line
+
+
+def _get_other_args_to_resolve_lines(other_args_to_resolve):
+    if not other_args_to_resolve:
+        return "{}"
+    indent = _get_indentation()
+    other_args_to_resolve_lines = []
+    for key, val in other_args_to_resolve.items():
+        if isinstance(val, DAGNode):
+            node_repr_lines = str(val).split("\n")
+            for index, node_repr_line in enumerate(node_repr_lines):
+                if index == 0:
+                    other_args_to_resolve_lines.append(
+                        f"{indent}{key}:"
+                        + f"{indent}"
+                        + "\n"
+                        + f"{indent}{indent}{indent}"
+                        + node_repr_line
+                    )
+                else:
+                    other_args_to_resolve_lines.append(
+                        f"{indent}{indent}" + node_repr_line
+                    )
+        else:
+            other_args_to_resolve_lines.append(f"{indent}{key}: " + str(val))
+
+    other_args_to_resolve_line = "{"
+    for line in other_args_to_resolve_lines:
+        other_args_to_resolve_line += f"\n{indent}{line}"
+    other_args_to_resolve_line += f"\n{indent}}}"
+    return other_args_to_resolve_line

lib/python3.12/site-packages/ray/dag/function_node.py

@@ -0,0 +1,59 @@
+from typing import Any, Dict, List
+
+import ray
+from ray.dag.dag_node import DAGNode
+from ray.dag.format_utils import get_dag_node_str
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class FunctionNode(DAGNode):
+    """Represents a bound task node in a Ray task DAG."""
+
+    def __init__(
+        self,
+        func_body,
+        func_args,
+        func_kwargs,
+        func_options,
+        other_args_to_resolve=None,
+    ):
+        self._body = func_body
+        super().__init__(
+            func_args,
+            func_kwargs,
+            func_options,
+            other_args_to_resolve=other_args_to_resolve,
+        )
+
+    def _copy_impl(
+        self,
+        new_args: List[Any],
+        new_kwargs: Dict[str, Any],
+        new_options: Dict[str, Any],
+        new_other_args_to_resolve: Dict[str, Any],
+    ):
+        return FunctionNode(
+            self._body,
+            new_args,
+            new_kwargs,
+            new_options,
+            other_args_to_resolve=new_other_args_to_resolve,
+        )
+
+    def _execute_impl(self, *args, **kwargs):
+        """Executor of FunctionNode by ray.remote().
+
+        Args and kwargs are to match base class signature, but not in the
+        implementation. All args and kwargs should be resolved and replaced
+        with value in bound_args and bound_kwargs via bottom-up recursion when
+        current node is executed.
+        """
+        return (
+            ray.remote(self._body)
+            .options(**self._bound_options)
+            .remote(*self._bound_args, **self._bound_kwargs)
+        )
+
+    def __str__(self) -> str:
+        return get_dag_node_str(self, str(self._body))

lib/python3.12/site-packages/ray/dag/input_node.py

@@ -0,0 +1,321 @@
+from typing import Any, Dict, List, Optional, Union
+
+from ray.dag import DAGNode
+from ray.dag.format_utils import get_dag_node_str
+from ray.experimental.gradio_utils import type_to_string
+from ray.util.annotations import DeveloperAPI
+
+IN_CONTEXT_MANAGER = "__in_context_manager__"
+
+
+@DeveloperAPI
+class InputNode(DAGNode):
+    r"""Ray dag node used in DAG building API to mark entrypoints of a DAG.
+
+    Should only be function or class method. A DAG can have multiple
+    entrypoints, but only one instance of InputNode exists per DAG, shared
+    among all DAGNodes.
+
+    Example:
+
+    .. code-block::
+
+                m1.forward
+                /       \
+        dag_input     ensemble -> dag_output
+                \       /
+                m2.forward
+
+    In this pipeline, each user input is broadcasted to both m1.forward and
+    m2.forward as first stop of the DAG, and authored like
+
+    .. code-block:: python
+
+        import ray
+
+        @ray.remote
+        class Model:
+            def __init__(self, val):
+                self.val = val
+            def forward(self, input):
+                return self.val * input
+
+        @ray.remote
+        def combine(a, b):
+            return a + b
+
+        with InputNode() as dag_input:
+            m1 = Model.bind(1)
+            m2 = Model.bind(2)
+            m1_output = m1.forward.bind(dag_input[0])
+            m2_output = m2.forward.bind(dag_input.x)
+            ray_dag = combine.bind(m1_output, m2_output)
+
+        # Pass mix of args and kwargs as input.
+        ray_dag.execute(1, x=2) # 1 sent to m1, 2 sent to m2
+
+        # Alternatively user can also pass single data object, list or dict
+        # and access them via list index, object attribute or dict key str.
+        ray_dag.execute(UserDataObject(m1=1, m2=2))
+        # dag_input.m1, dag_input.m2
+        ray_dag.execute([1, 2])
+        # dag_input[0], dag_input[1]
+        ray_dag.execute({"m1": 1, "m2": 2})
+        # dag_input["m1"], dag_input["m2"]
+    """
+
+    def __init__(
+        self,
+        *args,
+        input_type: Optional[Union[type, Dict[Union[int, str], type]]] = None,
+        _other_args_to_resolve=None,
+        **kwargs,
+    ):
+        """InputNode should only take attributes of validating and converting
+        input data rather than the input data itself. User input should be
+        provided via `ray_dag.execute(user_input)`.
+
+        Args:
+            input_type: Describes the data type of inputs user will be giving.
+                - if given through singular InputNode: type of InputNode
+                - if given through InputAttributeNodes: map of key -> type
+                Used when deciding what Gradio block to represent the input nodes with.
+            _other_args_to_resolve: Internal only to keep InputNode's execution
+                context throughput pickling, replacement and serialization.
+                User should not use or pass this field.
+        """
+        if len(args) != 0 or len(kwargs) != 0:
+            raise ValueError("InputNode should not take any args or kwargs.")
+
+        self.input_attribute_nodes = {}
+
+        self.input_type = input_type
+        if input_type is not None and isinstance(input_type, type):
+            if _other_args_to_resolve is None:
+                _other_args_to_resolve = {}
+            _other_args_to_resolve["result_type_string"] = type_to_string(input_type)
+
+        super().__init__([], {}, {}, other_args_to_resolve=_other_args_to_resolve)
+
+    def _copy_impl(
+        self,
+        new_args: List[Any],
+        new_kwargs: Dict[str, Any],
+        new_options: Dict[str, Any],
+        new_other_args_to_resolve: Dict[str, Any],
+    ):
+        return InputNode(_other_args_to_resolve=new_other_args_to_resolve)
+
+    def _execute_impl(self, *args, **kwargs):
+        """Executor of InputNode."""
+        # Catch and assert singleton context at dag execution time.
+        assert self._in_context_manager(), (
+            "InputNode is a singleton instance that should be only used in "
+            "context manager for dag building and execution. See the docstring "
+            "of class InputNode for examples."
+        )
+        # If user only passed in one value, for simplicity we just return it.
+        if len(args) == 1 and len(kwargs) == 0:
+            return args[0]
+
+        return DAGInputData(*args, **kwargs)
+
+    def _in_context_manager(self) -> bool:
+        """Return if InputNode is created in context manager."""
+        if (
+            not self._bound_other_args_to_resolve
+            or IN_CONTEXT_MANAGER not in self._bound_other_args_to_resolve
+        ):
+            return False
+        else:
+            return self._bound_other_args_to_resolve[IN_CONTEXT_MANAGER]
+
+    def set_context(self, key: str, val: Any):
+        """Set field in parent DAGNode attribute that can be resolved in both
+        pickle and JSON serialization
+        """
+        self._bound_other_args_to_resolve[key] = val
+
+    def __str__(self) -> str:
+        return get_dag_node_str(self, "__InputNode__")
+
+    def __getattr__(self, key: str):
+        assert isinstance(
+            key, str
+        ), "Please only access dag input attributes with str key."
+        if key not in self.input_attribute_nodes:
+            self.input_attribute_nodes[key] = InputAttributeNode(
+                self, key, "__getattr__"
+            )
+        return self.input_attribute_nodes[key]
+
+    def __getitem__(self, key: Union[int, str]) -> Any:
+        assert isinstance(key, (str, int)), (
+            "Please only use int index or str as first-level key to "
+            "access fields of dag input."
+        )
+
+        input_type = None
+        if self.input_type is not None and key in self.input_type:
+            input_type = type_to_string(self.input_type[key])
+
+        if key not in self.input_attribute_nodes:
+            self.input_attribute_nodes[key] = InputAttributeNode(
+                self, key, "__getitem__", input_type
+            )
+        return self.input_attribute_nodes[key]
+
+    def __enter__(self):
+        self.set_context(IN_CONTEXT_MANAGER, True)
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+    def get_result_type(self) -> str:
+        """Get type of the output of this DAGNode.
+
+        Generated by ray.experimental.gradio_utils.type_to_string().
+        """
+        if "result_type_string" in self._bound_other_args_to_resolve:
+            return self._bound_other_args_to_resolve["result_type_string"]
+
+
+@DeveloperAPI
+class InputAttributeNode(DAGNode):
+    """Represents partial access of user input based on an index (int),
+     object attribute or dict key (str).
+
+    Examples:
+
+        .. code-block:: python
+
+            with InputNode() as dag_input:
+                a = dag_input[0]
+                b = dag_input.x
+                ray_dag = add.bind(a, b)
+
+            # This makes a = 1 and b = 2
+            ray_dag.execute(1, x=2)
+
+            with InputNode() as dag_input:
+                a = dag_input[0]
+                b = dag_input[1]
+                ray_dag = add.bind(a, b)
+
+            # This makes a = 2 and b = 3
+            ray_dag.execute(2, 3)
+
+            # Alternatively, you can input a single object
+            # and the inputs are automatically indexed from the object:
+            # This makes a = 2 and b = 3
+            ray_dag.execute([2, 3])
+    """
+
+    def __init__(
+        self,
+        dag_input_node: InputNode,
+        key: Union[int, str],
+        accessor_method: str,
+        input_type: str = None,
+    ):
+        self._dag_input_node = dag_input_node
+        self._key = key
+        self._accessor_method = accessor_method
+        super().__init__(
+            [],
+            {},
+            {},
+            {
+                "dag_input_node": dag_input_node,
+                "key": key,
+                "accessor_method": accessor_method,
+                # Type of the input tied to this node. Used by
+                # gradio_visualize_graph.GraphVisualizer to determine which Gradio
+                # component should be used for this node.
+                "result_type_string": input_type,
+            },
+        )
+
+    def _copy_impl(
+        self,
+        new_args: List[Any],
+        new_kwargs: Dict[str, Any],
+        new_options: Dict[str, Any],
+        new_other_args_to_resolve: Dict[str, Any],
+    ):
+        return InputAttributeNode(
+            new_other_args_to_resolve["dag_input_node"],
+            new_other_args_to_resolve["key"],
+            new_other_args_to_resolve["accessor_method"],
+            new_other_args_to_resolve["result_type_string"],
+        )
+
+    def _execute_impl(self, *args, **kwargs):
+        """Executor of InputAttributeNode.
+
+        Args and kwargs are to match base class signature, but not in the
+        implementation. All args and kwargs should be resolved and replaced
+        with value in bound_args and bound_kwargs via bottom-up recursion when
+        current node is executed.
+        """
+
+        if isinstance(self._dag_input_node, DAGInputData):
+            return self._dag_input_node[self._key]
+        else:
+            # dag.execute() is called with only one arg, thus when an
+            # InputAttributeNode is executed, its dependent InputNode is
+            # resolved with original user input python object.
+            user_input_python_object = self._dag_input_node
+            if isinstance(self._key, str):
+                if self._accessor_method == "__getitem__":
+                    return user_input_python_object[self._key]
+                elif self._accessor_method == "__getattr__":
+                    return getattr(user_input_python_object, self._key)
+            elif isinstance(self._key, int):
+                return user_input_python_object[self._key]
+            else:
+                raise ValueError(
+                    "Please only use int index or str as first-level key to "
+                    "access fields of dag input."
+                )
+
+    def __str__(self) -> str:
+        return get_dag_node_str(self, f'["{self._key}"]')
+
+    def get_result_type(self) -> str:
+        """Get type of the output of this DAGNode.
+
+        Generated by ray.experimental.gradio_utils.type_to_string().
+        """
+        if "result_type_string" in self._bound_other_args_to_resolve:
+            return self._bound_other_args_to_resolve["result_type_string"]
+
+    @property
+    def key(self) -> Union[int, str]:
+        return self._key
+
+
+@DeveloperAPI
+class DAGInputData:
+    """If user passed multiple args and kwargs directly to dag.execute(), we
+    generate this wrapper for all user inputs as one object, accessible via
+    list index or object attribute key.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self._args = list(args)
+        self._kwargs = kwargs
+
+    def __getitem__(self, key: Union[int, str]) -> Any:
+        if isinstance(key, int):
+            # Access list args by index.
+            return self._args[key]
+        elif isinstance(key, str):
+            # Access kwarg by key.
+            return self._kwargs[key]
+        else:
+            raise ValueError(
+                "Please only use int index or str as first-level key to "
+                "access fields of dag input."
+            )

lib/python3.12/site-packages/ray/dag/output_node.py

@@ -0,0 +1,45 @@
+from typing import Any, Dict, List, Tuple, Union
+
+import ray
+from ray.dag import DAGNode
+from ray.dag.format_utils import get_dag_node_str
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class MultiOutputNode(DAGNode):
+    """Ray dag node used in DAG building API to mark the endpoint of DAG"""
+
+    def __init__(
+        self,
+        args: Union[List[DAGNode], Tuple[DAGNode]],
+        other_args_to_resolve: Dict[str, Any] = None,
+    ):
+        if isinstance(args, tuple):
+            args = list(args)
+        if not isinstance(args, list):
+            raise ValueError(f"Invalid input type for `args`, {type(args)}.")
+        super().__init__(
+            args,
+            {},
+            {},
+            other_args_to_resolve=other_args_to_resolve or {},
+        )
+
+    def _execute_impl(
+        self, *args, **kwargs
+    ) -> Union[ray.ObjectRef, "ray.actor.ActorHandle"]:
+        return self._bound_args
+
+    def _copy_impl(
+        self,
+        new_args: List[Any],
+        new_kwargs: Dict[str, Any],
+        new_options: Dict[str, Any],
+        new_other_args_to_resolve: Dict[str, Any],
+    ) -> "DAGNode":
+        """Return a copy of this node with the given new args."""
+        return MultiOutputNode(new_args, new_other_args_to_resolve)
+
+    def __str__(self) -> str:
+        return get_dag_node_str(self, "__MultiOutputNode__")

lib/python3.12/site-packages/ray/dag/py_obj_scanner.py

@@ -0,0 +1,103 @@
+import io
+import pickle  # noqa: F401
+from typing import Any, Dict, Generic, List, Tuple, Type, TypeVar, Union
+
+import ray
+from ray.dag.base import DAGNodeBase
+
+# Used in deserialization hooks to reference scanner instances.
+_instances: Dict[int, "_PyObjScanner"] = {}
+
+# Generic types for the scanner to transform from and to.
+SourceType = TypeVar("SourceType")
+TransformedType = TypeVar("TransformedType")
+
+
+def _get_node(instance_id: int, node_index: int) -> SourceType:
+    """Get the node instance.
+
+    Note: This function should be static and globally importable,
+    otherwise the serialization overhead would be very significant.
+    """
+    return _instances[instance_id]._replace_index(node_index)
+
+
+class _PyObjScanner(ray.cloudpickle.CloudPickler, Generic[SourceType, TransformedType]):
+    """Utility to find and replace the `source_type` in Python objects.
+
+    `source_type` can either be a single type or a tuple of multiple types.
+
+    The caller must first call `find_nodes()`, then compute a replacement table and
+    pass it to `replace_nodes`.
+
+    This uses cloudpickle under the hood, so all sub-objects that are not `source_type`
+    must be serializable.
+
+    Args:
+        source_type: the type(s) of object to find and replace. Default to DAGNodeBase.
+    """
+
+    def __init__(self, source_type: Union[Type, Tuple] = DAGNodeBase):
+        self.source_type = source_type
+        # Buffer to keep intermediate serialized state.
+        self._buf = io.BytesIO()
+        # List of top-level SourceType found during the serialization pass.
+        self._found = None
+        # List of other objects found during the serialization pass.
+        # This is used to store references to objects so they won't be
+        # serialized by cloudpickle.
+        self._objects = []
+        # Replacement table to consult during deserialization.
+        self._replace_table: Dict[SourceType, TransformedType] = None
+        _instances[id(self)] = self
+        super().__init__(self._buf)
+
+    def reducer_override(self, obj):
+        """Hook for reducing objects.
+
+        Objects of `self.source_type` are saved to `self._found` and a global map so
+        they can later be replaced.
+
+        All other objects fall back to the default `CloudPickler` serialization.
+        """
+        if isinstance(obj, self.source_type):
+            index = len(self._found)
+            self._found.append(obj)
+            return _get_node, (id(self), index)
+
+        return super().reducer_override(obj)
+
+    def find_nodes(self, obj: Any) -> List[SourceType]:
+        """
+        Serialize `obj` and store all instances of `source_type` found in `_found`.
+
+        Args:
+            obj: The object to scan for `source_type`.
+        Returns:
+            A list of all instances of `source_type` found in `obj`.
+        """
+        assert (
+            self._found is None
+        ), "find_nodes cannot be called twice on the same PyObjScanner instance."
+        self._found = []
+        self._objects = []
+        self.dump(obj)
+        return self._found
+
+    def replace_nodes(self, table: Dict[SourceType, TransformedType]) -> Any:
+        """Replace previously found DAGNodes per the given table."""
+        assert self._found is not None, "find_nodes must be called first"
+        self._replace_table = table
+        self._buf.seek(0)
+        return pickle.load(self._buf)
+
+    def _replace_index(self, i: int) -> SourceType:
+        return self._replace_table[self._found[i]]
+
+    def clear(self):
+        """Clear the scanner from the _instances"""
+        if id(self) in _instances:
+            del _instances[id(self)]
+
+    def __del__(self):
+        self.clear()

lib/python3.12/site-packages/ray/dag/utils.py

@@ -0,0 +1,66 @@
+from typing import Dict
+
+from ray.dag import (
+    ClassMethodNode,
+    ClassNode,
+    DAGNode,
+    FunctionNode,
+    InputAttributeNode,
+    InputNode,
+    MultiOutputNode,
+)
+
+
+class _DAGNodeNameGenerator(object):
+    """
+    Generate unique suffix for each given Node in the DAG.
+    Apply monotonic increasing id suffix for duplicated names.
+    """
+
+    def __init__(self):
+        self.name_to_suffix: Dict[str, int] = dict()
+
+    def get_node_name(self, node: DAGNode):
+        # InputNode should be unique.
+        if isinstance(node, InputNode):
+            return "INPUT_NODE"
+        if isinstance(node, MultiOutputNode):
+            return "MultiOutputNode"
+        # InputAttributeNode suffixes should match the user-defined key.
+        elif isinstance(node, InputAttributeNode):
+            return f"INPUT_ATTRIBUTE_NODE_{node._key}"
+
+        # As class, method, and function nodes may have duplicated names,
+        # generate unique suffixes for such nodes.
+        if isinstance(node, ClassMethodNode):
+            node_name = node.get_options().get("name", None) or node._method_name
+        elif isinstance(node, (ClassNode, FunctionNode)):
+            node_name = node.get_options().get("name", None) or node._body.__name__
+        # we use instance class name check here to avoid importing ServeNodes as
+        # serve components are not included in Ray Core.
+        elif type(node).__name__ in ("DeploymentNode", "DeploymentFunctionNode"):
+            node_name = node.get_deployment_name()
+        elif type(node).__name__ == "DeploymentFunctionExecutorNode":
+            node_name = node._deployment_function_handle.deployment_name
+        else:
+            raise ValueError(
+                "get_node_name() should only be called on DAGNode instances."
+            )
+
+        if node_name not in self.name_to_suffix:
+            self.name_to_suffix[node_name] = 0
+            return node_name
+        else:
+            self.name_to_suffix[node_name] += 1
+            suffix_num = self.name_to_suffix[node_name]
+
+            return f"{node_name}_{suffix_num}"
+
+    def reset(self):
+        self.name_to_suffix = dict()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.reset()

lib/python3.12/site-packages/ray/dag/vis_utils.py

@@ -0,0 +1,114 @@
+import os
+import tempfile
+
+from ray.dag import DAGNode
+from ray.dag.utils import _DAGNodeNameGenerator
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+def plot(dag: DAGNode, to_file=None):
+    if to_file is None:
+        tmp_file = tempfile.NamedTemporaryFile(suffix=".png")
+        to_file = tmp_file.name
+        extension = "png"
+    else:
+        _, extension = os.path.splitext(to_file)
+        if not extension:
+            extension = "png"
+        else:
+            extension = extension[1:]
+
+    graph = _dag_to_dot(dag)
+    graph.write(to_file, format=extension)
+
+    # Render the image directly if running inside a Jupyter notebook
+    try:
+        from IPython import display
+
+        return display.Image(filename=to_file)
+    except ImportError:
+        pass
+
+    # close temp file if needed
+    try:
+        tmp_file.close()
+    except NameError:
+        pass
+
+
+def _check_pydot_and_graphviz():
+    """Check if pydot and graphviz are installed.
+
+    pydot and graphviz are required for plotting. We check this
+    during runtime rather than adding them to Ray dependencies.
+
+    """
+    try:
+        import pydot
+    except ImportError:
+        raise ImportError(
+            "pydot is required to plot DAG, install it with `pip install pydot`."
+        )
+    try:
+        pydot.Dot.create(pydot.Dot())
+    except (OSError, pydot.InvocationException):
+        raise ImportError(
+            "graphviz is required to plot DAG, "
+            "download it from https://graphviz.gitlab.io/download/"
+        )
+
+
+def _get_nodes_and_edges(dag: DAGNode):
+    """Get all unique nodes and edges in the DAG.
+
+    A basic dfs with memorization to get all unique nodes
+    and edges in the DAG.
+    Unique nodes will be used to generate unique names,
+    while edges will be used to construct the graph.
+    """
+
+    edges = []
+    nodes = []
+
+    def _dfs(node):
+        nodes.append(node)
+        for child_node in node._get_all_child_nodes():
+            edges.append((child_node, node))
+        return node
+
+    dag.apply_recursive(_dfs)
+    return nodes, edges
+
+
+def _dag_to_dot(dag: DAGNode):
+    """Create a Dot graph from dag.
+
+    TODO(lchu):
+    1. add more Dot configs in kwargs,
+    e.g. rankdir, alignment, etc.
+    2. add more contents to graph,
+    e.g. args, kwargs and options of each node
+
+    """
+    # Step 0: check dependencies and init graph
+    _check_pydot_and_graphviz()
+    import pydot
+
+    graph = pydot.Dot(rankdir="LR")
+
+    # Step 1: generate unique name for each node in dag
+    nodes, edges = _get_nodes_and_edges(dag)
+    name_generator = _DAGNodeNameGenerator()
+    node_names = {}
+    for node in nodes:
+        node_names[node] = name_generator.get_node_name(node)
+
+    # Step 2: create graph with all the edges
+    for edge in edges:
+        graph.add_edge(pydot.Edge(node_names[edge[0]], node_names[edge[1]]))
+    # if there is only one node
+    if len(nodes) == 1 and len(edges) == 0:
+        graph.add_node(pydot.Node(node_names[nodes[0]]))
+
+    return graph

lib/python3.12/site-packages/ray/experimental/__init__.py

@@ -0,0 +1,11 @@
+from ray.experimental.dynamic_resources import set_resource
+from ray.experimental.gpu_object_manager import GPUObjectManager, wait_tensor_freed
+from ray.experimental.locations import get_local_object_locations, get_object_locations
+
+__all__ = [
+    "get_object_locations",
+    "get_local_object_locations",
+    "set_resource",
+    "GPUObjectManager",
+    "wait_tensor_freed",
+]

lib/python3.12/site-packages/ray/experimental/collective/__init__.py

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

lib/python3.12/site-packages/ray/experimental/collective/collective.py

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

lib/python3.12/site-packages/ray/experimental/collective/collective_tensor_transport.py

@@ -0,0 +1,199 @@
+from typing import TYPE_CHECKING, List, Optional
+
+import ray
+from ray.experimental.collective.tensor_transport_manager import (
+    TensorTransportManager,
+)
+from ray.util.collective.types import (
+    Backend,
+    CollectiveCommunicatorMetadata,
+    CollectiveTransportMetadata,
+)
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CollectiveTensorTransport(TensorTransportManager):
+    def __init__(self, tensor_transport_backend: Backend):
+        self._tensor_transport_backend = tensor_transport_backend
+
+    @property
+    def tensor_transport_backend(self) -> Backend:
+        return self._tensor_transport_backend
+
+    @staticmethod
+    def is_one_sided() -> bool:
+        return False
+
+    @staticmethod
+    def can_abort_transport() -> bool:
+        return False
+
+    def actor_has_tensor_transport(self, actor: "ray.actor.ActorHandle") -> bool:
+        from ray.experimental.collective import get_collective_groups
+
+        communicators = get_collective_groups(
+            [actor], backend=self.tensor_transport_backend
+        )
+        return len(communicators) > 0
+
+    @staticmethod
+    def extract_tensor_transport_metadata(
+        obj_id: str,
+        gpu_object: List["torch.Tensor"],
+    ) -> CollectiveTransportMetadata:
+        tensor_meta = []
+        device = None
+        if gpu_object:
+            device = gpu_object[0].device
+            for t in gpu_object:
+                if t.device.type != device.type:
+                    raise ValueError(
+                        "All tensors in an RDT object must have the same device type."
+                    )
+                tensor_meta.append((t.shape, t.dtype))
+        return CollectiveTransportMetadata(
+            tensor_meta=tensor_meta,
+            tensor_device=device,
+        )
+
+    @staticmethod
+    def get_tensor_transport_metadata(
+        src_actor: "ray.actor.ActorHandle",
+        obj_id: str,
+    ) -> CollectiveTransportMetadata:
+        def __ray_get_tensor_transport_metadata__(
+            self: "ray.actor.ActorHandle",
+            obj_id: str,
+        ) -> CollectiveTransportMetadata:
+
+            from ray._private.worker import global_worker
+
+            gpu_object_store = global_worker.gpu_object_manager.gpu_object_store
+            # NOTE: We do not specify a timeout here because the user task that returns
+            # it could take arbitrarily long and we don't want to trigger a spurious
+            # timeout.
+            gpu_object = gpu_object_store.wait_and_get_object(obj_id)
+            return CollectiveTensorTransport.extract_tensor_transport_metadata(
+                obj_id, gpu_object
+            )
+
+        # Submit a Ray actor task to the source actor to get the tensor metadata.
+        # The metadata is a list of tuples, where each tuple contains the shape and dtype
+        # of a tensor in the GPU object store. This function returns an ObjectRef that
+        # points to the tensor metadata.
+        # NOTE(swang): We put this task on the background thread to avoid tasks
+        # executing on the main thread blocking this task.
+
+        return src_actor.__ray_call__.options(concurrency_group="_ray_system").remote(
+            __ray_get_tensor_transport_metadata__, obj_id
+        )
+
+    @staticmethod
+    def get_communicator_metadata(
+        src_actor: "ray.actor.ActorHandle",
+        dst_actor: "ray.actor.ActorHandle",
+        backend: Optional[str] = None,
+    ) -> CollectiveCommunicatorMetadata:
+
+        from ray.experimental.collective import get_collective_groups
+
+        communicators = get_collective_groups(
+            [src_actor, dst_actor],
+            backend=backend,
+        )
+        # TODO(kevin85421): Support multiple communicators.
+        if len(communicators) == 0:
+            raise ValueError(
+                f"No communicators found for actors {src_actor} and {dst_actor}. "
+                "Create a communicator with "
+                "`ray.experimental.collective.create_collective_group` "
+                "before calling actor tasks. with non-default tensor_transport."
+            )
+        elif len(communicators) > 1:
+            raise ValueError(
+                f"There are {len(communicators)} possible communicators that contain actors {src_actor} and {dst_actor}. "
+                "Currently, RDT objects only support one communicator. Please make sure only "
+                "one communicator exists."
+            )
+        communicator = communicators[0]
+        src_rank = communicator.get_rank(src_actor)
+        if src_rank == -1:
+            raise ValueError(
+                f"Sender actor {src_actor} not found in communicator. "
+                "Please make sure the sender and receiver are in the same communicator."
+            )
+        dst_rank = communicator.get_rank(dst_actor)
+        if dst_rank == -1:
+            raise ValueError(
+                f"Receiver actor {dst_actor} not found in communicator. "
+                "Please make sure the sender and receiver are in the same communicator."
+            )
+
+        communicator_metadata = CollectiveCommunicatorMetadata(
+            communicator_name=communicator.name,
+            src_rank=src_rank,
+            dst_rank=dst_rank,
+        )
+        return communicator_metadata
+
+    @staticmethod
+    def recv_multiple_tensors(
+        tensors,
+        obj_id: str,
+        tensor_transport_metadata: CollectiveTransportMetadata,
+        communicator_metadata: CollectiveCommunicatorMetadata,
+    ):
+        from ray.util.collective import types
+        from ray.util.collective.collective import recv
+
+        assert isinstance(
+            tensor_transport_metadata, types.CollectiveTransportMetadata
+        ), "metadata must be a CollectiveTransportMetadata object for non-NIXL transport"
+        assert isinstance(
+            communicator_metadata, types.CollectiveCommunicatorMetadata
+        ), "metadata must be a CollectiveCommunicatorMetadata object for non-NIXL transport"
+
+        for tensor in tensors:
+            recv(
+                tensor,
+                communicator_metadata.src_rank,
+                communicator_metadata.communicator_name,
+            )
+
+    @staticmethod
+    def send_multiple_tensors(
+        tensors: List["torch.Tensor"],
+        tensor_transport_metadata: CollectiveTransportMetadata,
+        communicator_metadata: CollectiveCommunicatorMetadata,
+    ):
+        import ray.util.collective as collective
+
+        device = tensors[0].device if tensors else None
+
+        for tensor in tensors:
+            if tensor.device.type != device.type:
+                raise ValueError(
+                    f"tensor device {tensor.device} does not match device {device}"
+                )
+            collective.send(
+                tensor,
+                communicator_metadata.dst_rank,
+                communicator_metadata.communicator_name,
+            )
+
+    @staticmethod
+    def garbage_collect(
+        obj_id: str, tensor_transport_meta: CollectiveTransportMetadata
+    ):
+        pass
+
+    @staticmethod
+    def abort_transport(
+        obj_id: str,
+        communicator_metadata: CollectiveCommunicatorMetadata,
+    ):
+        raise NotImplementedError(
+            "Collective transport does not support abort_transport for now."
+        )