From 36104956c097ded4e9e55d046f78a94d4ca7b1dc Mon Sep 17 00:00:00 2001 From: Amin Saedi Date: Tue, 25 Aug 2026 13:44:25 +0000 Subject: [PATCH] fix(http_handler): defer finalizer client close while requests are in flight 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 now defers to EvictedClientCloser.close_or_defer, which closes an idle client immediately and queues a busy one until it reports no connection in flight and the grace window has passed. The handoff runs via loop.call_soon so no 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 | 17 +++ litellm/llms/custom_httpx/http_handler.py | 42 ++++--- .../caching/test_evicted_client_closer.py | 54 ++++++++ .../llms/custom_httpx/test_http_handler.py | 116 ++++++++++++------ 4 files changed, 171 insertions(+), 58 deletions(-) diff --git a/litellm/caching/evicted_client_closer.py b/litellm/caching/evicted_client_closer.py index eee7e2ea289..7ccf30a10c3 100644 --- a/litellm/caching/evicted_client_closer.py +++ b/litellm/caching/evicted_client_closer.py @@ -206,6 +206,23 @@ class EvictedClientCloser: ) ) + def close_or_defer(self, client: object) -> None: + """Close an unreferenced client now if idle, else queue it for a deferred close. + + 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. An idle client is + closed immediately, preserving the reclamation the finalizer used to do; a busy + one is queued and closed by ``reap`` once idle and out of grace. + """ + if _close_function(client) is None: + return + if not _has_connection_in_flight(client): + self._close(client) + return + self.mark_owned(client) + self.schedule(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 b6e93f590ca..6bae0f97977 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -9,7 +9,7 @@ import threading import time from collections.abc import AsyncIterable, Callable, Iterable, Mapping from http.cookiejar import CookieJar, DefaultCookiePolicy -from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, Optional, TypeAlias, TypedDict +from typing import TYPE_CHECKING, Any, Final, NoReturn, Optional, TypeAlias, TypedDict import certifi import httpx @@ -160,6 +160,26 @@ 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 _hand_off_client_to_evicted_closer(client: object) -> None: + """ + Route a finalized handler's client to the evicted-client closer instead of closing it. + + A request in flight holds references to the pooled connection, not to the client + object, so the sole-referrer refcount check cannot see it. Closing from ``__del__`` + therefore tears the pool down under live requests: a cache-evicted handler is + finalized the moment the cache drops it, and every SSE stream on its pool dies + mid-turn (one batch per handler-cache TTL per process). The closer closes an idle + client immediately (preserving reclamation) and defers a busy one until it reports + no connection in flight and a grace window has passed. + + Runs via ``loop.call_soon`` rather than inside the finalizer, so no lock is taken + in GC context. + """ + from litellm.caching.evicted_client_closer import default_evicted_client_closer + + default_evicted_client_closer.close_or_defer(client) + + def blocked_cookie_jar() -> CookieJar: """A jar that stores no response cookie and sends none, for httpx clients. @@ -933,21 +953,6 @@ class AsyncHTTPHandler: response.raise_for_status() return response - # Strong references to finalizer-scheduled client-close tasks. A bare - # create_task() result may be garbage-collected before it runs, leaving - # the underlying aiohttp session unclosed ("Unclosed client session"). - # Mirrors LiteLLMAiohttpTransport._background_close_tasks. - _finalizer_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes - - @classmethod - def _on_finalizer_close_done(cls, task: "asyncio.Task[None]") -> None: - cls._finalizer_close_tasks.discard(task) - if task.cancelled(): - return - exc: Final = task.exception() - if exc is not None: - verbose_logger.debug("Error closing client at finalization: %s", exc) - def _aiohttp_session_bound_elsewhere(self, loop: asyncio.AbstractEventLoop) -> bool: """True when the wrapped aiohttp session is bound to a loop other than ``loop`` — awaiting ``aclose()`` here would touch that loop's internals.""" @@ -1006,10 +1011,7 @@ 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()) - cls: Final = type(self) - cls._finalizer_close_tasks.add(task) - task.add_done_callback(cls._on_finalizer_close_done) + loop.call_soon(_hand_off_client_to_evicted_closer, self._client) except Exception: pass 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 16d57437043..718addaec32 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1,8 +1,6 @@ import asyncio import gc -import io import os -import pathlib import ssl import threading import weakref @@ -364,9 +362,11 @@ async def test_async_handler_with_shared_session(): async def test_get_async_httpx_client_with_shared_session(): """Test get_async_httpx_client with shared session""" from litellm.llms.custom_httpx.http_handler import ( - get_async_httpx_client, AsyncHTTPHandler as AsyncHTTPHandlerReload, ) + from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + ) from litellm.types.utils import LlmProviders # Create a mock shared session @@ -388,9 +388,11 @@ async def test_get_async_httpx_client_with_shared_session(): async def test_get_async_httpx_client_without_shared_session(): """Test get_async_httpx_client without shared session (backward compatibility)""" from litellm.llms.custom_httpx.http_handler import ( - get_async_httpx_client, AsyncHTTPHandler as AsyncHTTPHandlerReload, ) + from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + ) from litellm.types.utils import LlmProviders # Test without shared session @@ -426,6 +428,7 @@ async def test_session_reuse_chain(): def test_shared_session_parameter_in_acompletion(): """Test that acompletion function accepts shared_session parameter""" import inspect + from litellm.main import acompletion # Get the function signature @@ -443,6 +446,7 @@ def test_shared_session_parameter_in_acompletion(): def test_shared_session_parameter_in_completion(): """Test that completion function accepts shared_session parameter""" import inspect + from litellm.main import completion # Get the function signature @@ -461,9 +465,11 @@ def test_shared_session_parameter_in_completion(): async def test_session_reuse_integration(): """Integration test for session reuse functionality""" from litellm.llms.custom_httpx.http_handler import ( - get_async_httpx_client, AsyncHTTPHandler as AsyncHTTPHandlerReload, ) + from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + ) from litellm.types.utils import LlmProviders # Create a mock session @@ -1225,10 +1231,13 @@ def test_finalizer_without_running_loop_closes_dead_loop_session(): @pytest.mark.asyncio -async def test_finalizer_with_running_loop_schedules_close_and_holds_task_ref(): - """With a running loop, finalization schedules an async close and must keep - a strong reference to the task until it completes — a bare create_task() - result may be collected before it runs, leaving the session unclosed.""" +async def test_finalizer_with_running_loop_hands_idle_client_to_closer(): + """With a running loop, finalization defers to the evicted-client closer, + which closes an idle client immediately while holding a strong reference to + the close task — a bare create_task() result may be collected before it + runs, leaving the session unclosed.""" + from litellm.caching.evicted_client_closer import default_evicted_client_closer + handler = AsyncHTTPHandler(timeout=61.0) transport = handler.client._transport assert isinstance(transport, LiteLLMAiohttpTransport) @@ -1236,16 +1245,17 @@ async def test_finalizer_with_running_loop_schedules_close_and_holds_task_ref(): assert not session.closed del transport - baseline_tasks = set(AsyncHTTPHandler._finalizer_close_tasks) + baseline_tasks = set(default_evicted_client_closer._close_tasks) del handler gc.collect() + await asyncio.sleep(0) - scheduled = AsyncHTTPHandler._finalizer_close_tasks - baseline_tasks + scheduled = default_evicted_client_closer._close_tasks - baseline_tasks assert len(scheduled) == 1 await asyncio.gather(*scheduled) assert session.closed - assert not (AsyncHTTPHandler._finalizer_close_tasks & scheduled) + assert not (default_evicted_client_closer._close_tasks & scheduled) @pytest.mark.asyncio @@ -1277,40 +1287,70 @@ async def test_sync_close_helper_respects_session_ownership(): await owned_handler.close() -@pytest.mark.asyncio -async def test_finalizer_close_done_consumes_exception(): - """A failing finalizer close must have its exception retrieved by the done - callback, or asyncio emits "Task exception was never retrieved" at GC — - the same log noise the finalizer path exists to eliminate.""" - - async def failing_close() -> None: - raise RuntimeError("close failed") - - task = asyncio.get_running_loop().create_task(failing_close()) - AsyncHTTPHandler._finalizer_close_tasks.add(task) - await asyncio.sleep(0) - - AsyncHTTPHandler._on_finalizer_close_done(task) - assert task not in AsyncHTTPHandler._finalizer_close_tasks - - cancelled = asyncio.get_running_loop().create_task(asyncio.sleep(30)) - cancelled.cancel() - await asyncio.sleep(0) - AsyncHTTPHandler._on_finalizer_close_done(cancelled) - - @pytest.mark.asyncio async def test_finalizer_on_live_loop_disposes_foreign_loop_session_without_scheduling(): """GC on a live loop (e.g. the app's) of a handler whose session belongs to - another, dead loop must not schedule aclose() here — that is the cross-loop - path the transport refuses — and must still dispose the session.""" + another, dead loop must not hand the client to the closer — that is the + cross-loop path the transport refuses — and must still dispose the session.""" + from litellm.caching.evicted_client_closer import default_evicted_client_closer + handler = AsyncHTTPHandler(timeout=61.0) session = await asyncio.to_thread(_mint_session_on_dead_loop, handler) assert not session.closed - baseline_tasks = set(AsyncHTTPHandler._finalizer_close_tasks) + baseline_tasks = set(default_evicted_client_closer._close_tasks) + baseline_pending = default_evicted_client_closer.pending_count del handler gc.collect() + await asyncio.sleep(0) - assert AsyncHTTPHandler._finalizer_close_tasks == baseline_tasks + assert default_evicted_client_closer._close_tasks == baseline_tasks + assert default_evicted_client_closer.pending_count == baseline_pending assert session.closed + + +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()