From 40810f2a76631ed1c4975738d657b89344035537 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:34:58 -0700 Subject: [PATCH] fix(http): split the connection cap across HTTP/2 lanes Each lane had its own pool with the full connection cap, so an origin answering HTTP/1.1, or several origins on one client, could exceed max_connections once more than 100 requests were in flight. Lanes now open only while lanes times 100 streams stay under the cap, and each lane pool gets at most 100 of it. --- .../custom_httpx/http2_connection_spread.py | 27 +++++--- litellm/llms/custom_httpx/http_handler.py | 12 ++-- .../test_http2_connection_spread.py | 61 +++++++++++++++++-- 3 files changed, 81 insertions(+), 19 deletions(-) diff --git a/litellm/llms/custom_httpx/http2_connection_spread.py b/litellm/llms/custom_httpx/http2_connection_spread.py index 1fd7fb2c4a8..0fc1d8c38b9 100644 --- a/litellm/llms/custom_httpx/http2_connection_spread.py +++ b/litellm/llms/custom_httpx/http2_connection_spread.py @@ -23,13 +23,13 @@ class _Lane: class _Lanes: def __init__( self, - new_transport: Callable[[], httpx.AsyncBaseTransport], + new_transport: Callable[[httpx.Limits], httpx.AsyncBaseTransport], + limits: httpx.Limits, streams_per_connection: int, - max_connections: int | None, ) -> None: self._new_transport: Final = new_transport + self._limits: Final = limits self._streams_per_connection: Final = streams_per_connection - self._max_connections: Final = max_connections self.lanes: tuple[_Lane, ...] = () def acquire(self) -> _Lane: @@ -39,15 +39,26 @@ class _Lanes: return lane def _open_or_share(self) -> _Lane: - if self._max_connections is not None and len(self.lanes) >= self._max_connections: + cap: Final = self._limits.max_connections + if cap is not None and len(self.lanes) * self._streams_per_connection >= cap: return min(self.lanes, key=lambda lane: lane.in_flight) return self._open() def _open(self) -> _Lane: - lane: Final = _Lane(self._new_transport()) + lane: Final = _Lane(self._new_transport(self._next_lane_limits())) self.lanes = (*self.lanes, lane) return lane + def _next_lane_limits(self) -> httpx.Limits: + cap: Final = self._limits.max_connections + if cap is None: + return self._limits + return httpx.Limits( + max_connections=min(self._streams_per_connection, cap - len(self.lanes) * self._streams_per_connection), + max_keepalive_connections=self._limits.max_keepalive_connections, + keepalive_expiry=self._limits.keepalive_expiry, + ) + class _AsyncReleasingStream(httpx.AsyncByteStream): def __init__(self, stream: httpx.AsyncByteStream, release: Callable[[], None]) -> None: @@ -70,11 +81,11 @@ def _with_stream(response: httpx.Response, stream: httpx.AsyncByteStream) -> htt class AsyncHTTP2ConnectionSpread(httpx.AsyncBaseTransport): def __init__( self, - new_transport: Callable[[], httpx.AsyncBaseTransport], + new_transport: Callable[[httpx.Limits], httpx.AsyncBaseTransport], + limits: httpx.Limits, streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION, - max_connections: int | None = None, ) -> None: - self._lanes: Final = _Lanes(new_transport, streams_per_connection, max_connections) + self._lanes: Final = _Lanes(new_transport, limits, streams_per_connection) async def handle_async_request(self, request: httpx.Request) -> httpx.Response: lane: Final = self._lanes.acquire() diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 2ce6d1ea032..e00902d9047 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1442,10 +1442,10 @@ class AsyncHTTPHandler: local_address: Final = _IPV4_LOCAL_ADDRESS if litellm.force_ipv4 else None if http2 or http2_enabled(): return AsyncHTTP2ConnectionSpread( - lambda: AsyncHTTPTransport( - local_address=local_address, verify=verify, cert=cert, http2=True, limits=limits + lambda lane_limits: AsyncHTTPTransport( + local_address=local_address, verify=verify, cert=cert, http2=True, limits=lane_limits ), - max_connections=limits.max_connections, + limits, ) if litellm.force_ipv4: return AsyncHTTPTransport(local_address=local_address, verify=verify, cert=cert, limits=limits) @@ -1464,8 +1464,10 @@ class AsyncHTTPHandler: if http2 or http2_enabled(): return _environment_proxy_mounts( lambda proxy_url: AsyncHTTP2ConnectionSpread( - lambda: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True, limits=limits), - max_connections=limits.max_connections, + lambda lane_limits: AsyncHTTPTransport( + proxy=proxy_url, verify=verify, cert=cert, http2=True, limits=lane_limits + ), + limits, ) ) return _environment_proxy_mounts( diff --git a/tests/unit/llms/custom_httpx/test_http2_connection_spread.py b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py index d637631110f..a83c013c6a6 100644 --- a/tests/unit/llms/custom_httpx/test_http2_connection_spread.py +++ b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py @@ -12,6 +12,7 @@ from litellm.llms.custom_httpx.http2_connection_spread import ( URL: Final = "https://search.example/v1/search" BURST: Final = HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2 +UNCAPPED: Final = httpx.Limits(max_connections=None) class _GatedAsyncConnection(httpx.AsyncBaseTransport): @@ -39,7 +40,7 @@ class _AsyncConnections: self._gate: Final = gate self.opened: tuple[_GatedAsyncConnection, ...] = () - def __call__(self) -> _GatedAsyncConnection: + def __call__(self, _limits: httpx.Limits) -> _GatedAsyncConnection: connection: Final = _GatedAsyncConnection(self._gate) self.opened = (*self.opened, connection) return connection @@ -57,11 +58,39 @@ class _RefusedConnections: def __init__(self) -> None: self.opened = 0 - def __call__(self) -> _RefusedConnection: + def __call__(self, _limits: httpx.Limits) -> _RefusedConnection: self.opened += 1 return _RefusedConnection() +class _HTTP1Connections: + def __init__(self, gate: asyncio.Event) -> None: + self._gate: Final = gate + self.requested = 0 + self.open = 0 + + def __call__(self, limits: httpx.Limits) -> "_HTTP1Pool": + assert limits.max_connections is not None + return _HTTP1Pool(self, asyncio.Semaphore(limits.max_connections)) + + async def serve(self, pool_slots: asyncio.Semaphore) -> None: + self.requested += 1 + async with pool_slots: + self.open += 1 + await self._gate.wait() + self.open -= 1 + + +class _HTTP1Pool(httpx.AsyncBaseTransport): + def __init__(self, connections: _HTTP1Connections, slots: asyncio.Semaphore) -> None: + self._connections: Final = connections + self._slots: Final = slots + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + await self._connections.serve(self._slots) + return httpx.Response(200, stream=httpx.ByteStream(b"ok")) + + async def _until(condition: Callable[[], bool]) -> None: for _ in range(1000): if condition(): @@ -74,7 +103,7 @@ async def _until(condition: Callable[[], bool]) -> None: async def test_async_spread_opens_another_connection_only_once_streams_run_out(): gate: Final = asyncio.Event() connections: Final = _AsyncConnections(gate) - async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections)) as client: + async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections, UNCAPPED)) as client: pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(BURST)) await _until(lambda: connections.in_flight() == BURST) assert [connection.peak for connection in connections.opened] == [ @@ -97,7 +126,10 @@ async def test_async_spread_shares_the_least_busy_connection_once_the_connection gate: Final = asyncio.Event() connections: Final = _AsyncConnections(gate) burst: Final = 2 * HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2 - async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections, max_connections=2)) as client: + two_connections_of_streams: Final = httpx.Limits(max_connections=2 * HTTP2_STREAMS_PER_CONNECTION) + async with httpx.AsyncClient( + transport=AsyncHTTP2ConnectionSpread(connections, two_connections_of_streams) + ) as client: pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(burst)) await _until(lambda: connections.in_flight() == burst) assert [connection.peak for connection in connections.opened] == [burst // 2, burst // 2] @@ -106,12 +138,29 @@ async def test_async_spread_shares_the_least_busy_connection_once_the_connection assert {response.status_code for response in await asyncio.gather(*pending)} == {200} +@pytest.mark.asyncio +async def test_async_spread_keeps_http1_connections_within_the_connection_cap(): + gate: Final = asyncio.Event() + connections: Final = _HTTP1Connections(gate) + cap: Final = HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2 + burst: Final = 2 * cap + async with httpx.AsyncClient( + transport=AsyncHTTP2ConnectionSpread(connections, httpx.Limits(max_connections=cap)) + ) as client: + pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(burst)) + await _until(lambda: connections.requested == burst) + assert connections.open == cap + + gate.set() + assert {response.status_code for response in await asyncio.gather(*pending)} == {200} + + @pytest.mark.asyncio async def test_async_spread_holds_a_stream_until_its_response_is_closed(): gate: Final = asyncio.Event() gate.set() connections: Final = _AsyncConnections(gate) - async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections)) as client: + async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections, UNCAPPED)) as client: open_responses: Final = tuple( [ await client.send(client.build_request("GET", URL), stream=True) @@ -130,7 +179,7 @@ async def test_async_spread_holds_a_stream_until_its_response_is_closed(): @pytest.mark.asyncio async def test_async_spread_frees_the_stream_when_the_request_fails(): connections: Final = _RefusedConnections() - async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections)) as client: + async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections, UNCAPPED)) as client: for _ in range(BURST): with pytest.raises(httpx.ConnectError): await client.get(URL)