mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +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>
88 lines
2.9 KiB
Python
88 lines
2.9 KiB
Python
import json
|
|
from typing import TYPE_CHECKING, Any, Optional, Union
|
|
|
|
from .base_cache import BaseCache
|
|
|
|
if TYPE_CHECKING:
|
|
from opentelemetry.trace import Span as _Span
|
|
|
|
Span = Union[_Span, Any]
|
|
else:
|
|
Span = Any
|
|
|
|
|
|
class DiskCache(BaseCache):
|
|
def __init__(self, disk_cache_dir: Optional[str] = None):
|
|
try:
|
|
import diskcache as dc
|
|
except ModuleNotFoundError as e:
|
|
raise ModuleNotFoundError("Please install litellm with `litellm[caching]` to use disk caching.") from e
|
|
|
|
# if users don't provider one, use the default litellm cache
|
|
if disk_cache_dir is None:
|
|
self.disk_cache = dc.Cache(".litellm_cache")
|
|
else:
|
|
self.disk_cache = dc.Cache(disk_cache_dir)
|
|
|
|
def set_cache(self, key, value, **kwargs):
|
|
if "ttl" in kwargs:
|
|
self.disk_cache.set(key, value, expire=kwargs["ttl"])
|
|
else:
|
|
self.disk_cache.set(key, value)
|
|
|
|
async def async_set_cache(self, key, value, **kwargs):
|
|
self.set_cache(key=key, value=value, **kwargs)
|
|
|
|
async def async_set_cache_pipeline(self, cache_list, **kwargs):
|
|
for cache_key, cache_value in cache_list:
|
|
if "ttl" in kwargs:
|
|
self.set_cache(key=cache_key, value=cache_value, ttl=kwargs["ttl"])
|
|
else:
|
|
self.set_cache(key=cache_key, value=cache_value)
|
|
|
|
def get_cache(self, key, **kwargs):
|
|
original_cached_response = self.disk_cache.get(key)
|
|
if original_cached_response:
|
|
try:
|
|
cached_response = json.loads(original_cached_response) # type: ignore
|
|
except Exception:
|
|
cached_response = original_cached_response
|
|
return cached_response
|
|
return None
|
|
|
|
def batch_get_cache(self, keys: list, **kwargs):
|
|
return_val = []
|
|
for k in keys:
|
|
val = self.get_cache(key=k, **kwargs)
|
|
return_val.append(val)
|
|
return return_val
|
|
|
|
def increment_cache(self, key, value: int, **kwargs) -> int:
|
|
with self.disk_cache.transact():
|
|
cached_value = self.get_cache(key=key)
|
|
init_value = cached_value if isinstance(cached_value, int) else 0
|
|
new_value = init_value + value
|
|
self.set_cache(key, new_value, **kwargs)
|
|
return new_value
|
|
|
|
async def async_get_cache(self, key, **kwargs):
|
|
return self.get_cache(key=key, **kwargs)
|
|
|
|
async def async_batch_get_cache(self, keys: list, **kwargs):
|
|
return_val = []
|
|
for k in keys:
|
|
val = self.get_cache(key=k, **kwargs)
|
|
return_val.append(val)
|
|
return return_val
|
|
|
|
async def async_increment(self, key, value: int, **kwargs) -> int:
|
|
return self.increment_cache(key=key, value=value, **kwargs)
|
|
|
|
def flush_cache(self):
|
|
self.disk_cache.clear()
|
|
|
|
async def disconnect(self):
|
|
pass
|
|
|
|
def delete_cache(self, key):
|
|
self.disk_cache.pop(key)
|