diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 1cc3e106aa0..efa3caf3a4d 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -47,7 +47,7 @@ class InMemoryCache(BaseCache): self.max_size_per_item = max_size_per_item or MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB # 1MB = 1024KB # in-memory cache - self.cache_dict: dict = {} + self.cache_dict: dict[object, Any] = {} self.ttl_dict: dict = {} self.expiration_heap: list[tuple[float, str]] = [] self._increment_lock = threading.Lock() @@ -100,17 +100,17 @@ class InMemoryCache(BaseCache): except KeyError: return False - def _cache_keys(self) -> tuple[str, ...]: - return tuple(self.cache_dict) # pyright: ignore[reportUnknownVariableType] # cache_dict is untyped legacy state + def _cache_keys(self) -> tuple[object, ...]: + return tuple(self.cache_dict) - def _remove_key(self, key: str) -> None: + def _remove_key(self, key: object) -> None: """ Remove a key from both cache_dict and ttl_dict """ self.cache_dict.pop(key, None) self.ttl_dict.pop(key, None) - def evict_cache(self): + def evict_cache(self, *, reserve_space: bool = True) -> None: """ Eviction policy: 1. First, remove expired items from ttl_dict and cache_dict @@ -145,7 +145,7 @@ class InMemoryCache(BaseCache): heapq.heapify(self.expiration_heap) # Step 2: Evict if cache is still full - while self.expiration_heap and len(self.cache_dict) >= self.max_size_in_memory: + while self.expiration_heap and len(self.cache_dict) > self.max_size_in_memory - int(reserve_space): expiration_time, key = heapq.heappop(self.expiration_heap) # Skip if key was removed or updated if self.ttl_dict.get(key) == expiration_time: @@ -153,7 +153,7 @@ class InMemoryCache(BaseCache): if not self.expiration_heap: untracked_keys: Final = self._cache_keys() - overflow: Final = len(untracked_keys) - self.max_size_in_memory + 1 + overflow: Final = len(untracked_keys) - self.max_size_in_memory + int(reserve_space) for untracked_key in untracked_keys[: max(overflow, 0)]: self._remove_key(untracked_key) @@ -181,10 +181,13 @@ class InMemoryCache(BaseCache): # Always prune expired/outdated heap roots before inserting. # This keeps expiration_heap bounded even when the live cache stays # below max_size_in_memory and keys are reinserted after TTL expiry. - self.evict_cache() + self.evict_cache(reserve_space=key not in self.cache_dict) if not self.check_value_size(value): return + if key not in self.cache_dict and len(self.cache_dict) >= self.max_size_in_memory: + self.evict_cache() + self.cache_dict[key] = value if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl if "ttl" in kwargs and kwargs["ttl"] is not None: diff --git a/litellm/caching/llm_caching_handler.py b/litellm/caching/llm_caching_handler.py index a89e43b78b4..e7e7d9e6a90 100644 --- a/litellm/caching/llm_caching_handler.py +++ b/litellm/caching/llm_caching_handler.py @@ -37,7 +37,7 @@ class LLMClientCache(InMemoryCache): ) self.evicted_client_closer = evicted_client_closer or default_evicted_client_closer - def _remove_key(self, key: str) -> None: + def _remove_key(self, key: object) -> None: evicted: Final[object] = self.cache_dict.get(key) super()._remove_key(key) self.evicted_client_closer.schedule(evicted) diff --git a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py b/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py index 9361fc8dec4..f7a321e24eb 100644 --- a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py +++ b/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py @@ -27,7 +27,7 @@ class LangfuseInMemoryCache(InMemoryCache): (https://github.com/BerriAI/litellm/issues/11169). """ - def _remove_key(self, key: str) -> None: + def _remove_key(self, key: object) -> None: from litellm.integrations.langfuse.langfuse import LangFuseLogger evicted: Final = self.cache_dict.pop(key, None) diff --git a/litellm/proxy/agent_endpoints/databricks_oauth.py b/litellm/proxy/agent_endpoints/databricks_oauth.py index 7425c240d10..ac4fc021660 100644 --- a/litellm/proxy/agent_endpoints/databricks_oauth.py +++ b/litellm/proxy/agent_endpoints/databricks_oauth.py @@ -212,12 +212,12 @@ class DatabricksAppOAuthTokenCache(InMemoryCache): def __init__(self) -> None: super().__init__(default_ttl=_DEFAULT_TTL_SECONDS) - self._locks: dict[str, asyncio.Lock] = {} + self._locks: dict[object, asyncio.Lock] = {} def _get_lock(self, cache_key: str) -> asyncio.Lock: return self._locks.setdefault(cache_key, asyncio.Lock()) - def _remove_key(self, key: str) -> None: + def _remove_key(self, key: object) -> None: # Drop the per-key lock alongside the cached token so ``_locks`` stays # bounded by the live key set rather than growing for every key ever seen. super()._remove_key(key) diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 33633150711..e10598673c1 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -130,7 +130,7 @@ if os.environ.get("LITELLM_PROFILE", "false").lower() == "true": return {"top_50_memory_usage": result} -def _prompt_caching_pins(llm_router: "Router | None") -> tuple[Mapping[str, object], Mapping[str, float]]: +def _prompt_caching_pins(llm_router: "Router | None") -> tuple[Mapping[object, object], Mapping[str, float]]: if llm_router is None: return {}, {} pin_cache: Final = llm_router.prompt_caching_pin_cache.in_memory_cache diff --git a/tests/unit/caching/test_in_memory_cache.py b/tests/unit/caching/test_in_memory_cache.py index 75232476ecd..0cda0bf7fa2 100644 --- a/tests/unit/caching/test_in_memory_cache.py +++ b/tests/unit/caching/test_in_memory_cache.py @@ -1,21 +1,14 @@ -import asyncio -import json import random import sys import threading import time from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor -from typing import Final -from unittest.mock import MagicMock, patch +from typing import Final, Literal -import httpx import pytest -import respx -from fastapi.testclient import TestClient - -from unittest.mock import AsyncMock +from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache @@ -25,6 +18,17 @@ class _SlowInt(int): return _SlowInt(int(self) + value) +class _FakeClock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + def test_increment_cache_is_atomic_under_thread_concurrency(): cache = InMemoryCache() seed = 1000 @@ -139,6 +143,7 @@ def test_in_memory_cache_ttl(): assert new_ttl_time is not None time.sleep(1) cached_obj = in_memory_cache.get_cache(key="new-fake-key") + assert cached_obj is None new_ttl_time = in_memory_cache.ttl_dict.get("new-fake-key") assert new_ttl_time is None @@ -195,6 +200,48 @@ def test_in_memory_cache_max_size_with_ttl(): assert "key_3" in in_memory_cache.cache_dict +@pytest.mark.parametrize("updated_key", ["first", "second"]) +def test_updating_full_cache_preserves_entries_and_expiration(updated_key: Literal["first", "second"]) -> None: + clock: Final = _FakeClock() + cache: Final = InMemoryCache(max_size_in_memory=2, clock=clock) + cache.set_cache("first", "original-a", ttl=10) + cache.set_cache("second", "original-b", ttl=30) + + cache.set_cache(updated_key, "updated", ttl=60) + + assert cache.get_cache("first") == ("updated" if updated_key == "first" else "original-a") + assert cache.get_cache("second") == ("updated" if updated_key == "second" else "original-b") + + clock.advance(11.0) + assert cache.get_cache("first") is None + assert cache.get_cache("second") == ("updated" if updated_key == "second" else "original-b") + + clock.advance(20.0) + assert cache.get_cache("second") is None + + +@pytest.mark.parametrize( + ("updated_key", "expected_values"), + [ + ("first", ("updated", None, "original-c")), + ("third", (None, "original-b", "updated")), + ], +) +def test_updating_cache_after_capacity_reduction_respects_limit( + updated_key: Literal["first", "third"], expected_values: tuple[str | None, str | None, str | None] +) -> None: + cache: Final = InMemoryCache(max_size_in_memory=3, clock=_FakeClock()) + cache.set_cache("first", "original-a", ttl=10) + cache.set_cache("second", "original-b", ttl=20) + cache.set_cache("third", "original-c", ttl=30) + DualCache(in_memory_cache=cache).update_in_memory_max_size(2) + + cache.set_cache(updated_key, "updated", ttl=60) + + assert len(cache.cache_dict) == 2 + assert tuple(cache.get_cache(key) for key in ("first", "second", "third")) == expected_values + + def test_in_memory_cache_expired_items_evicted_first(): """ Test that expired items are evicted before non-expired items when cache is full. @@ -234,14 +281,9 @@ def test_in_memory_cache_eviction_order(): in_memory_cache = InMemoryCache(max_size_in_memory=2) # Add items with different TTLs - now = time.time() - in_memory_cache.set_cache( - key="early_expire", value="value_1", ttl=100 - ) # expires in 100 seconds + in_memory_cache.set_cache(key="early_expire", value="value_1", ttl=100) # expires in 100 seconds time.sleep(0.01) - in_memory_cache.set_cache( - key="late_expire", value="value_2", ttl=200 - ) # expires in 200 seconds + in_memory_cache.set_cache(key="late_expire", value="value_2", ttl=200) # expires in 200 seconds # Verify TTL order early_ttl = in_memory_cache.ttl_dict["early_expire"] @@ -337,6 +379,17 @@ def test_in_memory_cache_stays_bounded_when_entries_have_no_usable_expiry() -> N assert cache.get_cache("fresh") == "response" +def test_updating_full_cache_with_untracked_expiry_preserves_other_entries() -> None: + cache: Final = InMemoryCache(max_size_in_memory=2, clock=_FakeClock()) + cache.set_cache("first", "original-a", ttl="NaN") + cache.set_cache("second", "original-b", ttl="NaN") + + cache.set_cache("second", "updated", ttl=60) + + assert cache.get_cache("first") == "original-a" + assert cache.get_cache("second") == "updated" + + def test_in_memory_cache_prunes_expired_heap_entries_below_capacity(): """ Re-inserting expired keys below capacity should not grow expiration_heap