diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index cc37d1d9eca..05dff0cb9d8 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 io import BytesIO @@ -185,6 +186,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. @@ -778,6 +806,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): @@ -982,6 +1012,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): @@ -1451,6 +1483,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: @@ -1501,6 +1535,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: @@ -1551,6 +1587,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( @@ -1600,6 +1638,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 new file mode 100644 index 00000000000..1a6ab1b1827 --- /dev/null +++ b/tests/local_testing/test_handler_gc_does_not_close_client.py @@ -0,0 +1,315 @@ +""" +Collecting an HTTP handler must not abort a response that is still on the wire. + +``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. + +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. + +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. + +Related: https://github.com/BerriAI/litellm/issues/24929 +""" + +import asyncio +import gc +import threading +import time +import weakref +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import httpx +import pytest + +import litellm +from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + get_async_httpx_client, +) +from litellm.types.utils import LlmProviders + +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"]) + +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: + """In-process HTTP/1.1 server that answers every request with chunked SSE frames.""" + + def __init__(self, frame_count: int = FRAME_COUNT, frame_delay: float = 0.05) -> None: + self.frame_count = frame_count + self.frame_delay = frame_delay + parent = self + + class _Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def _stream(self): + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Transfer-Encoding", "chunked") + self.end_headers() + try: + for index in range(parent.frame_count): + frame = f"data: frame-{index}\n\n".encode() + self.wfile.write(b"%x\r\n" % len(frame) + frame + b"\r\n") + self.wfile.flush() + time.sleep(parent.frame_delay) + self.wfile.write(b"0\r\n\r\n") + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + pass + + do_GET = _stream + do_POST = _stream + + def log_message(self, *args): + pass + + self._server = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) + self.url = f"http://127.0.0.1:{self._server.server_address[1]}/stream" + + def __enter__(self): + threading.Thread(target=self._server.serve_forever, daemon=True).start() + return self + + def __exit__(self, *exc_info): + self._server.shutdown() + self._server.server_close() + + +def _select_transport(monkeypatch, disable_aiohttp_transport: bool) -> None: + monkeypatch.delenv("DISABLE_AIOHTTP_TRANSPORT", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", disable_aiohttp_transport) + monkeypatch.setattr(litellm, "force_ipv4", False) + + +async def _read_frames(response: httpx.Response) -> int: + """Count SSE frames, collecting garbage between chunks so a finalizer has every chance to fire. + + The body is joined before counting: a chunk boundary can fall inside the + marker, which a per-chunk count would miss. + """ + chunks = [] + async for chunk in response.aiter_bytes(): + chunks.append(chunk) + gc.collect() + 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 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)) + response = await handler.post(server.url, stream=True) + + ref = weakref.ref(handler) + del handler + gc.collect() + await asyncio.sleep(0) # let any close the finalizer scheduled run + + assert ref() is not None, STILL_PINNED + assert await asyncio.wait_for(_read_frames(response), timeout=READ_TIMEOUT_SECONDS) == FRAME_COUNT + + del response + gc.collect() + assert ref() is None, NOT_RELEASED + + +def test_sync_stream_survives_handler_collection(monkeypatch): + """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 + impossible" -- a workaround for this finalizer rather than a fix for it. + """ + monkeypatch.setattr(litellm, "force_ipv4", False) + + with _ChunkedSSEServer() as server: + handler = HTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0)) + response = handler.post(server.url, stream=True) + + 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 + + del response + gc.collect() + assert ref() is None, NOT_RELEASED + + +@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 +@BOTH_TRANSPORTS +async def test_cached_handler_eviction_does_not_abort_an_in_flight_stream(monkeypatch, disable_aiohttp_transport): + """Evicting a cached handler mid-stream leaves the stream alone. + + ``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. + """ + _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) + 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() + + assert ref() is not None, STILL_PINNED + assert await asyncio.wait_for(_read_frames(response), timeout=READ_TIMEOUT_SECONDS) == FRAME_COUNT + + del response + gc.collect() + assert ref() is None, NOT_RELEASED diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index a52bf58e944..33272a1a9e4 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1025,6 +1025,86 @@ def test_handed_out_sync_client_pool_survives_handler_collection(keepalive_serve consumer_client.close() +def _mock_transport() -> httpx.MockTransport: + """Answers anything with a short body, left unread when the caller asked to stream.""" + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, request=request, content=b"ab") + + return httpx.MockTransport(respond) + + +RELEASED_TOO_EARLY = "the handler was released while its response could still read" +NEVER_RELEASED = "the handler outlived the response that was holding it" + +# Every method that can hand back a body the caller has not read yet, which is +# every one that passes stream= down to send(). Parametrized so a method added +# later is covered here rather than being the one that forgets to anchor. +ASYNC_STREAMING_SENDS = ["post", "delete"] +SYNC_STREAMING_SENDS = ["post", "patch", "put", "delete"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ASYNC_STREAMING_SENDS) +async def test_a_streaming_response_holds_its_handler_until_it_is_released(method): + """The finalizer must not run while a body this handler issued can still arrive. + + ``_handler_may_close_client`` cannot see that body: it holds the connection it + reads from and never the client. Anchoring the handler to the response is what + withholds the close, and releasing the anchor is what still delivers one. + """ + handler = AsyncHTTPHandler() + handler.client._transport = _mock_transport() + ref = weakref.ref(handler) + response = await getattr(handler, method)("https://example.invalid/stream", stream=True) + + del handler + gc.collect() + assert ref() is not None, RELEASED_TOO_EARLY + + assert await response.aread() == b"ab" + del response + gc.collect() + assert ref() is None, NEVER_RELEASED + + +@pytest.mark.parametrize("method", SYNC_STREAMING_SENDS) +def test_a_sync_streaming_response_holds_its_handler_until_it_is_released(method): + """The sync finalizer closes inline, so the same anchor has to hold it off.""" + handler = HTTPHandler() + handler.client._transport = _mock_transport() + ref = weakref.ref(handler) + response = getattr(handler, method)("https://example.invalid/stream", stream=True) + + del handler + gc.collect() + assert ref() is not None, RELEASED_TOO_EARLY + + assert response.read() == b"ab" + del response + gc.collect() + assert ref() is None, NEVER_RELEASED + + +@pytest.mark.asyncio +async def test_a_fully_read_response_does_not_hold_its_handler(): + """A non-streaming response is complete when ``post`` returns, so it anchors nothing. + + Otherwise every client close would wait on whatever the caller does next with + a response it has already read. + """ + handler = AsyncHTTPHandler() + handler.client._transport = _mock_transport() + ref = weakref.ref(handler) + response = await handler.post("https://example.invalid/whole") + assert response.content == b"ab" + + del handler + gc.collect() + + assert ref() is None, "a fully-read response pinned its handler" + + def test_sync_close_leaves_caller_supplied_client_open(): supplied = httpx.Client() handler = HTTPHandler(client=supplied)