mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
refactor(caching): drop the RedisBatch pre-call SET, DEL and miss tracking
UserApiKeyCache never hands DualCache a Redis connection: with Redis attached every op goes through the native AuthCacheStore, and without Redis both inner DualCaches have redis_cache None, so DualCache.async_set_cache_pre_call and async_delete_cache_pre_call always returned None. The RedisBatch _Set/_Delete ops and the _misses set they fed were unreachable outside tests. UserApiKeyCache.async_delete_cache_pre_call becomes a sync declare_delete that forwards to the native store. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee
parent
df20f90d12
commit
bb2ac782dc
6 changed files with 9 additions and 189 deletions
|
|
@ -12,7 +12,7 @@ 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 +24,7 @@ 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_batch import active_post_call_redis_batch
|
||||
from .redis_cache import RedisCache, RedisCircuitBreakerOpenError, log_redis_failure
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -269,9 +269,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)
|
||||
|
||||
|
|
@ -493,28 +490,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
|
||||
|
|
|
|||
|
|
@ -11,7 +11,6 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import weakref
|
||||
|
|
@ -58,8 +57,6 @@ class _RedisPipeline(Protocol):
|
|||
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]: ...
|
||||
|
||||
|
||||
|
|
@ -221,55 +218,6 @@ class _Increment(_Op[float]):
|
|||
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."""
|
||||
|
||||
|
|
@ -306,23 +254,10 @@ class RedisBatch:
|
|||
_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
|
||||
return self._declare(_MGet(self.redis_cache, keys))
|
||||
|
||||
def script(
|
||||
self, source: str, run: RegisteredScript, keys: Sequence[str], args: Sequence[_ScriptArg]
|
||||
|
|
@ -332,14 +267,6 @@ class RedisBatch:
|
|||
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)
|
||||
|
|
@ -607,7 +534,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:
|
||||
|
|
|
|||
|
|
@ -3228,7 +3228,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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -65,14 +65,6 @@ class 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
|
||||
|
|
@ -127,19 +119,6 @@ class FakeRedisCache(RedisCache):
|
|||
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
|
||||
|
|
@ -156,10 +135,6 @@ def replies(command: tuple[Any, ...]) -> Any:
|
|||
return b"3.5"
|
||||
case "EXPIRE":
|
||||
return 1
|
||||
case "SET":
|
||||
return True
|
||||
case "DEL":
|
||||
return 1
|
||||
raise AssertionError(command)
|
||||
|
||||
|
||||
|
|
@ -324,11 +299,9 @@ async def test_a_reply_an_operation_cannot_decode_fails_only_that_operation() ->
|
|||
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
|
||||
|
||||
|
|
@ -418,54 +391,6 @@ def test_request_scope_hands_out_one_batch_per_backend_and_nests() -> None:
|
|||
assert active_request_redis_batch(cache_a) 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
|
||||
|
||||
|
||||
@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 == []
|
||||
|
||||
|
||||
@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
|
||||
|
|
|
|||
|
|
@ -6901,7 +6901,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()
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue