Merge ca25e3be82 into 0c98afa780

litellm/llms/custom_httpx/aiohttp_transport.py

@@ -4,6 +4,8 @@ import contextlib
 import os
 import ssl
 import sys
+import threading
+import time
 import typing
 import urllib.request
 from collections.abc import Callable, Generator
@@ -175,6 +177,11 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
     # 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
+    # A session finalized from a different running loop must stay alive until
+    # its thread-safe close future completes; the GC cycle is already being finalized.
+    _finalizer_sessions: ClassVar[set[ClientSession]] = set()  # mutable-ok: strong refs for foreign-loop finalizers
+    _FINALIZER_CLOSE_POLL_SECONDS: Final = 0.05
+    _FINALIZER_CLOSE_MAX_WAIT_SECONDS: Final = 5.0
 
     def __init__(
         self,
@@ -219,6 +226,43 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
         if exc is not None:
             verbose_logger.debug("Error closing recycled aiohttp session on its own loop: %s", exc)
 
+    @classmethod
+    def _on_finalizer_session_close_done(
+        cls, session: ClientSession, closing: object, future: "concurrent.futures.Future[None]"
+    ) -> None:
+        try:
+            if future.cancelled():
+                close_coroutine: Final = getattr(closing, "close", None)
+                if callable(close_coroutine):
+                    close_coroutine()
+                cls._mark_connector_closed(session)
+            elif future.exception() is not None:
+                cls._mark_connector_closed(session)
+        except Exception:
+            pass
+        finally:
+            cls._finalizer_sessions.discard(session)
+
+    @classmethod
+    def _watch_finalizer_session_close(
+        cls,
+        session: ClientSession,
+        session_loop: asyncio.AbstractEventLoop,
+        future: "concurrent.futures.Future[None]",
+    ) -> None:
+        deadline: Final = time.monotonic() + cls._FINALIZER_CLOSE_MAX_WAIT_SECONDS
+        while not future.done():
+            try:
+                loop_running = session_loop.is_running()
+            except Exception:
+                loop_running = False
+            if not loop_running or time.monotonic() >= deadline:
+                future.cancel()
+                cls._mark_connector_closed(session)
+                cls._finalizer_sessions.discard(session)
+                return
+            time.sleep(cls._FINALIZER_CLOSE_POLL_SECONDS)
+
     @staticmethod
     def _mark_connector_closed(session: ClientSession) -> None:
         """Synchronously dispose a session whose event loop is gone.
@@ -288,6 +332,53 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
         cls._background_close_tasks.add(task)
         task.add_done_callback(cls._on_close_task_done)
 
+    def _close_finalized_session(self, session: ClientSession) -> None:
+        """Dispose a session without letting its own finalizer race cleanup."""
+        session_loop: Final[asyncio.AbstractEventLoop | None] = getattr(session, "_loop", None)
+        try:
+            current_loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop()
+        except RuntimeError:
+            current_loop = None
+
+        if session_loop is not None and session_loop.is_running() and session_loop is not current_loop:
+            cls: Final = type(self)
+            cls._finalizer_sessions.add(session)
+            closing: Final = session.close()
+            try:
+                future: Final = asyncio.run_coroutine_threadsafe(closing, session_loop)
+            except RuntimeError as e:
+                closing.close()
+                cls._finalizer_sessions.discard(session)
+                verbose_logger.debug("Threadsafe finalizer session close failed: %s", e)
+                self._mark_connector_closed(session)
+            else:
+                future.add_done_callback(
+                    lambda completed: cls._on_finalizer_session_close_done(session, closing, completed)
+                )
+                threading.Thread(
+                    target=cls._watch_finalizer_session_close,
+                    args=(session, session_loop, future),
+                    daemon=True,
+                ).start()
+            return
+
+        # The session is on this loop, has no loop, or its loop is stopped or
+        # closed. Async close cannot run reliably during finalization; the
+        # connector's synchronous teardown flips the flags aiohttp checks.
+        self._mark_connector_closed(session)
+
+    def __del__(self) -> None:
+        """Best-effort cleanup for transports finalized outside an explicit close."""
+        try:
+            if not getattr(self, "_owns_session", False):
+                return
+            session: Final[object] = getattr(self, "client", None)
+            if isinstance(session, ClientSession) and not session.closed:
+                self._close_finalized_session(session)
+        except Exception:
+            # Finalizers must never surface errors during garbage collection or interpreter shutdown.
+            pass
+
     def _get_valid_client_session(self) -> ClientSession:
         """
         Helper to get a valid ClientSession for the current event loop.

tests/test_litellm/caching/test_llm_client_cache_e2e.py

@@ -14,7 +14,11 @@ See: https://github.com/BerriAI/litellm/pull/22247
 """
 
 import asyncio
+import gc
+import weakref
+from typing import Final
 
+import httpx
 import pytest
 
 import litellm
@@ -125,3 +129,38 @@ async def test_ttl_expired_openai_sdk_client_stays_usable():
         "'Cannot send a request, as the client has been closed' in production"
     )
     await client.close()
+
+
+@pytest.mark.asyncio
+async def test_evicted_openai_client_finalizes_aiohttp_session():
+    """Cache eviction must not leave the OpenAI path's aiohttp session unclosed."""
+    from openai import AsyncOpenAI
+
+    from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport
+    from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+
+    transport: Final = AsyncHTTPHandler._create_async_transport()
+    assert isinstance(transport, LiteLLMAiohttpTransport)
+    session_ref: Final = weakref.ref(transport._get_valid_client_session())
+    http_client: Final = httpx.AsyncClient(transport=transport)
+    client: Final = AsyncOpenAI(api_key="sk-test", http_client=http_client)
+    transport_ref: Final = weakref.ref(transport)
+    warning: Final[asyncio.Future[str]] = asyncio.get_running_loop().create_future()
+
+    def exception_handler(_loop: asyncio.AbstractEventLoop, context: dict[str, object]) -> None:
+        if not warning.done():
+            warning.set_result(str(context.get("message", "")))
+
+    asyncio.get_running_loop().set_exception_handler(exception_handler)
+
+    cache: Final = litellm.in_memory_llm_clients_cache
+    cache.set_cache("openai-aiohttp-client", client, ttl=600)
+    cache.set_cache("filler", "x", ttl=600)
+
+    del client, http_client, transport
+    gc.collect()
+    await asyncio.sleep(0.15)
+
+    assert transport_ref() is None
+    assert session_ref() is None
+    assert not warning.done(), "finalized OpenAI transport must not emit unclosed-session warnings"

tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py

@@ -1,8 +1,10 @@
 import asyncio
 import concurrent.futures
+import gc
+import queue
 import socket
 import sys
-from typing import Final
+from typing import Final, cast
 
 import aiohttp
 import aiohttp.abc
@@ -40,6 +42,22 @@ async def test_aclose_closes_owned_session():
     assert session.closed, "Owned session should be closed by transport"
 
 
+@pytest.mark.asyncio
+async def test_finalizer_does_not_close_shared_session():
+    """Finalization must preserve a session owned by the caller."""
+    session: Final = aiohttp.ClientSession()
+    transport: Final = LiteLLMAiohttpTransport(client=session, owns_session=False)
+
+    del transport
+    gc.collect()
+    await asyncio.sleep(0)
+
+    try:
+        assert not session.closed, "Finalizer must not close a shared session"
+    finally:
+        await session.close()
+
+
 @pytest.mark.asyncio
 async def test_owns_session_defaults_to_true():
     """Test that owns_session defaults to True for backwards compatibility."""
@@ -876,6 +894,90 @@ def _make_session_on_dead_loop() -> aiohttp.ClientSession:
     return result["session"]
 
 
+def test_finalizer_closes_owned_session_after_loop_closes():
+    """Finalization must synchronously dispose a session whose loop is already closed."""
+    session: Final = _make_session_on_dead_loop()
+    transport: Final = LiteLLMAiohttpTransport(client=session)
+
+    try:
+        del transport
+        gc.collect()
+        assert session.closed, "finalizer must dispose sessions bound to closed loops"
+    finally:
+        if not session.closed:
+            session._connector._close()
+
+
+def test_finalizer_close_failure_falls_back_to_sync_teardown():
+    """A failed foreign-loop close future must still dispose the session connector."""
+    session: Final = _make_session_on_dead_loop()
+    future: Final[concurrent.futures.Future[None]] = concurrent.futures.Future()
+    future.set_exception(RuntimeError("simulated close failure"))
+
+    try:
+        LiteLLMAiohttpTransport._on_finalizer_session_close_done(session, object(), future)
+        assert session.closed
+    finally:
+        if not session.closed:
+            session._connector._close()
+
+
+def test_finalizer_watcher_handles_loop_inspection_failure():
+    """A loop inspection error must not strand a finalizer session."""
+    from unittest.mock import Mock
+
+    session: Final = _make_session_on_dead_loop()
+    future: Final[concurrent.futures.Future[None]] = concurrent.futures.Future()
+    broken_loop: Final = cast(asyncio.AbstractEventLoop, Mock())
+    broken_loop.is_running.side_effect = RuntimeError("simulated loop inspection failure")
+
+    try:
+        LiteLLMAiohttpTransport._watch_finalizer_session_close(session, broken_loop, future)
+        assert future.cancelled()
+        assert session.closed
+    finally:
+        if not session.closed:
+            session._connector._close()
+
+
+@pytest.mark.asyncio
+async def test_finalizer_handoff_failure_closes_session_synchronously():
+    """A foreign-loop handoff race must close the session when scheduling fails."""
+    from unittest.mock import Mock, patch
+
+    session: Final = aiohttp.ClientSession()
+    original_loop: Final = session._loop
+    foreign_loop: Final = cast(asyncio.AbstractEventLoop, Mock())
+    foreign_loop.is_running.return_value = True
+    session._loop = foreign_loop
+    transport: Final = LiteLLMAiohttpTransport(client=session)
+
+    try:
+        with patch(  # test-quality-ok: inject handoff failure to verify finalizer fallback
+            "litellm.llms.custom_httpx.aiohttp_transport.asyncio.run_coroutine_threadsafe",
+            side_effect=RuntimeError("simulated loop shutdown"),
+        ):
+            transport._close_finalized_session(session)
+        assert session.closed
+    finally:
+        session._loop = original_loop
+        if not session.closed:
+            await session.close()
+
+
+def test_connector_close_failure_is_suppressed():
+    """Synchronous connector teardown must remain best effort in a finalizer."""
+    from unittest.mock import Mock
+
+    connector: Final = Mock()
+    connector._close.side_effect = RuntimeError("simulated connector failure")
+    session: Final = Mock()
+    session._connector = connector
+
+    LiteLLMAiohttpTransport._mark_connector_closed(cast(aiohttp.ClientSession, session))
+    connector._close.assert_called_once_with()
+
+
 def _flaky_get_running_loop_factory():
     """get_running_loop stand-in that fails once, then delegates.
 
@@ -1050,6 +1152,62 @@ async def test_session_from_other_running_loop_closed_threadsafe():
         await new_session.close()
 
 
+@pytest.mark.asyncio
+async def test_finalizer_closes_session_on_foreign_running_loop():
+    """Finalization must fall back when a foreign loop stops before close runs."""
+    import threading
+    import time
+
+    ready: Final = threading.Event()
+    blocked: Final = threading.Event()
+    release: Final = threading.Event()
+    state: Final[queue.Queue[tuple[asyncio.AbstractEventLoop, aiohttp.ClientSession]]] = queue.Queue()
+
+    def worker() -> None:
+        loop: Final = asyncio.new_event_loop()
+
+        async def make() -> None:
+            state.put((loop, aiohttp.ClientSession()))
+
+        def block_loop() -> None:
+            blocked.set()
+            release.wait(5)
+
+        loop.run_until_complete(make())
+        loop.call_soon(block_loop)
+        ready.set()
+        loop.run_forever()
+        loop.close()
+
+    thread: Final = threading.Thread(target=worker, daemon=True)
+    thread.start()
+    assert ready.wait(5), "worker loop failed to start"
+
+    loop, session = state.get(timeout=5)
+    assert blocked.wait(5), "worker loop failed to enter its blocking callback"
+    transport: Final = LiteLLMAiohttpTransport(client=session)
+
+    try:
+        transport.__del__()
+        loop.call_soon_threadsafe(loop.stop)
+        release.set()
+        deadline: Final = time.monotonic() + 5
+        while (
+            (not session.closed or LiteLLMAiohttpTransport._finalizer_sessions)
+            and time.monotonic() < deadline
+        ):
+            await asyncio.sleep(0.01)
+        assert session.closed, "foreign-loop session was never closed by finalization"
+        assert not LiteLLMAiohttpTransport._finalizer_sessions
+    finally:
+        release.set()
+        if loop.is_running():
+            loop.call_soon_threadsafe(loop.stop)
+        thread.join(5)
+        if not session.closed:
+            session._connector._close()
+
+
 def test_threadsafe_close_done_callback_tolerates_cancelled_future():
     """
     Regression test for #24230 (review finding): when the foreign loop stops