anemoi-utils 0.4.22__py3-none-any.whl → 0.4.24__py3-none-any.whl

anemoi/utils/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.22'
-__version_tuple__ = version_tuple = (0, 4, 22)
+__version__ = version = '0.4.24'
+__version_tuple__ = version_tuple = (0, 4, 24)

anemoi/utils/provenance.py

@@ -47,8 +47,11 @@ def lookup_git_repo(path: str) -> Optional[Any]:
     Repo, optional
         The git repository if found, otherwise None.
     """
-    from git import InvalidGitRepositoryError
-    from git import Repo
+    try:
+        from git import InvalidGitRepositoryError
+        from git import Repo
+    except ImportError:
+        return None
 
     while path != "/":
         try:

anemoi/utils/remote/s3.py

@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024-2025 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.
 
+
 """This module provides functions to upload, download, list and delete files and folders on S3.
 The functions of this package expect that the AWS credentials are set up in the environment
 typicaly by setting the `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` environment variables or
@@ -21,6 +24,7 @@ the `~/.config/anemoi/settings.toml`
 or `~/.config/anemoi/settings-secrets.toml` files.
 """
 
+import fnmatch
 import logging
 import os
 import threading
@@ -35,15 +39,15 @@ from ..humanize import bytes_to_human
 from . import BaseDownload
 from . import BaseUpload
 
-LOGGER = logging.getLogger(__name__)
-
+LOG = logging.getLogger(__name__)
+SECRETS = ["aws_access_key_id", "aws_secret_access_key"]
 
 # s3_clients are not thread-safe, so we need to create a new client for each thread
 
 thread_local = threading.local()
 
 
-def s3_client(bucket: str, region: str = None) -> Any:
+def s3_client(bucket: str, *, region: str = None, service: str = "s3") -> Any:
     """Get an S3 client for the specified bucket and region.
 
     Parameters
@@ -52,6 +56,8 @@ def s3_client(bucket: str, region: str = None) -> 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
     -------
@@ -65,14 +71,16 @@ def s3_client(bucket: str, region: str = None) -> Any:
     if not hasattr(thread_local, "s3_clients"):
         thread_local.s3_clients = {}
 
-    key = f"{bucket}-{region}"
-
-    boto3_config = dict(max_pool_connections=25)
+    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)
+    boto3_config = dict(
+        max_pool_connections=25,
+        request_checksum_calculation="when_required",
+        response_checksum_validation="when_required",
+    )
 
     if region:
         # This is using AWS
@@ -91,17 +99,27 @@ def s3_client(bucket: str, region: str = None) -> Any:
         # We may be accessing a different S3 compatible service
         # Use anemoi.config to get the configuration
 
-        options = {}
-        config = load_config(secrets=["aws_access_key_id", "aws_secret_access_key"])
+        region = "unknown-region"
+
+        options = {"region_name": region}
+        config = load_config(secrets=SECRETS)
 
         cfg = config.get("object-storage", {})
+        candidate = None
         for k, v in cfg.items():
             if isinstance(v, (str, int, float, bool)):
                 options[k] = v
 
-        for k, v in cfg.get(bucket, {}).items():
-            if isinstance(v, (str, int, float, bool)):
-                options[k] = v
+            if isinstance(v, dict):
+                if fnmatch.fnmatch(bucket, k):
+                    if candidate is not None:
+                        raise ValueError(f"Multiple object storage configurations match {bucket}: {candidate} and {k}")
+                    candidate = k
+
+        if candidate is not None:
+            for k, v in cfg.get(candidate, {}).items():
+                if isinstance(v, (str, int, float, bool)):
+                    options[k] = v
 
         type = options.pop("type", "s3")
         if type != "s3":
@@ -110,11 +128,27 @@ def s3_client(bucket: str, region: str = None) -> Any:
         if "config" in options:
             boto3_config.update(options["config"])
             del options["config"]
-            from botocore.client import Config
 
     options["config"] = Config(**boto3_config)
 
