scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(aiohttp): dispose recycled client sessions deterministically (#33428)

Browse source 
* fix(aiohttp): dispose recycled client sessions deterministically

LiteLLMAiohttpTransport replaced its cached aiohttp.ClientSession on
loop-mismatch, loop-inspection failure, and "Session is closed" retry
without reliably closing the previous session:

- the close task from asyncio.create_task() was never referenced, so
  it could be garbage-collected before running;
- the (RuntimeError, AttributeError) fallback branch replaced the
  session without closing it at all;
- sessions bound to a closed event loop were abandoned to the GC
  ("rely on GC"), and sessions bound to a loop running in another
  thread were closed from the wrong loop.

Replaced sessions surfaced as intermittent "Unclosed client session" /
"Unclosed connector" errors from the event-loop exception handler at
GC time.

_close_recycled_session() now covers the three lifecycles a recycled
session can be in: same-loop closes keep a strong task reference until
completion; sessions owned by a loop running elsewhere are closed on
their own loop via run_coroutine_threadsafe; sessions whose loop is
gone are disposed synchronously through the connector teardown that
aiohttp's own finalizer uses, which releases pooled connections and
silences the finalizer warnings.

Fixes #24230

* fix(aiohttp): guard threadsafe close callback against cancelled futures

---------

Co-authored-by: Anmol Jaiswal <68013660+anmolg1997@users.noreply.github.com>
This commit is contained in:
Yassin Kortam 2026-07-31 09:48:08 -07:00 • committed by GitHub
parent 416e398154
commit 16507f1174
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 416 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

123
litellm/llms/custom_httpx/aiohttp_transport.py
View file

@ -1,10 +1,11 @@
import asyncio
import concurrent.futures
import contextlib
import os
import ssl
import typing
import urllib.request
from typing import Any, Callable, Dict, Optional, Union
from typing import Any, Callable, ClassVar, Dict, Optional, Union
import aiohttp
import aiohttp.client_exceptions
@ -138,6 +139,11 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
Credit to: https://github.com/karpetrosyan/httpx-aiohttp for this implementation
"""
# Strong references to scheduled session-close tasks. A bare
# asyncio.create_task() result may be garbage-collected before it runs,
# leaving the recycled session unclosed ("Unclosed client session").
_background_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes
def __init__(
self,
client: Union[ClientSession, Callable[[], ClientSession]],
@ -164,6 +170,92 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
self._owns_session = True
return session
@classmethod
def _on_close_task_done(cls, task: "asyncio.Task[None]") -> None:
cls._background_close_tasks.discard(task)
if task.cancelled():
return
exc = task.exception()
if exc is not None:
verbose_logger.debug("Error closing recycled aiohttp session: %s", exc)
@staticmethod
def _on_threadsafe_close_done(future: "concurrent.futures.Future[None]") -> None:
if future.cancelled():
return
exc = future.exception()
if exc is not None:
verbose_logger.debug("Error closing recycled aiohttp session on its own loop: %s", exc)
@staticmethod
def _mark_connector_closed(session: ClientSession) -> None:
"""Synchronously dispose a session whose event loop is gone.
An async close can no longer run on a closed loop. BaseConnector._close
is the same synchronous teardown aiohttp's own finalizer (__del__)
uses: it is guarded for closed loops, releases pooled connections, and
flips the flags that ClientSession.closed / BaseConnector.closed read -
so no "Unclosed client session" / "Unclosed connector" warnings reach
the event-loop exception handler at garbage collection.
"""
connector = getattr(session, "_connector", None)
close_sync = getattr(connector, "_close", None)
if not callable(close_sync):
return
try:
close_sync()
except (RuntimeError, AttributeError, OSError) as e:
verbose_logger.debug("Best-effort connector close failed: %s", e)
def _close_recycled_session(self, session: ClientSession) -> None:
"""Deterministically dispose a ClientSession this transport is replacing.
Covers the three lifecycles a recycled session can be in:
- its loop is the current running loop: schedule an async close and keep
a strong reference to the task until it completes;
- its loop is still running elsewhere (e.g. another thread): hand the
close to that loop thread-safely;
- its loop is stopped or closed, or there is no running loop: fall
back to the synchronous finalizer-safe teardown.
"""
if session.closed:
return
session_loop = getattr(session, "_loop", None)
try:
current_loop: Optional[asyncio.AbstractEventLoop] = asyncio.get_running_loop()
except RuntimeError:
current_loop = None
if session_loop is not None and session_loop is not current_loop:
if not session_loop.is_closed() and session_loop.is_running():
# The session's loop is running somewhere else (e.g. another
# thread): closing from here would touch that loop's internals
# unsafely; hand the close to its own loop.
try:
future = asyncio.run_coroutine_threadsafe(session.close(), session_loop)
except RuntimeError as e: # loop shut down between the checks
verbose_logger.debug("Threadsafe session close failed: %s", e)
self._mark_connector_closed(session)
else:
future.add_done_callback(self._on_threadsafe_close_done)
return
# Foreign loop that is stopped or closed: an async close can no
# longer run there, and running it on the current loop would touch
# another loop's internals. Dispose synchronously instead.
self._mark_connector_closed(session)
return
if current_loop is None:
self._mark_connector_closed(session)
return
task = current_loop.create_task(session.close())
cls = type(self)
cls._background_close_tasks.add(task)
task.add_done_callback(cls._on_close_task_done)
def _get_valid_client_session(self) -> ClientSession:
"""
Helper to get a valid ClientSession for the current event loop.
@ -193,21 +285,25 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
# Close old session to prevent leaks
old_session = self.client
try:
if self._owns_session and not old_session.closed:
try:
asyncio.create_task(old_session.close())
except RuntimeError:
# Different event loop - can't schedule task, rely on GC
verbose_logger.debug("Old session from different loop, relying on GC")
if self._owns_session:
self._close_recycled_session(old_session)
except Exception as e:
verbose_logger.debug(f"Error closing old session: {e}")
# Create a new session in the current event loop
self.client = self._rebuild_session()
except (RuntimeError, AttributeError):
# If we can't check the loop or session is invalid, recreate it
except (RuntimeError, AttributeError) as e:
# If we can't check the loop or session is invalid, recreate it,
# but still dispose of the session being replaced.
old_session = self.client
if self._owns_session:
try:
self._close_recycled_session(old_session)
except (RuntimeError, AttributeError, OSError) as close_error:
verbose_logger.debug(f"Error closing old session: {close_error}")
self.client = self._rebuild_session()
verbose_logger.debug(f"Error checking session loop, created new session: {e}")
return self.client
@ -301,7 +397,14 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
# Handle the case where session was closed between our check and actual use
if "Session is closed" in str(e):
verbose_logger.debug(f"Session closed during request, retrying with new session: {e}")
# Force creation of a new session
# Dispose of the session that actually faulted. Do NOT read
# self.client here: a concurrent task may already have
# replaced it with a healthy session that must stay open.
# Guarded by isinstance: factory-injected sessions may be
# duck-typed test doubles without a close() coroutine.
# Read _owns_session before _rebuild_session() claims ownership.
if self._owns_session and isinstance(client_session, ClientSession):
self._close_recycled_session(client_session)
self.client = self._rebuild_session()
client_session = self.client

303
tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py
View file

@ -1,4 +1,5 @@
import asyncio
import concurrent.futures
import os
import sys
@ -827,3 +828,305 @@ async def test_stale_loop_rebuild_does_not_close_unowned_session():
shared_session._loop = running_loop
other_loop.close()
await shared_session.close()
# ---------------------------------------------------------------------------
# Recycled-session leak tests (#24230)
# ---------------------------------------------------------------------------
async def _new_session() -> aiohttp.ClientSession:
return aiohttp.ClientSession()
def _make_session_on_dead_loop() -> aiohttp.ClientSession:
"""Create a ClientSession bound to an event loop that is then closed.
Runs in a worker thread: the caller may already be inside a running
event loop, where a nested run_until_complete is forbidden.
"""
import threading
result: dict = {}
def build() -> None:
loop = asyncio.new_event_loop()
try:
result["session"] = loop.run_until_complete(_new_session())
finally:
loop.close()
thread = threading.Thread(target=build)
thread.start()
thread.join(5)
return result["session"]
def _flaky_get_running_loop_factory():
"""get_running_loop stand-in that fails once, then delegates.
Reproduces #24230: a transient loop-inspection failure sends
_get_valid_client_session into its (RuntimeError, AttributeError)
fallback branch.
"""
real_get_running_loop = asyncio.get_running_loop
calls = {"count": 0}
def flaky():
calls["count"] += 1
if calls["count"] == 1:
raise RuntimeError("simulated loop inspection failure")
return real_get_running_loop()
return flaky
@pytest.mark.asyncio
async def test_fallback_recreate_closes_previous_session():
"""
Regression test for #24230: when loop inspection fails and the fallback
branch recreates the session, the replaced session must still be closed -
not silently abandoned to the garbage collector.
"""
from unittest.mock import patch
old_session = aiohttp.ClientSession()
transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession())
transport.client = old_session
with patch(
"litellm.llms.custom_httpx.aiohttp_transport.asyncio.get_running_loop",
side_effect=_flaky_get_running_loop_factory(),
):
new_session = transport._get_valid_client_session()
try:
assert new_session is not old_session
for _ in range(3):
await asyncio.sleep(0)
assert old_session.closed, "replaced session must be closed, not leaked"
finally:
await new_session.close()
if not old_session.closed:
await old_session.close()
@pytest.mark.asyncio
async def test_replaced_session_emits_no_unclosed_warnings():
"""
Regression test for #24230: a session replaced by the fallback branch must
not surface "Unclosed client session" / "Unclosed connector" warnings when
the garbage collector finalizes it.
"""
import gc
import warnings as warnings_mod
from unittest.mock import patch
old_session = aiohttp.ClientSession()
transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession())
transport.client = old_session
with patch(
"litellm.llms.custom_httpx.aiohttp_transport.asyncio.get_running_loop",
side_effect=_flaky_get_running_loop_factory(),
):
new_session = transport._get_valid_client_session()
try:
for _ in range(3):
await asyncio.sleep(0)
del old_session
with warnings_mod.catch_warnings(record=True) as caught:
warnings_mod.simplefilter("always")
gc.collect()
unclosed = [
str(w.message)
for w in caught
if "Unclosed client session" in str(w.message) or "Unclosed connector" in str(w.message)
]
assert not unclosed, f"leaked session warnings: {unclosed}"
finally:
await new_session.close()
@pytest.mark.asyncio
async def test_dead_loop_session_closed_synchronously_on_recycle():
"""
Regression test for #24230: a session whose event loop is already closed
cannot run an async close anywhere. Recycling it must dispose of it
deterministically, the session reads closed as soon as the recycle
returns, so no finalizer warning window remains.
"""
old_session = _make_session_on_dead_loop()
transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession())
transport.client = old_session
new_session = transport._get_valid_client_session()
try:
assert new_session is not old_session
assert old_session.closed, "session from a closed loop must be disposed synchronously at recycle"
finally:
await new_session.close()
@pytest.mark.asyncio
async def test_close_task_strongly_referenced_until_done():
"""
Regression test for #24230: scheduled session-close tasks must be strongly
referenced (and pruned on completion) so they cannot be garbage-collected
before they run.
"""
old_session = aiohttp.ClientSession()
transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession())
transport._close_recycled_session(old_session)
assert LiteLLMAiohttpTransport._background_close_tasks, "close task must be strongly referenced while pending"
for _ in range(5):
await asyncio.sleep(0)
assert old_session.closed
assert not LiteLLMAiohttpTransport._background_close_tasks, "completed close tasks must be pruned from the registry"
@pytest.mark.asyncio
async def test_session_from_other_running_loop_closed_threadsafe():
"""
Regression test for #24230: a session that belongs to a loop still running
in another thread must be closed on its own loop (thread-safe), not driven
from the current loop.
"""
import threading
import time
ready = threading.Event()
holder: dict = {}
def worker() -> None:
loop = asyncio.new_event_loop()
holder["loop"] = loop
async def make() -> None:
holder["session"] = aiohttp.ClientSession()
loop.run_until_complete(make())
ready.set()
loop.run_forever()
loop.close()
thread = threading.Thread(target=worker, daemon=True)
thread.start()
assert ready.wait(5), "worker loop failed to start"
transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession())
transport.client = holder["session"]
new_session = transport._get_valid_client_session()
try:
deadline = time.monotonic() + 5
while not holder["session"].closed and time.monotonic() < deadline:
await asyncio.sleep(0.01)
assert holder["session"].closed, "foreign-loop session was never closed"
finally:
holder["loop"].call_soon_threadsafe(holder["loop"].stop)
thread.join(5)
await new_session.close()
def test_threadsafe_close_done_callback_tolerates_cancelled_future():
"""
Regression test for #24230 (review finding): when the foreign loop stops
before the handed-off close coroutine runs, asyncio cancels the
concurrent.futures.Future. The done-callback must return quietly instead
of letting future.exception() raise CancelledError (a BaseException that
escapes _invoke_callbacks and crashes the foreign loop's thread).
"""
future: "concurrent.futures.Future[None]" = concurrent.futures.Future()
future.cancel()
LiteLLMAiohttpTransport._on_threadsafe_close_done(future)
@pytest.mark.asyncio
async def test_session_closed_retry_does_not_close_concurrent_replacement():
"""
Regression test for #24230 (review finding): when the "Session is closed"
retry fires, the handler must dispose the session that actually faulted,
not self.client - a concurrent task may already have replaced self.client
with a healthy session, which must stay open.
"""
from unittest.mock import patch
faulted_session = aiohttp.ClientSession()
healthy_replacement = aiohttp.ClientSession()
transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession())
transport.client = faulted_session
calls = {"n": 0}
async def fake_make_request(*args, **kwargs):
calls["n"] += 1
if calls["n"] == 1:
# simulate a concurrent task replacing the shared session between
# the failed await and the exception handler
transport.client = healthy_replacement
raise RuntimeError("Session is closed")
raise StopAsyncIteration("stop after retry dispatch")
with patch.object(transport, "_make_aiohttp_request", side_effect=fake_make_request):
with pytest.raises(Exception):
await transport.handle_async_request(httpx.Request("GET", "http://example.com"))
try:
assert not healthy_replacement.closed, "concurrent replacement session must not be closed by the retry handler"
for _ in range(3):
await asyncio.sleep(0)
assert faulted_session.closed, "the faulted session must be disposed"
finally:
await faulted_session.close()
await healthy_replacement.close()
new_session = transport.client
if isinstance(new_session, aiohttp.ClientSession):
await new_session.close()
@pytest.mark.asyncio
async def test_stopped_loop_session_disposed_synchronously_on_recycle():
"""
Regression test for #24230 (review finding): a session whose loop is
stopped but not yet closed cannot safely run an async close on another
loop, and nothing will ever process a close handed to the stopped loop.
Recycling must dispose it synchronously, like the closed-loop case.
"""
import threading
result: dict = {}
def build() -> None:
loop = asyncio.new_event_loop()
async def make() -> None:
result["session"] = aiohttp.ClientSession()
loop.run_until_complete(make())
result["loop"] = loop # stopped, deliberately NOT closed
thread = threading.Thread(target=build)
thread.start()
thread.join(5)
old_session = result["session"]
transport = LiteLLMAiohttpTransport(client=lambda: aiohttp.ClientSession())
transport.client = old_session
new_session = transport._get_valid_client_session()
try:
assert new_session is not old_session
assert old_session.closed, "session from a stopped (not yet closed) loop must be disposed synchronously"
finally:
await new_session.close()
result["loop"].close()