anemoi-utils 0.4.23__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.23'
-__version_tuple__ = version_tuple = (0, 4, 23)
+__version__ = version = '0.4.24'
+__version_tuple__ = version_tuple = (0, 4, 24)

anemoi/utils/remote/s3.py

@@ -24,6 +24,7 @@ the `~/.config/anemoi/settings.toml`
 or `~/.config/anemoi/settings-secrets.toml` files.
 """
 
+import fnmatch
 import logging
 import os
 import threading
@@ -38,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
@@ -55,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
     -------
@@ -68,7 +71,7 @@ def s3_client(bucket: str, region: str = None) -> Any:
     if not hasattr(thread_local, "s3_clients"):
         thread_local.s3_clients = {}
 
-    key = f"{bucket}-{region}"
+    key = f"{bucket}-{region}-{service}"
 
     if key in thread_local.s3_clients:
         return thread_local.s3_clients[key]
@@ -96,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":
@@ -115,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]
 
@@ -215,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)
@@ -227,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:
@@ -400,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
@@ -409,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:
@@ -464,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
@@ -477,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
@@ -505,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:
@@ -525,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:
@@ -581,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
 
 
@@ -600,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)
 
@@ -643,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-0.4.23.dist-info → anemoi_utils-0.4.24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anemoi-utils
-Version: 0.4.23
+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

{anemoi_utils-0.4.23.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=AnB7uXIf9bJtYv1LZYemC9XOrzt9AeZnOcUDxAsfacA,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
@@ -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=dcXcgNddlgxwJ_OpgqOff8EWk-LT2mz20m7FcHHqz7w,17869
+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.23.dist-info/licenses/LICENSE,sha256=8HznKF1Vi2IvfLsKNE5A2iVyiri3pRjRPvPC9kxs6qk,11354
-anemoi_utils-0.4.23.dist-info/METADATA,sha256=jxTYBIvx6wzO0g8ftu2MoSPfeViRSTUrE38Zy3HUvlw,15439
-anemoi_utils-0.4.23.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-anemoi_utils-0.4.23.dist-info/entry_points.txt,sha256=LENOkn88xzFQo-V59AKoA_F_cfYQTJYtrNTtf37YgHY,60
-anemoi_utils-0.4.23.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
-anemoi_utils-0.4.23.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.23.dist-info → anemoi_utils-0.4.24.dist-info}/WHEEL

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