diff --git a/litellm/caching/evicted_client_closer.py b/litellm/caching/evicted_client_closer.py index 555f6a25324..862f9ddb82e 100644 --- a/litellm/caching/evicted_client_closer.py +++ b/litellm/caching/evicted_client_closer.py @@ -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]: diff --git a/litellm/caching/llm_caching_handler.py b/litellm/caching/llm_caching_handler.py index a89e43b78b4..ccc4b314e44 100644 --- a/litellm/caching/llm_caching_handler.py +++ b/litellm/caching/llm_caching_handler.py @@ -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) diff --git a/litellm/integrations/agentops/agentops.py b/litellm/integrations/agentops/agentops.py index 3364520a8ae..3666cd88a9d 100644 --- a/litellm/integrations/agentops/agentops.py +++ b/litellm/integrations/agentops/agentops.py @@ -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()) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index a8be217709f..5620971a8a2 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -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 diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index a9d2dcdffb8..6c416124d1a 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -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 diff --git a/tests/unit/caching/_event_loop_reuse.py b/tests/unit/caching/_event_loop_reuse.py new file mode 100644 index 00000000000..97cf8fe31f0 --- /dev/null +++ b/tests/unit/caching/_event_loop_reuse.py @@ -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) diff --git a/tests/unit/caching/test_evicted_client_closer.py b/tests/unit/caching/test_evicted_client_closer.py index 7679e276621..936eb37a733 100644 --- a/tests/unit/caching/test_evicted_client_closer.py +++ b/tests/unit/caching/test_evicted_client_closer.py @@ -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 diff --git a/tests/unit/caching/test_llm_caching_handler.py b/tests/unit/caching/test_llm_caching_handler.py index a3efd5d97aa..465ba2f9150 100644 --- a/tests/unit/caching/test_llm_caching_handler.py +++ b/tests/unit/caching/test_llm_caching_handler.py @@ -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(): """ diff --git a/tests/unit/integrations/test_agentops.py b/tests/unit/integrations/test_agentops.py index 5d4055ac75f..629722f8602 100644 --- a/tests/unit/integrations/test_agentops.py +++ b/tests/unit/integrations/test_agentops.py @@ -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 diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index cd3b377b943..0163cb9b26d 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -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 diff --git a/tests/unit/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py index 206fd13260c..46c14fcfd8f 100644 --- a/tests/unit/llms/openai/test_openai_common_utils.py +++ b/tests/unit/llms/openai/test_openai_common_utils.py @@ -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."""