From d6edd37810944752561ed49d9a8e85085f8dce81 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 27 Oct 2026 13:33:03 -0700 Subject: [PATCH] fix(aiohttp_transport): close an owned aiohttp session when its transport is collected (#48650) An OpenAI or Azure SDK client litellm built and evicted from the client cache is usually garbage collected inside the closer's 900 second grace window, so the closer's weak reference dies and the client's aiohttp session is never closed. LiteLLMAiohttpTransport.__del__ now hands an owned, still-open ClientSession to _close_recycled_session, which already covers the current-loop, foreign-loop and closed-loop cases. Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: BangerSoul --- .../llms/custom_httpx/aiohttp_transport.py | 7 ++ .../sdk/test_aiohttp_session_rebuild_wire.py | 86 ++++++++++++++++++- .../unit/caching/test_llm_caching_handler.py | 33 +++++++ .../custom_httpx/test_aiohttp_transport.py | 73 ++++++++++++++++ 4 files changed, 197 insertions(+), 2 deletions(-) diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index b8cb183c02f..0e4d97fd770 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -289,6 +289,13 @@ class LiteLLMAiohttpTransport(AiohttpTransport): cls._background_close_tasks.add(task) task.add_done_callback(cls._on_close_task_done) + def __del__(self) -> None: + session: Final = getattr(self, "client", None) + if not getattr(self, "_owns_session", False) or not isinstance(session, ClientSession): + return + with contextlib.suppress(Exception): + self._close_recycled_session(session) + def get_valid_client_session(self) -> ClientSession: """ Helper to get a valid ClientSession for the current event loop. diff --git a/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py b/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py index 1ec1f486bed..f1a18ba389a 100644 --- a/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py +++ b/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py @@ -11,11 +11,30 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Final import pytest -from pydantic import JsonValue, TypeAdapter +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter CONFIGURED_KEEPALIVE_SECONDS: Final = 1 IDLE_SECONDS: Final = 2 +EVICTED_CLIENTS: Final = 30 RESPONSES: Final = TypeAdapter(list[dict[str, JsonValue]]) +CHAT_COMPLETION: Final = { + "id": "chatcmpl-scripted", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-5.4-nano", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +class EvictionReport(BaseModel): + model_config = ConfigDict(frozen=True) + + cache_size: int + cached_clients: int + established_connections: int + unclosed_messages: tuple[str, ...] + REBUILT_SESSION_EXCHANGE: Final = textwrap.dedent( """ @@ -38,6 +57,47 @@ REBUILT_SESSION_EXCHANGE: Final = textwrap.dedent( ) +EVICTED_CLIENTS_EXCHANGE: Final = textwrap.dedent( + """ + import asyncio, gc, json, sys + import psutil + import litellm + from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport + + unclosed = [] + + def established(port: int) -> int: + connections = psutil.Process().net_connections("tcp") + return sum(1 for c in connections if c.status == "ESTABLISHED" and c.raddr and c.raddr.port == port) + + async def one_call(base_url: str, index: int) -> None: + await litellm.acompletion( + model="openai/gpt-5.4-nano", + api_key=f"sk-distinct-{index}", + api_base=base_url, + messages=[{"role": "user", "content": "hi"}], + ) + + async def main(base_url: str, port: int, evicted: int) -> None: + asyncio.get_running_loop().set_exception_handler(lambda loop, ctx: unclosed.append(ctx["message"])) + cache = litellm.in_memory_llm_clients_cache + for index in range(cache.max_size_in_memory + evicted): + await one_call(base_url, index) + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + await asyncio.sleep(0) + print(json.dumps({ + "cache_size": cache.max_size_in_memory, + "cached_clients": len(cache.cache_dict), + "established_connections": established(port), + "unclosed_messages": unclosed, + })) + + asyncio.run(main(sys.argv[1], int(sys.argv[2]), int(sys.argv[3]))) + """ +) + + class _ConnectionCountingPeer(ThreadingHTTPServer): daemon_threads = True @@ -62,7 +122,8 @@ class _ConnectionHandler(BaseHTTPRequestHandler): def do_POST(self) -> None: self.rfile.read(int(self.headers["Content-Length"])) - body: Final = json.dumps({"connection": self.connection_number}).encode() + payload: Final = CHAT_COMPLETION if self.path.endswith("/chat/completions") else {"connection": self.connection_number} + body: Final = json.dumps(payload).encode() self.send_response(200) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(body))) @@ -107,3 +168,24 @@ def test_rebuilt_shared_session_drops_idle_connection_after_configured_keepalive ) -> None: observed: Final = _rebuilt_session_exchange(connection_counting_peer) assert observed == [{"connection": 1}, {"connection": 2}], observed + + +def _evicted_clients_exchange(base_url: str) -> EvictionReport: + completed: Final = subprocess.run( + [sys.executable, "-P", "-c", EVICTED_CLIENTS_EXCHANGE, base_url, base_url.rsplit(":", 1)[1], str(EVICTED_CLIENTS)], + capture_output=True, + text=True, + timeout=300, + check=False, + ) + assert completed.returncode == 0, completed.stderr + return EvictionReport.model_validate_json(completed.stdout) + + +def test_sdk_clients_evicted_from_a_full_client_cache_leave_no_open_session_or_connection( + connection_counting_peer: str, +) -> None: + report: Final = _evicted_clients_exchange(connection_counting_peer) + assert report.cached_clients == report.cache_size, report + assert report.unclosed_messages == (), report + assert report.established_connections <= report.cache_size, report diff --git a/tests/unit/caching/test_llm_caching_handler.py b/tests/unit/caching/test_llm_caching_handler.py index 0c0fe8ca3c5..a3efd5d97aa 100644 --- a/tests/unit/caching/test_llm_caching_handler.py +++ b/tests/unit/caching/test_llm_caching_handler.py @@ -11,11 +11,14 @@ See: https://github.com/BerriAI/litellm/pull/22247 import asyncio, gc, httpx, importlib, litellm, os, threading, time, weakref import warnings +import aiohttp import pytest +from openai import AsyncOpenAI from litellm.caching.evicted_client_closer import EvictedClientCloser from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.custom_httpx.http_handler import( @@ -204,6 +207,36 @@ async def test_evicted_litellm_owned_client_is_closed_once_the_grace_window_elap assert client.closed is True +@pytest.mark.asyncio +async def test_an_evicted_owned_client_collected_inside_its_grace_window_still_closes_its_aiohttp_session(): + """ + The closer holds an evicted client weakly for the whole grace window. An SDK + process that exits, or a proxy that evicts and never calls the cache again, + collects the client first, so only its own finalizer can still close the + aiohttp session (LIT-9767). + """ + clock = _FakeClock() + cache = LLMClientCache( + max_size_in_memory=2, + evicted_client_closer=EvictedClientCloser(grace_seconds=900.0, clock=clock), + ) + session = aiohttp.ClientSession() + client = AsyncOpenAI( + api_key="sk-evicted", + http_client=httpx.AsyncClient(transport=LiteLLMAiohttpTransport(client=session)), + ) + cache.set_cache("client-key", client, litellm_owned_client=True, ttl=600) + + cache.ttl_dict = {key: 0 for key in cache.ttl_dict} + cache.expiration_heap = [(0, key) for _, key in cache.expiration_heap] + cache.evict_cache() + del client + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + + assert session.closed, "the evicted client was collected inside the grace window and left its session open" + + @pytest.mark.asyncio async def test_evicted_caller_supplied_client_is_never_closed(): """litellm does not own a client the caller passed in, so it must stay open.""" diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py index 0765371dc4a..c7d5518ef5b 100644 --- a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py @@ -1,8 +1,10 @@ import aiohttp as aiohttp_aiohttp_handler, asyncio, importlib, litellm import concurrent.futures +import gc import socket import sys import struct +import threading from typing import Final import aiohttp @@ -1310,3 +1312,74 @@ async def test_reset_after_request_is_not_replayed(monkeypatch: pytest.MonkeyPat assert len(received_requests) == 1 finally: await handler.close() + + +async def _new_session() -> aiohttp.ClientSession: + return aiohttp.ClientSession() + + +async def _until_closed(session: aiohttp.ClientSession) -> None: + while not session.closed: + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_a_collected_transport_closes_the_session_it_owns(): + """An SDK client evicted from the cache is collected long before the closer's grace window ends.""" + session = aiohttp.ClientSession() + transport = LiteLLMAiohttpTransport(client=session) + + del transport + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + + assert session.closed, "the owned session outlived the transport that was collected before closing it" + + +@pytest.mark.asyncio +async def test_a_collected_transport_leaves_a_shared_session_open(): + session = aiohttp.ClientSession() + try: + transport = LiteLLMAiohttpTransport(client=session, owns_session=False) + + del transport + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + + assert not session.closed, "a transport closed a session it does not own" + finally: + await session.close() + + +def test_a_transport_collected_after_its_loop_closed_still_closes_its_session(caplog: pytest.LogCaptureFixture): + """asyncio.run() closes the loop the session lives on before the evicted client is collected.""" + loop = asyncio.new_event_loop() + session = loop.run_until_complete(_new_session()) + loop.close() + transport = LiteLLMAiohttpTransport(client=session) + + del transport + gc.collect() + + assert session.closed, "the owned session stayed open after its transport was collected on a closed loop" + del session + gc.collect() + assert "Unclosed client session" not in caplog.text, caplog.text + + +def test_a_transport_collected_while_its_loop_runs_elsewhere_closes_its_session_on_that_loop(): + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + session = asyncio.run_coroutine_threadsafe(_new_session(), loop).result(timeout=10) + transport = LiteLLMAiohttpTransport(client=session) + try: + del transport + gc.collect() + asyncio.run_coroutine_threadsafe(_until_closed(session), loop).result(timeout=10) + + assert session.closed, "the owned session stayed open after its transport was collected from another thread" + finally: + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=10) + loop.close()