2023-02-14 00:04:46 +00:00
|
|
|
from logging import getLogger
|
2023-02-14 00:12:40 +00:00
|
|
|
from typing import Any, List, Tuple
|
2023-02-14 00:04:46 +00:00
|
|
|
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
|
|
|
|
|
|
|
|
class ModelCache:
|
2023-02-14 00:12:40 +00:00
|
|
|
cache: List[Tuple[str, Any, Any]]
|
2023-02-14 00:04:46 +00:00
|
|
|
limit: int
|
|
|
|
|
|
|
|
def __init__(self, limit: int) -> None:
|
2023-02-14 00:15:02 +00:00
|
|
|
self.cache = []
|
2023-02-14 00:04:46 +00:00
|
|
|
self.limit = limit
|
|
|
|
|
|
|
|
def drop(self, tag: str, key: Any) -> None:
|
2023-02-14 00:15:02 +00:00
|
|
|
self.cache[:] = [
|
2023-02-14 00:10:11 +00:00
|
|
|
model for model in self.cache if model[0] != tag and model[1] != key
|
|
|
|
]
|
2023-02-14 00:04:46 +00:00
|
|
|
|
|
|
|
def get(self, tag: str, key: Any) -> Any:
|
|
|
|
for t, k, v in self.cache:
|
|
|
|
if tag == t and key == k:
|
|
|
|
return v
|
|
|
|
|
|
|
|
return None
|
|
|
|
|
|
|
|
def set(self, tag: str, key: Any, value: Any) -> None:
|
2023-02-15 23:17:28 +00:00
|
|
|
if self.limit == 0:
|
2023-02-17 00:42:05 +00:00
|
|
|
logger.debug("cache limit set to 0, not caching model: %s", tag)
|
2023-02-15 23:17:28 +00:00
|
|
|
return
|
|
|
|
|
2023-02-14 00:04:46 +00:00
|
|
|
for i in range(len(self.cache)):
|
|
|
|
t, k, v = self.cache[i]
|
2023-02-19 13:53:20 +00:00
|
|
|
if tag == t and key != k:
|
|
|
|
logger.debug("updating model cache: %s", tag)
|
|
|
|
self.cache[i] = (tag, key, value)
|
|
|
|
return
|
2023-02-14 00:04:46 +00:00
|
|
|
|
2023-02-17 00:42:05 +00:00
|
|
|
logger.debug("adding new model to cache: %s", tag)
|
2023-02-14 00:04:46 +00:00
|
|
|
self.cache.append((tag, key, value))
|
|
|
|
self.prune()
|
|
|
|
|
|
|
|
def prune(self):
|
|
|
|
total = len(self.cache)
|
|
|
|
if total > self.limit:
|
2023-02-14 00:10:11 +00:00
|
|
|
logger.info(
|
|
|
|
"Removing models from cache, %s of %s", (total - self.limit), total
|
|
|
|
)
|
2023-02-14 00:33:06 +00:00
|
|
|
self.cache[:] = self.cache[-self.limit :]
|
2023-02-14 00:04:46 +00:00
|
|
|
else:
|
2023-02-17 00:42:05 +00:00
|
|
|
logger.debug("model cache below limit, %s of %s", total, self.limit)
|