diff --git a/litellm/redis/transport/contract.py b/litellm/redis/transport/contract.py index a7c8c2a0192..af8c86f1902 100644 --- a/litellm/redis/transport/contract.py +++ b/litellm/redis/transport/contract.py @@ -1,3 +1,4 @@ +import asyncio import math from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence from dataclasses import dataclass, field @@ -11,7 +12,6 @@ InfoValue: TypeAlias = InfoScalar | Mapping[str, InfoScalar] Popped: TypeAlias = bytes | tuple[bytes, ...] | None _T: Final = TypeVar("_T") -_T_co: Final = TypeVar("_T_co", covariant=True) class Admit(Protocol): @@ -42,14 +42,20 @@ def ttl_ms(ttl: timedelta) -> int: return max(1, math.ceil(ttl / timedelta(milliseconds=1))) -class BatchReply(Protocol[_T_co]): - def __await__(self) -> Generator[object, None, _T_co]: ... +class BatchReply(Protocol[_T]): + def __await__(self) -> Generator[object, None, _T]: ... @property def done(self) -> bool: ... + def on_settled(self, hook: Callable[[asyncio.Future[_T]], Awaitable[None] | None]) -> None: ... + class RedisCommandBatch(Protocol): + def add_flush_hook(self, hook: Callable[[], None]) -> None: ... + + def read_as_missing(self, key: str) -> bool: ... + def get(self, key: str) -> BatchReply[bytes | None | RedisFailure]: ... def mget(self, keys: Sequence[str]) -> BatchReply[tuple[bytes | None, ...] | RedisFailure]: ... diff --git a/litellm/redis/transport/redis_py.py b/litellm/redis/transport/redis_py.py index 99819448757..9641f5c2bca 100644 --- a/litellm/redis/transport/redis_py.py +++ b/litellm/redis/transport/redis_py.py @@ -215,6 +215,12 @@ class RedisPyCommandBatch: transport: RedisPyTransport engine: PipelineBatch + def add_flush_hook(self, hook: Callable[[], None]) -> None: + self.engine.add_flush_hook(hook) + + def read_as_missing(self, key: str) -> bool: + return self.engine.read_as_missing(key) + def get(self, key: str) -> BatchResult[bytes | None | RedisFailure]: name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own reply: Final = self._declare( diff --git a/tests/unit/redis/transport/test_redis_py.py b/tests/unit/redis/transport/test_redis_py.py index 633f4274be0..2d4390344ec 100644 --- a/tests/unit/redis/transport/test_redis_py.py +++ b/tests/unit/redis/transport/test_redis_py.py @@ -16,15 +16,93 @@ from redis.exceptions import ResponseError import litellm from litellm._internal_context import current_service_target, service_target from litellm._service_logger import ServiceLogging -from litellm.caching.redis_batch import active_request_redis_batch, request_redis_batch_scope +from litellm.caching.redis_batch import ( + active_request_redis_batch, + flush_post_call_redis_batches, + request_redis_batch_scope, +) from litellm.caching.redis_cluster_cache import ClusterConnectionGate, RedisClusterCache from litellm.redis.transport.batch import PipelineBatch -from litellm.redis.transport.contract import PubSubMessage, RedisFailure, raise_redis_failure +from litellm.redis.transport.contract import ( + BatchReply, + PubSubMessage, + RedisFailure, + RedisTransport, + raise_redis_failure, +) from litellm.redis.transport.redis_py import RedisPyCommandBatch, RedisPyTransport from litellm.types.services import ServiceTypes from tests._support.fake_redis import fake_redis_cache, fake_redis_transport +@dataclass +class _LazyRead: + transport: RedisTransport + reply: BatchReply[tuple[bytes | None, ...] | RedisFailure] | None = None + + def declare(self) -> None: + batch: Final = self.transport.request_batch() + if batch is not None and self.reply is None: + self.reply = batch.mget(("present", "absent")) + + +@pytest.mark.asyncio +async def test_transport_flush_hook_joins_a_legacy_request_read( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server, namespace="ns") + transport: Final[RedisTransport] = cache.transport() + await raw.set("ns:present", "3") + with request_redis_batch_scope(): + batch: Final = transport.request_batch() + legacy: Final = active_request_redis_batch(cache) + assert batch is not None and legacy is not None + reader: Final = _LazyRead(transport) + batch.add_flush_hook(reader.declare) + assert await legacy.mget(("present",)) == {"present": 3} + assert reader.reply is not None + assert await reader.reply == (b"3", None) + assert legacy.flushes == 1 + assert batch.read_as_missing("absent") is True + assert batch.read_as_missing("present") is False + assert await transport.set("absent", "new", ttl=None) is True + assert await batch.set("absent", "new", ttl=None) is True + assert batch.read_as_missing("absent") is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outage", (False, True)) +async def test_post_call_settled_hook_can_read_every_reply_and_rollback( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, outage: bool +) -> None: + transport: Final[RedisTransport] = fake_redis_cache(server, namespace="ns").transport() + with request_redis_batch_scope(): + batch: Final = transport.post_call_batch() + assert batch is not None + increment: Final = batch.incrbyfloat("counter", 2.5) + expired: Final = batch.expire("counter", timedelta(seconds=30)) + observed: Final[asyncio.Future[tuple[float | RedisFailure, bool | RedisFailure]]] = ( + asyncio.get_running_loop().create_future() + ) + + async def settle(future: asyncio.Future[float | RedisFailure]) -> None: + assert expired.done + observed.set_result((future.result(), await expired)) + server.connected = True + await transport.delete("counter") + + increment.on_settled(settle) + server.connected = not outage + await flush_post_call_redis_batches() + first, second = await observed + if outage: + assert isinstance(first, RedisFailure) and first.kind == "transport" + assert isinstance(second, RedisFailure) and second.kind == "transport" + else: + assert (first, second) == (2.5, True) + assert await raw.get("ns:counter") is None + + @pytest.mark.asyncio async def test_every_command_prefixes_the_namespace_once( server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis