diff --git a/litellm/caching/llm_caching_handler.py b/litellm/caching/llm_caching_handler.py index beb9301e570..ccc4b314e44 100644 --- a/litellm/caching/llm_caching_handler.py +++ b/litellm/caching/llm_caching_handler.py @@ -42,7 +42,7 @@ 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. @@ -50,12 +50,14 @@ class LLMClientCache(InMemoryCache): 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: @@ -63,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/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index e19a101b817..20fe521304c 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1123,6 +1123,8 @@ class AsyncHTTPHandler: 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: @@ -1816,7 +1818,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 @@ -1832,6 +1834,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 12ead582975..b12955e0152 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -258,6 +258,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 @@ -272,7 +276,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 @@ -313,6 +319,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/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 5d02acea327..7fcdd3947e4 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1388,7 +1388,6 @@ 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() @@ -1404,6 +1403,40 @@ def test_finalizer_close_does_not_keep_a_loop_closed_right_after_it_alive(): 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 831df30c98a..2764bff3042 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 @@ -164,6 +165,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."""