anemoi-utils 0.4.27__tar.gz → 0.4.28__tar.gz

{anemoi_utils-0.4.27 → anemoi_utils-0.4.28}/.github/pull_request_template.md

@@ -11,3 +11,5 @@
 <!-- Include any additional information, caveats, or considerations that the reviewer should be aware of. -->
 
 ***As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting  (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/***
+
+By opening this pull request, I affirm that all authors agree to the [Contributor License Agreement.](https://github.com/ecmwf/codex/blob/main/Legal/contributor_license_agreement.md)

anemoi_utils-0.4.28/.release-please-manifest.json

@@ -0,0 +1,3 @@
+{
+    ".": "0.4.28"
+}

{anemoi_utils-0.4.27 → anemoi_utils-0.4.28}/CHANGELOG.md

@@ -8,6 +8,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 Please add your functional changes to the appropriate section in the PR.
 Keep it human-readable, your future self will thank you!
 
+## [0.4.28](https://github.com/ecmwf/anemoi-utils/compare/0.4.27...0.4.28) (2025-07-03)
+
+
+### Features
+
+* Migrate mlflow utils from anemoi-training ([#174](https://github.com/ecmwf/anemoi-utils/issues/174)) ([0b7767b](https://github.com/ecmwf/anemoi-utils/commit/0b7767bc23486b140ad7423e3c5c7d5857cef71c))
+
+
+### Bug Fixes
+
+* Treat mlflow as an optional dependency ([#177](https://github.com/ecmwf/anemoi-utils/issues/177)) ([feb1088](https://github.com/ecmwf/anemoi-utils/commit/feb1088169a29f42032bf26d5c43f9817557bafc))
+
 ## [0.4.27](https://github.com/ecmwf/anemoi-utils/compare/0.4.26...0.4.27) (2025-06-27)
 
 

{anemoi_utils-0.4.27/src/anemoi_utils.egg-info → anemoi_utils-0.4.28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anemoi-utils
-Version: 0.4.27
+Version: 0.4.28
 Summary: A package to hold various functions to support training of ML models on ECMWF data.
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 License:                                  Apache License
@@ -236,7 +236,7 @@ Requires-Dist: rich
 Requires-Dist: tomli; python_version < "3.11"
 Requires-Dist: tqdm
 Provides-Extra: all
-Requires-Dist: anemoi-utils[grib,provenance,s3,text]; extra == "all"
+Requires-Dist: anemoi-utils[grib,mlflow,provenance,s3,text]; extra == "all"
 Provides-Extra: dev
 Requires-Dist: anemoi-utils[all,docs,tests]; extra == "dev"
 Provides-Extra: docs
@@ -250,6 +250,9 @@ Requires-Dist: sphinx-rtd-theme; extra == "docs"
 Requires-Dist: termcolor; extra == "docs"
 Provides-Extra: grib
 Requires-Dist: requests; extra == "grib"
+Provides-Extra: mlflow
+Requires-Dist: mlflow>=2.11.1; extra == "mlflow"
+Requires-Dist: requests; extra == "mlflow"
 Provides-Extra: provenance
 Requires-Dist: gitpython; extra == "provenance"
 Requires-Dist: nvsmi; extra == "provenance"
@@ -257,6 +260,7 @@ Provides-Extra: s3
 Requires-Dist: boto3>1.36; extra == "s3"
 Provides-Extra: tests
 Requires-Dist: pytest; extra == "tests"
+Requires-Dist: pytest-mock>=3; extra == "tests"
 Provides-Extra: text
 Requires-Dist: termcolor; extra == "text"
 Requires-Dist: wcwidth; extra == "text"

{anemoi_utils-0.4.27 → anemoi_utils-0.4.28}/pyproject.toml

@@ -53,7 +53,7 @@ dependencies = [
   "tqdm",
 ]
 
-optional-dependencies.all = [ "anemoi-utils[grib,provenance,text,s3]" ]
+optional-dependencies.all = [ "anemoi-utils[grib,provenance,text,s3,mlflow]" ]
 optional-dependencies.dev = [ "anemoi-utils[all,docs,tests]" ]
 
 optional-dependencies.docs = [
@@ -69,13 +69,15 @@ optional-dependencies.docs = [
 
 optional-dependencies.grib = [ "requests" ]
 
+optional-dependencies.mlflow = [ "mlflow>=2.11.1", "requests" ]
+
 optional-dependencies.provenance = [ "gitpython", "nvsmi" ]
 
 optional-dependencies.s3 = [
   "boto3>1.36",
 ]
 
-optional-dependencies.tests = [ "pytest" ]
+optional-dependencies.tests = [ "pytest", "pytest-mock>=3" ]
 
 optional-dependencies.text = [ "termcolor", "wcwidth" ]
 

{anemoi_utils-0.4.27 → anemoi_utils-0.4.28}/src/anemoi/utils/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.27'
-__version_tuple__ = version_tuple = (0, 4, 27)
+__version__ = version = '0.4.28'
+__version_tuple__ = version_tuple = (0, 4, 28)

anemoi_utils-0.4.28/src/anemoi/utils/mlflow/__init__.py

@@ -0,0 +1,8 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# 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.

anemoi_utils-0.4.28/src/anemoi/utils/mlflow/auth.py

@@ -0,0 +1,246 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# 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
+
+import logging
+import os
+import time
+from datetime import datetime
+from datetime import timezone
+from functools import wraps
+from getpass import getpass
+from typing import TYPE_CHECKING
+
+import requests
+from requests.exceptions import HTTPError
+
+from ..config import load_config
+from ..config import save_config
+from ..remote import robust
+from ..timer import Timer
+
+REFRESH_EXPIRE_DAYS = 29
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class TokenAuth:
+    """Manage authentication with a keycloak token server."""
+
+    config_file = "mlflow-token.json"
+
+    def __init__(
+        self,
+        url: str,
+        enabled: bool = True,
+        target_env_var: str = "MLFLOW_TRACKING_TOKEN",
+    ) -> None:
+        """Initialise the token authentication object.
+
+        Parameters
+        ----------
+        url : str
+            URL of the authentication server.
+        enabled : bool, optional
+            Set this to False to turn off authentication, by default True
+        target_env_var : str, optional
+            The environment variable to store the access token in after authenticating,
+            by default `MLFLOW_TRACKING_TOKEN`
+
+        """
+        self.url = url
+        self.target_env_var = target_env_var
+        self._enabled = enabled
+
+        config = self.load_config()
+
+        self._refresh_token = config.get("refresh_token")
+        self.refresh_expires = config.get("refresh_expires", 0)
+        self.access_token = None
+        self.access_expires = 0
+
+        # the command line tool adds a default handler to the root logger on runtime,
+        # so we init our logger here (on runtime, not on import) to avoid duplicate handlers
+        self.log = logging.getLogger(__name__)
+
+    def __call__(self) -> None:
+        self.authenticate()
+
+    @property
+    def refresh_token(self) -> str:
+        return self._refresh_token
+
+    @refresh_token.setter
+    def refresh_token(self, value: str) -> None:
+        self._refresh_token = value
+        self.refresh_expires = time.time() + (REFRESH_EXPIRE_DAYS * 86400)  # 86400 seconds in a day
+
+    @staticmethod
+    def load_config() -> dict:
+        return load_config(TokenAuth.config_file)
+
+    def enabled(fn: Callable) -> Callable:  # noqa: N805
+        """Decorator to call or ignore a function based on the `enabled` flag."""
+
+        @wraps(fn)
+        def _wrapper(self: TokenAuth, *args, **kwargs) -> Callable | None:
+            if self._enabled:
+                return fn(self, *args, **kwargs)
+            return None
+
+        return _wrapper
+
+    @enabled
+    def login(self, force_credentials: bool = False, **kwargs: dict) -> None:
+        """Acquire a new refresh token and save it to disk.
+
+        If an existing valid refresh token is already on disk it will be used.
+        If not, or the token has expired, the user will be asked to obtain one from the API.
+
+        Refresh token expiry time is set in the `REFRESH_EXPIRE_DAYS` constant (default 29 days).
+
+        This function should be called once, interactively, right before starting a training run.
+
+        Parameters
+        ----------
+        force_credentials : bool, optional
+            Force a credential login even if a refreh token is available, by default False.
+        kwargs : dict
+            Additional keyword arguments.
+
+        Raises
+        ------
+        RuntimeError
+            A new refresh token could not be acquired.
+
+        """
+        del kwargs  # unused
+        self.log.info("🌐 Logging in to %s", self.url)
+        new_refresh_token = None
+
+        if not force_credentials and self.refresh_token and self.refresh_expires > time.time():
+            new_refresh_token = self._token_request(ignore_exc=True).get("refresh_token")
+
+        if not new_refresh_token:
+            self.log.info("📝 Please obtain a seed refresh token from %s/seed", self.url)
+            self.log.info("📝 and paste it here (you will not see the output, just press enter after pasting):")
+            self.refresh_token = getpass("Refresh Token: ")
+
+            # perform a new refresh token request to check if the seed refresh token is valid
+            new_refresh_token = self._token_request().get("refresh_token")
+
+        if not new_refresh_token:
+            msg = "❌ Failed to log in. Please try again."
+            raise RuntimeError(msg)
+
+        self.refresh_token = new_refresh_token
+        self.save()
+
+        self.log.info("✅ Successfully logged in to MLflow. Happy logging!")
+
+    @enabled
+    def authenticate(self, **kwargs: dict) -> None:
+        """Check the access token and refresh it if necessary. A new refresh token will also be acquired upon refresh.
+
+        This requires a valid refresh token to be available, obtained from the `login` method.
+
+        The access token is stored in memory and in an environment variable.
+        If the access token is still valid, this function does nothing.
+
+        This function should be called before every MLflow API request.
+
+        Raises
+        ------
+        RuntimeError
+            No refresh token is available or the token request failed.
+
+        """
+        del kwargs  # unused
+        if self.access_expires > time.time():
+            return
+
+        if not self.refresh_token or self.refresh_expires < time.time():
+            msg = "You are not logged in to MLflow. Please log in first."
+            raise RuntimeError(msg)
+
+        with Timer("Access token refreshed", self.log):
+            response = self._token_request()
+
+        self.access_token = response.get("access_token")
+        self.access_expires = time.time() + (response.get("expires_in") * 0.7)  # bit of buffer
+        self.refresh_token = response.get("refresh_token")
+
+        os.environ[self.target_env_var] = self.access_token
+
+    @enabled
+    def save(self, **kwargs: dict) -> None:
+        """Save the latest refresh token to disk."""
+        del kwargs  # unused
+        if not self.refresh_token:
+            self.log.warning("No refresh token to save.")
+            return
+
+        config = {
+            "url": self.url,
+            "refresh_token": self.refresh_token,
+            "refresh_expires": self.refresh_expires,
+        }
+        save_config(self.config_file, config)
+
+        expire_date = datetime.fromtimestamp(self.refresh_expires, tz=timezone.utc)
+        self.log.info(
+            "Your MLflow login token is valid until %s UTC",
+            expire_date.strftime("%Y-%m-%d %H:%M:%S"),
+        )
+
+    def _token_request(
+        self,
+        ignore_exc: bool = False,
+    ) -> dict:
+        path = "refreshtoken"
+        payload = {"refresh_token": self.refresh_token}
+
+        try:
+            response = self._request(path, payload)
+        except Exception:
+            if ignore_exc:
+                return {}
+            raise
+
+        return response
+
+    def _request(self, path: str, payload: dict) -> dict:
+
+        headers = {
+            "Content-Type": "application/x-www-form-urlencoded",
+        }
+
+        try:
+            response = robust(requests.post)(
+                f"{self.url}/{path}",
+                headers=headers,
+                json=payload,
+                timeout=60,
+            )
+            response.raise_for_status()
+            response_json = response.json()
+
+            if response_json.get("status", "") != "OK":
+                error_description = response_json.get("response", "Error acquiring token.")
+                msg = f"❌ {error_description}"
+                raise RuntimeError(msg)
+
+            return response_json["response"]
+        except HTTPError:
+            self.log.exception("HTTP error occurred")
+            raise

anemoi_utils-0.4.28/src/anemoi/utils/mlflow/client.py

@@ -0,0 +1,76 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# 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
+
+try:
+    from mlflow import MlflowClient
+except ImportError:
+    raise ImportError(
+        "The `mlflow` package is required to use AnemoiMLflowclient. Please install it with `pip install mlflow`."
+    )
+
+from .auth import TokenAuth
+from .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
+
+    def login(self, force_credentials: bool = False, **kwargs) -> None:
+        """Explicitly log in to the MLflow server by acquiring or refreshing the token.
+
+        Parameters
+        ----------
+        force_credentials : bool, optional
+            Force a credential login even if a refresh token is available, by default False.
+        kwargs : dict
+            Additional keyword arguments passed to the underlying TokenAuth.login.
+        """
+        self.anemoi_auth.login(force_credentials=force_credentials, **kwargs)

anemoi_utils-0.4.28/src/anemoi/utils/mlflow/utils.py

@@ -0,0 +1,159 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# 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
+
+import functools
+import os
+from typing import Any
+
+import requests
+
+from ..remote import robust
+
+
+def health_check(tracking_uri: str) -> None:
+    """Query the health endpoint of an MLflow server.
+
+    If the server is not reachable, raise an error and remind the user that authentication may be required.
+
+    Raises
+    ------
+    ConnectionError
+        If the server is not reachable.
+
+    """
+    token = os.getenv("MLFLOW_TRACKING_TOKEN")
+
+    headers = {"Authorization": f"Bearer {token}"}
+    response = robust(requests.get, retry_after=30, maximum_tries=10)(
+        f"{tracking_uri}/health",
+        headers=headers,
+        timeout=60,
+    )
+
+    if response.text == "OK":
+        return
+
+    error_msg = f"Could not connect to MLflow server at {tracking_uri}. "
+    if not token:
+        error_msg += "The server may require authentication, did you forget to turn it on?"
+    raise ConnectionError(error_msg)
+
+
+def expand_iterables(
+    params: dict[str, Any],
+    *,
+    size_threshold: int | None = None,
+    recursive: bool = True,
+    delimiter: str = ".",
+) -> dict[str, Any]:
+    """Expand any iterable values to the form {key.i: value_i}.
+
+    If expanded will also add {key.all: [value_0, value_1, ...], key.length: len([value_0, value_1, ...])}.
+
+    If `size_threshold` is not None, expand the iterable only if the length of str(value) is
+    greater than `size_threshold`.
+
+    Parameters
+    ----------
+    params : dict[str, Any]
+        Parameters to be expanded.
+    size_threshold : int | None, optional
+        Threshold of str(value) to expand iterable at.
+        Default is None.
+    recursive : bool, optional
+        Expand nested dictionaries.
+        Default is True.
+    delimiter: str, optional
+        Delimiter to use for keys.
+        Default is ".".
+
+    Returns
+    -------
+    dict[str, Any]
+        Dictionary with all iterable values expanded.
+
+    Examples
+    --------
+        >>> expand_iterables({'a': ['a', 'b', 'c']})
+        {'a.0': 'a', 'a.1': 'b', 'a.2': 'c', 'a.all': ['a', 'b', 'c'], 'a.length': 3}
+        >>> expand_iterables({'a': {'b': ['a', 'b', 'c']}})
+        {'a': {'b.0': 'a', 'b.1': 'b', 'b.2': 'c', 'b.all': ['a', 'b', 'c'], 'b.length': 3}}
+        >>> expand_iterables({'a': ['a', 'b', 'c']}, size_threshold=100)
+        {'a': ['a', 'b', 'c']}
+        >>> expand_iterables({'a': [[0,1,2], 'b', 'c']})
+        {'a.0': {0: 0, 1: 1, 2: 2}, 'a.1': 'b', 'a.2': 'c', 'a.all': [[0, 1, 2], 'b', 'c'], 'a.length': 3}
+    """
+
+    def should_be_expanded(x: Any) -> bool:
+        return size_threshold is None or len(str(x)) > size_threshold
+
+    nested_func = functools.partial(expand_iterables, size_threshold=size_threshold, recursive=recursive)
+
+    def expand(val: dict | list) -> dict[str, Any]:
+        if not recursive:
+            return val
+        if isinstance(val, dict):
+            return nested_func(val)
+        if isinstance(val, list):
+            return nested_func(dict(enumerate(val)))
+        return val
+
+    expanded_params = {}
+
+    for key, value in params.items():
+        if isinstance(value, (list, tuple)):
+            if should_be_expanded(value):
+                for i, v in enumerate(value):
+                    expanded_params[f"{key}{delimiter}{i}"] = expand(v)
+
+                expanded_params[f"{key}{delimiter}all"] = value
+                expanded_params[f"{key}{delimiter}length"] = len(value)
+            else:
+                expanded_params[key] = value
+        else:
+            expanded_params[key] = expand(value)
+    return expanded_params
+
+
+def clean_config_params(params: dict[str, Any]) -> dict[str, Any]:
+    """Clean up params to avoid issues with mlflow.
+
+    Too many logged params will make the server take longer to render the
+    experiment.
+
+    Parameters
+    ----------
+    params : dict[str, Any]
+        Parameters to clean up.
+
+    Returns
+    -------
+    dict[str, Any]
+        Cleaned up params ready for MlFlow.
+    """
+    prefixes_to_remove = [
+        "hardware",
+        "data",
+        "dataloader",
+        "model",
+        "training",
+        "diagnostics",
+        "graph",
+        "metadata.config",
+        "config.dataset.sourcesmetadata.dataset.variables_metadata",
+        "metadata.dataset.sources",
+        "metadata.dataset.specific",
+        "metadata.dataset.variables_metadata",
+    ]
+
+    keys_to_remove = [key for key in params if any(key.startswith(prefix) for prefix in prefixes_to_remove)]
+    for key in keys_to_remove:
+        del params[key]
+    return params

{anemoi_utils-0.4.27 → anemoi_utils-0.4.28/src/anemoi_utils.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anemoi-utils
-Version: 0.4.27
+Version: 0.4.28
 Summary: A package to hold various functions to support training of ML models on ECMWF data.
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 License:                                  Apache License
@@ -236,7 +236,7 @@ Requires-Dist: rich
 Requires-Dist: tomli; python_version < "3.11"
 Requires-Dist: tqdm
 Provides-Extra: all
-Requires-Dist: anemoi-utils[grib,provenance,s3,text]; extra == "all"
+Requires-Dist: anemoi-utils[grib,mlflow,provenance,s3,text]; extra == "all"
 Provides-Extra: dev
 Requires-Dist: anemoi-utils[all,docs,tests]; extra == "dev"
 Provides-Extra: docs
@@ -250,6 +250,9 @@ Requires-Dist: sphinx-rtd-theme; extra == "docs"
 Requires-Dist: termcolor; extra == "docs"
 Provides-Extra: grib
 Requires-Dist: requests; extra == "grib"
+Provides-Extra: mlflow
+Requires-Dist: mlflow>=2.11.1; extra == "mlflow"
+Requires-Dist: requests; extra == "mlflow"
 Provides-Extra: provenance
 Requires-Dist: gitpython; extra == "provenance"
 Requires-Dist: nvsmi; extra == "provenance"
@@ -257,6 +260,7 @@ Provides-Extra: s3
 Requires-Dist: boto3>1.36; extra == "s3"
 Provides-Extra: tests
 Requires-Dist: pytest; extra == "tests"
+Requires-Dist: pytest-mock>=3; extra == "tests"
 Provides-Extra: text
 Requires-Dist: termcolor; extra == "text"
 Requires-Dist: wcwidth; extra == "text"

{anemoi_utils-0.4.27 → anemoi_utils-0.4.28}/src/anemoi_utils.egg-info/SOURCES.txt

@@ -73,6 +73,10 @@ src/anemoi/utils/commands/transfer.py
 src/anemoi/utils/mars/__init__.py
 src/anemoi/utils/mars/mars.yaml
 src/anemoi/utils/mars/requests.py
+src/anemoi/utils/mlflow/__init__.py
+src/anemoi/utils/mlflow/auth.py
+src/anemoi/utils/mlflow/client.py
+src/anemoi/utils/mlflow/utils.py
 src/anemoi/utils/remote/__init__.py
 src/anemoi/utils/remote/s3.py
 src/anemoi/utils/remote/ssh.py
@@ -88,6 +92,9 @@ tests/test_caching.py
 tests/test_compatibility.py
 tests/test_dates.py
 tests/test_frequency.py
+tests/test_mlflow_auth.py
+tests/test_mlflow_client.py
+tests/test_mlflow_utils.py
 tests/test_provenance.py
 tests/test_remote.py
 tests/test_sanetise.py

{anemoi_utils-0.4.27 → anemoi_utils-0.4.28}/src/anemoi_utils.egg-info/requires.txt

@@ -15,7 +15,7 @@ importlib-metadata
 tomli
 
 [all]
-anemoi-utils[grib,provenance,s3,text]
+anemoi-utils[grib,mlflow,provenance,s3,text]
 
 [dev]
 anemoi-utils[all,docs,tests]
@@ -33,6 +33,10 @@ termcolor
 [grib]
 requests
 
+[mlflow]
+mlflow>=2.11.1
+requests
+
 [provenance]
 gitpython
 nvsmi
@@ -42,6 +46,7 @@ boto3>1.36
 
 [tests]
 pytest
+pytest-mock>=3
 
 [text]
 termcolor