mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
feat(redis): cover every Python Redis use on the transport
The transport now has every data command Python sends, request-scoped and explicit batches on the one pipeline engine RedisBatch already used, MULTI/EXEC batches, pub/sub, and the admin and health commands. RedisBatch declares onto the same engine, so moved and unmoved callers in one request still share one round trip. Cluster mget goes through the cache's mget_nonatomic path instead of a cross-slot MGET. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee
Devin AI
parent
77129e4f92
commit
6a208ed271
12 changed files with 1432 additions and 312 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
255
litellm/redis/transport/batch.py
Normal file
255
litellm/redis/transport/batch.py
Normal file
|
|
@ -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)}
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
90
litellm/redis/transport/client.py
Normal file
90
litellm/redis/transport/client.py
Normal file
|
|
@ -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]: ...
|
||||
|
|
@ -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: ...
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,)), [])
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue