From 81003176d49b3f7425b022f4f273f7ad28a06556 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:01:53 -0700 Subject: [PATCH] fix(http): open another HTTP/2 connection once 100 streams are in use httpcore keeps every request to an origin on one HTTP/2 connection and caps it at 100 concurrent streams, so request 101 and later waited with no timeout. The handlers now spread requests over extra connections, opening one only when every open connection has 100 streams in flight. The force_ipv4 wire tests now cover ssl_verify=False and the client certificate. --- .../custom_httpx/http2_connection_spread.py | 130 ++++++++++++ litellm/llms/custom_httpx/http_handler.py | 87 +++++--- tests/integration/sdk/test_http2_wire.py | 98 ++++++++- .../test_http2_connection_spread.py | 190 ++++++++++++++++++ .../llms/custom_httpx/test_http_handler.py | 4 +- 5 files changed, 474 insertions(+), 35 deletions(-) create mode 100644 litellm/llms/custom_httpx/http2_connection_spread.py create mode 100644 tests/unit/llms/custom_httpx/test_http2_connection_spread.py diff --git a/litellm/llms/custom_httpx/http2_connection_spread.py b/litellm/llms/custom_httpx/http2_connection_spread.py new file mode 100644 index 00000000000..3803ccdf56e --- /dev/null +++ b/litellm/llms/custom_httpx/http2_connection_spread.py @@ -0,0 +1,130 @@ +import threading +from collections.abc import AsyncIterator, Callable, Iterator +from typing import Final, Generic, TypeVar + +import httpx + +HTTP2_STREAMS_PER_CONNECTION: Final = 100 + +_TransportT: Final = TypeVar("_TransportT", httpx.BaseTransport, httpx.AsyncBaseTransport) + + +class _Lane(Generic[_TransportT]): + __slots__ = ("in_flight", "transport") + + def __init__(self, transport: _TransportT) -> None: + self.transport: Final = transport + self.in_flight = 0 + + def acquire(self) -> None: + self.in_flight += 1 + + def release(self) -> None: + self.in_flight -= 1 + + +class _Lanes(Generic[_TransportT]): + def __init__(self, new_transport: Callable[[], _TransportT], streams_per_connection: int) -> None: + self._new_transport: Final = new_transport + self._streams_per_connection: Final = streams_per_connection + 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.acquire() + return lane + + def _open(self) -> _Lane[_TransportT]: + lane: Final = _Lane(self._new_transport()) + self.lanes = (*self.lanes, lane) + return lane + + +class _AsyncReleasingStream(httpx.AsyncByteStream): + def __init__(self, stream: httpx.AsyncByteStream, release: Callable[[], None]) -> None: + self._stream: Final = stream + self._release: Final = release + + async def __aiter__(self) -> AsyncIterator[bytes]: + async for chunk in self._stream: + yield chunk + + async def aclose(self) -> None: + self._release() + await self._stream.aclose() + + +class _ReleasingStream(httpx.SyncByteStream): + def __init__(self, stream: httpx.SyncByteStream, release: Callable[[], None]) -> None: + self._stream: Final = stream + self._release: Final = release + + def __iter__(self) -> Iterator[bytes]: + yield from self._stream + + def close(self) -> None: + self._release() + self._stream.close() + + +def _with_stream(response: httpx.Response, stream: httpx.SyncByteStream | httpx.AsyncByteStream) -> httpx.Response: + return httpx.Response(response.status_code, headers=response.headers, stream=stream, extensions=response.extensions) + + +class AsyncHTTP2ConnectionSpread(httpx.AsyncBaseTransport): + def __init__( + self, + new_transport: Callable[[], httpx.AsyncBaseTransport], + streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION, + ) -> None: + self._lanes: Final = _Lanes(new_transport, streams_per_connection) + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + lane: Final = self._lanes.acquire() + try: + response: Final = await lane.transport.handle_async_request(request) + except BaseException: + lane.release() + raise + stream: Final = response.stream + if not isinstance(stream, httpx.AsyncByteStream): + lane.release() + return response + return _with_stream(response, _AsyncReleasingStream(stream, lane.release)) + + async def aclose(self) -> None: + for lane in self._lanes.lanes: + await lane.transport.aclose() + + +class HTTP2ConnectionSpread(httpx.BaseTransport): + def __init__( + self, + new_transport: Callable[[], httpx.BaseTransport], + streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION, + ) -> None: + self._lanes: Final = _Lanes(new_transport, streams_per_connection) + self._lock: Final = threading.Lock() + + def _release(self, lane: _Lane[httpx.BaseTransport]) -> None: + with self._lock: + lane.release() + + def handle_request(self, request: httpx.Request) -> httpx.Response: + with self._lock: + lane: Final = self._lanes.acquire() + try: + response: Final = lane.transport.handle_request(request) + except BaseException: + self._release(lane) + raise + stream: Final = response.stream + if not isinstance(stream, httpx.SyncByteStream): + self._release(lane) + return response + return _with_stream(response, _ReleasingStream(stream, lambda: self._release(lane))) + + def close(self) -> None: + for lane in self._lanes.lanes: + lane.transport.close() diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index ac439ec1918..04e72bc7b28 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -43,6 +43,7 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) +from litellm.llms.custom_httpx.http2_connection_spread import AsyncHTTP2ConnectionSpread, HTTP2ConnectionSpread from litellm.types.llms.custom_http import * if TYPE_CHECKING: @@ -73,7 +74,7 @@ _RequestContent: TypeAlias = str | bytes | Iterable[bytes] | AsyncIterable[bytes _IPV4_LOCAL_ADDRESS: Final = "0.0.0.0" -_HttpxTransportT = TypeVar("_HttpxTransportT", HTTPTransport, AsyncHTTPTransport) +_HttpxTransportT: Final = TypeVar("_HttpxTransportT", httpx.BaseTransport, httpx.AsyncBaseTransport) HTTPEventHooks: TypeAlias = Mapping[str, list[Callable[..., object]]] @@ -1245,7 +1246,7 @@ class AsyncHTTPHandler: http2: bool = False, cert: CertTypes | None = None, limits: httpx.Limits = DEFAULT_LIMITS, - ) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | None: + ) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None: """ - Creates a transport for httpx.AsyncClient - if litellm.force_ipv4 is True, it will return AsyncHTTPTransport with local_address="0.0.0.0" @@ -1430,38 +1431,43 @@ class AsyncHTTPHandler: verify: ssl.SSLContext | bool = True, cert: CertTypes | None = None, limits: httpx.Limits = DEFAULT_LIMITS, - ) -> AsyncHTTPTransport | None: + ) -> AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None: """ Creates an AsyncHTTPTransport + - If HTTP/2 is on, it spreads requests over as many HTTP/2 connections as the stream limit needs - If force_ipv4 is True, it will create an AsyncHTTPTransport with local_address set to "0.0.0.0" - - [Default] If force_ipv4 is False, it will return None + - [Default] Otherwise it will return None """ - if litellm.force_ipv4: - return AsyncHTTPTransport( - local_address=_IPV4_LOCAL_ADDRESS, - verify=verify, - cert=cert, - http2=http2 or http2_enabled(), - limits=limits, + 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 + ) ) - else: - return None + if litellm.force_ipv4: + return AsyncHTTPTransport(local_address=local_address, verify=verify, cert=cert, limits=limits) + return None @staticmethod def create_httpx_proxy_mounts( - transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | None, + transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None, verify: VerifyTypes, cert: CertTypes | None, http2: bool = False, limits: httpx.Limits = DEFAULT_LIMITS, - ) -> Mapping[str, AsyncHTTPTransport | None] | None: - if not isinstance(transport, AsyncHTTPTransport): + ) -> Mapping[str, httpx.AsyncBaseTransport | None] | None: + if not isinstance(transport, (AsyncHTTPTransport, AsyncHTTP2ConnectionSpread)): return None - return _environment_proxy_mounts( - lambda proxy_url: AsyncHTTPTransport( - proxy=proxy_url, verify=verify, cert=cert, http2=http2 or http2_enabled(), limits=limits + 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) + ) ) + return _environment_proxy_mounts( + lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, limits=limits) ) _create_httpx_proxy_mounts = create_httpx_proxy_mounts @@ -1500,10 +1506,12 @@ class HTTPHandler: # Get default headers (User-Agent, overridable via LITELLM_USER_AGENT) default_headers: Final = get_default_headers() if not self.disable_default_headers else None + transport: Final = self._create_sync_transport(verify=ssl_config, cert=cert, http2=use_http2) + # Create a client with a connection pool return httpx.Client( - transport=self._create_sync_transport(verify=ssl_config, cert=cert, http2=use_http2), - mounts=self._create_sync_proxy_mounts(verify=ssl_config, cert=cert, http2=use_http2), + transport=transport, + mounts=self._create_sync_proxy_mounts(transport, verify=ssl_config, cert=cert, http2=use_http2), timeout=self.timeout if self.timeout is not None else _DEFAULT_TIMEOUT, verify=ssl_config, cert=cert, @@ -1786,29 +1794,44 @@ class HTTPHandler: pass @staticmethod - def _create_sync_transport(verify: VerifyTypes, cert: CertTypes | None, http2: bool) -> HTTPTransport | None: + def _create_sync_transport( + verify: VerifyTypes, cert: CertTypes | None, http2: bool + ) -> HTTPTransport | HTTP2ConnectionSpread | None: """ - Create an HTTP transport with IPv4 only if litellm.force_ipv4 is True. - Otherwise, return None. + Create an HTTP transport with IPv4 only if litellm.force_ipv4 is True, and one that spreads requests over + as many HTTP/2 connections as the stream limit needs if HTTP/2 is on. + Otherwise, return litellm.sync_transport or None. Some users have seen httpx ConnectionError when using ipv6 - forcing ipv4 resolves the issue for them """ + custom_transport: Final = None if litellm.force_ipv4 else getattr(litellm, "sync_transport", None) + if custom_transport is not None: + return custom_transport + 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) + ) if litellm.force_ipv4: - return HTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, verify=verify, cert=cert, http2=http2) - else: - return getattr(litellm, "sync_transport", None) + return HTTPTransport(local_address=local_address, verify=verify, cert=cert) + return None @staticmethod def _create_sync_proxy_mounts( + transport: httpx.BaseTransport | None, verify: VerifyTypes, cert: CertTypes | None, http2: bool, - ) -> Mapping[str, HTTPTransport | None] | None: - if not litellm.force_ipv4: + ) -> Mapping[str, httpx.BaseTransport | None] | None: + if transport is None or transport is getattr(litellm, "sync_transport", None): return None - return _environment_proxy_mounts( - lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2) - ) + if http2: + return _environment_proxy_mounts( + lambda proxy_url: HTTP2ConnectionSpread( + lambda: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True) + ) + ) + return _environment_proxy_mounts(lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert)) def get_async_httpx_client( diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index 527e43b10d3..4ff38b82837 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/tests/integration/sdk/test_http2_wire.py @@ -11,9 +11,11 @@ import ssl import threading import time from collections.abc import Callable, Generator, Iterator +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import Final import httpx @@ -35,6 +37,9 @@ from integration._support.ports import allocate_port from integration._support.vertex import service_account_json STREAM_CHUNKS: Final = 3 +HELD_SEARCHES: Final = 150 +HOLD_SECONDS: Final = 10.0 +_HOLD_ROOMS: Final = MappingProxyType({"hold-async": queue.SimpleQueue[int](), "hold-sync": queue.SimpleQueue[int]()}) VERTEX_SERVING_CONFIG_PATH: Final = ( "/v1/projects/p/locations/global/collections/default_collection/engines/e/servingConfigs/default_search" ) @@ -96,6 +101,13 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab if message["type"] == "http.request" and not message["more_body"]: break version: Final = scope["http_version"] + hold_room: Final = _HOLD_ROOMS.get(scope["path"].split("/")[1]) + if hold_room is not None and scope["path"].endswith(":search"): + client: Final = scope["client"] + await _send_json( + send, {"results": [{"id": await _hold_until_all_arrive(hold_room, client[1] if client else 0)}]} + ) + return if scope["path"] == "/_oauth/token": await _send_json(send, {"access_token": "scripted", "expires_in": 3600, "token_type": "Bearer"}) return @@ -119,6 +131,14 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab await _send_json(send, {"http_version": version}) +async def _hold_until_all_arrive(arrivals: queue.SimpleQueue[int], client_port: int) -> str: + arrivals.put(client_port) + deadline: Final = time.monotonic() + HOLD_SECONDS + while arrivals.qsize() < HELD_SEARCHES and time.monotonic() < deadline: + await asyncio.sleep(0.01) + return "held-together" if arrivals.qsize() >= HELD_SEARCHES else "released-alone" + + async def _send_json(send: ASGISendCallable, body: dict[str, object]) -> None: await send( HTTPResponseStartEvent(type="http.response.start", status=200, headers=[(b"content-type", b"application/json")]) @@ -454,11 +474,17 @@ def _answer_every_connection(context: ssl.SSLContext, raw: socket.socket, connec @contextmanager def _raw_tls_peer( - peer: TlsPeer, alpn: str, serve_connection: Callable[[ssl.SSLContext, socket.socket, int], None] + peer: TlsPeer, + alpn: str, + serve_connection: Callable[[ssl.SSLContext, socket.socket, int], None], + client_ca: Path | None = None, ) -> Generator[tuple[str, queue.SimpleQueue[int]]]: context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) context.load_cert_chain(peer.cert_file, peer.key_file) context.set_alpn_protocols([alpn]) + if client_ca is not None: + context.verify_mode = ssl.CERT_REQUIRED + context.load_verify_locations(client_ca) stop: Final = threading.Event() accepted: Final = queue.SimpleQueue[int]() @@ -532,3 +558,73 @@ def test_vertex_search_keeps_its_http2_connection_across_an_idle_gap( with _raw_tls_peer(http2_tls_peer, "h2", _answer_every_connection) as (search_url, accepted): assert asyncio.run(_search_twice(search_url)) == ["answered", "answered"] assert accepted.qsize() == 1 + + +def _held_burst(peer: TlsPeer, room: str) -> list[str]: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + + url: Final = f"{peer.url}/{room}/default_search:search" + if room == "hold-sync": + handler: Final = HTTPHandler(ssl_verify=str(peer.cert_file), http2=True) + try: + with ThreadPoolExecutor(max_workers=HELD_SEARCHES) as pool: + responses: Final = tuple(pool.map(lambda _: handler.post(url, json={}), range(HELD_SEARCHES))) + return [response.json()["results"][0]["id"] for response in responses] + finally: + handler.close() + + async def _burst() -> list[str]: + async_handler: Final = AsyncHTTPHandler(ssl_verify=str(peer.cert_file), http2=True) + try: + async_responses: Final = await asyncio.gather( + *(async_handler.post(url, json={}) for _ in range(HELD_SEARCHES)) + ) + return [response.json()["results"][0]["id"] for response in async_responses] + finally: + await async_handler.close() + + return asyncio.run(_burst()) + + +@pytest.mark.parametrize("room", ["hold-async", "hold-sync"]) +def test_http2_handlers_open_another_connection_once_the_streams_run_out( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer, room: str +) -> None: + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + + search_ids: Final = _held_burst(http2_tls_peer, room) + arrivals: Final = _HOLD_ROOMS[room] + client_ports: Final = frozenset(arrivals.get_nowait() for _ in range(arrivals.qsize())) + + assert search_ids == ["held-together"] * HELD_SEARCHES + assert len(client_ports) == 2 + + +def test_forced_ipv4_handlers_honour_ssl_verify_false(monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer) -> None: + import litellm + + monkeypatch.setattr(litellm, "force_ipv4", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + for enabled in (False, True): + _set_http2(monkeypatch, enabled) + _assert_negotiated(_async_exchange(http2_tls_peer.url), enabled) + _assert_negotiated(_sync_exchange(http2_tls_peer.url), enabled) + + +def test_vertex_search_over_forced_ipv4_presents_the_client_certificate( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer, tmp_path: Path +) -> None: + import litellm + + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + monkeypatch.setattr(litellm, "force_ipv4", True) + client_pem: Final = tmp_path / "client.pem" + client_pem.write_bytes(http2_tls_peer.cert_file.read_bytes() + http2_tls_peer.key_file.read_bytes()) + monkeypatch.setenv("SSL_CERTIFICATE", str(client_pem)) + + with _raw_tls_peer(http2_tls_peer, "h2", _answer_every_connection, client_ca=http2_tls_peer.cert_file) as ( + search_url, + _, + ): + assert _vertex_search_file_ids(http2_tls_peer, search_url=search_url) == ["answered"] + assert asyncio.run(_avertex_search_file_ids(http2_tls_peer, search_url=search_url)) == ["answered"] diff --git a/tests/unit/llms/custom_httpx/test_http2_connection_spread.py b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py new file mode 100644 index 00000000000..eddbc18f41a --- /dev/null +++ b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py @@ -0,0 +1,190 @@ +import asyncio +import threading +import time +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from typing import Final + +import httpx +import pytest + +from litellm.llms.custom_httpx.http2_connection_spread import ( + HTTP2_STREAMS_PER_CONNECTION, + AsyncHTTP2ConnectionSpread, + HTTP2ConnectionSpread, +) + +URL: Final = "https://search.example/v1/search" +BURST: Final = HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2 + + +class _GatedAsyncConnection(httpx.AsyncBaseTransport): + def __init__(self, gate: asyncio.Event) -> None: + self._gate: Final = gate + self.in_flight = 0 + self.peak = 0 + self.served = 0 + self.closed = False + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + self.in_flight += 1 + self.peak = max(self.peak, self.in_flight) + await self._gate.wait() + self.in_flight -= 1 + self.served += 1 + return httpx.Response(200, stream=httpx.ByteStream(b"ok")) + + async def aclose(self) -> None: + self.closed = True + + +class _AsyncConnections: + def __init__(self, gate: asyncio.Event) -> None: + self._gate: Final = gate + self.opened: tuple[_GatedAsyncConnection, ...] = () + + def __call__(self) -> _GatedAsyncConnection: + connection: Final = _GatedAsyncConnection(self._gate) + self.opened = (*self.opened, connection) + return connection + + def in_flight(self) -> int: + return sum(connection.in_flight for connection in self.opened) + + +class _RefusedConnection(httpx.AsyncBaseTransport): + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("refused", request=request) + + +class _RefusedConnections: + def __init__(self) -> None: + self.opened = 0 + + def __call__(self) -> _RefusedConnection: + self.opened += 1 + return _RefusedConnection() + + +class _GatedConnection(httpx.BaseTransport): + def __init__(self, gate: threading.Event) -> None: + self._gate: Final = gate + self._lock: Final = threading.Lock() + self.in_flight = 0 + self.peak = 0 + self.served = 0 + + def handle_request(self, request: httpx.Request) -> httpx.Response: + with self._lock: + self.in_flight += 1 + self.peak = max(self.peak, self.in_flight) + assert self._gate.wait(timeout=10) + with self._lock: + self.in_flight -= 1 + self.served += 1 + return httpx.Response(200, stream=httpx.ByteStream(b"ok")) + + +class _Connections: + def __init__(self, gate: threading.Event) -> None: + self._gate: Final = gate + self.opened: tuple[_GatedConnection, ...] = () + + def __call__(self) -> _GatedConnection: + connection: Final = _GatedConnection(self._gate) + self.opened = (*self.opened, connection) + return connection + + def in_flight(self) -> int: + return sum(connection.in_flight for connection in self.opened) + + +async def _until(condition: Callable[[], bool]) -> None: + for _ in range(1000): + if condition(): + return + await asyncio.sleep(0) + pytest.fail("condition never held") + + +def _wait_until(condition: Callable[[], bool]) -> None: + deadline: Final = time.monotonic() + 10 + while not condition(): + assert time.monotonic() < deadline, "condition never held" + time.sleep(0.01) + + +@pytest.mark.asyncio +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: + 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] == [ + HTTP2_STREAMS_PER_CONNECTION, + BURST - HTTP2_STREAMS_PER_CONNECTION, + ] + + gate.set() + assert {response.status_code for response in await asyncio.gather(*pending)} == {200} + assert (await client.get(URL)).text == "ok" + assert [connection.served for connection in connections.opened] == [ + HTTP2_STREAMS_PER_CONNECTION + 1, + BURST - HTTP2_STREAMS_PER_CONNECTION, + ] + assert all(connection.closed for connection in connections.opened) + + +@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: + open_responses: Final = tuple( + [ + await client.send(client.build_request("GET", URL), stream=True) + for _ in range(HTTP2_STREAMS_PER_CONNECTION) + ] + ) + await client.get(URL) + assert len(connections.opened) == 2 + + for response in open_responses: + await response.aclose() + await client.get(URL) + assert [connection.served for connection in connections.opened] == [HTTP2_STREAMS_PER_CONNECTION + 1, 1] + + +@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: + for _ in range(BURST): + with pytest.raises(httpx.ConnectError): + await client.get(URL) + assert connections.opened == 1 + + +def test_sync_spread_opens_another_connection_only_once_streams_run_out(): + gate: Final = threading.Event() + connections: Final = _Connections(gate) + with ( + httpx.Client(transport=HTTP2ConnectionSpread(connections)) as client, + ThreadPoolExecutor(max_workers=BURST) as pool, + ): + pending: Final = tuple(pool.submit(client.get, URL) for _ in range(BURST)) + _wait_until(lambda: connections.in_flight() == BURST) + assert [connection.peak for connection in connections.opened] == [ + HTTP2_STREAMS_PER_CONNECTION, + BURST - HTTP2_STREAMS_PER_CONNECTION, + ] + + gate.set() + assert {future.result().status_code for future in pending} == {200} + assert client.get(URL).text == "ok" + assert [connection.served for connection in connections.opened] == [ + HTTP2_STREAMS_PER_CONNECTION + 1, + BURST - HTTP2_STREAMS_PER_CONNECTION, + ] diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index bd4ed4c99cc..861b2bfa58d 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1917,12 +1917,12 @@ async def test_http2_flag_bypasses_aiohttp_transport(monkeypatch: pytest.MonkeyP monkeypatch.setattr(litellm, "http2", True) assert AsyncHTTPHandler.should_use_aiohttp_transport() is False - assert AsyncHTTPHandler.create_async_transport() is None + assert not isinstance(AsyncHTTPHandler.create_async_transport(), LiteLLMAiohttpTransport) monkeypatch.setattr(litellm, "http2", False) monkeypatch.setenv("LITELLM_HTTP2", "True") assert AsyncHTTPHandler.should_use_aiohttp_transport() is False - assert AsyncHTTPHandler.create_async_transport() is None + assert not isinstance(AsyncHTTPHandler.create_async_transport(), LiteLLMAiohttpTransport) @pytest.mark.asyncio