feat: add multiple save for redis registry (#281)

1) Option to multisave in redis registry
2) support caching for version calls.
3) In udfs/ load models using version call instead of latest calls.
4) call multiple_save in artifacts to save artifacts in a transaction.

---------

Signed-off-by: s0nicboOm <[email protected]>

.github/workflows/ci.yml

@@ -19,7 +19,7 @@ jobs:
  - uses: actions/checkout@v3
 
  - name: Install poetry
- run: pipx install poetry==1.5.1
+ run: pipx install poetry==1.6.1
 
  - name: Set up Python ${{ matrix.python-version }}
  uses: actions/setup-python@v4

.github/workflows/coverage.yml

@@ -19,7 +19,7 @@ jobs:
  - uses: actions/checkout@v3
 
  - name: Install poetry
- run: pipx install poetry==1.5.1
+ run: pipx install poetry==1.6.1
 
  - name: Set up Python ${{ matrix.python-version }}
  uses: actions/setup-python@v4

.github/workflows/pypi.yml

@@ -19,7 +19,7 @@ jobs:
  - uses: actions/checkout@v3
 
  - name: Install poetry
- run: pipx install poetry==1.5.1
+ run: pipx install poetry==1.6.1
 
  - name: Set up Python ${{ matrix.python-version }}
  uses: actions/setup-python@v4

numalogic/connectors/druid.py

@@ -52,7 +52,7 @@ def fetch(
 
  end_dt = datetime.now(pytz.utc)
  start_dt = end_dt - timedelta(hours=hours)
- intervals = f"{start_dt.isoformat()}/{end_dt.isoformat()}"
+ intervals = [f"{start_dt.isoformat()}/{end_dt.isoformat()}"]
 
  dimension_specs = map(lambda d: DimensionSpec(dimension=d, output_name=d), dimensions)
 

numalogic/registry/redis_registry.py

@@ -14,13 +14,14 @@
 from datetime import datetime, timedelta
 from typing import Optional
 import orjson
+import redis.client
 
 from redis.exceptions import RedisError
 
 from numalogic.registry.artifact import ArtifactManager, ArtifactData, ArtifactCache
 from numalogic.registry._serialize import loads, dumps
 from numalogic.tools.exceptions import ModelKeyNotFound, RedisRegistryError
-from numalogic.tools.types import artifact_t, redis_client_t, KEYS, META_T, META_VT
+from numalogic.tools.types import artifact_t, redis_client_t, KEYS, META_T, META_VT, KeyedArtifact
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -33,6 +34,8 @@ class RedisRegistry(ArtifactManager):
  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).
+ transactional: Flag to indicate if the registry should be transactional or
+ not (default = False).
 
  Examples
  --------
@@ -48,18 +51,20 @@ class RedisRegistry(ArtifactManager):
  >>> loaded_artifact = registry.load(skeys, dkeys)
  """
 
- __slots__ = ("client", "ttl", "cache_registry")
+ __slots__ = ("client", "ttl", "cache_registry", "transactional")
 
  def __init__(
  self,
  client: redis_client_t,
  ttl: int = 604800,
  cache_registry: Optional[ArtifactCache] = None,
+ transactional: bool = True,
  ):
  super().__init__("")
  self.client = client
  self.ttl = ttl
  self.cache_registry = cache_registry
+ self.transactional = transactional
 
  @staticmethod
  def construct_key(skeys: KEYS, dkeys: KEYS) -> str:
@@ -155,30 +160,37 @@ def __load_latest_artifact(self, key: str) -> tuple[ArtifactData, bool]:
  ------
  ModelKeyNotFound: If the model key is not found in the registry.
  """
- cached_artifact = self._load_from_cache(key)
+ latest_key = self.__construct_latest_key(key)
+ cached_artifact = self._load_from_cache(latest_key)
  if cached_artifact:
- _LOGGER.debug("Found cached artifact for key: %s", key)
+ _LOGGER.debug("Found cached artifact for key: %s", latest_key)
  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.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,
+ artifact, _ = self.__load_version_artifact(
+ version=self.get_version(model_key.decode()), key=key
  )
+ return artifact, False
 
