From ae8666b9d1a568db03f6aecfcd48de337381487e Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 31 Oct 2026 16:12:30 -0700 Subject: [PATCH] fix(http): diagnose aiohttp connection pool saturation (#49499) --- litellm/llms/custom_httpx/aiohttp_handler.py | 5 +- .../custom_httpx/aiohttp_pool_diagnostics.py | 130 ++++++++++ litellm/llms/custom_httpx/http_handler.py | 2 + litellm/proxy/proxy_server.py | 7 +- .../llms/custom_httpx/test_aiohttp_handler.py | 20 +- .../llms/custom_httpx/test_http_handler.py | 224 +++++++++++++++++- 6 files changed, 369 insertions(+), 19 deletions(-) create mode 100644 litellm/llms/custom_httpx/aiohttp_pool_diagnostics.py diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index e6544968d24..1d2b95625bb 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -14,6 +14,7 @@ from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.base_llm.image_variations.transformation import ( BaseImageVariationConfig, ) +from litellm.llms.custom_httpx.aiohttp_pool_diagnostics import build_pool_wait_trace_config from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -90,7 +91,7 @@ class BaseLLMAIOHTTPHandler: if self.transport is None: connector: Final = self._get_connector() if connector: - return aiohttp.ClientSession(connector=connector) + return aiohttp.ClientSession(connector=connector, trace_configs=[build_pool_wait_trace_config()]) transport: Final = self.transport or self._get_or_create_transport() if transport is not None and hasattr(transport, "_get_valid_client_session"): @@ -99,7 +100,7 @@ class BaseLLMAIOHTTPHandler: except RuntimeError: pass - return aiohttp.ClientSession() + return aiohttp.ClientSession(trace_configs=[build_pool_wait_trace_config()]) def _get_async_client_session(self, dynamic_client_session: ClientSession | None = None) -> ClientSession: if dynamic_client_session: diff --git a/litellm/llms/custom_httpx/aiohttp_pool_diagnostics.py b/litellm/llms/custom_httpx/aiohttp_pool_diagnostics.py new file mode 100644 index 00000000000..e0a4007033a --- /dev/null +++ b/litellm/llms/custom_httpx/aiohttp_pool_diagnostics.py @@ -0,0 +1,130 @@ +import asyncio +from types import SimpleNamespace +from typing import Final + +from aiohttp import ( + ClientSession, + TraceConfig, + TraceConnectionQueuedEndParams, + TraceConnectionQueuedStartParams, + TraceRequestExceptionParams, + TraceRequestRedirectParams, + TraceRequestStartParams, +) +from yarl import URL + +from litellm._logging import verbose_logger + +_POOL_WAIT_WARNING_SECONDS: Final = 1.0 + + +class _PoolWaitContext(SimpleNamespace): + def __init__(self, trace_request_ctx: object = None) -> None: + super().__init__() + self.host: str | None = None + self.port: int | None = None + self.started_at: float | None = None + self.waited: float = 0.0 + self.warning: asyncio.TimerHandle | None = None + self.warned: bool = False + + def set_destination(self, params: TraceRequestStartParams) -> None: + self.host = params.url.host + self.port = params.url.port + + def redirect(self, params: TraceRequestRedirectParams) -> None: + location: Final = params.response.headers.get("Location") or params.response.headers.get("URI") + if location is None: + return + try: + destination: Final = params.url.join(URL(location)) + host: Final = destination.host + port: Final = destination.port + except ValueError: + return + self.host = host + self.port = port + + def start(self, session: ClientSession) -> None: + loop: Final = asyncio.get_running_loop() + self.started_at = loop.time() + if not self.warned: + self.warning = loop.call_later(max(0.0, _POOL_WAIT_WARNING_SECONDS - self.waited), self.warn, session) + + def elapsed(self) -> float: + return self.waited + ( + asyncio.get_running_loop().time() - self.started_at if self.started_at is not None else 0.0 + ) + + def warn(self, session: ClientSession) -> None: + if self.warned or self.started_at is None: + return + self.warned = True + connector: Final = session.connector + if connector is None: + return + verbose_logger.warning( + "aiohttp connection pool saturated: host=%s port=%s waited=%.3fs for a connection slot " + "(acquired_total=%s, limit=%s, limit_per_host=%s). " + "Wait occurred before sending the request upstream.", + self.host, + self.port, + self.elapsed(), + len(connector._acquired), # pyright: ignore[reportPrivateUsage] # aiohttp exposes no public acquired count + connector.limit, + connector.limit_per_host, + ) + + def finish(self, session: ClientSession, failed: bool = False) -> None: + if self.warning is not None: + self.warning.cancel() + self.warning = None + if self.started_at is None: + return + if failed or self.elapsed() >= _POOL_WAIT_WARNING_SECONDS: + self.warn(session) + self.waited = self.elapsed() + self.started_at = None + + +async def _on_request_start(session: ClientSession, context: SimpleNamespace, params: TraceRequestStartParams) -> None: + if isinstance(context, _PoolWaitContext): + context.set_destination(params) + + +async def _on_queued_start( + session: ClientSession, context: SimpleNamespace, params: TraceConnectionQueuedStartParams +) -> None: + if isinstance(context, _PoolWaitContext): + context.start(session) + + +async def _on_queued_end( + session: ClientSession, context: SimpleNamespace, params: TraceConnectionQueuedEndParams +) -> None: + if isinstance(context, _PoolWaitContext): + context.finish(session) + + +async def _on_request_exception( + session: ClientSession, context: SimpleNamespace, params: TraceRequestExceptionParams +) -> None: + if isinstance(context, _PoolWaitContext): + context.finish(session, failed=True) + + +async def _on_request_redirect( + session: ClientSession, context: SimpleNamespace, params: TraceRequestRedirectParams +) -> None: + if isinstance(context, _PoolWaitContext): + context.redirect(params) + + +def build_pool_wait_trace_config() -> TraceConfig: + trace: Final = TraceConfig(trace_config_ctx_factory=_PoolWaitContext) + trace.on_request_start.append(_on_request_start) + trace.on_connection_queued_start.append(_on_queued_start) + trace.on_connection_queued_end.append(_on_queued_end) + trace.on_request_exception.append(_on_request_exception) + trace.on_request_redirect.append(_on_request_redirect) + return trace diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 5620971a8a2..3a6fa5f51da 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -44,6 +44,7 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) +from litellm.llms.custom_httpx.aiohttp_pool_diagnostics import build_pool_wait_trace_config from litellm.llms.custom_httpx.http2_connection_spread import AsyncHTTP2ConnectionSpread from litellm.types.llms.custom_http import * @@ -1456,6 +1457,7 @@ class AsyncHTTPHandler: connector=TCPConnector(**transport_connector_kwargs), cookie_jar=DummyCookieJar(), trust_env=trust_env, + trace_configs=[build_pool_wait_trace_config()], ) # Use shared session if provided and valid diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d8a57d3e9a1..2703d7d050c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -392,6 +392,7 @@ from litellm.llms.anthropic.messages.convertible_blocks import ( drop_non_string_tool_result_texts, malformed_tool_result_count_error, ) +from litellm.llms.custom_httpx.aiohttp_pool_diagnostics import build_pool_wait_trace_config from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.openai.live import parse_live_session_start from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS @@ -1460,7 +1461,11 @@ async def initialize_shared_aiohttp_session() -> "ClientSession | None": connector_kwargs["socket_factory"] = socket_factory connector: Final = TCPConnector(**connector_kwargs) - session: Final = ClientSession(connector=connector, cookie_jar=DummyCookieJar()) + session: Final = ClientSession( + connector=connector, + cookie_jar=DummyCookieJar(), + trace_configs=[build_pool_wait_trace_config()], + ) verbose_proxy_logger.info( "SESSION REUSE: Created shared aiohttp session for connection pooling (ID: %s, limit=%s, limit_per_host=%s)", diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_handler.py b/tests/unit/llms/custom_httpx/test_aiohttp_handler.py index bb66a70da84..2fa3a6a7db6 100644 --- a/tests/unit/llms/custom_httpx/test_aiohttp_handler.py +++ b/tests/unit/llms/custom_httpx/test_aiohttp_handler.py @@ -307,20 +307,12 @@ class TestBaseLLMAIOHTTPHandler: # Should not call aiohttp.ClientSession directly mock_client_session.assert_not_called() - @patch("aiohttp.ClientSession") - def test_create_client_session_with_connector(self, mock_client_session): - """Test session creation using connector""" - mock_connector = Mock(spec=aiohttp.BaseConnector) - mock_session_instance = Mock() - mock_client_session.return_value = mock_session_instance - - handler = BaseLLMAIOHTTPHandler(connector=mock_connector) - - result = handler._create_client_session_with_transport() - - # Should create session with connector - mock_client_session.assert_called_once_with(connector=mock_connector) - assert result is mock_session_instance + @pytest.mark.asyncio + async def test_create_client_session_with_connector(self) -> None: + connector: Final = aiohttp.TCPConnector() + handler: Final = BaseLLMAIOHTTPHandler(connector=connector) + async with handler._create_client_session_with_transport() as session: + assert session.connector is connector @pytest.mark.asyncio async def test_create_client_session_default_honors_global_ssl_verify_false( diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 0163cb9b26d..3e583bf8dff 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1,13 +1,16 @@ import asyncio, importlib, json, time import gc import io +import logging import os import pathlib import ssl import threading import weakref -from collections.abc import Callable, Mapping +from collections.abc import AsyncIterator, Callable, Mapping from concurrent.futures import ThreadPoolExecutor +from contextlib import asynccontextmanager +from contextvars import Context from typing import Final from unittest.mock import MagicMock, patch @@ -15,7 +18,10 @@ import certifi import httpx import pytest import respx -from aiohttp import ClientSession, TCPConnector +from aiohttp import ClientResponse, ClientSession, RequestInfo, TCPConnector +from aiohttp.tracing import Trace +from multidict import CIMultiDict, CIMultiDictProxy +from yarl import URL import litellm from litellm.caching.llm_caching_handler import LLMClientCache @@ -2304,3 +2310,217 @@ async def test_async_timeout_reports_the_per_call_timeout_as_a_float_when_one_is await handler.close() assert "Read timed out. Timeout passed=30.0, time taken=" in str(raised.value), str(raised.value) + + +class _PoolDiagnosticClock: + def __init__(self, loop: asyncio.AbstractEventLoop) -> None: + self.now = 1000.0 + self.loop = loop + self.schedule = loop.call_later + self.timers: list[asyncio.TimerHandle] = [] + + def time(self) -> float: + return self.now + + def call_later( + self, delay: float, callback: Callable[..., object], + *args: object, context: Context | None = None, + ) -> asyncio.TimerHandle: + timer: Final = self.schedule(delay, callback, *args, context=context) + self.timers.append(timer) + return timer + + async def advance(self, seconds: float) -> None: + self.now += seconds + checkpoint: Final[asyncio.Future[None]] = self.loop.create_future() + self.loop.call_soon(checkpoint.set_result, None) + await checkpoint + + +def _pool_diagnostic_clock(monkeypatch: pytest.MonkeyPatch) -> _PoolDiagnosticClock: + loop: Final = asyncio.get_running_loop() + clock: Final = _PoolDiagnosticClock(loop) + monkeypatch.setattr(loop, "time", clock.time) + monkeypatch.setattr(loop, "call_later", clock.call_later) + return clock + + +@asynccontextmanager +async def _pool_diagnostic_session(source: str) -> AsyncIterator[ClientSession]: + if source == "proxy": + from litellm.proxy.proxy_server import initialize_shared_aiohttp_session + + session: Final = await initialize_shared_aiohttp_session() + assert session is not None + async with session: + yield session + return + if source == "native-connector": + from litellm.llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler + + handler: Final = BaseLLMAIOHTTPHandler(connector=TCPConnector(limit=7, limit_per_host=2)) + try: + yield handler._get_async_client_session() + finally: + await handler.close() + return + transport: Final = AsyncHTTPHandler.create_aiohttp_transport() + if source == "rebuilt": + await transport.get_valid_client_session().close() + try: + yield transport.get_valid_client_session() + finally: + await transport.aclose() + + +async def _pool_diagnostic_request(session: ClientSession, url: URL) -> tuple[Trace, ...]: + traces: Final = tuple(Trace(session, config, config.trace_config_ctx()) for config in session.trace_configs) + for trace in traces: + await trace.send_request_start("POST", url, CIMultiDict()) + return traces + + +def _pool_diagnostic_messages(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: + return tuple(record.getMessage() for record in caplog.records if "connection pool" in record.getMessage()) + + +@pytest.mark.parametrize("source", ["sdk", "proxy", "native-connector", "rebuilt"]) +@pytest.mark.parametrize("error", [TimeoutError("private timeout"), asyncio.CancelledError("private cancellation")]) +async def test_aiohttp_pool_failure_identifies_queue_wait_and_cleans_up_timer( + source: str, error: BaseException, monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + pool_diagnostic_clock: Final = _pool_diagnostic_clock(monkeypatch) + caplog.set_level(logging.WARNING, logger="LiteLLM") + url: Final = URL("https://private-user:private-password@pool.test:8443/private-path?token=private-query") + async with _pool_diagnostic_session(source) as session: + connector: Final = session.connector + assert connector is not None + traces: Final = await _pool_diagnostic_request(session, url) + for trace in traces: + await trace.send_connection_queued_start() + await pool_diagnostic_clock.advance(0.25) + for trace in traces: + await trace.send_request_exception("POST", url, CIMultiDict(), error) + messages: Final = _pool_diagnostic_messages(caplog) + assert len(messages) == 1, messages + message: Final = messages[0] + assert "host=pool.test" in message + assert "port=8443" in message + assert "0.25" in message + assert "acquired_total=0" in message + assert f"limit={connector.limit}" in message + assert f"limit_per_host={connector.limit_per_host}" in message + assert "private" not in message + assert all(timer.cancelled() for timer in pool_diagnostic_clock.timers) + await pool_diagnostic_clock.advance(2.0) + assert _pool_diagnostic_messages(caplog) == messages + + +@pytest.mark.parametrize("queued", [False, True]) +async def test_aiohttp_pool_diagnostics_do_not_label_slow_upstream_as_pool_saturation( + queued: bool, monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + pool_diagnostic_clock: Final = _pool_diagnostic_clock(monkeypatch) + caplog.set_level(logging.WARNING, logger="LiteLLM") + url: Final = URL("https://healthy.test/completions") + async with _pool_diagnostic_session("sdk") as session: + traces: Final = await _pool_diagnostic_request(session, url) + if queued: + for trace in traces: + await trace.send_connection_queued_start() + await pool_diagnostic_clock.advance(0.25) + for trace in traces: + await trace.send_connection_queued_end() + await pool_diagnostic_clock.advance(2.0) + for trace in traces: + await trace.send_request_exception("POST", url, CIMultiDict(), TimeoutError("slow upstream")) + assert _pool_diagnostic_messages(caplog) == () + assert all(timer.cancelled() for timer in pool_diagnostic_clock.timers) + + +async def test_aiohttp_pool_warns_while_queued_once_per_request_with_isolated_concurrent_waits( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, +) -> None: + pool_diagnostic_clock: Final = _pool_diagnostic_clock(monkeypatch) + caplog.set_level(logging.WARNING, logger="LiteLLM") + slow_url: Final = URL("https://blocked.test/completions") + fast_url: Final = URL("https://available.test/completions") + async with _pool_diagnostic_session("sdk") as session: + blocked: Final = await _pool_diagnostic_request(session, slow_url) + available: Final = await _pool_diagnostic_request(session, fast_url) + for trace in blocked + available: + await trace.send_connection_queued_start() + await pool_diagnostic_clock.advance(0.25) + for trace in available: + await trace.send_connection_queued_end() + for trace in blocked: + await trace.send_connection_queued_end() + await trace.send_connection_queued_start() + await pool_diagnostic_clock.advance(0.8) + messages: Final = _pool_diagnostic_messages(caplog) + assert len(messages) == 1, messages + assert "host=blocked.test" in messages[0] + assert "1.05" in messages[0] + assert "available.test" not in messages[0] + for trace in blocked: + await trace.send_connection_queued_end() + await trace.send_connection_queued_start() + await pool_diagnostic_clock.advance(2.0) + for trace in blocked: + await trace.send_request_exception("POST", slow_url, CIMultiDict(), TimeoutError()) + for trace in available: + await trace.send_request_exception("POST", fast_url, CIMultiDict(), TimeoutError()) + await pool_diagnostic_clock.advance(2.0) + assert _pool_diagnostic_messages(caplog) == messages + assert all(timer.cancelled() for timer in pool_diagnostic_clock.timers) + + +@pytest.mark.parametrize(("header", "location", "host", "port"), [ + ("Location", "https://redirected.test:9443/private-path?token=private-query", "redirected.test", 9443), + ("Location", "/private-path?token=private-query", "initial.test", 8443), + ("URI", "https://redirected.test/private-path", "redirected.test", 443), + ("Location", "http://[invalid", None, None), +]) +async def test_aiohttp_pool_warning_tracks_redirect_destination_without_masking_invalid_redirects( + header: str, location: str, host: str | None, port: int | None, + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, +) -> None: + from aiohttp.client_proto import ResponseHandler + from aiohttp.helpers import TimerNoop + from aiohttp.http_writer import StreamWriter + + clock: Final = _pool_diagnostic_clock(monkeypatch) + caplog.set_level(logging.WARNING, logger="LiteLLM") + url: Final = URL("https://initial.test:8443/completions") + async with _pool_diagnostic_session("sdk") as session: + traces: Final = await _pool_diagnostic_request(session, url) + loop: Final = asyncio.get_running_loop() + response: Final = ClientResponse( + "POST", url, writer=None, continue100=None, timer=TimerNoop(), + request_info=RequestInfo(url, "POST", CIMultiDictProxy(CIMultiDict()), url), + traces=[], loop=loop, session=session, + stream_writer=StreamWriter(ResponseHandler(loop), loop), + ) + response._headers = CIMultiDictProxy(CIMultiDict({header: location})) + try: + for trace in traces: + await trace.send_request_redirect("POST", url, CIMultiDict(), response) + if host is None: + for trace in traces: + await trace.send_request_exception("POST", url, CIMultiDict(), ValueError("invalid redirect")) + assert _pool_diagnostic_messages(caplog) == () + return + for trace in traces: + await trace.send_connection_queued_start() + await clock.advance(0.25) + for trace in traces: + await trace.send_request_exception("POST", url, CIMultiDict(), TimeoutError()) + messages: Final = _pool_diagnostic_messages(caplog) + assert len(messages) == 1, messages + assert f"host={host}" in messages[0] + assert f"port={port}" in messages[0] + assert "private" not in messages[0] + finally: + response.close()