mlbricks-core 0.2.0__tar.gz → 0.2.1__tar.gz

{mlbricks_core-0.2.0 → mlbricks_core-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mlbricks-core
-Version: 0.2.0
+Version: 0.2.1
 Summary: Core ML tooling, including config resolution, registry, and storage.
 Author: Shivam Verma
 Author-email: Shivam Verma <shivam.vvm.verma@gmail.com>

{mlbricks_core-0.2.0 → mlbricks_core-0.2.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mlbricks-core"
-version = "0.2.0"
+version = "0.2.1"
 description = "Core ML tooling, including config resolution, registry, and storage."
 readme = "README.md"
 license = "MIT"

{mlbricks_core-0.2.0 → mlbricks_core-0.2.1}/pyproject.toml.orig

@@ -1,6 +1,6 @@
 [project]
 name = "mlbricks-core"
-version = "0.2.0"
+version = "0.2.1"
 description = "Core ML tooling, including config resolution, registry, and storage."
 readme = "README.md"
 license = "MIT"

{mlbricks_core-0.2.0 → mlbricks_core-0.2.1}/src/mlbricks/storage.py

@@ -6,18 +6,29 @@ Public API: ``StorageDriver`` only. ``Storage``, ``LocalStorage``, and
 
 import fnmatch
 import glob as glob_module
+import logging
 import shutil
+import time
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Any
 
 import boto3
 from botocore.exceptions import ClientError
+from s3transfer.exceptions import S3DownloadFailedError
 
 from mlbricks.utils import path_slug
 
 __all__ = ["StorageDriver"]
 
+logger = logging.getLogger(__name__)
+
+# Retry policy for the transient object-changed-mid-download race: another
+# writer overwrites the key between S3Storage.get's head_object and the
+# completion of its ranged GETs. See docs/storage.md for the full writeup.
+_GET_MAX_ATTEMPTS = 4
+_GET_RETRY_BASE_DELAY_SECONDS = 4.0
+
 
 def _validate_uri(uri: str) -> None:
     """Raise ValueError if uri is an absolute local path."""
@@ -160,14 +171,29 @@ class S3Storage(Storage):
 
         bucket, key = self._parse_uri(uri)
         client = self._get_client()
-
-        try:
-            staging.parent.mkdir(parents=True, exist_ok=True)
-            client.download_file(bucket, key, str(staging))
-        except ClientError as exc:
-            if exc.response["Error"]["Code"] in ("404", "NoSuchKey"):
-                return None
-            raise
+        staging.parent.mkdir(parents=True, exist_ok=True)
+
+        for attempt in range(1, _GET_MAX_ATTEMPTS + 1):
+            try:
+                client.download_file(bucket, key, str(staging))
+                return staging
+            except ClientError as exc:
+                if exc.response["Error"]["Code"] in ("404", "NoSuchKey"):
+                    return None
+                raise
+            except S3DownloadFailedError:
+                if attempt == _GET_MAX_ATTEMPTS:
+                    raise
+                delay = _GET_RETRY_BASE_DELAY_SECONDS * 2 ** (attempt - 1)
+                logger.warning(
+                    "S3 download of %r failed ETag validation (attempt %d/%d);"
+                    " retrying in %.1fs",
+                    uri,
+                    attempt,
+                    _GET_MAX_ATTEMPTS,
+                    delay,
+                )
+                time.sleep(delay)
 
         return staging