- def __load_version_artifact(self, version: str, key: str) -> ArtifactData:
- model_key = self.__construct_version_key(key, version)
- if not self.client.exists(model_key):
- raise ModelKeyNotFound("Could not find model key with key: %s" % model_key)
- return self.__get_artifact_data(
- model_key=model_key,
+ def __load_version_artifact(self, version: str, key: str) -> tuple[ArtifactData, bool]:
+ version_key = self.__construct_version_key(key, version)
+ cached_artifact = self._load_from_cache(version_key)
+ if cached_artifact:
+ _LOGGER.debug("Found cached version artifact for key: %s", version_key)
+ return cached_artifact, True
+ if not self.client.exists(version_key):
+ raise ModelKeyNotFound(f"Could not find model key with key: {version_key}")
+ return (
+ self.__get_artifact_data(
+ model_key=version_key,
+ ),
+ False,
  )
 
  def __save_artifact(
- self, pipe, artifact: artifact_t, metadata: META_T, key: KEYS, version: str
+ self, pipe, artifact: artifact_t, key: KEYS, version: str, **metadata: META_T
  ) -> str:
  new_version_key = self.__construct_version_key(key, version)
  latest_key = self.__construct_latest_key(key)
@@ -210,6 +222,8 @@ def load(
  is needed to load the respective artifact.
 
  If cache registry is provided, it will first check the cache registry for the artifact.
+ If latest is passed, latest key is saved otherwise version call saves the respective
+ version artifact in cache.
 
  Args:
  ----
@@ -230,24 +244,35 @@ def load(
  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, is_cached = self.__load_latest_artifact(key)
  else:
- artifact_data = self.__load_version_artifact(version, key)
+ artifact_data, is_cached = 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)
+ if not is_cached:
+ if latest:
+ _LOGGER.debug(
+ "Saving %s, in cache as %s", self.__construct_latest_key(key), key
+ )
+ self._save_in_cache(self.__construct_latest_key(key), artifact_data)
+ else:
+ _LOGGER.info(
+ "Saving %s, in cache as %s",
+ self.__construct_version_key(key, version),
+ key,
+ )
+ self._save_in_cache(self.__construct_version_key(key, version), artifact_data)
  return artifact_data
 
  def save(
  self,
  skeys: KEYS,
  dkeys: KEYS,
  artifact: artifact_t,
+ _pipe: Optional[redis.client.Pipeline] = None,
  **metadata: META_VT,
  ) -> Optional[str]:
  """Saves the artifact into redis registry and updates version.
@@ -257,6 +282,7 @@ def save(
  skeys: static key fields as list/tuple of strings
  dkeys: dynamic key fields as list/tuple of strings
  artifact: primary artifact to be saved
+ _pipe: RedisPipeline object
  metadata: additional metadata surrounding the artifact that needs to be saved.
 
  Returns
@@ -275,10 +301,15 @@ def save(
  _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)
- pipe.execute()
+ _redis_pipe = (
+ self.client.pipeline(transaction=self.transactional) if _pipe is None else _pipe
+ )
+ new_version_key = self.__save_artifact(
+ pipe=_redis_pipe, artifact=artifact, key=key, version=str(version), **metadata
+ )
+ _redis_pipe.expire(name=new_version_key, time=self.ttl)
+ if _pipe is None:
+ _redis_pipe.execute()
  except RedisError as err:
  raise RedisRegistryError(f"{err.__class__.__name__} raised") from err
  else:
@@ -306,7 +337,7 @@ def delete(self, skeys: KEYS, dkeys: KEYS, version: str) -> None:
  self.client.delete(del_key)
  else:
  raise ModelKeyNotFound(
- "Key to delete: %s, Not Found!" % del_key,
+ f"Key to delete: {del_key}, Not Found!",
  )
  except RedisError as err:
  raise RedisRegistryError(f"{err.__class__.__name__} raised") from err
@@ -338,3 +369,48 @@ def is_artifact_stale(artifact_data: ArtifactData, freq_hr: int) -> bool:
  raise RedisRegistryError("Error fetching timestamp information") from err
  stale_ts = (datetime.now() - timedelta(hours=freq_hr)).timestamp()
  return stale_ts > artifact_ts
+
+ def save_multiple(
+ self,
+ skeys: KEYS,
+ dict_artifacts: dict[str, KeyedArtifact],
+ **metadata: META_VT,
+ ):
+ """
+ Saves multiple artifacts into redis registry. The last save stores all the
+ artifact versions in the metadata.
+
+ Args:
+ ----
+ skeys: static key fields as list/tuple of strings
+ dict_artifacts: dict of artifacts to save
+ metadata: additional metadata surrounding the artifact that needs to be saved.
+ """
+ dict_model_ver = {}
+ try:
+ with self.client.pipeline(transaction=self.transactional) as pipe:
+ pipe.multi()
+ for key, value in dict_artifacts.items():
+ dict_model_ver[":".join(value.dkeys)] = self.save(
+ skeys=skeys,
+ dkeys=value.dkeys,
+ artifact=value.artifact,
+ _pipe=pipe,
+ **metadata,
+ )
+
+ if len(dict_artifacts) == len(dict_model_ver):
+ self.save(
+ skeys=skeys,
+ dkeys=value.dkeys,
+ artifact=value.artifact,
+ _pipe=pipe,
+ artifact_versions=dict_model_ver,
+ **metadata,
+ )
+ pipe.execute()
+ _LOGGER.info("Successfully saved all the artifacts with: %s", dict_model_ver)
+ except RedisError as err:
+ raise RedisRegistryError(f"{err.__class__.__name__} raised") from err
+ else:
+ return dict_model_ver

numalogic/tools/types.py

@@ -8,10 +8,8 @@
 # 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.
-
-
 from collections.abc import Sequence
-from typing import Union, TypeVar, ClassVar
+from typing import Union, TypeVar, ClassVar, NamedTuple
 
 from sklearn.base import BaseEstimator
 from torch import Tensor
@@ -40,6 +38,13 @@
 KEYS = TypeVar("KEYS", bound=Sequence[str], covariant=True)
 
 
+class KeyedArtifact(NamedTuple):
+ r"""namedtuple for artifacts."""
+
+ dkeys: KEYS
+ artifact: artifact_t
+
+
 class Singleton(type):
  r"""Helper metaclass to use as a Singleton class."""
 

numalogic/udfs/inference.py

@@ -12,17 +12,19 @@
 
 from numalogic.config import RegistryFactory
 from numalogic.registry import LocalLRUCache, ArtifactData
-from numalogic.tools.exceptions import RedisRegistryError, ModelKeyNotFound, ConfigNotFoundError
+from numalogic.tools.exceptions import ConfigNotFoundError
 from numalogic.tools.types import artifact_t, redis_client_t
 from numalogic.udfs._base import NumalogicUDF
 from numalogic.udfs._config import StreamConf, PipelineConf
 from numalogic.udfs.entities import StreamPayload, Header, Status
+from numalogic.udfs.tools import _load_artifact
 
 _LOGGER = logging.getLogger(__name__)
 
 # TODO: move to config
 LOCAL_CACHE_TTL = int(os.getenv("LOCAL_CACHE_TTL", "3600"))
 LOCAL_CACHE_SIZE = int(os.getenv("LOCAL_CACHE_SIZE", "10000"))
+LOAD_LATEST = os.getenv("LOAD_LATEST", "false").lower() == "true"
 
 
 class InferenceUDF(NumalogicUDF):
@@ -114,8 +116,13 @@ def exec(self, keys: list[str], datum: Datum) -> Messages:
  # Forward payload if a training request is tagged
  if payload.header == Header.TRAIN_REQUEST:
  return Messages(Message(keys=keys, value=payload.to_json()))
-
- artifact_data = self.load_artifact(keys, payload)
+ artifact_data, payload = _load_artifact(
+ skeys=keys,
+ dkeys=[self.get_conf(payload.config_id).numalogic_conf.model.name],
+ payload=payload,
+ model_registry=self.model_registry,
+ load_latest=LOAD_LATEST,
+ )
 
  # TODO: revisit retraining logic
  # Send training request if artifact loading is not successful
@@ -165,47 +172,6 @@ def exec(self, keys: list[str], datum: Datum) -> Messages:
  )
  return Messages(Message(keys=keys, value=payload.to_json()))
 
- def load_artifact(self, keys: list[str], payload: StreamPayload) -> Optional[ArtifactData]:
- """
- Load inference artifact from the registry.
-
- Args:
- keys: List of keys
- payload: StreamPayload object
-
- Returns
- -------
- ArtifactData instance
- """
- _conf = self.get_conf(payload.config_id).numalogic_conf
- try:
- artifact_data = self.model_registry.load(
- skeys=keys,
- dkeys=[_conf.model.name],
- )
- except ModelKeyNotFound:
- _LOGGER.warning(
- "%s - Model key not found for Keys: %s, Metric: %s",
- payload.uuid,
- payload.composite_keys,
- payload.metrics,
- )
- return None
- except RedisRegistryError:
- _LOGGER.exception(
- "%s - Error while fetching inference artifact, Keys: %s, Metric: %s",
- payload.uuid,
- payload.composite_keys,
- payload.metrics,
- )
- return None
- _LOGGER.info(
- "%s - Loaded artifact data from %s",
- payload.uuid,
- artifact_data.extras.get("source"),
- )
- return artifact_data
-
  def is_model_stale(self, artifact_data: ArtifactData, payload: StreamPayload) -> bool:
  """
  Check if the inference artifact is stale.