Adding locks to ARC and WTINYLFU caches and adding concurrency unit tests

Co-authored-by: omerdor001 <omerdo@post.bgu.ac.il>
Co-authored-by: adiaybgu <adiay@post.bgu.ac.il>

Author: 2 people

modelcache/cache.py

@@ -27,7 +27,7 @@
             #==================== Cache class definition =========================#
             #=====================================================================#
 
-executor = ThreadPoolExecutor(max_workers=6)
+executor = ThreadPoolExecutor(max_workers=2)
 
 def response_text(cache_resp):
     return cache_resp['data']
@@ -300,8 +300,8 @@ async def init(
                 config=vector_config,
                 metric_type=similarity_metric_type,
             ),
-            eviction='WTINYLFU',
-            max_size=100000,
+            eviction='ARC',
+            max_size=10000,
             normalize=normalize,
         )
 

modelcache/manager/data_manager.py

@@ -182,8 +182,7 @@ def __init__(
         self.eviction_base = MemoryCacheEviction(
             policy=policy,
             maxsize=max_size,
-            clean_size=clean_size,
-            on_evict=self._evict_ids)
+            clean_size=clean_size)
 
     def save(self, questions: List[any], answers: List[any], embedding_datas: List[any], **kwargs):
         model = kwargs.pop("model", None)
@@ -314,31 +313,6 @@ def truncate(self, model):
                     'ScalarDB': 'truncate scalar data failed, please check! e: {}'.format(e)}
         return {'status': 'success', 'VectorDB': 'rebuild', 'ScalarDB': 'delete_count: ' + str(delete_count)}
 
-    # added
-    def _evict_ids(self, ids, **kwargs):
-        model = kwargs.get("model")
-        if not ids or any(i is None for i in ids):
-            modelcache_log.warning("Skipping eviction for invalid IDs: %s", ids)
-            return
-
-        if isinstance(ids,str):
-            ids = [ids]
-
-        for _id in ids:
-            self.eviction_base.get_cache(model).pop(_id, None)
-
-        try:
-            self.s.mark_deleted(ids)
-            modelcache_log.info("Evicted from scalar storage: %s", ids)
-        except Exception as e:
-            modelcache_log.error("Failed to delete from scalar storage: %s", str(e))
-
-        try:
-            self.v.delete(ids, model=model)
-            modelcache_log.info("Evicted from vector storage (model=%s): %s", model, ids)
-        except Exception as e:
-            modelcache_log.error("Failed to delete from vector storage (model=%s): %s", model, str(e))
-
     def flush(self):
         self.s.flush()
         self.v.flush()

modelcache/manager/eviction/arc_cache.py

@@ -1,27 +1,18 @@
 from cachetools import Cache
 from collections import OrderedDict
+from readerwriterlock import rwlock
 
-class ARC(Cache):
-    """
-    Adaptive Replacement Cache (ARC) implementation with on_evict callback.
-    Balances recency and frequency via two active lists (T1, T2) and two ghost lists (B1, B2).
-    Calls on_evict([key]) whenever an item is evicted from the active cache.
-    """
+_sentinel = object()
 
-    def __init__(self, maxsize, getsizeof=None, on_evict=None):
-        """
-        Args:
-            maxsize (int): Maximum cache size.
-            getsizeof (callable, optional): Sizing function for items.
-            on_evict (callable, optional): Callback called as on_evict([key]) when a key is evicted.
-        """
+class ARC(Cache):
+    def __init__(self, maxsize, getsizeof=None):
         super().__init__(maxsize, getsizeof)
         self.t1 = OrderedDict()
         self.t2 = OrderedDict()
         self.b1 = OrderedDict()
         self.b2 = OrderedDict()
-        self.p = 0  # Adaptive target for T1 size.
-        self.on_evict = on_evict
+        self.p = 0
+        self._rw_lock = rwlock.RWLockWrite()
 
     def __len__(self):
         return len(self.t1) + len(self.t2)
@@ -30,96 +21,80 @@ def __contains__(self, key):
         return key in self.t1 or key in self.t2
 
     def _evict_internal(self):
-        """
-        Evicts items from T1 or T2 if cache is over capacity, and prunes ghost lists.
-        Calls on_evict for each evicted key.
-        """
-        # Evict from T1 or T2 if active cache > maxsize
         while len(self.t1) + len(self.t2) > self.maxsize:
             if len(self.t1) > self.p or (len(self.t1) == 0 and len(self.t2) > 0):
                 key, value = self.t1.popitem(last=False)
                 self.b1[key] = value
-                if self.on_evict:
-                    self.on_evict([key])
             else:
                 key, value = self.t2.popitem(last=False)
                 self.b2[key] = value
-                if self.on_evict:
-                    self.on_evict([key])
-        # Prune ghost lists to their max lengths
         while len(self.b1) > (self.maxsize - self.p):
             self.b1.popitem(last=False)
         while len(self.b2) > self.p:
             self.b2.popitem(last=False)
 
     def __setitem__(self, key, value):
-        # Remove from all lists before re-inserting
-        for l in (self.t1, self.t2, self.b1, self.b2):
-            l.pop(key, None)
-        self.t1[key] = value
-        self.t1.move_to_end(key)
-        self._evict_internal()
+        with self._rw_lock.gen_wlock():
+            for l in (self.t1, self.t2, self.b1, self.b2):
+                l.pop(key, None)
+            self.t1[key] = value
+            self.t1.move_to_end(key)
+            self._evict_internal()
 
     def __getitem__(self, key):
