Repository: Galileo-Galilei/kedro-mlflow
Pull request #370 (2022-10-09 20:57:39)

Add the ability to configure mlflow_registry_uri in mlflow.yml (#260)
## Description

First step for #260

## Development notes

- add an ``mlflow_registry_uri`` key in ``KedroMlflowConfig`` and ``mlflow.yml``

## 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.


codecov-commenter: # [Codecov](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/370?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9) Report
Base: **100.00**% // Head: **99.88**% // Decreases project coverage by **`-0.11%`** :warning:
> Coverage data is based on head [(`5dcb181`)](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/370?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9) compared to base [(`ec73e7e`)](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/commit/ec73e7e27771d19fed6c59507c0c606e123df0fa?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9).
> Patch coverage: 88.88% of modified lines in pull request are covered.

> :exclamation: Current head 5dcb181 differs from pull request most recent head 37235c2. Consider uploading reports for the commit 37235c2 to get more accurate results

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##            master     #370      +/-   ##
===========================================
- Coverage   100.00%   99.88%   -0.12%     
===========================================
  Files           19       27       +8     
  Lines          880      903      +23     
===========================================
+ Hits           880      902      +22     
- Misses           0        1       +1     
```


| [Impacted Files](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/370?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9) | Coverage Δ | |
|---|---|---|
| [kedro\_mlflow/config/kedro\_mlflow\_config.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/370/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9#diff-a2Vkcm9fbWxmbG93L2NvbmZpZy9rZWRyb19tbGZsb3dfY29uZmlnLnB5) | `99.00% <88.88%> (-1.00%)` | :arrow_down: |
| [kedro\_mlflow/io/metrics/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/370/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9#diff-a2Vkcm9fbWxmbG93L2lvL21ldHJpY3MvX19pbml0X18ucHk=) | `100.00% <0.00%> (ø)` | |
| [kedro\_mlflow/pipeline/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/370/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9#diff-a2Vkcm9fbWxmbG93L3BpcGVsaW5lL19faW5pdF9fLnB5) | `100.00% <0.00%> (ø)` | |
| [kedro\_mlflow/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/370/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9#diff-a2Vkcm9fbWxmbG93L19faW5pdF9fLnB5) | `100.00% <0.00%> (ø)` | |
| [kedro\_mlflow/utils.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/370/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9#diff-a2Vkcm9fbWxmbG93L3V0aWxzLnB5) | `100.00% <0.00%> (ø)` | |
| [kedro\_mlflow/framework/hooks/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/370/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9#diff-a2Vkcm9fbWxmbG93L2ZyYW1ld29yay9ob29rcy9fX2luaXRfXy5weQ==) | `100.00% <0.00%> (ø)` | |
| [kedro\_mlflow/io/models/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/370/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9#diff-a2Vkcm9fbWxmbG93L2lvL21vZGVscy9fX2luaXRfXy5weQ==) | `100.00% <0.00%> (ø)` | |
| [kedro\_mlflow/mlflow/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/370/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9#diff-a2Vkcm9fbWxmbG93L21sZmxvdy9fX2luaXRfXy5weQ==) | `100.00% <0.00%> (ø)` | |
| [kedro\_mlflow/io/artifacts/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/370/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9#diff-a2Vkcm9fbWxmbG93L2lvL2FydGlmYWN0cy9fX2luaXRfXy5weQ==) | `100.00% <0.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/370?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yolan+Honor%C3%A9-Roug%C3%A9).


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0a84cae..7eb5c82 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,13 +2,17 @@
 
 ## [Unreleased]
 
