scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #35870 from BerriAI/litellm_reland_evicted_client_closer

Browse source 
fix(caching): re-land evicted LLM client closing (#35492) atop self-healing handlers
This commit is contained in:
Mateo Wang 2026-08-05 12:22:12 -07:00 • committed by GitHub
commit 54f83b2614
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 1088 additions and 83 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

276
litellm/caching/evicted_client_closer.py Normal file
View file

@ -0,0 +1,276 @@
"""
Deferred close of HTTP/SDK clients that the LLM client cache has evicted.
Eviction only drops the cache's reference to a client. Every OpenAI/Azure SDK
client is a reference cycle (each resource namespace holds the client back), so
an evicted client and its pooled TCP connections survive until a generational
collection runs, which under load is thousands of requests later.
Closing at eviction time is not an option: a request that was handed the client
just before it was evicted is still using it, and closing it underneath that
request raises ``RuntimeError: Cannot send a request, as the client has been
closed.``
So an evicted client is closed once two conditions hold. A grace window must
have passed since its eviction, which covers a request that holds the client
but is momentarily not on the wire, and the client must report no connection in
flight. The second condition is what keeps the first honest: a request may run
for ``litellm.request_timeout`` seconds, 6000 by default, and a streaming
response is bounded only by how long the upstream keeps sending, so no deadline
on its own can promise that a request has finished.
Only clients litellm itself created are closed; a client the caller supplied is
left alone because litellm does not own its lifecycle.
A client that closes synchronously is closed from wherever the cache is next
used. One whose close is a coroutine needs the event loop it was evicted on, so
it waits for a call from that loop rather than having work scheduled onto a loop
it does not belong to. Queued clients are therefore bucketed by what it takes to
close them, and each bucket is ordered by deadline, so a reap walks the entries
that are due rather than the whole queue.
The queue holds its clients weakly, so waiting out a grace window never keeps
alive anything the collector would have reclaimed first.
"""
import asyncio
import contextlib
import inspect
import threading
import time
import weakref
from collections import deque
from collections.abc import Awaitable, Callable, Iterator
from dataclasses import dataclass, replace
from typing import Final
from litellm.constants import (
EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS,
EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING,
)
_CLOSABLE_ANYWHERE: Final = "closable-anywhere"
_CLOSABLE_ON_ANY_LOOP: Final = "closable-on-any-loop"
_BucketKey = str | int
@dataclass(frozen=True, slots=True)
class _PendingClose:
"""A queued close.
The client is held weakly, so queueing one never keeps alive anything the
collector would otherwise have reclaimed first.
``needs_loop`` is set for a client whose close is a coroutine; those can only
be closed from the event loop they were evicted on, recorded in ``loop_id``.
A client that closes synchronously carries neither constraint.
"""
client_ref: "weakref.ref[object]"
loop_id: int | None
needs_loop: bool
close_after: float
def _bucket_key(pending: _PendingClose) -> _BucketKey:
"""Which reaps can close this entry: any at all, any running a loop, or one loop's."""
if not pending.needs_loop:
return _CLOSABLE_ANYWHERE
if pending.loop_id is None:
return _CLOSABLE_ON_ANY_LOOP
return pending.loop_id
def _running_loop_id() -> int | None:
try:
return id(asyncio.get_running_loop())
except RuntimeError:
return None
def _close_function(client: object) -> Callable[[], object] | None:
close_fn: Final[Callable[[], object] | None] = getattr(client, "aclose", None) or getattr(client, "close", None)
return close_fn
def _transport_of(client: object) -> object:
"""The httpx transport behind an SDK wrapper, a litellm handler, or a bare client."""
for holder in (getattr(client, "_client", None), getattr(client, "client", None), client):
transport: object = getattr(holder, "_transport", None)
if transport is not None:
return transport
return None
def _connection_is_idle(connection: object) -> bool:
"""A pooled connection is idle unless it is servicing a request."""
is_idle: Final[object] = getattr(connection, "is_idle", None)
return bool(is_idle()) if callable(is_idle) else True
def _pool_has_busy_connection(transport: object) -> bool | None:
"""Whether the httpcore pool behind the transport is servicing a request.
``None`` when there is no such pool, so the caller can ask the other backend.
"""
pooled: Final[object] = getattr(getattr(transport, "_pool", None), "connections", None)
if not isinstance(pooled, (list, tuple)):
return None
return any(
not _connection_is_idle(connection) # pyright: ignore[reportUnknownArgumentType] # untyped pool list
for connection in pooled # pyright: ignore[reportUnknownVariableType] # untyped pool list
)
def _has_connection_in_flight(client: object) -> bool:
"""Whether the client is servicing a request right now.
Both connection backends litellm uses already account for the connections
they have handed out, so this reads the client's own lease accounting rather
than inferring it from elapsed time: httpcore reports a non-idle connection
for the whole of a response including a stream, and aiohttp holds the
connection in ``_acquired`` over the same span.
A client that cannot answer is reported as idle, which leaves the grace
window as the only guard, exactly as it was before this check existed.
"""
try:
transport: Final = _transport_of(client)
pooled_busy: Final = _pool_has_busy_connection(transport)
if pooled_busy is not None:
return pooled_busy
session: Final[object] = getattr(transport, "client", None)
return bool(getattr(getattr(session, "connector", None), "_acquired", None))
except Exception: # noqa: BLE001 - a client that cannot report its state is treated as idle
return False
async def _close_quietly(closing: Awaitable[object]) -> None:
with contextlib.suppress(Exception):
await closing
class EvictedClientCloser:
"""Closes evicted, litellm-owned clients once they are idle and out of grace."""
def __init__(
self,
grace_seconds: float = EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS,
max_pending: int = EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING,
clock: Callable[[], float] = time.monotonic,
) -> None:
self._grace_seconds = grace_seconds
self._max_pending = max_pending
self._clock = clock
self._owned: weakref.WeakSet[object] = weakref.WeakSet()
self._buckets: dict[_BucketKey, deque[_PendingClose]] = {} # mutable-ok: deadline-ordered queues
self._pending_count = 0
self._queue_lock = threading.Lock() # the cache is reachable from every worker thread's loop
self._close_tasks: set[asyncio.Task[None]] = set() # mutable-ok: strong refs to running closes
def mark_owned(self, client: object) -> None:
"""Record that litellm created this client, so it may be closed on eviction."""
try:
self._owned.add(client)
except TypeError:
pass # values that cannot be weak-referenced are never litellm clients
def _is_owned(self, client: object) -> bool:
try:
return client in self._owned
except TypeError:
return False # unhashable values are never litellm clients
def schedule(self, client: object) -> None:
"""Queue an evicted client for closing once it is idle and out of grace.
Past ``max_pending`` the client is left to the collector instead, so a
workload that churns the cache cannot grow this queue without bound.
Every queued entry comes due within one grace window, so the capacity it
occupies is returned within that window rather than held.
"""
if client is None or not self._is_owned(client):
return
close_fn: Final = _close_function(client)
if close_fn is None:
return
if self._pending_count >= self._max_pending:
return
self._enqueue(
_PendingClose(
client_ref=weakref.ref(client),
loop_id=_running_loop_id(),
needs_loop=inspect.iscoroutinefunction(close_fn),
close_after=self._clock() + self._grace_seconds,
)
)
def reap(self) -> None:
"""Close every queued client that is due, idle, and closable from here.
Called from the cache's read path, so the empty-queue exit comes first and
the work done past it is proportional to what is due, not to the queue.
"""
if not self._pending_count:
return
now: Final = self._clock()
for pending in self._take_due(_running_loop_id(), now):
client = pending.client_ref()
if client is None:
continue
if _has_connection_in_flight(client):
self._enqueue(replace(pending, close_after=now + self._grace_seconds))
continue
self._close(client)
@property
def pending_count(self) -> int:
return self._pending_count
def _enqueue(self, pending: _PendingClose) -> None:
"""Append to the entry's bucket, dropping any dead entries it queues behind.
Deadlines only ever move forward, so appending keeps each bucket ordered
by deadline, and entries whose client the collector already took sit at
the front rather than having to be searched for.
"""
with self._queue_lock:
bucket: Final = self._buckets.setdefault(_bucket_key(pending), deque()) # mutable-ok: FIFO by design
while bucket and bucket[0].client_ref() is None:
bucket.popleft()
self._pending_count -= 1
bucket.append(pending)
self._pending_count += 1
def _take_due(self, loop_id: int | None, now: float) -> tuple[_PendingClose, ...]:
buckets = (_CLOSABLE_ANYWHERE,) if loop_id is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_id)
with self._queue_lock:
return tuple(pending for key in buckets for pending in self._drain_locked(key, now))
def _drain_locked(self, key: _BucketKey, now: float) -> Iterator[_PendingClose]:
bucket: Final = self._buckets.get(key)
if bucket is None:
return
while bucket and bucket[0].close_after <= now:
self._pending_count -= 1
yield bucket.popleft()
if not bucket:
del self._buckets[key]
def _close(self, client: object) -> None:
close_fn: Final = _close_function(client)
if close_fn is None:
return
try:
closing: Final = close_fn()
except Exception: # noqa: BLE001 - a discarded client's close must never surface to callers
return
if not inspect.isawaitable(closing):
return
task: Final = asyncio.get_running_loop().create_task(_close_quietly(closing))
self._close_tasks.add(task)
task.add_done_callback(self._close_tasks.discard)
default_evicted_client_closer: Final = EvictedClientCloser()

