diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 146860d2a60..cb3e1371e42 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -13,7 +13,6 @@ from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast from pydantic import TypeAdapter -from redis.exceptions import RedisError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger @@ -601,6 +600,7 @@ class RedisUpdateBuffer: ) if len(rpush_list) == 0: return + from redis.exceptions import RedisError try: await self.redis_cache.async_rpush_pipeline(rpush_list=rpush_list) @@ -620,6 +620,8 @@ class RedisUpdateBuffer: """Move this worker's pending tool rows to Redis so only the leader writes LiteLLM_ToolTable.""" if self.redis_cache is None: return + from redis.exceptions import RedisError + items: Final = tool_discovery_queue.flush() if not items: return diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 798b1e5cc74..81b9c80a24b 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -11,7 +11,6 @@ from typing import TYPE_CHECKING, Final, Literal from fastapi import HTTPException from pydantic import BaseModel, ConfigDict, ValidationError -from redis.exceptions import RedisError import litellm from litellm import Router @@ -526,6 +525,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): member: Final = ( pool_member(priority_class) if user_api_key_dict.team_id is None else team_member(user_api_key_dict.team_id) ) + from redis.exceptions import RedisError + try: members: Final = await self.active_teams.touch(model, member, priority_class, self.v3_limiter.window_size) except (RedisError, RedisCircuitBreakerOpenError, ValidationError) as e: diff --git a/tests/unit/proxy/conftest.py b/tests/unit/proxy/conftest.py index 7ff922179ce..1aec7831b8e 100644 --- a/tests/unit/proxy/conftest.py +++ b/tests/unit/proxy/conftest.py @@ -4,13 +4,16 @@ import _hashlib import asyncio import copy import hashlib +import importlib import inspect import logging import os +import sys import tempfile import warnings from collections.abc import Awaitable, Callable, Iterator from contextlib import contextmanager +from types import ModuleType from typing import Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -570,3 +573,18 @@ def free_and_paid_router() -> "litellm.Router": _per_token_priced("paid-model", "openai/paid-backend", 1e-6, 2e-6), ] ) + + +@pytest.fixture +def import_without_redis_installed(monkeypatch: pytest.MonkeyPatch) -> Callable[[str], ModuleType]: + def _import(module_name: str) -> ModuleType: + package_name, _, attribute = module_name.rpartition(".") + monkeypatch.setattr(sys.modules[package_name], attribute, sys.modules[module_name]) + monkeypatch.setitem(sys.modules, "redis", None) + for name in tuple(sys.modules): + if name.startswith("redis."): + monkeypatch.setitem(sys.modules, name, None) + monkeypatch.delitem(sys.modules, module_name) + return importlib.import_module(module_name) + + return _import diff --git a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 7879adaed93..23ea7d2fedd 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -1,10 +1,13 @@ import json import logging +from collections.abc import Callable from datetime import datetime, timezone +from types import ModuleType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from redis.exceptions import RedisError from litellm._logging import verbose_proxy_logger from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY @@ -12,6 +15,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import DailyTagSpendTransaction from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue from litellm.proxy.proxy_server import ProxyStartupEvent @@ -572,8 +576,6 @@ async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer @pytest.mark.asyncio async def test_restore_transactions_to_redis_swallows_redis_error(redis_update_buffer, mock_redis_cache): """A Redis failure during restore must not propagate to the caller's finally block.""" - from redis.exceptions import RedisError - mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=RedisError("redis down")) await redis_update_buffer.restore_transactions_to_redis( @@ -909,11 +911,10 @@ async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_ @pytest.mark.asyncio -async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails(): - from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue - +@pytest.mark.parametrize("failure", [ConnectionError("redis down"), RedisError("redis down")]) +async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails(failure: Exception): redis_cache = MagicMock() - redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis down")) + redis_cache.async_rpush_pipeline = AsyncMock(side_effect=failure) queue = ToolDiscoveryQueue() queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) @@ -921,3 +922,18 @@ async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails() await RedisUpdateBuffer(redis_cache=redis_cache).store_tool_discoveries_in_redis(queue) assert queue.flush() == [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2}] + + +@pytest.mark.asyncio +async def test_buffer_without_redis_installed_skips_redis_work_when_no_cache_is_configured( + import_without_redis_installed: Callable[[str], ModuleType], +): + module: Final = import_without_redis_installed("litellm.proxy.db.db_transaction_queue.redis_update_buffer") + buffer: Final[RedisUpdateBuffer] = module.RedisUpdateBuffer(redis_cache=None) + queue: Final = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + + await buffer.restore_transactions_to_redis(db_spend_update_transactions={"key_list_transactions": {"key1": 1.0}}) + await buffer.store_tool_discoveries_in_redis(queue) + + assert queue.flush() == [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}] diff --git a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py index 87c88ae98f6..8a385b4a3fa 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -6,8 +6,9 @@ Core tests to validate that priority weights are respected (0.9/0.1) instead of import asyncio import time -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from datetime import datetime, timedelta +from types import ModuleType from typing import Final from unittest.mock import AsyncMock, patch @@ -3677,3 +3678,16 @@ async def test_a_worker_without_redis_reads_the_models_saturation_from_its_own_m group: Final = handler.llm_router.registry().group("shared-model") assert group is not None assert await handler._check_model_saturation("shared-model", group) == pytest.approx(0.7) + + +@pytest.mark.asyncio +async def test_the_limiter_loads_and_admits_requests_without_redis_installed( + monkeypatch: pytest.MonkeyPatch, import_without_redis_installed: Callable[[str], ModuleType] +): + module: Final = import_without_redis_installed("litellm.proxy.hooks.dynamic_rate_limiter_v3") + handler: Final[PROXY_DynamicRateLimitHandlerV3] = module.PROXY_DynamicRateLimitHandlerV3( + internal_usage_cache=DualCache() + ) + handler.update_variables(llm_router=_borrowing_handler(monkeypatch, borrows=False).llm_router) + + await _check(handler, _class_team("team-a"), saturation=0.0, priority="interactive")