Add files using upload-large-folder tool

.gitattributes

@@ -158,3 +158,4 @@ tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_inductor/_
 .venv/lib/python3.11/site-packages/ray/serve/_private/__pycache__/deployment_state.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text
 .venv/lib/python3.11/site-packages/xgrammar/xgrammar_bindings.cpython-311-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
 .venv/lib/python3.11/site-packages/ray/_raylet.so filter=lfs diff=lfs merge=lfs -text
+.venv/lib/python3.11/site-packages/ray/core/libjemalloc.so filter=lfs diff=lfs merge=lfs -text

.venv/lib/python3.11/site-packages/ray/core/libjemalloc.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0284919db23f95e692026039838aef89b5964b5cfec4a88acb9b3a9f4a226fd5
+size 885296

.venv/lib/python3.11/site-packages/ray/dag/__init__.py

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

.venv/lib/python3.11/site-packages/ray/dag/base.py

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

.venv/lib/python3.11/site-packages/ray/dag/class_node.py

@@ -0,0 +1,321 @@
+from weakref import ReferenceType
+
+import ray
+from ray.dag.dag_node import DAGNode
+from ray.dag.input_node import InputNode
+from ray.dag.format_utils import get_dag_node_str
+from ray.dag.constants import (
+    PARENT_CLASS_NODE_KEY,
+    PREV_CLASS_METHOD_CALL_KEY,
+    BIND_INDEX_KEY,
+    IS_CLASS_METHOD_OUTPUT_KEY,
+)
+from ray.util.annotations import DeveloperAPI
+
+from typing import Any, Dict, List, Union, Tuple, Optional
+
+
+@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 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

.venv/lib/python3.11/site-packages/ray/dag/collective_node.py

