diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index 39fe789160b..73e805cadcf 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -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.""" diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index cbb01cae603..f0d400b7b01 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -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"], diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index b3e6d2eb291..b8f14d7d711 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -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") diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 3f047e6001d..620b9de3eb3 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -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" diff --git a/litellm/integrations/otel/README.md b/litellm/integrations/otel/README.md index 70384e49f53..69a79419992 100644 --- a/litellm/integrations/otel/README.md +++ b/litellm/integrations/otel/README.md @@ -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 diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index 4e07d090edd..8c22381f031 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -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 diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index 8d186a2715a..02b713fdcdd 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -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`` diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8ecfce31076..fc818334f9c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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: diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 91529ee36fc..5119addaefc 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -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 diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f3d3c01b3a4..00b6f899c01 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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}") diff --git a/litellm/proxy/middleware/redis_request_batch_middleware.py b/litellm/proxy/middleware/redis_request_batch_middleware.py index bfb5f79a174..e1db6291d68 100644 --- a/litellm/proxy/middleware/redis_request_batch_middleware.py +++ b/litellm/proxy/middleware/redis_request_batch_middleware.py @@ -22,4 +22,4 @@ class RedisRequestBatchMiddleware: try: await self.app(scope, receive, send) finally: - await batches.flush_all() + await batches.flush_pre_call() diff --git a/tests/e2e/perf/cases/event_loop_stalls.py b/tests/e2e/perf/cases/event_loop_stalls.py index 8a5654cbf7b..97bd44cf144 100644 --- a/tests/e2e/perf/cases/event_loop_stalls.py +++ b/tests/e2e/perf/cases/event_loop_stalls.py @@ -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), diff --git a/tests/unit/caching/redis_batch_double.py b/tests/unit/caching/redis_batch_double.py new file mode 100644 index 00000000000..9b2349998b8 --- /dev/null +++ b/tests/unit/caching/redis_batch_double.py @@ -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) diff --git a/tests/unit/caching/test_dual_cache.py b/tests/unit/caching/test_dual_cache.py index bba0c215015..c51a429655e 100644 --- a/tests/unit/caching/test_dual_cache.py +++ b/tests/unit/caching/test_dual_cache.py @@ -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) diff --git a/tests/unit/caching/test_redis_batch.py b/tests/unit/caching/test_redis_batch.py index 158b9754309..ddd0da2df75 100644 --- a/tests/unit/caching/test_redis_batch.py +++ b/tests/unit/caching/test_redis_batch.py @@ -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 diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index d98b71f2ce3..47e8428bf33 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -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"] diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index e2ddda29944..85e3fd24cbe 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -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() diff --git a/tests/unit/caching/test_request_redis_batch_pre_call.py b/tests/unit/caching/test_request_redis_batch_pre_call.py index b8e7658ef24..b09c11b6920 100644 --- a/tests/unit/caching/test_request_redis_batch_pre_call.py +++ b/tests/unit/caching/test_request_redis_batch_pre_call.py @@ -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 == [] diff --git a/tests/unit/integrations/otel/test_otel_v2_logger.py b/tests/unit/integrations/otel/test_otel_v2_logger.py index 017eb745d90..02e3a343b97 100644 --- a/tests/unit/integrations/otel/test_otel_v2_logger.py +++ b/tests/unit/integrations/otel/test_otel_v2_logger.py @@ -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 diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index 9e006c20079..3ce44cbd154 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -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) diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 4cbe5320aa1..f6f744218ef 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -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() diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 58a852247e7..401c5173ea8 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -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() diff --git a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py index e9fad4ef374..e631cad2dfa 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py +++ b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py @@ -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 == [] diff --git a/tests/unit/redis/stores/router_usage_double.py b/tests/unit/redis/stores/router_usage_double.py index 1cf472776a3..2a1095c2cf7 100644 --- a/tests/unit/redis/stores/router_usage_double.py +++ b/tests/unit/redis/stores/router_usage_double.py @@ -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") diff --git a/tests/unit/redis/stores/test_auth_cache.py b/tests/unit/redis/stores/test_auth_cache.py index e03b7b81417..53e18845057 100644 --- a/tests/unit/redis/stores/test_auth_cache.py +++ b/tests/unit/redis/stores/test_auth_cache.py @@ -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 diff --git a/tests/unit/redis/stores/test_rate_limits.py b/tests/unit/redis/stores/test_rate_limits.py index 489604152c8..2a6df27ae25 100644 --- a/tests/unit/redis/stores/test_rate_limits.py +++ b/tests/unit/redis/stores/test_rate_limits.py @@ -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) diff --git a/tests/unit/redis/stores/test_router_cooldowns.py b/tests/unit/redis/stores/test_router_cooldowns.py index e86536d0a3d..c576d9ef562 100644 --- a/tests/unit/redis/stores/test_router_cooldowns.py +++ b/tests/unit/redis/stores/test_router_cooldowns.py @@ -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, diff --git a/tests/unit/redis/stores/test_router_usage.py b/tests/unit/redis/stores/test_router_usage.py index 54891c22a8f..22232c1b1aa 100644 --- a/tests/unit/redis/stores/test_router_usage.py +++ b/tests/unit/redis/stores/test_router_usage.py @@ -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") diff --git a/tests/unit/redis/stores/test_spend_counters.py b/tests/unit/redis/stores/test_spend_counters.py index 1d39cffc933..e7859873de7 100644 --- a/tests/unit/redis/stores/test_spend_counters.py +++ b/tests/unit/redis/stores/test_spend_counters.py @@ -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 diff --git a/tests/unit/router_utils/test_routing_read_batch.py b/tests/unit/router_utils/test_routing_read_batch.py index 96f46bc9e02..23e2d3fd11e 100644 --- a/tests/unit/router_utils/test_routing_read_batch.py +++ b/tests/unit/router_utils/test_routing_read_batch.py @@ -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 diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index cba28f075ac..4682aaffef3 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -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