rag: resilient embedder — rotate/split on endpoint errors; 4-GPU embed pool (#3)

rag/embeddings.py

@@ -19,6 +19,7 @@ from __future__ import annotations
 
 import os
 import logging
+import time
 from typing import Any
 
 import httpx
@@ -43,11 +44,18 @@ EMBED_DIM = int(os.environ.get("EMBED_DIM", "768"))
 
 
 class OllamaEmbeddings(EmbeddingFunction):
-    """Calls /api/embed across N Ollama endpoints, naive round-robin.
+    """Calls /api/embed across N Ollama endpoints, round-robin per batch.
 
     For indexing throughput on multiple GPUs, run one Ollama container
     per GPU (pinned via NVIDIA_VISIBLE_DEVICES) and pass all their URLs
     in OLLAMA_URL — the embedder picks the next endpoint per batch.
+
+    Resilient (ported from zerto-docs PR #45): a failed call rotates to
+    the next endpoint and retries with backoff instead of failing the
+    whole rebuild. HTTP status errors additionally halve the input —
+    the .0.125 Windows Ollama (4090) 400s when its model runner dies on
+    an oversized input array, and one endpoint rejecting a batch the
+    others accept shouldn't kill a multi-hour index build.
     """
 
     def __init__(self, urls: list[str] = OLLAMA_URLS, model: str = EMBED_MODEL):
@@ -56,14 +64,37 @@ class OllamaEmbeddings(EmbeddingFunction):
         self._next = 0
 
     def __call__(self, input: Documents) -> Embeddings:
+        return self._embed(list(input), attempt=1)
+
+    def _embed(self, texts: list, attempt: int) -> Embeddings:
         url = self.urls[self._next % len(self.urls)]
         self._next += 1
-        with httpx.Client(timeout=300) as c:
-            r = c.post(f"{url}/api/embed",
-                       json={"model": self.model, "input": list(input)})
-            r.raise_for_status()
-            data = r.json()
-        return data.get("embeddings") or []
+        try:
+            with httpx.Client(timeout=300) as c:
+                r = c.post(f"{url}/api/embed",
+                           json={"model": self.model, "input": texts})
+                r.raise_for_status()
+                return r.json().get("embeddings") or []
+        except (httpx.TransportError, httpx.HTTPStatusError) as e:
+            if isinstance(e, httpx.HTTPStatusError):
+                desc = f"HTTP {e.response.status_code} ({e.response.text[:200]})"
+            else:
+                desc = f"transport error {type(e).__name__}"
+            if attempt >= 5:
+                log.error("%s from %s (%d texts) — giving up after %d attempts",
+                          desc, url, len(texts), attempt)
+                raise
+            if isinstance(e, httpx.HTTPStatusError) and len(texts) > 16:
+                mid = len(texts) // 2
+                log.warning("%s from %s — splitting %d texts into %d+%d (attempt %d)",
+                            desc, url, len(texts), mid, len(texts) - mid, attempt)
+                return (self._embed(texts[:mid], attempt + 1)
+                        + self._embed(texts[mid:], attempt + 1))
+            backoff = 0.5 * (2 ** (attempt - 1))  # 0.5, 1, 2, 4
+            log.warning("%s (attempt %d, %s) — retrying in %.1fs",
+                        desc, attempt, url, backoff)
+            time.sleep(backoff)
+            return self._embed(texts, attempt + 1)
 
     def name(self) -> str:                  # newer chromadb requires this
         return f"ollama:{self.model}"