@@ -0,0 +1,191 @@
+from typing import Any, Dict, List, Union, Tuple, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch
+
+import ray
+from ray.dag import (
+    DAGNode,
+    ClassMethodNode,
+)
+from ray.dag.constants import COLLECTIVE_OPERATION_KEY
+from ray.experimental.channel import ChannelContext
+from ray.experimental.channel.torch_tensor_nccl_channel import _init_communicator
+from ray.experimental.channel.torch_tensor_type import Communicator, TorchTensorType
+from ray.experimental.util.types import _CollectiveOp, ReduceOp
+from ray.util.annotations import DeveloperAPI
+
+
+class _CollectiveOperation:
+    """
+    Represent metadata for a NCCL collective operation.
+
+    Args:
+        input_nodes: A list of input nodes to the collective operation.
+        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 NCCL group if specified.
+    """
+
+    def __init__(
+        self,
+        input_nodes: List[DAGNode],
+        op: _CollectiveOp,
+        transport: Optional[Union[str, Communicator]] = None,
+    ):
+        if len(input_nodes) == 0:
+            raise ValueError("Expected input nodes for a collective operation")
+        if len(set(input_nodes)) != len(input_nodes):
+            raise ValueError("Expected unique input nodes for a collective operation")
+
+        self._actor_handles: List["ray.actor.ActorHandle"] = []
+        for input_node in input_nodes:
+            actor_handle = input_node._get_actor_handle()
+            if actor_handle is None:
+                raise ValueError("Expected an actor handle from the input node")
+            self._actor_handles.append(actor_handle)
+        if len(set(self._actor_handles)) != len(self._actor_handles):
+            invalid_input_nodes = [
+                input_node
+                for input_node in input_nodes
+                if self._actor_handles.count(input_node._get_actor_handle()) > 1
+            ]
+            raise ValueError(
+                "Expected unique actor handles for a collective operation, "
+                "but found duplicate actor handles from input nodes: "
+                f"{invalid_input_nodes}"
+            )
+
+        self._op = op
+        if not isinstance(self._op, ReduceOp):
+            raise NotImplementedError("Only ReduceOp is implemented")
+        if transport is None:
+            transport = TorchTensorType.NCCL
+        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 NCCL group"
+                )
+
+    def __str__(self) -> str:
+        return (
+            f"CollectiveGroup("
+            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 init_communicator(self, communicator_id: Optional[str] = None) -> str:
+        """
+        Initialize the communicator if it has not been initialized yet. If
+        `communicator_id` is provided, it means the communicator has already
+        been initialized.
+        """
+        type_hint = self._type_hint
+        if type_hint.communicator_id is not None:
+            return type_hint.communicator_id
+        if communicator_id is None:
+            communicator_id = _init_communicator(
+                self._actor_handles, type_hint.get_custom_communicator()
+            )
+        type_hint.set_communicator_id(communicator_id)
+        return communicator_id
+
+    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 NCCL group")
+        return communicator
+
+    def execute(self, send_buf: "torch.Tensor") -> "torch.Tensor":
+        """
+        Call the collective operation on the input tensor. An output tensor is
+        allocated and returned.
+        """
+        import torch
+
+        if not isinstance(send_buf, torch.Tensor):
+            raise ValueError("Expected a torch tensor")
+        communicator = self.get_communicator()
+        recv_buf = torch.empty_like(send_buf)
+        communicator.allreduce(send_buf, recv_buf, self._op)
+        return recv_buf
+
+
+@DeveloperAPI
+class CollectiveOutputNode(ClassMethodNode):
+    """Represent an output node from a NCCL 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.
+        if not (
+            isinstance(method_args, tuple)
+            and len(method_args) == 1
+            and isinstance(method_args[0], DAGNode)
+        ):
+            raise ValueError("Expected a single input node")
+        self._input_node = method_args[0]
+        # Parse the collective operation.
+        self._collective_op: _CollectiveOperation = other_args_to_resolve.get(
+            COLLECTIVE_OPERATION_KEY, None
+        )
+        if self._collective_op is None:
+            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

.venv/lib/python3.11/site-packages/ray/dag/conftest.py

@@ -0,0 +1,16 @@
+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)

.venv/lib/python3.11/site-packages/ray/dag/constants.py

@@ -0,0 +1,33 @@
+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.
+RAY_CGRAPH_ENABLE_NVTX_PROFILING = (
+    os.environ.get("RAY_CGRAPH_ENABLE_NVTX_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"
+)

.venv/lib/python3.11/site-packages/ray/dag/context.py

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

.venv/lib/python3.11/site-packages/ray/dag/dag_node.py

@@ -0,0 +1,622 @@
+import copy
+from ray.experimental.channel.auto_transport_type import AutoTransportType
+from ray.experimental.channel.torch_tensor_type import TorchTensorType
+import ray
+from ray.dag.base import DAGNodeBase
+from ray.dag.py_obj_scanner import _PyObjScanner
+from ray.util.annotations import DeveloperAPI
+
+from itertools import chain
+
+from typing import (
+    Optional,
+    Union,
+    List,
+    Tuple,
+    Dict,
+    Any,
+    TypeVar,
+    Callable,
+)
+import uuid
+import asyncio
+
+from ray.dag.compiled_dag_node import build_compiled_dag_from_ray_dag
+from ray.experimental.channel import ChannelOutputType
+from ray.experimental.channel.communicator import Communicator
+
+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",
+        _static_shape: bool = False,
+        _direct_return: bool = False,
+    ):
+        if transport == "auto":
+            self._type_hint = AutoTransportType(
+                _static_shape=_static_shape,
+                _direct_return=_direct_return,
+            )
+        elif transport == "nccl":
+            self._type_hint = TorchTensorType(
+                transport=transport,
+                _static_shape=_static_shape,
+                _direct_return=_direct_return,
+            )
+        else:
+            if not isinstance(transport, Communicator):
+                raise ValueError(
+                    "transport must be 'auto', 'nccl' or a Communicator type"
+                )
+            self._type_hint = TorchTensorType(
+                transport=transport,
+                _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,
+        _overlap_gpu_communication: Optional[bool] = None,
+    ) -> "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.
+            _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.
+
+        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,
+            _overlap_gpu_communication,
+        )
+
+    def execute(
+        self, *args, _ray_cache_refs: bool = False, **kwargs
+    ) -> Union[ray.ObjectRef, "ray.actor.ActorHandle"]:
+        """Execute this DAG using the Ray default executor _execute_impl().
+
+        Args:
+            _ray_cache_refs: If true, stores the the default executor's return values
+                on each node in this DAG in a cache. These should be a mix of:
+                - ray.ObjectRefs pointing to the outputs of method and function nodes
+                - Serve handles for class nodes
+                - resolved values representing user input at runtime
+        """
+
+        def executor(node):
+            return node._execute_impl(*args, **kwargs)
+
+        result = self.apply_recursive(executor)
+        if _ray_cache_refs:
+            self.cache_from_last_execute = executor.cache
+        return result
+
+    def _get_toplevel_child_nodes(self) -> List["DAGNode"]:
+        """Return the list of nodes specified as top-level args.
+
+        For example, in `f.remote(a, [b])`, only `a` is a top-level arg.
+
+        This list of nodes are those that are typically resolved prior to
+        task execution in Ray. This does not include nodes nested within args.
+        For that, use ``_get_all_child_nodes()``.
+        """
+
+        # we use List instead of Set here because the hash key of the node
+        # object changes each time we create it. So if using Set here, the
+        # order of returned children can be different if we create the same
+        # nodes and dag one more time.
+        children = []
+        for a in self.get_args():
+            if isinstance(a, DAGNode):
+                if a not in children:
+                    children.append(a)
+        for a in self.get_kwargs().values():
+            if isinstance(a, DAGNode):
+                if a not in children:
+                    children.append(a)
+        for a in self.get_other_args_to_resolve().values():
+            if isinstance(a, DAGNode):
+                if a not in children:
+                    children.append(a)
+        return children
+
+    def _get_all_child_nodes(self) -> List["DAGNode"]:
+        """Return the list of nodes referenced by the args, kwargs, and
+        args_to_resolve in current node, even they're deeply nested.
+
+        Examples:
+            f.remote(a, [b]) -> [a, b]
+            f.remote(a, [b], key={"nested": [c]}) -> [a, b, c]
+        """
+
+        scanner = _PyObjScanner()
+        # we use List instead of Set here, reason explained
+        # in `_get_toplevel_child_nodes`.
+        children = []
+        for n in scanner.find_nodes(
+            [
+                self._bound_args,
+                self._bound_kwargs,
+                self._bound_other_args_to_resolve,
+            ]
+        ):
+            if n not in children:
+                children.append(n)
+        scanner.clear()
+        return children
+
+    def _apply_and_replace_all_child_nodes(
+        self, fn: "Callable[[DAGNode], T]"
+    ) -> "DAGNode":
+        """Apply and replace all immediate child nodes using a given function.
+
+        This is a shallow replacement only. To recursively transform nodes in
+        the DAG, use ``apply_recursive()``.
+
+        Args:
+            fn: Callable that will be applied once to each child of this node.
+
+        Returns:
+            New DAGNode after replacing all child nodes.
+        """
+
+        replace_table = {}
+        # CloudPickler scanner object for current layer of DAGNode. Same
+        # scanner should be use for a full find & replace cycle.
+        scanner = _PyObjScanner()
+        # Find all first-level nested DAGNode children in args.
+        # Update replacement table and execute the replace.
+        for node in scanner.find_nodes(
+            [
+                self._bound_args,
+                self._bound_kwargs,
+                self._bound_other_args_to_resolve,
+            ]
+        ):
+            if node not in replace_table:
+                replace_table[node] = fn(node)
+        new_args, new_kwargs, new_other_args_to_resolve = scanner.replace_nodes(
+            replace_table
+        )
+        scanner.clear()
+
+        # Return updated copy of self.
+        return self._copy(
+            new_args, new_kwargs, self.get_options(), new_other_args_to_resolve
+        )
+
+    def apply_recursive(self, fn: "Callable[[DAGNode], T]") -> T:
+        """Apply callable on each node in this DAG in a bottom-up tree walk.
+
+        Args:
+            fn: Callable that will be applied once to each node in the
+                DAG. It will be applied recursively bottom-up, so nodes can
+                assume the fn has been applied to their args already.
+
+        Returns:
+            Return type of the fn after application to the tree.
+        """
+
+        if not type(fn).__name__ == "_CachingFn":
+
+            class _CachingFn:
+                def __init__(self, fn):
+                    self.cache = {}
+                    self.fn = fn
+                    self.fn.cache = self.cache
+                    self.input_node_uuid = None
+
+                def __call__(self, node: "DAGNode"):
+                    from ray.dag.input_node import InputNode
+
+                    if node._stable_uuid not in self.cache:
+                        self.cache[node._stable_uuid] = self.fn(node)
+                    if isinstance(node, InputNode):
+                        if not self.input_node_uuid:
+                            self.input_node_uuid = node._stable_uuid
+                        elif self.input_node_uuid != node._stable_uuid:
+                            raise AssertionError(
+                                "Each DAG should only have one unique InputNode."
+                            )
+                    return self.cache[node._stable_uuid]
+
+            fn = _CachingFn(fn)
+        else:
+            if self._stable_uuid in fn.cache:
+                return fn.cache[self._stable_uuid]
+
+        return fn(
+            self._apply_and_replace_all_child_nodes(
+                lambda node: node.apply_recursive(fn)
+            )
+        )
+
+    def traverse_and_apply(self, fn: "Callable[[DAGNode], T]"):
+        """
+        Traverse all nodes in the connected component of the DAG that contains
+        the `self` node, and apply the given function to each node.
+        """
+        visited = set()
+        queue = [self]
+        cgraph_output_node: Optional[DAGNode] = None
+
+        while queue:
+            node = queue.pop(0)
+            if node._args_contain_nested_dag_node:
+                self._raise_nested_dag_node_error(node._bound_args)
+
+            if node not in visited:
+                if node.is_cgraph_output_node:
+                    # Validate whether there are multiple nodes that call
+                    # `experimental_compile`.
+                    if cgraph_output_node is not None:
+                        raise ValueError(
+                            "The DAG was compiled more than once. The following two "
+                            "nodes call `experimental_compile`: "
+                            f"(1) {cgraph_output_node}, (2) {node}"
+                        )
+                    cgraph_output_node = node
+                fn(node)
+                visited.add(node)
+                """
+                Add all unseen downstream and upstream nodes to the queue.
+                This function should be called by the root of the DAG. However,
+                in some invalid cases, some nodes may not be descendants of the
+                root. Therefore, we also add upstream nodes to the queue so that
+                a meaningful error message can be raised when the DAG is compiled.
+
+                ```
+                with InputNode() as inp:
+                    dag = MultiOutputNode([a1.inc.bind(inp), a2.inc.bind(1)])
+                ```
+
+                In the above example, `a2.inc` is not a descendant of inp. If we only
+                add downstream nodes to the queue, the `a2.inc` node will not be visited
+                , and the error message will be hard to understand, such as a key error
+                in the compiled DAG.
+                """
+                for neighbor in chain.from_iterable(
+                    [node._downstream_nodes, node._upstream_nodes]
+                ):
+                    if neighbor not in visited:
+                        queue.append(neighbor)
+
+    def _raise_nested_dag_node_error(self, args):
+        """
+        Raise an error for nested DAGNodes in Ray Compiled Graphs.
+
+        Args:
+            args: The arguments of the DAGNode.
+        """
+        for arg in args:
+            if isinstance(arg, DAGNode):
+                continue
+            else:
+                scanner = _PyObjScanner()
+                dag_nodes = scanner.find_nodes([arg])
+                scanner.clear()
+                if len(dag_nodes) > 0:
+                    raise ValueError(
+                        f"Found {len(dag_nodes)} DAGNodes from the arg {arg} "
+                        f"in {self}. Please ensure that the argument is a "
+                        "single DAGNode and that a DAGNode is not allowed to "
+                        "be placed inside any type of container."
+                    )
+        raise AssertionError(
+            "A DAGNode's args should contain nested DAGNodes as args, "
+            "but none were found during the compilation process. This is a "
+            "Ray internal error. Please report this issue to the Ray team."
+        )
+
+    def _find_root(self) -> "DAGNode":
+        """
+        Return the root node of the DAG. The root node must be an InputNode.
+        """
+        from ray.dag.input_node import InputNode
+
+        node = self
+        while not isinstance(node, InputNode):
+            if len(node._upstream_nodes) == 0:
+                raise ValueError(
+                    "No InputNode found in the DAG: when traversing upwards, "
+                    f"no upstream node was found for {node}."
+                )
+            node = node._upstream_nodes[0]
+        return node
+
+    def apply_functional(
+        self,
+        source_input_list: Any,
+        predictate_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.
+            predictate_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 appy 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 predictate_fn.
+        """
+        replace_table = {}
+        scanner = _PyObjScanner()
+        for node in scanner.find_nodes(source_input_list):
+            if predictate_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)

