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