Plugin/mlflow custom flavor

examples/hamilton_ui/run.py

@@ -41,6 +41,7 @@ def run(load_from_parquet: bool, username: str, project_id: int):
             "template": "machine_learning",
             "loading_data_from": "parquet" if load_from_parquet else "api",
             "TODO": "add_more_tags_to_find_your_run_later",
+            "custom_tag": "adding custom tag",
         },
     )
 

hamilton/plugins/mlflow_extensions.py

@@ -1,5 +1,6 @@
 import dataclasses
 import pathlib
+from types import ModuleType
 from typing import Any, Collection, Dict, Literal, Optional, Tuple, Type, Union
 
 try:
@@ -19,17 +20,17 @@ class MLFlowModelSaver(DataSaver):
     :param register_as: If not None, register the model under the specified name.
     :param flavor: Library format to save the model (sklearn, xgboost, etc.). Automatically inferred if None.
     :param run_id: Log model to a specific run. Leave to `None` if using the `MLFlowTracker`
-    :param kwargs: Arguments for `.log_model()`. Can be flavor-specific.
+    :param mlflow_kwargs: Arguments for `.log_model()`. Can be flavor-specific.
     """
 
     path: Union[str, pathlib.Path] = "model"
     register_as: Optional[str] = None
-    flavor: Optional[str] = None
+    flavor: Optional[Union[str, ModuleType]] = None
     run_id: Optional[str] = None
-    kwargs: Dict[str, Any] = None
+    mlflow_kwargs: Dict[str, Any] = None
 
     def __post_init__(self):
-        self.kwargs = self.kwargs if self.kwargs else {}
+        self.mlflow_kwargs = self.mlflow_kwargs if self.mlflow_kwargs else {}
 
     @classmethod
     def name(cls) -> str:
@@ -47,11 +48,16 @@ def save_data(self, data) -> Dict[str, Any]:
             # for example, extract `sklearn` from `sklearn.linear_model._base`
             flavor, _, _ = data.__module__.partition(".")
 
-        # retrieve the `mlflow.FLAVOR` submodule to `.log_model()`
-        try:
-            flavor_module = getattr(mlflow, flavor)
-        except ImportError:
-            raise ImportError(f"Flavor {flavor} is unsupported by MLFlow")
+        # pass flavor as module, this suports standard flavors `like mlflow.sklearn`
+        # but also supports custom flavors like `my_flavor` that implements save(), log(), load()
+        if isinstance(flavor, ModuleType):
+            flavor_module = flavor
+        else:
+            # retrieve the `mlflow.FLAVOR` submodule to `.log_model()`
+            try:
+                flavor_module = getattr(mlflow, flavor)
+            except ImportError:
+                raise ImportError(f"Flavor {flavor} is unsupported by MLFlow")
 
         # handle `run_id` and active run conflicts
         if mlflow.active_run() and self.run_id:
@@ -63,11 +69,11 @@ def save_data(self, data) -> Dict[str, Any]:
 
         # save to active run
         if mlflow.active_run():
-            model_info = flavor_module.log_model(data, self.path, **self.kwargs)
+            model_info = flavor_module.log_model(data, self.path, **self.mlflow_kwargs)
         # create a run with `run_id` and save to it
         else:
             with mlflow.start_run(run_id=self.run_id):
-                model_info = flavor_module.log_model(data, self.path, **self.kwargs)
+                model_info = flavor_module.log_model(data, self.path, **self.mlflow_kwargs)
 
         # create metadata from ModelInfo object
         metadata = {k.strip("_"): v for k, v in model_info.__dict__.items()}
@@ -99,7 +105,7 @@ class MLFlowModelLoader(DataLoader):
     :param version: Version of the registered model. Can pass as string `v1` or integer `1`
     :param version_alias: Version alias of the registered model. Specify either this or `version`
     :param flavor: Library format to load the model (sklearn, xgboost, etc.). Automatically inferred if None.
-    :param kwargs: Arguments for `.load_model()`. Can be flavor-specific.
+    :param mlflow_kwargs: Arguments for `.load_model()`. Can be flavor-specific.
     """
 
     model_uri: Optional[str] = None
@@ -109,15 +115,15 @@ class MLFlowModelLoader(DataLoader):
     model_name: Optional[str] = None
     version: Optional[Union[str, int]] = None
     version_alias: Optional[str] = None
-    flavor: Optional[str] = None
-    kwargs: Dict[str, Any] = None
+    flavor: Optional[Union[ModuleType, str]] = None
+    mlflow_kwargs: Dict[str, Any] = None
 
     # __post_init__ is required to set kwargs as empty dict because
     # can't set: kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)
     # otherwise raises `InvalidDecoratorException` because materializer factory check
     # for all params being set and `kwargs` would be unset until instantiation.
     def __post_init__(self):
