mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
fix(caching): key cached async clients by a per-loop token instead of id(loop) (#48877)
* fix(caching): key cached async clients by a per-loop token instead of id(loop) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(caching): drop deferred closes queued on a collected event loop Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(caching): pin that a close queued on a live loop waits for that loop * test(caching): run loop-reuse tests until a loop address repeats The collected-loop tests ran a fixed 50 loops and passed with the fix reverted whenever no loop reused an earlier address. They now stop at the first reuse and skip if none happens, through a shared helper. * test(caching): cover the cache key fallback for loops without weak references Also corrects the evicted client closer docstring: a dropped bucket's clients are left to the garbage collector. * fix(http): free short-lived event loops whose handler finalizer close finished as the loop stopped * fix(caching): share sync clients across event loops and catch a failed finalizer close * fix(agentops): stop closing the shared sync http client after the auth call --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
paveln83
GitHub
parent
d884167bf2
commit
4ce724e3df
11 changed files with 397 additions and 48 deletions
|
|
@ -29,6 +29,9 @@ 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.
|
||||
|
||||
A bucket whose loop has been collected is dropped, and its clients are left to the
|
||||
collector, as they are when the queue is full.
|
||||
|
||||
The queue holds its clients weakly, so waiting out a grace window never keeps
|
||||
alive anything the collector would have reclaimed first.
|
||||
"""
|
||||
|
|
@ -36,6 +39,7 @@ alive anything the collector would have reclaimed first.
|
|||
import asyncio
|
||||
import contextlib
|
||||
import inspect
|
||||
import itertools
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
|
|
@ -55,6 +59,9 @@ _CLOSABLE_ANYWHERE: Final = "closable-anywhere"
|
|||
_CLOSABLE_ON_ANY_LOOP: Final = "closable-on-any-loop"
|
||||
|
||||
_BucketKey = str | int
|
||||
_LOOP_TOKENS: Final[weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, int]] = weakref.WeakKeyDictionary()
|
||||
_LOOP_TOKEN_COUNTER: Final = itertools.count()
|
||||
_LOOP_TOKEN_LOCK: Final = threading.Lock()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -65,12 +72,12 @@ class _PendingClose:
|
|||
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``.
|
||||
be closed from the event loop they were evicted on, recorded in ``loop_token``.
|
||||
A client that closes synchronously carries neither constraint.
|
||||
"""
|
||||
|
||||
client_ref: "weakref.ref[object]"
|
||||
loop_id: int | None
|
||||
loop_token: int | None
|
||||
needs_loop: bool
|
||||
close_after: float
|
||||
|
||||
|
|
@ -79,16 +86,24 @@ 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:
|
||||
if pending.loop_token is None:
|
||||
return _CLOSABLE_ON_ANY_LOOP
|
||||
return pending.loop_id
|
||||
return pending.loop_token
|
||||
|
||||
|
||||
def _running_loop_id() -> int | None:
|
||||
def running_loop_token() -> int | None:
|
||||
try:
|
||||
return id(asyncio.get_running_loop())
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return None
|
||||
try:
|
||||
token: Final[int | None] = _LOOP_TOKENS.get(loop)
|
||||
if token is not None:
|
||||
return token
|
||||
with _LOOP_TOKEN_LOCK:
|
||||
return _LOOP_TOKENS.setdefault(loop, next(_LOOP_TOKEN_COUNTER))
|
||||
except TypeError:
|
||||
return id(loop)
|
||||
|
||||
|
||||
def _close_function(client: object) -> Callable[[], object] | None:
|
||||
|
|
@ -186,6 +201,8 @@ class EvictedClientCloser:
|
|||
self._clock = clock
|
||||
self._owned: weakref.WeakSet[object] = weakref.WeakSet()
|
||||
self._buckets: dict[_BucketKey, deque[_PendingClose]] = {} # mutable-ok: deadline-ordered queues
|
||||
self._dead_loop_tokens: deque[int] = deque() # mutable-ok: loop finalizers append without taking the queue lock
|
||||
self._watched_loop_tokens: set[int] = set() # mutable-ok: tracks tokens until their dead buckets are discarded
|
||||
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
|
||||
|
|
@ -216,13 +233,19 @@ class EvictedClientCloser:
|
|||
close_fn: Final = _close_function(client)
|
||||
if close_fn is None:
|
||||
return
|
||||
loop_token: Final = running_loop_token()
|
||||
needs_loop: Final = inspect.iscoroutinefunction(close_fn)
|
||||
if self._pending_count >= self._max_pending:
|
||||
return
|
||||
self._discard_dead_loops()
|
||||
if self._pending_count >= self._max_pending:
|
||||
return
|
||||
if needs_loop and loop_token is not None:
|
||||
self._watch_loop(loop_token)
|
||||
self._enqueue(
|
||||
_PendingClose(
|
||||
client_ref=weakref.ref(client),
|
||||
loop_id=_running_loop_id(),
|
||||
needs_loop=inspect.iscoroutinefunction(close_fn),
|
||||
loop_token=loop_token,
|
||||
needs_loop=needs_loop,
|
||||
close_after=self._clock() + self._grace_seconds,
|
||||
)
|
||||
)
|
||||
|
|
@ -236,7 +259,7 @@ class EvictedClientCloser:
|
|||
if not self._pending_count:
|
||||
return
|
||||
now: Final = self._clock()
|
||||
for pending in self._take_due(_running_loop_id(), now):
|
||||
for pending in self._take_due(running_loop_token(), now):
|
||||
client = pending.client_ref()
|
||||
if client is None:
|
||||
continue
|
||||
|
|
@ -249,6 +272,29 @@ class EvictedClientCloser:
|
|||
def pending_count(self) -> int:
|
||||
return self._pending_count
|
||||
|
||||
def _watch_loop(self, loop_token: int) -> None:
|
||||
if loop_token in self._watched_loop_tokens:
|
||||
return
|
||||
try:
|
||||
weakref.finalize(asyncio.get_running_loop(), self._dead_loop_tokens.append, loop_token)
|
||||
except TypeError:
|
||||
return
|
||||
self._watched_loop_tokens.add(loop_token)
|
||||
|
||||
def _discard_dead_loops(self) -> None:
|
||||
with self._queue_lock:
|
||||
self._discard_dead_loops_locked()
|
||||
|
||||
def _discard_dead_loops_locked(self) -> None:
|
||||
while self._dead_loop_tokens:
|
||||
self._discard_dead_loop_token_locked(self._dead_loop_tokens.popleft())
|
||||
|
||||
def _discard_dead_loop_token_locked(self, loop_token: int) -> None:
|
||||
self._watched_loop_tokens.discard(loop_token)
|
||||
bucket: Final = self._buckets.pop(loop_token, None)
|
||||
if bucket is not None:
|
||||
self._pending_count -= len(bucket)
|
||||
|
||||
def _enqueue(self, pending: _PendingClose) -> None:
|
||||
"""Append to the entry's bucket, dropping any dead entries it queues behind.
|
||||
|
||||
|
|
@ -257,6 +303,7 @@ class EvictedClientCloser:
|
|||
the front rather than having to be searched for.
|
||||
"""
|
||||
with self._queue_lock:
|
||||
self._discard_dead_loops_locked()
|
||||
bucket: Final = self._buckets.setdefault(_bucket_key(pending), deque())
|
||||
while bucket and bucket[0].client_ref() is None:
|
||||
bucket.popleft()
|
||||
|
|
@ -264,9 +311,12 @@ class EvictedClientCloser:
|
|||
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)
|
||||
def _take_due(self, loop_token: int | None, now: float) -> tuple[_PendingClose, ...]:
|
||||
buckets = (
|
||||
(_CLOSABLE_ANYWHERE,) if loop_token is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_token)
|
||||
)
|
||||
with self._queue_lock:
|
||||
self._discard_dead_loops_locked()
|
||||
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]:
|
||||
|
|
|
|||
|
|
@ -2,10 +2,9 @@
|
|||
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 .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer, running_loop_token
|
||||
from .in_memory_cache import InMemoryCache
|
||||
|
||||
|
||||
|
|
@ -43,24 +42,22 @@ class LLMClientCache(InMemoryCache):
|
|||
self.evicted_client_closer.schedule(evicted)
|
||||
self.evicted_client_closer.reap()
|
||||
|
||||
def update_cache_key_with_event_loop(self, key):
|
||||
def update_cache_key_with_event_loop(self, key: str) -> str:
|
||||
"""
|
||||
Add the event loop to the cache key, to prevent event loop closed errors.
|
||||
If none, use the key as is.
|
||||
"""
|
||||
try:
|
||||
event_loop: Final = asyncio.get_running_loop()
|
||||
stringified_event_loop: Final = str(id(event_loop))
|
||||
return f"{key}-{stringified_event_loop}"
|
||||
except RuntimeError: # handle no current running event loop
|
||||
return key
|
||||
token: Final = running_loop_token()
|
||||
return f"{key}-{token}" if token is not None else key
|
||||
|
||||
def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
|
||||
def _scoped_key(self, key: str, loop_bound: bool) -> str:
|
||||
return self.update_cache_key_with_event_loop(key) if loop_bound else key
|
||||
|
||||
def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, loop_bound: bool = True, **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)
|
||||
return super().set_cache(self._scoped_key(key, loop_bound), value, **kwargs)
|
||||
|
||||
async def async_set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
|
||||
if litellm_owned_client:
|
||||
|
|
@ -68,13 +65,13 @@ class LLMClientCache(InMemoryCache):
|
|||
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)
|
||||
def get_cache(self, key: str, loop_bound: bool = True, **kwargs):
|
||||
scoped_key: Final = self._scoped_key(key, loop_bound)
|
||||
self.evicted_client_closer.reap()
|
||||
|
||||
return super().get_cache(key, **kwargs)
|
||||
return super().get_cache(scoped_key, **kwargs)
|
||||
|
||||
async def async_get_cache(self, key, **kwargs):
|
||||
async def async_get_cache(self, key: str, **kwargs):
|
||||
key = self.update_cache_key_with_event_loop(key)
|
||||
|
||||
return await super().async_get_cache(key, **kwargs)
|
||||
|
|
|
|||
|
|
@ -102,18 +102,14 @@ class AgentOps(OpenTelemetry):
|
|||
"Connection": "keep-alive",
|
||||
}
|
||||
|
||||
client: Final = get_httpx_client()
|
||||
try:
|
||||
response: Final = client.post(
|
||||
url=auth_endpoint,
|
||||
headers=headers,
|
||||
json={"api_key": api_key},
|
||||
timeout=10,
|
||||
)
|
||||
response: Final = get_httpx_client().post(
|
||||
url=auth_endpoint,
|
||||
headers=headers,
|
||||
json={"api_key": api_key},
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"Failed to fetch auth token: {response.text}")
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"Failed to fetch auth token: {response.text}")
|
||||
|
||||
return _DECODED_JSON.validate_python(response.json())
|
||||
finally:
|
||||
client.close()
|
||||
return _DECODED_JSON.validate_python(response.json())
|
||||
|
|
|
|||
|
|
@ -1205,6 +1205,17 @@ class AsyncHTTPHandler:
|
|||
# Mirrors LiteLLMAiohttpTransport._background_close_tasks.
|
||||
_finalizer_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes
|
||||
|
||||
@classmethod
|
||||
async def _close_client_at_finalization(cls, client: httpx.AsyncClient) -> None:
|
||||
try:
|
||||
await client.aclose()
|
||||
except (OSError, RuntimeError) as exc:
|
||||
verbose_logger.debug("Error closing client at finalization: %s", exc)
|
||||
finally:
|
||||
task: Final = asyncio.current_task()
|
||||
if task is not None:
|
||||
cls._finalizer_close_tasks.discard(task)
|
||||
|
||||
@classmethod
|
||||
def _on_finalizer_close_done(cls, task: "asyncio.Task[None]") -> None:
|
||||
cls._finalizer_close_tasks.discard(task)
|
||||
|
|
@ -1272,8 +1283,8 @@ class AsyncHTTPHandler:
|
|||
# here is the cross-loop path the transport refuses.
|
||||
self._dispose_wrapped_aiohttp_session()
|
||||
return
|
||||
task: Final = loop.create_task(self._client.aclose())
|
||||
cls: Final = type(self)
|
||||
task: Final = loop.create_task(cls._close_client_at_finalization(self._client))
|
||||
cls._finalizer_close_tasks.add(task)
|
||||
task.add_done_callback(cls._on_finalizer_close_done)
|
||||
except Exception:
|
||||
|
|
@ -1973,7 +1984,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
|
|||
cache = LLMClientCache()
|
||||
setattr(litellm, "in_memory_llm_clients_cache", cache)
|
||||
|
||||
_cached_client: Final = cache.get_cache(_cache_key_name)
|
||||
_cached_client: Final = cache.get_cache(_cache_key_name, loop_bound=False)
|
||||
if _cached_client:
|
||||
return _cached_client
|
||||
|
||||
|
|
@ -1989,6 +2000,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
|
|||
value=_new_client,
|
||||
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
|
||||
litellm_owned_client=True,
|
||||
loop_bound=False,
|
||||
)
|
||||
return _new_client
|
||||
|
||||
|
|
|
|||
|
|
@ -275,6 +275,10 @@ def build_output_token_limit_response(
|
|||
return e.response.headers, (_aiter_once(chunk) if is_async else _iter_once(chunk))
|
||||
|
||||
|
||||
def _is_loop_bound(client_initialization_params: Mapping[str, object]) -> bool:
|
||||
return client_initialization_params.get("is_async") is not False
|
||||
|
||||
|
||||
class BaseOpenAILLM:
|
||||
"""
|
||||
Base class for OpenAI LLMs for getting their httpx clients and SSL verification settings
|
||||
|
|
@ -289,7 +293,9 @@ class BaseOpenAILLM:
|
|||
client_initialization_params=client_initialization_params,
|
||||
client_type=client_type,
|
||||
)
|
||||
_cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(_cache_key)
|
||||
_cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(
|
||||
_cache_key, loop_bound=_is_loop_bound(client_initialization_params)
|
||||
)
|
||||
return _cached_client
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -330,6 +336,7 @@ class BaseOpenAILLM:
|
|||
value=openai_client,
|
||||
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
|
||||
litellm_owned_client=litellm_owned_client,
|
||||
loop_bound=_is_loop_bound(client_initialization_params),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
20
tests/unit/caching/_event_loop_reuse.py
Normal file
20
tests/unit/caching/_event_loop_reuse.py
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
from collections.abc import Callable
|
||||
from typing import Final, TypeVar
|
||||
|
||||
import pytest
|
||||
|
||||
LoopResult = TypeVar("LoopResult")
|
||||
|
||||
|
||||
def runs_until_a_loop_address_repeats(
|
||||
run_on_new_loop: Callable[[], tuple[int, LoopResult]],
|
||||
earlier: tuple[tuple[int, LoopResult], ...] = (),
|
||||
limit: int = 50,
|
||||
) -> tuple[tuple[int, LoopResult], ...]:
|
||||
if len(earlier) == limit:
|
||||
pytest.skip(f"none of {limit} collected loops had its address handed to a later loop")
|
||||
latest: Final = run_on_new_loop()
|
||||
runs: Final = (*earlier, latest)
|
||||
if latest[0] in {loop_id for loop_id, _ in earlier}:
|
||||
return runs
|
||||
return runs_until_a_loop_address_repeats(run_on_new_loop, runs, limit)
|
||||
|
|
@ -10,6 +10,7 @@ never closed, because litellm does not own its lifecycle.
|
|||
import asyncio
|
||||
import gc
|
||||
import weakref
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import httpx
|
||||
|
|
@ -18,6 +19,7 @@ from redis.asyncio import ConnectionPool, Redis
|
|||
|
||||
from litellm.caching.evicted_client_closer import EvictedClientCloser
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats
|
||||
|
||||
|
||||
class FakeClock:
|
||||
|
|
@ -41,6 +43,14 @@ class AsyncClient:
|
|||
self.closed = True
|
||||
|
||||
|
||||
class RecordingAsyncClient:
|
||||
def __init__(self) -> None:
|
||||
self.close_calls = 0
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.close_calls += 1
|
||||
|
||||
|
||||
class SyncClient:
|
||||
def __init__(self) -> None:
|
||||
self.closed = False
|
||||
|
|
@ -266,6 +276,100 @@ def test_sync_client_evicted_outside_an_event_loop_is_still_closed():
|
|||
assert closer.pending_count == 0
|
||||
|
||||
|
||||
def test_async_client_from_collected_loop_is_not_closed_by_fresh_loops() -> None:
|
||||
clock: Final = FakeClock()
|
||||
closer: Final = EvictedClientCloser(grace_seconds=0, clock=clock)
|
||||
|
||||
async def reap_then_schedule(client: RecordingAsyncClient) -> None:
|
||||
closer.reap()
|
||||
await asyncio.sleep(0)
|
||||
closer.schedule(client)
|
||||
|
||||
def run_on_new_loop() -> tuple[int, RecordingAsyncClient]:
|
||||
client: Final = RecordingAsyncClient()
|
||||
closer.mark_owned(client)
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
loop_id: Final = id(loop)
|
||||
try:
|
||||
loop.run_until_complete(reap_then_schedule(client))
|
||||
finally:
|
||||
loop.close()
|
||||
del loop
|
||||
gc.collect()
|
||||
return loop_id, client
|
||||
|
||||
runs: Final = runs_until_a_loop_address_repeats(run_on_new_loop)
|
||||
|
||||
assert all(client.close_calls == 0 for _, client in runs)
|
||||
|
||||
|
||||
def test_async_client_queued_on_a_live_loop_is_closed_only_by_that_loop() -> None:
|
||||
clock: Final = FakeClock()
|
||||
closer: Final = EvictedClientCloser(grace_seconds=0, clock=clock)
|
||||
client: Final = RecordingAsyncClient()
|
||||
closer.mark_owned(client)
|
||||
owner_loop: Final = asyncio.new_event_loop()
|
||||
other_loop: Final = asyncio.new_event_loop()
|
||||
|
||||
async def schedule() -> None:
|
||||
closer.schedule(client)
|
||||
|
||||
async def reap() -> None:
|
||||
closer.reap()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
try:
|
||||
owner_loop.run_until_complete(schedule())
|
||||
other_loop.run_until_complete(reap())
|
||||
closes_after_other_loop_reaped: Final = client.close_calls
|
||||
pending_after_other_loop_reaped: Final = closer.pending_count
|
||||
owner_loop.run_until_complete(reap())
|
||||
finally:
|
||||
other_loop.close()
|
||||
owner_loop.close()
|
||||
|
||||
assert closes_after_other_loop_reaped == 0
|
||||
assert pending_after_other_loop_reaped == 1
|
||||
assert client.close_calls == 1
|
||||
assert closer.pending_count == 0
|
||||
|
||||
|
||||
def test_closes_queued_on_a_collected_loop_free_capacity() -> None:
|
||||
clock: Final = FakeClock()
|
||||
closer: Final = EvictedClientCloser(grace_seconds=0, max_pending=1, clock=clock)
|
||||
first: Final = RecordingAsyncClient()
|
||||
closer.mark_owned(first)
|
||||
loop_a: Final = asyncio.new_event_loop()
|
||||
|
||||
async def schedule_first() -> None:
|
||||
closer.schedule(first)
|
||||
|
||||
try:
|
||||
loop_a.run_until_complete(schedule_first())
|
||||
finally:
|
||||
loop_a.close()
|
||||
del loop_a
|
||||
gc.collect()
|
||||
|
||||
second: Final = RecordingAsyncClient()
|
||||
closer.mark_owned(second)
|
||||
loop_b: Final = asyncio.new_event_loop()
|
||||
|
||||
async def schedule_and_reap_second() -> None:
|
||||
closer.schedule(second)
|
||||
closer.reap()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
try:
|
||||
loop_b.run_until_complete(schedule_and_reap_second())
|
||||
finally:
|
||||
loop_b.close()
|
||||
|
||||
assert first.close_calls == 0
|
||||
assert second.close_calls == 1
|
||||
assert closer.pending_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_async_client_waits_for_a_loop_rather_than_being_dropped():
|
||||
clock = FakeClock()
|
||||
|
|
@ -289,8 +393,7 @@ async def test_an_async_client_waits_for_a_loop_rather_than_being_dropped():
|
|||
|
||||
|
||||
@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."""
|
||||
async def test_a_client_evicted_on_a_collected_loop_is_dropped():
|
||||
clock = FakeClock()
|
||||
closer = make_closer(clock)
|
||||
client = AsyncClient()
|
||||
|
|
@ -310,7 +413,7 @@ async def test_a_client_evicted_on_another_event_loop_is_left_alone():
|
|||
await asyncio.sleep(0.05)
|
||||
|
||||
assert client.closed is False
|
||||
assert closer.pending_count == 1
|
||||
assert closer.pending_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ See: https://github.com/BerriAI/litellm/pull/22247
|
|||
|
||||
import asyncio, gc, httpx, importlib, litellm, os, threading, time, weakref
|
||||
import warnings
|
||||
from typing import Final
|
||||
|
||||
import aiohttp
|
||||
import pytest
|
||||
|
|
@ -29,6 +30,41 @@ from litellm.llms.custom_httpx.http_handler import(
|
|||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import _invalidate_model_cost_lowercase_map
|
||||
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
||||
from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats
|
||||
|
||||
|
||||
async def _get_and_set_cache(cache: LLMClientCache) -> object | None:
|
||||
cached: Final = await cache.async_get_cache("client")
|
||||
await cache.async_set_cache("client", object())
|
||||
return cached
|
||||
|
||||
|
||||
def _cache_result_on_new_loop(cache: LLMClientCache) -> tuple[int, object | None]:
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
loop_id: Final = id(loop)
|
||||
try:
|
||||
result: Final = loop.run_until_complete(_get_and_set_cache(cache))
|
||||
finally:
|
||||
loop.close()
|
||||
del loop
|
||||
gc.collect()
|
||||
return loop_id, result
|
||||
|
||||
|
||||
async def _get_cache_key(cache: LLMClientCache) -> str:
|
||||
return cache.update_cache_key_with_event_loop("client")
|
||||
|
||||
|
||||
def _cache_key_on_new_loop(cache: LLMClientCache) -> tuple[int, str]:
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
loop_id: Final = id(loop)
|
||||
try:
|
||||
key: Final = loop.run_until_complete(_get_cache_key(cache))
|
||||
finally:
|
||||
loop.close()
|
||||
del loop
|
||||
gc.collect()
|
||||
return loop_id, key
|
||||
|
||||
|
||||
class MockAsyncClient:
|
||||
|
|
@ -51,6 +87,44 @@ class MockSyncClient:
|
|||
self.closed = True
|
||||
|
||||
|
||||
def test_async_client_cache_does_not_reuse_clients_between_collected_loops() -> None:
|
||||
cache: Final = LLMClientCache()
|
||||
runs: Final = runs_until_a_loop_address_repeats(lambda: _cache_result_on_new_loop(cache))
|
||||
|
||||
assert all(result is None for _, result in runs)
|
||||
|
||||
|
||||
def test_async_client_cache_keys_differ_for_every_collected_loop() -> None:
|
||||
cache: Final = LLMClientCache()
|
||||
runs: Final = runs_until_a_loop_address_repeats(lambda: _cache_key_on_new_loop(cache))
|
||||
keys: Final = tuple(key for _, key in runs)
|
||||
|
||||
assert len(set(keys)) == len(keys)
|
||||
|
||||
|
||||
class _LoopWithoutWeakReferences:
|
||||
__slots__ = ()
|
||||
|
||||
|
||||
def _cache_key_while_running(cache: LLMClientCache, loop: _LoopWithoutWeakReferences) -> str:
|
||||
asyncio._set_running_loop(loop)
|
||||
try:
|
||||
return cache.update_cache_key_with_event_loop("client")
|
||||
finally:
|
||||
asyncio._set_running_loop(None)
|
||||
|
||||
|
||||
def test_cache_keys_stay_per_loop_on_loops_without_weak_references() -> None:
|
||||
cache: Final = LLMClientCache()
|
||||
loop_a: Final = _LoopWithoutWeakReferences()
|
||||
loop_b: Final = _LoopWithoutWeakReferences()
|
||||
|
||||
first_key_on_a: Final = _cache_key_while_running(cache, loop_a)
|
||||
|
||||
assert _cache_key_while_running(cache, loop_a) == first_key_on_a
|
||||
assert _cache_key_while_running(cache, loop_b) != first_key_on_a
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remove_key_does_not_close_async_client():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,10 +1,14 @@
|
|||
import os
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.caching.llm_caching_handler import LLMClientCache
|
||||
from litellm.integrations.agentops.agentops import AgentOps, AgentOpsConfig
|
||||
from litellm.llms.custom_httpx.http_handler import get_httpx_client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -111,3 +115,17 @@ def test_agentops_initialization_no_auth():
|
|||
assert agentops.resource_attributes["deployment.environment"] == "test_env"
|
||||
assert agentops.resource_attributes["telemetry.sdk.name"] == "agentops"
|
||||
assert "project.id" not in agentops.resource_attributes
|
||||
|
||||
|
||||
def test_auth_fetch_leaves_the_shared_sync_client_open_for_other_callers(
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
|
||||
auth_endpoint: Final = "https://agentops.example/v3/auth/token"
|
||||
respx_mock.post(auth_endpoint).respond(json={"token": "jwt", "project_id": "project-1"})
|
||||
client_another_caller_is_using: Final = get_httpx_client().client
|
||||
|
||||
agentops: Final = AgentOps(config=AgentOpsConfig(api_key="key", auth_endpoint=auth_endpoint))
|
||||
|
||||
assert agentops.resource_attributes["project.id"] == "project-1"
|
||||
assert not client_another_caller_is_using.is_closed
|
||||
|
|
|
|||
|
|
@ -1388,6 +1388,55 @@ async def test_finalizer_with_running_loop_schedules_close_and_holds_task_ref():
|
|||
assert not (AsyncHTTPHandler._finalizer_close_tasks & scheduled)
|
||||
|
||||
|
||||
async def _drop_unused_handler() -> None:
|
||||
AsyncHTTPHandler(timeout=61.0)
|
||||
gc.collect()
|
||||
|
||||
|
||||
def test_finalizer_close_does_not_keep_a_loop_closed_right_after_it_alive():
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
loop.run_until_complete(_drop_unused_handler())
|
||||
loop.close()
|
||||
closed_loop: Final = weakref.ref(loop)
|
||||
del loop
|
||||
gc.collect()
|
||||
|
||||
assert closed_loop() is None
|
||||
|
||||
|
||||
class _TransportFailingToClose(httpx.AsyncBaseTransport):
|
||||
async def aclose(self) -> None:
|
||||
raise RuntimeError("close failed")
|
||||
|
||||
|
||||
async def _drop_handler_whose_close_fails() -> None:
|
||||
AsyncHTTPHandler(timeout=61.0, transport=_TransportFailingToClose())
|
||||
gc.collect()
|
||||
|
||||
|
||||
def test_failed_finalizer_close_in_a_loops_last_step_reports_no_unretrieved_exception():
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
exception_handler: Final = MagicMock()
|
||||
loop.set_exception_handler(exception_handler)
|
||||
loop.run_until_complete(_drop_handler_whose_close_fails())
|
||||
loop.close()
|
||||
del loop
|
||||
gc.collect()
|
||||
|
||||
exception_handler.assert_not_called()
|
||||
|
||||
|
||||
async def _sync_client_with_timeout_62() -> HTTPHandler:
|
||||
return get_httpx_client(params={"timeout": 62.0})
|
||||
|
||||
|
||||
def test_sync_client_is_shared_by_short_lived_loops_and_callers_outside_a_loop():
|
||||
built_in_first_loop: Final = asyncio.run(_sync_client_with_timeout_62())
|
||||
|
||||
assert asyncio.run(_sync_client_with_timeout_62()) is built_in_first_loop
|
||||
assert get_httpx_client(params={"timeout": 62.0}) is built_in_first_loop
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_close_helper_respects_session_ownership():
|
||||
"""The loop-less fallback closes only sessions the transport owns; a
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
|
@ -154,6 +155,28 @@ def test_get_openai_client_initialization_param_fields(client_type):
|
|||
assert "self" not in result
|
||||
|
||||
|
||||
def _loop_test_client_params(is_async: bool) -> dict[str, object]:
|
||||
return {"api_key": "sk-loop-scope-test", "is_async": is_async}
|
||||
|
||||
|
||||
async def _cache_openai_client(client: openai.OpenAI | openai.AsyncOpenAI, is_async: bool) -> None:
|
||||
BaseOpenAILLM.set_cached_openai_client(client, "openai", _loop_test_client_params(is_async))
|
||||
|
||||
|
||||
async def _cached_openai_client(is_async: bool) -> object:
|
||||
return BaseOpenAILLM.get_cached_openai_client(_loop_test_client_params(is_async), "openai")
|
||||
|
||||
|
||||
def test_sync_openai_client_cached_in_one_loop_is_reused_by_the_next_while_async_is_not():
|
||||
sync_client: Final = openai.OpenAI(api_key="sk-loop-scope-test")
|
||||
async_client: Final = openai.AsyncOpenAI(api_key="sk-loop-scope-test")
|
||||
asyncio.run(_cache_openai_client(sync_client, is_async=False))
|
||||
asyncio.run(_cache_openai_client(async_client, is_async=True))
|
||||
|
||||
assert asyncio.run(_cached_openai_client(is_async=False)) is sync_client
|
||||
assert asyncio.run(_cached_openai_client(is_async=True)) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client_type", ["openai", "azure"])
|
||||
def test_get_openai_client_cache_key(client_type):
|
||||
"""Verify get_openai_client_cache_key doesn't raise on tuple + tuple concatenation."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue