Repository: Galileo-Galilei/kedro-mlflow (python)

Issue:
Add `metadata` parameters to kedro-mlflow datasets
<!-- **_If you like the repo, please give it a :star:_** -->

Hi,
First of all, thanks a lot for this amazing plugin!

## Description
<!-- A clear and concise description of what you want to achieve. An image or a code example is worth thousand words! -->

It seems that kedro assumes its datasets can all take a `metadata` parameters. In previous versions of the kedro doc, I remember reading that `metadata` was meant to be exploited by orchestration-related kedro plugins, I could not find this statement anymore on the latest doc. 

However, there is now a hook documented that aims at showing how to access metadata: https://docs.kedro.org/en/stable/hooks/common_use_cases.html#use-hooks-to-read-metadata-from-datacatalog

## Context
<!-- Why is this change important to you? How would you use it? How can it benefit other users? -->

I am developing a kedro plugin for the dagster orchestrator: https://github.com/gtauzin/kedro-dagster and I would like to be able to associate `metadata` to kedro-mlflow datasets. I make use of the metadata arguments of datasets to display a description or other metadata on the Dagster UI.


Pull request:
Enable all datasets to take a `metadata` parameter
## Description

Fixes #625.
Make all kedro-mlflow datasets take a `metadata` parameter, thus aligning with kedro datasets.

## Development notes

Added a `metadata` parameters to all datasets and tested that it is properly set and it does not appear when `dataset._describe` is called.

## Checklist

