anemoi-utils 0.1.9__py3-none-any.whl → 0.2.0__py3-none-any.whl

anemoi/utils/__main__.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python
+# (C) Copyright 2024 ECMWF.
+#
+# 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 argparse
+import logging
+import sys
+import traceback
+
+from . import __version__
+from .commands import COMMANDS
+
+LOG = logging.getLogger(__name__)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+
+    parser.add_argument(
+        "--version",
+        "-V",
+        action="store_true",
+        help="show the version and exit",
+    )
+    parser.add_argument(
+        "--debug",
+        "-d",
+        action="store_true",
+        help="Debug mode",
+    )
+
+    subparsers = parser.add_subparsers(help="commands:", dest="command")
+    for name, command in COMMANDS.items():
+        command_parser = subparsers.add_parser(name, help=command.__doc__)
+        command.add_arguments(command_parser)
+
+    args = parser.parse_args()
+
+    if args.version:
+        print(__version__)
+        return
+
+    if args.command is None:
+        parser.print_help()
+        return
+
+    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,
+    )
+
+    try:
+        cmd.run(args)
+    except ValueError as e:
+        traceback.print_exc()
+        LOG.error("\n💣 %s", str(e).lstrip())
+        LOG.error("💣 Exiting")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

anemoi/utils/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.9'
-__version_tuple__ = version_tuple = (0, 1, 9)
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 0)

anemoi/utils/checkpoints.py

@@ -13,11 +13,38 @@ are zip archives containing the model weights.
 import json
 import logging
 import os
+import time
 import zipfile
+from tempfile import TemporaryDirectory
+
+import tqdm
 
 LOG = logging.getLogger(__name__)
 
-DEFAULT_NAME = "anemoi-metadata.json"
+DEFAULT_NAME = "ai-models.json"
+DEFAULT_FOLDER = "anemoi-metadata"
+
+
+def has_metadata(path: str, name: str = DEFAULT_NAME) -> bool:
+    """Check if a checkpoint file has a metadata file
+
+    Parameters
+    ----------
+    path : str
+        The path to the checkpoint file
+    name : str, optional
+        The name of the metadata file in the zip archive
+
+    Returns
+    -------
+    bool
+        True if the metadata file is found
+    """
+    with zipfile.ZipFile(path, "r") as f:
+        for b in f.namelist():
+            if os.path.basename(b) == name:
+                return True
+    return False
 
 
 def load_metadata(path: str, name: str = DEFAULT_NAME):
@@ -45,17 +72,17 @@ def load_metadata(path: str, name: str = DEFAULT_NAME):
         for b in f.namelist():
             if os.path.basename(b) == name:
                 if metadata is not None:
-                    LOG.warning(f"Found two '{name}' if {path}")
+                    raise ValueError(f"Found two or more '{name}' in {path}.")
                 metadata = b
 
     if metadata is not None:
         with zipfile.ZipFile(path, "r") as f:
             return json.load(f.open(metadata, "r"))
     else:
-        raise ValueError(f"Could not find {name} in {path}")
+        raise ValueError(f"Could not find '{name}' in {path}.")
 
 
-def save_metadata(path, metadata, name=DEFAULT_NAME):
+def save_metadata(path, metadata, name=DEFAULT_NAME, folder=DEFAULT_FOLDER):
     """Save metadata to a checkpoint file
 
     Parameters
@@ -68,8 +95,85 @@ def save_metadata(path, metadata, name=DEFAULT_NAME):
         The name of the metadata file in the zip archive
     """
     with zipfile.ZipFile(path, "a") as zipf:
-        base, _ = os.path.splitext(os.path.basename(path))
+
+        directories = set()
+
+        for b in zipf.namelist():
+            directory = os.path.dirname(b)
+            while os.path.dirname(directory) not in (".", ""):
+                directory = os.path.dirname(directory)
+            directories.add(directory)
+
+            if os.path.basename(b) == name:
+                raise ValueError(f"'{name}' already in {path}")
+
+        if len(directories) != 1:
+            # PyTorch checkpoints should have a single directory
+            # otherwise PyTorch will complain
+            raise ValueError(f"No or multiple directories in the checkpoint {path}, directories={directories}")
+
+        directory = list(directories)[0]
+
+        LOG.info("Saving metadata to %s/%s/%s", directory, folder, name)
+
         zipf.writestr(
-            f"{base}/{name}",
+            f"{directory}/{folder}/{name}",
             json.dumps(metadata),
         )
