From 5e9a1820da4ce3c6f96f87d2fe2364d27d1e3da7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:55:52 -0700 Subject: [PATCH] fix(http): keep sync vertex search on HTTP/1.1, since httpcore's sync HTTP/2 is not thread-safe httpcore 1.0.9 hands two threads the same HTTP/2 stream id, so threaded sync searches failed with protocol errors. Only the async client, which the proxy uses, now defaults to HTTP/2, and the sync connection spread is removed. --- .../base_llm/vector_store/transformation.py | 2 +- .../custom_httpx/http2_connection_spread.py | 74 ++++-------------- litellm/llms/custom_httpx/http_handler.py | 37 +++------ litellm/llms/custom_httpx/llm_http_handler.py | 9 +-- .../search_api/transformation.py | 2 +- tests/integration/sdk/test_http2_wire.py | 77 ++++++++----------- .../test_http2_connection_spread.py | 67 ---------------- 7 files changed, 57 insertions(+), 211 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index d15c30a8eb5..513ffc342a7 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -124,7 +124,7 @@ class BaseVectorStoreConfig: def validate_create_vector_store(self) -> None: return None - def get_httpx_client_params(self) -> Mapping[str, object]: + 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]: diff --git a/litellm/llms/custom_httpx/http2_connection_spread.py b/litellm/llms/custom_httpx/http2_connection_spread.py index cabdd5011cd..1fd7fb2c4a8 100644 --- a/litellm/llms/custom_httpx/http2_connection_spread.py +++ b/litellm/llms/custom_httpx/http2_connection_spread.py @@ -1,18 +1,15 @@ -import threading -from collections.abc import AsyncIterator, Callable, Iterator -from typing import Final, Generic, TypeVar +from collections.abc import AsyncIterator, Callable +from typing import Final import httpx HTTP2_STREAMS_PER_CONNECTION: Final = 100 -_TransportT: Final = TypeVar("_TransportT", httpx.BaseTransport, httpx.AsyncBaseTransport) - -class _Lane(Generic[_TransportT]): +class _Lane: __slots__ = ("in_flight", "transport") - def __init__(self, transport: _TransportT) -> None: + def __init__(self, transport: httpx.AsyncBaseTransport) -> None: self.transport: Final = transport self.in_flight = 0 @@ -23,27 +20,30 @@ class _Lane(Generic[_TransportT]): self.in_flight -= 1 -class _Lanes(Generic[_TransportT]): +class _Lanes: def __init__( - self, new_transport: Callable[[], _TransportT], streams_per_connection: int, max_connections: int | None + self, + new_transport: Callable[[], httpx.AsyncBaseTransport], + streams_per_connection: int, + max_connections: int | None, ) -> None: self._new_transport: Final = new_transport self._streams_per_connection: Final = streams_per_connection self._max_connections: Final = max_connections - self.lanes: tuple[_Lane[_TransportT], ...] = () + self.lanes: tuple[_Lane, ...] = () - def acquire(self) -> _Lane[_TransportT]: + 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[_TransportT]: + def _open_or_share(self) -> _Lane: if self._max_connections is not None and len(self.lanes) >= self._max_connections: return min(self.lanes, key=lambda lane: lane.in_flight) return self._open() - def _open(self) -> _Lane[_TransportT]: + def _open(self) -> _Lane: lane: Final = _Lane(self._new_transport()) self.lanes = (*self.lanes, lane) return lane @@ -63,20 +63,7 @@ class _AsyncReleasingStream(httpx.AsyncByteStream): 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: +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) @@ -105,36 +92,3 @@ class AsyncHTTP2ConnectionSpread(httpx.AsyncBaseTransport): 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, - max_connections: int | None = None, - ) -> None: - self._lanes: Final = _Lanes(new_transport, streams_per_connection, max_connections) - self._lock: Final = threading.Lock() - - def _release(self, lane: _Lane[httpx.BaseTransport]) -> None: - with self._lock: - lane.release() - - def handle_request(self, request: httpx.Request) -> httpx.Response: - with self._lock: - lane: Final = self._lanes.acquire() - try: - response: Final = lane.transport.handle_request(request) - except BaseException: - self._release(lane) - raise - stream: Final = response.stream - if not isinstance(stream, httpx.SyncByteStream): - self._release(lane) - return response - return _with_stream(response, _ReleasingStream(stream, lambda: self._release(lane))) - - def close(self) -> None: - for lane in self._lanes.lanes: - lane.transport.close() diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index e83d9e66758..2ce6d1ea032 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -43,7 +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.llms.custom_httpx.http2_connection_spread import AsyncHTTP2ConnectionSpread from litellm.types.llms.custom_http import * if TYPE_CHECKING: @@ -1485,19 +1485,17 @@ class HTTPHandler: disable_default_headers: bool | None = False, # arize phoenix returns different API responses when user agent header in request follow_redirects: bool = True, - http2: bool = False, ): self.timeout = timeout self.ssl_verify = ssl_verify self.disable_default_headers = disable_default_headers self.follow_redirects = follow_redirects - self.http2 = http2 self._owns_client = client is None self._heal_lock = threading.Lock() self._client = self.create_client() if client is None else client def create_client(self) -> httpx.Client: - use_http2: Final = _resolve_http2(self.http2, self.ssl_verify) + use_http2: Final = http2_enabled() # Get unified SSL configuration ssl_config: Final = get_ssl_configuration(self.ssl_verify, http2=use_http2) @@ -1796,28 +1794,16 @@ class HTTPHandler: pass @staticmethod - def _create_sync_transport( - verify: VerifyTypes, cert: CertTypes | None, http2: bool - ) -> HTTPTransport | HTTP2ConnectionSpread | None: + 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, and one that spreads requests over - as many HTTP/2 connections as the stream limit needs if HTTP/2 is on. + Create an HTTP transport with IPv4 only if litellm.force_ipv4 is True. 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), - max_connections=DEFAULT_LIMITS.max_connections, - ) if litellm.force_ipv4: - return HTTPTransport(local_address=local_address, verify=verify, cert=cert) - return 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( @@ -1828,14 +1814,9 @@ class HTTPHandler: ) -> Mapping[str, httpx.BaseTransport | None] | None: if transport is None or transport is getattr(litellm, "sync_transport", None): return None - if http2: - return _environment_proxy_mounts( - lambda proxy_url: HTTP2ConnectionSpread( - lambda: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=True), - max_connections=DEFAULT_LIMITS.max_connections, - ) - ) - return _environment_proxy_mounts(lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert)) + return _environment_proxy_mounts( + lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2) + ) def get_async_httpx_client( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 236674a4a8b..3895193d779 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -10349,7 +10349,7 @@ class BaseLLMHTTPHandler: llm_provider=litellm.LlmProviders(custom_llm_provider), params={ "ssl_verify": litellm_params.get("ssl_verify", None), - **vector_store_provider_config.get_httpx_client_params(), + **vector_store_provider_config.get_async_httpx_client_params(), }, ) else: @@ -10487,12 +10487,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = get_httpx_client( - params={ - "ssl_verify": litellm_params.get("ssl_verify", None), - **vector_store_provider_config.get_httpx_client_params(), - } - ) + sync_httpx_client = get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_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 d7676d54451..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,7 +234,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): def __init__(self): super().__init__() - def get_httpx_client_params(self) -> Mapping[str, object]: + def get_async_httpx_client_params(self) -> Mapping[str, object]: return MappingProxyType({"http2": True}) @staticmethod diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index 07b62a7d69e..b8f2ae9b76a 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/tests/integration/sdk/test_http2_wire.py @@ -39,7 +39,7 @@ 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]()}) +_HELD_ARRIVALS: Final = queue.SimpleQueue[int]() _TOKEN_REQUESTS: Final = queue.SimpleQueue[str]() _COUNTED_CLIENT_PORTS: Final = queue.SimpleQueue[int]() VERTEX_SERVING_CONFIG_PATH: Final = ( @@ -103,12 +103,9 @@ 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"): + if scope["path"] == "/hold:search": client: Final = scope["client"] - await _send_json( - send, {"results": [{"id": await _hold_until_all_arrive(hold_room, client[1] if client else 0)}]} - ) + 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"]) @@ -139,12 +136,12 @@ 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) +async def _hold_until_all_arrive(client_port: int) -> str: + _HELD_ARRIVALS.put(client_port) deadline: Final = time.monotonic() + HOLD_SECONDS - while arrivals.qsize() < HELD_SEARCHES and time.monotonic() < deadline: + while _HELD_ARRIVALS.qsize() < HELD_SEARCHES and time.monotonic() < deadline: await asyncio.sleep(0.01) - return "held-together" if arrivals.qsize() >= HELD_SEARCHES else "released-alone" + return "held-together" if _HELD_ARRIVALS.qsize() >= HELD_SEARCHES else "released-alone" async def _send_json(send: ASGISendCallable, body: dict[str, object]) -> None: @@ -302,11 +299,11 @@ async def _avertex_search_file_ids( return [result.get("file_id") for result in response.get("data") or ()] -def test_sync_vertex_search_negotiates_http2_without_global_flag( +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/2"] + assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"] def test_async_vertex_search_negotiates_http2_without_global_flag( @@ -372,7 +369,7 @@ def test_vertex_search_over_forced_ipv4_trusts_the_configured_ca( _trust_peer_without_global_http2(monkeypatch, http2_tls_peer) monkeypatch.setattr(litellm, "force_ipv4", True) - assert _vertex_search_file_ids(http2_tls_peer) == ["http/2"] + assert _vertex_search_file_ids(http2_tls_peer) == ["http/1.1"] assert asyncio.run(_avertex_search_file_ids(http2_tls_peer)) == ["http/2"] @@ -553,8 +550,9 @@ def test_vertex_search_retries_a_connection_the_server_drains_with_goaway( _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 _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"] + _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( @@ -597,41 +595,25 @@ def test_vertex_search_keeps_its_http2_connection_across_an_idle_gap( 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 +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) - 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())) + 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 @@ -685,5 +667,6 @@ def test_vertex_search_over_forced_ipv4_presents_the_client_certificate( 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"] + _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 index 36d79e0260b..d637631110f 100644 --- a/tests/unit/llms/custom_httpx/test_http2_connection_spread.py +++ b/tests/unit/llms/custom_httpx/test_http2_connection_spread.py @@ -1,8 +1,5 @@ import asyncio -import threading -import time from collections.abc import Callable -from concurrent.futures import ThreadPoolExecutor from typing import Final import httpx @@ -11,7 +8,6 @@ import pytest from litellm.llms.custom_httpx.http2_connection_spread import ( HTTP2_STREAMS_PER_CONNECTION, AsyncHTTP2ConnectionSpread, - HTTP2ConnectionSpread, ) URL: Final = "https://search.example/v1/search" @@ -66,39 +62,6 @@ class _RefusedConnections: 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(): @@ -107,13 +70,6 @@ async def _until(condition: Callable[[], bool]) -> None: 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() @@ -179,26 +135,3 @@ async def test_async_spread_frees_the_stream_when_the_request_fails(): 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, - ]