Skip to content

fix(workflow): avoid O(n^2) deep event comparison during rehydration - #6658

Open
lwangverizon wants to merge 1 commit into
google:mainfrom
lwangverizon:fix/rehydration-quadratic-scan
Open

fix(workflow): avoid O(n^2) deep event comparison during rehydration#6658
lwangverizon wants to merge 1 commit into
google:mainfrom
lwangverizon:fix/rehydration-quadratic-scan

Conversation

@lwangverizon

Copy link
Copy Markdown
Contributor

Link to Issue or Description of Change

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

Problem:

ReplayManager.get_events_for_rehydration merges top-level user prompts with a
value-based membership test over a list of pydantic models (_replay_manager.py:148):

user_prompts = [
    e for e in root_events if e.author == "user" and e not in node_events
]

e not in node_events invokes Event.__eq__ per probe — a recursive deep comparison of
the whole event including nested content/parts — so the merge costs
O(n_root x n_node x event_size).

Two problems follow:

  1. It blocks the event loop. get_events_for_rehydration and its caller
    _rehydrate_from_events (_dynamic_node_scheduler.py:333) are both plain def,
    reached from an async scheduler at _dynamic_node_scheduler.py:185 with no await
    in between — so this is synchronous CPU on the loop thread, stalling every other
    concurrent request on the process rather than only the invocation being rehydrated.

  2. It silently drops user prompts. Event.__eq__ is value-based, so a root-indexed
    user prompt that compares equal to an event already indexed under the node path is
    treated as already present and excluded. This is reachable through normal indexing:
    _build_event_index populates fc_to_parent incrementally, so a user
    function-response event arriving before its interrupt id is registered goes to
    root, while an equal-valued one arriving after is routed under the node path.

Eight lines below, the same method already keys membership by id() (line 156), so the
two halves of one function disagree on what membership means.

git log -S 'e not in node_events' dates the line to 64a7448f, "perf: Optimize
workflow rehydration performance with indexing — avoiding repeated linear scans of the
event log during rehydration"
. The quadratic scan arrived in the commit meant to remove
linear scans.

Solution:

Hoist {id(e) for e in node_events} above the comprehension and test id(e) not in it,
then reuse that set at line 156 instead of rebuilding it. Identity is the intended
semantics — both lists are views onto session.events, so an event is already indexed
under this node only if it is the same object.

This adopts the idiom already present in the method rather than introducing a new one, so
it is one changed line plus a comment explaining why identity (not value) is correct here
— the next reader would otherwise be tempted to "simplify" it back to in.

I deliberately did not restructure the indexing or change what gets merged; the goal
is to make this path do what it already intends, at linear cost.

Testing Plan

Unit Tests:

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

Two tests added to tests/unittests/workflow/utils/test_replay_manager.py, one per
symptom:

  • test_get_events_for_rehydration_merges_user_prompts_by_identity — a root-indexed
    user prompt that compares equal to a node-indexed event survives the merge, in session
    order. Built through the real indexing path, not by hand-populating the index.
  • test_get_events_for_rehydration_does_not_deep_compare_events — counts Event.__eq__
    invocations via a subclass and asserts zero. It also asserts the merge path actually
    ran, so eq_calls == 0 means "no deep comparison" rather than "nothing was compared".

Both fail on main and pass with the fix:

# fix reverted, tests kept
FAILED ...::test_get_events_for_rehydration_merges_user_prompts_by_identity
FAILED ...::test_get_events_for_rehydration_does_not_deep_compare_events
2 failed, 13 passed

# fix applied
15 passed

The second test's failure output is itself a readable statement of the bug — assert 20 == 0, i.e. 20 deep comparisons for 20 node events.

Suite results:

$ pytest tests/unittests/workflow -q
725 passed, 11 skipped, 10 xfailed

Wider sweep, compared against pristine main @ 4f583064 in a separate worktree
(tests/unittests/workflow, tests/unittests/agents, tests/unittests/test_runners.py):