+
+
+def _edit_metadata(path, name, callback):
+    new_path = f"{path}.anemoi-edit-{time.time()}-{os.getpid()}.tmp"
+
+    found = False
+
+    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)
+
+        if not found:
+            raise ValueError(f"Could not find '{name}' in {path}")
+
+        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)
+                        pbar.update(1)
+
+    os.rename(new_path, path)
+    LOG.info("Updated metadata in %s", path)
+
+
+def replace_metadata(path, metadata, name=DEFAULT_NAME):
+
+    if not isinstance(metadata, dict):
+        raise ValueError(f"metadata must be a dict, got {type(metadata)}")
+
+    if "version" not in metadata:
+        raise ValueError("metadata must have a 'version' key")
+
+    def callback(full):
+        with open(full, "w") as f:
+            json.dump(metadata, f)
+
+    _edit_metadata(path, name, callback)
+
+
+def remove_metadata(path, name=DEFAULT_NAME):
+
+    LOG.info("Removing metadata '%s' from %s", name, path)
+
+    def callback(full):
+        os.remove(full)
+
+    _edit_metadata(path, name, callback)

anemoi/utils/commands/__init__.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python
+# (C) Copyright 2024 ECMWF.
+#
+# 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 argparse
+import importlib
+import logging
+import os
+import sys
+
+LOG = logging.getLogger(__name__)
+
+
+def register(here, package, select, fail=None):
+    result = {}
+    not_available = {}
+
+    for p in os.listdir(here):
+        full = os.path.join(here, p)
+        if p.startswith("_"):
+            continue
+        if not (p.endswith(".py") or (os.path.isdir(full) and os.path.exists(os.path.join(full, "__init__.py")))):
+            continue
+
+        name, _ = os.path.splitext(p)
+
+        try:
+            imported = importlib.import_module(
+                f".{name}",
+                package=package,
+            )
+        except ImportError as e:
+            not_available[name] = e
+            continue
+
+        obj = select(imported)
+        if obj is not None:
+            result[name] = obj
+
+    for name, e in not_available.items():
+        if fail is None:
+            pass
+        if callable(fail):
+            result[name] = fail(name, e)
+
+    return result
+
+
+class Command:
+    def run(self, args):
+        raise NotImplementedError(f"Command not implemented: {args.command}")
+
+
+class Failed(Command):
+    def __init__(self, name, error):
+        self.name = name
+        self.error = error
+
+    def add_arguments(self, command_parser):
+        command_parser.add_argument("x", nargs=argparse.REMAINDER)
+
+    def run(self, args):
+        print(f"Command '{self.name}' not available: {self.error}")
+        sys.exit(1)
+
+
+COMMANDS = register(
+    os.path.dirname(__file__),
+    __name__,
+    lambda x: x.command(),
+    lambda name, error: Failed(name, error),
+)

anemoi/utils/commands/checkpoint.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python
+# (C) Copyright 2024 ECMWF.
+#
+# 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
+
+from . import Command
+
+
+def visit(x, path, name, value):
+    if isinstance(x, dict):
+        for k, v in x.items():
+            if k == name:
+                print(".".join(path), k, v)
+
+            if v == value:
+                print(".".join(path), k, v)
+
+            path.append(k)
+            visit(v, path, name, value)
+            path.pop()
+
+    if isinstance(x, list):
+        for i, v in enumerate(x):
+            path.append(str(i))
+            visit(v, path, name, value)
+            path.pop()
+
+
+class Checkpoint(Command):
+
+    def add_arguments(self, command_parser):
+        command_parser.add_argument("path", help="Path to the checkpoint.")
+        command_parser.add_argument("--name", help="Search for a specific name.")
+        command_parser.add_argument("--value", help="Search for a specific value.")
+
+    def run(self, args):
+        from anemoi.utils.checkpoints import load_metadata
+
+        checkpoint = load_metadata(args.path, "*.json")
+
+        if args.name or args.value:
+            visit(
+                checkpoint,
+                [],
+                args.name if args.name is not None else object(),
+                args.value if args.value is not None else object(),
+            )
+            return
+
+        print(json.dumps(checkpoint, sort_keys=True, indent=4))
+
+
+command = Checkpoint

