"""Bounded retry with exponential backoff for transient storage failures.
Two mechanisms recover from a failed write, and they answer different
questions. Re-running an operation recovers from a process that died: every
write is create-only and resumable, so the second run completes what the first
started. This module covers the smaller case -- a backend that failed while the
process is still alive and can simply try again.
Retrying is only safe because the operations underneath are idempotent. A retry
of a partially-succeeded publish re-enters the same body, discovers the objects
it already wrote, and continues. Without that property a retry would report a
collision instead of the transient error that caused it.
"""
from __future__ import annotations
import errno
import logging
import numbers
import random
import time
from collections.abc import Callable
from dataclasses import dataclass
from typing import TypeVar
from signal_dataset.errors import PublicationCollisionError, StorageError
logger = logging.getLogger(__name__)
T = TypeVar("T")
#: Jitter is drawn from a private generator rather than the `random` module.
#: Training pipelines seed the global generator, often identically across
#: workers -- which would make every worker retry in perfect lockstep, the one
#: thing jitter exists to prevent -- and a retry drawing from the global stream
#: would make a run's random sequence depend on how many times storage hiccuped.
rng = random.Random()
#: The largest exponent `delay` will raise 2 to. Past this the schedule is
#: pinned at `max_seconds` anyway, and 2**1024 overflows a float, which would
#: replace the caller's storage error with an `OverflowError` from in here.
MAX_BACKOFF_EXPONENT = 32
#: An upper bound on `attempts`, so `retry=10_000` cannot turn a failing
#: operation into a multi-day sleep.
MAX_ATTEMPTS = 64
#: Failures worth trying again. A backend that timed out, dropped a connection,
#: or returned a 5xx may well succeed on the next attempt.
TRANSIENT: tuple[type[BaseException], ...] = (
StorageError,
TimeoutError,
ConnectionError,
OSError,
)
#: Failures that must escape immediately, checked *before* `TRANSIENT`.
#:
#: Every entry here is an `OSError` subclass, which is the whole reason the list
#: exists: `PublicationCollisionError` derives from `FileExistsError`, so a bare
#: `except OSError` would swallow it. A collision is not a transient fault -- it
#: is how a resumed operation discovers its own earlier attempt, and retrying it
#: would hide a genuine conflict between two different writers.
NON_TRANSIENT: tuple[type[BaseException], ...] = (
PublicationCollisionError,
FileNotFoundError,
FileExistsError,
PermissionError,
IsADirectoryError,
NotADirectoryError,
)
#: Conditions that arrive as a plain `OSError` because Python maps only some
#: errnos to a named subclass. Every one of these describes the environment
#: rather than a passing fault: a full disk, a read-only mount, or a descriptor
#: bug in this package. Retrying them sleeps for tens of seconds, re-runs the
#: whole operation each time, and then reports the same failure.
NON_TRANSIENT_ERRNOS = frozenset(
{
errno.EACCES,
errno.EBADF,
errno.EDQUOT,
errno.EINVAL,
errno.EISDIR,
errno.ELOOP,
errno.EMFILE,
errno.ENAMETOOLONG,
errno.ENFILE,
errno.ENOSPC,
errno.ENOTDIR,
errno.EPERM,
errno.EROFS,
errno.EXDEV,
}
)
def transient(exc: BaseException) -> bool:
"""Whether trying `exc` again could plausibly produce a different result."""
if isinstance(exc, NON_TRANSIENT):
return False
if isinstance(exc, OSError) and exc.errno in NON_TRANSIENT_ERRNOS:
return False
return isinstance(exc, TRANSIENT)
[docs]
@dataclass(frozen=True, slots=True)
class RetryPolicy:
"""How often, and how patiently, to retry a transient failure.
The defaults are tuned for object storage rather than for local contention.
A bucket that just returned 503 is unlikely to be healthy a tenth of a
second later, so the first pause is a full second; five attempts span
roughly half a minute of trouble. Waiting costs only time, because the
operation underneath is idempotent.
"""
attempts: int = 5
base_seconds: float = 1.0
max_seconds: float = 60.0
jitter: bool = True
def __post_init__(self) -> None:
if isinstance(self.attempts, bool) or not isinstance(self.attempts, numbers.Integral):
raise ValueError("attempts must be an integer")
if not 1 <= self.attempts <= MAX_ATTEMPTS:
raise ValueError(f"attempts must be between 1 and {MAX_ATTEMPTS}")
for name in ("base_seconds", "max_seconds"):
value = getattr(self, name)
if isinstance(value, bool) or not isinstance(value, numbers.Real):
raise ValueError(f"{name} must be a number")
if value < 0:
raise ValueError(f"{name} must be nonnegative")
if self.max_seconds < self.base_seconds:
raise ValueError("max_seconds must be at least base_seconds")
if not isinstance(self.jitter, bool):
raise ValueError("jitter must be a bool")
[docs]
def delay(self, attempt: int) -> float:
"""Capped exponential delay before the attempt after `attempt`.
Jitter matters more than the exponent: several workers that fail at the
same moment and retry in lockstep reproduce whatever they collided with.
"""
scaled = self.base_seconds * (2 ** min(attempt, MAX_BACKOFF_EXPONENT))
capped = min(float(scaled), float(self.max_seconds))
return rng.uniform(0.5, 1.0) * capped if self.jitter else capped
[docs]
@classmethod
def coerce(cls, value: int | RetryPolicy | None) -> RetryPolicy:
"""Accept `retry=3`, `retry=RetryPolicy(...)`, or `retry=None`.
The integer shorthand is the spelling callers reach for first; the
dataclass is there when the delays need tuning too.
"""
if value is None:
return cls()
if isinstance(value, cls):
return value
if isinstance(value, bool) or not isinstance(value, numbers.Integral):
raise TypeError("retry must be an int, a RetryPolicy, or None")
return cls(attempts=int(value))
def retrying(
policy: RetryPolicy,
operation: Callable[[], T],
*,
description: str,
sleep: Callable[[float], None] | None = None,
) -> T:
"""Call `operation` until it succeeds, the attempts run out, or it fails
in a way that retrying cannot help.
`sleep` is injectable so tests can assert the delay sequence without
spending it. It is resolved when called rather than bound as a default, so
that patching `time.sleep` works too.
"""
pause_for = time.sleep if sleep is None else sleep
earlier: list[str] = []
for attempt in range(policy.attempts):
try:
return operation()
except BaseException as exc:
if not transient(exc):
raise
remaining = policy.attempts - attempt - 1
if remaining == 0:
for note in earlier:
exc.add_note(note)
logger.warning(
"Giving up on %s after %d attempts", description, policy.attempts
)
raise
pause = policy.delay(attempt)
logger.warning(
"Retrying %s after %s: %s (attempt %d of %d, sleeping %.2fs)",
description,
type(exc).__name__,
exc,
attempt + 1,
policy.attempts,
pause,
)
earlier.append(
f"attempt {attempt + 1} of {policy.attempts} failed with "
f"{type(exc).__name__}: {exc}"
)
pause_for(pause)
# `attempts` is validated to be at least 1, so the loop always returns or
# raises. This exists so the function has no implicit `None` return.
raise AssertionError("unreachable: attempts is at least 1") # pragma: no cover
__all__ = [
"MAX_ATTEMPTS",
"MAX_BACKOFF_EXPONENT",
"NON_TRANSIENT",
"NON_TRANSIENT_ERRNOS",
"TRANSIENT",
"RetryPolicy",
"retrying",
"transient",
]