diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 6a2a69f9199..328c68c8b5f 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -18,6 +18,7 @@ import certifi import httpx from aiohttp import ClientSession, DummyCookieJar, TCPConnector from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport +from httpx._config import DEFAULT_LIMITS from httpx._types import CertTypes, RequestFiles from httpx._utils import get_environment_proxies @@ -101,6 +102,13 @@ def _aiohttp_env_proxies_disabled() -> bool: return litellm.disable_aiohttp_trust_env or str_to_bool(os.getenv("DISABLE_AIOHTTP_TRUST_ENV", "False")) is True +_AIOHTTP_EQUIVALENT_LIMITS: Final = httpx.Limits( + max_connections=AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT > 0 else None, + max_keepalive_connections=None, + keepalive_expiry=AIOHTTP_KEEPALIVE_TIMEOUT, +) + + def _environment_proxy_mounts( build_proxy_transport: Callable[[str], _HttpxTransportT], ) -> Mapping[str, _HttpxTransportT | None]: @@ -705,6 +713,8 @@ class AsyncHTTPHandler: trust_env=False, ) use_http2: Final = _resolve_http2(self.http2, ssl_verify) + replaces_aiohttp: Final = use_http2 and AsyncHTTPHandler.should_use_aiohttp_transport() + limits: Final = _AIOHTTP_EQUIVALENT_LIMITS if replaces_aiohttp else DEFAULT_LIMITS # Get unified SSL configuration ssl_config: Final = get_ssl_configuration(ssl_verify, http2=use_http2) @@ -722,17 +732,21 @@ class AsyncHTTPHandler: shared_session=shared_session, http2=use_http2, cert=cert, + limits=limits, ) # Get default headers (User-Agent, overridable via LITELLM_USER_AGENT) default_headers: Final = get_default_headers() - trust_env: Final = not (use_http2 and _aiohttp_env_proxies_disabled()) + trust_env: Final = not (replaces_aiohttp and _aiohttp_env_proxies_disabled()) return httpx.AsyncClient( transport=transport, - mounts=AsyncHTTPHandler.create_httpx_proxy_mounts(transport, verify=ssl_config, cert=cert, http2=use_http2) + mounts=AsyncHTTPHandler.create_httpx_proxy_mounts( + transport, verify=ssl_config, cert=cert, http2=use_http2, limits=limits + ) if trust_env else None, + limits=limits, event_hooks=event_hooks, timeout=timeout, verify=ssl_config, @@ -1233,6 +1247,7 @@ class AsyncHTTPHandler: shared_session: Optional["ClientSession"] = None, http2: bool = False, cert: CertTypes | None = None, + limits: httpx.Limits = DEFAULT_LIMITS, ) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | None: """ - Creates a transport for httpx.AsyncClient @@ -1265,6 +1280,7 @@ class AsyncHTTPHandler: http2=http2, verify=ssl_context if ssl_context is not None else ssl_verify is not False, cert=cert, + limits=limits, ) _create_async_transport = create_async_transport @@ -1416,6 +1432,7 @@ class AsyncHTTPHandler: http2: bool = False, verify: ssl.SSLContext | bool = True, cert: CertTypes | None = None, + limits: httpx.Limits = DEFAULT_LIMITS, ) -> AsyncHTTPTransport | None: """ Creates an AsyncHTTPTransport @@ -1425,7 +1442,11 @@ class AsyncHTTPHandler: """ if litellm.force_ipv4: return AsyncHTTPTransport( - local_address=_IPV4_LOCAL_ADDRESS, verify=verify, cert=cert, http2=http2 or http2_enabled() + local_address=_IPV4_LOCAL_ADDRESS, + verify=verify, + cert=cert, + http2=http2 or http2_enabled(), + limits=limits, ) else: return None @@ -1436,12 +1457,13 @@ class AsyncHTTPHandler: verify: VerifyTypes, cert: CertTypes | None, http2: bool = False, + limits: httpx.Limits = DEFAULT_LIMITS, ) -> Mapping[str, AsyncHTTPTransport | None] | None: if not isinstance(transport, AsyncHTTPTransport): return None return _environment_proxy_mounts( lambda proxy_url: AsyncHTTPTransport( - proxy=proxy_url, verify=verify, cert=cert, http2=http2 or http2_enabled() + proxy=proxy_url, verify=verify, cert=cert, http2=http2 or http2_enabled(), limits=limits ) ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 1ddd5cdd61d..236674a4a8b 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -25,6 +25,7 @@ from typing import ( from urllib.parse import parse_qs, urlencode, urlparse, urlunparse import httpx +from h2.events import ConnectionTerminated from httpx import USE_CLIENT_DEFAULT from httpx._types import FileContent from openai.types.file_deleted import FileDeleted @@ -340,17 +341,31 @@ class _SignedBodyKwargs(TypedDict, total=False): json: ReadOnly[dict[str, object]] -async def _asend_idempotent(send: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response: +_GOAWAY_RETRY_MESSAGE: Final = "Vector store search connection was closed by an HTTP/2 GOAWAY, retrying once" + + +def _closed_by_http2_goaway(error: httpx.RemoteProtocolError) -> bool: + cause: Final = error.__cause__ + return cause is not None and any(isinstance(arg, ConnectionTerminated) for arg in cause.args) + + +async def _asend_with_goaway_retry(send: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response: try: return await send() - except httpx.RemoteProtocolError: + except httpx.RemoteProtocolError as error: + if not _closed_by_http2_goaway(error): + raise + verbose_logger.warning(_GOAWAY_RETRY_MESSAGE) return await send() -def _send_idempotent(send: Callable[[], httpx.Response]) -> httpx.Response: +def _send_with_goaway_retry(send: Callable[[], httpx.Response]) -> httpx.Response: try: return send() - except httpx.RemoteProtocolError: + except httpx.RemoteProtocolError as error: + if not _closed_by_http2_goaway(error): + raise + verbose_logger.warning(_GOAWAY_RETRY_MESSAGE) return send() @@ -10405,7 +10420,7 @@ class BaseLLMHTTPHandler: request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body try: - response: Final = await _asend_idempotent( + response: Final = await _asend_with_goaway_retry( partial(async_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout) ) except httpx.TimeoutException: @@ -10544,7 +10559,7 @@ class BaseLLMHTTPHandler: request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body try: - response: Final = _send_idempotent( + response: Final = _send_with_goaway_retry( partial(sync_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout) ) except httpx.TimeoutException: diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index b16f1699aaa..527e43b10d3 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/tests/integration/sdk/test_http2_wire.py @@ -4,18 +4,19 @@ import asyncio import contextlib import datetime import ipaddress -import itertools import json +import queue import socket import ssl import threading import time -from collections.abc import Generator, Iterator +from collections.abc import Callable, Generator, Iterator from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from typing import Final +import httpx import pytest from cryptography import x509 from cryptography.hazmat.primitives import hashes, serialization @@ -24,6 +25,7 @@ from cryptography.x509.oid import NameOID from h2.config import H2Configuration from h2.connection import H2Connection from h2.events import StreamEnded +from httpx._config import DEFAULT_LIMITS from hypercorn.asyncio import ( serve, # pyright: ignore[reportUnknownVariableType] # hypercorn's serve signature passes through untyped worker hooks ) @@ -356,29 +358,50 @@ def test_vertex_search_leaves_a_caller_ssl_context_on_http1( assert _alpn_offered_by(async_context, http2_tls_peer) != "h2" +def _route_https_through_a_dead_proxy(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("https_proxy", "HTTP_PROXY", "http_proxy", "ALL_PROXY", "all_proxy", "NO_PROXY", "no_proxy"): + monkeypatch.delenv(name, raising=False) + monkeypatch.delenv("DISABLE_AIOHTTP_TRANSPORT", raising=False) + monkeypatch.setenv("HTTPS_PROXY", f"http://127.0.0.1:{allocate_port()}") + monkeypatch.setenv("DISABLE_AIOHTTP_TRUST_ENV", "True") + + +async def _async_peer_http_version(peer: TlsPeer, http2: bool) -> str: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + handler: Final = AsyncHTTPHandler(ssl_verify=str(peer.cert_file), http2=http2) + try: + response: Final = await handler.client.post(f"{peer.url}/echo", json={"ping": "pong"}) + return response.http_version + finally: + await handler.close() + + def test_http2_client_skips_env_proxies_when_aiohttp_trust_env_is_disabled( monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer ) -> None: import litellm - from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler _set_http2(monkeypatch, False) - for name in ("https_proxy", "HTTP_PROXY", "http_proxy", "ALL_PROXY", "all_proxy", "NO_PROXY", "no_proxy"): - monkeypatch.delenv(name, raising=False) - monkeypatch.setenv("HTTPS_PROXY", f"http://127.0.0.1:{allocate_port()}") - monkeypatch.setenv("DISABLE_AIOHTTP_TRUST_ENV", "True") - - async def _peer_http_version() -> str: - handler: Final = AsyncHTTPHandler(ssl_verify=str(http2_tls_peer.cert_file), http2=True) - try: - response: Final = await handler.client.post(f"{http2_tls_peer.url}/echo", json={"ping": "pong"}) - return response.http_version - finally: - await handler.close() + _route_https_through_a_dead_proxy(monkeypatch) for force_ipv4 in (False, True): monkeypatch.setattr(litellm, "force_ipv4", force_ipv4) - assert asyncio.run(_peer_http_version()) == "HTTP/2" + assert asyncio.run(_async_peer_http_version(http2_tls_peer, http2=True)) == "HTTP/2" + + +def test_global_http2_toggle_keeps_env_proxies_when_aiohttp_trust_env_is_disabled( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + import litellm + + _set_http2(monkeypatch, True) + _route_https_through_a_dead_proxy(monkeypatch) + + for force_ipv4 in (False, True): + monkeypatch.setattr(litellm, "force_ipv4", force_ipv4) + with pytest.raises(httpx.ConnectError): + asyncio.run(_async_peer_http_version(http2_tls_peer, http2=False)) def _answer_or_drain(connection: H2Connection, stream_id: int, drain: bool) -> None: @@ -386,7 +409,7 @@ def _answer_or_drain(connection: H2Connection, stream_id: int, drain: bool) -> N connection.close_connection(last_stream_id=stream_id) return connection.send_headers(stream_id, [(":status", "200"), ("content-type", "application/json")]) - connection.send_data(stream_id, json.dumps({"results": [{"id": "after-goaway"}]}).encode(), end_stream=True) + connection.send_data(stream_id, json.dumps({"results": [{"id": "answered"}]}).encode(), end_stream=True) def _reply_to_ended_streams(connection: H2Connection, tls: ssl.SSLSocket, data: bytes, drain: bool) -> bool: @@ -407,28 +430,55 @@ def _serve_h2(context: ssl.SSLContext, raw: socket.socket, drain: bool) -> None: return +def _read_http1_request(tls: ssl.SSLSocket) -> bytes: + with tls.makefile("rb") as stream: + head: Final = tuple(iter(stream.readline, b"\r\n")) + length: Final = next( + (int(line.split(b":", 1)[1]) for line in head if line.lower().startswith(b"content-length:")), 0 + ) + return b"".join(head) + stream.read(length) + + +def _hang_up_after_request(context: ssl.SSLContext, raw: socket.socket, connection_number: int) -> None: + with contextlib.suppress(OSError), raw, context.wrap_socket(raw, server_side=True) as tls: + _ = _read_http1_request(tls) + + +def _drain_every_other_connection(context: ssl.SSLContext, raw: socket.socket, connection_number: int) -> None: + _serve_h2(context, raw, drain=connection_number % 2 == 1) + + +def _answer_every_connection(context: ssl.SSLContext, raw: socket.socket, connection_number: int) -> None: + _serve_h2(context, raw, drain=False) + + @contextmanager -def _goaway_peer(peer: TlsPeer) -> Generator[str]: +def _raw_tls_peer( + peer: TlsPeer, alpn: str, serve_connection: Callable[[ssl.SSLContext, socket.socket, int], 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(["h2"]) + context.set_alpn_protocols([alpn]) stop: Final = threading.Event() - connections: Final = itertools.count(1) + accepted: Final = queue.SimpleQueue[int]() with socket.create_server(("127.0.0.1", 0)) as listener: listener.settimeout(0.1) + def _accept_one() -> None: + raw: Final = listener.accept()[0] + accepted.put(accepted.qsize() + 1) + threading.Thread(target=serve_connection, args=(context, raw, accepted.qsize()), daemon=True).start() + def _accept() -> None: while not stop.is_set(): with contextlib.suppress(TimeoutError): - threading.Thread( - target=_serve_h2, args=(context, listener.accept()[0], next(connections) % 2 == 1), daemon=True - ).start() + _accept_one() acceptor: Final = threading.Thread(target=_accept, daemon=True) acceptor.start() try: - yield f"https://127.0.0.1:{listener.getsockname()[1]}" + yield f"https://127.0.0.1:{listener.getsockname()[1]}", accepted finally: stop.set() acceptor.join(timeout=5) @@ -439,6 +489,46 @@ def test_vertex_search_retries_a_connection_the_server_drains_with_goaway( ) -> None: _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) - with _goaway_peer(http2_tls_peer) as search_url: - assert _vertex_search_file_ids(http2_tls_peer, search_url=search_url) == ["after-goaway"] - assert asyncio.run(_avertex_search_file_ids(http2_tls_peer, search_url=search_url)) == ["after-goaway"] + with _raw_tls_peer(http2_tls_peer, "h2", _drain_every_other_connection) 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"] + + +def test_vertex_search_sends_once_when_an_http1_server_hangs_up( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + import litellm + + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + monkeypatch.setenv("LITELLM_HTTP2", "False") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + + with _raw_tls_peer(http2_tls_peer, "http/1.1", _hang_up_after_request) as (search_url, accepted): + with pytest.raises(litellm.InternalServerError, match="Server disconnected"): + _vertex_search_file_ids(http2_tls_peer, search_url=search_url) + assert accepted.qsize() == 1 + with pytest.raises(litellm.InternalServerError, match="Server disconnected"): + asyncio.run(_avertex_search_file_ids(http2_tls_peer, search_url=search_url)) + assert accepted.qsize() == 2 + + +@pytest.mark.parametrize("force_ipv4", [False, True]) +def test_vertex_search_keeps_its_http2_connection_across_an_idle_gap( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer, force_ipv4: bool +) -> None: + import litellm + from litellm.constants import AIOHTTP_KEEPALIVE_TIMEOUT + + idle_gap: Final = DEFAULT_LIMITS.keepalive_expiry + 1 + assert idle_gap < AIOHTTP_KEEPALIVE_TIMEOUT + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + monkeypatch.setattr(litellm, "force_ipv4", force_ipv4) + + async def _search_twice(search_url: str) -> list[str | None]: + first: Final = await _avertex_search_file_ids(http2_tls_peer, search_url=search_url) + await asyncio.sleep(idle_gap) + return [*first, *await _avertex_search_file_ids(http2_tls_peer, search_url=search_url)] + + 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