- [x] Read the [contributing](https://github.com/Galileo-Galilei/kedro-mlflow/blob/master/CONTRIBUTING.md) guidelines
- [x] Open this PR as a 'Draft Pull Request' if it is work-in-progress
- [x] Update the documentation to reflect the code changes
- [x] Add a description of this change and add your name to the list of supporting contributions in the [`CHANGELOG.md`](https://github.com/Galileo-Galilei/kedro-mlflow/blob/master/CHANGELOG.md) file. Please respect [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) guidelines.
- [x] Add tests to cover your changes

## Notice

- [x] I acknowledge and agree that, by checking this box and clicking "Submit Pull Request":

- I submit this contribution under the [Apache 2.0 license](https://www.apache.org/licenses/LICENSE-2.0.txt) and represent that I am entitled to do so on behalf of myself, my employer, or relevant third parties, as applicable.
- I certify that (a) this contribution is my original creation and / or (b) to the extent it is not my original creation, I am authorised to submit this contribution on behalf of the original creator(s) or their licensees.
- I certify that the use of this contribution as authorised by the Apache 2.0 license does not violate the intellectual property rights of anyone else.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7d9ec86..e3e8b55 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+
+-   :sparkles: All datasets can now take a metadata parameter ([#625](https://github.com/Galileo-Galilei/kedro-mlflow/issues/625), [#633](https://github.com/Galileo-Galilei/kedro-mlflow/pull/633))
+
 ## [0.14.2] - 2025-02-16
 
 ### Fixed
diff --git a/kedro_mlflow/io/artifacts/mlflow_artifact_dataset.py b/kedro_mlflow/io/artifacts/mlflow_artifact_dataset.py
index 67e1fb9..083bcb4 100644
--- a/kedro_mlflow/io/artifacts/mlflow_artifact_dataset.py
+++ b/kedro_mlflow/io/artifacts/mlflow_artifact_dataset.py
@@ -19,6 +19,7 @@ class MlflowArtifactDataset(AbstractVersionedDataset):
         run_id: str = None,
         artifact_path: str = None,
         credentials: Dict[str, Any] = None,
+        metadata: Dict[str, Any] | None = None,
     ):
         dataset_obj, dataset_args = parse_dataset_definition(config=dataset)
 
@@ -27,11 +28,12 @@ class MlflowArtifactDataset(AbstractVersionedDataset):
         # instead and since we can't modify the core package,
         # we create a subclass which inherits dynamically from the dataset class
         class MlflowArtifactDatasetChildren(dataset_obj):
-            def __init__(self, run_id, artifact_path):
+            def __init__(self, run_id, artifact_path, metadata):
                 super().__init__(**dataset_args)
                 self.run_id = run_id
                 self.artifact_path = artifact_path
                 self._logging_activated = True
+                self.metadata = metadata
 
             @property
             def _logging_activated(self):
@@ -147,7 +149,9 @@ class MlflowArtifactDataset(AbstractVersionedDataset):
         )
 
         mlflow_dataset_instance = MlflowArtifactDatasetChildren(
-            run_id=run_id, artifact_path=artifact_path
+            run_id=run_id,
+            artifact_path=artifact_path,
+            metadata=metadata,
         )
         return mlflow_dataset_instance
 
diff --git a/kedro_mlflow/io/metrics/mlflow_abstract_metric_dataset.py b/kedro_mlflow/io/metrics/mlflow_abstract_metric_dataset.py
index 6ac4b63..dca912e 100644
--- a/kedro_mlflow/io/metrics/mlflow_abstract_metric_dataset.py
+++ b/kedro_mlflow/io/metrics/mlflow_abstract_metric_dataset.py
@@ -12,6 +12,7 @@ class MlflowAbstractMetricDataset(AbstractDataset):
         run_id: str = None,
         load_args: Dict[str, Any] = None,
         save_args: Dict[str, Any] = None,
+        metadata: Dict[str, Any] | None = None,
     ):
         """Initialise MlflowMetricsHistoryDataset.
 
@@ -24,6 +25,7 @@ class MlflowAbstractMetricDataset(AbstractDataset):
         self._load_args = load_args or {}
         self._save_args = save_args or {}
         self._logging_activated = True  # by default, logging is activated!
+        self.metadata = metadata
 
     @property
     def run_id(self) -> Union[str, None]:
diff --git a/kedro_mlflow/io/metrics/mlflow_metric_dataset.py b/kedro_mlflow/io/metrics/mlflow_metric_dataset.py
index 23f5a85..a2f46f3 100644
--- a/kedro_mlflow/io/metrics/mlflow_metric_dataset.py
+++ b/kedro_mlflow/io/metrics/mlflow_metric_dataset.py
@@ -18,13 +18,14 @@ class MlflowMetricDataset(MlflowAbstractMetricDataset):
         run_id: str = None,
         load_args: Dict[str, Any] = None,
         save_args: Dict[str, Any] = None,
+        metadata: Dict[str, Any] | None = None,
     ):
         """Initialise MlflowMetricDataset.
         Args:
             run_id (str): The ID of the mlflow run where the metric should be logged
         """
 
-        super().__init__(key, run_id, load_args, save_args)
+        super().__init__(key, run_id, load_args, save_args, metadata)
 
         # We add an extra argument mode="overwrite" / "append" to enable logging update an existing metric
         # this is not an offical mlflow argument for log_metric, so we separate it from the others
diff --git a/kedro_mlflow/io/metrics/mlflow_metric_history_dataset.py b/kedro_mlflow/io/metrics/mlflow_metric_history_dataset.py
index 48fb734..39f4355 100644
--- a/kedro_mlflow/io/metrics/mlflow_metric_history_dataset.py
+++ b/kedro_mlflow/io/metrics/mlflow_metric_history_dataset.py
@@ -14,13 +14,14 @@ class MlflowMetricHistoryDataset(MlflowAbstractMetricDataset):
         run_id: str = None,
         load_args: Dict[str, Any] = None,
         save_args: Dict[str, Any] = None,
+        metadata: Dict[str, Any] | None = None,
     ):
         """Initialise MlflowMetricDataset.
         Args:
             run_id (str): The ID of the mlflow run where the metric should be logged
         """
 
-        super().__init__(key, run_id, load_args, save_args)
+        super().__init__(key, run_id, load_args, save_args, metadata)
 
     def _load(self):
         self._validate_run_id()
diff --git a/kedro_mlflow/io/metrics/mlflow_metrics_history_dataset.py b/kedro_mlflow/io/metrics/mlflow_metrics_history_dataset.py
index 6e9e84a..874c861 100644
--- a/kedro_mlflow/io/metrics/mlflow_metrics_history_dataset.py
+++ b/kedro_mlflow/io/metrics/mlflow_metrics_history_dataset.py
@@ -18,6 +18,7 @@ class MlflowMetricsHistoryDataset(AbstractDataset):
         self,
         run_id: str = None,
         prefix: Optional[str] = None,
+        metadata: Dict[str, Any] | None = None,
     ):
         """Initialise MlflowMetricsHistoryDataset.
 
@@ -28,6 +29,7 @@ class MlflowMetricsHistoryDataset(AbstractDataset):
         self._prefix = prefix
         self.run_id = run_id
         self._logging_activated = True  # by default, logging is activated!
+        self.metadata = metadata
 
     @property
     def run_id(self):
diff --git a/kedro_mlflow/io/models/mlflow_abstract_model_dataset.py b/kedro_mlflow/io/models/mlflow_abstract_model_dataset.py
index 13326da..ba4ac6f 100644
--- a/kedro_mlflow/io/models/mlflow_abstract_model_dataset.py
+++ b/kedro_mlflow/io/models/mlflow_abstract_model_dataset.py
@@ -20,6 +20,7 @@ class MlflowAbstractModelDataSet(AbstractVersionedDataset):
         load_args: Dict[str, Any] = None,
         save_args: Dict[str, Any] = None,
         version: Version = None,
+        metadata: Dict[str, Any] | None = None,
     ) -> None:
         """Initialize the Kedro MlflowAbstractModelDataSet.
 
@@ -39,6 +40,8 @@ class MlflowAbstractModelDataSet(AbstractVersionedDataset):
             save_args (Dict[str, Any], optional): Arguments to `log_model`
                 function from specified `flavor`. Defaults to {}.
             version (Version, optional): Specific version to load.
+            metadata: Any arbitrary metadata.
+                This is ignored by Kedro, but may be consumed by users or external plugins.
 
         Raises:
             DatasetError: When passed `flavor` does not exist.
@@ -61,6 +64,7 @@ class MlflowAbstractModelDataSet(AbstractVersionedDataset):
 
         self._load_args = load_args or {}
         self._save_args = save_args or {}
+        self.metadata = metadata
 
         try:
             self._mlflow_model_module
diff --git a/kedro_mlflow/io/models/mlflow_model_local_filesystem_dataset.py b/kedro_mlflow/io/models/mlflow_model_local_filesystem_dataset.py
index 5efe285..616cf72 100644
--- a/kedro_mlflow/io/models/mlflow_model_local_filesystem_dataset.py
+++ b/kedro_mlflow/io/models/mlflow_model_local_filesystem_dataset.py
@@ -21,6 +21,7 @@ class MlflowModelLocalFileSystemDataset(MlflowAbstractModelDataSet):
         save_args: Dict[str, Any] = None,
         log_args: Dict[str, Any] = None,
         version: Version = None,
+        metadata: Dict[str, Any] | None = None,
     ) -> None:
         """Initialize the Kedro MlflowModelDataSet.
 
@@ -40,6 +41,9 @@ class MlflowModelLocalFileSystemDataset(MlflowAbstractModelDataSet):
             save_args (Dict[str, Any], optional): Arguments to `save_model`
                 function from specified `flavor`. Defaults to None.
             version (Version, optional): Kedro version to use. Defaults to None.
+            metadata: Any arbitrary metadata.
+                This is ignored by Kedro, but may be consumed by users or external plugins.
+
         Raises:
             DatasetError: When passed `flavor` does not exist.
         """
@@ -50,6 +54,7 @@ class MlflowModelLocalFileSystemDataset(MlflowAbstractModelDataSet):
             load_args=load_args,
             save_args=save_args,
             version=version,
+            metadata=metadata,
         )
 
     def _load(self) -> Any:
diff --git a/kedro_mlflow/io/models/mlflow_model_registry_dataset.py b/kedro_mlflow/io/models/mlflow_model_registry_dataset.py
index 11374b3..5ee7b4b 100644
--- a/kedro_mlflow/io/models/mlflow_model_registry_dataset.py
+++ b/kedro_mlflow/io/models/mlflow_model_registry_dataset.py
@@ -19,6 +19,7 @@ class MlflowModelRegistryDataset(MlflowAbstractModelDataSet):
         flavor: Optional[str] = "mlflow.pyfunc",
         pyfunc_workflow: Optional[str] = "python_model",
         load_args: Optional[Dict[str, Any]] = None,
+        metadata: Dict[str, Any] | None = None,
     ) -> None:
         """Initialize the Kedro MlflowModelRegistryDataset.
 
@@ -37,6 +38,8 @@ class MlflowModelRegistryDataset(MlflowAbstractModelDataSet):
                 See https://www.mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#workflows.
             load_args (Dict[str, Any], optional): Arguments to `load_model`
                 function from specified `flavor`. Defaults to None.
+            metadata: Any arbitrary metadata.
+                This is ignored by Kedro, but may be consumed by users or external plugins.
 
         Raises:
             DatasetError: When passed `flavor` does not exist.
@@ -48,6 +51,7 @@ class MlflowModelRegistryDataset(MlflowAbstractModelDataSet):
             load_args=load_args,
             save_args={},
             version=None,
+            metadata=metadata,
         )
 
         if alias is None and stage_or_version is None:
