anemoi-utils 0.4.4__py3-none-any.whl → 0.4.5__py3-none-any.whl

anemoi/utils/__main__.py

@@ -1,12 +1,11 @@
-#!/usr/bin/env python
-# (C) Copyright 2024 ECMWF.
+# (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.
-#
 
 from anemoi.utils.cli import cli_main
 from anemoi.utils.cli import make_parser

anemoi/utils/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.4'
-__version_tuple__ = version_tuple = (0, 4, 4)
+__version__ = version = '0.4.5'
+__version_tuple__ = version_tuple = (0, 4, 5)

anemoi/utils/checkpoints.py

@@ -94,8 +94,8 @@ def load_metadata(path: str, *, supporting_arrays=False, name: str = DEFAULT_NAM
         with zipfile.ZipFile(path, "r") as f:
             metadata = json.load(f.open(metadata, "r"))
             if supporting_arrays:
-                metadata["supporting_arrays"] = load_supporting_arrays(f, metadata.get("supporting_arrays", {}))
-                return metadata, supporting_arrays
+                arrays = load_supporting_arrays(f, metadata.get("supporting_arrays_paths", {}))
+                return metadata, arrays
 
             return metadata
     else:

anemoi/utils/commands/__init__.py

@@ -1,12 +1,11 @@
-#!/usr/bin/env python
-# (C) Copyright 2024 ECMWF.
+# (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 os
 

anemoi/utils/commands/config.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 # (C) Copyright 2024 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0

anemoi/utils/compatibility.py

@@ -0,0 +1,76 @@
+# (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.
+
+from __future__ import annotations
+
+import functools
+from typing import Any
+from typing import Callable
+
+
+def aliases(
+    aliases: dict[str, str | list[str]] | None = None, **kwargs: str | list[str]
+) -> Callable[[Callable], Callable]:
+    """Alias keyword arguments in a function call.
+
+    Allows for dynamically renaming keyword arguments in a function call.
+
+    Parameters
+    ----------
+    aliases : dict[str, str  |  list[str]] | None, optional
+        Key, value pair of aliases, with keys being the true name, and value being a str or list of aliases,
+        by default None
+    **kwargs : str | list[str]
+        Kwargs form of aliases
+
+    Returns
+    -------
+    Callable
+        Decorator function that renames keyword arguments in a function call.
+
+    Raises
+    ------
+    ValueError
+        If the aliasing would result in duplicate keys.
+
+    Examples
+    --------
+    ```python
+    @aliases(a="b", c=["d", "e"])
+    def func(a, c):
+        return a, c
+
+    func(a=1, c=2)  # (1, 2)
+    func(b=1, d=2)  # (1, 2)
+    ```
+
+    """
+
+    if aliases is None:
+        aliases = {}
+    aliases.update(kwargs)
+
+    aliases = {v: k for k, vs in aliases.items() for v in (vs if isinstance(vs, list) else [vs])}
+
+    def decorator(func: Callable) -> Callable:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs) -> Any:
+            keys = kwargs.keys()
+            for k in set(keys).intersection(set(aliases.keys())):
+                if aliases[k] in keys:
+                    raise ValueError(
+                        f"When aliasing {k} with {aliases[k]} duplicate keys were present. Cannot include both."
+                    )
+                kwargs[aliases[k]] = kwargs.pop(k)
+
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

anemoi/utils/mars/__init__.py

@@ -1,6 +1,8 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.

anemoi/utils/registry.py

@@ -33,11 +33,12 @@ class Wrapper:
 class Registry:
     """A registry of factories"""
 
-    def __init__(self, package):
+    def __init__(self, package, key="_type"):
 
         self.package = package
         self.registered = {}
         self.kind = package.split(".")[-1]
+        self.key = key
 
     def register(self, name: str, factory: callable = None):
 
@@ -86,6 +87,8 @@ class Registry:
                 self.registered[name] = entry_point.load()
 
         if name not in self.registered:
+            for e in self.registered:
+                LOG.info(f"Registered: {e}")
             raise ValueError(f"Cannot load '{name}' from {self.package}")
 
         return self.registered[name]
@@ -96,3 +99,31 @@ class Registry:
 
     def __call__(self, name: str, *args, **kwargs):
         return self.create(name, *args, **kwargs)
+
+    def from_config(self, config, *args, **kwargs):
+        if isinstance(config, str):
+            config = {config: {}}
+
+        if not isinstance(config, dict):
+            raise ValueError(f"Invalid config: {config}")
+
+        if self.key in config:
+            config = config.copy()
+            key = config.pop(self.key)
+            return self.create(key, *args, **config, **kwargs)
+
+        if len(config) == 1:
+            key = list(config.keys())[0]
+            value = config[key]
+
+            if isinstance(value, dict):
+                return self.create(key, *args, **value, **kwargs)
+
+            if isinstance(value, list):
+                return self.create(key, *args, *value, **kwargs)
+
+            return self.create(key, *args, value, **kwargs)
+
+        raise ValueError(
+            f"Entry '{config}' must either be a string, a dictionray with a single entry, or a dictionary with a '{self.key}' key"
+        )

