mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
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).
This commit is contained in:
parent
5e4b3838aa
commit
36104956c0
4 changed files with 171 additions and 58 deletions
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue