diff --git a/rag/advanced_rag/knowlege_compile/_common.py b/rag/advanced_rag/knowlege_compile/_common.py index 83e833a37c..1adb0d9a53 100644 --- a/rag/advanced_rag/knowlege_compile/_common.py +++ b/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 diff --git a/rag/advanced_rag/knowlege_compile/structure.py b/rag/advanced_rag/knowlege_compile/structure.py index 4d9fe98308..7571020263 100644 --- a/rag/advanced_rag/knowlege_compile/structure.py +++ b/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) diff --git a/rag/svr/task_executor_refactor/dataset_wiki_generator.py b/rag/svr/task_executor_refactor/dataset_wiki_generator.py index 5823fd7254..26d817558d 100644 --- a/rag/svr/task_executor_refactor/dataset_wiki_generator.py +++ b/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)