diff --git a/hamilton/execution/graph_functions.py b/hamilton/execution/graph_functions.py index 0fbb3cb23..db57ede95 100644 --- a/hamilton/execution/graph_functions.py +++ b/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) diff --git a/tests/test_graph.py b/tests/test_graph.py index b20000184..1703424fe 100644 --- a/tests/test_graph.py +++ b/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.