Retry transient model download failures

Author: danias

scripts/real_backend_smoke.py

@@ -11,6 +11,9 @@
 from urllib import error, request
 
 
+SMOKE_RETRY_DELAYS_SECONDS = (10, 20, 40)
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(description="Run a real backend smoke test.")
     parser.add_argument(
@@ -21,11 +24,28 @@ def main() -> None:
     args = parser.parse_args()
 
     binary = args.binary
-    port = reserve_free_port()
+    run_with_retries("embed smoke", lambda: run_embed_smoke(binary))
+    run_with_retries("server smoke", lambda: run_server_smoke(binary, reserve_free_port()))
+    run_with_retries("daemon smoke", lambda: run_daemon_smoke(binary))
+
 
-    run_embed_smoke(binary)
-    run_server_smoke(binary, port)
-    run_daemon_smoke(binary)
+def run_with_retries(name: str, operation) -> None:
+    total_attempts = len(SMOKE_RETRY_DELAYS_SECONDS) + 1
+    for attempt in range(1, total_attempts + 1):
+        try:
+            operation()
+            return
+        except RuntimeError as exc:
+            if attempt >= total_attempts:
+                raise
+
+            delay_seconds = SMOKE_RETRY_DELAYS_SECONDS[attempt - 1]
+            print(
+                f"{name} failed on attempt {attempt}/{total_attempts}: {exc}\nRetrying in {delay_seconds}s...",
+                file=sys.stderr,
+                flush=True,
+            )
+            time.sleep(delay_seconds)
 
 
 def run_embed_smoke(binary: str) -> None:

src/bitloops_embeddings/backend/sentence_transformers_backend.py

@@ -1,10 +1,15 @@
 from __future__ import annotations
 
+import logging
+import time
 from pathlib import Path
 from typing import Any
 
 from bitloops_embeddings.errors import BackendLoadError, InferenceError
-from bitloops_embeddings.logging_utils import log_event
+from bitloops_embeddings.logging_utils import LOGGER_NAME, log_event
+
+
+MODEL_LOAD_RETRY_DELAYS_SECONDS = (5, 10, 20)
 
 
 class SentenceTransformersBackend:
@@ -56,19 +61,36 @@ def load(self) -> None:
             upstream_model_id=self._upstream_model_id,
             cache_dir=self._cache_dir,
         )
-        try:
-            self._model = SentenceTransformer(
-                self._upstream_model_id,
-                cache_folder=str(self._cache_dir),
-                device="cpu",
-            )
-            detected_dimensions = self._model.get_sentence_embedding_dimension()
-            if detected_dimensions is not None:
-                self._dimensions = int(detected_dimensions)
-        except Exception as exc:
-            raise BackendLoadError(
-                f"Failed to load model '{self.model_id}' from '{self._upstream_model_id}'."
-            ) from exc
+        max_attempts = len(MODEL_LOAD_RETRY_DELAYS_SECONDS) + 1
+        for attempt in range(1, max_attempts + 1):
+            try:
+                self._model = SentenceTransformer(
+                    self._upstream_model_id,
+                    cache_folder=str(self._cache_dir),
+                    device="cpu",
+                )
+                detected_dimensions = self._model.get_sentence_embedding_dimension()
+                if detected_dimensions is not None:
+                    self._dimensions = int(detected_dimensions)
+                break
+            except Exception as exc:
+                self._model = None
+                if attempt >= max_attempts or not _is_retryable_load_exception(exc):
+                    raise BackendLoadError(
+                        f"Failed to load model '{self.model_id}' from '{self._upstream_model_id}'."
+                    ) from exc
+
+                delay_seconds = MODEL_LOAD_RETRY_DELAYS_SECONDS[attempt - 1]
+                logging.getLogger(LOGGER_NAME).warning(
+                    "event=model_load_retry model_id=%s backend=%s attempt=%s max_attempts=%s delay_seconds=%s reason=%s",
+                    self.model_id,
+                    self.backend_name,
+                    attempt,
+                    max_attempts,
+                    delay_seconds,
+                    str(exc),
+                )
+                time.sleep(delay_seconds)
 
         log_event(
             "model_load_complete",
@@ -97,3 +119,23 @@ def embed(self, texts: list[str]) -> list[list[float]]:
 
     def close(self) -> None:
         self._model = None
+
+
+def _is_retryable_load_exception(exc: Exception) -> bool:
+    message = str(exc).lower()
+    retryable_markers = (
+        "http error 500",
+        "http error 502",
+        "http error 503",
+        "http error 504",
+        "connection error",
+        "connection aborted",
+        "connection reset",
+        "read timed out",
+        "timed out",
+        "temporarily unavailable",
+        "temporary failure",
+        "service unavailable",
+        "too many requests",
+    )
+    return any(marker in message for marker in retryable_markers)

tests/unit/test_sentence_transformers_backend.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import sys
+from types import ModuleType
+
+from bitloops_embeddings.backend.sentence_transformers_backend import SentenceTransformersBackend
+
+
+class FakeSentenceTransformer:
+    attempts = 0
+
+    def __init__(self, *args, **kwargs) -> None:
+        type(self).attempts += 1
+        if type(self).attempts < 3:
+            raise RuntimeError("HTTP Error 503 thrown while requesting HEAD https://huggingface.co/BAAI/bge-m3/resolve/main/config.json")
+
+    def get_sentence_embedding_dimension(self) -> int:
+        return 1024
+
+
+def test_sentence_transformers_backend_retries_transient_load_failures(
+    monkeypatch,
+    tmp_path,
+) -> None:
+    fake_module = ModuleType("sentence_transformers")
+    fake_module.SentenceTransformer = FakeSentenceTransformer
+    monkeypatch.setitem(sys.modules, "sentence_transformers", fake_module)
+    monkeypatch.setattr("bitloops_embeddings.backend.sentence_transformers_backend.time.sleep", lambda _: None)
+    FakeSentenceTransformer.attempts = 0
+
+    backend = SentenceTransformersBackend(
+        model_id="bge-m3",
+        upstream_model_id="BAAI/bge-m3",
+        cache_dir=tmp_path / "cache",
+        dimensions=1024,
+    )
+
+    backend.load()
+
+    assert backend.is_loaded is True
+    assert backend.dimensions == 1024
+    assert FakeSentenceTransformer.attempts == 3