+### Added
+
+- :sparkles: Added an extra ``server.mlflow_registry_uri`` key in ``mlflow.yml`` to set the mlflow registry uri. ([#260](https://github.com/Galileo-Galilei/kedro-mlflow/issues/260))
+
 ### Fixed
 
 -   :bug: `MlflowArtifactDataSet.load()` now correctly loads the artifact when both `artifact_path` and `run_id` arguments are specified. Previous fix in ``0.11.4`` did not work because when the file already exist locally, mlflow did not download it again so tests were incorrectly passing ([#362](https://github.com/Galileo-Galilei/kedro-mlflow/issues/362))
 
 ### Removed
 
--   :fire: :boom: Remove ``reload_kedro_mlflow`` line magic for notebook because kedro will deprecate the entrypoint in 0.18.3. It is still possible to access the mlflow client associated to the configuration in a notebook with ``context.mlflow.server._mlflow_client`` ([#349](https://github.com/Galileo-Galilei/kedro-mlflow/issues/349)). This is not considered as a breaking change since apparently no one uses it according to a [discussion with kedro's team](https://github.com/kedro-org/kedro/issues/878#issuecomment-1226545251)
+-   :fire: :boom: Remove ``reload_kedro_mlflow`` line magic for notebook because kedro will deprecate the entrypoint in 0.18.3. It is still possible to access the mlflow client associated to the configuration in a notebook with ``context.mlflow.server._mlflow_client`` ([#349](https://github.com/Galileo-Galilei/kedro-mlflow/issues/349)). This is not considered as a breaking change since apparently no one uses it according to a [discussion with kedro's team](https://github.com/kedro-org/kedro/issues/878#issuecomment-1226545251).
 
 ## [0.11.4] - 2022-10-04
 
diff --git a/docs/source/04_experimentation_tracking/01_configuration.md b/docs/source/04_experimentation_tracking/01_configuration.md
index 11391e6..0370a39 100644
--- a/docs/source/04_experimentation_tracking/01_configuration.md
+++ b/docs/source/04_experimentation_tracking/01_configuration.md
@@ -32,6 +32,13 @@ server:
 
 This is the **only mandatory key in the `mlflow.yml` file**, but there are many others described hereafter that provide fine-grained control on your mlflow setup.
 
+You can also specify the registry uri:
+
+```yaml
+server:
+  mlflow_registry_uri: registry.db
+```
+
 You can also specify some environment variables needed by mlflow (e.g `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`) in the credentials and specify them in the `mlflow.yml`. Any key specified will be automatically exported as environment variables.
 
 Your `credentials.yml` will look as follows:
diff --git a/kedro_mlflow/config/kedro_mlflow_config.py b/kedro_mlflow/config/kedro_mlflow_config.py
index 7a19a22..aa9f729 100644
--- a/kedro_mlflow/config/kedro_mlflow_config.py
+++ b/kedro_mlflow/config/kedro_mlflow_config.py
@@ -17,6 +17,7 @@ LOGGER = getLogger(__name__)
 class MlflowServerOptions(BaseModel):
     # mutable default is ok for pydantic : https://stackoverflow.com/questions/63793662/how-to-give-a-pydantic-list-field-a-default-value
     mlflow_tracking_uri: Optional[str] = None
+    mlflow_registry_uri: Optional[str] = None
     credentials: Optional[str] = None
     _mlflow_client: MlflowClient = PrivateAttr()
 
@@ -104,13 +105,28 @@ class KedroMlflowConfig(BaseModel):
     def setup(self, context):
         """Setup all the mlflow configuration"""
 
-        self.server.mlflow_tracking_uri = _validate_mlflow_tracking_uri(
-            project_path=context.project_path, uri=self.server.mlflow_tracking_uri
+        # Manage the tracking uri
+        mlflow_tracking_uri = self.server.mlflow_tracking_uri
+        if mlflow_tracking_uri is None:
+            # do not use mlflow.get_tracking_uri() because if there is no env var,
+            # it resolves to 'Path.cwd() / "mlruns"'
+            # but we want 'project_path / "mlruns"'
+            mlflow_tracking_uri = os.environ.get("MLFLOW_TRACKING_URI", "mlruns")
+
+        self.server.mlflow_tracking_uri = _validate_uri(
+            project_path=context.project_path, uri=mlflow_tracking_uri
         )
 
+        # Manage the registry uri: if None, it will use the tracking
+        if self.server.mlflow_registry_uri is not None:
+            self.server.mlflow_registry_uri = _validate_uri(
+                project_path=context.project_path, uri=self.server.mlflow_registry_uri
+            )
+
         # init after validating the uri, else mlflow creates a mlruns folder at the root
         self.server._mlflow_client = MlflowClient(
-            tracking_uri=self.server.mlflow_tracking_uri
+            tracking_uri=self.server.mlflow_tracking_uri,
+            registry_uri=self.server.mlflow_registry_uri,
         )
 
         self._export_credentials(context)
@@ -118,6 +134,7 @@ class KedroMlflowConfig(BaseModel):
         # we set the configuration now: it takes priority
         # if it has already be set in export_credentials
         mlflow.set_tracking_uri(self.server.mlflow_tracking_uri)
+        mlflow.set_registry_uri(self.server.mlflow_registry_uri)
 
         self._set_experiment()
 
@@ -134,7 +151,6 @@ class KedroMlflowConfig(BaseModel):
         Returns:
             mlflow.entities.Experiment -- [description]
         """
-
         # we retrieve the experiment manually to check if it exsits
         mlflow_experiment = self.server._mlflow_client.get_experiment_by_name(
             name=self.tracking.experiment.name
@@ -164,7 +180,7 @@ class KedroMlflowConfig(BaseModel):
         )
 
 
-def _validate_mlflow_tracking_uri(project_path: str, uri: Optional[str]) -> str:
+def _validate_uri(project_path: str, uri: Optional[str]) -> str:
     """Format the uri provided to match mlflow expectations.
 
     Arguments:
@@ -174,15 +190,9 @@ def _validate_mlflow_tracking_uri(project_path: str, uri: Optional[str]) -> str:
         str -- A valid mlflow_tracking_uri
     """
 
-    # this is a special reserved keyword for mlflow which should not be converted to a path
-    # se: https://mlflow.org/docs/latest/tracking.html#where-runs-are-recorded
-    if uri is None:
-        # do not use mlflow.get_tracking_uri() because if there is no env var,
-        # it resolves to 'Path.cwd() / "mlruns"'
-        # but we want 'project_path / "mlruns"'
-        uri = os.environ.get("MLFLOW_TRACKING_URI", "mlruns")
-
     if uri == "databricks":
+        # "databricks" is a special reserved keyword for mlflow which should not be converted to a path
+        # see: https://mlflow.org/docs/latest/tracking.html#where-runs-are-recorded
         return uri
 
     # if no tracking uri is provided, we register the runs locally at the root of the project
@@ -200,7 +210,7 @@ def _validate_mlflow_tracking_uri(project_path: str, uri: Optional[str]) -> str:
             # See : https://discuss.python.org/t/pathlib-absolute-vs-resolve/2573/6
             valid_uri = (Path(project_path) / uri).as_uri()
             LOGGER.info(
-                f"The 'mlflow_tracking_uri' key in mlflow.yml is relative ('server.mlflow_tracking_uri = {uri}'). It is converted to a valid uri: '{valid_uri}'"
+                f"The 'tracking_uri' key in mlflow.yml is relative ('server.mlflow_(tracking|registry)_uri = {uri}'). It is converted to a valid uri: '{valid_uri}'"
             )
         else:
             # else assume it is an uri
diff --git a/kedro_mlflow/template/project/mlflow.yml b/kedro_mlflow/template/project/mlflow.yml
index 38dd02f..ceffc9c 100644
--- a/kedro_mlflow/template/project/mlflow.yml
+++ b/kedro_mlflow/template/project/mlflow.yml
@@ -18,6 +18,7 @@
 
 server:
   mlflow_tracking_uri: null # if null, will use mlflow.get_tracking_uri() as a default
+  mlflow_registry_uri: null # if null, mlflow_tracking_uri will be used as mlflow default
   credentials: null  # must be a valid key in credentials.yml which refers to a dict of sensitive mlflow environment variables (password, tokens...). See top of the file.
 
 tracking:


Test patch:
diff --git a/tests/config/test_get_mlflow_config.py b/tests/config/test_get_mlflow_config.py
index 76bc5a7..4984881 100644
--- a/tests/config/test_get_mlflow_config.py
+++ b/tests/config/test_get_mlflow_config.py
@@ -27,6 +27,7 @@ def test_mlflow_config_default(kedro_project):
     dict_config = dict(
         server=dict(
             mlflow_tracking_uri="mlruns",
+            mlflow_registry_uri=None,
             credentials=None,
         ),
         tracking=dict(
@@ -64,6 +65,7 @@ def test_mlflow_config_in_uninitialized_project(kedro_project, package_name):
     assert context.mlflow.dict() == {
         "server": {
             "mlflow_tracking_uri": (kedro_project / "mlruns").as_uri(),
+            "mlflow_registry_uri": None,
             "credentials": None,
         },
         "tracking": {
@@ -92,6 +94,7 @@ def test_mlflow_config_with_no_experiment_name(kedro_project):
     assert context.mlflow.dict() == {
         "server": {
             "mlflow_tracking_uri": (kedro_project / "mlruns").as_uri(),
+            "mlflow_registry_uri": None,
             "credentials": None,
         },
         "tracking": {
@@ -196,6 +199,7 @@ def test_mlflow_config_with_templated_config_loader(fake_project):
     dict_config = dict(
         server=dict(
             mlflow_tracking_uri="${mlflow_tracking_uri}",
+            mlflow_registry_uri=None,
             credentials=None,
         ),
         tracking=dict(
diff --git a/tests/config/test_kedro_mlflow_config.py b/tests/config/test_kedro_mlflow_config.py
index bb0551c..8051591 100644
--- a/tests/config/test_kedro_mlflow_config.py
+++ b/tests/config/test_kedro_mlflow_config.py
@@ -6,10 +6,7 @@ from kedro.framework.session import KedroSession
 from kedro.framework.startup import bootstrap_project
 from mlflow.tracking import MlflowClient
 
-from kedro_mlflow.config.kedro_mlflow_config import (
-    KedroMlflowConfig,
-    _validate_mlflow_tracking_uri,
-)
+from kedro_mlflow.config.kedro_mlflow_config import KedroMlflowConfig, _validate_uri
 
 
 def test_kedro_mlflow_config_init():
@@ -19,6 +16,7 @@ def test_kedro_mlflow_config_init():
     assert config.dict() == dict(
         server=dict(
             mlflow_tracking_uri=None,  # not setup, not modified yet
+            mlflow_registry_uri=None,
             credentials=None,
         ),
         tracking=dict(
@@ -176,6 +174,26 @@ def test_kedro_mlflow_config_setup_set_tracking_uri(kedro_project_with_mlflow_co
     assert mlflow.get_tracking_uri() == mlflow_tracking_uri
 
 
+def test_kedro_mlflow_config_setup_set_registry_uri(kedro_project_with_mlflow_conf):
+
+    mlflow_registry_uri = (
+        kedro_project_with_mlflow_conf / "awesome_registry.db"
+    ).as_uri()
+
+    config = KedroMlflowConfig(
+        server=dict(
+            mlflow_registry_uri=mlflow_registry_uri,
+        ),
+    )
+
+    bootstrap_project(kedro_project_with_mlflow_conf)
+    with KedroSession.create(project_path=kedro_project_with_mlflow_conf) as session:
+        context = session.load_context()  # setup config
+        config.setup(context)
+
+    assert mlflow.get_registry_uri() == mlflow_registry_uri
+
+
 def test_kedro_mlflow_config_setup_export_credentials(kedro_project_with_mlflow_conf):
 
     (kedro_project_with_mlflow_conf / "conf/base/credentials.yml").write_text(
@@ -228,7 +246,7 @@ def test_kedro_mlflow_config_setup_tracking_priority(kedro_project_with_mlflow_c
 
 def test_validate_uri_local_relative_path(kedro_project_with_mlflow_conf):
 
-    validated_uri = _validate_mlflow_tracking_uri(
+    validated_uri = _validate_uri(
         uri=r"mlruns", project_path=kedro_project_with_mlflow_conf
     )
     assert validated_uri == (kedro_project_with_mlflow_conf / "mlruns").as_uri()
@@ -236,7 +254,7 @@ def test_validate_uri_local_relative_path(kedro_project_with_mlflow_conf):
 
 def test_validate_uri_local_absolute_posix(kedro_project_with_mlflow_conf, tmp_path):
 
-    validated_uri = _validate_mlflow_tracking_uri(
+    validated_uri = _validate_uri(
         uri=tmp_path.as_posix(), project_path=kedro_project_with_mlflow_conf
     )
     assert validated_uri == tmp_path.as_uri()
@@ -244,7 +262,7 @@ def test_validate_uri_local_absolute_posix(kedro_project_with_mlflow_conf, tmp_p
 
 def test_validate_uri_local_absolute_uri(kedro_project_with_mlflow_conf, tmp_path):
 
-    validated_uri = _validate_mlflow_tracking_uri(
+    validated_uri = _validate_uri(
         uri=tmp_path.as_uri(), project_path=kedro_project_with_mlflow_conf
     )
     assert validated_uri == tmp_path.as_uri()
@@ -252,7 +270,7 @@ def test_validate_uri_local_absolute_uri(kedro_project_with_mlflow_conf, tmp_pat
 
 def test_kedro_mlflow_config_validate_uri_databricks(kedro_project_with_mlflow_conf):
     # databricks is a reseved keyword which should not be modified
-    config_uri = _validate_mlflow_tracking_uri(
+    config_uri = _validate_uri(
         uri="databricks", project_path=kedro_project_with_mlflow_conf
     )
     assert config_uri == "databricks"