From 4db25169d26dc0e136272b935e77e43a93389426 Mon Sep 17 00:00:00 2001 From: Cole McIntosh <82463175+colesmcintosh@users.noreply.github.com> Date: Wed, 9 Jul 2025 10:25:17 -0600 Subject: [PATCH] Fix: Properly close aiohttp client sessions to prevent resource leaks (#12251) * Fix: Properly close aiohttp client sessions to prevent resource leaks (#12107) - Add close() method to BaseLLMAIOHTTPHandler to properly close aiohttp ClientSession - Create async_client_cleanup module with utility functions to close all cached async clients - Register automatic cleanup at exit via atexit hook - Export close_litellm_async_clients() function for manual cleanup - Add comprehensive tests to verify resource cleanup This fixes the "Unclosed client session" and "Unclosed connector" warnings when using acompletion with Gemini and other models that use aiohttp. Fixes #12107 * Fix: Remove unused import to satisfy linter * Fix: Extend cleanup to handle AsyncHTTPHandler instances used by Gemini The original implementation only cleaned up BaseLLMAIOHTTPHandler instances, but Gemini/Vertex AI providers use AsyncHTTPHandler objects which contain httpx clients with aiohttp transports. This commit extends the cleanup function to: - Handle AsyncHTTPHandler instances by accessing their internal client - Close both the aiohttp transport and httpx client - Add generic fallback for any objects with aclose method This properly fixes the resource leak warnings for all provider types. --- litellm/__init__.py | 5 + litellm/llms/custom_httpx/aiohttp_handler.py | 5 + .../llms/custom_httpx/async_client_cleanup.py | 83 +++++++++++++ tests/test_resource_cleanup.py | 116 ++++++++++++++++++ 4 files changed, 209 insertions(+) create mode 100644 litellm/llms/custom_httpx/async_client_cleanup.py create mode 100644 tests/test_resource_cleanup.py diff --git a/litellm/__init__.py b/litellm/__init__.py index e202f6dc4fb..c66bd25459d 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -68,11 +68,15 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager import httpx import dotenv +from litellm.llms.custom_httpx.async_client_cleanup import register_async_client_cleanup litellm_mode = os.getenv("LITELLM_MODE", "DEV") # "PRODUCTION", "DEV" if litellm_mode == "DEV": dotenv.load_dotenv() +# Register async client cleanup to prevent resource leaks +register_async_client_cleanup() + ################################################## if set_verbose == True: _turn_on_debug() @@ -1123,6 +1127,7 @@ from .llms.github_copilot.chat.transformation import GithubCopilotConfig from .llms.nebius.chat.transformation import NebiusConfig from .main import * # type: ignore from .integrations import * +from .llms.custom_httpx.async_client_cleanup import close_litellm_async_clients from .exceptions import ( AuthenticationError, InvalidRequestError, diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 5a1d4208656..d9fc85877c3 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -47,6 +47,11 @@ class BaseLLMAIOHTTPHandler: self.client_session = aiohttp.ClientSession() return self.client_session + async def close(self): + """Close the aiohttp client session if it exists.""" + if self.client_session and not self.client_session.closed: + await self.client_session.close() + async def _make_common_async_call( self, async_client_session: Optional[ClientSession], diff --git a/litellm/llms/custom_httpx/async_client_cleanup.py b/litellm/llms/custom_httpx/async_client_cleanup.py new file mode 100644 index 00000000000..45602576764 --- /dev/null +++ b/litellm/llms/custom_httpx/async_client_cleanup.py @@ -0,0 +1,83 @@ +""" +Utility functions for cleaning up async HTTP clients to prevent resource leaks. +""" +import asyncio + + +async def close_litellm_async_clients(): + """ + Close all cached async HTTP clients to prevent resource leaks. + + This function iterates through all cached clients in litellm's in-memory cache + and closes any aiohttp client sessions that are still open. + """ + # Import here to avoid circular import + import litellm + from litellm.llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler + + cache_dict = getattr(litellm.in_memory_llm_clients_cache, "cache_dict", {}) + + for key, handler in cache_dict.items(): + # Handle BaseLLMAIOHTTPHandler instances (aiohttp_openai provider) + if isinstance(handler, BaseLLMAIOHTTPHandler) and hasattr(handler, "close"): + try: + await handler.close() + except Exception: + # Silently ignore errors during cleanup + pass + + # Handle AsyncHTTPHandler instances (used by Gemini and other providers) + elif hasattr(handler, 'client'): + client = handler.client + # Check if the httpx client has an aiohttp transport + if hasattr(client, '_transport') and hasattr(client._transport, 'aclose'): + try: + await client._transport.aclose() + except Exception: + # Silently ignore errors during cleanup + pass + # Also close the httpx client itself + if hasattr(client, 'aclose') and not client.is_closed: + try: + await client.aclose() + except Exception: + # Silently ignore errors during cleanup + pass + + # Handle any other objects with aclose method + elif hasattr(handler, 'aclose'): + try: + await handler.aclose() + except Exception: + # Silently ignore errors during cleanup + pass + + +def register_async_client_cleanup(): + """ + Register the async client cleanup function to run at exit. + + This ensures that all async HTTP clients are properly closed when the program exits. + """ + import atexit + + def cleanup_wrapper(): + try: + loop = asyncio.get_event_loop() + if loop.is_running(): + # Schedule the cleanup coroutine + loop.create_task(close_litellm_async_clients()) + else: + # Run the cleanup coroutine + loop.run_until_complete(close_litellm_async_clients()) + except Exception: + # If we can't get an event loop or it's already closed, try creating a new one + try: + loop = asyncio.new_event_loop() + loop.run_until_complete(close_litellm_async_clients()) + loop.close() + except Exception: + # Silently ignore errors during cleanup + pass + + atexit.register(cleanup_wrapper) diff --git a/tests/test_resource_cleanup.py b/tests/test_resource_cleanup.py new file mode 100644 index 00000000000..41d56258be9 --- /dev/null +++ b/tests/test_resource_cleanup.py @@ -0,0 +1,116 @@ +""" +Test that async HTTP clients are properly cleaned up to prevent resource leaks. +Issue: https://github.com/BerriAI/litellm/issues/12107 +""" +import asyncio +import os +import warnings + +import pytest + +import litellm + + +@pytest.mark.asyncio +async def test_acompletion_resource_cleanup(): + """Test that acompletion doesn't leave unclosed client sessions.""" + # Suppress warnings to check for them later + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + # Make an async completion call + response = await litellm.acompletion( + model="gemini/gemini-2.0-flash-lite-001", + messages=[{"role": "user", "content": "Hello"}], + mock_response="Hi there! How can I help you today?", + ) + + # Check that response was received + assert ( + response.choices[0].message.content == "Hi there! How can I help you today?" + ) + + # Manually close async clients + await litellm.close_litellm_async_clients() + + # Give a small delay for any warnings to appear + await asyncio.sleep(0.1) + + # Check for resource warnings + resource_warnings = [ + warning + for warning in w + if "Unclosed" in str(warning.message) + and ( + "client session" in str(warning.message) + or "connector" in str(warning.message) + ) + ] + + # Should be no unclosed resource warnings + assert ( + len(resource_warnings) == 0 + ), f"Found unclosed resources: {[str(w.message) for w in resource_warnings]}" + + +@pytest.mark.asyncio +async def test_multiple_acompletion_calls_cleanup(): + """Test that multiple acompletion calls reuse clients and don't leak resources.""" + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + # Make multiple async completion calls + for i in range(3): + response = await litellm.acompletion( + model="gemini/gemini-2.0-flash-lite-001", + messages=[{"role": "user", "content": f"Hello {i}"}], + mock_response=f"Response {i}", + ) + assert response.choices[0].message.content == f"Response {i}" + + # Clean up + await litellm.close_litellm_async_clients() + + # Give a small delay for any warnings to appear + await asyncio.sleep(0.1) + + # Check for resource warnings + resource_warnings = [ + warning + for warning in w + if "Unclosed" in str(warning.message) + and ( + "client session" in str(warning.message) + or "connector" in str(warning.message) + ) + ] + + assert ( + len(resource_warnings) == 0 + ), f"Found unclosed resources: {[str(w.message) for w in resource_warnings]}" + + +@pytest.mark.asyncio +async def test_cleanup_function_is_safe_to_call_multiple_times(): + """Test that the cleanup function can be called multiple times safely.""" + # This should not raise any errors + await litellm.close_litellm_async_clients() + await litellm.close_litellm_async_clients() + await litellm.close_litellm_async_clients() + + # Should still work after multiple cleanups + response = await litellm.acompletion( + model="gemini/gemini-2.0-flash-lite-001", + messages=[{"role": "user", "content": "Hello"}], + mock_response="Hi!", + ) + assert response.choices[0].message.content == "Hi!" + + # Clean up again + await litellm.close_litellm_async_clients() + + +if __name__ == "__main__": + # Run the test + asyncio.run(test_acompletion_resource_cleanup()) + print("✅ All tests passed!")