diff --git a/litellm/llms/custom_httpx/http2_connection_spread.py b/litellm/llms/custom_httpx/http2_connection_spread.py index 3803ccdf56e..cabdd5011cd 100644 --- a/litellm/llms/custom_httpx/http2_connection_spread.py +++ b/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: diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 04e72bc7b28..e83d9e66758 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/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)) diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index 793dd8d490b..07b62a7d69e 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/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 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 eddbc18f41a..36d79e0260b 100644 --- a/tests/unit/llms/custom_httpx/test_http2_connection_spread.py +++ b/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()