anemoi-utils 0.4.36__tar.gz → 0.4.37__tar.gz

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

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

{anemoi_utils-0.4.36 → anemoi_utils-0.4.37}/CHANGELOG.md

@@ -8,6 +8,19 @@ 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.37](https://github.com/ecmwf/anemoi-utils/compare/0.4.36...0.4.37) (2025-09-30)
+
+
+### Features
+
+* **mlflow auth:** Support for multiple servers ([#217](https://github.com/ecmwf/anemoi-utils/issues/217)) ([8ccfb1a](https://github.com/ecmwf/anemoi-utils/commit/8ccfb1ab063cccfec5852c386580036286b097c6))
+
+
+### Bug Fixes
+
+* Update s3 chunk size to 10 MB ([#220](https://github.com/ecmwf/anemoi-utils/issues/220)) ([aa20fa8](https://github.com/ecmwf/anemoi-utils/commit/aa20fa8b0b572fb6fa510b2f28c2b8b8a2f76d7c))
+* Use `yaml` and `json` flag in metadata get command ([#222](https://github.com/ecmwf/anemoi-utils/issues/222)) ([6af46c4](https://github.com/ecmwf/anemoi-utils/commit/6af46c4e715fc55aca374d2112976aa7d1bac589))
+
 ## [0.4.36](https://github.com/ecmwf/anemoi-utils/compare/0.4.35...0.4.36) (2025-09-22)
 
 

{anemoi_utils-0.4.36/src/anemoi_utils.egg-info → anemoi_utils-0.4.37}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anemoi-utils
-Version: 0.4.36
+Version: 0.4.37
 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

{anemoi_utils-0.4.36 → anemoi_utils-0.4.37}/src/anemoi/utils/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.4.36'
-__version_tuple__ = version_tuple = (0, 4, 36)
+__version__ = version = '0.4.37'
+__version_tuple__ = version_tuple = (0, 4, 37)
 
-__commit_id__ = commit_id = 'g67b40f99c'
+__commit_id__ = commit_id = 'gceecf1ec8'

{anemoi_utils-0.4.36 → anemoi_utils-0.4.37}/src/anemoi/utils/commands/metadata.py

@@ -139,13 +139,13 @@ class Metadata(Command):
         command_parser.add_argument(
             "--json",
             action="store_true",
-            help="Use the JSON format with ``--dump``, ``--view`` and ``--edit``.",
+            help="Use the JSON format with ``--dump``, ``--view``, ``--get`` and ``--edit``.",
         )
 
         command_parser.add_argument(
             "--yaml",
             action="store_true",
-            help="Use the YAML format with ``--dump``, ``--view`` and ``--edit``.",
+            help="Use the YAML format with ``--dump``, ``--view``, ``--get`` and ``--edit``.",
         )
 
     def run(self, args: Namespace) -> None:
@@ -315,7 +315,6 @@ class Metadata(Command):
         args : Namespace
             The arguments passed to the command.
         """
-        from pprint import pprint
 
         from anemoi.utils.checkpoints import load_metadata
 
@@ -335,7 +334,10 @@ class Metadata(Command):
 
         print(f"Metadata values for {args.get}: ", end="\n" if isinstance(metadata, (dict, list)) else "")
         if isinstance(metadata, dict):
-            pprint(metadata, indent=2, compact=True)
+            if args.yaml:
+                print(yaml.dump(metadata, indent=2, sort_keys=True))
+                return
+            print(json.dumps(metadata, indent=2, sort_keys=True))
         else:
             print(metadata)
 

{anemoi_utils-0.4.36 → anemoi_utils-0.4.37}/src/anemoi/utils/mlflow/auth.py

@@ -13,6 +13,7 @@ from __future__ import annotations
 import logging
 import os
 import time
+import warnings
 from abc import ABC
 from abc import abstractmethod
 from datetime import datetime
@@ -22,10 +23,15 @@ from getpass import getpass
 from typing import TYPE_CHECKING
 
 import requests
+from pydantic import BaseModel
+from pydantic import RootModel
+from pydantic import field_validator
+from pydantic import model_validator
 from requests.exceptions import HTTPError
 
+from ..config import CONFIG_LOCK
 from ..config import config_path
-from ..config import load_config
+from ..config import load_raw_config
 from ..config import save_config
 from ..remote import robust
 from ..timer import Timer
@@ -37,6 +43,56 @@ if TYPE_CHECKING:
     from collections.abc import Callable
 
 
+class ServerConfig(BaseModel):
+    refresh_token: str | None = None
+    refresh_expires: int = 0
+
+    @field_validator("refresh_expires", mode="before")
+    def to_int(cls, value: float | int) -> int:
+        if not isinstance(value, int):
+            return int(value)
+        return value
+
+
+class ServerStore(RootModel):
+    root: dict[str, ServerConfig] = {}
+
+    def get(self, url: str) -> ServerConfig | None:
+        return self.root.get(url)
+
+    def __getitem__(self, url: str) -> ServerConfig:
+        return self.root[url]
+
+    def items(self):
+        return self.root.items()
+
+    def update(self, url, config: ServerConfig) -> None:
+        """Update the server configuration for a given URL."""
+        self.root[url] = config
+
+    @property
+    def servers(self) -> list[tuple[str, int]]:
+        """List of servers in the store, as a tuple (url, refresh_expires). Ordered most recently used first."""
+        return [
+            (url, cfg.refresh_expires)
+            for url, cfg in sorted(
+                self.root.items(),
+                key=lambda item: item[1].refresh_expires,
+                reverse=True,
+            )
+        ]
+
+    @model_validator(mode="before")
+    @classmethod
+    def load_legacy_format(cls, data: dict) -> dict:
+        """Convert legacy single-server config format to multi-server."""
+        if isinstance(data, dict) and "url" in data:
+            _data = data.copy()
+            _url = _data.pop("url")
+            data = {_url: ServerConfig(**_data)}
+        return data
+
+
 class AuthBase(ABC):
     """Base class for authentication implementations."""
 
@@ -76,7 +132,7 @@ class NoAuth(AuthBase):
 class TokenAuth(AuthBase):
     """Manage authentication with a keycloak token server."""
 
-    config_file = "mlflow-token.json"
+    _config_file = "mlflow-token.json"
 
     def __init__(
         self,
@@ -101,10 +157,16 @@ class TokenAuth(AuthBase):
         self.target_env_var = target_env_var
         self._enabled = enabled
 
-        config = self.load_config()
+        store = self._get_store()
+        config = store.get(self.url)
+
+        if config is not None:
+            self._refresh_token = config.refresh_token
+            self.refresh_expires = config.refresh_expires
+        else:
+            self._refresh_token = None
+            self.refresh_expires = 0
 
-        self._refresh_token = config.get("refresh_token")
-        self.refresh_expires = config.get("refresh_expires", 0)
         self.access_token = None
         self.access_expires = 0
 
@@ -124,17 +186,50 @@ class TokenAuth(AuthBase):
         self._refresh_token = value
         self.refresh_expires = time.time() + (REFRESH_EXPIRE_DAYS * 86400)  # 86400 seconds in a day
 
+    @staticmethod
+    def _get_store() -> ServerStore:
+        """Read the server store from disk."""
+        with CONFIG_LOCK:
+            file = TokenAuth._config_file
+            path = config_path(file)
+
+            if not os.path.exists(path):
+                save_config(file, {})
+
+            if os.path.exists(path) and os.stat(path).st_mode & 0o777 != 0o600:
+                os.chmod(path, 0o600)
+
+            return ServerStore(**load_raw_config(file))
+
+    @staticmethod
+    def get_servers() -> list[tuple[str, int]]:
+        """List of all saved servers, as a tuple (url, refresh_expires). Ordered most recently used first."""
+        return TokenAuth._get_store().servers
+
     @staticmethod
     def load_config() -> dict:
-        path = config_path(TokenAuth.config_file)
+        """Load the last used server configuration
 
-        if not os.path.exists(path):
-            save_config(TokenAuth.config_file, {})
+        Returns
+        -------
+        config : dict
+            Dictionary with the following keys: `url`, `refresh_token`, `refresh_expires`.
+            If no configuration is found, an empty dictionary is returned.
+        """
+        warnings.warn(
+            "TokenAuth.load_config() is deprecated and will be removed in a future release.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+        store = TokenAuth._get_store()
 
-        if os.path.exists(path) and os.stat(path).st_mode & 0o777 != 0o600:
-            os.chmod(path, 0o600)
+        last = {}
+        for url, cfg in store.items():
+            if cfg.refresh_expires > last.get("refresh_expires", 0):
+                last = dict(url=url, **cfg.model_dump())
 
-        return load_config(TokenAuth.config_file)
+        return last
 
     def enabled(fn: Callable) -> Callable:  # noqa: N805
         """Decorator to call or ignore a function based on the `enabled` flag."""
@@ -237,12 +332,15 @@ class TokenAuth(AuthBase):
             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)
+        server_config = ServerConfig(
+            refresh_token=self.refresh_token,
+            refresh_expires=self.refresh_expires,
+        )
+
+        with CONFIG_LOCK:
+            store = self._get_store()
+            store.update(self.url, server_config)
+            save_config(self._config_file, store.model_dump())
 
         expire_date = datetime.fromtimestamp(self.refresh_expires, tz=timezone.utc)
         self.log.info(

{anemoi_utils-0.4.36 → anemoi_utils-0.4.37}/src/anemoi/utils/remote/s3.py

@@ -262,7 +262,7 @@ def upload_file(source: str, target: str, overwrite: bool, resume: bool, verbosi
         leave=verbosity >= 2,
         delay=0 if verbosity > 0 else 10,
     ) as pbar:
-        chunk_size = 1024 * 1024
+        chunk_size = 1024 * 1024 * 10
         total = size
         with open(source, "rb") as f:
             with closing(obstore.open_writer(s3, obj.key, buffer_size=chunk_size)) as g:
@@ -331,7 +331,7 @@ def download_file(source: str, target: str, overwrite: bool, resume: bool, verbo
         leave=verbosity >= 2,
         delay=0 if verbosity > 0 else 10,
     ) as pbar:
-        chunk_size = 1024 * 1024
+        chunk_size = 1024 * 1024 * 10
         total = size
         with closing(obstore.open_reader(s3, obj.key, buffer_size=chunk_size)) as f:
             with open(target, "wb") as g:

{anemoi_utils-0.4.36 → anemoi_utils-0.4.37/src/anemoi_utils.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anemoi-utils
-Version: 0.4.36
+Version: 0.4.37
 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

{anemoi_utils-0.4.36 → anemoi_utils-0.4.37}/tests/test_mlflow_auth.py

@@ -15,6 +15,8 @@ import time
 import pytest
 
 from anemoi.utils.mlflow.auth import NoAuth
+from anemoi.utils.mlflow.auth import ServerConfig
+from anemoi.utils.mlflow.auth import ServerStore
 from anemoi.utils.mlflow.auth import TokenAuth
 
 
@@ -35,17 +37,19 @@ def mocks(
     response.update(token_request)
 
     config = {
-        "refresh_token": "old_refresh_token",
-        "refresh_expires": time.time() + 3600,
+        "https://test.url": {
+            "refresh_token": "old_refresh_token",
+            "refresh_expires": time.time() + 3600,
+        }
     }
-    config.update(load_config)
+    config["https://test.url"].update(load_config)
 
     mock_token_request = mocker.patch(
         "anemoi.utils.mlflow.auth.TokenAuth._token_request",
         return_value=response,
     )
     mocker.patch(
-        "anemoi.utils.mlflow.auth.load_config",
+        "anemoi.utils.mlflow.auth.load_raw_config",
         return_value=config,
     )
     mocker.patch(
@@ -177,3 +181,108 @@ def test_noauth_methods_do_nothing():
     assert auth.save() is None
     assert auth.login() is None
     assert auth.authenticate() is None
+
+
+def test_legacy_format(mocker: pytest.MockerFixture) -> None:
+    mocks(mocker)
+
+    legacy_config = {
+        "url": "https://test.url",
+        "refresh_token": "some_refresh_token",
+        "refresh_expires": 123,
+    }
+    new_config = {
+        legacy_config["url"]: {
+            "refresh_token": legacy_config["refresh_token"],
+            "refresh_expires": legacy_config["refresh_expires"],
+        }
+    }
+    mocker.patch(
+        "anemoi.utils.mlflow.auth.load_raw_config",
+        return_value=legacy_config,
+    )
+
+    # test backwards compatibility of deprecated load_config
+    # when this function is removed, also remove this assert
+    config = TokenAuth.load_config()
+    assert config == legacy_config
+
+    # test that the store can handle both formats and the outputs are identical
+    legacy_store = ServerStore(legacy_config)
+    new_store = ServerStore(new_config)
+    expected_config = new_config["https://test.url"]
+
+    assert (
+        legacy_store["https://test.url"].model_dump() == new_store["https://test.url"].model_dump() == expected_config
+    )
+    assert legacy_store.model_dump() == new_store.model_dump() == new_config
+
+
+multi_config = {
+    "https://server-1.url": {
+        "refresh_token": "refresh-token-1",
+        "refresh_expires": 1,
+    },
+    "https://server-3.url": {
+        "refresh_token": "refresh-token-3",
+        "refresh_expires": 3,
+    },
+    "https://server-2.url": {
+        "refresh_token": "refresh-token-2",
+        "refresh_expires": 2,
+    },
+}
+
+
+@pytest.mark.parametrize(
+    "url, unknown",
+    [
+        ("https://server-1.url", False),
+        ("https://server-2.url", False),
+        ("https://server-3.url", False),
+        ("https://unknown.url", True),
+    ],
+)
+def test_multi_server_format(mocker: pytest.MockerFixture, url: str, unknown: bool) -> None:
+    mocks(mocker)
+
+    mocker.patch(
+        "anemoi.utils.mlflow.auth.load_raw_config",
+        return_value=multi_config,
+    )
+
+    auth = TokenAuth(url)
+
+    if unknown:
+        assert auth.refresh_token is None
+        assert auth.refresh_expires == 0
+    else:
+        assert auth.refresh_token == multi_config[url]["refresh_token"]
+        assert auth.refresh_expires == multi_config[url]["refresh_expires"]
+
+
+def test_server_store() -> None:
+    store = ServerStore(multi_config)
+
+    config = store["https://server-2.url"]
+    assert isinstance(config, ServerConfig)
+    assert config.refresh_token == "refresh-token-2"
+    assert config.refresh_expires == 2
+
+    assert store.get("https://unknown.url") is None
+
+    # ordered by expiry time, highest first
+    assert store.servers == [("https://server-3.url", 3), ("https://server-2.url", 2), ("https://server-1.url", 1)]
+
+    assert ServerStore({}).model_dump() == {}
+
+
+def test_utils_interface():
+    """TokenAuth uses the utils CONFIG_LOCK when reading and writing the server store to ensure thread safety.
+    Ensure that CONFIG_LOCK stays a reentrant lock, if it were a normal lock it would deadlock itself.
+    """
+    from threading import RLock
+
+    from anemoi.utils.config import CONFIG_LOCK
+
+    assert isinstance(CONFIG_LOCK, type(RLock()))

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

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