address comments

Signed-off-by: Kai-Hsun Chen <kaihsun@anyscale.com>

Author: kevin85421

python/ray/dag/dag_node.py

@@ -4,6 +4,8 @@
 from ray.util.annotations import DeveloperAPI
 import copy
 
+from itertools import chain
+
 from typing import (
     Optional,
     Union,
@@ -60,28 +62,33 @@ def __init__(
         )
 
         # The list of nodes that use this DAG node as input.
-        self._downstream_nodes: List[DAGNode] = []
+        self._downstream_nodes: List["DAGNode"] = []
         # The list of nodes that this DAG node uses as input.
+        self._upstream_nodes: List["DAGNode"] = self._prepare_upstream_nodes()
+
+        # UUID that is not changed over copies of this node.
+        self._stable_uuid = uuid.uuid4().hex
+        # Cached values from last call to execute()
+        self.cache_from_last_execute = {}
+
+        self._type_hint: Optional[ChannelOutputType] = ChannelOutputType()
+        self.is_output_node = False
+
+    def _prepare_upstream_nodes(self) -> List["DAGNode"]:
+        """Retrieve upstream nodes and update their downstream dependencies."""
         scanner = _PyObjScanner()
-        self._upstream_nodes: List[DAGNode] = scanner.find_nodes(
+        upstream_nodes: List["DAGNode"] = scanner.find_nodes(
             [
                 self._bound_args,
                 self._bound_kwargs,
                 self._bound_other_args_to_resolve,
             ]
         )
         scanner.clear()
-        # Update upstream dependencies.
-        for upstream_node in self._upstream_nodes:
+        # Update dependencies.
+        for upstream_node in upstream_nodes:
             upstream_node._downstream_nodes.append(self)
-
-        # UUID that is not changed over copies of this node.
-        self._stable_uuid = uuid.uuid4().hex
-        # Cached values from last call to execute()
-        self.cache_from_last_execute = {}
-
-        self._type_hint: Optional[ChannelOutputType] = ChannelOutputType()
-        self.is_output_node = False
+        return upstream_nodes
 
     def with_type_hint(self, typ: ChannelOutputType):
         if typ.is_direct_return:
@@ -332,9 +339,11 @@ def __init__(self, fn):
                     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 type(node).__name__ == "InputNode":
+                    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:
@@ -368,14 +377,20 @@ def bfs(self, fn: "Callable[[DAGNode], T]"):
                 # 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.
-                for node in node._downstream_nodes + node._upstream_nodes:
+                for node in chain.from_iterable(
+                    [node._downstream_nodes, node._upstream_nodes]
+                ):
                     if node not in visited:
                         queue.append(node)
 
     def _find_root(self) -> "DAGNode":
-        """Return the root node of the DAG."""
+        """
+        Return the root node of the DAG. The root node must be an InputNode.
+        """
+        from ray.dag.input_node import InputNode
+
         root = self
-        while type(root).__name__ != "InputNode":
+        while not isinstance(root, InputNode):
             if len(root._upstream_nodes) == 0:
                 raise ValueError("No InputNode found in the DAG.")
             root = root._upstream_nodes[0]

python/ray/dag/py_obj_scanner.py

@@ -70,7 +70,14 @@ def reducer_override(self, obj):
         return super().reducer_override(obj)
 
     def find_nodes(self, obj: Any) -> List[SourceType]:
-        """Find top-level DAGNodes."""
+        """
+        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."