fix: training only when model loading from registry source (#159)

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

Author: s0nicboOm

numaprom/__init__.py

@@ -23,6 +23,8 @@ def emit(self, record):
 
 def __get_logger() -> logger:
     # Collect logs from logging library
+    logging.getLogger("pytorch_lightning.utilities.rank_zero").setLevel(logging.WARNING)
+    logging.getLogger("pytorch_lightning.accelerators.cuda").setLevel(logging.WARNING)
     logging.basicConfig(handlers=[InterceptHandler()], level=0)
     logger.remove()
 

numaprom/default-configs/numalogic_config.yaml

@@ -16,6 +16,8 @@ preprocess:
   - name: "StandardScaler"
 threshold:
   name: "StdDevThreshold"
+  conf:
+    min_threshold: 0.01
 postprocess:
   name: "TanhNorm"
   stateful: false

numaprom/udf/inference.py

@@ -97,8 +97,17 @@ def inference(_: list[str], datum: Datum) -> bytes:
         payload.set_status(Status.ARTIFACT_NOT_FOUND)
         return orjson.dumps(payload, option=orjson.OPT_SERIALIZE_NUMPY)
 
-    # Check if current model is stale
-    if RedisRegistry.is_artifact_stale(artifact_data, int(metric_config.retrain_freq_hr)):
+    LOGGER.info(
+        "{uuid} - Loaded artifact data from {source} ",
+        uuid=payload.uuid,
+        source=artifact_data.extras.get("source"),
+    )
+
+    # Check if current model is stale and source is 'registry'
+    if (
+        RedisRegistry.is_artifact_stale(artifact_data, int(metric_config.retrain_freq_hr))
+        and artifact_data.extras.get("source") == "registry"
+    ):
         payload.set_header(Header.MODEL_STALE)
 
     # Generate predictions