From 9ed8671f4cc40252e5a6d5155ba5d627b9d8c16e Mon Sep 17 00:00:00 2001 From: Octopus Date: Sun, 26 Apr 2026 20:37:20 +0800 Subject: [PATCH] fix(awel): BranchOperator incorrectly skips shared downstream nodes (fixes #2935) (#3035) Co-authored-by: octo-patch --- .../dbgpt/core/awel/runner/local_runner.py | 25 +++++++++- .../src/dbgpt/core/awel/tests/test_run_dag.py | 46 +++++++++++++++++++ 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/packages/dbgpt-core/src/dbgpt/core/awel/runner/local_runner.py b/packages/dbgpt-core/src/dbgpt/core/awel/runner/local_runner.py index a1eb33960..d5197a8e8 100644 --- a/packages/dbgpt-core/src/dbgpt/core/awel/runner/local_runner.py +++ b/packages/dbgpt-core/src/dbgpt/core/awel/runner/local_runner.py @@ -217,17 +217,40 @@ def _skip_current_downstream_by_node_name( ): if not skip_nodes: return + nodes_to_skip = [] for child in branch_node.downstream: child = cast(BaseOperator, child) if child.node_name in skip_nodes or child.node_id in skip_node_ids: logger.info(f"Skip node name {child.node_name}, node id {child.node_id}") + nodes_to_skip.append(child) + + # Pre-register all direct skip candidates so that shared downstream nodes + # (e.g. JoinOperator) can see the full set of skipped parents before deciding + # whether they themselves should be skipped. + for node in nodes_to_skip: + if node.can_skip_in_branch(): + skip_node_ids.add(node.node_id) + + # Now recurse into each candidate's downstream. + for node in nodes_to_skip: + for child in node.downstream: + child = cast(BaseOperator, child) _skip_downstream_by_id(child, skip_node_ids) def _skip_downstream_by_id(node: BaseOperator, skip_node_ids: Set[str]): if not node.can_skip_in_branch(): - # Current node can not skip, so skip its downstream + # Current node cannot be skipped, so leave it and its downstream intact. return + if node.node_id in skip_node_ids: + # Already marked for skipping; avoid redundant traversal. + return + # A node with multiple upstream parents (e.g. JoinOperator) should only be + # skipped when ALL of its parents are being skipped. If any parent is still + # active, this node must remain active to consume that parent's output. + for parent in node.upstream: + if isinstance(parent, BaseOperator) and parent.node_id not in skip_node_ids: + return skip_node_ids.add(node.node_id) for child in node.downstream: child = cast(BaseOperator, child) diff --git a/packages/dbgpt-core/src/dbgpt/core/awel/tests/test_run_dag.py b/packages/dbgpt-core/src/dbgpt/core/awel/tests/test_run_dag.py index f56533670..20c4577e5 100644 --- a/packages/dbgpt-core/src/dbgpt/core/awel/tests/test_run_dag.py +++ b/packages/dbgpt-core/src/dbgpt/core/awel/tests/test_run_dag.py @@ -138,3 +138,49 @@ async def test_branch_node( assert res.current_task_context.current_state == TaskState.SUCCESS expect_res = 999 if is_odd else 888 assert res.current_task_context.task_output.output == expect_res + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "input_node, is_odd", + [ + ({"outputs": [26]}, False), + ({"outputs": [3]}, True), + ], + indirect=["input_node"], +) +async def test_branch_node_shared_join_default_can_skip( + runner: WorkflowRunner, input_node: InputOperator, is_odd: bool +): + """Regression test for issue #2935. + + A JoinOperator shared by both branches of a BranchOperator must still execute + when only one branch is skipped, even when can_skip_in_branch is left at its + default value (True). Previously the skipped branch's downstream traversal + would incorrectly mark the shared JoinOperator as skipped. + """ + + def join_func(o1, o2) -> int: + return o1 or o2 + + with DAG("test_branch_shared_join") as _dag: + odd_node = MapOperator( + lambda x: 999, task_id="odd_node2", task_name="odd_node_name2" + ) + even_node = MapOperator( + lambda x: 888, task_id="even_node2", task_name="even_node_name2" + ) + # Intentionally omit can_skip_in_branch=False to reproduce the bug scenario. + join_node = JoinOperator(join_func) + branch_node = BranchOperator( + {lambda x: x % 2 == 1: odd_node, lambda x: x % 2 == 0: even_node} + ) + branch_node >> odd_node >> join_node + branch_node >> even_node >> join_node + + input_node >> branch_node + + res: DAGContext[int] = await runner.execute_workflow(join_node) + assert res.current_task_context.current_state == TaskState.SUCCESS + expect_res = 999 if is_odd else 888 + assert res.current_task_context.task_output.output == expect_res