Make execute_subdag and nodes_between traversals iterative

hamilton/execution/graph_functions.py

@@ -197,16 +197,33 @@ def execute_subdag(
     def dfs_traverse(
         node_: node.Node, dependency_type: node.DependencyType = node.DependencyType.REQUIRED
     ):
-        if node_.name in computed:
-            return
-        if node_.name in overrides:
-            computed[node_.name] = overrides[node_.name]
-            return
-        for n in node_.dependencies:
-            if n.name not in computed:
-                _, node_dependency_type = node_.input_types[n.name]
-                dfs_traverse(n, node_dependency_type)
+        # Explicit stack instead of recursion so chains deeper than the recursion limit work.
+        # Each entry holds a node, its dependency type, and its dependency iterator (None until entered).
+        stack = [(node_, dependency_type, None)]
+        in_progress = set()
+        while stack:
+            current, current_dependency_type, dependencies = stack.pop()
+            if dependencies is None:
+                if current.name in computed:
+                    continue
+                if current.name in overrides:
+                    computed[current.name] = overrides[current.name]
+                    continue
+                if current.name in in_progress:
+                    raise RecursionError(f"Cycle detected in the graph at node {current.name}.")
+                in_progress.add(current.name)
+                dependencies = iter(current.dependencies)
+            for n in dependencies:
+                if n.name not in computed:
+                    _, node_dependency_type = current.input_types[n.name]
+                    stack.append((current, current_dependency_type, dependencies))
+                    stack.append((n, node_dependency_type, None))
+                    break
+            else:
+                in_progress.remove(current.name)
+                compute_node(current, current_dependency_type)
 
+    def compute_node(node_: node.Node, dependency_type: node.DependencyType):
         logger.debug(f"Computing {node_.name}.")
         if node_.user_defined:
             if node_.name not in inputs:
@@ -376,24 +393,38 @@ def nodes_between(
     visited = set()
     search_node = None
 
-    def dfs_traverse(node_: node.Node):
-        if search_condition(node_):
-            # if we hit the end, we want to include all others in it
-            nonlocal search_node
-            search_node = node_
-            return True
-        if node_ in visited:
-            # if we've already seen it, we want to include it
-            return node_ in out
-        # now we mark that we've seen it
-        visited.add(node_)
-
-        any_deps_included = False
-        for n in node_.dependencies:
-            any_deps_included |= dfs_traverse(n)
-        if any_deps_included:
-            out.add(node_)
-        return any_deps_included
+    def dfs_traverse(start: node.Node):
+        nonlocal search_node
+        # Explicit stack instead of recursion so chains deeper than the recursion limit work.
+        # Each frame is [node, dependency iterator (None until entered), any_deps_included].
+        stack = [[start, None, False]]
+        while stack:
+            frame = stack[-1]
+            node_, dependencies, any_deps_included = frame
+            if dependencies is None:
+                if search_condition(node_):
+                    # if we hit the end, we want to include all others in it
+                    search_node = node_
+                    included = True
+                elif node_ in visited:
+                    # if we've already seen it, we want to include it
+                    included = node_ in out
+                else:
+                    # now we mark that we've seen it
+                    visited.add(node_)
+                    frame[1] = iter(node_.dependencies)
+                    continue
+            else:
+                n = next(dependencies, None)
+                if n is not None:
+                    stack.append([n, None, False])
+                    continue
+                if any_deps_included:
+                    out.add(node_)
+                included = any_deps_included
+            stack.pop()
+            if stack:
+                stack[-1][2] |= included
 
     for dep in end_node.dependencies:
         dfs_traverse(dep)

tests/test_graph.py

@@ -27,7 +27,7 @@
 
 import hamilton.graph_utils
 import hamilton.htypes
-from hamilton import ad_hoc_utils, base, graph, node
+from hamilton import ad_hoc_utils, base, driver, graph, node
 from hamilton import function_modifiers as fm
 from hamilton.execution import graph_functions
 from hamilton.function_modifiers import schema
@@ -643,6 +643,30 @@ def step(prev: float) -> float:
     assert all(fg.nodes[f"node_{i}"] in nodes for i in range(chain_size))
 
 
+def test_execute_and_nodes_between_large_chain_no_recursion_error(monkeypatch):
+    """Regression test for https://github.com/apache/hamilton/issues/1760.
+
+    Driver.execute (via execute_subdag) and nodes_between must not recurse once per
+    dependency level, so a chain deeper than the recursion limit has to work.
+    """
+    chain_size = 1200
+    assert chain_size > sys.getrecursionlimit()
+    module = ModuleType("large_exec_chain")
+    src = ["def node_0() -> float:\n    return 1.0\n"]
+    for i in range(1, chain_size):
+        src.append(f"def node_{i}(node_{i - 1}: float) -> float:\n    return node_{i - 1} + 1\n")
+    exec("\n".join(src), module.__dict__)
+    monkeypatch.setitem(sys.modules, module.__name__, module)
+    last = f"node_{chain_size - 1}"
+
+    dr = driver.Builder().with_modules(module).build()
+
+    assert dr.execute([last]) == {last: 1200.0}
+    assert dr.graph.nodes_between("node_0", last) == {
+        dr.graph.nodes[f"node_{i}"] for i in range(chain_size)
+    }
+
+
 def test_get_upstream_nodes_diamond_dag():
     """Tests that diamond-shaped DAGs don't produce duplicate visits.