From e5392836a463992886d3682d99f7255097c7dede Mon Sep 17 00:00:00 2001 From: mountain Date: Wed, 8 Jul 2026 10:31:37 +0800 Subject: [PATCH] fix(index): bound LLM concurrency safely, per-index and leak-free Cap concurrent in-flight LLM calls during indexing via a shared semaphore (bounded_gather), so a many-node document no longer schedules one socket per node and exhausts the process fd limit (Errno 24). Make the per-index max_concurrency override correct under concurrency: - Scope IndexConfig(max_concurrency=...) to the build_index call via a ContextVar (max_concurrency_scope) instead of mutating a process global. A one-off value no longer sticks as the new default, and concurrent indexing of other documents isn't affected. - Propagate the context through _run_async's worker-thread fallback so the override survives the sync-over-async thread hop. - set_max_concurrency() stays as the explicit process-wide setter. Also stop `from .utils import *` leaking a `config` name (SimpleNamespace alias) that shadowed the real pageindex.config submodule for the page_index modules; the alias is now `_config`. Adds regression tests for cap enforcement, scope stickiness/isolation, worker-thread propagation, and the config-namespace fix. Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS --- pageindex/config.py | 81 +++++++++++++ pageindex/index/page_index.py | 10 +- pageindex/index/page_index_md.py | 2 +- pageindex/index/pipeline.py | 73 +++++++----- pageindex/index/utils.py | 39 +++++- tests/test_concurrency.py | 197 +++++++++++++++++++++++++++++++ 6 files changed, 358 insertions(+), 44 deletions(-) create mode 100644 tests/test_concurrency.py diff --git a/pageindex/config.py b/pageindex/config.py index 2accaf8..1b6f060 100644 --- a/pageindex/config.py +++ b/pageindex/config.py @@ -2,6 +2,8 @@ from __future__ import annotations import os +from contextlib import contextmanager +from contextvars import ContextVar from pydantic import BaseModel @@ -23,6 +25,10 @@ class IndexConfig(BaseModel): if_add_node_summary: bool = True if_add_doc_description: bool = True if_add_node_text: bool = False + # Max concurrent in-flight LLM calls during indexing. None = use the global + # default (get_max_concurrency(), overridable via PAGEINDEX_MAX_CONCURRENCY). + # An explicit value here wins for this client. + max_concurrency: int | None = None def _env_drop_params_default() -> bool: @@ -45,6 +51,81 @@ _LLM_PARAMS: dict = {"temperature": 0, "drop_params": _env_drop_params_default() _RESERVED_LLM_PARAMS = ("model", "messages") +# Built-in fallback cap on concurrent in-flight LLM calls during indexing, used +# when PAGEINDEX_MAX_CONCURRENCY is unset or invalid. Kept conservative so a +# default run won't trip provider rate limits or the process fd ceiling; raise +# it via the env var / set_max_concurrency() / IndexConfig(max_concurrency=…). +_DEFAULT_MAX_CONCURRENCY = 5 + + +def _env_max_concurrency_default() -> int: + """Default max in-flight LLM calls, from PAGEINDEX_MAX_CONCURRENCY. + + A missing, non-integer, or non-positive value falls back to + ``_DEFAULT_MAX_CONCURRENCY``. Read once at import; change it at runtime via + set_max_concurrency() (a later env change doesn't apply). Bounding + concurrency keeps a many-node document from opening one socket per node all + at once and exhausting the process file-descriptor limit (Errno 24). + """ + raw = os.getenv("PAGEINDEX_MAX_CONCURRENCY", str(_DEFAULT_MAX_CONCURRENCY)).strip() + try: + value = int(raw) + except ValueError: + return _DEFAULT_MAX_CONCURRENCY + return value if value > 0 else _DEFAULT_MAX_CONCURRENCY + + +# Process-wide default for concurrent in-flight LLM completions during indexing. +# Overridable process-wide via set_max_concurrency() / the env var above, or +# per-index via max_concurrency_scope() (used by build_index for +# IndexConfig(max_concurrency=…)). Read through get_max_concurrency(). +_MAX_CONCURRENCY: int = _env_max_concurrency_default() + +# Per-index override, isolated per thread / async context so concurrent indexing +# of different documents never leaks one document's limit into another (and a +# one-off override never "sticks" as the new process default). None = no +# override -> fall back to the process-wide _MAX_CONCURRENCY. +_MAX_CONCURRENCY_OVERRIDE: ContextVar[int | None] = ContextVar( + "pageindex_max_concurrency_override", default=None +) + + +def get_max_concurrency() -> int: + """Return the effective cap on concurrent in-flight LLM calls during indexing. + + A per-index override (max_concurrency_scope) wins for the current context; + otherwise the process-wide default applies. + """ + override = _MAX_CONCURRENCY_OVERRIDE.get() + return override if override is not None else _MAX_CONCURRENCY + + +def set_max_concurrency(value: int) -> None: + """Set the process-wide default cap on concurrent in-flight LLM calls.""" + global _MAX_CONCURRENCY + if not isinstance(value, int) or value <= 0: + raise ValueError("max_concurrency must be a positive integer") + _MAX_CONCURRENCY = value + + +@contextmanager +def max_concurrency_scope(value: int | None): + """Scope a per-index max-concurrency override to the current context. + + ``value=None`` means "no override" (fall back to the process default). + Isolated per thread / async context and reset on exit, so concurrent + indexing doesn't leak across documents and a one-off value never becomes + the sticky new default. + """ + if value is not None and (not isinstance(value, int) or value <= 0): + raise ValueError("max_concurrency must be a positive integer") + token = _MAX_CONCURRENCY_OVERRIDE.set(value) + try: + yield + finally: + _MAX_CONCURRENCY_OVERRIDE.reset(token) + + def get_llm_params() -> dict: """Return a copy of the per-call kwargs PageIndex passes to litellm.""" return dict(_LLM_PARAMS) diff --git a/pageindex/index/page_index.py b/pageindex/index/page_index.py index a5bc337..862f358 100644 --- a/pageindex/index/page_index.py +++ b/pageindex/index/page_index.py @@ -89,7 +89,7 @@ async def check_title_appearance_in_start_concurrent(structure, page_list, model tasks.append(check_title_appearance_in_start(item['title'], page_text, model=model, logger=logger)) valid_items.append(item) - results = await asyncio.gather(*tasks, return_exceptions=True) + results = await bounded_gather(tasks, return_exceptions=True) for item, result in zip(valid_items, results): if isinstance(result, Exception): if logger: @@ -832,7 +832,7 @@ async def fix_incorrect_toc(toc_with_page_number, page_list, incorrect_results, process_and_check_item(item) for item in incorrect_results ] - results = await asyncio.gather(*tasks, return_exceptions=True) + results = await bounded_gather(tasks, return_exceptions=True) for item, result in zip(incorrect_results, results): if isinstance(result, Exception): print(f"Processing item {item} generated an exception: {result}") @@ -927,7 +927,7 @@ async def verify_toc(page_list, list_result, start_index=1, N=None, model=None): check_title_appearance(item, page_list, start_index, model) for item in indexed_sample_list ] - results = await asyncio.gather(*tasks) + results = await bounded_gather(tasks) # Process results correct_count = 0 @@ -1015,7 +1015,7 @@ async def process_large_node_recursively(node, page_list, opt=None, logger=None) process_large_node_recursively(child_node, page_list, opt, logger=logger) for child_node in node['nodes'] ] - await asyncio.gather(*tasks) + await bounded_gather(tasks) return node @@ -1051,7 +1051,7 @@ async def tree_parser(page_list, opt, doc=None, logger=None): process_large_node_recursively(node, page_list, opt, logger=logger) for node in toc_tree ] - await asyncio.gather(*tasks) + await bounded_gather(tasks) return toc_tree diff --git a/pageindex/index/page_index_md.py b/pageindex/index/page_index_md.py index f9e300a..a115b9a 100644 --- a/pageindex/index/page_index_md.py +++ b/pageindex/index/page_index_md.py @@ -16,7 +16,7 @@ async def get_node_summary(node, summary_token_threshold=200, model=None): async def generate_summaries_for_structure_md(structure, summary_token_threshold, model=None): nodes = structure_to_list(structure) tasks = [get_node_summary(node, summary_token_threshold=summary_token_threshold, model=model) for node in nodes] - summaries = await asyncio.gather(*tasks) + summaries = await bounded_gather(tasks) for node, summary in zip(nodes, summaries): if not node.get('nodes'): diff --git a/pageindex/index/pipeline.py b/pageindex/index/pipeline.py index 70d8c2f..6217a46 100644 --- a/pageindex/index/pipeline.py +++ b/pageindex/index/pipeline.py @@ -43,11 +43,16 @@ def _run_async(coro): """Run an async coroutine, handling the case where an event loop is already running.""" import asyncio import concurrent.futures + import contextvars try: asyncio.get_running_loop() - # Already inside an event loop -- run in a separate thread + # Already inside an event loop -- run in a separate thread. Copy the + # current context so ContextVar-based settings (e.g. the + # max_concurrency_scope override set by build_index) propagate into the + # worker thread instead of silently falling back to the process default. + ctx = contextvars.copy_context() with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: - return pool.submit(asyncio.run, coro).result() + return pool.submit(ctx.run, asyncio.run, coro).result() except RuntimeError: return asyncio.run(coro) @@ -58,48 +63,52 @@ def build_index(parsed: ParsedDocument, model: str = None, opt=None) -> dict: from .utils import (write_node_id, add_node_text, remove_structure_text, generate_summaries_for_structure, generate_doc_description, create_clean_structure_for_description) - from ..config import IndexConfig + from ..config import IndexConfig, max_concurrency_scope if opt is None: opt = IndexConfig(model=model) if model else IndexConfig() - nodes = parsed.nodes - strategy = detect_strategy(nodes) + # Scope the per-index concurrency cap to THIS call only (per thread/async + # context), so concurrent indexing of other documents isn't affected and a + # one-off value never sticks as the process default. + with max_concurrency_scope(getattr(opt, "max_concurrency", None)): + nodes = parsed.nodes + strategy = detect_strategy(nodes) - if strategy == "level_based": - structure = build_tree_from_levels(nodes) - # For level-based, text is already in the tree nodes - else: - # Strategies 1-3: convert ContentNode list to page_list format for existing pipeline - page_list = [(n.content, n.tokens) for n in nodes] - structure = _run_async(_content_based_pipeline(page_list, opt)) + if strategy == "level_based": + structure = build_tree_from_levels(nodes) + # For level-based, text is already in the tree nodes + else: + # Strategies 1-3: convert ContentNode list to page_list format for existing pipeline + page_list = [(n.content, n.tokens) for n in nodes] + structure = _run_async(_content_based_pipeline(page_list, opt)) - # Unified enhancement - if opt.if_add_node_id: - write_node_id(structure) + # Unified enhancement + if opt.if_add_node_id: + write_node_id(structure) - if strategy != "level_based": - if opt.if_add_node_text or opt.if_add_node_summary: - add_node_text(structure, page_list) + if strategy != "level_based": + if opt.if_add_node_text or opt.if_add_node_summary: + add_node_text(structure, page_list) - if opt.if_add_node_summary: - _run_async(generate_summaries_for_structure(structure, model=opt.model)) + if opt.if_add_node_summary: + _run_async(generate_summaries_for_structure(structure, model=opt.model)) - if not opt.if_add_node_text and strategy != "level_based": - remove_structure_text(structure) + if not opt.if_add_node_text and strategy != "level_based": + remove_structure_text(structure) - result = { - "doc_name": parsed.doc_name, - "structure": structure, - } + result = { + "doc_name": parsed.doc_name, + "structure": structure, + } - if opt.if_add_doc_description: - clean_structure = create_clean_structure_for_description(structure) - result["doc_description"] = generate_doc_description( - clean_structure, model=opt.model - ) + if opt.if_add_doc_description: + clean_structure = create_clean_structure_for_description(structure) + result["doc_description"] = generate_doc_description( + clean_structure, model=opt.model + ) - return result + return result class _NullLogger: diff --git a/pageindex/index/utils.py b/pageindex/index/utils.py index a2701e4..a850a23 100644 --- a/pageindex/index/utils.py +++ b/pageindex/index/utils.py @@ -14,14 +14,41 @@ from datetime import datetime from io import BytesIO from pathlib import Path from pprint import pprint -from types import SimpleNamespace as config +# Aliased with a leading underscore so `from .utils import *` (used by the +# page_index modules) doesn't export a name `config` that would shadow the real +# `pageindex.config` submodule for those modules. +from types import SimpleNamespace as _config -from ..config import get_llm_params +from ..config import get_llm_params, get_max_concurrency from ..tokens import count_tokens # re-exported for backward compat logger = logging.getLogger(__name__) +async def bounded_gather(coros, *, return_exceptions=False): + """``asyncio.gather`` with a cap on how many coroutines run concurrently. + + Each coroutine acquires a shared semaphore before running, so no more than + ``get_max_concurrency()`` LLM calls are ever in flight at once. Without this + a many-node document schedules every node's LLM call simultaneously, opening + one socket per node and exhausting the process file-descriptor limit + (Errno 24, "Too many open files"). + + The semaphore is created inside the running loop, so this stays correct when + the caller drives each document in its own ``asyncio.run()`` loop. Order of + results matches input order, mirroring ``asyncio.gather``. + """ + semaphore = asyncio.Semaphore(get_max_concurrency()) + + async def _run(coro): + async with semaphore: + return await coro + + return await asyncio.gather( + *(_run(c) for c in coros), return_exceptions=return_exceptions + ) + + def llm_completion(model, prompt, chat_history=None, return_finish_reason=False): if model: model = model.removeprefix("litellm/") @@ -213,7 +240,7 @@ async def generate_node_summary(node, model=None): async def generate_summaries_for_structure(structure, model=None): nodes = structure_to_list(structure) tasks = [generate_node_summary(node, model=model) for node in nodes] - summaries = await asyncio.gather(*tasks) + summaries = await bounded_gather(tasks) for node, summary in zip(nodes, summaries): node['summary'] = summary @@ -769,11 +796,11 @@ class ConfigLoader: if unknown_keys: raise ValueError(f"Unknown config keys: {unknown_keys}") - def load(self, user_opt=None) -> config: + def load(self, user_opt=None) -> _config: """Merge user options over IndexConfig defaults, returning a namespace.""" if user_opt is None: user_dict = {} - elif isinstance(user_opt, config): + elif isinstance(user_opt, _config): user_dict = vars(user_opt) elif isinstance(user_opt, dict): user_dict = user_opt @@ -782,7 +809,7 @@ class ConfigLoader: self._validate_keys(user_dict) merged = {**self._default_dict, **user_dict} - return config(**merged) + return _config(**merged) def create_node_mapping(tree, include_page_ranges=False, max_page=None): diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py new file mode 100644 index 0000000..73a8236 --- /dev/null +++ b/tests/test_concurrency.py @@ -0,0 +1,197 @@ +import asyncio +import threading + +import pytest + +from pageindex.config import ( + IndexConfig, + _env_max_concurrency_default, + get_max_concurrency, + max_concurrency_scope, + set_max_concurrency, +) +from pageindex.index.utils import bounded_gather + + +@pytest.fixture(autouse=True) +def _restore_max_concurrency(): + """Keep tests isolated — the concurrency setting is a module global.""" + prev = get_max_concurrency() + yield + set_max_concurrency(prev) + + +def test_bounded_gather_never_exceeds_the_cap(): + set_max_concurrency(5) + state = {"in_flight": 0, "peak": 0} + + async def worker(i): + state["in_flight"] += 1 + state["peak"] = max(state["peak"], state["in_flight"]) + await asyncio.sleep(0.01) + state["in_flight"] -= 1 + return i + + async def run(): + return await bounded_gather(worker(i) for i in range(30)) + + results = asyncio.run(run()) + + # Order is preserved (gather semantics) and the cap is respected: with 30 + # tasks and 5 slots, exactly 5 run at once — never the unbounded 30 that + # exhausted file descriptors. + assert results == list(range(30)) + assert state["peak"] == 5 + + +def test_bounded_gather_propagates_return_exceptions(): + async def ok(): + return "ok" + + async def boom(): + raise ValueError("boom") + + async def run(): + return await bounded_gather([ok(), boom()], return_exceptions=True) + + results = asyncio.run(run()) + assert results[0] == "ok" + assert isinstance(results[1], ValueError) + + +def test_set_get_max_concurrency_round_trip(): + set_max_concurrency(3) + assert get_max_concurrency() == 3 + + +def test_set_max_concurrency_rejects_non_positive(): + with pytest.raises(ValueError): + set_max_concurrency(0) + with pytest.raises(ValueError): + set_max_concurrency(-1) + + +def test_env_default_parsing(monkeypatch): + monkeypatch.delenv("PAGEINDEX_MAX_CONCURRENCY", raising=False) + assert _env_max_concurrency_default() == 5 + monkeypatch.setenv("PAGEINDEX_MAX_CONCURRENCY", "20") + assert _env_max_concurrency_default() == 20 + monkeypatch.setenv("PAGEINDEX_MAX_CONCURRENCY", "garbage") + assert _env_max_concurrency_default() == 5 + monkeypatch.setenv("PAGEINDEX_MAX_CONCURRENCY", "0") + assert _env_max_concurrency_default() == 5 + + +def test_index_config_max_concurrency_field(): + # Default is None → "use the global/env default"; explicit value overrides. + assert IndexConfig().max_concurrency is None + assert IndexConfig(max_concurrency=7).max_concurrency == 7 + + +def test_max_concurrency_scope_overrides_then_restores(): + # A per-index override applies inside the scope and, crucially, does NOT + # stick as the new process default afterwards (Finding A: no stickiness). + set_max_concurrency(10) + with max_concurrency_scope(3): + assert get_max_concurrency() == 3 + assert get_max_concurrency() == 10 + + +def test_max_concurrency_scope_none_is_a_no_op(): + set_max_concurrency(8) + with max_concurrency_scope(None): + assert get_max_concurrency() == 8 + assert get_max_concurrency() == 8 + + +def test_max_concurrency_scope_rejects_non_positive(): + with pytest.raises(ValueError): + with max_concurrency_scope(0): + pass + with pytest.raises(ValueError): + with max_concurrency_scope(-1): + pass + + +def test_max_concurrency_scope_is_isolated_across_threads(): + # A per-index override in one indexing thread must not leak into another + # thread indexing a different document concurrently (Finding B). The + # override is a ContextVar, so it's invisible outside its own context. + set_max_concurrency(10) + seen = {} + barrier = threading.Barrier(2) + + def worker(): + with max_concurrency_scope(2): + barrier.wait() # let main read while we're inside the scope + seen["worker"] = get_max_concurrency() + barrier.wait() + + t = threading.Thread(target=worker) + t.start() + barrier.wait() + seen["main"] = get_max_concurrency() + barrier.wait() + t.join() + + assert seen["worker"] == 2 # worker sees its own scoped override + assert seen["main"] == 10 # main is unaffected by the worker's scope + + +def test_bounded_gather_respects_scoped_override(): + # bounded_gather reads the cap at semaphore-creation time; a surrounding + # max_concurrency_scope must win and must not mutate the process default. + set_max_concurrency(10) + state = {"in_flight": 0, "peak": 0} + + async def worker(i): + state["in_flight"] += 1 + state["peak"] = max(state["peak"], state["in_flight"]) + await asyncio.sleep(0.01) + state["in_flight"] -= 1 + return i + + async def run(): + with max_concurrency_scope(4): + return await bounded_gather(worker(i) for i in range(20)) + + asyncio.run(run()) + assert state["peak"] == 4 + assert get_max_concurrency() == 10 + + +def test_run_async_propagates_scope_into_worker_thread(): + # When build_index runs inside an already-running loop, _run_async hops to a + # worker thread. The max_concurrency_scope override must ride along (copied + # context) instead of silently falling back to the process default. + from pageindex.index.pipeline import _run_async + + set_max_concurrency(10) + state = {"in_flight": 0, "peak": 0} + + async def worker(i): + state["in_flight"] += 1 + state["peak"] = max(state["peak"], state["in_flight"]) + await asyncio.sleep(0.01) + state["in_flight"] -= 1 + return i + + async def inner(): + return await bounded_gather(worker(i) for i in range(20)) + + async def outer(): + # We're inside a running loop -> _run_async uses the worker thread. + with max_concurrency_scope(3): + _run_async(inner()) + + asyncio.run(outer()) + assert state["peak"] == 3 + + +def test_utils_star_import_does_not_leak_config_name(): + # `from .utils import *` (used by the page_index modules) must not export a + # name `config` that would shadow the real pageindex.config submodule for + # those modules (Finding D). The SimpleNamespace alias is now `_config`. + ns = {} + exec("from pageindex.index.utils import *", ns) + assert "config" not in ns