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

Issue:
Using MLFlow>=2.18.0 with ThreadRunner creates several runs
<!-- **_If you like the repo, please give it a :star:_** -->

## Description
Updating to MLFlow 2.18.0 makes kedro-mlflow's datasets create and log to different runs (new run per artifact). It also breaks kedro-mlflow's test suite.

## Context
Cannot use MLFlow >= 2.18

## Steps to Reproduce

1. Log artifact(s) with any kedro-mlflow's dataset

## Expected Result

Artifact should be logged to the initial run

## Actual Result

Artifact is logged to a new run


Test suite failing on 2.18
```
===short test summary info ===
FAILED tests/mlflow/test_kedro_pipeline_model.py::test_kedro_pipeline_model_save_and_load[pipeline_inference_with_parameters-catalog_with_parameters-data-result3] - mlflow.exceptions.MlflowException: The provided MLFLOW_EXPERIMENT_ID environment variable value `514828615563594097` does not exist in the tracking server. Provide a valid experiment_id.
...

78 failed, 165 passed, 2 skipped, 1 xfailed, 41 warnings in 74.11s (0:01:14)
```

Initial run and subsequent runs with artifacts logged

<img width="1102" alt="image" src="https://github.com/user-attachments/assets/708556ca-b228-4a5d-aa9e-a544a4ecac6c">


## Your Environment

<!-- Include as many relevant details about the environment in which you experienced the bug: -->

* `kedro` and `kedro-mlflow` version used (`pip show kedro` and `pip show kedro-mlflow`): 0.19.8 and 0.13.3
* Python version used (`python -V`): 3.10.14
* Operating system and version: MacOS M2 15.1.1

## Does the bug also happen with the last version on master?

Yes
<!-- The plugin is still in early development and known bugs are fixed as soon as we can. If you are lucky, your bug is already fixed on the `master` branch which is the most up to date. This branch contains our more recent development unpublished on PyPI yet.

In your environment, please try:

```bash
pip install --upgrade git+https://github.com/Galileo-Galilei/kedro-mlflow
```

And check if you can to reproduce the error. If you can't, just wait for the next release or use the master branch at your own risk! -->


Pull request description:
Make artifacts be logged in the same run as parameters when using mlflow>=0.18 and ThreadRunner
## Description
> Why was this PR created?

Close #613. There are details in the issue, but in short : 
- Mlflow made the fluent API thread safe in this PR: https://github.com/mlflow/mlflow/pull/13456/files
- This causes EXPERIMENT_ID environment variable being reseted when a subprocess is launched, and it messes up all the tests. Apparently all tests pass while this is failing. 
- However, another bug arises : when we use a ``ThreadRunner`` (e.g. ``kedro run -r ThreadRunner``), all datasets (``MLflowArtifactDataset``, ``MLflowMetricsDataset``...)  are logging data in a *new* run, not the one launched at the beginning of the pipeline run. There is way to circumvent this at mlflow level. We can however modify the catalog in the ``after_context_created`` hook to ensure proper propagation of the run id. This may introduce some edge cases behaviour in interactive mode, bu CLI experience is definitely more important because that's how most users use kedro-mlflow. 

## Development notes
What have you changed, and how has this been tested?

- Change the ``conftest.py`` auto fixture to "cleanup" mlflow after each test (including env vars and so on), based on what is done here: https://github.com/mlflow/mlflow/pull/13456/files
- Update the ``MlflowHook`` to add the ``run_id`` of the run started at the beginning of the pipeline to the catalog configuration of ``MlflowArtifactDataset``s. This forces to use the ``MlflowClient`` and logs in the correct run. 

⚠️ ~All old tests are now green again, but the new bug for ``ThreadRunner`` is not fixed yet.~ Should be ok, at least for artifacts.


