Skip to content

fix(workflow): make JoinNode in a loop wait for each iteration's predecessors - #7499

Closed
AkshayShah03 wants to merge 1 commit into
google:mainfrom
AkshayShah03:fix/workflow-join-loop-iteration
Closed

AkshayShah03 wants to merge 1 commit into
google:mainfrom
AkshayShah03:fix/workflow-join-loop-iteration

Conversation

@AkshayShah03

@AkshayShah03 AkshayShah03 commented Oct 10, 2026 •

Copy link
Copy Markdown
Contributor

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

Problem:
Workflow._buffer_barrier_trigger considered a wait-for-all node (JoinNode) ready when every predecessor's NodeState.status was COMPLETED. That status is never consumed when the join fires, so inside a loop a predecessor that has not re-run yet still counts as done. The join fires as soon as the first predecessor of the next iteration completes, with the others' outputs from the previous iteration, and then fires again when they arrive, so everything downstream runs an extra time. With human-in-the-loop, the join even fires while the slower branch is still waiting for user input.

Solution:
Add _LoopState.join_arrivals: for each wait-for-all node, the set of predecessors that completed since it last fired. _handle_completion records the arrival (_record_join_arrivals), and _buffer_barrier_trigger fires only when every predecessor (other than START, which stays always-satisfied as before) is in the set, then clears it.

Notes:

  • An arrival is recorded for every completed run of a predecessor, regardless of its route, matching the previous COMPLETED check. Only the "since last fire" scoping changes.
  • Nothing new is persisted: _LoopState is rebuilt on resume by replaying recorded completions through _handle_completion, which rebuilds join_arrivals too. The resume test below covers this for both resumable and non-resumable apps.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Added to tests/unittests/workflow/test_join_node.py:

  • test_join_node_in_loop_waits_for_each_iteration: two-iteration loop where one branch is a hop longer. Expects join inputs [{A1, B1}, {A2, B2}]; main gives [{A1, B1}, {A2, B1}, {A2, B2}].
  • test_join_node_in_loop_waits_for_predecessor_paused_for_input[False|True]: same loop, but the slower branch pauses for RequestInput on iteration 2. Expects no join firing while it waits, then exactly {A2, B2} after resume. On main the join fires with {A2, B1} before the user answers.

All three fail on main and pass with this change.

$ pytest tests/unittests/workflow
1062 passed, 3 xfailed

$ pytest tests/unittests -n auto
18386 passed, 90 skipped, 22 xfailed, 2 xpassed

pyink and isort report no changes.

Manual End-to-End (E2E) Tests:

Ran the reproduction script from #7498 through a real Runner with InMemorySessionService:

  • on main: the join fires 3 times, the second time with {'a': 'A2', 'b': 'B1'}
  • with this change: the join fires twice, {'a': 'A1', 'b': 'B1'} then {'a': 'A2', 'b': 'B2'}

Also checked an existing non-loop join + RequestInput resume scenario (resumable and non-resumable): output unchanged.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

Additional context

A related, separate issue in the same area: _handle_completion only overwrites node_outputs[name] when the output is not None, so a predecessor that returns None on a later iteration passes its previous output to the join. I left that out to keep this PR to one concern and can follow up separately.

…ecessors

_buffer_barrier_trigger treated a wait-for-all node as ready when every
predecessor's status was COMPLETED. That status is never consumed when the
join fires, so in a loop the join fired as soon as the first predecessor of
the next iteration completed, using the other predecessors' outputs from the
previous iteration, and then fired again when they arrived.

Track the predecessors that completed since each wait-for-all node last fired
and clear them when it fires. Resume replays recorded completions through the
same path, so the tracking is rebuilt without persisting anything new.

Fixes google#7498
copybara-service Bot pushed a commit that referenced this pull request Oct 17, 2026
Merge #7499

- A predecessor that routes away from the JoinNode no longer counts as arrived; the join waits until every incoming edge has fired for the current iteration. Previously this was order-dependent (the join fired if the routed-away predecessor happened to complete first), so some existing graphs with conditional edges into a JoinNode will see the join stop firing.
- A node that completes without an output no longer leaves its previous iteration's output in node_outputs; downstream nodes (and the workflow's terminal output) now observe None for that run.
- Known limitation: if an iteration restarts before the JoinNode fires (for example, when one predecessor arrives at the join while another routes back to restart the loop), the partial arrival carries over and a later firing can combine outputs from different iterations.

Fixes #7498
Fixes #7515

PiperOrigin-RevId: 1001343757
@adk-bot

adk-bot commented Oct 17, 2026

Copy link
Copy Markdown
Collaborator

Thank you @AkshayShah03 for your contribution! 🎉

Your changes have been successfully imported and merged via Copybara in commit a6ed859.

Closing this PR as the changes are now in the main branch.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

merged [Status] This PR is merged

Projects

None yet

Development

Successfully merging this pull request may close these issues.

JoinNode in a loop fires early with the previous iteration's predecessor outputs, then fires again

3 participants