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()