Fix sync LLM concurrency limit

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)
 
 

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"

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