.venv/lib/python3.11/site-packages/ray/dag/dag_node_operation.py

@@ -0,0 +1,789 @@
+from functools import total_ordering
+from enum import Enum
+from typing import Set, Tuple, List, Dict, Optional
+import copy
+import logging
+import ray
+import heapq
+from collections import defaultdict
+
+
+logger = logging.getLogger(__name__)
+
+
+class _DAGNodeOperationType(Enum):
+    """
+    There are three types of operations that a DAG node can perform:
+    1. READ: Read from an input channel.
+    2. COMPUTE: Execute the method corresponding to the node.
+    3. WRITE: Write to an output channel.
+    """
+
+    READ = "READ"
+    COMPUTE = "COMPUTE"
+    WRITE = "WRITE"
+
+    def viz_str(self):
+        """
+        A string representation of the operation type to be used in visualization.
+
+        The result string is a single character because conciseness is preferred.
+        """
+        if self == _DAGNodeOperationType.READ:
+            return "R"
+        elif self == _DAGNodeOperationType.COMPUTE:
+            return "C"
+        elif self == _DAGNodeOperationType.WRITE:
+            return "W"
+        assert False, f"Unknown operation type: {self}"
+
+
+class _DAGNodeOperation:
+    def __init__(
+        self,
+        exec_task_idx: int,
+        operation_type: _DAGNodeOperationType,
+        method_name: Optional[str] = None,
+    ):
+        """
+        Args:
+            exec_task_idx: The index of the task that this operation belongs to
+                in the actor's ExecutableTask list. The index is not the same
+                as bind_index because there may be more tasks bound to an actor
+                than tasks that appear in the current compiled DAG.
+            operation_type: The type of operation to perform.
+            method_name: The name of the method that this operation originates
+                from. This is only for visualization and debugging purposes.
+        """
+        self.exec_task_idx = exec_task_idx
+        self.type = operation_type
+        self.method_name = method_name
+
+    def __repr__(self):
+        return (
+            f"_DAGNodeOperation("
+            f"exec_task_idx: {self.exec_task_idx}, "
+            f" type: {self.type})"
+        )
+
+    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_nccl: 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_nccl: Whether this operation requires NCCL.
+        """
+        self.operation = operation
+        self.task_idx = task_idx
+        self.actor_handle = actor_handle
+        self.requires_nccl = requires_nccl
+        # 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 collective nodes are the nodes that belong to the same collective
+        # operation. Each node is represented by a tuple of its task idx and type.
+        self.collective_idxs: Set[Tuple[int, _DAGNodeOperationType]] = set()
+        # The ready collective nodes are the nodes that are ready to be executed,
+        # i.e., their in-degrees are zero. When a collective node is ready, it
+        # will be added to the ready collective nodes of all the nodes in its
+        # collective operation.
+        self.ready_collective_idxs: Set[Tuple[int, _DAGNodeOperationType]] = set()
+
+    def __repr__(self):
+        return (
+            f"_DAGOperationGraphNode("
+            f"operation: {self.operation}, "
+            f"task_idx: {self.task_idx}, "
+            f"actor_handle: {self.actor_handle}, "
+            f"requires_nccl: {self.requires_nccl})"
+        )
+
+    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.
+        """
+
+        def compare(lhs: "_DAGOperationGraphNode", rhs: "_DAGOperationGraphNode"):
+            # If both nodes belong to the same actor, the node with the smaller
+            # `exec_task_idx` is prioritized. If two nodes belong to different
+            # actors, it approximates balancing the scheduled tasks across actors,
+            # by prioritizing the node with the smaller `exec_task_idx`. The tie
+            # is broken by the `task_idx`.
+            if lhs.operation.exec_task_idx != rhs.operation.exec_task_idx:
+                return lhs.operation.exec_task_idx < rhs.operation.exec_task_idx
+            return lhs.task_idx < rhs.task_idx
+
+        if self.actor_handle == other.actor_handle:
+            # When both nodes belong to the same actor, use the default comparison.
+            return compare(self, other)
+        elif self.is_nccl_op != other.is_nccl_op:
+            # When one node is a NCCL operation and the other is not, prioritize
+            # the non-NCCL operation.
+            return not self.is_nccl_op
+        else:
+            # When either both nodes are NCCL operations or both nodes are not
+            # NCCL operations, use the default comparison.
+            return compare(self, other)
+
+    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 a NCCL collective, it is ready when it has a zero
+        in-degree. If it is a NCCL collective, it is ready when all the nodes
+        in its collective operation have zero in-degrees.
+        """
+        return self.in_degree == 0 and (
+            len(self.ready_collective_idxs) == len(self.collective_idxs)
+        )
+
+    @property
+    def is_read(self) -> bool:
+        return self.operation.type == _DAGNodeOperationType.READ
+
+    @property
+    def is_nccl_collective(self) -> bool:
+        """
+        A node is a NCCL collective if it is a compute node and requires NCCL.
+        """
+        return (
+            self.operation.type == _DAGNodeOperationType.COMPUTE and self.requires_nccl
+        )
+
+    @property
+    def is_nccl_write(self) -> bool:
+        """
+        A node is a NCCL write if it is a write node and requires NCCL.
+        """
+        return self.operation.type == _DAGNodeOperationType.WRITE and self.requires_nccl
+
+    @property
+    def is_nccl_op(self) -> bool:
+        return self.is_nccl_collective or self.is_nccl_write
+
+    def viz_str(self):
+        """
+        A string representation of the node to be used in visualization.
+        """
+        return self.operation.viz_str()
+
+    @property
+    def _actor_id(self):
+        return self.actor_handle._ray_actor_id.hex()
+
+
+def _add_edge(
+    from_node: _DAGOperationGraphNode,
+    to_node: _DAGOperationGraphNode,
+    label: str = "",
+    control_dependency: bool = False,
+):
+    """
+    Add an edge from `from_node` to `to_node`.
+
+    Args:
+        from_node: The node from which the edge originates.
+        to_node: The node to which the edge points.
+        label: The label of the edge. This will be used to annotate the edge
+            in the visualization of the execution schedule.
+    """
+    from_node.out_edges[(to_node.task_idx, to_node.operation.type)] = (
+        label,
+        control_dependency,
+    )
+    to_node.in_edges[(from_node.task_idx, from_node.operation.type)] = (
+        label,
+        control_dependency,
+    )
+
+
+def _push_candidate_node_if_ready(
+    actor_to_candidates: Dict["ray._raylet.ActorID", List[_DAGOperationGraphNode]],
+    graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]],
+    node: _DAGOperationGraphNode,
+) -> None:
+    # Collective operations are ready when all the collective nodes have zero
+    # in-degrees. Only one node per collective will be added as ready.
+    if node.is_nccl_collective:
+        for collective_node_metadata in node.collective_idxs:
+            task_idx, op_type = collective_node_metadata
+            collective_node = graph[task_idx][op_type]
+            collective_node.ready_collective_idxs.add(
+                (node.task_idx, node.operation.type)
+            )
+    if node.is_ready:
+        heapq.heappush(
+            actor_to_candidates[node.actor_handle._actor_id],
+            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 a NCCL collective node, it is ready to be
+    executed if it has a zero in-degree. For a NCCL collective node, it is ready to
+    be executed when all the nodes in its collective operation have zero in-degrees.
+
+    If a node is a NCCL 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 NCCL collective node from its collective operation is
+    selected from the candidate list.
+
+    If the selected node is a NCCL write node, select all the downstream NCCL
+    read nodes. If the selected node is a NCCL collective node, select all the NCCL
+    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.items():
+        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 = [
+        heapq.heappop(actor_to_candidates[top_priority_node.actor_handle._actor_id])
+    ]
+
+    if not top_priority_node.is_nccl_op:
+        # A non-NCCL operation node is picked.
+        assert len(next_nodes) == 1
+    elif top_priority_node.is_nccl_write:
+        # a NCCL write node is picked. NCCL is a blocking operation, so we need
+        # to pick all the corresponding NCCL read nodes to avoid a deadlock.
+        for downstream_node_metadata in top_priority_node.out_edges:
+            task_idx, op_type = downstream_node_metadata
+            downstream_node = graph[task_idx][op_type]
+            assert downstream_node.is_read
+            next_nodes.append(downstream_node)
+        assert len(next_nodes) == 1 + len(top_priority_node.out_edges)
+    elif top_priority_node.is_nccl_collective:
+        # a NCCL collective node is picked. NCCL is a blocking operation, so we need
+        # to pick all the corresponding NCCL collective nodes in its collective
+        # operation to avoid a deadlock.
+        for collective_node_metadata in top_priority_node.collective_idxs:
+            task_idx, op_type = collective_node_metadata
+            collective_node = graph[task_idx][op_type]
+            assert collective_node.is_nccl_collective and collective_node.is_ready
+            if collective_node != top_priority_node:
+                next_nodes.append(collective_node)
+        assert len(next_nodes) == len(top_priority_node.collective_idxs)
+
+    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
+
+    # Add an edge from WRITE of the writer task to READ of the reader task.
+    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
+            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:
+                        _add_edge(
+                            graph[task_idx][_DAGNodeOperationType.WRITE],
+                            graph[consumer_idx][_DAGNodeOperationType.READ],
+                            "nccl"
+                            if graph[task_idx][
+                                _DAGNodeOperationType.WRITE
+                            ].requires_nccl
+                            else "shm",
+                        )
+                continue
+            _add_edge(
+                graph[task_idx][_DAGNodeOperationType.WRITE],
+                graph[downstream_task_idx][_DAGNodeOperationType.READ],
+                "nccl"
+                if graph[task_idx][_DAGNodeOperationType.WRITE].requires_nccl
+                else "shm",
+            )
+
+    return graph
+
+
+def _actor_viz_label(actor: "ray.actor.ActorHandle"):
+    """
+    Returns the label of an actor in the visualization of the execution schedule.
+
+    Args:
+        actor: The actor to be represented.
+    """
+    class_name = actor._ray_actor_creation_function_descriptor.class_name
+    actor_id = actor._ray_actor_id.hex()
+    return f"Actor class name: {class_name}\nActor ID: {actor_id}"
+
+
+def _node_viz_id_and_label(
+    node: _DAGOperationGraphNode, idx: int, optimized_index: int
+):
+    """
+    Returns the visualization id and label of a node. The visualization id is unique
+    across all nodes.
+
+    Args:
+        node: The node to be represented.
+        idx: The index of the node in the execution schedule.
+        optimized_index: The index of the node in the optimized execution schedule.
+    """
+    node_viz_label = node.viz_str() + f" {idx},{optimized_index}"
+    node_viz_id = f"{node._actor_id}_{node_viz_label}"
+    return node_viz_id, node_viz_label
+
+
+def _visualize_execution_schedule(
+    actor_to_execution_schedule: Dict[
+        "ray.actor.ActorHandle", List[_DAGOperationGraphNode]
+    ],
+    actor_to_overlapped_schedule: Optional[
+        Dict["ray.actor.ActorHandle", List[_DAGOperationGraphNode]]
+    ],
+    graph: Dict[int, Dict[_DAGNodeOperationType, _DAGOperationGraphNode]],
+):
+    """
+    Visualize the execution schedule for each actor.
+
+    The visualization will be saved as a PNG file named `compiled_graph_schedule.png`.
+    Details of the visualization: # noqa
+
+        Node description format:
+            [<task_index>] <method_name> <operation> <orig_index>, <overlap_index>
+
+        Node description fields:
+            operation: is R(READ), C(COMPUTE), or W(WRITE)
+            orig_index: the index in the original execution schedule
+            overlap_index: the index in the overlap-communication optimized execution schedule
+            If this is different from orig_index, the node is highlighted in red color
+
+        Node grouping:
+            The nodes belonging to the same actor are grouped in the same rectangle
+            The actor class name and the actor id are shown in the rectangle
+
+        Edges:
+            black color (without label): data dependency
+            black color (annotated with "shm"): shared memory channel
+            blue color (annotated with "nccl): NCCL 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 == "nccl" 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 "nccl): NCCL 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 a NCCL operation, only itself is returned.
+        # 2. If a selected node is a NCCL write operation, the corresponding NCCL
+        #    read operations are also returned.
+        # 3. If a selected node is a NCCL collective operation, all the nodes in
+        #    its collective operation are returned.
+        # In cases 1 and 3, all the selected nodes are ready. In case 2, the NCCL
+        # write node is ready, while the NCCL read nodes are not ready until their
+        # in-degrees are updated.
+        nodes = _select_next_nodes(actor_to_candidates, graph)
+        if nodes is None:
+            break
+        # Filter out the visited nodes.
+        nodes = [node for node in nodes if node not in visited_nodes]
+        # Add the selected nodes to the execution schedule.
+        for node in nodes:
+            actor_to_execution_schedule[node.actor_handle].append(node)
+            visited_nodes.add(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]
+                out_node.in_edges.pop((node.task_idx, node.operation.type))
+                if out_node.in_degree == 0 and out_node not in visited_nodes:
+                    # If the downstream node is already visited, it has been added
+                    # to the execution schedule. They are the NCCL read nodes in
+                    # case 2.
+                    _push_candidate_node_if_ready(actor_to_candidates, graph, out_node)
+    assert len(visited_nodes) == len(graph) * 3, "Expected all nodes to be visited"
+    for node in visited_nodes:
+        assert node.is_ready, f"Expected {node} to be ready"
+    for _, candidates in actor_to_candidates.items():
+        assert len(candidates) == 0, "Expected all candidates to be empty"
+
+    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 NCCL read operation (i.e., recv), scan backwards to find the nearest
+    compute node to swap with so that the NCCL 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_nccl
+            ):
+                # For each NCCL read operation (i.e., recv), scan backwards
+                # to find the nearest compute node to swap with so that
+                # the NCCL 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
+                        nccl_read_op = overlapped_schedule[i]
+                        prev_ops = overlapped_schedule[j:i]
+                        overlapped_schedule[j + 1 : i + 1] = prev_ops
+                        overlapped_schedule[j] = nccl_read_op
+                        break
+                    if (
+                        overlapped_schedule[j].operation.type
+                        == _DAGNodeOperationType.READ
+                        or overlapped_schedule[j].operation.type
+                        == _DAGNodeOperationType.WRITE
+                    ) and overlapped_schedule[j].requires_nccl:
+                        # Found a NCCL read/write operation, skip the overlap
+                        # optimization to keep relative order of NCCL 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()
+    }

