refa: optimize wiki compilation concurrency (#17192)

rag/advanced_rag/knowlege_compile/_common.py

@@ -493,27 +493,39 @@ async def run_chunked_pipeline(
         return aggregate([]) if aggregate else []
 
     total = len(batches)
-    semaphore = asyncio.Semaphore(max_workers) if max_workers and max_workers > 0 else None
+    worker_count = max_workers if max_workers and max_workers > 0 else 6
+    work_queue: asyncio.Queue[tuple[int, list[dict]] | None] = asyncio.Queue(maxsize=worker_count)
+    results: list[Any] = [None] * total
+    completed: list[bool] = [False] * total
 
-    async def _one(idx: int, entries: list[dict]) -> Any:
-        async def _do() -> Any:
-            return await process_batch(entries, idx, total)
+    async def _producer() -> None:
+        for idx, entries in enumerate(batches):
+            await work_queue.put((idx, entries))
+        for _ in range(worker_count):
+            await work_queue.put(None)
 
-        if semaphore is not None:
-            async with semaphore:
-                return await _do()
-        return await _do()
-
-    tasks = [asyncio.create_task(_one(i, b)) for i, b in enumerate(batches) if b]
-    if not tasks:
-        return aggregate([]) if aggregate else []
+    async def _worker() -> None:
+        while True:
+            item = await work_queue.get()
+            if item is None:
+                work_queue.task_done()
+                return
+            idx, entries = item
+            try:
+                results[idx] = await process_batch(entries, idx, total)
+                completed[idx] = True
+            finally:
+                work_queue.task_done()
 
+    producer = asyncio.create_task(_producer())
+    workers = [asyncio.create_task(_worker()) for _ in range(worker_count)]
     try:
-        results = await asyncio.gather(*tasks, return_exceptions=False)
-    except Exception:
-        for t in tasks:
+        await asyncio.gather(producer, *workers)
+    except BaseException:
+        producer.cancel()
+        for t in workers:
             t.cancel()
-        await asyncio.gather(*tasks, return_exceptions=True)
+        await asyncio.gather(producer, *workers, return_exceptions=True)
         raise
 
     if callback:
@@ -522,7 +534,8 @@ async def run_chunked_pipeline(
         except Exception:
             logging.debug("%s: completion callback failed", log_prefix, exc_info=True)
 
-    return aggregate(results) if aggregate else results
+    ordered_results = [results[idx] for idx in range(total) if completed[idx]]
+    return aggregate(ordered_results) if aggregate else ordered_results
 
 
 # ---------------------------------------------------------------------------
@@ -633,7 +646,7 @@ async def _embedding_dedup(
     merge_threshold: float = 0.90,
     ambiguous_low: float = 0.75,
 ) -> tuple[dict[int, int], list[tuple[int, int]], Optional[list]]:
-    """Vectorised pairwise cosine; same-type-only when ``type_key`` given.
+    """Blockwise pairwise cosine; same-type-only when ``type_key`` given.
 
     Returns ``(merged_into, ambiguous_pairs, vectors)``. ``merged_into``
     is a union-find map ``index → parent_index``. ``ambiguous_pairs`` is the
@@ -655,11 +668,17 @@ async def _embedding_dedup(
         return {}, [], None
 
     try:
-        from sklearn.metrics.pairwise import cosine_similarity
         import numpy as np
 
         matrix = np.asarray([list(v) for v in vectors], dtype=float)
-        sims = cosine_similarity(matrix)
+        if matrix.ndim != 2 or matrix.shape[0] != n:
+            raise ValueError(f"invalid embedding matrix shape: {matrix.shape}")
+
+        # Normalize once, then calculate only bounded blocks of the upper
+        # triangle. The old implementation materialized an N x N matrix,
+        # which becomes a significant memory spike for large KBs.
+        norms = np.linalg.norm(matrix, axis=1, keepdims=True)
+        matrix = np.divide(matrix, norms, out=np.zeros_like(matrix), where=norms > 0)
     except Exception:
         logging.exception("bulk_dedup: pairwise cosine failed; skipping")
         return {}, [], vectors
@@ -673,16 +692,29 @@ async def _embedding_dedup(
 
     auto_pairs: list[tuple[int, int]] = []
     ambiguous_pairs: list[tuple[int, int]] = []
+    block_size = 1024
+    groups: dict[Any, list[int]] = {}
+    for index, item in enumerate(canonical):
+        groups.setdefault(item.get(type_key) if type_key else None, []).append(index)
 
-    for i in range(n):
-        for j in range(i + 1, n):
-            if type_key and canonical[i].get(type_key) != canonical[j].get(type_key):
-                continue
-            s = float(sims[i, j])
-            if s >= merge_threshold:
-                auto_pairs.append((i, j))
-            elif s >= ambiguous_low:
-                ambiguous_pairs.append((i, j))
+    for group_indices in groups.values():
+        for left_start in range(0, len(group_indices), block_size):
+            left_indices = group_indices[left_start : left_start + block_size]
+            left_vectors = matrix[left_indices]
+            for right_start in range(left_start, len(group_indices), block_size):
+                right_indices = group_indices[right_start : right_start + block_size]
+                sims = left_vectors @ matrix[right_indices].T
+                if right_start == left_start:
+                    candidate_mask = np.triu(sims >= ambiguous_low, k=1)
+                else:
+                    candidate_mask = sims >= ambiguous_low
+                rows, cols = np.nonzero(candidate_mask)
+                for row, col in zip(rows.tolist(), cols.tolist(), strict=True):
+                    score = float(sims[row, col])
+                    if score >= merge_threshold:
+                        auto_pairs.append((left_indices[row], right_indices[col]))
+                    else:
+                        ambiguous_pairs.append((left_indices[row], right_indices[col]))
 
     for i, j in auto_pairs:
         ri, rj = _root(i), _root(j)
@@ -709,7 +741,13 @@ async def _resolve_ambiguous_pairs(
     llm_timeout: int = 60,
     system_prompt: str = DEFAULT_DISAMBIGUATE_SYSTEM,
 ) -> dict[int, int]:
-    """LLM-judged disambiguation in batches; returns updated ``merged_into``."""
+    """LLM-judged disambiguation in bounded concurrent waves.
+
+    Each wave filters pairs using the latest union-find roots, then runs at
+    most one batch per available pooled LLM slot. Results are applied in one
+    serial pass before the next wave starts, avoiding redundant calls for
+    pairs merged by an earlier wave.
+    """
     if not ambiguous_pairs:
         return merged_into
 
@@ -718,12 +756,7 @@ async def _resolve_ambiguous_pairs(
             i = merged_into[i]
         return i
 
-    for start in range(0, len(ambiguous_pairs), batch_size):
-        batch = ambiguous_pairs[start : start + batch_size]
-        batch = [(i, j) for i, j in batch if _root(i) != _root(j)]
-        if not batch:
-            continue
-
+    async def _resolve_batch(batch: list[tuple[int, int]]) -> tuple[list[tuple[int, int]], Optional[list]]:
         lines: list[str] = []
         for k, (i, j) in enumerate(batch):
             a_type = f" ({canonical[i].get(type_key, '')})" if type_key else ""
@@ -744,10 +777,10 @@ async def _resolve_ambiguous_pairs(
             )
         except asyncio.TimeoutError:
             logging.warning("bulk_dedup: disambiguation timed out (%d pairs)", len(batch))
-            continue
+            return batch, None
         except Exception:
             logging.exception("bulk_dedup: disambiguation call failed (%d pairs)", len(batch))
-            continue
+            return batch, None
 
         decisions = None
         if isinstance(res, list):
@@ -759,19 +792,34 @@ async def _resolve_ambiguous_pairs(
                     break
         if not isinstance(decisions, list):
             logging.warning("bulk_dedup: disambiguation returned unexpected shape: %r", type(res))
-            continue
+            return batch, None
+        return batch, decisions
 
-        for k, (i, j) in enumerate(batch):
-            verdict = decisions[k] if k < len(decisions) else False
-            if not verdict:
+    pool = getattr(chat_mdl, "_pool", None)
+    wave_batch_count = max(1, int(getattr(pool, "max_concurrency", 10)))
+    remaining = list(ambiguous_pairs)
+    while remaining:
+        eligible = [(i, j) for i, j in remaining if _root(i) != _root(j)]
+        if not eligible:
+            break
+        wave_pairs = eligible[: wave_batch_count * batch_size]
+        batches = [wave_pairs[start : start + batch_size] for start in range(0, len(wave_pairs), batch_size)]
+        results = await asyncio.gather(*(_resolve_batch(batch) for batch in batches))
+        for batch, decisions in results:
+            if decisions is None:
                 continue
-            ri, rj = _root(i), _root(j)
-            if ri == rj:
-                continue
-            if canonical[ri].get("mention_count", 0) >= canonical[rj].get("mention_count", 0):
-                merged_into[rj] = ri
-            else:
-                merged_into[ri] = rj
+            for k, (i, j) in enumerate(batch):
+                verdict = decisions[k] if k < len(decisions) else False
+                if not verdict:
+                    continue
+                ri, rj = _root(i), _root(j)
+                if ri == rj:
+                    continue
+                if canonical[ri].get("mention_count", 0) >= canonical[rj].get("mention_count", 0):
+                    merged_into[rj] = ri
+                else:
+                    merged_into[ri] = rj
+        remaining = eligible[len(wave_pairs) :]
 
     return merged_into
 
@@ -837,7 +885,7 @@ async def bulk_dedup_items(
     ``aggregate_extra(group)``.
 
     Phase 2 (when ``embd_mdl`` is provided AND ``len(canonical) > 1``):
-    vectorised pairwise cosine over the canonical ``name_key`` values.
+    blockwise pairwise cosine over the canonical ``name_key`` values.
     Pairs at similarity ≥ ``merge_threshold`` auto-merge; pairs in
     ``[ambiguous_low, merge_threshold)`` move to phase 3. When ``type_key``
     is given, pairs are only considered when both endpoints share the same

rag/advanced_rag/knowlege_compile/structure.py

@@ -50,18 +50,29 @@ _ES_DEDUP_INSERT_BATCH_SIZE = 256
 class LLMCallPool:
     """Task-scoped priority scheduler for actual chat-model calls."""
 
-    def __init__(self, max_concurrency: int = 10):
+    def __init__(self, max_concurrency: int = 10, max_pending: int | None = None):
         self.max_concurrency = max(1, int(max_concurrency))
+        self.max_pending = max(self.max_concurrency, int(max_pending or self.max_concurrency))
         self._active = 0
         self._ticket = 0
         self._waiting: list[tuple[int, int]] = []
         self._condition = asyncio.Condition()
 
+    @property
+    def active_count(self) -> int:
+        return self._active
+
+    @property
+    def pending_count(self) -> int:
+        return self._active + len(self._waiting)
+
     def wrap(self, chat_mdl, *, priority: int, label: str, context: str | None = None):
         return PooledChatModel(self, chat_mdl, priority=priority, label=label, context=context)
 
     async def call(self, fn, *, priority: int, label: str, context: str | None = None):
         async with self._condition:
+            while self.pending_count >= self.max_pending:
+                await self._condition.wait()
             ticket = (int(priority), self._ticket)
             self._ticket += 1
             heapq.heappush(self._waiting, ticket)

rag/svr/task_executor_refactor/dataset_wiki_generator.py

@@ -55,6 +55,7 @@ from common import settings
 from common.constants import LLMType
 from common.misc_utils import thread_pool_exec
 from rag.nlp import search
+from rag.advanced_rag.knowlege_compile.structure import LLMCallPool
 from rag.advanced_rag.knowlege_compile.wiki import (
     wiki_map_from_chunks,
     wiki_plan_from_reduction,
@@ -72,6 +73,22 @@ from rag.svr.task_executor_refactor.task_context import TaskContext
 # room for the function's internal split_chunks packing to do real work.
 WIKI_MAP_BATCH_CHUNKS = 64
 
+# The pool limits actual MAP LLM calls rather than the surrounding batch
+# tasks. This lets the next waiting batch start as soon as a model call
+# finishes, without waiting for the previous batch's ES persistence work.
+WIKI_MAP_LLM_POOL_SIZE = 20
+
+# Global MAP admission limit: active calls plus calls waiting in the pool.
+WIKI_MAP_MAX_PENDING = 25
+
+# Keep only a small number of outer batches buffered. With 20 workers this
+# bounds the in-memory MAP work to roughly 25 batches (20 active + 5 waiting).
+WIKI_MAP_QUEUE_SIZE = 5
+
+# REFINE pages are independent. Keep enough page workers to feed the shared
+# pool; the pool itself still caps actual LLM requests at 20.
+WIKI_REFINE_WORKERS = WIKI_MAP_LLM_POOL_SIZE
+
 # Per-node cap on ``source_chunk_ids`` carried by the canvas graph blob.
 # Pages can accumulate hundreds of source chunks; the graph response is
 # meant for fast canvas rendering, not full provenance audit, so we trim
@@ -744,61 +761,86 @@ async def run_wiki(
     # were already processed in a prior run — this is the incremental
     # behavior the user asked for.
     #
-    # ``kb_chat_llm_id`` is captured from the first eligible template
-    # and used as the canonical chat model for REDUCE/PLAN/REFINE.
-    # ``kb_writer_example`` follows the same first-template-wins rule:
-    # the REFINE writer's page-structure section is pulled from the
-    # first eligible artifact template's ``parser_config.example``
-    # override (None falls back to the built-in ``WIKI_TEMPLATE_EXAMPLE``).
-    kb_chat_llm_id: Optional[str] = None
-    kb_writer_example: Optional[str] = None
-    n_docs = len(eligible)
-    for i, (doc, template_id) in enumerate(eligible):
-        doc_id = doc["id"]
-        progress(
-            0.05 + 0.6 * (i / n_docs),
-            f"MAP {i + 1}/{n_docs}: {doc.get('name', doc_id)}",
-        )
-
+    # Resolve templates before starting workers so the first eligible
+    # template remains the deterministic source for KB-wide REDUCE/PLAN/
+    # REFINE, independent of which document worker happens to finish first.
+    resolved_eligible: list[tuple[dict, str, dict]] = []
+    for doc, template_id in eligible:
         template = CompilationTemplateService.get_saved(template_id, ctx.tenant_id)
         if not template:
             logging.warning(
                 "artifact: template %s not found for doc %s; skipping",
                 template_id,
-                doc_id,
+                doc["id"],
             )
             continue
-        parser_cfg = template.get("config") or {}
+        resolved_eligible.append((doc, template_id, template.get("config") or {}))
 
-        map_llm_id = (parser_cfg.get("llm_id") or "").strip() if isinstance(parser_cfg, dict) else ""
-        map_chat_mdl = _bundle_for(map_llm_id)
-        if kb_chat_llm_id is None:
-            # First eligible template wins — canonical for KB-wide
-            # REDUCE/PLAN/REFINE further down.
-            kb_chat_llm_id = map_llm_id or None
-            tmpl_example = parser_cfg.get("example") if isinstance(parser_cfg, dict) else None
-            if isinstance(tmpl_example, str) and tmpl_example.strip():
-                kb_writer_example = tmpl_example
+    if not resolved_eligible:
+        progress(1.0, "No valid templates resolved for wiki compilation.")
+        return
 
-        # Stream the doc's chunks in batches and call MAP per batch so
-        # peak memory stays bounded for long docs.
-        agg = {"entities": 0, "concepts": 0, "claims": 0, "relations": 0}
-        agg_delta = {"new": 0, "changed": 0, "unchanged": 0, "deleted": 0}
-        doc_had_delta = False
-        saw_any = False
-        batch_no = 0
-        async for batch in load_chunks_for_doc(
-            ctx.tenant_id,
-            ctx.kb_id,
-            doc_id,
-            batch_size=WIKI_MAP_BATCH_CHUNKS,
-        ):
-            saw_any = True
-            batch_no += 1
+    # ``kb_chat_llm_id`` is captured from the first eligible template and
+    # used as the canonical chat model for KB-wide REDUCE/PLAN/REFINE.
+    # ``kb_writer_example`` follows the same first-template-wins rule.
+    first_parser_cfg = resolved_eligible[0][2]
+    first_parser_cfg = first_parser_cfg if isinstance(first_parser_cfg, dict) else {}
+    kb_chat_llm_id: Optional[str] = (first_parser_cfg.get("llm_id") or "").strip() or None
+    first_example = first_parser_cfg.get("example")
+    kb_writer_example: Optional[str] = first_example if isinstance(first_example, str) and first_example.strip() else None
+    map_llm_pool = LLMCallPool(WIKI_MAP_LLM_POOL_SIZE, max_pending=WIKI_MAP_MAX_PENDING)
+    n_docs = len(resolved_eligible)
+    map_queue: asyncio.Queue = asyncio.Queue(maxsize=WIKI_MAP_QUEUE_SIZE)
+    doc_stats = {
+        i: {
+            "doc": doc,
+            "batch_count": 0,
+            "saw_any": False,
+            "status": "ok",
+            "agg": {"entities": 0, "concepts": 0, "claims": 0, "relations": 0},
+            "delta": {"new": 0, "changed": 0, "unchanged": 0, "deleted": 0},
+            "had_delta": False,
+        }
+        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
+        doc_id = doc["id"]
+        stats = doc_stats[i]
+        progress(0.05 + 0.6 * (i / n_docs), f"MAP {i + 1}/{n_docs}: {doc.get('name', doc_id)}")
+        try:
+            async for batch in load_chunks_for_doc(
+                ctx.tenant_id,
+                ctx.kb_id,
+                doc_id,
+                batch_size=WIKI_MAP_BATCH_CHUNKS,
+            ):
+                stats["saw_any"] = True
+                stats["batch_count"] += 1
+                await map_queue.put((i, doc, template_id, parser_cfg, batch, stats["batch_count"]))
+        except Exception:
+            logging.exception("wiki: loading MAP chunks failed for doc %s", doc_id)
+            stats["status"] = "error"
+
+    async def _map_worker() -> None:
+        while True:
+            item = await map_queue.get()
             try:
+                if item is None:
+                    return
+                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 ""
                 phase1 = await wiki_map_from_chunks(
                     chunks=batch,
-                    chat_mdl=map_chat_mdl,
+                    chat_mdl=map_llm_pool.wrap(
+                        _bundle_for(map_llm_id),
+                        priority=30,
+                        label=f"wiki-map:{doc_id}",
+                        context=f"{ctx.kb_id}:{doc_id}:map",
+                    ),
                     embd_mdl=embedding_model,
                     doc_id=doc_id,
                     tenant_id=ctx.tenant_id,
@@ -808,26 +850,44 @@ async def run_wiki(
                     parser_config=parser_cfg,
                     batch_size_cap=8,
                     window_fraction=0.5,
+                    # Keep a bounded internal worker queue. The shared pool
+                    # globally limits active + admitted waiting calls to
+                    # WIKI_MAP_MAX_PENDING, while this prevents every outer
+                    # batch from creating all of its sub-batch tasks at once.
+                    max_workers=6,
                 )
+                for key in stats["agg"]:
+                    stats["agg"][key] += len(phase1.get(key) or [])
+                meta = phase1.get("_meta") or {}
+                if isinstance(meta, dict):
+                    for key in stats["delta"]:
+                        stats["delta"][key] += int(meta.get(key, 0) or 0)
+                    stats["had_delta"] |= bool(meta.get("had_delta"))
             except Exception:
-                logging.exception(
-                    "wiki: MAP failed for doc %s batch %d",
-                    doc_id,
-                    batch_no,
-                )
-                continue
-            for k in agg.keys():
-                agg[k] += len(phase1.get(k) or [])
-            meta = phase1.get("_meta") or {}
-            if isinstance(meta, dict):
-                for k in agg_delta.keys():
-                    agg_delta[k] += int(meta.get(k, 0) or 0)
-                if meta.get("had_delta"):
-                    doc_had_delta = True
+                logging.exception("wiki: MAP failed for doc %s batch %d", doc_id, batch_no)
+                stats["status"] = "error"
+            finally:
+                map_queue.task_done()
 
-        if not saw_any:
+    producers = [asyncio.create_task(_produce_document(i, job)) for i, job in enumerate(resolved_eligible)]
+    workers = [asyncio.create_task(_map_worker()) for _ in range(WIKI_MAP_LLM_POOL_SIZE)]
+    try:
+        await asyncio.gather(*producers)
+        await map_queue.join()
+    finally:
+        for task in producers + workers:
+            if not task.done():
+                task.cancel()
+        await asyncio.gather(*producers, *workers, return_exceptions=True)
+
+    for i, stats in doc_stats.items():
+        doc = stats["doc"]
+        doc_id = doc["id"]
+        if not stats["saw_any"] and stats["status"] == "ok":
+            stats["status"] = "empty"
             logging.info("wiki: no chunks for doc %s; skipping", doc_id)
-            continue
+        agg = stats["agg"]
+        delta = stats["delta"]
         logging.info(
             "wiki: MAP doc=%s entities=%d concepts=%d claims=%d relations=%d (batches=%d, new=%d changed=%d unchanged=%d deleted=%d, delta=%s)",
             doc_id,
@@ -835,23 +895,28 @@ async def run_wiki(
             agg["concepts"],
             agg["claims"],
             agg["relations"],
-            batch_no,
-            agg_delta["new"],
-            agg_delta["changed"],
-            agg_delta["unchanged"],
-            agg_delta["deleted"],
-            doc_had_delta,
+            stats["batch_count"],
+            delta["new"],
+            delta["changed"],
+            delta["unchanged"],
+            delta["deleted"],
+            stats["had_delta"],
         )
-
     # 5. REDUCE / PLAN / REFINE KB-wide. Each phase has its own
     # input_hash gate (REDUCE keys off the MAP-state hash, PLAN off
     # REDUCE's hash, REFINE off PLAN's hash) so re-runs without an
     # upstream delta short-circuit at the cache layer.
     kb_chat_mdl = _bundle_for(kb_chat_llm_id)
+
     try:
         progress(0.65, "Reducing extracts KB-wide...")
         await wiki_reduce_from_extracts(
-            chat_mdl=kb_chat_mdl,
+            chat_mdl=map_llm_pool.wrap(
+                kb_chat_mdl,
+                priority=20,
+                label="wiki-reduce",
+                context=f"{ctx.kb_id}:reduce",
+            ),
             embd_mdl=embedding_model,
             tenant_id=ctx.tenant_id,
             kb_id=ctx.kb_id,
@@ -871,10 +936,16 @@ async def run_wiki(
 
         progress(0.85, "Refining pages...")
         pages = await wiki_refine_from_plan(
-            chat_mdl=kb_chat_mdl,
+            chat_mdl=map_llm_pool.wrap(
+                kb_chat_mdl,
+                priority=30,
+                label="wiki-refine",
+                context=f"{ctx.kb_id}:refine",
+            ),
             embd_mdl=embedding_model,
             tenant_id=ctx.tenant_id,
             kb_id=ctx.kb_id,
+            max_workers=WIKI_REFINE_WORKERS,
             callback=_stage_cb("[wiki REFINE]"),
             example=kb_writer_example,
         )
@@ -885,13 +956,13 @@ async def run_wiki(
 
     # 6. Persist searchable artifact_page rows.
     try:
-        await persist_wiki_pages_to_es(ctx, pages or [], embedding_model)
+        await persist_wiki_pages_to_es(ctx=ctx, pages=pages or [], embd_mdl=embedding_model)
     except Exception:
         logging.exception("wiki: ES persist failed for kb %s", ctx.kb_id)
 
     # 7. Materialize the canvas graph from the refined pages.
     try:
-        await persist_wiki_page_graph_to_es(ctx, pages or [])
+        await persist_wiki_page_graph_to_es(ctx=ctx, pages=pages or [])
     except Exception:
         logging.exception("wiki: page-graph persist failed for kb %s", ctx.kb_id)