Merge pull request #289 from Maximilian-Winter/main

Diskcache implementation for llama state.

llama_cpp/llama.py

@@ -4,6 +4,7 @@ import uuid
 import time
 import math
 import multiprocessing
+from abc import ABC
 from typing import (
     List,
     Optional,
@@ -17,6 +18,8 @@ from typing import (
 )
 from collections import deque, OrderedDict
 
+import diskcache
+
 from . import llama_cpp
 from .llama_types import *
 
@@ -24,12 +27,39 @@ import numpy as np
 import numpy.typing as npt
 
 
-class LlamaCache:
+class LlamaCache(ABC):
-    """Cache for a llama.cpp model."""
+    """Base cache class for a llama.cpp model."""
 
     def __init__(self, capacity_bytes: int = (2 << 30)):
-        self.cache_state: OrderedDict[Tuple[int, ...], "LlamaState"] = OrderedDict()
+        pass
+
+    @property
+    def cache_size(self):
+        return 0
+
+    def _find_longest_prefix_key(
+            self,
+            key: Tuple[int, ...],
+    ) -> Optional[Tuple[int, ...]]:
+        pass
+
+    def __getitem__(self, key: Sequence[int]) -> "LlamaState":
+        pass
+
+    def __contains__(self, key: Sequence[int]) -> bool:
+        pass
+
+    def __setitem__(self, key: Sequence[int], value: "LlamaState"):
+        pass
+
+
+class LlamaRAMCache(LlamaCache):
+    """Cache for a llama.cpp model using RAM."""
+
+    def __init__(self, capacity_bytes: int = (2 << 30)):
+        super().__init__(capacity_bytes)
         self.capacity_bytes = capacity_bytes
+        self.cache_state: OrderedDict[Tuple[int, ...], "LlamaState"] = OrderedDict()
 
     @property
     def cache_size(self):
@@ -54,7 +84,7 @@ class LlamaCache:
         key = tuple(key)
         _key = self._find_longest_prefix_key(key)
         if _key is None:
-            raise KeyError(f"Key not found")
+            raise KeyError("Key not found")
         value = self.cache_state[_key]
         self.cache_state.move_to_end(_key)
         return value
@@ -71,6 +101,49 @@ class LlamaCache:
             self.cache_state.popitem(last=False)
 
 
+class LlamaDiskCache(LlamaCache):
+    """Cache for a llama.cpp model using disk."""
+
+    def __init__(self, cache_dir="./llama_cache", capacity_bytes: int = (2 << 30)):
+        super().__init__(capacity_bytes)
+        self.cache = diskcache.Cache(cache_dir)
+
+    @property
+    def cache_size(self):
+        return self.cache.volume()
+
+    def _find_longest_prefix_key(
+            self,
+            key: Tuple[int, ...],
+    ) -> Optional[Tuple[int, ...]]:
+        min_len = 0
+        min_key = None
+        for k in self.cache.iterkeys():
+            prefix_len = Llama.longest_token_prefix(k, key)
+            if prefix_len > min_len:
+                min_len = prefix_len
+                min_key = k
+        return min_key
+
+    def __getitem__(self, key: Sequence[int]) -> "LlamaState":
+        key = tuple(key)
+        _key = self._find_longest_prefix_key(key)
+        if _key is None:
+            raise KeyError("Key not found")
+        value = self.cache.pop(_key)
+        self.cache.push(_key)
+        return value
+
+    def __setitem__(self, key: Sequence[int], value: "LlamaState"):
+        key = tuple(key)
+        if key in self.cache:
+            del self.cache[key]
+        self.cache[key] = value
+        while self.cache_size > self.capacity_bytes:
+            key_to_remove = next(iter(self.cache))
+            del self.cache[key_to_remove]
+
+
 class LlamaState:
     def __init__(
             self,

pyproject.toml

@@ -16,6 +16,7 @@ include = [
 python = "^3.8.1"
 typing-extensions = "^4.6.3"
 numpy = "^1.20.0"
+diskcache = "^5.6.1"
 uvicorn = { version = "^0.22.0", optional = true }
 fastapi = { version = "^0.96.0", optional = true }
 sse-starlette = { version = "^1.6.1", optional = true }

setup.py

@@ -16,7 +16,7 @@ setup(
     license="MIT",
     package_dir={"llama_cpp": "llama_cpp", "llama_cpp.server": "llama_cpp/server"},
     packages=["llama_cpp", "llama_cpp.server"],
-    install_requires=["typing-extensions>=4.5.0", "numpy>=1.20.0"],
+    install_requires=["typing-extensions>=4.5.0", "numpy>=1.20.0", "diskcache>=5.6.1"],
     extras_require={
         "server": ["uvicorn>=0.21.1", "fastapi>=0.95.0", "sse-starlette>=1.3.3"],
     },