# main @ 4f583064   2 failed, 1559 passed, 11 skipped, 12 xfailed, 10 errors
# this branch       2 failed, 1561 passed, 11 skipped, 12 xfailed, 10 errors   (+2 = the new tests)

The 2 failures and 10 errors are pre-existing on main and unrelated to this change; I
diffed the two FAILED lists and they are identical, so this branch introduces no
regressions. test_base_agent.py and test_output_key_visibility.py are excluded from
both runs — they fail to collect on main for a missing pytest_mock dependency.

pyink --check reports both changed files unchanged.

Manual End-to-End (E2E) Tests:

Ran the repro from #6657 against pristine main and against this branch. It constructs
Event objects directly and drives ReplayManager, so it needs no network, API key, or
model access. Part 1 checks the dropped-prompt behavior; part 2 scales n_root and
n_node together, which is how a multi-turn session grows both.

Before — main @ 4f583064:

PART 1 - behavior: is a value-equal user prompt preserved?
  e_user_root == e_user_routed : True
  e_user_root is e_user_routed : False
  session events               : 3
  events returned              : 2
  root user prompt preserved?  : False
  session order preserved?     : False
  RESULT: FAIL (prompt dropped)

PART 2 - performance
   n_root  n_node    merge time   vs previous
      100     100      16.17ms             -
      200     200      65.44ms          4.0x
      400     400     258.24ms          3.9x
      800     800    1039.47ms          4.0x

After — this branch:

PART 1 - behavior: is a value-equal user prompt preserved?
  session events               : 3
  events returned              : 3
  root user prompt preserved?  : True
  session order preserved?     : True
  RESULT: PASS

PART 2 - performance
   n_root  n_node    merge time   vs previous
      100     100       0.21ms             -
      200     200       0.38ms          1.8x
      400     400       0.93ms          2.4x
      800     800       1.95ms          2.1x

Scaling goes from ~4x per doubling to ~2x, and the n=800 case drops from 1039 ms to
1.95 ms
— about 533x. The absolute numbers are machine-specific; the exponent is the
part that matters.

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. — N/A, no dependent changes.

Additional context

Reported and root-caused by @wangge, who traced it from production faulthandler stacks
captured while a process was frozen — the blocking frame was pydantic/main.py __eq__
called from _replay_manager.py:148 via _rehydrate_from_events. Please credit them if
this lands.

Scope note: identity semantics slightly change which events the merge keeps, in the one
case where two distinct-but-equal events exist. I believe that direction is correct (the
current behavior loses a real user prompt) and the new test pins it, but flagging it
explicitly rather than presenting this as a pure no-op performance change — if you'd
prefer the perf fix without that behavior change, say so and I'll split it.

`ReplayManager.get_events_for_rehydration` merged top-level user prompts
with `e not in node_events`, a linear scan over a list of pydantic models.
Every probe invoked `Event.__eq__`, a recursive deep comparison of the whole
event including nested content and parts, so the merge cost scaled as
O(n_root x n_node x event_size).

The call chain is fully synchronous - `_rehydrate_from_events` and
`get_events_for_rehydration` are both plain `def`, reached from an `async`
scheduler - so this runs on the event loop thread with no `await`, and long
sessions can stall unrelated concurrent work.

Eight lines below, the same method already keys membership by `id()` when
rebuilding the chronological result. Use that idiom for the merge too, and
reuse the resulting set instead of rebuilding it.

This is also a behavior fix, not only a performance one: `Event.__eq__` is
value-based, so a root-indexed user prompt that happened to compare equal to
an event already indexed under the node path was silently dropped from
rehydration. Identity is the intended test - both lists are views onto
`session.events`.

Measured on real `Event` objects, doubling n quadruples the old cost:

  n=50   6.5ms -> 0.008ms     n=200  108.0ms -> 0.055ms
  n=100 27.2ms -> 0.020ms     n=400  445.8ms -> 0.125ms
@adk-bot adk-bot added the core [Component] This issue is related to the core interface and implementation label Aug 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core [Component] This issue is related to the core interface and implementation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Workflow rehydration does an O(n^2) deep event comparison that blocks the event loop and drops user prompts

3 participants