Distributed cache feature using Redis (#518)

* [add] initial draft for distributed caching

Signed-off-by: Anurag Wagh <[email protected]>

* [mod] Simplify Model schema for distributed cache management.

Signed-off-by: Anurag Wagh <[email protected]>

* [mod] update expiry time for keys when accessed or saved

Signed-off-by: Anurag Wagh <[email protected]>

* [add] NoOpEviction mechanism to avoid redundant calls to redis

Signed-off-by: Anurag Wagh <[email protected]>

* [add] unit tests for distributed cache

Signed-off-by: Anurag Wagh <[email protected]>

* [add] temporary directory for sqlite

Signed-off-by: Anurag Wagh <[email protected]>

* [mod] skip adding scalar ids to cache if NoOpEviction is used

Signed-off-by: Anurag Wagh <[email protected]>

* [add] documentation, example code and refactoring

Signed-off-by: Anurag Wagh <[email protected]>

* [refactor] added docstrings, consistent quotes, removed literal comparison

Signed-off-by: Anurag Wagh <[email protected]>

* [refactor] removed unused imports

Signed-off-by: Anurag Wagh <[email protected]>

* [fix] add lazy import for redis in distributed_cache.py

Signed-off-by: Anurag Wagh <[email protected]>

* [add] `pylint: disable=wrong-import-position`

Signed-off-by: Anurag Wagh <[email protected]>

* [refactor] grouped import statements for distributed cache

Signed-off-by: Anurag Wagh <[email protected]>

* [refactor] revert `get_data_manager` and `manager_factory` signatures to include memory cache config

Signed-off-by: Anurag Wagh <[email protected]>

* [mod] add default values to memory cache params

Signed-off-by: Anurag Wagh <[email protected]>

* [mod] add param to set `mammemory-samples` config for redis cache

Signed-off-by: Anurag Wagh <[email protected]>

* [mod] updated cache config for testing lru cache

Signed-off-by: Anurag Wagh <[email protected]>

* [del] removed `test-lru-cache`

Signed-off-by: Anurag Wagh <[email protected]>

* [add] description for `maxmemory_samples` param

Signed-off-by: Anurag Wagh <[email protected]>

* [add] unit test for validating ttl configuration

Signed-off-by: Anurag Wagh <[email protected]>

* [add] unit tests for validating cache configuration, ttl access

Signed-off-by: Anurag Wagh <[email protected]>

* [add] unit test to validate str based access for eviction base in `get_data_manager`

Signed-off-by: Anurag Wagh <[email protected]>

---------

Signed-off-by: Anurag Wagh <[email protected]>

examples/eviction/distributed_eviction.py

@@ -0,0 +1,73 @@
+from gptcache import Cache
+from gptcache.embedding import Onnx
+
+from gptcache.manager.eviction import EvictionBase
+
+from gptcache.manager import get_data_manager, CacheBase, VectorBase, manager_factory
+
+
+def get_data_manager_example():
+    """
+    This example shows how to create a data manager with a mongo as a scalar storage, faiss vector base,
+    and redis eviction base.
+    This type of configuration can be used to scale GPTCache horizontally.
+    Where keys will be maintained in redis key-value store instead of in-memory.
+    The eviction of the keys will be handled based on the eviction policy of redis.
+    """
+    onnx = Onnx()
+    data_manager = get_data_manager(cache_base=CacheBase("mongo", url="mongodb://localhost:27017/"),
+                                    vector_base=VectorBase("faiss", dimension=onnx.dimension),
+                                    eviction_base=EvictionBase("redis",
+                                                               maxmemory="100mb",
+                                                               policy="allkeys-lru",
+                                                               ttl=100))
+
+    cache = Cache()
+    cache.init(data_manager=data_manager)
+    question = "What is github?"
+    answer = "Online platform for version control and code collaboration."
+    embedding = onnx.to_embeddings(question)
+    cache.import_data([question], [answer], [embedding])
+
+
+def get_manager_example_redis_only():
+    """
+    Note: Since, `RedisScalarStorage` can be configured to internally handle the ttl of the keys and their eviction.
+    In this scenario, `no_op_eviction` is used as the eviction base. It will not add any keys or update their ttls.
+
+    This example shows how to create a data manager with a redis as a scalar storage, as well as eviction base.
+    This type of configuration can be used to scale GPTCache horizontally.
+    Where keys will be maintained in redis key-value store instead of in-memory.
+    The eviction of the keys will be handled based on the eviction policy of redis.
+
+    """
+    onnx = Onnx()
+    data_manager = get_data_manager(cache_base=CacheBase("redis", maxmemory="100mb", policy="allkeys-lru", ttl=100),
+                                    vector_base=VectorBase("faiss", dimension=onnx.dimension),
+                                    eviction_base=EvictionBase("no_op_eviction"))
+
+    cache = Cache()
+    cache.init(data_manager=data_manager)
+    question = "What is github?"
+    answer = "Online platform for version control and code collaboration."
+    embedding = onnx.to_embeddings(question)
+    cache.import_data([question], [answer], [embedding])
+
+
+def manager_factory_example():
+    onnx = Onnx()
+    data_manager = manager_factory("redis,faiss",
+                                   eviction_manager="redis",
+                                   scalar_params={"url": "redis://localhost:6379"},
+                                   vector_params={"dimension": onnx.dimension},
+                                   eviction_params={"maxmemory": "100mb",
+                                                    "policy": "allkeys-lru",
+                                                    "ttl": 1}
+                                   )
+
+    cache = Cache()
+    cache.init(data_manager=data_manager)
+    question = "What is github?"
+    answer = "Online platform for version control and code collaboration."
+    embedding = onnx.to_embeddings(question)
+    cache.import_data([question], [answer], [embedding])

gptcache/manager/data_manager.py

@@ -7,6 +7,7 @@
 import requests
 
 from gptcache.manager.eviction import EvictionBase
+from gptcache.manager.eviction.distributed_cache import NoOpEviction
 from gptcache.manager.eviction_manager import EvictionManager
 from gptcache.manager.object_data.base import ObjectBase
 from gptcache.manager.scalar_data.base import (
@@ -30,11 +31,11 @@ def save(self, question, answer, embedding_data, **kwargs):
 
     @abstractmethod
     def import_data(
-        self,
-        questions: List[Any],
-        answers: List[Any],
-        embedding_datas: List[Any],
-        session_ids: List[Optional[str]],
+            self,
+            questions: List[Any],
+            answers: List[Any],
+            embedding_datas: List[Any],
+            session_ids: List[Optional[str]],
     ):
         pass
 
@@ -213,37 +214,37 @@ class SSDataManager(DataManager):
     :type s: CacheStorage
     :param v: VectorBase to manager the vector data, it can be generated with :meth:`gptcache.manager.VectorBase`.
     :type v:  VectorBase
-    :param max_size: the max size for the cache, defaults to 1000.
-    :type max_size: int
-    :param clean_size: the size to clean up, defaults to `max_size * 0.2`.
-    :type clean_size: int
-    :param eviction: The eviction policy, it is support "LRU" and "FIFO" now, and defaults to "LRU".
-    :type eviction:  str
+    :param o: ObjectBase to manager the object data, it can be generated with :meth:`gptcache.manager.ObjectBase`.
+    :type o:  ObjectBase
+    :param e: EvictionBase to manager the eviction data, it can be generated with :meth:`gptcache.manager.EvictionBase`.
+    :type e:  EvictionBase
     """
 
     def __init__(
         self,
         s: CacheStorage,
         v: VectorBase,
         o: Optional[ObjectBase],
+        e: Optional[EvictionBase],
         max_size,
         clean_size,
-        policy="LRU",
+        policy="LRU"
     ):
-        self.max_size = max_size
-        self.clean_size = clean_size
         self.s = s
         self.v = v
         self.o = o
         self.eviction_manager = EvictionManager(self.s, self.v)
-        self.eviction_base = EvictionBase(
-            name="memory",
-            policy=policy,
-            maxsize=max_size,
-            clean_size=clean_size,
-            on_evict=self._clear,
-        )
-        self.eviction_base.put(self.s.get_ids(deleted=False))
+        if e is None:
+            e = EvictionBase(name="memory",
+                             maxsize=max_size,
+                             clean_size=clean_size,
+                             policy=policy,
+                             on_evict=self._clear)
+        self.eviction_base = e
+
+        if not isinstance(self.eviction_base, NoOpEviction):
+            # if eviction manager is no op redis, we don't need to put data into eviction base
+            self.eviction_base.put(self.s.get_ids(deleted=False))
 
     def _clear(self, marked_keys):
         self.eviction_manager.soft_evict(marked_keys)

gptcache/manager/eviction/distributed_cache.py

@@ -0,0 +1,121 @@
+# pylint: disable=wrong-import-position
+from abc import ABC, abstractmethod
+from typing import List
+
+from gptcache.utils import import_redis
+from gptcache.manager.eviction.base import EvictionBase
+
+import_redis()
+import redis
+from redis_om import get_redis_connection
+
+
+class DistributedEviction(EvictionBase, ABC):
+    """
+    Base class for Distributed Eviction Strategy.
+    """
+
+    @abstractmethod
+    def put(self, objs: List[str]):
+        pass
+
+    @abstractmethod
+    def get(self, obj: str):
+        pass
+
+    @property
+    @abstractmethod
+    def policy(self) -> str:
+        pass
+
+
+class RedisCacheEviction(DistributedEviction, ABC):
+    """eviction: Distributed Cache Eviction Strategy using Redis.
+
+    :param host: the host of redis
+    :type host: str
+    :param port: the port of redis
+    :type port: int
+    :param policy: eviction strategy policy of redis such as allkeys-lru, volatile-lru, allkeys-random, volatile-random, etc.
+    refer https://redis.io/docs/reference/eviction/ for more information.
+    :type policy: str
+    :param maxsize: the maxsize of cache data
+    :type maxsize: int
+    :param on_evict: the function for cleaning the data in the store
+    :type  on_evict: Callable[[List[Any]], None]
+    :param maxmemory: the maxmemory of redis
+    :type maxmemory: str
+    :param global_key_prefix: the global key prefix
+    :type global_key_prefix: str
+    :param ttl: the ttl of the cache data
+    :type ttl: int
+    :param maxmemory_samples: Number of keys to sample when evicting keys
+    :type maxmemory_samples: int
+    :param kwargs: the kwargs
+    :type kwargs: Any
+    """
+
+    def __init__(self,
+                 host="localhost",
+                 port=6379,
+                 maxmemory: str = None,
+                 policy: str = None,
+                 global_key_prefix="gptcache",
+                 ttl: int = None,
+                 maxmemory_samples: int = None,
+                 **kwargs):
+        self._redis = get_redis_connection(host=host, port=port, **kwargs)
+        if maxmemory:
+            self._redis.config_set("maxmemory", maxmemory)
+        if maxmemory_samples:
+            self._redis.config_set("maxmemory-samples", maxmemory_samples)
+        if policy:
+            self._redis.config_set("maxmemory-policy", policy)
+            self._policy = policy.lower()
+
+        self._global_key_prefix = global_key_prefix
+        self._ttl = ttl
+
+    def _create_key(self, key: str) -> str:
+        return f"{self._global_key_prefix}:evict:{key}"
+
+    def put(self, objs: List[str], expire=False):
+        ttl = self._ttl if expire else None
+        for key in objs:
+            self._redis.set(self._create_key(key), "True", ex=ttl)
+
+    def get(self, obj: str):
+
+        try:
+            value = self._redis.get(self._create_key(obj))
+            # update key expire time when accessed
+            if self._ttl:
+                self._redis.expire(self._create_key(obj), self._ttl)
+            return value
+        except redis.RedisError:
+            print(f"Error getting key {obj} from cache")
+            return None
+
+    @property
+    def policy(self) -> str:
+        return self._policy
+
+
+class NoOpEviction(EvictionBase):
+    """eviction: No Op Eviction Strategy. This is used when Eviction is managed internally
+    by the Databases such as Redis or memcached and no eviction is required to perform.
+
+    """
+
+    @property
+    def policy(self) -> str:
+        return ""
+
+    def __init__(self, **kwargs):
+        pass
+
+    def put(self, objs: List[str]):
+        pass
+
+    def get(self, obj: str):
+        pass

gptcache/manager/eviction/manager.py

@@ -18,8 +18,8 @@ def __init__(self):
     @staticmethod
     def get(
         name: str,
-        policy: str,
-        maxsize: int,
+        policy: str = "LRU",
+        maxsize: int = 1000,
         clean_size: int = 0,
         on_evict: Callable[[List[Any]], None] = None,
         **kwargs
@@ -32,6 +32,17 @@ def get(
             eviction_base = MemoryCacheEviction(
                 policy, maxsize, clean_size, on_evict, **kwargs
             )
+            return eviction_base
+        if name == "redis":
+            from gptcache.manager.eviction.distributed_cache import RedisCacheEviction
+            if policy == "LRU":
+                policy = None
+            eviction_base = RedisCacheEviction(policy=policy, **kwargs)
+            return eviction_base
+        if name == "no_op_eviction":
+            from gptcache. manager.eviction.distributed_cache import NoOpEviction
+            eviction_base = NoOpEviction()
+            return eviction_base
+
         else:
             raise NotFoundError("eviction base", name)
-        return eviction_base

gptcache/manager/eviction/memory_cache.py

@@ -33,12 +33,12 @@ class MemoryCacheEviction(EvictionBase):
     """
 
     def __init__(
-        self,
-        policy: str,
-        maxsize: int,
-        clean_size: int = 0,
-        on_evict: Callable[[List[Any]], None] = None,
-        **kwargs,
+            self,
+            policy: str = "LRU",
+            maxsize: int = 1000,
+            clean_size: int = 0,
+            on_evict: Callable[[List[Any]], None] = None,
+            **kwargs,
     ):
         self._policy = policy.upper()
         if self._policy == "LRU":