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

{anemoi_utils-0.4.27 → anemoi_utils-0.4.29}/.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.27 → anemoi_utils-0.4.29}/.pre-commit-config.yaml

@@ -41,7 +41,7 @@ repos:
     - --profile black
     - --project anemoi
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.11.12
+  rev: v0.12.2
   hooks:
   - id: ruff
     args:
@@ -69,7 +69,7 @@ repos:
   hooks:
   - id: pyproject-fmt
 -   repo: https://github.com/jshwi/docsig # Check docstrings against function sig
-    rev: v0.69.3
+    rev: v0.70.0
     hooks:
     -   id: docsig
         args:

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

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

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

@@ -8,6 +8,34 @@ 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.29](https://github.com/ecmwf/anemoi-utils/compare/0.4.28...0.4.29) (2025-07-22)
+
+
+### Features
+
+* Better support for negative timedeltas ([#180](https://github.com/ecmwf/anemoi-utils/issues/180)) ([3f8041a](https://github.com/ecmwf/anemoi-utils/commit/3f8041a46b525b6fcbe6171cd8a8a40ec30b2c1f))
+* **deps:** Use mlflow-skinny instead of mlflow ([#184](https://github.com/ecmwf/anemoi-utils/issues/184)) ([82e5c30](https://github.com/ecmwf/anemoi-utils/commit/82e5c3053962cd8e1e8f6a1ea9e8f92492e497b4))
+* Protect mlflow token file ([#183](https://github.com/ecmwf/anemoi-utils/issues/183)) ([fdf0fc8](https://github.com/ecmwf/anemoi-utils/commit/fdf0fc84ee3e8076928f6c888374cd3aa008023b))
+* **sanitise:** Sanitation level ([#175](https://github.com/ecmwf/anemoi-utils/issues/175)) ([8d85d8f](https://github.com/ecmwf/anemoi-utils/commit/8d85d8fd889bf72b8066cc021d4d7b329a360848))
+* Support negative timedelta ([#178](https://github.com/ecmwf/anemoi-utils/issues/178)) ([546f6ec](https://github.com/ecmwf/anemoi-utils/commit/546f6ec76534cd39094957ce3b57b34f14f7a000))
+
+
+### Bug Fixes
+
+* Clean utils ([#185](https://github.com/ecmwf/anemoi-utils/issues/185)) ([de3c7a4](https://github.com/ecmwf/anemoi-utils/commit/de3c7a47f14c258997942564717c480caa124ee6))
+
+## [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.29}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anemoi-utils
-Version: 0.4.27
+Version: 0.4.29
 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,13 +250,18 @@ 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-skinny>=2.11.1; extra == "mlflow"
+Requires-Dist: requests; extra == "mlflow"
 Provides-Extra: provenance
 Requires-Dist: gitpython; extra == "provenance"
 Requires-Dist: nvsmi; extra == "provenance"
 Provides-Extra: s3
 Requires-Dist: boto3>1.36; extra == "s3"
 Provides-Extra: tests
+Requires-Dist: anemoi-utils[mlflow]; 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.29}/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-skinny>=2.11.1", "requests" ]
+
 optional-dependencies.provenance = [ "gitpython", "nvsmi" ]
 
 optional-dependencies.s3 = [
   "boto3>1.36",
 ]
 
-optional-dependencies.tests = [ "pytest" ]
+optional-dependencies.tests = [ "anemoi-utils[mlflow]", "pytest", "pytest-mock>=3" ]
 
 optional-dependencies.text = [ "termcolor", "wcwidth" ]
 

{anemoi_utils-0.4.27 → anemoi_utils-0.4.29}/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.29'
+__version_tuple__ = version_tuple = (0, 4, 29)

{anemoi_utils-0.4.27 → anemoi_utils-0.4.29}/src/anemoi/utils/dates.py

@@ -199,6 +199,15 @@ def as_timedelta(frequency: Union[int, str, datetime.timedelta]) -> datetime.tim
     except ValueError:
         pass
 
+    if frequency.startswith(" ") or frequency.startswith(" "):
+        frequency = frequency.strip()
+
+    if frequency.startswith("-"):
+        return -as_timedelta(frequency[1:])
+
+    if frequency.startswith("+"):
+        return as_timedelta(frequency[1:])
+
     if re.match(r"^\d+[hdms]$", frequency, re.IGNORECASE):
         unit = frequency[-1].lower()
         v = int(frequency[:-1])
@@ -261,6 +270,8 @@ def frequency_to_string(frequency: datetime.timedelta) -> str:
     frequency = frequency_to_timedelta(frequency)
 
     total_seconds = frequency.total_seconds()
+    if total_seconds < 0:
+        return f"-{frequency_to_string(-frequency)}"
     assert int(total_seconds) == total_seconds, total_seconds
     total_seconds = int(total_seconds)
 

anemoi_utils-0.4.29/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.29/src/anemoi/utils/mlflow/auth.py

@@ -0,0 +1,255 @@
+# (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 config_path
+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:
+        path = config_path(TokenAuth.config_file)
+
+        if not os.path.exists(path):
+            save_config(TokenAuth.config_file, {})
+
+        if os.path.exists(path) and os.stat(path).st_mode & 0o777 != 0o600:
+            os.chmod(path, 0o600)
+
+        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.29/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.29/src/anemoi/utils/mlflow/utils.py

@@ -0,0 +1,44 @@
+# (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 os
+
+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)