mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-18 21:21:05 +02:00
fix: scoped concurrency semaphore over-release on cancellation
The sync-concurrency fix marked the scoped-override semaphore for release before it was actually acquired, so a coroutine cancelled while polling for a scoped permit (or a sync acquire interrupted mid-wait) ran the finally and released a permit it never held — inflating the scoped cap for later calls (the mirror of the ceiling-leak fix). Only bind the release guard after the acquire succeeds, in both the async and sync semaphores. Regression test included: cancelling a waiter leaves the scoped permit count at 1, not 2.
This commit is contained in:
parent
56590c63d5
commit
d3ea9b9f2e
2 changed files with 49 additions and 9 deletions
|
|
@ -90,18 +90,20 @@ async def _llm_semaphore():
|
||||||
# we've already given up on it.
|
# we've already given up on it.
|
||||||
while not ceiling_sem.acquire(False):
|
while not ceiling_sem.acquire(False):
|
||||||
await asyncio.sleep(0.05)
|
await asyncio.sleep(0.05)
|
||||||
|
# Only set once the permit is actually held, so a cancellation while polling
|
||||||
|
# for it doesn't make the finally release a permit we never acquired (which
|
||||||
|
# would inflate the scoped cap — the mirror of the ceiling leak fixed above).
|
||||||
scoped_sem = None
|
scoped_sem = None
|
||||||
try:
|
try:
|
||||||
effective = get_max_concurrency()
|
effective = get_max_concurrency()
|
||||||
ceiling = _process_wide_max_concurrency()
|
ceiling = _process_wide_max_concurrency()
|
||||||
if effective < ceiling:
|
if effective < ceiling:
|
||||||
scoped_sem = _max_concurrency_scope_semaphore()
|
candidate = _max_concurrency_scope_semaphore()
|
||||||
if scoped_sem is not None:
|
if candidate is not None:
|
||||||
while not scoped_sem.acquire(False):
|
while not candidate.acquire(False):
|
||||||
await asyncio.sleep(0.05)
|
await asyncio.sleep(0.05)
|
||||||
yield
|
scoped_sem = candidate
|
||||||
else:
|
yield
|
||||||
yield
|
|
||||||
finally:
|
finally:
|
||||||
if scoped_sem is not None:
|
if scoped_sem is not None:
|
||||||
scoped_sem.release()
|
scoped_sem.release()
|
||||||
|
|
@ -117,14 +119,17 @@ def _sync_llm_semaphore():
|
||||||
"""
|
"""
|
||||||
ceiling_sem = _process_ceiling_semaphore()
|
ceiling_sem = _process_ceiling_semaphore()
|
||||||
ceiling_sem.acquire()
|
ceiling_sem.acquire()
|
||||||
|
# Only set once the permit is actually held (mirrors _llm_semaphore): guards
|
||||||
|
# against releasing a permit we never acquired if acquire() is interrupted.
|
||||||
scoped_sem = None
|
scoped_sem = None
|
||||||
try:
|
try:
|
||||||
effective = get_max_concurrency()
|
effective = get_max_concurrency()
|
||||||
ceiling = _process_wide_max_concurrency()
|
ceiling = _process_wide_max_concurrency()
|
||||||
if effective < ceiling:
|
if effective < ceiling:
|
||||||
scoped_sem = _max_concurrency_scope_semaphore()
|
candidate = _max_concurrency_scope_semaphore()
|
||||||
if scoped_sem is not None:
|
if candidate is not None:
|
||||||
scoped_sem.acquire()
|
candidate.acquire()
|
||||||
|
scoped_sem = candidate
|
||||||
yield
|
yield
|
||||||
finally:
|
finally:
|
||||||
if scoped_sem is not None:
|
if scoped_sem is not None:
|
||||||
|
|
|
||||||
|
|
@ -11,6 +11,7 @@ from pageindex.config import (
|
||||||
IndexConfig,
|
IndexConfig,
|
||||||
_env_llm_timeout_default,
|
_env_llm_timeout_default,
|
||||||
_env_max_concurrency_default,
|
_env_max_concurrency_default,
|
||||||
|
_max_concurrency_scope_semaphore,
|
||||||
get_llm_params,
|
get_llm_params,
|
||||||
get_max_concurrency,
|
get_max_concurrency,
|
||||||
llm_params_scope,
|
llm_params_scope,
|
||||||
|
|
@ -133,6 +134,40 @@ def test_llm_semaphore_cancellation_while_waiting_does_not_leak_a_permit():
|
||||||
asyncio.run(run())
|
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():
|
def test_llm_semaphore_uses_scoped_override():
|
||||||
# A per-index max_concurrency_scope active when the loop's semaphore is first
|
# 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.
|
# created must set its size, and must not mutate the process default.
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue