danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add support for retrying docling requests

Browse source
This commit is contained in:
bryan davis 2026-06-25 12:11:25 -05:00
parent 73beae5a4c
commit b58f2f5496
No known key found for this signature in database
GPG key ID: D11B4A4C0C731E5E
2 changed files with 235 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

122
haiku_rag_slim/haiku/rag/providers/docling_serve.py
View file

@ -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)

126
tests/test_docling_serve_client.py
View file

@ -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(