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

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

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

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

@@ -8,6 +8,21 @@ 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.27](https://github.com/ecmwf/anemoi-utils/compare/0.4.26...0.4.27) (2025-06-27)
+
+
+### Features
+
+* Split s3 config from s3 client code ([#170](https://github.com/ecmwf/anemoi-utils/issues/170)) ([56dacb1](https://github.com/ecmwf/anemoi-utils/commit/56dacb19efa0979acd72edb72a95f058b69d612a))
+
+## [0.4.26](https://github.com/ecmwf/anemoi-utils/compare/0.4.25...0.4.26) (2025-06-25)
+
+
+### Features
+
+* Fixtures for temp dir handling for test data ([#166](https://github.com/ecmwf/anemoi-utils/issues/166)) ([2b9677f](https://github.com/ecmwf/anemoi-utils/commit/2b9677fffc5eba84876f974001b87b73c7e542af))
+* Move anemoi-inference metadata command to this package, add metadata removal options ([#167](https://github.com/ecmwf/anemoi-utils/issues/167)) ([cabb989](https://github.com/ecmwf/anemoi-utils/commit/cabb989bdd4154a0476acf48e1ac44099c91c6db))
+
 ## [0.4.25](https://github.com/ecmwf/anemoi-utils/compare/0.4.24...0.4.25) (2025-06-24)
 
 

{anemoi_utils-0.4.25/src/anemoi_utils.egg-info → anemoi_utils-0.4.27}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anemoi-utils
-Version: 0.4.25
+Version: 0.4.27
 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.25 → anemoi_utils-0.4.27}/src/anemoi/utils/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.25'
-__version_tuple__ = version_tuple = (0, 4, 25)
+__version__ = version = '0.4.27'
+__version_tuple__ = version_tuple = (0, 4, 27)

{anemoi_utils-0.4.25 → anemoi_utils-0.4.27}/src/anemoi/utils/checkpoints.py

@@ -308,7 +308,6 @@ def remove_metadata(path: str, *, name: str = DEFAULT_NAME) -> None:
     name : str, optional
         The name of the metadata file in the zip archive
     """
-    LOG.info("Removing metadata '%s' from %s", name, path)
 
     def callback(full):
         os.remove(full)

anemoi_utils-0.4.27/src/anemoi/utils/commands/metadata.py

@@ -0,0 +1,391 @@
+# (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.
+
+
+import json
+import logging
+import os
+import shutil
+import subprocess
+from argparse import ArgumentParser
+from argparse import Namespace
+from tempfile import TemporaryDirectory
+from typing import Any
+from typing import Dict
+
+import yaml
+
+from . import Command
+
+LOG = logging.getLogger(__name__)
+
+EDITOR_OPTIONS = {"code": ["--wait"]}
+
+
+class Metadata(Command):
+    """Edit, remove, dump or load metadata from a checkpoint file."""
+
+    def add_arguments(self, command_parser: ArgumentParser) -> None:
+        """Add command line arguments to the parser.
+
+        Parameters
+        ----------
+        command_parser : ArgumentParser
+            The argument parser to which the arguments will be added.
+        """
+        from anemoi.utils.checkpoints import DEFAULT_NAME
+
+        command_parser.add_argument("path", help="Path to the checkpoint.")
+
+        group = command_parser.add_mutually_exclusive_group(required=True)
+
+        group.add_argument(
+            "--dump",
+            action="store_true",
+            help=(
+                "Extract the metadata from the checkpoint and print it to the standard output"
+                " or the file specified by ``--output``, in JSON or YAML format."
+            ),
+        )
+        group.add_argument(
+            "--load",
+            action="store_true",
+            help=(
+                "Set the metadata in the checkpoint from the content"
+                " of a file specified by the ``--input`` argument."
+            ),
+        )
+
+        group.add_argument(
+            "--edit",
+            action="store_true",
+            help=(
+                "Edit the metadata in place, using the specified editor."
+                " See the ``--editor`` argument for more information."
+            ),
+        )
+
+        group.add_argument(
+            "--view",
+            action="store_true",
+            help=(
+                "View the metadata in place, using the specified pager."
+                " See the ``--pager`` argument for more information."
+            ),
+        )
+
+        group.add_argument(
+            "--remove",
+            action="store_true",
+            help="Remove the metadata from the checkpoint.",
+        )
+
+        group.add_argument(
+            "--supporting-arrays",
+            action="store_true",
+            help="Print the supporting arrays.",
+        )
+
+        group.add_argument(
+            "--get",
+            help="Navigate the metadata via dot-separated path.",
+        )
+
+        group.add_argument(
+            "--pytest",
+            action="store_true",
+            help=("Extract the metadata from the checkpoint so it can be added to the test suite."),
+        )
+
+        command_parser.add_argument(
+            "--name",
+            default=DEFAULT_NAME,
+            help="Name of metadata record to be used with the actions above.",
+        )
+
+        command_parser.add_argument(
+            "--input",
+            help="The output file name to be used by the ``--load`` option.",
+        )
+
+        command_parser.add_argument(
+            "--output",
+            help="The output file name to be used by the ``--dump`` option.",
+        )
+
+        command_parser.add_argument(
+            "--inplace",
+            action="store_true",
+            help="If set, update the source file in place instead of writing to a separate target.",
+        )
+
+        command_parser.add_argument(
+            "--editor",
+            help="Editor to use for the ``--edit`` option. Default to ``$EDITOR`` if defined, else ``vi``.",
+            default=os.environ.get("EDITOR", "vi"),
+        )
+
+        command_parser.add_argument(
+            "--pager",
+            help="Editor to use for the ``--view`` option. Default to ``$PAGER`` if defined, else ``less``.",
+            default=os.environ.get("PAGER", "less"),
+        )
+
+        command_parser.add_argument(
+            "--json",
+            action="store_true",
+            help="Use the JSON format with ``--dump``, ``--view`` and ``--edit``.",
+        )
+
+        command_parser.add_argument(
+            "--yaml",
+            action="store_true",
+            help="Use the YAML format with ``--dump``, ``--view`` and ``--edit``.",
+        )
+
+    def run(self, args: Namespace) -> None:
+        """Execute the command based on the provided arguments.
+
+        Parameters
+        ----------
+        args : Namespace
+            The arguments passed to the command.
+        """
+        if args.edit:
+            return self.edit(args)
+
+        if args.view:
+            return self.view(args)
+
+        if args.get:
+            return self.get(args)
+
+        if args.remove:
+            return self.remove(args)
+
+        if args.dump or args.pytest:
+            return self.dump(args)
+
+        if args.load:
+            return self.load(args)
+
+        if args.supporting_arrays:
+            return self.supporting_arrays(args)
+
+    def edit(self, args: Namespace) -> None:
+        """Edit the metadata in place using the specified editor.
+
+        Parameters
+        ----------
+        args : Namespace
+            The arguments passed to the command.
+        """
+        return self._edit(args, view=False, cmd=args.editor)
+
+    def view(self, args: Namespace) -> None:
+        """View the metadata in place using the specified pager.
+
+        Parameters
+        ----------
+        args : Namespace
+            The arguments passed to the command.
+        """
+        return self._edit(args, view=True, cmd=args.pager)
+
+    def _edit(self, args: Namespace, view: bool, cmd: str) -> None:
+        """Internal method to edit or view the metadata.
+
+        Parameters
+        ----------
+        args : Namespace
+            The arguments passed to the command.
+        view : bool
+            If True, view the metadata; otherwise, edit it.
+        cmd : str
+            The command to use for editing or viewing.
+        """
+        from anemoi.utils.checkpoints import load_metadata
+        from anemoi.utils.checkpoints import replace_metadata
+
+        kwargs: Dict[str, Any] = {}
+
+        if args.json:
+            ext = "json"
+            dump = json.dump
+            load = json.load
+            if args.test:
+                kwargs = {"sort_keys": True}
+            else:
+                kwargs = {"indent": 4, "sort_keys": True}
+        else:
+            ext = "yaml"
+            dump = yaml.dump
+            load = yaml.safe_load
+            kwargs = {"default_flow_style": False}
+
+        with TemporaryDirectory() as temp_dir:
+
+            path = os.path.join(temp_dir, f"checkpoint.{ext}")
+            metadata = load_metadata(args.path)
+
+            with open(path, "w") as f:
+                dump(metadata, f, **kwargs)
+
+            subprocess.check_call([cmd, *EDITOR_OPTIONS.get(cmd, []), path])
+
+            if not view:
+                with open(path) as f:
+                    edited = load(f)
+
+                if edited != metadata:
+                    replace_metadata(args.path, edited)
+                else:
+                    LOG.info("No changes made.")
+
+    def remove(self, args: Namespace) -> None:
+        """Remove the metadata from the checkpoint.
+
+        Parameters
+        ----------
+        args : Namespace
+            The arguments passed to the command.
+        """
+        from anemoi.utils.checkpoints import remove_metadata
+
+        if args.inplace and args.output:
+            raise ValueError("Only choose one of --inplace and --output")
+
+        LOG.info("Removing metadata from %s", args.path)
+
+        if args.inplace:
+            output = args.path
+        else:
+            if not args.output:
+                raise ValueError("Argument --output is required unless --inplace is set")
+
+            shutil.copy2(args.path, args.output)
+            output = args.output
+
+        LOG.info("Writing checkpoint at %s", output)
+        remove_metadata(output)
+
+    def dump(self, args: Namespace) -> None:
+        """Dump the metadata from the checkpoint to a file or standard output.
+
+        Parameters
+        ----------
+        args : Namespace
+            The arguments passed to the command.
+        """
+        from anemoi.utils.checkpoints import load_metadata
+
+        if args.output:
+            file = open(args.output, "w")
+        else:
+            file = None
+
+        metadata = load_metadata(args.path)
+        if args.pytest:
+            from anemoi.inference.testing.mock_checkpoint import minimum_mock_checkpoint
+
+            # We remove all unessential metadata for testing purposes
+            metadata = minimum_mock_checkpoint(metadata)
+
+        if args.yaml:
+            print(yaml.dump(metadata, indent=2, sort_keys=True), file=file)
+            return
+
+        if args.json or True:
+            if args.pytest:
+                print(json.dumps(metadata, sort_keys=True), file=file)
+            else:
+                print(json.dumps(metadata, indent=4, sort_keys=True), file=file)
+            return
+
+    def get(self, args: Namespace) -> None:
+        """Navigate and print the metadata via a dot-separated path.
+
+        Parameters
+        ----------
+        args : Namespace
+            The arguments passed to the command.
+        """
+        from pprint import pprint
+
+        from anemoi.utils.checkpoints import load_metadata
+
+        metadata = load_metadata(args.path, name=args.name)
+
+        if args.get == ".":
+            print("Metadata from root: ", list(metadata.keys()))
+            return
+
+        for key in args.get.split("."):
+            if key == "":
+                keys = list(metadata.keys())
+                print(f"Metadata keys from {args.get[:-1]}: ", keys)
+                return
+            else:
+                metadata = metadata[key]
+
+        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)
+        else:
+            print(metadata)
+
+    def load(self, args: Namespace) -> None:
+        """Load metadata into the checkpoint from a specified file.
+
+        Parameters
+        ----------
+        args : Namespace
+            The arguments passed to the command.
+        """
+        from anemoi.utils.checkpoints import has_metadata
+        from anemoi.utils.checkpoints import replace_metadata
+        from anemoi.utils.checkpoints import save_metadata
+
+        if args.input is None:
+            raise ValueError("Please specify a value for --input")
+
+        _, ext = os.path.splitext(args.input)
+        if ext == ".json" or args.json:
+            with open(args.input) as f:
+                metadata = json.load(f)
+
+        elif ext in (".yaml", ".yml") or args.yaml:
+            with open(args.input) as f:
+                metadata = yaml.safe_load(f)
+
+        else:
+            raise ValueError(f"Unknown file extension {ext}. Please specify --json or --yaml")
+
+        if has_metadata(args.path, name=args.name):
+            replace_metadata(args.path, metadata)
+        else:
+            save_metadata(args.path, metadata, name=args.name)
+
+    def supporting_arrays(self, args: Namespace) -> None:
+        """Print the supporting arrays from the metadata.
+
+        Parameters
+        ----------
+        args : Namespace
+            The arguments passed to the command.
+        """
+        from anemoi.utils.checkpoints import load_metadata
+
+        _, supporting_arrays = load_metadata(args.path, supporting_arrays=True)
+
+        for name, array in supporting_arrays.items():
+            print(f"{name}: shape={array.shape} dtype={array.dtype}")
+
+
+command = Metadata

{anemoi_utils-0.4.25 → anemoi_utils-0.4.27}/src/anemoi/utils/commands/transfer.py

@@ -17,7 +17,7 @@ from . import Command
 
 
 class Transfer(Command):
-    """Handle configuration related commands."""
+    """Transfer files or folders from the source to the target location."""
 
     def add_arguments(self, command_parser: ArgumentParser) -> None:
         """Add arguments to the command parser.

{anemoi_utils-0.4.25 → anemoi_utils-0.4.27}/src/anemoi/utils/remote/s3.py

@@ -47,8 +47,8 @@ SECRETS = ["aws_access_key_id", "aws_secret_access_key"]
 thread_local = threading.local()
 
 
-def s3_client(bucket: str, *, region: str = None, service: str = "s3") -> Any:
-    """Get an S3 client for the specified bucket and region.
+def _s3_config(bucket: str, *, region: str = None) -> Any:
+    """Get an S3 client config for the specified bucket and region.
 
     Parameters
     ----------
@@ -56,31 +56,15 @@ def s3_client(bucket: str, *, region: str = None, service: str = "s3") -> Any:
         The name of the S3 bucket.
     region : str, optional
         The AWS region of the S3 bucket.
-    service : str, optional
-        The AWS service to use, default is "s3".
 
     Returns
     -------
     Any
         The S3 client.
     """
-    import boto3
     from botocore import UNSIGNED
-    from botocore.client import Config
 
-    if not hasattr(thread_local, "s3_clients"):
-        thread_local.s3_clients = {}
-
-    key = f"{bucket}-{region}-{service}"
-
-    if key in thread_local.s3_clients:
-        return thread_local.s3_clients[key]
-
-    boto3_config = dict(
-        max_pool_connections=25,
-        request_checksum_calculation="when_required",
-        response_checksum_validation="when_required",
-    )
+    boto3_config = {}
 
     if region:
         # This is using AWS
@@ -129,6 +113,85 @@ def s3_client(bucket: str, *, region: str = None, service: str = "s3") -> Any:
             boto3_config.update(options["config"])
             del options["config"]
 
+    def _(options):
+
+        def __(k, v):
+            if k in SECRETS:
+                return "***"
+            return v
+
+        if isinstance(options, dict):
+            return {k: __(k, v) for k, v in options.items()}
+
+        if isinstance(options, list):
+            return [_(o) for o in options]
+
+        return options
+
+    LOG.debug(f"Using S3 options: {_(options)}")
+
+    return boto3_config, options
+
+
+def s3_options(bucket: str, *, region: str = None, service: str = "s3") -> dict:
+    """Get the S3 configuration for the specified bucket and region.
+
+    Parameters
+    ----------
+    bucket : str
+        The name of the S3 bucket.
+    region : str, optional
+        The AWS region of the S3 bucket.
+    service : str, optional
+        The AWS service to use, default is "s3".
+
+    Returns
+    -------
+    dict
+        The S3 configuration.
+    """
+    _, options = _s3_config(bucket, region=region)
+    return options
+
+
+def s3_client(bucket: str, *, region: str = None, service: str = "s3") -> Any:
+    """Get an S3 client for the specified bucket and region.
+
+    Parameters
+    ----------
+    bucket : str
+        The name of the S3 bucket.
+    region : str, optional
+        The AWS region of the S3 bucket.
+    service : str, optional
+        The AWS service to use, default is "s3".
+
+    Returns
+    -------
+    Any
+        The S3 client.
+    """
+    import boto3
+    from botocore.client import Config
+
+    if not hasattr(thread_local, "s3_clients"):
+        thread_local.s3_clients = {}
+
+    key = f"{bucket}-{region}-{service}"
+
+    if key in thread_local.s3_clients:
+        return thread_local.s3_clients[key]
+
+    boto3_config, options = _s3_config(bucket, region=region)
+
+    boto3_config.update(
+        dict(
+            max_pool_connections=25,
+            request_checksum_calculation="when_required",
+            response_checksum_validation="when_required",
+        )
+    )
+
     options["config"] = Config(**boto3_config)
 
     def _(options):

{anemoi_utils-0.4.25 → anemoi_utils-0.4.27}/src/anemoi/utils/testing.py

@@ -7,14 +7,12 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-import atexit
 import logging
 import os
 import shutil
-import tempfile
-import threading
 import warnings
 from functools import lru_cache
+from pathlib import Path
 
 import pytest
 from multiurl import download
@@ -25,30 +23,6 @@ LOG = logging.getLogger(__name__)
 
 TEST_DATA_URL = "https://object-store.os-api.cci1.ecmwf.int/ml-tests/test-data/samples/"
 
-lock = threading.RLock()
-TEMPORARY_DIRECTORY = None
-
-
-def _temporary_directory() -> str:
-    """Return a temporary directory in which to download test data.
-
-    Returns
-    -------
-    str
-        The path to the temporary directory.
-    """
-    global TEMPORARY_DIRECTORY
-    with lock:
-        if TEMPORARY_DIRECTORY is not None:
-            return TEMPORARY_DIRECTORY
-
-        TEMPORARY_DIRECTORY = tempfile.mkdtemp()
-
-        # Register a cleanup function to remove the directory at exit
-        atexit.register(shutil.rmtree, TEMPORARY_DIRECTORY)
-
-        return TEMPORARY_DIRECTORY
-
 
 def _check_path(path: str) -> None:
     """Check if the given path is normalized, not absolute, and does not start with a dot.
@@ -68,21 +42,17 @@ def _check_path(path: str) -> None:
     assert not path.startswith("."), f"Path '{path}' should not start with '.'"
 
 
-def _temporary_directory_for_test_data(path: str) -> str:
-    """Get the temporary directory for a test dataset.
+@pytest.fixture(scope="session")
+def temporary_directory_for_test_data(tmp_path_factory) -> callable:
+    base_dir = tmp_path_factory.mktemp("test_data_base")
 
-    Parameters
-    ----------
-    path : str
-        The relative path to the test data in the object store.
+    def _temporary_directory_for_test_data(path: str = "", archive: bool = False) -> str:
+        if path == "":
+            return str(base_dir)
+        _check_path(path)
+        return str(base_dir.joinpath(*Path(path).parts)) + (".extracted" if archive else "")
 
-    Returns
-    -------
-    str
-        The path to the temporary directory.
-    """
-    _check_path(path)
-    return os.path.normpath(os.path.join(_temporary_directory(), path))
+    return _temporary_directory_for_test_data
 
 
 def url_for_test_data(path: str) -> str:
@@ -103,27 +73,29 @@ def url_for_test_data(path: str) -> str:
     return f"{TEST_DATA_URL}{path}"
 
 
-def get_test_data(path: str, gzipped=False) -> str:
-    """Download the test data to a temporary directory and return the local path.
+@pytest.fixture()
+def get_test_data(temporary_directory_for_test_data):
+    def _get_test_data(path: str, gzipped=False) -> callable:
+        """Download the test data to a temporary directory and return the local path.
 
-    Parameters
-    ----------
-    path : str
-        The relative path to the test data.
-    gzipped : bool, optional
-        Flag indicating if the remote file is gzipped, by default False. The local file will be gunzipped.
+        Parameters
+        ----------
+        path : str
+            The relative path to the test data.
+        gzipped : bool, optional
+            Flag indicating if the remote file is gzipped, by default False. The local file will be gunzipped.
 
-    Returns
-    -------
-    str
-        The local path to the downloaded test data.
-    """
+        Returns
+        -------
+        str
+            The local path to the downloaded test data.
+        """
+
+        if _offline():
+            raise RuntimeError("Offline mode: cannot download test data, add @pytest.mark.skipif(not offline(),...)")
 
-    if _offline():
-        raise RuntimeError("Offline mode: cannot download test data, add @pytest.mark.skipif(not offline(),...)")
+        target = temporary_directory_for_test_data(path)
 
-    target = _temporary_directory_for_test_data(path)
-    with lock:
         if os.path.exists(target):
             return target
 
@@ -149,27 +121,29 @@ def get_test_data(path: str, gzipped=False) -> str:
 
         return target
 
+    return _get_test_data
 
-def get_test_archive(path: str, extension=".extracted") -> str:
-    """Download an archive file (.zip, .tar, .tar.gz, .tar.bz2, .tar.xz) to a temporary directory
-    unpack it, and return the local path to the directory containing the extracted files.
 
-    Parameters
-    ----------
-    path : str
-        The relative path to the test data.
-    extension : str, optional
-        The extension to add to the extracted directory, by default '.extracted'
+@pytest.fixture()
+def get_test_archive(temporary_directory_for_test_data, get_test_data) -> callable:
+    def _get_test_archive(path: str) -> str:
+        """Download an archive file (.zip, .tar, .tar.gz, .tar.bz2, .tar.xz) to a temporary directory
+        unpack it, and return the local path to the directory containing the extracted files.
 
-    Returns
-    -------
-    str
-        The local path to the downloaded test data.
-    """
+        Parameters
+        ----------
+        path : str
+            The relative path to the test data.
+        extension : str, optional
+            The extension to add to the extracted directory, by default '.extracted'
 
-    with lock:
+        Returns
+        -------
+        str
+            The local path to the downloaded test data.
+        """
 
-        target = _temporary_directory_for_test_data(path) + extension
+        target = Path(temporary_directory_for_test_data(path, archive=True))
 
         if os.path.exists(target):
             return target
@@ -183,6 +157,8 @@ def get_test_archive(path: str, extension=".extracted") -> str:
 
         return target
 
+    return _get_test_archive
+
 
 def packages_installed(*names: str) -> bool:
     """Check if all the given packages are installed.

{anemoi_utils-0.4.25 → anemoi_utils-0.4.27/src/anemoi_utils.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anemoi-utils
-Version: 0.4.25
+Version: 0.4.27
 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.25 → anemoi_utils-0.4.27}/src/anemoi_utils.egg-info/SOURCES.txt

@@ -67,6 +67,7 @@ src/anemoi/utils/text.py
 src/anemoi/utils/timer.py
 src/anemoi/utils/commands/__init__.py
 src/anemoi/utils/commands/config.py
+src/anemoi/utils/commands/metadata.py
 src/anemoi/utils/commands/requests.py
 src/anemoi/utils/commands/transfer.py
 src/anemoi/utils/mars/__init__.py

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

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