diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index e1f0fc9e7d3..612d5997ec5 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -7,6 +7,7 @@ import ssl import sys import threading import time +import weakref from collections.abc import AsyncIterable, Callable, Iterable, Mapping from http.cookiejar import CookieJar, DefaultCookiePolicy from types import MappingProxyType @@ -178,6 +179,33 @@ def _handler_may_close_client(client_refcount: int, owns_client: bool) -> bool: return owns_client and client_refcount <= _CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER +def _drop_streaming_anchor(_handler: object) -> None: + """Release a handler anchored to a streaming response. See ``_anchor_handler_to``. + + The work is the reference held until this point, so there is nothing to do here. + """ + + +def _anchor_handler_to(response: httpx.Response, handler: object) -> None: + """Keep the handler alive for as long as a streaming response can still read. + + A body still arriving reads through the handler's connection pool, and closing + the client tears that pool down. The refcount ``_handler_may_close_client`` + reads cannot see that body: the reference graph runs response -> stream -> + connection and stops there, so a client carrying one looks exactly like an + unreferenced client, and the finalizer closes it mid-body. + + ``weakref.finalize`` holds the handler in its own registry rather than on the + response, which matters twice. The handler stays out of the response's + reference cycle, so it is finalized by refcount once the anchor drops and can + still schedule an async close, instead of being finalized inside a cyclic + collection that reaps its aiohttp session in the same pass. And a handler + serving several streams collects only once every one of them is done, because + each anchor holds it separately. + """ + weakref.finalize(response, _drop_streaming_anchor, handler) + + def blocked_cookie_jar() -> CookieJar: """A jar that stores no response cookie and sends none, for httpx clients. @@ -704,6 +732,8 @@ class AsyncHTTPHandler: content=request_content, ) response: Final = await self.client.send(req, stream=stream) + if stream: + _anchor_handler_to(response, self) response.raise_for_status() return response except (httpx.RemoteProtocolError, httpx.ConnectError): @@ -898,6 +928,8 @@ class AsyncHTTPHandler: content=request_content, ) response: Final = await self.client.send(req, stream=stream) + if stream: + _anchor_handler_to(response, self) response.raise_for_status() return response except (httpx.RemoteProtocolError, httpx.ConnectError): @@ -1362,6 +1394,8 @@ class HTTPHandler: content=request_content, ) response: Final = self.client.send(req, stream=stream) + if stream: + _anchor_handler_to(response, self) response.raise_for_status() return response except httpx.TimeoutException: @@ -1412,6 +1446,8 @@ class HTTPHandler: content=request_content, ) response: Final = self.client.send(req, stream=stream) + if stream: + _anchor_handler_to(response, self) response.raise_for_status() return response except httpx.TimeoutException: @@ -1462,6 +1498,8 @@ class HTTPHandler: content=request_content, ) response: Final = self.client.send(req, stream=stream) + if stream: + _anchor_handler_to(response, self) return response except httpx.TimeoutException: raise litellm.Timeout( @@ -1511,6 +1549,8 @@ class HTTPHandler: content=request_content, ) response: Final = self.client.send(req, stream=stream) + if stream: + _anchor_handler_to(response, self) response.raise_for_status() return response except httpx.TimeoutException: diff --git a/tests/local_testing/conftest.py b/tests/local_testing/conftest.py index 5535a62bb81..228457f4d55 100644 --- a/tests/local_testing/conftest.py +++ b/tests/local_testing/conftest.py @@ -75,6 +75,9 @@ _VCR_INCOMPATIBLE_FILES = frozenset( "test_router_caching.py", # Hits the local fake OpenAI endpoint on 127.0.0.1; nothing to record. "test_fake_openai_endpoint.py", + # Needs the real connection pool a collected handler tears down; vcrpy + # patches the transport that pool lives in. + "test_handler_gc_does_not_close_client.py", } ) diff --git a/tests/local_testing/test_handler_gc_does_not_close_client.py b/tests/local_testing/test_handler_gc_does_not_close_client.py index 3874405a261..1a6ab1b1827 100644 --- a/tests/local_testing/test_handler_gc_does_not_close_client.py +++ b/tests/local_testing/test_handler_gc_does_not_close_client.py @@ -1,36 +1,39 @@ """ -Garbage-collecting an HTTP handler must not close the httpx client it holds. +Collecting an HTTP handler must not abort a response that is still on the wire. -``HTTPHandler`` and ``AsyncHTTPHandler`` used to close their client from -``__del__``. Closing an httpx client tears down the connection pool that every -in-flight response is streaming through, and it permanently invalidates the -client for future requests -- including for callers who only ever borrowed -``handler.client``. Since nothing in a response's reference graph points back -at the handler, and since litellm caches handlers behind a one-hour TTL, the -handler routinely became collectable while its client was still in use. +``HTTPHandler`` and ``AsyncHTTPHandler`` close their client from ``__del__``. +Closing a client tears down the connection pool, which aborts every response +still streaming through it. ``_handler_may_close_client`` already withholds the +close from a client someone else holds, but a streaming response holds the +connection it is reading from and never the client, so the refcount it reads +says "sole referrer" for exactly the client that is busiest. The handler is +routinely collectable at that moment: a provider's streaming call returns the +response and drops the handler, and ``get_async_httpx_client`` caches handlers +behind a one-hour TTL and then lets them go. -``LLMClientCache`` documents the invariant this broke: evicted clients "may -still be in use by in-flight requests", so they are left to normal garbage -collection rather than closed eagerly. A finalizer that closes on collection -defeats exactly that. +The fix anchors the handler to the streaming response, so these tests turn on +*when* the handler is collected rather than on whether it is: pinned while the +body can still arrive, released once the caller is done with the response. -Each test below is one shape the finalizers broke; all of them fail if either -``__del__`` comes back. Async cases run on both transports, because litellm -defaults to aiohttp and only uses httpcore when aiohttp is disabled. +Nothing here re-tests the shapes ``_handler_may_close_client`` covers -- a +borrowed ``handler.client``, a caller-supplied client, an evicted-but-held +client. Those are pinned in ``tests/test_litellm/llms/custom_httpx/ +test_http_handler.py``. What is uncovered there is the in-flight response, so no +test here may keep the client in a local: that inflates the very refcount under +test, and the test then passes on a broken handler. They hold weak references +instead, which the refcount does not count. These live here rather than under ``tests/test_litellm/`` because they need a real connection pool: a mocked transport goes on yielding chunks after its client is closed, so the very teardown under test is what a mock cannot -reproduce. The server is a hermetic, credential-free ``ThreadingHTTPServer`` -on an ephemeral loopback port, and needs no network access beyond it. +reproduce. The server is a hermetic, credential-free ``ThreadingHTTPServer`` on +an ephemeral loopback port, and needs no network access beyond it. Related: https://github.com/BerriAI/litellm/issues/24929 """ import asyncio import gc -import os -import sys import threading import time import weakref @@ -39,8 +42,6 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import httpx import pytest -sys.path.insert(0, os.path.abspath("../..")) - import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import ( @@ -54,15 +55,12 @@ FRAME_COUNT = 6 # Generous: the server emits all frames in ~0.3s. A client whose pool was torn # down mid-stream can stall silently instead of raising, so reads are bounded. READ_TIMEOUT_SECONDS = 15.0 +RELEASE_TIMEOUT_SECONDS = 3.0 BOTH_TRANSPORTS = pytest.mark.parametrize("disable_aiohttp_transport", [False, True], ids=["aiohttp", "httpcore"]) -# Every test here rests on the handler actually being collected at the ``del``. -# If something ever pins it, the test would pass while guarding nothing, so each -# one checks the premise. The check comes last: a reintroduced finalizer also -# fails it, by resurrecting the handler into the task it creates for ``close()``, -# and the transport error is the more useful thing to see first. -HANDLER_NOT_COLLECTED = "handler was not collected; this test no longer exercises the finalizer path" +STILL_PINNED = "the handler was released while its response could still read" +NOT_RELEASED = "the handler outlived the response that was holding it" class _ChunkedSSEServer: @@ -129,93 +127,44 @@ async def _read_frames(response: httpx.Response) -> int: return b"".join(chunks).count(b"data: frame-") +async def _wait_until(is_done, failure: str) -> None: + deadline = time.monotonic() + RELEASE_TIMEOUT_SECONDS + while time.monotonic() < deadline: + if is_done(): + return + await asyncio.sleep(0.05) + pytest.fail(failure) + + @pytest.mark.asyncio @BOTH_TRANSPORTS async def test_async_stream_survives_handler_collection(monkeypatch, disable_aiohttp_transport): - """A response being streamed keeps working after its handler is collected.""" + """A response still streaming keeps working after its handler goes out of scope. + + The caller holds the response and nothing else, which is what a provider's + streaming path is left with once ``post(..., stream=True)`` has returned. + """ _select_transport(monkeypatch, disable_aiohttp_transport) with _ChunkedSSEServer() as server: handler = AsyncHTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0)) - client = handler.client - try: - response = await client.send(client.build_request("GET", server.url), stream=True) + response = await handler.post(server.url, stream=True) - # The handler loses its last reference while the body is still streaming. - ref = weakref.ref(handler) - del handler - gc.collect() + ref = weakref.ref(handler) + del handler + gc.collect() + await asyncio.sleep(0) # let any close the finalizer scheduled run - frames = await asyncio.wait_for(_read_frames(response), timeout=READ_TIMEOUT_SECONDS) - assert frames == FRAME_COUNT - assert client.is_closed is False - assert ref() is None, HANDLER_NOT_COLLECTED - finally: - await client.aclose() + assert ref() is not None, STILL_PINNED + assert await asyncio.wait_for(_read_frames(response), timeout=READ_TIMEOUT_SECONDS) == FRAME_COUNT - -@pytest.mark.asyncio -@BOTH_TRANSPORTS -async def test_borrowed_async_client_outlives_its_handler(monkeypatch, disable_aiohttp_transport): - """A caller that keeps only ``handler.client`` can still send requests once the handler is gone. - - This is the shape at litellm/a2a_protocol/main.py (``httpx_client = - _async_handler.client``, handed to the a2a SDK) and at - litellm/proxy/pass_through_endpoints/pass_through_endpoints.py (``async_client - = async_client_obj.client``). Both take the handler from - ``get_async_httpx_client``, so the cache pins it for - ``_DEFAULT_TTL_FOR_HTTPX_CLIENTS`` (one hour) and then lets it go on eviction, - at which point it is collected while the borrowed client is still serving a - longer-lived consumer: ``create_a2a_client`` hands its client to the a2a SDK - and documents it as "create client once, reuse for multiple requests". - """ - _select_transport(monkeypatch, disable_aiohttp_transport) - - with _ChunkedSSEServer(frame_count=1, frame_delay=0.0) as server: - handler = AsyncHTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0)) - client = handler.client - try: - ref = weakref.ref(handler) - del handler - gc.collect() - # A finalizer would close the client from a task, so let the loop turn. - await asyncio.sleep(0.05) - - assert client.is_closed is False - response = await client.get(server.url) - assert response.status_code == 200 - assert ref() is None, HANDLER_NOT_COLLECTED - finally: - await client.aclose() - - -def test_sync_handler_collection_does_not_close_a_caller_owned_client(monkeypatch): - """A throwaway handler wrapped around someone else's client must not close it. - - litellm/llms/azure/azure.py builds ``HTTPHandler(client=litellm.client_session)`` - for a single image generation and drops it. With a finalizer, that one call - left the user's shared session closed for the rest of the process. - """ - monkeypatch.setattr(litellm, "force_ipv4", False) - - with _ChunkedSSEServer(frame_count=1, frame_delay=0.0) as server: - caller_client = httpx.Client(timeout=httpx.Timeout(10.0, connect=5.0)) - monkeypatch.setattr(litellm, "client_session", caller_client) - try: - handler = HTTPHandler(client=litellm.client_session) - ref = weakref.ref(handler) - del handler - gc.collect() - - assert caller_client.is_closed is False - assert caller_client.get(server.url).status_code == 200 - assert ref() is None, HANDLER_NOT_COLLECTED - finally: - caller_client.close() + del response + gc.collect() + assert ref() is None, NOT_RELEASED def test_sync_stream_survives_handler_collection(monkeypatch): - """A sync response being streamed keeps working after its handler is collected. + """The sync handler closes inline from its finalizer, so a stream must hold it off. litellm/main.py builds a sync handler only for non-streaming calls, commented "Keep this here, otherwise, the httpx.client closes and streaming is @@ -225,26 +174,115 @@ def test_sync_stream_survives_handler_collection(monkeypatch): with _ChunkedSSEServer() as server: handler = HTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0)) - client = handler.client - try: - response = client.send(client.build_request("GET", server.url), stream=True) + response = handler.post(server.url, stream=True) - # The handler loses its last reference while the body is still streaming. - ref = weakref.ref(handler) - del handler + ref = weakref.ref(handler) + del handler + gc.collect() + assert ref() is not None, STILL_PINNED + + # Joined before counting, as in ``_read_frames``. + chunks = [] + for chunk in response.iter_bytes(): + chunks.append(chunk) gc.collect() + assert b"".join(chunks).count(b"data: frame-") == FRAME_COUNT - # Joined before counting, as in ``_read_frames``. - chunks = [] - for chunk in response.iter_bytes(): - chunks.append(chunk) - gc.collect() + del response + gc.collect() + assert ref() is None, NOT_RELEASED - assert b"".join(chunks).count(b"data: frame-") == FRAME_COUNT - assert client.is_closed is False - assert ref() is None, HANDLER_NOT_COLLECTED - finally: - client.close() + +@pytest.mark.asyncio +@BOTH_TRANSPORTS +async def test_an_abandoned_stream_still_releases_its_handler(monkeypatch, disable_aiohttp_transport): + """A caller that drops a stream unread must not pin the handler for good. + + Tying the handler to the response's own lifetime is what bounds this. No + deadline, and no poll of the connection's state, can tell an abandoned body + from one the upstream is merely slow to finish: httpx leaves the connection + checked out until the response is read or closed, and a legitimate stream is + bounded only by how long the upstream keeps sending. + """ + _select_transport(monkeypatch, disable_aiohttp_transport) + + with _ChunkedSSEServer() as server: + handler = AsyncHTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0)) + client_ref = weakref.ref(handler.client) + response = await handler.post(server.url, stream=True) + + ref = weakref.ref(handler) + del handler, response + gc.collect() + + assert ref() is None, NOT_RELEASED + await _wait_until( + lambda: client_ref() is None or client_ref().is_closed, + "the client outlived the abandoned stream without being closed", + ) + + +@pytest.mark.asyncio +@BOTH_TRANSPORTS +async def test_the_pool_is_released_once_the_stream_it_carried_ends(monkeypatch, disable_aiohttp_transport): + """Holding the finalizer off must defer the close, not drop it. + + Otherwise a collected handler leaks its pool for every streaming request it + was carrying, and on aiohttp warns "Unclosed client session" when the + collector eventually takes it. The pool and the session are children of the + client, so keeping one here does not inflate the refcount the finalizer + reads, the way keeping the client would. + """ + _select_transport(monkeypatch, disable_aiohttp_transport) + + with _ChunkedSSEServer() as server: + handler = AsyncHTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0)) + transport = handler.client._transport + if disable_aiohttp_transport: + pool = transport._pool + + def is_released() -> bool: + return pool.connections == [] + else: + session = transport._get_valid_client_session() + + def is_released() -> bool: + return session.closed + + response = await handler.post(server.url, stream=True) + + del handler, transport + gc.collect() + assert not is_released(), "the pool was torn down while it was still carrying a body" + + assert await asyncio.wait_for(_read_frames(response), timeout=READ_TIMEOUT_SECONDS) == FRAME_COUNT + del response + gc.collect() + + await _wait_until(is_released, "the pool outlived the stream it carried, unclosed") + + +@pytest.mark.asyncio +@BOTH_TRANSPORTS +async def test_a_non_streaming_response_does_not_pin_its_handler(monkeypatch, disable_aiohttp_transport): + """Only a body that can still arrive holds the handler. + + A non-streaming response has been read in full by the time ``post`` returns, + so pinning the handler to it would delay every client close behind whatever + the caller goes on to do with the response. + """ + _select_transport(monkeypatch, disable_aiohttp_transport) + + with _ChunkedSSEServer(frame_count=1, frame_delay=0.0) as server: + handler = AsyncHTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0)) + response = await handler.post(server.url) + assert response.status_code == 200 + + ref = weakref.ref(handler) + del handler + gc.collect() + + assert ref() is None, "a fully-read response pinned its handler" @pytest.mark.asyncio @@ -254,30 +292,24 @@ async def test_cached_handler_eviction_does_not_abort_an_in_flight_stream(monkey ``get_async_httpx_client`` caches handlers for an hour. When that TTL expires the cache drops the only reference to a handler whose client is - still streaming -- the production shape of #24929, and the case - ``LLMClientCache`` documents as "may still be in use by in-flight requests". + still streaming -- the production shape of #24929. """ _select_transport(monkeypatch, disable_aiohttp_transport) monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) with _ChunkedSSEServer() as server: handler = get_async_httpx_client(llm_provider=LlmProviders.OPENAI) - client = handler.client - try: - response = await client.send(client.build_request("GET", server.url), stream=True) + response = await handler.post(server.url, stream=True) - # An hour passes: the TTL expires and the cache lets the handler go. - ref = weakref.ref(handler) - litellm.in_memory_llm_clients_cache.flush_cache() - del handler - gc.collect() + # An hour passes: the TTL expires and the cache lets the handler go. + ref = weakref.ref(handler) + litellm.in_memory_llm_clients_cache.flush_cache() + del handler + gc.collect() - frames = await asyncio.wait_for(_read_frames(response), timeout=READ_TIMEOUT_SECONDS) - assert frames == FRAME_COUNT - assert client.is_closed is False + assert ref() is not None, STILL_PINNED + assert await asyncio.wait_for(_read_frames(response), timeout=READ_TIMEOUT_SECONDS) == FRAME_COUNT - # And the evicted client is still usable for the next request. - assert (await client.get(server.url)).status_code == 200 - assert ref() is None, HANDLER_NOT_COLLECTED - finally: - await client.aclose() + del response + gc.collect() + assert ref() is None, NOT_RELEASED