Merge pull request #2693 from BerriAI/litellm_batch_write_redis_cache

[FEAT] batch write redis cache

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()

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?"},
             ],

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