diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 5c9f74e4c18..594417173ed 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -15,6 +15,7 @@ import hashlib import inspect import json import time +from collections.abc import Awaitable, Callable, Sequence from datetime import timedelta from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union, cast @@ -535,7 +536,7 @@ class RedisCache(BaseCache): ) raise e - def async_register_script(self, script: str) -> Any: + def async_register_script(self, script: str) -> Callable[..., Awaitable[Any]]: """ Register a Lua script with Redis asynchronously. Works with both standalone Redis and Redis Cluster. @@ -545,39 +546,87 @@ class RedisCache(BaseCache): scripts would operate on raw keys while the rest of the cache uses the namespace, leaving rate-limit and lock keys outside the configured prefix. + Registration is deferred to call time and cached per running event loop + (via in_memory_llm_clients_cache, which keys its entries on the loop). A + registered script is bound to the connection of the loop it was created + on; awaiting it from another loop raises "got Future attached to a + different loop". Binding lazily on the calling loop gives the script the + same per-loop scoping init_async_client already gives the clients, so a + script registered once at startup is never reused across loops. + Args: script (str): The Lua script to register Returns: - Any: A script object that can be called with keys and args + A callable ``(keys, args, client=None)`` that runs the script + against the calling loop's Redis client. """ - try: - _redis_client = self.init_async_client() - # For standalone Redis - if hasattr(_redis_client, "register_script"): - registered_script = _redis_client.register_script(script) # type: ignore + # Keyed by connection params and namespace as well as the script, so + # two RedisCache instances pointing at different servers or using + # different key prefixes never share an executor; in_memory_llm_clients_cache + # then adds the running loop, completing the per-(client, namespace, loop) + # scoping. + script_cache_key = ( + f"redis-registered-script-{self._get_async_client_cache_key()}-" + f"{self.namespace}-{hashlib.sha256(script.encode()).hexdigest()[:16]}" + ) - async def namespaced_script( - keys: list[str], args: list[Any], client: Any = None - ) -> Any: - keys = [self.check_and_fix_namespace(key=key) for key in keys] - return await registered_script(keys=keys, args=args, client=client) + async def run_script( + keys: Sequence[str], args: Sequence[Any], client: Any = None + ) -> Any: + executor: Optional[Callable[..., Awaitable[Any]]] = ( + litellm.in_memory_llm_clients_cache.get_cache(key=script_cache_key) + ) + if executor is None: + executor = self._register_script_for_current_loop(script) + litellm.in_memory_llm_clients_cache.set_cache( + key=script_cache_key, value=executor + ) + return await executor(keys=keys, args=args, client=client) - return namespaced_script - # For Redis Cluster - elif hasattr(_redis_client, "script_load"): - # Load the script and get its SHA - script_sha = _redis_client.script_load(script) # type: ignore + return run_script - # Return a callable that uses evalsha - async def script_callable(keys: List[str], args: List[Any]) -> Any: - keys = [self.check_and_fix_namespace(key=key) for key in keys] - return _redis_client.evalsha(script_sha, len(keys), *keys, *args) # type: ignore + def _register_script_for_current_loop( + self, script: str + ) -> Callable[..., Awaitable[Any]]: + """ + Register the script against the current event loop's Redis client. - return script_callable - except Exception as e: - verbose_logger.error(f"Error registering Redis script: {str(e)}") - raise e + Kept separate from async_register_script so each loop caches its own + executor; see that method for why the binding must be per loop. + """ + _redis_client: Any = self.init_async_client() + if hasattr(_redis_client, "register_script"): + registered_script = _redis_client.register_script(script) + + async def standalone_executor( + keys: Sequence[str], args: Sequence[Any], client: Any = None + ) -> Any: + namespaced_keys = tuple( + self.check_and_fix_namespace(key=key) for key in keys + ) + return await registered_script( + keys=namespaced_keys, args=args, client=client + ) + + return standalone_executor + + if hasattr(_redis_client, "script_load"): + script_sha = _redis_client.script_load(script) + + async def cluster_executor( + keys: Sequence[str], args: Sequence[Any], client: Any = None + ) -> Any: + namespaced_keys = tuple( + self.check_and_fix_namespace(key=key) for key in keys + ) + return await _redis_client.evalsha( + script_sha, len(namespaced_keys), *namespaced_keys, *args + ) + + return cluster_executor + + raise ValueError("Redis client does not support Lua script registration") @_redis_circuit_breaker_guard async def async_set_cache(self, key, value, **kwargs): diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 37c938bb38d..de5b3b32105 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1,3 +1,4 @@ +import asyncio import os import sys from unittest.mock import MagicMock, patch @@ -593,10 +594,136 @@ async def test_async_register_script_namespaces_keys( assert result == "ok" registered_script.assert_awaited_once_with( - keys=expected_keys, args=[60], client=None + keys=tuple(expected_keys), args=[60], client=None ) +# LIT-3298: rate limits tripped at ~40M instead of 80M. async_register_script +# registered the Lua script once at startup and stored the object on the +# limiter, so a request running on a different event loop awaited a script bound +# to the startup loop's connection -> "got Future attached to a different loop". +# The limiter then fell back to a pipeline that reset the window TTL, so two +# minutes of tokens piled into one window. The script must instead be registered +# lazily against the calling loop's client and cached per loop. + + +@pytest.mark.parametrize("namespace", [None, "litellm_sandbox"]) +def test_async_register_script_binds_per_event_loop(namespace, monkeypatch): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + + clients_built = [] + + def make_client(): + client = MagicMock() + client.register_script = MagicMock(return_value=AsyncMock(return_value="ok")) + clients_built.append(client) + return client + + unique_script = "return 'lit3298'" + + with patch.object(redis_cache, "init_async_client", side_effect=make_client): + script = redis_cache.async_register_script(unique_script) + + # Registration is deferred: no client is touched until the script runs. + assert clients_built == [] + + # Two loops kept alive at once so their ids can't be recycled into one + # cache key. The buggy version reuses the first loop's bound object. + loop_a = asyncio.new_event_loop() + loop_b = asyncio.new_event_loop() + try: + result_a = loop_a.run_until_complete( + script(keys=["{k:v}:tokens"], args=[60]) + ) + result_b = loop_b.run_until_complete( + script(keys=["{k:v}:tokens"], args=[60]) + ) + finally: + loop_a.close() + loop_b.close() + + assert result_a == "ok" + assert result_b == "ok" + assert len(clients_built) == 2 + for client in clients_built: + client.register_script.assert_called_once_with(unique_script) + + +@pytest.mark.asyncio +async def test_async_register_script_not_shared_across_namespaces( + monkeypatch, redis_no_ping +): + """Two caches with different namespaces registering the SAME script must + each run against their own client and key prefix. A content-only executor + cache would let the second cache reuse the first's executor and namespace.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + cache_a = RedisCache(namespace="ns_a") + cache_b = RedisCache(namespace="ns_b") + + reg_a = AsyncMock(return_value="a") + client_a = MagicMock() + client_a.register_script = MagicMock(return_value=reg_a) + reg_b = AsyncMock(return_value="b") + client_b = MagicMock() + client_b.register_script = MagicMock(return_value=reg_b) + + same_script = "return redis.call('GET', KEYS[1])" + with patch.object( + cache_a, "init_async_client", return_value=client_a + ), patch.object(cache_b, "init_async_client", return_value=client_b): + script_a = cache_a.async_register_script(same_script) + script_b = cache_b.async_register_script(same_script) + result_a = await script_a(keys=["k"], args=[]) + result_b = await script_b(keys=["k"], args=[]) + + assert (result_a, result_b) == ("a", "b") + reg_a.assert_awaited_once_with(keys=("ns_a:k",), args=[], client=None) + reg_b.assert_awaited_once_with(keys=("ns_b:k",), args=[], client=None) + + +@pytest.mark.asyncio +async def test_async_register_script_cluster_path_uses_evalsha( + monkeypatch, redis_no_ping +): + """Redis Cluster exposes script_load/evalsha rather than register_script. + The script is loaded once and invoked via evalsha with namespaced keys.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="ns") + + cluster_client = MagicMock(spec=["script_load", "evalsha"]) + cluster_client.script_load = MagicMock(return_value="sha123") + cluster_client.evalsha = AsyncMock(return_value="cluster-ok") + + with patch.object( + redis_cache, "init_async_client", return_value=cluster_client + ): + script = redis_cache.async_register_script("return 'cluster'") + result = await script(keys=["{k:v}:tokens"], args=[5, 60]) + + assert result == "cluster-ok" + cluster_client.script_load.assert_called_once_with("return 'cluster'") + cluster_client.evalsha.assert_awaited_once_with( + "sha123", 1, "ns:{k:v}:tokens", 5, 60 + ) + + +@pytest.mark.asyncio +async def test_async_register_script_raises_for_unsupported_client( + monkeypatch, redis_no_ping +): + """A client exposing neither register_script nor script_load fails loudly + rather than silently returning a no-op callable.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + bad_client = MagicMock(spec=[]) + + with patch.object(redis_cache, "init_async_client", return_value=bad_client): + script = redis_cache.async_register_script("return 'x'") + with pytest.raises(ValueError, match="does not support Lua script"): + await script(keys=["k"], args=[1]) + + @pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) @pytest.mark.asyncio async def test_async_delete_cache_namespaces_key(