fix: loading secondary artifacts (#16)

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

numalogic/registry/mlflow_registry.py

@@ -155,12 +155,14 @@ def load(
  run_data = self.client.get_run(run_id).data.to_dictionary()
  if run_data["params"]:
  data = run_data["params"]
- secondary_artifacts = pickle.loads(
- codecs.decode(data["secondary_artifacts"].encode(), "base64")
- )
- _LOGGER.info("Successfully loaded secondary_artifacts from Mlflow")
- metadata = pickle.loads(codecs.decode(data["metadata"].encode(), "base64"))
- _LOGGER.info("Successfully loaded model metadata from Mlflow")
+ if "secondary_artifacts" in data:
+ secondary_artifacts = pickle.loads(
+ codecs.decode(data["secondary_artifacts"].encode(), "base64")
+ )
+ _LOGGER.info("Successfully loaded secondary_artifacts from Mlflow")
+ if "metadata" in data:
+ metadata = pickle.loads(codecs.decode(data["metadata"].encode(), "base64"))
+ _LOGGER.info("Successfully loaded model metadata from Mlflow")
  return self.__as_dict(model, secondary_artifacts, metadata, model_properties)
  except Exception as ex:
  _LOGGER.exception("Error when loading a model with key: %s: %r", model_key, ex)

numalogic/tests/registry/test_mlflow_registry.py

@@ -64,8 +64,7 @@ def test_insert_model(self):
  dkeys=dkeys,
  primary_artifact=model,
  secondary_artifacts=[make_pipeline(return_scaler)],
- artifact_state_dict=model.state_dict(),
- models_to_retain=2,
+ **model.state_dict(),
  )
  mock_status = "READY"
  self.assertEqual(mock_status, status.status)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "numalogic"
-version = "0.2.3"
+version = "0.2.4"
 description = "Collection of operational Machine Learning models and tools."
 authors = ["Numalogic Developers"]
 packages = [{ include = "numalogic" }]