fix(workflow): avoid O(n^2) deep event comparison during rehydration - #6658
Open
lwangverizon wants to merge 1 commit into
Open
fix(workflow): avoid O(n^2) deep event comparison during rehydration#6658lwangverizon wants to merge 1 commit into
lwangverizon wants to merge 1 commit into
Conversation
`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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
Problem:
ReplayManager.get_events_for_rehydrationmerges top-level user prompts with avalue-based membership test over a list of pydantic models (
_replay_manager.py:148):e not in node_eventsinvokesEvent.__eq__per probe — a recursive deep comparison ofthe whole event including nested
content/parts— so the merge costsO(n_root x n_node x event_size).Two problems follow:
It blocks the event loop.
get_events_for_rehydrationand its caller_rehydrate_from_events(_dynamic_node_scheduler.py:333) are both plaindef,reached from an
asyncscheduler at_dynamic_node_scheduler.py:185with noawaitin 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.
It silently drops user prompts.
Event.__eq__is value-based, so a root-indexeduser 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_indexpopulatesfc_to_parentincrementally, so a userfunction-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 thetwo halves of one function disagree on what membership means.
git log -S 'e not in node_events'dates the line to64a7448f, "perf: Optimizeworkflow 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 testid(e) not init,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 indexedunder 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:
Two tests added to
tests/unittests/workflow/utils/test_replay_manager.py, one persymptom:
test_get_events_for_rehydration_merges_user_prompts_by_identity— a root-indexeduser 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— countsEvent.__eq__invocations via a subclass and asserts zero. It also asserts the merge path actually
ran, so
eq_calls == 0means "no deep comparison" rather than "nothing was compared".Both fail on
mainand pass with the fix: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:
Wider sweep, compared against pristine
main@4f583064in a separate worktree(
tests/unittests/workflow,tests/unittests/agents,tests/unittests/test_runners.py):The 2 failures and 10 errors are pre-existing on
mainand unrelated to this change; Idiffed the two
FAILEDlists and they are identical, so this branch introduces noregressions.
test_base_agent.pyandtest_output_key_visibility.pyare excluded fromboth runs — they fail to collect on
mainfor a missingpytest_mockdependency.pyink --checkreports both changed files unchanged.Manual End-to-End (E2E) Tests:
Ran the repro from #6657 against pristine
mainand against this branch. It constructsEventobjects directly and drivesReplayManager, so it needs no network, API key, ormodel access. Part 1 checks the dropped-prompt behavior; part 2 scales
n_rootandn_nodetogether, which is how a multi-turn session grows both.Before —
main@4f583064:After — this branch:
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
Additional context
Reported and root-caused by @wangge, who traced it from production
faulthandlerstackscaptured while a process was frozen — the blocking frame was
pydantic/main.py __eq__called from
_replay_manager.py:148via_rehydrate_from_events. Please credit them ifthis 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.