anemoi/utils/remote/__init__.py

@@ -0,0 +1,328 @@
+# (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.
+
+import concurrent.futures
+import logging
+import os
+import shutil
+from abc import abstractmethod
+
+import tqdm
+
+from ..humanize import bytes_to_human
+
+LOGGER = logging.getLogger(__name__)
+
+
+def _ignore(number_of_files, total_size, total_transferred, transfering):
+    pass
+
+
+class Loader:
+
+    def transfer_folder(self, *, source, target, overwrite=False, resume=False, verbosity=1, threads=1, progress=None):
+        assert verbosity == 1, verbosity
+
+        if progress is None:
+            progress = _ignore
+
+        # from boto3.s3.transfer import TransferConfig
+        # config = TransferConfig(use_threads=False)
+        config = None
+        with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor:
+            try:
+                if verbosity > 0:
+                    LOGGER.info(f"{self.action} {source} to {target}")
+
+                total_size = 0
+                total_transferred = 0
+
+                futures = []
+                for name in self.list_source(source):
+
+                    futures.append(
+                        executor.submit(
+                            self.transfer_file,
+                            source=self.source_path(name, source),
+                            target=self.target_path(name, source, target),
+                            overwrite=overwrite,
+                            resume=resume,
+                            verbosity=verbosity - 1,
+                            config=config,
+                        )
+                    )
+                    total_size += self.source_size(name)
+
+                    if len(futures) % 10000 == 0:
+
+                        progress(len(futures), total_size, 0, False)
+
+                        if verbosity > 0:
+                            LOGGER.info(f"Preparing transfer, {len(futures):,} files... ({bytes_to_human(total_size)})")
+                        done, _ = concurrent.futures.wait(
+                            futures,
+                            timeout=0.001,
+                            return_when=concurrent.futures.FIRST_EXCEPTION,
+                        )
+                        # Trigger exceptions if any
+                        for future in done:
+                            future.result()
+
+                number_of_files = len(futures)
+                progress(number_of_files, total_size, 0, True)
+
+                if verbosity > 0:
+                    LOGGER.info(f"{self.action} {number_of_files:,} files ({bytes_to_human(total_size)})")
+                    with tqdm.tqdm(total=total_size, unit="B", unit_scale=True, unit_divisor=1024) as pbar:
+                        for future in concurrent.futures.as_completed(futures):
+                            size = future.result()
+                            pbar.update(size)
+                            total_transferred += size
+                            progress(number_of_files, total_size, total_transferred, True)
+                else:
+                    for future in concurrent.futures.as_completed(futures):
+                        size = future.result()
+                        total_transferred += size
+                        progress(number_of_files, total_size, total_transferred, True)
+
+            except Exception:
+                executor.shutdown(wait=False, cancel_futures=True)
+                raise
+
+    def transfer_file(self, source, target, overwrite, resume, verbosity, threads=1, progress=None, config=None):
+        try:
+            return self._transfer_file(source, target, overwrite, resume, verbosity, threads=threads, config=config)
+        except Exception as e:
+            LOGGER.exception(f"Error transferring {source} to {target}")
+            LOGGER.error(e)
+            raise
+
+    @abstractmethod
+    def list_source(self, source):
+        raise NotImplementedError
+
+    @abstractmethod
+    def source_path(self, local_path, source):
+        raise NotImplementedError
+
+    @abstractmethod
+    def target_path(self, source_path, source, target):
+        raise NotImplementedError
+
+    @abstractmethod
+    def source_size(self, local_path):
+        raise NotImplementedError
+
+    @abstractmethod
+    def copy(self, source, target, **kwargs):
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_temporary_target(self, target, pattern):
+        raise NotImplementedError
+
+    @abstractmethod
+    def rename_target(self, target, temporary_target):
+        raise NotImplementedError
+
+
+class BaseDownload(Loader):
+    action = "Downloading"
+
+    @abstractmethod
+    def copy(self, source, target, **kwargs):
+        raise NotImplementedError
+
+    def get_temporary_target(self, target, pattern):
+        dirname, basename = os.path.split(target)
+        return pattern.format(dirname=dirname, basename=basename)
+
+    def rename_target(self, target, new_target):
+        os.rename(target, new_target)
+
+    def delete_target(self, target):
+        if os.path.exists(target):
+            shutil.rmtree(target)
+
+
+class BaseUpload(Loader):
+    action = "Uploading"
+
+    def copy(self, source, target, **kwargs):
+        if os.path.isdir(source):
+            self.transfer_folder(source=source, target=target, **kwargs)
+        else:
+            self.transfer_file(source=source, target=target, **kwargs)
+
+    def list_source(self, source):
+        for root, _, files in os.walk(source):
+            for file in files:
+                yield os.path.join(root, file)
+
+    def source_path(self, local_path, source):
+        return local_path
+
+    def target_path(self, source_path, source, target):
+        relative_path = os.path.relpath(source_path, source)
+        path = os.path.join(target, relative_path)
+        return path
+
+    def source_size(self, local_path):
+        return os.path.getsize(local_path)
+
+
+class TransferMethodNotImplementedError(NotImplementedError):
+    pass
+
+
+class Transfer:
+    """This is the internal API and should not be used directly. Use the transfer function instead."""
+
+    TransferMethodNotImplementedError = TransferMethodNotImplementedError
+
+    def __init__(
+        self,
+        source,
+        target,
+        overwrite=False,
+        resume=False,
+        verbosity=1,
+        threads=1,
+        progress=None,
+        temporary_target=False,
+    ):
+        if target == ".":
+            target = os.path.basename(source)
+
+        temporary_target = {
+            False: "{dirname}/{basename}",
+            True: "{dirname}-downloading/{basename}",
+            "-tmp/*": "{dirname}-tmp/{basename}",
+            "*-tmp": "{dirname}/{basename}-tmp",
+            "tmp-*": "{dirname}/tmp-{basename}",
+        }.get(temporary_target, temporary_target)
+        assert isinstance(temporary_target, str), (type(temporary_target), temporary_target)
+
+        self.source = source
+        self.target = target
+        self.overwrite = overwrite
+        self.resume = resume
+        self.verbosity = verbosity
+        self.threads = threads
+        self.progress = progress
+        self.temporary_target = temporary_target
+
+        cls = _find_transfer_class(self.source, self.target)
+        self.loader = cls()
+
+    def run(self):
+
+        target = self.loader.get_temporary_target(self.target, self.temporary_target)
+        if target != self.target:
+            LOGGER.info(f"Using temporary target {target} to copy to {self.target}")
+
+        if self.overwrite:
+            # delete the target if it exists
+            LOGGER.info(f"Deleting {self.target}")
+            self.delete_target(target)
+
+            # carefully delete the temporary target if it exists
+            head, tail = os.path.split(self.target)
+            head_, tail_ = os.path.split(target)
+            if not head_.startswith(head) or tail not in tail_:
+                LOGGER.info(f"{target} is too different from {self.target} to delete it automatically.")
+            else:
+                self.delete_target(target)
+
+        self.loader.copy(
+            self.source,
+            target,
+            overwrite=self.overwrite,
+            resume=self.resume,
+            verbosity=self.verbosity,
+            threads=self.threads,
+            progress=self.progress,
+        )
+
+        self.rename_target(target, self.target)
+
+        return self
+
+    def rename_target(self, target, new_target):
+        if target != new_target:
+            LOGGER.info(f"Renaming temporary target {target} into {self.target}")
+            return self.loader.rename_target(target, new_target)
+
+    def delete_target(self, target):
+        return self.loader.delete_target(target)
+
+
+def _find_transfer_class(source, target):
+    from_ssh = source.startswith("ssh://")
+    into_ssh = target.startswith("ssh://")
+
+    from_s3 = source.startswith("s3://")
+    into_s3 = target.startswith("s3://")
+
+    from_local = not from_ssh and not from_s3
+    into_local = not into_ssh and not into_s3
+
+    # check that exactly one source type and one target type is specified
+    assert sum([into_ssh, into_local, into_s3]) == 1, (into_ssh, into_local, into_s3)
+    assert sum([from_ssh, from_local, from_s3]) == 1, (from_ssh, from_local, from_s3)
+
+    if from_local and into_ssh:  # local -> ssh
+        from .ssh import RsyncUpload
+
+        return RsyncUpload
+
+    if from_s3 and into_local:  # local <- S3
+        from .s3 import S3Download
+
+        return S3Download
+
+    if from_local and into_s3:  # local -> S3
+        from .s3 import S3Upload
+
+        return S3Upload
+
+    raise TransferMethodNotImplementedError(f"Transfer from {source} to {target} is not implemented")
+
+
+# this is the public API
+def transfer(*args, **kwargs) -> Loader:
+    """Parameters
+    ----------
+    source : str
+        A path to a local file or folder or a URL to a file or a folder on S3.
+        The url should start with 's3://'.
+    target : str
+        A path to a local file or folder or a URL to a file or a folder on S3 or a remote folder.
+        The url should start with 's3://' or 'ssh://'.
+    overwrite : bool, optional
+        If the data is alreay on in the target location it will be overwritten.
+        By default False
+    resume : bool, optional
+        If the data is alreay on S3 it will not be uploaded, unless the remote file has a different size
+        Ignored if the target is an SSH remote folder (ssh://).
+        By default False
+    verbosity : int, optional
+        The level of verbosity, by default 1
+    progress: callable, optional
+        A callable that will be called with the number of files, the total size of the files, the total size
+        transferred and a boolean indicating if the transfer has started. By default None
+    threads : int, optional
+        The number of threads to use when uploading a directory, by default 1
+    temporary_target : bool, optional
+        Experimental feature
+        If True and if the target location supports it, the data will be uploaded to a temporary location
+        then renamed to the final location. Supported by SSH and local targets, not supported by S3.
+        By default False
+    """
+    copier = Transfer(*args, **kwargs)
+    copier.run()
+    return copier