mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-15 21:11:05 +02:00
fix(index): bound LLM concurrency at the leaf, not per gather call
The previous approach (bounded_gather building a fresh semaphore per call) did NOT compose: the indexing call graph nests gathers (tree_parser -> process_large_node_recursively -> recurse, plus each node's check_title gather), and each level got its own independent cap. Peak in-flight LLM calls grew ~N^depth, so a deep/wide document still exhausted file descriptors (Errno 24) — the exact failure the cap was meant to prevent — while the flat summary phase was over-serialized. Naively sharing one semaphore across gather levels would instead deadlock (a parent holds a slot while awaiting children that need slots). Move the throttle to the single chokepoint every LLM call funnels through, llm_acompletion: one shared semaphore per event loop, acquired only around the litellm.acompletion network call. This gives a true global cap that composes across any nesting and can't deadlock (a parent awaiting children holds no slot). bounded_gather is gone; the call sites revert to plain asyncio.gather. Also: - Reject bool in max_concurrency validation (bool is an int subclass, so set_max_concurrency(True) / IndexConfig(max_concurrency=True) previously became Semaphore(1) and silently serialized). Shared _validate_max_concurrency + a pydantic field_validator. - Guard check_title_appearance_in_start_concurrent against an out-of-range or 0 physical_index (LLM can emit one): it was dereferenced during task construction, outside the gather's return_exceptions protection, aborting the whole build; 0 silently wrapped to the last page. Now marked 'no'. - Propagate contextvars into agent.py's worker-thread run (mirrors pipeline._run_async) so ContextVar settings stay consistent. - Tests rewritten to cover the nested case the old flat tests missed, the leaf-level throttle in llm_acompletion, bool rejection, and the out-of-range physical_index guard. Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
This commit is contained in:
mountain
parent
e5392836a4
commit
2d46d68052
7 changed files with 203 additions and 130 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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'):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue