import random import sys import threading import time from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor from typing import Final, Literal import pytest from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache class _SlowInt(int): def __add__(self, value: int) -> "_SlowInt": time.sleep(0.05) 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 cache.set_cache("counter", _SlowInt(seed)) thread_count = 8 barrier = threading.Barrier(thread_count) def increment(_: int) -> float: barrier.wait() return cache.increment_cache("counter", 1) with ThreadPoolExecutor(max_workers=thread_count) as executor: tuple(executor.map(increment, range(thread_count))) assert cache.get_cache("counter") == seed + thread_count async def test_async_increment_delegates_to_locked_sync_path(): cache = InMemoryCache() assert await cache.async_increment("counter", 2) == 2 assert await cache.async_increment("counter", 3) == 5 assert cache.get_cache("counter") == 5 class _AnotherReaderExpiresTheEntryFirstCache(InMemoryCache): other_reader_started = False other_reader_result: object = "not run" def evict_element_if_expired(self, key: str) -> bool: if not self.other_reader_started: self.other_reader_started = True self.other_reader_result = self.get_cache(key) return super().evict_element_if_expired(key) def test_get_cache_reads_an_entry_another_reader_just_expired_as_a_miss() -> None: cache: Final = _AnotherReaderExpiresTheEntryFirstCache() cache.set_cache("tokenizer_failure", "HF unavailable", ttl=-1) assert cache.get_cache("tokenizer_failure") is None assert cache.other_reader_result is None class _ClockThatLetsAnotherReaderRunFirst: cache: InMemoryCache | None = None other_reader_result: object = "not run" def __call__(self) -> float: cache: Final = self.cache if cache is not None: self.cache = None self.other_reader_result = cache.get_cache("tokenizer_failure") return 100.0 def test_get_cache_reads_an_entry_another_reader_expires_during_the_expiry_check_as_a_miss() -> None: clock: Final = _ClockThatLetsAnotherReaderRunFirst() cache: Final = InMemoryCache(clock=clock) cache.set_cache("tokenizer_failure", "HF unavailable", ttl=-1) clock.cache = cache assert cache.get_cache("tokenizer_failure") is None assert clock.other_reader_result is None def test_in_memory_openai_obj_cache(): from openai import OpenAI openai_obj = OpenAI(api_key="my-fake-key") in_memory_cache = InMemoryCache() in_memory_cache.set_cache(key="my-fake-key", value=openai_obj) cached_obj = in_memory_cache.get_cache(key="my-fake-key") assert cached_obj is not None assert cached_obj == openai_obj def test_in_memory_cache_max_size_per_item(): """ Test that the cache will not store items larger than the max size per item """ in_memory_cache = InMemoryCache(max_size_per_item=100) result = in_memory_cache.check_value_size("a" * 100000000) assert result is False def test_in_memory_cache_ttl(): """ Check that - if ttl is not set, it will be set to default ttl - if object expires, the ttl is also removed """ in_memory_cache = InMemoryCache() in_memory_cache.set_cache(key="my-fake-key", value="my-fake-value", ttl=10) initial_ttl_time = in_memory_cache.ttl_dict["my-fake-key"] assert initial_ttl_time is not None in_memory_cache.set_cache(key="my-fake-key", value="my-fake-value-2", ttl=10) new_ttl_time = in_memory_cache.ttl_dict["my-fake-key"] assert new_ttl_time == initial_ttl_time # ttl should not be updated ## On object expiration, the ttl should be removed in_memory_cache.set_cache(key="new-fake-key", value="new-fake-value", ttl=1) new_ttl_time = in_memory_cache.ttl_dict["new-fake-key"] 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 def test_in_memory_cache_ttl_allow_override(): """ Check that - if ttl is not set, it will be set to default ttl - if object expires, the ttl is also removed """ in_memory_cache = InMemoryCache() ## On object expiration, but no get_cache, the override should be allowed in_memory_cache.set_cache(key="new-fake-key", value="new-fake-value", ttl=1) initial_ttl_time = in_memory_cache.ttl_dict["new-fake-key"] assert initial_ttl_time is not None time.sleep(1) in_memory_cache.set_cache(key="new-fake-key", value="new-fake-value-2", ttl=1) new_ttl_time = in_memory_cache.ttl_dict["new-fake-key"] assert new_ttl_time is not None assert new_ttl_time != initial_ttl_time def test_in_memory_cache_max_size_with_ttl(): """ Test that max_size_in_memory is respected even when all items have long TTLs. This tests the fix for the unbounded growth issue. """ in_memory_cache = InMemoryCache(max_size_in_memory=3) long_ttl = 86400 # 1 day # Fill the cache to max capacity for i in range(3): in_memory_cache.set_cache(key=f"key_{i}", value=f"value_{i}", ttl=long_ttl) time.sleep(0.01) # Small delay to ensure different timestamps assert len(in_memory_cache.cache_dict) == 3 assert len(in_memory_cache.ttl_dict) == 3 # Add another item - should evict the earliest item in_memory_cache.set_cache(key="key_3", value="value_3", ttl=long_ttl) # Cache should still be at max size, not larger assert len(in_memory_cache.cache_dict) == 3 assert len(in_memory_cache.ttl_dict) == 3 # key_0 should have been evicted (it was added first) assert "key_0" not in in_memory_cache.cache_dict assert "key_0" not in in_memory_cache.ttl_dict # Other keys should still be present assert "key_1" in in_memory_cache.cache_dict assert "key_2" in in_memory_cache.cache_dict 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. """ in_memory_cache = InMemoryCache(max_size_in_memory=3) # Add items with short TTL that will expire in_memory_cache.set_cache(key="expired_1", value="value_1", ttl=1) in_memory_cache.set_cache(key="expired_2", value="value_2", ttl=1) # Add item with long TTL in_memory_cache.set_cache(key="long_lived", value="value_long", ttl=86400) assert len(in_memory_cache.cache_dict) == 3 # Wait for short TTL items to expire time.sleep(2) # Add new item - should evict expired items first, not the long-lived one in_memory_cache.set_cache(key="new_item", value="new_value", ttl=86400) # Long-lived item should still be present assert "long_lived" in in_memory_cache.cache_dict assert "new_item" in in_memory_cache.cache_dict # Expired items should be gone assert "expired_1" not in in_memory_cache.cache_dict assert "expired_2" not in in_memory_cache.cache_dict assert "expired_1" not in in_memory_cache.ttl_dict assert "expired_2" not in in_memory_cache.ttl_dict def test_in_memory_cache_eviction_order(): """ Test that when non-expired items need to be evicted, those with earliest expiration times are evicted first. """ in_memory_cache = InMemoryCache(max_size_in_memory=2) # Add items with different TTLs 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 # Verify TTL order early_ttl = in_memory_cache.ttl_dict["early_expire"] late_ttl = in_memory_cache.ttl_dict["late_expire"] assert early_ttl < late_ttl, "early_expire should have earlier expiration time" assert len(in_memory_cache.cache_dict) == 2 # Add third item - should evict the one with earliest expiration time in_memory_cache.set_cache(key="new_item", value="value_3", ttl=300) assert len(in_memory_cache.cache_dict) == 2 # Item with earliest expiration should be evicted assert "early_expire" not in in_memory_cache.cache_dict assert "early_expire" not in in_memory_cache.ttl_dict # Items with later expiration should remain assert "late_expire" in in_memory_cache.cache_dict assert "new_item" in in_memory_cache.cache_dict def test_in_memory_cache_heap_size_staus_bounded(): """ Test that the expiration_heap does not grow unbounded when the same key is updated repeaatedly. """ in_memory_cache = InMemoryCache(max_size_in_memory=10) for i in range(1_000): in_memory_cache.set_cache(key="hot_key", value=f"value_{i}", ttl=60) # Expiration heap should only have 1 entry assert len(in_memory_cache.expiration_heap) == 1 def test_in_memory_cache_heap_stays_bounded_when_a_live_key_is_rewritten_with_a_new_ttl() -> None: in_memory_cache = InMemoryCache(max_size_in_memory=10) in_memory_cache.set_cache(key="expires_first", value="pin", ttl=300) for i in range(1_000): in_memory_cache.delete_cache("hot_key") in_memory_cache.set_cache(key="hot_key", value=f"value_{i}", ttl=3600) assert in_memory_cache.get_cache("hot_key") == "value_999" assert len(in_memory_cache.expiration_heap) <= 2 * in_memory_cache.max_size_in_memory + 1 @pytest.fixture def fine_grained_thread_switching() -> Iterator[None]: previous_interval: Final = sys.getswitchinterval() sys.setswitchinterval(1e-6) yield sys.setswitchinterval(previous_interval) def _keys_until(deadline: float, rng: random.Random) -> Iterator[str]: while time.monotonic() < deadline: yield f"hot_{rng.randrange(5)}" if rng.random() < 0.7 else f"cold_{rng.randrange(40)}" def _rewrite_keys_until(cache: InMemoryCache, deadline: float, seed: int) -> None: rng: Final = random.Random(seed) for key in _keys_until(deadline, rng): cache.delete_cache(key) cache.set_cache(key, "pin", ttl=rng.choice((300, 3600))) @pytest.mark.usefixtures("fine_grained_thread_switching") def test_in_memory_cache_survives_threads_rewriting_keys_in_a_full_cache() -> None: cache: Final = InMemoryCache(max_size_in_memory=5) deadline: Final = time.monotonic() + 0.5 thread_count: Final = 8 with ThreadPoolExecutor(max_workers=thread_count) as executor: futures: Final = tuple( executor.submit(_rewrite_keys_until, cache, deadline, seed) for seed in range(thread_count) ) assert [future.exception() for future in futures] == [None] * thread_count keys: Final = (*(f"hot_{i}" for i in range(5)), *(f"cold_{i}" for i in range(40))) assert sum(cache.get_cache(key) is not None for key in keys) <= cache.max_size_in_memory def test_in_memory_cache_stays_bounded_when_entries_have_no_usable_expiry() -> None: cache: Final = InMemoryCache(max_size_in_memory=10) for i in range(30): cache.set_cache(key=f"nan_{i}", value="response", ttl="NaN") cache.set_cache(key="fresh", value="response", ttl=60) keys: Final = (*(f"nan_{i}" for i in range(30)), "fresh") retained: Final = tuple(key for key in keys if cache.get_cache(key) is not None) assert len(retained) <= cache.max_size_in_memory 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 without bound. """ in_memory_cache = InMemoryCache(max_size_in_memory=200, default_ttl=1) for cycle in range(3): for i in range(5): in_memory_cache.set_cache(key=f"key_{i}", value=f"value_{cycle}_{i}", ttl=1) time.sleep(1.1) for i in range(5): in_memory_cache.set_cache(key=f"key_{i}", value=f"value_final_{i}", ttl=1) assert len(in_memory_cache.cache_dict) == 5 assert len(in_memory_cache.ttl_dict) == 5 assert len(in_memory_cache.expiration_heap) == 5 def test_in_memory_cache_injected_clock_controls_expiry_and_eviction() -> None: class Clock: now = 0.0 def __call__(self) -> float: return self.now clock = Clock() cache = InMemoryCache(max_size_in_memory=2, default_ttl=60, clock=clock) cache.set_cache("first", "original", ttl=10) clock.now = 9.0 cache.set_cache("second", "survivor") assert cache.get_cache("first") == "original" clock.now = 10.001 assert cache.get_cache("first") is None cache.set_cache("third", "replacement") assert cache.get_cache("second") == "survivor" clock.now = 69.001 cache.set_cache("fourth", "new") assert cache.get_cache("second") is None assert cache.get_cache("third") == "replacement" assert cache.get_cache("fourth") == "new"