## 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
- [ ] 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/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 0c559d1..239594d 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,7 +1,7 @@
 exclude: ^kedro_mlflow/template/project/run.py$
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.5.6
+    rev: v0.8.3
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]
diff --git a/CHANGELOG.md b/CHANGELOG.md
index c814423..a66b9a6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Fixed
+
+-   :bug: :ambulance: Ensure `MlflowArtifactDataset` logs in the same run that parameters to when using ``mlflow>=2.18`` in combination with ``ThreadRunner`` [#613](https://github.com/Galileo-Galilei/kedro-mlflow/issues/613))
+
 ## [0.13.3] - 2024-10-29
 
 ### Added
diff --git a/kedro_mlflow/framework/hooks/mlflow_hook.py b/kedro_mlflow/framework/hooks/mlflow_hook.py
index 7aed8cc..9db324b 100644
--- a/kedro_mlflow/framework/hooks/mlflow_hook.py
+++ b/kedro_mlflow/framework/hooks/mlflow_hook.py
@@ -27,6 +27,9 @@ from kedro_mlflow.framework.hooks.utils import (
     _flatten_dict,
     _generate_kedro_command,
 )
+from kedro_mlflow.io.catalog.add_run_id_to_artifact_datasets import (
+    add_run_id_to_artifact_datasets,
+)
 from kedro_mlflow.io.catalog.switch_catalog_logging import switch_catalog_logging
 from kedro_mlflow.io.metrics import (
     MlflowMetricDataset,
@@ -270,6 +273,13 @@ class MlflowHook:
                     pipeline_name=run_params["pipeline_name"],
                 ),
             )
