diff --git a/litellm/caching/evicted_client_closer.py b/litellm/caching/evicted_client_closer.py deleted file mode 100644 index c895669be2b..00000000000 --- a/litellm/caching/evicted_client_closer.py +++ /dev/null @@ -1,277 +0,0 @@ -""" -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 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: - try: - await closing - except Exception: # noqa: BLE001 - a discarded client's close must never surface to callers - pass - - -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() diff --git a/litellm/caching/llm_caching_handler.py b/litellm/caching/llm_caching_handler.py index 6fa5963c99b..7d072a40195 100644 --- a/litellm/caching/llm_caching_handler.py +++ b/litellm/caching/llm_caching_handler.py @@ -5,44 +5,21 @@ 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.). - 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.`` + 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. - 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. + Clients that are no longer referenced will be garbage-collected normally. + For explicit shutdown cleanup, use ``close_litellm_async_clients()``. """ - 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, - ): - 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. @@ -55,22 +32,16 @@ class LLMClientCache(InMemoryCache): except RuntimeError: # handle no current running event loop return key - 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) + def set_cache(self, key, value, **kwargs): key = self.update_cache_key_with_event_loop(key) return super().set_cache(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) + async def async_set_cache(self, key, value, **kwargs): 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) diff --git a/litellm/constants.py b/litellm/constants.py index 0c7316455d6..264f595027f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -197,16 +197,6 @@ 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)) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 25dd9698624..9e613ae4eb4 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -509,8 +509,6 @@ class BaseAzureLLM(BaseOpenAILLM): openai_client=openai_client, client_initialization_params=client_initialization_params, client_type="azure", - litellm_owned_client=client is None - and self.owns_wrapped_http_client(azure_client_params.get("http_client")), ) return openai_client diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index e09fd48743a..619341be62b 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1408,7 +1408,6 @@ 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 @@ -1454,6 +1453,5 @@ 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 diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 527f44b930f..5c5e78c062d 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -128,33 +128,13 @@ 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 - - ``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. - """ + """Stores the OpenAI client in the in-memory cache for _DEFAULT_TTL_FOR_HTTPX_CLIENTS SECONDS""" _cache_key: Final = BaseOpenAILLM.get_openai_client_cache_key( client_initialization_params=client_initialization_params, client_type=client_type, @@ -163,7 +143,6 @@ class BaseOpenAILLM: key=_cache_key, value=openai_client, ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS, - litellm_owned_client=litellm_owned_client, ) @staticmethod diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 3c6846823e2..998319f3e85 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -360,16 +360,11 @@ 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=http_client, + http_client=OpenAIChatCompletion._get_async_http_client(shared_session=shared_session), timeout=timeout, max_retries=max_retries, organization=organization, @@ -378,7 +373,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): _new_client = OpenAI( api_key=api_key, base_url=api_base, - http_client=http_client, + http_client=OpenAIChatCompletion._get_sync_http_client(), timeout=timeout, max_retries=max_retries, organization=organization, @@ -389,7 +384,6 @@ 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 diff --git a/tests/test_litellm/caching/test_evicted_client_closer.py b/tests/test_litellm/caching/test_evicted_client_closer.py deleted file mode 100644 index a08fd58079d..00000000000 --- a/tests/test_litellm/caching/test_evicted_client_closer.py +++ /dev/null @@ -1,409 +0,0 @@ -""" -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() - - 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 handler.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 handler.client.is_closed is True - 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" - ) diff --git a/tests/test_litellm/caching/test_llm_caching_handler.py b/tests/test_litellm/caching/test_llm_caching_handler.py index 5f0e82dbb80..8e6a94945b0 100644 --- a/tests/test_litellm/caching/test_llm_caching_handler.py +++ b/tests/test_litellm/caching/test_llm_caching_handler.py @@ -19,7 +19,6 @@ 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 @@ -157,71 +156,6 @@ 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.). diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index 85db11fdb24..c0446a6cfba 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -2034,74 +2034,3 @@ 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 diff --git a/tests/test_litellm/llms/openai/test_openai_common_utils.py b/tests/test_litellm/llms/openai/test_openai_common_utils.py index a099b5c659f..ce25f7e9af6 100644 --- a/tests/test_litellm/llms/openai/test_openai_common_utils.py +++ b/tests/test_litellm/llms/openai/test_openai_common_utils.py @@ -175,75 +175,3 @@ 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