diff --git a/.agents/skills/redis-access/SKILL.md b/.agents/skills/redis-access/SKILL.md index 9e8b319b56c..a3e38f0b944 100644 --- a/.agents/skills/redis-access/SKILL.md +++ b/.agents/skills/redis-access/SKILL.md @@ -24,18 +24,25 @@ metadata: ## The transport -`RedisTransport` in `litellm/redis/transport/contract.py` is how Python keyspace code reaches redis-py. It is a Protocol of plain methods, one per Redis command the Python code uses (`get`, `mget`, `set`, `delete`, `eval`), command-shaped and never redis-py shaped +`RedisTransport` in `litellm/redis/transport/contract.py` is how Python code reaches redis-py. It is a Protocol of plain methods, command-shaped and never redis-py shaped, and it covers every Redis use in Python, so a keyspace move never waits on the transport + +- Data commands `get`, `mget`, `set`, `delete`, `incrby`, `incrbyfloat`, `expire`, `ttl`, `rpush`, `lpop`, `ltrim`, `sadd`, `scan` and `eval` +- Batches, `batch()` for one caller's commands, `batch(transaction=True)` for MULTI/EXEC, and `request_batch()` / `post_call_batch()` for the request scope +- Pub/sub, `publish`, `publish_all` and `subscribe` +- Admin and health, `ping`, `info`, `client_list` and `flushall` - Every method returns its value or a `RedisFailure`. Nothing raises for a Redis problem - `RedisFailure.kind` is `transport` (timeout or lost connection, the command may still have applied, so a failed write is an unknown outcome, never "not applied"), `server` (Redis answered with an error, the connection is fine and the breaker does not count it), `decode` (the reply was not the shape expected) or `breaker_open` (refused without touching the socket) -- The namespace is applied once, inside the transport, to every key including Lua `KEYS`. A key that already starts with `{namespace}:` is left alone. Stores and callers pass logical keys +- The namespace is applied once, inside the transport, to every key including Lua `KEYS` and pub/sub channels. A key that already starts with `{namespace}:` is left alone. Stores and callers pass logical keys - Every write takes its TTL from the caller as a `timedelta`. `None` means no expiry, there is no default - Replies are undecoded bytes. `eval` returns `int | bytes | None | tuple[...]`. Decoding belongs to the store that owns the keyspace - `eval` takes script text every call and does `EVALSHA` with `EVAL` fallback, so nothing registers scripts - These rules are shared with `crates/redis-client`, whose `AGENTS.md` has the Rust side. Change both together, since Python and Rust pods share keys in one Redis -- `RedisCache.transport()` builds the transport over that cache's own client, namespace, circuit breaker, `_admit_async_call` and `ServiceLogging`. No new connection, no second breaker, and `RedisClusterCache`'s connection gate still applies +- `RedisCache.transport()` builds the transport over that cache's own client, pub/sub client, namespace, circuit breaker, `_admit_async_call`, cluster-aware `mget` and `ServiceLogging`. No new connection, no second breaker, and `RedisClusterCache`'s connection gate still applies - Every method runs its command through two steps, in this order. `run_under_breaker` wraps the existing `_run_under_circuit_breaker` from `litellm/caching/redis_cache.py`, so admission, the generation check, the health failure classification and the timeout flag are the ones legacy traffic uses, and maps the outcome to a `RedisFailure` (`breaker_open`, `transport`, `server`) instead of raising. `RedisServiceEvents.observe` times the call, logs a failure and emits the service event +- `ping`, the admin commands and pub/sub skip the breaker, like the `RedisCache` methods they replace, so health checks still reach a Redis the breaker took out. They still go through `RedisServiceEvents.observe`, except `get_message` polling + ```python @dataclass(frozen=True, slots=True) class RedisPyTransport: @@ -44,6 +51,7 @@ class RedisPyTransport: breaker: RedisCircuitBreaker admit: Admit events: RedisServiceEvents + ... async def get(self, key: str) -> bytes | None | RedisFailure: return _bytes_reply(await self._run("get", lambda: self.client().get(self._key(key)))) @@ -51,9 +59,20 @@ class RedisPyTransport: async def _run(self, name: str, call: Callable[[], Awaitable[_T]]) -> _T | RedisFailure: return await self.events.observe(name, lambda: run_under_breaker(self.breaker, name, self.admit, call)) ``` -- Adding a command means a method on the Protocol, an implementation in `redis_py.py`, and a fakeredis test that the namespace is applied once, the TTL lands, a rejected command is `server` and an outage is `transport`. Before adding one, check whether the operation is really one `eval`. A counter with a TTL on first write or an owner-checked release is one script, not three round trips +- A Redis command Python starts using is a method on the Protocol and on the batch, an implementation in `redis_py.py`, and a fakeredis test that the namespace is applied once, the TTL lands, a rejected command is `server` and an outage is `transport` +- A store still picks the fewest round trips. A counter with a TTL on first write or an owner-checked release is one `eval`, not three commands - Multi-key commands and scripts keep their keys in one slot, or they break on cluster depending on slot ownership +## Batches + +The pipeline engine is `PipelineBatch` in `litellm/redis/transport/batch.py`, the one `RedisBatch` has always had. Transport batches and `RedisBatch` (`litellm/caching/redis_batch.py`) declare onto it, so a request keeps one round trip per backend while callers move over one at a time + +- Every batch method returns an awaitable reply, and awaiting any reply flushes everything declared so far. Each reply is its value or a `RedisFailure`, so one rejected command fails only its own reply, and an outage fails every reply and counts once on the breaker +- `request_batch()` and `post_call_batch()` return the engine `request_redis_batch_scope` holds for this backend, or `None` outside a request. Never build a second request scope +- A script Redis has not cached reruns on its own after the pipeline, so do not rely on its effect being visible to later commands in the same batch. A transaction batch sends `EVAL` instead, so its scripts run inside the MULTI +- On a cluster client a request batch runs each command on its own, as `RedisBatch` always has. Explicit and transaction batches still pipeline +- One REDIS event per flush, named by the batch and targeted by the key families that shared it, `mixed` with the sorted families when there are several + ## Stores A store owns one keyspace. Python and Rust pods share one Redis, so each keyspace needs one definition of its keys and values, and the store is that definition @@ -73,7 +92,7 @@ A store owns one keyspace. Python and Rust pods share one Redis, so each keyspac - Delete the caller's own timing, service hooks and Redis health logging. The transport emits for it now - Move one keyspace per change. An extraction is not a Rust port - `RedisCache` is a connection object wearing a cache name. Its own cache verbs move onto the transport one method at a time, each conversion deleting that method's inline logging, until what is left is the connection, the breaker and a thin facade -- Not on the transport yet, so leave them on `RedisCache` until the transport grows them: pub/sub, admin commands, pipelines and `RedisBatch`. Vector and semantic search keep their own clients +- Pipelined and batched access moves onto a transport batch, pub/sub onto `publish` / `subscribe`, and admin onto the admin methods. Vector and semantic search keep their own clients ## Observability diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index aa458c9c926..82942dcf9c5 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -1,10 +1,10 @@ -"""One Redis pipeline for several independent operations, each with its own result and its own failure. +"""``RedisCache``-shaped operations on the transport's pipeline engine, and the request scope that hands each +backend one engine per request. -A ``RedisBatch`` collects MGETs, Lua scripts and increments declared by unrelated callers and sends them -in one ``pipeline(transaction=False)`` round trip. Every declaration returns an awaitable; awaiting one -flushes whatever has been declared so far, so callers keep their existing ``await`` shape and their own -error handling while sharing the wire. Redis Cluster clients run each operation on its own, as before: -a cluster pipeline is per node anyway and the existing per-operation paths already group by slot. +A ``RedisBatch`` collects MGETs, Lua scripts and increments declared by unrelated callers and sends them in one +round trip, together with any transport commands declared on the same engine. Every declaration returns an +awaitable; awaiting one flushes whatever has been declared so far, so callers keep their existing ``await`` shape +and their own error handling while sharing the wire. Redis Cluster clients run each operation on its own. """ from __future__ import annotations @@ -12,131 +12,58 @@ from __future__ import annotations import asyncio import hashlib import json -import logging -import time import weakref -from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from contextvars import ContextVar, Token -from dataclasses import dataclass, field from datetime import timedelta from types import MappingProxyType, TracebackType -from typing import Final, Generic, Protocol, TypeVar +from typing import Final, Protocol -from litellm._internal_context import ( - REDIS_FAMILIES_METADATA_KEY, - current_service_target, - service_caller, - service_target, -) -from litellm._logging import verbose_logger -from litellm.caching.redis_cache import ( - RedisCache, - _get_call_stack_info, # pyright: ignore[reportPrivateUsage] # same caller chain every RedisCache method reports - _run_under_circuit_breaker, # pyright: ignore[reportPrivateUsage] # same health signal as every RedisCache method - log_redis_failure, -) +from litellm.caching.redis_cache import RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache -from litellm.types.services import ServiceTypes +from litellm.redis.transport.batch import ( + MIXED_PIPELINE_TARGET, + BatchResult, + PipelineBackend, + PipelineBatch, + SettledHook, +) +from litellm.redis.transport.client import RedisPipeline +from litellm.redis.transport.contract import Arg + +__all__ = ( + "MIXED_PIPELINE_TARGET", + "POST_CALL_FLUSH_DEADLINE_SECONDS", + "BatchResult", + "RedisBatch", + "RegisteredScript", + "RequestRedisBatches", + "SettledHook", + "active_post_call_redis_batch", + "active_request_redis_batch", + "active_request_redis_batches", + "drain_post_call_redis_batches", + "flush_post_call_redis_batches", + "post_call_batch_engine", + "request_batch_engine", + "request_redis_batch_scope", +) -_T = TypeVar("_T") -_ScriptArg = str | bytes | int | float -SettledHook = Callable[[asyncio.Future[_T]], Awaitable[None] | None] POST_CALL_FLUSH_DEADLINE_SECONDS: Final = 1.0 class RegisteredScript(Protocol): - def __call__(self, keys: Sequence[str], args: Sequence[_ScriptArg]) -> Awaitable[object]: ... + def __call__(self, keys: Sequence[str], args: Sequence[Arg]) -> Awaitable[object]: ... -class _RedisPipeline(Protocol): - def mget(self, keys: Sequence[str]) -> object: ... - def evalsha(self, sha: str, numkeys: int, *keys_and_args: _ScriptArg) -> 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: ... - def delete(self, *names: str) -> object: ... - async def execute(self, raise_on_error: bool = True) -> list[object]: ... - - -class _Op(Generic[_T]): - """One declared operation: how many pipeline replies it consumes, how to turn them into a result, and - how to run on its own when the batch cannot pipeline (cluster client, or a reply the pipeline cannot - settle, like NOSCRIPT).""" - - __slots__ = ("caller", "future", "settled_hooks", "target") - - def __init__(self) -> None: - self.future: Final[asyncio.Future[_T]] = asyncio.get_running_loop().create_future() - self.future.add_done_callback(_mark_retrieved) - self.settled_hooks: Final[list[SettledHook[_T]]] = [] # mutable-ok: append-only registry - self.target: Final = current_service_target() - self.caller: Final = _get_call_stack_info() - - async def run_settled_hooks(self) -> None: - for hook in self.settled_hooks: - await self._run_settled_hook(hook) - - async def _run_settled_hook(self, hook: SettledHook[_T]) -> None: - try: - follow_up: Final = hook(self.future) - if follow_up is not None: - await follow_up - except Exception as e: # noqa: BLE001 # one owner's follow-up must not stop the others - verbose_logger.warning("redis batch settled hook failed: %s", e) - - def enqueue(self, pipe: _RedisPipeline) -> int: - raise NotImplementedError - - def resolve(self, replies: Sequence[object]) -> _T: - raise NotImplementedError - - 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.""" - 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 - if _is_missing_script(failure): - return self._settle_alone() - self.future.set_exception(failure) - return None - - async def _settle_alone(self) -> None: - try: - with service_target(self.target), service_caller(self.caller): - self.future.set_result(await self.run_alone()) - except Exception as e: # noqa: BLE001 # the declaring caller owns the failure of its own operation - self.future.set_exception(e) - - -def _is_missing_script(failure: Exception) -> bool: - """Imported lazily: this module is reachable from a base ``import litellm`` while redis is not a base dependency.""" - from redis.exceptions import NoScriptError - - return isinstance(failure, NoScriptError) - - -def _mark_retrieved(future: asyncio.Future[object]) -> None: - """A caller that stops awaiting (cancelled request) must not leave an 'exception never retrieved' log.""" - if not future.cancelled(): - future.exception() - - -class _MGet(_Op[Mapping[str, object]]): +class _MGet: __slots__ = ("_keys", "_redis_cache") def __init__(self, redis_cache: RedisCache, keys: Sequence[str]) -> None: - super().__init__() self._redis_cache: Final = redis_cache self._keys: Final[tuple[str, ...]] = tuple(dict.fromkeys(keys)) - def enqueue(self, pipe: _RedisPipeline) -> int: + def enqueue(self, pipe: RedisPipeline) -> int: pipe.mget(tuple(self._redis_cache.check_and_fix_namespace(key=key) for key in self._keys)) return 1 @@ -148,6 +75,9 @@ class _MGet(_Op[Mapping[str, object]]): {key: self._redis_cache._get_cache_logic(value) for key, value in zip(self._keys, values)} # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownArgumentType] # shared decode with async_batch_get_cache ) + def recover(self, error: Exception) -> Mapping[str, object]: + raise error + async def run_alone(self) -> Mapping[str, object]: found: Mapping[str, object] = await self._redis_cache.async_batch_get_cache(key_list=list(self._keys)) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped cache API if any(key not in found for key in self._keys): @@ -155,7 +85,7 @@ class _MGet(_Op[Mapping[str, object]]): return found -class _Script(_Op[object]): +class _Script: __slots__ = ("_args", "_keys", "_redis_cache", "_run", "_sha") def __init__( @@ -164,16 +94,15 @@ class _Script(_Op[object]): source: str, run: RegisteredScript, keys: Sequence[str], - args: Sequence[_ScriptArg], + args: Sequence[Arg], ) -> None: - super().__init__() 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) + self._args: Final[tuple[Arg, ...]] = tuple(args) - def enqueue(self, pipe: _RedisPipeline) -> int: + 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) return 1 @@ -181,21 +110,23 @@ class _Script(_Op[object]): def resolve(self, replies: Sequence[object]) -> object: return replies[0] + def recover(self, error: Exception) -> object: + raise error + async def run_alone(self) -> object: return await self._run(keys=self._keys, args=self._args) -class _Increment(_Op[float]): +class _Increment: __slots__ = ("_key", "_redis_cache", "_ttl", "_value") def __init__(self, redis_cache: RedisCache, key: str, value: float, ttl: int | None) -> None: - super().__init__() self._redis_cache: Final = redis_cache self._key: Final = key self._value: Final = value self._ttl: Final = ttl - def enqueue(self, pipe: _RedisPipeline) -> int: + def enqueue(self, pipe: RedisPipeline) -> int: name: Final = self._redis_cache.check_and_fix_namespace(key=self._key) pipe.incrbyfloat(name, self._value) if self._ttl is None: @@ -209,6 +140,9 @@ class _Increment(_Op[float]): raise TypeError(f"INCRBYFLOAT reply is not numeric: {type(reply).__name__}") return float(reply) + def recover(self, error: Exception) -> float: + raise error + async def run_alone(self) -> float: value: object = await self._redis_cache.async_increment(key=self._key, value=self._value, ttl=self._ttl) # pyright: ignore[reportUnknownMemberType] # untyped cache API if not isinstance(value, (int, float)): @@ -216,19 +150,18 @@ class _Increment(_Op[float]): return float(value) -class _Set(_Op[None]): +class _Set: """SET with the cache's TTL rules, same encoding as ``async_set_cache_pipeline_with_ttls``.""" __slots__ = ("_key", "_redis_cache", "_ttl", "_value") def __init__(self, redis_cache: RedisCache, key: str, value: object, ttl: float | None) -> None: - super().__init__() self._redis_cache: Final = redis_cache self._key: Final = key self._value: Final = value self._ttl: Final = ttl - def enqueue(self, pipe: _RedisPipeline) -> int: + def enqueue(self, pipe: RedisPipeline) -> int: ttl: Final = self._redis_cache.get_ttl(ttl=self._ttl) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped cache API pipe.set( self._redis_cache.check_and_fix_namespace(key=self._key), @@ -240,195 +173,104 @@ class _Set(_Op[None]): def resolve(self, replies: Sequence[object]) -> None: return None + def recover(self, error: Exception) -> None: + raise error + async def run_alone(self) -> None: await self._redis_cache.async_set_cache_pipeline_with_ttls(((self._key, self._value, self._ttl),)) -class _Delete(_Op[None]): +class _Delete: """DEL of one key, the pipelined twin of ``async_delete_cache``.""" __slots__ = ("_key", "_redis_cache") def __init__(self, redis_cache: RedisCache, key: str) -> None: - super().__init__() self._redis_cache: Final = redis_cache self._key: Final = key - def enqueue(self, pipe: _RedisPipeline) -> int: + def enqueue(self, pipe: RedisPipeline) -> int: pipe.delete(self._redis_cache.check_and_fix_namespace(key=self._key)) return 1 def resolve(self, replies: Sequence[object]) -> None: return None + def recover(self, error: Exception) -> None: + raise error + async def run_alone(self) -> None: await self._redis_cache.async_delete_cache(self._key) -class BatchResult(Generic[_T]): - """Awaitable handle for one declared operation; awaiting it flushes the batch it belongs to.""" - - __slots__ = ("_batch", "_op") - - def __init__(self, batch: RedisBatch, op: _Op[_T]) -> None: - self._batch: Final = batch - self._op: Final = op - - def __await__(self) -> Generator[object, None, _T]: - return self._wait().__await__() - - async def _wait(self) -> _T: - if not self._op.future.done(): - await self._batch.flush() - return self._op.future.result() - - @property - def done(self) -> bool: - return self._op.future.done() - - def on_settled(self, hook: SettledHook[_T]) -> None: - """For an owner that does not await: runs inside the flush once this operation has its result or - failure (or was cancelled with the pipeline), so the flush completes with the follow-up done.""" - self._op.settled_hooks.append(hook) +def pipeline_backend(redis_cache: RedisCache) -> PipelineBackend: + return PipelineBackend( + client=redis_cache.init_async_client, + breaker=redis_cache._circuit_breaker, # pyright: ignore[reportPrivateUsage] # same breaker as the cache's own methods + admit=redis_cache._admit_async_call, # pyright: ignore[reportPrivateUsage] # same connection gate as the cache's own methods + service_logging=redis_cache.service_logger_obj, + runs_alone=isinstance(redis_cache, RedisClusterCache), + ) -@dataclass(slots=True) class RedisBatch: """Operations declared here go out in one pipeline the next time any of them is awaited or ``flush`` runs.""" - redis_cache: RedisCache - name: str = "redis_batch" - _pending: list[_Op[object]] = field(default_factory=list) # mutable-ok: drained by flush - _flush_hooks: list[Callable[[], None]] = field(default_factory=list) # mutable-ok: append-only registry - _lock: asyncio.Lock = field(default_factory=asyncio.Lock) - _misses: set[str] = field(default_factory=set) # mutable-ok: keys an MGET of this request read as absent - flushes: int = 0 + __slots__ = ("engine", "redis_cache") + + def __init__(self, redis_cache: RedisCache, name: str = "redis_batch") -> None: + self.redis_cache: Final = redis_cache + self.engine: Final = PipelineBatch(pipeline_backend(redis_cache), name=name) + + @property + def name(self) -> str: + return self.engine.name + + @property + def flushes(self) -> int: + return self.engine.flushes + + @property + def pending(self) -> int: + return self.engine.pending def mget(self, keys: Sequence[str]) -> BatchResult[Mapping[str, object]]: - op: Final = _MGet(self.redis_cache, keys) - op.future.add_done_callback(self._note_misses) - return self._declare(op) + result: Final = self.engine.declare(_MGet(self.redis_cache, keys)) + result.on_settled(self._note_misses) + return result def _note_misses(self, future: asyncio.Future[Mapping[str, object]]) -> None: if future.cancelled() or future.exception() is not None: return - self._misses.update(key for key, value in future.result().items() if value is None) + self.engine.note_missing(key for key, value in future.result().items() if value is None) def read_as_missing(self, key: str) -> bool: """True when an MGET on this batch already found no value under ``key`` and nothing has set it since, so a per-key GET later in the same request can be answered without another round trip.""" - return key in self._misses + return self.engine.read_as_missing(key) def script( - self, source: str, run: RegisteredScript, keys: Sequence[str], args: Sequence[_ScriptArg] + self, source: str, run: RegisteredScript, keys: Sequence[str], args: Sequence[Arg] ) -> BatchResult[object]: - return self._declare(_Script(self.redis_cache, source, run, keys, args)) + return self.engine.declare(_Script(self.redis_cache, source, run, keys, args)) def increment(self, key: str, value: float, ttl: int | None = None) -> BatchResult[float]: - return self._declare(_Increment(self.redis_cache, key, value, ttl)) + return self.engine.declare(_Increment(self.redis_cache, key, value, ttl)) def set(self, key: str, value: object, ttl: float | None = None) -> BatchResult[None]: - self._misses.discard(key) - return self._declare(_Set(self.redis_cache, key, value, ttl)) + self.engine.note_present(key) + return self.engine.declare(_Set(self.redis_cache, key, value, ttl)) def delete(self, key: str) -> BatchResult[None]: - self._misses.add(key) - return self._declare(_Delete(self.redis_cache, key)) + self.engine.note_missing((key,)) + return self.engine.declare(_Delete(self.redis_cache, key)) def add_flush_hook(self, hook: Callable[[], None]) -> None: """Called at the start of every flush so lazily bound readers can declare their keys into the same trip.""" - self._flush_hooks.append(hook) - - @property - def pending(self) -> int: - return len(self._pending) - - def _declare(self, op: _Op[_T]) -> BatchResult[_T]: - self._pending.append(op) # pyright: ignore[reportArgumentType] # heterogeneous ops share the flush loop - return BatchResult(self, op) + self.engine.add_flush_hook(hook) async def flush(self) -> None: - async with self._lock: - for hook in self._flush_hooks: - hook() - ops: Final = tuple(self._pending) - self._pending.clear() - if not ops: - return - self.flushes += 1 - try: - if isinstance(self.redis_cache, RedisClusterCache): - await asyncio.gather(*(op._settle_alone() for op in ops)) # pyright: ignore[reportPrivateUsage] # batch owns its ops - else: - await self._flush_pipeline(ops) - finally: - for op in ops: - if not op.future.done(): - op.future.cancel() - await asyncio.gather(*(op.run_settled_hooks() for op in ops)) - - async def _flush_pipeline(self, ops: Sequence[_Op[object]]) -> 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]: - 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) - - 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 - except Exception as e: # noqa: BLE001 # each declaring caller applies its own Redis fallback - log_redis_failure(verbose_logger, logging.WARNING, f"{self.name}: pipeline of {len(ops)} ops failed", e) - with service_target(target): - asyncio.create_task( - self.redis_cache.service_logger_obj.async_service_failure_hook( - service=ServiceTypes.REDIS, - duration=time.time() - start_time, - error=e, - call_type=self.name, - start_time=start_time, - end_time=time.time(), - event_metadata=metadata, - ) - ) - for op in ops: - op.future.set_exception(e) - return - with service_target(target): - asyncio.create_task( - self.redis_cache.service_logger_obj.async_service_success_hook( - service=ServiceTypes.REDIS, - duration=time.time() - start_time, - call_type=self.name, - start_time=start_time, - end_time=time.time(), - 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) - - -MIXED_PIPELINE_TARGET: Final = "mixed" - - -def _pipeline_service_event(ops: Sequence[_Op[object]]) -> tuple[str | None, dict[str, int | str]]: - """The target and metadata of one pipeline flush: the one key family every op was declared under, or - ``"mixed"`` plus the sorted families when owners of several families share the trip.""" - families: Final = sorted({op.target for op in ops if op.target is not None}) - if len(families) > 1: - return MIXED_PIPELINE_TARGET, {"op_count": len(ops), REDIS_FAMILIES_METADATA_KEY: ",".join(families)} - return next(iter(families), None), {"op_count": len(ops)} + await self.engine.flush() def _backend_key(redis_cache: RedisCache) -> object: @@ -576,3 +418,14 @@ class request_redis_batch_scope: ) -> None: if self._token is not None: _active_request_batches.reset(self._token) + + +def request_batch_engine(redis_cache: RedisCache) -> PipelineBatch | None: + """The engine the transport declares request-scoped commands on, shared with this backend's ``RedisBatch``.""" + batch: Final = active_request_redis_batch(redis_cache) + return None if batch is None else batch.engine + + +def post_call_batch_engine(redis_cache: RedisCache) -> PipelineBatch | None: + batch: Final = active_post_call_redis_batch(redis_cache) + return None if batch is None else batch.engine diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 2efd979f462..def36fa5891 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -42,6 +42,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.coroutine_checker import coroutine_checker +from litellm.redis.transport.client import AsyncPubSubClient, AsyncRedisCommands from litellm.redis.transport.contract import RedisFailure, RedisTransport from litellm.types.caching import ( RedisPipelineIncrementOperation, @@ -73,7 +74,7 @@ else: Span = Any -class _AsyncRedisCommands(Protocol): +class _AsyncRedisCommands(AsyncRedisCommands, Protocol): """Async redis commands this cache issues. redis-py's type stubs omit these methods on RedisCluster, so the union returned by @@ -82,12 +83,6 @@ class _AsyncRedisCommands(Protocol): def ping(self) -> Awaitable[bool]: ... - def get(self, name: str) -> Awaitable[object]: ... - - def mget(self, keys: Sequence[str]) -> Awaitable[object]: ... - - def set(self, name: str, value: str | bytes | float, *, nx: bool, px: int | None) -> Awaitable[object]: ... - def delete(self, *names: str) -> Awaitable[int]: ... def ttl(self, name: str) -> Awaitable[int]: ... @@ -96,14 +91,8 @@ class _AsyncRedisCommands(Protocol): def rpush(self, name: str, *values: str | bytes | float) -> Awaitable[int]: ... - def lpop(self, name: str, count: int | None = None) -> Awaitable[object]: ... - def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ... - def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... - - def evalsha(self, sha: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... - _GENERIC_CALLER_MODULES: Final = frozenset( { @@ -868,7 +857,11 @@ class RedisCache(BaseCache): def _async_commands(self) -> _AsyncRedisCommands: return self.init_async_client() + def _pubsub_commands(self) -> AsyncPubSubClient: + return self.init_pubsub_client() # pyright: ignore[reportUnknownVariableType] # redis generics + def transport(self) -> RedisTransport: + from litellm.caching.redis_batch import post_call_batch_engine, request_batch_engine from litellm.redis.transport.events import RedisServiceEvents from litellm.redis.transport.redis_py import RedisPyTransport @@ -878,6 +871,11 @@ class RedisCache(BaseCache): breaker=self._circuit_breaker, admit=self._admit_async_call, events=RedisServiceEvents(self.service_logger_obj), + multi_get=lambda keys: self._async_run_redis_mget_operation(list(keys)), + server_version=self._parse_redis_major_version, + pubsub_client=self._pubsub_commands, + request_engine=lambda: request_batch_engine(self), + post_call_engine=lambda: post_call_batch_engine(self), ) def check_and_fix_namespace(self, key: str) -> str: diff --git a/litellm/redis/AGENTS.md b/litellm/redis/AGENTS.md index 2db518f00c8..1285bdc57d5 100644 --- a/litellm/redis/AGENTS.md +++ b/litellm/redis/AGENTS.md @@ -1,8 +1,9 @@ For adding Redis access or moving raw Redis access here, follow [.agents/skills/redis-access/SKILL.md](../../.agents/skills/redis-access/SKILL.md) -- `transport/contract.py` is the Python seam to redis-py. `RedisTransport` is plain methods, one per Redis command the Python code uses, command shaped and never redis-py shaped. It is not a seam for the Rust client, native stores sit behind a store `Protocol` instead +- `transport/contract.py` is the Python seam to redis-py. `RedisTransport` is plain methods covering every Redis use in Python, data commands, batches, pub/sub and admin, command shaped and never redis-py shaped. It is not a seam for the Rust client, native stores sit behind a store `Protocol` instead +- `transport/client.py` holds the redis-py client shapes the transport drives, as Protocols with no redis import. `transport/batch.py` is the one pipeline engine, shared by transport batches and `RedisBatch` - Rules shared with `crates/redis-client`, so Python and Rust pods agree on one Redis - - The namespace is applied here, once, including Lua `KEYS`. A key that already starts with `{namespace}:` is left alone + - The namespace is applied here, once, including Lua `KEYS` and pub/sub channels. A key that already starts with `{namespace}:` is left alone - Every write takes its TTL from the caller, `None` means no expiry, there is no default - Replies are undecoded bytes, `mget` and `eval` arrays are tuples - Failures are values. `transport` covers timeouts and connection loss, where the command may still have applied. `server` means Redis answered with an error and the connection is fine diff --git a/litellm/redis/transport/batch.py b/litellm/redis/transport/batch.py new file mode 100644 index 00000000000..b62de930cf1 --- /dev/null +++ b/litellm/redis/transport/batch.py @@ -0,0 +1,255 @@ +import asyncio +import time +from collections.abc import Awaitable, Callable, Generator, Iterable, Sequence +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Final, Generic, Protocol, TypeVar + +from litellm._internal_context import ( + REDIS_FAMILIES_METADATA_KEY, + current_service_target, + service_caller, + service_target, +) +from litellm._logging import verbose_logger +from litellm.caching.redis_cache import ( + RedisCircuitBreaker, + _get_call_stack_info, # pyright: ignore[reportPrivateUsage] # same caller chain every RedisCache method reports + _run_under_circuit_breaker, # pyright: ignore[reportPrivateUsage] # same health signal as every RedisCache method +) +from litellm.redis.transport.client import PipelineClient, RedisPipeline +from litellm.redis.transport.contract import Admit + +if TYPE_CHECKING: + from litellm._service_logger import ServiceLogging + +_T = TypeVar("_T") +_T_co = TypeVar("_T_co", covariant=True) +SettledHook = Callable[[asyncio.Future[_T]], Awaitable[None] | None] +MIXED_PIPELINE_TARGET: Final = "mixed" + + +class OpSpec(Protocol[_T_co]): + """How one operation rides a pipeline: the commands it queues, how its replies become its result, what a + failed pipeline or command makes of it (raise to fail the op, return to settle it), and how it runs alone.""" + + def enqueue(self, pipe: RedisPipeline) -> int: ... + + def resolve(self, replies: Sequence[object]) -> _T_co: ... + + def recover(self, error: Exception) -> _T_co: ... + + def run_alone(self) -> Awaitable[_T_co]: ... + + +@dataclass(frozen=True, slots=True) +class PipelineBackend: + client: Callable[[], PipelineClient] + breaker: RedisCircuitBreaker + admit: Admit + service_logging: "ServiceLogging" + runs_alone: bool + + +def _mark_retrieved(future: asyncio.Future[object]) -> None: + if not future.cancelled(): + future.exception() + + +def _is_missing_script(failure: Exception) -> bool: + from litellm.redis.transport.redis_py import is_missing_script + + return is_missing_script(failure) + + +class _Op(Generic[_T]): + __slots__ = ("caller", "future", "settled_hooks", "spec", "target") + + def __init__(self, spec: OpSpec[_T]) -> None: + self.spec: Final = spec + self.future: Final[asyncio.Future[_T]] = asyncio.get_running_loop().create_future() + self.future.add_done_callback(_mark_retrieved) + self.settled_hooks: Final[list[SettledHook[_T]]] = [] # mutable-ok: append-only registry + self.target: Final = current_service_target() + self.caller: Final = _get_call_stack_info() + + async def run_settled_hooks(self) -> None: + for hook in self.settled_hooks: + await self._run_settled_hook(hook) + + async def _run_settled_hook(self, hook: SettledHook[_T]) -> None: + try: + follow_up: Final = hook(self.future) + if follow_up is not None: + await follow_up + except Exception as e: # noqa: BLE001 # one owner's follow-up must not stop the others + verbose_logger.warning("redis batch settled hook failed: %s", e) + + def settle(self, replies: Sequence[object]) -> Awaitable[None] | None: + failure: Final = next((reply for reply in replies if isinstance(reply, Exception)), None) + if failure is None: + try: + self.future.set_result(self.spec.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 + if _is_missing_script(failure): + return self.settle_alone() + self.recover(failure) + return None + + def recover(self, error: Exception) -> None: + try: + self.future.set_result(self.spec.recover(error)) + except Exception as e: # noqa: BLE001 # the declaring caller owns the failure of its own operation + self.future.set_exception(e) + + async def settle_alone(self) -> None: + try: + with service_target(self.target), service_caller(self.caller): + self.future.set_result(await self.spec.run_alone()) + except Exception as e: # noqa: BLE001 # the declaring caller owns the failure of its own operation + self.future.set_exception(e) + + +class BatchResult(Generic[_T]): + """Awaitable handle for one declared operation; awaiting it flushes the batch it belongs to.""" + + __slots__ = ("_batch", "_op") + + def __init__(self, batch: "PipelineBatch", op: _Op[_T]) -> None: + self._batch: Final = batch + self._op: Final = op + + def __await__(self) -> Generator[object, None, _T]: + return self._wait().__await__() + + async def _wait(self) -> _T: + if not self._op.future.done(): + await self._batch.flush() + return self._op.future.result() + + @property + def done(self) -> bool: + return self._op.future.done() + + def on_settled(self, hook: SettledHook[_T]) -> None: + """For an owner that does not await: runs inside the flush once this operation has its result or + failure (or was cancelled with the pipeline), so the flush completes with the follow-up done.""" + self._op.settled_hooks.append(hook) + + +@dataclass(slots=True) +class PipelineBatch: + """Operations declared here go out in one pipeline the next time any of them is awaited or ``flush`` runs. + A backend that cannot pipeline them (a cluster client) runs each on its own, unless the batch is a + transaction, which always goes out as one MULTI/EXEC.""" + + backend: PipelineBackend + name: str = "redis_batch" + transaction: bool = False + _pending: list[_Op[object]] = field(default_factory=list) # mutable-ok: drained by flush + _flush_hooks: list[Callable[[], None]] = field(default_factory=list) # mutable-ok: append-only registry + _lock: asyncio.Lock = field(default_factory=asyncio.Lock) + _misses: set[str] = field(default_factory=set) # mutable-ok: keys a read of this batch found absent + flushes: int = 0 + + def declare(self, spec: OpSpec[_T]) -> BatchResult[_T]: + op: Final = _Op(spec) + self._pending.append(op) # pyright: ignore[reportArgumentType] # heterogeneous ops share the flush loop + return BatchResult(self, op) + + def note_missing(self, keys: Iterable[str]) -> None: + self._misses.update(keys) + + def note_present(self, key: str) -> None: + self._misses.discard(key) + + def read_as_missing(self, key: str) -> bool: + """True when a read on this batch already found no value under ``key`` and nothing has set it since, + so a per-key GET later in the same request can be answered without another round trip.""" + return key in self._misses + + def add_flush_hook(self, hook: Callable[[], None]) -> None: + """Called at the start of every flush so lazily bound readers can declare their keys into the same trip.""" + self._flush_hooks.append(hook) + + @property + def pending(self) -> int: + return len(self._pending) + + async def flush(self) -> None: + async with self._lock: + for hook in self._flush_hooks: + hook() + ops: Final = tuple(self._pending) + self._pending.clear() + if not ops: + return + self.flushes += 1 + try: + if self.backend.runs_alone and not self.transaction: + await asyncio.gather(*(op.settle_alone() for op in ops)) + else: + await self._flush_pipeline(ops) + finally: + for op in ops: + if not op.future.done(): + op.future.cancel() + await asyncio.gather(*(op.run_settled_hooks() for op in ops)) + + async def _flush_pipeline(self, ops: Sequence[_Op[object]]) -> None: + from litellm.redis.transport.events import report_pipeline + + 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 with self.backend.client().pipeline(transaction=self.transaction) as pipe: + widths.extend(op.spec.enqueue(pipe) for op in ops) + return await pipe.execute(raise_on_error=False) + + try: + replies: Final = await _run_under_circuit_breaker( + self.backend.breaker, self.name, lambda: self.backend.admit(run) + ) + except Exception as e: # noqa: BLE001 # each declaring caller applies its own Redis fallback + report_pipeline( + self.backend.service_logging, + self.name, + op_count=len(ops), + target=target, + metadata=metadata, + start_time=start_time, + error=e, + ) + for op in ops: + op.recover(e) + return + report_pipeline( + self.backend.service_logging, + self.name, + op_count=len(ops), + target=target, + metadata=metadata, + start_time=start_time, + error=None, + ) + 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) + + +def _pipeline_service_event(ops: Sequence[_Op[object]]) -> tuple[str | None, dict[str, int | str]]: + """The target and metadata of one pipeline flush: the one key family every op was declared under, or + ``"mixed"`` plus the sorted families when owners of several families share the trip.""" + families: Final = sorted({op.target for op in ops if op.target is not None}) + if len(families) > 1: + return MIXED_PIPELINE_TARGET, {"op_count": len(ops), REDIS_FAMILIES_METADATA_KEY: ",".join(families)} + return next(iter(families), None), {"op_count": len(ops)} diff --git a/litellm/redis/transport/breaker.py b/litellm/redis/transport/breaker.py index e9ad1819676..c10ef1a732b 100644 --- a/litellm/redis/transport/breaker.py +++ b/litellm/redis/transport/breaker.py @@ -19,16 +19,16 @@ async def run_under_breaker( ) -> _T | RedisFailure: try: return await _run_under_circuit_breaker(breaker, name, lambda: admit(call)) - except RedisCircuitBreakerOpenError as e: - return RedisFailure(kind="breaker_open", message=str(e)) except Exception as e: - failure: Final = _redis_failure(e) + failure: Final = failure_from(e) if failure is None: raise return failure -def _redis_failure(exc: BaseException) -> RedisFailure | None: +def failure_from(exc: BaseException) -> RedisFailure | None: + if isinstance(exc, RedisCircuitBreakerOpenError): + return RedisFailure(kind="breaker_open", message=str(exc)) if _is_redis_health_failure(exc): return RedisFailure(kind="transport", message=str(exc), timeout=is_redis_timeout_failure(exc)) if is_redis_error(exc): diff --git a/litellm/redis/transport/client.py b/litellm/redis/transport/client.py new file mode 100644 index 00000000000..23f455c7ed9 --- /dev/null +++ b/litellm/redis/transport/client.py @@ -0,0 +1,90 @@ +from collections.abc import AsyncIterator, Awaitable, Sequence +from contextlib import AbstractAsyncContextManager +from datetime import timedelta +from typing import Protocol + +from litellm.redis.transport.contract import Arg + + +class RedisPipeline(Protocol): + def get(self, name: str) -> object: ... + def mget(self, keys: Sequence[str]) -> object: ... + def set( + self, name: str, value: Arg, ex: timedelta | None = None, px: int | None = None, nx: bool = False + ) -> object: ... + def delete(self, *names: str) -> object: ... + def incrby(self, name: str, amount: int) -> object: ... + def incrbyfloat(self, name: str, amount: float) -> object: ... + def expire(self, name: str, time: timedelta) -> object: ... + def pexpire(self, name: str, time: int) -> object: ... + def ttl(self, name: str) -> object: ... + def rpush(self, name: str, *values: Arg) -> object: ... + def lpop(self, name: str, count: int | None = None) -> object: ... + def ltrim(self, name: str, start: int, end: int) -> object: ... + def sadd(self, name: str, *values: Arg) -> object: ... + def eval(self, script: str, numkeys: int, *keys_and_args: Arg) -> object: ... + def evalsha(self, sha: str, numkeys: int, *keys_and_args: Arg) -> object: ... + def publish(self, channel: str, message: str | bytes) -> object: ... + async def execute(self, raise_on_error: bool = True) -> list[object]: ... + + +class PipelineClient(Protocol): + def pipeline(self, transaction: bool = True) -> AbstractAsyncContextManager[RedisPipeline]: ... + + +class AsyncRedisCommands(Protocol): + def get(self, name: str) -> Awaitable[object]: ... + + def mget(self, keys: Sequence[str]) -> Awaitable[object]: ... + + def set(self, name: str, value: Arg, *, nx: bool, px: int | None) -> Awaitable[object]: ... + + def delete(self, *names: str) -> Awaitable[object]: ... + + def incrby(self, name: str, amount: int) -> Awaitable[object]: ... + + def incrbyfloat(self, name: str, amount: float) -> Awaitable[object]: ... + + def pexpire(self, name: str, time: int) -> Awaitable[object]: ... + + def ttl(self, name: str) -> Awaitable[object]: ... + + def rpush(self, name: str, *values: Arg) -> Awaitable[object]: ... + + def lpop(self, name: str, count: int | None = None) -> Awaitable[object]: ... + + def ltrim(self, name: str, start: int, end: int) -> Awaitable[object]: ... + + def sadd(self, name: str, *values: Arg) -> Awaitable[object]: ... + + def scan_iter(self, *, match: str, count: int) -> AsyncIterator[object]: ... + + def eval(self, script: str, numkeys: int, *keys_and_args: Arg) -> Awaitable[object]: ... + + def evalsha(self, sha: str, numkeys: int, *keys_and_args: Arg) -> Awaitable[object]: ... + + def pipeline(self, transaction: bool = True) -> AbstractAsyncContextManager[RedisPipeline]: ... + + def ping(self) -> Awaitable[object]: ... + + def info(self) -> Awaitable[object]: ... + + def client_list(self) -> Awaitable[object]: ... + + def flushall(self) -> Awaitable[object]: ... + + +class AsyncPubSub(Protocol): + def subscribe(self, *channels: str) -> Awaitable[object]: ... + + def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Awaitable[object]: ... + + def reset(self) -> Awaitable[object]: ... + + +class AsyncPubSubClient(Protocol): + def publish(self, channel: str, message: str | bytes) -> Awaitable[object]: ... + + def pubsub(self) -> AsyncPubSub: ... + + def pipeline(self, transaction: bool = True) -> AbstractAsyncContextManager[RedisPipeline]: ... diff --git a/litellm/redis/transport/contract.py b/litellm/redis/transport/contract.py index 1a39e50c502..83c4e2b0664 100644 --- a/litellm/redis/transport/contract.py +++ b/litellm/redis/transport/contract.py @@ -1,13 +1,17 @@ import math -from collections.abc import Awaitable, Callable, Sequence +from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence from dataclasses import dataclass from datetime import timedelta from typing import Final, Literal, Protocol, TypeAlias, TypeVar Arg: TypeAlias = str | bytes | int | float EvalReply: TypeAlias = int | bytes | None | tuple["EvalReply", ...] +InfoScalar: TypeAlias = str | int | float +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): @@ -21,10 +25,61 @@ class RedisFailure: timeout: bool = False +@dataclass(frozen=True, slots=True) +class PubSubMessage: + channel: bytes + data: bytes + + 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]: ... + + @property + def done(self) -> bool: ... + + +class RedisCommandBatch(Protocol): + def get(self, key: str) -> BatchReply[bytes | None | RedisFailure]: ... + + def mget(self, keys: Sequence[str]) -> BatchReply[tuple[bytes | None, ...] | RedisFailure]: ... + + def set( + self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False + ) -> BatchReply[bool | RedisFailure]: ... + + def delete(self, *keys: str) -> BatchReply[int | RedisFailure]: ... + + def incrby(self, key: str, amount: int) -> BatchReply[int | RedisFailure]: ... + + def incrbyfloat(self, key: str, amount: float) -> BatchReply[float | RedisFailure]: ... + + def expire(self, key: str, ttl: timedelta) -> BatchReply[bool | RedisFailure]: ... + + def ttl(self, key: str) -> BatchReply[int | RedisFailure]: ... + + def rpush(self, key: str, *values: Arg) -> BatchReply[int | RedisFailure]: ... + + def lpop(self, key: str, count: int | None = None) -> BatchReply[Popped | RedisFailure]: ... + + def ltrim(self, key: str, start: int, stop: int) -> BatchReply[bool | RedisFailure]: ... + + def sadd(self, key: str, *members: Arg) -> BatchReply[int | RedisFailure]: ... + + def eval(self, script: str, keys: Sequence[str], args: Sequence[Arg]) -> BatchReply[EvalReply | RedisFailure]: ... + + async def flush(self) -> None: ... + + +class RedisSubscription(Protocol): + async def get_message(self, *, timeout: float) -> PubSubMessage | None | RedisFailure: ... + + async def close(self) -> None | RedisFailure: ... + + class RedisTransport(Protocol): async def get(self, key: str) -> bytes | None | RedisFailure: ... @@ -34,4 +89,42 @@ class RedisTransport(Protocol): async def delete(self, *keys: str) -> int | RedisFailure: ... + async def incrby(self, key: str, amount: int) -> int | RedisFailure: ... + + async def incrbyfloat(self, key: str, amount: float) -> float | RedisFailure: ... + + async def expire(self, key: str, ttl: timedelta) -> bool | RedisFailure: ... + + async def ttl(self, key: str) -> int | RedisFailure: ... + + async def rpush(self, key: str, *values: Arg) -> int | RedisFailure: ... + + async def lpop(self, key: str, count: int | None = None) -> Popped | RedisFailure: ... + + async def ltrim(self, key: str, start: int, stop: int) -> bool | RedisFailure: ... + + async def sadd(self, key: str, *members: Arg) -> int | RedisFailure: ... + + async def scan(self, prefix: str, *, limit: int) -> tuple[bytes, ...] | RedisFailure: ... + async def eval(self, script: str, keys: Sequence[str], args: Sequence[Arg]) -> EvalReply | RedisFailure: ... + + def batch(self, *, transaction: bool = False) -> RedisCommandBatch: ... + + def request_batch(self) -> RedisCommandBatch | None: ... + + def post_call_batch(self) -> RedisCommandBatch | None: ... + + async def publish(self, channel: str, message: str | bytes) -> int | RedisFailure: ... + + async def publish_all(self, channel: str, messages: Sequence[str | bytes]) -> tuple[int, ...] | RedisFailure: ... + + async def subscribe(self, *channels: str) -> RedisSubscription | RedisFailure: ... + + async def ping(self) -> bool | RedisFailure: ... + + async def info(self) -> Mapping[str, InfoValue] | RedisFailure: ... + + async def client_list(self) -> tuple[Mapping[str, str], ...] | RedisFailure: ... + + async def flushall(self) -> bool | RedisFailure: ... diff --git a/litellm/redis/transport/events.py b/litellm/redis/transport/events.py index d1ae7e81eaa..e6e3fc35b41 100644 --- a/litellm/redis/transport/events.py +++ b/litellm/redis/transport/events.py @@ -1,14 +1,16 @@ +import asyncio import logging import time -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from typing import Final, TypeVar -from litellm._internal_context import current_service_target +from litellm._internal_context import current_service_target, service_target from litellm._logging import verbose_logger from litellm._service_logger import ServiceLogging from litellm.caching.redis_cache import ( _get_call_stack_info, # pyright: ignore[reportPrivateUsage] # same caller attribution as RedisCache events + log_redis_failure, log_redis_failure_value, ) from litellm.redis.transport.contract import RedisFailure @@ -69,3 +71,39 @@ def _log_failure(call_type: str, result: object) -> None: log_redis_failure_value( verbose_logger, logging.WARNING, f"Redis {call_type} failed" + (f" for {target}" if target else ""), result ) + + +def report_pipeline( + service_logging: ServiceLogging, + call_type: str, + *, + op_count: int, + target: str | None, + metadata: Mapping[str, int | str], + start_time: float, + error: Exception | None, +) -> None: + end_time: Final = time.time() + if error is not None: + log_redis_failure(verbose_logger, logging.WARNING, f"{call_type}: pipeline of {op_count} ops failed", error) + with service_target(target): + asyncio.create_task( + service_logging.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=end_time - start_time, + call_type=call_type, + start_time=start_time, + end_time=end_time, + event_metadata=dict(metadata), + ) + if error is None + else service_logging.async_service_failure_hook( + service=ServiceTypes.REDIS, + duration=end_time - start_time, + error=error, + call_type=call_type, + start_time=start_time, + end_time=end_time, + event_metadata=dict(metadata), + ) + ) diff --git a/litellm/redis/transport/redis_py.py b/litellm/redis/transport/redis_py.py index 8a65413be6e..227996ed6ff 100644 --- a/litellm/redis/transport/redis_py.py +++ b/litellm/redis/transport/redis_py.py @@ -1,34 +1,40 @@ +import asyncio import hashlib -from collections.abc import Awaitable, Callable, Iterator, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from dataclasses import dataclass from datetime import timedelta -from typing import Final, Protocol, TypeVar +from types import MappingProxyType +from typing import Final, Generic, TypeVar from pydantic import TypeAdapter, ValidationError from redis.exceptions import NoScriptError from litellm.caching.redis_cache import RedisCircuitBreaker -from litellm.redis.transport.breaker import run_under_breaker -from litellm.redis.transport.contract import Admit, Arg, EvalReply, RedisFailure, ttl_ms +from litellm.redis.transport.batch import BatchResult, PipelineBackend, PipelineBatch +from litellm.redis.transport.breaker import failure_from, run_under_breaker +from litellm.redis.transport.client import AsyncPubSub, AsyncPubSubClient, AsyncRedisCommands, RedisPipeline +from litellm.redis.transport.contract import ( + Admit, + Arg, + EvalReply, + InfoValue, + Popped, + PubSubMessage, + RedisFailure, + ttl_ms, +) from litellm.redis.transport.events import RedisServiceEvents _T: Final = TypeVar("_T") _ARRAY: Final = TypeAdapter(tuple[object, ...]) +_INFO: Final[TypeAdapter[Mapping[str, InfoValue]]] = TypeAdapter(Mapping[str, InfoValue]) +_CLIENTS: Final = TypeAdapter(tuple[Mapping[str, str], ...]) +_MESSAGE: Final = TypeAdapter(Mapping[str, object]) -class AsyncRedisCommands(Protocol): - def get(self, name: str) -> Awaitable[object]: ... - - def mget(self, keys: Sequence[str]) -> Awaitable[object]: ... - - def set(self, name: str, value: Arg, *, nx: bool, px: int | None) -> Awaitable[object]: ... - - def delete(self, *names: str) -> Awaitable[int]: ... - - def eval(self, script: str, numkeys: int, *keys_and_args: Arg) -> Awaitable[object]: ... - - def evalsha(self, sha: str, numkeys: int, *keys_and_args: Arg) -> Awaitable[object]: ... +def is_missing_script(failure: Exception) -> bool: + return isinstance(failure, NoScriptError) @dataclass(frozen=True, slots=True) @@ -38,13 +44,18 @@ class RedisPyTransport: breaker: RedisCircuitBreaker admit: Admit events: RedisServiceEvents + multi_get: Callable[[Sequence[str]], Awaitable[object]] + server_version: Callable[[], int] + pubsub_client: Callable[[], AsyncPubSubClient] + request_engine: Callable[[], PipelineBatch | None] + post_call_engine: Callable[[], PipelineBatch | None] async def get(self, key: str) -> bytes | None | RedisFailure: return _bytes_reply(await self._run("get", lambda: self.client().get(self._key(key)))) async def mget(self, keys: Sequence[str]) -> tuple[bytes | None, ...] | RedisFailure: namespaced: Final = tuple(self._key(k) for k in keys) - return _mget_reply(await self._run("mget", lambda: self.client().mget(namespaced)), len(namespaced)) + return _mget_reply(await self._run("mget", lambda: self.multi_get(namespaced)), len(namespaced)) async def set(self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False) -> bool | RedisFailure: px: Final = None if ttl is None else ttl_ms(ttl) @@ -54,21 +65,345 @@ class RedisPyTransport: namespaced: Final = tuple(self._key(k) for k in keys) return _int_reply(await self._run("delete", lambda: self.client().delete(*namespaced))) + async def incrby(self, key: str, amount: int) -> int | RedisFailure: + return _int_reply(await self._run("incrby", lambda: self.client().incrby(self._key(key), amount))) + + async def incrbyfloat(self, key: str, amount: float) -> float | RedisFailure: + return _float_reply(await self._run("incrbyfloat", lambda: self.client().incrbyfloat(self._key(key), amount))) + + async def expire(self, key: str, ttl: timedelta) -> bool | RedisFailure: + return _bool_reply(await self._run("expire", lambda: self.client().pexpire(self._key(key), ttl_ms(ttl)))) + + async def ttl(self, key: str) -> int | RedisFailure: + return _int_reply(await self._run("ttl", lambda: self.client().ttl(self._key(key)))) + + async def rpush(self, key: str, *values: Arg) -> int | RedisFailure: + return _int_reply(await self._run("rpush", lambda: self.client().rpush(self._key(key), *values))) + + async def lpop(self, key: str, count: int | None = None) -> Popped | RedisFailure: + name: Final = self._key(key) + if count is None or self.server_version() >= 7: + return _popped_reply(await self._run("lpop", lambda: self.client().lpop(name, count))) + return _popped_each(await self._run("lpop", lambda: _lpop_each_alone(self.client(), name, count))) + + async def ltrim(self, key: str, start: int, stop: int) -> bool | RedisFailure: + return _bool_reply(await self._run("ltrim", lambda: self.client().ltrim(self._key(key), start, stop))) + + async def sadd(self, key: str, *members: Arg) -> int | RedisFailure: + return _int_reply(await self._run("sadd", lambda: self.client().sadd(self._key(key), *members))) + + async def scan(self, prefix: str, *, limit: int) -> tuple[bytes, ...] | RedisFailure: + match: Final = f"{self._key(prefix)}*" + return _keys_reply(await self._run("scan", lambda: _scan(self.client(), match, limit))) + async def eval(self, script: str, keys: Sequence[str], args: Sequence[Arg]) -> EvalReply | RedisFailure: - sha: Final = hashlib.sha1(script.encode(), usedforsecurity=False).hexdigest() + sha: Final = _sha(script) namespaced: Final = tuple(self._key(k) for k in keys) return _eval_reply( await self._run("eval", lambda: _evalsha_or_eval(self.client(), script, sha, namespaced, args)) ) + def batch(self, *, transaction: bool = False) -> "RedisPyCommandBatch": + return RedisPyCommandBatch(self, PipelineBatch(self._backend(runs_alone=False), transaction=transaction)) + + def request_batch(self) -> "RedisPyCommandBatch | None": + engine: Final = self.request_engine() + return None if engine is None else RedisPyCommandBatch(self, engine) + + def post_call_batch(self) -> "RedisPyCommandBatch | None": + engine: Final = self.post_call_engine() + return None if engine is None else RedisPyCommandBatch(self, engine) + + async def publish(self, channel: str, message: str | bytes) -> int | RedisFailure: + return _int_reply( + await self._run_unguarded("publish", lambda: self.pubsub_client().publish(self._key(channel), message)) + ) + + async def publish_all(self, channel: str, messages: Sequence[str | bytes]) -> tuple[int, ...] | RedisFailure: + name: Final = self._key(channel) + return _ints_reply( + await self._run_unguarded("publish", lambda: _publish_each(self.pubsub_client(), name, messages)), + len(messages), + ) + + async def subscribe(self, *channels: str) -> "RedisPySubscription | RedisFailure": + names: Final = tuple(self._key(channel) for channel in channels) + return await self._run_unguarded("subscribe", lambda: _subscribe(self.pubsub_client(), names)) + + async def ping(self) -> bool | RedisFailure: + return _bool_reply(await self._run_unguarded("ping", lambda: self.client().ping())) + + async def info(self) -> Mapping[str, InfoValue] | RedisFailure: + return _info_reply(await self._run_unguarded("info", lambda: self.client().info())) + + async def client_list(self) -> tuple[Mapping[str, str], ...] | RedisFailure: + return _clients_reply(await self._run_unguarded("client_list", lambda: self.client().client_list())) + + async def flushall(self) -> bool | RedisFailure: + return _bool_reply(await self._run_unguarded("flushall", lambda: self.client().flushall())) + def _key(self, key: str) -> str: if not self.namespace or key.startswith(f"{self.namespace}:"): return key return f"{self.namespace}:{key}" + def _backend(self, *, runs_alone: bool) -> PipelineBackend: + return PipelineBackend( + client=self.client, + breaker=self.breaker, + admit=self.admit, + service_logging=self.events.service_logging, + runs_alone=runs_alone, + ) + async def _run(self, name: str, call: Callable[[], Awaitable[_T]]) -> _T | RedisFailure: return await self.events.observe(name, lambda: run_under_breaker(self.breaker, name, self.admit, call)) + async def _run_unguarded(self, name: str, call: Callable[[], Awaitable[_T]]) -> _T | RedisFailure: + return await self.events.observe(name, lambda: _without_breaker(call)) + + +@dataclass(frozen=True, slots=True) +class _Command(Generic[_T]): + queue: Callable[[RedisPipeline], int] + decode: Callable[[Sequence[object]], _T | RedisFailure] + alone: Callable[[], Awaitable[_T | RedisFailure]] + + def enqueue(self, pipe: RedisPipeline) -> int: + return self.queue(pipe) + + def resolve(self, replies: Sequence[object]) -> _T | RedisFailure: + return self.decode(replies) + + def recover(self, error: Exception) -> _T | RedisFailure: + failure: Final = failure_from(error) + if failure is None: + raise error + return failure + + def run_alone(self) -> Awaitable[_T | RedisFailure]: + return self.alone() + + +@dataclass(frozen=True, slots=True) +class RedisPyCommandBatch: + transport: RedisPyTransport + engine: PipelineBatch + + 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( + lambda pipe: _one(pipe.get(name)), lambda replies: _bytes_reply(replies[0]), lambda: self.transport.get(key) + ) + reply.on_settled(lambda future: self._note_read((key,), future, lambda value: (value,))) + return reply + + def mget(self, keys: Sequence[str]) -> BatchResult[tuple[bytes | None, ...] | RedisFailure]: + logical: Final = tuple(keys) + names: Final = tuple(self.transport._key(k) for k in logical) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + reply: Final = self._declare( + lambda pipe: _one(pipe.mget(names)), + lambda replies: _mget_reply(replies[0], len(names)), + lambda: self.transport.mget(logical), + ) + reply.on_settled(lambda future: self._note_read(logical, future, lambda values: values)) + return reply + + def set(self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False) -> BatchResult[bool | RedisFailure]: + name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + px: Final = None if ttl is None else ttl_ms(ttl) + self.engine.note_present(key) + return self._declare( + lambda pipe: _one(pipe.set(name, value, px=px, nx=nx)), + lambda replies: _bool_reply(replies[0]), + lambda: self.transport.set(key, value, ttl=ttl, nx=nx), + ) + + def delete(self, *keys: str) -> BatchResult[int | RedisFailure]: + names: Final = tuple(self.transport._key(k) for k in keys) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + self.engine.note_missing(keys) + return self._declare( + lambda pipe: _one(pipe.delete(*names)), + lambda replies: _int_reply(replies[0]), + lambda: self.transport.delete(*keys), + ) + + def incrby(self, key: str, amount: int) -> BatchResult[int | RedisFailure]: + name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + self.engine.note_present(key) + return self._declare( + lambda pipe: _one(pipe.incrby(name, amount)), + lambda replies: _int_reply(replies[0]), + lambda: self.transport.incrby(key, amount), + ) + + def incrbyfloat(self, key: str, amount: float) -> BatchResult[float | RedisFailure]: + name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + self.engine.note_present(key) + return self._declare( + lambda pipe: _one(pipe.incrbyfloat(name, amount)), + lambda replies: _float_reply(replies[0]), + lambda: self.transport.incrbyfloat(key, amount), + ) + + def expire(self, key: str, ttl: timedelta) -> BatchResult[bool | RedisFailure]: + name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + return self._declare( + lambda pipe: _one(pipe.pexpire(name, ttl_ms(ttl))), + lambda replies: _bool_reply(replies[0]), + lambda: self.transport.expire(key, ttl), + ) + + def ttl(self, key: str) -> BatchResult[int | RedisFailure]: + name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + return self._declare( + lambda pipe: _one(pipe.ttl(name)), lambda replies: _int_reply(replies[0]), lambda: self.transport.ttl(key) + ) + + def rpush(self, key: str, *values: Arg) -> BatchResult[int | RedisFailure]: + name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + self.engine.note_present(key) + return self._declare( + lambda pipe: _one(pipe.rpush(name, *values)), + lambda replies: _int_reply(replies[0]), + lambda: self.transport.rpush(key, *values), + ) + + def lpop(self, key: str, count: int | None = None) -> BatchResult[Popped | RedisFailure]: + name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + if count is None or self.transport.server_version() >= 7: + return self._declare( + lambda pipe: _one(pipe.lpop(name, count)), + lambda replies: _popped_reply(replies[0]), + lambda: self.transport.lpop(key, count), + ) + return self._declare( + lambda pipe: _queue_lpop_each(pipe, name, count), + _popped_each, + lambda: self.transport.lpop(key, count), + ) + + def ltrim(self, key: str, start: int, stop: int) -> BatchResult[bool | RedisFailure]: + name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + return self._declare( + lambda pipe: _one(pipe.ltrim(name, start, stop)), + lambda replies: _bool_reply(replies[0]), + lambda: self.transport.ltrim(key, start, stop), + ) + + def sadd(self, key: str, *members: Arg) -> BatchResult[int | RedisFailure]: + name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + self.engine.note_present(key) + return self._declare( + lambda pipe: _one(pipe.sadd(name, *members)), + lambda replies: _int_reply(replies[0]), + lambda: self.transport.sadd(key, *members), + ) + + def eval(self, script: str, keys: Sequence[str], args: Sequence[Arg]) -> BatchResult[EvalReply | RedisFailure]: + names: Final = tuple(self.transport._key(k) for k in keys) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + sha: Final = _sha(script) + transaction: Final = self.engine.transaction + return self._declare( + lambda pipe: _one( + pipe.eval(script, len(names), *names, *args) + if transaction + else pipe.evalsha(sha, len(names), *names, *args) + ), + lambda replies: _eval_reply(replies[0]), + lambda: self.transport.eval(script, keys, args), + ) + + async def flush(self) -> None: + await self.engine.flush() + + def _declare( + self, + queue: Callable[[RedisPipeline], int], + decode: Callable[[Sequence[object]], _T | RedisFailure], + alone: Callable[[], Awaitable[_T | RedisFailure]], + ) -> BatchResult[_T | RedisFailure]: + return self.engine.declare(_Command(queue, decode, alone)) + + def _note_read( + self, + keys: Sequence[str], + future: "asyncio.Future[_T | RedisFailure]", + values: Callable[[_T], Sequence[bytes | None]], + ) -> None: + if future.cancelled() or future.exception() is not None: + return + result: Final = future.result() + if isinstance(result, RedisFailure): + return + self.engine.note_missing(key for key, value in zip(keys, values(result)) if value is None) + + +@dataclass(frozen=True, slots=True) +class RedisPySubscription: + pubsub: AsyncPubSub + + async def get_message(self, *, timeout: float) -> PubSubMessage | None | RedisFailure: + return _message_reply( + await _without_breaker(lambda: self.pubsub.get_message(ignore_subscribe_messages=True, timeout=timeout)) + ) + + async def close(self) -> None | RedisFailure: + closed: Final = await _without_breaker(self.pubsub.reset) + return closed if isinstance(closed, RedisFailure) else None + + +def _one(_queued: object) -> int: + return 1 + + +def _sha(script: str) -> str: + return hashlib.sha1(script.encode(), usedforsecurity=False).hexdigest() + + +def _queue_lpop_each(pipe: RedisPipeline, name: str, count: int) -> int: + for _ in range(count): + pipe.lpop(name) + return count + + +async def _lpop_each_alone(client: AsyncRedisCommands, name: str, count: int) -> Sequence[object]: + async with client.pipeline(transaction=False) as pipe: + _queue_lpop_each(pipe, name, count) + return await pipe.execute() + + +async def _publish_each(client: AsyncPubSubClient, channel: str, messages: Sequence[str | bytes]) -> Sequence[object]: + async with client.pipeline(transaction=False) as pipe: + for message in messages: + pipe.publish(channel, message) + return await pipe.execute() + + +async def _subscribe(client: AsyncPubSubClient, channels: Sequence[str]) -> RedisPySubscription: + pubsub: Final = client.pubsub() + await pubsub.subscribe(*channels) + return RedisPySubscription(pubsub) + + +async def _scan(client: AsyncRedisCommands, match: str, limit: int) -> tuple[object, ...]: + keys: Final[list[object]] = [] # mutable-ok: filled while the cursor walks + async for key in client.scan_iter(match=match, count=limit): + keys.append(key) + if len(keys) >= limit: + break + return tuple(keys) + + +async def _without_breaker(call: Callable[[], Awaitable[_T]]) -> _T | RedisFailure: + try: + return await call() + except Exception as e: + failure: Final = failure_from(e) + if failure is None: + raise + return failure + async def _evalsha_or_eval( client: AsyncRedisCommands, script: str, sha: str, keys: Sequence[str], args: Sequence[Arg] @@ -96,7 +431,7 @@ def _bytes_reply(reply: object) -> bytes | None | RedisFailure: def _array(reply: object) -> tuple[object, ...] | None: - if not isinstance(reply, list): + if not isinstance(reply, (list, tuple)): return None try: return _ARRAY.validate_python(reply) @@ -104,6 +439,15 @@ def _array(reply: object) -> tuple[object, ...] | None: return None +def _all_bytes(values: Sequence[object]) -> tuple[bytes, ...] | RedisFailure: + decoded: Final = tuple(_optional_bytes(value) for value in values) + failure: Final = next((value for value in decoded if isinstance(value, RedisFailure)), None) + if failure is not None: + return failure + present: Final = tuple(value for value in decoded if isinstance(value, bytes)) + return present if len(present) == len(decoded) else _decode_failure(values) + + def _mget_reply(reply: object, expected: int) -> tuple[bytes | None, ...] | RedisFailure: if isinstance(reply, RedisFailure): return reply @@ -115,10 +459,36 @@ def _mget_reply(reply: object, expected: int) -> tuple[bytes | None, ...] | Redi return failure if failure is not None else tuple(value for value in values if not isinstance(value, RedisFailure)) +def _keys_reply(reply: object) -> tuple[bytes, ...] | RedisFailure: + if isinstance(reply, RedisFailure): + return reply + array: Final = _array(reply) + return _decode_failure(reply) if array is None else _all_bytes(array) + + +def _popped_reply(reply: object) -> Popped | RedisFailure: + if isinstance(reply, RedisFailure) or reply is None: + return reply + array: Final = _array(reply) + if array is None: + return _bytes_reply(reply) + return _all_bytes(array) + + +def _popped_each(replies: object) -> Popped | RedisFailure: + if isinstance(replies, RedisFailure): + return replies + array: Final = _array(replies) + if array is None: + return _decode_failure(replies) + popped: Final = _all_bytes(tuple(value for value in array if value is not None)) + return None if popped == () else popped + + def _bool_reply(reply: object) -> bool | RedisFailure: if isinstance(reply, RedisFailure): return reply - return reply is not None and reply is not False + return reply is not None and reply is not False and reply != 0 def _int_reply(reply: object) -> int | RedisFailure: @@ -127,6 +497,60 @@ def _int_reply(reply: object) -> int | RedisFailure: return reply if isinstance(reply, int) and not isinstance(reply, bool) else _decode_failure(reply) +def _ints_reply(reply: object, expected: int) -> tuple[int, ...] | RedisFailure: + if isinstance(reply, RedisFailure): + return reply + array: Final = _array(reply) + if array is None or len(array) != expected: + return _decode_failure(reply) + counts: Final = tuple(_int_reply(value) for value in array) + failure: Final = next((value for value in counts if isinstance(value, RedisFailure)), None) + return failure if failure is not None else tuple(value for value in counts if not isinstance(value, RedisFailure)) + + +def _float_reply(reply: object) -> float | RedisFailure: + if isinstance(reply, RedisFailure): + return reply + if isinstance(reply, bool) or not isinstance(reply, (int, float, str, bytes)): + return _decode_failure(reply) + try: + return float(reply) + except ValueError: + return _decode_failure(reply) + + +def _info_reply(reply: object) -> Mapping[str, InfoValue] | RedisFailure: + if isinstance(reply, RedisFailure): + return reply + try: + return MappingProxyType(_INFO.validate_python(reply)) + except ValidationError: + return _decode_failure(reply) + + +def _clients_reply(reply: object) -> tuple[Mapping[str, str], ...] | RedisFailure: + if isinstance(reply, RedisFailure): + return reply + try: + return tuple(MappingProxyType(client) for client in _CLIENTS.validate_python(reply)) + except ValidationError: + return _decode_failure(reply) + + +def _message_reply(reply: object) -> PubSubMessage | None | RedisFailure: + if isinstance(reply, RedisFailure) or reply is None: + return reply + try: + message: Final = _MESSAGE.validate_python(reply) + except ValidationError: + return _decode_failure(reply) + channel: Final = _optional_bytes(message.get("channel")) + data: Final = _optional_bytes(message.get("data")) + if not isinstance(channel, bytes) or not isinstance(data, bytes): + return _decode_failure(reply) + return PubSubMessage(channel=channel, data=data) + + def _eval_reply(reply: object) -> EvalReply | RedisFailure: frames: Final[list[tuple[Iterator[object], list[EvalReply]]]] = [ # mutable-ok: linear-time traversal stack (iter((reply,)), []) diff --git a/tests/_support/fake_redis.py b/tests/_support/fake_redis.py index a12b2ed4212..8af6cf88a14 100644 --- a/tests/_support/fake_redis.py +++ b/tests/_support/fake_redis.py @@ -46,4 +46,9 @@ def fake_redis_transport( breaker=RedisCircuitBreaker(failure_threshold=failure_threshold, recovery_timeout=3600), admit=_admit_directly, events=RedisServiceEvents(events or ServiceLogging()), + multi_get=lambda keys: client.mget(keys), + server_version=lambda: 7, + pubsub_client=lambda: client, + request_engine=lambda: None, + post_call_engine=lambda: None, ) diff --git a/tests/unit/redis/transport/test_redis_py.py b/tests/unit/redis/transport/test_redis_py.py index a3a0b687f8a..7cd0de37958 100644 --- a/tests/unit/redis/transport/test_redis_py.py +++ b/tests/unit/redis/transport/test_redis_py.py @@ -1,20 +1,23 @@ import asyncio import logging -from dataclasses import dataclass +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass, replace from datetime import datetime, timedelta from typing import Final, Literal - from unittest.mock import patch import fakeredis import pytest +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_cluster_cache import ClusterConnectionGate, RedisClusterCache -from litellm.redis.transport.contract import RedisFailure +from litellm.redis.transport.batch import PipelineBatch +from litellm.redis.transport.contract import PubSubMessage, RedisFailure +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 @@ -372,3 +375,344 @@ async def test_a_failure_is_logged_once_with_the_command_and_target( assert [record.getMessage() for record in caplog.records] == [ f"Redis get failed for response_polling: {failure.message}" ] + + +_Command = Callable[[RedisPyTransport], Awaitable[object]] + + +_WRONG_TYPE_COMMANDS: Final[tuple[_Command, ...]] = ( + lambda t: t.incrby("k", 1), + lambda t: t.incrbyfloat("k", 1.0), + lambda t: t.rpush("k", "x"), + lambda t: t.sadd("k", "x"), +) +_OUTAGE_COMMANDS: Final[tuple[_Command, ...]] = ( + lambda t: t.incrby("k", 1), + lambda t: t.expire("k", timedelta(seconds=1)), + lambda t: t.ttl("k"), + lambda t: t.lpop("k"), + lambda t: t.scan("k", limit=1), + lambda t: t.publish("c", "m"), + lambda t: t.ping(), +) + + +@pytest.mark.asyncio +async def test_counters_add_to_the_stored_value_under_the_namespace( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + transport = fake_redis_cache(server, namespace="ns").transport() + + assert await transport.incrby("n", 2) == 2 + assert await transport.incrby("ns:n", 3) == 5 + assert await transport.incrbyfloat("f", 1.5) == 1.5 + assert await transport.incrbyfloat("f", 0.25) == 1.75 + + assert await raw.get("ns:n") == b"5" + assert await raw.get("ns:f") == b"1.75" + + +@pytest.mark.asyncio +async def test_expire_sets_the_ttl_in_milliseconds_and_ttl_reads_it_back( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + transport = fake_redis_cache(server).transport() + await raw.set("k", "v") + + assert await transport.ttl("k") == -1 + assert await transport.expire("k", timedelta(milliseconds=1500)) is True + assert await transport.expire("absent", timedelta(seconds=5)) is False + + assert 1000 < await raw.pttl("k") <= 1500 + assert await transport.ttl("k") == 1 + assert await transport.ttl("absent") == -2 + + +@pytest.mark.asyncio +async def test_list_commands_push_trim_and_pop_from_the_head( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + transport = fake_redis_cache(server, namespace="ns").transport() + + assert await transport.rpush("l", "a", "b", "c", "d") == 4 + assert await transport.ltrim("l", -3, -1) is True + assert await transport.lpop("l") == b"b" + assert await transport.lpop("l", 5) == (b"c", b"d") + assert await transport.lpop("l") is None + assert await transport.lpop("l", 2) is None + + assert await raw.exists("ns:l") == 0 + + +@pytest.mark.asyncio +async def test_lpop_with_a_count_before_redis_7_pops_the_same_items( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + transport = replace(fake_redis_cache(server).transport(), server_version=lambda: 6) + await raw.rpush("l", "a", "b", "c") + + assert await transport.lpop("l", 2) == (b"a", b"b") + assert await transport.lpop("l", 5) == (b"c",) + assert await transport.lpop("l", 5) is None + + +@pytest.mark.asyncio +async def test_sadd_counts_only_new_members(server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis) -> None: + transport = fake_redis_cache(server, namespace="ns").transport() + + assert await transport.sadd("s", "a", "b") == 2 + assert await transport.sadd("s", "b", "c") == 1 + + assert await raw.smembers("ns:s") == {b"a", b"b", b"c"} + + +@pytest.mark.asyncio +async def test_scan_finds_keys_under_the_namespaced_prefix_up_to_the_limit( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + transport = fake_redis_cache(server, namespace="ns").transport() + for key in ("ns:team:1", "ns:team:2", "ns:team:3", "ns:user:1", "team:4"): + await raw.set(key, "v") + + assert sorted(await transport.scan("team:", limit=10)) == [b"ns:team:1", b"ns:team:2", b"ns:team:3"] + assert len(await transport.scan("team:", limit=2)) == 2 + + +@pytest.mark.parametrize("command", _WRONG_TYPE_COMMANDS) +@pytest.mark.asyncio +async def test_a_command_on_the_wrong_type_is_a_server_failure( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, command: _Command +) -> None: + transport = fake_redis_cache(server).transport() + await raw.hset("k", "f", "v") + + failure = await command(transport) + + assert isinstance(failure, RedisFailure) and failure.kind == "server" + + +@pytest.mark.parametrize("command", _OUTAGE_COMMANDS) +@pytest.mark.asyncio +async def test_an_outage_is_a_transport_failure_for_every_command( + server: fakeredis.FakeServer, command: _Command +) -> None: + transport = fake_redis_cache(server).transport() + server.connected = False + + failure = await command(transport) + + assert isinstance(failure, RedisFailure) and failure.kind == "transport" + + +class _ClusterClient: + def __init__(self, inner: fakeredis.FakeAsyncRedis) -> None: + self.inner = inner + + async def mget(self, keys: Sequence[str]) -> object: + raise ResponseError("CROSSSLOT Keys in request don't hash to the same slot") + + async def mget_nonatomic(self, keys: Sequence[str]) -> object: + return [await self.inner.get(key) for key in keys] + + +@pytest.mark.asyncio +async def test_mget_on_a_cluster_cache_reads_keys_from_different_slots(raw: fakeredis.FakeAsyncRedis) -> None: + with patch("litellm._redis.init_redis_cluster"): + cache = RedisClusterCache(startup_nodes=[{"host": "localhost", "port": 7000}]) + litellm.in_memory_llm_clients_cache.set_cache( + key=cache._get_async_client_cache_key(), # pyright: ignore[reportPrivateUsage] # seeds the per-loop client slot RedisCache reads + value=_ClusterClient(raw), + ) + await raw.set("a", "1") + await raw.set("b", "2") + + assert await cache.transport().mget(("a", "b", "c")) == (b"1", b"2", None) + + +@pytest.mark.asyncio +async def test_a_batch_sends_every_declared_command_in_one_round_trip( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + events = _RecordingEvents() + transport = fake_redis_cache(server, namespace="ns", service_logger=events).transport() + await _settle() + await raw.set("ns:a", "1") + before = len(events.seen) + batch = transport.batch() + + got = batch.get("a") + written = batch.set("b", "2", ttl=timedelta(seconds=30)) + counted = batch.incrbyfloat("n", 1.5) + pushed = batch.rpush("l", "x", "y") + read = batch.mget(("a", "b", "missing")) + assert not got.done + + assert await got == b"1" + assert (await written, await counted, await pushed) == (True, 1.5, 2) + assert await read == (b"1", b"2", None) + await _settle() + assert [event.call_type for event in events.seen[before:]] == ["redis_batch"] + assert 0 < await raw.pttl("ns:b") <= 30_000 + + +@pytest.mark.asyncio +async def test_a_failing_command_in_a_batch_fails_only_its_own_reply( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + transport = fake_redis_cache(server).transport() + await raw.hset("h", "f", "v") + batch = transport.batch() + + bad = batch.incrby("h", 1) + good = batch.incrby("n", 1) + + failure = await bad + assert isinstance(failure, RedisFailure) and failure.kind == "server" + assert await good == 1 + + +@pytest.mark.asyncio +async def test_an_outage_fails_every_reply_in_the_batch_and_counts_once_on_the_breaker( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + transport = fake_redis_transport(raw, failure_threshold=1) + batch = transport.batch() + server.connected = False + + declared = (batch.get("a"), batch.incrby("n", 1)) + replies = [await reply for reply in declared] + server.connected = True + refused = await transport.get("a") + + assert all(isinstance(reply, RedisFailure) and reply.kind == "transport" for reply in replies) + assert isinstance(refused, RedisFailure) and refused.kind == "breaker_open" + + +@pytest.mark.asyncio +async def test_a_script_redis_does_not_have_cached_still_runs_in_a_batch( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + transport = fake_redis_cache(server).transport() + await raw.script_flush() + batch = transport.batch() + + scripted = batch.eval("return redis.call('INCRBY', KEYS[1], ARGV[1])", ("n",), (4,)) + counted = batch.incrby("other", 1) + + assert await scripted == 4 + assert await counted == 1 + assert await raw.get("n") == b"4" + + +@pytest.mark.asyncio +async def test_a_transaction_batch_runs_its_commands_atomically_with_uncached_scripts( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + transport = fake_redis_cache(server).transport() + await raw.script_flush() + batch = transport.batch(transaction=True) + + pushed = batch.rpush("l", "a", "b", "c") + trimmed = batch.ltrim("l", -2, -1) + scripted = batch.eval("return redis.call('LLEN', KEYS[1])", ("l",), ()) + + assert (await pushed, await trimmed, await scripted) == (3, True, 2) + assert await raw.lrange("l", 0, -1) == [b"b", b"c"] + + +@pytest.mark.asyncio +async def test_a_transaction_stays_one_round_trip_on_a_backend_that_runs_batched_commands_alone( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + events = _RecordingEvents() + transport = fake_redis_transport(raw, failure_threshold=1, events=events) + alone_backend = replace(transport.batch().engine.backend, runs_alone=True) + batch = RedisPyCommandBatch(transport, PipelineBatch(alone_backend, name="tx", transaction=True)) + + pushed = batch.rpush("l", "a", "b") + trimmed = batch.ltrim("l", -1, -1) + + assert (await pushed, await trimmed) == (2, True) + await _settle() + assert [event.call_type for event in events.seen] == ["tx"] + + +@pytest.mark.asyncio +async def test_request_batch_commands_share_the_round_trip_of_redis_batch_callers( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache = fake_redis_cache(server, namespace="ns") + await raw.set("ns:legacy", '"old"') + await raw.set("ns:moved", "new") + transport = cache.transport() + assert transport.request_batch() is None + + with request_redis_batch_scope(): + legacy = active_request_redis_batch(cache) + commands = transport.request_batch() + assert legacy is not None and commands is not None + old_read = legacy.mget(["legacy"]) + new_read = commands.get("moved") + missing = commands.get("absent") + + assert await old_read == {"legacy": "old"} + assert await new_read == b"new" + assert await missing is None + assert legacy.flushes == 1 + assert legacy.read_as_missing("absent") is True + assert legacy.read_as_missing("moved") is False + + +@pytest.mark.asyncio +async def test_publish_reaches_a_subscriber_on_the_namespaced_channel( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + transport = fake_redis_cache(server, namespace="ns").transport() + subscription = await transport.subscribe("config") + assert not isinstance(subscription, RedisFailure) + + assert await transport.publish("config", "one") == 1 + assert await transport.publish_all("config", ("two", "three")) == (1, 1) + assert await raw.publish("config", "elsewhere") == 0 + + received = [await subscription.get_message(timeout=1) for _ in range(4)] + messages = [message for message in received if message is not None] + assert messages == [ + PubSubMessage(channel=b"ns:config", data=b"one"), + PubSubMessage(channel=b"ns:config", data=b"two"), + PubSubMessage(channel=b"ns:config", data=b"three"), + ] + assert await subscription.close() is None + + +@pytest.mark.asyncio +async def test_health_and_admin_commands_still_run_while_the_breaker_is_open( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + transport = fake_redis_transport(raw, failure_threshold=1) + await raw.set("k", "v") + server.connected = False + await transport.get("k") + server.connected = True + assert isinstance(refused := await transport.get("k"), RedisFailure) and refused.kind == "breaker_open" + + assert await transport.ping() is True + clients = await transport.client_list() + assert not isinstance(clients, RedisFailure) and len(clients) >= 1 + assert await transport.flushall() is True + assert await raw.exists("k") == 0 + + +class _InfoClient: + async def info(self) -> object: + return {"redis_version": "7.4.0", "connected_clients": 3, "db0": {"keys": 2, "expires": 0}} + + +@pytest.mark.asyncio +async def test_info_returns_the_parsed_sections(raw: fakeredis.FakeAsyncRedis) -> None: + transport = replace(fake_redis_transport(raw, failure_threshold=1), client=lambda: _InfoClient()) + + info = await transport.info() + + assert not isinstance(info, RedisFailure) + assert (info["redis_version"], info["connected_clients"], info["db0"]) == ("7.4.0", 3, {"keys": 2, "expires": 0})