45
litellm/caching/llm_caching_handler.py
View file

@ -5,21 +5,44 @@ Add the event loop to the cache key, to prevent event loop closed errors.
import asyncio
from typing import Final
from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer
from .in_memory_cache import InMemoryCache
class LLMClientCache(InMemoryCache):
"""Cache for LLM HTTP clients (OpenAI, Azure, httpx, etc.).
IMPORTANT: This cache intentionally does NOT close clients on eviction.
Evicted clients may still be in use by in-flight requests. Closing them
eagerly causes ``RuntimeError: Cannot send a request, as the client has
been closed.`` errors in production after the TTL (1 hour) expires.
An evicted client is never closed on the spot: a request handed the client
just before eviction is still using it, and closing it there raises
``RuntimeError: Cannot send a request, as the client has been closed.``
Clients that are no longer referenced will be garbage-collected normally.
For explicit shutdown cleanup, use ``close_litellm_async_clients()``.
Nor can eviction be left to rely on garbage collection. The SDK clients are
reference cycles, so an evicted client and its open TCP connections survive
until a generational collection runs. Instead a client litellm created is
handed to ``EvictedClientCloser``, which closes it once a grace window has
passed. Clients the caller supplied are left untouched.
"""
def __init__(
self,
max_size_in_memory: int | None = 200,
default_ttl: int | None = 600,
max_size_per_item: int | None = 1024,
evicted_client_closer: EvictedClientCloser | None = None,
) -> None:
super().__init__(
max_size_in_memory=max_size_in_memory,
default_ttl=default_ttl,
max_size_per_item=max_size_per_item,
)
self.evicted_client_closer = evicted_client_closer or default_evicted_client_closer
def _remove_key(self, key: str) -> None:
evicted: Final[object] = self.cache_dict.get(key)
super()._remove_key(key)
self.evicted_client_closer.schedule(evicted)
self.evicted_client_closer.reap()
def update_cache_key_with_event_loop(self, key):
"""
Add the event loop to the cache key, to prevent event loop closed errors.
@ -32,16 +55,22 @@ class LLMClientCache(InMemoryCache):
except RuntimeError: # handle no current running event loop
return key
def set_cache(self, key, value, **kwargs):
def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
"""``litellm_owned_client`` marks a client litellm built, so it may be closed once evicted."""
if litellm_owned_client:
self.evicted_client_closer.mark_owned(value)
key = self.update_cache_key_with_event_loop(key)
return super().set_cache(key, value, **kwargs)
async def async_set_cache(self, key, value, **kwargs):
async def async_set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
if litellm_owned_client:
self.evicted_client_closer.mark_owned(value)
key = self.update_cache_key_with_event_loop(key)
return await super().async_set_cache(key, value, **kwargs)
def get_cache(self, key, **kwargs):
key = self.update_cache_key_with_event_loop(key)
self.evicted_client_closer.reap()
return super().get_cache(key, **kwargs)

10
litellm/constants.py
View file

@ -197,6 +197,16 @@ RUNWAYML_POLLING_TIMEOUT = int(os.getenv("RUNWAYML_POLLING_TIMEOUT", 600)) # 10
########## Networking constants ##############################################################
_DEFAULT_TTL_FOR_HTTPX_CLIENTS: Final = 3600 # 1 hour, re-use the same httpx client for 1 hour
# The earliest an evicted, litellm-created client may be closed. A request handed the
# client just before eviction is still using it, so nothing is closed inside this window;
# past it, the client is closed once it reports no connection in flight.
EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS: Final = 900
# How many evicted clients may be queued for closing at once. Past this, an evicted client
# is left to the collector rather than letting a cache-churning workload grow the queue
# without bound. Each queued entry is ~100 bytes and comes due within one grace window.
EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING: Final = 10_000
# Aiohttp connection pooling - prevents memory leaks from unbounded connection growth
# Set to 0 for unlimited (not recommended for production)
AIOHTTP_CONNECTOR_LIMIT: Final = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 1000))

151
litellm/llms/azure/common_utils.py
View file

@ -427,88 +427,95 @@ class BaseAzureLLM(BaseOpenAILLM):
f"|azure_password={hashlib.sha256(_azure_password.encode()).hexdigest() if isinstance(_azure_password, str) else None}"
f"|azure_scope={_lp.get('azure_scope')}"
)
if client is None:
cached_client: Final = self.get_cached_openai_client(
client_initialization_params=client_initialization_params,
client_type="azure",
)
if cached_client:
if isinstance(cached_client, (AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI)):
return cached_client
azure_client_params: Final = self.initialize_azure_sdk_client(
litellm_params=litellm_params or {},
api_key=api_key,
api_base=api_base,
model_name=model,
api_version=api_version,
is_async=_is_async,
)
# For Azure v1 API, use standard OpenAI client instead of AzureOpenAI
# See: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#api-specs
if self._is_azure_v1_api_version(api_version):
# Extract only params that OpenAI client accepts
# Always use /openai/v1/ regardless of whether user passed "v1", "latest", or "preview"
# The OpenAI client accepts a callable for `api_key` and re-invokes it
# on every request (via `_refresh_api_key`), so passing
# `azure_ad_token_provider` directly preserves Azure AD token refresh
# behavior that the regular AzureOpenAI client provides.
v1_api_key: str | Callable[[], Any] | None = (
azure_client_params.get("api_key")
or azure_client_params.get("azure_ad_token_provider")
or azure_client_params.get("azure_ad_token")
)
if _is_async is True and callable(v1_api_key):
# AsyncOpenAI expects an async provider; wrap the sync provider
# returned by azure-identity. Offload to a thread so a token
# refresh (blocking HTTP call to AAD on cache miss) does not
# stall the event loop.
_sync_provider: Final = v1_api_key
async def _async_v1_api_key() -> str:
return await asyncio.to_thread(_sync_provider)
v1_api_key = _async_v1_api_key
v1_params: Final[dict[str, Any]] = {
"api_key": v1_api_key,
"base_url": f"{api_base}/openai/v1/",
}
if "timeout" in azure_client_params:
v1_params["timeout"] = azure_client_params["timeout"]
if "max_retries" in azure_client_params:
v1_params["max_retries"] = azure_client_params["max_retries"]
if "http_client" in azure_client_params:
v1_params["http_client"] = azure_client_params["http_client"]
verbose_logger.debug("Using Azure v1 API with base_url: %s", v1_params["base_url"])
if _is_async is True:
openai_client = AsyncOpenAI(**v1_params)
else:
openai_client = OpenAI(**v1_params)
else:
# Traditional Azure API uses AzureOpenAI client
if _is_async is True:
openai_client = AsyncAzureOpenAI(**azure_client_params)
else:
openai_client = AzureOpenAI(**azure_client_params)
else:
openai_client = client
if client is not None:
if (
api_version is not None
and isinstance(openai_client, (AzureOpenAI, AsyncAzureOpenAI))
and isinstance(openai_client._custom_query, dict)
and isinstance(client, (AzureOpenAI, AsyncAzureOpenAI))
and isinstance(client._custom_query, dict)
):
# set api_version to version passed by user
openai_client._custom_query.setdefault("api-version", api_version)
client._custom_query.setdefault("api-version", api_version)
self.set_cached_openai_client(
openai_client=client,
client_initialization_params=client_initialization_params,
client_type="azure",
litellm_owned_client=False,
)
return client
cached_client: Final = self.get_cached_openai_client(
client_initialization_params=client_initialization_params,
client_type="azure",
)
if cached_client:
if isinstance(cached_client, (AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI)):
return cached_client
azure_client_params: Final = self.initialize_azure_sdk_client(
litellm_params=litellm_params or {},
api_key=api_key,
api_base=api_base,
model_name=model,
api_version=api_version,
is_async=_is_async,
)
# For Azure v1 API, use standard OpenAI client instead of AzureOpenAI
# See: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#api-specs
if self._is_azure_v1_api_version(api_version):
# Extract only params that OpenAI client accepts
# Always use /openai/v1/ regardless of whether user passed "v1", "latest", or "preview"
# The OpenAI client accepts a callable for `api_key` and re-invokes it
# on every request (via `_refresh_api_key`), so passing
# `azure_ad_token_provider` directly preserves Azure AD token refresh
# behavior that the regular AzureOpenAI client provides.
v1_api_key: str | Callable[[], Any] | None = (
azure_client_params.get("api_key")
or azure_client_params.get("azure_ad_token_provider")
or azure_client_params.get("azure_ad_token")
)
if _is_async is True and callable(v1_api_key):
# AsyncOpenAI expects an async provider; wrap the sync provider
# returned by azure-identity. Offload to a thread so a token
# refresh (blocking HTTP call to AAD on cache miss) does not
# stall the event loop.
_sync_provider: Final = v1_api_key
async def _async_v1_api_key() -> str:
return await asyncio.to_thread(_sync_provider)
v1_api_key = _async_v1_api_key
v1_params: Final[dict[str, Any]] = {
"api_key": v1_api_key,
"base_url": f"{api_base}/openai/v1/",
}
if "timeout" in azure_client_params:
v1_params["timeout"] = azure_client_params["timeout"]
if "max_retries" in azure_client_params:
v1_params["max_retries"] = azure_client_params["max_retries"]
if "http_client" in azure_client_params:
v1_params["http_client"] = azure_client_params["http_client"]
verbose_logger.debug("Using Azure v1 API with base_url: %s", v1_params["base_url"])
if _is_async is True:
openai_client = AsyncOpenAI(**v1_params)
else:
openai_client = OpenAI(**v1_params)
else:
# Traditional Azure API uses AzureOpenAI client
if _is_async is True:
openai_client = AsyncAzureOpenAI(**azure_client_params)
else:
openai_client = AzureOpenAI(**azure_client_params)
# save client in-memory cache
self.set_cached_openai_client(
openai_client=openai_client,
client_initialization_params=client_initialization_params,
client_type="azure",
litellm_owned_client=self.owns_wrapped_http_client(azure_client_params.get("http_client")),
)
return openai_client

2
litellm/llms/custom_httpx/http_handler.py
View file

@ -1441,6 +1441,7 @@ def get_async_httpx_client(
key=_cache_key_name,
value=_new_client,
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
litellm_owned_client=True,
)
return _new_client
@ -1486,5 +1487,6 @@ def _get_httpx_client(params: dict | None = None) -> HTTPHandler:
key=_cache_key_name,
value=_new_client,
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
litellm_owned_client=True,
)
return _new_client

23
litellm/llms/openai/common_utils.py
View file

@ -128,13 +128,33 @@ class BaseOpenAILLM:
_cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(_cache_key)
return _cached_client
@staticmethod
def owns_wrapped_http_client(http_client: httpx.Client | httpx.AsyncClient | None) -> bool:
"""Whether litellm may close an SDK client built around ``http_client``.
``_get_async_http_client`` / ``_get_sync_http_client`` hand back
``litellm.aclient_session`` / ``litellm.client_session`` when the caller
configured one. The SDK's ``close()`` closes whatever http client it was
given, so an SDK client wrapping one of those shared sessions must never be
closed on eviction; the caller goes on using the session. ``None`` means the
SDK built its own http client, which litellm does own.
"""
if http_client is None:
return True
return http_client is not litellm.aclient_session and http_client is not litellm.client_session
@staticmethod
def set_cached_openai_client(
openai_client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI,
client_type: Literal["openai", "azure"],
client_initialization_params: dict,
litellm_owned_client: bool = False,
):
"""Stores the OpenAI client in the in-memory cache for _DEFAULT_TTL_FOR_HTTPX_CLIENTS SECONDS"""
"""Stores the OpenAI client in the in-memory cache for _DEFAULT_TTL_FOR_HTTPX_CLIENTS SECONDS
``litellm_owned_client`` says litellm built this client, so the cache may close it once it
is evicted. A client the caller supplied stays open, since litellm does not own it.
"""
_cache_key: Final = BaseOpenAILLM.get_openai_client_cache_key(
client_initialization_params=client_initialization_params,
client_type=client_type,
@ -143,6 +163,7 @@ class BaseOpenAILLM:
key=_cache_key,
value=openai_client,
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
litellm_owned_client=litellm_owned_client,
)
@staticmethod

10
litellm/llms/openai/openai.py
View file

@ -360,11 +360,16 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
if cached_client:
if isinstance(cached_client, OpenAI) or isinstance(cached_client, AsyncOpenAI):
return cached_client
http_client: Final[httpx.Client | httpx.AsyncClient | None] = (
OpenAIChatCompletion._get_async_http_client(shared_session=shared_session)
if is_async
else OpenAIChatCompletion._get_sync_http_client()
)
if is_async:
_new_client: OpenAI | AsyncOpenAI = AsyncOpenAI(
api_key=api_key,
base_url=api_base,
http_client=OpenAIChatCompletion._get_async_http_client(shared_session=shared_session),
http_client=http_client,
timeout=timeout,
max_retries=max_retries,
organization=organization,
@ -373,7 +378,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
_new_client = OpenAI(
api_key=api_key,
base_url=api_base,
http_client=OpenAIChatCompletion._get_sync_http_client(),
http_client=http_client,
timeout=timeout,
max_retries=max_retries,
organization=organization,
@ -384,6 +389,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
openai_client=_new_client,
client_initialization_params=client_initialization_params,
client_type="openai",
litellm_owned_client=self.owns_wrapped_http_client(http_client),
)
return _new_client

411
tests/test_litellm/caching/test_evicted_client_closer.py Normal file
View file

@ -0,0 +1,411 @@
"""
Tests for EvictedClientCloser.
An evicted client must stay open long enough for a request that already holds it
to finish, and must then actually be closed, otherwise its connection pool is
retained until a generational collection runs. A client the caller supplied is
never closed, because litellm does not own its lifecycle.
"""
import asyncio
import gc
import weakref
import httpx
import pytest
from litellm.caching.evicted_client_closer import EvictedClientCloser
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
class FakeClock:
"""Hand-advanced monotonic clock, so grace windows need no real waiting."""
def __init__(self) -> None:
self.now = 1000.0
def __call__(self) -> float:
return self.now
def advance(self, seconds: float) -> None:
self.now += seconds
class AsyncClient:
def __init__(self) -> None:
self.closed = False
async def close(self) -> None:
self.closed = True
class SyncClient:
def __init__(self) -> None:
self.closed = False
def close(self) -> None:
self.closed = True
class CountingDeadline(float):
"""A clock reading that tallies every deadline comparison made against it.
Deadline comparisons are the work a reap does, so counting them says whether
that work tracks the entries that are due or the size of the whole queue.
"""
comparisons = 0
def __add__(self, other: float) -> "CountingDeadline":
return CountingDeadline(float(self) + other)
def __le__(self, other: float) -> bool:
CountingDeadline.comparisons += 1
return float(self) <= float(other)
def __gt__(self, other: float) -> bool:
CountingDeadline.comparisons += 1
return float(self) > float(other)
def make_closer(clock: FakeClock, grace_seconds: float = 60.0) -> EvictedClientCloser:
return EvictedClientCloser(grace_seconds=grace_seconds, clock=clock)
async def _trickling_upstream(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
"""Serves a chunked body slowly, so a request stays on the wire long enough to observe."""
await reader.read(4096)
writer.write(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n")
await writer.drain()
for _ in range(6):
writer.write(b"5\r\nhello\r\n")
await writer.drain()
await asyncio.sleep(0.1)
writer.write(b"0\r\n\r\n")
await writer.drain()
@pytest.mark.asyncio
async def test_owned_client_is_closed_once_the_grace_window_elapses():
clock = FakeClock()
closer = make_closer(clock)
client = AsyncClient()
closer.mark_owned(client)
closer.schedule(client)
clock.advance(61.0)
closer.reap()
await asyncio.sleep(0.05)
assert client.closed is True
assert closer.pending_count == 0
@pytest.mark.asyncio
async def test_owned_client_stays_open_inside_the_grace_window():
"""A request handed the client just before eviction is still using it."""
clock = FakeClock()
closer = make_closer(clock)
client = AsyncClient()
closer.mark_owned(client)
closer.schedule(client)
clock.advance(59.0)
closer.reap()
await asyncio.sleep(0.05)
assert client.closed is False
assert closer.pending_count == 1
@pytest.mark.asyncio
async def test_caller_supplied_client_is_never_closed():
clock = FakeClock()
closer = make_closer(clock)
client = AsyncClient()
closer.schedule(client)
clock.advance(3600.0)
closer.reap()
await asyncio.sleep(0.05)
assert client.closed is False
assert closer.pending_count == 0
@pytest.mark.asyncio
async def test_sync_client_is_closed_once_the_grace_window_elapses():
clock = FakeClock()
closer = make_closer(clock)
client = SyncClient()
closer.mark_owned(client)
closer.schedule(client)
clock.advance(61.0)
closer.reap()
assert client.closed is True
@pytest.mark.asyncio
async def test_a_failing_close_does_not_propagate_or_block_the_others():
class ExplodingClient:
async def close(self) -> None:
raise RuntimeError("connection already gone")
clock = FakeClock()
closer = make_closer(clock)
exploding, healthy = ExplodingClient(), AsyncClient()
for client in (exploding, healthy):
closer.mark_owned(client)
closer.schedule(client)
clock.advance(61.0)
closer.reap()
await asyncio.sleep(0.05)
assert healthy.closed is True
@pytest.mark.asyncio
async def test_an_unhashable_cached_value_does_not_break_eviction():
"""The cache holds arbitrary values; an ownership test must never raise on one."""
class Unhashable:
__hash__ = None # pyright: ignore[reportAssignmentType] # unhashable by construction
clock = FakeClock()
closer = make_closer(clock)
closer.mark_owned(Unhashable())
closer.schedule(Unhashable())
assert closer.pending_count == 0
@pytest.mark.asyncio
async def test_values_with_nothing_to_close_are_never_queued():
"""The cache holds plain values too; those have nothing to reclaim."""
class NotAClient:
pass
clock = FakeClock()
closer = make_closer(clock)
value = NotAClient()
closer.mark_owned(value)
closer.schedule(value)
assert closer.pending_count == 0
@pytest.mark.asyncio
async def test_a_queued_client_is_not_kept_alive_by_the_queue():
"""Waiting out a grace window must not retain what the collector would free first."""
clock = FakeClock()
closer = make_closer(clock)
client = AsyncClient()
gone = weakref.ref(client)
closer.mark_owned(client)
closer.schedule(client)
del client
gc.collect()
assert gone() is None, "the pending queue is holding the client alive"
clock.advance(61.0)
closer.reap()
assert closer.pending_count == 0
def test_sync_client_evicted_outside_an_event_loop_is_still_closed():
"""The sync httpx handler is cached and evicted from call sites with no loop."""
clock = FakeClock()
closer = make_closer(clock)
client = SyncClient()
closer.mark_owned(client)
closer.schedule(client)
assert closer.pending_count == 1
clock.advance(61.0)
closer.reap()
assert client.closed is True
assert closer.pending_count == 0
@pytest.mark.asyncio
async def test_an_async_client_waits_for_a_loop_rather_than_being_dropped():
clock = FakeClock()
closer = make_closer(clock)
client = AsyncClient()
closer.mark_owned(client)
def schedule_outside_a_loop() -> None:
closer.schedule(client)
clock.advance(61.0)
closer.reap()
await asyncio.to_thread(schedule_outside_a_loop)
assert client.closed is False, "no loop was running, so it could not have been closed"
assert closer.pending_count == 1
closer.reap()
await asyncio.sleep(0.05)
assert client.closed is True
@pytest.mark.asyncio
async def test_a_client_evicted_on_another_event_loop_is_left_alone():
"""Closing a client bound to a different loop would schedule work on that loop."""
clock = FakeClock()
closer = make_closer(clock)
client = AsyncClient()
closer.mark_owned(client)
def schedule_on_its_own_loop() -> None:
asyncio.run(_schedule())
async def _schedule() -> None:
closer.schedule(client)
await asyncio.to_thread(schedule_on_its_own_loop)
assert closer.pending_count == 1
clock.advance(61.0)
closer.reap()
await asyncio.sleep(0.05)
assert client.closed is False
assert closer.pending_count == 1
@pytest.mark.asyncio
async def test_a_client_serving_a_request_is_not_closed_when_its_grace_window_ends():
"""The grace window on its own cannot promise that a request has finished.
``litellm.request_timeout`` defaults to 6000 seconds and a streaming response
is bounded only by how long the upstream keeps sending, so a client past its
deadline is closed only once its own pool reports nothing in flight.
"""
server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
clock = FakeClock()
closer = make_closer(clock)
client = httpx.AsyncClient()
closer.mark_owned(client)
closer.schedule(client)
async def read_the_stream() -> int:
received = 0
async with client.stream("GET", f"http://127.0.0.1:{port}/") as response:
async for chunk in response.aiter_bytes():
received += len(chunk)
return received
streaming = asyncio.create_task(read_the_stream())
await asyncio.sleep(0.25) # the request is on the wire
clock.advance(3600.0) # and its grace window is long gone
closer.reap()
await asyncio.sleep(0.05)
assert client.is_closed is False, "closed a client that was serving a request"
assert await streaming > 0, "the in-flight request did not survive the reap"
clock.advance(3600.0)
closer.reap()
await asyncio.sleep(0.05)
assert client.is_closed is True, "an idle client past its grace window must be closed"
assert closer.pending_count == 0
server.close()
@pytest.mark.asyncio
async def test_the_aiohttp_backed_handler_is_not_closed_mid_request():
"""The default async path is aiohttp-backed, whose pool accounts for its own leases."""
server = await asyncio.start_server(_trickling_upstream, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
clock = FakeClock()
closer = make_closer(clock)
handler = AsyncHTTPHandler()
held_client = handler.client
closer.mark_owned(handler)
closer.schedule(handler)
request = asyncio.create_task(handler.get(f"http://127.0.0.1:{port}/"))
await asyncio.sleep(0.25)
clock.advance(3600.0)
closer.reap()
await asyncio.sleep(0.05)
assert held_client.is_closed is False, "closed a handler that was serving a request"
assert (await request).status_code == 200
clock.advance(3600.0)
closer.reap()
await asyncio.sleep(0.05)
assert held_client.is_closed is True
assert handler.client.is_closed is False, "a held handler must self-heal after its evicted client is closed"
server.close()
def test_the_pending_queue_cannot_grow_past_its_bound():
"""A caller that churns the client cache must not be able to grow this queue."""
clock = FakeClock()
closer = EvictedClientCloser(grace_seconds=60.0, max_pending=8, clock=clock)
clients = tuple(SyncClient() for _ in range(50))
for client in clients:
closer.mark_owned(client)
closer.schedule(client)
assert closer.pending_count == 8, "the queue grew past max_pending"
clock.advance(61.0)
closer.reap()
assert closer.pending_count == 0
assert sum(client.closed for client in clients) == 8, "everything queued should have been closed"
def test_a_reap_looks_at_what_is_due_rather_than_at_the_whole_queue():
"""Sustained churn evicts a client per request, and every read of the cache reaps.
So the cost of a reap has to track the entries that are due, not the length of
the queue; a reap that filters the whole queue makes the pair quadratic. Each
bucket is ordered by deadline, so an up-to-date reap compares one entry per
bucket and stops. Counting the comparisons measures that directly, where a
wall-clock budget would only measure the machine.
"""
evictions = 1_000
clock = FakeClock()
closer = EvictedClientCloser(
grace_seconds=60.0,
max_pending=evictions,
clock=lambda: CountingDeadline(clock.now),
)
clients = tuple(SyncClient() for _ in range(evictions))
for client in clients:
closer.mark_owned(client)
CountingDeadline.comparisons = 0
for client in clients:
closer.schedule(client)
closer.reap() # nothing is due yet, which is the hot path
clock.advance(61.0)
closer.reap()
assert closer.pending_count == 0
assert all(client.closed for client in clients)
assert CountingDeadline.comparisons < 10 * evictions, (
f"{CountingDeadline.comparisons} deadline comparisons for {evictions} evictions; "
"a reap is walking the whole queue"
)

66
tests/test_litellm/caching/test_llm_caching_handler.py
View file

@ -19,6 +19,7 @@ sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from litellm.caching.evicted_client_closer import EvictedClientCloser
from litellm.caching.llm_caching_handler import LLMClientCache
@ -156,6 +157,71 @@ def test_remove_key_no_event_loop():
assert "test-key" not in cache.cache_dict
class _FakeClock:
"""Hand-advanced monotonic clock, so grace windows need no real waiting."""
def __init__(self) -> None:
self.now = 1000.0
def __call__(self) -> float:
return self.now
def advance(self, seconds: float) -> None:
self.now += seconds
@pytest.mark.asyncio
async def test_evicted_litellm_owned_client_is_closed_once_the_grace_window_elapses():
"""
Eviction only drops the cache's reference. The SDK clients are reference
cycles, so without an explicit close the client keeps its connection pool
open until a generational collection runs.
"""
clock = _FakeClock()
cache = LLMClientCache(
max_size_in_memory=2,
evicted_client_closer=EvictedClientCloser(grace_seconds=60.0, clock=clock),
)
client = MockAsyncClient()
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()
await asyncio.sleep(0.1)
assert client.closed is False, "an in-flight request may still hold the client"
clock.advance(61.0)
cache.get_cache("any-key")
await asyncio.sleep(0.1)
assert client.closed is True
@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."""
clock = _FakeClock()
cache = LLMClientCache(
max_size_in_memory=2,
evicted_client_closer=EvictedClientCloser(grace_seconds=60.0, clock=clock),
)
client = MockAsyncClient()
cache.set_cache("client-key", client, 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()
clock.advance(3600.0)
cache.get_cache("any-key")
await asyncio.sleep(0.1)
assert client.closed is False
def test_remove_key_removes_plain_values():
"""
_remove_key correctly removes non-client values (strings, dicts, etc.).

71
tests/test_litellm/llms/azure/test_azure_common_utils.py
View file

@ -2034,3 +2034,74 @@ def test_azure_traditional_api_uses_azure_openai_client():
assert isinstance(
async_client, AsyncAzureOpenAI
), f"Expected AsyncAzureOpenAI client for api_version={api_version}"
def test_evicting_an_azure_client_built_on_the_callers_session_leaves_it_open(monkeypatch):
"""`initialize_azure_sdk_client` puts `litellm.aclient_session` on the SDK client.
That session belongs to the caller. `AsyncAzureOpenAI.close()` closes whatever
http client it was handed, so treating the wrapper as litellm's to close would
close the caller's shared session out from under them.
"""
import httpx
from litellm.caching.evicted_client_closer import EvictedClientCloser
from litellm.caching.llm_caching_handler import LLMClientCache
shared_session = httpx.AsyncClient()
closer = EvictedClientCloser(grace_seconds=0.0)
monkeypatch.setattr(litellm, "aclient_session", shared_session)
monkeypatch.setattr(
litellm,
"in_memory_llm_clients_cache",
LLMClientCache(evicted_client_closer=closer),
)
wrapper = BaseAzureLLM().get_azure_openai_client(
api_key="not-a-real-key",
api_base="https://litellm.openai.azure.com",
api_version="2024-02-01",
litellm_params={},
_is_async=True,
)
assert wrapper is not None
assert wrapper._client is shared_session, "the wrapper should be built on the caller's session"
closer.schedule(wrapper)
closer.reap()
assert closer.pending_count == 0, "a wrapper around the caller's session must never be queued"
assert shared_session.is_closed is False, "closed the session the caller configured"
def test_an_azure_client_litellm_built_its_own_http_client_for_is_still_closed(monkeypatch):
"""The ownership check must not turn the reclaim off for the ordinary case."""
from litellm.caching.evicted_client_closer import EvictedClientCloser
from litellm.caching.llm_caching_handler import LLMClientCache
closer = EvictedClientCloser(grace_seconds=0.0)
monkeypatch.setattr(litellm, "aclient_session", None)
monkeypatch.setattr(litellm, "client_session", None)
monkeypatch.setattr(
litellm,
"in_memory_llm_clients_cache",
LLMClientCache(evicted_client_closer=closer),
)
wrapper = BaseAzureLLM().get_azure_openai_client(
api_key="not-a-real-key",
api_base="https://litellm.openai.azure.com",
api_version="2024-02-01",
litellm_params={},
_is_async=False,
)
assert wrapper is not None
closer.schedule(wrapper)
assert closer.pending_count == 1, "litellm built this client's http client, so it owns it"
closer.reap()
assert wrapper.is_closed() is True

34
tests/test_litellm/llms/custom_httpx/test_http_handler.py
View file

@ -838,6 +838,40 @@ async def test_init_held_async_handler_survives_external_client_close():
await handler.close()
@pytest.mark.asyncio
async def test_init_held_async_handler_survives_evicted_client_close():
from litellm.caching.evicted_client_closer import EvictedClientCloser
from litellm.caching.llm_caching_handler import LLMClientCache
cache = LLMClientCache(evicted_client_closer=EvictedClientCloser(grace_seconds=0))
handler = AsyncHTTPHandler(timeout=42.5)
held_client = handler.client
cache.set_cache("init-held-handler", handler, litellm_owned_client=True, ttl=0)
await asyncio.sleep(0.02)
assert cache.get_cache("init-held-handler") is None
await asyncio.sleep(0.05)
assert held_client.is_closed
async def respond(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
await _read_http_request(reader)
writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")
await writer.drain()
writer.close()
server = await asyncio.start_server(respond, "127.0.0.1", 0)
port = server.sockets[0].getsockname()[1]
try:
response = await handler.post(f"http://127.0.0.1:{port}/v1/compress", json={"messages": []})
finally:
server.close()
await server.wait_closed()
assert response.status_code == 200
assert handler.client is not held_client
assert handler.client.timeout == httpx.Timeout(42.5)
await handler.close()
def test_init_held_sync_handler_recreates_closed_client():
from http.server import BaseHTTPRequestHandler, HTTPServer

72
tests/test_litellm/llms/openai/test_openai_common_utils.py
View file

@ -175,3 +175,75 @@ def test_get_openai_client_cache_key(client_type):
)
assert isinstance(key, str)
assert "api_key=sk-test" in key
def test_evicting_a_client_built_on_the_callers_session_leaves_that_session_open(monkeypatch):
"""`litellm.aclient_session` belongs to the caller, who goes on using it.
`_get_async_http_client` hands that session straight back, so the SDK client
litellm builds around it is only a wrapper. The SDK's `close()` closes
whatever http client it was given, so treating the wrapper as litellm's to
close would close the caller's shared session out from under them.
"""
import httpx
from litellm.caching.evicted_client_closer import EvictedClientCloser
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.llms.openai.openai import OpenAIChatCompletion
shared_session = httpx.AsyncClient()
closer = EvictedClientCloser(grace_seconds=0.0)
monkeypatch.setattr(litellm, "aclient_session", shared_session)
monkeypatch.setattr(
litellm,
"in_memory_llm_clients_cache",
LLMClientCache(evicted_client_closer=closer),
)
wrapper = OpenAIChatCompletion()._get_openai_client(
is_async=True,
api_key="sk-not-a-real-key",
api_base="https://api.openai.com/v1",
max_retries=2,
)
assert wrapper is not None
assert wrapper._client is shared_session, "the wrapper should be built on the caller's session"
closer.schedule(wrapper)
closer.reap()
assert closer.pending_count == 0, "a wrapper around the caller's session must never be queued"
assert shared_session.is_closed is False, "closed the session the caller configured"
def test_a_client_litellm_built_its_own_http_client_for_is_still_closed(monkeypatch):
"""The ownership check must not turn the reclaim off for the ordinary case."""
from litellm.caching.evicted_client_closer import EvictedClientCloser
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.llms.openai.openai import OpenAIChatCompletion
closer = EvictedClientCloser(grace_seconds=0.0)
monkeypatch.setattr(litellm, "aclient_session", None)
monkeypatch.setattr(litellm, "client_session", None)
monkeypatch.setattr(
litellm,
"in_memory_llm_clients_cache",
LLMClientCache(evicted_client_closer=closer),
)
wrapper = OpenAIChatCompletion()._get_openai_client(
is_async=False,
api_key="sk-not-a-real-key",
api_base="https://api.openai.com/v1",
max_retries=2,
)
assert wrapper is not None
closer.schedule(wrapper)
assert closer.pending_count == 1, "litellm built this client's http client, so it owns it"
closer.reap()
assert wrapper.is_closed() is True