.venv/lib/python3.11/site-packages/ray/dag/dag_operation_future.py

@@ -0,0 +1,95 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar
+from ray.util.annotations import DeveloperAPI
+
+
+if TYPE_CHECKING:
+    import cupy as cp
+
+T = TypeVar("T")
+
+
+@DeveloperAPI
+class DAGOperationFuture(ABC, Generic[T]):
+    """
+    A future representing the result of a DAG operation.
+
+    This is an abstraction that is internal to each actor,
+    and is not exposed to the DAG caller.
+    """
+
+    @abstractmethod
+    def wait(self):
+        """
+        Wait for the future and return the result of the operation.
+        """
+        raise NotImplementedError
+
+
+@DeveloperAPI
+class ResolvedFuture(DAGOperationFuture):
+    """
+    A future that is already resolved. Calling `wait()` on this will
+    immediately return the result without blocking.
+    """
+
+    def __init__(self, result):
+        """
+        Initialize a resolved future.
+
+        Args:
+            result: The result of the future.
+        """
+        self._result = result
+
+    def wait(self):
+        """
+        Wait and immediately return the result. This operation will not block.
+        """
+        return self._result
+
+
+@DeveloperAPI
+class GPUFuture(DAGOperationFuture[Any]):
+    """
+    A future for a GPU event on a CUDA stream.
+
+    This future wraps a buffer, and records an event on the given stream
+    when it is created. When the future is waited on, it makes the current
+    CUDA stream wait on the event, then returns the buffer.
+
+    The buffer must be a GPU tensor produced by an earlier operation launched
+    on the given stream, or it could be CPU data. Then the future guarantees
+    that when the wait() returns, the buffer is ready on the current stream.
+
+    The `wait()` does not block CPU.
+    """
+
+    def __init__(self, buf: Any, stream: Optional["cp.cuda.Stream"] = None):
+        """
+        Initialize a GPU future on the given stream.
+
+        Args:
+            buf: The buffer to return when the future is resolved.
+            stream: The CUDA stream to record the event on, this event is waited
+                on when the future is resolved. If None, the current stream is used.
+        """
+        import cupy as cp
+
+        if stream is None:
+            stream = cp.cuda.get_current_stream()
+
+        self._buf = buf
+        self._event = cp.cuda.Event()
+        self._event.record(stream)
+
+    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.
+        """
+        import cupy as cp
+
+        current_stream = cp.cuda.get_current_stream()
+        current_stream.wait_event(self._event)
+        return self._buf

.venv/lib/python3.11/site-packages/ray/dag/format_utils.py

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

.venv/lib/python3.11/site-packages/ray/dag/function_node.py

@@ -0,0 +1,60 @@
+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))

.venv/lib/python3.11/site-packages/ray/dag/input_node.py

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

.venv/lib/python3.11/site-packages/ray/dag/output_node.py

@@ -0,0 +1,45 @@
+import ray
+from typing import Any, Dict, List, Union, Tuple
+
+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__")

.venv/lib/python3.11/site-packages/ray/dag/py_obj_scanner.py

@@ -0,0 +1,105 @@
+import io
+from typing import Any, Dict, Generic, List, Tuple, Type, TypeVar, Union
+
+import pickle  # noqa: F401
+
+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()

.venv/lib/python3.11/site-packages/ray/dag/utils.py

@@ -0,0 +1,66 @@
+from typing import Dict
+
+from ray.dag import (
+    DAGNode,
+    InputNode,
+    InputAttributeNode,
+    FunctionNode,
+    ClassNode,
+    ClassMethodNode,
+    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()

.venv/lib/python3.11/site-packages/ray/dag/vis_utils.py

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

.venv/lib/python3.11/site-packages/ray/experimental/channel/__init__.py

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