-        # Case 1: Hit in T1 → promote to T2
-        if key in self.t1:
-            value = self.t1.pop(key)
-            self.t2[key] = value
-            self.t2.move_to_end(key)
-            self.p = max(0, self.p - 1)
-            self._evict_internal()
-            return value
-        # Case 2: Hit in T2 → refresh in T2
-        if key in self.t2:
-            value = self.t2.pop(key)
-            self.t2[key] = value
-            self.t2.move_to_end(key)
-            self.p = min(self.maxsize, self.p + 1)
-            self._evict_internal()
-            return value
-        # Case 3: Hit in B1 (ghost) → fetch and promote to T2
-        if key in self.b1:
-            self.b1.pop(key)
-            self.p = min(self.maxsize, self.p + 1)
-            self._evict_internal()
-            value = super().__missing__(key)
-            self.t2[key] = value
-            self.t2.move_to_end(key)
-            return value
-        # Case 4: Hit in B2 (ghost) → fetch and promote to T2
-        if key in self.b2:
-            self.b2.pop(key)
-            self.p = max(0, self.p - 1)
-            self._evict_internal()
-            value = super().__missing__(key)
-            self.t2[key] = value
-            self.t2.move_to_end(key)
-            return value
-        # Case 5: Cold miss → handled by Cache base class (calls __setitem__ after __missing__)
-        return super().__getitem__(key)
+        with self._rw_lock.gen_wlock():
+            if key in self.t1:
+                value = self.t1.pop(key)
+                self.t2[key] = value
+                self.t2.move_to_end(key)
+                self.p = max(0, self.p - 1)
+                self._evict_internal()
+                return value
+            if key in self.t2:
+                value = self.t2.pop(key)
+                self.t2[key] = value
+                self.t2.move_to_end(key)
+                self.p = min(self.maxsize, self.p + 1)
+                self._evict_internal()
+                return value
+            if key in self.b1:
+                self.b1.pop(key)
+                self.p = min(self.maxsize, self.p + 1)
+                self._evict_internal()
+                value = super().__missing__(key)
+                self.t2[key] = value
+                self.t2.move_to_end(key)
+                return value
+            if key in self.b2:
+                self.b2.pop(key)
+                self.p = max(0, self.p - 1)
+                self._evict_internal()
+                value = super().__missing__(key)
+                self.t2[key] = value
+                self.t2.move_to_end(key)
+                return value
+            return super().__getitem__(key)
 
     def __missing__(self, key):
-        """
-        Override this in a subclass, or rely on direct assignment (cache[key] = value).
-        """
         raise KeyError(key)
 
-    def pop(self, key, default=None):
-        """
-        Remove key from all lists.
-        """
-        for l in (self.t1, self.t2, self.b1, self.b2):
-            if key in l:
-                return l.pop(key)
-        return default
+    def pop(self, key, default=_sentinel):
+        with self._rw_lock.gen_wlock():
+            for l in (self.t1, self.t2, self.b1, self.b2):
+                if key in l:
+                    return l.pop(key)
+            if default is _sentinel:
+                raise KeyError(key)
+            return default
 
     def clear(self):
-        self.t1.clear()
-        self.t2.clear()
-        self.b1.clear()
-        self.b2.clear()
-        self.p = 0
-        super().clear()
+        with self._rw_lock.gen_wlock():
+            self.t1.clear()
+            self.t2.clear()
+            self.b1.clear()
+            self.b2.clear()
+            self.p = 0
+            super().clear()
 
     def __iter__(self):
         yield from self.t1

modelcache/manager/eviction/memory_cache.py

@@ -19,17 +19,14 @@ def wrapper(*args, **kwargs):
 
 
 class MemoryCacheEviction(EvictionBase):
-    def __init__(self, policy: str, maxsize: int, clean_size: int, on_evict: Callable[[List[Any]], None], **kwargs):
+    def __init__(self, policy: str, maxsize: int, clean_size: int, **kwargs):
         self._policy = policy.upper()
         self.model_to_cache = dict()
         self.maxsize = maxsize
         self.clean_size = clean_size
-        self.on_evict = on_evict
         self.kwargs = kwargs
 
     def create_cache(self, model: str):
-
-        ################# Not integrated with on_evict yet #######################
         if self._policy == "LRU":
             cache = cachetools.LRUCache(maxsize=self.maxsize, **self.kwargs)
         elif self._policy == "LFU":
@@ -38,18 +35,14 @@ def create_cache(self, model: str):
             cache = cachetools.FIFOCache(maxsize=self.maxsize, **self.kwargs)
         elif self._policy == "RR":
             cache = cachetools.RRCache(maxsize=self.maxsize, **self.kwargs)
-        ###########################################################################
-
         elif self._policy == "WTINYLFU":
-            cache = W2TinyLFU(maxsize=self.maxsize, on_evict=lambda x: self.on_evict(x,model=model))
+            cache = W2TinyLFU(maxsize=self.maxsize)
         elif self._policy == "ARC":
-            cache = ARC(maxsize=self.maxsize, on_evict=lambda x: self.on_evict(x,model=model))
+            cache = ARC(maxsize=self.maxsize)
         else:
             raise ValueError(f"Unknown policy {self.policy}")
-        cache.popitem = popitem_wrapper(cache.popitem, self.on_evict, self.clean_size)
         return cache
 
-
     def put(self, objs: List[Tuple[Any, Any]], model: str):
         cache = self.get_cache(model)
         for key, value in objs: