add support for retrying docling requests

haiku_rag_slim/haiku/rag/providers/docling_serve.py

@@ -2,10 +2,34 @@
 
 import asyncio
 import itertools
-from typing import Any
+import logging
+from collections.abc import Awaitable, Callable
+from typing import Any, TypeVar
 
 import httpx
 
+logger = logging.getLogger(__name__)
+
+_T = TypeVar("_T")
+
+# Statuses worth retrying on another instance: 408/429 are transient overload,
+# 5xx covers a crashed or restarting docling-serve worker (the OOM-leak failure
+# mode). Other 4xx are the caller's fault and won't succeed elsewhere.
+_RETRYABLE_STATUS = frozenset({408, 429})
+
+
+def _is_retryable(exc: BaseException) -> bool:
+    """Whether a failed docling-serve request should be retried on another
+    instance. True for transport-level failures (connection reset / timeout —
+    an instance that crashed or went unresponsive) and server-side 5xx/overload;
+    False for other 4xx and task-level ``ValueError``\\ s, which won't succeed on
+    a retry (a deterministically bad document is handled upstream, not here)."""
+    if isinstance(exc, httpx.HTTPStatusError):
+        status = exc.response.status_code
+        return status in _RETRYABLE_STATUS or status >= 500
+    return isinstance(exc, httpx.TransportError)
+
+
 # Process-global registry of round-robin rotators over docling-serve
 # instances, keyed by the sorted tuple of base URLs. Clients are constructed
 # per-job by `get_converter` / `get_chunker`, so the rotation cursor has to
@@ -31,6 +55,9 @@ class DoclingServeClient:
         api_key: str | None = None,
         timeout: float = 300,
         transport: httpx.AsyncBaseTransport | None = None,
+        max_attempts: int = 3,
+        retry_base_delay: float = 0.5,
+        retry_max_delay: float = 8.0,
     ):
         urls = [base_urls] if isinstance(base_urls, str) else list(base_urls)
         if not urls:
@@ -40,6 +67,12 @@ class DoclingServeClient:
         self.timeout = timeout
         # transport is for testing — production callers leave it None.
         self._transport = transport
+        # Bounded retry with failover: docling-serve instances crash (memory
+        # leaks), so a request that hits a dying instance is retried, preferring
+        # an instance that hasn't already failed this request.
+        self._max_attempts = max(1, max_attempts)
+        self._retry_base_delay = retry_base_delay
+        self._retry_max_delay = retry_max_delay
         # setdefault is atomic under the GIL — concurrent constructors with
         # the same instance set will end up sharing one rotator.
         key = tuple(sorted(self.base_urls))
@@ -56,8 +89,24 @@ class DoclingServeClient:
         callers should let `_pick_url` round-robin per request."""
         return self.base_urls[0]
 
-    def _pick_url(self) -> str:
+    def _pick_url(self, exclude: frozenset[str] = frozenset()) -> str:
-        return next(self._instance_rotator)
+        """Next instance in the round-robin, preferring one not in ``exclude``
+        (instances that already failed this request). Falls back to a
+        possibly-excluded instance when every instance is excluded — a
+        single-instance fleet, or one where all instances failed, still gets
+        retried after backoff in case the instance has since restarted."""
+        url = next(self._instance_rotator)
+        if url not in exclude:
+            return url
+        for _ in range(len(self.base_urls) - 1):
+            url = next(self._instance_rotator)
+            if url not in exclude:
+                return url
+        return url
+
+    def _retry_delay(self, attempt_no: int) -> float:
+        """Capped exponential backoff between retry attempts."""
+        return min(self._retry_base_delay * (2**attempt_no), self._retry_max_delay)
 
     def _get_headers(self) -> dict[str, str]:
         """Get headers for API requests."""
@@ -66,6 +115,47 @@ class DoclingServeClient:
             headers["X-Api-Key"] = self.api_key
         return headers
 
+    async def _run_with_retry(
+        self,
+        attempt: Callable[[httpx.AsyncClient, str], Awaitable[_T]],
+        name: str,
+    ) -> _T:
+        """Run ``attempt(client, base_url)`` with bounded retry + failover.
+
+        Each attempt runs the full submit → poll → fetch trio against one
+        instance (task IDs are instance-local, so a trio can't be split across
+        instances). On a retryable failure — a transport error (crashed /
+        unresponsive instance) or 5xx/overload — the next attempt prefers an
+        instance that hasn't already failed this request. Non-retryable errors
+        (4xx, task-level ``ValueError``) propagate immediately.
+        """
+        tried: set[str] = set()
+        last_exc: BaseException | None = None
+        for attempt_no in range(self._max_attempts):
+            base_url = self._pick_url(exclude=frozenset(tried))
+            try:
+                async with self._httpx_client() as client:
+                    return await attempt(client, base_url)
+            except Exception as exc:
+                if not _is_retryable(exc):
+                    raise
+                last_exc = exc
+                tried.add(base_url)
+                remaining = self._max_attempts - attempt_no - 1
+                if remaining <= 0:
+                    raise
+                logger.warning(
+                    "docling-serve request for %s failed on %s (%s); retrying "
+                    "on another instance (%d attempt(s) left)",
+                    name,
+                    base_url,
+                    exc,
+                    remaining,
+                )
+                await asyncio.sleep(self._retry_delay(attempt_no))
+        # Unreachable: the final iteration either returns or re-raises.
+        raise last_exc or RuntimeError("retry loop exited without a result")
+
     async def _submit_and_wait(
         self,
         client: httpx.AsyncClient,
@@ -118,15 +208,16 @@ class DoclingServeClient:
     ) -> dict[str, Any]:
         """Submit a task and poll until completion; fetch result as JSON.
 
