mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23: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-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).
This commit is contained in:
parent
9a715df212
commit
389c4befca
4 changed files with 153 additions and 1 deletions
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue