diff --git a/pageindex/agent.py b/pageindex/agent.py index a186fa4..2592ec0 100644 --- a/pageindex/agent.py +++ b/pageindex/agent.py @@ -159,6 +159,10 @@ class AgentRunner: result = Runner.run_sync(agent, question) else: import concurrent.futures + import contextvars + # Copy the current context into the worker thread so ContextVar-based + # settings propagate (mirrors pipeline._run_async). + ctx = contextvars.copy_context() with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: - result = pool.submit(asyncio.run, Runner.run(agent, question)).result() + result = pool.submit(ctx.run, asyncio.run, Runner.run(agent, question)).result() return result.final_output diff --git a/pageindex/config.py b/pageindex/config.py index 1b6f060..e1d48a8 100644 --- a/pageindex/config.py +++ b/pageindex/config.py @@ -5,7 +5,7 @@ import os from contextlib import contextmanager from contextvars import ContextVar -from pydantic import BaseModel +from pydantic import BaseModel, field_validator class IndexConfig(BaseModel): @@ -30,6 +30,16 @@ class IndexConfig(BaseModel): # An explicit value here wins for this client. max_concurrency: int | None = None + @field_validator("max_concurrency", mode="before") + @classmethod + def _validate_max_concurrency_field(cls, v): + # Reject bool before pydantic coerces True->1 / False->0, and reject + # non-positive ints, so a bad value fails loudly instead of silently + # serializing (Semaphore(1)) or crashing (Semaphore(0)). + if v is not None: + _validate_max_concurrency(v) + return v + def _env_drop_params_default() -> bool: return os.getenv("PAGEINDEX_DROP_PARAMS", "true").strip().lower() not in ( @@ -90,6 +100,17 @@ _MAX_CONCURRENCY_OVERRIDE: ContextVar[int | None] = ContextVar( ) +def _validate_max_concurrency(value) -> None: + """Raise ValueError unless ``value`` is a positive int. + + ``bool`` is an ``int`` subclass, so it's rejected explicitly — otherwise + ``set_max_concurrency(True)`` would pass and become ``Semaphore(1)``, + silently serializing all indexing instead of failing loudly. + """ + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError("max_concurrency must be a positive integer") + + def get_max_concurrency() -> int: """Return the effective cap on concurrent in-flight LLM calls during indexing. @@ -103,8 +124,7 @@ def get_max_concurrency() -> int: 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") + _validate_max_concurrency(value) _MAX_CONCURRENCY = value @@ -117,8 +137,8 @@ def max_concurrency_scope(value: int | None): 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") + if value is not None: + _validate_max_concurrency(value) token = _MAX_CONCURRENCY_OVERRIDE.set(value) try: yield diff --git a/pageindex/index/page_index.py b/pageindex/index/page_index.py index 862f358..9687a28 100644 --- a/pageindex/index/page_index.py +++ b/pageindex/index/page_index.py @@ -75,21 +75,28 @@ async def check_title_appearance_in_start_concurrent(structure, page_list, model if logger: logger.info("Checking title appearance in start concurrently") - # skip items without physical_index + # Mark items we can't check as 'no' up front: missing physical_index, or one + # out of range for page_list. An out-of-range index (the LLM can emit one) + # would otherwise raise IndexError below — during task-list construction, + # outside the gather's return_exceptions protection — and abort the build. + def _valid_physical_index(item): + idx = item.get('physical_index') + return idx is not None and 1 <= idx <= len(page_list) + for item in structure: - if item.get('physical_index') is None: + if not _valid_physical_index(item): item['appear_start'] = 'no' - # only for items with valid physical_index + # only for items with a valid, in-range physical_index tasks = [] valid_items = [] for item in structure: - if item.get('physical_index') is not None: + if _valid_physical_index(item): page_text = page_list[item['physical_index'] - 1][0] tasks.append(check_title_appearance_in_start(item['title'], page_text, model=model, logger=logger)) valid_items.append(item) - results = await bounded_gather(tasks, return_exceptions=True) + results = await asyncio.gather(*tasks, return_exceptions=True) for item, result in zip(valid_items, results): if isinstance(result, Exception): if logger: @@ -832,7 +839,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 bounded_gather(tasks, return_exceptions=True) + results = await asyncio.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 +934,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 bounded_gather(tasks) + results = await asyncio.gather(*tasks) # Process results correct_count = 0 @@ -1015,7 +1022,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 bounded_gather(tasks) + await asyncio.gather(*tasks) return node @@ -1051,7 +1058,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 bounded_gather(tasks) + await asyncio.gather(*tasks) return toc_tree diff --git a/pageindex/index/page_index_md.py b/pageindex/index/page_index_md.py index a115b9a..f9e300a 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 bounded_gather(tasks) + summaries = await asyncio.gather(*tasks) for node, summary in zip(nodes, summaries): if not node.get('nodes'): diff --git a/pageindex/index/utils.py b/pageindex/index/utils.py index a850a23..e8be39e 100644 --- a/pageindex/index/utils.py +++ b/pageindex/index/utils.py @@ -7,6 +7,8 @@ import json import copy import re import asyncio +import threading +import weakref import PyPDF2 import pymupdf import yaml @@ -25,28 +27,41 @@ 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. +# One shared semaphore per event loop, bounding concurrent in-flight LLM calls. +# Keyed by the loop object (WeakKeyDictionary drops the entry once the loop is +# closed and garbage-collected) so each asyncio.run() gets its own, correctly +# loop-bound semaphore. The lock only guards the tiny get-or-create against two +# threads (each driving its own loop) racing to insert; within a single loop +# everything is single-threaded, so no lock is needed on the hot path. +_LLM_SEMAPHORES: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +_LLM_SEMAPHORES_LOCK = threading.Lock() - 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``. +def _llm_semaphore() -> asyncio.Semaphore: + """Shared per-loop cap on concurrent in-flight LLM calls. + + Acquired only around the leaf ``litellm.acompletion`` call in + ``llm_acompletion`` — the single point every LLM request funnels through — + so the cap is a TRUE global bound no matter how deeply the indexing gathers + nest (``tree_parser`` → ``process_large_node_recursively`` → …). Bounding at + the leaf rather than at each gather call site is also deadlock-free: a parent + coroutine awaiting its children holds no slot, so children can always + acquire one. + + Sized from ``get_max_concurrency()`` the first time it's needed in a loop, so + a per-index ``max_concurrency_scope`` override in effect at that moment is + honored. Without this bound a many-node document opens one socket per node at + once and exhausts the process file-descriptor limit (Errno 24). """ - 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 - ) + loop = asyncio.get_running_loop() + sem = _LLM_SEMAPHORES.get(loop) + if sem is None: + with _LLM_SEMAPHORES_LOCK: + sem = _LLM_SEMAPHORES.get(loop) + if sem is None: + sem = asyncio.Semaphore(get_max_concurrency()) + _LLM_SEMAPHORES[loop] = sem + return sem def llm_completion(model, prompt, chat_history=None, return_finish_reason=False): @@ -86,11 +101,14 @@ async def llm_acompletion(model, prompt): messages = [{"role": "user", "content": prompt}] for i in range(max_retries): try: - response = await litellm.acompletion( - model=model, - messages=messages, - **get_llm_params(), # per-call kwargs; never the litellm global - ) + # Hold a concurrency slot only around the actual network call — not + # across retry backoff — so the cap counts real in-flight requests. + async with _llm_semaphore(): + response = await litellm.acompletion( + model=model, + messages=messages, + **get_llm_params(), # per-call kwargs; never the litellm global + ) return response.choices[0].message.content except Exception as e: logger.warning("Retrying async LLM completion (%d/%d)", i + 1, max_retries) @@ -240,7 +258,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 bounded_gather(tasks) + summaries = await asyncio.gather(*tasks) for node, summary in zip(nodes, summaries): node['summary'] = summary diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py index 73a8236..a9234ad 100644 --- a/tests/test_concurrency.py +++ b/tests/test_concurrency.py @@ -1,6 +1,8 @@ import asyncio import threading +from types import SimpleNamespace +import pydantic import pytest from pageindex.config import ( @@ -10,7 +12,7 @@ from pageindex.config import ( max_concurrency_scope, set_max_concurrency, ) -from pageindex.index.utils import bounded_gather +from pageindex.index.utils import _llm_semaphore, llm_acompletion @pytest.fixture(autouse=True) @@ -21,42 +23,90 @@ def _restore_max_concurrency(): set_max_concurrency(prev) -def test_bounded_gather_never_exceeds_the_cap(): - set_max_concurrency(5) +async def _nested_llm_load(state, *, branches=5, leaves=5): + """Drive branches*leaves leaf calls, nested two levels deep, each holding + the shared per-loop LLM semaphore — the exact shape of the indexing pipeline + (tree_parser gather -> per-node gather -> leaf LLM call).""" + + async def leaf(): + async with _llm_semaphore(): + state["in_flight"] += 1 + state["peak"] = max(state["peak"], state["in_flight"]) + await asyncio.sleep(0.01) + state["in_flight"] -= 1 + + async def branch(): + await asyncio.gather(*(leaf() for _ in range(leaves))) + + await asyncio.gather(*(branch() for _ in range(branches))) + + +def test_llm_semaphore_bounds_concurrency_even_when_nested(): + # The core fix: the cap is a TRUE global bound even when acquired from + # deeply nested gathers. 25 leaf calls nested two levels, cap 3 -> peak 3. + # A per-gather-call semaphore (the previous design) would let this reach + # branches*leaves and blow past the cap. + set_max_concurrency(3) + state = {"in_flight": 0, "peak": 0} + asyncio.run(_nested_llm_load(state, branches=5, leaves=5)) + assert state["peak"] == 3 + + +def test_llm_semaphore_uses_scoped_override(): + # A per-index max_concurrency_scope active when the loop's semaphore is first + # created must set its size, and must not mutate the process default. + set_max_concurrency(10) state = {"in_flight": 0, "peak": 0} - async def worker(i): + async def run(): + with max_concurrency_scope(2): + await _nested_llm_load(state, branches=4, leaves=4) + + asyncio.run(run()) + assert state["peak"] == 2 + assert get_max_concurrency() == 10 + + +def test_llm_acompletion_holds_the_shared_semaphore(monkeypatch): + # Prove llm_acompletion (the single chokepoint every LLM call funnels + # through) actually acquires the shared cap around the network call. + set_max_concurrency(3) + state = {"in_flight": 0, "peak": 0} + + async def fake_acompletion(**kwargs): state["in_flight"] += 1 state["peak"] = max(state["peak"], state["in_flight"]) await asyncio.sleep(0.01) state["in_flight"] -= 1 - return i + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))] + ) + + monkeypatch.setattr("litellm.acompletion", fake_acompletion) async def run(): - return await bounded_gather(worker(i) for i in range(30)) + await asyncio.gather(*(llm_acompletion("gpt-x", f"p{i}") for i in range(20))) - 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 + asyncio.run(run()) + assert state["peak"] == 3 -def test_bounded_gather_propagates_return_exceptions(): - async def ok(): - return "ok" +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) and still bound the (nested) LLM load in that worker loop. + from pageindex.index.pipeline import _run_async - async def boom(): - raise ValueError("boom") + set_max_concurrency(10) + state = {"in_flight": 0, "peak": 0} - async def run(): - return await bounded_gather([ok(), boom()], return_exceptions=True) + async def outer(): + # We're inside a running loop -> _run_async uses the worker thread. + with max_concurrency_scope(3): + _run_async(_nested_llm_load(state, branches=4, leaves=4)) - results = asyncio.run(run()) - assert results[0] == "ok" - assert isinstance(results[1], ValueError) + asyncio.run(outer()) + assert state["peak"] == 3 def test_set_get_max_concurrency_round_trip(): @@ -64,11 +114,11 @@ def test_set_get_max_concurrency_round_trip(): 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_set_max_concurrency_rejects_invalid(): + # bool is an int subclass -> must be rejected, not silently -> Semaphore(1). + for bad in (0, -1, True, False, 2.5, "3", None): + with pytest.raises(ValueError): + set_max_concurrency(bad) def test_env_default_parsing(monkeypatch): @@ -88,9 +138,16 @@ def test_index_config_max_concurrency_field(): assert IndexConfig(max_concurrency=7).max_concurrency == 7 +def test_index_config_rejects_bool_and_non_positive_max_concurrency(): + # bool would otherwise be coerced by pydantic to 1/0; both must be rejected. + for bad in (True, False, 0, -1): + with pytest.raises(pydantic.ValidationError): + IndexConfig(max_concurrency=bad) + + 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). + # stick as the new process default afterwards (no stickiness). set_max_concurrency(10) with max_concurrency_scope(3): assert get_max_concurrency() == 3 @@ -104,19 +161,17 @@ def test_max_concurrency_scope_none_is_a_no_op(): 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_rejects_invalid(): + for bad in (0, -1, True, False): + with pytest.raises(ValueError): + with max_concurrency_scope(bad): + 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. + # thread indexing a different document concurrently. The override is a + # ContextVar, so it's invisible outside its own context. set_max_concurrency(10) seen = {} barrier = threading.Barrier(2) @@ -138,60 +193,10 @@ def test_max_concurrency_scope_is_isolated_across_threads(): 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`. + # those modules. The SimpleNamespace alias is now `_config`. ns = {} exec("from pageindex.index.utils import *", ns) assert "config" not in ns diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 9e1e54e..ce66b63 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -93,3 +93,22 @@ def test_null_logger_methods(): logger.error("test error") logger.debug("test debug") logger.info({"key": "value"}) + + +def test_check_title_appearance_tolerates_out_of_range_physical_index(): + """An LLM-emitted physical_index outside page_list must be marked 'no', not + raise IndexError (which happens during task construction, outside the + gather's return_exceptions protection, and would abort the whole build).""" + from pageindex.index.page_index import check_title_appearance_in_start_concurrent + + page_list = [("only page text", 3)] # length 1 + structure = [ + {"title": "A", "physical_index": 5}, # out of range -> would IndexError + {"title": "B", "physical_index": 0}, # 0 -> would wrap to page_list[-1] + {"title": "C", "physical_index": None}, # missing + {"title": "D"}, # no physical_index key at all + ] + result = asyncio.run( + check_title_appearance_in_start_concurrent(structure, page_list) + ) + assert all(item["appear_start"] == "no" for item in result)