Skip to content

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

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

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

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

2 participants