ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/plugins/datatools/s3/s3.py

@@ -1,3 +1,4 @@
+import errno
 import json
 import os
 import re
@@ -17,11 +18,12 @@ from metaflow.metaflow_config import (
     DATATOOLS_S3ROOT,
     S3_RETRY_COUNT,
     S3_TRANSIENT_RETRY_COUNT,
+    S3_LOG_TRANSIENT_RETRIES,
     S3_SERVER_SIDE_ENCRYPTION,
+    S3_WORKER_COUNT,
     TEMPDIR,
 )
 from metaflow.util import (
-    namedtuple_with_defaults,
     is_stringish,
     to_bytes,
     to_unicode,
@@ -29,6 +31,7 @@ from metaflow.util import (
     url_quote,
     url_unquote,
 )
+from metaflow.tuple_util import namedtuple_with_defaults
 from metaflow.exception import MetaflowException
 from metaflow.debug import debug
 import metaflow.tracing as tracing
@@ -136,6 +139,10 @@ class MetaflowS3InvalidRange(MetaflowException):
     headline = "S3 invalid range"
 
 
+class MetaflowS3InsufficientDiskSpace(MetaflowException):
+    headline = "Insufficient disk space"
+
+
 class S3Object(object):
     """
     This object represents a path or an object in S3,
@@ -492,18 +499,22 @@ class S3(object):
 
     Parameters
     ----------
-    tmproot : str, default: '.'
+    tmproot : str, default '.'
         Where to store the temporary directory.
-    bucket : str, optional
+    bucket : str, optional, default None
         Override the bucket from `DATATOOLS_S3ROOT` when `run` is specified.
-    prefix : str, optional
+    prefix : str, optional, default None
         Override the path from `DATATOOLS_S3ROOT` when `run` is specified.
-    run : FlowSpec or Run, optional
+    run : FlowSpec or Run, optional, default None
         Derive path prefix from the current or a past run ID, e.g. S3(run=self).
-    s3root : str, optional
+    s3root : str, optional, default None
         If `run` is not specified, use this as the S3 prefix.
+    encryption : str, optional, default None
+        Server-side encryption to use when uploading objects to S3.
     """
 
+    TYPE = "s3"
+
     @classmethod
     def get_root_from_config(cls, echo, create_on_absent=True):
         return DATATOOLS_S3ROOT
@@ -570,7 +581,13 @@ class S3(object):
         self._s3_inject_failures = kwargs.get(
             "inject_failure_rate", TEST_INJECT_RETRYABLE_FAILURES
         )
-        self._tmpdir = mkdtemp(dir=tmproot, prefix="metaflow.s3.")
+        # Storing tmproot, bucket, ... as members to allow easier reconstruction
+        # during JSON deserialization
+        self._tmproot = tmproot
+        self._bucket = bucket
+        self._prefix = prefix
+        self._run = run
+        self._tmpdir = mkdtemp(dir=self._tmproot, prefix="metaflow.s3.")
         self._encryption = encryption
 
     def __enter__(self) -> "S3":
@@ -598,7 +615,9 @@ class S3(object):
         # returned are Unicode.
         key = getattr(key_value, "key", key_value)
         if self._s3root is None:
-            parsed = urlparse(to_unicode(key))
+            # NOTE: S3 allows fragments as part of object names, e.g. /dataset #1/data.txt
+            # Without allow_fragments=False the parsed.path for an object name with fragments is incomplete.
+            parsed = urlparse(to_unicode(key), allow_fragments=False)
             if parsed.scheme == "s3" and parsed.path:
                 return key
             else:
@@ -619,7 +638,9 @@ class S3(object):
                     "Don't use absolute S3 URLs when the S3 client is "
                     "initialized with a prefix. URL: %s" % key
                 )
-            return os.path.join(self._s3root, key)
+            # Strip leading slashes to ensure os.path.join works correctly
+            # os.path.join discards the first argument if the second starts with '/'
+            return os.path.join(self._s3root, key.lstrip("/"))
         else:
             return self._s3root
 
