refa: optimize wiki compilation concurrency#17192
Conversation
📝 WalkthroughWalkthroughEmbedding deduplication now uses bounded normalized cosine comparisons and concurrent ambiguity resolution. ChangesWiki compilation pipeline
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant run_wiki
participant MAPQueue
participant MAPWorker
participant LLMCallPool
participant wiki_map_from_chunks
run_wiki->>MAPQueue: enqueue chunk batches
MAPQueue->>MAPWorker: deliver batch jobs
MAPWorker->>LLMCallPool: submit pooled MAP call
LLMCallPool->>wiki_map_from_chunks: admit bounded LLM work
wiki_map_from_chunks-->>MAPWorker: return batch results
MAPWorker-->>run_wiki: aggregate document statistics
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
dosubot
Bot
added
size:L
wangq8
left a comment
wangq8
left a comment
There was a problem hiding this comment.
Review: refa: optimize wiki compilation concurrency (#17192)
Thanks for this PR — the producer-consumer pattern and shared LLM pool are a solid architectural improvement. I left detailed notes below. The main callout is the concurrency explosion risk from max_workers=0 and the lost pre-filtering in disambiguation.
1. (Important) max_workers=0 can create unbounded concurrent tasks
In dataset_wiki_generator.py the MAP worker now calls wiki_map_from_chunks(..., max_workers=0). When max_workers is 0, _run_chunked_pipeline creates all sub-batches as concurrent asyncio tasks with no semaphore:
# _common.py, run_chunked_pipeline
semaphore = asyncio.Semaphore(max_workers) if max_workers and max_workers > 0 else NoneIf split_chunks produces e.g. 20 sub-batches per batch, and 20 workers are active, that is up to 400 asyncio tasks launched simultaneously. The shared LLMCallPool (max_pending=25) limits the actual LLM calls, but having hundreds of tasks pending is wasteful and risks excessive memory pressure from all the suspended coroutine frames.
Suggestion: Either set a reasonable non-zero max_workers (e.g. 6 like the default) or add an explicit bound inside _run_chunked_pipeline for the max_workers=0 path. At minimum, add a comment documenting the reasoning so future readers are not surprised.
2. Disambiguation pre-filtering is lost
Old code in _resolve_ambiguous_pairs:
batch = [(i, j) for i, j in batch if _root(i) != _root(j)]This filtered pairs that were already redundant before firing the LLM call. The new code runs all batches concurrently via asyncio.gather and only filters at apply time. This means the LLM may be asked to judge pairs that will already be merged by an earlier-applied batch result.
This is a functional tradeoff (correctness is preserved by the final _root(i) != _root(j) guard in the apply loop), but it wastes LLM calls and tokens. Consider adding the pre-filtering back after gathering results, or at least batching the filtering with the concurrent dispatch.
3. Blockwise embedding dedup looks correct but worth verifying edge case
The _embedding_dedup change from sklearn.metrics.pairwise.cosine_similarity to blockwise matrix multiply is a good memory optimization. One thing to double-check in testing:
The DIY normalization:
norms = np.linalg.norm(matrix, axis=1, keepdims=True)
matrix = np.divide(matrix, norms, out=np.zeros_like(matrix), where=norms > 0)The out=np.zeros_like(matrix) means zero-norm vectors become all-zero rows. Their cosine similarity with any other vector is then 0 (well below both thresholds). This is consistent with the old sklearn behavior (which would return NaN/0 for zero-vectors), but worth an explicit test case to confirm.
4. Minor: WIKI_MAP_QUEUE_SIZE = 5 may be too conservative
With 20 workers and potentially large documents, a queue depth of 5 means producers are throttled aggressively. If a document has many batches, the producer will put→block→put→block repeatedly, reducing throughput. Consider raising this to 2 * WIKI_MAP_LLM_POOL_SIZE (40) or similar, since the pool is the true concurrency limiter and the queue is just a buffer.
5. LLMCallPool pending gate is correct
The max_pending bound in LLMCallPool.call() gates on pending_count >= max_pending using async Condition.wait(). This correctly back-pressures without deadlocking since callers are pure async (not holding locks). ✅
6. Docs are solid
The docstring updates in _embedding_dedup, _resolve_ambiguous_pairs, and bulk_dedup_items clearly explain the changed semantics. The comments in dataset_wiki_generator.py about the producer-consumer design are helpful. ✅
Summary
The overall architecture is sound: shared pool bounding actual LLM calls, async queue for producer-consumer MAP, concurrent REFINE. My main concerns are (1) the unbounded task explosion from max_workers=0 and (2) the lost disambiguation pre-filtering. Neither is a correctness bug, but both could hurt production performance/latency.
Consider addressing #1 and #2 before merging. The rest looks good.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rag/advanced_rag/knowlege_compile/_common.py`:
- Around line 694-703: Optimize the blockwise comparison around the
upper-triangle branch by applying the ambiguous_low similarity mask to the
upper-triangle indices before iterating, so only relevant pairs reach the Python
loop while preserving existing merge and ambiguous classification. Update the
rows/cols iteration in this block to use zip(..., strict=True).
- Around line 787-790: Update the batch resolution flow around _resolve_batch
and asyncio.gather to use a local semaphore that limits concurrent batch tasks
to the available LLM capacity. Acquire the semaphore before creating or starting
each timed resolution, release it after completion, and preserve the existing
handling of None decisions so deduplications are not silently dropped.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 6dd7ff91-5657-4029-8e4f-d29932ae8e75
📒 Files selected for processing (3)
rag/advanced_rag/knowlege_compile/_common.pyrag/advanced_rag/knowlege_compile/structure.pyrag/svr/task_executor_refactor/dataset_wiki_generator.py
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
rag/svr/task_executor_refactor/dataset_wiki_generator.py (2)
875-880: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winCancel background workers to prevent deadlocks and orphaned tasks.
If
await asyncio.gather(*producers)orawait map_queue.join()raises an exception (or is cancelled), the function exits and leaves orphanedworkersrunning in the background, causing resource leaks.
Additionally, if any worker dies prematurely,await map_queue.put(None)can hang indefinitely because there aren't enough alive workers to consume the poison pills, deadlocking the cleanup phase.Since
await map_queue.join()guarantees all work is done and workers are idle, you can safely and robustly shut them down by cancelling them in afinallyblock, removing the need for poison pills entirely.🛡️ Proposed robust cleanup
- await asyncio.gather(*producers) - await map_queue.join() - for _ in workers: - await map_queue.put(None) - await asyncio.gather(*workers) + try: + await asyncio.gather(*producers) + await map_queue.join() + finally: + for w in workers: + if not w.done(): + w.cancel() + # Await the cancelled workers to suppress "Task was destroyed but it is pending" warnings + await asyncio.gather(*workers, return_exceptions=True)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rag/svr/task_executor_refactor/dataset_wiki_generator.py` around lines 875 - 880, Update the producer/worker coordination around asyncio.gather and map_queue.join to use a finally block that cancels every task in workers and awaits their completion, including cancellation/error paths. Remove the poison-pill map_queue.put(None) shutdown loop, while preserving normal producer completion and queue draining behavior.
832-834: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMove item unpacking inside the
tryblock to prevent silent worker crashes and deadlocks.If an exception occurs during unpacking or dictionary access (e.g., if
docis missing an"id"), the worker will crash without executing thefinallyblock. This preventsmap_queue.task_done()from being called, which causesawait map_queue.join()to hang indefinitely and deadlock the entire pipeline.🛡️ Proposed fix to ensure queue completion
- i, doc, template_id, parser_cfg, batch, batch_no = item - doc_id = doc["id"] - stats = doc_stats[i] try: + i, doc, template_id, parser_cfg, batch, batch_no = item + doc_id = doc["id"] + stats = doc_stats[i] map_llm_id = (parser_cfg.get("llm_id") or "").strip() if isinstance(parser_cfg, dict) else ""🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rag/svr/task_executor_refactor/dataset_wiki_generator.py` around lines 832 - 834, Move the item unpacking and subsequent doc access in the worker flow, including the assignments to i, doc, template_id, parser_cfg, batch, batch_no, doc_id, and stats, inside the existing try block so exceptions still reach the finally block and map_queue.task_done() is always called.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rag/advanced_rag/knowlege_compile/_common.py`:
- Around line 502-504: Preserve 1:1 batch-to-result cardinality in the batch
processing flow: in rag/advanced_rag/knowlege_compile/_common.py lines 502-504,
update the loop around process_batch to enqueue every entries value, including
empty batches; in lines 521-522, remove the early return for batches containing
no entries so process_batch produces the corresponding safe defaults. Both sites
require direct changes.
---
Outside diff comments:
In `@rag/svr/task_executor_refactor/dataset_wiki_generator.py`:
- Around line 875-880: Update the producer/worker coordination around
asyncio.gather and map_queue.join to use a finally block that cancels every task
in workers and awaits their completion, including cancellation/error paths.
Remove the poison-pill map_queue.put(None) shutdown loop, while preserving
normal producer completion and queue draining behavior.
- Around line 832-834: Move the item unpacking and subsequent doc access in the
worker flow, including the assignments to i, doc, template_id, parser_cfg,
batch, batch_no, doc_id, and stats, inside the existing try block so exceptions
still reach the finally block and map_queue.task_done() is always called.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 2add3c4a-1aa8-4814-be25-f188e083223f
📒 Files selected for processing (2)
rag/advanced_rag/knowlege_compile/_common.pyrag/svr/task_executor_refactor/dataset_wiki_generator.py
wangq8
left a comment
wangq8
left a comment
There was a problem hiding this comment.
PR Review: refa: optimize wiki compilation concurrency
Thanks for this PR. The overall direction — shared LLM pool, blockwise embedding dedup, and concurrent disambiguation / REFINE — is solid and should meaningfully improve Wiki compilation throughput. A few observations:
✅ What looks good
-
run_chunked_pipeline— Producer-consumer pattern viaasyncio.Queueis cleaner than the oldSemaphore+gatherapproach. It avoids creating all tasks upfront and the sentinel-based worker termination is correct. Exception propagation is also handled properly (worker crash →gatherraises → cancel + re-raise). -
Blockwise cosine similarity — Replacing
sklearn.cosine_similarity(N×N materialization) with blockwisematrix @ block.Tis a strong improvement for large KBs. The grouping bytype_keyand chunk-by-chunktriu_indices/np.nonzerologic correctly covers all pairs without duplication. -
LLMCallPoolbackpressure — Addingmax_pendingwith thepending_countguard keeps queue growth bounded. Since the existing code usesnotify_all()everywhere, no deadlock risk; thewhile pending_count >= max_pendingre-check after wakeup is correct. -
Concurrent
_resolve_ambiguous_pairs— Deferring union-find mutations to a single serial pass afterasyncio.gatheravoids race conditions onmerged_into. The docstring clearly explains the intent. -
Template resolution upfront — Moving template resolution before the MAP worker pool avoids subtle ordering dependencies on which document happens to finish first.
⚠️ Items worth addressing
1. WIKI_MAP_QUEUE_SIZE = 5 with 20 workers may starve workers
With 20 map workers but a queue capacity of only 5 batches, workers will frequently block on map_queue.get(). Since each batch triggers expensive LLM calls, this may be acceptable in practice, but the ratio feels off. Consider at least matching WIKI_MAP_LLM_POOL_SIZE (20), or document why 5 was chosen (e.g. ES fetch latency is the bottleneck, not the queue).
2. _resolve_ambiguous_pairs sends all pairs to LLM without pre-filtering
The old code filtered _root(i) != _root(j) before each batch's LLM call, avoiding wasted LLM queries for already-merged transitive pairs. The new concurrent approach can't do that pre-filter (no ordering guarantee), which is correct — but it means every ambiguous pair always triggers an LLM call, even if transitive merges in concurrent batches would have made it redundant. This is a compute-cost trade-off vs. latency. Worth a comment in the code acknowledging it.
3. _embedding_dedup normalization change for zero-magnitude vectors
np.divide(matrix, norms, out=np.zeros_like(matrix), where=norms > 0) produces zero vectors for zero-magnitude embeddings, which then yield similarity 0. The old cosine_similarity would produce NaN. In practice zero embeddings are rare, so this is low risk, but the behavioral diff is worth noting.
4. Progress granularity reduced
Progress is now reported only when each document starts (inside _produce_document), not per-batch. For documents with many batches, the progress bar will stall at a single value for a long time. Consider passing per-batch progress info back from workers via a shared counter, or accept the coarser granularity (MAP is only ~60% of the total progress range anyway).
5. No test additions
This PR touches 3 files with significant concurrency / algorithmic changes and no test coverage. At minimum, adding a basic test for the LLMCallPool max_pending backpressure behavior and the blockwise dedup logic would increase confidence.
🔍 Minor / style
- The
_call_modelinner function inside_resolve_batchjust wrapsasyncio.wait_for(gen_json(...), timeout=llm_timeout)— it could be inlined. persist_wiki_pages_to_es/persist_wiki_page_graph_to_esswitch to keyword args: cosmetic, parameter names match, so it's compatible._pool_managed = TrueonPooledChatModelappears unused in this PR — is it for a follow-up?
Verdict: Approve with suggestions. The core design is sound; the queue sizing and the ambiguity pre-filtering trade-off are the main things to discuss or document.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #17192 +/- ##
==========================================
- Coverage 94.56% 90.65% -3.91%
==========================================
Files 10 10
Lines 717 717
Branches 118 118
==========================================
- Hits 678 650 -28
- Misses 25 39 +14
- Partials 14 28 +14 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rag/svr/task_executor_refactor/dataset_wiki_generator.py (1)
767-777: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winRemove the unused
template_idfrom MAP job tuples.After resolving
parser_cfg,template_idis never consumed; Ruff reports RUF059 at Line 832. Remove it fromresolved_eligible, producer arguments, and queue items instead of carrying a dead field.Proposed fix
- resolved_eligible: list[tuple[dict, str, dict]] = [] + resolved_eligible: list[tuple[dict, dict]] = [] ... - resolved_eligible.append((doc, template_id, template.get("config") or {})) + resolved_eligible.append((doc, template.get("config") or {})) ... - for i, (doc, _, _) in enumerate(resolved_eligible) + for i, (doc, _) in enumerate(resolved_eligible) ... - async def _produce_document(i: int, job: tuple[dict, str, dict]) -> None: - doc, template_id, parser_cfg = job + async def _produce_document(i: int, job: tuple[dict, dict]) -> None: + doc, parser_cfg = job ... - await map_queue.put((i, doc, template_id, parser_cfg, batch, stats["batch_count"])) + await map_queue.put((i, doc, parser_cfg, batch, stats["batch_count"])) ... - i, doc, template_id, parser_cfg, batch, batch_no = item + i, doc, parser_cfg, batch, batch_no = itemAlso applies to: 807-832
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rag/svr/task_executor_refactor/dataset_wiki_generator.py` around lines 767 - 777, Remove the unused template_id field from the resolved_eligible tuple and all downstream MAP job data. Update the producer logic around CompilationTemplateService.get_saved and the queue-item construction/consumption near the parser_cfg handling so tuples and arguments contain only the values that are actually used, including adjusting unpacking consistently.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@rag/svr/task_executor_refactor/dataset_wiki_generator.py`:
- Around line 767-777: Remove the unused template_id field from the
resolved_eligible tuple and all downstream MAP job data. Update the producer
logic around CompilationTemplateService.get_saved and the queue-item
construction/consumption near the parser_cfg handling so tuples and arguments
contain only the values that are actually used, including adjusting unpacking
consistently.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 34828ca2-3b24-43e1-ac49-2615a7e81cfe
📒 Files selected for processing (3)
rag/advanced_rag/knowlege_compile/_common.pyrag/advanced_rag/knowlege_compile/structure.pyrag/svr/task_executor_refactor/dataset_wiki_generator.py
💤 Files with no reviewable changes (1)
- rag/advanced_rag/knowlege_compile/structure.py
🚧 Files skipped from review as they are similar to previous changes (1)
- rag/advanced_rag/knowlege_compile/_common.py
3 participants
What problem does this PR solve?
Improve Wiki compilation throughput while keeping LLM concurrency and memory usage bounded.
Type of change