scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

refactor(redis): delete the Python RedisBatch now that every request op is native (#51383)

Browse source 
The request scope keeps only the native batches. RedisBatch, its ops, the cluster fallback, the Python post-call
deadline and the DualCache post-call helpers had no production caller left, and the pre-call SET, DEL and miss
tracking that UserApiKeyCache never reached go with them. The mixed-families pipeline metadata had no producer left
This commit is contained in:
yujonglee 2026-11-13 12:22:28 +09:00 • committed by GitHub
parent 9a1acee4a7
commit 2be96aed13
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
31 changed files with 270 additions and 1605 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
litellm/_internal_context.py
View file

@ -36,11 +36,6 @@ _billing_time: Final[ContextVar[datetime | None]] = ContextVar("billing_time", d
_post_response: Final[ContextVar[bool]] = ContextVar("post_response", default=False)
_service_target: Final[ContextVar[str | None]] = ContextVar("service_target", default=None)
# Event-metadata key under which a Redis pipeline reports the sorted, comma-joined families its ops were
# declared under when they span more than one.
REDIS_FAMILIES_METADATA_KEY: Final = "families"
_service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", default=None)
_mirrored_call: Final[ContextVar[bool]] = ContextVar("mirrored_call", default=False)
@ -140,22 +135,6 @@ def with_service_target(
return decorate
@contextmanager
def service_caller(caller: str | None) -> Generator[None]:
"""Name the litellm code a datastore call was issued for when its own frames cannot: an operation
declared in one task and run in another (a batch op retried on the flush) carries the chain captured
where it was declared."""
token: Final = _service_caller.set(caller)
try:
yield
finally:
_restore(_service_caller, token)
def current_service_caller() -> str | None:
return _service_caller.get()
@contextmanager
def mirrored_call() -> Generator[None]:
"""A silent_model mirror of a request: every litellm call inside, its retries and fallbacks included."""

86
litellm/caching/dual_cache.py
View file

@ -8,11 +8,10 @@ Has 4 primary methods:
- async_get_cache
"""
import asyncio
import itertools
import logging
import time
from collections.abc import Awaitable, Mapping, Sequence
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from threading import Lock
from typing import TYPE_CHECKING, Any, Final
@ -24,7 +23,6 @@ from litellm.types.caching import RedisPipelineIncrementOperation
from .base_cache import BaseCache
from .in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache
from .redis_batch import BatchResult, RedisBatch, active_post_call_redis_batch, active_request_redis_batch
from .redis_cache import RedisCache, RedisCircuitBreakerOpenError, log_redis_failure
if TYPE_CHECKING:
@ -59,14 +57,6 @@ class PendingBatchRead:
previous_access_times: dict[str, float | None]
def _log_deferred_increment_failure(future: asyncio.Future[float]) -> None:
if future.cancelled():
return
failure: Final = future.exception()
if failure is not None:
log_redis_failure(verbose_logger, logging.WARNING, "post-call Redis increment failed", failure)
class DualCache(BaseCache):
"""
DualCache is a cache implementation that updates both Redis and an in-memory cache simultaneously.
@ -269,9 +259,6 @@ class DualCache(BaseCache):
result = in_memory_result
if result is None and self.redis_cache is not None and local_only is False:
request_batch: Final = active_request_redis_batch(self.redis_cache)
if request_batch is not None and request_batch.read_as_missing(key):
return None
# If not found in in-memory cache, try fetching from Redis
redis_result: Final = await self.redis_cache.async_get_cache(key, parent_otel_span=parent_otel_span)
@ -316,20 +303,6 @@ class DualCache(BaseCache):
return sublist_keys, previous_access_times
def reserve_redis_batch_reads(self, keys: Sequence[str]) -> tuple[list[str], dict[str, float | None]]:
"""Reserve the memory-missed keys whose throttled Redis reads are due, as a batch read would."""
if self.redis_cache is None:
return [], {}
key_list: Final = list(keys)
memory: Final = self.in_memory_cache
in_memory_result: Final = (
None
if memory is None # pyright: ignore[reportUnnecessaryComparison] # handle an absent in-memory tier
else memory.batch_get_cache(key_list)
)
result: Final = in_memory_result if in_memory_result is not None else tuple(None for _ in key_list)
return self._reserve_redis_batch_keys(time.time(), key_list, result)
def _rollback_redis_batch_key_reservations(self, previous_access_times: dict[str, float | None]) -> None:
with self._last_redis_batch_access_time_lock:
for key, previous_time in previous_access_times.items():
@ -493,28 +466,6 @@ class DualCache(BaseCache):
verbose_logger, logging.ERROR, "LiteLLM Cache: exception in async add_cache", e, with_traceback=True
)
async def async_set_cache_pre_call(self, key: str, value: object, ttl: float | None) -> BatchResult[None] | None:
"""Memory now, the Redis SET on the request's pipeline, sent with the next read any caller awaits; None
when no pipeline is open, so the caller takes its direct path."""
batch: Final = None if self.redis_cache is None else active_request_redis_batch(self.redis_cache)
return None if batch is None else await self._set_on_batch(batch, key, value, ttl)
async def async_delete_cache_pre_call(self, key: str) -> Awaitable[None] | None:
"""Memory now, the Redis DEL on the request's pipeline; None when no pipeline is open, so the caller
takes its direct path."""
batch: Final = None if self.redis_cache is None else active_request_redis_batch(self.redis_cache)
if batch is None:
return None
if self.in_memory_cache is not None:
self.in_memory_cache.delete_cache(key)
return batch.delete(key)
async def _set_on_batch(self, batch: RedisBatch, key: str, value: object, ttl: float | None) -> BatchResult[None]:
effective_ttl: Final = self.default_in_memory_ttl if ttl is None else ttl
if self.in_memory_cache is not None:
await self.in_memory_cache.async_set_cache(key, value, ttl=effective_ttl)
return batch.set(key, value, effective_ttl)
# async_batch_set_cache
async def async_set_cache_pipeline(
self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs
@ -582,41 +533,6 @@ class DualCache(BaseCache):
)
return result
async def async_increment_cache_post_call(
self,
key: str,
value: float,
ttl: int | None,
parent_otel_span: Span | None = None,
) -> None:
"""Memory is incremented now; the Redis increment rides the request's post-call pipeline when one is
open, and runs on its own as ``async_increment_cache`` otherwise."""
await self.async_increment_cache_pipeline_post_call(
(RedisPipelineIncrementOperation(key=key, increment_value=value, ttl=ttl),), parent_otel_span
)
async def async_increment_cache_pipeline_post_call(
self,
increment_list: Sequence["RedisPipelineIncrementOperation"],
parent_otel_span: Span | None = None,
) -> None:
batch: Final = None if self.redis_cache is None else active_post_call_redis_batch(self.redis_cache)
operations: Final = list(increment_list)
if batch is None:
await self.async_increment_cache_pipeline(operations, parent_otel_span=parent_otel_span)
return
try:
if self.in_memory_cache is not None:
await self.in_memory_cache.async_increment_pipeline(
increment_list=operations, parent_otel_span=parent_otel_span
)
except Exception as e: # noqa: BLE001 # same tolerance as async_increment_cache_pipeline
log_redis_failure(verbose_logger, logging.WARNING, "in-memory increment failed", e)
for operation in increment_list:
batch.increment(operation["key"], operation["increment_value"], operation["ttl"]).on_settled(
_log_deferred_increment_failure
)
async def async_increment_cache_pipeline(
self,
increment_list: list["RedisPipelineIncrementOperation"],

535
litellm/caching/redis_batch.py
View file

@ -1,453 +1,24 @@
"""One Redis pipeline for several independent operations, each with its own result and its own failure.
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.
"""
"""The current request's native Redis batches: a pre-call batch for the reads it needs before answering and a
post-call batch for the writes nobody waits on, both owned by the native ``NativeRequestRedis``."""
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, Generator, Sequence
from contextvars import ContextVar, Token
from dataclasses import dataclass, field
from datetime import timedelta
from enum import Enum, auto
from types import MappingProxyType, TracebackType
from types import TracebackType
from typing import TYPE_CHECKING, Final, Generic, Literal, Protocol, TypeAlias, 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 (
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_cluster_cache import RedisClusterCache
from litellm.types.services import ServiceTypes
if TYPE_CHECKING:
from litellm.rust_bridge._native import NativeRequestRedis, NativeSend
_T = TypeVar("_T")
_T_co: Final = TypeVar("_T_co", covariant=True)
_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]: ...
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]]):
__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:
pipe.mget(tuple(self._redis_cache.check_and_fix_namespace(key=key) for key in self._keys))
return 1
def resolve(self, replies: Sequence[object]) -> Mapping[str, object]:
values: Final = replies[0]
if not isinstance(values, (list, tuple)):
raise TypeError(f"MGET reply is not a list: {type(values).__name__}")
return MappingProxyType(
{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
)
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):
raise ConnectionError("batch get did not return every key")
return found
class _Script(_Op[object]):
__slots__ = ("_args", "_keys", "_redis_cache", "_run", "_sha")
def __init__(
self,
redis_cache: RedisCache,
source: str,
run: RegisteredScript,
keys: Sequence[str],
args: Sequence[_ScriptArg],
) -> 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)
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
def resolve(self, replies: Sequence[object]) -> object:
return replies[0]
async def run_alone(self) -> object:
return await self._run(keys=self._keys, args=self._args)
class _Increment(_Op[float]):
__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:
name: Final = self._redis_cache.check_and_fix_namespace(key=self._key)
pipe.incrbyfloat(name, self._value)
if self._ttl is None:
return 1
pipe.expire(name, timedelta(seconds=self._ttl))
return 2
def resolve(self, replies: Sequence[object]) -> float:
reply: Final = replies[0]
if not isinstance(reply, (int, float, str, bytes)):
raise TypeError(f"INCRBYFLOAT reply is not numeric: {type(reply).__name__}")
return float(reply)
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)):
raise TypeError(f"increment did not return a number: {type(value).__name__}")
return float(value)
class _Set(_Op[None]):
"""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:
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),
json.dumps(self._value),
ex=None if ttl is None else timedelta(seconds=ttl),
)
return 1
def resolve(self, replies: Sequence[object]) -> None:
return None
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]):
"""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:
pipe.delete(self._redis_cache.check_and_fix_namespace(key=self._key))
return 1
def resolve(self, replies: Sequence[object]) -> None:
return None
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)
@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
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)
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)
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
def script(
self, source: str, run: RegisteredScript, keys: Sequence[str], args: Sequence[_ScriptArg]
) -> BatchResult[object]:
return self._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))
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))
def delete(self, key: str) -> BatchResult[None]:
self._misses.add(key)
return self._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)
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._health, 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)}
def _backend_key(redis_cache: RedisCache) -> object:
"""Two ``RedisCache`` instances built from the same connection settings and namespace talk to the same server
under the same key prefix, so the proxy's cache and the router's cache share one pipeline (the router gets its
port as a string, hence the ``str`` comparison); a cache whose settings cannot be compared (a test double) gets
its own."""
try:
settings: Final = tuple(sorted((str(k), str(v)) for k, v in redis_cache.redis_kwargs.items() if v is not None)) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportUnknownArgumentType] # untyped cache API
except AttributeError:
return ("instance", id(redis_cache))
return (type(redis_cache), redis_cache.namespace, settings)
class NativePending(Protocol[_T_co]):
def resolve(self) -> Awaitable[_T_co]: ...
@ -511,29 +82,22 @@ def _native_send(phase: RequestPhase, request: NativeRequestRedis) -> NativeSend
_open_post_call: Final[weakref.WeakSet[RequestRedisBatches]] = weakref.WeakSet()
"""Requests whose post-call batch still holds declared ops, so a shutdown can send them before Redis goes away."""
"""Requests whose post-call batch may still hold declared ops, so a shutdown can send them before Redis goes away."""
class RequestRedisBatches:
"""One ``RedisBatch`` per Redis backend for the current request, so readers of different caches that
share a server (the proxy's and the router's) share the pipeline.
The post-call batches hold the writes nothing waits on (counters, token scripts, the response cache).
They flush once, when the success or failure callbacks have all run, or at ``post_call_deadline``
seconds after the first declaration when no callback phase closes them."""
"""The request's native batches. The post-call batch goes out once, when the success or failure callbacks have
all run, or ``post_call_deadline`` seconds after its first declaration when no callback phase closes it."""
__slots__ = (
"__weakref__",
"_batches",
"_deadline",
"_deadline_flush",
"_native",
"_native_factory",
"_native_flush_hooks",
"_native_misses",
"_native_send_factory",
"_post_call",
"_post_call_settlements",
"_post_call_unsent",
"post_call_deadline",
"prefetched",
)
@ -544,41 +108,17 @@ class RequestRedisBatches:
native_factory: NativeRequestFactory = _native_request,
native_send_factory: NativeSendFactory = _native_send,
) -> None:
self._batches: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend
self._post_call: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend
self._post_call_settlements: tuple[PostCallSettlement, ...] = ()
self._post_call_unsent = False
self.post_call_deadline: Final = post_call_deadline
self._native_factory: Final = native_factory
self._native_send_factory: Final = native_send_factory
self._native: NativeRequestRedis | None = None
self._native_flush_hooks: Final[list[Callable[[], None]]] = [] # mutable-ok: append-only registry
self._native_misses: frozenset[str] = frozenset()
self._deadline: asyncio.TimerHandle | None = None
self._deadline_flush: asyncio.Task[None] | None = None
# Reads declared early for a consumer that runs later in the request, keyed by consumer name.
self.prefetched: Final[dict[str, object]] = {} # mutable-ok: armed pre-admission, taken at use
def batch(self, redis_cache: RedisCache) -> RedisBatch:
key: Final = _backend_key(redis_cache)
batch = self._batches.get(key)
if batch is None:
batch = RedisBatch(redis_cache, name="request_redis_batch")
self._batches[key] = batch
return batch
def post_call(self, redis_cache: RedisCache) -> RedisBatch:
key: Final = _backend_key(redis_cache)
existing: Final = self._post_call.get(key)
batch: Final = (
existing
if existing is not None
else self._post_call.setdefault(key, RedisBatch(redis_cache, name="post_call_redis_batch"))
)
if self._deadline is None:
self._deadline = asyncio.get_running_loop().call_later(self.post_call_deadline, self._flush_on_deadline)
_open_post_call.add(self)
return batch
def native(self) -> NativeRequestRedis:
"""The request's native batches, built on first use so requests that touch no native store pay nothing."""
if self._native is None:
@ -588,6 +128,7 @@ class RequestRedisBatches:
def native_post_call(self) -> NativeRequestRedis:
"""``native`` for a post-call declaration, so a shutdown drain also flushes it."""
_open_post_call.add(self)
self._post_call_unsent = True
return self.native()
def native_send(self, phase: RequestPhase) -> NativeSend:
@ -607,7 +148,7 @@ class RequestRedisBatches:
self._native_flush_hooks.append(hook)
def note_native_misses(self, keys: Sequence[str]) -> None:
"""Keys a native pre-call read found absent (or a native DEL removed), like ``RedisBatch.read_as_missing``."""
"""Keys a native pre-call read found absent, or a native DEL removed, so a later GET can skip Redis."""
self._native_misses = self._native_misses | frozenset(keys)
def forget_native_miss(self, key: str) -> None:
@ -623,31 +164,21 @@ class RequestRedisBatches:
hook()
finally:
_active_request_batches.reset(token)
await self.flush_pre_call()
async def flush_pre_call(self) -> None:
"""Send the pre-call ops declared so far, such as write-backs nobody awaits before the request scope closes.
The pre-call flush hooks are skipped: they declare reads for a reader about to await, and none is waiting."""
if self._native is not None:
await self._native.flush_pre()
def _flush_on_deadline(self) -> None:
self._deadline = None
self._deadline_flush = asyncio.ensure_future(self.flush_post_call())
async def flush_all(self) -> None:
"""Batch whatever is still declared (write-backs nobody awaits) before the request scope closes. The pre-call
flush hooks are skipped: they declare reads for later readers, and none are left."""
await asyncio.gather(
*(batch.flush() for batch in self._batches.values() if batch.pending),
*(() if self._native is None else (self._native.flush_pre(),)),
)
async def flush_post_call(self) -> None:
"""One pipeline per backend for the post-call writes; the deadline is disarmed since this is that flush."""
if self._deadline is not None:
self._deadline.cancel()
self._deadline = None
await asyncio.gather(
*(batch.flush() for batch in self._post_call.values() if batch.pending),
*(() if self._native is None else (self._native.flush_post(),)),
)
if not any(batch.pending for batch in self._post_call.values()):
"""Send the post-call writes. The request stays registered with the shutdown drain until they are back, so a
drain that overlaps this flush waits for it, and an op declared while it runs keeps the request registered."""
self._post_call_unsent = False
if self._native is not None:
await self._native.flush_post()
if not self._post_call_unsent:
_open_post_call.discard(self)
def track_post_call_settlement(self, settlement: PostCallSettlement) -> None:
@ -658,36 +189,16 @@ class RequestRedisBatches:
def post_call_settlements(self) -> tuple[PostCallSettlement, ...]:
return self._post_call_settlements
@property
def batches(self) -> tuple[RedisBatch, ...]:
return tuple(self._batches.values())
_active_request_batches: Final[ContextVar[RequestRedisBatches | None]] = ContextVar(
"request_redis_batches", default=None
)
def active_request_redis_batch(redis_cache: RedisCache) -> RedisBatch | None:
"""The request's batch for this backend, or None outside a ``request_redis_batch_scope``."""
batches: Final = _active_request_batches.get()
if batches is None:
return None
return batches.batch(redis_cache)
def active_request_redis_batches() -> RequestRedisBatches | None:
return _active_request_batches.get()
def active_post_call_redis_batch(redis_cache: RedisCache) -> RedisBatch | None:
"""The request's post-call batch for this backend, or None outside a ``request_redis_batch_scope``."""
batches: Final = _active_request_batches.get()
if batches is None:
return None
return batches.post_call(redis_cache)
async def flush_post_call_redis_batches() -> None:
"""Called where the success and failure callbacks of a request have all run."""
batches: Final = _active_request_batches.get()
@ -701,7 +212,7 @@ async def drain_post_call_redis_batches() -> None:
class request_redis_batch_scope:
"""Redis reads declared inside share one pipeline per backend; nested scopes join the outer one."""
"""Redis ops declared inside share the request's native batches; nested scopes join the outer one."""
__slots__ = ("_native_factory", "_native_send_factory", "_post_call_deadline", "_token")

