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)