mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
fix: close cached AsyncOpenAI clients in close_litellm_async_clients()
litellm caches the openai SDK client as an AsyncOpenAI in in_memory_llm_clients_cache (set_cached_openai_client), built on litellm's aiohttp transport, so each cached client owns an aiohttp ClientSession. close_litellm_async_clients() dispatched only on BaseLLMAIOHTTPHandler, objects with .client, and objects with .aclose(); AsyncOpenAI has none of these -- its async teardown is named .close() -- so it was never closed and its session leaked as an "Unclosed client session" error at loop teardown (one per distinct cached OpenAI client config). Widen the final fallback branch to await whichever of aclose()/close() a cached client exposes, guarded on inspect.iscoroutinefunction so sync clients (sync OpenAI / HTTPHandler) are left untouched. Additive: objects already handled by the aclose() branch keep working. Adds a regression test under tests/test_litellm/. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
1b81148f2a
commit
673d69f094
2 changed files with 67 additions and 7 deletions
|
|
@ -3,6 +3,7 @@ Utility functions for cleaning up async HTTP clients to prevent resource leaks.
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
|
||||
|
||||
async def close_litellm_async_clients():
|
||||
|
|
@ -46,13 +47,22 @@ async def close_litellm_async_clients():
|
|||
# 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
|
||||
# Handle any other cached async client. Most expose aclose(), but some
|
||||
# SDK clients -- notably the openai SDK's AsyncOpenAI / AsyncAzureOpenAI,
|
||||
# which are cached here (set_cached_openai_client) and built on litellm's
|
||||
# aiohttp transport -- name the async teardown close() instead. Accept
|
||||
# either, guarding on iscoroutinefunction so sync close() methods
|
||||
# (sync OpenAI / HTTPHandler) are left untouched.
|
||||
else:
|
||||
closer = getattr(handler, "aclose", None) or getattr(
|
||||
handler, "close", None
|
||||
)
|
||||
if closer is not None and inspect.iscoroutinefunction(closer):
|
||||
try:
|
||||
await closer()
|
||||
except Exception:
|
||||
# Silently ignore errors during cleanup
|
||||
pass
|
||||
|
||||
# Close the global base_llm_aiohttp_handler instance (issue #12443)
|
||||
# This is used by Gemini and other providers that use aiohttp
|
||||
|
|
|
|||
|
|
@ -0,0 +1,50 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test for the AsyncOpenAI cleanup gap: close_litellm_async_clients() must close
|
||||
cached clients whose async teardown is named close() (the openai SDK's
|
||||
AsyncOpenAI / AsyncAzureOpenAI), not only those exposing aclose().
|
||||
|
||||
Proposed path: tests/test_litellm/llms/custom_httpx/test_openai_session_leak.py
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.async_client_cleanup import close_litellm_async_clients
|
||||
|
||||
|
||||
class _AsyncCloseOnlyClient:
|
||||
"""Mirrors openai.AsyncOpenAI: async close(), no aclose(), no .client."""
|
||||
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _SyncClient:
|
||||
"""A sync client (close() is not a coroutine) -- must be left untouched."""
|
||||
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_closes_async_close_only_clients():
|
||||
async_client = _AsyncCloseOnlyClient()
|
||||
sync_client = _SyncClient()
|
||||
litellm.in_memory_llm_clients_cache.set_cache(
|
||||
key="test_async_close_only", value=async_client, ttl=60
|
||||
)
|
||||
litellm.in_memory_llm_clients_cache.set_cache(
|
||||
key="test_sync_close", value=sync_client, ttl=60
|
||||
)
|
||||
|
||||
await close_litellm_async_clients()
|
||||
|
||||
assert async_client.closed, "async close()-only client should be closed"
|
||||
assert not sync_client.closed, "sync client should be left untouched"
|
||||
Loading…
Add table
Reference in a new issue