16
litellm/caching/redis_cache.py
View file

@ -29,7 +29,6 @@ from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast
from pydantic import TypeAdapter
import litellm
from litellm._internal_context import current_service_caller
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import (
DEFAULT_REDIS_MAJOR_VERSION,
@ -112,7 +111,6 @@ class _SupportsScriptLoad(Protocol):
_GENERIC_CALLER_MODULES: Final = frozenset(
{
__name__,
"litellm.caching.redis_batch",
"litellm.caching.dual_cache",
"litellm.caching.caching",
"litellm.rust_bridge.execution",
@ -125,8 +123,6 @@ _GENERIC_CALLER_FRAME_NAMES: Final = frozenset(
"wrapper",
"_run_under_circuit_breaker",
"_run_under_circuit_breaker_sync",
"run_alone",
"_settle_alone",
"get_cache",
"set_cache",
"async_get_cache",
@ -184,11 +180,9 @@ def _get_call_stack_info(num_frames: int = 2) -> str:
Get the function names of the nearest meaningful callers of the cache method.
Frames that merely forward the call (this module's circuit-breaker guards, the
cache facades, the batch pipeline's retry path, generic cache verbs) are
skipped, and the walk stops at the event loop, so the chain names the litellm
code that wanted the call. When nothing but forwarding frames is found (the call
runs in a task of its own, like a batch op retried on the flush) the chain the
declaring code threaded through ``service_caller`` is reported, else ``unknown``.
cache facades, generic cache verbs) are skipped, and the walk stops at the event
loop, so the chain names the litellm code that wanted the call, or ``unknown``
when nothing but forwarding frames is found.
Args:
num_frames: Number of previous frames to include (default: 2)
@ -219,9 +213,7 @@ def _get_call_stack_info(num_frames: int = 2) -> str:
)
function_names: Final = leading_names + further_names
if function_names:
return " <- ".join(function_names)
return current_service_caller() or "unknown"
return " <- ".join(function_names) if function_names else "unknown"
except Exception:
return "unknown"

23
litellm/integrations/otel/README.md
View file

@ -89,23 +89,13 @@ shows as `redis.incr router_cooldowns` followed by `redis.ttl router_cooldowns`
way (`pod_lock`, `budget_reset`, `spend_queue`, `health_check`, `scheduler_queue`,
`managed_files`, `mcp_servers`, ...), so a job tick renders `redis.set pod_lock`
rather than a bare `redis.set`; `tests/unit/test_internal_context.py` scans every
module under `litellm/` and `enterprise/` that calls a shared cache or declares a
read or write on the request batch (`reserve_redis_batch_reads`,
`batch.mget`, `batch.set`, `batch.script`) and fails when
module under `litellm/` and `enterprise/` that calls a shared cache and fails when
one has no declared family, with the process-local `InMemoryCache` callers listed
as the only exemptions. A batch op carries the family that was active when it was
declared, so the routing prefetch armed before deployment selection
(`RoutingPrefetch.arm`) is `router_cooldowns` when only cooldown keys go out,
`router_usage` when only usage counters do, and `router_cooldowns_usage` when both
ride the same MGET, whichever pipeline or standalone read later settles it. A
per-request pipeline (`RedisBatch`) that carries ops of
one family is `"redis.pipeline auth_objects"`; one that carries several owners'
ops is `"redis.pipeline mixed"` with the sorted family list on
`litellm.redis.families` and the op count on `litellm.metadata.op_count` (an int,
never stringified). A cluster client cannot pipeline across slots, so there every
batch op settles on its own and one write-back of three auth objects shows as three
parallel `redis.set auth_objects` spans with the same caller, not one pipeline
span. Every `call_type` the Redis cache layer emits maps to a verb,
ride the same MGET, whichever pipeline or standalone read later settles it. Every `call_type` the Redis cache layer emits maps to a verb,
so the `{service} {call_type}` fallback is unreachable for Redis (a test asserts
it). Postgres helpers are `postgres.{verb} {table}`
(`"postgres.select LiteLLM_VerificationToken"`, `"postgres.update LiteLLM_TeamTable"`,
@ -130,12 +120,9 @@ Redis it is also `db.operation.name`), the target lands on
`litellm.service.target`, and the litellm call chain that issued the call
(`_retrieve_from_cache <- _async_get_cache`) travels as
`ServiceLoggerPayload.caller` onto `litellm.service.caller`, with the forwarding
frames (cache facades, circuit-breaker guards, batch retry wrappers, the native
execution's `lifecycle`/`streams` drivers) skipped so it names the code that wanted
the call. A call whose own frames are all forwarders
(a batch op settled in a task of its own, on a cluster client or a NOSCRIPT retry)
reports the chain its declaring code captured and threaded through
`service_caller(...)`, never the forwarders, and `unknown` when there is none.
frames (cache facades, circuit-breaker guards, the native execution's
`lifecycle`/`streams` drivers) skipped so it names the code that wanted the call,
and `unknown` when its own frames are all forwarders.
The cache key itself is never on the span: it is unbounded and carries key hashes
and session ids, and the span is already named by key family. Like every other
span they parent to the

10
litellm/integrations/otel/mappers/genai.py
View file

@ -10,7 +10,6 @@ table: one lambda per mapping operation, applied against the typed span data.
from collections.abc import Callable
from typing import Final
from litellm._internal_context import REDIS_FAMILIES_METADATA_KEY
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
from litellm.integrations.otel.mappers.utils import (
MAX_TOOL_DEFINITION_ATTRS_PER_SPAN,
@ -200,12 +199,5 @@ class GenAIMapper:
# semconv naming the server it reached. Internal services (router, budget
# jobs, …) have no db.system, so they get only the litellm.service.* keys.
attrs.update(db_span_attributes(data.service_name, data.call_type, postgres_operation(data)))
attrs.update(
{
LiteLLM.REDIS_FAMILIES
if key == REDIS_FAMILIES_METADATA_KEY
else f"{LiteLLM.METADATA_PREFIX}{key}": value
for key, value in data.event_metadata.items()
}
)
attrs.update({f"{LiteLLM.METADATA_PREFIX}{key}": value for key, value in data.event_metadata.items()})
return attrs

2
litellm/integrations/otel/model/semconv.py
View file

@ -334,8 +334,6 @@ class LiteLLM:
SERVICE_CALL_TYPE: Final = "litellm.service.call_type"
SERVICE_CALLER: Final = "litellm.service.caller"
SERVICE_TARGET: Final = "litellm.service.target"
# The sorted, comma-joined key families one Redis pipeline carried ops for; bounded, unlike the keys.
REDIS_FAMILIES: Final = "litellm.redis.families"
PREPROCESSING_MS: Final = "litellm.preprocessing.duration_ms"
# The logical name of the MCP server a tool call was routed to. There is no
# semconv key for an MCP server's *name* (the convention uses ``server.address``

2
litellm/proxy/auth/auth_checks.py
View file

@ -3248,7 +3248,7 @@ async def cache_team_object(
if team_table.team_alias:
alias_key: Final = f"team_alias:{team_table.team_alias}"
try:
pipelined_delete: Final = await user_api_key_cache.async_delete_cache_pre_call(alias_key)
pipelined_delete: Final = user_api_key_cache.declare_delete(alias_key)
if pipelined_delete is None:
await user_api_key_cache.async_delete_cache(key=alias_key)
else:

13
litellm/proxy/common_utils/user_api_key_cache.py
View file

@ -327,11 +327,7 @@ class UserApiKeyCache(DualCache):
await self._set_shared(shared, key, payload, local_only=local_only, pipelined=pipelined, kwargs=kwargs)
return None
if key is not None and is_user_key_cache_key(key):
if pipelined and await self.key_object_cache.async_set_cache_pre_call(key, payload, ttl) is not None:
return None
return await self.key_object_cache.async_set_cache(key=key, value=payload, local_only=local_only, **kwargs)
if pipelined and await super().async_set_cache_pre_call(key, payload, ttl) is not None:
return None
return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs)
@staticmethod
@ -371,12 +367,9 @@ class UserApiKeyCache(DualCache):
return
await super().async_delete_cache(key)
async def async_delete_cache_pre_call(self, key: str) -> Awaitable[None] | None:
if self._shared is not None:
return self._shared.declare_delete(key)
if is_user_key_cache_key(key):
return await self.key_object_cache.async_delete_cache_pre_call(key)
return await super().async_delete_cache_pre_call(key)
def declare_delete(self, key: str) -> Awaitable[None] | None:
"""The DEL on the request's pre-call batch; None without Redis or outside a request, so the caller deletes."""
return None if self._shared is None else self._shared.declare_delete(key)
async def async_delete_cache_keys(self, keys: Sequence[str]) -> None:
"""Batch twin of ``async_delete_cache``, partitioned like

2
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -621,7 +621,7 @@ async def refresh_cached_team(
)
batches: Final = active_request_redis_batches()
if batches is not None:
await batches.flush_all()
await batches.flush_pre_call()
await publish_auth_cache_invalidation(cache_key=f"team_id:{team_row.team_id}")
if team_table.team_alias is not None:
await publish_auth_cache_invalidation(cache_key=f"team_alias:{team_table.team_alias}")

2
litellm/proxy/middleware/redis_request_batch_middleware.py
View file

@ -22,4 +22,4 @@ class RedisRequestBatchMiddleware:
try:
await self.app(scope, receive, send)
finally:
await batches.flush_all()
await batches.flush_pre_call()

1
tests/e2e/perf/cases/event_loop_stalls.py
View file

@ -40,7 +40,6 @@ CASE: Final = replace(
KnownStall("prisma/actions.py", max_count=2),
KnownStall("SpendCounterReseed", max_count=2),
KnownStall("Subscriber._run", max_count=2),
KnownStall("RedisBatch.flush", max_count=2),
KnownStall("run_coroutine_job", max_count=2),
KnownStall("user_api_key_auth.py", max_count=2),
KnownStall("route_streaming_logging_to_handler", max_count=2),

105
tests/unit/caching/redis_batch_double.py Normal file
View file

@ -0,0 +1,105 @@
"""Test doubles for the request's native Redis batches and for a ``RedisCache`` that never reaches a server."""
from __future__ import annotations
import uuid
from collections.abc import Callable, Sequence
from dataclasses import dataclass, field
from typing import Any, Final, Generic, NoReturn, TypeVar
from litellm._service_logger import ServiceLogging
from litellm.caching.redis_cache import RedisCache, RedisHealth
from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings
from litellm.rust_bridge._native import NativeRedisBreaker
_Answer = TypeVar("_Answer")
class FakeRedisCache(RedisCache):
"""A ``RedisCache`` the native stores connect to by its settings; Python itself must never open a client on it."""
def __init__(self, namespace: str | None = None, failure_threshold: int = 5) -> None: # super().__init__ needs a server
self.namespace = namespace
self.redis_kwargs = {"host": f"batch-{uuid.uuid4().hex}.invalid", "port": 1, "socket_timeout": 0.2}
self.service_logger_obj = ServiceLogging()
breaker: Final = RedisBreakerSettings(
failure_threshold=failure_threshold, recovery_timeout=30.0, timeout_min_duration=5.0
)
self._health = RedisHealth( # pyright: ignore[reportGeneralTypeIssues] # the fake skips RedisCache.__init__, which sets it
NativeRedisBreaker.shared(self.redis_kwargs, breaker), self.service_logger_obj
)
self.store_settings = RedisStoreSettings(
redis_kwargs=self.redis_kwargs, namespace=namespace, default_ttl=60, breaker=breaker
)
self.default_ttl = None
self.alone: list[tuple[str, Any]] = []
self.store: dict[str, Any] = {}
def init_async_client(self) -> NoReturn: # pyright: ignore[reportIncompatibleMethodOverride] # no server behind the fake
raise AssertionError("a Python Redis client was opened; Redis traffic goes through the native stores")
async def async_batch_get_cache(self, key_list: Sequence[str], **kwargs: object) -> dict[str, Any]: # pyright: ignore[reportIncompatibleMethodOverride] # records the direct read
self.alone.append(("MGET", tuple(key_list)))
return {key: self.store.get(key) for key in key_list}
async def async_increment(self, key: str, value: float, ttl: int | None = None, **kwargs: object) -> float: # pyright: ignore[reportIncompatibleMethodOverride] # records the direct write
self.alone.append(("INCRBYFLOAT", key, value))
self.store[key] = float(self.store.get(key, 0.0)) + value
return self.store[key]
@dataclass
class FakeNativeRequest:
"""Stands in for ``NativeRequestRedis``: records each flush in ``events`` next to what the test's pendings log."""
post_call_deadline: float
events: list[str] = field(default_factory=list)
async def flush_pre(self) -> None:
self.events.append("flush_pre")
async def flush_post(self) -> None:
self.events.append("flush_post")
@dataclass
class FakeNatives:
built: list[FakeNativeRequest] = field(default_factory=list)
def __call__(self, post_call_deadline: float) -> FakeNativeRequest:
self.built.append(FakeNativeRequest(post_call_deadline))
return self.built[-1]
@dataclass
class FakeNativePending(Generic[_Answer]):
"""A read declared on a ``FakeNativeRequest``; it only answers once that request's pre-call batch was flushed."""
request: FakeNativeRequest
answer: Callable[[], _Answer]
async def resolve(self) -> _Answer:
if "flush_pre" not in self.request.events:
raise ConnectionError("resolved before its batch was flushed")
self.request.events.append("resolve")
return self.answer()
@dataclass
class ReadyPending(Generic[_Answer]):
"""An op declared with ``NativeSend.Now()``: it answers whenever it is resolved."""
answer: Callable[[], _Answer]
async def resolve(self) -> _Answer:
return self.answer()
def fake_send(_phase: object, request: FakeNativeRequest) -> FakeNativeRequest:
"""A ``native_send_factory`` for ``FakeNatives``: the fake request stands in for its ``NativeSend``."""
return request
def fake_pending(send: object, answer: Callable[[], _Answer]) -> FakeNativePending[_Answer] | ReadyPending[_Answer]:
"""What a fake handle declares for ``send``: on the fake request's batch, or on its own for ``Now()``."""
return FakeNativePending(send, answer) if isinstance(send, FakeNativeRequest) else ReadyPending(answer)

29
tests/unit/caching/test_dual_cache.py
View file

@ -137,35 +137,6 @@ def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error():
assert "shared_a" not in dual_cache.last_redis_batch_access_time
def test_reserve_redis_batch_reads_reserves_memory_misses_and_can_be_rolled_back():
mock_redis: Final = MagicMock(spec=RedisCache)
dual_cache: Final = DualCache(
in_memory_cache=InMemoryCache(),
redis_cache=mock_redis,
default_redis_batch_cache_expiry=10,
)
dual_cache.in_memory_cache.set_cache("memory_key", "memory_value")
reserved, previous_access_times = dual_cache.reserve_redis_batch_reads(["memory_key", "missing_key"])
assert reserved == ["missing_key"]
assert previous_access_times == {"missing_key": None}
assert dual_cache.reserve_redis_batch_reads(["memory_key", "missing_key"]) == ([], {})
dual_cache._rollback_redis_batch_key_reservations(previous_access_times)
assert dual_cache.reserve_redis_batch_reads(["memory_key", "missing_key"]) == (
["missing_key"],
{"missing_key": None},
)
def test_reserve_redis_batch_reads_returns_empty_without_redis():
dual_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=None)
assert dual_cache.reserve_redis_batch_reads(["missing_key"]) == ([], {})
def test_dual_cache_batch_get_cache_returns_memory_only_when_redis_read_is_throttled():
mock_redis = _redis_mock_for_sync_batch({"throttled_key": "redis_value"})
dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10)

615
tests/unit/caching/test_redis_batch.py
View file

@ -1,424 +1,35 @@
"""RedisBatch: independent operations share one pipeline, each keeps its own result and failure."""
"""The request scope: one set of native batches per request, nested scopes join it, and the shutdown drain sends
every post-call batch still waiting."""
from __future__ import annotations
import asyncio
import gc
import hashlib
import json
import uuid
import weakref
from collections.abc import Awaitable, Callable, Sequence
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import timedelta
from typing import Any, Final, Generic, TypeVar
from typing import Final
import pytest
from redis.exceptions import NoScriptError
from litellm._internal_context import current_service_target, service_target
from litellm._service_logger import ServiceLogging
from litellm.caching.redis_batch import (
MIXED_PIPELINE_TARGET,
RedisBatch,
RequestRedisBatches,
active_request_redis_batch,
active_request_redis_batches,
drain_post_call_redis_batches,
request_redis_batch_scope,
)
from litellm.caching.redis_cache import (
RedisCache,
RedisHealth,
_get_call_stack_info, # pyright: ignore[reportPrivateUsage] # the chain the service hook reports
)
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings
from litellm.rust_bridge._native import NativeRedisBreaker
SCRIPT = "return redis.call('GET', KEYS[1])"
_Answer = TypeVar("_Answer")
SHA = hashlib.sha1(SCRIPT.encode()).hexdigest() # noqa: S324
from .redis_batch_double import FakeNativeRequest, fake_send
class FakePipeline:
def __init__(self, reply_for: Callable[[tuple[object, ...]], object], fail: Exception | None) -> None:
self.commands: list[tuple[Any, ...]] = []
self.reply_for = reply_for
self.fail = fail
self.executed = False
async def __aenter__(self) -> FakePipeline:
return self
async def __aexit__(self, *exc: object) -> None:
return None
def mget(self, keys: Sequence[str]) -> FakePipeline:
self.commands.append(("MGET", *keys))
return self
def evalsha(self, sha: str, numkeys: int, *keys_and_args: object) -> FakePipeline:
self.commands.append(("EVALSHA", sha, numkeys, *keys_and_args))
return self
def incrbyfloat(self, name: str, amount: float) -> FakePipeline:
self.commands.append(("INCRBYFLOAT", name, amount))
return self
def expire(self, name: str, time: timedelta) -> FakePipeline:
self.commands.append(("EXPIRE", name, int(time.total_seconds())))
return self
def set(self, name: str, value: str, ex: timedelta | None = None) -> FakePipeline:
self.commands.append(("SET", name, value, None if ex is None else int(ex.total_seconds())))
return self
def delete(self, *names: str) -> FakePipeline:
self.commands.append(("DEL", *names))
return self
async def execute(self, raise_on_error: bool = True) -> list[Any]:
assert raise_on_error is False
self.executed = True
if self.fail is not None:
raise self.fail
return [self.reply_for(command) for command in self.commands]
class FakeClient:
def __init__(self, reply_for: Callable[[tuple[object, ...]], object], fail: Exception | None = None) -> None:
self.pipelines: list[FakePipeline] = []
self.reply_for = reply_for
self.fail = fail
def pipeline(self, transaction: bool = True) -> FakePipeline:
assert transaction is False
pipe = FakePipeline(self.reply_for, self.fail)
self.pipelines.append(pipe)
return pipe
class FakeRedisCache(RedisCache):
def __init__(
self, client: FakeClient, namespace: str | None = None, failure_threshold: int = 5
) -> None: # super().__init__ needs a server
self.client = client
self.namespace = namespace
self.redis_kwargs = {"host": f"batch-{uuid.uuid4().hex}.invalid", "port": 1, "socket_timeout": 0.2}
self.service_logger_obj = ServiceLogging()
breaker: Final = RedisBreakerSettings(
failure_threshold=failure_threshold, recovery_timeout=30.0, timeout_min_duration=5.0
)
self._health = RedisHealth( # pyright: ignore[reportGeneralTypeIssues] # the fake skips RedisCache.__init__, which sets it
NativeRedisBreaker.shared(self.redis_kwargs, breaker), self.service_logger_obj
)
self.store_settings = RedisStoreSettings(
redis_kwargs=self.redis_kwargs, namespace=namespace, default_ttl=60, breaker=breaker
)
self.default_ttl = None
self.alone: list[tuple[str, Any]] = []
self.store: dict[str, Any] = {}
def init_async_client(self) -> FakeClient: # pyright: ignore[reportIncompatibleMethodOverride] # fake client, no server
return self.client
async def async_batch_get_cache(self, key_list: Sequence[str], **kwargs: object) -> dict[str, Any]: # pyright: ignore[reportIncompatibleMethodOverride] # records the direct read
self.alone.append(("MGET", tuple(key_list)))
return {key: self.store.get(key) for key in key_list}
async def async_increment(self, key: str, value: float, ttl: int | None = None, **kwargs: object) -> float: # pyright: ignore[reportIncompatibleMethodOverride] # records the direct write
self.alone.append(("INCRBYFLOAT", key, value))
self.store[key] = float(self.store.get(key, 0.0)) + value
return self.store[key]
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # fake, no server
self.alone.append(("SET", key, value))
self.store[key] = value
async def async_delete_cache(self, key: str) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # records the direct delete
self.alone.append(("DEL", key))
self.store.pop(key, None)
async def async_set_cache_pipeline_with_ttls(self, cache_list: Sequence[tuple[str, object, float | None]]) -> None:
self.alone.append(("SET_PIPELINE", tuple(cache_list)))
for key, value, _ttl in cache_list:
self.store[key] = value
class FakeClusterCache(RedisClusterCache, FakeRedisCache):
def __init__(self, client: FakeClient) -> None: # super().__init__ needs a server
FakeRedisCache.__init__(self, client)
def replies(command: tuple[Any, ...]) -> Any:
match command[0]:
case "MGET":
return [json.dumps({"k": key}) if key.endswith("hit") else None for key in command[1:]]
case "EVALSHA":
return [1, 2]
case "INCRBYFLOAT":
return b"3.5"
case "EXPIRE":
return 1
case "SET":
return True
case "DEL":
return 1
raise AssertionError(command)
@dataclass
class FakeNativeRequest:
"""Stands in for ``NativeRequestRedis``: records each flush in ``events`` next to what the test's pendings log."""
post_call_deadline: float
events: list[str] = field(default_factory=list)
async def flush_pre(self) -> None:
self.events.append("flush_pre")
async def flush_post(self) -> None:
self.events.append("flush_post")
@dataclass
class FakeNatives:
built: list[FakeNativeRequest] = field(default_factory=list)
def __call__(self, post_call_deadline: float) -> FakeNativeRequest:
self.built.append(FakeNativeRequest(post_call_deadline))
return self.built[-1]
@dataclass
class FakeNativePending(Generic[_Answer]):
"""A read declared on a ``FakeNativeRequest``; it only answers once that request's pre-call batch was flushed."""
request: FakeNativeRequest
answer: Callable[[], _Answer]
async def resolve(self) -> _Answer:
if "flush_pre" not in self.request.events:
raise ConnectionError("resolved before its batch was flushed")
self.request.events.append("resolve")
return self.answer()
@dataclass
class ReadyPending(Generic[_Answer]):
"""An op declared with ``NativeSend.Now()``: it answers whenever it is resolved."""
answer: Callable[[], _Answer]
async def resolve(self) -> _Answer:
return self.answer()
def fake_send(_phase: object, request: FakeNativeRequest) -> FakeNativeRequest:
"""A ``native_send_factory`` for ``FakeNatives``: the fake request stands in for its ``NativeSend``."""
return request
def fake_pending(send: object, answer: Callable[[], _Answer]) -> FakeNativePending[_Answer] | ReadyPending[_Answer]:
"""What a fake handle declares for ``send``: on the fake request's batch, or on its own for ``Now()``."""
return FakeNativePending(send, answer) if isinstance(send, FakeNativeRequest) else ReadyPending(answer)
def make(
fail: Exception | None = None, namespace: str | None = None, failure_threshold: int = 5
) -> tuple[FakeRedisCache, FakeClient]:
client = FakeClient(replies, fail)
return FakeRedisCache(client, namespace, failure_threshold), client
async def run_alone_script(keys: Sequence[str], args: Sequence[Any]) -> object:
return ["alone", *keys, *args]
@pytest.mark.asyncio
async def test_pipeline_flush_reports_its_name_as_the_call_type_and_the_op_count_as_metadata() -> None:
"""The service event is ``request_redis_batch`` with ``op_count`` on the metadata, not
``request_redis_batch[3]``: the span renders as ``redis.pipeline`` and the metrics label
stays one value per batch name instead of one per batch size."""
cache, _client = make()
events: list[dict[str, Any]] = []
async def record(**kwargs: Any) -> None:
events.append(kwargs)
cache.service_logger_obj.async_service_success_hook = record # pyright: ignore[reportAttributeAccessIssue] # fake, records the hook call
batch = RedisBatch(cache, name="request_redis_batch")
got = batch.mget(["a:hit"])
incr = batch.increment("cnt", 1)
await got
await incr
await asyncio.gather(*(t for t in asyncio.all_tasks() if t is not asyncio.current_task()))
(event,) = events
assert event["call_type"] == "request_redis_batch"
assert event["event_metadata"] == {"op_count": 2}
@pytest.mark.asyncio
async def test_one_pipeline_carries_every_declared_operation_and_awaiting_one_flushes_all() -> None:
cache, client = make(namespace="ns")
batch = RedisBatch(cache)
got = batch.mget(["a:hit", "b", "a:hit"])
script = batch.script(SCRIPT, run_alone_script, ["w"], [7, "x"])
incr = batch.increment("cnt", 2.5, ttl=60)
plain = batch.increment("cnt2", 1)
assert client.pipelines == []
assert await got == {"a:hit": {"k": "ns:a:hit"}, "b": None}
assert script.done and incr.done and plain.done
assert await script == [1, 2]
assert await incr == 3.5
assert await plain == 3.5
assert batch.flushes == 1
assert [pipe.commands for pipe in client.pipelines] == [
[
("MGET", "ns:a:hit", "ns:b"),
("EVALSHA", SHA, 1, "ns:w", 7, "x"),
("INCRBYFLOAT", "ns:cnt", 2.5),
("EXPIRE", "ns:cnt", 60),
("INCRBYFLOAT", "ns:cnt2", 1),
]
]
assert cache.alone == []
@pytest.mark.asyncio
async def test_operations_declared_after_a_flush_go_out_in_the_next_pipeline() -> None:
cache, client = make()
batch = RedisBatch(cache)
await batch.mget(["a"])
later = batch.increment("cnt", 1)
assert not later.done
assert await later == 3.5
assert batch.flushes == 2
assert [pipe.commands for pipe in client.pipelines] == [[("MGET", "a")], [("INCRBYFLOAT", "cnt", 1)]]
@pytest.mark.asyncio
async def test_a_failing_reply_fails_only_its_own_operation() -> None:
def reply_for(command: tuple[Any, ...]) -> Any:
if command[0] == "EVALSHA":
return ValueError("script blew up")
return replies(command)
client = FakeClient(reply_for)
cache = FakeRedisCache(client)
batch = RedisBatch(cache)
got = batch.mget(["a:hit"])
script = batch.script(SCRIPT, run_alone_script, ["w"], [])
assert await got == {"a:hit": {"k": "a:hit"}}
with pytest.raises(ValueError, match="script blew up"):
await script
assert cache.alone == []
@pytest.mark.asyncio
async def test_a_reply_an_operation_cannot_decode_fails_only_that_operation() -> None:
def reply_for(command: tuple[Any, ...]) -> Any:
if command[0] == "MGET":
return "not-a-list"
return replies(command)
client = FakeClient(reply_for)
cache = FakeRedisCache(client)
batch = RedisBatch(cache)
got = batch.mget(["a:hit"])
written = batch.set("w", {"k": 1})
script = batch.script(SCRIPT, run_alone_script, ["w"], [])
with pytest.raises(TypeError, match="MGET reply is not a list"):
await got
assert await written is None
assert await script == [1, 2]
assert len(client.pipelines) == 1
@pytest.mark.asyncio
async def test_pipeline_failure_fails_every_operation_and_trips_the_breaker() -> None:
cache, _client = make(fail=ConnectionError("redis down"), failure_threshold=1)
batch = RedisBatch(cache)
got = batch.mget(["a"])
incr = batch.increment("cnt", 1)
with pytest.raises(ConnectionError):
await got
with pytest.raises(ConnectionError):
await incr
assert cache._health.breaker.state == "open" # pyright: ignore[reportPrivateUsage] # the breaker the pipeline fed
@pytest.mark.asyncio
async def test_noscript_reply_reruns_that_script_through_the_registered_executor() -> None:
def reply_for(command: tuple[Any, ...]) -> Any:
if command[0] == "EVALSHA":
return NoScriptError("NOSCRIPT No matching script")
return replies(command)
client = FakeClient(reply_for)
cache = FakeRedisCache(client)
batch = RedisBatch(cache)
script = batch.script(SCRIPT, run_alone_script, ["w"], [1])
incr = batch.increment("cnt", 1)
assert await script == ["alone", "w", 1]
assert await incr == 3.5
assert batch.flushes == 1
@pytest.mark.asyncio
async def test_cluster_cache_runs_each_operation_on_its_own_path() -> None:
client = FakeClient(replies)
cache = FakeClusterCache(client)
cache.store["a"] = 4
batch = RedisBatch(cache)
got = batch.mget(["a", "b"])
incr = batch.increment("cnt", 2)
assert await got == {"a": 4, "b": None}
assert await incr == 2.0
assert client.pipelines == []
assert cache.alone == [("MGET", ("a", "b")), ("INCRBYFLOAT", "cnt", 2)]
@pytest.mark.asyncio
async def test_flush_hook_lets_a_lazy_reader_join_the_pipeline_that_is_going_out() -> None:
cache, client = make()
batch = RedisBatch(cache)
joined: list[Any] = []
batch.add_flush_hook(lambda: joined.append(batch.mget(["late"])))
await batch.mget(["early"])
assert len(joined) == 1 and joined[0].done
assert await joined[0] == {"late": None}
assert [pipe.commands for pipe in client.pipelines] == [[("MGET", "early"), ("MGET", "late")]]
@pytest.mark.asyncio
async def test_concurrent_awaiters_share_one_flush() -> None:
cache, client = make()
batch = RedisBatch(cache)
first = batch.mget(["a"])
second = batch.mget(["b"])
results = await asyncio.gather(first._wait(), second._wait()) # pyright: ignore[reportPrivateUsage]
assert results == [{"a": None}, {"b": None}]
assert batch.flushes == 1
assert len(client.pipelines) == 1
def test_request_scope_hands_out_one_batch_per_backend_and_nests() -> None:
cache_a, _ = make()
cache_b, _ = make()
assert active_request_redis_batch(cache_a) is None
with request_redis_batch_scope() as batches:
first = active_request_redis_batch(cache_a)
assert first is not None
assert active_request_redis_batch(cache_a) is first
assert active_request_redis_batch(cache_b) is not first
def test_nested_request_scopes_share_the_outer_request_and_close_with_it() -> None:
assert active_request_redis_batches() is None
with request_redis_batch_scope() as outer:
assert active_request_redis_batches() is outer
with request_redis_batch_scope() as inner:
assert inner is batches
assert active_request_redis_batch(cache_a) is first
assert active_request_redis_batch(cache_a) is first
assert len(batches.batches) == 2
assert active_request_redis_batch(cache_a) is None
assert inner is outer
assert active_request_redis_batches() is outer
assert active_request_redis_batches() is None
def test_a_request_scope_that_built_its_native_batches_is_freed_once_it_closes() -> None:
@ -432,168 +43,48 @@ def test_a_request_scope_that_built_its_native_batches_is_freed_once_it_closes()
assert closed() is None
@pytest.mark.asyncio
async def test_a_key_an_mget_read_as_absent_stays_known_missing_until_something_sets_it() -> None:
cache, client = make()
batch = RedisBatch(cache)
values = await batch.mget(["a-hit", "b-miss"])
assert values == {"a-hit": {"k": "a-hit"}, "b-miss": None}
assert batch.read_as_missing("b-miss") is True
assert batch.read_as_missing("a-hit") is False
assert batch.read_as_missing("never-read") is False
batch.set("b-miss", "now-present")
assert batch.read_as_missing("b-miss") is False
@dataclass
class _DeclaresWhileFlushingPost(FakeNativeRequest):
declare_once: list[Callable[[], object]] = field(default_factory=list)
async def flush_post(self) -> None:
await super().flush_post()
while self.declare_once:
self.declare_once.pop()()
@pytest.mark.asyncio
async def test_a_delete_rides_the_pipeline_under_the_namespace_and_reads_as_missing_afterwards() -> None:
cache, client = make(namespace="ns")
batch = RedisBatch(cache)
gone = batch.delete("team_alias:x")
got = batch.mget(["a-hit"])
assert await gone is None
assert await got == {"a-hit": {"k": "ns:a-hit"}}
assert len(client.pipelines) == 1
assert client.pipelines[0].commands[0] == ("DEL", "ns:team_alias:x")
assert batch.read_as_missing("team_alias:x") is True
assert cache.alone == []
async def test_a_post_call_write_declared_while_the_post_call_batch_goes_out_is_still_sent_by_the_shutdown_drain() -> (
None
):
native: Final = _DeclaresWhileFlushingPost(post_call_deadline=60.0)
with request_redis_batch_scope(native_factory=lambda _deadline: native, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
request.native_post_call()
native.declare_once.append(request.native_post_call)
await request.flush_post_call()
await drain_post_call_redis_batches()
assert native.events == ["flush_post", "flush_post"]
@dataclass
class _HeldPostFlush(FakeNativeRequest):
sent: asyncio.Event = field(default_factory=asyncio.Event)
async def flush_post(self) -> None:
await self.sent.wait()
await super().flush_post()
@pytest.mark.asyncio
async def test_a_delete_on_a_cluster_cache_runs_as_its_own_del() -> None:
client = FakeClient(replies)
cache = FakeClusterCache(client)
cache.store["team_alias:x"] = "stale"
batch = RedisBatch(cache)
assert await batch.delete("team_alias:x") is None
assert cache.alone == [("DEL", "team_alias:x")]
assert "team_alias:x" not in cache.store
assert client.pipelines == []
@pytest.mark.asyncio
async def test_a_failed_mget_marks_nothing_as_missing() -> None:
cache, client = make(fail=ConnectionError("down"))
batch = RedisBatch(cache)
with pytest.raises(ConnectionError):
await batch.mget(["b-miss"])
assert batch.read_as_missing("b-miss") is False
@pytest.mark.asyncio
async def test_an_operation_retried_alone_keeps_the_target_it_was_declared_under() -> None:
"""The retry runs on the flush, outside the declaring caller's block, so the op carries
the target it was declared under and the retried call is still named by its purpose."""
seen: list[str | None] = []
async def record_target(keys: Sequence[str], args: Sequence[Any]) -> object:
seen.append(current_service_target())
return ["alone", *keys]
def reply_for(command: tuple[Any, ...]) -> Any:
if command[0] == "EVALSHA":
return NoScriptError("NOSCRIPT")
return replies(command)
cache = FakeRedisCache(FakeClient(reply_for))
batch = RedisBatch(cache)
with service_target("spend_counters"):
script = batch.script(SCRIPT, record_target, ["w"], [])
assert current_service_target() is None
assert await script == ["alone", "w"]
assert seen == ["spend_counters"]
assert current_service_target() is None
class CallerRecordingClusterCache(FakeClusterCache):
def __init__(self, client: FakeClient) -> None:
super().__init__(client)
self.callers: list[str] = []
async def async_batch_get_cache(self, key_list: Sequence[str], **kwargs: object) -> dict[str, Any]: # pyright: ignore[reportIncompatibleMethodOverride] # records what the service hook would report
self.callers.append(_get_call_stack_info())
return await super().async_batch_get_cache(key_list, **kwargs)
def _prefetch_auth_objects(batch: RedisBatch) -> Awaitable[Sequence[Any]]:
return batch.mget(["team", "user"])
@pytest.mark.asyncio
async def test_a_cluster_op_names_the_code_that_declared_it_not_its_wrappers() -> None:
"""On a cluster client every op runs alone, in a task driven by the flush, so above its
wrappers there is only the event loop. Production reported ``_run_under_circuit_breaker <-
wrapper``; the op carries the chain captured where it was declared and reports that."""
cache = CallerRecordingClusterCache(FakeClient(replies))
batch = RedisBatch(cache)
with service_target("auth_objects"):
pending = _prefetch_auth_objects(batch)
assert await pending == {"team": None, "user": None}
assert cache.callers == [
"_prefetch_auth_objects <- test_a_cluster_op_names_the_code_that_declared_it_not_its_wrappers"
]
async def _flush_and_record_service_events(
cache: FakeRedisCache, *results: Awaitable[object]
) -> list[dict[str, object]]:
events: list[dict[str, object]] = [] # mutable-ok: filled by the recording hooks
async def record(**kwargs: object) -> None:
events.append({**kwargs, "target": current_service_target()})
cache.service_logger_obj.async_service_success_hook = record # pyright: ignore[reportAttributeAccessIssue] # fake, records the hook call
cache.service_logger_obj.async_service_failure_hook = record # pyright: ignore[reportAttributeAccessIssue] # fake, records the hook call
await asyncio.gather(*results, return_exceptions=True)
await asyncio.gather(*(t for t in asyncio.all_tasks() if t is not asyncio.current_task()))
return events
@pytest.mark.asyncio
async def test_pipeline_of_one_key_family_is_targeted_by_that_family() -> None:
"""Every op in the flush was declared under ``auth_objects``, so the span is
``redis.pipeline auth_objects`` and carries only the op count."""
cache, _client = make()
batch = RedisBatch(cache, name="request_redis_batch")
with service_target("auth_objects"):
first = batch.mget(["a:hit"])
second = batch.mget(["b:hit"])
(event,) = await _flush_and_record_service_events(cache, first, second)
assert (event["target"], event["event_metadata"]) == ("auth_objects", {"op_count": 2})
@pytest.mark.asyncio
async def test_pipeline_of_several_key_families_is_mixed_and_lists_the_families_sorted() -> None:
"""Owners of different families sharing one round trip render as ``redis.pipeline mixed``
with the sorted family list beside the op count, never as a bare ``redis.pipeline``."""
cache, _client = make()
batch = RedisBatch(cache, name="request_redis_batch")
with service_target("spend_counters"):
incr = batch.increment("cnt", 1)
with service_target("auth_objects"):
auth = batch.mget(["a:hit"])
with service_target("router_cooldowns"):
cooldown = batch.mget(["c:hit"])
(event,) = await _flush_and_record_service_events(cache, incr, auth, cooldown)
assert event["target"] == MIXED_PIPELINE_TARGET
assert event["event_metadata"] == {"op_count": 3, "families": "auth_objects,router_cooldowns,spend_counters"}
assert current_service_target() is None
@pytest.mark.asyncio
async def test_failed_pipeline_reports_the_same_family_target_as_a_successful_one() -> None:
"""The failure event names the pipeline the same way, so the error span lines up with the
success spans of the same flush shape in a trace search."""
cache, _client = make(fail=ConnectionError("redis down"))
batch = RedisBatch(cache, name="post_call_redis_batch")
with service_target("spend_counters"):
incr = batch.increment("cnt", 1)
with service_target("auth_objects"):
auth = batch.mget(["a:hit"])
(event,) = await _flush_and_record_service_events(cache, incr, auth)
assert isinstance(event["error"], ConnectionError)
assert (event["call_type"], event["target"]) == ("post_call_redis_batch", MIXED_PIPELINE_TARGET)
assert event["event_metadata"] == {"op_count": 2, "families": "auth_objects,spend_counters"}
async def test_a_shutdown_drain_that_overlaps_a_post_call_flush_waits_for_its_writes() -> None:
native: Final = _HeldPostFlush(post_call_deadline=60.0)
with request_redis_batch_scope(native_factory=lambda _deadline: native, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
request.native_post_call()
flush: Final = asyncio.create_task(request.flush_post_call())
await asyncio.sleep(0)
drain: Final = asyncio.create_task(drain_post_call_redis_batches())
await asyncio.sleep(0)
drained_while_in_flight: Final = drain.done()
native.sent.set()
await asyncio.gather(flush, drain)
assert not drained_while_in_flight

40
tests/unit/caching/test_redis_cache.py
View file

@ -1595,7 +1595,7 @@ async def test_async_rpush_and_trim_runs_push_and_trim_in_one_transaction(monkey
def test_call_stack_info_skips_generic_cache_facade_frames():
"""A read through ``DualCache.async_get_cache`` -> ``RedisCache.async_get_cache`` used to
report ``async_get_cache <- async_get_cache``; the chain names the code that wanted the
read, skipping the facade verbs and the batch retry wrappers in between."""
read, skipping the facade verbs and the circuit-breaker guard in between."""
from litellm.caching.redis_cache import _get_call_stack_info
def probe(): # the RedisCache method that sets call_type
@ -1604,11 +1604,11 @@ def test_call_stack_info_skips_generic_cache_facade_frames():
def async_get_cache(): # a facade's generic verb
return probe()
def run_alone(): # the batch retry wrapper
def _run_under_circuit_breaker(): # the breaker guard
return async_get_cache()
def _retrieve_from_cache():
return run_alone()
return _run_under_circuit_breaker()
def _async_get_cache():
return _retrieve_from_cache()
@ -1630,31 +1630,7 @@ def test_call_stack_info_stops_at_the_event_loop():
assert asyncio.run(_lookup()) == "_lookup"
def test_call_stack_info_reports_the_threaded_caller_when_only_wrappers_are_found():
"""A batch op retried on the flush runs in a task of its own, so above its wrappers there
is only the event loop; the chain is the one its declaring code threaded through
``service_caller``, never the wrapper names (``run_alone <- _settle_alone`` says nothing)."""
from litellm._internal_context import service_caller
from litellm.caching.redis_cache import _get_call_stack_info
def probe():
return _get_call_stack_info()
def run_alone():
return probe()
async def _settle_alone():
return run_alone()
async def flush():
with service_caller("prefetch_auth_objects <- user_api_key_auth"):
task = asyncio.create_task(_settle_alone())
return await task
assert asyncio.run(flush()) == "prefetch_auth_objects <- user_api_key_auth"
def test_call_stack_info_is_unknown_when_only_wrappers_are_found_and_nothing_was_threaded():
def test_call_stack_info_is_unknown_when_only_wrappers_are_found():
import threading
from litellm.caching.redis_cache import _get_call_stack_info
@ -1662,14 +1638,14 @@ def test_call_stack_info_is_unknown_when_only_wrappers_are_found_and_nothing_was
def probe():
return _get_call_stack_info()
def run_alone():
def wrapper():
return probe()
def _settle_alone():
return run_alone()
def _run_under_circuit_breaker():
return wrapper()
seen: list[str] = []
worker = threading.Thread(target=lambda: seen.append(_settle_alone()))
worker = threading.Thread(target=lambda: seen.append(_run_under_circuit_breaker()))
worker.start()
worker.join()
assert seen == ["unknown"]

174
tests/unit/caching/test_request_redis_batch_post_call.py
View file

@ -1,14 +1,11 @@
"""One Redis pipeline per backend for the post-call writes of a request: spend counters, rate-limit token
scripts and slot releases and deployment TPM all ride the post-call batch, which goes out once the success/failure
callbacks have run (or at the deadline when no callback phase closes it). The response-cache SET stays direct so the
next identical request can hit it while the callbacks are still running."""
"""The post-call writes of a request: spend counters, rate-limit token scripts and slot releases and deployment TPM
all ride the native post-call batch, which goes out once the success/failure callbacks have run (or at the deadline
when no callback phase closes it). The response-cache SET stays direct so the next identical request can hit it while
the callbacks are still running."""
from __future__ import annotations
import asyncio
import datetime
import hashlib
import json
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Final
@ -21,7 +18,7 @@ from litellm.caching.caching import Cache
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_batch import (
active_post_call_redis_batch,
active_request_redis_batches,
drain_post_call_redis_batches,
flush_post_call_redis_batches,
request_redis_batch_scope,
@ -38,12 +35,11 @@ from litellm.proxy.utils import InternalUsageCache
from litellm.redis.stores.rate_limits import CounterIncrement, RateLimitScope
from litellm.redis.stores.router_usage import CounterIncrement as UsageIncrement
from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import ModelResponse
from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits
from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, usage_store
from .test_redis_batch import FakeClient, FakeNatives, FakeRedisCache, fake_send
from .redis_batch_double import FakeNatives, FakeRedisCache, fake_send
async def _script_outside_the_pipeline(keys: Sequence[str], args: Sequence[object]) -> object:
@ -56,11 +52,6 @@ class PostCallFakeRedisCache(FakeRedisCache):
def async_register_script(self, script: str) -> Callable[..., Awaitable[object]]:
return _script_outside_the_pipeline
async def async_increment_pipeline(
self, increment_list: list[RedisPipelineIncrementOperation], **kwargs: object
) -> list[float]:
return [await self.async_increment(op["key"], op["increment_value"]) for op in increment_list]
async def async_delete_cache(self, key: str, **kwargs: object) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # the fake drops RedisCache's unused kwargs
self.alone.append(("DEL", key))
self.store.pop(key, None)
@ -70,36 +61,6 @@ class PostCallFakeRedisCache(FakeRedisCache):
self.store[key] = value
def sha_of(script: str) -> str:
return hashlib.sha1(script.encode()).hexdigest() # noqa: S324
def _ok_replies(command: tuple[object, ...]) -> object:
match command[0]:
case "INCRBYFLOAT":
return b"7.5"
case "EXPIRE":
return 1
case "SET":
return True
case "EVALSHA":
return [3, 0]
case "MGET":
return [json.dumps({"spend": 1.0}) for _ in command[1:]]
raise AssertionError(command)
async def _run_ready_callbacks(client: FakeClient) -> None:
for _ in range(20):
if client.pipelines:
return
await asyncio.sleep(0)
def _names(client: FakeClient, index: int = 0) -> list[str]:
return [command[0] for command in client.pipelines[index].commands]
def _limiter(fakes: FakeRateLimits) -> PROXY_MaxParallelRequestsHandler_v3:
return PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=InternalUsageCache(dual_cache=DualCache()), rate_limit_store=fakes.store()
@ -154,8 +115,7 @@ def _tpm_kwargs() -> Mapping[str, object]:
@pytest.mark.asyncio
async def test_post_call_owners_declare_on_the_native_post_call_batch_until_the_callbacks_finish():
client = FakeClient(_ok_replies)
redis_cache = PostCallFakeRedisCache(client)
redis_cache = PostCallFakeRedisCache()
calls: list[tuple[str, list[str], list[object]]] = [] # mutable-ok: filled by the fake scripts
fakes = FakeRateLimits(token_increment_script=_scripts(calls, [10.0]), parallel_release_script=_scripts(calls, [0]))
limiter = _limiter(fakes)
@ -177,7 +137,6 @@ async def test_post_call_owners_declare_on_the_native_post_call_batch_until_the_
assert usage.calls == [("post_call", (UsageIncrement(tpm_counter, 42, 60),))]
assert (tpm_counter.deployment_id, tpm_counter.model, tpm_counter.kind) == ("dep-a", "openai/gpt-4o-mini", "tpm")
assert usage.local == {tpm_counter: 42}
assert client.pipelines == []
assert redis_cache.alone == []
@ -193,8 +152,7 @@ async def test_outside_a_request_the_tpm_increment_goes_out_on_its_own_pipeline(
@pytest.mark.asyncio
async def test_the_response_cache_set_reaches_redis_before_the_post_call_pipeline_goes_out():
client = FakeClient(_ok_replies)
redis_cache = PostCallFakeRedisCache(client)
redis_cache = PostCallFakeRedisCache()
response_cache = _response_cache(redis_cache)
kwargs = {"messages": [{"role": "user", "content": "hi"}], "model": "gpt", "ttl": 120}
cache_key = response_cache.get_cache_key(**kwargs)
@ -204,15 +162,13 @@ async def test_the_response_cache_set_reaches_redis_before_the_post_call_pipelin
assert redis_cache.store[cache_key]["response"] == {"id": "resp"}
await flush_post_call_redis_batches()
assert client.pipelines == []
(direct_set,) = redis_cache.alone
assert (direct_set[0], direct_set[1], direct_set[2]["ttl"]) == ("SET", cache_key, 120)
@pytest.mark.asyncio
async def test_a_chat_response_written_through_the_handler_dual_cache_is_in_memory_and_redis_at_once():
client = FakeClient(_ok_replies)
redis_cache = PostCallFakeRedisCache(client)
redis_cache = PostCallFakeRedisCache()
response_cache = _response_cache(redis_cache)
handler_cache = DualCache(redis_cache=redis_cache, in_memory_cache=InMemoryCache())
kwargs = {"messages": [{"role": "user", "content": "hi"}], "model": "gpt", "ttl": 120}
@ -225,7 +181,6 @@ async def test_a_chat_response_written_through_the_handler_dual_cache_is_in_memo
assert redis_cache.store[cache_key]["response"] == '{"id": "resp"}'
await flush_post_call_redis_batches()
assert client.pipelines == []
(direct_set,) = redis_cache.alone
assert (direct_set[0], direct_set[1], direct_set[2]["ttl"]) == ("SET", cache_key, 120)
@ -279,41 +234,19 @@ async def test_a_released_slot_is_free_locally_at_once_and_the_older_redis_count
@pytest.mark.asyncio
async def test_failure_refunds_ride_the_post_call_pipeline_and_count_in_memory_at_once():
client = FakeClient(_ok_replies)
dual_cache = DualCache()
dual_cache.attach_redis_cache(PostCallFakeRedisCache(client))
refund = [RedisPipelineIncrementOperation(key="{api_key:k1}:tokens", increment_value=-500, ttl=60)]
with request_redis_batch_scope():
await dual_cache.async_increment_cache_pipeline_post_call(refund)
assert await dual_cache.in_memory_cache.async_get_cache("{api_key:k1}:tokens") == -500
assert client.pipelines == []
await flush_post_call_redis_batches()
assert client.pipelines[0].commands[0] == ("INCRBYFLOAT", "{api_key:k1}:tokens", -500)
@pytest.mark.asyncio
async def test_outside_a_request_scope_owners_write_directly_as_before():
client = FakeClient(_ok_replies)
redis_cache = PostCallFakeRedisCache(client)
dual_cache = DualCache()
dual_cache.attach_redis_cache(redis_cache)
async def test_outside_a_request_scope_the_response_cache_writes_directly():
redis_cache = PostCallFakeRedisCache()
response_cache = _response_cache(redis_cache)
await dual_cache.async_increment_cache_post_call("dep:tpm", 42, ttl=60)
await response_cache.async_add_cache({"id": "resp"}, messages=[{"role": "user", "content": "hi"}], model="gpt")
assert client.pipelines == []
assert redis_cache.alone[0] == ("INCRBYFLOAT", "dep:tpm", 42)
assert active_post_call_redis_batch(redis_cache) is None
(direct_set,) = redis_cache.alone
assert direct_set[0] == "SET"
@pytest.mark.asyncio
async def test_a_set_with_options_keeps_the_direct_path():
client = FakeClient(_ok_replies)
redis_cache = PostCallFakeRedisCache(client)
redis_cache = PostCallFakeRedisCache()
response_cache = _response_cache(redis_cache)
with request_redis_batch_scope():
@ -322,78 +255,21 @@ async def test_a_set_with_options_keeps_the_direct_path():
)
await flush_post_call_redis_batches()
assert client.pipelines == []
(direct_set,) = redis_cache.alone
assert direct_set[0] == "SET" and direct_set[2]["nx"] is True
@pytest.mark.asyncio
async def test_two_backends_get_one_post_call_pipeline_each():
a_client, b_client = FakeClient(_ok_replies), FakeClient(_ok_replies)
a, b = DualCache(), DualCache()
a.attach_redis_cache(PostCallFakeRedisCache(a_client))
b.attach_redis_cache(PostCallFakeRedisCache(b_client))
with request_redis_batch_scope():
await a.async_increment_cache_post_call("x", 1, ttl=None)
await b.async_increment_cache_post_call("y", 1, ttl=None)
await a.async_increment_cache_post_call("z", 1, ttl=None)
await flush_post_call_redis_batches()
assert len(a_client.pipelines) == 1 and len(b_client.pipelines) == 1
assert [c[1] for c in a_client.pipelines[0].commands if c[0] == "INCRBYFLOAT"] == ["x", "z"]
@pytest.mark.asyncio
async def test_post_call_writes_still_waiting_on_their_callbacks_are_drained_at_shutdown():
client = FakeClient(_ok_replies)
dual_cache = DualCache()
dual_cache.attach_redis_cache(PostCallFakeRedisCache(client))
with request_redis_batch_scope(post_call_deadline=60) as request:
await dual_cache.async_increment_cache_post_call("x", 1, ttl=None)
await request.flush_all()
assert client.pipelines == []
await drain_post_call_redis_batches()
assert len(client.pipelines) == 1 and _names(client) == ["INCRBYFLOAT"]
await drain_post_call_redis_batches()
assert len(client.pipelines) == 1
@pytest.mark.asyncio
async def test_a_post_call_batch_nobody_closes_goes_out_at_the_deadline(monkeypatch: pytest.MonkeyPatch):
client = FakeClient(_ok_replies)
dual_cache = DualCache()
dual_cache.attach_redis_cache(PostCallFakeRedisCache(client))
loop = asyncio.get_running_loop()
armed_at = loop.time()
with request_redis_batch_scope(post_call_deadline=60) as request:
await dual_cache.async_increment_cache_post_call("x", 1, ttl=None)
await request.flush_all()
await _run_ready_callbacks(client)
assert client.pipelines == [], "the request boundary drains the immediate batch, not the post-call one"
monkeypatch.setattr(loop, "time", lambda: armed_at + 61)
await _run_ready_callbacks(client)
assert len(client.pipelines) == 1 and _names(client) == ["INCRBYFLOAT"]
@pytest.mark.asyncio
async def test_the_success_handler_closes_the_post_call_batch_after_the_last_callback(monkeypatch):
client = FakeClient(_ok_replies)
dual_cache = DualCache()
dual_cache.attach_redis_cache(PostCallFakeRedisCache(client))
pipelines_seen_by_callbacks: list[int] = []
natives: Final = FakeNatives()
flushes_seen_by_callbacks: list[int] = []
class Counter(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
await dual_cache.async_increment_cache_post_call("counted", 1, ttl=None)
pipelines_seen_by_callbacks.append(len(client.pipelines))
request = active_request_redis_batches()
assert request is not None
request.native_post_call()
flushes_seen_by_callbacks.append(natives.built[0].events.count("flush_post"))
monkeypatch.setattr(litellm, "_async_success_callback", [])
logging_obj = LitellmLogging(
@ -415,11 +291,11 @@ async def test_the_success_handler_closes_the_post_call_batch_after_the_last_cal
"model_parameters": {},
}
with request_redis_batch_scope():
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double
await logging_obj.async_success_handler(result=ModelResponse(), standard_logging_object=payload)
assert pipelines_seen_by_callbacks == [0, 0]
assert len(client.pipelines) == 1 and _names(client) == ["INCRBYFLOAT", "INCRBYFLOAT"]
assert flushes_seen_by_callbacks == [0, 0]
assert natives.built[0].events == ["flush_post"]
@pytest.mark.asyncio
@ -469,7 +345,7 @@ async def test_native_post_call_ops_go_out_with_the_post_call_flush_not_the_requ
with request_redis_batch_scope(post_call_deadline=60, native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
request.native_post_call()
await request.flush_all()
await request.flush_pre_call()
at_the_boundary = list(natives.built[0].events)
await flush_post_call_redis_batches()
await drain_post_call_redis_batches()
@ -484,7 +360,7 @@ async def test_native_post_call_ops_still_waiting_on_their_callbacks_are_drained
with request_redis_batch_scope(post_call_deadline=60, native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
request.native_post_call()
await request.flush_all()
await request.flush_pre_call()
await drain_post_call_redis_batches()
await drain_post_call_redis_batches()

97
tests/unit/caching/test_request_redis_batch_pre_call.py
View file

@ -38,8 +38,7 @@ from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cach
from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits
from tests.unit.redis.stores.router_usage_double import FakeUsageRead
from .test_redis_batch import (
FakeClient,
from .redis_batch_double import (
FakeNativePending,
FakeNativeRequest,
FakeNatives,
@ -107,14 +106,6 @@ def _declared_before_the_first_flush(natives: FakeNatives) -> list[str]:
return events[: events.index("flush_pre")]
def _lua_ok_replies(command: tuple[Any, ...]) -> Any:
if command[0] == "MGET":
return [None for _ in command[1:]]
if command[0] == "SET":
return True
raise AssertionError(command)
@pytest.mark.asyncio
async def test_descriptor_lua_calls_share_one_native_pipeline_and_each_keeps_its_result():
limiter, _fakes = _limiter()
@ -329,16 +320,10 @@ _PING: Final = [{"role": "user", "content": "ping"}]
_PINNED_MINUTE: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc)
def _python_cooldown_reads(client: FakeClient, redis_cache: FakeRedisCache) -> tuple[str, ...]:
pipelined: Final = tuple(
key
for pipeline in client.pipelines
for command in pipeline.commands
if command[0] == "MGET"
for key in command[1:]
def _python_cooldown_reads(redis_cache: FakeRedisCache) -> tuple[str, ...]:
return tuple(
key for command, keys in redis_cache.alone if command == "MGET" for key in keys if key.endswith(":cooldown")
)
alone: Final = tuple(key for command, keys in redis_cache.alone if command == "MGET" for key in keys)
return tuple(key for key in pipelined + alone if key.endswith(":cooldown"))
async def _route(router: Router) -> str:
@ -488,8 +473,7 @@ async def test_a_usage_based_pick_reads_the_counters_the_prefetch_armed():
async def test_a_cooldown_recorded_through_the_store_excludes_its_deployment_without_a_python_cooldown_read(
routing_strategy: str,
):
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
redis_cache = FakeRedisCache()
router = _router(redis_cache, routing_strategy=routing_strategy)
router.cooldown_cache.cooldown_store.put(
"dep-a",
@ -502,12 +486,12 @@ async def test_a_cooldown_recorded_through_the_store_excludes_its_deployment_wit
picks = {await _route(router) for _ in range(5)}
assert picks == {"dep-b"}
assert _python_cooldown_reads(client, redis_cache) == ()
assert _python_cooldown_reads(redis_cache) == ()
@pytest.mark.asyncio
async def test_arming_outside_a_request_scope_is_a_no_op():
redis_cache = FakeRedisCache(FakeClient(_lua_ok_replies))
redis_cache = FakeRedisCache()
router = _router(redis_cache)
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
assert active_request_redis_batches() is None
@ -515,8 +499,7 @@ async def test_arming_outside_a_request_scope_is_a_no_op():
@pytest.mark.asyncio
async def test_simple_shuffle_cooldown_prefetch_does_not_arm_usage_reads_or_read_eagerly():
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
redis_cache = FakeRedisCache()
router = _router(redis_cache, routing_strategy="simple-shuffle")
with request_redis_batch_scope() as request:
@ -525,57 +508,9 @@ async def test_simple_shuffle_cooldown_prefetch_does_not_arm_usage_reads_or_read
deployment = await _route(router)
assert deployment in {"dep-a", "dep-b"}
assert client.pipelines == []
assert redis_cache.alone == []
@pytest.mark.asyncio
async def test_two_backends_flush_concurrently_one_pipeline_each():
a_client, b_client = FakeClient(_lua_ok_replies), FakeClient(_lua_ok_replies)
a, b = FakeRedisCache(a_client), FakeRedisCache(b_client)
with request_redis_batch_scope() as request:
ra = request.batch(a).mget(["x", "y"])
rb = request.batch(b).mget(["x"])
await asyncio.gather(ra, rb)
assert len(a_client.pipelines) == 1 and len(b_client.pipelines) == 1
class _SameServerCache(FakeRedisCache):
def __init__(self, client: FakeClient, namespace: str | None = None, **redis_kwargs: object) -> None:
super().__init__(client, namespace)
self.redis_kwargs = redis_kwargs
@pytest.mark.asyncio
async def test_caches_built_from_the_same_connection_settings_share_the_request_pipeline():
client = FakeClient(_lua_ok_replies)
proxy_cache = _SameServerCache(client, host="r", port=6379, db=0)
router_cache = _SameServerCache(FakeClient(_lua_ok_replies), port="6379", host="r", db=0, password=None)
other_cache = _SameServerCache(FakeClient(_lua_ok_replies), host="r", port=6380, db=0)
with request_redis_batch_scope() as request:
assert request.batch(proxy_cache) is request.batch(router_cache)
assert request.batch(proxy_cache) is not request.batch(other_cache)
a = request.batch(proxy_cache).mget(["a"])
b = request.batch(router_cache).mget(["b"])
await asyncio.gather(a, b)
assert len(client.pipelines) == 1
assert [c[0] for c in client.pipelines[0].commands] == ["MGET", "MGET"]
@pytest.mark.asyncio
async def test_caches_on_one_server_with_different_namespaces_keep_their_own_key_prefix():
proxy_client, router_client = FakeClient(_lua_ok_replies), FakeClient(_lua_ok_replies)
proxy_cache = _SameServerCache(proxy_client, namespace="proxy", host="r", port=6379, db=0)
router_cache = _SameServerCache(router_client, namespace="router", host="r", port=6379, db=0)
with request_redis_batch_scope() as request:
await asyncio.gather(request.batch(proxy_cache).mget(["a"]), request.batch(router_cache).mget(["b"]))
sent: Final = tuple(
tuple(command for pipe in client.pipelines for command in pipe.commands)
for client in (proxy_client, router_client)
)
assert sent == ((("MGET", "proxy:a"),), (("MGET", "router:b"),)), "each cache reads under its own namespace"
def _user_entry() -> tuple[_CacheEntry, LiteLLM_UserTable]:
entry = _CacheEntry("user-1", "user_row", LiteLLM_UserTable, 42)
return entry, LiteLLM_UserTable(user_id="user-1", max_budget=None, spend=0.0)
@ -589,7 +524,7 @@ async def test_auth_write_back_rides_the_native_pre_call_batch_and_the_scope_dra
await _write_back([_user_entry()], cache)
assert cache.get_local("user-1")["user_id"] == "user-1"
assert natives.built[0].events == ["declare write_back (('user-1', 42),)"]
await request.flush_all()
await request.flush_pre_call()
assert natives.built[0].events[-1] == "flush_pre"
assert fake.commands == []
@ -609,7 +544,7 @@ async def test_a_key_the_request_read_as_absent_is_not_read_again_by_a_per_key_g
assert await cache.async_get_cache("absent-key") is None
await cache.async_set_cache("absent-key", {"v": 1}, ttl=5)
assert await cache.async_get_cache("absent-key") == {"v": 1}
await request.flush_all()
await request.flush_pre_call()
assert fake.commands == []
@ -622,7 +557,7 @@ async def test_management_object_writes_inside_a_request_ride_its_pipeline_and_w
await cache.async_set_cache("hashed-key-object", {"token": "hashed-key-object"}, ttl=60)
assert cache.get_local("team_id:t1") == {"team_id": "t1"}
assert await cache.async_get_cache("hashed-key-object") == {"token": "hashed-key-object"}
await request.flush_all()
await request.flush_pre_call()
assert natives.built[0].events == [
"declare set team_id:t1 ttl=60.0",
"declare set hashed-key-object ttl=60.0",
@ -634,7 +569,7 @@ async def test_management_object_writes_inside_a_request_ride_its_pipeline_and_w
@pytest.mark.asyncio
async def test_a_team_refresh_inside_a_request_sends_its_set_and_alias_del_in_one_pipeline_before_returning():
redis_cache = FakeRedisCache(FakeClient(_lua_ok_replies))
redis_cache = FakeRedisCache()
cache, fake = shared_cache(connection=redis_cache)
usage_cache = DualCache(redis_cache=redis_cache)
usage_cache.in_memory_cache.set_cache("team_id:t1", "stale team")
@ -655,7 +590,7 @@ async def test_a_team_refresh_inside_a_request_sends_its_set_and_alias_del_in_on
assert usage_cache.in_memory_cache.get_cache("team_alias:alpha") is None
assert cache.get_local("team_alias:alpha") is None
assert cache.get_local("team_id:t1")["team_id"] == "t1"
await request.flush_all()
await request.flush_pre_call()
assert fake.commands == [] and redis_cache.alone == []
@ -689,7 +624,7 @@ async def test_the_auth_prefetch_and_the_armed_cooldown_read_go_out_in_one_nativ
await prefetch_identity_keys(["key-hit"], cache, before_read=lambda: arm_cooldowns(cooldowns, ["dep-a"]))
arm_cooldowns(cooldowns, ["dep-a"])
active = await read_cooldowns(cooldowns, ["dep-a"], None)
await request.flush_all()
await request.flush_pre_call()
assert active == ()
assert len(natives.built) == 1
@ -721,7 +656,7 @@ async def test_awaiting_a_native_read_runs_the_native_hooks_then_flushes_the_nat
result: Final = request.native_result(FakeNativePending(native, lambda: "answer"))
declared_only: Final = list(native.events)
answer: Final = await result
await request.flush_all()
await request.flush_pre_call()
assert request.native() is native
assert [built.post_call_deadline for built in natives.built] == [7]
@ -751,7 +686,7 @@ async def test_a_request_that_no_native_store_asks_for_builds_no_native_batches(
natives: Final = FakeNatives()
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
await request.flush_all()
await request.flush_pre_call()
await request.flush_post_call()
assert natives.built == []

32
tests/unit/integrations/otel/test_otel_v2_logger.py
View file

@ -3537,20 +3537,21 @@ _REDIS_CACHE_MODULES = (
"litellm/caching/redis_cluster_cache.py",
"litellm/caching/redis_semantic_cache.py",
"litellm/caching/dual_cache.py",
"litellm/caching/redis_batch.py",
)
_NATIVE_REQUEST_BATCHES = "litellm/rust/crates/python-bridge/src/redis/store/request.rs"
def _redis_call_types_emitted_by_the_cache_layer():
"""Every ``call_type`` literal the Redis cache layer hands to the service logger, plus the
two ``RedisBatch`` names it passes as ``call_type=self.name``."""
names the native request batches report their flushes under."""
import re
from pathlib import Path
repo = Path(__file__).resolve().parents[4]
sources = "\n".join((repo / module).read_text() for module in _REDIS_CACHE_MODULES)
literal = frozenset(re.findall(r'call_type="([a-z_]+)"', sources))
batch_names = frozenset(re.findall(r'RedisBatch\([^)]*name="([a-z_]+)"', sources))
native_request = (repo / _NATIVE_REQUEST_BATCHES).read_text()
batch_names = frozenset(re.findall(r'call\(py, "([a-z_]+)"\)', native_request))
return sorted(literal | batch_names)
@ -3573,31 +3574,6 @@ def test_no_call_type_the_redis_cache_layer_emits_can_fall_back_to_the_raw_metho
assert fallbacks == []
def test_mixed_pipeline_families_land_on_their_own_bounded_attribute():
"""A flush that carried several owners' ops is ``redis.pipeline mixed``; the sorted family
list goes to ``litellm.redis.families`` (not under ``litellm.metadata.``) while ``op_count``
stays an int on ``litellm.metadata.op_count``, and the legacy vocabulary keeps the bare keys."""
logger, exporter = _logger()
parent = _service_parent(logger)
try:
asyncio.run(
logger.async_service_success_hook(
payload=_ServicePayload("redis", "request_redis_batch", target="mixed"),
parent_otel_span=parent,
event_metadata={"op_count": 3, "families": "auth_objects,spend_counters"},
)
)
finally:
parent.end()
(span,) = [s for s in exporter.get_finished_spans() if s.name.startswith("redis")]
assert span.name == "redis.pipeline mixed"
assert span.attributes[LiteLLM.REDIS_FAMILIES] == "auth_objects,spend_counters"
assert span.attributes["families"] == "auth_objects,spend_counters"
assert f"{LiteLLM.METADATA_PREFIX}families" not in span.attributes
assert span.attributes[f"{LiteLLM.METADATA_PREFIX}op_count"] == 3
assert type(span.attributes[f"{LiteLLM.METADATA_PREFIX}op_count"]) is int
def test_deferred_close_starts_at_the_provider_handoff_not_the_logging_objects_birth():
"""A destination logger never sees ``pre_call``, so its copy of ``chat`` is created at
close. Starting it at the logging object's ``start_time`` makes it span routing and the

12
tests/unit/integrations/test_prometheus_carried_budget_state.py
View file

@ -13,11 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from prometheus_client import REGISTRY
from litellm.caching.redis_batch import (
active_post_call_redis_batch,
flush_post_call_redis_batches,
request_redis_batch_scope,
)
from litellm.caching.redis_batch import flush_post_call_redis_batches, request_redis_batch_scope
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy import proxy_server
from litellm.proxy._types import (
@ -41,7 +37,7 @@ from litellm.types.proxy.carried_budget_state import (
TeamBudgetSnapshot,
UserBudgetSnapshot,
)
from tests.unit.caching.test_redis_batch import FakeNatives, fake_send, make
from tests.unit.caching.redis_batch_double import FakeNatives, fake_send
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
TEAM_RESET_AT = datetime(2026, 10, 1, tzinfo=timezone.utc)
@ -478,15 +474,13 @@ async def test_gauges_do_not_add_the_own_charge_twice_once_the_pipeline_went_out
async def test_an_unrelated_post_call_flush_does_not_make_the_gauge_drop_its_own_unsent_charge(
prometheus_logger, getters, redis_spend_counters
):
"""An unrelated write went out first, then another request's 2.0 reached the counters. This request's 2.0 counts
"""A post-call flush went out first, then another request's 2.0 reached the counters. This request's 2.0 counts
once whether it is undeclared, declared but unsent, or settled."""
_stub_gauges(prometheus_logger)
_store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0)
unrelated_cache, _ = make()
metadata: Final = _request_metadata(_token_committed_at_auth())
with _request_scope():
active_post_call_redis_batch(unrelated_cache).increment("unrelated-write", 1.0)
await flush_post_call_redis_batches()
await _emit(prometheus_logger, metadata)
undeclared: Final = _remaining(prometheus_logger)

2
tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
View file

@ -6966,7 +6966,7 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias():
cache = MagicMock()
cache.async_set_cache = AsyncMock()
cache.async_delete_cache = AsyncMock()
cache.async_delete_cache_pre_call = AsyncMock(return_value=None) # no request pipeline open
cache.declare_delete = MagicMock(return_value=None)
logging_obj = MagicMock()
logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock()

2
tests/unit/proxy/management_endpoints/test_team_endpoints.py
View file

@ -2525,7 +2525,7 @@ async def test_refresh_cached_team_broadcasts_team_and_alias_invalidations(
async def test_refresh_cached_team_flushes_redis_set_before_broadcast(monkeypatch) -> None:
from litellm.caching.redis_batch import request_redis_batch_scope
from litellm.proxy.management_endpoints import team_endpoints
from tests.unit.caching.test_redis_batch import FakeNatives, fake_send
from tests.unit.caching.redis_batch_double import FakeNatives, fake_send
from tests.unit.redis.stores.auth_cache_double import shared_cache
natives: Final = FakeNatives()

4
tests/unit/proxy/spend_tracking/test_spend_counter_batch.py
View file

@ -26,7 +26,7 @@ from litellm.proxy.spend_tracking.spend_counter_batch import (
)
from litellm.rust_bridge._native import NativeRedisStoreError
from litellm.caching.redis_batch import request_redis_batch_scope
from tests.unit.caching.test_redis_batch import FakeNativePending, FakeNatives, fake_send
from tests.unit.caching.redis_batch_double import FakeNativePending, FakeNatives, fake_send
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
TOKEN = UserAPIKeyAuth(token="hashed", team_id="team", user_id="user", org_id="org")
@ -147,7 +147,7 @@ async def test_counters_bound_at_auth_cost_no_redis_read_when_the_request_ends_b
spend_counter_batch_scope(spend.store()),
):
bind_admission_counter_keys(TOKEN, end_user_id=None)
await request.flush_all()
await request.flush_pre_call()
assert spend.declared == []
assert spend.commands == []

2
tests/unit/redis/stores/router_usage_double.py
View file

@ -12,7 +12,7 @@ from typing import Final, TypeVar
from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter
from litellm.router_python_host import PythonHost
from litellm.rust_bridge._native import NativeRedisClient
from tests.unit.caching.test_redis_batch import fake_pending
from tests.unit.caching.redis_batch_double import fake_pending
T = TypeVar("T")

4
tests/unit/redis/stores/test_auth_cache.py
View file

@ -11,7 +11,7 @@ from litellm.proxy.common_utils.user_api_key_cache import is_user_key_cache_key
from litellm.redis.stores.auth_cache import AuthCacheSettings, AuthCacheStore
from litellm.redis.stores.connection import connect
from litellm.rust_bridge._native import LocalSizes, NativeAuthCache
from tests.unit.caching.test_redis_batch import FakeNatives, fake_send
from tests.unit.caching.redis_batch_double import FakeNatives, fake_send
from .auth_cache_double import FakeAuthCache
@ -134,7 +134,7 @@ async def test_a_key_the_request_read_found_absent_is_answered_without_another_r
refilled = await store.get("team_id:absent")
store.declare_delete("team_id:hit")
deleted = await store.get("team_id:hit")
await request.flush_all()
await request.flush_pre_call()
assert found == {"team_id:absent": None, "team_id:hit": {"team_id": "hit"}}
assert absent is None

2
tests/unit/redis/stores/test_rate_limits.py
View file

@ -12,7 +12,7 @@ from litellm.redis.stores.rate_limits import (
RateLimitScope,
RateLimitStore,
)
from tests.unit.caching.test_redis_batch import FakeNatives, fake_send
from tests.unit.caching.redis_batch_double import FakeNatives, fake_send
@dataclass(frozen=True, slots=True)

2
tests/unit/redis/stores/test_router_cooldowns.py
View file

@ -22,7 +22,7 @@ from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSetting
from litellm.rust_bridge._native import NativeRedisClient, NativeRouterCooldowns
from litellm.rust_bridge.host import settings
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
from tests.unit.caching.test_redis_batch import (
from tests.unit.caching.redis_batch_double import (
FakeNativePending,
FakeNatives,
ReadyPending,

2
tests/unit/redis/stores/test_router_usage.py
View file

@ -13,7 +13,7 @@ from litellm.redis.stores.connection import connect
from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter
from litellm.rust_bridge._native import NativeRedisClient
from litellm.types.utils import ModelResponse
from tests.unit.caching.test_redis_batch import FakeNatives, fake_send
from tests.unit.caching.redis_batch_double import FakeNatives, fake_send
T = TypeVar("T")

2
tests/unit/redis/stores/test_spend_counters.py
View file

@ -27,7 +27,7 @@ from litellm.redis.stores.spend_counters import (
)
from litellm.redis.store_settings import RedisStoreSettings
from litellm.rust_bridge._native import NativeRedisClient, NativeRedisStoreError, NativeSpendCounters, SpendRead
from tests.unit.caching.test_redis_batch import FakeNatives, fake_send
from tests.unit.caching.redis_batch_double import FakeNatives, fake_send
from .spend_counters_double import FakeSettlement, FakeSpendCounters

2
tests/unit/router_utils/test_routing_read_batch.py
View file

@ -19,7 +19,7 @@ from litellm.router_config import build_router_config
from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2
from litellm.router_utils import routing_prefetch
from litellm.utils import get_utc_datetime
from tests.unit.caching.test_redis_batch import FakeNativePending, FakeNatives, ReadyPending, fake_pending, fake_send
from tests.unit.caching.redis_batch_double import FakeNativePending, FakeNatives, ReadyPending, fake_pending, fake_send
from tests.unit.conftest import redis_cache_double
from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, UsageHost
from tests.unit.redis.stores.test_router_cooldowns import FakeCooldowns

34
tests/unit/test_internal_context.py
View file

@ -1,6 +1,6 @@
"""``with_service_target`` and ``service_caller`` carry the purpose and the caller of a datastore call
to code that cannot see them from its own frames, and every Redis producer on the proxy request path
declares a key family so no request-path span renders as a bare ``redis.get``."""
"""``with_service_target`` carries the purpose of a datastore call to code that cannot see it from its own
frames, and every Redis producer on the proxy request path declares a key family so no request-path span renders
as a bare ``redis.get``."""
import ast
import asyncio
@ -11,14 +11,10 @@ from dataclasses import dataclass
from pathlib import Path
from typing import Final
import pytest
from litellm._internal_context import (
current_request_parent_otel_span,
current_service_caller,
current_service_target,
request_parent_otel_span,
service_caller,
service_target,
set_request_parent_otel_span,
with_service_target,
@ -100,8 +96,7 @@ _CACHE_CALL: Final = re.compile(
r"\.(?:async_)?(?:get_cache|set_cache|batch_get_cache|batch_get_cache_shared|increment_cache|increment_if_present|increment"
r"|set_cache_pipeline|set_cache_pipeline_with_ttls|set_cache_sadd|delete_cache|batch_set_cache|increment_pipeline"
r"|rpush|lpop|scan_iter|get_ttl|mget)\("
r"|\b(?:reserve_redis_batch_reads|_prepare_batch_get)\("
r"|\bbatch\.(?:set|delete|script|increment)\("
r"|\b_prepare_batch_get\("
)
_DECLARES_TARGET: Final = re.compile(r"\b(?:with_service_target|service_target|response_cache_phase)\(")
_BUILDS_A_REDIS_CACHE: Final = re.compile(r"\bRedisCache\(|\bredis_cache=(?!None\b)")
@ -173,27 +168,6 @@ def test_with_service_target_keeps_the_wrapped_signature_and_coroutine_ness() ->
assert hook.__name__ == "hook"
def test_service_caller_is_inherited_by_a_task_spawned_inside_it_and_cleared_after() -> None:
async def spawned() -> str | None:
return current_service_caller()
async def main() -> tuple[str | None, str | None]:
with service_caller("prefetch <- auth"):
task = asyncio.create_task(spawned())
return await task, current_service_caller()
assert asyncio.run(main()) == ("prefetch <- auth", None)
@pytest.mark.parametrize("value", [None, "x"])
def test_service_caller_restores_the_outer_value(value: str | None) -> None:
with service_caller(value):
with service_caller("inner"):
assert current_service_caller() == "inner"
assert current_service_caller() == value
assert current_service_caller() is None
class _Suspend:
def __await__(self) -> Generator[None]:
yield