-    thread_local.s3_clients[key] = boto3.client("s3", **options)
+    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.info(f"Using S3 options: {_(options)}")
+
+    thread_local.s3_clients[key] = boto3.client(service, **options)
 
     return thread_local.s3_clients[key]
 
@@ -162,7 +196,14 @@ class S3Upload(BaseUpload):
         # delete(target)
 
     def _transfer_file(
-        self, source: str, target: str, overwrite: bool, resume: bool, verbosity: int, threads: int, config: dict = None
+        self,
+        source: str,
+        target: str,
+        overwrite: bool,
+        resume: bool,
+        verbosity: int,
+        threads: int,
+        config: dict = None,
     ) -> int:
         """Transfer a file to S3.
 
@@ -203,7 +244,7 @@ class S3Upload(BaseUpload):
         size = os.path.getsize(source)
 
         if verbosity > 0:
-            LOGGER.info(f"{self.action} {source} to {target} ({bytes_to_human(size)})")
+            LOG.info(f"{self.action} {source} to {target} ({bytes_to_human(size)})")
 
         try:
             results = s3.head_object(Bucket=bucket, Key=key)
@@ -215,7 +256,7 @@ class S3Upload(BaseUpload):
 
         if remote_size is not None:
             if remote_size != size:
-                LOGGER.warning(
+                LOG.warning(
                     f"{target} already exists, but with different size, re-uploading (remote={remote_size}, local={size})"
                 )
             elif resume:
@@ -227,7 +268,13 @@ class S3Upload(BaseUpload):
 
         if verbosity > 0:
             with tqdm.tqdm(total=size, unit="B", unit_scale=True, unit_divisor=1024, leave=False) as pbar:
-                s3.upload_file(source, bucket, key, Callback=lambda x: pbar.update(x), Config=config)
+                s3.upload_file(
+                    source,
+                    bucket,
+                    key,
+                    Callback=lambda x: pbar.update(x),
+                    Config=config,
+                )
         else:
             s3.upload_file(source, bucket, key, Config=config)
 
@@ -326,7 +373,14 @@ class S3Download(BaseDownload):
         return s3_object["Size"]
 
     def _transfer_file(
-        self, source: str, target: str, overwrite: bool, resume: bool, verbosity: int, threads: int, config: dict = None
+        self,
+        source: str,
+        target: str,
+        overwrite: bool,
+        resume: bool,
+        verbosity: int,
+        threads: int,
+        config: dict = None,
     ) -> int:
         """Transfer a file from S3 to the local filesystem.
 
@@ -375,7 +429,7 @@ class S3Download(BaseDownload):
         size = int(response["ContentLength"])
 
         if verbosity > 0:
-            LOGGER.info(f"{self.action} {source} to {target} ({bytes_to_human(size)})")
+            LOG.info(f"{self.action} {source} to {target} ({bytes_to_human(size)})")
 
         if overwrite:
             resume = False
@@ -384,7 +438,7 @@ class S3Download(BaseDownload):
             if os.path.exists(target):
                 local_size = os.path.getsize(target)
                 if local_size != size:
