diff --git a/litellm/caching.py b/litellm/caching.py index fcebeb91da..2c2a3f8bc1 100644 --- a/litellm/caching.py +++ b/litellm/caching.py @@ -38,6 +38,9 @@ class BaseCache: async def async_get_cache(self, key, **kwargs): raise NotImplementedError + async def batch_cache_write(self, result, *args, **kwargs): + raise NotImplementedError + async def disconnect(self): raise NotImplementedError @@ -96,7 +99,9 @@ class InMemoryCache(BaseCache): class RedisCache(BaseCache): # if users don't provider one, use the default litellm cache - def __init__(self, host=None, port=None, password=None, **kwargs): + def __init__( + self, host=None, port=None, password=None, redis_flush_size=100, **kwargs + ): from ._redis import get_redis_client, get_redis_connection_pool redis_kwargs = {} @@ -111,6 +116,10 @@ class RedisCache(BaseCache): self.redis_client = get_redis_client(**redis_kwargs) self.redis_kwargs = redis_kwargs self.async_redis_conn_pool = get_redis_connection_pool(**redis_kwargs) + + # for high traffic, we store the redis results in memory and then batch write to redis + self.redis_batch_writing_buffer = [] + self.redis_flush_size = redis_flush_size self.redis_version = "Unknown" try: self.redis_version = self.redis_client.info()["redis_version"] @@ -193,6 +202,21 @@ class RedisCache(BaseCache): except Exception as e: print_verbose(f"Error occurred in pipeline write - {str(e)}") + async def batch_cache_write(self, key, value, **kwargs): + print_verbose( + f"in batch cache writing for redis buffer size={len(self.redis_batch_writing_buffer)}", + ) + self.redis_batch_writing_buffer.append((key, value)) + if len(self.redis_batch_writing_buffer) >= self.redis_flush_size: + await self.flush_cache_buffer() + + async def flush_cache_buffer(self): + print_verbose( + f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}" + ) + await self.async_set_cache_pipeline(self.redis_batch_writing_buffer) + self.redis_batch_writing_buffer = [] + def _get_cache_logic(self, cached_response: Any): """ Common 'get_cache_logic' across sync + async redis client implementations @@ -909,6 +933,7 @@ class Cache: s3_path: Optional[str] = None, redis_semantic_cache_use_async=False, redis_semantic_cache_embedding_model="text-embedding-ada-002", + redis_flush_size=None, **kwargs, ): """ @@ -931,7 +956,9 @@ class Cache: None. Cache is set as a litellm param """ if type == "redis": - self.cache: BaseCache = RedisCache(host, port, password, **kwargs) + self.cache: BaseCache = RedisCache( + host, port, password, redis_flush_size, **kwargs + ) elif type == "redis-semantic": self.cache = RedisSemanticCache( host, @@ -968,6 +995,7 @@ class Cache: self.supported_call_types = supported_call_types # default to ["completion", "acompletion", "embedding", "aembedding"] self.type = type self.namespace = namespace + self.redis_flush_size = redis_flush_size self.ttl = ttl def get_cache_key(self, *args, **kwargs): @@ -1252,10 +1280,14 @@ class Cache: Async implementation of add_cache """ try: - cache_key, cached_data, kwargs = self._add_cache_logic( - result=result, *args, **kwargs - ) - await self.cache.async_set_cache(cache_key, cached_data, **kwargs) + if self.type == "redis" and self.redis_flush_size is not None: + # high traffic - fill in results in memory and then flush + await self.batch_cache_write(result, *args, **kwargs) + else: + cache_key, cached_data, kwargs = self._add_cache_logic( + result=result, *args, **kwargs + ) + await self.cache.async_set_cache(cache_key, cached_data, **kwargs) except Exception as e: print_verbose(f"LiteLLM Cache: Excepton add_cache: {str(e)}") traceback.print_exc() @@ -1293,6 +1325,12 @@ class Cache: print_verbose(f"LiteLLM Cache: Excepton add_cache: {str(e)}") traceback.print_exc() + async def batch_cache_write(self, result, *args, **kwargs): + cache_key, cached_data, kwargs = self._add_cache_logic( + result=result, *args, **kwargs + ) + await self.cache.batch_cache_write(cache_key, cached_data, **kwargs) + async def ping(self): if hasattr(self.cache, "ping"): return await self.cache.ping() diff --git a/litellm/proxy/proxy_load_test/locustfile.py b/litellm/proxy/proxy_load_test/locustfile.py index ebfca80ad3..5778cb916f 100644 --- a/litellm/proxy/proxy_load_test/locustfile.py +++ b/litellm/proxy/proxy_load_test/locustfile.py @@ -1,6 +1,7 @@ from locust import HttpUser, task, between, events import json import time +import uuid class MyUser(HttpUser): @@ -20,7 +21,8 @@ class MyUser(HttpUser): "messages": [ { "role": "system", - "content": "this is a very sweet test message from ishaan", + "content": f"{uuid.uuid4()} this is a very sweet test message from ishaan" + * 100, }, {"role": "user", "content": "Hello, how are you?"}, ], diff --git a/litellm/tests/test_caching.py b/litellm/tests/test_caching.py index 4bc5d1d5ae..21ddf1659d 100644 --- a/litellm/tests/test_caching.py +++ b/litellm/tests/test_caching.py @@ -386,6 +386,48 @@ async def test_redis_cache_basic(): assert stored_val["id"] == response1.id +@pytest.mark.asyncio +async def test_redis_batch_cache_write(): + """ + Init redis client + - write to client + - read from client + """ + litellm.set_verbose = True + import uuid + + messages = [ + {"role": "user", "content": f"write a one sentence poem about: {uuid.uuid4()}"}, + ] + litellm.cache = Cache( + type="redis", + host=os.environ["REDIS_HOST"], + port=os.environ["REDIS_PORT"], + password=os.environ["REDIS_PASSWORD"], + redis_flush_size=2, + ) + response1 = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=messages, + ) + + response2 = await litellm.acompletion( + model="anthropic/claude-3-opus-20240229", + messages=messages, + mock_response="good morning from this test", + ) + + # we hit the flush size, this will now send to redis + await asyncio.sleep(2) + + response4 = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=messages, + ) + + assert response1.id == response4.id + + def test_redis_cache_completion(): litellm.set_verbose = False