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

Issue:
Add support for authorization with expiring tokens
## Description
In many of our integrations we deploy MLflow instances behind some authorization methods, is most cases - an OAuth2.0. Secured MLflow instances require HTTP requests to have the `Authorization: Bearer <token>` header.
Now, your plugin does a great job with the `server.credentials`, but it does not play well with tokens that are short-lived (e.g. 1h), which might expire during the execution (e.g. long running training job).

MLflow itself does expose the `mlflow.request_header_provider` plugin entrypoint (https://github.com/mlflow/mlflow/blob/master/mlflow/tracking/request_header/abstract_request_header_provider.py#L4), but it slightly collides with the usual use case - the ML project is not a Python package directly and it does not have `setup.py` which is required for the MLflow's plugin system to work.

## Context
Commercial deployments of MLflow in almost 100% cases are secured by some kind of authorization mechanism.

## Possible Implementation
Right now, we can attach our custom `request_header_provider` by accessing private field in MLflow and calling `register` with our implementation. We need to call it before kedro-mlflow hook, so we ended up with "disabling" hook for kedro-mlflow and then manually adding it in the project's `settings` after our `request_header_provider`.

```python
 from mlflow.tracking.request_header.registry import (
     _request_header_provider_registry,
 )
_request_header_provider_registry.register(OurCustomClass())
```

but it seems really "hacky".

I was wondering if you have any other ideas how this can be integrated into Kedro-MLflow? Maybe the plugin could expose some interface to dynamically (configurable in-project) provide the authorization headers to the requests? In a way, kedro-mlflow could have "always attached" `request_header_provider` which would act as a proxy for in-project configuration.

## Possible Alternatives
Ask `mlflow` to expose the header provider plugin in "dynamic" way, rather than only `setup.py`'s entrypoint.


Pull request description:
:sparkles: Add custom request_header_provider support for authorization with expiring tokens
## Description

Close #357

## Development notes

- add an extra key in ``mlflow.yml`` and ``KedroMlflwoConfig`` to specifuy a path to custom class which inherits from ``mlflow.tracking.request_header.abstract_request_header_provider.RequestHeaderProvider``
- automatically resolve and add this class to the mlflow private global ``mlflow.tracking.request_header.registry._request_header_provider_registry``

## 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 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


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