diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index cbb01cae603..68c008fc1a0 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -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 diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index b3e6d2eb291..26afc43d94d 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -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: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e546c093473..1c7d64023bf 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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: 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/tests/unit/caching/test_redis_batch.py b/tests/unit/caching/test_redis_batch.py index 536d510f6d6..a58a9dddc1a 100644 --- a/tests/unit/caching/test_redis_batch.py +++ b/tests/unit/caching/test_redis_batch.py @@ -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 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 78ce9728f3b..caddd4ec06c 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 @@ -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()