diff --git a/pageindex/config.py b/pageindex/config.py index dad0fc5..b40e285 100644 --- a/pageindex/config.py +++ b/pageindex/config.py @@ -2,6 +2,7 @@ from __future__ import annotations import os +import threading from contextlib import contextmanager from contextvars import ContextVar @@ -142,6 +143,9 @@ _MAX_CONCURRENCY: int = _env_max_concurrency_default() _MAX_CONCURRENCY_OVERRIDE: ContextVar[int | None] = ContextVar( "pageindex_max_concurrency_override", default=None ) +_MAX_CONCURRENCY_SCOPE_SEMAPHORE: ContextVar[threading.Semaphore | None] = ContextVar( + "pageindex_max_concurrency_scope_semaphore", default=None +) def _validate_max_concurrency(value) -> None: @@ -175,6 +179,15 @@ def _process_wide_max_concurrency() -> int: return _MAX_CONCURRENCY +def _max_concurrency_scope_semaphore() -> threading.Semaphore | None: + """Return the semaphore backing the active scoped override, if any. + + Internal helper used by the LLM call sites so sync and async completions + share the same scoped cap when a context is copied across threads. + """ + return _MAX_CONCURRENCY_SCOPE_SEMAPHORE.get() + + def set_max_concurrency(value: int) -> None: """Set the process-wide default cap on concurrent in-flight LLM calls.""" global _MAX_CONCURRENCY @@ -193,10 +206,13 @@ def max_concurrency_scope(value: int | None): """ if value is not None: _validate_max_concurrency(value) + scoped_sem = threading.Semaphore(value) if value is not None else None token = _MAX_CONCURRENCY_OVERRIDE.set(value) + sem_token = _MAX_CONCURRENCY_SCOPE_SEMAPHORE.set(scoped_sem) try: yield finally: + _MAX_CONCURRENCY_SCOPE_SEMAPHORE.reset(sem_token) _MAX_CONCURRENCY_OVERRIDE.reset(token) diff --git a/pageindex/index/utils.py b/pageindex/index/utils.py index 7267fbc..1156086 100644 --- a/pageindex/index/utils.py +++ b/pageindex/index/utils.py @@ -8,7 +8,6 @@ import copy import re import asyncio import threading -import weakref import PyPDF2 import pymupdf import yaml @@ -21,9 +20,14 @@ from pprint import pprint # `pageindex.config` submodule for those modules. from types import SimpleNamespace as _config -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, contextmanager -from ..config import get_llm_params, get_max_concurrency, _process_wide_max_concurrency +from ..config import ( + get_llm_params, + get_max_concurrency, + _max_concurrency_scope_semaphore, + _process_wide_max_concurrency, +) from ..tokens import count_tokens # re-exported for backward compat logger = logging.getLogger(__name__) @@ -41,16 +45,6 @@ _PROCESS_LLM_SEMAPHORE: threading.Semaphore | None = None _PROCESS_LLM_SEMAPHORE_SIZE: int | None = None _PROCESS_LLM_SEMAPHORE_LOCK = threading.Lock() -# Per-loop, per-size semaphores for a max_concurrency_scope() override that's -# narrower than the process ceiling — isolates one call's own subtree to a -# tighter self-imposed limit without needing to be cross-thread itself (it can -# never let MORE calls through than the process ceiling above already allows, -# since both are held simultaneously; see _llm_semaphore). Keyed by (loop, size) -# rather than just loop so a later scope with a different size in the same loop -# isn't silently ignored. -_SCOPED_LLM_SEMAPHORES: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() -_SCOPED_LLM_SEMAPHORES_LOCK = threading.Lock() - def _process_ceiling_semaphore() -> threading.Semaphore: global _PROCESS_LLM_SEMAPHORE, _PROCESS_LLM_SEMAPHORE_SIZE @@ -62,33 +56,14 @@ def _process_ceiling_semaphore() -> threading.Semaphore: return _PROCESS_LLM_SEMAPHORE -def _scoped_llm_semaphore(size: int) -> asyncio.Semaphore: - loop = asyncio.get_running_loop() - per_loop = _SCOPED_LLM_SEMAPHORES.get(loop) - if per_loop is None: - with _SCOPED_LLM_SEMAPHORES_LOCK: - per_loop = _SCOPED_LLM_SEMAPHORES.get(loop) - if per_loop is None: - per_loop = {} - _SCOPED_LLM_SEMAPHORES[loop] = per_loop - sem = per_loop.get(size) - if sem is None: - with _SCOPED_LLM_SEMAPHORES_LOCK: - sem = per_loop.get(size) - if sem is None: - sem = asyncio.Semaphore(size) - per_loop[size] = sem - return sem - - @asynccontextmanager async def _llm_semaphore(): """Bound concurrent in-flight LLM calls to a TRUE process-wide ceiling, optionally narrowed further by an active max_concurrency_scope() override. Acquired only around the leaf ``litellm.acompletion`` call in - ``llm_acompletion`` — the single point every LLM request funnels through — - so the cap holds no matter how deeply the indexing gathers nest + ``llm_acompletion`` — sync calls use ``_sync_llm_semaphore`` below — so the + cap holds no matter how deeply the indexing gathers nest (``tree_parser`` → ``process_large_node_recursively`` → …) AND no matter how many threads are each running their own indexing job concurrently. Bounding at the leaf rather than at each gather call site is also deadlock-free: a @@ -97,7 +72,7 @@ async def _llm_semaphore(): The process ceiling (threading.Semaphore, shared cross-thread) is sized from the process-wide default only; a narrower max_concurrency_scope() override - is enforced as a second, nested, per-loop restriction — it can only + is enforced as a second, nested context-local restriction — it can only *tighten* the effective cap for its own call tree, never widen it past the ceiling. Without the outer bound a many-node document opens one socket per node at once and exhausts the process file-descriptor limit (Errno 24). @@ -115,15 +90,45 @@ async def _llm_semaphore(): # we've already given up on it. while not ceiling_sem.acquire(False): await asyncio.sleep(0.05) + scoped_sem = None try: effective = get_max_concurrency() ceiling = _process_wide_max_concurrency() if effective < ceiling: - async with _scoped_llm_semaphore(effective): - yield + scoped_sem = _max_concurrency_scope_semaphore() + if scoped_sem is not None: + while not scoped_sem.acquire(False): + await asyncio.sleep(0.05) + yield else: yield finally: + if scoped_sem is not None: + scoped_sem.release() + ceiling_sem.release() + + +@contextmanager +def _sync_llm_semaphore(): + """Synchronous companion to ``_llm_semaphore`` for ``llm_completion``. + + It uses the same process-wide ceiling so sync and async LLM calls share one + real cap. A scoped override can only narrow that cap for the active context. + """ + ceiling_sem = _process_ceiling_semaphore() + ceiling_sem.acquire() + scoped_sem = None + try: + effective = get_max_concurrency() + ceiling = _process_wide_max_concurrency() + if effective < ceiling: + scoped_sem = _max_concurrency_scope_semaphore() + if scoped_sem is not None: + scoped_sem.acquire() + yield + finally: + if scoped_sem is not None: + scoped_sem.release() ceiling_sem.release() @@ -134,13 +139,16 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}] for i in range(max_retries): try: - response = litellm.completion( - model=model, - messages=messages, - # Per-call litellm kwargs (default temperature=0, drop_params=True); - # configure via config.set_llm_params(...) — never the litellm global. - **get_llm_params(), - ) + # Hold a concurrency slot only around the actual network call, not + # retry backoff, so sync completions obey the same cap as async ones. + with _sync_llm_semaphore(): + response = litellm.completion( + model=model, + messages=messages, + # Per-call litellm kwargs (default temperature=0, drop_params=True); + # configure via config.set_llm_params(...) — never the litellm global. + **get_llm_params(), + ) content = response.choices[0].message.content if return_finish_reason: finish_reason = "max_output_reached" if response.choices[0].finish_reason == "length" else "finished" diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py index 391249b..aa697e2 100644 --- a/tests/test_concurrency.py +++ b/tests/test_concurrency.py @@ -1,5 +1,7 @@ import asyncio import threading +import time +from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace import pydantic @@ -16,7 +18,12 @@ from pageindex.config import ( set_llm_params, set_max_concurrency, ) -from pageindex.index.utils import _llm_semaphore, _process_ceiling_semaphore, llm_acompletion +from pageindex.index.utils import ( + _llm_semaphore, + _process_ceiling_semaphore, + llm_acompletion, + llm_completion, +) @pytest.fixture(autouse=True) @@ -142,8 +149,8 @@ def test_llm_semaphore_uses_scoped_override(): 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. + # Prove llm_acompletion actually acquires the shared cap around the async + # network call. set_max_concurrency(3) state = {"in_flight": 0, "peak": 0} @@ -182,6 +189,39 @@ def test_llm_acompletion_passes_a_timeout_to_litellm(monkeypatch): assert "timeout" in seen and seen["timeout"] == get_llm_params()["timeout"] +def test_llm_completion_holds_the_shared_semaphore(monkeypatch): + # Sync litellm.completion calls must share the same process-wide cap as the + # async path; otherwise concurrent indexing threads can exceed + # set_max_concurrency(). + set_max_concurrency(1) + state = {"in_flight": 0, "peak": 0} + lock = threading.Lock() + + def fake_completion(**kwargs): + with lock: + state["in_flight"] += 1 + state["peak"] = max(state["peak"], state["in_flight"]) + time.sleep(0.02) + with lock: + state["in_flight"] -= 1 + return SimpleNamespace( + choices=[ + SimpleNamespace( + message=SimpleNamespace(content="ok"), + finish_reason="stop", + ) + ] + ) + + monkeypatch.setattr("litellm.completion", fake_completion) + + with ThreadPoolExecutor(max_workers=3) as pool: + results = list(pool.map(lambda i: llm_completion("gpt-x", f"p{i}"), range(6))) + + assert results == ["ok"] * 6 + assert state["peak"] == 1 + + 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