@@ -763,7 +784,9 @@ class S3(object):
         """
 
         url = self._url(key)
-        src = urlparse(url)
+        # NOTE: S3 allows fragments as part of object names, e.g. /dataset #1/data.txt
+        # Without allow_fragments=False the parsed src.path for an object name with fragments is incomplete.
+        src = urlparse(url, allow_fragments=False)
 
         def _info(s3, tmp):
             resp = s3.head_object(Bucket=src.netloc, Key=src.path.lstrip('/"'))
@@ -889,7 +912,9 @@ class S3(object):
         DOWNLOAD_MAX_CHUNK = 2 * 1024 * 1024 * 1024 - 1
 
         url, r = self._url_and_range(key)
-        src = urlparse(url)
+        # NOTE: S3 allows fragments as part of object names, e.g. /dataset #1/data.txt
+        # Without allow_fragments=False the parsed src.path for an object name with fragments is incomplete.
+        src = urlparse(url, allow_fragments=False)
 
         def _download(s3, tmp):
             if r:
@@ -1171,7 +1196,9 @@ class S3(object):
         blob.close = lambda: None
 
         url = self._url(key)
-        src = urlparse(url)
+        # NOTE: S3 allows fragments as part of object names, e.g. /dataset #1/data.txt
+        # Without allow_fragments=False the parsed src.path for an object name with fragments is incomplete.
+        src = urlparse(url, allow_fragments=False)
         extra_args = None
         if content_type or metadata or self._encryption:
             extra_args = {}
@@ -1245,12 +1272,12 @@ class S3(object):
 
         def _store():
             for key_obj in key_objs:
-                if isinstance(key_obj, tuple):
-                    key = key_obj[0]
-                    obj = key_obj[1]
-                else:
+                if isinstance(key_obj, S3PutObject):
                     key = key_obj.key
                     obj = key_obj.value
+                else:
+                    key = key_obj[0]
+                    obj = key_obj[1]
                 store_info = {
                     "key": key,
                     "content_type": getattr(key_obj, "content_type", None),
@@ -1319,12 +1346,12 @@ class S3(object):
 
         def _check():
             for key_path in key_paths:
-                if isinstance(key_path, tuple):
-                    key = key_path[0]
-                    path = key_path[1]
-                else:
+                if isinstance(key_path, S3PutObject):
                     key = key_path.key
                     path = key_path.path
+                else:
+                    key = key_path[0]
+                    path = key_path[1]
                 store_info = {
                     "key": key,
                     "content_type": getattr(key_path, "content_type", None),
@@ -1366,8 +1393,13 @@ class S3(object):
                 elif error_code == "NoSuchBucket":
                     raise MetaflowS3URLException("Specified S3 bucket doesn't exist.")
                 error = str(err)
+            except OSError as e:
+                if e.errno == errno.ENOSPC:
+                    raise MetaflowS3InsufficientDiskSpace(str(e))
+            except MetaflowException as ex:
+                # Re-raise Metaflow exceptions (including TimeoutException)
+                raise
             except Exception as ex:
-                # TODO specific error message for out of disk space
                 error = str(ex)
             if tmp:
                 os.unlink(tmp.name)
@@ -1380,9 +1412,31 @@ class S3(object):
         )
 
     # add some jitter to make sure retries are not synchronized
-    def _jitter_sleep(self, trynum, multiplier=2):
-        interval = multiplier**trynum + random.randint(0, 10)
-        time.sleep(interval)
+    def _jitter_sleep(
+        self, trynum: int, base: int = 2, cap: int = 360, jitter: float = 0.1
+    ) -> None:
+        """
+        Sleep for an exponentially increasing interval with added jitter.
+
+        Parameters
+        ----------
+        trynum: The current retry attempt number.
+        base: The base multiplier for the exponential backoff.
+        cap: The maximum interval to sleep.
+        jitter: The maximum jitter percentage to add to the interval.
+        """
+        # Calculate the exponential backoff interval
+        interval = min(cap, base**trynum)
+
+        # Add random jitter
+        jitter_value = interval * jitter * random.uniform(-1, 1)
+        interval_with_jitter = interval + jitter_value
+
+        # Ensure the interval is not negative
+        interval_with_jitter = max(0, interval_with_jitter)
+
+        # Sleep for the calculated interval
+        time.sleep(interval_with_jitter)
 
     # NOTE: re: _read_many_files and _put_many_files
     # All file IO is through binary files - we write bytes, we read
@@ -1470,20 +1524,17 @@ class S3(object):
         #    - a known transient failure (SlowDown for example) in which case we will
         #      retry *only* the inputs that have this transient failure.
         #    - an unknown failure (something went wrong but we cannot say if it was
-        #      a known permanent failure or something else). In this case, we retry
-        #      the operation completely.
-        #
-        # There are therefore two retry counts:
-        #  - the transient failure retry count: how many times do we try on known
-        #    transient errors
-        #  - the top-level retry count: how many times do we try on unknown failures
+        #      a known permanent failure or something else). In this case, we assume
+        #      it's a transient failure and retry only those inputs (same as above).
         #
-        # Note that, if the operation runs out of transient failure retries, it will
-        # count as an "unknown" failure (ie: it will be retried according to the
-        # outer top-level retry count). In other words, you can potentially have
-        # transient_retry_count * retry_count tries).
-        # Finally, if on transient failures, we make NO progress (ie: no input is
-        # successfully processed), that counts as an "unknown" failure.
+        # NOTES(npow): 2025-05-13
+        # Previously, this code would also retry the fatal failures, including no_progress
+        # and unknown failures, from the beginning. This is not ideal because:
+        # 1. Fatal errors are not supposed to be retried.
+        # 2. Retrying from the beginning does not improve the situation, and is
+        #    wasteful since we have already uploaded some files.
+        # 3. The number of transient errors is far more than fatal errors, so we
+        #    can be optimistic and assume the unknown errors are transient.
         cmdline = [sys.executable, os.path.abspath(s3op.__file__), mode]
         recursive_get = False
         for key, value in options.items():
@@ -1518,7 +1569,6 @@ class S3(object):
             # Otherwise, we cap the failure rate at 90%
             return min(90, self._s3_inject_failures)
 
-        retry_count = 0  # Number of retries (excluding transient failures)
         transient_retry_count = 0  # Number of transient retries (per top-level retry)
         inject_failures = _inject_failure_rate()
         out_lines = []  # List to contain the lines returned by _s3op_with_retries
@@ -1585,7 +1635,12 @@ class S3(object):
                         # things, this will shrink more and more until we are doing a
                         # single operation at a time. If things start going better, it
                         # will increase by 20% every round.
-                        max_count = min(int(last_ok_count * 1.2), len(pending_retries))
+                        #
+                        # If we made no progress (last_ok_count == 0) we retry at most
+                        # 2*S3_WORKER_COUNT from whatever is left in `pending_retries`
+                        max_count = min(
+                            int(last_ok_count * 1.2), len(pending_retries)
+                        ) or min(2 * S3_WORKER_COUNT, len(pending_retries))
                         tmp_input.writelines(pending_retries[:max_count])
                         tmp_input.flush()
                         debug.s3client_exec(
@@ -1702,53 +1757,49 @@ class S3(object):
                         _update_out_lines(out_lines, ok_lines, resize=loop_count == 0)
                         return 0, 0, inject_failures, err_out
 
-        while retry_count <= S3_RETRY_COUNT:
+        while transient_retry_count <= S3_TRANSIENT_RETRY_COUNT:
             (
                 last_ok_count,
                 last_retry_count,
                 inject_failures,
                 err_out,
             ) = try_s3_op(last_ok_count, pending_retries, out_lines, inject_failures)
-            if err_out or (
-                last_retry_count != 0
-                and (
-                    last_ok_count == 0
-                    or transient_retry_count > S3_TRANSIENT_RETRY_COUNT
-                )
-            ):
-                # We had a fatal failure (err_out is not None)
-                # or we made no progress (last_ok_count is 0)
-                # or we are out of transient retries
-                # so we will restart from scratch (being very conservative)
-                retry_count += 1
-                err_msg = err_out
-                if err_msg is None and last_ok_count == 0:
-                    err_msg = "No progress"
-                if err_msg is None:
-                    err_msg = "Too many transient errors"
-                print(
-                    "S3 non-transient error (attempt #%d): %s" % (retry_count, err_msg)
-                )
-                _reset()
-                if retry_count <= S3_RETRY_COUNT:
-                    self._jitter_sleep(retry_count)
-                continue
-            elif last_retry_count != 0:
+            if err_out:
+                break
+            if last_retry_count != 0:
                 # During our last try, we did not manage to process everything we wanted
                 # due to a transient failure so we try again.
                 transient_retry_count += 1
                 total_ok_count += last_ok_count
-                print(
-                    "Transient S3 failure (attempt #%d) -- total success: %d, "
-                    "last attempt %d/%d -- remaining: %d"
-                    % (
-                        transient_retry_count,
-                        total_ok_count,
-                        last_ok_count,
-                        last_ok_count + last_retry_count,
-                        len(pending_retries),
+
+                if S3_LOG_TRANSIENT_RETRIES:
+                    # Extract transient error type from pending retry lines
+                    error_info = ""
+                    if pending_retries:
+                        try:
+                            # Parse the first line to get transient error type
+                            first_retry = json.loads(
+                                pending_retries[0].decode("utf-8").strip()
+                            )
+                            if "transient_error_type" in first_retry:
+                                error_info = (
+                                    " (%s)" % first_retry["transient_error_type"]
+                                )
+                        except (json.JSONDecodeError, IndexError, KeyError):
+                            pass
+
+                    print(
+                        "Transient S3 failure (attempt #%d) -- total success: %d, "
+                        "last attempt %d/%d -- remaining: %d%s"
+                        % (
+                            transient_retry_count,
+                            total_ok_count,
+                            last_ok_count,
+                            last_ok_count + last_retry_count,
+                            len(pending_retries),
+                            error_info,
+                        )
                     )
-                )
                 if inject_failures == 0:
                     # Don't sleep when we are "faking" the failures
                     self._jitter_sleep(transient_retry_count)