diff --git a/src/skillspector/llm_analyzer_base.py b/src/skillspector/llm_analyzer_base.py index 908dc25a..b64b6923 100644 --- a/src/skillspector/llm_analyzer_base.py +++ b/src/skillspector/llm_analyzer_base.py @@ -29,8 +29,9 @@ import asyncio import os +import threading import time -from collections import defaultdict +from collections import defaultdict, deque from collections.abc import Callable from dataclasses import dataclass, field from typing import Any, Literal, cast @@ -134,12 +135,109 @@ def _uses_native_connection_retries( return False +# ONE BUDGET FOR THE PROCESS, AND IT CANNOT BE AN asyncio.Semaphore. +# +# The analyzers are separate graph nodes and each node is a *synchronous* function: it reaches +# ``run_async()``, which calls ``asyncio.run()`` — and, when a loop is already running, does so on +# a fresh thread. Every analyzer therefore gets its own event loop. An ``asyncio.Semaphore`` is +# bound to the loop that created it, so one semaphore per loop is one semaphore per analyzer, and +# the process still puts N x limit requests on the wire. That is the burst +# ``SKILLSPECTOR_MAX_LLM_CONCURRENCY`` exists to prevent: users set it to 1 on a free tier +# precisely because a burst guarantees 429s, and 429'd batches are dropped from the result. +# +# So the permits live outside asyncio, under a plain lock, and each waiter parks on a future +# belonging to *its own* loop. Releasing hands the permit to the next waiter through that loop's +# ``call_soon_threadsafe``, which is the one asyncio primitive that is safe to call from another +# thread. Nothing blocks a worker thread while it waits, which rules out the obvious alternative +# of wrapping a ``threading.Semaphore`` in ``run_in_executor``: that pins one thread per queued +# request, and the queue is exactly as long as the fan-out this is meant to bound. +class _GlobalLLMLimiter: + """A permit counter shared by every event loop and thread in the process.""" + + def __init__(self, limit: int) -> None: + self._limit = max(1, limit) + self._lock = threading.Lock() + self._in_flight = 0 + # (loop, future) in arrival order. FIFO, so a late analyzer is not starved by an early + # one that keeps re-acquiring. + self._waiters: deque[tuple[asyncio.AbstractEventLoop, asyncio.Future]] = deque() + + async def acquire(self) -> None: + loop = asyncio.get_running_loop() + with self._lock: + if self._in_flight < self._limit: + self._in_flight += 1 + return + waiter: asyncio.Future = loop.create_future() + self._waiters.append((loop, waiter)) + try: + await waiter + except BaseException: + # Cancelled or timed out while queued. Leaving the future in the deque would let a + # later release() hand a permit to a coroutine that is already gone, and the permit + # would never come back. + with self._lock: + try: + self._waiters.remove((loop, waiter)) + except ValueError: + pass # already taken out of the queue: the handover was in flight + if waiter.done() and not waiter.cancelled() and waiter.exception() is None: + # The permit arrived in the same instant the wait was abandoned. It is ours and + # nobody will use it: pass it on rather than leak it. + self.release() + raise + + def release(self) -> None: + with self._lock: + while self._waiters: + loop, waiter = self._waiters.popleft() + if waiter.cancelled(): + continue + try: + # The permit is transferred, not returned: _in_flight stays as it is. + loop.call_soon_threadsafe(_settle, waiter) + except RuntimeError: + # That loop is closed — its waiter can never run. Try the next one. + continue + return + self._in_flight = max(0, self._in_flight - 1) + + async def __aenter__(self) -> _GlobalLLMLimiter: + await self.acquire() + return self + + async def __aexit__(self, *_exc: object) -> None: + self.release() + + +def _settle(waiter: asyncio.Future) -> None: + """Resolve a queued waiter, unless it went away between the handover and the callback.""" + if not waiter.done(): + waiter.set_result(None) + + +# Keyed by the resolved limit, not shared blindly: a caller that resolves a different value gets +# its own budget instead of silently resizing one that other coroutines are holding permits from. +_limiters: dict[int, _GlobalLLMLimiter] = {} +_limiters_lock = threading.Lock() + + +def _shared_limiter(limit: int) -> _GlobalLLMLimiter: + """Return the process-wide limiter for *limit*.""" + with _limiters_lock: + limiter = _limiters.get(limit) + if limiter is None: + limiter = _GlobalLLMLimiter(limit) + _limiters[limit] = limiter + return limiter + + def resolve_max_concurrency() -> int: """Resolve the LLM fan-out concurrency from ``SKILLSPECTOR_MAX_LLM_CONCURRENCY``. Defaults to :data:`DEFAULT_MAX_LLM_CONCURRENCY`. Users on rate-limited providers (free tiers with a low RPM) can set it to ``1`` to serialize - requests instead of bursting up to 10 in parallel — a burst that otherwise + requests across every analyzer instead of bursting up to 10 in parallel — a burst that otherwise guarantees 429s, and 429'd batches are dropped from the result (see the analyzer fan-out below). Invalid values fall back to the default; values below 1 are clamped to 1. @@ -930,9 +1028,15 @@ async def arun_batches_detailed( **kwargs: object, ) -> BatchExecutionResult: """Execute batches concurrently and retain sanitized per-batch failures.""" + # Resolved from the environment: one budget for the whole process, because that is what + # the variable promises — "requests in parallel", not "requests in parallel per analyzer". + # An explicit argument keeps its documented meaning and stays local to this call: a caller + # that passes a number is asking for a fan-out width, not for a share of the budget. + sem: _GlobalLLMLimiter | asyncio.Semaphore if max_concurrency is None: - max_concurrency = resolve_max_concurrency() - sem = asyncio.Semaphore(max_concurrency) + sem = _shared_limiter(resolve_max_concurrency()) + else: + sem = asyncio.Semaphore(max_concurrency) async def _process(batch: Batch) -> tuple[Batch, list]: async with sem: diff --git a/tests/nodes/test_llm_analyzer_base.py b/tests/nodes/test_llm_analyzer_base.py index 0e1aca1c..f94fc9db 100644 --- a/tests/nodes/test_llm_analyzer_base.py +++ b/tests/nodes/test_llm_analyzer_base.py @@ -17,7 +17,11 @@ from __future__ import annotations +import asyncio +import concurrent.futures +import contextlib import json +import threading from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -39,6 +43,7 @@ LLMAnalyzerBase, LLMFinding, LLMRuntimeLimitError, + _shared_limiter, append_output_language_instruction, chunk_file_by_lines, estimate_tokens, @@ -48,7 +53,7 @@ resolve_max_concurrency, resolve_output_language, ) -from skillspector.llm_utils import AgentCLIChatModel, StructuredOutputParseError +from skillspector.llm_utils import AgentCLIChatModel, StructuredOutputParseError, run_async from skillspector.models import Finding from skillspector.nodes.meta_analyzer import ( LLMMetaAnalyzer, @@ -2736,3 +2741,189 @@ def test_unknown_model_uses_default(self) -> None: out = get_max_output_tokens("unknown/model") assert inp == int(mocked_ctx * 0.75) assert out == int(mocked_ctx * 0.25) + + +class TestConcurrencyIsGlobal: + """`SKILLSPECTOR_MAX_LLM_CONCURRENCY` must bound the whole process, not one analyzer. + + The analyzers are separate LangGraph nodes and the graph fans out to them in parallel + (``workflow.add_edge("build_context", analyzer_id)`` for each). A semaphore created inside + one analyzer's fan-out therefore bounds only that analyzer: setting the variable to 1 still + puts N requests on the wire at once, which is exactly what a user on a rate-limited endpoint + set it to 1 to avoid. + """ + + MODEL = "gpt-4o-mini" + + @staticmethod + def _counting_analyzer(peak: list[int], live: list[int]) -> LLMAnalyzerBase: + analyzer = LLMAnalyzerBase(base_prompt="test", model=TestConcurrencyIsGlobal.MODEL) + + async def _invoke(*_args, **_kwargs): + live[0] += 1 + peak[0] = max(peak[0], live[0]) + # Yield control so a second in-flight request has the chance to be observed. Without + # this the coroutine could complete before any other one starts, and the assertion + # would pass on a serialization the code does not actually provide. + await asyncio.sleep(0.02) + live[0] -= 1 + return LLMAnalysisResult(findings=[]) + + analyzer._structured_llm.ainvoke = AsyncMock(side_effect=_invoke) + return analyzer + + @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) + async def test_two_analyzers_respect_one_global_slot(self, monkeypatch) -> None: + monkeypatch.setenv("SKILLSPECTOR_MAX_LLM_CONCURRENCY", "1") + peak, live = [0], [0] + first = self._counting_analyzer(peak, live) + second = self._counting_analyzer(peak, live) + batches = [Batch(file_path="a.py", content="a"), Batch(file_path="b.py", content="b")] + + await asyncio.gather( + first.arun_batches_detailed(list(batches)), + second.arun_batches_detailed(list(batches)), + ) + + assert peak[0] == 1, ( + f"{peak[0]} requests were in flight with the limit set to 1: the bound is per " + "analyzer, so N analyzers issue N simultaneous requests" + ) + + @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) + async def test_limit_of_two_allows_two_across_analyzers(self, monkeypatch) -> None: + """The bound must be the ceiling, not a serialization: a knob that only ever means 1 is + as wrong as one that means nothing.""" + monkeypatch.setenv("SKILLSPECTOR_MAX_LLM_CONCURRENCY", "2") + peak, live = [0], [0] + first = self._counting_analyzer(peak, live) + second = self._counting_analyzer(peak, live) + batches = [Batch(file_path="a.py", content="a"), Batch(file_path="b.py", content="b")] + + await asyncio.gather( + first.arun_batches_detailed(list(batches)), + second.arun_batches_detailed(list(batches)), + ) + + assert peak[0] == 2 + + @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) + async def test_explicit_argument_still_bounds_only_its_own_call(self, monkeypatch) -> None: + """An explicit ``max_concurrency`` keeps its documented meaning: it wins for that call. + + Callers that pass a number are asking for a local fan-out width, not for a share of the + process-wide budget, and tests rely on that isolation. + """ + monkeypatch.setenv("SKILLSPECTOR_MAX_LLM_CONCURRENCY", "1") + peak, live = [0], [0] + analyzer = self._counting_analyzer(peak, live) + batches = [Batch(file_path="a.py", content="a"), Batch(file_path="b.py", content="b")] + + await analyzer.arun_batches_detailed(batches, max_concurrency=2) + + assert peak[0] == 2 + + +class TestConcurrencyIsGlobalAcrossLoops: + """The same bound, exercised the way the graph actually runs. + + The tests above gather two analyzers on one artificial event loop, and that is not the shape + of production: every analyzer node is a *synchronous* function that reaches ``run_async()``, + which calls ``asyncio.run()`` — a brand-new loop each time, on a brand-new thread when a loop + is already running. Anything keyed by the running loop is therefore keyed by the analyzer, + and a same-loop test cannot tell the two apart. This class reproduces the real execution + model: N analyzers, N threads, N loops, one process-wide counter. + """ + + MODEL = "gpt-4o-mini" + + @staticmethod + def _counting_analyzer(state: dict, lock: threading.Lock) -> LLMAnalyzerBase: + analyzer = LLMAnalyzerBase( + base_prompt="test", model=TestConcurrencyIsGlobalAcrossLoops.MODEL + ) + + async def _invoke(*_args, **_kwargs): + # The counter is shared across threads now, so it needs a real lock: reading a peak + # through a data race would make this test lie in whichever direction was convenient. + with lock: + state["live"] += 1 + state["peak"] = max(state["peak"], state["live"]) + # Long enough that every thread is inside a request at the same time if the bound + # does not hold. Without the sleep a serial execution and a bounded one look alike. + await asyncio.sleep(0.05) + with lock: + state["live"] -= 1 + return LLMAnalysisResult(findings=[]) + + analyzer._structured_llm.ainvoke = AsyncMock(side_effect=_invoke) + return analyzer + + @staticmethod + def _run_like_a_node(analyzer: LLMAnalyzerBase, batches: list[Batch]) -> None: + """Exactly what an analyzer node does: a sync call into ``run_async``.""" + run_async(analyzer.arun_batches_detailed(batches)) + + @pytest.mark.parametrize("limit", [1, 2]) + @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) + def test_three_nodes_on_three_loops_share_one_budget(self, limit, monkeypatch) -> None: + monkeypatch.setenv("SKILLSPECTOR_MAX_LLM_CONCURRENCY", str(limit)) + state = {"live": 0, "peak": 0} + lock = threading.Lock() + analyzers = [self._counting_analyzer(state, lock) for _ in range(3)] + batches = [Batch(file_path="a.py", content="a"), Batch(file_path="b.py", content="b")] + + with concurrent.futures.ThreadPoolExecutor(max_workers=len(analyzers)) as pool: + futures = [pool.submit(self._run_like_a_node, a, list(batches)) for a in analyzers] + for f in futures: + f.result() + + assert state["peak"] <= limit, ( + f"{state['peak']} requests were in flight with the limit set to {limit}: each node " + "runs on its own event loop, so a per-loop bound is a per-analyzer bound" + ) + assert state["live"] == 0, "a permit was never released" + + @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) + def test_the_bound_is_a_ceiling_not_a_serialization(self, monkeypatch) -> None: + """A knob that always means 1 is as wrong as one that means nothing.""" + monkeypatch.setenv("SKILLSPECTOR_MAX_LLM_CONCURRENCY", "3") + state = {"live": 0, "peak": 0} + lock = threading.Lock() + analyzers = [self._counting_analyzer(state, lock) for _ in range(3)] + batches = [Batch(file_path="a.py", content="a")] + + with concurrent.futures.ThreadPoolExecutor(max_workers=len(analyzers)) as pool: + for f in [pool.submit(self._run_like_a_node, a, list(batches)) for a in analyzers]: + f.result() + + assert state["peak"] == 3, ( + f"peak was {state['peak']} with a limit of 3 across three nodes: the bound has " + "become a queue" + ) + + @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) + def test_a_cancelled_waiter_does_not_strand_its_permit(self, monkeypatch) -> None: + """A permit handed to a coroutine that has gone away must come back. + + This is the failure mode that does not announce itself: the count never recovers, and + every later request waits for a slot that no longer exists. It looks like a hang, hours + after the cancellation that caused it. + """ + monkeypatch.setenv("SKILLSPECTOR_MAX_LLM_CONCURRENCY", "1") + + async def _exercise() -> int: + limiter = _shared_limiter(1) + await limiter.acquire() + queued = asyncio.ensure_future(limiter.acquire()) + await asyncio.sleep(0) + queued.cancel() + with contextlib.suppress(asyncio.CancelledError): + await queued + limiter.release() + # The permit must be free again: a second acquire has to succeed immediately. + await asyncio.wait_for(limiter.acquire(), timeout=1.0) + limiter.release() + return limiter._in_flight + + assert asyncio.run(_exercise()) == 0, "the limiter leaked a permit"