-        httpx exceptions (ConnectError, HTTPStatusError, TimeoutException,
+        Retries on a different instance when an instance crashes or returns
-        etc.) propagate so the ingester's pipeline classifier can route
+        5xx/overload (see ``_run_with_retry``). Non-retryable httpx exceptions
-        4xx → PermanentError and 5xx/network → TransientError. ValueError
+        (4xx other than 408/429) and the ValueError raised by
-        is raised by `_submit_and_wait` when docling-serve reports a task
+        ``_submit_and_wait`` (task failure / missing task_id) propagate so the
-        failure or returns no task_id.
+        ingester's pipeline classifier can route them (4xx → PermanentError,
+        ValueError → TransientError for a whole-document retry).
         """
         headers = self._get_headers()
-        base_url = self._pick_url()
+
-        async with self._httpx_client() as client:
+        async def _attempt(client: httpx.AsyncClient, base_url: str) -> dict[str, Any]:
             task_id = await self._submit_and_wait(
                 client, base_url, endpoint, files, data, headers, name
             )
@@ -135,6 +226,8 @@ class DoclingServeClient:
             result_response.raise_for_status()
             return result_response.json()
 
+        return await self._run_with_retry(_attempt, name)
+
     async def submit_and_poll_zip(
         self,
         endpoint: str,
@@ -147,11 +240,12 @@ class DoclingServeClient:
         Used when the caller requested ``target_type=zip`` (e.g. to retrieve
         picture image bytes that docling-serve only emits as referenced files
         bundled into a zip archive). The submit/poll flow is identical to
-        ``submit_and_poll``; only the result-fetching step differs.
+        ``submit_and_poll`` (including the retry + failover); only the
+        result-fetching step differs.
         """
         headers = self._get_headers()
-        base_url = self._pick_url()
+
-        async with self._httpx_client() as client:
+        async def _attempt(client: httpx.AsyncClient, base_url: str) -> bytes:
             task_id = await self._submit_and_wait(
                 client, base_url, endpoint, files, data, headers, name
             )
@@ -159,3 +253,5 @@ class DoclingServeClient:
             result_response = await client.get(result_url, headers=headers)
             result_response.raise_for_status()
             return result_response.content
+
+        return await self._run_with_retry(_attempt, name)

tests/test_docling_serve_client.py

@@ -110,6 +110,132 @@ def test_round_robin_shared_across_fresh_clients():
     assert picks == [urls[0], urls[1], urls[2], urls[0]]
 
 
