chore: import upstream snapshot with attribution

python/ray/dag/BUILD.bazel

@@ -0,0 +1,221 @@
+load("@rules_python//python:defs.bzl", "py_library", "py_test")
+load("//bazel:python.bzl", "doctest", "py_test_module_list")
+
+doctest(
+    files = glob(
+        ["**/*.py"],
+        exclude = ["**/experimental/**/*.py"],
+    ),
+    tags = ["team:core"],
+    deps = [":dag_lib"],
+)
+
+# This is a dummy test dependency that causes the above tests to be
+# re-run if any of these files changes.
+py_library(
+    name = "dag_lib",
+    srcs = glob(
+        ["**/*.py"],
+        exclude = ["tests/**/*.py"],
+    ),
+    visibility = [
+        "//python/ray/dag:__pkg__",
+        "//python/ray/dag:__subpackages__",
+        "//release:__pkg__",
+    ],
+)
+
+dag_tests_srcs = glob(["tests/**/*.py"])
+
+py_test(
+    name = "test_function_dag",
+    size = "small",
+    srcs = dag_tests_srcs,
+    tags = [
+        "exclusive",
+        "ray_dag_tests",
+        "team:core",
+    ],
+    deps = [":dag_lib"],
+)
+
+py_test(
+    name = "test_class_dag",
+    size = "small",
+    srcs = dag_tests_srcs,
+    tags = [
+        "exclusive",
+        "ray_dag_tests",
+        "team:core",
+    ],
+    deps = [":dag_lib"],
+)
+
+py_test(
+    name = "test_input_node",
+    size = "small",
+    srcs = dag_tests_srcs,
+    tags = [
+        "exclusive",
+        "ray_dag_tests",
+        "team:core",
+    ],
+    deps = [":dag_lib"],
+)
+
+py_test(
+    name = "test_output_node",
+    size = "small",
+    srcs = dag_tests_srcs,
+    tags = [
+        "exclusive",
+        "ray_dag_tests",
+        "team:core",
+    ],
+    deps = [":dag_lib"],
+)
+
+py_test(
+    name = "test_plot",
+    size = "small",
+    srcs = dag_tests_srcs,
+    tags = [
+        "exclusive",
+        "ray_dag_tests",
+        "team:core",
+    ],
+    deps = [":dag_lib"],
+)
+
+py_test(
+    name = "test_py_obj_scanner",
+    size = "small",
+    srcs = dag_tests_srcs,
+    tags = [
+        "exclusive",
+        "ray_dag_tests",
+        "team:core",
+    ],
+    deps = [":dag_lib"],
+)
+
+py_test_module_list(
+    size = "medium",
+    files = [
+        "tests/experimental/test_collective_dag.py",
+        "tests/experimental/test_dag_error_handling.py",
+        "tests/experimental/test_dag_visualization.py",
+        "tests/experimental/test_execution_schedule.py",
+        "tests/experimental/test_mocked_nccl_dag.py",
+        "tests/experimental/test_torch_tensor_dag.py",
+    ],
+    tags = [
+        "compiled_graphs",
+        "exclusive",
+        "no_windows",
+        "team:core",
+    ],
+    deps = ["//:ray_lib"],
+)
+
+py_test_module_list(
+    size = "medium",
+    files = [
+        "tests/experimental/test_multi_node_dag.py",
+    ],
+    tags = [
+        "compiled_graphs",
+        "exclusive",
+        # Disabled as this is a flaky compiled graphs test.
+        "manual",
+        "no_windows",
+        "team:core",
+    ],
+    deps = ["//:ray_lib"],
+)
+
+py_test_module_list(
+    size = "enormous",
+    files = [
+        "tests/experimental/test_compiled_graphs.py",
+    ],
+    tags = [
+        "compiled_graphs",
+        "exclusive",
+        "no_windows",
+        "team:core",
+    ],
+    deps = ["//:ray_lib"],
+)
+
+py_test(
+    name = "test_torch_tensor_dag_gpu",
+    size = "enormous",
+    srcs = [
+        "tests/experimental/test_torch_tensor_dag.py",
+    ],
+    env = {"RAY_PYTEST_USE_GPU": "1"},
+    main = "tests/experimental/test_torch_tensor_dag.py",
+    tags = [
+        "compiled_graphs",
+        "custom_setup",
+        "exclusive",
+        # Disabled as this test is consistently failing in CI and cgraphs will be deprecated/removed soon.
+        "manual",
+        "multi_gpu",
+        "no_windows",
+        "team:core",
+    ],
+    deps = ["//:ray_lib"],
+)
+
+py_test(
+    name = "test_torch_tensor_transport_gpu",
+    size = "enormous",
+    srcs = [
+        "tests/experimental/test_torch_tensor_transport.py",
+    ],
+    env = {"RAY_PYTEST_USE_GPU": "1"},
+    main = "tests/experimental/test_torch_tensor_transport.py",
+    tags = [
+        "compiled_graphs",
+        "custom_setup",
+        "exclusive",
+        "multi_gpu",
+        "no_windows",
+        "team:core",
+    ],
+    deps = ["//:ray_lib"],
+)
+
+# TODO(ruisearch42): Add this test once issues are fixed.
+# py_test(
+#     name = "test_execution_schedule_gpu",
+#     size = "enormous",
+#     srcs = [
+#         "tests/experimental/test_execution_schedule_gpu.py",
+#     ],
+#     env = {"RAY_PYTEST_USE_GPU": "1"},
+#     main = "tests/experimental/test_execution_schedule_gpu.py",
+#     tags = [
+#         "compiled_graphs",
+#         "exclusive",
+#         "multi_gpu",
+#         "custom_setup",
+#         "no_windows",
+#         "team:core",
+#     ],
+#     deps = ["//:ray_lib"],
+# )
+
+py_test(
+    name = "test_cpu_communicator_dag",
+    size = "medium",
+    srcs = dag_tests_srcs,
+    tags = [
+        "exclusive",
+        "ray_dag_tests",
+        "team:core",
+    ],
+    deps = [":dag_lib"],
+)

python/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",
+]

python/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."""

python/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

python/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

python/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)

python/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"
+)

python/ray/dag/context.py

@@ -0,0 +1,109 @@
+import os
+import threading
+from dataclasses import dataclass
+from typing import Optional
+
+from ray._common.utils import env_bool
+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 = env_bool(
+    "RAY_CGRAPH_overlap_gpu_communication", False
+)
+
+
+@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

python/ray/dag/dag_node.py

@@ -0,0 +1,724 @@
+import asyncio
+import copy
+import uuid
+import warnings
+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, RayDeprecationWarning
+
+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.
+
+        Returns:
+            This DAG node with the configured tensor transport.
+        """
+        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: Any, _ray_cache_refs: bool = False, **kwargs: Any
+    ) -> Union[ray.ObjectRef, "ray.actor.ActorHandle"]:
+        """Execute this DAG using the Ray default executor _execute_impl().
+
+        Args:
+            *args: Positional arguments forwarded to ``_execute_impl`` on each node.
+            _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
+            **kwargs: Keyword arguments forwarded to ``_execute_impl`` on each node.
+
+        Returns:
+            The result of executing the DAG (an ``ObjectRef`` or an
+            ``ActorHandle`` depending on the root node type).
+        """
+        warnings.warn(
+            "DAGNode.execute() is deprecated and will be removed in a future release.",
+            RayDeprecationWarning,
+            stacklevel=2,
+        )
+
+        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]
+
+        Returns:
+            All child DAGNodes referenced (transitively) by this node's
+            args, kwargs, and other_args_to_resolve.
+        """
+
+        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: Tuple[Any, ...]) -> None:
+        """
+        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)

python/ray/dag/dag_node_operation.py

@@ -0,0 +1,862 @@
+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,
+    ):
+        """Initialize a _DAGNodeOperation.
+
+        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.
+        control_dependency: If True, the edge represents a control dependency
+            (used for visualization) rather than a data dependency.
+    """
+    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") -> str:
+    """
+    Returns the label of an actor in the visualization of the execution schedule.
+
+    Args:
+        actor: The actor to be represented.
+
+    Returns:
+        A human-readable label combining the actor's class name and ID.
+    """
+    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
+) -> Tuple[str, str]:
+    """
+    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.
+
+    Returns:
+        A ``(node_viz_id, node_viz_label)`` tuple suitable for visualization.
+    """
+    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()
+    }

