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

Issue:
MLFlow conflict when using Databricks
<!-- **_If you like the repo, please give it a :star:_** -->

## Description
When running kedro-mlflow on Databricks, occasionally a new run of the experiment might be triggered when running parallelized code. This is because Databricks enables autologging (at least in recent runtimes), and the new runs might be due to an mlflow bug.

Proposed solution: Add a new hook to disable autolog, or include it in the current hook.

```
class DisableMLFlowAutoLogger:    
    @hook_impl(tryfirst=True)
    def after_context_created(self, context) -> None:    
        mlflow.autolog(disable=True)
```

Although I encountered this because of Databricks, I can't imagine a context where you'd like to enable autolog together with the plugin. Could be a parameter of mlflow.yml if you want to be flexible.

## Context

See conversation for context:
https://kedro-org.slack.com/archives/C03RKP2LW64/p1732141412790889

## Steps to Reproduce

1. Start a a Kedro pipeline using kedro-mlflow in a Databricks interactive notebook
3. Use some parallelized code to trigger a new run. Minimal example with Optuna:
```
    study = optuna.create_study()
    study.optimize(lambda trial: objective(my_data,trial),n_trials=100,n_jobs=-1)
```
This will trigger maybe 4-6 new runs when using LightGBM in your objective.

## Expected Result

Results should be in the run started by kedro-mlflow.

## Actual Result

New runs are triggered.

## Your Environment

Databricks Runtime 15.4 ML
Kedro 19.9
kedro-mlflow 0.13.3

## Does the bug also happen with the last version on master?

Yes


Pull request:
:sparkles: Add configuration to disable autologging for Databricks
## Description
Close #610

## Development notes
What have you changed, and how has this been tested?

- Add a configuration in ``mlflow.yml`` to disable autologging globally and set it to true by default. 
- Trigger it in ``KedroMlflowConfig.setup()``

