improve: read write lock and less blocking

Author: samedii

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "shared-lru-cache"
-version = "0.1.1"
+version = "0.1.2"
 description = ""
 authors = ["Richard Löwenström <[email protected]>"]
 readme = "README.md"

pyproject.toml

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from multiprocessing import Condition, Lock, Manager, Value
+
+
+class ReadWriteLock:
+    readers: Value[int]
+    readers_lock: Lock
+    writers_lock: Lock
+    readers_ok: Condition
+
+    def __init__(self, manager: Manager):
+        self.readers = manager.Value("i", 0)
+        self.readers_lock = manager.Lock()
+        self.writers_lock = manager.Lock()
+        self.readers_ok = manager.Condition(self.readers_lock)
+
+    def read_lock(self):
+        return ReadLock(self)
+
+    def write_lock(self):
+        return WriteLock(self)
+
+
+class ReadLock:
+    rw_lock: ReadWriteLock
+
+    def __init__(self, rw_lock):
+        self.rw_lock = rw_lock
+
+    def __enter__(self):
+        with self.rw_lock.readers_lock:
+            self.rw_lock.readers.value += 1
+            if self.rw_lock.readers.value == 1:
+                self.rw_lock.writers_lock.acquire()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        with self.rw_lock.readers_lock:
+            self.rw_lock.readers.value -= 1
+            if self.rw_lock.readers.value == 0:
+                self.rw_lock.writers_lock.release()
+
+
+class WriteLock:
+    rw_lock: ReadWriteLock
+
+    def __init__(self, rw_lock):
+        self.rw_lock = rw_lock
+
+    def __enter__(self):
+        self.rw_lock.writers_lock.acquire()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.rw_lock.writers_lock.release()

shared_lru_cache/read_write_lock.py

@@ -2,14 +2,18 @@
 import pickle
 from multiprocessing import Manager
 
+import numpy as np
+import torch
+
+from .read_write_lock import ReadWriteLock
+
 
 class SharedLRUCache:
     def __init__(self, maxsize=128):
         self.maxsize = maxsize
         self.manager = Manager()
         self.cache = self.manager.dict()
-        self.read_lock = self.manager.RLock()
-        self.write_lock = self.manager.Lock()
+        self.lock = ReadWriteLock(self.manager)
         self.order = self.manager.list()
         self.data_store = self.manager.dict()
 
@@ -18,26 +22,26 @@ def __call__(self, func):
         def wrapper(*args, **kwargs):
             key = str((args, frozenset(kwargs.items())))
 
-            with self.read_lock:
+            with self.lock.read_lock():
                 if key in self.cache:
-                    with self.write_lock:
-                        if (
-                            key in self.order
-                        ):  # Check if key is in order before removing
-                            self.order.remove(key)
-                        self.order.append(key)
-                    return pickle.loads(self.data_store[key])
+                    hit = True
+                    serialized_result, obj_info = self.data_store[key]
+                else:
+                    hit = False
+
+            if hit:
+                return self.deserialize(serialized_result, obj_info)
 
             result = func(*args, **kwargs)
-            serialized_result = pickle.dumps(result)
+            serialized_result, obj_info = self.serialize(result)
 
-            with self.write_lock:
-                # Check again in case another process has updated the cache
+            with self.lock.read_lock():
                 if key in self.cache:
                     return result
 
+            with self.lock.write_lock():
                 self.cache[key] = key
-                self.data_store[key] = serialized_result
+                self.data_store[key] = (serialized_result, obj_info)
                 self.order.append(key)
 
                 while len(self.order) > self.maxsize:
@@ -52,6 +56,32 @@ def wrapper(*args, **kwargs):
         wrapper.data_store = self.data_store
         return wrapper
 
+    def serialize(self, obj):
+        if isinstance(obj, np.ndarray):
+            obj_info = ("numpy", obj.shape, obj.dtype.str)
+            return obj.tobytes(), obj_info
+        elif isinstance(obj, torch.Tensor):
+            obj.byte()
+            numpy_array = obj.cpu().numpy()
+            obj_info = ("torch", numpy_array.shape, numpy_array.dtype.str)
+            return numpy_array.tobytes(), obj_info
+        else:
+            obj_info = ("other",)
+            return pickle.dumps(obj), obj_info
+
+    def deserialize(self, data, obj_info):
+        obj_type, *info = obj_info
+        if obj_type == "numpy":
+            shape, dtype = info
+            return np.frombuffer(data, dtype=np.dtype(dtype)).reshape(shape)
+        elif obj_type == "torch":
+            shape, dtype = info
+            dtype = np.dtype(dtype) if isinstance(dtype, str) else dtype
+            numpy_array = np.frombuffer(data, dtype=dtype).reshape(shape)
+            return torch.from_numpy(numpy_array)
+        else:
+            return pickle.loads(data)
+
 
 def shared_lru_cache(maxsize=128):
     return SharedLRUCache(maxsize)

shared_lru_cache/shared_lru_cache.py

@@ -7,12 +7,12 @@
 
 from shared_lru_cache import shared_lru_cache
 
-MAX_INDEX = 20
+MAX_INDEX = 10
 
 
 def load_image(idx):
-    time.sleep(0.2)  # Simulate some data loading time
-    return torch.randn(1024, 1024)
+    time.sleep(5)  # Simulate some data loading time
+    return torch.zeros((30000, 2500), dtype=torch.uint8)
 
 
 lru_cached_load_image = lru_cache(maxsize=128)(load_image)