perf: summaries run deepest-first and start while expand is still deciding - #432
Open
rejojer wants to merge 18 commits into
Open
perf: summaries run deepest-first and start while expand is still deciding#432rejojer wants to merge 18 commits into
rejojer wants to merge 18 commits into
Conversation
summarize_tree visits the tree level by level (asyncio's ready queue is FIFO), so shallow leaves reach the semaphore first and the deepest leaves last — yet the deepest leaves are the ones still owed a chain of parent calls, which then run one after another with the pool nearly idle. On fed-2023 (182 calls) the 64-wide phase lasted ~22 s and the tail ~33 s, in-flight decaying 36→15→10→6→3→1. The semaphore becomes a priority gate: waiters are served deepest first (FIFO within a depth), a ready parent included. Depth counts the calls left on a node's path to the root, its own included, so this is longest-remaining-path-first scheduling; a prompt's content does not depend on admission order, so the summaries are unchanged. A permit granted to a waiter that is cancelled before it runs is passed on, as asyncio.Semaphore does; a cancelled waiter still queued is skipped at the next release. Fewer than one permit is refused at construction, as asyncio.Semaphore refused a negative count. Measured on the stored fed-2023 tree, FIFO and priority alternating: FIFO 62.3 / 57.1 s, priority 54.9 / 48.7 s (-13%); the landed code re-run 52.2 s. Mean start time by depth flipped from d3 7.8 s < d4 15.1 s < d5 19.5 s to d5 7.1 s < d4 10.4 s. Claude-Session: https://claude.ai/code/session_01MbNxisC33qUpuwq6ifvXwK
The priority gate orders only the calls that wait; the first SUMMARY_CONCURRENCY calls walk straight through in arrival order, and arrival order was the recursive visit's level-by-level expansion, shallow leaves first. So whenever the gate had free permits, the deep leaves whose parents the run ends on still started last. Now every node's task is created up front, deepest first: children always exist before their parent, which awaits their tasks instead of spawning them, so the deep leaves reach the gate first and a parent's own call no longer queues behind leaves that merely arrived earlier. Same 182 calls on fed, prompts unchanged. Summary stage at 64 wide, mirror-ordered: gate only 51.1/50.0 s, this 46.8/40.0 s; the landed code re-runs at 45.4 s against 57.4/58.7 s for the gate alone measured the same afternoon. Error handling is untouched: parents still gather their children with return_exceptions and re-raise the unrecoverable ones, roots the same. Claude-Session: https://claude.ai/code/session_01MbNxisC33qUpuwq6ifvXwK
… settle summarize_tree built every task up front from a finished tree, so summaries could only begin once expand had returned it. SummaryScheduler keeps the same machinery (one task per node, parents awaiting their children, the priority gate, deepest-first task creation) behind two calls: mark_final(nodes) says those nodes will not gain, lose or swap children and spawns tasks for their subtrees; finish() awaits the roots, applies the every-call-failed check and strips the bookkeeping keys. A node's task first awaits its own mark, so a subtree can be handed over while the rest of the tree is still being decided, and a node is never summarized twice however often it is marked: tasks are memoized per node and a parent composes whatever its children already wrote. summarize_tree marks the whole tree and finishes, which reproduces the old behavior exactly: same task order, same gate, same error semantics. The fed summary stage re-runs at 44.2 s against 45.4 s before. Claude-Session: https://claude.ai/code/session_01MbNxisC33qUpuwq6ifvXwK
page_index_flash ran optimize and the summaries as two asyncio.run calls, one after the other: the whole summary wave waited for the expand chain, a handful of dependent calls that take 18-67 s while the summary channels sat idle. Now the two share one loop. optimize(on_final=...) reports which nodes will not change any more, and the SummaryScheduler starts those at once; a parent's task is already waiting on its children, so it follows the moment its last child is done. A node is final when its children cannot change: a collapsed node under the trigger, a collapsed node expand has judged (frozen), or any node with children. That rule holds throughout the run, not only at the start. Expand touches only collapsed nodes over the trigger. The cost merge never fires on a surviving node after the first round: keeping X expanded means expand_cost < S(X), and tree_cost(X) right after attaching the children equals expand_cost term for term, only falling as children expand further; ancestors' tree_cost can only fall with it while their S stays put, because assign_ends pins the last child at the parent's old end. Same-page fusion was the one later mutation, and it is now done where the duplicates arise, right after a merge collapses a subtree onto a sibling's exact pages and right after expand attaches children, instead of at the next round's start, so no node waits a round for it. Reports go out after each round's merges, at each candidate's decision (kept collapsed, nothing found, or expanded together with what it grew), and for the whole tree at the end. Merge-only trees for the nine corpus PDFs are identical before and after; the only visible change is SpaceX ending after two rounds instead of a third that did nothing. Same-hour A/B, end to end: fed 97.9 s -> 72.6 s, PRML 174.3 s -> 136.8 s. Peak in flight is now the expand cap plus the summary cap. Claude-Session: https://claude.ai/code/session_01MbNxisC33qUpuwq6ifvXwK
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| gate.release() # granted to `waiter`, not yet resumed | ||
| waiter.cancel() | ||
| with pytest.raises(asyncio.CancelledError): | ||
| await waiter |
Member
Author
Code reviewNo issues found. Checked for bugs and CLAUDE.md compliance. 🤖 Generated with Claude Code - If this code review was useful, please react with 👍. Otherwise, react with 👎. |
… contract Four fixes from the PR #432 review round: - expand's in-flight same-page fusion now obeys do_merge: a caller who disabled merging (--no-merge, optimize_tree(do_merge=False)) keeps every proposed child and every document title. - _leaf_summary consults the tokenizer only when the length leaves "small" possible: no text averages 8+ chars per token, so long leaves skip the synchronous count that ran back-to-back on the loop before either lane's first model call (first-call latency measured 2.6 s -> ~30 ms on a 760-page synthetic, total 7.5 s -> 4.9 s). - finish() verifies the promise mark_final rests on (a final node stays in the tree and keeps its children) and refuses to wait on tasked nodes that were never marked: a broken promise now fails loud with a diagnosis instead of a silently wrong tree or an indefinite, symptomless hang. - the deepest-first test now pins the deep leaf against the shallow leaf document order would start first; deleting the sort in mark_final turns it red (mutation-verified). Claude-Session: https://claude.ai/code/session_01FtuLuuUh5f7Yy3W6Kx84xf
…e overlap Both still described the pre-overlap order: merge_same_page "runs first" / "before merge() ... before expand()" (it now also runs right after merge and on freshly attached children inside expand), and _optimize_async sitting "between extraction and summaries" (summaries now run alongside it). Claude-Session: https://claude.ai/code/session_01FtuLuuUh5f7Yy3W6Kx84xf
The settled promise (a final node is never touched again) rests on this inequality staying strict: expansion then strictly lowers tree_cost while spans never change, so merge's span <= cost can never newly hold in later rounds. At <=, a zero-gain expand would create a node merge folds right back. finish()'s contract check catches a break at runtime; this line warns before the edit. Claude-Session: https://claude.ai/code/session_01FtuLuuUh5f7Yy3W6Kx84xf
A same-page fusion during expansion bumped the run counters while the round entry still said same_page=False. The flag is now derived from the round's own log slice, so it counts every fusion and cannot disagree with the counters. Loop exit is unchanged: an in-expand fusion only happens after an attach, which already sets expanded. Claude-Session: https://claude.ai/code/session_01FtuLuuUh5f7Yy3W6Kx84xf
Member
Author
Code reviewNo issues found. Checked for bugs and CLAUDE.md compliance. 🤖 Generated with Claude Code - If this code review was useful, please react with 👍. Otherwise, react with 👎. |
The length gate bb6ee89 put in front of count_tokens rested on "no text averages 8+ chars per token", which is false: dot-leader contents pages run 9-10 chars per token, rule and whitespace runs 60-120. Because the two tests were and-ed, a leaf over 1600 chars but under the 200-token floor paid a model call where it was documented to reuse its raw text. The 2.6 s the prefilter was measured to save was import litellm, paid by the first count_tokens in a harness whose faked model never imported it. On the nine example PDFs with litellm already loaded, as a real run's first model call leaves it, the prefilter moves expand's first call by 10-200 ms per document (PRML 0.36 s -> 0.17 s, about 0.1% of the run). The tokenizer now takes the summary model like every other count_tokens site, instead of a hardcoded gpt-4o. litellm falls back to cl100k_base for names tiktoken does not know and swallows HF tokenizer failures, so no model name raises; names outside the gpt-4o family count CJK text about 2x higher, which narrows the raw-text floor for those runs.
- finish()'s tasked-but-unmarked check ran over self._tasks before the root tasks existed, so a root that never received mark_final slipped past it and finish() parked forever - the symptomless hang the check was added to prevent. It now counts every node in the tree. - _optimize_and_summarize took the expand model as `model`, the name the neighbouring _summarize uses for the summary model, with both passed positionally as adjacent strings. Renamed to optimize_model and called by keyword; the overlap test now runs the two lanes on different names and pins which model each receives. - expand's keep_collapsed settle had no coverage: the existing test exits through no_children. A second test forces the cost-rejected branch with min_gain_ratio and checks the node is reported final before the closing sweep. - merge_same_page's docstrings claimed every seam; merge_tree() is a bare merge(). They now name optimize()'s seams. Claude-Session: https://claude.ai/code/session_018aNxBG6NTxK7hVwFCJK7Cx
Member
Author
Code reviewNo issues found. Checked for bugs and CLAUDE.md compliance. 🤖 Generated with Claude Code - If this code review was useful, please react with 👍. Otherwise, react with 👎. |
…ords
The two summary prompts asked for a "points" list next to the summary and
parse_summary threw the list away: about 70% of every reply was generated
to be discarded. Dropping the list alone is not enough. The model then
pours the same content into the summary (656 -> ~2000 chars) and parents,
which read their children's summaries, slow down more than the leaves
gained (fed wall +6%). With a word cap the reply shrinks for real.
Measured on gpt-5.6-luna, mirror A/B, summary stage only:
per-call median 9.7 s -> 5.3 s (-45%)
fed-2023 wall 47.5 s -> 30.7 s (-35%)
PRML wall 71.1 s -> 38.1 s (-46%)
reply chars ~3340 -> ~1175 (-65% output tokens)
summary length ~670 -> ~1160 chars (the model writes cap x 1.25)
Blinded pairwise judge (claude-sonnet-5, source in view): 150-word
summaries beat the current ones 21-1-0 on fed + attention leaves. The
specifics that used to sit in the discarded list (law names, rates, terms)
now land in the summary. A full prompt rewrite lost to this two-line edit
8-11-3, so the edit stays minimal.
The cap is a parameter, default 150: page_index_flash(summary_max_words=),
PageIndexLocalClient(summary_max_words=) / index={"summary_max_words": n},
and run_pageindex.py --summary-max-words. The classic pipeline's summary
prompt is a different one without a points field and is untouched.
The CLI could set all three; the SDK could not. They join the local index
side the way summary_max_words did: flat arguments on both clients, keys
in the index= slot, LocalAPI hands them to page_index_flash. optimize
takes the CLI's words ("full" / "merge" / "off"; "off" is no optimize
pass at all), so a wrong value refuses in the constructor like every
other slot. The slot check's empty-value rule now lets False through,
since use_embedded_toc is the first bool in the vocabulary.
summary_concurrency also gets a CLI flag.
Lowering summary_concurrency for a tight quota used to lower only the summary lane: expand kept its own 32 and still hit the API in bursts. The one knob now bounds both lanes: expand's gate is min(32, the cap), so 8 means 8 in each lane (the lanes overlap, so up to cap + min(32, cap) calls are in flight), while raising it past 32 leaves expand at its measured plateau. Defaults are unchanged (64 and 32).
Belongs with 06f9c87; the file lost the test on disk between the test run and that commit.
- max_words goes after concurrency in summarize_tree and SummaryScheduler: concurrency shipped in 0.2.11 as the sixth positional of summarize_tree, so inserting before it would have turned a positional concurrency into a word cap. - _resolve_index_slot maps slot keys with a two-entry rename instead of a table that repeated every local key; conf is annotated dict[str, Any] so the comprehension types cleanly. - run_pageindex.py checks the flash-only flags in one loop, like the standard-only ones. - summary_concurrency's docs say what actually runs: per lane, and the lanes overlap, so up to cap + min(32, cap) calls at once. - test_summary_concurrency_caps_expand_too keeps an uncapped control run (peak 3) so the capped assertion cannot pass by accident; the knob test uses the default store under tmp_path instead of per-run directories. - page_index_flash's docstring lists summary_max_words in signature order.
…e order is tested - count_tokens(model=...) picks the model's tokenizer through litellm, and the HuggingFace tokenizers it selects for llama-family names reject a lone surrogate with TypeError where tiktoken counts it. _leaf_summary's raw-text gate runs inside _visit's except, so that raise blanked the leaf's summary silently. On any failure count_tokens now re-counts with litellm's bundled default encoding (cl100k: offline, surrogate-safe). - get_page_tokens (both parsers) and _index_standard route through count_tokens instead of calling litellm.token_counter themselves, so the standard path shares the fallback; a page whose extract_text() is None now counts 0 instead of raising. - expand settles a node only after merge_same_page has fused its new children: mark_final snapshots the children and finish() rejects a later change. The order was right but nothing held it; the fusion test now drives a real SummaryScheduler through on_final and awaits finish(), so swapping the two lines fails it.
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
1 participant
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.
Flash indexing spends most of its wall time in summaries, and until now that stage waited for expand to finish and then ran its calls in whatever order the tree recursion produced. This branch makes the summary stage run deepest node first and start while expand is still deciding, so the LLM channels never sit idle waiting on the expand chain.
What changes
_PriorityGate: the summary semaphore admits the queued call with the most work still above it (depth = calls left on the node's path to the root, its own included), FIFO within a depth. Cancellation-safe likeasyncio.Semaphore.summarize_treebecomes a thin wrapper overSummaryScheduler:mark_final(nodes)says those nodes will not gain, lose or swap children and starts their subtrees;finish()awaits the roots. Same task order, gate and error semantics as before.optimize(on_final=...)reports which nodes are final as it goes: after each round's merges, at each expand candidate's decision (together with what it grew), and for the whole tree at the end. A node is final when it is collapsed under the trigger, collapsed and already judged by expand, or has children — the cost merge cannot fire on a surviving node after the first round (see the commit message for the argument).page_index_flashruns expand and summaries on one event loop when both are on; every other combination keeps the old path.Measured (same hour, end to end via
submit_document)Summary-stage only (fed, 182 calls, 64 wide): FIFO 58–62 s → gate 50–57 s → gate + deepest-first 45 s.
Same calls, same prompts; outputs are order-independent. Peak in flight is now the expand cap plus the summary cap (32 + 64).
Tests: 459 passing, including the ordering, cancellation, scheduler, final-node reporting, immediate-fusion and one-loop overlap cases; green on the without-frameworks leg and on Python 3.10 / 3.13.
Summary prompt and indexing knobs
The summary prompts no longer ask for the
pointslist thatparse_summarydiscarded, and cap the summary atsummary_max_words(default 150). Measured on gpt-5.6-luna, mirror A/B, summary stage only: per-call latency 9.7 → 5.3 s (−45%), fed-2023 47.5 → 30.7 s (−35%), PRML 71.1 → 38.1 s (−46%), output tokens −65%. Summaries come out ~1160 chars instead of ~670 and carry the specifics that used to sit in the discarded list; a blinded pairwise judge (claude-sonnet-5, source in view) prefers them 21-1-0 over the old ones. Deleting the list without a cap is not enough: the model then pours it into the summary (3× longer) and parents slow down more than the leaves gain.Four indexing knobs are settable from the SDK (flat arguments or the
index=slot) and the CLI:summary_max_words,summary_concurrency,use_embedded_toc,optimize("full"/"merge"/"off").summary_concurrencybounds both lanes: expand's gate becomes min(32, the cap), so one knob lowers the whole indexing lane on a tight quota (the lanes overlap, so up to cap + min(32, cap) calls run at once). Defaults are unchanged.https://claude.ai/code/session_01MbNxisC33qUpuwq6ifvXwK