diff --git a/kedro_mlflow/io/models/mlflow_model_tracking_dataset.py b/kedro_mlflow/io/models/mlflow_model_tracking_dataset.py
index 0250685..127f2ff 100644
--- a/kedro_mlflow/io/models/mlflow_model_tracking_dataset.py
+++ b/kedro_mlflow/io/models/mlflow_model_tracking_dataset.py
@@ -19,6 +19,7 @@ class MlflowModelTrackingDataset(MlflowAbstractModelDataSet):
         pyfunc_workflow: Optional[str] = None,
         load_args: Optional[Dict[str, Any]] = None,
         save_args: Optional[Dict[str, Any]] = None,
+        metadata: Dict[str, Any] | None = None,
     ) -> None:
         """Initialize the Kedro MlflowModelDataSet.
 
@@ -40,6 +41,8 @@ class MlflowModelTrackingDataset(MlflowAbstractModelDataSet):
                 function from specified `flavor`. Defaults to None.
             save_args (Dict[str, Any], optional): Arguments to `log_model`
                 function from specified `flavor`. Defaults to None.
+            metadata: Any arbitrary metadata.
+                This is ignored by Kedro, but may be consumed by users or external plugins.
 
         Raises:
             DatasetError: When passed `flavor` does not exist.
@@ -51,6 +54,7 @@ class MlflowModelTrackingDataset(MlflowAbstractModelDataSet):
             load_args=load_args,
             save_args=save_args,
             version=None,
+            metadata=metadata,
         )
 
         self._run_id = run_id


Test patch:
diff --git a/tests/io/artifacts/test_mlflow_artifact_dataset.py b/tests/io/artifacts/test_mlflow_artifact_dataset.py
index 8d67ee6..1590731 100644
--- a/tests/io/artifacts/test_mlflow_artifact_dataset.py
+++ b/tests/io/artifacts/test_mlflow_artifact_dataset.py
@@ -376,3 +376,23 @@ def test_artifact_dataset_legacy_dataset(tmp_path, mlflow_client, df1):
     remote_path = (Path("artifact_dir") / filepath.name).as_posix()
     assert remote_path in run_artifacts
     assert df1.equals(mlflow_dataset.load())
+
+
+@pytest.mark.parametrize(
+    "metadata",
+    (
+        None,
+        {"description": "My awsome dataset"},
+        {"string": "bbb", "int": 0},
+    ),
+)
+def test_artifact_dataset_with_metadata(metadata):
+    mlflow_csv_dataset = MlflowArtifactDataset(
+        dataset=dict(type=CSVDataset, filepath="/my/file/path"),
+        metadata=metadata,
+    )
+
+    assert mlflow_csv_dataset.metadata == metadata
+
+    # Metadata should not show in _describe
+    assert "metadata" not in mlflow_csv_dataset._describe()
diff --git a/tests/io/metrics/test_mlflow_metric_dataset.py b/tests/io/metrics/test_mlflow_metric_dataset.py
index 7be0d47..5f596dd 100644
--- a/tests/io/metrics/test_mlflow_metric_dataset.py
+++ b/tests/io/metrics/test_mlflow_metric_dataset.py
@@ -193,3 +193,23 @@ def test_mlflow_metric_logging_deactivation_is_bool():
 
     with pytest.raises(ValueError, match="_logging_activated must be a boolean"):
         mlflow_metric_dataset._logging_activated = "hello"
+
+
+@pytest.mark.parametrize(
+    "metadata",
+    (
+        None,
+        {"description": "My awsome dataset"},
+        {"string": "bbb", "int": 0},
+    ),
+)
+def test_metric_dataset_with_metadata(tmp_path, metadata):
+    mlflow_metric_dataset = MlflowMetricDataset(
+        key="hello",
+        metadata=metadata,
+    )
+
+    assert mlflow_metric_dataset.metadata == metadata
+
+    # Metadata should not show in _describe
+    assert "metadata" not in mlflow_metric_dataset._describe()
diff --git a/tests/io/metrics/test_mlflow_metric_history_dataset.py b/tests/io/metrics/test_mlflow_metric_history_dataset.py
index fe16f1e..088d821 100644
--- a/tests/io/metrics/test_mlflow_metric_history_dataset.py
+++ b/tests/io/metrics/test_mlflow_metric_history_dataset.py
@@ -68,3 +68,23 @@ def test_mlflow_metric_history_dataset_logging_deactivation(mlflow_tracking_uri)
     with mlflow.start_run():
         metric_ds.save([0.1])
         assert metric_ds._exists() is False
+
+
+@pytest.mark.parametrize(
+    "metadata",
+    (
+        None,
+        {"description": "My awsome dataset"},
+        {"string": "bbb", "int": 0},
+    ),
+)
+def test_metric_history_dataset_with_metadata(tmp_path, metadata):
+    metric_ds = MlflowMetricHistoryDataset(
+        key="hello",
+        metadata=metadata,
+    )
+
+    assert metric_ds.metadata == metadata
+
+    # Metadata should not show in _describe
+    assert "metadata" not in metric_ds._describe()
diff --git a/tests/io/metrics/test_mlflow_metrics_dataset.py b/tests/io/metrics/test_mlflow_metrics_dataset.py
index d47120f..dd89f3c 100644
--- a/tests/io/metrics/test_mlflow_metrics_dataset.py
+++ b/tests/io/metrics/test_mlflow_metrics_dataset.py
@@ -189,3 +189,23 @@ def test_mlflow_metrics_logging_deactivation_is_bool():
 
     with pytest.raises(ValueError, match="_logging_activated must be a boolean"):
         mlflow_metrics_dataset._logging_activated = "hello"
+
+
+@pytest.mark.parametrize(
+    "metadata",
+    (
+        None,
+        {"description": "My awsome dataset"},
+        {"string": "bbb", "int": 0},
+    ),
+)
+def test_metrics_history_dataset_with_metadata(tmp_path, metadata):
+    mlflow_metrics_dataset = MlflowMetricsHistoryDataset(
+        prefix="hello",
+        metadata=metadata,
+    )
+
+    assert mlflow_metrics_dataset.metadata == metadata
+
+    # Metadata should not show in _describe
+    assert "metadata" not in mlflow_metrics_dataset._describe()
diff --git a/tests/io/models/test_mlflow_model_local_filesystem_dataset.py b/tests/io/models/test_mlflow_model_local_filesystem_dataset.py
index 757f9df..7c4a283 100644
--- a/tests/io/models/test_mlflow_model_local_filesystem_dataset.py
+++ b/tests/io/models/test_mlflow_model_local_filesystem_dataset.py
@@ -193,3 +193,24 @@ def test_pyfunc_flavor_python_model_save_and_load(
     loaded_model.predict(pd.DataFrame(data=[1], columns=["a"])) == pd.DataFrame(
         data=[2], columns=["a"]
     )
+
+
+@pytest.mark.parametrize(
+    "metadata",
+    (
+        None,
+        {"description": "My awsome dataset"},
+        {"string": "bbb", "int": 0},
+    ),
+)
+def test_metrics_history_dataset_with_metadata(metadata):
+    mlflow_model_ds = MlflowModelLocalFileSystemDataset(
+        flavor="mlflow.sklearn",
+        filepath="/my/file/path",
+        metadata=metadata,
+    )
+
+    assert mlflow_model_ds.metadata == metadata
+
+    # Metadata should not show in _describe
+    assert "metadata" not in mlflow_model_ds._describe()
diff --git a/tests/io/models/test_mlflow_model_registry_dataset.py b/tests/io/models/test_mlflow_model_registry_dataset.py
index bb5f47c..52c9edf 100644
--- a/tests/io/models/test_mlflow_model_registry_dataset.py
+++ b/tests/io/models/test_mlflow_model_registry_dataset.py
@@ -156,3 +156,23 @@ def test_mlflow_model_registry_load_given_alias(tmp_path, monkeypatch):
     ml_ds = MlflowModelRegistryDataset(model_name="demo_model", alias="champion")
     loaded_model = ml_ds.load()
     assert loaded_model.metadata.run_id == runs[1]
+
+
+@pytest.mark.parametrize(
+    "metadata",
+    (
+        None,
+        {"description": "My awsome dataset"},
+        {"string": "bbb", "int": 0},
+    ),
+)
+def test_metrics_history_dataset_with_metadata(tmp_path, metadata):
+    mlflow_model_ds = MlflowModelRegistryDataset(
+        model_name="demo_model",
+        metadata=metadata,
+    )
+
+    assert mlflow_model_ds.metadata == metadata
+
+    # Metadata should not show in _describe
+    assert "metadata" not in mlflow_model_ds._describe()
diff --git a/tests/io/models/test_mlflow_model_tracking_dataset.py b/tests/io/models/test_mlflow_model_tracking_dataset.py
index 4e57ca1..e7f8f36 100644
--- a/tests/io/models/test_mlflow_model_tracking_dataset.py
+++ b/tests/io/models/test_mlflow_model_tracking_dataset.py
@@ -364,3 +364,23 @@ def test_mlflow_model_tracking_logging_deactivation_is_bool():
 
     with pytest.raises(ValueError, match="_logging_activated must be a boolean"):
         mlflow_model_tracking_dataset._logging_activated = "hello"
+
+
+@pytest.mark.parametrize(
+    "metadata",
+    (
+        None,
+        {"description": "My awsome dataset"},
+        {"string": "bbb", "int": 0},
+    ),
+)
+def test_metrics_history_dataset_with_metadata(metadata):
+    mlflow_model_ds = MlflowModelTrackingDataset(
+        flavor="mlflow.sklearn",
+        metadata=metadata,
+    )
+
+    assert mlflow_model_ds.metadata == metadata
+
+    # Metadata should not show in _describe
+    assert "metadata" not in mlflow_model_ds._describe()