fix(python): stop retrying permanent embedding errors and cap the backoff delay (#4542)

Permanent embedding errors (HTTP 400/404/422 and other 4xx) were retried
the full `max_retries` times with an
uncapped, multiplicatively jittered backoff, so a typo'd model name
looked like a hang of about an hour
(median 50 min of sleeping with a 404; see the issue for the repro and
numbers).

Changes (`python/python/lancedb/embeddings/utils.py`,
`embeddings/base.py`):
- An HTTP status in 4xx other than 408/409/429 is raised immediately.
The status is read from `status_code`,
`http_status`, `response.status_code`, an integer `code`
(google-api-core, urllib) or a botocore
`response["ResponseMetadata"]["HTTPStatusCode"]`. AWS throttling
reported as 400 stays retryable (matched by
  error code). `AuthenticationError` by name is unchanged.
- Still retried: 408, 409, 429, >=500, and anything without a
recognisable status (connection errors, timeouts,
custom exceptions), so existing custom embedding functions behave as
before.
- New keyword `max_delay=60` (last in the signature, default keeps call
sites working). `delay = min(delay * base,
max_delay)` with equal jitter `delay/2 + random() * delay/2`. With
defaults the waits are [1,2], [2,4], [4,8],
[8,16], [16,32], [30,60], [30,60] s, worst case about 182 s instead of
up to 2.5 h.
- The "Maximum number of retries" exception now has `from e`. The
`max_retries` comment in `base.py` describes the
  rules.

Not changed: `embeddings/openai.py` (it swallows `BadRequestError`; open
#4154/#4005/#3152 cover it), Retry-After
handling (the code never used it), and the "C Data interface error"
wrapping from the native layer (the original
message is now included in it).

Behaviour changes to be aware of: every 4xx except 408/409/429 fails
fast, including 401/403 coming through
`response.status_code` (requests/httpx), which used to be retried; a
provider that signals transient failures with
a 4xx and no AWS-style code is no longer retried; the first retry waits
1-2 s instead of 2-4 s.

Tests (`python/python/tests/test_embeddings.py`): 48 new cases (23 test
functions): permanent errors are not
retried, transient errors are, botocore-style errors, the delay is
capped, exact exponential growth without jitter,
jitter bounds, no float overflow for a large `max_retries`, the original
error is attached, and end to end through
`table.add` and `table.search` with a stub function (one call and no
sleep). On `main` 23 of them fail with the
change stashed; the rest are guards that must not change.
`test_embeddings.py`: 61 passed, 4 skipped.
`ruff format`/`ruff check` are clean on the touched files.

Verification caveat: `main`'s Python layer needs a newer native module
than the 0.40.0 wheel, so the tests were
run with the wheel's native module preloaded through a small shim (the
same shim ran the baseline on `main`); a
wider run of embeddings/pydantic/query/table tests has the same 36
failures before and after (that mismatch).
Not run: real provider APIs (a fake OpenAI-style server was used), Node
or Rust code.

Fixes #4541

Co-authored-by: Xuanwo <github@xuanwo.io>

python/python/lancedb/embeddings/base.py

@@ -38,9 +38,11 @@ class EmbeddingFunction(BaseModel, ABC):
     3. ndims() which returns the number of dimensions of the vector column
     """
 
-    max_retries: int = (
-        7  # Setting 0 disables retires. Maybe this should not be enabled by default,
-    )
+    # Maximum number of retries for transient errors (rate limits, 5xx responses,
+    # connection errors, ...), with exponential backoff capped at 60 seconds between
+    # attempts. Permanent errors (HTTP 4xx other than 408/409/429) are never retried.
+    # Setting 0 disables retries.
+    max_retries: int = 7
     _ndims: int = PrivateAttr()
     _original_args: dict = PrivateAttr()
 

python/python/lancedb/embeddings/utils.py

@@ -241,15 +241,30 @@ def retry_with_exponential_backoff(
     exponential_base: float = 2,
     jitter: bool = True,
     max_retries: int = 7,
+    max_delay: float = 60,
 ):
     """Retry a function with exponential backoff.
 
+    Only errors that may succeed when repeated are retried. Errors that carry an
+    HTTP status code are retried when the status is 408, 409, 429 or >= 500; any
+    other 4xx status (for example 400, 401, 403, 404 or 422) is a permanent failure
+    that is raised immediately. Errors without a recognisable status code (for
+    example connection errors, timeouts, or exceptions raised by a custom embedding
+    function) are retried.
+
+    The delay before the n-th retry is ``initial_delay * exponential_base ** n``,
+    never more than ``max_delay``. With ``jitter`` enabled, a random amount of up to
+    half of that delay is removed ("equal jitter"), so the actual wait is always in
+    the range ``[delay / 2, delay]``.
+
     Args:
         func (function): The function to be retried.
         initial_delay (float): Initial delay in seconds (default is 1).
         exponential_base (float): The base for exponential backoff (default is 2).
         jitter (bool): Whether to add jitter to the delay (default is True).
         max_retries (int): Maximum number of retries (default is 7).
+        max_delay (float): Maximum delay between two attempts, in seconds
+            (default is 60).
 
     Returns:
         function: The decorated function.
@@ -265,13 +280,12 @@ def retry_with_exponential_backoff(
             try:
                 return func(*args, **kwargs)
 
-            # Currently retrying on all exceptions as there is no way to know the
-            # format of the error msgs used by different APIs. We'll log the error
-            # and say that it is assumed that if this portion errors out, it's due
-            # to rate limit but the user should check the error message to be sure.
+            # Errors are retried unless they are known to be permanent (see
+            # `_is_non_retryable_error`). Custom embedding functions can raise
+            # arbitrary exceptions, so unknown errors are assumed to be transient.
             except Exception as e:  # noqa: PERF203
-                # Don't retry on authentication errors (e.g., OpenAI 401)
-                # These are permanent failures that won't be fixed by retrying
+                # Don't retry permanent failures (e.g., OpenAI 401 or 404). They
+                # won't be fixed by retrying and would only delay the real error.
                 if _is_non_retryable_error(e):
                     raise
 
@@ -280,24 +294,72 @@ def retry_with_exponential_backoff(
                 if num_retries > max_retries:
                     raise Exception(
                         f"Maximum number of retries ({max_retries}) exceeded.", e
-                    )
+                    ) from e
 
-                delay *= exponential_base * (1 + jitter * random.random())
+                delay = min(delay * exponential_base, max_delay)
+                sleep_for = delay
+                if jitter:
+                    sleep_for = delay / 2 + random.random() * delay / 2
                 logging.warning(
                     "Error occurred: %s \n Retrying in %s seconds (retry %s of %s) \n",
                     e,
-                    delay,
+                    sleep_for,
                     num_retries,
                     max_retries,
                 )
-                time.sleep(delay)
+                time.sleep(sleep_for)
 
     return wrapper
 
 
+# HTTP statuses in the 4xx range that are worth retrying: request timeout,
+# conflict (the OpenAI SDK treats it as transient) and rate limiting.
+_RETRYABLE_CLIENT_STATUSES = (408, 409, 429)
+
+# AWS error codes that are reported with a 4xx status but are transient.
+_RETRYABLE_AWS_ERROR_CODES = (
+    "Throttl",
+    "TooManyRequests",
+    "RequestLimitExceeded",
+    "ProvisionedThroughputExceeded",
+    "ModelNotReady",
+    "ModelTimeout",
+)
+
+
+def _get_http_status(error: Exception):
+    """Best-effort extraction of an HTTP status code from an exception.
+
+    Provider SDKs expose it in different ways: ``status_code`` (openai, cohere),
+    ``http_status`` (voyageai), ``response.status_code`` (requests, httpx),
+    ``code`` (google-api-core, urllib) or ``response["ResponseMetadata"]``
+    (botocore). Returns None if no status can be found.
+    """
+    candidates = [
+        getattr(error, "status_code", None),
+        getattr(error, "http_status", None),
+    ]
+    response = getattr(error, "response", None)
+    if isinstance(response, dict):
+        candidates.append(response.get("ResponseMetadata", {}).get("HTTPStatusCode"))
+    else:
+        candidates.append(getattr(response, "status_code", None))
+    candidates.append(getattr(error, "code", None))
+
+    for status in candidates:
+        if isinstance(status, int) and not isinstance(status, bool):
+            if 100 <= status <= 599:
+                return status
+    return None
+
+
 def _is_non_retryable_error(error: Exception) -> bool:
     """Check if an error should not be retried.
 
+    An error is non-retryable if it is an authentication error or if it carries an
+    HTTP client-error status (4xx) other than 408, 409 and 429. Errors without a
+    recognisable status are treated as retryable.
+
     Args:
         error: The exception to check
 
@@ -314,18 +376,25 @@ def _is_non_retryable_error(error: Exception) -> bool:
     if error_type == "AuthenticationError":
         return True
 
-    # Check for other common non-retryable HTTP status codes
-    # 401 Unauthorized, 403 Forbidden
-    if hasattr(error, "status_code"):
-        if error.status_code in (401, 403):
-            return True
-
     # A missing API key or a missing optional dependency fails the same way on
     # every attempt
     if isinstance(error, (MissingApiKeyError, ImportError)):
         return True
 
-    return False
+    status = _get_http_status(error)
+    if status is None or not 400 <= status < 500:
+        return False
+    if status in _RETRYABLE_CLIENT_STATUSES:
+        return False
+
+    # AWS reports throttling with a 400 status; the error code tells them apart.
+    response = getattr(error, "response", None)
+    if isinstance(response, dict):
+        aws_code = str(response.get("Error", {}).get("Code", ""))
+        if aws_code.startswith(_RETRYABLE_AWS_ERROR_CODES):
+            return False
+
+    return True
 
 
 def url_retrieve(url: str):

python/python/tests/test_embeddings.py

@@ -633,6 +633,269 @@ def test_openai_no_retry_on_401(mock_sleep):
     assert mock_sleep.call_count == 0
 
 
+class _HttpStatusError(Exception):
+    """Stand-in for provider SDK errors (openai, cohere, ...) with a status code."""
+
+    def __init__(self, status_code, message="error"):
+        super().__init__(message)
+        self.status_code = status_code
+
+
+class _ResponseStatusError(Exception):
+    """Stand-in for requests/httpx errors, which expose `response.status_code`."""
+
+    def __init__(self, status_code):
+        super().__init__(f"HTTP {status_code}")
+        self.response = MagicMock(status_code=status_code)
+
+
+def _run_with_retry(side_effect, **kwargs):
+    """Run a function through the retry helper with a mocked `time.sleep`.
+
+    Returns (mock_func, sleeps, error).
+    """
+    from lancedb.embeddings.utils import retry_with_exponential_backoff
+
+    mock_func = MagicMock(side_effect=side_effect)
+    sleeps = []
+    error = None
+    with patch("time.sleep", side_effect=sleeps.append):
+        try:
+            retry_with_exponential_backoff(mock_func, **kwargs)()
+        except Exception as e:
+            error = e
+    return mock_func, sleeps, error
+
+
+@pytest.mark.parametrize("status", [400, 401, 403, 404, 405, 410, 413, 422])
+@pytest.mark.parametrize("make_error", [_HttpStatusError, _ResponseStatusError])
+def test_retry_does_not_retry_permanent_http_errors(status, make_error):
+    err = make_error(status)
+    mock_func, sleeps, error = _run_with_retry(err, max_retries=7)
+
+    assert error is err
+    assert mock_func.call_count == 1
+    assert sleeps == []
+
+
+@pytest.mark.parametrize("status", [408, 409, 429, 500, 502, 503, 504])
+@pytest.mark.parametrize("make_error", [_HttpStatusError, _ResponseStatusError])
+def test_retry_retries_transient_http_errors(status, make_error):
+    mock_func, sleeps, error = _run_with_retry(
+        [make_error(status), make_error(status), "result"], max_retries=7
+    )
+
+    assert error is None
+    assert mock_func.call_count == 3
+    assert len(sleeps) == 2
+
+
+@pytest.mark.parametrize(
+    "exc",
+    [
+        ConnectionError("connection reset"),
+        TimeoutError("timed out"),
+        OSError("network is unreachable"),
+        ValueError("custom embedding function failure"),
+        RuntimeError("boom"),
+        Exception("plain exception"),
+    ],
+)
+def test_retry_retries_errors_without_http_status(exc):
+    mock_func, sleeps, error = _run_with_retry([exc, exc, "result"], max_retries=7)
+
+    assert error is None
+    assert mock_func.call_count == 3
+    assert len(sleeps) == 2
+
+
+def test_retry_error_status_attributes_of_other_types_are_ignored():
+    # `code` is a string for openai errors ("model_not_found"), and arbitrary
+    # exceptions may carry unrelated numeric attributes. Neither is an HTTP status.
+    class WeirdError(Exception):
+        code = "model_not_found"
+        status_code = None
+
+    class ExitCodeError(Exception):
+        code = 2
+
+    for exc in (WeirdError("a"), ExitCodeError("b")):
+        mock_func, sleeps, error = _run_with_retry(exc, max_retries=2)
+        assert mock_func.call_count == 3
+        assert len(sleeps) == 2
+
+
+def test_retry_botocore_style_errors():
+    # botocore ClientError exposes `response["ResponseMetadata"]["HTTPStatusCode"]`.
+    # Bedrock reports throttling with a 400, which must stay retryable.
+    class ClientError(Exception):
+        def __init__(self, status, code):
+            super().__init__(code)
+            self.response = {
+                "ResponseMetadata": {"HTTPStatusCode": status},
+                "Error": {"Code": code},
+            }
+
+    throttled = ClientError(400, "ThrottlingException")
+    mock_func, sleeps, error = _run_with_retry([throttled, "result"], max_retries=3)
+    assert error is None
+    assert mock_func.call_count == 2
+
+    invalid = ClientError(400, "ValidationException")
+    mock_func, sleeps, error = _run_with_retry(invalid, max_retries=3)
+    assert error is invalid
+    assert mock_func.call_count == 1
+    assert sleeps == []
+
+
+def test_retry_authentication_error_by_name_not_retried():
+    class AuthenticationError(Exception):
+        pass
+
+    mock_func, sleeps, error = _run_with_retry(AuthenticationError("bad key"))
+    assert isinstance(error, AuthenticationError)
+    assert mock_func.call_count == 1
+    assert sleeps == []
+
+
+def test_retry_gives_up_after_max_retries_with_original_error_attached():
+    original = _HttpStatusError(503, "overloaded")
+    mock_func, sleeps, error = _run_with_retry(original, max_retries=3)
+
+    assert mock_func.call_count == 4
+    assert len(sleeps) == 3
+    assert "Maximum number of retries (3) exceeded" in str(error)
+    assert error.args[1] is original
+    assert error.__cause__ is original
+
+
+def test_retry_delay_is_capped_by_max_delay():
+    # Default settings (7 retries, 60s cap): the old implementation could sleep for
+    # up to ~47 minutes in a single retry and hours in total.
+    exc = _HttpStatusError(503)
+    for seed in range(50):
+        import random
+
+        random.seed(seed)
+        _, sleeps, _ = _run_with_retry(exc)
+        assert len(sleeps) == 7
+        assert max(sleeps) <= 60
+        assert sum(sleeps) <= 7 * 60
+
+    # Custom cap
+    _, sleeps, _ = _run_with_retry(exc, max_retries=10, max_delay=5)
+    assert len(sleeps) == 10
+    assert max(sleeps) <= 5
+
+
+def test_retry_delay_without_jitter_is_exact_exponential_then_capped():
+    _, sleeps, _ = _run_with_retry(
+        _HttpStatusError(500),
+        initial_delay=1,
+        exponential_base=2,
+        jitter=False,
+        max_retries=8,
+        max_delay=30,
+    )
+    assert sleeps == [2, 4, 8, 16, 30, 30, 30, 30]
+
+
+def test_retry_jitter_is_bounded():
+    # Each wait is within [delay / 2, delay], where delay grows by exponential_base
+    # per retry (no random factor is compounded across retries).
+    exc = _HttpStatusError(500)
+    for seed in range(50):
+        import random
+
+        random.seed(seed)
+        _, sleeps, _ = _run_with_retry(exc, max_retries=8, max_delay=100)
+        for i, slept in enumerate(sleeps):
+            delay = min(2 ** (i + 1), 100)
+            assert delay / 2 <= slept <= delay
+
+    # Extremes of the random source map to the ends of the range
+    with patch("random.random", return_value=0.0):
+        _, sleeps, _ = _run_with_retry(exc, max_retries=3)
+    assert sleeps == [1, 2, 4]
+    with patch("random.random", return_value=1.0):
+        _, sleeps, _ = _run_with_retry(exc, max_retries=3)
+    assert sleeps == [2, 4, 8]
+
+
+def test_retry_large_max_retries_does_not_overflow():
+    _, sleeps, _ = _run_with_retry(_HttpStatusError(503), max_retries=1500)
+    assert len(sleeps) == 1500
+    assert max(sleeps) <= 60
+
+
+def _register_failing_embedding(name, error):
+    """Register an embedding function that always raises `error`.
+
+    Returns the class and a list that records every call. The function is
+    re-created from table metadata, so the call count has to live outside of it.
+    """
+    calls = []
+
+    @register(name)
+    class FailingEmbeddingFunction(TextEmbeddingFunction):
+        def ndims(self):
+            return 4
+
+        def generate_embeddings(self, texts):
+            calls.append(len(texts))
+            raise error
+
+    return FailingEmbeddingFunction, calls
+
+
+@pytest.mark.parametrize("status", [400, 404, 422])
+def test_permanent_embedding_error_is_not_retried(tmp_path, status):
+    """A permanent provider error (e.g. a misspelled model name) must surface
+    immediately from both `table.add` and `table.search` instead of being retried
+    with minutes of backoff."""
+    error = _HttpStatusError(status, "The model `typo-embed` does not exist")
+    cls, calls = _register_failing_embedding(f"failing-embedding-{status}", error)
+    func = cls.create()  # default max_retries=7
+
+    class Schema(LanceModel):
+        text: str = func.SourceField()
+        vector: Vector(func.ndims()) = func.VectorField()
+
+    db = lancedb.connect(tmp_path)
+    table = db.create_table("t", schema=Schema)
+
+    with patch("time.sleep") as mock_sleep:
+        with pytest.raises(Exception, match="The model `typo-embed` does not exist"):
+            table.add([{"text": "hello"}])
+        assert len(calls) == 1
+
+        with pytest.raises(_HttpStatusError, match="does not exist"):
+            table.search("hello").limit(1).to_list()
+        assert len(calls) == 2
+
+    mock_sleep.assert_not_called()
+
+
+def test_transient_embedding_error_is_retried_in_table_add(tmp_path):
+    cls, calls = _register_failing_embedding(
+        "failing-embedding-503", _HttpStatusError(503)
+    )
+    func = cls.create(max_retries=2)
+
+    class Schema(LanceModel):
+        text: str = func.SourceField()
+        vector: Vector(func.ndims()) = func.VectorField()
+
+    db = lancedb.connect(tmp_path)
+    table = db.create_table("t", schema=Schema)
+
+    with patch("time.sleep") as mock_sleep:
+        with pytest.raises(Exception, match="Maximum number of retries"):
+            table.add([{"text": "hello"}])
+
+    assert len(calls) == 3
+    assert mock_sleep.call_count == 2
+    assert all(call.args[0] <= 60 for call in mock_sleep.call_args_list)
 @patch("time.sleep")
 @pytest.mark.parametrize(
     "error",