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

anemoi_utils-0.4.37/.github/workflows/pr-label-ats.yml

@@ -0,0 +1,66 @@
+# This workflow checks that the appropriate ATS labels are applied to a pull request:
+# - "ATS Approval Not Needed": pass
+# - "ATS Approved": pass
+# - "ATS Approval Needed": fail
+# - Missing ATS label: fail
+
+name: "[PR] ATS labels"
+
+on:
+  pull_request:
+    types:
+      - opened
+      - edited
+      - reopened
+      - labeled
+      - unlabeled
+      - synchronize
+
+permissions:
+  pull-requests: read
+
+jobs:
+  check:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Get PR details
+        uses: actions/github-script@v7
+        id: check-pr
+        with:
+          script: |
+            const pr = await github.rest.pulls.get({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              pull_number: context.payload.pull_request.number,
+            });
+            const labels = pr.data.labels.map(label => label.name);
+            core.setOutput('labels', JSON.stringify(labels));
+            core.setOutput('author', context.payload.pull_request.user.login);
+
+      - name: Evaluate ATS labels
+        run: |
+          AUTHOR='${{ steps.check-pr.outputs.author }}'
+          if [ "$AUTHOR" == "DeployDuck" ] || [ "$AUTHOR" == "pre-commit-ci[bot]" ]; then
+            echo "Bot PR, skipping."
+            exit 0
+          fi
+          LABELS=$(echo '${{ steps.check-pr.outputs.labels }}' | jq -r '.[]')
+          echo "Labels found:"
+          echo -e "$LABELS\n"
+          echo "Result:"
+          if echo "$LABELS" | grep -qi "ATS approval not needed"; then
+            echo "ATS approval not needed. Passing."
+            EXIT_CODE=0
+          elif echo "$LABELS" | grep -qi "ATS approved"; then
+            echo "ATS Approved. Passing."
+            EXIT_CODE=0
+          elif echo "$LABELS" | grep -qi "ATS approval needed"; then
+            echo "ATS Approval Needed. Failing."
+            EXIT_CODE=1
+          else
+            echo "No ATS approval labels found. Please assign the appropriate ATS label. Failing."
+            EXIT_CODE=1
+          fi
+          echo -e "\nFor more information on ATS labels, see:"
+          echo "https://anemoi.readthedocs.io/en/latest/contributing/guidelines.html#labelling-guidelines"
+          exit $EXIT_CODE

{anemoi_utils-0.4.35 → anemoi_utils-0.4.37}/.pre-commit-config.yaml

@@ -10,7 +10,7 @@ repos:
     entry: jupyter nbconvert --ClearOutputPreprocessor.enabled=True --inplace
     additional_dependencies: [jupyter]
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v5.0.0
+  rev: v6.0.0
   hooks:
   - id: check-yaml # Check YAML files for syntax errors only
     args: [--unsafe, --allow-multiple-documents]
@@ -41,7 +41,7 @@ repos:
     - --profile black
     - --project anemoi
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.12.7
+  rev: v0.12.11
   hooks:
   - id: ruff
     args:

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

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

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

@@ -8,6 +8,36 @@ 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)
+
+
+### Features
+
+* Add aliases to registry ([#219](https://github.com/ecmwf/anemoi-utils/issues/219)) ([37267b5](https://github.com/ecmwf/anemoi-utils/commit/37267b548556a796a01b43abb908011eeec85454))
+* Debug imports ([#182](https://github.com/ecmwf/anemoi-utils/issues/182)) ([1eaa615](https://github.com/ecmwf/anemoi-utils/commit/1eaa61540dc9ac3d5fe82f2c91b7fc98c8bb10af))
+* NoAuth for AML mlflow Logging ([#200](https://github.com/ecmwf/anemoi-utils/issues/200)) ([732182e](https://github.com/ecmwf/anemoi-utils/commit/732182ea5d255ba69ea2ed0a23b307d6f64aaf84))
+* Rich logging ([#209](https://github.com/ecmwf/anemoi-utils/issues/209)) ([3c762a5](https://github.com/ecmwf/anemoi-utils/commit/3c762a593ba2dc734becc54b92984d6dc62967ac))
+* Speedup checkpoint editing - remove compression ([#218](https://github.com/ecmwf/anemoi-utils/issues/218)) ([b49120f](https://github.com/ecmwf/anemoi-utils/commit/b49120f763b0b6ee10c805bab2aa7b973047f963))
+* Use obstore to access s3 buckets ([#210](https://github.com/ecmwf/anemoi-utils/issues/210)) ([da380be](https://github.com/ecmwf/anemoi-utils/commit/da380be71d78274d72bd0a3859ef00b1c80e9469))
+
+
+### Bug Fixes
+
+* Add missing s3 function used by datasets ([#212](https://github.com/ecmwf/anemoi-utils/issues/212)) ([30589e8](https://github.com/ecmwf/anemoi-utils/commit/30589e891fbdb1cff205f0350c63e93a725c7242))
+
 ## [0.4.35](https://github.com/ecmwf/anemoi-utils/compare/0.4.34...0.4.35) (2025-08-12)
 
 

{anemoi_utils-0.4.35/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.35
+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
@@ -225,6 +225,7 @@ Requires-Python: >=3.10
 License-File: LICENSE
 Requires-Dist: aniso8601
 Requires-Dist: deprecation
+Requires-Dist: entrypoints
 Requires-Dist: importlib-metadata; python_version < "3.10"
 Requires-Dist: multiurl
 Requires-Dist: numpy
@@ -256,7 +257,7 @@ Provides-Extra: provenance
 Requires-Dist: gitpython; extra == "provenance"
 Requires-Dist: nvsmi; extra == "provenance"
 Provides-Extra: s3
-Requires-Dist: boto3>1.36; extra == "s3"
+Requires-Dist: obstore; extra == "s3"
 Provides-Extra: tests
 Requires-Dist: anemoi-utils[mlflow]; extra == "tests"
 Requires-Dist: pytest; extra == "tests"

{anemoi_utils-0.4.35 → anemoi_utils-0.4.37}/pyproject.toml

@@ -41,6 +41,7 @@ dynamic = [ "version" ]
 dependencies = [
   "aniso8601",
   "deprecation",
+  "entrypoints",
   "importlib-metadata; python_version<'3.10'",
   "multiurl",
   "numpy",
@@ -73,7 +74,7 @@ optional-dependencies.mlflow = [ "mlflow-skinny>=2.11.1", "requests" ]
 optional-dependencies.provenance = [ "gitpython", "nvsmi" ]
 
 optional-dependencies.s3 = [
-  "boto3>1.36",
+  "obstore",
 ]
 
 optional-dependencies.tests = [ "anemoi-utils[mlflow]", "pytest", "pytest-mock>=3" ]

{anemoi_utils-0.4.35 → anemoi_utils-0.4.37}/src/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-0.4.35 → 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.35'
-__version_tuple__ = version_tuple = (0, 4, 35)
+__version__ = version = '0.4.37'
+__version_tuple__ = version_tuple = (0, 4, 37)
 
-__commit_id__ = commit_id = 'gd34987af9'
+__commit_id__ = commit_id = 'gceecf1ec8'

{anemoi_utils-0.4.35 → anemoi_utils-0.4.37}/src/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-0.4.35 → anemoi_utils-0.4.37}/src/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-0.4.35 → 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.35 → anemoi_utils-0.4.37}/src/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-0.4.35 → anemoi_utils-0.4.37}/src/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-0.4.35 → anemoi_utils-0.4.37}/src/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]")