Skip to content

Commit adc2433

Browse files
inamdarmihircopybara-github
authored andcommitted
fix: exclude replayed nodes from workflow max_concurrency count
Merge #7758 Nodes fast-forwarded from session history on resume no longer take a `max_concurrency` slot, so a replayed node waiting for its recorded turn cannot block the node it is waiting for. Nodes that actually execute are still limited. A recorded agent transfer whose interrupts are answered starts once its recorded turn comes and then takes a slot. Fixes #7618 PiperOrigin-RevId: 1009119531
1 parent 82fe029 commit adc2433

3 files changed

Lines changed: 490 additions & 13 deletions

File tree

‎.agents/skills/adk-architecture/references/interface-workflow.md‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -376,6 +376,13 @@ stages:
376376
a checkpoint if a replayed node comes back `WAITING`) and seeding
377377
`loop_state.runs[node_path]` with
378378
`DynamicNodeRun(state=NodeState(run_id=run_id), recovered_state=recovered, is_static=True)`.
379+
`_at_concurrency_limit` excludes `replayed_nodes` from the
380+
`max_concurrency` count. A resolved-transfer fast-forward is not in that
381+
set because it runs its transfer target once the replay barrier releases
382+
it, so under `max_concurrency` `_schedule_ready_nodes` keeps its trigger
383+
buffered (`_is_held_transfer`) until its barrier key is released and then
384+
counts it like any running node. If nothing else is pending, it starts
385+
one held node anyway so the barrier reports the divergence.
379386
- `_start_node_task` then dispatches `ctx._run_node_internal` to
380387
`DynamicNodeScheduler._check_existing_run`, which runs
381388
`check_interception(node=curr_node, recovered=run.recovered_state, current_run=run)`:

‎src/google/adk/workflow/_workflow.py‎

Lines changed: 72 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -528,6 +528,7 @@ async def _schedule_ready_nodes(
528528
# loop_state.trigger_buffer is a dict, and Python 3.7+ dicts preserve insertion order.
529529
# Therefore, nodes are processed strictly in the order their triggers arrived,
530530
# ensuring deterministic scheduling order for parallel branches.
531+
held: list[str] = []
531532
for node_name in list(loop_state.trigger_buffer.keys()):
532533
if node_name in loop_state.pending_tasks:
533534
continue
@@ -544,26 +545,84 @@ async def _schedule_ready_nodes(
544545
if node_state.status == NodeStatus.WAITING and node_state.interrupts:
545546
continue
546547

548+
if self._is_held_transfer(loop_state, ctx, node_name):
549+
held.append(node_name)
550+
continue
551+
547552
if self._at_concurrency_limit(loop_state):
548553
break
549554

550-
trigger = self._pop_trigger(loop_state, node_name)
551-
if trigger is None:
552-
continue
555+
await self._start_buffered_node(loop_state, ctx, node_name)
553556

554-
self._prepare_node_state_for_starting(loop_state, node_name, trigger)
555-
started = self._start_node_task(loop_state, ctx, node_name, trigger)
556-
# Only a genuinely-executing node marks a new step; a replayed node
557-
# being fast-forwarded from history should not re-announce itself.
558-
if started:
559-
await self._emit_node_checkpoint(loop_state, ctx)
557+
# With nothing else in flight, nothing can release a held transfer, so
558+
# start one and let the replay barrier report the divergence.
559+
if held and not loop_state.pending_tasks:
560+
await self._start_buffered_node(loop_state, ctx, held[0])
561+
562+
async def _start_buffered_node(
563+
self, loop_state: _LoopState, ctx: Context, node_name: str
564+
) -> None:
565+
"""Pop the next trigger for a node and start it."""
566+
trigger = self._pop_trigger(loop_state, node_name)
567+
if trigger is None:
568+
return
569+
570+
self._prepare_node_state_for_starting(loop_state, node_name, trigger)
571+
started = self._start_node_task(loop_state, ctx, node_name, trigger)
572+
# Only a genuinely-executing node marks a new step; a replayed node
573+
# being fast-forwarded from history should not re-announce itself.
574+
if started:
575+
await self._emit_node_checkpoint(loop_state, ctx)
560576

561577
def _at_concurrency_limit(self, loop_state: _LoopState) -> bool:
562-
"""Check if max_concurrency has been reached."""
578+
"""Check if max_concurrency has been reached.
579+
580+
Replayed nodes are history fast-forwards that may block on the replay
581+
sequence barrier waiting for another node to be scheduled, so they must
582+
not hold a concurrency slot.
583+
"""
584+
if self.max_concurrency is None or self.max_concurrency <= 0:
585+
return False
586+
active = sum(
587+
1
588+
for name in loop_state.pending_tasks
589+
if name not in loop_state.replayed_nodes
590+
)
591+
return active >= self.max_concurrency
592+
593+
def _is_held_transfer(
594+
self, loop_state: _LoopState, ctx: Context, node_name: str
595+
) -> bool:
596+
"""Check if a node must wait for its replay turn before starting.
597+
598+
A recorded transfer with resolved interrupts runs its transfer target once
599+
the replay sequence barrier releases it. Under max_concurrency it is
600+
started only after that release and counted like any other node once it
601+
runs. The one exception: with nothing else pending, nothing can release
602+
it, so _schedule_ready_nodes starts one anyway and the barrier reports the
603+
divergence.
604+
"""
605+
barrier = loop_state.replay_manager.sequence_barrier
606+
if (
607+
self.max_concurrency is None
608+
or self.max_concurrency <= 0
609+
or barrier is None
610+
):
611+
return False
612+
counters = loop_state.run_counters.get(ctx.node_path or "", {})
613+
key = f"{node_name}@{counters.get(node_name, 0) + 1}"
614+
event = barrier.events.get(key)
615+
recovered = loop_state.recovered_executions.get(key)
616+
if event is None or event.is_set() or recovered is None:
617+
return False
618+
result = check_interception(
619+
node=self._get_static_node_by_name(node_name), recovered=recovered
620+
)
563621
return (
564-
self.max_concurrency is not None
565-
and self.max_concurrency > 0
566-
and len(loop_state.pending_tasks) >= self.max_concurrency
622+
not result.should_run
623+
and result.transfer_to_agent is not None
624+
and bool(recovered.interrupt_ids)
625+
and not recovered.interrupt_ids - recovered.resolved_ids
567626
)
568627

569628
def _pop_trigger(

0 commit comments

Comments
 (0)