From 6df30258b2306432a0d5a45e13767875bd691f0c Mon Sep 17 00:00:00 2001 From: YMuskrat <101849520+YMuskrat@users.noreply.github.com> Date: Tue, 6 Oct 2026 00:55:30 +0100 Subject: [PATCH 1/2] fix(batch): release pooled model slots on failures and cancellation Release reserved slots in finally blocks and keep asynchronous capacity waiting cancellable. Signed-off-by: YMuskrat <101849520+YMuskrat@users.noreply.github.com> --- contrib/batch_scan/api_pool.py | 28 ++++--- tests/test_batch_scan_pool_cleanup.py | 109 ++++++++++++++++++++++++++ 2 files changed, 126 insertions(+), 11 deletions(-) create mode 100644 tests/test_batch_scan_pool_cleanup.py diff --git a/contrib/batch_scan/api_pool.py b/contrib/batch_scan/api_pool.py index 6960eab9b..820f4f516 100644 --- a/contrib/batch_scan/api_pool.py +++ b/contrib/batch_scan/api_pool.py @@ -459,19 +459,19 @@ def _invoke_with_retry( for attempt in range(self._max_retries + 1): key = self._pool.acquire() - llm = self._build_llm(key) + release_success = True try: + llm = self._build_llm(key) if callbacks is None: result = llm.invoke(prompt) else: result = llm.invoke(prompt, config={"callbacks": callbacks}) - self._pool.release(key, success=True) if attempt > 0: self._pool.record_retry_success() return result except Exception as exc: if self._is_rate_limit(exc) and attempt < self._max_retries: - self._pool.release(key, success=False) + release_success = False logger.debug( "PooledChatModel: rate-limited, retrying " "(attempt %d/%d)", @@ -479,9 +479,10 @@ def _invoke_with_retry( self._max_retries, ) continue - self._pool.release(key, success=True) last_exception = exc raise + finally: + self._pool.release(key, success=release_success) raise RuntimeError( f"PooledChatModel: exhausted {self._max_retries} retries " @@ -494,27 +495,31 @@ async def _ainvoke_with_retry( *, callbacks: list[object] | None = None, ) -> object: - """Async retry loop — non-blocking acquire first, block only if full.""" + """Async retry loop with cancellation-safe slot acquisition and release.""" import asyncio last_exception: Exception | None = None for attempt in range(self._max_retries + 1): key = self._pool.try_acquire() - if key is None: - key = await asyncio.to_thread(self._pool.acquire) - llm = self._build_llm(key) + while key is None: + # A blocking acquire in a worker thread can reserve a slot + # after its awaiting task has been cancelled. Keep waiting in + # the task so cancellation leaves no orphaned pool waiter. + await asyncio.sleep(0.05) + key = self._pool.try_acquire() + release_success = True try: + llm = self._build_llm(key) if callbacks is None: result = await llm.ainvoke(prompt) else: result = await llm.ainvoke(prompt, config={"callbacks": callbacks}) - self._pool.release(key, success=True) if attempt > 0: self._pool.record_retry_success() return result except Exception as exc: if self._is_rate_limit(exc) and attempt < self._max_retries: - self._pool.release(key, success=False) + release_success = False logger.debug( "PooledChatModel: rate-limited, retrying " "(attempt %d/%d)", @@ -522,9 +527,10 @@ async def _ainvoke_with_retry( self._max_retries, ) continue - self._pool.release(key, success=True) last_exception = exc raise + finally: + self._pool.release(key, success=release_success) raise RuntimeError( f"PooledChatModel: exhausted {self._max_retries} retries " diff --git a/tests/test_batch_scan_pool_cleanup.py b/tests/test_batch_scan_pool_cleanup.py new file mode 100644 index 000000000..794d771a3 --- /dev/null +++ b/tests/test_batch_scan_pool_cleanup.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Regression coverage for pooled model slot cleanup without provider calls.""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, Mock + +import pytest + +from contrib.batch_scan.api_pool import ApiKey, ApiKeyPool, PooledChatModel + + +def _model() -> PooledChatModel: + return PooledChatModel( + ApiKeyPool([ApiKey(key="test", model="test", base_url=None, max_concurrent=1)]) + ) + + +@pytest.mark.parametrize("async_call", [False, True]) +async def test_model_construction_failure_releases_slot(monkeypatch, async_call: bool) -> None: + model = _model() + monkeypatch.setattr(model, "_build_llm", Mock(side_effect=ValueError("invalid client"))) + with pytest.raises(ValueError, match="invalid client"): + if async_call: + await model.ainvoke("test") + else: + model.invoke("test") + assert model._pool.active_requests == 0 + key = model._pool.acquire(timeout=0.05) + model._pool.release(key) + + +async def test_cancelling_inflight_call_releases_slot(monkeypatch) -> None: + model = _model() + started = asyncio.Event() + + async def blocked(*args, **kwargs): + started.set() + await asyncio.Event().wait() + + monkeypatch.setattr(model, "_build_llm", Mock(return_value=Mock(ainvoke=blocked))) + task = asyncio.create_task(model.ainvoke("test")) + await asyncio.wait_for(started.wait(), timeout=1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert model._pool.active_requests == 0 + + +async def test_cancelling_waiter_does_not_reserve_a_slot_later(monkeypatch) -> None: + model = _model() + held = model._pool.acquire() + tried = asyncio.Event() + original_try = model._pool.try_acquire + + def try_acquire(): + tried.set() + return original_try() + + monkeypatch.setattr(model._pool, "try_acquire", try_acquire) + # The old thread-based acquire can outlive cancellation. Bound it so the + # regression test cleans up its executor even when run against old code. + original_acquire = model._pool.acquire + monkeypatch.setattr(model._pool, "acquire", lambda: original_acquire(timeout=0.1)) + task = asyncio.create_task(model.ainvoke("test")) + await asyncio.wait_for(tried.wait(), timeout=1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + model._pool.release(held) + await asyncio.sleep(0.15) + assert model._pool.active_requests == 0 + + +async def test_async_waiter_runs_after_capacity_is_released(monkeypatch) -> None: + model = _model() + held = model._pool.acquire() + invoke = AsyncMock(return_value="result") + monkeypatch.setattr(model, "_build_llm", Mock(return_value=Mock(ainvoke=invoke))) + task = asyncio.create_task(model.ainvoke("test")) + await asyncio.sleep(0) + assert not task.done() + model._pool.release(held) + assert await asyncio.wait_for(task, timeout=1) == "result" + assert model._pool.active_requests == 0 + + +@pytest.mark.parametrize("async_call", [False, True]) +async def test_rate_limit_retry_releases_each_slot_once(monkeypatch, async_call: bool) -> None: + pool = ApiKeyPool( + [ApiKey(key=key, model="test", base_url=None, max_concurrent=1) for key in ("one", "two")] + ) + model = PooledChatModel(pool, max_retries=1) + release = Mock(wraps=pool.release) + monkeypatch.setattr(pool, "release", release) + invoke = AsyncMock if async_call else Mock + clients = [Mock(), Mock()] + for client, outcome in zip(clients, [RuntimeError("429 rate limit"), "result"], strict=True): + setattr(client, "ainvoke" if async_call else "invoke", invoke(side_effect=[outcome])) + monkeypatch.setattr(model, "_build_llm", Mock(side_effect=clients)) + result = await model.ainvoke("test") if async_call else model.invoke("test") + assert result == "result" + assert [call.kwargs["success"] for call in release.call_args_list] == [False, True] + assert pool.active_requests == 0 + assert pool.rate_limits_hit == 1 + assert pool.retry_successes == 1 From d0a2a427535615d430cb4c51a91be3539f258eac Mon Sep 17 00:00:00 2001 From: YMuskrat <101849520+YMuskrat@users.noreply.github.com> Date: Wed, 14 Oct 2026 14:28:16 +0100 Subject: [PATCH 2/2] test(batch): cover pooled model cancellation and slot recovery Keep cancellation, capacity recovery and one-release-per-retry coverage against main; drop duplicate client-construction cases. Signed-off-by: YMuskrat <101849520+YMuskrat@users.noreply.github.com> --- tests/test_batch_scan_pool_cleanup.py | 20 +++----------------- 1 file changed, 3 insertions(+), 17 deletions(-) diff --git a/tests/test_batch_scan_pool_cleanup.py b/tests/test_batch_scan_pool_cleanup.py index 794d771a3..d25501e73 100644 --- a/tests/test_batch_scan_pool_cleanup.py +++ b/tests/test_batch_scan_pool_cleanup.py @@ -19,20 +19,6 @@ def _model() -> PooledChatModel: ) -@pytest.mark.parametrize("async_call", [False, True]) -async def test_model_construction_failure_releases_slot(monkeypatch, async_call: bool) -> None: - model = _model() - monkeypatch.setattr(model, "_build_llm", Mock(side_effect=ValueError("invalid client"))) - with pytest.raises(ValueError, match="invalid client"): - if async_call: - await model.ainvoke("test") - else: - model.invoke("test") - assert model._pool.active_requests == 0 - key = model._pool.acquire(timeout=0.05) - model._pool.release(key) - - async def test_cancelling_inflight_call_releases_slot(monkeypatch) -> None: model = _model() started = asyncio.Event() @@ -61,10 +47,10 @@ def try_acquire(): return original_try() monkeypatch.setattr(model._pool, "try_acquire", try_acquire) - # The old thread-based acquire can outlive cancellation. Bound it so the - # regression test cleans up its executor even when run against old code. + # Bound any blocking fallback so the test cannot leave a waiting worker + # behind if a regression reintroduces threaded slot acquisition. original_acquire = model._pool.acquire - monkeypatch.setattr(model._pool, "acquire", lambda: original_acquire(timeout=0.1)) + monkeypatch.setattr(model._pool, "acquire", lambda **kwargs: original_acquire(timeout=0.1)) task = asyncio.create_task(model.ainvoke("test")) await asyncio.wait_for(tried.wait(), timeout=1) task.cancel()