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
This commit is contained in:
mountain 2026-07-08 10:31:37 +08:00
parent 890b520b1c
commit e5392836a4
6 changed files with 358 additions and 44 deletions

197
tests/test_concurrency.py Normal file
View file

@ -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