Merge 5cb8b9b63f into 39e5b0c2d1

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.
 

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, Optional, TypeAlias, TypedDict
+from typing import TYPE_CHECKING, Any, Final, 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
 

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

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()