Extract retry_with_backoff to rigging for reuse outside RPC (#4499)

Add retry_with_backoff() to rigging.timing alongside ExponentialBackoff.
It accepts a retryable predicate and on_retry(exc, attempt) callback,
making it usable for any exception type—not just gRPC ConnectErrors.
Refactor iris.call_with_retry and marin.call_with_hf_backoff to delegate
to it. Delete fn_utils.with_retries (fixed delay, no backoff) and update
its two callers in hf_upload.py.

Author: rjpower

lib/iris/src/iris/rpc/errors.py

@@ -16,7 +16,7 @@
 
 from iris.rpc import errors_pb2
 from iris.time_proto import timestamp_to_proto
-from rigging.timing import Deadline, ExponentialBackoff, Timestamp
+from rigging.timing import Deadline, ExponentialBackoff, Timestamp, retry_with_backoff
 
 logger = logging.getLogger(__name__)
 
@@ -171,56 +171,22 @@ def call_with_retry(
     Raises:
         Exception from call_fn if all retries exhausted or error is not retryable
     """
-    if backoff is None:
-        backoff = ExponentialBackoff(initial=0.5, maximum=10.0, factor=2.0)
-    else:
-        backoff = backoff.copy()
-    last_exception = None
-    start_time = time.monotonic()
-
-    for attempt in range(max_attempts):
-        try:
-            return call_fn()
-        except Exception as e:
-            last_exception = e
-            if not is_retryable_error(e):
-                raise
-
-            if on_retry is not None:
-                on_retry(e)
-
-            elapsed = time.monotonic() - start_time
-            attempts_exhausted = attempt + 1 >= max_attempts
-            time_exhausted = max_elapsed is not None and elapsed >= max_elapsed
-
-            if attempts_exhausted or time_exhausted:
-                logger.exception(
-                    "Operation %s failed after %d attempts (%.1fs elapsed): %s",
-                    operation,
-                    attempt + 1,
-                    elapsed,
-                    e,
-                )
-                raise
-
-            delay = backoff.next_interval()
-            if max_elapsed is not None:
-                remaining = max_elapsed - elapsed
-                delay = min(delay, max(0, remaining))
-
-            logger.exception(
-                "Operation %s failed (attempt %d/%d, %.1fs elapsed), retrying in %.2fs: %s",
-                operation,
-                attempt + 1,
-                max_attempts,
-                elapsed,
-                delay,
-                e,
-            )
-            time.sleep(delay)
-
-    assert last_exception is not None
-    raise last_exception
+    wrapped_on_retry: Callable[[Exception, int], None] | None = None
+    if on_retry is not None:
+
+        def wrapped_on_retry(exc: Exception, _attempt: int) -> None:
+            assert on_retry is not None
+            on_retry(exc)
+
+    return retry_with_backoff(
+        call_fn,
+        retryable=is_retryable_error,
+        max_attempts=max_attempts,
+        max_elapsed=max_elapsed,
+        backoff=backoff,
+        on_retry=wrapped_on_retry,
+        operation=operation,
+    )
 
 
 def poll_with_retries(

lib/marin/src/marin/export/hf_upload.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import dataclasses
-import functools
 import io
 import logging
 import os
@@ -14,11 +13,11 @@
 import humanfriendly
 from fsspec.implementations.local import LocalFileSystem
 from rigging.filesystem import open_url
+from rigging.timing import ExponentialBackoff, retry_with_backoff
 from huggingface_hub import create_commit, upload_folder
 from tqdm_loggable.auto import tqdm
 
 from marin.execution import ExecutorStep, InputName
-from marin.utilities.fn_utils import with_retries
 from marin.utils import fsspec_glob
 
 logger = logging.getLogger(__name__)
@@ -222,16 +221,22 @@ def _actually_upload_to_hf(config: UploadToHfConfig):
                 )
 
 
-@functools.wraps(upload_folder)
-@with_retries()
 def retrying_upload_folder(*args, **kwargs):
-    return upload_folder(*args, **kwargs)
+    return retry_with_backoff(
+        lambda: upload_folder(*args, **kwargs),
+        max_attempts=3,
+        backoff=ExponentialBackoff(initial=2.0, maximum=30.0, factor=2.0),
+        operation="upload_folder",
+    )
 
 
-@functools.wraps(create_commit)
-@with_retries()
 def retrying_create_commit(*args, **kwargs):
-    return create_commit(*args, **kwargs)
+    return retry_with_backoff(
+        lambda: create_commit(*args, **kwargs),
+        max_attempts=3,
+        backoff=ExponentialBackoff(initial=2.0, maximum=30.0, factor=2.0),
+        operation="create_commit",
+    )
 
 
 def _wrap_in_buffered_base(fileobj):

lib/marin/src/marin/utilities/fn_utils.py

@@ -4,8 +4,6 @@
 import functools
 import logging
 import os
-import random
-import time
 from collections.abc import Callable
 from contextlib import contextmanager
 from dataclasses import fields, is_dataclass
@@ -17,6 +15,7 @@
 import fsspec
 import requests
 from rigging.filesystem import url_to_fs
+from rigging.timing import ExponentialBackoff, retry_with_backoff
 from huggingface_hub.utils import HfHubHTTPError
 
 logger = logging.getLogger(__name__)
@@ -153,47 +152,22 @@ def _hf_should_retry(exc: Exception) -> bool:
     return any(keyword in message for keyword in _HF_RETRY_KEYWORDS)
 
 
-def _hf_sleep_with_jitter(delay: float, max_delay: float) -> tuple[float, float]:
-    jitter = random.uniform(0.5, 1.5)
-    sleep_seconds = min(delay * jitter, max_delay)
-    time.sleep(sleep_seconds)
-    next_delay = min(delay * 2, max_delay)
-    return sleep_seconds, next_delay
-
-
 def call_with_hf_backoff(
     fn: Callable[[], T],
     *,
     context: str,
     max_attempts: int = 6,
     initial_delay: float = 2.0,
     max_delay: float = 60.0,
-    logger: logging.Logger | None = None,
 ) -> T:
     """Call ``fn`` with exponential backoff tuned for HF rate limits."""
-
-    log_obj = logger or logging.getLogger(__name__)
-    delay = initial_delay
-
-    for attempt in range(1, max_attempts + 1):
-        try:
-            return fn()
-        except Exception as exc:  # pragma: no cover - network failure
-            retryable = _hf_should_retry(exc)
-            if not retryable or attempt == max_attempts:
-                raise
-
-            sleep_seconds, delay = _hf_sleep_with_jitter(delay, max_delay)
-            log_obj.warning(
-                "HF request failed for %s (attempt %s/%s): %s. Retrying in %.1fs",
-                context,
-                attempt,
-                max_attempts,
-                exc,
-                sleep_seconds,
-            )
-
-    raise RuntimeError(f"Exceeded max attempts ({max_attempts}) for HF request: {context}")
+    return retry_with_backoff(
+        fn,
+        retryable=_hf_should_retry,
+        max_attempts=max_attempts,
+        backoff=ExponentialBackoff(initial=initial_delay, maximum=max_delay, factor=2.0, jitter=0.25),
+        operation=context,
+    )
 
 
 def load_dataset_with_backoff(
@@ -202,7 +176,6 @@ def load_dataset_with_backoff(
     max_attempts: int = 6,
     initial_delay: float = 2.0,
     max_delay: float = 120.0,
-    logger: logging.Logger | None = None,
     **dataset_kwargs: Any,
 ):
     return call_with_hf_backoff(
@@ -211,7 +184,6 @@ def load_dataset_with_backoff(
         max_attempts=max_attempts,
         initial_delay=initial_delay,
         max_delay=max_delay,
-        logger=logger,
     )
 
 
@@ -222,7 +194,6 @@ def load_tokenizer_with_backoff(
     max_attempts: int = 6,
     initial_delay: float = 2.0,
     max_delay: float = 60.0,
-    logger: logging.Logger | None = None,
 ):
     from levanter.tokenizers import load_tokenizer
 
@@ -233,7 +204,6 @@ def load_tokenizer_with_backoff(
         max_attempts=max_attempts,
         initial_delay=initial_delay,
         max_delay=max_delay,
-        logger=logger,
     )
 
 

lib/marin/src/marin/utils.py

@@ -9,6 +9,7 @@
 import time
 from collections.abc import Callable, Iterator
 from datetime import datetime, timedelta, timezone
+from typing import TypeVar
 
 logger = logging.getLogger(__name__)
 
@@ -554,3 +555,73 @@ def wait_until_or_raise(
     ) -> None:
         if not self.wait_until(condition, timeout):
             raise TimeoutError(error_message)
+
+
+T = TypeVar("T")
+
+
+def retry_with_backoff(
+    call_fn: Callable[[], T],
+    *,
+    retryable: Callable[[Exception], bool] = lambda _: True,
+    max_attempts: int = 10,
+    max_elapsed: float | None = None,
+    backoff: ExponentialBackoff | None = None,
+    on_retry: Callable[[Exception, int], None] | None = None,
+    operation: str = "",
+) -> T:
+    """Execute call_fn with exponential backoff retry.
+
+    Args:
+        call_fn: Function to call and potentially retry.
+        retryable: Returns True to retry this exception, False to re-raise immediately.
+            Defaults to retrying all exceptions.
+        max_attempts: Maximum total attempts (default 10).
+        max_elapsed: Wall-clock budget in seconds; None means no time limit.
+        backoff: Backoff schedule; a fresh copy is made internally so the caller's
+            instance is not mutated. Defaults to
+            ExponentialBackoff(initial=0.5, maximum=10.0, factor=2.0).
+        on_retry: Called with (exception, attempt_index) before each sleep.
+        operation: Description used in log messages.
+    """
+    if backoff is None:
+        backoff = ExponentialBackoff(initial=0.5, maximum=10.0, factor=2.0)
+    else:
+        backoff = backoff.copy()
+
+    start_time = time.monotonic()
+
+    for attempt in range(max_attempts):
+        try:
+            return call_fn()
+        except Exception as e:
+            if not retryable(e):
+                raise
+
+            if on_retry is not None:
+                on_retry(e, attempt)
+
+            elapsed = time.monotonic() - start_time
+            attempts_exhausted = attempt + 1 >= max_attempts
+            time_exhausted = max_elapsed is not None and elapsed >= max_elapsed
+
+            if attempts_exhausted or time_exhausted:
+                raise
+
+            delay = backoff.next_interval()
+            if max_elapsed is not None:
+                remaining = max_elapsed - elapsed
+                delay = min(delay, max(0, remaining))
+
+            logger.warning(
+                "Operation %s failed (attempt %d/%d, %.1fs elapsed), retrying in %.2fs: %s",
+                operation,
+                attempt + 1,
+                max_attempts,
+                elapsed,
+                delay,
+                e,
+            )
+            time.sleep(delay)
+
+    raise AssertionError("unreachable")