mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix: prevent httpx client closed errors from cache eviction
Root cause: LLMClientCache._remove_key() was closing httpx clients when
they were evicted from cache (TTL or size-based). But other code still
held references to those clients (e.g., litellm.module_level_aclient
stored in module __dict__, in-flight requests). This caused
RuntimeError('Cannot send a request, as the client has been closed.')
for all subsequent users of those client references.
The issue manifested after ~1 hour (TTL=3600s) of uptime, as clients
were evicted from the LLMClientCache. It persisted until pod restart
because litellm.module_level_aclient in the module __dict__ was never
re-created (Python's __getattr__ is not called when the attribute
already exists in __dict__).
Fix (3 layers of defense):
1. LLMClientCache._remove_key() no longer closes clients on eviction.
Client cleanup is deferred to GC (__del__) and atexit handlers,
which only run when no references remain.
2. get_async_httpx_client() and _get_httpx_client() now check if a
cached client's underlying httpx client is closed before returning
it. If closed, they create a new one.
3. AsyncHTTPHandler.post/put/patch/delete/get catch RuntimeError with
'client has been closed' and transparently retry with a fresh
httpx.AsyncClient.
Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>
This commit is contained in:
parent
ae13a40c01
commit
a06404451d
4 changed files with 453 additions and 29 deletions
|
|
@ -9,23 +9,23 @@ from .in_memory_cache import InMemoryCache
|
|||
|
||||
class LLMClientCache(InMemoryCache):
|
||||
def _remove_key(self, key: str) -> None:
|
||||
"""Close async clients before evicting them to prevent connection pool leaks."""
|
||||
value = self.cache_dict.get(key)
|
||||
"""
|
||||
Remove the key from cache WITHOUT closing the client.
|
||||
|
||||
Closing clients on eviction is unsafe because other parts of the code
|
||||
may still hold references to the evicted client (e.g.,
|
||||
litellm.module_level_aclient stored in the module __dict__, or
|
||||
in-flight requests that obtained the client before eviction).
|
||||
|
||||
Closing such clients causes RuntimeError("Cannot send a request, as
|
||||
the client has been closed.") for all subsequent or in-flight users
|
||||
of that client reference.
|
||||
|
||||
Client cleanup is handled by:
|
||||
- AsyncHTTPHandler.__del__ / HTTPHandler.__del__ (GC-triggered)
|
||||
- atexit handler registered by register_async_client_cleanup()
|
||||
"""
|
||||
super()._remove_key(key)
|
||||
if value is not None:
|
||||
close_fn = getattr(value, "aclose", None) or getattr(
|
||||
value, "close", None
|
||||
)
|
||||
if close_fn and asyncio.iscoroutinefunction(close_fn):
|
||||
try:
|
||||
asyncio.get_running_loop().create_task(close_fn())
|
||||
except RuntimeError:
|
||||
pass
|
||||
elif close_fn and callable(close_fn):
|
||||
try:
|
||||
close_fn()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def update_cache_key_with_event_loop(self, key):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -423,10 +423,21 @@ class AsyncHTTPHandler:
|
|||
params = params or {}
|
||||
params.update(HTTPHandler.extract_query_params(url))
|
||||
|
||||
response = await self.client.get(
|
||||
url, params=params, headers=headers, follow_redirects=_follow_redirects # type: ignore
|
||||
)
|
||||
return response
|
||||
try:
|
||||
response = await self.client.get(
|
||||
url, params=params, headers=headers, follow_redirects=_follow_redirects # type: ignore
|
||||
)
|
||||
return response
|
||||
except RuntimeError as e:
|
||||
if "client has been closed" not in str(e):
|
||||
raise
|
||||
self.client = self.create_client(
|
||||
timeout=self.timeout, event_hooks=self.event_hooks
|
||||
)
|
||||
response = await self.client.get(
|
||||
url, params=params, headers=headers, follow_redirects=_follow_redirects # type: ignore
|
||||
)
|
||||
return response
|
||||
|
||||
@track_llm_api_timing()
|
||||
async def post(
|
||||
|
|
@ -483,6 +494,21 @@ class AsyncHTTPHandler:
|
|||
)
|
||||
finally:
|
||||
await new_client.aclose()
|
||||
except RuntimeError as e:
|
||||
if "client has been closed" not in str(e):
|
||||
raise
|
||||
self.client = self.create_client(
|
||||
timeout=timeout, event_hooks=self.event_hooks
|
||||
)
|
||||
return await self.single_connection_post_request(
|
||||
url=url,
|
||||
client=self.client,
|
||||
data=data,
|
||||
json=json,
|
||||
params=params,
|
||||
headers=headers,
|
||||
stream=stream,
|
||||
)
|
||||
except httpx.TimeoutException as e:
|
||||
end_time = time.time()
|
||||
time_delta = round(end_time - start_time, 3)
|
||||
|
|
@ -555,6 +581,21 @@ class AsyncHTTPHandler:
|
|||
)
|
||||
finally:
|
||||
await new_client.aclose()
|
||||
except RuntimeError as e:
|
||||
if "client has been closed" not in str(e):
|
||||
raise
|
||||
self.client = self.create_client(
|
||||
timeout=timeout, event_hooks=self.event_hooks
|
||||
)
|
||||
return await self.single_connection_post_request(
|
||||
url=url,
|
||||
client=self.client,
|
||||
data=data,
|
||||
json=json,
|
||||
params=params,
|
||||
headers=headers,
|
||||
stream=stream,
|
||||
)
|
||||
except httpx.TimeoutException as e:
|
||||
headers = {}
|
||||
error_response = getattr(e, "response", None)
|
||||
|
|
@ -621,6 +662,21 @@ class AsyncHTTPHandler:
|
|||
)
|
||||
finally:
|
||||
await new_client.aclose()
|
||||
except RuntimeError as e:
|
||||
if "client has been closed" not in str(e):
|
||||
raise
|
||||
self.client = self.create_client(
|
||||
timeout=timeout, event_hooks=self.event_hooks
|
||||
)
|
||||
return await self.single_connection_post_request(
|
||||
url=url,
|
||||
client=self.client,
|
||||
data=data,
|
||||
json=json,
|
||||
params=params,
|
||||
headers=headers,
|
||||
stream=stream,
|
||||
)
|
||||
except httpx.TimeoutException as e:
|
||||
headers = {}
|
||||
error_response = getattr(e, "response", None)
|
||||
|
|
@ -687,6 +743,21 @@ class AsyncHTTPHandler:
|
|||
)
|
||||
finally:
|
||||
await new_client.aclose()
|
||||
except RuntimeError as e:
|
||||
if "client has been closed" not in str(e):
|
||||
raise
|
||||
self.client = self.create_client(
|
||||
timeout=timeout, event_hooks=self.event_hooks
|
||||
)
|
||||
return await self.single_connection_post_request(
|
||||
url=url,
|
||||
client=self.client,
|
||||
data=data,
|
||||
json=json,
|
||||
params=params,
|
||||
headers=headers,
|
||||
stream=stream,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
setattr(e, "status_code", e.response.status_code)
|
||||
if stream is True:
|
||||
|
|
@ -1231,10 +1302,10 @@ def get_async_httpx_client(
|
|||
|
||||
_cached_client = cache.get_cache(_cache_key_name)
|
||||
if _cached_client:
|
||||
return _cached_client
|
||||
if not _is_async_client_closed(_cached_client):
|
||||
return _cached_client
|
||||
|
||||
if params is not None:
|
||||
# Filter out params that are only used for cache key, not for AsyncHTTPHandler.__init__
|
||||
handler_params = {k: v for k, v in params.items() if k != "disable_aiohttp_transport"}
|
||||
handler_params["shared_session"] = shared_session
|
||||
_new_client = AsyncHTTPHandler(**handler_params)
|
||||
|
|
@ -1280,10 +1351,10 @@ def _get_httpx_client(params: Optional[dict] = None) -> HTTPHandler:
|
|||
|
||||
_cached_client = cache.get_cache(_cache_key_name)
|
||||
if _cached_client:
|
||||
return _cached_client
|
||||
if not _is_sync_client_closed(_cached_client):
|
||||
return _cached_client
|
||||
|
||||
if params is not None:
|
||||
# Filter out params that are only used for cache key, not for HTTPHandler.__init__
|
||||
handler_params = {k: v for k, v in params.items() if k != "disable_aiohttp_transport"}
|
||||
_new_client = HTTPHandler(**handler_params)
|
||||
else:
|
||||
|
|
@ -1295,3 +1366,19 @@ def _get_httpx_client(params: Optional[dict] = None) -> HTTPHandler:
|
|||
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
|
||||
)
|
||||
return _new_client
|
||||
|
||||
|
||||
def _is_async_client_closed(handler: AsyncHTTPHandler) -> bool:
|
||||
"""Check if an AsyncHTTPHandler's underlying httpx client is closed."""
|
||||
try:
|
||||
return handler.client.is_closed
|
||||
except Exception:
|
||||
return True
|
||||
|
||||
|
||||
def _is_sync_client_closed(handler: HTTPHandler) -> bool:
|
||||
"""Check if an HTTPHandler's underlying httpx client is closed."""
|
||||
try:
|
||||
return handler.client.is_closed
|
||||
except Exception:
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -132,9 +132,10 @@ async def test_disconnect_idempotent():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_eviction_calls_aclose():
|
||||
"""When an async client is evicted from LLMClientCache, its aclose()
|
||||
should be scheduled via create_task."""
|
||||
async def test_eviction_does_not_close_client():
|
||||
"""When an async client is evicted from LLMClientCache, it should NOT
|
||||
be closed because other code may still hold references to it.
|
||||
Cleanup is deferred to GC (__del__) and atexit handlers."""
|
||||
cache = LLMClientCache(max_size_in_memory=2, default_ttl=600)
|
||||
|
||||
client = AsyncMock()
|
||||
|
|
@ -145,10 +146,9 @@ async def test_eviction_calls_aclose():
|
|||
# Third insert triggers eviction of client-0
|
||||
cache.set_cache(key="trigger", value="y")
|
||||
|
||||
# Let the scheduled task run
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
assert client.aclose.await_count > 0
|
||||
assert client.aclose.await_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -0,0 +1,337 @@
|
|||
"""
|
||||
Regression tests for httpx "Cannot send a request, as the client has been closed" error.
|
||||
|
||||
Root cause: LLMClientCache._remove_key() used to close httpx clients on cache
|
||||
eviction (TTL or size-based). But other code still held references to those
|
||||
clients (e.g. litellm.module_level_aclient, in-flight requests), causing
|
||||
RuntimeError when they tried to use the now-closed client.
|
||||
|
||||
Fix:
|
||||
1. LLMClientCache._remove_key() no longer closes clients on eviction.
|
||||
2. get_async_httpx_client() / _get_httpx_client() check if a cached client
|
||||
is closed before returning it; if so, create a new one.
|
||||
3. AsyncHTTPHandler HTTP methods catch RuntimeError("client has been closed")
|
||||
and retry with a fresh underlying httpx.AsyncClient.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
import litellm
|
||||
from litellm.caching.llm_caching_handler import LLMClientCache
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
_is_async_client_closed,
|
||||
_is_sync_client_closed,
|
||||
get_async_httpx_client,
|
||||
_get_httpx_client,
|
||||
)
|
||||
|
||||
|
||||
class TestLLMClientCacheNoCloseOnEviction:
|
||||
"""Verify that LLMClientCache does NOT close clients when evicting them."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_not_close_async_client_on_ttl_eviction(self):
|
||||
"""Client evicted due to TTL expiry must not be closed."""
|
||||
cache = LLMClientCache(max_size_in_memory=10, default_ttl=1)
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.aclose = AsyncMock()
|
||||
cache.set_cache(key="test-client", value=mock_client, ttl=0.01)
|
||||
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
cache.get_cache(key="test-client")
|
||||
|
||||
await asyncio.sleep(0.05)
|
||||
assert mock_client.aclose.await_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_not_close_async_client_on_size_eviction(self):
|
||||
"""Client evicted due to cache size limit must not be closed."""
|
||||
cache = LLMClientCache(max_size_in_memory=2, default_ttl=600)
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.aclose = AsyncMock()
|
||||
|
||||
cache.set_cache(key="client-a", value=mock_client)
|
||||
cache.set_cache(key="filler", value="x")
|
||||
cache.set_cache(key="trigger", value="y")
|
||||
|
||||
await asyncio.sleep(0.05)
|
||||
assert mock_client.aclose.await_count == 0
|
||||
|
||||
def test_should_not_close_sync_client_on_eviction(self):
|
||||
"""Sync client evicted from cache must not have close() called."""
|
||||
cache = LLMClientCache(max_size_in_memory=2, default_ttl=600)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.close = MagicMock()
|
||||
|
||||
cache.set_cache(key="sync-client", value=mock_client)
|
||||
cache.set_cache(key="filler", value="x")
|
||||
cache.set_cache(key="trigger", value="y")
|
||||
|
||||
assert mock_client.close.call_count == 0
|
||||
|
||||
|
||||
class TestIsClientClosedHelpers:
|
||||
"""Verify the is_closed helper functions."""
|
||||
|
||||
def test_should_detect_closed_async_client(self):
|
||||
handler = AsyncHTTPHandler()
|
||||
assert _is_async_client_closed(handler) is False
|
||||
handler.client._state = httpx._client.ClientState.CLOSED
|
||||
assert _is_async_client_closed(handler) is True
|
||||
|
||||
def test_should_detect_closed_sync_client(self):
|
||||
handler = HTTPHandler()
|
||||
assert _is_sync_client_closed(handler) is False
|
||||
handler.client._state = httpx._client.ClientState.CLOSED
|
||||
assert _is_sync_client_closed(handler) is True
|
||||
|
||||
def test_should_return_true_for_broken_handler(self):
|
||||
handler = AsyncHTTPHandler()
|
||||
del handler.client
|
||||
assert _is_async_client_closed(handler) is True
|
||||
|
||||
|
||||
class TestGetAsyncHttpxClientClosedCheck:
|
||||
"""get_async_httpx_client should create a new client if cached one is closed."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_return_new_client_when_cached_is_closed(self):
|
||||
cache = LLMClientCache(max_size_in_memory=200, default_ttl=3600)
|
||||
original_cache = getattr(litellm, "in_memory_llm_clients_cache", None)
|
||||
|
||||
try:
|
||||
litellm.in_memory_llm_clients_cache = cache
|
||||
|
||||
client_1 = get_async_httpx_client(
|
||||
llm_provider="test_closed_provider",
|
||||
)
|
||||
|
||||
await client_1.client.aclose()
|
||||
assert client_1.client.is_closed is True
|
||||
|
||||
client_2 = get_async_httpx_client(
|
||||
llm_provider="test_closed_provider",
|
||||
)
|
||||
|
||||
assert client_2 is not client_1
|
||||
assert client_2.client.is_closed is False
|
||||
finally:
|
||||
if original_cache is not None:
|
||||
litellm.in_memory_llm_clients_cache = original_cache
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_reuse_client_when_cached_is_open(self):
|
||||
cache = LLMClientCache(max_size_in_memory=200, default_ttl=3600)
|
||||
original_cache = getattr(litellm, "in_memory_llm_clients_cache", None)
|
||||
|
||||
try:
|
||||
litellm.in_memory_llm_clients_cache = cache
|
||||
|
||||
client_1 = get_async_httpx_client(
|
||||
llm_provider="test_open_provider",
|
||||
)
|
||||
|
||||
client_2 = get_async_httpx_client(
|
||||
llm_provider="test_open_provider",
|
||||
)
|
||||
|
||||
assert client_2 is client_1
|
||||
finally:
|
||||
if original_cache is not None:
|
||||
litellm.in_memory_llm_clients_cache = original_cache
|
||||
|
||||
|
||||
class TestGetSyncHttpxClientClosedCheck:
|
||||
"""_get_httpx_client should create a new client if cached one is closed."""
|
||||
|
||||
def test_should_return_new_client_when_cached_is_closed(self):
|
||||
cache = LLMClientCache(max_size_in_memory=200, default_ttl=3600)
|
||||
original_cache = getattr(litellm, "in_memory_llm_clients_cache", None)
|
||||
|
||||
try:
|
||||
litellm.in_memory_llm_clients_cache = cache
|
||||
|
||||
client_1 = _get_httpx_client(
|
||||
params={"client_alias": "test_closed_sync"},
|
||||
)
|
||||
|
||||
client_1.client.close()
|
||||
assert client_1.client.is_closed is True
|
||||
|
||||
client_2 = _get_httpx_client(
|
||||
params={"client_alias": "test_closed_sync"},
|
||||
)
|
||||
|
||||
assert client_2 is not client_1
|
||||
assert client_2.client.is_closed is False
|
||||
finally:
|
||||
if original_cache is not None:
|
||||
litellm.in_memory_llm_clients_cache = original_cache
|
||||
|
||||
|
||||
class TestAsyncHTTPHandlerClosedClientRecovery:
|
||||
"""AsyncHTTPHandler should recover from 'client has been closed' by creating a new client."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_recover_post_when_client_closed(self):
|
||||
handler = AsyncHTTPHandler(timeout=httpx.Timeout(5.0))
|
||||
|
||||
await handler.client.aclose()
|
||||
assert handler.client.is_closed is True
|
||||
|
||||
mock_response = httpx.Response(200, text="ok")
|
||||
with patch.object(
|
||||
AsyncHTTPHandler,
|
||||
"single_connection_post_request",
|
||||
return_value=mock_response,
|
||||
) as mock_post:
|
||||
response = await handler.post(
|
||||
url="https://example.com/api",
|
||||
json={"test": True},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
mock_post.assert_called_once()
|
||||
|
||||
assert handler.client.is_closed is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_recover_get_when_client_closed(self):
|
||||
handler = AsyncHTTPHandler(timeout=httpx.Timeout(5.0))
|
||||
|
||||
await handler.client.aclose()
|
||||
assert handler.client.is_closed is True
|
||||
|
||||
mock_response = httpx.Response(200, text="ok")
|
||||
with patch.object(
|
||||
httpx.AsyncClient,
|
||||
"get",
|
||||
return_value=mock_response,
|
||||
):
|
||||
response = await handler.get(url="https://example.com/api")
|
||||
assert response.status_code == 200
|
||||
|
||||
assert handler.client.is_closed is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_recover_put_when_client_closed(self):
|
||||
handler = AsyncHTTPHandler(timeout=httpx.Timeout(5.0))
|
||||
|
||||
await handler.client.aclose()
|
||||
|
||||
mock_response = httpx.Response(200, text="ok")
|
||||
with patch.object(
|
||||
AsyncHTTPHandler,
|
||||
"single_connection_post_request",
|
||||
return_value=mock_response,
|
||||
):
|
||||
response = await handler.put(
|
||||
url="https://example.com/api",
|
||||
json={"test": True},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
assert handler.client.is_closed is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_recover_delete_when_client_closed(self):
|
||||
handler = AsyncHTTPHandler(timeout=httpx.Timeout(5.0))
|
||||
|
||||
await handler.client.aclose()
|
||||
|
||||
mock_response = httpx.Response(200, text="ok")
|
||||
with patch.object(
|
||||
AsyncHTTPHandler,
|
||||
"single_connection_post_request",
|
||||
return_value=mock_response,
|
||||
):
|
||||
response = await handler.delete(
|
||||
url="https://example.com/api",
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
assert handler.client.is_closed is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_recover_patch_when_client_closed(self):
|
||||
handler = AsyncHTTPHandler(timeout=httpx.Timeout(5.0))
|
||||
|
||||
await handler.client.aclose()
|
||||
|
||||
mock_response = httpx.Response(200, text="ok")
|
||||
with patch.object(
|
||||
AsyncHTTPHandler,
|
||||
"single_connection_post_request",
|
||||
return_value=mock_response,
|
||||
):
|
||||
response = await handler.patch(
|
||||
url="https://example.com/api",
|
||||
json={"test": True},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
assert handler.client.is_closed is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_reraise_unrelated_runtime_error(self):
|
||||
"""RuntimeErrors not related to closed client should propagate."""
|
||||
handler = AsyncHTTPHandler(timeout=httpx.Timeout(5.0))
|
||||
|
||||
with patch.object(
|
||||
httpx.AsyncClient,
|
||||
"send",
|
||||
side_effect=RuntimeError("some other error"),
|
||||
):
|
||||
with pytest.raises(RuntimeError, match="some other error"):
|
||||
await handler.post(
|
||||
url="https://example.com/api",
|
||||
json={"test": True},
|
||||
)
|
||||
|
||||
|
||||
class TestCacheTTLExpiryDoesNotBreakClients:
|
||||
"""End-to-end: simulate cache TTL expiry and verify client still works."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_not_break_reference_after_ttl_expiry(self):
|
||||
"""Simulate the exact production scenario:
|
||||
1. Client is cached and a reference is held externally
|
||||
2. Cache TTL expires, client is evicted
|
||||
3. External reference should still be usable (not closed)
|
||||
"""
|
||||
cache = LLMClientCache(max_size_in_memory=200, default_ttl=0.01)
|
||||
original_cache = getattr(litellm, "in_memory_llm_clients_cache", None)
|
||||
|
||||
try:
|
||||
litellm.in_memory_llm_clients_cache = cache
|
||||
|
||||
client = get_async_httpx_client(
|
||||
llm_provider="test_ttl_provider",
|
||||
)
|
||||
|
||||
external_ref = client
|
||||
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
_result = cache.get_cache("anything")
|
||||
|
||||
assert external_ref.client.is_closed is False, (
|
||||
"Client should NOT be closed after cache eviction"
|
||||
)
|
||||
finally:
|
||||
if original_cache is not None:
|
||||
litellm.in_memory_llm_clients_cache = original_cache
|
||||
Loading…
Add table
Reference in a new issue