## 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/CHANGELOG.md b/CHANGELOG.md
index 3111f49..cbffdc7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+
+- :sparkles: Add the ``tracking.disable_tracking.disable_autologging`` configuration option in ``mlflow.yml ``to disable autologging by default. This simplify the workflow for Databricks users who have autologging activated by default, which conflicts with the plugin([[#610](https://github.com/Galileo-Galilei/kedro-mlflow/issues/610)]).
+
 ## [0.14.3] - 2025-02-17
 
 ### Added
diff --git a/kedro_mlflow/config/kedro_mlflow_config.py b/kedro_mlflow/config/kedro_mlflow_config.py
index 7b7331b..1e90ca2 100644
--- a/kedro_mlflow/config/kedro_mlflow_config.py
+++ b/kedro_mlflow/config/kedro_mlflow_config.py
@@ -48,6 +48,7 @@ class MlflowServerOptions(BaseModel):
 class DisableTrackingOptions(BaseModel):
     # mutable default is ok for pydantic : https://stackoverflow.com/questions/63793662/how-to-give-a-pydantic-list-field-a-default-value
     pipelines: List[str] = []
+    disable_autologging: bool = True
 
     class Config:
         extra = "forbid"
@@ -158,6 +159,12 @@ class KedroMlflowConfig(BaseModel):
 
         self._set_experiment()
 
+        if self.tracking.disable_tracking.disable_autologging is True:
+            # notice that we dont't pass 'self.tracking.disable_tracking.disable_autologging' directly
+            # because if it's false we just keep the default behaviour of the platform which is often True by default,
+            # except on Databricks
+            mlflow.autolog(disable=True)
+
     def _register_request_header_provider(self, context: KedroContext):
         # this is a specific trick to deal with expiring tokens for authentication, see https://github.com/Galileo-Galilei/kedro-mlflow/issues/357
         if self.server.request_header_provider.type is not None:
diff --git a/kedro_mlflow/template/project/mlflow.yml b/kedro_mlflow/template/project/mlflow.yml
index 51c2c64..1166dc6 100644
--- a/kedro_mlflow/template/project/mlflow.yml
+++ b/kedro_mlflow/template/project/mlflow.yml
@@ -31,6 +31,7 @@ tracking:
   # in a new mlflow run
 
   disable_tracking:
+    disable_autologging: True  # If True, we force autologging to be disabled. This is useful on databricks with autologging by default which conflicts with the plugin. If False, we keep the default behaviour which is disable by default anayway.
     pipelines: []
 
   experiment:
@@ -41,6 +42,7 @@ tracking:
     id: null # if `id` is None, a new run will be created
     name: null # if `name` is None, pipeline name will be used for the run name. You can use "${km.random_name:}" to generate a random name (mlflow's default)
     nested: True  # if `nested` is False, you won't be able to launch sub-runs inside your nodes
+
   params:
     dict_params:
       flatten: False  # if True, parameter which are dictionary will be splitted in multiple parameters when logged in mlflow, one for each key.


Test patch:
diff --git a/tests/config/test_get_mlflow_config.py b/tests/config/test_get_mlflow_config.py
index e3a2415..0fbe9e0 100644
--- a/tests/config/test_get_mlflow_config.py
+++ b/tests/config/test_get_mlflow_config.py
@@ -35,7 +35,10 @@ def test_mlflow_config_default(kedro_project):
             request_header_provider=dict(type=None, pass_context=False, init_kwargs={}),
         ),
         tracking=dict(
-            disable_tracking=dict(pipelines=["my_disabled_pipeline"]),
+            disable_tracking=dict(
+                pipelines=["my_disabled_pipeline"],
+                disable_autologging=True,
+            ),
             experiment=dict(name="fake_package", restore_if_deleted=True),
             run=dict(id="123456789", name="my_run", nested=True),
             params=dict(
@@ -74,7 +77,10 @@ def test_mlflow_config_in_uninitialized_project(kedro_project, package_name):
             request_header_provider=dict(type=None, pass_context=False, init_kwargs={}),
         ),
         tracking=dict(
-            disable_tracking=dict(pipelines=[]),
+            disable_tracking=dict(
+                pipelines=[],
+                disable_autologging=True,
+            ),
             experiment=dict(name="fake_project", restore_if_deleted=True),
             run=dict(id=None, name=None, nested=True),
             params=dict(
@@ -101,7 +107,10 @@ def test_mlflow_config_with_no_experiment_name(kedro_project):
             request_header_provider=dict(type=None, pass_context=False, init_kwargs={}),
         ),
         tracking=dict(
-            disable_tracking=dict(pipelines=[]),
+            disable_tracking=dict(
+                pipelines=[],
+                disable_autologging=True,
+            ),
             experiment=dict(name="fake_project", restore_if_deleted=True),
             run=dict(id=None, name=None, nested=True),
             params=dict(
@@ -215,7 +224,7 @@ def test_mlflow_config_correctly_set(kedro_project, project_settings):
             request_header_provider=dict(type=None, pass_context=False, init_kwargs={}),
         ),
         tracking=dict(
-            disable_tracking=dict(pipelines=[]),
+            disable_tracking=dict(pipelines=[], disable_autologging=True),
             experiment=dict(name="fake_project", restore_if_deleted=True),
             run=dict(id=None, name=None, nested=True),
             params=dict(
@@ -238,7 +247,10 @@ def test_mlflow_config_interpolated_with_globals_resolver(monkeypatch, fake_proj
             request_header_provider=dict(type=None, pass_context=False, init_kwargs={}),
         ),
         tracking=dict(
-            disable_tracking=dict(pipelines=["my_disabled_pipeline"]),
+            disable_tracking=dict(
+                pipelines=["my_disabled_pipeline"],
+                disable_autologging=True,
+            ),
             experiment=dict(name="fake_package", restore_if_deleted=True),
             run=dict(id="123456789", name="my_run", nested=True),
             params=dict(
diff --git a/tests/config/test_kedro_mlflow_config.py b/tests/config/test_kedro_mlflow_config.py
index 42be077..e8566e9 100644
--- a/tests/config/test_kedro_mlflow_config.py
+++ b/tests/config/test_kedro_mlflow_config.py
@@ -19,7 +19,7 @@ def test_kedro_mlflow_config_init():
             request_header_provider=dict(type=None, pass_context=False, init_kwargs={}),
         ),
         tracking=dict(
-            disable_tracking=dict(pipelines=[]),
+            disable_tracking=dict(pipelines=[], disable_autologging=True),
             experiment=dict(name="Default", restore_if_deleted=True),
             run=dict(id=None, name=None, nested=True),
             params=dict(
@@ -270,3 +270,21 @@ def test_from_dict_to_dict_idempotent(kedro_project_with_mlflow_conf):
     # modify config
     reloaded_config = KedroMlflowConfig.parse_obj(original_config_dict)
     assert config == reloaded_config
+
+
+def test_config_disable_autologging(kedro_project_with_mlflow_conf):
+    config = KedroMlflowConfig(
+        tracking=dict(disable_tracking=dict(disable_autologging=True))
+    )
+
+    mlflow.autolog()
+    from mlflow.utils.autologging_utils import autologging_is_disabled
+
+    assert not autologging_is_disabled(integration_name="mlflow")
+
+    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 autologging_is_disabled(integration_name="mlflow")
diff --git a/tests/template/project/test_mlflow_yml.py b/tests/template/project/test_mlflow_yml.py
index 3a067bb..8507bbc 100644
--- a/tests/template/project/test_mlflow_yml.py
+++ b/tests/template/project/test_mlflow_yml.py
@@ -33,7 +33,7 @@ def test_mlflow_yml_rendering(template_mlflowyml):
     expected_config = KedroMlflowConfig.construct(
         project_path="fake/path",
         tracking=dict(
-            disable_tracking=dict(pipelines=[]),
+            disable_tracking=dict(pipelines=[], disable_autologging=True),
             experiment=dict(name="fake_project", restore_if_deleted=True),
             params=dict(
                 dict_params=dict(flatten=False, recursive=True, sep="."),