-        self.kwargs = self.kwargs if self.kwargs else {}
+        self.mlflow_kwargs = self.mlflow_kwargs if self.mlflow_kwargs else {}
 
         if self.model_uri:
             return
@@ -163,13 +169,18 @@ def load_data(self, type_: Type) -> Tuple[Any, Dict[str, Any]]:
             except StopIteration:
                 flavor = "pyfunc"
 
-        # retrieve the `mlflow.FLAVOR` submodule to `.log_model()`
-        try:
-            flavor_module = getattr(mlflow, flavor)
-        except ImportError:
-            raise ImportError(f"Flavor {flavor} is unsupported by MLFlow")
+        # pass flavor as module, this suports standard flavors `like mlflow.sklearn`
+        # but also supports custom flavors like `my_flavor` that implements save(), log(), load()
+        if isinstance(flavor, ModuleType):
+            flavor_module = flavor
+        else:
+            # retrieve the `mlflow.FLAVOR` submodule to `.log_model()`
+            try:
+                flavor_module = getattr(mlflow, flavor)
+            except ImportError:
+                raise ImportError(f"Flavor {flavor} is unsupported by MLFlow")
 
-        model = flavor_module.load_model(model_uri=self.model_uri)
+        model = flavor_module.load_model(model_uri=self.model_uri, **self.mlflow_kwargs)
         return model, metadata
 
 

tests/plugins/test_mlflow_extension.py

@@ -17,6 +17,8 @@
 from sklearn.base import BaseEstimator
 from sklearn.linear_model import LinearRegression
 
+from hamilton.io.materialization import from_, to
+
 # TODO move these tests to `plugin_tests` because the required read-writes can get
 # complicated and tests are time consuming.
 
@@ -145,22 +147,71 @@ def test_mlflow_load_registry_model(fitted_sklearn_model: BaseEstimator, tmp_pat
     assert coefficients_are_equal(fitted_sklearn_model, loaded_model)
 
 
-def test_mlflow_infer_flavor(fitted_sklearn_model: BaseEstimator, tmp_path: Path):
+def test_mlflow_infer_flavor(fitted_sklearn_model: BaseEstimator):
     saver = MLFlowModelSaver(path="model")
 
     metadata = saver.save_data(fitted_sklearn_model)
 
     assert "sklearn" in metadata["flavors"].keys()
 
 
+def test_mlflow_specify_flavor_using_module(fitted_sklearn_model: BaseEstimator, tmp_path):
+    model_path = tmp_path / "sklearn_model"
+    saver = MLFlowModelSaver(flavor=mlflow.sklearn)
+
+    mlflow.set_tracking_uri(model_path.as_uri())
+    with mlflow.start_run():
+        # save model
+        metadata = saver.save_data(fitted_sklearn_model)
+    # reload model
+    loaded_model = mlflow.sklearn.load_model(metadata["model_uri"])
+
+    assert coefficients_are_equal(fitted_sklearn_model, loaded_model)
+
+
 def test_mlflow_handle_saver_kwargs():
     path = "tmp/path"
     flavor = "sklearn"
-    saver = MLFlowModelSaver(path=path, flavor=flavor, kwargs=dict(unknown_kwarg=True))
+    saver = MLFlowModelSaver(path=path, flavor=flavor, mlflow_kwargs=dict(unknown_kwarg=True))
 
     assert saver.path == path
     assert saver.flavor == flavor
-    assert saver.kwargs.get("unknown_kwarg") is True
+    assert saver.mlflow_kwargs.get("unknown_kwarg") is True
+
+
+def test_io_to_mlflow_handle_saver_kwargs():
+    path = "tmp/path"
+    flavor = "sklearn"
+    id = "saver_id"
+    dependencies = ["tmp_node"]
+    saver = to.mlflow(
+        path=path,
+        flavor=flavor,
+        id=id,
+        dependencies=dependencies,
+        mlflow_kwargs=dict(unknown_kwarg=True),
+    )
+    mlflow_saver = vars(saver)["data_saver_kwargs"]
+
+    assert mlflow_saver["path"].value == path
+    assert mlflow_saver["flavor"].value == flavor
+    assert mlflow_saver["mlflow_kwargs"].value.get("unknown_kwarg") is True
+
+
+def test_io_to_mlflow_handle_loader_kwargs():
+    path = "tmp/path"
+    flavor = "sklearn"
+    loader = from_.mlflow(
+        target="test_node",
+        model_uri=path,
+        flavor=flavor,
+        mlflow_kwargs=dict(unknown_kwarg=True),
+    )
+    mlflow_loader = vars(loader)["data_loader_kwargs"]
+
+    assert mlflow_loader["model_uri"].value == path
+    assert mlflow_loader["flavor"].value == flavor
+    assert mlflow_loader["mlflow_kwargs"].value.get("unknown_kwarg") is True
 
 
 def test_mlflow_registered_model_metadata(fitted_sklearn_model: BaseEstimator, tmp_path: Path):