(perf) Litellm redis router fix - ~100ms improvement (#6483)

* docs(exception_mapping.md): add missing exception types

Fixes https://github.com/Aider-AI/aider/issues/2120#issuecomment-2438971183

* fix(main.py): register custom model pricing with specific key

Ensure custom model pricing is registered to the specific model+provider key combination

* test: make testing more robust for custom pricing

* fix(redis_cache.py): instrument otel logging for sync redis calls

ensures complete coverage for all redis cache calls

* refactor: pass parent_otel_span for redis caching calls in router

allows for more observability into what calls are causing latency issues

* test: update tests with new params

* refactor: ensure e2e otel tracing for router

* refactor(router.py): add more otel tracing acrosss router

catch all latency issues for router requests

* fix: fix linting error

* fix(router.py): fix linting error

* fix: fix test

* test: fix tests

* fix(dual_cache.py): pass ttl to redis cache

* fix: fix param

* perf(cooldown_cache.py): improve cooldown cache, to store cache results in memory for 5s, prevents redis call from being made on each request

reduces 100ms latency per call with caching enabled on router

* fix: fix test

* fix(cooldown_cache.py): handle if a result is None

* fix(cooldown_cache.py): add debug statements

* refactor(dual_cache.py): move to using an in-memory check for batch get cache, to prevent redis from being hit for every call

* fix(cooldown_cache.py): fix linting erropr

tests/local_testing/test_acooldowns_router.py

@@ -17,6 +17,7 @@ import concurrent
 from dotenv import load_dotenv
 
 import litellm
+
 from litellm import Router
 
 load_dotenv()
@@ -130,6 +131,7 @@ def test_multiple_deployments_parallel():
 @pytest.mark.parametrize("sync_mode", [True, False])
 @pytest.mark.asyncio
 async def test_cooldown_same_model_name(sync_mode):
+    litellm._turn_on_debug()
     # users could have the same model with different api_base
     # example
     # azure/chatgpt, api_base: 1234

tests/local_testing/test_caching.py

@@ -681,6 +681,7 @@ async def test_redis_cache_basic():
 
 
 @pytest.mark.asyncio
+@pytest.mark.flaky(retries=3, delay=1)
 async def test_redis_batch_cache_write():
     """
     Init redis client
@@ -2477,3 +2478,30 @@ async def test_redis_caching_ttl_sadd():
         )
         print(f"expected_timedelta: {expected_timedelta}")
         assert mock_expire.call_args.args[1] == expected_timedelta
+
+
+@pytest.mark.asyncio()
+async def test_dual_cache_caching_batch_get_cache():
+    """
+    - check redis cache called for initial batch get cache
+    - check redis cache not called for consecutive batch get cache with same keys
+    """
+    from litellm.caching.dual_cache import DualCache
+    from litellm.caching.redis_cache import RedisCache
+
+    dc = DualCache(redis_cache=MagicMock(spec=RedisCache))
+
+    with patch.object(
+        dc.redis_cache,
+        "async_batch_get_cache",
+        new=AsyncMock(
+            return_value={"test_key1": "test_value1", "test_key2": "test_value2"}
+        ),
+    ) as mock_async_get_cache:
+        await dc.async_batch_get_cache(keys=["test_key1", "test_key2"])
+
+        assert mock_async_get_cache.call_count == 1
+
+        await dc.async_batch_get_cache(keys=["test_key1", "test_key2"])
+
+        assert mock_async_get_cache.call_count == 1

tests/local_testing/test_router.py

@@ -2445,6 +2445,8 @@ async def test_aaarouter_dynamic_cooldown_message_retry_time(sync_mode):
             except litellm.RateLimitError:
                 pass
 
+        await asyncio.sleep(2)
+
         if sync_mode:
             cooldown_deployments = _get_cooldown_deployments(
                 litellm_router_instance=router, parent_otel_span=None

tests/local_testing/test_unit_test_caching.py

@@ -135,7 +135,7 @@ def test_get_cache_key_text_completion():
 def test_get_hashed_cache_key():
     cache = Cache()
     cache_key = "model:gpt-3.5-turbo,messages:Hello world"
-    hashed_key = cache._get_hashed_cache_key(cache_key)
+    hashed_key = Cache._get_hashed_cache_key(cache_key)
     assert len(hashed_key) == 64  # SHA-256 produces a 64-character hex string
 
 

tests/router_unit_tests/test_router_cooldown_utils.py

@@ -11,7 +11,7 @@ from litellm.router import Deployment, LiteLLM_Params, ModelInfo
 from concurrent.futures import ThreadPoolExecutor
 from collections import defaultdict
 from dotenv import load_dotenv
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 from litellm.integrations.prometheus import PrometheusLogger
 from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback
 from litellm.router_utils.cooldown_handlers import (