fix(http): diagnose aiohttp connection pool saturation (#49499)

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:

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

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

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)",

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(

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()