-                    LOGGER.warning(
+                    LOG.warning(
                         f"{target} already with different size, re-downloading (remote={size}, local={local_size})"
                     )
                 else:
@@ -397,7 +451,13 @@ class S3Download(BaseDownload):
 
         if verbosity > 0:
             with tqdm.tqdm(total=size, unit="B", unit_scale=True, unit_divisor=1024, leave=False) as pbar:
-                s3.download_file(bucket, key, target, Callback=lambda x: pbar.update(x), Config=config)
+                s3.download_file(
+                    bucket,
+                    key,
+                    target,
+                    Callback=lambda x: pbar.update(x),
+                    Config=config,
+                )
         else:
             s3.download_file(bucket, key, target, Config=config)
 
@@ -433,7 +493,7 @@ def _list_objects(target: str, batch: bool = False) -> Iterable:
                 yield from objects
 
 
-def _delete_folder(target: str) -> None:
+def delete_folder(target: str) -> None:
     """Delete a folder from S3.
 
     Parameters
@@ -446,13 +506,13 @@ def _delete_folder(target: str) -> None:
 
     total = 0
     for batch in _list_objects(target, batch=True):
-        LOGGER.info(f"Deleting {len(batch):,} objects from {target}")
+        LOG.info(f"Deleting {len(batch):,} objects from {target}")
         s3.delete_objects(Bucket=bucket, Delete={"Objects": [{"Key": o["Key"]} for o in batch]})
         total += len(batch)
-        LOGGER.info(f"Deleted {len(batch):,} objects (total={total:,})")
+        LOG.info(f"Deleted {len(batch):,} objects (total={total:,})")
 
 
-def _delete_file(target: str) -> None:
+def delete_file(target: str) -> None:
     """Delete a file from S3.
 
     Parameters
@@ -474,12 +534,12 @@ def _delete_file(target: str) -> None:
         exits = False
 
     if not exits:
-        LOGGER.warning(f"{target} does not exist. Did you mean to delete a folder? Then add a trailing '/'")
+        LOG.warning(f"{target} does not exist. Did you mean to delete a folder? Then add a trailing '/'")
         return
 
-    LOGGER.info(f"Deleting {target}")
+    LOG.info(f"Deleting {target}")
     s3.delete_object(Bucket=bucket, Key=key)
-    LOGGER.info(f"{target} is deleted")
+    LOG.info(f"{target} is deleted")
 
 
 def delete(target: str) -> None:
@@ -494,9 +554,9 @@ def delete(target: str) -> None:
     assert target.startswith("s3://")
 
     if target.endswith("/"):
-        _delete_folder(target)
+        delete_folder(target)
     else:
-        _delete_file(target)
+        delete_file(target)
 
 
 def list_folder(folder: str) -> Iterable:
@@ -524,7 +584,9 @@ def list_folder(folder: str) -> Iterable:
 
     for page in paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter="/"):
         if "CommonPrefixes" in page:
-            yield from [folder + _["Prefix"] for _ in page.get("CommonPrefixes")]
+            yield from [folder + _["Prefix"] for _ in page.get("CommonPrefixes") if _["Prefix"] != "/"]
+        if "Contents" in page:
+            yield from [folder + _["Key"] for _ in page.get("Contents")]
 
 
 def object_info(target: str) -> dict:
@@ -548,7 +610,33 @@ def object_info(target: str) -> dict:
         return s3.head_object(Bucket=bucket, Key=key)
     except s3.exceptions.ClientError as e:
         if e.response["Error"]["Code"] == "404":
-            raise ValueError(f"{target} does not exist")
+            raise FileNotFoundError(f"{target} does not exist")
+        raise
+
+
+def object_exists(target: str) -> bool:
+    """Check if an object exists.
+
+    Parameters
+    ----------
+    target : str
+        The URL of a file or a folder on S3. The URL should start with 's3://'.
+
+    Returns
+    -------
+    bool
+        True if the object exists, False otherwise.
+    """
+
+    _, _, bucket, key = target.split("/", 3)
+    s3 = s3_client(bucket)
+
+    try:
+        s3.head_object(Bucket=bucket, Key=key)
+        return True
+    except s3.exceptions.ClientError as e:
+        if e.response["Error"]["Code"] == "404":
+            return False
         raise
 
 
