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.

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

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(

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)