fix(http): keep the HTTP/2 connection spread within the connection cap

litellm/llms/custom_httpx/http2_connection_spread.py

@@ -24,17 +24,25 @@ class _Lane(Generic[_TransportT]):
 
 
 class _Lanes(Generic[_TransportT]):
-    def __init__(self, new_transport: Callable[[], _TransportT], streams_per_connection: int) -> None:
+    def __init__(
+        self, new_transport: Callable[[], _TransportT], streams_per_connection: int, max_connections: int | None
+    ) -> None:
         self._new_transport: Final = new_transport
         self._streams_per_connection: Final = streams_per_connection
+        self._max_connections: Final = max_connections
         self.lanes: tuple[_Lane[_TransportT], ...] = ()
 
     def acquire(self) -> _Lane[_TransportT]:
         free: Final = next((lane for lane in self.lanes if lane.in_flight < self._streams_per_connection), None)
-        lane: Final = free if free is not None else self._open()
+        lane: Final = free if free is not None else self._open_or_share()
         lane.acquire()
         return lane
 
+    def _open_or_share(self) -> _Lane[_TransportT]:
+        if self._max_connections is not None and len(self.lanes) >= self._max_connections:
+            return min(self.lanes, key=lambda lane: lane.in_flight)
+        return self._open()
+
     def _open(self) -> _Lane[_TransportT]:
         lane: Final = _Lane(self._new_transport())
         self.lanes = (*self.lanes, lane)
@@ -77,8 +85,9 @@ class AsyncHTTP2ConnectionSpread(httpx.AsyncBaseTransport):
         self,
         new_transport: Callable[[], httpx.AsyncBaseTransport],
         streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION,
+        max_connections: int | None = None,
     ) -> None:
-        self._lanes: Final = _Lanes(new_transport, streams_per_connection)
+        self._lanes: Final = _Lanes(new_transport, streams_per_connection, max_connections)
 
     async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
         lane: Final = self._lanes.acquire()
@@ -103,8 +112,9 @@ class HTTP2ConnectionSpread(httpx.BaseTransport):
         self,
         new_transport: Callable[[], httpx.BaseTransport],
         streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION,
+        max_connections: int | None = None,
     ) -> None:
-        self._lanes: Final = _Lanes(new_transport, streams_per_connection)
+        self._lanes: Final = _Lanes(new_transport, streams_per_connection, max_connections)
         self._lock: Final = threading.Lock()
 
     def _release(self, lane: _Lane[httpx.BaseTransport]) -> None:

litellm/llms/custom_httpx/http_handler.py

@@ -1444,7 +1444,8 @@ class AsyncHTTPHandler:
             return AsyncHTTP2ConnectionSpread(
                 lambda: AsyncHTTPTransport(
                     local_address=local_address, verify=verify, cert=cert, http2=True, limits=limits
-                )
+                ),
+                max_connections=limits.max_connections,
             )
         if litellm.force_ipv4:
             return AsyncHTTPTransport(local_address=local_address, verify=verify, cert=cert, limits=limits)
@@ -1463,7 +1464,8 @@ 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)
+                    lambda: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True, limits=limits),
+                    max_connections=limits.max_connections,
                 )
             )
         return _environment_proxy_mounts(
@@ -1810,7 +1812,8 @@ class HTTPHandler:
         local_address: Final = _IPV4_LOCAL_ADDRESS if litellm.force_ipv4 else None
         if http2:
             return HTTP2ConnectionSpread(
-                lambda: HTTPTransport(local_address=local_address, verify=verify, cert=cert, http2=True)
+                lambda: HTTPTransport(local_address=local_address, verify=verify, cert=cert, http2=True),
+                max_connections=DEFAULT_LIMITS.max_connections,
             )
         if litellm.force_ipv4:
             return HTTPTransport(local_address=local_address, verify=verify, cert=cert)
@@ -1828,7 +1831,8 @@ class HTTPHandler:
         if http2:
             return _environment_proxy_mounts(
                 lambda proxy_url: HTTP2ConnectionSpread(
-                    lambda: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True)
+                    lambda: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True),
+                    max_connections=DEFAULT_LIMITS.max_connections,
                 )
             )
         return _environment_proxy_mounts(lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert))

tests/integration/sdk/test_http2_wire.py

@@ -41,6 +41,7 @@ HELD_SEARCHES: Final = 150
 HOLD_SECONDS: Final = 10.0
 _HOLD_ROOMS: Final = MappingProxyType({"hold-async": queue.SimpleQueue[int](), "hold-sync": queue.SimpleQueue[int]()})
 _TOKEN_REQUESTS: Final = queue.SimpleQueue[str]()
+_COUNTED_CLIENT_PORTS: Final = queue.SimpleQueue[int]()
 VERTEX_SERVING_CONFIG_PATH: Final = (
     "/v1/projects/p/locations/global/collections/default_collection/engines/e/servingConfigs/default_search"
 )
@@ -113,6 +114,11 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab
         _TOKEN_REQUESTS.put(scope["path"])
         await _send_json(send, {"access_token": "scripted", "expires_in": 3600, "token_type": "Bearer"})
         return
+    if scope["path"] == "/counted:search":
+        counted_client: Final = scope["client"]
+        _COUNTED_CLIENT_PORTS.put(counted_client[1] if counted_client else 0)
+        await _send_json(send, {"results": [{"id": "counted"}]})
+        return
     if scope["path"].endswith(":search"):
         await _send_json(send, {"results": [{"id": f"http/{version}"}]})
         return
@@ -631,6 +637,28 @@ def test_http2_handlers_open_another_connection_once_the_streams_run_out(
     assert len(client_ports) == 2
 
 
+def test_http2_transport_opens_no_more_connections_than_its_cap(http2_tls_peer: TlsPeer) -> None:
+    from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+
+    async def _burst() -> frozenset[str]:
+        transport: Final = AsyncHTTPHandler.create_async_transport(
+            ssl_context=ssl.create_default_context(cafile=str(http2_tls_peer.cert_file)),
+            http2=True,
+            limits=httpx.Limits(max_connections=1),
+        )
+        async with httpx.AsyncClient(transport=transport) as client:
+            responses: Final = await asyncio.gather(
+                *(client.post(f"{http2_tls_peer.url}/counted:search", json={}) for _ in range(HELD_SEARCHES))
+            )
+        return frozenset(response.json()["results"][0]["id"] for response in responses)
+
+    assert asyncio.run(_burst()) == {"counted"}
+    client_ports: Final = frozenset(
+        iter(lambda: None if _COUNTED_CLIENT_PORTS.empty() else _COUNTED_CLIENT_PORTS.get(), None)
+    )
+    assert len(client_ports) == 1
+
+
 def test_forced_ipv4_handlers_honour_ssl_verify_false(monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer) -> None:
     import litellm
 

tests/unit/llms/custom_httpx/test_http2_connection_spread.py

@@ -136,6 +136,20 @@ async def test_async_spread_opens_another_connection_only_once_streams_run_out()
     assert all(connection.closed for connection in connections.opened)
 
 
+@pytest.mark.asyncio
+async def test_async_spread_shares_the_least_busy_connection_once_the_connection_cap_is_reached():
+    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:
+        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]
+
+        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()