mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
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.
This commit is contained in:
mateo-berri
parent
5894a0c46b
commit
81003176d4
5 changed files with 474 additions and 35 deletions
130
litellm/llms/custom_httpx/http2_connection_spread.py
Normal file
130
litellm/llms/custom_httpx/http2_connection_spread.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
190
tests/unit/llms/custom_httpx/test_http2_connection_spread.py
Normal file
190
tests/unit/llms/custom_httpx/test_http2_connection_spread.py
Normal file
|
|
@ -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,
|
||||
]
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue