From 389c4befcaef9c8434a3bf9c4114763b65096158 Mon Sep 17 00:00:00 2001 From: Amin Saedi Date: Thu, 10 Sep 2026 18:05:26 +0000 Subject: [PATCH] fix(http_handler): defer finalizer client close while requests are in flight MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A cache-evicted AsyncHTTPHandler is finalized the moment the cache drops it. The finalizer's sole-referrer refcount guard proves nothing else holds the client object, but a request in flight references only the pooled connection, so the guard cannot see it: closing there tears the pool down under every live SSE stream, one batch per handler-cache TTL per process. The finalizer-scheduled close task now consults the evicted-client closer first: a busy client (connection in flight) is queued and closed by the closer once idle and out of grace (EvictedClientCloser.defer_if_busy), an idle one is closed in the task as before — the existing task-reference machinery, no-loop and cross-loop disposal branches are unchanged. The in-flight check runs inside the task, not in __del__, so no closer lock is taken in GC context. Resolves the hourly batched mid-stream stream deaths observed in a production proxy deployment (streams died in same-second batches at a fixed phase after process start, matching the client-cache TTL). --- litellm/caching/evicted_client_closer.py | 30 +++++++++++ litellm/llms/custom_httpx/http_handler.py | 23 +++++++- .../caching/test_evicted_client_closer.py | 54 +++++++++++++++++++ .../llms/custom_httpx/test_http_handler.py | 47 ++++++++++++++++ 4 files changed, 153 insertions(+), 1 deletion(-) diff --git a/litellm/caching/evicted_client_closer.py b/litellm/caching/evicted_client_closer.py index eee7e2ea289..334de189eff 100644 --- a/litellm/caching/evicted_client_closer.py +++ b/litellm/caching/evicted_client_closer.py @@ -206,6 +206,36 @@ class EvictedClientCloser: ) ) + def defer_if_busy(self, client: object) -> bool: + """Queue a deferred close for a client with a request in flight; no-op when idle. + + For a finalized handler's client the sole-referrer refcount check has already + proven nothing else holds the client object, but a request in flight references + only the pooled connection, so it is invisible to that check. A busy client is + queued and closed by ``reap`` once idle and out of grace; returns True when the + close was deferred (the caller must NOT close the client itself), False when the + client is idle and the caller may close it directly. + """ + if _close_function(client) is None: + return False + if not _has_connection_in_flight(client): + return False + self.mark_owned(client) + self.schedule(client) + return True + + def close_or_defer(self, client: object) -> None: + """Close an unreferenced client now if idle, else queue it for a deferred close. + + Same in-flight rule as ``defer_if_busy``; an idle client is closed immediately, + preserving the reclamation the finalizer used to do. + """ + if _close_function(client) is None: + return + if self.defer_if_busy(client): + return + self._close(client) + def reap(self) -> None: """Close every queued client that is due, idle, and closable from here. diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index f4883b57fbc..396b53c2c62 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1102,13 +1102,34 @@ class AsyncHTTPHandler: # here is the cross-loop path the transport refuses. self._dispose_wrapped_aiohttp_session() return - task: Final = loop.create_task(self._client.aclose()) + task: Final = loop.create_task(self._finalizer_close_client(self._client)) cls: Final = type(self) cls._finalizer_close_tasks.add(task) task.add_done_callback(cls._on_finalizer_close_done) except Exception: pass + @staticmethod + async def _finalizer_close_client(client: httpx.AsyncClient) -> None: + """Close a finalized handler's client, unless a request is still in flight. + + The sole-referrer refcount guard in ``__del__`` proves nothing else holds the + client *object*, but a request in flight references only the pooled connection, + so it is invisible to that check: closing here would tear the pool down under + every live SSE stream (a cache-evicted handler is finalized the moment the cache + drops it — one batch of mid-turn stream deaths per handler-cache TTL per + process). A busy client is handed to the evicted-client closer, which closes it + once it reports no connection in flight and a grace window has passed; an idle + one is closed here so the finalizer keeps its reclamation. The in-flight check + runs inside this task, not in ``__del__``, so no closer lock is taken in GC + context. + """ + from litellm.caching.evicted_client_closer import default_evicted_client_closer + + if default_evicted_client_closer.defer_if_busy(client): + return + await client.aclose() + @staticmethod def _create_async_transport( ssl_context: ssl.SSLContext | None = None, diff --git a/tests/test_litellm/caching/test_evicted_client_closer.py b/tests/test_litellm/caching/test_evicted_client_closer.py index 939be5f3d6b..4d2383b02ee 100644 --- a/tests/test_litellm/caching/test_evicted_client_closer.py +++ b/tests/test_litellm/caching/test_evicted_client_closer.py @@ -409,3 +409,57 @@ def test_a_reap_looks_at_what_is_due_rather_than_at_the_whole_queue(): f"{CountingDeadline.comparisons} deadline comparisons for {evictions} evictions; " "a reap is walking the whole queue" ) + + +@pytest.mark.asyncio +async def test_close_or_defer_closes_an_idle_client_immediately(): + clock = FakeClock() + closer = make_closer(clock) + client = AsyncClient() + + closer.close_or_defer(client) + await asyncio.sleep(0.05) + + assert client.closed is True + assert closer.pending_count == 0 + + +@pytest.mark.asyncio +async def test_close_or_defer_defers_a_client_with_a_request_on_the_wire(): + server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + clock = FakeClock() + closer = make_closer(clock) + client = httpx.AsyncClient() + + async with asyncio.timeout(30): + async with client.stream("GET", f"http://127.0.0.1:{port}/") as response: + body_iter = response.aiter_raw() + await body_iter.__anext__() + + closer.close_or_defer(client) + + assert not client.is_closed + assert closer.pending_count == 1 + + remainder = b"".join([chunk async for chunk in body_iter]) + assert b"hello" in remainder + + clock.advance(61.0) + closer.reap() + await asyncio.sleep(0.05) + + assert client.is_closed + assert closer.pending_count == 0 + # No wait_closed(): on Python >= 3.12.1 it waits for every client + # transport, and a pooled keepalive connection would park it forever. + server.close() + + +def test_close_or_defer_ignores_a_value_without_a_close_function(): + clock = FakeClock() + closer = make_closer(clock) + + closer.close_or_defer(object()) + + assert closer.pending_count == 0 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 f8868cfaf83..92fe128d3ac 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1675,3 +1675,50 @@ async def test_bounded_get_closes_stream_on_cancellation(respx_mock, monkeypatch finally: await handler.close() assert closed.is_set() + + +async def _slow_chunked_upstream(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + """Serves a chunked body in two installments, so a stream is on the wire while the handler dies.""" + await reader.read(4096) + writer.write(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n") + writer.write(b"5\r\nfirst\r\n") + await writer.drain() + await asyncio.sleep(0.4) + writer.write(b"4\r\nlast\r\n0\r\n\r\n") + await writer.drain() + + +@pytest.mark.asyncio +async def test_collected_handler_never_kills_a_stream_in_flight(monkeypatch): + """ + Regression: a cache-evicted (hence collected) handler's finalizer used to close the + owned client while its pool still served live SSE streams, killing every one of them + mid-turn once per handler-cache TTL. The finalizer must defer to the evicted-client + closer while a connection is in flight, so the stream reads to completion. + """ + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "force_ipv4", False) + + server = await asyncio.start_server(_slow_chunked_upstream, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + + handler = AsyncHTTPHandler() + async with asyncio.timeout(30): + request = handler.client.build_request("GET", f"http://127.0.0.1:{port}/") + response = await handler.client.send(request, stream=True) + body_iter = response.aiter_raw() + first = await body_iter.__anext__() + assert b"first" in first + + del handler + gc.collect() + await asyncio.sleep(0.1) + + remainder = b"".join([chunk async for chunk in body_iter]) + assert b"last" in remainder + + await response.aclose() + # No wait_closed(): the surviving client's pooled keepalive connection is + # the point of this test, and on Python >= 3.12.1 wait_closed() waits for + # every client transport, parking the suite forever. + server.close()