(Redis fix) - use mget_non_atomic (#8682)

* use mget_nonatomic

* redis cluster override MGET op

* fix redis cluster + MGET

* test redis cluster

litellm/caching/redis_cache.py

@@ -637,6 +637,23 @@ class RedisCache(BaseCache):
                 "litellm.caching.caching: get() - Got exception from REDIS: ", e
             )
 
+    def _run_redis_mget_operation(self, keys: List[str]) -> List[Any]:
+        """
+        Wrapper to call `mget` on the redis client
+
+        We use a wrapper so RedisCluster can override this method
+        """
+        return self.redis_client.mget(keys=keys)  # type: ignore
+
+    async def _async_run_redis_mget_operation(self, keys: List[str]) -> List[Any]:
+        """
+        Wrapper to call `mget` on the redis client
+
+        We use a wrapper so RedisCluster can override this method
+        """
+        async_redis_client = self.init_async_client()
+        return await async_redis_client.mget(keys=keys)  # type: ignore
+
     def batch_get_cache(
         self,
         key_list: Union[List[str], List[Optional[str]]],
@@ -661,7 +678,7 @@ class RedisCache(BaseCache):
                 cache_key = self.check_and_fix_namespace(key=cache_key or "")
                 _keys.append(cache_key)
             start_time = time.time()
-            results: List = self.redis_client.mget(keys=_keys)  # type: ignore
+            results: List = self._run_redis_mget_operation(keys=_keys)
             end_time = time.time()
             _duration = end_time - start_time
             self.service_logger_obj.service_success_hook(
@@ -757,7 +774,6 @@ class RedisCache(BaseCache):
         `.mget` does not support None keys. This will filter out None keys.
         """
         # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `mget`
-        _redis_client: Any = self.init_async_client()
         key_value_dict = {}
         start_time = time.time()
         _key_list = [key for key in key_list if key is not None]
@@ -766,7 +782,7 @@ class RedisCache(BaseCache):
             for cache_key in _key_list:
                 cache_key = self.check_and_fix_namespace(key=cache_key)
                 _keys.append(cache_key)
-            results = await _redis_client.mget(keys=_keys)
+            results = await self._async_run_redis_mget_operation(keys=_keys)
             ## LOGGING ##
             end_time = time.time()
             _duration = end_time - start_time

litellm/caching/redis_cluster_cache.py

@@ -5,7 +5,7 @@ Key differences:
 - RedisClient NEEDs to be re-used across requests, adds 3000ms latency if it's re-created
 """
 
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, List, Optional
 
 from litellm.caching.redis_cache import RedisCache
 
@@ -26,19 +26,34 @@ else:
 class RedisClusterCache(RedisCache):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.redis_cluster_client: Optional[RedisCluster] = None
+        self.redis_async_redis_cluster_client: Optional[RedisCluster] = None
+        self.redis_sync_redis_cluster_client: Optional[RedisCluster] = None
 
     def init_async_client(self):
         from redis.asyncio import RedisCluster
 
         from .._redis import get_redis_async_client
 
-        if self.redis_cluster_client:
-            return self.redis_cluster_client
+        if self.redis_async_redis_cluster_client:
+            return self.redis_async_redis_cluster_client
 
         _redis_client = get_redis_async_client(
             connection_pool=self.async_redis_conn_pool, **self.redis_kwargs
         )
         if isinstance(_redis_client, RedisCluster):
-            self.redis_cluster_client = _redis_client
+            self.redis_async_redis_cluster_client = _redis_client
+
         return _redis_client
+
+    def _run_redis_mget_operation(self, keys: List[str]) -> List[Any]:
+        """
+        Overrides `_run_redis_mget_operation` in redis_cache.py
+        """
+        return self.redis_client.mget_nonatomic(keys=keys)  # type: ignore
+
+    async def _async_run_redis_mget_operation(self, keys: List[str]) -> List[Any]:
+        """
+        Overrides `_async_run_redis_mget_operation` in redis_cache.py
+        """
+        async_redis_cluster_client = self.init_async_client()
+        return await async_redis_cluster_client.mget_nonatomic(keys=keys)  # type: ignore