fix (RedisRegistry): avoid overwriting cache with the same key during…

… load (#223)

- fix cache key overwrite issue which caused a stale model to persist
for a long time
- Save sourcetype in artifact data extras to be used for other logic
- Improve docs

---------

Signed-off-by: Avik Basu <[email protected]>

numalogic/registry/artifact.py

@@ -48,6 +48,8 @@ class ArtifactManager(Generic[KEYS, A_D]):
  uri: server/connection uri
  """
 
+ _STORETYPE = "registry"
+
  __slots__ = ("uri",)
 
  def __init__(self, uri: str):
@@ -137,6 +139,8 @@ class ArtifactCache(Generic[M_K, A_D]):
  ttl: time to live for each item in the cache
  """
 
+ _STORETYPE = "cache"
+
  __slots__ = ("_cachesize", "_ttl")
 
  def __init__(self, cachesize: int, ttl: int):

numalogic/registry/dynamodb_registry.py

@@ -1,3 +1,14 @@
+# Copyright 2022 The Numaproj Authors.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# http:https://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import logging
 import time
 from typing import Any, Optional

numalogic/registry/localcache.py

@@ -9,6 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from copy import deepcopy
 from typing import Optional
 
 from cachetools import TTLCache
@@ -34,14 +35,55 @@ def __init__(self, cachesize: int = 512, ttl: int = 300):
  if not self.__cache:
  self.__cache = TTLCache(maxsize=cachesize, ttl=ttl)
 
- def load(self, artifact_key: str) -> ArtifactData:
+ def __contains__(self, artifact_key: str) -> bool:
+ """Check if an artifact is in the cache."""
+ return artifact_key in self.__cache
+
+ def load(self, artifact_key: str) -> Optional[ArtifactData]:
+ """
+ Load an artifact from the cache.
+
+ Args:
+ ----
+ artifact_key: The key of the artifact to load.
+
+ Returns
+ -------
+ The artifact data instance if found, None otherwise.
+ """
  return self.__cache.get(artifact_key)
 
  def save(self, key: str, artifact: ArtifactData) -> None:
+ """
+ Save an artifact to the cache.
+
+ Args:
+ ----
+ key: The key of the artifact to save.
+ artifact: The artifact data instance to save.
+ """
+ artifact = deepcopy(artifact)
+ artifact.extras["source"] = self._STORETYPE
  self.__cache[key] = artifact
 
  def delete(self, key: str) -> Optional[ArtifactData]:
+ """
+ Delete an artifact from the cache.
+
+ Args:
+ ----
+ key: The key of the artifact to delete.
+
+ Returns
+ -------
+ The deleted artifact data instance if found, None otherwise.
+ """
  return self.__cache.pop(key, default=None)
 
  def clear(self) -> None:
+ """Clears the whole cache."""
  self.__cache.clear()
+
+ def keys(self) -> list[str]:
+ """Returns the current keys of the cache."""
+ return list(_key for _key in self.__cache)

numalogic/registry/redis_registry.py

@@ -1,3 +1,14 @@
+# Copyright 2022 The Numaproj Authors.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# http:https://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import logging
 import time
 from datetime import datetime, timedelta
@@ -18,7 +29,7 @@ class RedisRegistry(ArtifactManager):
 
  Args:
  ----
- client: Take in the reids client already established/created
+ client: Take in the redis client already established/created
  ttl: Total Time to Live (in seconds) for the key when saving in redis (dafault = 604800)
  cache_registry: Cache registry to use (default = None).
 
@@ -94,6 +105,7 @@ def _load_from_cache(self, key: str) -> Optional[ArtifactData]:
 
  def _save_in_cache(self, key: str, artifact_data: ArtifactData) -> None:
  if self.cache_registry:
+ _LOGGER.debug("Saving artifact in cache with key: %s", key)
  self.cache_registry.save(key, artifact_data)
 
  def _clear_cache(self, key: Optional[str] = None) -> Optional[ArtifactData]:
@@ -123,20 +135,38 @@ def __get_artifact_data(
  extras={
  "timestamp": float(artifact_timestamp.decode()),
  "version": artifact_version.decode(),
+ "source": self._STORETYPE,
  },
  )
 
- def __load_latest_artifact(self, key: str) -> ArtifactData:
+ def __load_latest_artifact(self, key: str) -> tuple[ArtifactData, bool]:
+ """
+ Load the latest artifact from the registry.
+
+ Args:
+ key: full model key.
+
+ Returns
+ -------
+ ArtifactData and a boolean flag indicating if the artifact was loaded from cache.
+
+ Raises
+ ------
+ ModelKeyNotFound: If the model key is not found in the registry.
+ """
  cached_artifact = self._load_from_cache(key)
  if cached_artifact:
  _LOGGER.debug("Found cached artifact for key: %s", key)
- return cached_artifact
+ return cached_artifact, True
  latest_key = self.__construct_latest_key(key)
  if not self.client.exists(latest_key):
  raise ModelKeyNotFound(f"latest key: {latest_key}, Not Found !!!")
  model_key = self.client.get(latest_key)
- _LOGGER.info("latest key, %s, is pointing to the key : %s", latest_key, model_key)
- return self.__load_version_artifact(version=self.get_version(model_key.decode()), key=key)
+ _LOGGER.debug("latest key, %s, is pointing to the key : %s", latest_key, model_key)
+ return (
+ self.__load_version_artifact(version=self.get_version(model_key.decode()), key=key),
+ False,
+ )
 
  def __load_version_artifact(self, version: str, key: str) -> ArtifactData:
  model_key = self.__construct_version_key(key, version)
@@ -152,7 +182,7 @@ def __save_artifact(
  new_version_key = self.__construct_version_key(key, version)
  latest_key = self.__construct_latest_key(key)
  pipe.set(name=latest_key, value=new_version_key)
- _LOGGER.info("Setting latest key : %s ,to this new key = %s", latest_key, new_version_key)
+ _LOGGER.debug("Setting latest key : %s ,to this new key = %s", latest_key, new_version_key)
  serialized_metadata = ""
  if metadata:
  serialized_metadata = dumps(deserialized_object=metadata)
@@ -178,6 +208,8 @@ def load(
  """Loads the artifact from redis registry. Either latest or version (one of the arguments)
  is needed to load the respective artifact.
 
+ If cache registry is provided, it will first check the cache registry for the artifact.
+
  Args:
  ----
  skeys: static key fields as list/tuple of strings
@@ -188,19 +220,26 @@ def load(
  Returns
  -------
  ArtifactData instance
+
+ Raises
+ ------
+ ValueError: If both latest and version are provided or none of them are provided.
+ RedisRegistryError: If any redis error occurs.
  """
  if (latest and version) or (not latest and not version):
  raise ValueError("Either One of 'latest' or 'version' needed in load method call")
  key = self.construct_key(skeys, dkeys)
+ is_cached = False
  try:
  if latest:
- artifact_data = self.__load_latest_artifact(key)
- self._save_in_cache(key, artifact_data)
+ artifact_data, is_cached = self.__load_latest_artifact(key)
  else:
  artifact_data = self.__load_version_artifact(version, key)
  except RedisError as err:
  raise RedisRegistryError(f"{err.__class__.__name__} raised") from err
  else:
+ if (not is_cached) and latest:
+ self._save_in_cache(key, artifact_data)
  return artifact_data
 
  def save(
@@ -221,24 +260,28 @@ def save(
 
  Returns
  -------
- model version
+ Model version (str)
+
+ Raises
+ ------
+ RedisRegistryError: If there is any RedisError while saving the artifact.
  """
  key = self.construct_key(skeys, dkeys)
  latest_key = self.__construct_latest_key(key)
  version = 0
  try:
  if self.client.exists(latest_key):
- _LOGGER.debug("latest key exists for the model")
+ _LOGGER.debug("Latest key: %s exists for the model", latest_key)
  version_key = self.client.get(name=latest_key)
  version = int(self.get_version(version_key.decode())) + 1
  with self.client.pipeline() as pipe:
  new_version_key = self.__save_artifact(pipe, artifact, metadata, key, str(version))
  pipe.expire(name=new_version_key, time=self.ttl)
- _LOGGER.info("Model with the key = %s, loaded successfully.", new_version_key)
  pipe.execute()
  except RedisError as err:
  raise RedisRegistryError(f"{err.__class__.__name__} raised") from err
  else:
+ _LOGGER.info("Model with the key = %s, saved successfully.", new_version_key)
  return str(version)
 
  def delete(self, skeys: KEYS, dkeys: KEYS, version: str) -> None:
@@ -249,21 +292,25 @@ def delete(self, skeys: KEYS, dkeys: KEYS, version: str) -> None:
  skeys: static key fields as list/tuple of strings
  dkeys: dynamic key fields as list/tuple of strings
  version: model version to delete.
+
+ Raises
+ ------
+ ModelKeyNotFound: If the model version is not found in registry.
+ RedisRegistryError: If there is any RedisError while deleting the artifact.
  """
  key = self.construct_key(skeys, dkeys)
  del_key = self.__construct_version_key(key, version)
  try:
  if self.client.exists(del_key):
  self.client.delete(del_key)
- _LOGGER.info("Model with the key = %s, deleted successfully", del_key)
  else:
- _LOGGER.debug("Key to delete: %s, Not Found !!!\n Exiting.....", del_key)
  raise ModelKeyNotFound(
- "Key to delete: %s, Not Found !!!\n Exiting....." % del_key,
+ "Key to delete: %s, Not Found!" % del_key,
  )
  except RedisError as err:
  raise RedisRegistryError(f"{err.__class__.__name__} raised") from err
  else:
+ _LOGGER.info("Model with the key = %s, deleted successfully", del_key)
  self._clear_cache(del_key)
 
  @staticmethod
@@ -276,6 +323,13 @@ def is_artifact_stale(artifact_data: ArtifactData, freq_hr: int) -> bool:
  artifact_data: ArtifactData object to look into
  freq_hr: Frequency of retraining in hours.
 
+ Returns
+ -------
+ True if artifact is stale, False otherwise.
+
+ Raises
+ ------
+ RedisRegistryError: If there is any error while fetching timestamp information.
  """
  try:
  artifact_ts = float(artifact_data.extras["timestamp"])

tests/registry/test_cache.py

@@ -27,9 +27,11 @@ def test_cache_size(self):
 
  self.assertIsNone(cache_registry.load("m1"))
  self.assertIsInstance(cache_registry.load("m2"), ArtifactData)
- self.assertIsInstance(cache_registry.load("m3"), ArtifactData)
  self.assertEqual(2, cache_registry.cachesize)
  self.assertEqual(1, cache_registry.ttl)
+ self.assertTrue("m2" in cache_registry)
+ self.assertTrue("m3" in cache_registry)
+ self.assertListEqual(["m2", "m3"], cache_registry.keys())
 
  def test_cache_overwrite(self):
  cache_registry = LocalLRUCache(cachesize=2, ttl=1)
@@ -41,7 +43,7 @@ def test_cache_overwrite(self):
  )
 
  loaded_artifact = cache_registry.load("m1")
- self.assertDictEqual({"version": "2"}, loaded_artifact.extras)
+ self.assertDictEqual({"version": "2", "source": "cache"}, loaded_artifact.extras)
 
  def test_cache_ttl(self):
  cache_registry = LocalLRUCache(cachesize=2, ttl=1)

tests/registry/test_redis_registry.py

@@ -1,3 +1,5 @@
+import logging
+import time
 import unittest
 from datetime import datetime, timedelta
 from unittest.mock import Mock, patch
@@ -12,6 +14,8 @@
 from numalogic.registry import RedisRegistry, LocalLRUCache, ArtifactData
 from numalogic.tools.exceptions import ModelKeyNotFound, RedisRegistryError
 
+logging.basicConfig(level=logging.DEBUG)
+
 
 class TestRedisRegistry(unittest.TestCase):
  @classmethod
@@ -25,7 +29,7 @@ def setUpClass(cls) -> None:
  cls.redis_client = fakeredis.FakeStrictRedis(server=server, decode_responses=False)
 
  def setUp(self):
- self.cache = LocalLRUCache(cachesize=4, ttl=300)
+ self.cache = LocalLRUCache(cachesize=4, ttl=1)
  self.registry = RedisRegistry(
  client=self.redis_client,
  cache_registry=self.cache,
@@ -162,19 +166,44 @@ def test_load_model_when_no_model(self):
  with self.assertRaises(ModelKeyNotFound):
  self.registry.load(skeys=self.skeys, dkeys=self.dkeys)
 
- def test_load_model_when_model_stale(self):
- with self.assertRaises(ModelKeyNotFound):
- version = self.registry.save(
- skeys=self.skeys, dkeys=self.dkeys, artifact=self.pytorch_model
- )
- self.registry.delete(skeys=self.skeys, dkeys=self.dkeys, version=str(version))
- self.registry.load(skeys=self.skeys, dkeys=self.dkeys)
+ def test_load_latest_model_twice(self):
+ with freeze_time(datetime.today() - timedelta(days=5)):
+ self.registry.save(skeys=self.skeys, dkeys=self.dkeys, artifact=self.pytorch_model)
+
+ artifact_data_1 = self.registry.load(skeys=self.skeys, dkeys=self.dkeys)
+ artifact_data_2 = self.registry.load(skeys=self.skeys, dkeys=self.dkeys)
+ self.assertTrue(self.registry.is_artifact_stale(artifact_data_1, 4))
+ self.assertEqual("registry", artifact_data_1.extras["source"])
+ self.assertEqual("cache", artifact_data_2.extras["source"])
+
+ def test_load_latest_cache_ttl_expire(self):
+ self.registry.save(skeys=self.skeys, dkeys=self.dkeys, artifact=self.pytorch_model)
+ artifact_data_1 = self.registry.load(skeys=self.skeys, dkeys=self.dkeys)
+ time.sleep(1)
+ artifact_data_2 = self.registry.load(skeys=self.skeys, dkeys=self.dkeys)
+ self.assertEqual("registry", artifact_data_1.extras["source"])
+ self.assertEqual("registry", artifact_data_2.extras["source"])
+
+ def test_load_non_latest_model_twice(self):
+ old_version = self.registry.save(
+ skeys=self.skeys, dkeys=self.dkeys, artifact=self.pytorch_model
+ )
+ self.registry.save(skeys=self.skeys, dkeys=self.dkeys, artifact=self.pytorch_model)
+
+ artifact_data_1 = self.registry.load(
+ skeys=self.skeys, dkeys=self.dkeys, latest=False, version=old_version
+ )
+ artifact_data_2 = self.registry.load(
+ skeys=self.skeys, dkeys=self.dkeys, latest=False, version=old_version
+ )
+ self.assertEqual("registry", artifact_data_1.extras["source"])
+ self.assertEqual("registry", artifact_data_2.extras["source"])
 
  def test_delete_version(self):
+ version = self.registry.save(
+ skeys=self.skeys, dkeys=self.dkeys, artifact=self.pytorch_model
+ )
  with self.assertRaises(ModelKeyNotFound):
- version = self.registry.save(
- skeys=self.skeys, dkeys=self.dkeys, artifact=self.pytorch_model
- )
  self.registry.delete(skeys=self.skeys, dkeys=self.dkeys, version=str(version))
  self.registry.load(skeys=self.skeys, dkeys=self.dkeys)