fix: stale check; conf lazy imports (#178)

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

.pre-commit-config.yaml

@@ -8,11 +8,9 @@ repos:
  language_version: python3.9
  args: [--config=pyproject.toml, --diff, --color ]
 - repo: https://github.com/charliermarsh/ruff-pre-commit
- # Ruff version.
  rev: 'v0.0.264'
  hooks:
  - id: ruff
- args: [ --fix, --exit-non-zero-on-fix ]
 - repo: https://github.com/adamchainz/blacken-docs
  rev: "1.13.0"
  hooks:

numalogic/config/factory.py

@@ -10,23 +10,7 @@
 # limitations under the License.
 from typing import Union
 
-from sklearn.preprocessing import StandardScaler, MinMaxScaler, MaxAbsScaler, RobustScaler
-
 from numalogic.config._config import ModelInfo, RegistryInfo
-from numalogic.models.autoencoder.variants import (
- VanillaAE,
- SparseVanillaAE,
- Conv1dAE,
- SparseConv1dAE,
- LSTMAE,
- SparseLSTMAE,
- TransformerAE,
- SparseTransformerAE,
-)
-from numalogic.models.threshold import StdDevThreshold, StaticThreshold, SigmoidThreshold
-from numalogic.postprocess import TanhNorm, ExpMovingAverage
-from numalogic.preprocess import LogTransformer, StaticPowerTransformer, TanhScaler
-from numalogic.registry import MLflowRegistry, RedisRegistry
 from numalogic.tools.exceptions import UnknownConfigArgsError
 
 
@@ -52,6 +36,9 @@ def get_cls(self, object_info: Union[ModelInfo, RegistryInfo]):
 
 
 class PreprocessFactory(_ObjectFactory):
+ from sklearn.preprocessing import StandardScaler, MinMaxScaler, MaxAbsScaler, RobustScaler
+ from numalogic.preprocess import LogTransformer, StaticPowerTransformer, TanhScaler
+
  _CLS_MAP = {
  "StandardScaler": StandardScaler,
  "MinMaxScaler": MinMaxScaler,
@@ -64,10 +51,14 @@ class PreprocessFactory(_ObjectFactory):
 
 
 class PostprocessFactory(_ObjectFactory):
+ from numalogic.postprocess import TanhNorm, ExpMovingAverage
+
  _CLS_MAP = {"TanhNorm": TanhNorm, "ExpMovingAverage": ExpMovingAverage}
 
 
 class ThresholdFactory(_ObjectFactory):
+ from numalogic.models.threshold import StdDevThreshold, StaticThreshold, SigmoidThreshold
+
  _CLS_MAP = {
  "StdDevThreshold": StdDevThreshold,
  "StaticThreshold": StaticThreshold,
@@ -76,6 +67,17 @@ class ThresholdFactory(_ObjectFactory):
 
 
 class ModelFactory(_ObjectFactory):
+ from numalogic.models.autoencoder.variants import (
+ VanillaAE,
+ SparseVanillaAE,
+ Conv1dAE,
+ SparseConv1dAE,
+ LSTMAE,
+ SparseLSTMAE,
+ TransformerAE,
+ SparseTransformerAE,
+ )
+
  _CLS_MAP = {
  "VanillaAE": VanillaAE,
  "SparseVanillaAE": SparseVanillaAE,
@@ -89,7 +91,9 @@ class ModelFactory(_ObjectFactory):
 
 
 class RegistryFactory(_ObjectFactory):
+ import numalogic.registry as reg
+
  _CLS_MAP = {
- "RedisRegistry": RedisRegistry,
- "MLflowRegistry": MLflowRegistry,
+ "RedisRegistry": getattr(reg, "RedisRegistry"),
+ "MLflowRegistry": getattr(reg, "MLflowRegistry"),
  }

numalogic/registry/artifact.py

@@ -13,7 +13,7 @@
 from dataclasses import dataclass
 from typing import Any, Generic, TypeVar
 
-from numalogic.tools.types import artifact_t, KEYS, META_T, EXTRA_T
+from numalogic.tools.types import artifact_t, KEYS, META_T, META_VT, EXTRA_T
 
 
 @dataclass
@@ -54,7 +54,7 @@ def load(
  """
  raise NotImplementedError("Please implement this method!")
 
- def save(self, skeys: KEYS, dkeys: KEYS, artifact: artifact_t, **metadata: META_T) -> Any:
+ def save(self, skeys: KEYS, dkeys: KEYS, artifact: artifact_t, **metadata: META_VT) -> Any:
  r"""
  Saves the artifact into mlflow registry and updates version.
  Args:
@@ -75,6 +75,18 @@ def delete(self, skeys: KEYS, dkeys: KEYS, version: str) -> None:
  """
  raise NotImplementedError("Please implement this method!")
 
+ @staticmethod
+ def is_artifact_stale(artifact_data: ArtifactData, freq_hr: int) -> bool:
+ """
+ Returns whether the given artifact is stale or not, i.e. if
+ more time has elasped since it was last retrained.
+ Args:
+ artifact_data: ArtifactData object to look into
+ freq_hr: Frequency of retraining in hours
+
+ """
+ raise NotImplementedError("Please implement this method!")
+
  @staticmethod
  def construct_key(skeys: KEYS, dkeys: KEYS) -> str:
  """

numalogic/registry/mlflow_registry.py

@@ -11,6 +11,7 @@
 
 
 import logging
+from datetime import datetime, timedelta
 from enum import Enum
 from typing import Optional, Any
 
@@ -24,7 +25,7 @@
 from numalogic.registry import ArtifactManager, ArtifactData
 from numalogic.registry.artifact import ArtifactCache
 from numalogic.tools.exceptions import ModelVersionError
-from numalogic.tools.types import artifact_t, KEYS, META_T
+from numalogic.tools.types import artifact_t, KEYS, META_VT
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -183,7 +184,7 @@ def save(
  dkeys: KEYS,
  artifact: artifact_t,
  run_id: str = None,
- **metadata: META_T,
+ **metadata: META_VT,
  ) -> Optional[ModelVersion]:
  """
  Saves the artifact into mlflow registry and updates version.
@@ -213,6 +214,20 @@ def save(
  finally:
  mlflow.end_run()
 
+ @staticmethod
+ def is_artifact_stale(artifact_data: ArtifactData, freq_hr: int) -> bool:
+ """
+ Returns whether the given artifact is stale or not, i.e. if
+ more time has elasped since it was last retrained.
+ Args:
+ artifact_data: ArtifactData object to look into
+ freq_hr: Frequency of retraining in hours
+
+ """
+ date_updated = artifact_data.extras["last_updated_timestamp"] / 1000
+ stale_date = (datetime.now() - timedelta(hours=freq_hr)).timestamp()
+ return date_updated < stale_date
+
  def delete(self, skeys: KEYS, dkeys: KEYS, version: str) -> None:
  """
  Deletes the artifact with a specified version from mlflow registry.

numalogic/registry/redis_registry.py

@@ -1,13 +1,14 @@
 import logging
 import time
+from datetime import datetime, timedelta
 from typing import Optional
 
 from redis.exceptions import RedisError
 
 from numalogic.registry import ArtifactManager, ArtifactData
 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
+from numalogic.tools.types import artifact_t, redis_client_t, KEYS, META_T, META_VT
 
 _LOGGER = logging.getLogger()
 
@@ -24,14 +25,13 @@ class RedisRegistry(ArtifactManager):
  >>> import redis
  >>> from numalogic.models.autoencoder.variants import VanillaAE
  >>> from numalogic.registry.redis_registry import RedisRegistry
+ >>> ...
  >>> r = redis.StrictRedis(host='127.0.0.1', port=6379)
- >>> cli = r.client()
- >>> registry = RedisRegistry(client=cli)
- >>> skeys = ['c', 'a']
- >>> dkeys = ['d', 'a']
- >>> model = VanillaAE(10)
+ >>> registry = RedisRegistry(client=r)
+ >>> skeys, dkeys = ("mymetric", "ae"), ("vanilla", "seq10")
+ >>> model = VanillaAE(seq_len=10)
  >>> registry.save(skeys, dkeys, artifact=model, **{'lr': 0.01})
- >>> registry.load(skeys, dkeys)
+ >>> loaded_artifact = registry.load(skeys, dkeys)
  """
 
  __slots__ = ("client", "ttl")
@@ -177,7 +177,7 @@ def save(
  skeys: KEYS,
  dkeys: KEYS,
  artifact: artifact_t,
- **metadata: META_T,
+ **metadata: META_VT,
  ) -> Optional[str]:
  """
  Saves the artifact into redis registry and updates version.
@@ -229,3 +229,20 @@ def delete(self, skeys: KEYS, dkeys: KEYS, version: str) -> None:
  )
  except RedisError as err:
  raise RedisRegistryError(f"{err.__class__.__name__} raised") from err
+
+ @staticmethod
+ def is_artifact_stale(artifact_data: ArtifactData, freq_hr: int) -> bool:
+ """
+ Returns whether the given artifact is stale or not, i.e. if
+ more time has elasped since it was last retrained.
+ Args:
+ artifact_data: ArtifactData object to look into
+ freq_hr: Frequency of retraining in hours
+
+ """
+ try:
+ artifact_ts = float(artifact_data.extras["timestamp"])
+ except (KeyError, TypeError) as err:
+ raise RedisRegistryError("Error fetching timestamp information") from err
+ stale_ts = (datetime.now() - timedelta(hours=freq_hr)).timestamp()
+ return stale_ts > artifact_ts

numalogic/tools/types.py

@@ -19,6 +19,7 @@
 
 artifact_t = TypeVar("artifact_t", bound=Union[nn.Module, BaseEstimator])
 META_T = TypeVar("META_T", bound=dict[str, Union[str, list, dict]])
+META_VT = TypeVar("META_VT", bound=Union[str, list, dict])
 EXTRA_T = TypeVar("EXTRA_T", bound=dict[str, Union[str, list, dict]])
 redis_client_t = TypeVar("redis_client_t", bound=AbstractRedis, covariant=True)
 KEYS = TypeVar("KEYS", bound=Sequence[str], covariant=True)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "numalogic"
-version = "0.4.dev3"
+version = "0.4.dev4"
 description = "Collection of operational Machine Learning models and tools."
 authors = ["Numalogic Developers"]
 packages = [{ include = "numalogic" }]
@@ -26,7 +26,7 @@ homepage = "https://numalogic.numaproj.io/"
 python = ">=3.9, <3.11"
 numpy = "^1.23"
 pandas = "^2.0"
-scikit-learn = "^1.0"
+scikit-learn = "^1.2"
 mlflow-skinny = { version = ">2.0, <2.3", optional = true }
 protobuf = "~3.20" # needed to support pytorch-lightning
 omegaconf = "^2.3.0"
@@ -55,6 +55,7 @@ torchinfo = "^1.7.2"
 ruff = "^0.0.264"
 pre-commit = "^3.3.1"
 fakeredis = "^2.11.2"
+freezegun = "^1.2.2"
 
 [tool.poetry.group.jupyter]
 optional = true

tests/registry/test_mlflow_registry.py

@@ -2,6 +2,7 @@
 from contextlib import contextmanager
 from unittest.mock import patch, Mock
 
+from freezegun import freeze_time
 from mlflow import ActiveRun
 from mlflow.exceptions import RestException
 from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST, ErrorCode, RESOURCE_LIMIT_EXCEEDED
@@ -290,6 +291,37 @@ def test_load_other_mlflow_err(self):
  dkeys = self.dkeys
  self.assertIsNone(ml.load(skeys=skeys, dkeys=dkeys))
 
+ @patch("mlflow.pytorch.log_model", mock_log_model_pytorch())
+ @patch("mlflow.start_run", Mock(return_value=ActiveRun(return_pytorch_rundata_dict())))
+ @patch("mlflow.active_run", Mock(return_value=return_pytorch_rundata_dict()))
+ @patch("mlflow.log_params", {"lr": 0.01})
+ @patch("mlflow.tracking.MlflowClient.transition_model_version_stage", mock_transition_stage)
+ @patch("mlflow.tracking.MlflowClient.get_latest_versions", mock_get_model_version)
+ @patch("mlflow.pytorch.load_model", Mock(return_value=VanillaAE(10)))
+ @patch("mlflow.tracking.MlflowClient.get_run", Mock(return_value=return_pytorch_rundata_dict()))
+ def test_is_model_stale_true(self):
+ model = self.model
+ ml = MLflowRegistry(TRACKING_URI, artifact_type="pytorch")
+ ml.save(skeys=self.skeys, dkeys=self.dkeys, artifact=model, **{"lr": 0.01})
+ data = ml.load(skeys=self.skeys, dkeys=self.dkeys)
+ self.assertTrue(ml.is_artifact_stale(data, 12))
+
+ @patch("mlflow.pytorch.log_model", mock_log_model_pytorch())
+ @patch("mlflow.start_run", Mock(return_value=ActiveRun(return_pytorch_rundata_dict())))
+ @patch("mlflow.active_run", Mock(return_value=return_pytorch_rundata_dict()))
+ @patch("mlflow.log_params", {"lr": 0.01})
+ @patch("mlflow.tracking.MlflowClient.transition_model_version_stage", mock_transition_stage)
+ @patch("mlflow.tracking.MlflowClient.get_latest_versions", mock_get_model_version)
+ @patch("mlflow.pytorch.load_model", Mock(return_value=VanillaAE(10)))
+ @patch("mlflow.tracking.MlflowClient.get_run", Mock(return_value=return_pytorch_rundata_dict()))
+ def test_is_model_stale_false(self):
+ model = self.model
+ ml = MLflowRegistry(TRACKING_URI, artifact_type="pytorch")
+ ml.save(skeys=self.skeys, dkeys=self.dkeys, artifact=model, **{"lr": 0.01})
+ data = ml.load(skeys=self.skeys, dkeys=self.dkeys)
+ with freeze_time("2022-05-24 10:30:00"):
+ self.assertFalse(ml.is_artifact_stale(data, 12))
+
 
 if __name__ == "__main__":
  unittest.main()