diff --git a/litellm/__init__.py b/litellm/__init__.py index 1784deeb797..c33fda86bf8 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -518,7 +518,7 @@ aiohttp_trust_env: bool = False # set to true to use HTTP_ Proxy settings disable_aiohttp_transport: bool = False # Set this to true to use httpx instead disable_aiohttp_trust_env: bool = False # When False, aiohttp will respect HTTP(S)_PROXY env vars force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. -http2: bool = False +http2: bool | None = None network_mock: bool = False # When True, use mock transport — no real network calls ####### STOP SEQUENCE LIMIT ####### diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 28f6348e4d9..513ffc342a7 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -124,6 +124,9 @@ class BaseVectorStoreConfig: def validate_create_vector_store(self) -> None: return None + def get_async_httpx_client_params(self) -> Mapping[str, object]: + return MappingProxyType({}) + def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]: return [] 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..0fc1d8c38b9 --- /dev/null +++ b/litellm/llms/custom_httpx/http2_connection_spread.py @@ -0,0 +1,105 @@ +from collections.abc import AsyncIterator, Callable +from typing import Final + +import httpx + +HTTP2_STREAMS_PER_CONNECTION: Final = 100 + + +class _Lane: + __slots__ = ("in_flight", "transport") + + def __init__(self, transport: httpx.AsyncBaseTransport) -> 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: + def __init__( + self, + new_transport: Callable[[httpx.Limits], httpx.AsyncBaseTransport], + limits: httpx.Limits, + streams_per_connection: int, + ) -> None: + self._new_transport: Final = new_transport + self._limits: Final = limits + self._streams_per_connection: Final = streams_per_connection + self.lanes: tuple[_Lane, ...] = () + + def acquire(self) -> _Lane: + 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_or_share() + lane.acquire() + return lane + + def _open_or_share(self) -> _Lane: + cap: Final = self._limits.max_connections + if cap is not None and len(self.lanes) * self._streams_per_connection >= cap: + return min(self.lanes, key=lambda lane: lane.in_flight) + return self._open() + + def _open(self) -> _Lane: + lane: Final = _Lane(self._new_transport(self._next_lane_limits())) + self.lanes = (*self.lanes, lane) + return lane + + def _next_lane_limits(self) -> httpx.Limits: + cap: Final = self._limits.max_connections + if cap is None: + return self._limits + return httpx.Limits( + max_connections=min(self._streams_per_connection, cap - len(self.lanes) * self._streams_per_connection), + max_keepalive_connections=self._limits.max_keepalive_connections, + keepalive_expiry=self._limits.keepalive_expiry, + ) + + +class _AsyncReleasingStream(httpx.AsyncByteStream): + def __init__(self, stream: httpx.AsyncByteStream, release: Callable[[], None]) -> None: + 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() + + +def _with_stream(response: httpx.Response, stream: 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.Limits], httpx.AsyncBaseTransport], + limits: httpx.Limits, + streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION, + ) -> None: + self._lanes: Final = _Lanes(new_transport, limits, streams_per_connection) + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + lane: Final = self._lanes.acquire() + 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() diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index cffeb44f12c..e00902d9047 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 @@ -42,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 from litellm.types.llms.custom_http import * if TYPE_CHECKING: @@ -72,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]]] @@ -83,6 +85,31 @@ def http2_enabled() -> bool: return litellm.http2 is True or str_to_bool(os.getenv("LITELLM_HTTP2", "False")) is True +def http2_disabled() -> bool: + from litellm.secret_managers.main import str_to_bool + + return litellm.http2 is False or str_to_bool(os.getenv("LITELLM_HTTP2")) is False + + +def _resolve_http2(provider_default: bool, ssl_verify: VerifyTypes | None) -> bool: + if http2_enabled(): + return True + return provider_default and not http2_disabled() and not isinstance(ssl_verify, ssl.SSLContext) + + +def _aiohttp_env_proxies_disabled() -> bool: + from litellm.secret_managers.main import str_to_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]: @@ -279,9 +306,7 @@ def _prepare_request_data_and_content( # Cache for SSL contexts to avoid creating duplicate contexts with the same configuration -# Key: tuple of (cafile, ssl_security_level, ssl_ecdh_curve) -# Value: ssl.SSLContext -_ssl_context_cache: Final[dict[tuple[str | None, str | None, str | None], ssl.SSLContext]] = {} +_ssl_context_cache: Final[dict[tuple[str | None, str | None, str | None, bool], ssl.SSLContext]] = {} def _create_ssl_context( @@ -374,6 +399,7 @@ def get_ssl_verify( def get_ssl_configuration( ssl_verify: VerifyTypes | None = None, + http2: bool = False, ) -> bool | str | ssl.SSLContext: """ Unified SSL configuration function that handles ssl_context and ssl_verify logic. @@ -422,7 +448,7 @@ def get_ssl_configuration( if ssl_verify is not False: # Create cache key from configuration parameters - cache_key: Final = (cafile, ssl_security_level, ssl_ecdh_curve) + cache_key: Final = (cafile, ssl_security_level, ssl_ecdh_curve, http2) # Check if we have a cached SSL context for this configuration if cache_key not in _ssl_context_cache: @@ -633,6 +659,7 @@ class AsyncHTTPHandler: shared_session: Optional["ClientSession"] = None, transport: httpx.AsyncBaseTransport | None = None, follow_redirects: bool = True, + http2: bool = False, ): self.timeout = timeout self.event_hooks = event_hooks @@ -640,6 +667,7 @@ class AsyncHTTPHandler: self.shared_session = shared_session self.transport = transport self.follow_redirects = follow_redirects + self.http2 = http2 self._owns_client = True self._client = self.create_client( timeout=timeout, @@ -682,8 +710,11 @@ class AsyncHTTPHandler: follow_redirects=self.follow_redirects, 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) + ssl_config: Final = get_ssl_configuration(ssl_verify, http2=use_http2) # An SSL certificate used by the requested host to authenticate the client. # /path/to/client.pem @@ -697,14 +728,23 @@ class AsyncHTTPHandler: ssl_context=ssl_config if isinstance(ssl_config, ssl.SSLContext) else None, ssl_verify=ssl_config if isinstance(ssl_config, bool) else None, 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 (replaces_aiohttp and _aiohttp_env_proxies_disabled()) return httpx.AsyncClient( transport=transport, - mounts=AsyncHTTPHandler.create_httpx_proxy_mounts(transport, verify=ssl_config, cert=cert), + 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, @@ -712,7 +752,8 @@ class AsyncHTTPHandler: headers=default_headers, cookies=blocked_cookie_jar(), follow_redirects=self.follow_redirects, - http2=http2_enabled(), + http2=use_http2, + trust_env=trust_env, ) async def close(self): @@ -1202,7 +1243,10 @@ class AsyncHTTPHandler: ssl_context: ssl.SSLContext | None = None, ssl_verify: bool | None = None, shared_session: Optional["ClientSession"] = None, - ) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | None: + http2: bool = False, + cert: CertTypes | None = None, + limits: httpx.Limits = DEFAULT_LIMITS, + ) -> 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" @@ -1220,7 +1264,7 @@ class AsyncHTTPHandler: ######################################################### # AIOHTTP TRANSPORT is off by default ######################################################### - if AsyncHTTPHandler.should_use_aiohttp_transport(): + if not http2 and AsyncHTTPHandler.should_use_aiohttp_transport(): return AsyncHTTPHandler.create_aiohttp_transport( ssl_context=ssl_context, ssl_verify=ssl_verify, @@ -1230,7 +1274,12 @@ class AsyncHTTPHandler: ######################################################### # HTTPX TRANSPORT is used when aiohttp is not installed ######################################################### - return AsyncHTTPHandler._create_httpx_transport() + return AsyncHTTPHandler._create_httpx_transport( + 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 @@ -1377,28 +1426,52 @@ class AsyncHTTPHandler: _create_aiohttp_transport = create_aiohttp_transport @staticmethod - def _create_httpx_transport() -> AsyncHTTPTransport | None: + def _create_httpx_transport( + http2: bool = False, + verify: ssl.SSLContext | bool = True, + cert: CertTypes | None = None, + limits: httpx.Limits = DEFAULT_LIMITS, + ) -> 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 """ + local_address: Final = _IPV4_LOCAL_ADDRESS if litellm.force_ipv4 else None + if http2 or http2_enabled(): + return AsyncHTTP2ConnectionSpread( + lambda lane_limits: AsyncHTTPTransport( + local_address=local_address, verify=verify, cert=cert, http2=True, limits=lane_limits + ), + limits, + ) if litellm.force_ipv4: - return AsyncHTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, http2=http2_enabled()) - else: - return None + 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, - ) -> Mapping[str, AsyncHTTPTransport | None] | None: - if not isinstance(transport, AsyncHTTPTransport): + http2: bool = False, + limits: httpx.Limits = DEFAULT_LIMITS, + ) -> Mapping[str, httpx.AsyncBaseTransport | None] | None: + if not isinstance(transport, (AsyncHTTPTransport, AsyncHTTP2ConnectionSpread)): return None + if http2 or http2_enabled(): + return _environment_proxy_mounts( + lambda proxy_url: AsyncHTTP2ConnectionSpread( + lambda lane_limits: AsyncHTTPTransport( + proxy=proxy_url, verify=verify, cert=cert, http2=True, limits=lane_limits + ), + limits, + ) + ) return _environment_proxy_mounts( - lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2_enabled()) + lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, limits=limits) ) _create_httpx_proxy_mounts = create_httpx_proxy_mounts @@ -1424,8 +1497,9 @@ class HTTPHandler: self._client = self.create_client() if client is None else client def create_client(self) -> httpx.Client: + use_http2: Final = http2_enabled() # Get unified SSL configuration - ssl_config: Final = get_ssl_configuration(self.ssl_verify) + ssl_config: Final = get_ssl_configuration(self.ssl_verify, http2=use_http2) # An SSL certificate used by the requested host to authenticate the client. # /path/to/client.pem @@ -1434,17 +1508,19 @@ 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(), - mounts=self._create_sync_proxy_mounts(verify=ssl_config, cert=cert), + 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, headers=default_headers, cookies=blocked_cookie_jar(), follow_redirects=self.follow_redirects, - http2=http2_enabled(), + http2=use_http2, ) @property @@ -1719,27 +1795,29 @@ class HTTPHandler: except Exception: pass - def _create_sync_transport(self) -> HTTPTransport | None: + @staticmethod + def _create_sync_transport(verify: VerifyTypes, cert: CertTypes | None, http2: bool) -> HTTPTransport | None: """ Create an HTTP transport with IPv4 only if litellm.force_ipv4 is True. - Otherwise, return None. + Otherwise, return litellm.sync_transport or None. Some users have seen httpx ConnectionError when using ipv6 - forcing ipv4 resolves the issue for them """ if litellm.force_ipv4: - return HTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, http2=http2_enabled()) - else: - return getattr(litellm, "sync_transport", None) + return HTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, verify=verify, cert=cert, http2=http2) + return getattr(litellm, "sync_transport", None) @staticmethod def _create_sync_proxy_mounts( + transport: httpx.BaseTransport | None, verify: VerifyTypes, cert: CertTypes | None, - ) -> Mapping[str, HTTPTransport | None] | None: - if not litellm.force_ipv4: + http2: bool, + ) -> 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_enabled()) + lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2) ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a7ca0d5f49e..32bc726aa4d 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 @@ -342,6 +343,34 @@ class _SignedBodyKwargs(TypedDict, total=False): json: ReadOnly[dict[str, object]] +_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 as error: + if not _closed_by_http2_goaway(error): + raise + verbose_logger.warning(_GOAWAY_RETRY_MESSAGE) + return await send() + + +def _send_with_goaway_retry(send: Callable[[], httpx.Response]) -> httpx.Response: + try: + return send() + except httpx.RemoteProtocolError as error: + if not _closed_by_http2_goaway(error): + raise + verbose_logger.warning(_GOAWAY_RETRY_MESSAGE) + return send() + + def _signed_body_kwargs(*, signed_body: bytes | None, data: dict[str, object]) -> _SignedBodyKwargs: if signed_body is not None: return {"data": signed_body} @@ -10458,7 +10487,10 @@ class BaseLLMHTTPHandler: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( llm_provider=litellm.LlmProviders(custom_llm_provider), - params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + params={ + "ssl_verify": litellm_params.get("ssl_verify", None), + **vector_store_provider_config.get_async_httpx_client_params(), + }, ) else: async_httpx_client = client @@ -10530,11 +10562,8 @@ class BaseLLMHTTPHandler: request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body try: - response: Final = await async_httpx_client.post( - url=url, - headers=headers, - data=request_data, - timeout=timeout, + response: Final = await _asend_with_goaway_retry( + partial(async_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout) ) except httpx.TimeoutException: raise vector_store_provider_config.get_error_class( @@ -10667,11 +10696,8 @@ class BaseLLMHTTPHandler: request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body try: - response: Final = sync_httpx_client.post( - url=url, - headers=headers, - data=request_data, - timeout=timeout, + response: Final = _send_with_goaway_retry( + partial(sync_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout) ) except httpx.TimeoutException: raise vector_store_provider_config.get_error_class( diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 5bff43828d6..351436c813c 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -418,12 +418,7 @@ class ContextCachingEndpoints(VertexBase): headers.update(extra_headers) if client is None or not isinstance(client, HTTPHandler): - _params: Final = {} - if timeout is not None: - if isinstance(timeout, float) or isinstance(timeout, int): - timeout = httpx.Timeout(timeout) - _params["timeout"] = timeout - client = HTTPHandler(**_params) + client = HTTPHandler(timeout=httpx.Timeout(timeout) if isinstance(timeout, float | int) else timeout) else: client = client diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index aa2cc8575f9..dfd34c706fd 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -234,6 +234,9 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): def __init__(self): super().__init__() + def get_async_httpx_client_params(self) -> Mapping[str, object]: + return MappingProxyType({"http2": True}) + @staticmethod def get_supported_extra_body_fields(is_engine: bool = False) -> frozenset[str]: """ @@ -458,3 +461,6 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): input_cost_per_query: Final = model_info.get("input_cost_per_query") or 0.0 return input_cost_per_query, 0.0 + + +vertex_search_api_vector_store_config: Final = VertexSearchAPIVectorStoreConfig() diff --git a/litellm/utils.py b/litellm/utils.py index 7de7ce6188d..11c1f73d65f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10234,10 +10234,10 @@ class ProviderConfigManager: return VertexVectorStoreConfig() elif api_type == "search_api": from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( - VertexSearchAPIVectorStoreConfig, + vertex_search_api_vector_store_config, ) - return VertexSearchAPIVectorStoreConfig() + return vertex_search_api_vector_store_config elif litellm.LlmProviders.BEDROCK == provider: from litellm.llms.bedrock.vector_stores.transformation import ( BedrockVectorStoreConfig, diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index 4fcf590f0eb..b8f2ae9b76a 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/tests/integration/sdk/test_http2_wire.py @@ -1,30 +1,57 @@ from __future__ import annotations import asyncio +import contextlib import datetime import ipaddress import json +import queue import socket +import ssl import threading import time -from collections.abc import Iterator +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 import pytest from cryptography import x509 from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import rsa 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 ) from hypercorn.config import Config from hypercorn.typing import ASGIReceiveCallable, ASGISendCallable, HTTPResponseBodyEvent, HTTPResponseStartEvent, Scope 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 +_HELD_ARRIVALS: Final = queue.SimpleQueue[int]() +_TOKEN_REQUESTS: Final = queue.SimpleQueue[str]() +_COUNTED_CLIENT_PORTS: Final = queue.SimpleQueue[int]() +VERTEX_SERVING_CONFIG_PATH: Final = ( + "/v1/projects/p/locations/global/collections/default_collection/engines/e/servingConfigs/default_search" +) + + +@dataclass(frozen=True, slots=True) +class TlsPeer: + url: str + cert_file: Path + key_file: Path @dataclass(frozen=True, slots=True) @@ -76,6 +103,22 @@ 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"] + if scope["path"] == "/hold:search": + client: Final = scope["client"] + await _send_json(send, {"results": [{"id": await _hold_until_all_arrive(client[1] if client else 0)}]}) + return + if scope["path"].endswith("/_oauth/token"): + _TOKEN_REQUESTS.put(scope["path"]) + await _send_json(send, {"access_token": "scripted", "expires_in": 3600, "token_type": "Bearer"}) + return + if scope["path"] == "/counted:search": + counted_client: Final = scope["client"] + _COUNTED_CLIENT_PORTS.put(counted_client[1] if counted_client else 0) + await _send_json(send, {"results": [{"id": "counted"}]}) + return + if scope["path"].endswith(":search"): + await _send_json(send, {"results": [{"id": f"http/{version}"}]}) + return if scope["path"] == "/stream": await send( HTTPResponseStartEvent( @@ -90,18 +133,26 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab ) await send(HTTPResponseBodyEvent(type="http.response.body", body=b"", more_body=False)) return + await _send_json(send, {"http_version": version}) + + +async def _hold_until_all_arrive(client_port: int) -> str: + _HELD_ARRIVALS.put(client_port) + deadline: Final = time.monotonic() + HOLD_SECONDS + while _HELD_ARRIVALS.qsize() < HELD_SEARCHES and time.monotonic() < deadline: + await asyncio.sleep(0.01) + return "held-together" if _HELD_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")]) ) - await send( - HTTPResponseBodyEvent( - type="http.response.body", body=json.dumps({"http_version": version}).encode(), more_body=False - ) - ) + await send(HTTPResponseBodyEvent(type="http.response.body", body=json.dumps(body).encode(), more_body=False)) @pytest.fixture(scope="module") -def http2_tls_peer(tmp_path_factory: pytest.TempPathFactory) -> Iterator[str]: +def http2_tls_peer(tmp_path_factory: pytest.TempPathFactory) -> Iterator[TlsPeer]: cert_file, key_file = _write_self_signed_cert(tmp_path_factory.mktemp("h2certs")) port: Final = allocate_port() shutdown: Final = threading.Event() @@ -126,7 +177,7 @@ def http2_tls_peer(tmp_path_factory: pytest.TempPathFactory) -> Iterator[str]: time.sleep(0.05) else: pytest.fail("hypercorn peer did not start") - yield f"https://127.0.0.1:{port}" + yield TlsPeer(url=f"https://127.0.0.1:{port}", cert_file=cert_file, key_file=key_file) shutdown.set() thread.join(timeout=10) @@ -185,15 +236,437 @@ def _assert_negotiated(observed: Observed, enabled: bool) -> None: @pytest.mark.covers("other.sdk_wire.http2.async_handler_negotiates_h2_only_when_enabled") -def test_async_handler_negotiates_http2_only_when_enabled(monkeypatch: pytest.MonkeyPatch, http2_tls_peer: str) -> None: +def test_async_handler_negotiates_http2_only_when_enabled( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") for enabled in (False, True): _set_http2(monkeypatch, enabled) - _assert_negotiated(_async_exchange(http2_tls_peer), enabled) + _assert_negotiated(_async_exchange(http2_tls_peer.url), enabled) @pytest.mark.covers("other.sdk_wire.http2.sync_handler_negotiates_h2_only_when_enabled") -def test_sync_handler_negotiates_http2_only_when_enabled(monkeypatch: pytest.MonkeyPatch, http2_tls_peer: str) -> None: +def test_sync_handler_negotiates_http2_only_when_enabled( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: for enabled in (False, True): _set_http2(monkeypatch, enabled) - _assert_negotiated(_sync_exchange(http2_tls_peer), enabled) + _assert_negotiated(_sync_exchange(http2_tls_peer.url), enabled) + + +def _trust_peer_without_global_http2(monkeypatch: pytest.MonkeyPatch, peer: TlsPeer) -> None: + import litellm + + _set_http2(monkeypatch, False) + monkeypatch.delenv("DISABLE_AIOHTTP_TRANSPORT", raising=False) + monkeypatch.setenv("SSL_VERIFY", str(peer.cert_file)) + monkeypatch.setenv("REQUESTS_CA_BUNDLE", str(peer.cert_file)) + litellm.in_memory_llm_clients_cache.flush_cache() + + +def _vertex_search_file_ids( + peer: TlsPeer, search_url: str | None = None, ssl_verify: ssl.SSLContext | None = None +) -> list[str | None]: + import litellm + + response: Final = litellm.vector_stores.search( + vector_store_id="default_search", + query="ping", + custom_llm_provider="vertex_ai/search_api", + api_base=f"{search_url or peer.url}{VERTEX_SERVING_CONFIG_PATH}", + vertex_credentials=service_account_json("p", peer.url), + vertex_project="p", + ssl_verify=ssl_verify, + ) + assert isinstance(response, dict) + return [result.get("file_id") for result in response.get("data") or ()] + + +async def _avertex_search_file_ids( + peer: TlsPeer, search_url: str | None = None, ssl_verify: ssl.SSLContext | None = None +) -> list[str | None]: + import litellm + + response: Final = await litellm.vector_stores.asearch( + vector_store_id="default_search", + query="ping", + custom_llm_provider="vertex_ai/search_api", + api_base=f"{search_url or peer.url}{VERTEX_SERVING_CONFIG_PATH}", + vertex_credentials=service_account_json("p", peer.url), + vertex_project="p", + ssl_verify=ssl_verify, + ) + return [result.get("file_id") for result in response.get("data") or ()] + + +def test_sync_vertex_search_stays_on_http1_without_global_flag( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"] + + +def test_async_vertex_search_negotiates_http2_without_global_flag( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/2"] + + +def test_default_async_client_stays_http1_after_vertex_search_used_http2( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + + async def _run() -> tuple[list[str | None], str]: + file_ids: Final = await _avertex_search_file_ids(http2_tls_peer) + handler: Final = AsyncHTTPHandler(ssl_verify=str(http2_tls_peer.cert_file)) + try: + response: Final = await handler.post(f"{http2_tls_peer.url}/echo", json={"ping": "pong"}) + return file_ids, response.json()["http_version"] + finally: + await handler.close() + + assert asyncio.run(_run()) == (["http/2"], "1.1") + + +def test_vertex_search_fetches_one_token_for_many_searches( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + import litellm + + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + credentials: Final = service_account_json("p", f"{http2_tls_peer.url}/token-reuse") + search: Final = MappingProxyType( + { + "vector_store_id": "default_search", + "query": "ping", + "custom_llm_provider": "vertex_ai/search_api", + "api_base": f"{http2_tls_peer.url}{VERTEX_SERVING_CONFIG_PATH}", + "vertex_credentials": credentials, + "vertex_project": "p", + } + ) + + async def _searches() -> None: + await litellm.vector_stores.asearch(**search) + await asyncio.gather(*(litellm.vector_stores.asearch(**search) for _ in range(5))) + + asyncio.run(_searches()) + litellm.vector_stores.search(**search) + + token_paths: Final = tuple(iter(lambda: None if _TOKEN_REQUESTS.empty() else _TOKEN_REQUESTS.get(), None)) + assert token_paths.count("/token-reuse/_oauth/token") == 1 + + +def test_vertex_search_over_forced_ipv4_trusts_the_configured_ca( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + import litellm + + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + monkeypatch.setattr(litellm, "force_ipv4", True) + + assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"] + assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/2"] + + +def test_explicit_http2_off_keeps_vertex_search_on_http1( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + import litellm + + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + monkeypatch.setenv("LITELLM_HTTP2", "False") + assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"] + assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/1.1"] + + monkeypatch.delenv("LITELLM_HTTP2") + monkeypatch.setattr(litellm, "http2", False) + litellm.in_memory_llm_clients_cache.flush_cache() + assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"] + assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/1.1"] + + +def _alpn_offered_by(context: ssl.SSLContext, peer: TlsPeer) -> str | None: + host, port = peer.url.removeprefix("https://").split(":") + with ( + socket.create_connection((host, int(port)), timeout=5) as raw, + context.wrap_socket(raw, server_hostname=host) as tls, + ): + return tls.selected_alpn_protocol() + + +def test_vertex_search_leaves_a_caller_ssl_context_on_http1( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + sync_context: Final = ssl.create_default_context(cafile=str(http2_tls_peer.cert_file)) + async_context: Final = ssl.create_default_context(cafile=str(http2_tls_peer.cert_file)) + + assert _vertex_search_file_ids(http2_tls_peer, ssl_verify=sync_context) == ["http/1.1"] + assert asyncio.run(_avertex_search_file_ids(http2_tls_peer, ssl_verify=async_context)) == ["http/1.1"] + assert _alpn_offered_by(sync_context, http2_tls_peer) != "h2" + 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 + + _set_http2(monkeypatch, False) + _route_https_through_a_dead_proxy(monkeypatch) + + for force_ipv4 in (False, True): + monkeypatch.setattr(litellm, "force_ipv4", force_ipv4) + 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: + if drain: + 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": "answered"}]}).encode(), end_stream=True) + + +def _reply_to_ended_streams(connection: H2Connection, tls: ssl.SSLSocket, data: bytes, drain: bool) -> bool: + ended: Final = tuple(event.stream_id for event in connection.receive_data(data) if isinstance(event, StreamEnded)) + for stream_id in ended: + _answer_or_drain(connection, stream_id, drain) + tls.sendall(connection.data_to_send()) + return bool(ended) + + +def _serve_h2(context: ssl.SSLContext, raw: socket.socket, drain: bool) -> None: + with contextlib.suppress(OSError), raw, context.wrap_socket(raw, server_side=True) as tls: + connection: Final = H2Connection(config=H2Configuration(client_side=False)) + connection.initiate_connection() + tls.sendall(connection.data_to_send()) + while data := tls.recv(65535): + if _reply_to_ended_streams(connection, tls, data, drain) and drain: + 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 _raw_tls_peer( + 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]() + + 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): + _accept_one() + + acceptor: Final = threading.Thread(target=_accept, daemon=True) + acceptor.start() + try: + yield f"https://127.0.0.1:{listener.getsockname()[1]}", accepted + finally: + stop.set() + acceptor.join(timeout=5) + + +def test_vertex_search_retries_a_connection_the_server_drains_with_goaway( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + + with _raw_tls_peer(http2_tls_peer, "h2", _drain_every_other_connection) as (search_url, _): + assert asyncio.run(_avertex_search_file_ids(http2_tls_peer, search_url=search_url)) == ["answered"] + _set_http2(monkeypatch, True) + assert _vertex_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 + + +def test_async_http2_handler_opens_another_connection_once_the_streams_run_out( + monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer +) -> None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) + + async def _burst() -> list[str]: + handler: Final = AsyncHTTPHandler(ssl_verify=str(http2_tls_peer.cert_file), http2=True) + try: + responses: Final = await asyncio.gather( + *(handler.post(f"{http2_tls_peer.url}/hold:search", json={}) for _ in range(HELD_SEARCHES)) + ) + return [response.json()["results"][0]["id"] for response in responses] + finally: + await handler.close() + + search_ids: Final = asyncio.run(_burst()) + client_ports: Final = frozenset(_HELD_ARRIVALS.get_nowait() for _ in range(_HELD_ARRIVALS.qsize())) + + assert search_ids == ["held-together"] * HELD_SEARCHES + assert len(client_ports) == 2 + + +def test_http2_transport_opens_no_more_connections_than_its_cap(http2_tls_peer: TlsPeer) -> None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + async def _burst() -> frozenset[str]: + transport: Final = AsyncHTTPHandler.create_async_transport( + ssl_context=ssl.create_default_context(cafile=str(http2_tls_peer.cert_file)), + http2=True, + limits=httpx.Limits(max_connections=1), + ) + async with httpx.AsyncClient(transport=transport) as client: + responses: Final = await asyncio.gather( + *(client.post(f"{http2_tls_peer.url}/counted:search", json={}) for _ in range(HELD_SEARCHES)) + ) + return frozenset(response.json()["results"][0]["id"] for response in responses) + + assert asyncio.run(_burst()) == {"counted"} + client_ports: Final = frozenset( + iter(lambda: None if _COUNTED_CLIENT_PORTS.empty() else _COUNTED_CLIENT_PORTS.get(), None) + ) + assert len(client_ports) == 1 + + +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 asyncio.run(_avertex_search_file_ids(http2_tls_peer, search_url=search_url)) == ["answered"] + _set_http2(monkeypatch, True) + assert _vertex_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..a83c013c6a6 --- /dev/null +++ b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py @@ -0,0 +1,186 @@ +import asyncio +from collections.abc import Callable +from typing import Final + +import httpx +import pytest + +from litellm.llms.custom_httpx.http2_connection_spread import ( + HTTP2_STREAMS_PER_CONNECTION, + AsyncHTTP2ConnectionSpread, +) + +URL: Final = "https://search.example/v1/search" +BURST: Final = HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2 +UNCAPPED: Final = httpx.Limits(max_connections=None) + + +class _GatedAsyncConnection(httpx.AsyncBaseTransport): + 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, _limits: httpx.Limits) -> _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, _limits: httpx.Limits) -> _RefusedConnection: + self.opened += 1 + return _RefusedConnection() + + +class _HTTP1Connections: + def __init__(self, gate: asyncio.Event) -> None: + self._gate: Final = gate + self.requested = 0 + self.open = 0 + + def __call__(self, limits: httpx.Limits) -> "_HTTP1Pool": + assert limits.max_connections is not None + return _HTTP1Pool(self, asyncio.Semaphore(limits.max_connections)) + + async def serve(self, pool_slots: asyncio.Semaphore) -> None: + self.requested += 1 + async with pool_slots: + self.open += 1 + await self._gate.wait() + self.open -= 1 + + +class _HTTP1Pool(httpx.AsyncBaseTransport): + def __init__(self, connections: _HTTP1Connections, slots: asyncio.Semaphore) -> None: + self._connections: Final = connections + self._slots: Final = slots + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + await self._connections.serve(self._slots) + return httpx.Response(200, stream=httpx.ByteStream(b"ok")) + + +async def _until(condition: Callable[[], bool]) -> None: + for _ in range(1000): + if condition(): + return + await asyncio.sleep(0) + pytest.fail("condition never held") + + +@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, UNCAPPED)) as client: + pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(BURST)) + await _until(lambda: connections.in_flight() == BURST) + assert [connection.peak for connection in connections.opened] == [ + 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_shares_the_least_busy_connection_once_the_connection_cap_is_reached(): + gate: Final = asyncio.Event() + connections: Final = _AsyncConnections(gate) + burst: Final = 2 * HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2 + two_connections_of_streams: Final = httpx.Limits(max_connections=2 * HTTP2_STREAMS_PER_CONNECTION) + async with httpx.AsyncClient( + transport=AsyncHTTP2ConnectionSpread(connections, two_connections_of_streams) + ) as client: + pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(burst)) + await _until(lambda: connections.in_flight() == burst) + assert [connection.peak for connection in connections.opened] == [burst // 2, burst // 2] + + gate.set() + assert {response.status_code for response in await asyncio.gather(*pending)} == {200} + + +@pytest.mark.asyncio +async def test_async_spread_keeps_http1_connections_within_the_connection_cap(): + gate: Final = asyncio.Event() + connections: Final = _HTTP1Connections(gate) + cap: Final = HTTP2_STREAMS_PER_CONNECTION + HTTP2_STREAMS_PER_CONNECTION // 2 + burst: Final = 2 * cap + async with httpx.AsyncClient( + transport=AsyncHTTP2ConnectionSpread(connections, httpx.Limits(max_connections=cap)) + ) as client: + pending: Final = tuple(asyncio.create_task(client.get(URL)) for _ in range(burst)) + await _until(lambda: connections.requested == burst) + assert connections.open == cap + + gate.set() + assert {response.status_code for response in await asyncio.gather(*pending)} == {200} + + +@pytest.mark.asyncio +async def test_async_spread_holds_a_stream_until_its_response_is_closed(): + gate: Final = asyncio.Event() + gate.set() + connections: Final = _AsyncConnections(gate) + async with httpx.AsyncClient(transport=AsyncHTTP2ConnectionSpread(connections, UNCAPPED)) 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, UNCAPPED)) as client: + for _ in range(BURST): + with pytest.raises(httpx.ConnectError): + await client.get(URL) + assert connections.opened == 1 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 diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 63cc45b29df..cac7caeb6cd 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -10,10 +10,12 @@ from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, Mock, patch +import httpcore import httpx import pytest import respx from botocore.credentials import RefreshableCredentials +from h2.events import ConnectionTerminated import litellm from litellm._logging import verbose_logger @@ -5542,3 +5544,54 @@ async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_ sent: Final = received[0] assert "safeguards" not in json.loads(sent.content) assert "dangerous-tool-use-2026-09-03" not in sent.headers.get("anthropic-beta", ""), dict(sent.headers) + + +_EMPTY_SEARCH_PAGE: Final = { + "object": "vector_store.search_results.page", + "search_query": "q", + "data": [], + "has_more": False, + "next_page": None, +} + + +def _close_with_goaway(request: httpx.Request) -> httpx.Response: + raise httpx.RemoteProtocolError("GOAWAY", request=request) from httpcore.RemoteProtocolError(ConnectionTerminated()) + + +def _hang_up(request: httpx.Request) -> httpx.Response: + raise httpx.RemoteProtocolError("Server disconnected without sending a response.", request=request) + + +def _answer_search(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=_EMPTY_SEARCH_PAGE) + + +async def _search_through(transport: httpx.MockTransport, is_async: bool) -> object: + search_kwargs: Final = {"vector_store_id": "vs_1", "query": "q", "custom_llm_provider": "openai", "api_key": "sk-test"} + if not is_async: + with httpx.Client(transport=transport) as sync_client: + return litellm.vector_stores.search(**search_kwargs, client=HTTPHandler(client=sync_client)) + async with httpx.AsyncClient(transport=transport) as async_client: + client: Final = AsyncHTTPHandler() + await client.close() + client.client = async_client + return await litellm.vector_stores.asearch(**search_kwargs, client=client) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True)) +async def test_vector_store_search_sends_again_once_after_an_http2_goaway(is_async: bool) -> None: + sends: Final = iter((_close_with_goaway, _answer_search)) + response: Final = await _search_through(httpx.MockTransport(lambda request: next(sends)(request)), is_async) + assert response["search_query"] == "q" + assert next(sends, None) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True)) +async def test_vector_store_search_sends_once_when_the_connection_drops_without_a_goaway(is_async: bool) -> None: + sends: Final = iter((_hang_up, _answer_search)) + with pytest.raises(litellm.InternalServerError, match="Server disconnected"): + await _search_through(httpx.MockTransport(lambda request: next(sends)(request)), is_async) + assert next(sends) is _answer_search