mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-15 21:11:05 +02:00
9ad54122bb
- index/utils.py: fix an asyncio deadlock in _sync_llm_semaphore. Sync LLM calls (check_toc / process_no_toc / toc_transformer) run on the event-loop thread nested inside the async meta_processor, and its blocking ceiling acquire could wait forever for a permit held by async _llm_semaphore holders that can only release it once the (now-frozen) loop runs. Take the slot non-blocking when on a running loop; keep the blocking acquire off-loop. - index/utils.py: bound parse_pages ranges before materializing range() into the list, so a huge span like '1-2000000000' is rejected up front instead of exhausting memory before the 1000-page cap is ever checked (DoS). - storage/sqlite.py: bump a generation counter on close() so a thread that cached a connection in thread-local storage reconnects on its next call instead of reusing a closed handle (ProgrammingError). - backend/cloud.py: after connect, bail out of the SSE background thread if the consumer already abandoned the stream, instead of draining it in the background. Adds regression tests for each fix.
528 lines
19 KiB
Python
528 lines
19 KiB
Python
import asyncio
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from types import SimpleNamespace
|
|
|
|
import pydantic
|
|
import pytest
|
|
|
|
from pageindex.config import (
|
|
IndexConfig,
|
|
_env_llm_timeout_default,
|
|
_env_max_concurrency_default,
|
|
_max_concurrency_scope_semaphore,
|
|
get_llm_params,
|
|
get_max_concurrency,
|
|
llm_params_scope,
|
|
max_concurrency_scope,
|
|
set_llm_params,
|
|
set_max_concurrency,
|
|
)
|
|
from pageindex.index.utils import (
|
|
_llm_semaphore,
|
|
_process_ceiling_semaphore,
|
|
llm_acompletion,
|
|
llm_completion,
|
|
)
|
|
|
|
|
|
@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)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _restore_llm_params():
|
|
"""Keep tests isolated — llm params are a module global too."""
|
|
prev = get_llm_params()
|
|
yield
|
|
set_llm_params(**prev)
|
|
|
|
|
|
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_is_a_true_process_wide_ceiling_across_threads():
|
|
# The bug this fixes: each asyncio.run() (its own event loop) used to get
|
|
# an independent full-size semaphore, so N concurrently-indexing threads
|
|
# multiplied the effective cap by N. 2 threads, cap=3 -> combined peak
|
|
# must stay at 3, not 6.
|
|
set_max_concurrency(3)
|
|
state = {"in_flight": 0, "peak": 0}
|
|
lock = threading.Lock()
|
|
|
|
async def leaf():
|
|
async with _llm_semaphore():
|
|
with lock:
|
|
state["in_flight"] += 1
|
|
state["peak"] = max(state["peak"], state["in_flight"])
|
|
await asyncio.sleep(0.05)
|
|
with lock:
|
|
state["in_flight"] -= 1
|
|
|
|
async def load():
|
|
await asyncio.gather(*(leaf() for _ in range(5)))
|
|
|
|
threads = [threading.Thread(target=lambda: asyncio.run(load())) for _ in range(2)]
|
|
[t.start() for t in threads]
|
|
[t.join() for t in threads]
|
|
|
|
assert state["peak"] == 3
|
|
|
|
|
|
def test_llm_semaphore_cancellation_while_waiting_does_not_leak_a_permit():
|
|
# A blocking ceiling_sem.acquire() run via asyncio.to_thread() would leak a
|
|
# permit under cancellation: the worker thread can't be interrupted, so if
|
|
# the awaiting coroutine is cancelled while the thread is still parked
|
|
# inside acquire(), the thread can go on to actually acquire the permit
|
|
# *after* the coroutine already unwound, and the matching finally:
|
|
# release() never runs for that attempt. Cancel a task waiting on an
|
|
# already-exhausted ceiling and confirm the permit count fully recovers.
|
|
set_max_concurrency(1)
|
|
|
|
async def run():
|
|
async def hold():
|
|
async with _llm_semaphore():
|
|
await asyncio.sleep(10)
|
|
|
|
holder = asyncio.create_task(hold())
|
|
await asyncio.sleep(0.1) # let it acquire the single permit
|
|
|
|
waiter = asyncio.create_task(_llm_semaphore().__aenter__())
|
|
await asyncio.sleep(0.1) # let it start waiting for the permit
|
|
waiter.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await waiter
|
|
|
|
holder.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await holder
|
|
|
|
await asyncio.sleep(0.2) # give any orphaned acquire a chance to land
|
|
assert _process_ceiling_semaphore()._value == 1
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_scoped_semaphore_cancellation_while_waiting_does_not_over_release():
|
|
# Mirror of the ceiling test for the scoped override: a coroutine cancelled
|
|
# while polling for a scoped permit must NOT let the finally release a permit
|
|
# it never acquired, which would inflate the scoped cap for later calls.
|
|
set_max_concurrency(5) # high ceiling so the scope is the narrower cap
|
|
|
|
async def run():
|
|
with max_concurrency_scope(1):
|
|
sem = _max_concurrency_scope_semaphore()
|
|
assert sem._value == 1
|
|
|
|
async def hold():
|
|
async with _llm_semaphore():
|
|
await asyncio.sleep(10)
|
|
|
|
holder = asyncio.create_task(hold())
|
|
await asyncio.sleep(0.1) # holder takes the single scoped permit
|
|
|
|
waiter = asyncio.create_task(_llm_semaphore().__aenter__())
|
|
await asyncio.sleep(0.1) # waiter is now polling for the scoped permit
|
|
waiter.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await waiter
|
|
|
|
holder.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await holder
|
|
|
|
await asyncio.sleep(0.2)
|
|
assert sem._value == 1 # fully recovered, not inflated to 2
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
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 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 actually acquires the shared cap around the async
|
|
# 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 SimpleNamespace(
|
|
choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))]
|
|
)
|
|
|
|
monkeypatch.setattr("litellm.acompletion", fake_acompletion)
|
|
|
|
async def run():
|
|
await asyncio.gather(*(llm_acompletion("gpt-x", f"p{i}") for i in range(20)))
|
|
|
|
asyncio.run(run())
|
|
assert state["peak"] == 3
|
|
|
|
|
|
def test_llm_acompletion_passes_a_timeout_to_litellm(monkeypatch):
|
|
# A per-request timeout must reach litellm so a hung / half-open connection
|
|
# fails fast instead of stalling indexing forever. It rides get_llm_params(),
|
|
# so both the default and an override flow through automatically.
|
|
seen = {}
|
|
|
|
async def fake_acompletion(**kwargs):
|
|
seen.update(kwargs)
|
|
return SimpleNamespace(
|
|
choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))]
|
|
)
|
|
|
|
monkeypatch.setattr("litellm.acompletion", fake_acompletion)
|
|
asyncio.run(llm_acompletion("gpt-x", "hi"))
|
|
assert "timeout" in seen and seen["timeout"] == get_llm_params()["timeout"]
|
|
|
|
|
|
def test_llm_completion_degrades_to_empty_on_exhaustion(monkeypatch, caplog):
|
|
# A persistently-failing LLM call must NOT abort the whole index: it returns
|
|
# an empty result (logged at WARNING, not silent) so callers degrade
|
|
# (extract_json('') -> {} -> .get(default)) and the rest still indexes.
|
|
import logging
|
|
|
|
def boom(**kwargs):
|
|
raise RuntimeError("provider down")
|
|
|
|
monkeypatch.setattr("litellm.completion", boom)
|
|
monkeypatch.setattr("pageindex.index.utils.time.sleep", lambda *a: None) # skip retry backoff
|
|
|
|
with caplog.at_level(logging.WARNING, logger="pageindex.index.utils"):
|
|
assert llm_completion("gpt-x", "hi") == ""
|
|
assert llm_completion("gpt-x", "hi", return_finish_reason=True) == ("", "error")
|
|
assert any("failed after" in r.message and "degrading" in r.message for r in caplog.records)
|
|
|
|
|
|
def test_llm_acompletion_degrades_to_empty_on_exhaustion(monkeypatch):
|
|
# Async counterpart: exhausted retries return "" instead of raising, so the
|
|
# return_exceptions gathers see a plain empty result and callers degrade.
|
|
async def boom(**kwargs):
|
|
raise RuntimeError("provider down")
|
|
|
|
async def _instant_sleep(*a):
|
|
return None
|
|
|
|
monkeypatch.setattr("litellm.acompletion", boom)
|
|
monkeypatch.setattr("pageindex.index.utils.asyncio.sleep", _instant_sleep) # skip retry backoff
|
|
|
|
assert asyncio.run(llm_acompletion("gpt-x", "hi")) == ""
|
|
|
|
|
|
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_sync_llm_completion_on_event_loop_does_not_deadlock(monkeypatch):
|
|
# Regression: _sync_llm_semaphore used a BLOCKING ceiling acquire. Sync LLM
|
|
# helpers (check_toc, process_no_toc, toc_transformer, …) run synchronously
|
|
# ON the event loop (nested inside the async meta_processor). If async
|
|
# llm_acompletion holders occupy every ceiling permit across their awaits,
|
|
# a blocking acquire froze the loop -> the holders could never resume to
|
|
# release their permits -> permanent deadlock. The sync path must never
|
|
# block the running loop.
|
|
set_max_concurrency(2)
|
|
|
|
async def fake_acompletion(**kwargs):
|
|
await asyncio.sleep(0.3) # hold a ceiling permit across the await
|
|
return SimpleNamespace(
|
|
choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))]
|
|
)
|
|
|
|
def fake_completion(**kwargs):
|
|
return SimpleNamespace(
|
|
choices=[SimpleNamespace(
|
|
message=SimpleNamespace(content="sync-ok"), finish_reason="stop")]
|
|
)
|
|
|
|
monkeypatch.setattr("litellm.acompletion", fake_acompletion)
|
|
monkeypatch.setattr("litellm.completion", fake_completion)
|
|
|
|
async def run():
|
|
# Both ceiling permits taken by async holders, held across their await.
|
|
holders = [asyncio.create_task(llm_acompletion("m", f"p{i}")) for i in range(2)]
|
|
await asyncio.sleep(0.05) # let them acquire the permits
|
|
# Sync call on the loop thread: pre-fix this blocks forever waiting for a
|
|
# permit the holders own and can't release (loop is frozen).
|
|
result = llm_completion("m", "sync")
|
|
await asyncio.gather(*holders)
|
|
return result
|
|
|
|
# Run in a thread with a join timeout so a regression FAILS instead of
|
|
# hanging CI: a real deadlock freezes the loop, so asyncio.wait_for can't
|
|
# cancel it (its timeout callback never runs on the frozen loop).
|
|
box = {}
|
|
|
|
def target():
|
|
box["result"] = asyncio.run(run())
|
|
|
|
t = threading.Thread(target=target, daemon=True)
|
|
t.start()
|
|
t.join(timeout=8)
|
|
assert not t.is_alive(), "deadlock: sync llm_completion blocked the event loop"
|
|
assert box["result"] == "sync-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
|
|
|
|
set_max_concurrency(10)
|
|
state = {"in_flight": 0, "peak": 0}
|
|
|
|
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))
|
|
|
|
asyncio.run(outer())
|
|
assert state["peak"] == 3
|
|
|
|
|
|
def test_set_get_max_concurrency_round_trip():
|
|
set_max_concurrency(3)
|
|
assert get_max_concurrency() == 3
|
|
|
|
|
|
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):
|
|
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_env_llm_timeout_parsing(monkeypatch):
|
|
monkeypatch.delenv("PAGEINDEX_LLM_TIMEOUT", raising=False)
|
|
assert _env_llm_timeout_default() == 120
|
|
monkeypatch.setenv("PAGEINDEX_LLM_TIMEOUT", "45")
|
|
assert _env_llm_timeout_default() == 45
|
|
monkeypatch.setenv("PAGEINDEX_LLM_TIMEOUT", "garbage")
|
|
assert _env_llm_timeout_default() == 120
|
|
monkeypatch.setenv("PAGEINDEX_LLM_TIMEOUT", "0") # <=0 opts out of the timeout
|
|
assert _env_llm_timeout_default() is None
|
|
|
|
|
|
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_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 (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_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. 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_llm_params_scope_overrides_then_restores():
|
|
set_llm_params(temperature=0)
|
|
with llm_params_scope({"temperature": 1}):
|
|
assert get_llm_params()["temperature"] == 1
|
|
assert get_llm_params()["temperature"] == 0
|
|
|
|
|
|
def test_llm_params_scope_none_is_a_no_op():
|
|
set_llm_params(temperature=0)
|
|
with llm_params_scope(None):
|
|
assert get_llm_params()["temperature"] == 0
|
|
assert get_llm_params()["temperature"] == 0
|
|
|
|
|
|
def test_llm_params_scope_rejects_reserved_keys():
|
|
with pytest.raises(ValueError):
|
|
with llm_params_scope({"model": "x"}):
|
|
pass
|
|
|
|
|
|
def test_llm_params_scope_is_isolated_across_threads():
|
|
set_llm_params(temperature=0)
|
|
seen = {}
|
|
barrier = threading.Barrier(2)
|
|
|
|
def worker():
|
|
with llm_params_scope({"temperature": 1}):
|
|
barrier.wait()
|
|
seen["worker"] = get_llm_params()["temperature"]
|
|
barrier.wait()
|
|
|
|
t = threading.Thread(target=worker)
|
|
t.start()
|
|
barrier.wait()
|
|
seen["main"] = get_llm_params()["temperature"]
|
|
barrier.wait()
|
|
t.join()
|
|
|
|
assert seen["worker"] == 1
|
|
assert seen["main"] == 0
|
|
|
|
|
|
def test_llm_params_scope_does_not_leak_across_concurrent_indexing():
|
|
# The bug this fixes: set_llm_params() mutates a bare process-wide dict, so
|
|
# two documents indexed concurrently with different llm_params_scope()
|
|
# overrides must not see each other's temperature.
|
|
set_llm_params(temperature=0)
|
|
seen = {"a": None, "b": None}
|
|
|
|
async def job(name, temperature, delay_before, delay_after):
|
|
with llm_params_scope({"temperature": temperature}):
|
|
await asyncio.sleep(delay_before)
|
|
seen[name] = get_llm_params()["temperature"]
|
|
await asyncio.sleep(delay_after)
|
|
|
|
async def run():
|
|
await asyncio.gather(
|
|
job("a", 1, 0.0, 0.05),
|
|
job("b", 2, 0.02, 0.0),
|
|
)
|
|
|
|
asyncio.run(run())
|
|
assert seen["a"] == 1
|
|
assert seen["b"] == 2
|
|
|
|
|
|
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. The SimpleNamespace alias is now `_config`.
|
|
ns = {}
|
|
exec("from pageindex.index.utils import *", ns)
|
|
assert "config" not in ns
|