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