anemoi-utils 0.4.35__py3-none-any.whl → 0.4.37__py3-none-any.whl

anemoi/utils/_environment.py

@@ -17,6 +17,9 @@ class Environment:
     ANEMOI_CONFIG_OVERRIDE_PATH: str
     """Path to the configuration override file for Anemoi."""
 
+    ANEMOI_DEBUG_IMPORTS: bool
+    """Enable debug imports to trace module loading."""
+
     def __setattr__(self, name: str, value: Any) -> None:
         raise AttributeError("Cannot set attributes on Environment class. Use environment variables instead.")
 

anemoi/utils/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.4.35'
-__version_tuple__ = version_tuple = (0, 4, 35)
+__version__ = version = '0.4.37'
+__version_tuple__ = version_tuple = (0, 4, 37)
 
 __commit_id__ = commit_id = None

anemoi/utils/checkpoints.py

@@ -214,7 +214,7 @@ def save_metadata(
             zipf.writestr(entry["path"], value.tobytes())
 
 
-def _edit_metadata(path: str, name: str, callback: Callable, supporting_arrays: dict = None) -> None:
+def _edit_metadata(path: str, name: str, callback: Callable, supporting_arrays: dict | None = None) -> None:
     """Edit metadata in a checkpoint file.
 
     Parameters
@@ -230,41 +230,56 @@ def _edit_metadata(path: str, name: str, callback: Callable, supporting_arrays:
     """
     new_path = f"{path}.anemoi-edit-{time.time()}-{os.getpid()}.tmp"
 
-    found = False
-
-    directory = None
-    with TemporaryDirectory() as temp_dir:
-        zipfile.ZipFile(path, "r").extractall(temp_dir)
-        total = 0
-        for root, dirs, files in os.walk(temp_dir):
-            for f in files:
-                total += 1
-                full = os.path.join(root, f)
-                if f == name:
-                    found = True
-                    callback(full)
-                    directory = os.path.dirname(full)
-
-        if not found:
+    with zipfile.ZipFile(path, "r") as source_zip:
+        file_list = source_zip.namelist()
+
+        # Find the target file and its directory
+        target_file = None
+        directory = None
+        for file_path in file_list:
+            if os.path.basename(file_path) == name:
+                target_file = file_path
+                directory = os.path.dirname(file_path)
+                break
+
+        if target_file is None:
             raise ValueError(f"Could not find '{name}' in {path}")
 
+        # Calculate total files for progress bar
+        total_files = len(file_list)
         if supporting_arrays is not None:
+            total_files += len(supporting_arrays)
+
+        with zipfile.ZipFile(new_path, "w", zipfile.ZIP_STORED) as new_zip:
+            with tqdm.tqdm(total=total_files, desc="Rebuilding checkpoint") as pbar:
+
+                # Copy all files except the target file
+                for file_path in file_list:
+                    if file_path != target_file:
+                        with source_zip.open(file_path) as source_file:
+                            data = source_file.read()
+                            new_zip.writestr(file_path, data)
+                        pbar.update(1)
 
-            for key, entry in supporting_arrays.items():
-                value = entry.tobytes()
-                fname = os.path.join(directory, f"{key}.numpy")
-                os.makedirs(os.path.dirname(fname), exist_ok=True)
-                with open(fname, "wb") as f:
-                    f.write(value)
-                    total += 1
-
-        with zipfile.ZipFile(new_path, "w", zipfile.ZIP_DEFLATED) as zipf:
-            with tqdm.tqdm(total=total, desc="Rebuilding checkpoint") as pbar:
-                for root, dirs, files in os.walk(temp_dir):
-                    for f in files:
-                        full = os.path.join(root, f)
-                        rel = os.path.relpath(full, temp_dir)
-                        zipf.write(full, rel)
+                # Handle the target file with callback
+                with TemporaryDirectory() as temp_dir:
+                    # Extract only the target file
+                    source_zip.extract(target_file, temp_dir)
+                    target_full_path = os.path.join(temp_dir, target_file)
+
+                    # Apply the callback
+                    callback(target_full_path)
+
+                    # Add the modified file to the new zip (if it still exists)
+                    if os.path.exists(target_full_path):
+                        new_zip.write(target_full_path, target_file)
+                    pbar.update(1)
+
+                # Add supporting arrays if provided
+                if supporting_arrays is not None:
+                    for key, entry in supporting_arrays.items():
+                        array_path = os.path.join(directory, f"{key}.numpy") if directory else f"{key}.numpy"
+                        new_zip.writestr(array_path, entry.tobytes())
                         pbar.update(1)
 
     os.rename(new_path, path)

anemoi/utils/cli.py

@@ -16,11 +16,26 @@ import sys
 import traceback
 from collections.abc import Callable
 
+from anemoi.utils import ENV
+
 try:
     import argcomplete
 except ImportError:
     argcomplete = None
 
+
+if ENV.ANEMOI_DEBUG_IMPORTS:
+    from datetime import datetime
+    from importlib.abc import MetaPathFinder
+
+    class ImportTracer(MetaPathFinder):
+        def find_spec(self, fullname, path, target=None):
+            now = datetime.now().isoformat(timespec="milliseconds")
+            print(f"[{now}] Importing {fullname} from {path}")
+            return None  # allow normal import processing to continue
+
+    sys.meta_path.insert(0, ImportTracer())
+
 LOG = logging.getLogger(__name__)
 
 
@@ -29,6 +44,10 @@ class Command:
 
     accept_unknown_args = False
 
+    def check(self, parser: argparse.ArgumentParser, args: argparse.Namespace) -> None:
+        """Check the command arguments."""
+        pass
+
     def run(self, args: argparse.Namespace) -> None:
         """Run the command.
 
@@ -72,6 +91,11 @@ def make_parser(description: str, commands: dict[str, Command]) -> argparse.Argu
         action="store_true",
         help="Debug mode",
     )
+    parser.add_argument(
+        "--rich",
+        action="store_true",
+        help="Use rich for logging",
+    )
 
     subparsers = parser.add_subparsers(help="commands:", dest="command")
     for name, command in commands.items():
@@ -216,16 +240,25 @@ def cli_main(
 
     cmd = commands[args.command]
 
-    logging.basicConfig(
-        format="%(asctime)s %(levelname)s %(message)s",
-        datefmt="%Y-%m-%d %H:%M:%S",
-        level=logging.DEBUG if args.debug else logging.INFO,
-    )
+    if args.rich:
+        from .logs import get_rich_handler
+
+        logging.basicConfig(
+            format="%(message)s", level=logging.DEBUG if args.debug else logging.INFO, handlers=[get_rich_handler()]
+        )
+    else:
+        logging.basicConfig(
+            format="%(asctime)s %(levelname)s %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+            level=logging.DEBUG if args.debug else logging.INFO,
+        )
 
     if unknown and not cmd.accept_unknown_args:
         # This should trigger an error
         parser.parse_args(test_arguments)
 
+    cmd.check(parser, args)
+
     try:
         if unknown:
             cmd.run(args, unknown)

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/commands/transfer.py

@@ -28,10 +28,14 @@ class Transfer(Command):
             The argument parser to which the arguments will be added.
         """
         command_parser.add_argument(
-            "--source", help="A path to a local file or folder or a URL to a file or a folder on S3."
+            "--source",
+            help="A path to a local file or folder or a URL to a file or a folder on S3.",
+            required=True,
         )
         command_parser.add_argument(
-            "--target", help="A path to a local file or folder or a URL to a file or a folder on S3 or a remote folder."
+            "--target",
+            help="A path to a local file or folder or a URL to a file or a folder on S3 or a remote folder.",
+            required=True,
         )
         command_parser.add_argument(
             "--overwrite",

anemoi/utils/config.py

@@ -15,7 +15,6 @@ import json
 import logging
 import os
 import threading
-import warnings
 from typing import Any
 
 import yaml
@@ -187,7 +186,6 @@ class DotDict(dict):
             The attribute value.
         """
 
-        self.warn_on_mutation(attr)
         value = self.convert_to_nested_dot_dict(value)
         super().__setitem__(attr, value)
 
@@ -201,13 +199,9 @@ class DotDict(dict):
         value : Any
             The value to set.
         """
-        self.warn_on_mutation(key)
         value = self.convert_to_nested_dot_dict(value)
         super().__setitem__(key, value)
 
-    def warn_on_mutation(self, key):
-        warnings.warn("Modifying an instance of DotDict(). This class is intended to be immutable.")
-
     def __repr__(self) -> str:
         """Return a string representation of the DotDict.
 

anemoi/utils/logs.py

@@ -10,10 +10,10 @@
 
 """Logging utilities."""
 
+import contextvars
 import logging
-import threading
 
-thread_local = threading.local()
+LOGGING_NAME = contextvars.ContextVar("logging_name", default="main")
 
 
 LOGGER = logging.getLogger(__name__)
@@ -27,7 +27,7 @@ def set_logging_name(name: str) -> None:
     name : str
         The name to set for logging.
     """
-    thread_local.logging_name = name
+    LOGGING_NAME.set(name)
 
 
 class ThreadCustomFormatter(logging.Formatter):
@@ -46,7 +46,7 @@ class ThreadCustomFormatter(logging.Formatter):
         str
             The formatted log record.
         """
-        record.logging_name = thread_local.logging_name
+        record.logging_name = LOGGING_NAME.get()
         return super().format(record)
 
 
@@ -58,11 +58,39 @@ def enable_logging_name(name: str = "main") -> None:
     name : str, optional
         The default logging name to set, by default "main".
     """
-    thread_local.logging_name = name
 
-    formatter = ThreadCustomFormatter("%(asctime)s - %(logging_name)s - %(levelname)s - %(message)s")
+    logger = logging.getLogger()
+    is_rich = any(handler.__class__.__name__ == "CustomRichHandler" for handler in logger.handlers)
+
+    set_logging_name(name)
+
+    if is_rich:
+        formatter = ThreadCustomFormatter("%(message)s")
+    else:
+        formatter = ThreadCustomFormatter("%(asctime)s - [%(logging_name)s] - %(levelname)s - %(message)s")
 
     logger = logging.getLogger()
 
     for handler in logger.handlers:
         handler.setFormatter(formatter)
+
+
+def get_rich_handler() -> logging.Handler:
+    """Return a RichHandler with custom formatting for logging."""
+
+    from rich.logging import RichHandler
+    from rich.text import Text
+
+    class CustomRichHandler(RichHandler):
+        def render_message(self, record, message):
+            global width
+
+            text = super().render_message(record, message)
+
+            if hasattr(record, "logging_name"):
+                name = record.logging_name
+                text = Text.assemble(f"[{name}]", (" → ", "dim"), text)
+
+            return text
+
+    return CustomRichHandler(log_time_format="[%X]")

anemoi/utils/mlflow/auth.py

@@ -13,6 +13,9 @@ from __future__ import annotations
 import logging
 import os
 import time
+import warnings
+from abc import ABC
+from abc import abstractmethod
 from datetime import datetime
 from datetime import timezone
 from functools import wraps
@@ -20,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
@@ -35,10 +43,96 @@ if TYPE_CHECKING:
     from collections.abc import Callable
 
 
-class TokenAuth:
+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."""
+
+    @abstractmethod
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def save(self, **kwargs):
+        pass
+
+    @abstractmethod
+    def login(self, force_credentials: bool = False, **kwargs):
+        pass
+
+    @abstractmethod
+    def authenticate(self, **kwargs):
+        pass
+
+
+class NoAuth(AuthBase):
+    """No-op authentication class."""
+
+    def __init__(self, *args, **kwargs):
+        self._enabled = False
+
+    def save(self, **kwargs):
+        pass
+
+    def login(self, force_credentials: bool = False, **kwargs):
+        pass
+
+    def authenticate(self, **kwargs):
+        pass
+
+
+class TokenAuth(AuthBase):
     """Manage authentication with a keycloak token server."""
 
-    config_file = "mlflow-token.json"
+    _config_file = "mlflow-token.json"
 
     def __init__(
         self,
@@ -63,10 +157,16 @@ class TokenAuth:
         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
 
@@ -86,17 +186,50 @@ class TokenAuth:
         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,
+        )
 
-        if os.path.exists(path) and os.stat(path).st_mode & 0o777 != 0o600:
-            os.chmod(path, 0o600)
+        store = TokenAuth._get_store()
 
-        return load_config(TokenAuth.config_file)
+        last = {}
+        for url, cfg in store.items():
+            if cfg.refresh_expires > last.get("refresh_expires", 0):
+                last = dict(url=url, **cfg.model_dump())
+
+        return last
 
     def enabled(fn: Callable) -> Callable:  # noqa: N805
         """Decorator to call or ignore a function based on the `enabled` flag."""
@@ -199,12 +332,15 @@ class TokenAuth:
             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(