python/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: Any):
+        """
+        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

python/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

python/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))

python/ray/dag/input_node.py

@@ -0,0 +1,336 @@
+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: Any,
+        input_type: Optional[Union[type, Dict[Union[int, str], type]]] = None,
+        _other_args_to_resolve: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """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:
+            *args: Reserved; passing any positional argument raises ``ValueError``.
+            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.
+            **kwargs: Reserved; passing any keyword argument raises ``ValueError``.
+        """
+        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,
+    ):
+        """Initialize an InputAttributeNode.
+
+        Args:
+            dag_input_node: The parent ``InputNode`` this attribute access
+                derives from.
+            key: The index, attribute name, or dict key used to access a
+                value of the user input.
+            accessor_method: The accessor method used to extract the value
+                from the user input (e.g., ``"__getitem__"`` or
+                ``"__getattr__"``).
+            input_type: Type hint for the extracted value, used by the
+                Gradio visualizer to pick a UI component.
+        """
+        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."
+            )

python/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__")

python/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()

python/ray/dag/tests/conftest.py

@@ -0,0 +1,2 @@
+# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
+from ray.tests.conftest import pytest_runtest_makereport  # noqa

python/ray/dag/tests/experimental/actor_defs.py

@@ -0,0 +1,102 @@
+import os
+import random
+import time
+
+import ray
+
+
+@ray.remote
+class Actor:
+    def __init__(self, init_value, fail_after=None, sys_exit=False):
+        self.i = init_value
+        self.fail_after = fail_after
+        self.sys_exit = sys_exit
+
+        self.count = 0
+
+    def _fail_if_needed(self):
+        if self.fail_after and self.count > self.fail_after:
+            # Randomize the failures to better cover multi actor scenarios.
+            if random.random() > 0.5:
+                if self.sys_exit:
+                    os._exit(1)
+                else:
+                    raise RuntimeError("injected fault")
+
+    def inc(self, x):
+        self.i += x
+        self.count += 1
+        self._fail_if_needed()
+        return self.i
+
+    def double_and_inc(self, x):
+        self.i *= 2
+        self.i += x
+        return self.i
+
+    def echo(self, x):
+        self.count += 1
+        self._fail_if_needed()
+        return x
+
+    def append_to(self, lst):
+        lst.append(self.i)
+        return lst
+
+    def inc_two(self, x, y):
+        self.i += x
+        self.i += y
+        return self.i
+
+    def sleep(self, x):
+        time.sleep(x)
+        return x
+
+    @ray.method(num_returns=2)
+    def return_two(self, x):
+        return x, x + 1
+
+    def read_input(self, x):
+        return x
+
+    @ray.method(num_returns=2)
+    def inc_and_return_two(self, x):
+        self.i += x
+        return self.i, self.i + 1
+
+    @ray.method(num_returns=1)
+    def return_two_as_one(self, x):
+        return x, x + 1
+
+    @ray.method(num_returns=2)
+    def return_two_from_three(self, x):
+        return x, x + 1, x + 2
+
+    @ray.method(num_returns=2)
+    def return_two_but_raise_exception(self, x):
+        raise RuntimeError
+        return 1, 2
+
+    def get_events(self):
+        return getattr(self, "__ray_cgraph_events", [])
+
+
+@ray.remote
+class Collector:
+    def __init__(self):
+        self.results = []
+
+    def collect(self, x):
+        self.results.append(x)
+        return self.results
+
+    def collect_two(self, x, y):
+        self.results.append(x)
+        self.results.append(y)
+        return self.results
+
+    def collect_three(self, x, y, z):
+        self.results.append(x)
+        self.results.append(y)
+        self.results.append(z)
+        return self.results

python/ray/dag/tests/experimental/test_collective_dag.py

@@ -0,0 +1,497 @@
+# coding: utf-8
+import logging
+import os
+import sys
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
+
+import pytest
+
+import ray
+import ray.experimental.collective as collective
+from ray.dag import InputNode, MultiOutputNode
+from ray.experimental.channel import CPUCommunicator
+from ray.experimental.collective.conftest import (
+    AbstractNcclGroup,
+    CPUTorchTensorWorker,
+    check_nccl_group_init,
+    check_nccl_group_teardown,
+)
+from ray.experimental.util.types import ReduceOp
+from ray.tests.conftest import *  # noqa
+
+if TYPE_CHECKING:
+    import cupy as cp
+    import torch
+
+
+logger = logging.getLogger(__name__)
+
+if sys.platform != "linux" and sys.platform != "darwin":
+    pytest.skip("Skipping, requires Linux or Mac.", allow_module_level=True)
+
+
+class MockCommunicator(CPUCommunicator):
+    """
+    Use a mock communicator to test the actor schedules.
+    """
+
+    def __init__(self, world_size: int, actor_handles: List["ray.actor.ActorHandle"]):
+        self._world_size = world_size
+        self._actor_handles = actor_handles
+
+    def send(self, value: "torch.Tensor", peer_rank: int) -> None:
+        raise NotImplementedError
+
+    def recv(
+        self,
+        shape: Tuple[int],
+        dtype: "torch.dtype",
+        peer_rank: int,
+        allocator: Optional[
+            Callable[[Tuple[int], "torch.dtype"], "torch.Tensor"]
+        ] = None,
+    ) -> "torch.Tensor":
+        raise NotImplementedError
+
+    def allgather(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+    ) -> None:
+        raise NotImplementedError
+
+    def allreduce(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp,
+    ) -> None:
+        raise NotImplementedError
+
+    def reducescatter(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp,
+    ) -> None:
+        raise NotImplementedError
+
+    @property
+    def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+        raise NotImplementedError
+
+    @property
+    def send_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+        raise NotImplementedError
+
+    def destroy(self) -> None:
+        raise NotImplementedError
+
+
+@ray.remote
+class DDPWorker:
+    def __init__(self):
+        return
+
+    def backward(self, _):
+        return 0
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
+def test_all_reduce_duplicate_actors(ray_start_regular):
+    """
+    Test an error is thrown when two input nodes from the same actor bind to
+    an all-reduce.
+    """
+    actor_cls = CPUTorchTensorWorker.options()
+    worker = actor_cls.remote()
+
+    with InputNode() as inp:
+        computes = [worker.return_tensor.bind(inp) for _ in range(2)]
+        with pytest.raises(
+            ValueError,
+            match="Expected unique actor handles, but found duplicate actor handles from input nodes",
+        ):
+            collective.allreduce.bind(computes)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
+def test_all_reduce_custom_comm_wrong_actors(ray_start_regular):
+    """
+    Test an error is thrown when an all-reduce binds to a custom NCCL group and
+    a wrong set of actors.
+    """
+    actor_cls = CPUTorchTensorWorker.options()
+
+    num_workers = 2
+    workers = [actor_cls.remote() for _ in range(num_workers)]
+
+    nccl_group = AbstractNcclGroup([workers[0]])
+    with InputNode() as inp:
+        computes = [worker.return_tensor.bind(inp) for worker in workers]
+        with pytest.raises(
+            ValueError,
+            match="Expected actor handles to match the custom communicator group",
+        ):
+            collective.allreduce.bind(computes, transport=nccl_group)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
+def test_all_reduce_bind_list_of_nodes_duplicate_nodes(ray_start_regular):
+    """
+    Test an error is thrown when an all-reduce binds to lists of nodes
+    that are duplicated.
+    """
+    actor_cls = CPUTorchTensorWorker.options()
+
+    num_workers = 2
+    workers = [actor_cls.remote() for _ in range(num_workers)]
+
+    nccl_group = AbstractNcclGroup([workers[0]])
+    with InputNode() as inp:
+        computes_0 = [worker.return_tensor.bind(inp) for worker in workers]
+        computes_1 = [workers[0].return_tensor.bind(inp) for _ in range(2)]
+        with pytest.raises(
+            ValueError,
+            match="Expected unique actor handles at list at index",
+        ):
+            collective.allreduce.bind([computes_0, computes_1], transport=nccl_group)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
+def test_all_reduce_bind_list_of_nodes_unequal_number_of_nodes(ray_start_regular):
+    """
+    Test an error is thrown when an all-reduce binds to lists of nodes
+    of different number of nodes across actors.
+    """
+    actor_cls = CPUTorchTensorWorker.options()
+
+    num_workers = 2
+    workers = [actor_cls.remote() for _ in range(num_workers)]
+
+    nccl_group = AbstractNcclGroup([workers[0]])
+    with InputNode() as inp:
+        computes_0 = [worker.return_tensor.bind(inp) for worker in workers]
+        computes_1 = [worker.return_tensor.bind(inp) for worker in workers[1:]]
+        with pytest.raises(
+            ValueError,
+            match="Expected all input lists to have the same number of nodes",
+        ):
+            collective.allreduce.bind([computes_0, computes_1], transport=nccl_group)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
+def test_all_reduce_bind_list_of_nodes_different_actors(ray_start_regular):
+    """
+    Test an error is thrown when an all-reduce binds to a list of nodes
+    from different set of actors.
+    """
+    actor_cls = CPUTorchTensorWorker.options()
+
+    num_workers = 3
+    workers = [actor_cls.remote() for _ in range(num_workers)]
+
+    nccl_group = AbstractNcclGroup([workers[0]])
+    with InputNode() as inp:
+        computes_0 = [worker.return_tensor.bind(inp) for worker in workers[:2]]
+        computes_1 = [worker.return_tensor.bind(inp) for worker in workers[1:]]
+        with pytest.raises(
+            ValueError,
+            match="Expected all input lists to have the same set of actor handles",
+        ):
+            collective.allreduce.bind([computes_0, computes_1], transport=nccl_group)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
+def test_all_reduce_bind_list_of_nodes_different_dtypes(ray_start_regular):
+    """
+    Test an error is thrown when an all-reduce binds to a list of nodes
+    that execute with tensors of different dtypes.
+    """
+    actor_cls = CPUTorchTensorWorker.options()
+
+    num_workers = 3
+    workers = [actor_cls.remote() for _ in range(num_workers)]
+
+    comm = MockCommunicator(num_workers, workers)
+    with InputNode() as inp:
+        computes_0 = [worker.return_tensor.bind(inp[0], inp[1]) for worker in workers]
+        computes_1 = [worker.return_tensor.bind(inp[0], inp[2]) for worker in workers]
+        collectives = collective.allreduce.bind(
+            [computes_0, computes_1], transport=comm
+        )
+        recvs = [
+            worker.recv_tensors.bind(*collective)
+            for worker, collective in zip(workers, collectives)
+        ]
+        dag = MultiOutputNode(recvs)
+
+    compiled_dag = dag.experimental_compile()
+    with pytest.raises(
+        ValueError,
+        match="Expected all input tensors to have the same dtype",
+    ):
+        import torch
+
+        ray.get(compiled_dag.execute(1, torch.float16, torch.float32))
+
+
+@pytest.mark.parametrize(
+    "ray_start_regular", [{"num_cpus": 4, "num_gpus": 4}], indirect=True
+)
+def test_comm_all_reduces(ray_start_regular, monkeypatch):
+    """
+    Test different communicators are used for different all-reduce calls of
+    different sets of actors.
+    """
+    actor_cls = CPUTorchTensorWorker.options(num_cpus=0, num_gpus=1)
+
+    num_workers = 2
+    workers = [actor_cls.remote() for _ in range(num_workers)]
+
+    with InputNode() as inp:
+        computes = [worker.return_tensor.bind(inp) for worker in workers]
+        # There are two all-reduces, each on one actor.
+        collectives = [collective.allreduce.bind([compute]) for compute in computes]
+        # collective[0] is the only CollectiveOutputNode for each all-reduce.
+        dag = MultiOutputNode([collective[0] for collective in collectives])
+
+    compiled_dag, mock_nccl_group_set = check_nccl_group_init(
+        monkeypatch,
+        dag,
+        {
+            (frozenset([workers[0]]), None),
+            (frozenset([workers[1]]), None),
+        },
+    )
+
+    check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
+
+
+@pytest.mark.parametrize(
+    "ray_start_regular", [{"num_cpus": 4, "num_gpus": 4}], indirect=True
+)
+def test_comm_deduplicate_all_reduces(ray_start_regular, monkeypatch):
+    """
+    Test communicators are deduplicated when all-reduces are called on the same
+    group of actors more than once.
+    """
+    actor_cls = CPUTorchTensorWorker.options(num_cpus=0, num_gpus=1)
+
+    num_workers = 2
+    workers = [actor_cls.remote() for _ in range(num_workers)]
+
+    with InputNode() as inp:
+        tensors = [worker.return_tensor.bind(inp) for worker in workers]
+        collectives = collective.allreduce.bind(tensors)
+        collectives = collective.allreduce.bind(collectives)
+        dag = MultiOutputNode(collectives)
+
+    compiled_dag, mock_nccl_group_set = check_nccl_group_init(
+        monkeypatch,
+        dag,
+        {(frozenset(workers), None)},
+    )
+
+    check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
+
+
+@pytest.mark.parametrize(
+    "ray_start_regular", [{"num_cpus": 4, "num_gpus": 4}], indirect=True
+)
+def test_comm_deduplicate_p2p_and_collective(ray_start_regular, monkeypatch):
+    """
+    Test communicators are deduplicated when the collective and the P2P are on
+    the same set of actors.
+    """
+    actor_cls = CPUTorchTensorWorker.options(num_cpus=0, num_gpus=1)
+
+    num_workers = 2
+    workers = [actor_cls.remote() for _ in range(num_workers)]
+
+    with InputNode() as inp:
+        computes = [worker.return_tensor.bind(inp) for worker in workers]
+        collectives = collective.allreduce.bind(computes)
+        recvs = [
+            # Each of the 2 workers receives from the other.
+            workers[0].recv.bind(
+                collectives[1].with_tensor_transport(transport="nccl")
+            ),
+            workers[1].recv.bind(
+                collectives[0].with_tensor_transport(transport="nccl")
+            ),
+        ]
+        dag = MultiOutputNode(recvs)
+
+    compiled_dag, mock_nccl_group_set = check_nccl_group_init(
+        monkeypatch,
+        dag,
+        {(frozenset(workers), None)},
+    )
+
+    check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
+
+    with InputNode() as inp:
+        computes = [worker.return_tensor.bind(inp) for worker in workers]
+        collectives = collective.allreduce.bind(computes)
+        # Sender is workers[0] and receiver is workers[1].
+        dag = workers[1].recv.bind(
+            collectives[0].with_tensor_transport(transport="nccl")
+        )
+        dag = MultiOutputNode([dag, collectives[1]])
+
+    compiled_dag, mock_nccl_group_set = check_nccl_group_init(
+        monkeypatch,
+        dag,
+        {(frozenset(workers), None)},
+    )
+
+    check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
+
+
+@pytest.mark.parametrize(
+    "ray_start_regular", [{"num_cpus": 4, "num_gpus": 4}], indirect=True
+)
+def test_custom_comm(ray_start_regular, monkeypatch):
+    """
+    Test a custom GPU communicator is used when specified and a default
+    communicator is used otherwise.
+    """
+    actor_cls = CPUTorchTensorWorker.options(num_cpus=0, num_gpus=1)
+
+    num_workers = 2
+    workers = [actor_cls.remote() for _ in range(num_workers)]
+
+    comm = AbstractNcclGroup(workers)
+    with InputNode() as inp:
+        computes = [worker.return_tensor.bind(inp) for worker in workers]
+        collectives = collective.allreduce.bind(computes, transport=comm)
+        collectives = collective.allreduce.bind(collectives)
+        dag = workers[0].recv.bind(
+            collectives[1].with_tensor_transport(transport="nccl")
+        )
+        dag = MultiOutputNode([dag, collectives[0]])
+
+    compiled_dag, mock_nccl_group_set = check_nccl_group_init(
+        monkeypatch,
+        dag,
+        {
+            (frozenset(workers), comm),
+            (frozenset(workers), None),
+        },
+    )
+
+    check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
+
+    comm = AbstractNcclGroup(workers)
+    with InputNode() as inp:
+        computes = [worker.return_tensor.bind(inp) for worker in workers]
+        collectives = collective.allreduce.bind(computes)
+        collectives = collective.allreduce.bind(collectives)
+        dag = workers[0].recv.bind(collectives[1].with_tensor_transport(transport=comm))
+        dag = MultiOutputNode([dag, collectives[0]])
+
+    compiled_dag, mock_nccl_group_set = check_nccl_group_init(
+        monkeypatch,
+        dag,
+        {
+            (frozenset(workers), comm),
+            (frozenset(workers), None),
+        },
+    )
+
+    check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
+
+
+@pytest.mark.parametrize(
+    "ray_start_regular", [{"num_cpus": 4, "num_gpus": 4}], indirect=True
+)
+def test_custom_comm_init_teardown(ray_start_regular, monkeypatch):
+    """
+    Test custom NCCL groups are properly initialized and destroyed.
+    1. Test when multiple type hints have the same `transport=custom_nccl_group`,
+    the `custom_nccl_group` is initialized only once.
+    2. Test all initialized NCCL groups are destroyed during teardown.
+    """
+    actor_cls = CPUTorchTensorWorker.options(num_cpus=0, num_gpus=1)
+
+    num_workers = 2
+    workers = [actor_cls.remote() for _ in range(num_workers)]
+
+    comm = AbstractNcclGroup(workers)
+
+    with InputNode() as inp:
+        tensors = [worker.return_tensor.bind(inp) for worker in workers]
+        allreduce = collective.allreduce.bind(tensors, transport=comm)
+        dag = workers[0].recv.bind(allreduce[1].with_tensor_transport(transport=comm))
+        dag = MultiOutputNode([dag, allreduce[0]])
+
+    compiled_dag, mock_nccl_group_set = check_nccl_group_init(
+        monkeypatch,
+        dag,
+        {(frozenset(workers), comm)},
+    )
+
+    check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
+
+    comm_1 = AbstractNcclGroup(workers)
+    comm_2 = AbstractNcclGroup(workers)
+    comm_3 = AbstractNcclGroup(workers)
+
+    with InputNode() as inp:
+        tensors = [worker.return_tensor.bind(inp) for worker in workers]
+        allreduce1 = collective.allreduce.bind(tensors, transport=comm_1)
+        allreduce2 = collective.allreduce.bind(allreduce1, transport=comm_2)
+        dag = workers[0].recv.bind(
+            allreduce2[1].with_tensor_transport(transport=comm_3)
+        )
+        dag = MultiOutputNode([dag, allreduce2[0]])
+
+    compiled_dag, mock_nccl_group_set = check_nccl_group_init(
+        monkeypatch,
+        dag,
+        {
+            (frozenset(workers), comm_1),
+            (frozenset(workers), comm_2),
+            (frozenset(workers), comm_3),
+        },
+    )
+
+    check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
+@pytest.mark.parametrize("num_workers", [2, 4])
+def test_exec_schedules_ddp(ray_start_regular, num_workers):
+    """
+    Test the execution schedules for the DDP strategy. Each worker should have
+    identical schedules.
+    """
+    actor_cls = DDPWorker.options(num_cpus=1)
+    workers = [actor_cls.remote() for _ in range(num_workers)]
+    comm = MockCommunicator(num_workers, workers)
+
+    outputs = []
+    with InputNode() as inp:
+        grads = [worker.backward.bind(inp) for worker in workers]
+        grads_reduced = collective.allreduce.bind(grads, transport=comm)
+        outputs.extend(grads_reduced)
+        grads = [worker.backward.bind(grad) for worker, grad in zip(workers, grads)]
+        grads_reduced = collective.allreduce.bind(grads, transport=comm)
+        outputs.extend(grads_reduced)
+        dag = MultiOutputNode(outputs)
+
+    compiled_dag = dag.experimental_compile(_default_communicator=comm)
+    actor_to_execution_schedule = list(
+        compiled_dag.actor_to_execution_schedule.values()
+    )
+    expected_schedule = actor_to_execution_schedule[0]
+    for schedule in actor_to_execution_schedule[1:]:
+        assert schedule == expected_schedule
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/dag/tests/experimental/test_cpu_communicator_dag.py

@@ -0,0 +1,361 @@
+import os
+import sys
+
+import pytest
+import torch
+
+import ray
+import ray.cluster_utils
+import ray.experimental.collective as collective
+from ray.dag import InputNode
+from ray.dag.output_node import MultiOutputNode
+from ray.exceptions import RayChannelError, RayTaskError
+from ray.experimental.channel.cpu_communicator import CPUCommunicator
+from ray.tests.conftest import *  # noqa
+
+
+@ray.remote
+class CPUTorchTensorWorker:
+    def __init__(self):
+        self.device = torch.device(type="cpu")
+
+    def send(self, shape, dtype, value: int, send_tensor=True):
+        if not send_tensor:
+            return 1
+        return torch.ones(shape, dtype=dtype) * value
+
+    def send_dict(self, entries):
+        results = {}
+        for key, entry in entries.items():
+            value, shape, dtype = entry
+            results[key] = torch.ones(shape, dtype=dtype) * value
+        return results
+
+    def send_or_raise(self, shape, dtype, value: int, raise_exception=False):
+        if raise_exception:
+            raise RuntimeError()
+        return torch.ones(shape, dtype=dtype) * value
+
+    def recv(self, tensor):
+        assert tensor.device == self.device
+        return (tensor[0].item(), tensor.shape, tensor.dtype)
+
+    def recv_dict(self, tensor_dict):
+        vals = {}
+        for i, tensor in tensor_dict.items():
+            assert tensor.device == self.device
+            vals[i] = self.recv(tensor)
+        return vals
+
+    def compute_with_tuple_args(self, args, i: int):
+        shape, dtype, value = args[i]
+        tensor = torch.ones(shape, dtype=dtype) * value
+        return tensor
+
+    def recv_tensor(self, tensor):
+        assert tensor.device == self.device
+        return tensor
+
+    def return_tensor(self, size: int) -> torch.Tensor:
+        return torch.ones(size)
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 2,
+            "num_gpus": 0,
+            "num_nodes": 1,
+        }
+    ],
+    indirect=True,
+)
+def test_p2p_basic(ray_start_cluster):
+    sender = CPUTorchTensorWorker.remote()
+    receiver = CPUTorchTensorWorker.remote()
+
+    cpu_group = CPUCommunicator(2, [sender, receiver])
+
+    shape = (10,)
+    dtype = torch.float16
+
+    with InputNode() as inp:
+        dag = sender.send.bind(inp.shape, inp.dtype, inp[0])
+        dag = dag.with_tensor_transport(transport=cpu_group)
+        dag = receiver.recv.bind(dag)
+
+    compiled_dag = dag.experimental_compile()
+    for i in range(3):
+        ref = compiled_dag.execute(i, shape=shape, dtype=dtype)
+        assert ray.get(ref) == (i, shape, dtype)
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 2,
+            "num_gpus": 0,
+            "num_nodes": 1,
+        }
+    ],
+    indirect=True,
+)
+def test_allreduce_basic(ray_start_cluster):
+    num_workers = 2
+    workers = [CPUTorchTensorWorker.remote() for _ in range(num_workers)]
+
+    cpu_group = CPUCommunicator(num_workers, workers)
+
+    with InputNode() as inp:
+        computes = [
+            worker.compute_with_tuple_args.bind(inp, i)
+            for i, worker in enumerate(workers)
+        ]
+        collectives = collective.allreduce.bind(computes, transport=cpu_group)
+        recvs = [
+            worker.recv.bind(collective)
+            for worker, collective in zip(workers, collectives)
+        ]
+        dag = MultiOutputNode(recvs)
+
+    compiled_dag = dag.experimental_compile()
+
+    for i in range(3):
+        i += 1
+        shape = (i * 10,)
+        dtype = torch.float16
+        ref = compiled_dag.execute(
+            [(shape, dtype, i + idx) for idx in range(num_workers)]
+        )
+        result = ray.get(ref)
+        reduced_val = sum(i + idx for idx in range(num_workers))
+        assert result == [(reduced_val, shape, dtype) for _ in workers]
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 2,
+            "num_gpus": 0,
+            "num_nodes": 1,
+        }
+    ],
+    indirect=True,
+)
+def test_allreduce_get_partial(ray_start_cluster):
+    num_workers = 2
+    workers = [CPUTorchTensorWorker.remote() for _ in range(num_workers)]
+
+    cpu_group = CPUCommunicator(num_workers, workers)
+
+    shape = (10,)
+    dtype = torch.float16
+
+    with InputNode() as inp:
+        computes = [
+            worker.compute_with_tuple_args.bind(inp, i)
+            for i, worker in enumerate(workers)
+        ]
+        collectives = collective.allreduce.bind(computes, transport=cpu_group)
+        recv = workers[0].recv.bind(collectives[0])
+        tensor = workers[1].recv_tensor.bind(collectives[0])
+        dag = MultiOutputNode([recv, tensor, collectives[1]])
+
+    compiled_dag = dag.experimental_compile()
+
+    for i in range(3):
+        ref = compiled_dag.execute(
+            [(shape, dtype, i + idx + 1) for idx in range(num_workers)]
+        )
+        result = ray.get(ref)
+        metadata, tensor, _ = result
+        reduced_val = sum(i + idx + 1 for idx in range(num_workers))
+        assert metadata == (reduced_val, shape, dtype)
+        expected_tensor_val = torch.ones(shape, dtype=dtype) * reduced_val
+        assert torch.equal(tensor, expected_tensor_val)
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 2,
+            "num_gpus": 0,
+            "num_nodes": 1,
+        }
+    ],
+    indirect=True,
+)
+def test_allreduce_wrong_shape(ray_start_cluster):
+    """
+    Test an error is thrown when the tensors in an all-reduce have different shapes.
+    """
+    num_workers = 2
+    workers = [CPUTorchTensorWorker.remote() for _ in range(num_workers)]
+
+    cpu_group = CPUCommunicator(num_workers, workers)
+
+    dtype = torch.float16
+
+    with InputNode() as inp:
+        computes = [
+            worker.compute_with_tuple_args.bind(inp, i)
+            for i, worker in enumerate(workers)
+        ]
+        collectives = collective.allreduce.bind(computes, transport=cpu_group)
+        recvs = [
+            worker.recv.bind(collective)
+            for worker, collective in zip(workers, collectives)
+        ]
+        dag = MultiOutputNode(recvs)
+
+    compiled_dag = dag.experimental_compile()
+
+    ref = compiled_dag.execute([((20,), dtype, idx + 1) for idx in range(num_workers)])
+    reduced_val = (1 + num_workers) * num_workers / 2
+    assert ray.get(ref) == [(reduced_val, (20,), dtype) for _ in range(num_workers)]
+
+    ref = compiled_dag.execute(
+        [((10 * (idx + 1),), dtype, idx + 1) for idx in range(num_workers)]
+    )
+    # Execution hangs because of shape mismatch and a task error is raised.
+    with pytest.raises(RayTaskError):
+        ray.get(ref)
+
+    # Since we have buffered channels, the execution should not error, but the
+    # get should error, as the dag should no longer work after the application-
+    # level exception.
+    ref = compiled_dag.execute([((20,), dtype, 1) for _ in workers])
+    with pytest.raises(RayChannelError):
+        ray.get(ref)
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 2,
+            "num_gpus": 0,
+            "num_nodes": 1,
+        }
+    ],
+    indirect=True,
+)
+def test_allreduce_scheduling(ray_start_cluster):
+    """
+    Test scheduling avoids potential deadlocks that arise from all-reduce operations.
+
+    inp --> x(0) --> +------------+
+        |            | all-reduce |
+        --> y(1) --> +------------+
+        |
+        --> t(0) --> recv(1)
+
+    In the above graph, x, y, t are tensors, and the numbers inside parentheses
+    identify the actors. If actor 1 launches an all-reduce with tensor y while
+    actor 0 starts sending t, then actor 1 waits for actor 0 to join the all-reduce
+    while actor 1 waits for actor 0 to receive t.
+    """
+    num_workers = 2
+    workers = [CPUTorchTensorWorker.remote() for _ in range(num_workers)]
+
+    cpu_group = CPUCommunicator(num_workers, workers)
+
+    shape = (10,)
+    dtype = torch.float16
+
+    with InputNode() as inp:
+        # Tensors in the all-reduce.
+        x = workers[0].send.bind(shape, dtype, inp)
+        y = workers[1].send.bind(shape, dtype, inp)
+
+        # Tensor to be sent from workers[0] to workers[1].
+        t = workers[0].send.bind(shape, dtype, inp)
+        t = t.with_tensor_transport(transport=cpu_group)
+
+        collectives = collective.allreduce.bind([x, y], transport=cpu_group)
+        recv = workers[1].recv.bind(t)
+        dag = MultiOutputNode([collectives[0], collectives[1], recv])
+
+    compiled_dag = dag.experimental_compile()
+
+    value = 10
+    ref = compiled_dag.execute(value)
+    result = ray.get(ref)
+    reduced_value = value * 2
+    expected_tensor_val = torch.ones(shape, dtype=dtype) * reduced_value
+    assert torch.equal(result[0], expected_tensor_val)
+    assert torch.equal(result[1], expected_tensor_val)
+    assert result[2] == (value, shape, dtype)
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 2,
+            "num_gpus": 0,
+            "num_nodes": 1,
+        }
+    ],
+    indirect=True,
+)
+def test_allreduce_duplicate_actors(ray_start_cluster):
+    """
+    Test an error is thrown when two input nodes from the same actor bind to
+    an all-reduce.
+    """
+    num_workers = 2
+    worker = CPUTorchTensorWorker.remote()
+
+    cpu_group = CPUCommunicator(num_workers, [worker, worker])
+
+    with InputNode() as inp:
+        computes = [worker.return_tensor.bind(inp) for _ in range(2)]
+        with pytest.raises(
+            ValueError,
+            match=(
+                "Expected unique actor handles, but found duplicate actor handles "
+                "from input nodes"
+            ),
+        ):
+            collective.allreduce.bind(computes, transport=cpu_group)
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 2,
+            "num_gpus": 0,
+            "num_nodes": 1,
+        }
+    ],
+    indirect=True,
+)
+def test_allreduce_wrong_actors(ray_start_cluster):
+    """
+    Test an error is thrown when an all-reduce binds to a wrong set of actors.
+    """
+    num_workers = 2
+    workers = [CPUTorchTensorWorker.remote() for _ in range(num_workers * 2)]
+
+    cpu_group = CPUCommunicator(num_workers, workers[:2])
+
+    with InputNode() as inp:
+        computes = [worker.return_tensor.bind(inp) for worker in workers[2:]]
+        with pytest.raises(
+            ValueError,
+            match="Expected actor handles to match the custom communicator group",
+        ):
+            collective.allreduce.bind(computes, transport=cpu_group)
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/dag/tests/experimental/test_dag_error_handling.py

@@ -0,0 +1,814 @@
+# coding: utf-8
+import copy
+import logging
+import pickle
+import re
+import signal
+import sys
+import time
+
+import pytest
+
+import ray
+import ray._private
+import ray.cluster_utils
+from ray._common.test_utils import SignalActor
+from ray._common.utils import (
+    get_or_create_event_loop,
+)
+from ray._private.test_utils import (
+    run_string_as_driver_nonblocking,
+    wait_for_pid_to_exit,
+)
+from ray.dag import DAGContext, InputNode, MultiOutputNode
+from ray.dag.tests.experimental.actor_defs import Actor
+from ray.exceptions import ActorDiedError, RayChannelError, RayChannelTimeoutError
+from ray.tests.conftest import *  # noqa
+
+logger = logging.getLogger(__name__)
+
+
+pytestmark = [
+    pytest.mark.skipif(
+        sys.platform != "linux" and sys.platform != "darwin",
+        reason="Requires Linux or MacOS",
+    ),
+    pytest.mark.timeout(500),
+]
+
+
+@pytest.fixture
+def temporary_change_timeout(request):
+    ctx = DAGContext.get_current()
+    original = ctx.submit_timeout
+    ctx.submit_timeout = request.param
+    yield ctx.submit_timeout
+    ctx.submit_timeout = original
+
+
+@pytest.fixture
+def zero_teardown_timeout(request):
+    ctx = DAGContext.get_current()
+    original = ctx.teardown_timeout
+    ctx.teardown_timeout = 0
+    yield ctx.teardown_timeout
+    ctx.teardown_timeout = original
+
+
+def test_kwargs_not_supported(ray_start_regular):
+    a = Actor.remote(0)
+
+    # Binding InputNode as kwarg is not supported.
+    with InputNode() as i:
+        dag = a.inc_two.bind(x=i, y=1)
+    with pytest.raises(
+        ValueError,
+        match=r"Compiled DAG currently does not support binding to other DAG "
+        "nodes as kwargs",
+    ):
+        compiled_dag = dag.experimental_compile()
+
+    # Binding another DAG node as kwarg is not supported.
+    with InputNode() as i:
+        dag = a.inc.bind(i)
+        dag = a.inc_two.bind(x=dag, y=1)
+    with pytest.raises(
+        ValueError,
+        match=r"Compiled DAG currently does not support binding to other DAG "
+        "nodes as kwargs",
+    ):
+        compiled_dag = dag.experimental_compile()
+
+    # Binding normal Python value as a kwarg is supported.
+    with InputNode() as i:
+        dag = a.inc_two.bind(i, y=1)
+    compiled_dag = dag.experimental_compile()
+    assert ray.get(compiled_dag.execute(2)) == 3
+
+
+def test_dag_exception_basic(ray_start_regular, capsys):
+    # Test application throwing exceptions with a single task.
+    a = Actor.remote(0)
+    with InputNode() as inp:
+        dag = a.inc.bind(inp)
+
+    # Can throw an error.
+    compiled_dag = dag.experimental_compile()
+    ref = compiled_dag.execute("hello")
+    with pytest.raises(TypeError) as exc_info:
+        ray.get(ref)
+    # Traceback should match the original actor class definition.
+    assert "self.i += x" in str(exc_info.value)
+
+    # Can throw an error multiple times.
+    ref = compiled_dag.execute("hello")
+    with pytest.raises(TypeError) as exc_info:
+        ray.get(ref)
+    # Traceback should match the original actor class definition.
+    assert "self.i += x" in str(exc_info.value)
+
+    # Can use the DAG after exceptions are thrown.
+    assert ray.get(compiled_dag.execute(1)) == 1
+
+
+def test_dag_exception_chained(ray_start_regular, capsys):
+    # Test application throwing exceptions with a task that depends on another
+    # task.
+    a = Actor.remote(0)
+    with InputNode() as inp:
+        dag = a.inc.bind(inp)
+        dag = a.inc.bind(dag)
+
+    # Can throw an error.
+    compiled_dag = dag.experimental_compile()
+    ref = compiled_dag.execute("hello")
+    with pytest.raises(TypeError) as exc_info:
+        ray.get(ref)
+    # Traceback should match the original actor class definition.
+    assert "self.i += x" in str(exc_info.value)
+
+    # Can throw an error multiple times.
+    ref = compiled_dag.execute("hello")
+    with pytest.raises(TypeError) as exc_info:
+        ray.get(ref)
+    # Traceback should match the original actor class definition.
+    assert "self.i += x" in str(exc_info.value)
+
+    # Can use the DAG after exceptions are thrown.
+    assert ray.get(compiled_dag.execute(1)) == 2
+
+
+def test_dag_exception_multi_output(ray_start_regular, capsys):
+    # Test application throwing exceptions with a DAG with multiple outputs.
+    a = Actor.remote(0)
+    b = Actor.remote(0)
+    with InputNode() as inp:
+        dag = MultiOutputNode([a.inc.bind(inp), b.inc.bind(inp)])
+
+    compiled_dag = dag.experimental_compile()
+
+    # Verify that fetching each output individually raises the error.
+    refs = compiled_dag.execute("hello")
+    for ref in refs:
+        with pytest.raises(TypeError) as exc_info:
+            ray.get(ref)
+        # Traceback should match the original actor class definition.
+        assert "self.i += x" in str(exc_info.value)
+
+    # Verify that another bad input exhibits the same behavior.
+    refs = compiled_dag.execute("hello")
+    for ref in refs:
+        with pytest.raises(TypeError) as exc_info:
+            ray.get(ref)
+        # Traceback should match the original actor class definition.
+        assert "self.i += x" in str(exc_info.value)
+
+    # Verify that the DAG can be used after the errors.
+    assert ray.get(compiled_dag.execute(1)) == [1, 1]
+
+
+def test_dag_errors(ray_start_regular):
+    a = Actor.remote(0)
+    dag = a.inc.bind(1)
+    with pytest.raises(
+        ValueError,
+        match="No InputNode found in the DAG: when traversing upwards, "
+        "no upstream node was found for",
+    ):
+        dag.experimental_compile()
+
+    a2 = Actor.remote(0)
+    with InputNode() as inp:
+        dag = MultiOutputNode([a.inc.bind(inp), a2.inc.bind(1)])
+    with pytest.raises(
+        ValueError,
+        match="Compiled DAGs require each task to take a ray.dag.InputNode or "
+        "at least one other DAGNode as an input",
+    ):
+        dag.experimental_compile()
+
+    @ray.remote
+    def f(x):
+        return x
+
+    with InputNode() as inp:
+        dag = f.bind(inp)
+    with pytest.raises(
+        NotImplementedError,
+        match="Compiled DAGs currently only support actor method nodes",
+    ):
+        dag.experimental_compile()
+
+    with InputNode() as inp:
+        dag = a.inc.bind(inp)
+    compiled_dag = dag.experimental_compile()
+    ref = compiled_dag.execute(1)
+    with pytest.raises(
+        TypeError,
+        match=(
+            re.escape(
+                "wait() expected a list of ray.ObjectRef or ray.ObjectRefGenerator, "
+                "got <class 'ray.experimental.compiled_dag_ref.CompiledDAGRef'>"
+            )
+        ),
+    ):
+        ray.wait(ref)
+
+    with pytest.raises(
+        TypeError,
+        match=(
+            re.escape(
+                "wait() expected a list of ray.ObjectRef or ray.ObjectRefGenerator, "
+                "got list containing "
+                "<class 'ray.experimental.compiled_dag_ref.CompiledDAGRef'>"
+            )
+        ),
+    ):
+        ray.wait([ref])
+
+    with pytest.raises(TypeError, match=r".*was found to be non-serializable.*"):
+        ray.put([ref])
+
+    with pytest.raises(ValueError, match="CompiledDAGRef cannot be copied."):
+        copy.copy(ref)
+
+    with pytest.raises(ValueError, match="CompiledDAGRef cannot be deep copied."):
+        copy.deepcopy(ref)
+
+    with pytest.raises(ValueError, match="CompiledDAGRef cannot be pickled."):
+        pickle.dumps(ref)
+
+    with pytest.raises(
+        TypeError, match="CompiledDAGRef cannot be used as Ray task/actor argument."
+    ):
+        f.remote(ref)
+
+    with pytest.raises(
+        TypeError, match="CompiledDAGRef cannot be used as Ray task/actor argument."
+    ):
+        a2.inc.remote(ref)
+
+    result = ray.get(ref)
+    assert result == 1
+
+    with pytest.raises(
+        ValueError,
+        match=(
+            r"ray.get\(\) can only be called once "
+            r"on a CompiledDAGRef, and it was already called."
+        ),
+    ):
+        ray.get(ref)
+
+
+def test_get_timeout(ray_start_regular, zero_teardown_timeout):
+    a = Actor.remote(0)
+    with InputNode() as inp:
+        dag = a.sleep.bind(inp)
+
+    compiled_dag = dag.experimental_compile()
+    ref = compiled_dag.execute(5)
+
+    timed_out = False
+    epsilon = 0.1  # Allow for some slack in the timeout checking
+    try:
+        start_time = time.monotonic()
+        ray.get(ref, timeout=1)
+    except RayChannelTimeoutError:
+        duration = time.monotonic() - start_time
+        assert duration > 1 - epsilon
+        assert duration < 1 + epsilon
+        timed_out = True
+    assert timed_out
+
+    compiled_dag.teardown(kill_actors=True)
+
+
+def test_buffered_get_timeout(ray_start_regular, zero_teardown_timeout):
+    a = Actor.remote(0)
+    with InputNode() as inp:
+        dag = a.sleep.bind(inp)
+
+    compiled_dag = dag.experimental_compile()
+    # The tasks will execute in order and sleep 1s, 1s, then 0s, respectively.
+    refs = [
+        compiled_dag.execute(1),
+        compiled_dag.execute(1),
+        compiled_dag.execute(0),
+    ]
+
+    with pytest.raises(RayChannelTimeoutError):
+        # The final task takes <1s on its own, but because it's queued behind the
+        # other two that take 1s each, this should time out.
+        ray.get(refs[-1], timeout=1)
+
+    compiled_dag.teardown(kill_actors=True)
+
+
+def test_get_with_zero_timeout(ray_start_regular):
+    @ray.remote
+    class Actor:
+        def __init__(self, signal_actor):
+            self.signal_actor = signal_actor
+
+        def send(self, x):
+            self.signal_actor.send.remote()
+            return x
+
+    signal_actor = SignalActor.remote()
+    a = Actor.remote(signal_actor)
+    with InputNode() as inp:
+        dag = a.send.bind(inp)
+
+    compiled_dag = dag.experimental_compile()
+    ref = compiled_dag.execute(1)
+    # Give enough time for DAG execution result to be ready
+    ray.get(signal_actor.wait.remote())
+    time.sleep(0.1)
+    # Use timeout=0 to either get result immediately or raise an exception
+    result = ray.get(ref, timeout=0)
+    assert result == 1
+
+
+class TestDAGExceptionCompileMultipleTimes:
+    @pytest.mark.parametrize("use_multi_output_node", [False, True])
+    def test_compile_twice_fails(self, ray_start_regular, use_multi_output_node: bool):
+        a = Actor.remote(0)
+        with InputNode() as i:
+            if use_multi_output_node:
+                dag = MultiOutputNode([a.echo.bind(i)])
+            else:
+                dag = a.echo.bind(i)
+        compiled_dag = dag.experimental_compile()
+
+        # Trying to compile again should fail.
+        expected_err = (
+            "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."
+        )
+        with pytest.raises(
+            ValueError,
+            match=expected_err,
+        ):
+            compiled_dag = dag.experimental_compile()
+
+        # Even if we teardown the DAG, trying to compile again should still fail.
+        compiled_dag.teardown()
+        with pytest.raises(
+            ValueError,
+            match=expected_err,
+        ):
+            compiled_dag = dag.experimental_compile()
+
+    def test_compile_twice_with_different_nodes(self, ray_start_regular):
+        a = Actor.remote(0)
+        b = Actor.remote(0)
+        with InputNode() as i:
+            branch1 = a.echo.bind(i)
+            branch2 = b.echo.bind(i)
+            dag = MultiOutputNode([branch1, branch2])
+        compiled_dag = dag.experimental_compile()
+        compiled_dag.teardown()
+        with pytest.raises(
+            ValueError,
+            match="The DAG was compiled more than once. The following two "
+            "nodes call `experimental_compile`: ",
+        ):
+            branch2.experimental_compile()
+
+
+def test_exceed_max_buffered_results(ray_start_regular):
+    a = Actor.remote(0)
+    with InputNode() as i:
+        dag = a.inc.bind(i)
+
+    compiled_dag = dag.experimental_compile(_max_buffered_results=1)
+
+    refs = []
+    for i in range(2):
+        ref = compiled_dag.execute(1)
+        # Hold the refs to avoid get() being called on the ref
+        # when it goes out of scope
+        refs.append(ref)
+
+    # ray.get() on the 2nd ref fails because the DAG cannot buffer 2 results.
+    with pytest.raises(
+        ray.exceptions.RayCgraphCapacityExceeded,
+        match=(
+            "The compiled graph can't have more than 1 buffered results, "
+            r"and you currently have 1 buffered results. Call `ray.get\(\)` on "
+            r"CompiledDAGRef's \(or await on CompiledDAGFuture's\) to retrieve "
+            "results, or increase `_max_buffered_results` if buffering is "
+            "desired, note that this will increase driver memory usage."
+        ),
+    ):
+        ray.get(ref)
+
+
+def test_exceed_max_buffered_results_multi_output(ray_start_regular):
+    a = Actor.remote(0)
+    b = Actor.remote(0)
+    with InputNode() as inp:
+        dag = MultiOutputNode([a.inc.bind(inp), b.inc.bind(inp)])
+
+    compiled_dag = dag.experimental_compile(_max_buffered_results=1)
+
+    refs = []
+    for _ in range(2):
+        ref = compiled_dag.execute(1)
+        # Hold the refs to avoid get() being called on the ref
+        # when it goes out of scope
+        refs.append(ref)
+
+    # If there are results not fetched from an execution, that execution
+    # still counts towards the number of buffered results.
+    ray.get(refs[0][0])
+
+    # ray.get() on the 2nd ref fails because the DAG cannot buffer 2 results.
+    with pytest.raises(
+        ray.exceptions.RayCgraphCapacityExceeded,
+        match=(
+            "The compiled graph can't have more than 1 buffered results, "
+            r"and you currently have 1 buffered results. Call `ray.get\(\)` on "
+            r"CompiledDAGRef's \(or await on CompiledDAGFuture's\) to retrieve "
+            "results, or increase `_max_buffered_results` if buffering is "
+            "desired, note that this will increase driver memory usage."
+        ),
+    ):
+        ray.get(ref[0])
+
+
+def test_dag_fault_tolerance_chain(ray_start_regular):
+    actors = [
+        Actor.remote(0, fail_after=10 if i == 0 else None, sys_exit=False)
+        for i in range(4)
+    ]
+    with InputNode() as i:
+        dag = i
+        for a in actors:
+            dag = a.echo.bind(dag)
+
+    compiled_dag = dag.experimental_compile()
+
+    for i in range(9):
+        ref = compiled_dag.execute(i)
+        results = ray.get(ref)
+
+    with pytest.raises(RuntimeError):
+        for i in range(9):
+            ref = compiled_dag.execute(i)
+            results = ray.get(ref)
+            assert results == i
+
+    compiled_dag.teardown()
+
+    # All actors are still alive.
+    ray.get([actor.sleep.remote(0) for actor in actors])
+
+    # Remaining actors can be reused.
+    actors.pop(0)
+    with InputNode() as i:
+        dag = i
+        for a in actors:
+            dag = a.echo.bind(dag)
+
+    compiled_dag = dag.experimental_compile()
+    for i in range(10):
+        ref = compiled_dag.execute(i)
+        results = ray.get(ref)
+        assert results == i
+
+
+def test_dag_fault_tolerance(ray_start_regular):
+    actors = [
+        Actor.remote(0, fail_after=10 if i == 0 else None, sys_exit=False)
+        for i in range(4)
+    ]
+    with InputNode() as i:
+        out = [a.inc.bind(i) for a in actors]
+        dag = MultiOutputNode(out)
+
+    compiled_dag = dag.experimental_compile()
+
+    for i in range(9):
+        refs = compiled_dag.execute(1)
+        assert ray.get(refs) == [i + 1] * len(actors)
+
+    with pytest.raises(RuntimeError):
+        for i in range(9, 20):
+            refs = compiled_dag.execute(1)
+            assert ray.get(refs) == [i + 1] * len(actors)
+
+    compiled_dag.teardown()
+
+    # All actors are still alive.
+    ray.get([actor.sleep.remote(0) for actor in actors])
+
+    # Remaining actors can be reused.
+    actors.pop(0)
+    with InputNode() as i:
+        out = [a.inc.bind(i) for a in actors]
+        dag = MultiOutputNode(out)
+
+    compiled_dag = dag.experimental_compile()
+    for i in range(10):
+        ray.get(compiled_dag.execute(1))
+
+
+def test_dag_fault_tolerance_sys_exit(ray_start_regular):
+    actors = [
+        Actor.remote(0, fail_after=10 if i == 0 else None, sys_exit=True)
+        for i in range(4)
+    ]
+    with InputNode() as i:
+        out = [a.inc.bind(i) for a in actors]
+        dag = MultiOutputNode(out)
+
+    compiled_dag = dag.experimental_compile()
+
+    for i in range(9):
+        refs = compiled_dag.execute(1)
+        assert ray.get(refs) == [i + 1] * len(actors)
+
+    with pytest.raises(
+        ActorDiedError, match="The actor died unexpectedly before finishing this task."
+    ):
+        for i in range(9):
+            refs = compiled_dag.execute(1)
+            ray.get(refs)
+
+    # Remaining actors are still alive.
+    with pytest.raises(ray.exceptions.RayActorError):
+        ray.get(actors[0].echo.remote("hello"))
+    actors.pop(0)
+    ray.get([actor.echo.remote("hello") for actor in actors])
+
+    # Remaining actors can be reused.
+    with InputNode() as i:
+        out = [a.inc.bind(i) for a in actors]
+        dag = MultiOutputNode(out)
+
+    compiled_dag = dag.experimental_compile()
+    for i in range(10):
+        refs = compiled_dag.execute(1)
+        ray.get(refs)
+
+
+def test_dag_teardown_while_running(ray_start_regular):
+    a = Actor.remote(0)
+
+    with InputNode() as inp:
+        dag = a.sleep.bind(inp)
+
+    compiled_dag = dag.experimental_compile()
+    ref = compiled_dag.execute(3)  # 3-second slow task running async
+    compiled_dag.teardown()
+    try:
+        ray.get(ref)  # Sanity check the channel doesn't block.
+    except Exception:
+        pass
+
+    # Check we can still use the actor after first DAG teardown.
+    with InputNode() as inp:
+        dag = a.sleep.bind(inp)
+
+    compiled_dag = dag.experimental_compile()
+    ref = compiled_dag.execute(0.1)
+    result = ray.get(ref)
+    assert result == 0.1
+
+
+def test_asyncio_exceptions(ray_start_regular):
+    a = Actor.remote(0)
+    with InputNode() as i:
+        dag = a.inc.bind(i)
+
+    loop = get_or_create_event_loop()
+    compiled_dag = dag.experimental_compile(enable_asyncio=True)
+
+    async def main():
+        fut = await compiled_dag.execute_async(1)
+        result = await fut
+        assert result == 1
+
+        fut = await compiled_dag.execute_async("hello")
+        with pytest.raises(TypeError) as exc_info:
+            await fut
+        # Traceback should match the original actor class definition.
+        assert "self.i += x" in str(exc_info.value)
+
+        # Can throw an error multiple times.
+        fut = await compiled_dag.execute_async("hello")
+        with pytest.raises(TypeError) as exc_info:
+            await fut
+        # Traceback should match the original actor class definition.
+        assert "self.i += x" in str(exc_info.value)
+
+        # Can use the DAG after exceptions are thrown.
+        fut = await compiled_dag.execute_async(1)
+        result = await fut
+        assert result == 2
+
+    loop.run_until_complete(main())
+
+
+def test_channel_read_after_close(ray_start_regular):
+    # Tests that read to a channel after Compiled Graph teardown raises a
+    # RayChannelError exception as the channel is closed (see issue #46284).
+    @ray.remote
+    class Actor:
+        def foo(self, arg):
+            return arg
+
+    a = Actor.remote()
+    with InputNode() as inp:
+        dag = a.foo.bind(inp)
+
+    dag = dag.experimental_compile()
+    ref = dag.execute(1)
+    dag.teardown()
+
+    with pytest.raises(RayChannelError, match="Channel closed."):
+        ray.get(ref)
+
+
+def test_channel_write_after_close(ray_start_regular):
+    # Tests that write to a channel after Compiled Graph teardown raises a
+    # RayChannelError exception as the channel is closed.
+    @ray.remote
+    class Actor:
+        def foo(self, arg):
+            return arg
+
+    a = Actor.remote()
+    with InputNode() as inp:
+        dag = a.foo.bind(inp)
+
+    dag = dag.experimental_compile()
+    dag.teardown()
+
+    with pytest.raises(RayChannelError, match="Channel closed."):
+        dag.execute(1)
+
+
+def test_multi_arg_exception(shutdown_only):
+    a = Actor.remote(0)
+    with InputNode() as i:
+        o1, o2 = a.return_two_but_raise_exception.bind(i)
+        dag = MultiOutputNode([o1, o2])
+
+    compiled_dag = dag.experimental_compile()
+    for _ in range(3):
+        x, y = compiled_dag.execute(1)
+        with pytest.raises(RuntimeError):
+            ray.get(x)
+        with pytest.raises(RuntimeError):
+            ray.get(y)
+
+
+def test_multi_arg_exception_async(shutdown_only):
+    a = Actor.remote(0)
+    with InputNode() as i:
+        o1, o2 = a.return_two_but_raise_exception.bind(i)
+        dag = MultiOutputNode([o1, o2])
+
+    compiled_dag = dag.experimental_compile(enable_asyncio=True)
+
+    async def main():
+        for _ in range(3):
+            x, y = await compiled_dag.execute_async(1)
+            with pytest.raises(RuntimeError):
+                await x
+            with pytest.raises(RuntimeError):
+                await y
+
+    loop = get_or_create_event_loop()
+    loop.run_until_complete(main())
+
+
+def test_signature_mismatch(shutdown_only):
+    @ray.remote
+    class Worker:
+        def w(self, x):
+            return 1
+
+        def f(self, x, *, y):
+            pass
+
+        def g(self, x, y, z=1):
+            pass
+
+    worker = Worker.remote()
+    with pytest.raises(
+        TypeError,
+        match=(
+            r"got an unexpected keyword argument 'y'\. The function `w` has a "
+            r"signature `\(x\)`, but the given arguments to `bind` doesn't match\. "
+            r".*args:.*kwargs:.*"
+        ),
+    ):
+        with InputNode() as inp:
+            _ = worker.w.bind(inp, y=inp)
+
+    with pytest.raises(
+        TypeError,
+        match=(
+            r"too many positional arguments\. The function `w` has a signature "
+            r"`\(x\)`, but the given arguments to `bind` doesn't match\. "
+            r"args:.*kwargs:.*"
+        ),
+    ):
+        with InputNode() as inp:
+            _ = worker.w.bind(inp, inp)
+
+    with pytest.raises(
+        TypeError,
+        # Starting from Python 3.12, the error message includes "keyword-only."
+        # Therefore, we need to match both "required keyword-only argument" and
+        # "required argument."
+        match=(
+            r"missing a required (keyword-only )?argument: 'y'\. "
+            r"The function `f` has a signature `\(x, \*, y\)`, "
+            r"but the given arguments to `bind` doesn't match\. "
+            r"args:.*kwargs:.*"
+        ),
+    ):
+        with InputNode() as inp:
+            _ = worker.f.bind(inp)
+
+    with pytest.raises(
+        TypeError,
+        match=(
+            r"missing a required argument: 'y'\. The function `g` has a signature "
+            r"`\(x, y, z=1\)`, but the given arguments to `bind` doesn't match\. "
+            r"args:.*kwargs:.*"
+        ),
+    ):
+        with InputNode() as inp:
+            _ = worker.g.bind(inp)
+
+
+def test_missing_input_node():
+    @ray.remote
+    class Actor:
+        def __init__(self):
+            pass
+
+        def f(self, input):
+            return input
+
+        def add(self, a, b):
+            return a + b
+
+    actor = Actor.remote()
+
+    with ray.dag.InputNode() as dag_input:
+        input0, input1, input2 = dag_input[0], dag_input[1], dag_input[2]
+        _ = actor.f.bind(input1)
+        dag = actor.add.bind(input0, input2)
+
+    with pytest.raises(
+        ValueError,
+        match="Compiled Graph expects input to be accessed "
+        "using all of attributes 0, 1, 2, "
+        "but 1 is unused. "
+        "Ensure all input attributes are used and contribute "
+        "to the computation of the Compiled Graph output.",
+    ):
+        dag.experimental_compile()
+
+
+def test_sigint_get_dagref(ray_start_cluster):
+    driver_script = """
+import ray
+from ray.dag import InputNode
+import time
+
+ray.init()
+
+@ray.remote
+class Actor:
+    def sleep(self, x):
+        time.sleep(x)
+
+a = Actor.remote()
+with InputNode() as inp:
+    dag = a.sleep.bind(inp)
+compiled_dag = dag.experimental_compile()
+ref = compiled_dag.execute(100)
+print("executing", flush=True)
+ray.get(ref)
+"""
+    driver_proc = run_string_as_driver_nonblocking(
+        driver_script, env={"RAY_CGRAPH_teardown_timeout": "0"}
+    )
+    # wait for graph execution to start
+    assert driver_proc.stdout.readline() == b"executing\n"
+    driver_proc.send_signal(signal.SIGINT)  # ctrl+c
+    # teardown will kill actors after timeout
+    wait_for_pid_to_exit(driver_proc.pid, 10)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/dag/tests/experimental/test_dag_visualization.py

@@ -0,0 +1,712 @@
+import os
+import sys
+
+import pydot
+import pytest
+
+import ray
+from ray.dag import InputNode, MultiOutputNode
+from ray.tests.conftest import *  # noqa
+
+
+@pytest.fixture
+def cleanup_files():
+    """Clean up files generated during the test."""
+
+    def _cleanup_files(filename: str):
+        for ext in ["", ".png", ".pdf", ".jpeg", ".dot"]:
+            file_path = filename + ext
+            if os.path.exists(file_path):
+                os.remove(file_path)
+
+    return _cleanup_files
+
+
+def test_visualize_basic(ray_start_regular, cleanup_files):
+    """
+    Expect output or dot_source:
+        MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
+            0 -> 1 [label=SharedMemoryType]
+            1 -> 2 [label=SharedMemoryType]
+    """
+
+    @ray.remote
+    class Actor:
+        def echo(self, x):
+            return x
+
+    actor = Actor.remote()
+
+    with InputNode() as i:
+        dag = actor.echo.bind(i)
+
+    compiled_dag = dag.experimental_compile()
+
+    # Call the visualize method
+    dot_source = compiled_dag.visualize()
+
+    graphs = pydot.graph_from_dot_data(dot_source)
+    graph = graphs[0]
+
+    node_names = {node.get_name() for node in graph.get_nodes()}
+    edge_pairs = {
+        (edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
+    }
+
+    expected_nodes = {"0", "1", "2"}
+    assert expected_nodes.issubset(
+        node_names
+    ), f"Expected nodes {expected_nodes} not found."
+
+    expected_edges = {("0", "1"), ("1", "2")}
+    assert expected_edges.issubset(
+        edge_pairs
+    ), f"Expected edges {expected_edges} not found."
+
+    cleanup_files("compiled_graph")
+
+
+def test_visualize_multi_return(ray_start_regular, cleanup_files):
+    """
+    Expect output or dot_source:
+        MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
+            0 -> 1 [label=SharedMemoryType]
+            1 -> 2 [label=SharedMemoryType]
+            1 -> 3 [label=SharedMemoryType]
+            2 -> 4 [label=SharedMemoryType]
+            3 -> 4 [label=SharedMemoryType]
+    """
+
+    @ray.remote
+    class Actor:
+        @ray.method(num_returns=2)
+        def return_two(self, x):
+            return x, x + 1
+
+    actor = Actor.remote()
+
+    with InputNode() as i:
+        o1, o2 = actor.return_two.bind(i)
+        dag = MultiOutputNode([o1, o2])
+
+    compiled_dag = dag.experimental_compile()
+
+    # Get the DOT source
+    dot_source = compiled_dag.visualize()
+
+    graphs = pydot.graph_from_dot_data(dot_source)
+    graph = graphs[0]
+
+    node_names = {node.get_name() for node in graph.get_nodes()}
+    edge_pairs = {
+        (edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
+    }
+
+    expected_nodes = {"0", "1", "2", "3", "4"}
+    assert expected_nodes.issubset(
+        node_names
+    ), f"Expected nodes {expected_nodes} not found."
+
+    expected_edges = {("0", "1"), ("1", "2"), ("1", "3"), ("2", "4"), ("3", "4")}
+    assert expected_edges.issubset(
+        edge_pairs
+    ), f"Expected edges {expected_edges} not found."
+    cleanup_files("compiled_graph")
+
+
+def test_visualize_multi_return2(ray_start_regular, cleanup_files):
+    """
+    Expect output or dot_source:
+        MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
+            0 -> 1 [label=SharedMemoryType]
+            1 -> 2 [label=SharedMemoryType]
+            1 -> 3 [label=SharedMemoryType]
+            2 -> 4 [label=SharedMemoryType]
+            3 -> 5 [label=SharedMemoryType]
+            4 -> 6 [label=SharedMemoryType]
+            5 -> 6 [label=SharedMemoryType]
+    """
+
+    @ray.remote
+    class Actor:
+        @ray.method(num_returns=2)
+        def return_two(self, x):
+            return x, x + 1
+
+        def echo(self, x):
+            return x
+
+    a = Actor.remote()
+    b = Actor.remote()
+    with InputNode() as i:
+        o1, o2 = a.return_two.bind(i)
+        o3 = b.echo.bind(o1)
+        o4 = b.echo.bind(o2)
+        dag = MultiOutputNode([o3, o4])
+
+    compiled_dag = dag.experimental_compile()
+
+    # Get the DOT source
+    dot_source = compiled_dag.visualize()
+
+    graphs = pydot.graph_from_dot_data(dot_source)
+    graph = graphs[0]
+
+    node_names = {node.get_name() for node in graph.get_nodes()}
+    edge_pairs = {
+        (edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
+    }
+
+    expected_nodes = {"0", "1", "2", "3", "4", "5", "6"}
+    assert expected_nodes.issubset(
+        node_names
+    ), f"Expected nodes {expected_nodes} not found."
+
+    expected_edges = {
+        ("0", "1"),
+        ("1", "2"),
+        ("1", "3"),
+        ("2", "4"),
+        ("3", "5"),
+        ("4", "6"),
+        ("5", "6"),
+    }
+    assert expected_edges.issubset(
+        edge_pairs
+    ), f"Expected edges {expected_edges} not found."
+    cleanup_files("compiled_graph")
+
+
+def test_visualize_multi_input_nodes(ray_start_regular, cleanup_files):
+    """
+    Expect output or dot_source:
+        MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
+            0 -> 1
+            0 -> 2
+            0 -> 3
+            1 -> 4
+            2 -> 5
+            3 -> 6
+            4 -> 7
+            5 -> 7
+            6 -> 7
+    """
+
+    @ray.remote
+    class Actor:
+        def echo(self, x):
+            return x
+
+    actor = Actor.remote()
+
+    with InputNode() as inp:
+        o1 = actor.echo.bind(inp.x)
+        o2 = actor.echo.bind(inp.y)
+        o3 = actor.echo.bind(inp.z)
+        dag = MultiOutputNode([o1, o2, o3])
+
+    compiled_dag = dag.experimental_compile()
+
+    # Get the DOT source
+    dot_source = compiled_dag.visualize()
+
+    graphs = pydot.graph_from_dot_data(dot_source)
+    graph = graphs[0]
+
+    node_names = {node.get_name() for node in graph.get_nodes()}
+    edge_pairs = {
+        (edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
+    }
+
+    expected_nodes = {"0", "1", "2", "3", "4", "5", "6", "7"}
+    assert expected_nodes.issubset(
+        node_names
+    ), f"Expected nodes {expected_nodes} not found."
+
+    expected_edges = {
+        ("0", "1"),
+        ("0", "2"),
+        ("0", "3"),
+        ("1", "4"),
+        ("2", "5"),
+        ("3", "6"),
+        ("4", "7"),
+        ("5", "7"),
+        ("6", "7"),
+    }
+    assert expected_edges.issubset(
+        edge_pairs
+    ), f"Expected edges {expected_edges} not found."
+    cleanup_files("compiled_graph")
+
+
+class TestVisualizationAscii:
+
+    """Tests for the visualize_ascii method of compiled DAGs."""
+
+    @staticmethod
+    def parse_ascii_visualization(ascii_visualization: str):
+        """
+        Parses the ASCII visualization output to extract node names and edge pairs.
+
+        Args:
+            ascii_visualization: The ASCII visualization
+                output generated by the `visualize` function.
+
+        Returns:
+            tuple: A tuple containing:
+                - node_names: A set of strings representing node names.
+                - edge_pairs: A set of tuples representing edge
+                pairs with type hints.
+        """
+        import re
+
+        # Sets to store unique nodes and edges
+        node_names = set()
+        edge_pairs = set()
+
+        # Extract nodes from "Nodes Information" section
+        node_pattern = re.compile(r'^(\d+) \[label="Task \d+')
+        edge_pattern = re.compile(r"^(\d+) (--->|\+\+\+>) (\d+)")
+
+        lines = ascii_visualization.splitlines()
+        in_nodes_section = False
+        in_edges_section = False
+
+        for line in lines:
+            line = line.strip()
+
+            # Check for nodes section
+            if line.startswith("Nodes Information:"):
+                in_nodes_section = True
+                in_edges_section = False
+                continue
+
+            # Check for edges section
+            if line.startswith("Edges Information:"):
+                in_edges_section = True
+                in_nodes_section = False
+                continue
+
+            # Collect nodes
+            if in_nodes_section:
+                node_match = node_pattern.match(line)
+                if node_match:
+                    node_id = node_match.group(1)
+                    node_names.add(node_id)
+
+            # Collect edges
+            if in_edges_section:
+                edge_match = edge_pattern.match(line)
+                if edge_match:
+                    from_node, _, to_node = edge_match.groups()
+                    edge_pairs.add((from_node, to_node))
+
+        return node_names, edge_pairs
+
+    def test_visualize_ascii_basic(self, ray_start_regular):
+        """
+        Expect output:
+            Nodes Information:
+            0 [label="Task 0  InputNode"]
+            1 [label="Task 1  Actor: d6c5c4... Method: echo"]
+            2 [label="Task 2  MultiOutputNode"]
+
+            Edges Information:
+            0 ---> 1
+            1 ---> 2
+
+            Legend:
+            +++> : Represents Nccl-type data channels
+            ---> : Represents Shared Memory data channels
+
+            Experimental Graph:
+            0:InputNode
+            |
+            1:Actor_d6c5c4:echo
+            |
+            2:MultiOutputNode
+        """
+
+        @ray.remote
+        class Actor:
+            def echo(self, x):
+                return x
+
+        actor = Actor.remote()
+
+        with InputNode() as i:
+            dag = actor.echo.bind(i)
+
+        compiled_dag = dag.experimental_compile()
+
+        # Call the visualize method
+        ascii_visualization = compiled_dag.visualize(format="ascii")
+
+        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
+            ascii_visualization
+        )
+        print(node_names, edge_pairs)
+        expected_nodes = {"0", "1", "2"}
+        assert expected_nodes.issubset(
+            node_names
+        ), f"Expected nodes {expected_nodes} not found."
+
+        expected_edges = {("0", "1"), ("1", "2")}
+        assert expected_edges.issubset(
+            edge_pairs
+        ), f"Expected edges {expected_edges} not found."
+
+    def test_visualize_ascii_multi_return(self, ray_start_regular):
+        """
+        Expect output:
+            Nodes Information:
+            0 [label="Task 0  InputNode"]
+            1 [label="Task 1  Actor: 885f1d... Method: return_two"]
+            2 [label="Task 2  ClassMethodOutputNode[0]"]
+            3 [label="Task 3  ClassMethodOutputNode[1]"]
+            4 [label="Task 4  MultiOutputNode"]
+
+            Edges Information:
+            0 ---> 1
+            1 ---> 2
+            1 ---> 3
+            2 ---> 4
+            3 ---> 4
+
+            Legend:
+            +++> : Represents Nccl-type data channels
+            ---> : Represents Shared Memory data channels
+
+            Graph Built:
+            0:InputNode
+            |
+            1:Actor_885f1d:return_two
+            |---------------------------->|
+            2:Output[0]                   3:Output[1]
+            |<----------------------------|
+            4:MultiOutputNode
+        """
+
+        @ray.remote
+        class Actor:
+            @ray.method(num_returns=2)
+            def return_two(self, x):
+                return x, x + 1
+
+        actor = Actor.remote()
+
+        with InputNode() as i:
+            o1, o2 = actor.return_two.bind(i)
+            dag = MultiOutputNode([o1, o2])
+
+        compiled_dag = dag.experimental_compile()
+
+        ascii_visualization = compiled_dag.visualize(format="ascii")
+
+        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
+            ascii_visualization
+        )
+
+        expected_nodes = {"0", "1", "2", "3", "4"}
+        assert expected_nodes.issubset(
+            node_names
+        ), f"Expected nodes {expected_nodes} not found."
+
+        expected_edges = {("0", "1"), ("1", "2"), ("1", "3"), ("2", "4"), ("3", "4")}
+        assert expected_edges.issubset(
+            edge_pairs
+        ), f"Expected edges {expected_edges} not found."
+
+    def test_visualize_ascii_multi_return2(self, ray_start_regular):
+        """
+        Expect output:
+            Nodes Information:
+            0 [label="Task 0  InputNode"]
+            1 [label="Task 1  Actor: f3e919... Method: return_two"]
+            2 [label="Task 2  ClassMethodOutputNode[0]"]
+            3 [label="Task 3  ClassMethodOutputNode[1]"]
+            4 [label="Task 4  Actor: 15ec69... Method: echo"]
+            5 [label="Task 5  Actor: 15ec69... Method: echo"]
+            6 [label="Task 6  MultiOutputNode"]
+
+            Edges Information:
+            0 ---> 1
+            1 ---> 2
+            1 ---> 3
+            2 ---> 4
+            3 ---> 5
+            4 ---> 6
+            5 ---> 6
+
+            Legend:
+            +++> : Represents Nccl-type data channels
+            ---> : Represents Shared Memory data channels
+
+            Graph Built:
+            0:InputNode
+            |
+            1:Actor_f3e919:return_two
+            |---------------------------->|
+            2:Output[0]                   3:Output[1]
+            |                             |
+            4:Actor_15ec69:echo           5:Actor_15ec69:echo
+            |<----------------------------|
+            6:MultiOutputNode
+        """
+
+        @ray.remote
+        class Actor:
+            @ray.method(num_returns=2)
+            def return_two(self, x):
+                return x, x + 1
+
+            def echo(self, x):
+                return x
+
+        a = Actor.remote()
+        b = Actor.remote()
+        with InputNode() as i:
+            o1, o2 = a.return_two.bind(i)
+            o3 = b.echo.bind(o1)
+            o4 = b.echo.bind(o2)
+            dag = MultiOutputNode([o3, o4])
+
+        compiled_dag = dag.experimental_compile()
+
+        ascii_visualization = compiled_dag.visualize(format="ascii")
+
+        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
+            ascii_visualization
+        )
+
+        expected_nodes = {"0", "1", "2", "3", "4", "5", "6"}
+        assert expected_nodes.issubset(
+            node_names
+        ), f"Expected nodes {expected_nodes} not found."
+
+        expected_edges = {
+            ("0", "1"),
+            ("1", "2"),
+            ("1", "3"),
+            ("2", "4"),
+            ("3", "5"),
+            ("4", "6"),
+            ("5", "6"),
+        }
+        assert expected_edges.issubset(
+            edge_pairs
+        ), f"Expected edges {expected_edges} not found."
+
+    def test_visualize_ascii_complicate(self, ray_start_regular):
+        """
+        Expect output:
+            Nodes Information:
+            0 [label="Task 0  InputNode"]
+            1 [label="Task 1  Actor: 54777d... Method: return_three"]
+            2 [label="Task 2  ClassMethodOutputNode[0]"]
+            3 [label="Task 3  ClassMethodOutputNode[1]"]
+            4 [label="Task 4  ClassMethodOutputNode[2]"]
+            5 [label="Task 5  Actor: c927c9... Method: echo"]
+            6 [label="Task 6  Actor: c927c9... Method: echo"]
+            7 [label="Task 7  Actor: c927c9... Method: return_two"]
+            8 [label="Task 8  MultiOutputNode"]
+            9 [label="Task 9  ClassMethodOutputNode[0]"]
+            10 [label="Task 10  ClassMethodOutputNode[1]"]
+
+            Edges Information:
+            0 ---> 1
+            1 ---> 2
+            1 ---> 3
+            1 ---> 4
+            2 ---> 5
+            3 ---> 6
+            4 ---> 7
+            5 ---> 8
+            6 ---> 8
+            9 ---> 8
+            10 ---> 8
+            7 ---> 9
+            7 ---> 10
+
+            Legend:
+            +++> : Represents Nccl-type data channels
+            ---> : Represents Shared Memory data channels
+
+            Graph Built:
+            0:InputNode
+            |
+            1:Actor_54777d:return_three
+            |---------------------------->|---------------------------->|                                                  # noqa
+            2:Output[0]                   3:Output[1]                   4:Output[2]                                        # noqa
+            |                             |                             |                                                  # noqa
+            5:Actor_c927c9:echo           6:Actor_c927c9:echo           7:Actor_c927c9:return_two                          # noqa
+            |                             |                             |---------------------------->|                    # noqa
+            |                             |                             9:Output[0]                   10:Output[1]         # noqa
+            |<----------------------------|-----------------------------|-----------------------------|                    # noqa
+            8:MultiOutputNode
+        """
+
+        @ray.remote
+        class Actor:
+            @ray.method(num_returns=3)
+            def return_three(self, x):
+                return x, x + 1, x + 2
+
+            def echo(self, x):
+                return x
+
+            @ray.method(num_returns=2)
+            def return_two(self, x):
+                return x, x + 1
+
+        a = Actor.remote()
+        b = Actor.remote()
+        with InputNode() as i:
+            o1, o2, o3 = a.return_three.bind(i)
+            o4 = b.echo.bind(o1)
+            o5 = b.echo.bind(o2)
+            o6, o7 = b.return_two.bind(o3)
+            dag = MultiOutputNode([o4, o5, o6, o7])
+
+        compiled_dag = dag.experimental_compile()
+
+        ascii_visualization = compiled_dag.visualize(format="ascii")
+
+        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
+            ascii_visualization
+        )
+
+        expected_nodes = {"0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10"}
+        assert expected_nodes.issubset(
+            node_names
+        ), f"Expected nodes {expected_nodes} not found."
+
+        expected_edges = {
+            ("0", "1"),
+            ("1", "2"),
+            ("1", "3"),
+            ("1", "4"),
+            ("2", "5"),
+            ("3", "6"),
+            ("4", "7"),
+            ("5", "8"),
+            ("6", "8"),
+            ("9", "8"),
+            ("10", "8"),
+            ("7", "9"),
+            ("7", "10"),
+        }
+        assert expected_edges.issubset(
+            edge_pairs
+        ), f"Expected edges {expected_edges} not found."
+
+    def test_visualize_ascii_cross_line(self, ray_start_regular):
+        """
+        Expect output:
+            Nodes Information:
+            0 [label="Task 0  InputNode"]
+            1 [label="Task 1  Actor: 84835a... Method: return_three"]
+            2 [label="Task 2  ClassMethodOutputNode[0]"]
+            3 [label="Task 3  ClassMethodOutputNode[1]"]
+            4 [label="Task 4  ClassMethodOutputNode[2]"]
+            5 [label="Task 5  Actor: 02a6a1... Method: echo"]
+            6 [label="Task 6  Actor: 02a6a1... Method: return_two"]
+            7 [label="Task 7  Actor: 02a6a1... Method: echo"]
+            8 [label="Task 8  MultiOutputNode"]
+            9 [label="Task 9  ClassMethodOutputNode[0]"]
+            10 [label="Task 10  ClassMethodOutputNode[1]"]
+
+            Edges Information:
+            0 ---> 1
+            1 ---> 2
+            1 ---> 3
+            1 ---> 4
+            2 ---> 5
+            3 ---> 6
+            4 ---> 7
+            5 ---> 8
+            7 ---> 8
+            9 ---> 8
+            10 ---> 8
+            6 ---> 9
+            6 ---> 10
+
+            Legend:
+            +++> : Represents Nccl-type data channels
+            ---> : Represents Shared Memory data channels
+
+            Graph Built:
+            0:InputNode
+            |
+            1:Actor_84835a:return_three
+            |---------------------------->|---------------------------->|                            # noqa
+            2:Output[0]                   3:Output[1]                   4:Output[2]                  # noqa
+            |                             |                             |                            # noqa
+            5:Actor_02a6a1:echo           6:Actor_02a6a1:return_two     7:Actor_02a6a1:echo          # noqa
+            |                             |---------------------------->|                            # noqa
+            |                             9:Output[0]                   10:Output[1]                 # noqa
+            |<----------------------------------------------------------|                            # noqa
+            8:MultiOutputNod
+        """
+
+        @ray.remote
+        class Actor:
+            @ray.method(num_returns=3)
+            def return_three(self, x):
+                return x, x + 1, x + 2
+
+            def echo(self, x):
+                return x
+
+            @ray.method(num_returns=2)
+            def return_two(self, x):
+                return x, x + 1
+
+        a = Actor.remote()
+        b = Actor.remote()
+        with InputNode() as i:
+            o1, o2, o3 = a.return_three.bind(i)
+            o4 = b.echo.bind(o1)
+            o5 = b.echo.bind(o3)
+            o6, o7 = b.return_two.bind(o2)
+            dag = MultiOutputNode([o4, o5, o6, o7])
+
+        compiled_dag = dag.experimental_compile()
+
+        ascii_visualization = compiled_dag.visualize(format="ascii")
+
+        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
+            ascii_visualization
+        )
+
+        expected_nodes = {"0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10"}
+        assert expected_nodes.issubset(
+            node_names
+        ), f"Expected nodes {expected_nodes} not found."
+
+        expected_edges = {
+            ("0", "1"),
+            ("1", "2"),
+            ("1", "3"),
+            ("1", "4"),
+            ("2", "5"),
+            ("3", "6"),
+            ("4", "7"),
+            ("5", "8"),
+            ("7", "8"),
+            ("9", "8"),
+            ("10", "8"),
+            ("6", "9"),
+            ("6", "10"),
+        }
+        assert expected_edges.issubset(
+            edge_pairs
+        ), f"Expected edges {expected_edges} not found."
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/dag/tests/experimental/test_execution_schedule_gpu.py

@@ -0,0 +1,469 @@
+# coding: utf-8
+import os
+import sys
+from typing import Optional
+
+import pytest
+import torch
+
+import ray
+import ray.cluster_utils
+from ray.dag import InputNode, MultiOutputNode
+from ray.dag.compiled_dag_node import CompiledDAG
+from ray.dag.dag_node_operation import _DAGNodeOperationType
+from ray.tests.conftest import *  # noqa
+
+if sys.platform != "linux" and sys.platform != "darwin":
+    pytest.skip("Skipping, requires Linux or Mac.", allow_module_level=True)
+
+USE_GPU = os.environ.get("RAY_PYTEST_USE_GPU") == "1"
+
+if not USE_GPU:
+    pytest.skip("Skipping, these tests require GPUs.", allow_module_level=True)
+
+
+@ray.remote(num_cpus=0, num_gpus=1)
+class Worker:
+    def __init__(self, rank: Optional[int] = None):
+        self.rank = rank
+        self.trace = []
+
+    def fwd(self, value):
+        self.trace.append(("FWD", self.rank))
+        return value
+
+    def bwd(self, value):
+        self.trace.append(("BWD", self.rank))
+        return value
+
+    def pop_trace(self):
+        trace = self.trace
+        self.trace = []
+        return trace
+
+    def read_input(self, input):
+        return input
+
+    def send(self, shape, dtype, value: int, send_tensor=True):
+        if not send_tensor:
+            return 1
+        return torch.ones(shape, dtype=dtype, device=self.device) * value
+
+    def recv(self, tensor):
+        # Check that tensor got loaded to the correct device.
+        assert tensor.device == self.device
+        return (tensor[0].item(), tensor.shape, tensor.dtype)
+
+    def no_op(self, value):
+        return value
+
+    def no_op_two(self, value1, value2):
+        return value1, value2
+
+
+def generate_1f1b_dag(
+    num_workers: int, num_microbatches: int, num_lead_microbatches: int
+) -> CompiledDAG:
+    workers = [Worker.remote(rank) for rank in range(num_workers)]
+
+    with ray.dag.InputNode() as inp:
+        fwd_queues = [[] for _ in range(num_workers)]
+        bwd_queues = [[] for _ in range(num_workers)]
+        # Once a worker's counter reaches 0, it cannot execute another fwd until it
+        # executes a bwd first.
+        fwd_counter = [num_lead_microbatches - i for i in range(num_workers)]
+        # All of the done batches.
+        done = []
+
+        # FWD on worker 0.
+        input_data = workers[0].read_input.bind(inp)
+        for i in range(num_microbatches):
+            fwd_queues[0].append(input_data)
+
+        while len(done) < num_microbatches:
+            for i, worker in enumerate(workers):
+                if fwd_counter[i] > 0 and fwd_queues[i]:
+                    b = fwd_queues[i].pop(0)
+                    b = worker.fwd.bind(b)
+                    if i < num_workers - 1:
+                        fwd_queues[i + 1].append(b)
+                        # Use NCCL channel for communication between workers.
+                        b.with_tensor_transport(transport="nccl")
+                    else:
+                        bwd_queues[i].append(b)
+                    fwd_counter[i] -= 1
+                elif bwd_queues[i]:
+                    b = bwd_queues[i].pop(0)
+                    b = worker.bwd.bind(b)
+                    if i > 0:
+                        bwd_queues[i - 1].append(b)
+                        # Use NCCL channel for communication between workers.
+                        b.with_tensor_transport(transport="nccl")
+                    else:
+                        done.append(b)
+                    fwd_counter[i] += 1
+        dag = ray.dag.MultiOutputNode(done)
+    compiled_dag = dag.experimental_compile()
+    return compiled_dag
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 2}], indirect=True)
+@pytest.mark.parametrize("single_fetch", [True, False])
+def test_simulate_pp_2workers_2batches_1f1b(
+    ray_start_regular, single_fetch, monkeypatch
+):
+    """
+    This test simulates a simple 1F1B pipeline parallelism for training with
+    2 workers and 2 batches.
+
+    w1: fwd_b1  fwd_b2          bwd_b1          bwd_b2
+    w2:         fwd_b1  bwd_b1  fwd_b2  bwd_b2
+
+    The communication between workers is done using NCCL. The communication
+    within the worker actor is done using IntraProcessChannel.
+    """
+    if not USE_GPU:
+        pytest.skip("NCCL tests require GPUs")
+
+    w1 = Worker.remote()
+    w2 = Worker.remote()
+
+    with InputNode() as inp:
+        w1_input = w1.read_input.bind(inp)
+        batch_1 = w1.fwd.bind(w1_input)
+        batch_1.with_tensor_transport(transport="nccl")
+        batch_2 = w1.fwd.bind(w1_input)
+        batch_2.with_tensor_transport(transport="nccl")
+        batch_1 = w2.fwd.bind(batch_1)
+        batch_1 = w2.bwd.bind(batch_1)
+        batch_1.with_tensor_transport(transport="nccl")
+        batch_2 = w2.fwd.bind(batch_2)
+        batch_1 = w1.bwd.bind(batch_1)
+        batch_2 = w2.bwd.bind(batch_2)
+        batch_2.with_tensor_transport(transport="nccl")
+        batch_2 = w1.bwd.bind(batch_2)
+        dag = MultiOutputNode([batch_1, batch_2])
+    compiled_dag = dag.experimental_compile()
+
+    w1_expected_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.READ),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+        # `w1 (3, READ)` (P2P recv) is scheduled together with
+        # `w2 (1, WRITE)` (P2P send).
+        (3, _DAGNodeOperationType.READ),
+        (2, _DAGNodeOperationType.READ),
+        (2, _DAGNodeOperationType.COMPUTE),
+        (2, _DAGNodeOperationType.WRITE),
+        # `w1 (4, READ)` (P2P recv) is scheduled together with
+        # `w2 (3, WRITE)` (P2P send).
+        (4, _DAGNodeOperationType.READ),
+        (3, _DAGNodeOperationType.COMPUTE),
+        (3, _DAGNodeOperationType.WRITE),
+        (4, _DAGNodeOperationType.COMPUTE),
+        (4, _DAGNodeOperationType.WRITE),
+    ]
+    w2_expected_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.READ),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+        (2, _DAGNodeOperationType.READ),
+        (2, _DAGNodeOperationType.COMPUTE),
+        (2, _DAGNodeOperationType.WRITE),
+        (3, _DAGNodeOperationType.READ),
+        (3, _DAGNodeOperationType.COMPUTE),
+        (3, _DAGNodeOperationType.WRITE),
+    ]
+    w1_schedule = compiled_dag.actor_to_execution_schedule[w1]
+    w2_schedule = compiled_dag.actor_to_execution_schedule[w2]
+
+    for schedule, expected_schedule in zip(
+        [w1_schedule, w2_schedule], [w1_expected_schedule, w2_expected_schedule]
+    ):
+        assert len(schedule) == len(expected_schedule)
+        for i, operation in enumerate(schedule):
+            assert operation.exec_task_idx == expected_schedule[i][0]
+            assert operation.type == expected_schedule[i][1]
+
+    tensor_cpu = torch.zeros(10, 10)
+    tensor_cuda = tensor_cpu.to("cuda:0")
+    refs = compiled_dag.execute(tensor_cuda)
+
+    if single_fetch:
+        assert len(refs) == 2
+        for ref in refs:
+            tensor = ray.get(ref)
+            assert torch.equal(tensor.cpu(), tensor_cpu)
+    else:
+        tensors = ray.get(refs)
+        assert len(tensors) == 2
+        for tensor in tensors:
+            assert torch.equal(tensor.cpu(), tensor_cpu)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 4}], indirect=True)
+def test_simulate_pp_4workers_8batches_1f1b(ray_start_regular, monkeypatch):
+    """
+    This test simulates a 1F1B pipeline parallelism for training with
+    4 workers and 8 batches.
+    """
+    if not USE_GPU:
+        pytest.skip("NCCL tests require GPUs")
+
+    num_workers, num_microbatches, num_lead_microbatches = 4, 8, 4
+    compiled_dag = generate_1f1b_dag(
+        num_workers, num_microbatches, num_lead_microbatches
+    )
+
+    tensor_cpu = torch.zeros(10, 10)
+    tensor_cuda = tensor_cpu.to("cuda:0")
+    tensors = ray.get(compiled_dag.execute(tensor_cuda))
+
+    assert len(tensors) == num_microbatches
+    for t in tensors:
+        assert torch.equal(t.cpu(), tensor_cpu)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 3}], indirect=True)
+def test_three_actors_with_nccl_1(ray_start_regular):
+    """
+    Driver -> a.no_op -> b.no_op -> a.no_op_two -> Driver
+                      |          |
+                      -> c.no_op -
+    """
+    if not USE_GPU:
+        pytest.skip("NCCL tests require GPUs")
+
+    a = Worker.remote()
+    b = Worker.remote()
+    c = Worker.remote()
+
+    with InputNode() as inp:
+        dag = a.no_op.bind(inp)
+        dag.with_tensor_transport(transport="nccl")
+        branch1 = b.no_op.bind(dag)
+        branch1.with_tensor_transport(transport="nccl")
+        branch2 = c.no_op.bind(dag)
+        branch2.with_tensor_transport(transport="nccl")
+        dag = a.no_op_two.bind(branch1, branch2)
+
+    compiled_dag = dag.experimental_compile()
+
+    a_expected_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.READ),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+    ]
+    b_expected_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+    ]
+    c_expected_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+    ]
+    a_schedule = compiled_dag.actor_to_execution_schedule[a]
+    b_schedule = compiled_dag.actor_to_execution_schedule[b]
+    c_schedule = compiled_dag.actor_to_execution_schedule[c]
+
+    for schedule, expected_schedule in zip(
+        [a_schedule, b_schedule, c_schedule],
+        [a_expected_schedule, b_expected_schedule, c_expected_schedule],
+    ):
+        assert len(schedule) == len(expected_schedule)
+        for i, operation in enumerate(schedule):
+            assert operation.exec_task_idx == expected_schedule[i][0]
+            assert operation.type == expected_schedule[i][1]
+
+    tensor_cpu = torch.zeros(10, 10)
+    tensor_cuda = tensor_cpu.to("cuda:0")
+    ref = compiled_dag.execute(tensor_cuda)
+    tensors = ray.get(ref)
+
+    assert len(tensors) == 2
+    for t in tensors:
+        assert torch.equal(t.cpu(), tensor_cpu)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 3}], indirect=True)
+@pytest.mark.parametrize("single_fetch", [True, False])
+def test_three_actors_with_nccl_2(ray_start_regular, single_fetch, monkeypatch):
+    """
+    Driver --> a.no_op -> b.no_op --> Driver
+            |                     |
+            -> b.no_op -> c.no_op -
+            |                     |
+            -> c.no_op -> a.no_op -
+    """
+    if not USE_GPU:
+        pytest.skip("NCCL tests require GPUs")
+
+    a = Worker.remote()
+    b = Worker.remote()
+    c = Worker.remote()
+
+    with InputNode() as inp:
+        branch1 = a.no_op.bind(inp)
+        branch1.with_tensor_transport(transport="nccl")
+        branch2 = b.no_op.bind(inp)
+        branch2.with_tensor_transport(transport="nccl")
+        branch3 = c.no_op.bind(inp)
+        branch3.with_tensor_transport(transport="nccl")
+        dag = MultiOutputNode(
+            [
+                a.no_op.bind(branch3),
+                b.no_op.bind(branch1),
+                c.no_op.bind(branch2),
+            ]
+        )
+
+    compiled_dag = dag.experimental_compile()
+
+    a_expected_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.READ),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+    ]
+    b_expected_schedule = [
+        # `b (1, READ)` (P2P recv) is scheduled together with
+        # `a (0, WRITE)` (P2P send).
+        (1, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+    ]
+    c_expected_schedule = [
+        # `c (1, READ)` (P2P recv) is scheduled together with
+        # `a (0, WRITE)` (P2P send).
+        (1, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+    ]
+
+    a_schedule = compiled_dag.actor_to_execution_schedule[a]
+    b_schedule = compiled_dag.actor_to_execution_schedule[b]
+    c_schedule = compiled_dag.actor_to_execution_schedule[c]
+
+    for schedule, expected_schedule in zip(
+        [a_schedule, b_schedule, c_schedule],
+        [a_expected_schedule, b_expected_schedule, c_expected_schedule],
+    ):
+        assert len(schedule) == len(expected_schedule)
+        for i, operation in enumerate(schedule):
+            assert operation.exec_task_idx == expected_schedule[i][0]
+            assert operation.type == expected_schedule[i][1]
+
+    tensor_cpu = torch.zeros(10, 10)
+    tensor_cuda = tensor_cpu.to("cuda:0")
+    refs = compiled_dag.execute(tensor_cuda)
+
+    if single_fetch:
+        assert len(refs) == 3
+        for ref in refs:
+            tensor = ray.get(ref)
+            assert torch.equal(tensor.cpu(), tensor_cpu)
+    else:
+        tensors = ray.get(refs)
+        assert len(tensors) == 3
+        for tensor in tensors:
+            assert torch.equal(tensor.cpu(), tensor_cpu)
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 3}], indirect=True)
+@pytest.mark.parametrize("overlap_gpu_communication", [True, False])
+def test_overlap_gpu_communication(ray_start_regular, overlap_gpu_communication):
+    """
+    Driver --> sender1.send -> receiver.recv --> Driver
+            |                                |
+            -> sender2.send -> receiver.recv -
+    """
+    if not USE_GPU:
+        pytest.skip("NCCL tests require GPUs")
+
+    sender1 = Worker.remote()
+    sender2 = Worker.remote()
+    receiver = Worker.remote()
+
+    shape = (10000,)
+    dtype = torch.float16
+
+    with InputNode() as inp:
+        branch1 = sender1.send.bind(shape, dtype, inp)
+
+        branch1 = branch1.with_tensor_transport(
+            transport="nccl", _static_shape=True, _direct_return=True
+        )
+        branch1 = receiver.recv.bind(branch1)
+
+        branch2 = sender2.send.bind(shape, dtype, inp)
+        branch2 = branch2.with_tensor_transport(
+            transport="nccl", _static_shape=True, _direct_return=True
+        )
+        branch2 = receiver.recv.bind(branch2)
+        dag = MultiOutputNode([branch1, branch2])
+
+    # Test normal execution.
+    compiled_dag = dag.experimental_compile(
+        _overlap_gpu_communication=overlap_gpu_communication
+    )
+
+    # Check receiver schedule
+    expected_no_overlap_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        # `receiver (1, READ)` (P2P recv) is scheduled together with
+        # `sender2 (0, WRITE)` (P2P send).
+        (1, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+    ]
+    expected_overlap_schedule = [
+        (0, _DAGNodeOperationType.READ),
+        # `receiver (1, READ)` (P2P recv) is scheduled together with
+        # `sender2 (0, WRITE)` (P2P send).
+        (1, _DAGNodeOperationType.READ),
+        (0, _DAGNodeOperationType.COMPUTE),
+        (0, _DAGNodeOperationType.WRITE),
+        (1, _DAGNodeOperationType.COMPUTE),
+        (1, _DAGNodeOperationType.WRITE),
+    ]
+    if overlap_gpu_communication:
+        expected_receiver_schedule = expected_overlap_schedule
+    else:
+        expected_receiver_schedule = expected_no_overlap_schedule
+
+    receiver_schedule = compiled_dag.actor_to_execution_schedule[receiver]
+
+    assert len(receiver_schedule) == len(expected_receiver_schedule)
+    for i, operation in enumerate(receiver_schedule):
+        assert operation.exec_task_idx == expected_receiver_schedule[i][0]
+        assert operation.type == expected_receiver_schedule[i][1]
+
+    compiled_dag.teardown()
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/dag/tests/experimental/test_mocked_nccl_dag.py

@@ -0,0 +1,408 @@
+# coding: utf-8
+import os
+import sys
+
+import pytest
+import torch
+
+import ray
+import ray.cluster_utils
+from ray._common.test_utils import wait_for_condition
+from ray.dag import InputNode
+from ray.exceptions import RayChannelError, RayTaskError
+from ray.experimental.channel.conftest import (
+    Barrier,
+    start_nccl_mock,
+)
+from ray.tests.conftest import *  # noqa
+
+
+def error_logged(capsys, msg):
+    out, err = capsys.readouterr()
+    # Write captured back to stdout, stderr for easier test debugging.
+    sys.stdout.write(out)
+    sys.stderr.write(err)
+    return msg in err
+
+
+@ray.remote(num_cpus=0, num_gpus=1)
+class MockedWorker:
+    def __init__(self):
+        self.chan = None
+
+    def start_mock(self):
+        """
+        Patch methods that require CUDA.
+        """
+        start_nccl_mock()
+
+    def send(self, shape, dtype, value: int, send_as_dict=False):
+        if send_as_dict:
+            return self.send_dict([(value, value, shape, dtype)])
+
+        return torch.ones(shape, dtype=dtype) * value
+
+    def recv(self, tensor):
+        if isinstance(tensor, dict):
+            assert len(tensor) == 1
+            tensor = list(tensor.values())[0]
+
+        return (tensor[0].item(), tensor.shape, tensor.dtype)
+
+    def send_dict(self, entries):
+        results = {}
+        for key, value, shape, dtype in entries:
+            results[key] = torch.ones(shape, dtype=dtype) * value
+        return results
+
+    def recv_dict(self, tensor_dict):
+        results = []
+        for key in sorted(tensor_dict.keys()):
+            tensor = tensor_dict[key]
+            results.append((key, tensor[0].item(), tensor.shape, tensor.dtype))
+        return results
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 2,
+            "num_gpus": 2,
+            "num_nodes": 1,
+        }
+    ],
+    indirect=True,
+)
+def test_p2p(ray_start_cluster):
+    """
+    Test simple sender -> receiver pattern. Check that receiver receives
+    correct results.
+    """
+    # Barrier name should be barrier-{lower rank}-{higher rank}.
+    barrier = Barrier.options(name="barrier-0-1").remote()  # noqa
+
+    sender = MockedWorker.remote()
+    receiver = MockedWorker.remote()
+
+    ray.get([sender.start_mock.remote(), receiver.start_mock.remote()])
+
+    shape = (10,)
+    dtype = torch.float16
+
+    # Test torch.Tensor sent between actors.
+    with InputNode() as inp:
+        dag = sender.send.bind(inp.shape, inp.dtype, inp[0], inp.send_as_dict)
+        dag = dag.with_tensor_transport(transport="nccl")
+        dag = receiver.recv.bind(dag)
+
+    compiled_dag = dag.experimental_compile()
+    for i in range(3):
+        ref = compiled_dag.execute(i, shape=shape, dtype=dtype, send_as_dict=False)
+        assert ray.get(ref) == (i, shape, dtype)
+
+    # Sending tensors of different shape also works.
+    for i in range(3):
+        ref = compiled_dag.execute(i, shape=(20,), dtype=dtype, send_as_dict=False)
+        assert ray.get(ref) == (i, (20,), dtype)
+
+    # Sending tensors inside a dictionary also works.
+    for i in range(3):
+        ref = compiled_dag.execute(i, shape=shape, dtype=dtype, send_as_dict=True)
+        assert ray.get(ref) == (i, shape, dtype)
+
+    compiled_dag.teardown()
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 2,
+            "num_gpus": 2,
+            "num_nodes": 1,
+        }
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("send_as_dict", [True, False])
+def test_p2p_static_shape(ray_start_cluster, send_as_dict):
+    """
+    Test simple send -> recv pattern with
+    _static_shape=True. If sender always sends tensors of
+    the same shape, then it works.
+    """
+    # Barrier name should be barrier-{lower rank}-{higher rank}.
+    barrier = Barrier.options(name="barrier-0-1").remote()  # noqa
+
+    sender = MockedWorker.remote()
+    receiver = MockedWorker.remote()
+
+    ray.get([sender.start_mock.remote(), receiver.start_mock.remote()])
+
+    shape = (10,)
+    dtype = torch.float16
+
+    # Test torch.Tensor sent between actors.
+    with InputNode() as inp:
+        dag = sender.send.bind(inp.shape, inp.dtype, inp[0], send_as_dict=send_as_dict)
+        dag = dag.with_tensor_transport(transport="nccl", _static_shape=True)
+        dag = receiver.recv.bind(dag)
+
+    compiled_dag = dag.experimental_compile()
+    for i in range(3):
+        ref = compiled_dag.execute(i, shape=shape, dtype=dtype)
+        assert ray.get(ref) == (i, shape, dtype)
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 2,
+            "num_gpus": 2,
+            "num_nodes": 1,
+        }
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("send_as_dict", [True, False])
+def test_p2p_static_shape_error(capsys, ray_start_cluster, send_as_dict):
+    """
+    Test that when static_shape=True, an error is thrown when a tensor with a
+    different shape or dtype is found.
+    """
+    # Barrier name should be barrier-{lower rank}-{higher rank}.
+    barrier = Barrier.options(name="barrier-0-1").remote()  # noqa
+
+    sender = MockedWorker.remote()
+    receiver = MockedWorker.remote()
+
+    ray.get([sender.start_mock.remote(), receiver.start_mock.remote()])
+
+    shape = (10,)
+    dtype = torch.float16
+
+    # Test torch.Tensor sent between actors.
+    with InputNode() as inp:
+        dag = sender.send.bind(inp.shape, inp.dtype, inp[0], send_as_dict=send_as_dict)
+        dag = dag.with_tensor_transport(transport="nccl", _static_shape=True)
+        dag = receiver.recv.bind(dag)
+
+    compiled_dag = dag.experimental_compile()
+    for i in range(3):
+        ref = compiled_dag.execute(i, shape=shape, dtype=dtype)
+        assert ray.get(ref) == (i, shape, dtype)
+
+    # Sending wrong shape errors.
+    ref = compiled_dag.execute(i, shape=(20,), dtype=dtype)
+    with pytest.raises(RayTaskError):
+        ray.get(ref)
+
+    # Sending correct shape still errors because the DAG has already been torn
+    # down after the previous error.
+    with pytest.raises(RayChannelError):
+        ref = compiled_dag.execute(i, shape=shape, dtype=dtype)
+
+    wait_for_condition(
+        lambda: error_logged(
+            capsys,
+            "ValueError: Expected torch.Tensors with shapes and dtypes: "
+            "[(shape=torch.Size([10]), dtype=torch.float16)], found: "
+            "[(shape=torch.Size([20]), dtype=torch.float16)]",
+        )
+    )
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 2,
+            "num_gpus": 2,
+            "num_nodes": 1,
+        }
+    ],
+    indirect=True,
+)
+def test_p2p_direct_return(ray_start_cluster):
+    """
+    Test simple sender -> receiver pattern with _direct_return=True
+    """
+    # Barrier name should be barrier-{lower rank}-{higher rank}.
+    barrier = Barrier.options(name="barrier-0-1").remote()  # noqa
+
+    sender = MockedWorker.remote()
+    receiver = MockedWorker.remote()
+
+    ray.get([sender.start_mock.remote(), receiver.start_mock.remote()])
+
+    # Test torch.Tensor sent between actors.
+    with InputNode() as inp:
+        dag = sender.send.bind(inp.shape, inp.dtype, inp.value, inp.send_as_dict)
+        dag = dag.with_tensor_transport(
+            transport="nccl",
+            _direct_return=True,
+        )
+        dag = receiver.recv.bind(dag)
+
+    compiled_dag = dag.experimental_compile()
+
+    dtype = torch.float16
+    for i in range(3):
+        shape = (10 * (i + 1),)
+        ref = compiled_dag.execute(
+            shape=shape, dtype=dtype, value=i, send_as_dict=False
+        )
+        assert ray.get(ref) == (i, shape, dtype)
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 2,
+            "num_gpus": 2,
+            "num_nodes": 1,
+        }
+    ],
+    indirect=True,
+)
+def test_p2p_direct_return_error(capsys, ray_start_cluster):
+    """
+    Test simple sender -> receiver pattern with
+    _direct_return=True. Test that error is thrown when
+    actor task does not return a tensor directly.
+    """
+    # Barrier name should be barrier-{lower rank}-{higher rank}.
+    barrier = Barrier.options(name="barrier-0-1").remote()  # noqa
+
+    sender = MockedWorker.remote()
+    receiver = MockedWorker.remote()
+
+    ray.get([sender.start_mock.remote(), receiver.start_mock.remote()])
+
+    # Test torch.Tensor sent between actors.
+    with InputNode() as inp:
+        dag = sender.send.bind(inp.shape, inp.dtype, inp.value, inp.send_as_dict)
+        dag = dag.with_tensor_transport(
+            transport="nccl",
+            _direct_return=True,
+        )
+        dag = receiver.recv.bind(dag)
+
+    compiled_dag = dag.experimental_compile()
+
+    dtype = torch.float16
+    for i in range(3):
+        shape = (10 * (i + 1),)
+        ref = compiled_dag.execute(
+            shape=shape, dtype=dtype, value=i, send_as_dict=False
+        )
+        assert ray.get(ref) == (i, shape, dtype)
+
+    # Error is thrown if we do not send a tensor.
+    ref = compiled_dag.execute(shape=shape, dtype=dtype, value=1, send_as_dict=True)
+    with pytest.raises(RayTaskError):
+        ray.get(ref)
+
+    # Currently the receiver cannot catch the exception so the DAG cannot be
+    # used again.
+    with pytest.raises(RayChannelError):
+        ref = compiled_dag.execute(
+            shape=shape, dtype=dtype, value=1, send_as_dict=False
+        )
+
+    wait_for_condition(
+        lambda: error_logged(
+            capsys,
+            "Task annotated with _direct_return=True must "
+            "return a CUDA torch.Tensor",
+        )
+    )
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 2,
+            "num_gpus": 2,
+            "num_nodes": 1,
+        }
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("check_static_shape", [True, False])
+def test_p2p_static_shape_and_direct_return(
+    capsys, ray_start_cluster, check_static_shape
+):
+    """
+    Test simple sender -> receiver pattern with both _static_shape=True and
+    _direct_return=True. Check errors are thrown if tensors with wrong shape
+    are passed (check_static_shape=True) OR if non-tensor value is returned
+    (check_static_shape=False).
+    """
+    # Barrier name should be barrier-{lower rank}-{higher rank}.
+    barrier = Barrier.options(name="barrier-0-1").remote()  # noqa
+
+    sender = MockedWorker.remote()
+    receiver = MockedWorker.remote()
+
+    ray.get([sender.start_mock.remote(), receiver.start_mock.remote()])
+
+    # Test torch.Tensor sent between actors.
+    with InputNode() as inp:
+        dag = sender.send.bind(inp.shape, inp.dtype, inp.value, inp.send_as_dict)
+        dag = dag.with_tensor_transport(
+            transport="nccl",
+            _static_shape=True,
+            _direct_return=True,
+        )
+        dag = receiver.recv.bind(dag)
+
+    compiled_dag = dag.experimental_compile()
+
+    shape = (10,)
+    dtype = torch.float16
+    for i in range(3):
+        ref = compiled_dag.execute(
+            shape=shape, dtype=dtype, value=i, send_as_dict=False
+        )
+        assert ray.get(ref) == (i, shape, dtype)
+
+    if check_static_shape:
+        # Error is thrown if we send the wrong shape.
+        ref = compiled_dag.execute(
+            shape=(20,), dtype=dtype, value=1, send_as_dict=False
+        )
+    else:
+        # Error is thrown if we do not send a tensor.
+        ref = compiled_dag.execute(shape=shape, dtype=dtype, value=1, send_as_dict=True)
+
+    with pytest.raises(RayTaskError):
+        ray.get(ref)
+
+    # Currently the receiver cannot catch either kind of
+    # exception so the DAG cannot be used again.
+    with pytest.raises(RayChannelError):
+        ref = compiled_dag.execute(
+            shape=shape, dtype=dtype, value=1, send_as_dict=False
+        )
+
+    if check_static_shape:
+        msg = (
+            "ValueError: Expected torch.Tensors with shapes and dtypes: "
+            "[(shape=torch.Size([10]), dtype=torch.float16)], found: "
+            "[(shape=torch.Size([20]), dtype=torch.float16)]"
+        )
+    else:
+        msg = "Task annotated with _direct_return=True must return a CUDA torch.Tensor"
+    wait_for_condition(lambda: error_logged(capsys, msg))
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/dag/tests/experimental/test_multi_args_gpu.py

@@ -0,0 +1,75 @@
+# coding: utf-8
+import os
+import sys
+
+import pytest
+import torch
+
+import ray
+import ray.cluster_utils
+from ray.dag import InputNode, MultiOutputNode
+from ray.tests.conftest import *  # noqa
+
+if sys.platform != "linux" and sys.platform != "darwin":
+    pytest.skip("Skipping, requires Linux or Mac.", allow_module_level=True)
+
+USE_GPU = os.environ.get("RAY_PYTEST_USE_GPU") == "1"
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 2}], indirect=True)
+def test_multi_args_simulate_pp(ray_start_regular):
+    if not USE_GPU:
+        pytest.skip("NCCL tests require GPUs")
+
+    @ray.remote(num_cpus=0, num_gpus=1)
+    class Worker:
+        def __init__(self):
+            pass
+
+        def forward(self, data):
+            return data
+
+        def backward(self, data):
+            return data
+
+    NUM_MICROBATCHES = 2
+    w0 = Worker.remote()
+    w1 = Worker.remote()
+    with InputNode() as dag_input:
+        dag_outs = []
+        for microbatch_idx in range(NUM_MICROBATCHES):
+            microbatch = dag_input[microbatch_idx]
+            stage_fwd_out = w0.forward.bind(microbatch)
+            stage_fwd_out.with_tensor_transport(transport="nccl")
+            stage_fwd_out = w1.forward.bind(stage_fwd_out)
+            dag_outs.append(stage_fwd_out)
+
+        grad_out = dag_input[NUM_MICROBATCHES]
+        for _ in range(NUM_MICROBATCHES):
+            stage_bwd_out = w1.backward.bind(grad_out)
+            stage_bwd_out.with_tensor_transport(transport="nccl")
+            stage_bwd_out = w0.backward.bind(stage_bwd_out)
+            dag_outs.append(stage_bwd_out)
+
+        dag = MultiOutputNode(dag_outs)
+    compiled_dag = dag.experimental_compile()
+
+    tensor_cpu_list = [torch.zeros(1, i + 1) for i in range(3)]
+    tensor_cuda_list = [t.to("cuda:0") for t in tensor_cpu_list]
+    ref = compiled_dag.execute(
+        tensor_cuda_list[0], tensor_cuda_list[1], tensor_cuda_list[2]
+    )
+    tensors = ray.get(ref)
+
+    assert len(tensors) == 4
+    assert torch.equal(tensors[0], tensor_cpu_list[0])
+    assert torch.equal(tensors[1], tensor_cpu_list[1])
+    assert torch.equal(tensors[2], tensor_cpu_list[2])
+    assert torch.equal(tensors[3], tensor_cpu_list[2])
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/dag/tests/experimental/test_multi_node_dag.py

@@ -0,0 +1,414 @@
+import os
+import random
+import sys
+import time
+
+import pytest
+
+import ray
+import ray.remote_function
+from ray._common.test_utils import wait_for_condition
+from ray.dag import InputNode, MultiOutputNode
+from ray.tests.conftest import *  # noqa
+
+if sys.platform != "linux" and sys.platform != "darwin":
+    pytest.skip("Skipping, requires Linux or Mac.", allow_module_level=True)
+
+
+@ray.remote
+class Actor:
+    def __init__(self, init_value, fail_after=None, sys_exit=False):
+        self.i = init_value
+        self.fail_after = fail_after
+        self.sys_exit = sys_exit
+
+        self.count = 0
+
+    def _fail_if_needed(self):
+        if self.fail_after and self.count > self.fail_after:
+            # Randomize the failures to better cover multi actor scenarios.
+            if random.random() > 0.5:
+                if self.sys_exit:
+                    os._exit(1)
+                else:
+                    raise RuntimeError("injected fault")
+
+    def inc(self, x):
+        self.i += x
+        self.count += 1
+        self._fail_if_needed()
+        return self.i
+
+    def double_and_inc(self, x):
+        self.i *= 2
+        self.i += x
+        return self.i
+
+    def echo(self, x):
+        print("ECHO!")
+        self.count += 1
+        self._fail_if_needed()
+        return x
+
+    def append_to(self, lst):
+        lst.append(self.i)
+        return lst
+
+    def inc_two(self, x, y):
+        self.i += x
+        self.i += y
+        return self.i
+
+    def sleep(self, x):
+        time.sleep(x)
+        return x
+
+    @ray.method(num_returns=2)
+    def return_two(self, x):
+        return x, x + 1
+
+
+def test_readers_on_different_nodes(ray_start_cluster):
+    cluster = ray_start_cluster
+    # This node is for the driver (including the CompiledDAG.DAGDriverProxyActor) and
+    # one of the readers.
+    cluster.add_node(num_cpus=1)
+    ray.init(address=cluster.address)
+    # 2 more nodes for other readers.
+    cluster.add_node(num_cpus=1)
+    cluster.add_node(num_cpus=1)
+    cluster.wait_for_nodes()
+    # Wait until nodes actually start, otherwise the code below will fail.
+    wait_for_condition(lambda: len(ray.nodes()) == 3)
+
+    a = Actor.options(num_cpus=1).remote(0)
+    b = Actor.options(num_cpus=1).remote(0)
+    c = Actor.options(num_cpus=1).remote(0)
+    actors = [a, b, c]
+
+    def _get_node_id(self) -> "ray.NodeID":
+        return ray.get_runtime_context().get_node_id()
+
+    node_ids = ray.get([act.__ray_call__.remote(_get_node_id) for act in actors])
+    assert len(set(node_ids)) == 3
+
+    with InputNode() as inp:
+        x = a.inc.bind(inp)
+        y = b.inc.bind(inp)
+        z = c.inc.bind(inp)
+        dag = MultiOutputNode([x, y, z])
+
+    cdag = dag.experimental_compile()
+
+    for i in range(1, 10):
+        assert ray.get(cdag.execute(1)) == [i, i, i]
+
+
+def test_bunch_readers_on_different_nodes(ray_start_cluster):
+    cluster = ray_start_cluster
+    ACTORS_PER_NODE = 2
+    NUM_REMOTE_NODES = 2
+    # driver node
+    cluster.add_node(num_cpus=ACTORS_PER_NODE)
+    ray.init(address=cluster.address)
+    # additional nodes for multi readers in multi nodes
+    for _ in range(NUM_REMOTE_NODES):
+        cluster.add_node(num_cpus=ACTORS_PER_NODE)
+    cluster.wait_for_nodes()
+
+    wait_for_condition(lambda: len(ray.nodes()) == NUM_REMOTE_NODES + 1)
+
+    actors = [
+        Actor.options(num_cpus=1).remote(0)
+        for _ in range(ACTORS_PER_NODE * (NUM_REMOTE_NODES + 1))
+    ]
+
+    def _get_node_id(self) -> "ray.NodeID":
+        return ray.get_runtime_context().get_node_id()
+
+    node_ids = ray.get([act.__ray_call__.remote(_get_node_id) for act in actors])
+    assert len(set(node_ids)) == NUM_REMOTE_NODES + 1
+
+    with InputNode() as inp:
+        outputs = []
+        for actor in actors:
+            outputs.append(actor.inc.bind(inp))
+        dag = MultiOutputNode(outputs)
+
+    cdag = dag.experimental_compile()
+
+    for i in range(1, 10):
+        assert ray.get(cdag.execute(1)) == [
+            i for _ in range(ACTORS_PER_NODE * (NUM_REMOTE_NODES + 1))
+        ]
+
+
+@pytest.mark.parametrize("single_fetch", [True, False])
+def test_pp(ray_start_cluster, single_fetch):
+    cluster = ray_start_cluster
+    # This node is for the driver.
+    cluster.add_node(num_cpus=0)
+    ray.init(address=cluster.address)
+    TP = 2
+    # This node is for the PP stage 1.
+    cluster.add_node(resources={"pp1": TP})
+    # This node is for the PP stage 2.
+    cluster.add_node(resources={"pp2": TP})
+
+    @ray.remote
+    class Worker:
+        def __init__(self):
+            pass
+
+        def execute_model(self, val):
+            return val
+
+    pp1_workers = [
+        Worker.options(num_cpus=0, resources={"pp1": 1}).remote() for _ in range(TP)
+    ]
+    pp2_workers = [
+        Worker.options(num_cpus=0, resources={"pp2": 1}).remote() for _ in range(TP)
+    ]
+
+    with InputNode() as inp:
+        outputs = [inp for _ in range(TP)]
+        outputs = [pp1_workers[i].execute_model.bind(outputs[i]) for i in range(TP)]
+        outputs = [pp2_workers[i].execute_model.bind(outputs[i]) for i in range(TP)]
+        dag = MultiOutputNode(outputs)
+
+    compiled_dag = dag.experimental_compile()
+    refs = compiled_dag.execute(1)
+    if single_fetch:
+        for i in range(TP):
+            assert ray.get(refs[i]) == 1
+    else:
+        assert ray.get(refs) == [1] * TP
+
+    # So that raylets' error messages are printed to the driver
+    time.sleep(2)
+
+
+@pytest.mark.parametrize("single_fetch", [True, False])
+def test_pp_exception(ray_start_cluster, single_fetch):
+    """
+    This test is to verify that the exception can be passed properly
+    through pipeline parallel workers on different nodes.
+    """
+    cluster = ray_start_cluster
+    # This node is for the driver.
+    cluster.add_node(num_cpus=0)
+    ray.init(address=cluster.address)
+    TP = 2
+    # This node is for the PP stage 1.
+    cluster.add_node(resources={"pp1": TP})
+    # This node is for the PP stage 2.
+    cluster.add_node(resources={"pp2": TP})
+    # This node is for the PP stage 3.
+    cluster.add_node(resources={"pp3": TP})
+
+    # Simulate a large error message (e.g., those with a long stack trace)
+    large_error_message = "Model execution failed" * 10000
+
+    @ray.remote
+    class Worker:
+        def __init__(self):
+            pass
+
+        def execute_model(self, val):
+            if val == "exception_trigger":
+                # Simulate an exception happened during model execution
+                raise RuntimeError(large_error_message)
+            return val
+
+    pp1_workers = [
+        Worker.options(num_cpus=0, resources={"pp1": 1}).remote() for _ in range(TP)
+    ]
+    pp2_workers = [
+        Worker.options(num_cpus=0, resources={"pp2": 1}).remote() for _ in range(TP)
+    ]
+    pp3_workers = [
+        Worker.options(num_cpus=0, resources={"pp3": 1}).remote() for _ in range(TP)
+    ]
+
+    with InputNode() as inp:
+        outputs = [inp for _ in range(TP)]
+        outputs = [pp1_workers[i].execute_model.bind(outputs[i]) for i in range(TP)]
+        outputs = [pp2_workers[i].execute_model.bind(outputs[i]) for i in range(TP)]
+        outputs = [pp3_workers[i].execute_model.bind(outputs[i]) for i in range(TP)]
+        dag = MultiOutputNode(outputs)
+
+    compiled_dag = dag.experimental_compile()
+    refs = compiled_dag.execute("exception_trigger")
+
+    # Without the fix in this PR, we will encounter the following exception:
+    # File "/Users/ruiqiao/repos2/ray/python/ray/_private/serialization.py",
+    # line 460, in deserialize_objects
+    #     obj = self._deserialize_object(data, metadata, object_ref)
+    #     raise Exception(
+    #     Exception: Can't deserialize object:
+    #       ObjectRef(00a33d534c5b0ce51bdf175790467da3114801680100000002e1f505), metadata: b'\x00'
+    # With this fix, the original exception will be propagated.
+    if single_fetch:
+        for i in range(TP):
+            with pytest.raises(RuntimeError) as exc_info:
+                ray.get(refs[i])
+            assert "Can't deserialize object" not in str(exc_info.value)
+            assert large_error_message in str(exc_info.value)
+    else:
+        with pytest.raises(RuntimeError) as exc_info:
+            ray.get(refs)
+        assert "Can't deserialize object" not in str(exc_info.value)
+        assert large_error_message in str(exc_info.value)
+
+
+def test_payload_large(ray_start_cluster, monkeypatch):
+    GRPC_MAX_SIZE = 1024 * 1024 * 5
+    monkeypatch.setenv("RAY_max_grpc_message_size", str(GRPC_MAX_SIZE))
+    cluster = ray_start_cluster
+    # This node is for the driver (including the CompiledDAG.DAGDriverProxyActor).
+    first_node_handle = cluster.add_node(num_cpus=1)
+    # This node is for the reader.
+    second_node_handle = cluster.add_node(num_cpus=1)
+    ray.init(address=cluster.address)
+    cluster.wait_for_nodes()
+
+    nodes = [first_node_handle.node_id, second_node_handle.node_id]
+    # We want to check that there are two nodes. Thus, we convert `nodes` to a set and
+    # then back to a list to remove duplicates. Then we check that the length of `nodes`
+    # is 2.
+    nodes = list(set(nodes))
+    assert len(nodes) == 2
+
+    def create_actor(node):
+        return Actor.options(label_selector={ray._raylet.RAY_NODE_ID_KEY: node}).remote(
+            0
+        )
+
+    def get_node_id(self):
+        return ray.get_runtime_context().get_node_id()
+
+    driver_node = get_node_id(None)
+    nodes.remove(driver_node)
+
+    a = create_actor(nodes[0])
+    a_node = ray.get(a.__ray_call__.remote(get_node_id))
+    assert a_node == nodes[0]
+    # Check that the driver and actor are on different nodes.
+    assert driver_node != a_node
+
+    with InputNode() as i:
+        dag = a.echo.bind(i)
+
+    compiled_dag = dag.experimental_compile()
+
+    size = GRPC_MAX_SIZE + (1024 * 1024 * 2)
+    val = b"x" * size
+
+    for i in range(3):
+        ref = compiled_dag.execute(val)
+        result = ray.get(ref)
+        assert result == val
+
+
+@pytest.mark.parametrize("num_actors", [1, 4])
+@pytest.mark.parametrize("num_nodes", [1, 4])
+def test_multi_node_multi_reader_large_payload(
+    ray_start_cluster, num_actors, num_nodes, monkeypatch
+):
+    # Set max grpc size to 5mb.
+    GRPC_MAX_SIZE = 1024 * 1024 * 5
+    monkeypatch.setenv("RAY_max_grpc_message_size", str(GRPC_MAX_SIZE))
+    cluster = ray_start_cluster
+    ACTORS_PER_NODE = num_actors
+    NUM_REMOTE_NODES = num_nodes
+    cluster.add_node(num_cpus=ACTORS_PER_NODE)
+    ray.init(address=cluster.address)
+    # This node is for the other two readers.
+    for _ in range(NUM_REMOTE_NODES):
+        cluster.add_node(num_cpus=ACTORS_PER_NODE)
+    cluster.wait_for_nodes()
+
+    wait_for_condition(lambda: len(ray.nodes()) == NUM_REMOTE_NODES + 1)
+
+    actors = [
+        Actor.options(num_cpus=1).remote(0)
+        for _ in range(ACTORS_PER_NODE * (NUM_REMOTE_NODES + 1))
+    ]
+
+    def _get_node_id(self) -> "ray.NodeID":
+        return ray.get_runtime_context().get_node_id()
+
+    node_ids = ray.get([act.__ray_call__.remote(_get_node_id) for act in actors])
+    assert len(set(node_ids)) == NUM_REMOTE_NODES + 1
+
+    with InputNode() as inp:
+        outputs = []
+        for actor in actors:
+            outputs.append(actor.echo.bind(inp))
+        dag = MultiOutputNode(outputs)
+
+    compiled_dag = dag.experimental_compile()
+
+    # Set the object size a little bigger than the gRPC size so that
+    # it triggers chunking and resizing.
+    size = GRPC_MAX_SIZE + (1024 * 1024 * 2)
+    val = b"x" * size
+
+    for _ in range(3):
+        ref = compiled_dag.execute(val)
+        # In the CI environment, the object store may use /tmp instead of /dev/shm
+        # due to limited size of /tmp/shm and therefore has degraded performance.
+        # Therefore, we use a longer timeout to avoid flakiness.
+        result = ray.get(ref, timeout=50)
+        assert result == [val for _ in range(ACTORS_PER_NODE * (NUM_REMOTE_NODES + 1))]
+
+
+def test_multi_node_dag_from_actor(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=1)
+    ray.init()
+    cluster.add_node(num_cpus=1)
+
+    @ray.remote(num_cpus=0)
+    class SameNodeActor:
+        def predict(self, x: str):
+            return x
+
+    @ray.remote(num_cpus=1)
+    class RemoteNodeActor:
+        def predict(self, x: str, y: str):
+            return y
+
+    @ray.remote(num_cpus=1)
+    class DriverActor:
+        def __init__(self):
+            self._base_actor = SameNodeActor.options(
+                label_selector={
+                    ray._raylet.RAY_NODE_ID_KEY: ray.get_runtime_context().get_node_id()
+                }
+            ).remote()
+            self._refiner_actor = RemoteNodeActor.remote()
+
+            with InputNode() as inp:
+                x = self._base_actor.predict.bind(inp)
+                dag = self._refiner_actor.predict.bind(
+                    inp,
+                    x,
+                )
+
+            self._cdag = dag.experimental_compile(
+                _submit_timeout=120,
+            )
+
+        def call(self, prompt: str) -> bytes:
+            return ray.get(self._cdag.execute(prompt))
+
+    parallel = DriverActor.remote()
+    assert ray.get(parallel.call.remote("abc")) == "abc"
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/dag/tests/experimental/test_torch_tensor_transport.py

@@ -0,0 +1,627 @@
+import os
+import sys
+from typing import Any, Dict
+
+import pytest
+import torch
+
+import ray
+from ray.dag import InputNode
+from ray.exceptions import RaySystemError, RayTaskError
+from ray.tests.conftest import *  # noqa
+
+if sys.platform != "linux" and sys.platform != "darwin":
+    pytest.skip("Skipping, requires Linux or Mac.", allow_module_level=True)
+
+USE_GPU = os.environ.get("RAY_PYTEST_USE_GPU") == "1"
+
+
+@ray.remote
+class Actor:
+    def echo_device(self, tensor: torch.Tensor) -> str:
+        if isinstance(tensor, RaySystemError):
+            raise tensor
+        return str(tensor.device)
+
+    def echo_dict_device(
+        self, dict_of_tensors: Dict[str, torch.Tensor]
+    ) -> Dict[str, str]:
+        if isinstance(dict_of_tensors, RaySystemError):
+            raise dict_of_tensors
+        return {k: str(v.device) for k, v in dict_of_tensors.items()}
+
+    def send(self, device: str) -> torch.Tensor:
+        return torch.ones((100,), device=device)
+
+    def send_dict(self, name_device_pairs: Dict[str, str]) -> Dict[str, torch.Tensor]:
+        tensor_dict = {}
+        for name, device in name_device_pairs.items():
+            tensor_dict[name] = torch.ones((100,), device=device)
+        return tensor_dict
+
+
+def run_driver_to_worker_dag(
+    actor: "ray.actor.ActorHandle",
+    device: str,
+    tensor_input: Any,
+    is_dict: bool = False,
+):
+    """Create and execute a DAG with tensor transport for driver to worker tests.
+
+    Args:
+        actor: Ray actor to use
+        device: Target device ("cpu", "cuda", or "default")
+        tensor_input: Input tensor(s) to execute with
+        is_dict: Whether to use dict version of the method
+
+    Returns:
+        ray.ObjectRef: Result reference
+    """
+    with InputNode() as inp:
+        method = actor.echo_dict_device if is_dict else actor.echo_device
+        dag = method.bind(inp.with_tensor_transport(device=device))
+    compiled_dag = dag.experimental_compile()
+    return compiled_dag.execute(tensor_input)
+
+
+def run_worker_to_worker_dag(
+    sender: "ray.actor.ActorHandle",
+    receiver: "ray.actor.ActorHandle",
+    device: str,
+    input_device: str,
+    is_dict: bool = False,
+):
+    """Create and execute a DAG with tensor transport for worker to worker tests.
+
+    Args:
+        sender: Sender Ray actor
+        receiver: Receiver Ray actor
+        device: Target device for tensor transport
+        input_device: Device string to pass to sender
+        is_dict: Whether to use dict version of the methods
+
+    Returns:
+        ray.ObjectRef: Result reference or ValueError for compilation errors
+    """
+    with InputNode() as inp:
+        if is_dict:
+            tensor = sender.send_dict.bind(inp)
+            dag = receiver.echo_dict_device.bind(
+                tensor.with_tensor_transport(device=device)
+            )
+        else:
+            tensor = sender.send.bind(inp)
+            dag = receiver.echo_device.bind(tensor.with_tensor_transport(device=device))
+    compiled_dag = dag.experimental_compile()
+    return compiled_dag.execute(input_device)
+
+
+def run_worker_to_driver_dag(
+    actor: "ray.actor.ActorHandle",
+    device: str,
+    input_device: str,
+    is_dict: bool = False,
+):
+    """Create and execute a DAG with tensor transport for worker to driver tests.
+
+    Args:
+        actor: Ray actor to use
+        device: Target device for tensor transport
+        input_device: Device string to pass to actor
+        is_dict: Whether to use dict version of the method
+
+    Returns:
+        ray.ObjectRef: Result reference
+    """
+    with InputNode() as inp:
+        if is_dict:
+            dag = actor.send_dict.bind(inp).with_tensor_transport(device=device)
+        else:
+            dag = actor.send.bind(inp).with_tensor_transport(device=device)
+    compiled_dag = dag.experimental_compile()
+    return compiled_dag.execute(input_device)
+
+
+class TestDriverToWorkerDeviceCPU:
+    """Tests driver to worker tensor transport with CPU device."""
+
+    def create_and_execute_dag(self, actor, device, tensor_input, is_dict=False):
+        """Create a DAG with tensor transport and execute it."""
+        with InputNode() as inp:
+            method = actor.echo_dict_device if is_dict else actor.echo_device
+            dag = method.bind(inp.with_tensor_transport(device=device))
+        compiled_dag = dag.experimental_compile()
+        return compiled_dag.execute(tensor_input)
+
+    def test_src_cpu_tensor_dst_cpu_node(self, ray_start_regular):
+        actor = Actor.remote()
+        ref = run_driver_to_worker_dag(actor, "cpu", torch.tensor([1]))
+        assert ray.get(ref) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_gpu_tensor_dst_cpu_node(self, ray_start_regular):
+        actor = Actor.remote()
+        ref = run_driver_to_worker_dag(actor, "cpu", torch.tensor([1], device="cuda"))
+        assert ray.get(ref) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_cpu_tensor_dst_gpu_node(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        ref = run_driver_to_worker_dag(actor, "cpu", torch.tensor([1]))
+        assert ray.get(ref) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_gpu_tensor_dst_gpu_node(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        ref = run_driver_to_worker_dag(actor, "cpu", torch.tensor([1], device="cuda"))
+        assert ray.get(ref) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_mix_tensors_dst_cpu_node(self, ray_start_regular):
+        actor = Actor.remote()
+        tensor_dict = {
+            "cpu_tensor": torch.tensor([1]),
+            "gpu_tensor": torch.tensor([1], device="cuda"),
+        }
+        ref = run_driver_to_worker_dag(actor, "cpu", tensor_dict, is_dict=True)
+        assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cpu"}
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_mix_tensors_dst_gpu_node(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        tensor_dict = {
+            "cpu_tensor": torch.tensor([1]),
+            "gpu_tensor": torch.tensor([1], device="cuda"),
+        }
+        ref = run_driver_to_worker_dag(actor, "cpu", tensor_dict, is_dict=True)
+        assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cpu"}
+
+
+class TestDriverToWorkerDeviceGPU:
+    """Tests driver to worker tensor transport with GPU device."""
+
+    def create_and_execute_dag(self, actor, device, tensor_input, is_dict=False):
+        """Create a DAG with tensor transport and execute it."""
+        with InputNode() as inp:
+            method = actor.echo_dict_device if is_dict else actor.echo_device
+            dag = method.bind(inp.with_tensor_transport(device=device))
+        compiled_dag = dag.experimental_compile()
+        return compiled_dag.execute(tensor_input)
+
+    def test_src_cpu_tensor_dst_cpu_node(self, ray_start_regular):
+        actor = Actor.remote()
+        ref = run_driver_to_worker_dag(actor, "cuda", torch.tensor([1]))
+        if torch.cuda.is_available():
+            assert ray.get(ref) == "cuda:0"
+        else:
+            with pytest.raises(
+                RayTaskError, match="RuntimeError: No CUDA GPUs are available"
+            ):
+                ray.get(ref)
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_gpu_tensor_dst_cpu_node(self, ray_start_regular):
+        actor = Actor.remote()
+        ref = run_driver_to_worker_dag(actor, "cuda", torch.tensor([1], device="cuda"))
+        assert ray.get(ref) == "cuda:0"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_cpu_tensor_dst_gpu_node(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        ref = run_driver_to_worker_dag(actor, "cuda", torch.tensor([1]))
+        assert ray.get(ref) == "cuda:0"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_gpu_tensor_dst_gpu_node(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        ref = run_driver_to_worker_dag(actor, "cuda", torch.tensor([1], device="cuda"))
+        assert ray.get(ref) == "cuda:0"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_mix_tensors_dst_cpu_node(self, ray_start_regular):
+        actor = Actor.remote()
+        tensor_dict = {
+            "cpu_tensor": torch.tensor([1]),
+            "gpu_tensor": torch.tensor([1], device="cuda"),
+        }
+        ref = run_driver_to_worker_dag(actor, "cuda", tensor_dict, is_dict=True)
+        assert ray.get(ref) == {"cpu_tensor": "cuda:0", "gpu_tensor": "cuda:0"}
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_mix_tensors_dst_gpu_node(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        tensor_dict = {
+            "cpu_tensor": torch.tensor([1]),
+            "gpu_tensor": torch.tensor([1], device="cuda"),
+        }
+        ref = run_driver_to_worker_dag(actor, "cuda", tensor_dict, is_dict=True)
+        assert ray.get(ref) == {"cpu_tensor": "cuda:0", "gpu_tensor": "cuda:0"}
+
+
+class TestDriverToWorkerDeviceDefault:
+    """Tests driver to worker tensor transport with default device."""
+
+    def create_and_execute_dag(self, actor, device, tensor_input, is_dict=False):
+        """Create a DAG with tensor transport and execute it."""
+        with InputNode() as inp:
+            method = actor.echo_dict_device if is_dict else actor.echo_device
+            dag = method.bind(inp.with_tensor_transport(device=device))
+        compiled_dag = dag.experimental_compile()
+        return compiled_dag.execute(tensor_input)
+
+    def test_src_cpu_tensor_dst_cpu_node(self, ray_start_regular):
+        actor = Actor.remote()
+        ref = run_driver_to_worker_dag(actor, "default", torch.tensor([1]))
+        assert ray.get(ref) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_gpu_tensor_dst_cpu_node(self, ray_start_regular):
+        actor = Actor.remote()
+        ref = run_driver_to_worker_dag(
+            actor, "default", torch.tensor([1], device="cuda")
+        )
+        assert ray.get(ref) == "cuda:0"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_cpu_tensor_dst_gpu_node(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        ref = run_driver_to_worker_dag(actor, "default", torch.tensor([1]))
+        assert ray.get(ref) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_gpu_tensor_dst_gpu_node(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        ref = run_driver_to_worker_dag(
+            actor, "default", torch.tensor([1], device="cuda")
+        )
+        assert ray.get(ref) == "cuda:0"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_mix_tensors_dst_cpu_node(self, ray_start_regular):
+        actor = Actor.remote()
+        tensor_dict = {
+            "cpu_tensor": torch.tensor([1]),
+            "gpu_tensor": torch.tensor([1], device="cuda"),
+        }
+        ref = run_driver_to_worker_dag(actor, "default", tensor_dict, is_dict=True)
+        assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cuda:0"}
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_mix_tensors_dst_gpu_node(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        tensor_dict = {
+            "cpu_tensor": torch.tensor([1]),
+            "gpu_tensor": torch.tensor([1], device="cuda"),
+        }
+        ref = run_driver_to_worker_dag(actor, "default", tensor_dict, is_dict=True)
+        assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cuda:0"}
+
+
+class TestWorkerToWorkerDeviceCPU:
+    """Tests worker to worker tensor transport with CPU device."""
+
+    def test_src_cpu_tensor_dst_cpu_node(self, ray_start_regular):
+        sender = Actor.remote()
+        receiver = Actor.remote()
+        ref = run_worker_to_worker_dag(sender, receiver, "cpu", "cpu")
+        assert ray.get(ref) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_cpu_tensor_dst_gpu_node(self, ray_start_regular):
+        sender = Actor.remote()
+        receiver = Actor.options(num_gpus=1).remote()
+        ref = run_worker_to_worker_dag(sender, receiver, "cpu", "cpu")
+        assert ray.get(ref) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_gpu_tensor_dst_cpu_node(self, ray_start_regular):
+        sender = Actor.options(num_gpus=1).remote()
+        receiver = Actor.remote()
+        ref = run_worker_to_worker_dag(sender, receiver, "cpu", "cpu")
+        assert ray.get(ref) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 2}], indirect=True)
+    def test_src_gpu_tensor_dst_gpu_node(self, ray_start_regular):
+        sender = Actor.options(num_gpus=1).remote()
+        receiver = Actor.options(num_gpus=1).remote()
+
+        with pytest.raises(
+            ValueError,
+            match="accelerator transport is not supported with CPU target device.",
+        ):
+            run_worker_to_worker_dag(sender, receiver, "cpu", "cpu")
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_mix_tensors_dst_cpu_node(self, ray_start_regular):
+        sender = Actor.options(num_gpus=1).remote()
+        receiver = Actor.options().remote()
+        ref = run_worker_to_worker_dag(
+            sender,
+            receiver,
+            "cpu",
+            {"cpu_tensor": "cpu", "gpu_tensor": "cuda"},
+            is_dict=True,
+        )
+        assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cpu"}
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 2}], indirect=True)
+    def test_src_mix_tensors_dst_gpu_node(self, ray_start_regular):
+        sender = Actor.options(num_gpus=1).remote()
+        receiver = Actor.options(num_gpus=1).remote()
+
+        with pytest.raises(
+            ValueError,
+            match="accelerator transport is not supported with CPU target device.",
+        ):
+            run_worker_to_worker_dag(
+                sender,
+                receiver,
+                "cpu",
+                {"cpu_tensor": "cpu", "gpu_tensor": "cuda"},
+                is_dict=True,
+            )
+
+
+class TestWorkerToWorkerDeviceGPU:
+    """Tests worker to worker tensor transport with GPU device."""
+
+    @pytest.mark.parametrize("gpu_device", ["gpu", "cuda"])
+    def test_src_cpu_tensor_dst_cpu_node(self, ray_start_regular, gpu_device):
+        sender = Actor.remote()
+        receiver = Actor.remote()
+        ref = run_worker_to_worker_dag(sender, receiver, gpu_device, "cpu")
+        if torch.cuda.is_available():
+            assert ray.get(ref) == "cuda:0"
+        else:
+            with pytest.raises(
+                RayTaskError, match="RuntimeError: No CUDA GPUs are available"
+            ):
+                ray.get(ref)
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_cpu_tensor_dst_gpu_node(self, ray_start_regular):
+        sender = Actor.remote()
+        receiver = Actor.options(num_gpus=1).remote()
+        ref = run_worker_to_worker_dag(sender, receiver, "cuda", "cpu")
+        assert ray.get(ref) == "cuda:0"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_gpu_tensor_dst_cpu_node(self, ray_start_regular):
+        sender = Actor.options(num_gpus=1).remote()
+        receiver = Actor.remote()
+        ref = run_worker_to_worker_dag(sender, receiver, "cuda", "cuda")
+        assert ray.get(ref) == "cuda:0"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 2}], indirect=True)
+    def test_src_gpu_tensor_dst_gpu_node(self, ray_start_regular):
+        sender = Actor.options(num_gpus=1).remote()
+        receiver = Actor.options(num_gpus=1).remote()
+
+        with pytest.raises(
+            ValueError,
+            match="accelerator transport is not supported with CPU target device.",
+        ):
+            run_worker_to_worker_dag(sender, receiver, "cpu", "cpu")
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_mix_tensors_dst_cpu_node(self, ray_start_regular):
+        sender = Actor.options(num_gpus=1).remote()
+        receiver = Actor.options().remote()
+        ref = run_worker_to_worker_dag(
+            sender,
+            receiver,
+            "cuda",
+            {"cpu_tensor": "cpu", "gpu_tensor": "cuda"},
+            is_dict=True,
+        )
+        assert ray.get(ref) == {"cpu_tensor": "cuda:0", "gpu_tensor": "cuda:0"}
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 2}], indirect=True)
+    @pytest.mark.parametrize("gpu_device", ["gpu", "cuda"])
+    def test_src_mix_tensors_dst_gpu_node(self, ray_start_regular, gpu_device):
+        sender = Actor.options(num_gpus=1).remote()
+        receiver = Actor.options(num_gpus=1).remote()
+
+        ref = run_worker_to_worker_dag(
+            sender,
+            receiver,
+            gpu_device,
+            {"cpu_tensor": "cpu", "gpu_tensor": "cuda"},
+            is_dict=True,
+        )
+        assert ray.get(ref) == {"cpu_tensor": "cuda:0", "gpu_tensor": "cuda:0"}
+
+
+class TestWorkerToWorkerDeviceDefault:
+    """Tests worker to worker tensor transport with default device."""
+
+    def test_src_cpu_tensor_dst_cpu_node(self, ray_start_regular):
+        sender = Actor.remote()
+        receiver = Actor.remote()
+        ref = run_worker_to_worker_dag(sender, receiver, "default", "cpu")
+        assert ray.get(ref) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_cpu_tensor_dst_gpu_node(self, ray_start_regular):
+        sender = Actor.remote()
+        receiver = Actor.options(num_gpus=1).remote()
+        ref = run_worker_to_worker_dag(sender, receiver, "default", "cpu")
+        assert ray.get(ref) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_gpu_tensor_dst_cpu_node(self, ray_start_regular):
+        sender = Actor.options(num_gpus=1).remote()
+        receiver = Actor.remote()
+        ref = run_worker_to_worker_dag(sender, receiver, "default", "cuda")
+        assert ray.get(ref) == "cuda:0"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 2}], indirect=True)
+    def test_src_gpu_tensor_dst_gpu_node(self, ray_start_regular):
+        sender = Actor.options(num_gpus=1).remote()
+        receiver = Actor.options(num_gpus=1).remote()
+
+        with pytest.raises(
+            ValueError,
+            match="accelerator transport is not supported with CPU target device.",
+        ):
+            run_worker_to_worker_dag(sender, receiver, "cpu", "cpu")
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_mix_tensors_dst_cpu_node(self, ray_start_regular):
+        sender = Actor.options(num_gpus=1).remote()
+        receiver = Actor.options().remote()
+        ref = run_worker_to_worker_dag(
+            sender,
+            receiver,
+            "default",
+            {"cpu_tensor": "cpu", "gpu_tensor": "cuda"},
+            is_dict=True,
+        )
+        assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cuda:0"}
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 2}], indirect=True)
+    def test_src_mix_tensors_dst_gpu_node(self, ray_start_regular):
+        sender = Actor.options(num_gpus=1).remote()
+        receiver = Actor.options(num_gpus=1).remote()
+        ref = run_worker_to_worker_dag(
+            sender,
+            receiver,
+            "default",
+            {"cpu_tensor": "cpu", "gpu_tensor": "cuda"},
+            is_dict=True,
+        )
+        assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cuda:0"}
+
+
+class TestWorkerToDriverDeviceCPU:
+    """Tests worker to driver tensor transport with CPU device."""
+
+    def test_src_cpu_tensor(self, ray_start_regular):
+        actor = Actor.remote()
+        ref = run_worker_to_driver_dag(actor, "cpu", "cpu")
+        tensor = ray.get(ref)
+        assert str(tensor.device) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_gpu_tensor(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        ref = run_worker_to_driver_dag(actor, "cpu", "cuda")
+        tensor = ray.get(ref)
+        assert str(tensor.device) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_mix_tensors(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        ref = run_worker_to_driver_dag(
+            actor, "cpu", {"cpu_tensor": "cpu", "gpu_tensor": "cuda"}, is_dict=True
+        )
+        tensor = ray.get(ref)
+        assert str(tensor["cpu_tensor"].device) == "cpu"
+        assert str(tensor["gpu_tensor"].device) == "cpu"
+
+
+class TestWorkerToDriverDeviceGPU:
+    """Tests worker to driver tensor transport with GPU device."""
+
+    def test_src_cpu_tensor(self, ray_start_regular):
+        actor = Actor.remote()
+        ref = run_worker_to_driver_dag(actor, "cuda", "cpu")
+
+        # different behavior between a driver node with GPU and without GPU
+        if torch.cuda.is_available():
+            tensor = ray.get(ref)
+            assert str(tensor.device) == "cuda:0"
+        else:
+            with pytest.raises(
+                RayTaskError, match="RuntimeError: No CUDA GPUs are available"
+            ):
+                ray.get(ref)
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_gpu_tensor(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        ref = run_worker_to_driver_dag(actor, "cuda", "cuda")
+
+        # different behavior between a driver node with GPU and without GPU
+        if torch.cuda.is_available():
+            tensor = ray.get(ref)
+            assert str(tensor.device) == "cuda:0"
+        else:
+            with pytest.raises(
+                RayTaskError, match="RuntimeError: No CUDA GPUs are available"
+            ):
+                ray.get(ref)
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_mix_tensors(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        ref = run_worker_to_driver_dag(
+            actor, "cuda", {"cpu_tensor": "cpu", "gpu_tensor": "cuda"}, is_dict=True
+        )
+
+        # different behavior between a driver node with GPU and without GPU
+        if torch.cuda.is_available():
+            tensor = ray.get(ref)
+            assert str(tensor["cpu_tensor"].device) == "cuda:0"
+            assert str(tensor["gpu_tensor"].device) == "cuda:0"
+        else:
+            with pytest.raises(
+                RayTaskError, match="RuntimeError: No CUDA GPUs are available"
+            ):
+                ray.get(ref)
+
+
+class TestWorkerToDriverDeviceDefault:
+    """Tests worker to driver tensor transport with default device."""
+
+    def test_src_cpu_tensor(self, ray_start_regular):
+        actor = Actor.remote()
+        ref = run_worker_to_driver_dag(actor, "default", "cpu")
+        tensor = ray.get(ref)
+        assert str(tensor.device) == "cpu"
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_gpu_tensor(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        ref = run_worker_to_driver_dag(actor, "default", "cuda")
+
+        # different behavior between a driver node with GPU and without GPU
+        if torch.cuda.is_available():
+            tensor = ray.get(ref)
+            assert str(tensor.device) == "cuda:0"
+        else:
+            with pytest.raises(
+                RayTaskError, match="RuntimeError: No CUDA GPUs are available"
+            ):
+                ray.get(ref)
+
+    @pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
+    def test_src_mix_tensors(self, ray_start_regular):
+        actor = Actor.options(num_gpus=1).remote()
+        ref = run_worker_to_driver_dag(
+            actor, "default", {"cpu_tensor": "cpu", "gpu_tensor": "cuda"}, is_dict=True
+        )
+
+        # different behavior between a driver node with GPU and without GPU
+        if torch.cuda.is_available():
+            tensor = ray.get(ref)
+            assert str(tensor["cpu_tensor"].device) == "cpu"
+            assert str(tensor["gpu_tensor"].device) == "cuda:0"
+        else:
+            with pytest.raises(
+                RayTaskError, match="RuntimeError: No CUDA GPUs are available"
+            ):
+                ray.get(ref)
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/dag/tests/test_class_dag.py

@@ -0,0 +1,362 @@
+import pytest
+
+import ray
+from ray.dag import (
+    PARENT_CLASS_NODE_KEY,
+    PREV_CLASS_METHOD_CALL_KEY,
+    InputNode,
+    MultiOutputNode,
+)
+
+
+@ray.remote
+class Counter:
+    def __init__(self, init_value=0):
+        self.i = init_value
+
+    def inc(self):
+        self.i += 1
+
+    def get(self):
+        return self.i
+
+
+@ray.remote
+class Actor:
+    def __init__(self, init_value):
+        self.i = init_value
+
+    def inc(self, x):
+        self.i += x
+
+    def get(self):
+        return self.i
+
+    def echo(self, x):
+        return x
+
+    @ray.method(num_returns=2)
+    def return_two(self, x):
+        return x, x + 1
+
+    @ray.method(num_returns=2)
+    def inc_and_return_two(self, x):
+        self.i += x
+        return self.i, self.i + 1
+
+    @ray.method(num_returns=1)
+    def return_two_as_one(self, x):
+        return x, x + 1
+
+    @ray.method(num_returns=2)
+    def return_two_from_three(self, x):
+        return x, x + 1, x + 2
+
+
+def test_basic_actor_dag(shared_ray_instance):
+    @ray.remote
+    def combine(x, y):
+        return x + y
+
+    a1 = Actor.bind(10)
+    res = a1.get.bind()
+    print(res)
+    assert ray.get(res.execute()) == 10
+
+    a2 = Actor.bind(10)
+    a1.inc.bind(2)
+    a1.inc.bind(4)
+    a2.inc.bind(6)
+    dag = combine.bind(a1.get.bind(), a2.get.bind())
+
+    print(dag)
+    assert ray.get(dag.execute()) == 32
+
+
+def test_class_as_class_constructor_arg(shared_ray_instance):
+    @ray.remote
+    class OuterActor:
+        def __init__(self, inner_actor):
+            self.inner_actor = inner_actor
+
+        def inc(self, x):
+            self.inner_actor.inc.remote(x)
+
+        def get(self):
+            return ray.get(self.inner_actor.get.remote())
+
+    outer = OuterActor.bind(Actor.bind(10))
+    outer.inc.bind(2)
+    dag = outer.get.bind()
+    print(dag)
+    assert ray.get(dag.execute()) == 12
+
+
+def test_class_as_function_constructor_arg(shared_ray_instance):
+    @ray.remote
+    def f(actor_handle):
+        return ray.get(actor_handle.get.remote())
+
+    dag = f.bind(Actor.bind(10))
+    print(dag)
+    assert ray.get(dag.execute()) == 10
+
+
+def test_basic_actor_dag_constructor_options(shared_ray_instance):
+    a1 = Actor.bind(10)
+    dag = a1.get.bind()
+    print(dag)
+    assert ray.get(dag.execute()) == 10
+
+    a1 = Actor.options(name="Actor", namespace="test", max_pending_calls=10).bind(10)
+    dag = a1.get.bind()
+    print(dag)
+    # Ensure execution result is identical with .options() in init()
+    assert ray.get(dag.execute()) == 10
+    # Ensure options are passed in
+    assert a1.get_options().get("name") == "Actor"
+    assert a1.get_options().get("namespace") == "test"
+    assert a1.get_options().get("max_pending_calls") == 10
+
+
+def test_actor_method_options(shared_ray_instance):
+    a1 = Actor.bind(10)
+    dag = a1.get.options(name="actor_method_options").bind()
+    print(dag)
+    assert ray.get(dag.execute()) == 10
+    assert dag.get_options().get("name") == "actor_method_options"
+
+
+def test_basic_actor_dag_constructor_invalid_options(shared_ray_instance):
+    with pytest.raises(
+        ValueError, match=r".*quantity of resource num_cpus cannot be negative.*"
+    ):
+        a1 = Actor.options(num_cpus=-1).bind(10)
+        invalid_dag = a1.get.bind()
+        ray.get(invalid_dag.execute())
+
+
+def test_actor_options_complicated(shared_ray_instance):
+    """Test a more complicated setup where we apply .options() in both
+    constructor and method call with overlapping keys, and ensure end to end
+    options correctness.
+    """
+
+    @ray.remote
+    def combine(x, y):
+        return x + y
+
+    a1 = Actor.options(name="a1_v0").bind(10)
+    res = a1.get.options(name="v1").bind()
+    print(res)
+    assert ray.get(res.execute()) == 10
+    assert a1.get_options().get("name") == "a1_v0"
+    assert res.get_options().get("name") == "v1"
+
+    a1 = Actor.options(name="a1_v1").bind(10)  # Cannot
+    a2 = Actor.options(name="a2_v0").bind(10)
+    a1.inc.options(name="v1").bind(2)
+    a1.inc.options(name="v2").bind(4)
+    a2.inc.options(name="v3").bind(6)
+    dag = combine.options(name="v4").bind(a1.get.bind(), a2.get.bind())
+
+    print(dag)
+    assert ray.get(dag.execute()) == 32
+    test_a1 = dag.get_args()[0]  # call graph for a1.get.bind()
+    test_a2 = dag.get_args()[1]  # call graph for a2.get.bind()
+    assert test_a2.get_options() == {}  # No .options() at outer call
+    # refer to a2 constructor .options() call
+    assert (
+        test_a2.get_other_args_to_resolve()[PARENT_CLASS_NODE_KEY]
+        .get_options()
+        .get("name")
+        == "a2_v0"
+    )
+    # refer to actor method a2.inc.options() call
+    assert (
+        test_a2.get_other_args_to_resolve()[PREV_CLASS_METHOD_CALL_KEY]
+        .get_options()
+        .get("name")
+        == "v3"
+    )
+    # refer to a1 constructor .options() call
+    assert (
+        test_a1.get_other_args_to_resolve()[PARENT_CLASS_NODE_KEY]
+        .get_options()
+        .get("name")
+        == "a1_v1"
+    )
+    # refer to latest actor method a1.inc.options() call
+    assert (
+        test_a1.get_other_args_to_resolve()[PREV_CLASS_METHOD_CALL_KEY]
+        .get_options()
+        .get("name")
+        == "v2"
+    )
+    # refer to first bound actor method a1.inc.options() call
+    assert (
+        test_a1.get_other_args_to_resolve()[PREV_CLASS_METHOD_CALL_KEY]
+        .get_other_args_to_resolve()[PREV_CLASS_METHOD_CALL_KEY]
+        .get_options()
+        .get("name")
+        == "v1"
+    )
+
+
+def test_pass_actor_handle(shared_ray_instance):
+    @ray.remote
+    class Actor:
+        def ping(self):
+            return "hello"
+
+    @ray.remote
+    def caller(handle):
+        assert isinstance(handle, ray.actor.ActorHandle), handle
+        return ray.get(handle.ping.remote())
+
+    a1 = Actor.bind()
+    dag = caller.bind(a1)
+    print(dag)
+    assert ray.get(dag.execute()) == "hello"
+
+
+def test_dynamic_pipeline(shared_ray_instance):
+    @ray.remote
+    class Model:
+        def __init__(self, arg):
+            self.arg = arg
+
+        def forward(self, x):
+            return self.arg + str(x)
+
+    @ray.remote
+    class ModelSelection:
+        def is_even(self, x):
+            return x % 2 == 0
+
+    @ray.remote
+    def pipeline(x, m1, m2, selection):
+        sel = selection.is_even.remote(x)
+        if ray.get(sel):
+            result = m1.forward.remote(x)
+        else:
+            result = m2.forward.remote(x)
+        return ray.get(result)
+
+    m1 = Model.bind("Even: ")
+    m2 = Model.bind("Odd: ")
+    selection = ModelSelection.bind()
+
+    even_input = pipeline.bind(20, m1, m2, selection)
+    print(even_input)
+    assert ray.get(even_input.execute()) == "Even: 20"
+
+    odd_input = pipeline.bind(21, m1, m2, selection)
+    print(odd_input)
+    assert ray.get(odd_input.execute()) == "Odd: 21"
+
+
+def test_unsupported_bind():
+    @ray.remote
+    class Actor:
+        def ping(self):
+            return "hello"
+
+    with pytest.raises(
+        AttributeError,
+        match=r"\.bind\(\) cannot be used again on",
+    ):
+        actor = Actor.bind()
+        _ = actor.bind()
+
+    with pytest.raises(
+        AttributeError,
+        match=r"\.remote\(\) cannot be used on ClassMethodNodes",
+    ):
+        actor = Actor.bind()
+        _ = actor.ping.remote()
+
+
+def test_unsupported_remote():
+    @ray.remote
+    class Actor:
+        def ping(self):
+            return "hello"
+
+    with pytest.raises(AttributeError, match="'Actor' has no attribute 'remote'"):
+        _ = Actor.bind().remote()
+
+    @ray.remote
+    def func():
+        return 1
+
+    with pytest.raises(AttributeError, match=r"\.remote\(\) cannot be used on"):
+        _ = func.bind().remote()
+
+
+def test_two_returns_first():
+    a = Actor.remote(0)
+    with InputNode() as i:
+        o1, o2 = a.return_two.bind(i)
+        dag = o1
+
+    for _ in range(3):
+        res = ray.get(dag.execute(1))
+        assert res == 1
+
+
+def test_two_returns_second():
+    a = Actor.remote(0)
+    with InputNode() as i:
+        o1, o2 = a.return_two.bind(i)
+        dag = o2
+
+    for _ in range(3):
+        res = ray.get(dag.execute(1))
+        assert res == 2
+
+
+def test_two_returns_one_reader_multi_times():
+    a = Actor.remote(0)
+    b = Actor.remote(0)
+    with InputNode() as i:
+        o1, o2 = a.return_two.bind(i)
+        o3 = b.echo.bind(o1)
+        o4 = b.echo.bind(o2)
+        dag = MultiOutputNode([o3, o4])
+
+    for _ in range(3):
+        res = ray.get(dag.execute(1))
+        assert res == [1, 2]
+
+
+def test_two_returns_two_readers():
+    a = Actor.remote(0)
+    b = Actor.remote(0)
+    c = Actor.remote(0)
+    with InputNode() as i:
+        o1, o2 = a.return_two.bind(i)
+        o3 = b.echo.bind(o1)
+        o4 = c.echo.bind(o2)
+        dag = MultiOutputNode([o3, o4])
+
+    for _ in range(3):
+        res = ray.get(dag.execute(1))
+        assert res == [1, 2]
+
+
+def test_inc_two_returns():
+    a = Actor.remote(0)
+    with InputNode() as i:
+        o1, o2 = a.inc_and_return_two.bind(i)
+        dag = MultiOutputNode([o1, o2])
+
+    for i in range(3):
+        res = ray.get(dag.execute(1))
+        assert res == [i + 1, i + 2]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/dag/tests/test_function_dag.py

@@ -0,0 +1,215 @@
+import pytest
+
+import ray
+
+
+@ray.remote
+class Counter:
+    def __init__(self, init_value=0):
+        self.i = init_value
+
+    def inc(self):
+        self.i += 1
+
+    def get(self):
+        return self.i
+
+
+def test_basic_task_dag(shared_ray_instance):
+    ct = Counter.remote()
+
+    @ray.remote
+    def a():
+        ray.get(ct.inc.remote())
+        return 2
+
+    @ray.remote
+    def b(x):
+        ray.get(ct.inc.remote())
+        return x * 2
+
+    @ray.remote
+    def c(x):
+        ray.get(ct.inc.remote())
+        return x + 1
+
+    @ray.remote
+    def d(x, y):
+        ray.get(ct.inc.remote())
+        return x + y
+
+    a_ref = a.bind()
+    b_ref = b.bind(a_ref)
+    c_ref = c.bind(a_ref)
+    d_ref = d.bind(b_ref, c_ref)
+    d1_ref = d.bind(d_ref, d_ref)
+    d2_ref = d.bind(d1_ref, d_ref)
+    dag = d.bind(d2_ref, d_ref)
+    print(dag)
+
+    assert ray.get(dag.execute()) == 28
+    assert ray.get(ct.get.remote()) == 7
+
+
+def test_basic_task_dag_with_options(shared_ray_instance):
+    ct = Counter.remote()
+
+    @ray.remote
+    def a():
+        ray.get(ct.inc.remote())
+        return 2
+
+    @ray.remote
+    def b(x):
+        ray.get(ct.inc.remote())
+        return x * 2
+
+    @ray.remote
+    def c(x):
+        ray.get(ct.inc.remote())
+        return x + 1
+
+    @ray.remote
+    def d(x, y):
+        ray.get(ct.inc.remote())
+        return x + y
+
+    a_ref = a.bind()
+    b_ref = b.options(name="b", num_returns=1).bind(a_ref)
+    c_ref = c.options(name="c", max_retries=3).bind(a_ref)
+    dag = d.options(name="d", num_cpus=2).bind(b_ref, c_ref)
+
+    print(dag)
+
+    assert ray.get(dag.execute()) == 7
+    assert ray.get(ct.get.remote()) == 4
+
+    assert b_ref.get_options().get("name") == "b"
+    assert b_ref.get_options().get("num_returns") == 1
+    assert c_ref.get_options().get("name") == "c"
+    assert c_ref.get_options().get("max_retries") == 3
+    assert dag.get_options().get("name") == "d"
+    assert dag.get_options().get("num_cpus") == 2
+
+
+def test_invalid_task_options(shared_ray_instance):
+    """
+    Test to ensure options used in DAG binding are applied, and will throw
+    as expected even given invalid values.
+    """
+
+    @ray.remote
+    def a():
+        return 2
+
+    @ray.remote
+    def b(x):
+        return x * 2
+
+    a_ref = a.bind()
+    dag = b.bind(a_ref)
+
+    # Ensure current DAG is executable
+    assert ray.get(dag.execute()) == 4
+    with pytest.raises(
+        ValueError, match=r".*quantity of resource num_cpus cannot be negative.*"
+    ):
+        invalid_dag = b.options(num_cpus=-1).bind(a_ref)
+        ray.get(invalid_dag.execute())
+
+
+def test_node_accessors(shared_ray_instance):
+    @ray.remote
+    def a(*a, **kw):
+        pass
+
+    tmp1 = a.bind()
+    tmp2 = a.bind()
+    tmp3 = a.bind()
+    node = a.bind(1, tmp1, x=tmp2, y={"foo": tmp3})
+    assert node.get_args() == (1, tmp1)
+    assert node.get_kwargs() == {"x": tmp2, "y": {"foo": tmp3}}
+    assert node._get_toplevel_child_nodes() == [tmp1, tmp2]
+    assert node._get_all_child_nodes() == [tmp1, tmp2, tmp3]
+
+    tmp4 = a.bind()
+    tmp5 = a.bind()
+    replace = {tmp1: tmp4, tmp2: tmp4, tmp3: tmp5}
+    n2 = node._apply_and_replace_all_child_nodes(lambda x: replace[x])
+    assert n2._get_all_child_nodes() == [tmp4, tmp5]
+
+
+def test_nested_args(shared_ray_instance):
+    ct = Counter.remote()
+
+    @ray.remote
+    def a():
+        ray.get(ct.inc.remote())
+        return 2
+
+    @ray.remote
+    def b(**kwargs):
+        ray.get(ct.inc.remote())
+        return kwargs["x"] * 2
+
+    @ray.remote
+    def c(**kwargs):
+        ray.get(ct.inc.remote())
+        return kwargs["x"] + 1
+
+    @ray.remote
+    def d(nested):
+        ray.get(ct.inc.remote())
+        return ray.get(nested["x"]) + ray.get(nested["y"])
+
+    a_ref = a.bind()
+    b_ref = b.bind(x=a_ref)
+    c_ref = c.bind(x=a_ref)
+    dag = d.bind({"x": b_ref, "y": c_ref})
+    print(dag)
+
+    assert ray.get(dag.execute()) == 7
+    assert ray.get(ct.get.remote()) == 4
+
+
+def test_dag_options(shared_ray_instance):
+    @ray.remote(num_gpus=100)
+    def foo():
+        pass
+
+    assert foo.bind().get_options() == {"max_calls": 1, "num_gpus": 100}
+    assert foo.options(num_gpus=300).bind().get_options() == {"num_gpus": 300}
+    assert foo.options(num_cpus=500).bind().get_options() == {
+        "num_gpus": 100,
+        "num_cpus": 500,
+    }
+
+    @ray.remote
+    def bar():
+        pass
+
+    assert bar.bind().get_options() == {}
+    assert bar.options(num_gpus=100).bind().get_options() == {"num_gpus": 100}
+
+    @ray.remote(num_gpus=100)
+    class Foo:
+        pass
+
+    assert Foo.bind().get_options() == {"num_gpus": 100}
+    assert Foo.options(num_gpus=300).bind().get_options() == {"num_gpus": 300}
+    assert Foo.options(num_cpus=500).bind().get_options() == {
+        "num_gpus": 100,
+        "num_cpus": 500,
+    }
+
+    @ray.remote
+    class Bar:
+        pass
+
+    assert Bar.bind().get_options() == {}
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/dag/tests/test_input_node.py

@@ -0,0 +1,386 @@
+"""
+Tests to ensure ray DAG can correctly mark its input(s) to take user
+request, for all DAGNode types.
+"""
+
+from typing import Any, TypeVar
+
+import pytest
+
+import ray
+from ray.dag.dag_node import DAGNode
+from ray.dag.input_node import InputNode
+
+RayHandleLike = TypeVar("RayHandleLike")
+
+
+def test_no_args_to_input_node(shared_ray_instance):
+    @ray.remote
+    def f(input):
+        return input
+
+    with pytest.raises(
+        ValueError, match="InputNode should not take any args or kwargs"
+    ):
+        with InputNode(0) as dag_input:
+            f.bind(dag_input)
+    with pytest.raises(
+        ValueError,
+        match="InputNode should not take any args or kwargs",
+    ):
+        with InputNode(key=1) as dag_input:
+            f.bind(dag_input)
+
+
+def test_simple_func(shared_ray_instance):
+    @ray.remote
+    def a(input: str):
+        return f"{input} -> a"
+
+    @ray.remote
+    def b(a: "RayHandleLike"):
+        # At runtime, a is replaced with execution result of a.
+        return f"{a} -> b"
+
+    # input -> a - > b -> ouput
+    with InputNode() as dag_input:
+        a_node = a.bind(dag_input)
+        dag = b.bind(a_node)
+
+    assert ray.get(dag.execute("input")) == "input -> a -> b"
+    assert ray.get(dag.execute("test")) == "test -> a -> b"
+
+
+def test_func_dag(shared_ray_instance):
+    @ray.remote
+    def a(user_input):
+        return user_input
+
+    @ray.remote
+    def b(x):
+        return x * 2
+
+    @ray.remote
+    def c(x):
+        return x + 1
+
+    @ray.remote
+    def d(x, y):
+        return x + y
+
+    with InputNode() as dag_input:
+        a_ref = a.bind(dag_input)
+        b_ref = b.bind(a_ref)
+        c_ref = c.bind(a_ref)
+        d_ref = d.bind(b_ref, c_ref)
+        d1_ref = d.bind(d_ref, d_ref)
+        d2_ref = d.bind(d1_ref, d_ref)
+        dag = d.bind(d2_ref, d_ref)
+
+    # [(2*2 + 2+1) + (2*2 + 2+1)] + [(2*2 + 2+1) + (2*2 + 2+1)]
+    assert ray.get(dag.execute(2)) == 28
+    # [(3*2 + 3+1) + (3*2 + 3+1)] + [(3*2 + 3+1) + (3*2 + 3+1)]
+    assert ray.get(dag.execute(3)) == 40
+
+
+def test_multi_input_func_dag(shared_ray_instance):
+    @ray.remote
+    def a(user_input):
+        return user_input * 2
+
+    @ray.remote
+    def b(user_input):
+        return user_input + 1
+
+    @ray.remote
+    def c(x, y):
+        return x + y
+
+    with InputNode() as dag_input:
+        a_ref = a.bind(dag_input)
+        b_ref = b.bind(dag_input)
+        dag = c.bind(a_ref, b_ref)
+
+    # (2*2) + (2*1)
+    assert ray.get(dag.execute(2)) == 7
+    # (3*2) + (3*1)
+    assert ray.get(dag.execute(3)) == 10
+
+
+def test_invalid_input_node_as_class_constructor(shared_ray_instance):
+    @ray.remote
+    class Actor:
+        def __init__(self, val):
+            self.val = val
+
+        def get(self):
+            return self.val
+
+    with pytest.raises(
+        ValueError,
+        match=(
+            "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."
+        ),
+    ):
+        with InputNode() as dag_input:
+            Actor.bind(dag_input)
+
+
+def test_class_method_input(shared_ray_instance):
+    @ray.remote
+    class Model:
+        def __init__(self, weight: int):
+            self.weight = weight
+
+        def forward(self, input: "RayHandleLike"):
+            return self.weight * input
+
+    @ray.remote
+    class FeatureProcessor:
+        def __init__(self, scale):
+            self.scale = scale
+
+        def process(self, input: int):
+            return input * self.scale
+
+    with InputNode() as dag_input:
+        preprocess = FeatureProcessor.bind(0.5)
+        feature = preprocess.process.bind(dag_input)
+        model = Model.bind(4)
+        dag = model.forward.bind(feature)
+
+    # 2 * 0.5 * 4
+    assert ray.get(dag.execute(2)) == 4
+    # 6 * 0.5 * 4
+    assert ray.get(dag.execute(6)) == 12
+
+
+def test_multi_class_method_input(shared_ray_instance):
+    """
+    Test a multiple class methods can all be used as inputs in a dag.
+    """
+
+    @ray.remote
+    class Model:
+        def __init__(self, weight: int):
+            self.weight = weight
+
+        def forward(self, input: int):
+            return self.weight * input
+
+    @ray.remote
+    def combine(m1: "RayHandleLike", m2: "RayHandleLike"):
+        return m1 + m2
+
+    with InputNode() as dag_input:
+        m1 = Model.bind(2)
+        m2 = Model.bind(3)
+
+        m1_output = m1.forward.bind(dag_input)
+        m2_output = m2.forward.bind(dag_input)
+
+        dag = combine.bind(m1_output, m2_output)
+
+    # 1*2 + 1*3
+    assert ray.get(dag.execute(1)) == 5
+    # 2*2 + 2*3
+    assert ray.get(dag.execute(2)) == 10
+
+
+def test_func_class_mixed_input(shared_ray_instance):
+    """
+    Test both class method and function are used as input in the
+    same dag.
+    """
+
+    @ray.remote
+    class Model:
+        def __init__(self, weight: int):
+            self.weight = weight
+
+        def forward(self, input: int):
+            return self.weight * input
+
+    @ray.remote
+    def model_func(input: int):
+        return input * 2
+
+    @ray.remote
+    def combine(m1: "RayHandleLike", m2: "RayHandleLike"):
+        return m1 + m2
+
+    with InputNode() as dag_input:
+        m1 = Model.bind(3)
+        m1_output = m1.forward.bind(dag_input)
+        m2_output = model_func.bind(dag_input)
+
+    dag = combine.bind(m1_output, m2_output)
+    # 2*3 + 2*2
+    assert ray.get(dag.execute(2)) == 10
+    # 3*3 + 3*2
+    assert ray.get(dag.execute(3)) == 15
+
+
+def test_input_attr_partial_access(shared_ray_instance):
+    @ray.remote
+    class Model:
+        def __init__(self, weight: int):
+            self.weight = weight
+
+        def forward(self, input: int):
+            return self.weight * input
+
+    @ray.remote
+    def combine(a, b, c, d=None):
+        if not d:
+            return a + b + c
+        else:
+            return a + b + c + d["deep"]["nested"]
+
+    # 1) Test default wrapping of args and kwargs into internal python object
+    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[1])
+        dag = combine.bind(m1_output, m2_output, dag_input.m3, dag_input.m4)
+    # 1*1 + 2*2 + 3 + 4 = 12
+    assert ray.get(dag.execute(1, 2, m3=3, m4={"deep": {"nested": 4}})) == 12
+
+    # 2) Test user passed data object as only input to the dag.execute()
+    class UserDataObj:
+        user_object_field_0: Any
+        user_object_field_1: Any
+        field_3: Any
+
+        def __init__(
+            self, user_object_field_0: Any, user_object_field_1: Any, field_3: Any
+        ) -> None:
+            self.user_object_field_0 = user_object_field_0
+            self.user_object_field_1 = user_object_field_1
+            self.field_3 = field_3
+
+    with InputNode() as dag_input:
+        m1 = Model.bind(1)
+        m2 = Model.bind(2)
+        m1_output = m1.forward.bind(dag_input.user_object_field_0)
+        m2_output = m2.forward.bind(dag_input.user_object_field_1)
+        dag = combine.bind(m1_output, m2_output, dag_input.field_3)
+
+    # 1*1 + 2*2 + 3
+    assert ray.get(dag.execute(UserDataObj(1, 2, 3))) == 8
+
+    # 3) Test user passed only one list object with regular list index accessor
+    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[1])
+        dag = combine.bind(m1_output, m2_output, dag_input[2])
+    # 1*1 + 2*2 + 3 + 4 = 12
+    assert ray.get(dag.execute([1, 2, 3])) == 8
+
+    # 4) Test user passed only one dict object with key str accessor
+    with InputNode() as dag_input:
+        m1 = Model.bind(1)
+        m2 = Model.bind(2)
+        m1_output = m1.forward.bind(dag_input["m1"])
+        m2_output = m2.forward.bind(dag_input["m2"])
+        dag = combine.bind(m1_output, m2_output, dag_input["m3"])
+    # 1*1 + 2*2 + 3 + 4 = 12
+    assert ray.get(dag.execute({"m1": 1, "m2": 2, "m3": 3})) == 8
+
+    with pytest.raises(
+        AssertionError,
+        match="Please only use int index or str as first-level key",
+    ):
+        with InputNode() as dag_input:
+            m1 = Model.bind(1)
+            dag = m1.forward.bind(dag_input[(1, 2)])
+
+
+def test_ensure_in_context_manager(shared_ray_instance):
+    # No enforcement on creation given __enter__ executes after __init__
+    input = InputNode()
+    with pytest.raises(
+        AssertionError,
+        match=(
+            "InputNode is a singleton instance that should be only used "
+            "in context manager"
+        ),
+    ):
+        input.execute()
+
+    @ray.remote
+    def f(input):
+        return input
+
+    # No enforcement on creation given __enter__ executes after __init__
+    dag = f.bind(InputNode())
+    with pytest.raises(
+        AssertionError,
+        match=(
+            "InputNode is a singleton instance that should be only used "
+            "in context manager"
+        ),
+    ):
+        dag.execute()
+
+
+def test_ensure_input_node_singleton(shared_ray_instance):
+    @ray.remote
+    def f(input):
+        return input
+
+    @ray.remote
+    def combine(a, b):
+        return a + b
+
+    with InputNode() as input_1:
+        a = f.bind(input_1)
+    with InputNode() as input_2:
+        b = f.bind(input_2)
+        dag = combine.bind(a, b)
+
+    with pytest.raises(
+        AssertionError, match="Each DAG should only have one unique InputNode"
+    ):
+        _ = ray.get(dag.execute(2))
+
+
+def test_apply_recursive_caching(shared_ray_instance):
+    @ray.remote
+    def f(input):
+        return input
+
+    input = InputNode()
+    f_node = f.bind(input)
+
+    a, b = input, f_node
+    for _ in range(10):
+        a, b = f.bind(a, b), f.bind(a, b)
+
+    counter = 0
+    original_apply_recursive = DAGNode.apply_recursive
+
+    def _apply_recursive_with_counter(self, fn):
+        nonlocal counter
+        counter += 1
+        return original_apply_recursive(self, fn)
+
+    DAGNode.apply_recursive = _apply_recursive_with_counter
+
+    a.apply_recursive(lambda node: node)
+
+    # Prior to #40337; count was 2559
+    assert counter == 40
+    DAGNode.apply_recursive = original_apply_recursive
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/dag/tests/test_output_node.py

@@ -0,0 +1,203 @@
+import pytest
+
+import ray
+from ray._common.test_utils import wait_for_condition
+from ray.dag.input_node import InputNode
+from ray.dag.output_node import MultiOutputNode
+from ray.util.state import list_tasks
+
+
+def test_output_node(shared_ray_instance):
+    @ray.remote
+    def f(input):
+        return input
+
+    with pytest.raises(ValueError):
+        with InputNode() as input_data:
+            dag = MultiOutputNode(f.bind(input_data))
+
+    with InputNode() as input_data:
+        dag = MultiOutputNode([f.bind(input_data)])
+
+    assert ray.get(dag.execute(1)) == [1]
+    assert ray.get(dag.execute(2)) == [2]
+
+    with InputNode() as input_data:
+        dag = MultiOutputNode([f.bind(input_data["x"]), f.bind(input_data["y"])])
+
+    refs = dag.execute({"x": 1, "y": 2})
+    assert len(refs) == 2
+    assert ray.get(refs) == [1, 2]
+
+    with InputNode() as input_data:
+        dag = MultiOutputNode(
+            [f.bind(input_data["x"]), f.bind(input_data["y"]), f.bind(input_data["x"])]
+        )
+
+    refs = dag.execute({"x": 1, "y": 2})
+    assert len(refs) == 3
+    assert ray.get(refs) == [1, 2, 1]
+
+
+def test_dag_with_actor_handle(shared_ray_instance):
+    """Verify DAG API works with actor created by .remote"""
+
+    @ray.remote
+    class Worker:
+        def __init__(self):
+            self.forward_called = 0
+            self.init_called = 0
+
+        def forward(self, input):
+            print("forward")
+            self.forward_called += 1
+            return input
+
+        def initialize(self, input):
+            print("initialize")
+            self.init_called += 1
+            return input
+
+        def get(self):
+            return (self.forward_called, self.init_called)
+
+    worker = Worker.remote()
+    with InputNode() as input_node:
+        init_dag = worker.initialize.bind(input_node)
+    with InputNode() as input_node:
+        forward_dag = worker.forward.bind(input_node)
+
+    assert ray.get(init_dag.execute(1)) == 1
+    assert ray.get(forward_dag.execute(2)) == 2
+
+    # Make sure both forward/initialize called only once
+    assert ray.get(worker.get.remote()) == (1, 1)
+
+    # Double check the actor is resued.
+    assert ray.get(init_dag.execute(1)) == 1
+    assert ray.get(worker.get.remote()) == (1, 2)
+
+
+def test_dag_with_alive_actors_chained(shared_ray_instance):
+    """Verify we can have multiple DAGs to the
+    same actor that are chained.
+    """
+
+    @ray.remote
+    class Actor:
+        def __init__(self, init_value):
+            self.i = init_value
+
+        def add(self, x):
+            return self.i + x
+
+    @ray.remote
+    def combine(x, y):
+        return x + y
+
+    a1 = Actor.remote(10)
+    a1_dag = a1.add.bind(a1.add.bind(2))  # 22
+    a1_dag_2 = a1.add.bind(a1.add.bind(6))  # 26
+    dag = combine.bind(a1_dag, a1_dag_2)
+
+    assert ray.get(dag.execute()) == 48
+
+
+def test_tensor_parallel_dag(shared_ray_instance):
+    """Simulate the TP DAG with N workers.
+    Input -> forward -> MultiOutput
+    """
+
+    @ray.remote
+    class Worker:
+        def __init__(self, rank):
+            self.rank = rank
+            self.forwarded = 0
+
+        def forward(self, input_data: int):
+            print(input_data)
+            self.forwarded += 1
+            return self.rank + input_data
+
+        def initialize(self):
+            pass
+
+        def get_forwarded(self):
+            return self.forwarded
+
+    NUM_WORKERS = 4
+    workers = [Worker.remote(i) for i in range(NUM_WORKERS)]
+    # Init multiple times.
+    for _ in range(4):
+        ray.get([worker.initialize.remote() for worker in workers])
+
+    with InputNode() as input_data:
+        dag = MultiOutputNode([worker.forward.bind(input_data) for worker in workers])
+
+    # Run DAG repetitively.
+    ITER = 4
+    assert ITER > 1
+    for i in range(ITER):
+        ref = dag.execute(i)
+        all_outputs = ray.get(ref)
+        assert len(all_outputs) == NUM_WORKERS
+        assert all_outputs == [i + j for j in range(NUM_WORKERS)]
+
+    forwarded = ray.get([worker.get_forwarded.remote() for worker in workers])
+    assert forwarded == [ITER for _ in range(NUM_WORKERS)]
+
+
+def test_shared_output(shared_ray_instance):
+    """Verify when an upstream task output is shared by
+    multi output, the upstream task runs only once.
+    """
+
+    @ray.remote
+    def shared_f():
+        return 1
+
+    @ray.remote
+    def g(input):
+        return input + 1
+
+    @ray.remote
+    def h(input):
+        return input + 2
+
+    x = shared_f.bind()
+    dag = MultiOutputNode([g.bind(x), h.bind(x)])
+
+    assert ray.get(dag.execute()) == [2, 3]
+
+    # Verify f ran only once.
+    def verify():
+        tasks = list_tasks(filters=[("name", "=", "shared_f")])
+        return len(tasks) == 1
+
+    wait_for_condition(verify)
+
+
+def test_bind_survives_handle_deletion(shared_ray_instance):
+    """Verify that .bind().execute() still works even if the original handle was dropped."""
+
+    @ray.remote
+    class A:
+        def f(self):
+            return 1
+
+    # Grab the handle and the bound method node
+    actor = A.remote()
+    method_node = actor.f.bind()
+
+    # Destroy the only Python variable reference and force collection
+    del actor
+
+    # Executing should now succeed because the node holds the ref
+    result = ray.get(method_node.execute())
+    assert result == 1
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/dag/tests/test_plot.py

@@ -0,0 +1,67 @@
+import os
+import sys
+import tempfile
+
+import pytest
+
+import ray
+
+
+def test_basic_dag_with_names_plot():
+    @ray.remote
+    def a(*args, **kwargs):
+        pass
+
+    tmp1 = a.options(name="tmp1").bind()
+    tmp2 = a.options(name="tmp2").bind()
+    tmp3 = a.options(name="tmp3").bind(tmp1, tmp2)
+    tmp4 = a.options(name="tmp4").bind()
+    tmp5 = a.options(name="tmp5").bind(tmp4)
+    tmp6 = a.options(name="tmp6").bind()
+    dag = a.bind(tmp3, tmp5, tmp6)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        to_file = os.path.join(tmpdir, "tmp.png")
+        ray.dag.plot(dag, to_file)
+        assert os.path.isfile(to_file)
+
+    graph = ray.dag.vis_utils._dag_to_dot(dag)
+    to_string = graph.to_string()
+    assert "tmp1 -> tmp3" in to_string
+    assert "tmp2 -> tmp3" in to_string
+    assert "tmp4 -> tmp5" in to_string
+    assert "tmp3 -> a" in to_string
+    assert "tmp5 -> a" in to_string
+    assert "tmp6 -> a" in to_string
+
+
+def test_basic_dag_without_names_plot():
+    @ray.remote
+    def a(*args, **kwargs):
+        pass
+
+    tmp1 = a.bind()
+    tmp2 = a.bind()
+    tmp3 = a.bind(tmp1, tmp2)
+    tmp4 = a.bind()
+    tmp5 = a.bind(tmp4)
+    tmp6 = a.bind()
+    dag = a.bind(tmp3, tmp5, tmp6)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        to_file = os.path.join(tmpdir, "tmp.png")
+        ray.dag.plot(dag, to_file)
+        assert os.path.isfile(to_file)
+
+    graph = ray.dag.vis_utils._dag_to_dot(dag)
+    to_string = graph.to_string()
+    assert "a -> a_2" in to_string
+    assert "a_1 -> a_2" in to_string
+    assert "a_3 -> a_4" in to_string
+    assert "a_2 -> a_6" in to_string
+    assert "a_4 -> a_6" in to_string
+    assert "a_5 -> a_6" in to_string
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))

python/ray/dag/tests/test_py_obj_scanner.py

@@ -0,0 +1,88 @@
+from typing import Any
+
+import pytest
+
+from ray.dag.py_obj_scanner import _instances, _PyObjScanner
+
+
+class Source:
+    pass
+
+
+def test_simple_replace():
+    scanner = _PyObjScanner(source_type=Source)
+    my_objs = [Source(), [Source(), {"key": Source()}]]
+
+    found = scanner.find_nodes(my_objs)
+    assert len(found) == 3
+
+    replaced = scanner.replace_nodes({obj: 1 for obj in found})
+    assert replaced == [1, [1, {"key": 1}]]
+
+
+def test_replace_multiple_types():
+    class OtherSource:
+        pass
+
+    scanner = _PyObjScanner(source_type=(Source, OtherSource))
+    my_objs = [Source(), [Source(), {"key": Source(), "key2": OtherSource()}]]
+
+    found = scanner.find_nodes(my_objs)
+    assert len(found) == 4
+
+    replaced = scanner.replace_nodes(
+        {obj: 1 if isinstance(obj, Source) else 2 for obj in found}
+    )
+    assert replaced == [1, [1, {"key": 1, "key2": 2}]]
+
+
+def test_replace_nested_in_obj():
+    """Test that the source can be nested in arbitrary objects."""
+    scanner = _PyObjScanner(source_type=Source)
+
+    class Outer:
+        def __init__(self, inner: Any):
+            self._inner = inner
+
+        def __eq__(self, other):
+            return self._inner == other._inner
+
+    my_objs = [Outer(Source()), Outer(Outer(Source())), Outer((Source(),))]
+
+    found = scanner.find_nodes(my_objs)
+    assert len(found) == 3
+
+    replaced = scanner.replace_nodes({obj: 1 for obj in found})
+    assert replaced == [Outer(1), Outer(Outer(1)), Outer((1,))]
+
+
+def test_scanner_clear():
+    """Test scanner clear to make the scanner GCable"""
+    prev_len = len(_instances)
+
+    def call_find_nodes():
+        scanner = _PyObjScanner(source_type=Source)
+        my_objs = [Source(), [Source(), {"key": Source()}]]
+        scanner.find_nodes(my_objs)
+        scanner.clear()
+        assert id(scanner) not in _instances
+
+    call_find_nodes()
+    assert prev_len == len(_instances)
+
+    def call_find_and_replace_nodes():
+        scanner = _PyObjScanner(source_type=Source)
+        my_objs = [Source(), [Source(), {"key": Source()}]]
+        found = scanner.find_nodes(my_objs)
+        scanner.replace_nodes({obj: 1 for obj in found})
+        scanner.clear()
+        assert id(scanner) not in _instances
+
+    call_find_and_replace_nodes()
+    assert prev_len == len(_instances)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/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()

python/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 memoization 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