diff --git a/haiku_rag_slim/haiku/rag/providers/docling_serve.py b/haiku_rag_slim/haiku/rag/providers/docling_serve.py index 2d32e98c..7416ccde 100644 --- a/haiku_rag_slim/haiku/rag/providers/docling_serve.py +++ b/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: - return next(self._instance_rotator) + def _pick_url(self, exclude: frozenset[str] = frozenset()) -> str: + """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, - etc.) propagate so the ingester's pipeline classifier can route - 4xx → PermanentError and 5xx/network → TransientError. ValueError - is raised by `_submit_and_wait` when docling-serve reports a task - failure or returns no task_id. + Retries on a different instance when an instance crashes or returns + 5xx/overload (see ``_run_with_retry``). Non-retryable httpx exceptions + (4xx other than 408/429) and the ValueError raised by + ``_submit_and_wait`` (task failure / missing task_id) propagate so the + 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) diff --git a/tests/test_docling_serve_client.py b/tests/test_docling_serve_client.py index b19360d4..d9c943d1 100644 --- a/tests/test_docling_serve_client.py +++ b/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(