From 2476e2fbfa697b3c85d9fcac8ddc8fe8052fb80d Mon Sep 17 00:00:00 2001 From: buua436 Date: Fri, 4 Sep 2026 17:05:22 +0800 Subject: [PATCH] feat: configure knowledge compilation runtime settings (#19254) --- docker/.env | 31 ++++++ .../runtime_configuration.md | 94 +++++++++++++++++++ rag/advanced_rag/knowlege_compile/_common.py | 44 +++++++++ rag/advanced_rag/knowlege_compile/runner.py | 19 ++-- .../knowlege_compile/structure.py | 32 ++++++- rag/advanced_rag/knowlege_compile/wiki.py | 15 +-- .../chunk_post_processor.py | 8 +- .../dataset_wiki_generator.py | 17 ++-- .../rag/advanced_rag/test_llm_call_pool.py | 50 ++++++++++ 9 files changed, 284 insertions(+), 26 deletions(-) create mode 100644 docs/guides/knowledge_compilation/runtime_configuration.md diff --git a/docker/.env b/docker/.env index 188862cbb9..02dc08e93c 100644 --- a/docker/.env +++ b/docker/.env @@ -407,6 +407,37 @@ DOTNET_SYSTEM_GLOBALIZATION_INVARIANT=1 # Used for ThreadPoolExecutor THREAD_POOL_MAX_WORKERS=128 +# ----------------------------------------------------------------------------- +# Knowledge compilation +# ----------------------------------------------------------------------------- +# Wiki LLM pool and phase timeouts. Values are read when the Python service +# starts; restart the service after changing them. +WIKI_MAP_LLM_POOL_SIZE=20 +WIKI_MAP_MAX_PENDING=25 +WIKI_REFINE_WORKERS=4 +WIKI_MAP_WORKERS=20 +WIKI_MAP_TIMEOUT=600 +WIKI_REDUCE_TIMEOUT=60 +WIKI_PLAN_TIMEOUT=600 +WIKI_REFINE_TIMEOUT=300 +WIKI_MERGE_TIMEOUT=600 + +# Structure Compile batching, pool, and timeout settings. +DOC_STRUCTURE_COMPILE_BATCH_CHUNKS=4 +DOC_STRUCTURE_COMPILE_MAX_IN_FLIGHT=15 +DOC_STRUCTURE_LLM_POOL_SIZE=20 +STRUCTURE_CONTEXT_FRACTION=0.5 +STRUCTURE_DEFAULT_CONTEXT=100000 +KNOWLEDGE_GRAPH_CONTEXT_FRACTION=0.1 +KNOWLEDGE_GRAPH_MIN_BATCH_TOKENS=2048 +KNOWLEDGE_GRAPH_MAX_BATCH_TOKENS=4096 +STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S=120 + +# Shared LLM pool rate-limit retry settings. +LLM_POOL_RATE_LIMIT_RETRIES=3 +LLM_POOL_RATE_LIMIT_RETRY_BASE_DELAY=1.0 +LLM_POOL_RATE_LIMIT_RETRY_MAX_DELAY=30.0 + #Option to disable login form for SSO DISABLE_PASSWORD_LOGIN=false diff --git a/docs/guides/knowledge_compilation/runtime_configuration.md b/docs/guides/knowledge_compilation/runtime_configuration.md new file mode 100644 index 0000000000..1588fabfa1 --- /dev/null +++ b/docs/guides/knowledge_compilation/runtime_configuration.md @@ -0,0 +1,94 @@ +--- +sidebar_position: 6 +title: Knowledge Compilation Runtime Configuration +sidebar_label: Runtime Configuration +slug: /knowledge_compilation/runtime_configuration +sidebar_custom_props: { + categoryIcon: LucideWandSparkles +} +--- + +# Knowledge Compilation Runtime Configuration + +Knowledge compilation runtime parameters can be configured with environment +variables. In the Docker deployment, set them in `docker/.env` and restart the +RAGFlow service for the changes to take effect. + +The values below are the defaults. If an environment variable is not set, its +default value is used. Invalid values are replaced with the default value. +Values below a configured minimum or above a configured maximum are clamped to +the valid range and logged. + +## Wiki Configuration + +| Environment variable | Default | Description | +| --- | ---: | --- | +| `WIKI_MAP_LLM_POOL_SIZE` | `20` | Maximum number of concurrent LLM calls allowed by the Wiki task's shared LLM pool. | +| `WIKI_MAP_MAX_PENDING` | `25` | Maximum number of active and waiting calls admitted by the Wiki LLM pool. The effective value is never lower than `WIKI_MAP_LLM_POOL_SIZE`. | +| `WIKI_REFINE_WORKERS` | `4` | Number of page-refinement workers. Actual LLM concurrency is still limited by the shared pool. | +| `WIKI_MAP_WORKERS` | `20` | Default worker count used by direct Wiki MAP calls. | +| `WIKI_MAP_TIMEOUT` | `600` seconds | Timeout for one Wiki MAP extraction call. | +| `WIKI_REDUCE_TIMEOUT` | `60` seconds | Timeout for one Wiki REDUCE disambiguation call. | +| `WIKI_PLAN_TIMEOUT` | `600` seconds | Timeout for Wiki PLAN calls, including page planning and MAYBE resolution. | +| `WIKI_REFINE_TIMEOUT` | `300` seconds | Timeout for one Wiki page-writing call. | +| `WIKI_MERGE_TIMEOUT` | `600` seconds | Timeout for merging an existing Wiki page with newly generated content. | + +`WIKI_MAP_LLM_POOL_SIZE` is a hard ceiling for the task. The pool can reduce +the effective concurrency for a model after rate limiting, so setting this +value does not force every provider to receive that many concurrent requests. + +## Structure Compile Configuration + +| Environment variable | Default | Description | +| --- | ---: | --- | +| `DOC_STRUCTURE_LLM_POOL_SIZE` | `20` | Maximum number of concurrent LLM calls for a document Structure Compile task. | +| `DOC_STRUCTURE_COMPILE_MAX_IN_FLIGHT` | `15` | Maximum number of structure batch/template operations in flight. | +| `DOC_STRUCTURE_COMPILE_BATCH_CHUNKS` | `4` | Number of source chunks passed to one outer Structure Compile batch. | +| `STRUCTURE_CONTEXT_FRACTION` | `0.5` | Fraction of the model context window used when packing structure batches, clamped to `(0, 1]`. | +| `STRUCTURE_DEFAULT_CONTEXT` | `100000` tokens | Fallback model context size when the model does not provide one. | +| `KNOWLEDGE_GRAPH_CONTEXT_FRACTION` | `0.1` | Fraction of the model context window used for Knowledge Graph batches, clamped to `(0, 1]`. | +| `KNOWLEDGE_GRAPH_MIN_BATCH_TOKENS` | `2048` tokens | Minimum Knowledge Graph batch size. | +| `KNOWLEDGE_GRAPH_MAX_BATCH_TOKENS` | `4096` tokens | Maximum Knowledge Graph batch size. The effective value is never lower than `KNOWLEDGE_GRAPH_MIN_BATCH_TOKENS`. | +| `STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S` | `120` seconds | Time limit for the Structure Compile chain-correction LLM step. | + +The regular Structure Compile entity/relation extraction path does not define +an independent application-level timeout. Its request timeout is determined by +the configured LLM provider/client. + +## LLM Pool Rate-Limit Handling + +These variables apply to the shared adaptive LLM pool: + +| Environment variable | Default | Description | +| --- | ---: | --- | +| `LLM_POOL_RATE_LIMIT_RETRIES` | `3` | Number of retries after a rate-limit response. | +| `LLM_POOL_RATE_LIMIT_RETRY_BASE_DELAY` | `1.0` second | Initial delay used by exponential backoff. | +| `LLM_POOL_RATE_LIMIT_RETRY_MAX_DELAY` | `30.0` seconds | Maximum delay between rate-limit retries. | + +After a rate-limit response, the pool lowers the effective concurrency for the +affected model and retries the request. Successful calls gradually restore +the model's concurrency up to the configured pool limit. + +The pool normalizes `max_pending` to at least the configured pool size, and +normalizes the retry maximum delay to at least the retry base delay. + +When all retries are exhausted, the backend records the detailed failure in +its service log and the frontend task log receives only the stage, context, and +error type. Provider response bodies are not forwarded to the frontend. + +## Example + +The following configuration lowers concurrency for providers with a smaller +request limit and shortens the Wiki MAP timeout: + +```dotenv +WIKI_MAP_LLM_POOL_SIZE=8 +WIKI_MAP_MAX_PENDING=10 +DOC_STRUCTURE_LLM_POOL_SIZE=8 +LLM_POOL_RATE_LIMIT_RETRIES=5 +WIKI_MAP_TIMEOUT=300 +``` + +Restart the RAGFlow backend after changing `docker/.env`. The settings are +read when the Python modules are loaded and are not changed for already-running +tasks. diff --git a/rag/advanced_rag/knowlege_compile/_common.py b/rag/advanced_rag/knowlege_compile/_common.py index 77393d6427..9cd36157c4 100644 --- a/rag/advanced_rag/knowlege_compile/_common.py +++ b/rag/advanced_rag/knowlege_compile/_common.py @@ -37,6 +37,8 @@ from __future__ import annotations import asyncio import logging +import math +import os import string from typing import Any, Awaitable, Callable, Iterable, Optional @@ -48,6 +50,48 @@ from rag.nlp import rag_tokenizer from rag.prompts.generator import INPUT_UTILIZATION, gen_json, split_chunks +def env_int(name: str, default: int, *, minimum: int | None = None) -> int: + """Read a non-empty integer environment override, or return ``default``.""" + raw = os.getenv(name) + if raw is None or not raw.strip(): + value = default + else: + try: + value = int(raw) + except ValueError: + logging.warning("Invalid integer environment variable %s=%r; using %d", name, raw, default) + value = default + if minimum is not None and value < minimum: + logging.warning("Environment variable %s=%d is below minimum %d; using %d", name, value, minimum, minimum) + value = minimum + logging.debug("Resolved environment variable %s=%d", name, value) + return value + + +def env_float(name: str, default: float, *, minimum: float | None = None, maximum: float | None = None) -> float: + """Read a non-empty float environment override, or return ``default``.""" + raw = os.getenv(name) + if raw is None or not raw.strip(): + value = default + else: + try: + value = float(raw) + except ValueError: + logging.warning("Invalid float environment variable %s=%r; using %s", name, raw, default) + value = default + if not math.isfinite(value): + logging.warning("Environment variable %s=%r is not finite; using %s", name, value, default) + value = default + if minimum is not None and value < minimum: + logging.warning("Environment variable %s=%s is below minimum %s; using %s", name, value, minimum, minimum) + value = minimum + if maximum is not None and value > maximum: + logging.warning("Environment variable %s=%s is above maximum %s; using %s", name, value, maximum, maximum) + value = maximum + logging.debug("Resolved environment variable %s=%s", name, value) + return value + + def knowledge_compile_gen_conf(chat_mdl, gen_conf: Optional[dict] = None) -> dict: """Add model-specific reasoning controls for knowledge compilation only.""" conf = dict(gen_conf or {}) diff --git a/rag/advanced_rag/knowlege_compile/runner.py b/rag/advanced_rag/knowlege_compile/runner.py index 692936aaf7..2f4ad78607 100644 --- a/rag/advanced_rag/knowlege_compile/runner.py +++ b/rag/advanced_rag/knowlege_compile/runner.py @@ -53,30 +53,31 @@ from rag.advanced_rag.knowlege_compile.structure import ( rebuild_dataset_structure_graph_json, rebuild_structure_graph_json, ) +from rag.advanced_rag.knowlege_compile._common import env_float, env_int # ----- tunables ------------------------------------------------------ # Bound how many source chunks are handed to a single # ``compile_structure_from_text`` invocation for regular templates. -DOC_STRUCTURE_COMPILE_BATCH_CHUNKS = 4 +DOC_STRUCTURE_COMPILE_BATCH_CHUNKS = env_int("DOC_STRUCTURE_COMPILE_BATCH_CHUNKS", 4, minimum=1) # Structure compilation packs chunks up to half of the model context. # ``compile_structure_from_text`` applies the exact prompt-aware packing again # before the LLM call. -STRUCTURE_CONTEXT_FRACTION = 0.5 -STRUCTURE_DEFAULT_CONTEXT = 100_000 -KNOWLEDGE_GRAPH_CONTEXT_FRACTION = 0.1 -KNOWLEDGE_GRAPH_MIN_BATCH_TOKENS = 2048 -KNOWLEDGE_GRAPH_MAX_BATCH_TOKENS = 4096 +STRUCTURE_CONTEXT_FRACTION = env_float("STRUCTURE_CONTEXT_FRACTION", 0.5, minimum=0.01, maximum=1.0) +STRUCTURE_DEFAULT_CONTEXT = env_int("STRUCTURE_DEFAULT_CONTEXT", 100_000, minimum=1) +KNOWLEDGE_GRAPH_CONTEXT_FRACTION = env_float("KNOWLEDGE_GRAPH_CONTEXT_FRACTION", 0.1, minimum=0.01, maximum=1.0) +KNOWLEDGE_GRAPH_MIN_BATCH_TOKENS = env_int("KNOWLEDGE_GRAPH_MIN_BATCH_TOKENS", 2048, minimum=1) +KNOWLEDGE_GRAPH_MAX_BATCH_TOKENS = env_int("KNOWLEDGE_GRAPH_MAX_BATCH_TOKENS", 4096, minimum=KNOWLEDGE_GRAPH_MIN_BATCH_TOKENS) # Bound the number of batch/template extraction calls in flight. Results are # committed in submission order so accumulator updates and merge flushes stay # deterministic while the LLM calls run concurrently. -DOC_STRUCTURE_COMPILE_MAX_IN_FLIGHT = 15 +DOC_STRUCTURE_COMPILE_MAX_IN_FLIGHT = env_int("DOC_STRUCTURE_COMPILE_MAX_IN_FLIGHT", 15, minimum=1) # Total task-scoped Chat LLM capacity shared by compile, chain validation and # merge decisions. A request waits in the priority queue when all slots are busy. -DOC_STRUCTURE_LLM_POOL_SIZE = 20 +DOC_STRUCTURE_LLM_POOL_SIZE = env_int("DOC_STRUCTURE_LLM_POOL_SIZE", 20, minimum=1) # Bound how many compiled ES-ready docs may accumulate before we flush # them through ``merge_compiled_structures``. The merger does pairwise @@ -85,7 +86,7 @@ DOC_STRUCTURE_LLM_POOL_SIZE = 20 DOC_STRUCTURE_MERGE_MAX_DOCS = 512 # Hard wall on the chain-validator LLM correction step. -STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S = 120.0 +STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S = env_float("STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S", 120.0, minimum=0.1) # ----- template resolution ------------------------------------------- diff --git a/rag/advanced_rag/knowlege_compile/structure.py b/rag/advanced_rag/knowlege_compile/structure.py index e9ed583526..638f99c699 100644 --- a/rag/advanced_rag/knowlege_compile/structure.py +++ b/rag/advanced_rag/knowlege_compile/structure.py @@ -39,6 +39,8 @@ from ._common import ( union_ordered as _union_ordered, run_chunked_pipeline as _run_chunked_pipeline, knowledge_compile_gen_conf as _knowledge_compile_gen_conf, + env_float as _env_float, + env_int as _env_int, ) @@ -70,6 +72,10 @@ MERGE_SCOPE_DATASET = "dataset" _STRUCT_MERGE_LOCK_TIMEOUT_S = 60 _STRUCT_MERGE_LOCK_BLOCKING_TIMEOUT_S = 5 +LLM_POOL_RATE_LIMIT_RETRIES = _env_int("LLM_POOL_RATE_LIMIT_RETRIES", 3, minimum=0) +LLM_POOL_RATE_LIMIT_RETRY_BASE_DELAY = _env_float("LLM_POOL_RATE_LIMIT_RETRY_BASE_DELAY", 1.0, minimum=0.0) +LLM_POOL_RATE_LIMIT_RETRY_MAX_DELAY = _env_float("LLM_POOL_RATE_LIMIT_RETRY_MAX_DELAY", 30.0, minimum=0.0) + class _RechunkedDocs(list): """Compiled structure rows plus the formal chunks created by rechunking.""" @@ -123,11 +129,12 @@ class LLMCallPool: decrease_cooldown: float = 5.0, recovery_successes: int = 20, recovery_cooldown: float = 30.0, - rate_limit_retries: int = 3, - rate_limit_retry_base_delay: float = 1.0, - rate_limit_retry_max_delay: float = 30.0, + rate_limit_retries: int = LLM_POOL_RATE_LIMIT_RETRIES, + rate_limit_retry_base_delay: float = LLM_POOL_RATE_LIMIT_RETRY_BASE_DELAY, + rate_limit_retry_max_delay: float = LLM_POOL_RATE_LIMIT_RETRY_MAX_DELAY, clock: Callable[[], float] = time.monotonic, on_concurrency_change: Callable[[int, int, str], None] | None = None, + on_error: Callable[[str, str | None, str], None] | None = None, ): self.max_concurrency = max(1, int(max_concurrency)) self.max_pending = max(self.max_concurrency, int(max_pending or self.max_concurrency)) @@ -141,6 +148,7 @@ class LLMCallPool: self.rate_limit_retry_max_delay = max(self.rate_limit_retry_base_delay, float(rate_limit_retry_max_delay)) self._clock = clock self._on_concurrency_change = on_concurrency_change + self._on_error = on_error self._active = 0 self._ticket = 0 self._waiting: list[tuple[int, int, str]] = [] @@ -266,6 +274,21 @@ class LLMCallPool: except Exception: logging.exception("LLM pool concurrency change callback failed") + def _notify_error(self, label: str, context: str | None, error: BaseException | str) -> None: + error_type = type(error).__name__ if isinstance(error, BaseException) else "ProviderErrorResult" + logging.error( + "LLM pool terminal call failure label=%s context=%s error_type=%s", + label, + context, + error_type, + ) + if self._on_error is None: + return + try: + self._on_error(label, context, error_type) + except Exception: + logging.exception("LLM pool error callback failed") + async def _acquire(self, model_key: str, priority: int) -> None: async with self._condition: while self.pending_count >= self.max_pending: @@ -311,6 +334,7 @@ class LLMCallPool: rate_limited = self._is_rate_limited(exc) await self._release(model_key, "rate_limited" if rate_limited else "failed", label=label, context=context) if not rate_limited or retry >= self.rate_limit_retries: + self._notify_error(label, context, exc) raise else: error_result = self._is_error_result(result) @@ -322,6 +346,8 @@ class LLMCallPool: context=context, ) if not rate_limited or retry >= self.rate_limit_retries: + if error_result: + self._notify_error(label, context, result) return result delay = self._rate_limit_retry_delay(retry + 1) diff --git a/rag/advanced_rag/knowlege_compile/wiki.py b/rag/advanced_rag/knowlege_compile/wiki.py index cf7b9c182b..b9126b5d15 100644 --- a/rag/advanced_rag/knowlege_compile/wiki.py +++ b/rag/advanced_rag/knowlege_compile/wiki.py @@ -54,6 +54,7 @@ from ._common import ( knowledge_compile_gen_conf as _knowledge_compile_gen_conf, run_chunked_pipeline as _run_chunked_pipeline, stable_row_id as _stable_row_id, + env_int as _env_int, ) @@ -90,8 +91,8 @@ from .structure import ( WIKI_MAP_COMPILE_KWD = "wiki_map_extract" WIKI_MAP_STATE_COMPILE_KWD = "wiki_map_state" WIKI_MAP_STATE_META_COMPILE_KWD = "wiki_map_state_meta" -DEFAULT_WIKI_MAP_WORKERS = 20 -DEFAULT_WIKI_MAP_TIMEOUT = 600 +DEFAULT_WIKI_MAP_WORKERS = _env_int("WIKI_MAP_WORKERS", 20, minimum=1) +DEFAULT_WIKI_MAP_TIMEOUT = _env_int("WIKI_MAP_TIMEOUT", 600, minimum=1) async def _wiki_disabled_doc_ids(kb_id: str) -> set[str]: @@ -1355,7 +1356,7 @@ WIKI_REDUCE_COMPILE_KWD = "wiki_reduce_result" DEFAULT_WIKI_REDUCE_MERGE_THRESHOLD = 0.95 DEFAULT_WIKI_REDUCE_AMBIGUOUS_LOW = 0.75 DEFAULT_WIKI_REDUCE_AMBIGUOUS_BATCH = 50 -DEFAULT_WIKI_REDUCE_TIMEOUT = 60 +DEFAULT_WIKI_REDUCE_TIMEOUT = _env_int("WIKI_REDUCE_TIMEOUT", 60, minimum=1) # System prompt for the LLM disambiguation batch. The shared engine @@ -1863,7 +1864,7 @@ WIKI_PLAN_COMPILE_KWD = "wiki_compilation_plan" WIKI_PAGE_COMPILE_KWD = "wiki_page" DEFAULT_WIKI_PLAN_UPDATE_THRESHOLD = 0.95 DEFAULT_WIKI_PLAN_MAYBE_THRESHOLD = 0.60 -DEFAULT_WIKI_PLAN_TIMEOUT = 600 # ~10 min — the planning call emits one big +DEFAULT_WIKI_PLAN_TIMEOUT = _env_int("WIKI_PLAN_TIMEOUT", 600, minimum=1) # ~10 min — the planning call emits one big # JSON plan and reasoning models can spend a # long time thinking before emitting tokens. # Override via the ``llm_timeout`` arg to @@ -2799,11 +2800,11 @@ async def wiki_plan_from_reduction( WIKI_DRAFT_COMPILE_KWD = "wiki_page_draft" -DEFAULT_WIKI_REFINE_WORKERS = 4 -DEFAULT_WIKI_REFINE_TIMEOUT = 300 +DEFAULT_WIKI_REFINE_WORKERS = _env_int("WIKI_REFINE_WORKERS", 4, minimum=1) +DEFAULT_WIKI_REFINE_TIMEOUT = _env_int("WIKI_REFINE_TIMEOUT", 300, minimum=1) WIKI_REFINE_SOURCE_BUDGET_CHARS = 60_000 WIKI_MERGE_BODY_SHRINK_THRESHOLD = 0.7 -WIKI_MERGE_TIMEOUT = 600 +WIKI_MERGE_TIMEOUT = _env_int("WIKI_MERGE_TIMEOUT", 600, minimum=1) WIKI_TEMPLATE_EXAMPLE = ( diff --git a/rag/svr/task_executor_refactor/chunk_post_processor.py b/rag/svr/task_executor_refactor/chunk_post_processor.py index 1dc17eba60..d4678ba9bf 100644 --- a/rag/svr/task_executor_refactor/chunk_post_processor.py +++ b/rag/svr/task_executor_refactor/chunk_post_processor.py @@ -1110,7 +1110,13 @@ async def run_document_structure_compile(handler, embedding_model: LLMBundle) -> chat_mdl_by_tid = {template_id: chat_mdl for template_id, _ in active_templates} from rag.advanced_rag.knowlege_compile.structure import LLMCallPool - llm_pool = LLMCallPool(DOC_STRUCTURE_LLM_POOL_SIZE) + def _on_llm_error(label: str, context: str | None, error_type: str) -> None: + ctx.progress_cb(msg=f"LLM call failed ({label}, {context or 'no context'}): {error_type}") + + llm_pool = LLMCallPool( + DOC_STRUCTURE_LLM_POOL_SIZE, + on_error=_on_llm_error, + ) tree_templates: list[tuple[str, dict]] = [] non_tree_templates: list[tuple[str, dict]] = [] diff --git a/rag/svr/task_executor_refactor/dataset_wiki_generator.py b/rag/svr/task_executor_refactor/dataset_wiki_generator.py index 1330547998..de6ca7f6aa 100644 --- a/rag/svr/task_executor_refactor/dataset_wiki_generator.py +++ b/rag/svr/task_executor_refactor/dataset_wiki_generator.py @@ -56,6 +56,7 @@ 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._common import env_int from rag.advanced_rag.knowlege_compile.wiki import ( WIKI_MAP_STATE_COMPILE_KWD, WIKI_MAP_STATE_META_COMPILE_KWD, @@ -83,28 +84,32 @@ WIKI_MAP_BATCH_CHUNKS = 64 # The pool limits actual LLM calls rather than the surrounding batch tasks and # adapts each model's admission limit when it reports rate limiting. This lets # high-capacity and low-concurrency models share one Wiki task safely. -WIKI_MAP_LLM_POOL_SIZE = 20 +WIKI_MAP_LLM_POOL_SIZE = env_int("WIKI_MAP_LLM_POOL_SIZE", 20, minimum=1) # Global MAP admission limit: active calls plus calls waiting in the pool. -WIKI_MAP_MAX_PENDING = 25 +WIKI_MAP_MAX_PENDING = env_int("WIKI_MAP_MAX_PENDING", 25, minimum=WIKI_MAP_LLM_POOL_SIZE) -# 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). +# Keep only a small number of outer batches buffered. Together with the +# configured worker and pending limits, this bounds in-memory MAP work. 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 +# pool; the pool itself still caps actual LLM requests at the configured size. +WIKI_REFINE_WORKERS = env_int("WIKI_REFINE_WORKERS", 4, minimum=1) def _create_wiki_llm_pool(progress: Callable) -> LLMCallPool: def _on_concurrency_change(old: int, new: int, reason: str) -> None: progress(msg=f"LLM pool concurrency {old} -> {new} ({reason}).") + def _on_error(label: str, context: str | None, error_type: str) -> None: + progress(msg=f"LLM call failed ({label}, {context or 'no context'}): {error_type}") + pool = LLMCallPool( WIKI_MAP_LLM_POOL_SIZE, max_pending=WIKI_MAP_MAX_PENDING, on_concurrency_change=_on_concurrency_change, + on_error=_on_error, ) logging.info("Wiki LLM pool initialized max_concurrency=%d max_pending=%d", pool.max_concurrency, pool.max_pending) progress(0.0, f"LLM pool max {pool.max_concurrency}.") diff --git a/test/unit_test/rag/advanced_rag/test_llm_call_pool.py b/test/unit_test/rag/advanced_rag/test_llm_call_pool.py index cfd327c510..0ca929cb1d 100644 --- a/test/unit_test/rag/advanced_rag/test_llm_call_pool.py +++ b/test/unit_test/rag/advanced_rag/test_llm_call_pool.py @@ -196,6 +196,56 @@ def test_llm_call_pool_reports_concurrency_changes(): asyncio.run(exercise()) +def test_llm_call_pool_reports_final_call_error(): + async def exercise(): + errors = [] + pool = LLMCallPool( + 2, + rate_limit_retries=0, + on_error=lambda label, context, error: errors.append((label, context, str(error))), + ) + model_key = "provider/model" + + async def fail(): + raise RuntimeError("provider unavailable") + + with pytest.raises(RuntimeError, match="provider unavailable"): + await pool.call( + fail, + model_key=model_key, + priority=1, + label="wiki-map", + context="kb:doc", + ) + + assert errors == [("wiki-map", "kb:doc", "RuntimeError")] + + asyncio.run(exercise()) + + +def test_llm_call_pool_reports_terminal_error_result(): + async def exercise(): + errors = [] + pool = LLMCallPool( + 2, + rate_limit_retries=0, + on_error=lambda label, context, error_type: errors.append((label, context, error_type)), + ) + + result = await pool.call( + lambda: _return("**ERROR**: provider unavailable"), + model_key="provider/model", + priority=1, + label="wiki-plan", + context="kb:plan", + ) + + assert result == "**ERROR**: provider unavailable" + assert errors == [("wiki-plan", "kb:plan", "ProviderErrorResult")] + + asyncio.run(exercise()) + + def test_llm_call_pool_isolates_models_and_avoids_head_of_line_blocking(): async def exercise(): pool = LLMCallPool(2, max_pending=4, decrease_cooldown=0, rate_limit_retries=0)