chore!: make mlflow as an optional dependency

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

numalogic/registry/__init__.py

@@ -1,5 +1,8 @@
 from numalogic.registry.artifact import ArtifactManager
-from numalogic.registry.mlflow_registry import MLflowRegistrar
 
-
-__all__ = ["ArtifactManager", "MLflowRegistrar"]
+try:
+ from numalogic.registry.mlflow_registry import MLflowRegistrar
+except ImportError:
+ __all__ = ["ArtifactManager"]
+else:
+ __all__ = ["ArtifactManager", "MLflowRegistrar"]

numalogic/tests/registry/_mlflow_utils.py

@@ -110,6 +110,7 @@ def mock_log_model_pytorch(*_, **__):
  signature_dict=None,
  utc_time_created="2022-05-23 22:35:59.557372",
  mlflow_version="1.26.0",
+ signature=None
  )
 
 
@@ -133,6 +134,7 @@ def mock_log_model_sklearn(*_, **__):
  signature_dict=None,
  utc_time_created="2022-05-23 22:35:59.557372",
  mlflow_version="1.26.0",
+ signature=None
  )
 
 

numalogic/tools/types.py

@@ -1,15 +1,12 @@
 from typing import Union, Dict, NewType, TypeVar, Sequence, Optional
 
-from mlflow.entities.model_registry import ModelVersion
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.pipeline import Pipeline
 from torch import nn
 
 Artifact = NewType("Artifact", Union[nn.Module, BaseEstimator, TransformerMixin, Pipeline])
 ArtifactDict = NewType(
  "ArtifactDict",
- Optional[
- Dict[str, Union[Sequence[Artifact], Dict[str, Artifact], Artifact, None, ModelVersion]]
- ],
+ Optional[Dict[str, Union[Sequence[Artifact], Dict[str, Artifact], Artifact, None]]],
 )
 AutoencoderModel = TypeVar("AutoencoderModel", bound="TorchAE")