mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
432954a2ab
* fix(cache): make in-memory and disk increments atomic * refactor(cache): narrow in-memory increment lock scope * fix(cache): address follow-up review on increment tests/types * fix(cache): refresh atomic increment coverage * test(cache): widen increment race window with non-zero _SlowInt seed The zero seed was falsy, so InMemoryCache.increment_cache's `get_cache(...) or 0` and DiskCache.get_cache's truthiness guard both discarded the _SlowInt before __add__ could run, leaving the sleep-based window-widening inert. Seed a non-zero value and return _SlowInt from __add__ so the sleep fires on every read-modify-write in both backends, making the concurrency regression deterministic. * test(cache): cover InMemoryCache.async_increment delegation Add a focused async test asserting async_increment accumulates through the locked sync path, exercising the previously uncovered delegation line. --------- Co-authored-by: Emerson Gomes <emerson.gomes@thalesgroup.com>
67 lines
1.9 KiB
Python
67 lines
1.9 KiB
Python
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("diskcache")
|
|
|
|
from litellm.caching.disk_cache import DiskCache
|
|
|
|
|
|
class _SlowInt(int):
|
|
def __add__(self, value: int) -> "_SlowInt":
|
|
time.sleep(0.05)
|
|
return _SlowInt(int(self) + value)
|
|
|
|
|
|
@pytest.fixture
|
|
def cache(tmp_path):
|
|
return DiskCache(disk_cache_dir=str(tmp_path))
|
|
|
|
|
|
def test_increment_cache_starts_from_zero_when_key_missing(cache):
|
|
assert cache.increment_cache("counter", 3) == 3
|
|
assert cache.get_cache("counter") == 3
|
|
|
|
|
|
def test_increment_cache_adds_to_existing_int(cache):
|
|
cache.set_cache("counter", 7)
|
|
assert cache.increment_cache("counter", 5) == 12
|
|
assert cache.get_cache("counter") == 12
|
|
|
|
|
|
def test_increment_cache_treats_non_int_cached_value_as_zero(cache):
|
|
cache.set_cache("counter", "not-a-number")
|
|
assert cache.increment_cache("counter", 4) == 4
|
|
assert cache.get_cache("counter") == 4
|
|
|
|
|
|
def test_increment_cache_is_atomic_under_thread_concurrency(cache):
|
|
seed = 1000
|
|
cache.set_cache("counter", _SlowInt(seed))
|
|
thread_count = 8
|
|
barrier = threading.Barrier(thread_count)
|
|
|
|
def increment(_: int) -> int:
|
|
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_starts_from_zero_when_key_missing(cache):
|
|
assert await cache.async_increment("counter", 2) == 2
|
|
|
|
|
|
async def test_async_increment_adds_to_existing_int(cache):
|
|
await cache.async_set_cache("counter", 10)
|
|
assert await cache.async_increment("counter", 5) == 15
|
|
|
|
|
|
async def test_async_increment_treats_non_int_cached_value_as_zero(cache):
|
|
await cache.async_set_cache("counter", "corrupt")
|
|
assert await cache.async_increment("counter", 9) == 9
|