anemoi/utils/timer.py

@@ -0,0 +1,32 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+"""Logging utilities."""
+
+import logging
+import time
+
+from .humanize import seconds
+
+LOGGER = logging.getLogger(__name__)
+
+
+class Timer:
+    def __init__(self, title, logger=LOGGER):
+        self.title = title
+        self.start = time.time()
+        self.logger = logger
+
+    def __enter__(self):
+        return self
+
+    @property
+    def elapsed(self):
+        return time.time() - self.start
+
+    def __exit__(self, *args):
+        self.logger.info("%s: %s.", self.title, seconds(self.elapsed))

{anemoi_utils-0.1.9.dist-info → anemoi_utils-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-utils
-Version: 0.1.9
+Version: 0.2.0
 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
@@ -224,6 +224,7 @@ Requires-Python: >=3.9
 License-File: LICENSE
 Requires-Dist: tomli
 Requires-Dist: pyyaml
+Requires-Dist: tqdm
 Provides-Extra: all
 Requires-Dist: tomli ; extra == 'all'
 Requires-Dist: GitPython ; extra == 'all'
@@ -236,14 +237,15 @@ Requires-Dist: GitPython ; extra == 'dev'
 Requires-Dist: nvsmi ; extra == 'dev'
 Requires-Dist: termcolor ; extra == 'dev'
 Requires-Dist: requests ; extra == 'dev'
-Requires-Dist: sphinx ; extra == 'dev'
-Requires-Dist: sphinx-rtd-theme ; extra == 'dev'
-Requires-Dist: nbsphinx ; extra == 'dev'
-Requires-Dist: pandoc ; extra == 'dev'
 Provides-Extra: docs
 Requires-Dist: tomli ; extra == 'docs'
 Requires-Dist: termcolor ; extra == 'docs'
 Requires-Dist: requests ; extra == 'docs'
+Requires-Dist: sphinx ; extra == 'docs'
+Requires-Dist: sphinx-rtd-theme ; extra == 'docs'
+Requires-Dist: nbsphinx ; extra == 'docs'
+Requires-Dist: pandoc ; extra == 'docs'
+Requires-Dist: sphinx-argparse ; extra == 'docs'
 Provides-Extra: grib
 Requires-Dist: requests ; extra == 'grib'
 Provides-Extra: provenance

anemoi_utils-0.2.0.dist-info/RECORD

@@ -0,0 +1,22 @@
+anemoi/utils/__init__.py,sha256=zZZpbKIoGWwdCOuo6YSruLR7C0GzvzI1Wzhyqaa0K7M,456
+anemoi/utils/__main__.py,sha256=CGl8WF7rWMx9EoArysla0-ThjUFtEZUEGM58LbdU488,1798
+anemoi/utils/_version.py,sha256=H-qsvrxCpdhaQzyddR-yajEqI71hPxLa4KxzpP3uS1g,411
+anemoi/utils/caching.py,sha256=HrC9aFHlcCTaM2Z5u0ivGIXz7eFu35UQQhUuwwuG2pk,1743
+anemoi/utils/checkpoints.py,sha256=1_3mg4B-ykTVfIvIUEv7IxGyREx_ZcilVbB3U-V6O6I,5165
+anemoi/utils/config.py,sha256=XEesqODvkuE3ZA7dnEnZ-ooBRtU6ecPmkfP65FtialA,2147
+anemoi/utils/dates.py,sha256=Ot9OTY1uFvHxW1EU4DPv3oUqmzvkXTwKuwhlfVlY788,8426
+anemoi/utils/grib.py,sha256=gVfo4KYQv31iRyoqRDwk5tiqZDUgOIvhag_kO0qjYD0,3067
+anemoi/utils/humanize.py,sha256=LD6dGnqChxA5j3tMhSybsAGRQzi33d_qS9pUoUHubkc,10330
+anemoi/utils/provenance.py,sha256=v54L9jF1JgYcclOhg3iojRl1v3ajbiWz_oc289xTgO4,9574
+anemoi/utils/text.py,sha256=pGWtDvRFoDxAnSuZJiA-GOGJOJLHsw2dAm0tfVvPKno,8599
+anemoi/utils/timer.py,sha256=5aNdcxVmiCijRmqp0URmsqsDypLUJgME0GSn9bk8zxo,920
+anemoi/utils/commands/__init__.py,sha256=Pc5bhVgW92ox1lMR5WUOLuhiY2HT6PsadSHclyw99Vc,1983
+anemoi/utils/commands/checkpoint.py,sha256=SEnAizU3WklqMXUjmIh4eNrgBVwmheKG9gEBS90zwYU,1741
+anemoi/utils/mars/__init__.py,sha256=RAeY8gJ7ZvsPlcIvrQ4fy9xVHs3SphTAPw_XJDtNIKo,1750
+anemoi/utils/mars/mars.yaml,sha256=R0dujp75lLA4wCWhPeOQnzJ45WZAYLT8gpx509cBFlc,66
+anemoi_utils-0.2.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+anemoi_utils-0.2.0.dist-info/METADATA,sha256=7IO_KqlBHcAKCmdszjmE5Zxe5G_ox2QOwKn16mgBIa4,15174
+anemoi_utils-0.2.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+anemoi_utils-0.2.0.dist-info/entry_points.txt,sha256=LENOkn88xzFQo-V59AKoA_F_cfYQTJYtrNTtf37YgHY,60
+anemoi_utils-0.2.0.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
+anemoi_utils-0.2.0.dist-info/RECORD,,

anemoi_utils-0.2.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+anemoi-utils = anemoi.utils.__main__:main

anemoi_utils-0.1.9.dist-info/RECORD

@@ -1,17 +0,0 @@
-anemoi/utils/__init__.py,sha256=zZZpbKIoGWwdCOuo6YSruLR7C0GzvzI1Wzhyqaa0K7M,456
-anemoi/utils/_version.py,sha256=NWmu2cvzOcqY9v-ee-qFLmtXRczssdN-cFGZ9qMNSmY,411
-anemoi/utils/caching.py,sha256=HrC9aFHlcCTaM2Z5u0ivGIXz7eFu35UQQhUuwwuG2pk,1743
-anemoi/utils/checkpoints.py,sha256=IR86FFNh5JR_uQVlgybnZG74PyU0CNLhyocqARwZIrs,2069
-anemoi/utils/config.py,sha256=XEesqODvkuE3ZA7dnEnZ-ooBRtU6ecPmkfP65FtialA,2147
-anemoi/utils/dates.py,sha256=Ot9OTY1uFvHxW1EU4DPv3oUqmzvkXTwKuwhlfVlY788,8426
-anemoi/utils/grib.py,sha256=gVfo4KYQv31iRyoqRDwk5tiqZDUgOIvhag_kO0qjYD0,3067
-anemoi/utils/humanize.py,sha256=LD6dGnqChxA5j3tMhSybsAGRQzi33d_qS9pUoUHubkc,10330
-anemoi/utils/provenance.py,sha256=v54L9jF1JgYcclOhg3iojRl1v3ajbiWz_oc289xTgO4,9574
-anemoi/utils/text.py,sha256=pGWtDvRFoDxAnSuZJiA-GOGJOJLHsw2dAm0tfVvPKno,8599
-anemoi/utils/mars/__init__.py,sha256=RAeY8gJ7ZvsPlcIvrQ4fy9xVHs3SphTAPw_XJDtNIKo,1750
-anemoi/utils/mars/mars.yaml,sha256=R0dujp75lLA4wCWhPeOQnzJ45WZAYLT8gpx509cBFlc,66
-anemoi_utils-0.1.9.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-anemoi_utils-0.1.9.dist-info/METADATA,sha256=wVXW4E6hpkTjCEHUCWDiQXm-pZfJl81Mrzdy6kwKbfA,15101
-anemoi_utils-0.1.9.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-anemoi_utils-0.1.9.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
-anemoi_utils-0.1.9.dist-info/RECORD,,