@@ -567,7 +655,7 @@ def object_acl(target: str) -> dict:
     """
 
     _, _, bucket, key = target.split("/", 3)
-    s3 = s3_client()
+    s3 = s3_client(bucket)
 
     return s3.get_object_acl(Bucket=bucket, Key=key)
 
@@ -610,3 +698,29 @@ def upload(source: str, target: str, *args, **kwargs) -> None:
 
     assert target.startswith("s3://"), f"target {target} should start with 's3://'"
     return transfer(source, target, *args, **kwargs)
+
+
+def quotas(target: str) -> dict:
+    """Get the quotas for an S3 bucket.
+
+    Parameters
+    ----------
+    target : str
+        The URL of a file or a folder on S3. The URL should start with 's3://'.
+
+    Returns
+    -------
+    dict
+        A dictionary with the quotas for the bucket.
+    """
+    from botocore.exceptions import ClientError
+
+    _, _, bucket, _ = target.split("/", 3)
+    s3 = s3_client(bucket, service="service-quotas")
+
+    try:
+        return s3.list_service_quotas(ServiceCode="ec2")
+    except ClientError as e:
+        if e.response["Error"]["Code"] == "404":
+            raise ValueError(f"{target} does not exist")
+        raise

anemoi/utils/testing.py

@@ -261,12 +261,12 @@ def _run_slow_tests() -> bool:
 @lru_cache(maxsize=None)
 def _offline() -> bool:
     """Check if we are offline."""
-
-    import socket
+    from urllib import request
 
     try:
-        socket.create_connection(("anemoi.ecmwf.int", 443), timeout=5)
-    except OSError:
+        request.urlopen("https://anemoi.ecmwf.int", timeout=1)
+        return False
+    except request.URLError:
         return True
 
     return False

{anemoi_utils-0.4.22.dist-info → anemoi_utils-0.4.24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anemoi-utils
-Version: 0.4.22
+Version: 0.4.24
 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
@@ -252,7 +252,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: boto3>1.36; extra == "s3"
 Provides-Extra: tests
 Requires-Dist: pytest; extra == "tests"
 Provides-Extra: text

{anemoi_utils-0.4.22.dist-info → anemoi_utils-0.4.24.dist-info}/RECORD

@@ -1,6 +1,6 @@
 anemoi/utils/__init__.py,sha256=uVhpF-VjIl_4mMywOVtgTutgsdIsqz-xdkwxeMhzuag,730
 anemoi/utils/__main__.py,sha256=6LlE4MYrPvqqrykxXh7XMi50UZteUY59NeM8P9Zs2dU,910
-anemoi/utils/_version.py,sha256=ncPL58_Lzy93I55YUff3cwMUk7Kbpd4bhXC7ujoaqGw,513
+anemoi/utils/_version.py,sha256=eqPTx1mit5QVY_CIZPeIYTsnPPxmLZEDtDv2l58DdBE,513
 anemoi/utils/caching.py,sha256=rXbeAmpBcMbbfN4EVblaHWKicsrtx1otER84FEBtz98,6183
 anemoi/utils/checkpoints.py,sha256=N4WpAZXa4etrpSEKhHqUUtG2-x9w3FJMHcLO-dDAXPY,9600
 anemoi/utils/cli.py,sha256=IyZfnSw0u0yYnrjOrzvm2RuuKvDk4cVb8pf8BkaChgA,6209
@@ -13,13 +13,13 @@ anemoi/utils/grids.py,sha256=uYgkU_KIg8FTUiuKV0Pho2swMMeXcSQ9CQe0MFlRr_I,5262
 anemoi/utils/hindcasts.py,sha256=iYVIxSNFL2HJcc_k1abCFLkpJFGHT8WKRIR4wcAwA3s,2144
 anemoi/utils/humanize.py,sha256=pjnFJAKHbEAOfcvn8c48kt-8eFy6FGW_U2ruJvfamrA,25189
 anemoi/utils/logs.py,sha256=naTgrmPwWHD4eekFttXftS4gtcAGYHpCqG4iwYprNDA,1804
-anemoi/utils/provenance.py,sha256=xC6mTstF7f_asqtPSrulC7c34xjOSuAxWhkwc3yKhHg,14629
+anemoi/utils/provenance.py,sha256=iTsn4r-VPq2D8tSHPSuAIqG077_opkqMT42G03DRWJg,14690
 anemoi/utils/registry.py,sha256=e3nOIRyMYQ-mpEvaHAv5tuvMYNbkJ5yz94ns7BnvkjM,9717
 anemoi/utils/rules.py,sha256=VspUoPmw7tijrs6l_wl4vDjr_zVQsFjx9ITiBSvxgc8,6972
 anemoi/utils/s3.py,sha256=xMT48kbcelcjjqsaU567WI3oZ5eqo88Rlgyx5ECszAU,4074
 anemoi/utils/sanitise.py,sha256=ZYGdSX6qihQANr3pHZjbKnoapnzP1KcrWdW1Ul1mOGk,3668
 anemoi/utils/sanitize.py,sha256=43ZKDcfVpeXSsJ9TFEc9aZnD6oe2cUh151XnDspM98M,462
-anemoi/utils/testing.py,sha256=psfHfluNqXa-cXSDF4xD7xIQCVPvBGDNCWWZ-yVh_24,7750
+anemoi/utils/testing.py,sha256=kwgAgLh3exYOTZSaX4xcPFjiMOyQDz-vcAlPcJqMiZk,7784
 anemoi/utils/text.py,sha256=HkzIvi24obDceFLpJEwBJ9PmPrJUkQN2TrElJ-A87gU,14441
 anemoi/utils/timer.py,sha256=_leKMYza2faM7JKlGE7LCNy13rbdPnwaCF7PSrI_NmI,3895
 anemoi/utils/commands/__init__.py,sha256=5u_6EwdqYczIAgJfCwRSyQAYFEqh2ZuHHT57g9g7sdI,808
@@ -29,13 +29,13 @@ anemoi/utils/mars/__init__.py,sha256=b-Lc3L1TAQd9ODs0Z1YSJzgZCO1K_M3DSgx_yd2qXvM
 anemoi/utils/mars/mars.yaml,sha256=R0dujp75lLA4wCWhPeOQnzJ45WZAYLT8gpx509cBFlc,66
 anemoi/utils/mars/requests.py,sha256=VFMHBVAAl0_2lOcMBa1lvaKHctN0lDJsI6_U4BucGew,1142
 anemoi/utils/remote/__init__.py,sha256=swPWHQoh-B6Xq9R489tPw0FykMue7f-bJ8enneFYSYE,20776
-anemoi/utils/remote/s3.py,sha256=spQ8l0rwQjLZh9dZu5cOsYIvNwKihQfCJ6YsFYegeqI,17339
+anemoi/utils/remote/s3.py,sha256=lhLxwBU-AKERzYURXJ-cOlaccEYmjoXqlvX3zYu5h-8,19977
 anemoi/utils/remote/ssh.py,sha256=xNtsawh8okytCKRehkRCVExbHZj-CRUQNormEHglfuw,8088
 anemoi/utils/schemas/__init__.py,sha256=nkinKlsPLPXEjfTYQT1mpKC4cvs-14w_zBkDRxakwxw,698
 anemoi/utils/schemas/errors.py,sha256=lgOXzVTYzAE0qWQf3OZ42vCWixv8lilSqLLhzARBmvI,1831
-anemoi_utils-0.4.22.dist-info/licenses/LICENSE,sha256=8HznKF1Vi2IvfLsKNE5A2iVyiri3pRjRPvPC9kxs6qk,11354
-anemoi_utils-0.4.22.dist-info/METADATA,sha256=DOdoppZLABKmpmK5DAgYtEl5BoKs7zMkWmquG6qFfgo,15439
-anemoi_utils-0.4.22.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-anemoi_utils-0.4.22.dist-info/entry_points.txt,sha256=LENOkn88xzFQo-V59AKoA_F_cfYQTJYtrNTtf37YgHY,60
-anemoi_utils-0.4.22.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
-anemoi_utils-0.4.22.dist-info/RECORD,,
+anemoi_utils-0.4.24.dist-info/licenses/LICENSE,sha256=8HznKF1Vi2IvfLsKNE5A2iVyiri3pRjRPvPC9kxs6qk,11354
+anemoi_utils-0.4.24.dist-info/METADATA,sha256=EsIP40bMdXsSEAp1v4iFkKMrzdVaE4Cw892SLNlzuWs,15439
+anemoi_utils-0.4.24.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+anemoi_utils-0.4.24.dist-info/entry_points.txt,sha256=LENOkn88xzFQo-V59AKoA_F_cfYQTJYtrNTtf37YgHY,60
+anemoi_utils-0.4.24.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
+anemoi_utils-0.4.24.dist-info/RECORD,,

{anemoi_utils-0.4.22.dist-info → anemoi_utils-0.4.24.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any