From 6c6db971dd289db2acbbe4d4c8704493acb339e6 Mon Sep 17 00:00:00 2001 From: yuneng Date: Sat, 10 Oct 2026 21:32:40 +0000 Subject: [PATCH] fix(caching): reload missing Lua scripts in one pipeline so concurrent reservations stay whole --- litellm/caching/redis_batch.py | 75 +++++++++++++++++------- tests/unit/caching/test_redis_batch.py | 81 +++++++++++++++++++++++++- 2 files changed, 135 insertions(+), 21 deletions(-) diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index aa458c9c926..e300bf946a3 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -19,8 +19,15 @@ from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence from contextvars import ContextVar, Token from dataclasses import dataclass, field from datetime import timedelta +from itertools import accumulate from types import MappingProxyType, TracebackType -from typing import Final, Generic, Protocol, TypeVar +from typing import ( + Final, + Generic, + Protocol, + TypeVar, + cast, # noqa: TID251 # Redis pipeline stubs require a typed boundary +) from litellm._internal_context import ( REDIS_FAMILIES_METADATA_KEY, @@ -51,6 +58,7 @@ class RegisteredScript(Protocol): class _RedisPipeline(Protocol): def mget(self, keys: Sequence[str]) -> object: ... def evalsha(self, sha: str, numkeys: int, *keys_and_args: _ScriptArg) -> object: ... + def script_load(self, script: str) -> object: ... def incrbyfloat(self, name: str, amount: float) -> object: ... def expire(self, name: str, time: timedelta) -> object: ... def set(self, name: str, value: str, ex: timedelta | None = None) -> object: ... @@ -90,22 +98,25 @@ class _Op(Generic[_T]): def resolve(self, replies: Sequence[object]) -> _T: raise NotImplementedError + @property + def script_source(self) -> str | None: + return None + async def run_alone(self) -> _T: raise NotImplementedError - def settle(self, replies: Sequence[object]) -> Awaitable[None] | None: - """Resolve from pipeline replies; return a coroutine when the op has to be retried on its own.""" + def settle(self, replies: Sequence[object]) -> bool: failure: Final = next((reply for reply in replies if isinstance(reply, Exception)), None) if failure is None: try: self.future.set_result(self.resolve(replies)) except Exception as e: # noqa: BLE001 # a reply this op cannot decode fails this op alone self.future.set_exception(e) - return None + return False if _is_missing_script(failure): - return self._settle_alone() + return True self.future.set_exception(failure) - return None + return False async def _settle_alone(self) -> None: try: @@ -156,7 +167,7 @@ class _MGet(_Op[Mapping[str, object]]): class _Script(_Op[object]): - __slots__ = ("_args", "_keys", "_redis_cache", "_run", "_sha") + __slots__ = ("_args", "_keys", "_redis_cache", "_run", "_sha", "_source") def __init__( self, @@ -167,12 +178,17 @@ class _Script(_Op[object]): args: Sequence[_ScriptArg], ) -> None: super().__init__() + self._source: Final = source self._redis_cache: Final = redis_cache self._sha: Final = hashlib.sha1(source.encode()).hexdigest() # noqa: S324 # EVALSHA identifies scripts by SHA-1 self._run: Final = run self._keys: Final[tuple[str, ...]] = tuple(keys) self._args: Final[tuple[_ScriptArg, ...]] = tuple(args) + @property + def script_source(self) -> str: + return self._source + def enqueue(self, pipe: _RedisPipeline) -> int: namespaced: Final = tuple(self._redis_cache.check_and_fix_namespace(key=key) for key in self._keys) pipe.evalsha(self._sha, len(namespaced), *namespaced, *self._args) @@ -367,16 +383,27 @@ class RedisBatch: op.future.cancel() await asyncio.gather(*(op.run_settled_hooks() for op in ops)) - async def _flush_pipeline(self, ops: Sequence[_Op[object]]) -> None: + async def _flush_pipeline( + self, + ops: Sequence[_Op[object]], + scripts_to_load: tuple[str, ...] = (), + reloaded: bool = False, + ) -> None: start_time: Final = time.time() target, metadata = _pipeline_service_event(ops) widths: list[int] = [] # mutable-ok: filled while enqueuing - async def run() -> list[object]: + async def run() -> Sequence[object]: client: Final = self.redis_cache.init_async_client() async with client.pipeline(transaction=False) as pipe: - widths.extend(op.enqueue(pipe) for op in ops) - return await pipe.execute(raise_on_error=False) + pipeline: Final = cast( # cast-ok: Redis Cluster flushes take the separate path + _RedisPipeline, pipe + ) + for source in scripts_to_load: + pipeline.script_load(source) + widths.extend(op.enqueue(pipeline) for op in ops) + replies: Final = await pipeline.execute(raise_on_error=False) + return replies[len(scripts_to_load) :] try: replies: Final = await _run_under_circuit_breaker(self.redis_cache._circuit_breaker, self.name, run) # pyright: ignore[reportPrivateUsage] # same breaker as the cache's own methods @@ -408,15 +435,23 @@ class RedisBatch: event_metadata=metadata, ) ) - retries: list[Awaitable[None]] = [] # mutable-ok: collected while slicing replies - offset = 0 - for op, width in zip(ops, widths): - retry = op.settle(replies[offset : offset + width]) - offset += width - if retry is not None: - retries.append(retry) - if retries: - await asyncio.gather(*retries) + reply_offsets: Final = tuple(accumulate(widths, initial=0)) + missing_ops: Final = tuple( + op + for index, op in enumerate(ops) + if op.settle(replies[reply_offsets[index] : reply_offsets[index + 1]]) + ) + if not missing_ops: + return + if reloaded: + await asyncio.gather( + *(op._settle_alone() for op in missing_ops) # pyright: ignore[reportPrivateUsage] # batch owns its ops + ) + return + sources_to_load: Final = tuple( + dict.fromkeys(source for op in missing_ops if (source := op.script_source) is not None) + ) + await self._flush_pipeline(missing_ops, sources_to_load, reloaded=True) MIXED_PIPELINE_TARGET: Final = "mixed" diff --git a/tests/unit/caching/test_redis_batch.py b/tests/unit/caching/test_redis_batch.py index cd270035416..c3a5ce54c39 100644 --- a/tests/unit/caching/test_redis_batch.py +++ b/tests/unit/caching/test_redis_batch.py @@ -7,7 +7,7 @@ import hashlib import json from collections.abc import Awaitable, Callable, Sequence from datetime import timedelta -from typing import Any +from typing import Any, Final import pytest from redis.exceptions import NoScriptError @@ -52,6 +52,10 @@ class FakePipeline: self.commands.append(("EVALSHA", sha, numkeys, *keys_and_args)) return self + def script_load(self, script: str) -> FakePipeline: + self.commands.append(("SCRIPT LOAD", script)) + return self + def incrbyfloat(self, name: str, amount: float) -> FakePipeline: self.commands.append(("INCRBYFLOAT", name, amount)) return self @@ -134,6 +138,8 @@ def replies(command: tuple[Any, ...]) -> Any: match command[0]: case "MGET": return [json.dumps({"k": key}) if key.endswith("hit") else None for key in command[1:]] + case "SCRIPT LOAD": + return hashlib.sha1(str(command[1]).encode()).hexdigest() case "EVALSHA": return [1, 2] case "INCRBYFLOAT": @@ -288,6 +294,79 @@ async def test_noscript_reply_reruns_that_script_through_the_registered_executor assert batch.flushes == 1 +@pytest.mark.asyncio +async def test_scripts_missing_on_the_server_are_loaded_and_rerun_together_so_concurrent_reservations_stay_whole() -> None: + class ReservationServer: + def __init__(self) -> None: + self.loaded: set[str] = set() + self.counters: dict[str, int] = {"key": 0, "team": 0} + self.arrivals: list[ + tuple[ + tuple[str, ...], + tuple[str | bytes | int | float, ...], + asyncio.Future[object], + ] + ] = [] + + def reserve(self, key: str, increment: int, limit: int) -> int: + current: Final = self.counters[key] + if current + increment > limit: + return 0 + self.counters[key] = current + increment + return 1 + + def reply(self, command: tuple[object, ...]) -> object: + match command[0]: + case "SCRIPT LOAD": + source: Final = command[1] + assert isinstance(source, str) + loaded_sha: Final = hashlib.sha1(source.encode()).hexdigest() + self.loaded.add(loaded_sha) + return loaded_sha + case "EVALSHA": + script_sha: Final = command[1] + key: Final = command[3] + increment: Final = command[4] + limit: Final = command[5] + assert isinstance(script_sha, str) + assert isinstance(key, str) + assert isinstance(increment, int) + assert isinstance(limit, int) + if script_sha not in self.loaded: + return NoScriptError("NOSCRIPT No matching script") + return self.reserve(key, increment, limit) + case _: + return replies(command) + + async def run( + self, keys: Sequence[str], args: Sequence[str | bytes | int | float] + ) -> object: + self.loaded.add(SHA) + result: Final[asyncio.Future[object]] = asyncio.get_running_loop().create_future() + self.arrivals.append((tuple(keys), tuple(args), result)) + if len(self.arrivals) == 4: + for index in (0, 3, 1, 2): + arrival_keys, arrival_args, pending = self.arrivals[index] + pending.set_result( + self.reserve(arrival_keys[0], int(arrival_args[0]), int(arrival_args[1])) + ) + return await result + + server: Final = ReservationServer() + cache: Final = FakeRedisCache(FakeClient(server.reply)) + batch_a: Final = RedisBatch(cache) + batch_b: Final = RedisBatch(cache) + a_key: Final = batch_a.script(SCRIPT, server.run, ["key"], [72, 100]) + a_team: Final = batch_a.script(SCRIPT, server.run, ["team"], [72, 100]) + b_key: Final = batch_b.script(SCRIPT, server.run, ["key"], [72, 100]) + b_team: Final = batch_b.script(SCRIPT, server.run, ["team"], [72, 100]) + + async with asyncio.timeout(5): + await asyncio.gather(batch_a.flush(), batch_b.flush()) + + assert sorted(((await a_key, await a_team), (await b_key, await b_team))) == [(0, 0), (1, 1)] + + @pytest.mark.asyncio async def test_cluster_cache_runs_each_operation_on_its_own_path() -> None: client = FakeClient(replies)