Repository: Galileo-Galilei/kedro-mlflow
Pull request #379 (2022-12-06 21:01:12)

⬆️ Update sphinx-click requirement from <4.4,>=3.1 to >=3.1,<4.5
Updates the requirements on [sphinx-click](https://github.com/click-contrib/sphinx-click) to permit the latest version.
<details>
<summary>Commits</summary>
<ul>
<li><a href="https://github.com/click-contrib/sphinx-click/commit/0b81d33f0cdd444c5f520e94bca2543d7bbe335f"><code>0b81d33</code></a> Add release note for Python 3.11 support</li>
<li><a href="https://github.com/click-contrib/sphinx-click/commit/0f73ff495889204683d6fefc7b61a58335fc92bf"><code>0f73ff4</code></a> actions: Test Python 3.11</li>
<li><a href="https://github.com/click-contrib/sphinx-click/commit/7698cac2c52482e31b1b805b89299f250ca2c134"><code>7698cac</code></a> Points to alternative to complement Click-based CLI documentation.</li>
<li><a href="https://github.com/click-contrib/sphinx-click/commit/2ce4c2d9dfbe26f59efb20954eaedbccfe42b5b7"><code>2ce4c2d</code></a> Add Sphinx extension classifier</li>
<li><a href="https://github.com/click-contrib/sphinx-click/commit/f8f1a5a05654ac381dec47b0999df83816019fcd"><code>f8f1a5a</code></a> Add pull request template</li>
<li><a href="https://github.com/click-contrib/sphinx-click/commit/09a71dcc5052b9e8882eaa606b06826b950ee05b"><code>09a71dc</code></a> Trivial setup.py, setup.cfg changes</li>
<li><a href="https://github.com/click-contrib/sphinx-click/commit/30516a4a74fbd2e62faf252af108bbac0f0b3306"><code>30516a4</code></a> Render option defaults as literals</li>
<li><a href="https://github.com/click-contrib/sphinx-click/commit/ffceac79df71c9db7add514e42214afaab6016ca"><code>ffceac7</code></a> Change target Python version used for black</li>
<li><a href="https://github.com/click-contrib/sphinx-click/commit/f2c35ccbd965da5b781d275c39a9ed7aa3deacc4"><code>f2c35cc</code></a> actions: Update to latest versions of actions</li>
<li><a href="https://github.com/click-contrib/sphinx-click/commit/6114f8d3ae1941057f743e6778cfa35a82477477"><code>6114f8d</code></a> Use deep clone for GitHub actions</li>
<li>Additional commits viewable in <a href="https://github.com/click-contrib/sphinx-click/compare/3.1.0...4.4.0">compare view</a></li>
</ul>
</details>
<br />


Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`.

[//]: # (dependabot-automerge-start)
[//]: # (dependabot-automerge-end)

---

<details>
<summary>Dependabot commands and options</summary>
<br />

You can trigger Dependabot actions by commenting on this PR:
- `@dependabot rebase` will rebase this PR
- `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it
- `@dependabot merge` will merge this PR after your CI passes on it
- `@dependabot squash and merge` will squash and merge this PR after your CI passes on it
- `@dependabot cancel merge` will cancel a previously requested merge and block automerging
- `@dependabot reopen` will reopen this PR if it is closed
- `@dependabot close` will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
- `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)


</details>

codecov-commenter: # [Codecov](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/379?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: **100.00**% // No change to project coverage :thumbsup:
> Coverage data is based on head [(`a9eab1b`)](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/379?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 [(`b3126f8`)](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/commit/b3126f830ca2d9b513565e140380c54df60449b3?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 has no changes to coverable lines.

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


```diff
@@            Coverage Diff            @@
##            master      #379   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           27        19    -8     
  Lines          903       885   -18     
=========================================
- Hits           903       885   -18     
```


| [Impacted Files](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/379?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/mlflow/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/379/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==) | | |
| [kedro\_mlflow/io/artifacts/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/379/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==) | | |
| [kedro\_mlflow/framework/hooks/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/379/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==) | | |
| [kedro\_mlflow/utils.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/379/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) | | |
| [kedro\_mlflow/pipeline/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/379/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) | | |
| [kedro\_mlflow/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/379/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) | | |
| [kedro\_mlflow/io/models/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/379/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==) | | |
| [kedro\_mlflow/io/metrics/\_\_init\_\_.py](https://codecov.io/gh/Galileo-Galilei/kedro-mlflow/pull/379/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=) | | |

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/379?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 7eb5c82..025f769 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@
 ### 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))
+- :sparkles: Add support for authorization with expiring tokens by adding an extra ``server.request_header_provider`` entry in ``mlflow.yml`` ([#357](https://github.com/Galileo-Galilei/kedro-mlflow/issues/357))
 
 ### Fixed
 
diff --git a/docs/source/03_getting_started/01_example_project.md b/docs/source/03_getting_started/01_example_project.md
index 3587ccc..00cccf2 100644
--- a/docs/source/03_getting_started/01_example_project.md
+++ b/docs/source/03_getting_started/01_example_project.md
@@ -12,7 +12,7 @@ pip install kedro-mlflow==0.11.4
 
 ## Install the toy project
 
-For this end to end example, we will use the [kedro starter](https://kedro.readthedocs.io/en/latest/get_started/starters.html) with the [iris dataset](https://github.com/quantumblacklabs/kedro-starter-pandas-iris).
+For this end to end example, we will use the [kedro starter](https://kedro.readthedocs.io/en/stable/get_started/starters.html) with the [iris dataset](https://github.com/quantumblacklabs/kedro-starter-pandas-iris).
 
 We use this project because:
 
diff --git a/docs/source/04_experimentation_tracking/01_configuration.md b/docs/source/04_experimentation_tracking/01_configuration.md
index 7454f2e..06e4770 100644
--- a/docs/source/04_experimentation_tracking/01_configuration.md
+++ b/docs/source/04_experimentation_tracking/01_configuration.md
@@ -21,6 +21,8 @@ The ``mlflow.yml`` file contains all configuration you can pass either to kedro
 
 ### Configure the tracking server
 
+#### Configure the tracking and registry uri
+
 ``kedro-mlflow`` needs the tracking uri of your mlflow tracking server to operate properly. The ``mlflow.yml`` file must have the ``mlflow_tracking_uri`` key with a [valid mlflow_tracking_uri associated](https://mlflow.org/docs/latest/tracking.html#where-runs-are-recorded) value. The ``mlflow.yml`` default have this keys set to ``mlruns``. This will create a ``mlruns`` folder locally at the root of your kedro project and enable you to use the plugin without any setup of a mlflow tracking server.
 
 Unlike mlflow, `kedro-mlflow` allows the `mlflow_tracking_uri` to be a relative path. It will convert it to an absolute uri automatically.
@@ -39,6 +41,10 @@ server:
   mlflow_registry_uri: registry.db
 ```
 
+#### Configure the credentials
+
+##### Default credentials with environment variables
+
 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:
@@ -56,7 +62,25 @@ server:
   credentials: my_mlflow_credentials
 ```
 
-For safety reasons, the credentials will not be accessible within `KedroMlflowConfig` objects. They will be exported as environment variables *on the fly* when running the pipeline.
+```{note}
+For security reasons, the credentials will not be accessible within `KedroMlflowConfig` objects. They will be exported as environment variables *on the fly* when running the pipeline.
+```
+
+##### Authentication with expiring tokens
+
+Mlflow can be deployed with OAuth2.0. authentication method : in this case, secured MLflow instances require HTTP requests to have the ``Authorization: Bearer <token>`` header. Mlflow exposes a [``RequestHeaderProvider`` abstract class](https://github.com/mlflow/mlflow/blob/master/mlflow/tracking/request_header/abstract_request_header_provider.py#L4) to mange this use case. If you need kedro-mlflow to use a *custom* header provider, you can configure your ``mlflow.yml`` as follow:  
+
+```yaml
+# mlflow.yml
+server:
+    request_header_provider:
+        type: path.to.your.class.CustomRequestHeaderProvider
+        pass_context: True # if you want to pass context. it must be named ``kedro_context`` in the ``__init__`` method of your custom ``request_header_provider``)
+        init_kwargs:
+            my_kwarg: 1
+```
+
+This will automatically register in the mlflow entrypoint the ``CustomRequestHeaderProvider(kedro_context=<kedro-context>, my_kwarg=1)`` request header when running a kedro pipeline.
 
 ### Deactivate tracking under conditions
 
diff --git a/docs/source/06_interactive_use/01_notebook_use.md b/docs/source/06_interactive_use/01_notebook_use.md
index f1bf458..ddd6f1c 100644
--- a/docs/source/06_interactive_use/01_notebook_use.md
+++ b/docs/source/06_interactive_use/01_notebook_use.md
@@ -29,7 +29,7 @@ Or if you are on JupyterLab,
 %load_ext kedro.ipython
 ```
 
-Kedro [creates a bunch of global variables](https://kedro.readthedocs.io/en/latest/tools_integration/ipython.html#kedro-and-jupyter), including a `session`, a ``context`` and a ``catalog`` which are automatically accessible.
+Kedro [creates a bunch of global variables](https://kedro.readthedocs.io/en/stable/tools_integration/ipython.html#use-kedro-with-ipython-and-jupyter), including a `session`, a ``context`` and a ``catalog`` which are automatically accessible.
 
 When the context was created, ``kedro-mlflow`` automatically:
 - loaded and setup (create the tracking uri, export credentials...) the mlflow configuration of your `mlflow.yml`
diff --git a/kedro_mlflow/config/kedro_mlflow_config.py b/kedro_mlflow/config/kedro_mlflow_config.py
index aa9f729..968ec42 100644
--- a/kedro_mlflow/config/kedro_mlflow_config.py
+++ b/kedro_mlflow/config/kedro_mlflow_config.py
@@ -1,24 +1,44 @@
 import os
 from logging import getLogger
 from pathlib import Path, PurePath
-from typing import List, Optional
+from typing import Dict, List, Optional
 from urllib.parse import urlparse
 
 import mlflow
+import mlflow.tracking.request_header.registry as mtrr  # necessary to access the global variable '_request_header_provider_registry' of the namespace
 from kedro.framework.context import KedroContext
+from kedro.utils import load_obj
 from mlflow.entities import Experiment
 from mlflow.tracking.client import MlflowClient
+from mlflow.tracking.request_header.abstract_request_header_provider import (
+    RequestHeaderProvider,
+)
 from pydantic import BaseModel, PrivateAttr, StrictBool
 from typing_extensions import Literal
 
 LOGGER = getLogger(__name__)
 
 
+class RequestHeaderProviderOptions(BaseModel):
+    # mutable default is ok for pydantic : https://stackoverflow.com/questions/63793662/how-to-give-a-pydantic-list-field-a-default-value
+    type: Optional[str] = None
+    pass_context: bool = False
+    init_kwargs: Dict[str, str] = {}
+
+    class Config:
+        extra = "forbid"
+        # necessary to check instance of RequestHeaderProvider:
+        arbitrary_types_allowed = "allowed"
+
+
 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
+    request_header_provider: RequestHeaderProviderOptions = (
+        RequestHeaderProviderOptions()
+    )
     _mlflow_client: MlflowClient = PrivateAttr()
 
     class Config:
@@ -131,6 +151,7 @@ class KedroMlflowConfig(BaseModel):
 
         self._export_credentials(context)
 
+        self._register_request_header_provider(context)
         # 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)
@@ -138,6 +159,34 @@ class KedroMlflowConfig(BaseModel):
 
         self._set_experiment()
 
+    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:
+            # it type is none, there is nothing to register
+            request_header_provider_class = load_obj(
+                self.server.request_header_provider.type
+            )
+
+            if not issubclass(request_header_provider_class, RequestHeaderProvider):
+                raise ValueError(
+                    f"request_header_provider.type='{self.server.request_header_provider.type}' should be a sublass of 'mlflow.tracking.request_header.abstract_request_header_provider', got {request_header_provider_class} instead."
+                )
+
+            # we allow the user to instantiate the class with parameters, including the kedro context for advanded retrieval of kedro objects
+            init_kwargs = (
+                {
+                    **self.server.request_header_provider.init_kwargs,
+                    "kedro_context": context,
+                }
+                if self.server.request_header_provider.pass_context
+                else self.server.request_header_provider.init_kwargs
+            )
+
+            # the "register" method because expects a callable class with no arguments so we tricked it with a lambda
+            mtrr._request_header_provider_registry.register(
+                lambda: request_header_provider_class(**init_kwargs)
+            )
+
     def _export_credentials(self, context: KedroContext):
         conf_creds = context._get_config_credentials()
         mlflow_creds = conf_creds.get(self.server.credentials, {})
diff --git a/kedro_mlflow/template/project/mlflow.yml b/kedro_mlflow/template/project/mlflow.yml
index ceffc9c..377c63c 100644
--- a/kedro_mlflow/template/project/mlflow.yml
+++ b/kedro_mlflow/template/project/mlflow.yml
@@ -20,6 +20,10 @@ 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.
+  request_header_provider: # this is only useful to deal with expiring token, see https://github.com/Galileo-Galilei/kedro-mlflow/issues/357
+    type: null # The path to a class : my_project.pipelines.module.MyClass. Should inherit from https://github.com/mlflow/mlflow/blob/master/mlflow/tracking/request_header/abstract_request_header_provider.py#L4
+    pass_context: False # should the class be instantiated with "kedro_context" argument?
+    init_kwargs: {} # any kwargs to pass to the class when it is instantiated
 
 tracking:
   # You can specify a list of pipeline names for which tracking will be disabled
diff --git a/setup.py b/setup.py
index a8793a3..54a710b 100644
--- a/setup.py
+++ b/setup.py
@@ -41,7 +41,7 @@ setup(
             "sphinx>=4.5.0,<6.0.0",
             "sphinx_rtd_theme~=1.0.0",
             "sphinx-markdown-tables~=0.0.15",
-            "sphinx-click>=3.1,<4.4",
+            "sphinx-click>=3.1,<4.5",
             "sphinx_copybutton~=0.5.0",
             "myst-parser>=0.17.2,<0.19.0",
         ],
@@ -50,9 +50,9 @@ setup(
             "pytest-cov>=2.8.0, <4.0.0",
             "pytest-lazy-fixture>=0.6.0, <1.0.0",
             "pytest-mock>=3.1.0, <4.0.0",
-            "scikit-learn>=0.23.0, <1.2.0",
+            "scikit-learn>=0.23.0, <1.3.0",
             "flake8==5.0.4",  # ensure consistency with pre-commit
-            "black==22.10.0",  # pin black version because it is not compatible with a pip range (because of non semver version number)
+            "black==22.12.0",  # pin black version because it is not compatible with a pip range (because of non semver version number)
             "isort==5.10.1",  # ensure consistency with pre-commit
         ],
         "dev": [


Test patch:
diff --git a/tests/config/test_get_mlflow_config.py b/tests/config/test_get_mlflow_config.py
index 4984881..a0033e3 100644
--- a/tests/config/test_get_mlflow_config.py
+++ b/tests/config/test_get_mlflow_config.py
@@ -1,12 +1,15 @@
 import json
+import sys
 from pathlib import Path
 
+import mlflow.tracking.request_header.registry as mtrr  # necessary to access the global variable '_request_header_provider_registry' of the namespace
 import pytest
 import toml
 import yaml
 from dynaconf.validator import Validator
 from kedro import __version__ as kedro_version
 from kedro.config import TemplatedConfigLoader
+from kedro.framework.context import KedroContext
 from kedro.framework.project import _IsSubclassValidator, _ProjectSettings
 from kedro.framework.session import KedroSession
 from kedro.framework.startup import bootstrap_project
@@ -29,6 +32,7 @@ def test_mlflow_config_default(kedro_project):
             mlflow_tracking_uri="mlruns",
             mlflow_registry_uri=None,
             credentials=None,
+            request_header_provider=dict(type=None, pass_context=False, init_kwargs={}),
         ),
         tracking=dict(
             disable_tracking=dict(pipelines=["my_disabled_pipeline"]),
@@ -62,23 +66,24 @@ def test_mlflow_config_in_uninitialized_project(kedro_project, package_name):
     bootstrap_project(kedro_project)
     session = KedroSession.create(project_path=kedro_project, package_name=package_name)
     context = session.load_context()
-    assert context.mlflow.dict() == {
-        "server": {
-            "mlflow_tracking_uri": (kedro_project / "mlruns").as_uri(),
-            "mlflow_registry_uri": None,
-            "credentials": None,
-        },
-        "tracking": {
-            "disable_tracking": {"pipelines": []},
-            "experiment": {"name": "fake_project", "restore_if_deleted": True},
-            "run": {"id": None, "name": None, "nested": True},
-            "params": {
-                "dict_params": {"flatten": False, "recursive": True, "sep": "."},
-                "long_params_strategy": "fail",
-            },
-        },
-        "ui": {"port": "5000", "host": "127.0.0.1"},
-    }
+    assert context.mlflow.dict() == dict(
+        server=dict(
+            mlflow_tracking_uri=(kedro_project / "mlruns").as_uri(),
+            mlflow_registry_uri=None,
+            credentials=None,
+            request_header_provider=dict(type=None, pass_context=False, init_kwargs={}),
+        ),
+        tracking=dict(
+            disable_tracking=dict(pipelines=[]),
+            experiment=dict(name="fake_project", restore_if_deleted=True),
+            run=dict(id=None, name=None, nested=True),
+            params=dict(
+                dict_params=dict(flatten=False, recursive=True, sep="."),
+                long_params_strategy="fail",
+            ),
+        ),
+        ui=dict(port="5000", host="127.0.0.1"),
+    )
 
 
 def test_mlflow_config_with_no_experiment_name(kedro_project):
@@ -91,23 +96,24 @@ def test_mlflow_config_with_no_experiment_name(kedro_project):
         project_path=kedro_project, package_name="fake_project"
     )
     context = session.load_context()
-    assert context.mlflow.dict() == {
-        "server": {
-            "mlflow_tracking_uri": (kedro_project / "mlruns").as_uri(),
-            "mlflow_registry_uri": None,
-            "credentials": None,
-        },
-        "tracking": {
-            "disable_tracking": {"pipelines": []},
-            "experiment": {"name": "fake_project", "restore_if_deleted": True},
-            "run": {"id": None, "name": None, "nested": True},
-            "params": {
-                "dict_params": {"flatten": False, "recursive": True, "sep": "."},
-                "long_params_strategy": "fail",
-            },
-        },
-        "ui": {"port": "5000", "host": "127.0.0.1"},
-    }
+    assert context.mlflow.dict() == dict(
+        server=dict(
+            mlflow_tracking_uri=(kedro_project / "mlruns").as_uri(),
+            mlflow_registry_uri=None,
+            credentials=None,
+            request_header_provider=dict(type=None, pass_context=False, init_kwargs={}),
+        ),
+        tracking=dict(
+            disable_tracking=dict(pipelines=[]),
+            experiment=dict(name="fake_project", restore_if_deleted=True),
+            run=dict(id=None, name=None, nested=True),
+            params=dict(
+                dict_params=dict(flatten=False, recursive=True, sep="."),
+                long_params_strategy="fail",
+            ),
+        ),
+        ui=dict(port="5000", host="127.0.0.1"),
+    )
 
 
 @pytest.fixture(autouse=True)
@@ -201,6 +207,7 @@ def test_mlflow_config_with_templated_config_loader(fake_project):
             mlflow_tracking_uri="${mlflow_tracking_uri}",
             mlflow_registry_uri=None,
             credentials=None,
+            request_header_provider=dict(type=None, pass_context=False, init_kwargs={}),
         ),
         tracking=dict(
             disable_tracking=dict(pipelines=["my_disabled_pipeline"]),
@@ -234,3 +241,245 @@ def test_mlflow_config_with_templated_config_loader(fake_project):
     with KedroSession.create("fake_package", fake_project) as session:
         context = session.load_context()
         assert context.mlflow.dict(exclude={"project_path"}) == expected
+
+
+@pytest.fixture
+def request_header_provider_cleaner(fake_project):
+    sys.path.append(fake_project.as_posix())
+    yield
+    # cleanup test specific setup
+    (fake_project / "custom_rhp.py").unlink()
+    sys.path.pop()
+    del sys.modules["custom_rhp"]
+    mtrr._request_header_provider_registry._registry.pop()
+
+
+@pytest.mark.usefixtures("request_header_provider_cleaner")
+def test_mlflow_config_with_request_header_provider(fake_project):
+
+    # emulate import of custom request header class
+    custom_rhp_txt = """
+from mlflow.tracking.request_header.abstract_request_header_provider import RequestHeaderProvider
+
+class CustomRequestHeaderProvider(RequestHeaderProvider):
+    def in_context(self):
+        pass
+    def request_headers(self):
+        pass
+"""
+
+    with open(fake_project / "custom_rhp.py", "w") as fhandler:
+        fhandler.write(custom_rhp_txt)
+
+    dict_config = dict(
+        server=dict(
+            mlflow_tracking_uri=None,  # not setup, not modified yet
+            credentials=None,
+            request_header_provider={"type": "custom_rhp.CustomRequestHeaderProvider"},
+        ),
+        tracking=dict(
+            disable_tracking=dict(pipelines=[]),
+            experiment=dict(name="Default", restore_if_deleted=True),
+            run=dict(id=None, name=None, nested=True),
+            params=dict(
+                dict_params=dict(
+                    flatten=False,
+                    recursive=True,
+                    sep=".",
+                ),
+                long_params_strategy="fail",
+            ),
+        ),
+        ui=dict(port="5000", host="127.0.0.1"),
+    )
+
+    _write_yaml(fake_project / "conf" / "local" / "mlflow.yml", dict_config)
+
+    bootstrap_project(fake_project)
+    with KedroSession.create("fake_package", fake_project) as session:
+        session.load_context()  # trigger setup and request_header_provider registration
+
+        assert (
+            mtrr._request_header_provider_registry._registry[-1].__class__.__name__
+            == "CustomRequestHeaderProvider"
+        )
+
+
+@pytest.mark.usefixtures("request_header_provider_cleaner")
+def test_mlflow_config_with_request_header_provider_with_init_kwargs(
+    fake_project,
+):
+
+    # emulate import of custom request header class
+    custom_rhp_txt = """
+from mlflow.tracking.request_header.abstract_request_header_provider import RequestHeaderProvider
+
+class CustomRequestHeaderProviderInitKwargs(RequestHeaderProvider):
+    def __init__(self, a):
+        super().__init__()
+        self.a=a
+
+    def in_context(self):
+        pass
+    def request_headers(self):
+        pass
+"""
+
+    with open(fake_project / "custom_rhp.py", "w") as fhandler:
+        fhandler.write(custom_rhp_txt)
+
+    dict_config = dict(
+        server=dict(
+            mlflow_tracking_uri=None,  # not setup, not modified yet
+            credentials=None,
+            request_header_provider=dict(
+                type="custom_rhp.CustomRequestHeaderProviderInitKwargs",
+                init_kwargs=dict(a=1),
+            ),
+        ),
+        tracking=dict(
+            disable_tracking=dict(pipelines=[]),
+            experiment=dict(name="Default", restore_if_deleted=True),
+            run=dict(id=None, name=None, nested=True),
+            params=dict(
+                dict_params=dict(
+                    flatten=False,
+                    recursive=True,
+                    sep=".",
+                ),
+                long_params_strategy="fail",
+            ),
+        ),
+        ui=dict(port="5000", host="127.0.0.1"),
+    )
+
+    _write_yaml(fake_project / "conf" / "local" / "mlflow.yml", dict_config)
+
+    bootstrap_project(fake_project)
+    with KedroSession.create("fake_package", fake_project) as session:
+        session.load_context()  # trigger setup and request_header_provider registration
+
+        assert (
+            mtrr._request_header_provider_registry._registry[-1].__class__.__name__
+            == "CustomRequestHeaderProviderInitKwargs"
+        )
+    assert mtrr._request_header_provider_registry._registry[-1].a == "1"
+    assert not hasattr(mtrr._request_header_provider_registry._registry[-1], "context")
+
+
+@pytest.mark.usefixtures("request_header_provider_cleaner")
+def test_mlflow_config_with_request_header_provider_with_with_context(
+    fake_project,
+):
+
+    # emulate import of custom request header class
+    custom_rhp_txt = """
+from mlflow.tracking.request_header.abstract_request_header_provider import RequestHeaderProvider
+
+class CustomRequestHeaderProviderInitKwargsKedroContext(RequestHeaderProvider):
+    def __init__(self, kedro_context, b):
+        super().__init__()
+        self.context=kedro_context
+        self.b=b
+
+    def in_context(self):
+        pass
+    def request_headers(self):
+        pass
+"""
+
+    with open(fake_project / "custom_rhp.py", "w") as fhandler:
+        fhandler.write(custom_rhp_txt)
+
+    dict_config = dict(
+        server=dict(
+            mlflow_tracking_uri=None,  # not setup, not modified yet
+            credentials=None,
+            request_header_provider=dict(
+                type="custom_rhp.CustomRequestHeaderProviderInitKwargsKedroContext",
+                pass_context=True,
+                init_kwargs=dict(b=2),
+            ),
+        ),
+        tracking=dict(
+            disable_tracking=dict(pipelines=[]),
+            experiment=dict(name="Default", restore_if_deleted=True),
+            run=dict(id=None, name=None, nested=True),
+            params=dict(
+                dict_params=dict(
+                    flatten=False,
+                    recursive=True,
+                    sep=".",
+                ),
+                long_params_strategy="fail",
+            ),
+        ),
+        ui=dict(port="5000", host="127.0.0.1"),
+    )
+
+    _write_yaml(fake_project / "conf" / "local" / "mlflow.yml", dict_config)
+
+    bootstrap_project(fake_project)
+    with KedroSession.create("fake_package", fake_project) as session:
+        session.load_context()  # trigger setup and request_header_provider registration
+
+    assert (
+        mtrr._request_header_provider_registry._registry[-1].__class__.__name__
+        == "CustomRequestHeaderProviderInitKwargsKedroContext"
+    )
+    assert mtrr._request_header_provider_registry._registry[-1].b == "2"
+    assert hasattr(mtrr._request_header_provider_registry._registry[-1], "context")
+    assert isinstance(
+        mtrr._request_header_provider_registry._registry[-1].context, KedroContext
+    )
+
+
+def test_mlflow_config_with_bad_request_header_provider(fake_project):
+
+    # emulate import of custom request header class
+    # same as before, except CustomRequestHeaderProvider inherits from object
+    custom_rhp_txt = """
+class BadCustomRequestHeaderProvider():
+    def in_context(self):
+        pass
+    def request_headers(self):
+        pass
+"""
+    sys.path.append(fake_project.as_posix())
+    with open(fake_project / "bad_custom_rhp.py", "w") as fhandler:
+        fhandler.write(custom_rhp_txt)
+
+    dict_config = dict(
+        server=dict(
+            mlflow_tracking_uri=None,  # not setup, not modified yet
+            credentials=None,
+            request_header_provider=dict(
+                type="bad_custom_rhp.BadCustomRequestHeaderProvider"
+            ),
+        ),
+        tracking=dict(
+            disable_tracking=dict(pipelines=[]),
+            experiment=dict(name="Default", restore_if_deleted=True),
+            run=dict(id=None, name=None, nested=True),
+            params=dict(
+                dict_params=dict(
+                    flatten=False,
+                    recursive=True,
+                    sep=".",
+                ),
+                long_params_strategy="fail",
+            ),
+        ),
+        ui=dict(port="5000", host="127.0.0.1"),
+    )
+
+    _write_yaml(fake_project / "conf" / "local" / "mlflow.yml", dict_config)
+
+    bootstrap_project(fake_project)
+    with KedroSession.create("fake_package", fake_project) as session:
+        with pytest.raises(ValueError, match=r"should be a sublass of"):
+            session.load_context()  # trigger setup and request_header_provider registration
+
+    # cleanup test specific setup
+    (fake_project / "bad_custom_rhp.py").unlink()
+    sys.path.pop()
diff --git a/tests/config/test_kedro_mlflow_config.py b/tests/config/test_kedro_mlflow_config.py
index 8051591..f492803 100644
--- a/tests/config/test_kedro_mlflow_config.py
+++ b/tests/config/test_kedro_mlflow_config.py
@@ -10,14 +10,13 @@ from kedro_mlflow.config.kedro_mlflow_config import KedroMlflowConfig, _validate
 
 
 def test_kedro_mlflow_config_init():
-    # kedro_project_with_mlflow_conf is a global fixture in conftest
-
     config = KedroMlflowConfig()
     assert config.dict() == dict(
         server=dict(
             mlflow_tracking_uri=None,  # not setup, not modified yet
             mlflow_registry_uri=None,
             credentials=None,
+            request_header_provider=dict(type=None, pass_context=False, init_kwargs={}),
         ),
         tracking=dict(
             disable_tracking=dict(pipelines=[]),
diff --git a/tests/io/artifacts/test_mlflow_artifact_dataset.py b/tests/io/artifacts/test_mlflow_artifact_dataset.py
index fddb7ed..2a11b11 100644
--- a/tests/io/artifacts/test_mlflow_artifact_dataset.py
+++ b/tests/io/artifacts/test_mlflow_artifact_dataset.py
@@ -232,7 +232,6 @@ def test_artifact_dataset_load_with_run_id(tmp_path, tracking_uri, df1, df2):
 def test_artifact_dataset_load_with_run_id_and_artifact_path(
     tmp_path, tracking_uri, df1, artifact_path
 ):
-    print("artifact_path", artifact_path)
     mlflow.set_tracking_uri(tracking_uri.as_uri())
 
     # save first and retrieve run id