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

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

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

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

@@ -8,6 +8,33 @@ 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.38](https://github.com/ecmwf/anemoi-utils/compare/0.4.37...0.4.38) (2025-10-22)
+
+
+### Features
+
+* **testing:** Add download timeout ([#230](https://github.com/ecmwf/anemoi-utils/issues/230)) ([721d114](https://github.com/ecmwf/anemoi-utils/commit/721d114f2702985d9fbabf68f384f7ccedb7cfb3))
+* **testing:** Sane test data download retries ([#227](https://github.com/ecmwf/anemoi-utils/issues/227)) ([1e08996](https://github.com/ecmwf/anemoi-utils/commit/1e089962b14f05f4aa56eb66a884bfe50ce60dcc))
+
+
+### Bug Fixes
+
+* Fix frequency_to_string outputing 108000s ([#216](https://github.com/ecmwf/anemoi-utils/issues/216)) ([5806a0c](https://github.com/ecmwf/anemoi-utils/commit/5806a0c996235fb3a19d34eeb25813280c5c989f))
+* Support dicts of supporting_arrays ([#229](https://github.com/ecmwf/anemoi-utils/issues/229)) ([9badbad](https://github.com/ecmwf/anemoi-utils/commit/9badbad360609d254717e0d32c6171beb903eb21))
+
+## [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.38}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anemoi-utils
-Version: 0.4.36
+Version: 0.4.38
 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.38}/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.38'
+__version_tuple__ = version_tuple = (0, 4, 38)
 
-__commit_id__ = commit_id = 'g67b40f99c'
+__commit_id__ = commit_id = 'g30b2d7452'

{anemoi_utils-0.4.36 → anemoi_utils-0.4.38}/src/anemoi/utils/checkpoints.py

@@ -20,6 +20,7 @@ import zipfile
 from collections.abc import Callable
 from tempfile import TemporaryDirectory
 
+import numpy as np
 import tqdm
 
 LOG = logging.getLogger(__name__)
@@ -147,6 +148,41 @@ def load_supporting_arrays(zipf: zipfile.ZipFile, entries: dict) -> dict:
     return supporting_arrays
 
 
+def _get_supporting_arrays_paths(directory: str, folder: str, supporting_arrays: dict | np.ndarray) -> dict:
+    """Get the paths of supporting arrays."""
+    if supporting_arrays is None:
+        return {}
+
+    if isinstance(supporting_arrays, dict):
+        return {
+            new_key: _get_supporting_arrays_paths(f"{directory}/{folder}", new_key, new_value)
+            for new_key, new_value in supporting_arrays.items()
+        }
+
+    return dict(
+        path=f"{directory}/{folder}.numpy",
+        shape=supporting_arrays.shape,
+        dtype=str(supporting_arrays.dtype),
+    )
+
+
+def _write_array_to_bytes(array: dict | np.ndarray, name: str, entry: dict, zipf: zipfile.ZipFile) -> None:
+    """Write a supporting array to bytes in a zip file."""
+    if isinstance(array, dict):
+        for sub_name, sub_array in array.items():
+            _write_array_to_bytes(sub_array, sub_name, entry[sub_name], zipf)
+        return None
+
+    LOG.info(
+        "Saving supporting array `%s` to %s (shape=%s, dtype=%s)",
+        name,
+        entry["path"],
+        entry["shape"],
+        entry["dtype"],
+    )
+    zipf.writestr(entry["path"], array.tobytes())
+
+
 def save_metadata(
     path: str, metadata: dict, *, supporting_arrays: dict = None, name: str = DEFAULT_NAME, folder: str = DEFAULT_FOLDER
 ) -> None:
@@ -189,29 +225,14 @@ def save_metadata(
         LOG.info("Saving metadata to %s/%s/%s", directory, folder, name)
 
         metadata = metadata.copy()
-        if supporting_arrays is not None:
-            metadata["supporting_arrays_paths"] = {
-                key: dict(path=f"{directory}/{folder}/{key}.numpy", shape=value.shape, dtype=str(value.dtype))
-                for key, value in supporting_arrays.items()
-            }
-        else:
-            metadata["supporting_arrays_paths"] = {}
+        metadata["supporting_arrays_paths"] = _get_supporting_arrays_paths(directory, folder, supporting_arrays)
 
         zipf.writestr(
             f"{directory}/{folder}/{name}",
             json.dumps(metadata),
         )
 
-        for name, entry in metadata["supporting_arrays_paths"].items():
-            value = supporting_arrays[name]
-            LOG.info(
-                "Saving supporting array `%s` to %s (shape=%s, dtype=%s)",
-                name,
-                entry["path"],
-                entry["shape"],
-                entry["dtype"],
-            )
-            zipf.writestr(entry["path"], value.tobytes())
+        _write_array_to_bytes(supporting_arrays, "", metadata["supporting_arrays_paths"], zipf)
 
 
 def _edit_metadata(path: str, name: str, callback: Callable, supporting_arrays: dict | None = None) -> None:

{anemoi_utils-0.4.36 → anemoi_utils-0.4.38}/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.38}/src/anemoi/utils/dates.py

@@ -260,39 +260,21 @@ def frequency_to_string(frequency: datetime.timedelta) -> str:
         A string representation of the frequency.
     """
 
-    frequency = frequency_to_timedelta(frequency)
+    total_seconds = int(frequency.total_seconds())
 
-    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)
 
-    seconds = total_seconds
+    if total_seconds % (24 * 3600) == 0:
+        return f"{total_seconds // (24 * 3600)}d"
 
-    days = seconds // (24 * 3600)
-    seconds %= 24 * 3600
-    hours = seconds // 3600
-    seconds %= 3600
-    minutes = seconds // 60
-    seconds %= 60
+    if total_seconds % 3600 == 0:
+        return f"{total_seconds // 3600}h"
 
-    if days > 0 and hours == 0 and minutes == 0 and seconds == 0:
-        return f"{days}d"
+    if total_seconds % 60 == 0:
+        return f"{total_seconds // 60}m"
 
-    if days == 0 and hours > 0 and minutes == 0 and seconds == 0:
-        return f"{hours}h"
-
-    if days == 0 and hours == 0 and minutes > 0 and seconds == 0:
-        return f"{minutes}m"
-
-    if days == 0 and hours == 0 and minutes == 0 and seconds > 0:
-        return f"{seconds}s"
-
-    if days > 0:
-        return f"{total_seconds}s"
-
-    return str(frequency)
+    return f"{total_seconds}s"
 
 
 def frequency_to_seconds(frequency: int | str | datetime.timedelta) -> int:

{anemoi_utils-0.4.36 → anemoi_utils-0.4.38}/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.38}/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.38}/src/anemoi/utils/testing.py

@@ -114,7 +114,7 @@ class GetTestData:
 
         LOG.info(f"Downloading test data from {url} to {target}")
 
-        download(url, target)
+        download(url, target, maximum_retries=5, retry_after=60, timeout=60)
 
         if gzipped:
             import gzip

{anemoi_utils-0.4.36 → anemoi_utils-0.4.38/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.38
 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.38}/tests/test_frequency.py

@@ -17,10 +17,11 @@ from anemoi.utils.dates import frequency_to_timedelta
 def test_frequency_to_string() -> None:
     """Test the frequency_to_string function for converting timedelta to string."""
     assert frequency_to_string(datetime.timedelta(hours=1)) == "1h"
-    assert frequency_to_string(datetime.timedelta(hours=1, minutes=30)) == "1:30:00"
+    assert frequency_to_string(datetime.timedelta(hours=1, minutes=30)) == "90m"
     assert frequency_to_string(datetime.timedelta(days=10)) == "10d"
+    assert frequency_to_string(datetime.timedelta(hours=30)) == "30h"
     assert frequency_to_string(datetime.timedelta(minutes=10)) == "10m"
-    assert frequency_to_string(datetime.timedelta(minutes=90)) == "1:30:00"
+    assert frequency_to_string(datetime.timedelta(minutes=90)) == "90m"
 
 
 def test_frequency_to_timedelta() -> None:

{anemoi_utils-0.4.36 → anemoi_utils-0.4.38}/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"
-}