fix(redis): loop-scope async Lua script registration (#31501)

* fix(redis): loop-scope async Lua script registration

async_register_script registered the Lua script eagerly and returned a
callable bound to the Redis client of the event loop running at
registration time. The v3 parallel request limiter registers its three
scripts once in __init__ at proxy startup and stores them, so a request
or logging callback on another loop awaited a script bound to the startup
loop and hit "got Future attached to a different loop". The limiter then
fell back to a pipeline that reset the window TTL every increment, so
counters never expired and an 80M TPM model rate-limited around 40M.

Defer registration to call time and cache the per-loop executor in
in_memory_llm_clients_cache (which already keys on the running loop), so
each loop runs the script against its own client. Covers all five
consumers of the primitive.

Resolves LIT-3298

* fix(redis): await evalsha on the cluster Lua script path

The cluster branch returned the evalsha coroutine without awaiting it, so
callers received a coroutine instead of the script result. Await it, which
also addresses the cluster path called out in review.

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):

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(