+
+            # This function ensures the run_id started at the beginning of the pipeline
+            # is associated to all the datasets. This is necessary because to make mlflow thread safe
+            # each call to the "active run" now creates a new run when started in a new thread. See
+            # https://github.com/Galileo-Galilei/kedro-mlflow/issues/613 and https://github.com/Galileo-Galilei/kedro-mlflow/pull/615
+            add_run_id_to_artifact_datasets(catalog, mlflow.active_run().info.run_id)
+
         else:
             self._logger.info(
                 "kedro-mlflow logging is deactivated for this pipeline in the configuration. This includes DataSets and parameters."
diff --git a/kedro_mlflow/io/artifacts/mlflow_artifact_dataset.py b/kedro_mlflow/io/artifacts/mlflow_artifact_dataset.py
index 9e1d42f..67e1fb9 100644
--- a/kedro_mlflow/io/artifacts/mlflow_artifact_dataset.py
+++ b/kedro_mlflow/io/artifacts/mlflow_artifact_dataset.py
@@ -171,3 +171,9 @@ class MlflowArtifactDataset(AbstractVersionedDataset):
         and consequently does not implements abtracts methods
         """
         pass
+
+
+def _is_instance_mlflow_artifact_dataset(dataset_instance):
+    parent_classname = dataset_instance.__class__.__bases__[0].__name__
+    instance_classname = f"Mlflow{parent_classname}"
+    return type(dataset_instance).__name__ == instance_classname
diff --git a/kedro_mlflow/io/catalog/add_run_id_to_artifact_datasets.py b/kedro_mlflow/io/catalog/add_run_id_to_artifact_datasets.py
new file mode 100644
index 0000000..05f2735
--- /dev/null
+++ b/kedro_mlflow/io/catalog/add_run_id_to_artifact_datasets.py
@@ -0,0 +1,9 @@
+from kedro_mlflow.io.artifacts.mlflow_artifact_dataset import (
+    _is_instance_mlflow_artifact_dataset,
+)
+
+
+def add_run_id_to_artifact_datasets(catalog, run_id: str):
+    for name, dataset in catalog._datasets.items():
+        if _is_instance_mlflow_artifact_dataset(dataset):
+            catalog._datasets[name].run_id = run_id
diff --git a/setup.py b/setup.py
index 278cc0f..9c589be 100644
--- a/setup.py
+++ b/setup.py
@@ -50,7 +50,7 @@ setup(
             "pytest-cov>=2.8.0, <7.0.0",
             "pytest-lazy-fixtures>=1.0.0, <2.0.0",
             "pytest-mock>=3.1.0, <4.0.0",
-            "ruff>=0.5.0,<0.8.0",  # ensure consistency with pre-commit
+            "ruff>=0.5.0,<0.9.0",  # ensure consistency with pre-commit
             "scikit-learn>=0.23.0, <1.7.0",
             "kedro-datasets[pandas.CSVDataSet]",
         ],


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 460c095..10a9d3f 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -29,22 +29,24 @@ def mlflow_client(tracking_uri):
 
 @pytest.fixture(autouse=True)
 def cleanup_mlflow_after_runs():
-    # A test function will be run at this point
-    yield
+    yield  # A test function will be run at this point
     while mlflow.active_run():
         mlflow.end_run()
 
-    import mlflow.tracking.fluent as mtf
-
     # if set_experiment has been called before, it stores the experiment_id
     # as a global variable, so if we change the tracking_uri afterwards
     # mlflow is completly lost because the experiment id no longer exists
     # we just reset it after a test, like in a brand new session
-    if hasattr(mtf, "_active_experiment_id"):
-        mtf._active_experiment_id = None
 
-    if "MLFLOW_TRACKING_URI" in os.environ:
-        os.environ.pop("MLFLOW_TRACKING_URI")
+    # CAVEAT 1 : do not import from "mlflow.tracking.fluent import _active_experiment_id"
+    # because due to python namespacing import, it will not change the global variable accessed by mlflow
+
+    # CAVEAT 2 : Since this PR: https://github.com/mlflow/mlflow/pull/13456/files
+    # we need to reset experiment ID too because its now resetted in each thread
+    mlflow.tracking.fluent._active_experiment_id = None
+    os.environ.pop("MLFLOW_EXPERIMENT_ID", None)
+    os.environ.pop("MLFLOW_TRACKING_URI", None)
+    os.environ.pop("MLFLOW_REGISTRY_URI", None)
 
     # see https://github.com/kedro-org/kedro/blob/859f98217eed12208a922b771a97cbfb82ba7e80/tests/framework/session/test_session.py#L173
 
diff --git a/tests/framework/hooks/test_hook_log_artifact.py b/tests/framework/hooks/test_hook_log_artifact.py
new file mode 100644
index 0000000..9cd3205
--- /dev/null
+++ b/tests/framework/hooks/test_hook_log_artifact.py
@@ -0,0 +1,159 @@
+import mlflow
+import pandas as pd
+import pytest
+from kedro.framework.session import KedroSession
+from kedro.framework.startup import bootstrap_project
+from kedro.io import DataCatalog, MemoryDataset
+from kedro.pipeline import Pipeline, node
+from kedro.runner import ThreadRunner
+from kedro_datasets.pickle import PickleDataset
+
+from kedro_mlflow.framework.hooks.mlflow_hook import MlflowHook
+from kedro_mlflow.io.artifacts import MlflowArtifactDataset
+
+
+@pytest.fixture
+def dummy_pipeline():
+    def preprocess_fun(data):
+        return data
+
+    def train_fun(data):
+        return 2
+
+    dummy_pipeline = Pipeline(
+        [
+            node(
+                func=preprocess_fun,
+                inputs="raw_data",
+                outputs="data",
+            ),
+            node(
+                func=train_fun,
+                inputs=["data"],
+                outputs="model",
+            ),
+        ]
+    )
+    return dummy_pipeline
+
+
+@pytest.fixture
+def dummy_catalog(tmp_path):
+    dummy_catalog = DataCatalog(
+        {
+            "raw_data": MemoryDataset(pd.DataFrame(data=[1], columns=["a"])),
+            "data": MemoryDataset(),
+            "model": MlflowArtifactDataset(
+                dataset=dict(
+                    type=PickleDataset, filepath=(tmp_path / "model.csv").as_posix()
+                )
+            ),
+        }
+    )
+    return dummy_catalog
+
+
+@pytest.fixture
+def dummy_run_params(tmp_path):
+    dummy_run_params = {
+        "project_path": tmp_path.as_posix(),
+        "env": "local",
+        "kedro_version": "0.16.0",
+        "tags": [],
+        "from_nodes": [],
+        "to_nodes": [],
+        "node_names": [],
+        "from_inputs": [],
+        "load_versions": [],
+        "pipeline_name": "my_cool_pipeline",
+        "extra_params": [],
+    }
+    return dummy_run_params
+
+
+def test_mlflow_hook_automatically_update_artifact_run_id(
+    kedro_project, dummy_run_params, dummy_pipeline, dummy_catalog
+):
+    # since mlflow>=2.18, the fluent API create a new run for each thread
+    # hence for thread runner we need to prefix the catalog with the run id
+
+    bootstrap_project(kedro_project)
+    with KedroSession.create(project_path=kedro_project) as session:
+        context = session.load_context()  # triggers conf setup
+
+        mlflow_hook = MlflowHook()
+        mlflow_hook.after_context_created(context)  # setup mlflow config
+
+        mlflow_hook.after_catalog_created(
+            catalog=dummy_catalog,
+            # `after_catalog_created` is not using any of below arguments,
+            # so we are setting them to empty values.
+            conf_catalog={},
+            conf_creds={},
+            feed_dict={},
+            save_version="",
+            load_versions="",
+        )
+
+        mlflow_hook.before_pipeline_run(
+            run_params=dummy_run_params, pipeline=dummy_pipeline, catalog=dummy_catalog
+        )
+
+        run_id = mlflow.active_run().info.run_id
+        # Check if metrics datasets have prefix with its names.
+        # for metric
+        assert dummy_catalog._datasets["model"].run_id == run_id
+
+
+def test_mlflow_hook_log_artifacts_within_same_run_with_thread_runner(
+    kedro_project, dummy_run_params, dummy_pipeline, dummy_catalog
+):
+    # this test is very specific to a new design introduced in mlflow 2.18 to make it htread safe
+    # see https://github.com/Galileo-Galilei/kedro-mlflow/issues/613
+    bootstrap_project(kedro_project)
+
+    with KedroSession.create(project_path=kedro_project) as session:
+        context = session.load_context()  # setup mlflow
+
+        mlflow_hook = MlflowHook()
+        runner = ThreadRunner()  # this is what we want to test
+
+        mlflow_hook.after_context_created(context)
+        mlflow_hook.after_catalog_created(
+            catalog=dummy_catalog,
+            # `after_catalog_created` is not using any of arguments bellow,
+            # so we are setting them to empty values.
+            conf_catalog={},
+            conf_creds={},
+            feed_dict={},
+            save_version="",
+            load_versions="",
+        )
+        mlflow_hook.before_pipeline_run(
+            run_params=dummy_run_params,
+            pipeline=dummy_pipeline,
+            catalog=dummy_catalog,
+        )
+
+        # we get the run id BEFORE running the pipeline because it was modified in different thread
+        run_id_before_run = mlflow.active_run().info.run_id
+
+        runner.run(dummy_pipeline, dummy_catalog, session._hook_manager)
+
+        run_id_after_run = mlflow.active_run().info.run_id
+
+        # CHECK 1: check that we are not on the second id created by the thread.lock()
+        assert run_id_before_run == run_id_after_run
+
+        mlflow_hook.after_pipeline_run(
+            run_params=dummy_run_params,
+            pipeline=dummy_pipeline,
+            catalog=dummy_catalog,
+        )
+
+        mlflow_client = context.mlflow.server._mlflow_client
+
+        # check that the artifact is assocaied to the initial run:
+
+        artifacts_list = mlflow_client.list_artifacts(run_id_before_run)
+        assert len(artifacts_list) == 1