+def _failover_transport(
+    down_hosts: set[str], task_id: str, result: dict
+) -> tuple[httpx.MockTransport, list[str]]:
+    """MockTransport where any request to a host in ``down_hosts`` raises a
+    ConnectError (simulating a crashed instance); other hosts serve a normal
+    submit/poll/result trio. Records every host attempted."""
+    seen_hosts: list[str] = []
+    routes = _success_routes(task_id, result)
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        seen_hosts.append(request.url.host)
+        if request.url.host in down_hosts:
+            raise httpx.ConnectError("connection refused", request=request)
+        return routes.get(request.url.path, httpx.Response(404))
+
+    return httpx.MockTransport(handler), seen_hosts
+
+
+@pytest.mark.asyncio
+async def test_retry_fails_over_to_healthy_instance():
+    """A crashed instance (connection error) is retried on the next instance,
+    and the call succeeds without surfacing the failure."""
+    transport, seen = _failover_transport(
+        down_hosts={"down-a"}, task_id="t", result={"ok": True}
+    )
+    client = DoclingServeClient(
+        base_urls=["http://down-a:5001", "http://up-a:5001"],
+        transport=transport,
+        retry_base_delay=0.0,  # don't sleep in tests
+    )
+
+    result = await client.submit_and_poll(
+        endpoint="/v1/convert/file/async",
+        files={"file": ("x.md", b"x", "text/markdown")},
+        data={},
+    )
+
+    assert result == {"ok": True}
+    # First attempt hit the crashed instance; failover moved to the healthy one.
+    assert seen[0] == "down-a"
+    assert "up-a" in seen
+    # The successful trio (submit/poll/result) all landed on the healthy host.
+    assert seen[-3:] == ["up-a", "up-a", "up-a"]
+
+
+@pytest.mark.asyncio
+async def test_retry_5xx_fails_over():
+    """A 5xx from a struggling instance is retried elsewhere (status-based
+    retryability, distinct from the transport-error path)."""
+    seen: list[str] = []
+    ok = _success_routes("t", {"ok": True})
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        seen.append(request.url.host)
+        if request.url.host == "sad-b":
+            return httpx.Response(503)
+        return ok.get(request.url.path, httpx.Response(404))
+
+    client = DoclingServeClient(
+        base_urls=["http://sad-b:5001", "http://ok-b:5001"],
+        transport=httpx.MockTransport(handler),
+        retry_base_delay=0.0,
+    )
+
+    result = await client.submit_and_poll(
+        endpoint="/v1/convert/file/async",
+        files={"file": ("x.md", b"x", "text/markdown")},
+        data={},
+    )
+    assert result == {"ok": True}
+    assert seen[0] == "sad-b"
+    assert "ok-b" in seen
+
+
+@pytest.mark.asyncio
+async def test_retry_exhausts_all_instances_then_raises():
+    """When every instance is down, the call retries up to max_attempts and
+    then surfaces the transport error."""
+    transport, seen = _failover_transport(
+        down_hosts={"down-c", "down-d"}, task_id="t", result={}
+    )
+    client = DoclingServeClient(
+        base_urls=["http://down-c:5001", "http://down-d:5001"],
+        transport=transport,
+        max_attempts=2,
+        retry_base_delay=0.0,
+    )
+
+    with pytest.raises(httpx.ConnectError):
+        await client.submit_and_poll(
+            endpoint="/v1/convert/file/async",
+            files={"file": ("x.md", b"x", "text/markdown")},
+            data={},
+        )
+
+    # Two attempts, each preferring a not-yet-failed instance.
+    assert seen == ["down-c", "down-d"]
+
+
+@pytest.mark.asyncio
+async def test_4xx_is_not_retried():
+    """A 4xx (other than 408/429) is the caller's fault — it must not be
+    retried on another instance; it propagates after a single attempt."""
+    seen: list[str] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        seen.append(request.url.host)
+        return httpx.Response(400, json={"detail": "bad request"})
+
+    client = DoclingServeClient(
+        base_urls=["http://e:5001", "http://f:5001"],
+        transport=httpx.MockTransport(handler),
+        retry_base_delay=0.0,
+    )
+
+    with pytest.raises(httpx.HTTPStatusError):
+        await client.submit_and_poll(
+            endpoint="/v1/convert/file/async",
+            files={"file": ("x.md", b"x", "text/markdown")},
+            data={},
+        )
+
+    # Single attempt — no failover on a non-retryable status.
+    assert seen == ["e"]
+
+
 @pytest.mark.asyncio
 async def test_zip_endpoint_uses_round_robin_too():
     transport, seen = _scripted_transport(