Repository: ecmwf/anemoi-training (language: python)

Issue:
Add an AnemoiMlflowClient with authentication support
### Is your feature request related to a problem? Please describe.

If you want to use the mlflow python API on a server with token authentication, you currently have to manually insert the authentication calls before every mlflow call, like this:

```python 
import mlflow
from anemoi.training.diagnostics.mlflow.auth import TokenAuth

tracking_uri = "https://mlflow.url"

auth = TokenAuth(tracking_uri)
client = mlflow.tracking.MlflowClient(tracking_uri)

# check if we already have a login token on disk
# this can also be done from the command line with `anemoi-training mlflow login`
auth.login()

# get an access token
auth.authenticate()
# do mlflow stuff
client.search_experiments()

# authenticate again to be sure the token didn't expire since last time
auth.authenticate()
# do more mlflow stuff
client.log_metric(...)
```


### Describe the solution you'd like

A custom `AnemoiMlflowClient` that inherits from `mlflow.tracking.MlflowClient` , that injects authentication calls before every mlflow call.

```python
from anemoi.training.somewhere import AnemoiMlflowClient

client = AnemoiMlflowClient(tracking_uri, authentication=True)

# do mlflow stuff with login + authentication built in
client.search_experiments()
client.log_metric(...)
```



Discussion / hints:
One possible downside with anemoi.training is that it has dependencies on torch lighting (and potentially other packages) that are not required for the code.
If we need to remove this dependency on torch lightning, we could move the mlflow code into [yet another] anemoi package. If we need to, in the future.
Thanks for thinking about this Gert! I have some very simple scripts for cleaning out old runs, which rely on getting run info from mlflow. Would be great if there was a way to do this without having to install the full set of anemoi-training dependencies into my environment. Thanks!
Thanks for the feedback, dependencies are an issue. 

You could do `pip install --no-deps anemoi-training` (sort of fine for one-off scripts, but very ugly)

Like Florian said we can move all the code elsewhere, utils or yet another package, but at the same time it makes sense to me that MLflow code lives here, because it is directly related to training. 

I think we need to solve this problem at the root, which is "we have training code that is related to training, and it makes sense that this lives in anemoi-training, but it doesn't need the training dependencies". MLflow will probably not be the only thing that falls in this category.


Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8826b89a..3ca30a81 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -26,6 +26,7 @@ Keep it human-readable, your future self will thank you!
 - Feature: Add configurable models [#50](https://github.com/ecmwf/anemoi-training/pulls/50)
 - Feature: Authentication support for mlflow sync - [#51](https://github.com/ecmwf/anemoi-training/pull/51)
 - Feature: Support training for datasets with missing time steps [#48](https://github.com/ecmwf/anemoi-training/pulls/48)
+- Feature: `AnemoiMlflowClient`, an mlflow client with authentication support [#86](https://github.com/ecmwf/anemoi-training/pull/86)
 - Long Rollout Plots
 
 ### Fixed
diff --git a/src/anemoi/training/diagnostics/mlflow/client.py b/src/anemoi/training/diagnostics/mlflow/client.py
new file mode 100644
index 00000000..5c49f929
--- /dev/null
+++ b/src/anemoi/training/diagnostics/mlflow/client.py
@@ -0,0 +1,56 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+from __future__ import annotations
+
+from typing import Any
+
+from mlflow import MlflowClient
+
+from anemoi.training.diagnostics.mlflow.auth import TokenAuth
+from anemoi.training.diagnostics.mlflow.utils import health_check
+
+
+class AnemoiMlflowClient(MlflowClient):
+    """Anemoi extension of the MLflow client with token authentication support."""
+
+    def __init__(
+        self,
+        tracking_uri: str,
+        *args,
+        authentication: bool = False,
+        check_health: bool = True,
+        **kwargs,
+    ) -> None:
+        """Behaves like a normal `mlflow.MlflowClient` but with token authentication injected on every call.
+
+        Parameters
+        ----------
+        tracking_uri : str
+            The URI of the MLflow tracking server.
+        authentication : bool, optional
+            Enable token authentication, by default False
+        check_health : bool, optional
+            Check the health of the MLflow server on init, by default True
+        *args : Any
+            Additional arguments to pass to the MLflow client.
+        **kwargs : Any
+            Additional keyword arguments to pass to the MLflow client.
+
+        """
+        self.anemoi_auth = TokenAuth(tracking_uri, enabled=authentication)
+        if check_health:
+            super().__getattribute__("anemoi_auth").authenticate()
+            health_check(tracking_uri)
+        super().__init__(tracking_uri, *args, **kwargs)
+
+    def __getattribute__(self, name: str) -> Any:
+        """Intercept attribute access and inject authentication."""
+        attr = super().__getattribute__(name)
+        if callable(attr) and name != "anemoi_auth":
+            super().__getattribute__("anemoi_auth").authenticate()
+        return attr


Test patch:
diff --git a/tests/diagnostics/mlflow/test_client.py b/tests/diagnostics/mlflow/test_client.py
new file mode 100644
index 00000000..93c0aeeb
--- /dev/null
+++ b/tests/diagnostics/mlflow/test_client.py
@@ -0,0 +1,38 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+if TYPE_CHECKING:
+    import pytest_mock
+
+from anemoi.training.diagnostics.mlflow.client import AnemoiMlflowClient
+
+
+@pytest.fixture(autouse=True)
+def mocks(mocker: pytest_mock.MockerFixture) -> None:
+    mocker.patch("anemoi.training.diagnostics.mlflow.client.TokenAuth")
+    mocker.patch("anemoi.training.diagnostics.mlflow.client.health_check")
+    mocker.patch("anemoi.training.diagnostics.mlflow.client.AnemoiMlflowClient.search_experiments")
+
+
+def test_auth_injected() -> None:
+    client = AnemoiMlflowClient("http://localhost:5000", authentication=True, check_health=False)
+    client.search_experiments()
+    client.search_experiments()
+
+    assert client.anemoi_auth.authenticate.call_count == 2
+
+
+def test_health_check() -> None:
+    # the internal health check will trigger an authenticate call
+    client = AnemoiMlflowClient("http://localhost:5000", authentication=True, check_health=True)
+
+    client.anemoi_auth.authenticate.assert_called_once()