Merge pull request #486 from ggozad/feat/breakers-per-instance

docling-serve client: fail over to another instance with per-instance circuit breaking
This commit is contained in:
Yiorgis Gozadinos 2026-07-08 11:24:49 +03:00 committed by GitHub
commit 87548770a3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 528 additions and 60 deletions

View file

@ -4,6 +4,7 @@
### Added ### Added
- `update_document` accepts a `uri` argument to change a document's URI. - `update_document` accepts a `uri` argument to change a document's URI.
- docling-serve requests fail over to another instance on transport/5xx errors and skip instances whose circuit breaker is open; tune via `providers.docling_serve.max_attempts` and `providers.docling_serve.circuit_breaker`.
## [0.63.2] - 2026-07-03 ## [0.63.2] - 2026-07-03

View file

@ -96,13 +96,21 @@ providers:
- http://gpu-1:5001 - http://gpu-1:5001
- http://cpu-1:5001 - http://cpu-1:5001
- http://cpu-2:5001 - http://cpu-2:5001
max_attempts: 3
circuit_breaker:
failure_threshold: 3
cooldown_s: 30.0
``` ```
The round-robin counter is per-process — multiple concurrent ingester or The round-robin counter is per-process — multiple concurrent ingester or
client processes pick independently, so the distribution evens out over many client processes pick independently, so the distribution evens out over many
jobs without coordination. For tight load balancing, failover, or health jobs without coordination. When a listed instance crashes or returns 5xx, the
checks, run a real load balancer (nginx `least_conn`, etc.) in front and client fails the request over to another instance (up to `max_attempts`) and
configure a single URL here. opens that instance's circuit breaker so subsequent jobs skip it until its
`cooldown_s` elapses. An external load balancer can only front docling-serve in
RQ mode (shared Redis task state); with the default standalone instances the
submit / poll / result trio is instance-pinned, so the failover and health
checks live in the client.
**Tuning `ingester.workers.worker_count` for docling-serve users**: convert **Tuning `ingester.workers.worker_count` for docling-serve users**: convert
is usually the throughput ceiling — a default docling-serve instance is usually the throughput ceiling — a default docling-serve instance

View file

@ -56,10 +56,7 @@ class DoclingServeChunker(DocumentChunker):
def __init__(self, config: AppConfig = Config): def __init__(self, config: AppConfig = Config):
self.config = config self.config = config
self.client = DoclingServeClient( self.client = DoclingServeClient.from_config(config.providers.docling_serve)
base_urls=config.providers.docling_serve.base_urls,
api_key=config.providers.docling_serve.api_key,
)
self.chunker_type = config.processing.chunker_type self.chunker_type = config.processing.chunker_type
def _build_chunking_data(self) -> dict[str, str | list[str]]: def _build_chunking_data(self) -> dict[str, str | list[str]]:

View file

@ -5,7 +5,7 @@ from haiku.rag.config import CircuitBreakerConfig
class CircuitBreaker: class CircuitBreaker:
"""Three-state breaker over discover() failures. """Three-state breaker over repeated failures of a single target.
- closed: failures are counted; threshold flips to open. - closed: failures are counted; threshold flips to open.
- open: probes are blocked until cooldown elapses, then a single probe - open: probes are blocked until cooldown elapses, then a single probe

View file

@ -11,6 +11,7 @@ from haiku.rag.config.models import (
AppConfig, AppConfig,
CircuitBreakerConfig, CircuitBreakerConfig,
ConversionOptions, ConversionOptions,
DoclingServeConfig,
EmbeddingModelConfig, EmbeddingModelConfig,
EmbeddingsConfig, EmbeddingsConfig,
FSSourceConfig, FSSourceConfig,
@ -40,6 +41,7 @@ __all__ = [
"AppConfig", "AppConfig",
"CircuitBreakerConfig", "CircuitBreakerConfig",
"ConversionOptions", "ConversionOptions",
"DoclingServeConfig",
"EmbeddingModelConfig", "EmbeddingModelConfig",
"EmbeddingsConfig", "EmbeddingsConfig",
"FSSourceConfig", "FSSourceConfig",

View file

@ -237,15 +237,44 @@ class OllamaConfig(BaseModel):
) )
class CircuitBreakerConfig(BaseModel):
"""Breaker over repeated failures of a single target. Stops callers from
hammering a target that's persistently failing (an ingester source, a
docling-serve instance)."""
failure_threshold: int = Field(
default=5, description="Consecutive failures before the breaker opens."
)
cooldown_s: float = Field(
default=600.0,
description="How long the breaker stays open before allowing a probe.",
)
class DoclingServeConfig(BaseModel): class DoclingServeConfig(BaseModel):
"""docling-serve endpoints. Accepts a single URL or a list — when a list """docling-serve endpoints. Accepts a single URL or a list — when a list is
is given, the client round-robins jobs across the URLs. Each job's given, the client round-robins jobs across the URLs, fails a request over to
submit/poll/result trio stays on the same instance (task IDs are another instance when one crashes or returns 5xx, and trips a per-instance
instance-local). The round-robin counter is per-process; for true load circuit breaker so repeated failures route around a dead instance. Each
balancing or failover, put an LB in front.""" job's submit/poll/result trio stays on the same instance (task IDs are
instance-local).
An external load balancer can only front docling-serve in RQ mode (shared
Redis task state); in the default standalone LocalOrchestrator mode the trio
is instance-pinned, so this client does the balancing/failover itself."""
base_url: str | list[str] = "http://localhost:5001" base_url: str | list[str] = "http://localhost:5001"
api_key: str = "" api_key: str = ""
max_attempts: int = Field(
default=3,
description="Max attempts per request across the fleet before giving up; "
"each retry fails over to another instance.",
)
circuit_breaker: CircuitBreakerConfig = Field(
default_factory=lambda: CircuitBreakerConfig(
failure_threshold=3, cooldown_s=30.0
)
)
@property @property
def base_urls(self) -> list[str]: def base_urls(self) -> list[str]:
@ -314,19 +343,6 @@ class RetryPolicyConfig(BaseModel):
jitter: float = Field(default=0.25, ge=0.0, le=1.0) jitter: float = Field(default=0.25, ge=0.0, le=1.0)
class CircuitBreakerConfig(BaseModel):
"""Per-source breaker over discover() failures. Stops the ingester from
hammering a source that's persistently failing."""
failure_threshold: int = Field(
default=5, description="Consecutive failures before the breaker opens."
)
cooldown_s: float = Field(
default=600.0,
description="How long the breaker stays open before allowing a probe.",
)
class WorkerConfig(BaseModel): class WorkerConfig(BaseModel):
# Reject unknown keys so a renamed/removed setting (e.g. the former # Reject unknown keys so a renamed/removed setting (e.g. the former
# claim_timeout_s) fails loudly instead of being silently ignored. # claim_timeout_s) fails loudly instead of being silently ignored.

View file

@ -57,10 +57,7 @@ class DoclingServeConverter(DocumentConverter):
config: Application configuration containing docling-serve settings. config: Application configuration containing docling-serve settings.
""" """
self.config = config self.config = config
self.client = DoclingServeClient( self.client = DoclingServeClient.from_config(config.providers.docling_serve)
base_urls=config.providers.docling_serve.base_urls,
api_key=config.providers.docling_serve.api_key,
)
@property @property
def supported_extensions(self) -> list[str]: def supported_extensions(self) -> list[str]:

View file

@ -1,4 +1,4 @@
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker from haiku.rag.circuit_breaker import CircuitBreaker
from haiku.rag.ingester.pollers.factory import build_source from haiku.rag.ingester.pollers.factory import build_source
from haiku.rag.ingester.pollers.fs import FSPoller from haiku.rag.ingester.pollers.fs import FSPoller
from haiku.rag.ingester.pollers.manager import PollerManager from haiku.rag.ingester.pollers.manager import PollerManager

View file

@ -3,9 +3,9 @@ import logging
import random import random
from datetime import UTC, datetime from datetime import UTC, datetime
from haiku.rag.circuit_breaker import CircuitBreaker
from haiku.rag.config import SourceConfig from haiku.rag.config import SourceConfig
from haiku.rag.ingester.batch import BatchChange, BatchSourceSummary from haiku.rag.ingester.batch import BatchChange, BatchSourceSummary
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
from haiku.rag.ingester.queue.models import JobOp, SyncRow from haiku.rag.ingester.queue.models import JobOp, SyncRow
from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
from haiku.rag.ingester.sources.base import ( from haiku.rag.ingester.sources.base import (

View file

@ -11,8 +11,8 @@ from haiku.rag.ingester.sources.filter import FileFilter
from haiku.rag.telemetry import logfire from haiku.rag.telemetry import logfire
if TYPE_CHECKING: if TYPE_CHECKING:
from haiku.rag.circuit_breaker import CircuitBreaker
from haiku.rag.config import FSSourceConfig from haiku.rag.config import FSSourceConfig
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
from haiku.rag.ingester.sources.fs import FSSource from haiku.rag.ingester.sources.fs import FSSource
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)

View file

@ -4,10 +4,10 @@ from collections.abc import Sequence
from datetime import UTC, datetime from datetime import UTC, datetime
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from haiku.rag.circuit_breaker import CircuitBreaker
from haiku.rag.config import FSSourceConfig, SourceConfig from haiku.rag.config import FSSourceConfig, SourceConfig
from haiku.rag.ingester.batch import BatchManifest from haiku.rag.ingester.batch import BatchManifest
from haiku.rag.ingester.pollers.base import BasePoller from haiku.rag.ingester.pollers.base import BasePoller
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
from haiku.rag.ingester.pollers.factory import build_source from haiku.rag.ingester.pollers.factory import build_source
from haiku.rag.ingester.pollers.fs import FSPoller from haiku.rag.ingester.pollers.fs import FSPoller
from haiku.rag.ingester.pollers.periodic import PeriodicPoller from haiku.rag.ingester.pollers.periodic import PeriodicPoller

View file

@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
from haiku.rag.ingester.pollers.base import BasePoller from haiku.rag.ingester.pollers.base import BasePoller
if TYPE_CHECKING: if TYPE_CHECKING:
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker from haiku.rag.circuit_breaker import CircuitBreaker
class PeriodicPoller(BasePoller): class PeriodicPoller(BasePoller):

View file

@ -5,9 +5,9 @@ import time
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from uuid import uuid4 from uuid import uuid4
from haiku.rag.circuit_breaker import CircuitBreaker
from haiku.rag.config import CircuitBreakerConfig from haiku.rag.config import CircuitBreakerConfig
from haiku.rag.ingester.exceptions import PermanentError, TransientError from haiku.rag.ingester.exceptions import PermanentError, TransientError
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
from haiku.rag.ingester.queue.models import Job, JobOp from haiku.rag.ingester.queue.models import Job, JobOp
from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
from haiku.rag.ingester.workers.pipeline import run_job from haiku.rag.ingester.workers.pipeline import run_job

View file

@ -2,27 +2,61 @@
import asyncio import asyncio
import itertools import itertools
from typing import Any import logging
import time
from collections.abc import Awaitable, Callable
from typing import Any, TypeVar
import httpx import httpx
# Process-global registry of round-robin rotators over docling-serve from haiku.rag.circuit_breaker import CircuitBreaker
# instances, keyed by the sorted tuple of base URLs. Clients are constructed from haiku.rag.config import CircuitBreakerConfig, DoclingServeConfig
# per-job by `get_converter` / `get_chunker`, so the rotation cursor has to
# live OUTSIDE the instance to actually advance across jobs. Two clients logger = logging.getLogger(__name__)
# pointing at the same instance set share one rotator; clients with
# different sets get independent rotators. _T = TypeVar("_T")
# Statuses that mean "try again elsewhere" rather than "this request is bad".
_RETRYABLE_STATUS = frozenset({408, 429})
def _is_retryable(exc: BaseException) -> bool:
"""Whether a failure reflects an unhealthy docling-serve instance.
True for transport errors (a crashed / unresponsive instance) and
server-side 5xx / 408 / 429. Such failures are retried on another instance
and counted against the instance's circuit breaker. Other 4xx and the
task-level ValueError from `_submit_and_wait` (a bad document, not the
instance's fault) won't succeed on a retry, so they propagate untouched.
"""
if isinstance(exc, httpx.HTTPStatusError):
status = exc.response.status_code
return status >= 500 or status in _RETRYABLE_STATUS
return isinstance(exc, httpx.TransportError)
# Process-global registries over docling-serve instances. Clients are built
# per-job by `get_converter` / `get_chunker`, so this shared state lives OUTSIDE
# the instance to persist across jobs: the round-robin cursor (keyed by the
# sorted URL set) advances across jobs, and the per-URL breakers keep a crashed
# instance skipped across jobs until its cooldown elapses. The first client to
# touch a URL fixes that breaker's config; in practice all clients share one.
_instance_rotators: dict[tuple[str, ...], "itertools.cycle[str]"] = {} _instance_rotators: dict[tuple[str, ...], "itertools.cycle[str]"] = {}
_instance_breakers: dict[str, CircuitBreaker] = {}
class DoclingServeClient: class DoclingServeClient:
"""Client for docling-serve async workflow. """Client for docling-serve async workflow.
Handles the submit poll fetch pattern used by both conversion and Handles the submit poll fetch pattern used by both conversion and
chunking. Accepts a list of base URLs and round-robins jobs across chunking. Accepts a list of base URLs and round-robins jobs across them via
them via a process-wide rotator keyed by the URL set. Each job's a process-wide rotator keyed by the URL set. Each job's submit/poll/result
submit/poll/result trio stays on the same instance task IDs are trio stays on the same instance task IDs are instance-local, so picking a
instance-local, so picking a different URL mid-job would 404. different URL mid-job would 404.
When an instance crashes or returns 5xx, the request fails over to another
instance (up to `max_attempts`) and the failure trips that instance's
circuit breaker, so subsequent `_pick_url` calls skip it while it recovers.
""" """
def __init__( def __init__(
@ -31,6 +65,11 @@ class DoclingServeClient:
api_key: str | None = None, api_key: str | None = None,
timeout: float = 300, timeout: float = 300,
transport: httpx.AsyncBaseTransport | None = None, transport: httpx.AsyncBaseTransport | None = None,
circuit_breaker: CircuitBreakerConfig | None = None,
max_attempts: int = 3,
retry_base_delay: float = 0.5,
retry_max_delay: float = 8.0,
now_fn: Callable[[], float] = time.monotonic,
): ):
urls = [base_urls] if isinstance(base_urls, str) else list(base_urls) urls = [base_urls] if isinstance(base_urls, str) else list(base_urls)
if not urls: if not urls:
@ -40,6 +79,11 @@ class DoclingServeClient:
self.timeout = timeout self.timeout = timeout
# transport is for testing — production callers leave it None. # transport is for testing — production callers leave it None.
self._transport = transport self._transport = transport
self._breaker_config = circuit_breaker or CircuitBreakerConfig()
self._max_attempts = max(1, max_attempts)
self._retry_base_delay = retry_base_delay
self._retry_max_delay = retry_max_delay
self._now = now_fn
# setdefault is atomic under the GIL — concurrent constructors with # setdefault is atomic under the GIL — concurrent constructors with
# the same instance set will end up sharing one rotator. # the same instance set will end up sharing one rotator.
key = tuple(sorted(self.base_urls)) key = tuple(sorted(self.base_urls))
@ -47,6 +91,15 @@ class DoclingServeClient:
key, itertools.cycle(self.base_urls) key, itertools.cycle(self.base_urls)
) )
@classmethod
def from_config(cls, config: DoclingServeConfig) -> "DoclingServeClient":
return cls(
base_urls=config.base_urls,
api_key=config.api_key,
circuit_breaker=config.circuit_breaker,
max_attempts=config.max_attempts,
)
def _httpx_client(self) -> httpx.AsyncClient: def _httpx_client(self) -> httpx.AsyncClient:
return httpx.AsyncClient(timeout=self.timeout, transport=self._transport) return httpx.AsyncClient(timeout=self.timeout, transport=self._transport)
@ -56,9 +109,80 @@ class DoclingServeClient:
callers should let `_pick_url` round-robin per request.""" callers should let `_pick_url` round-robin per request."""
return self.base_urls[0] return self.base_urls[0]
def _pick_url(self) -> str: def _breaker_for(self, url: str) -> CircuitBreaker:
breaker = _instance_breakers.get(url)
if breaker is None:
breaker = CircuitBreaker(self._breaker_config, now_fn=self._now)
_instance_breakers[url] = breaker
return breaker
def _pick_url(self, exclude: frozenset[str] = frozenset()) -> str:
"""Next instance in the round-robin, skipping instances that already
failed this request (`exclude`) or whose circuit breaker is open.
When no healthy, not-excluded instance exists, probe one anyway (the
loop consumed a full rotation, so one more step rotates across calls
rather than pinning one node) a single-instance or fully-down fleet
still gets an attempt instead of failing with nothing to pick.
"""
for _ in range(len(self.base_urls)):
url = next(self._instance_rotator)
if url in exclude:
continue
if not self._breaker_for(url).is_open:
return url
return next(self._instance_rotator) return next(self._instance_rotator)
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)
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). A retryable failure counts against that instance's breaker
and fails over to another instance; a successful attempt closes the
breaker. Non-retryable errors 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))
breaker = self._breaker_for(base_url)
try:
async with self._httpx_client() as client:
result = await attempt(client, base_url)
except Exception as exc:
if not _is_retryable(exc):
raise
breaker.record_failure()
tried.add(base_url)
last_exc = exc
if attempt_no + 1 >= self._max_attempts:
raise
logger.warning(
"docling-serve request for %s failed on %s "
"(attempt %d/%d), retrying on another instance: %s",
name,
base_url,
attempt_no + 1,
self._max_attempts,
exc,
)
await asyncio.sleep(self._retry_delay(attempt_no))
else:
breaker.record_success()
return result
raise last_exc or RuntimeError( # pragma: no cover
"retry loop exited without a result"
)
def _get_headers(self) -> dict[str, str]: def _get_headers(self) -> dict[str, str]:
"""Get headers for API requests.""" """Get headers for API requests."""
headers: dict[str, str] = {} headers: dict[str, str] = {}
@ -118,15 +242,16 @@ class DoclingServeClient:
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Submit a task and poll until completion; fetch result as JSON. """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 `_submit_and_wait`
is raised by `_submit_and_wait` when docling-serve reports a task (task failure / missing task_id) propagate so the ingester's pipeline
failure or returns no task_id. classifier can route them (4xx PermanentError, ValueError
TransientError for a whole-document retry).
""" """
headers = self._get_headers() 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( task_id = await self._submit_and_wait(
client, base_url, endpoint, files, data, headers, name client, base_url, endpoint, files, data, headers, name
) )
@ -135,6 +260,8 @@ class DoclingServeClient:
result_response.raise_for_status() result_response.raise_for_status()
return result_response.json() return result_response.json()
return await self._run_with_retry(_attempt, name)
async def submit_and_poll_zip( async def submit_and_poll_zip(
self, self,
endpoint: str, endpoint: str,
@ -147,11 +274,12 @@ class DoclingServeClient:
Used when the caller requested ``target_type=zip`` (e.g. to retrieve Used when the caller requested ``target_type=zip`` (e.g. to retrieve
picture image bytes that docling-serve only emits as referenced files picture image bytes that docling-serve only emits as referenced files
bundled into a zip archive). The submit/poll flow is identical to 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() 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( task_id = await self._submit_and_wait(
client, base_url, endpoint, files, data, headers, name client, base_url, endpoint, files, data, headers, name
) )
@ -159,3 +287,5 @@ class DoclingServeClient:
result_response = await client.get(result_url, headers=headers) result_response = await client.get(result_url, headers=headers)
result_response.raise_for_status() result_response.raise_for_status()
return result_response.content return result_response.content
return await self._run_with_retry(_attempt, name)

View file

@ -1,5 +1,5 @@
from haiku.rag.circuit_breaker import CircuitBreaker
from haiku.rag.config import CircuitBreakerConfig from haiku.rag.config import CircuitBreakerConfig
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
class _Clock: class _Clock:

View file

@ -4,6 +4,7 @@ from pathlib import Path
import pytest import pytest
from haiku.rag.circuit_breaker import CircuitBreaker
from haiku.rag.config import ( from haiku.rag.config import (
CircuitBreakerConfig, CircuitBreakerConfig,
FSSourceConfig, FSSourceConfig,
@ -11,7 +12,6 @@ from haiku.rag.config import (
S3SourceConfig, S3SourceConfig,
WebDAVSourceConfig, WebDAVSourceConfig,
) )
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
from haiku.rag.ingester.pollers.manager import PollerManager from haiku.rag.ingester.pollers.manager import PollerManager
from haiku.rag.ingester.pollers.periodic import PeriodicPoller from haiku.rag.ingester.pollers.periodic import PeriodicPoller
from haiku.rag.ingester.queue.models import JobOp, JobStatus from haiku.rag.ingester.queue.models import JobOp, JobStatus

View file

@ -9,6 +9,7 @@ resume mid-rotation, breaking specific-order assertions.
import httpx import httpx
import pytest import pytest
from haiku.rag.config import CircuitBreakerConfig
from haiku.rag.providers.docling_serve import DoclingServeClient from haiku.rag.providers.docling_serve import DoclingServeClient
@ -134,3 +135,319 @@ async def test_zip_endpoint_uses_round_robin_too():
) )
per_call = [seen[i : i + 3] for i in range(0, 6, 3)] per_call = [seen[i : i + 3] for i in range(0, 6, 3)]
assert [triple[0] for triple in per_call] == ["a", "b"] assert [triple[0] for triple in per_call] == ["a", "b"]
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 (a crashed instance); other hosts serve a normal
submit/poll/result trio. Mutate ``down_hosts`` mid-test to flip an
instance's health. Records every host attempted."""
seen: list[str] = []
routes = _success_routes(task_id, result)
def handler(request: httpx.Request) -> httpx.Response:
seen.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
async def _poll(client: DoclingServeClient):
return await client.submit_and_poll(
endpoint="/v1/convert/file/async",
files={"file": ("x.md", b"x", "text/markdown")},
data={},
)
@pytest.mark.asyncio
async def test_retry_fails_over_to_healthy_instance():
"""A crashed instance (connection error) is retried on another instance and
the call succeeds without surfacing the failure."""
transport, seen = _failover_transport({"down-a"}, "t", {"ok": True})
client = DoclingServeClient(
base_urls=["http://down-a:5001", "http://up-a:5001"],
transport=transport,
retry_base_delay=0.0,
)
result = await _poll(client)
assert result == {"ok": True}
assert seen[0] == "down-a"
# The successful trio 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 _poll(client)
assert result == {"ok": True}
assert seen[0] == "sad-b"
assert "ok-b" in seen
@pytest.mark.asyncio
async def test_retry_429_fails_over():
"""A 429 (transient overload) is retried elsewhere — the status-set
membership branch of retryability, distinct from 5xx."""
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 == "busy-g":
return httpx.Response(429)
return ok.get(request.url.path, httpx.Response(404))
client = DoclingServeClient(
base_urls=["http://busy-g:5001", "http://free-g:5001"],
transport=httpx.MockTransport(handler),
retry_base_delay=0.0,
)
result = await _poll(client)
assert result == {"ok": True}
assert seen[0] == "busy-g"
assert "free-g" 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-c", "down-d"}, "t", {})
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 _poll(client)
# 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 — not 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 _poll(client)
assert seen == ["e"]
@pytest.mark.asyncio
async def test_task_failure_is_not_retried():
"""A docling-serve task 'failure' status raises ValueError and is NOT
retried a document problem, not an instance one."""
seen: list[str] = []
def handler(request: httpx.Request) -> httpx.Response:
seen.append(request.url.host)
if request.url.path == "/v1/convert/file/async":
return httpx.Response(200, json={"task_id": "t"})
if request.url.path == "/v1/status/poll/t":
return httpx.Response(200, json={"task_status": "failure", "detail": "x"})
return httpx.Response(404)
client = DoclingServeClient(
base_urls=["http://h:5001", "http://i:5001"],
transport=httpx.MockTransport(handler),
retry_base_delay=0.0,
)
with pytest.raises(ValueError, match="task failed"):
await _poll(client)
# Only the first instance was attempted (submit + poll) — no failover.
assert set(seen) == {"h"}
def test_pick_url_skips_excluded_instances():
"""On retry, _pick_url advances past every excluded instance."""
urls = ["http://p1:5001", "http://p2:5001", "http://p3:5001"]
client = DoclingServeClient(base_urls=urls)
assert client._pick_url(exclude=frozenset({urls[0]})) == urls[1]
assert client._pick_url(exclude=frozenset({urls[1], urls[2]})) == urls[0]
@pytest.mark.asyncio
async def test_retryable_failure_fails_over_and_trips_breaker():
"""A retryable failure does both jobs at once: the request fails over to a
healthy instance AND the failure counts against the crashed instance's
breaker."""
trip_a, trip_b = "http://trip-a:5001", "http://trip-b:5001"
transport, seen = _failover_transport({"trip-a"}, "t", {"ok": True})
client = DoclingServeClient(
base_urls=[trip_a, trip_b],
transport=transport,
circuit_breaker=CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
max_attempts=2,
retry_base_delay=0.0,
)
result = await _poll(client)
assert result == {"ok": True}
assert seen[0] == "trip-a"
assert client._breaker_for(trip_a).is_open
assert not client._breaker_for(trip_b).is_open
@pytest.mark.asyncio
async def test_open_breaker_skips_crashed_instance():
"""Once an instance's breaker has opened, later requests route straight to a
healthy instance without even attempting the dead one."""
crash, live = "http://crash-x:5001", "http://live-x:5001"
transport, seen = _failover_transport({"crash-x"}, "t", {"ok": True})
client = DoclingServeClient(
base_urls=[crash, live],
transport=transport,
circuit_breaker=CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
max_attempts=2,
retry_base_delay=0.0,
)
await _poll(client) # crash-x fails, opens its breaker, fails over to live-x
assert client._breaker_for(crash).is_open
seen.clear()
for _ in range(3):
await _poll(client)
assert "crash-x" not in seen
assert set(seen) == {"live-x"}
@pytest.mark.asyncio
async def test_breaker_recovers_after_cooldown():
"""An open breaker auto-probes after its cooldown; once the instance is
healthy again a successful request closes the breaker."""
clock = [1000.0]
down = {"flip-y"}
flip, spare = "http://flip-y:5001", "http://spare-y:5001"
transport, seen = _failover_transport(down, "t", {"ok": True})
client = DoclingServeClient(
base_urls=[flip, spare],
transport=transport,
circuit_breaker=CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
max_attempts=2,
retry_base_delay=0.0,
now_fn=lambda: clock[0],
)
await _poll(client) # flip-y fails, opens; fails over to spare-y
assert client._breaker_for(flip).is_open
# Instance recovers, but within the cooldown it's still treated as open.
down.clear()
assert client._breaker_for(flip).is_open
# After the cooldown the breaker allows a probe; traffic returns and a
# success closes it.
clock[0] += 31.0
seen.clear()
await _poll(client)
assert "flip-y" in seen
assert not client._breaker_for(flip).is_open
@pytest.mark.asyncio
async def test_4xx_does_not_trip_breaker():
"""A 4xx is the caller's fault — it must not count against instance health,
even at a 1-failure threshold."""
bad = "http://bad-z:5001"
def handler(_request: httpx.Request) -> httpx.Response:
return httpx.Response(400, json={"detail": "bad request"})
client = DoclingServeClient(
base_urls=[bad],
transport=httpx.MockTransport(handler),
circuit_breaker=CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
retry_base_delay=0.0,
)
for _ in range(3):
with pytest.raises(httpx.HTTPStatusError):
await _poll(client)
assert not client._breaker_for(bad).is_open
def test_all_open_breakers_rotate_instead_of_pinning_one():
"""When every breaker is open, successive picks rotate across the fleet
rather than pinning the first instance an all-429 overload shouldn't pile
every retry on one node."""
urls = [
"http://allopen-a:5001",
"http://allopen-b:5001",
"http://allopen-c:5001",
]
client = DoclingServeClient(
base_urls=urls,
circuit_breaker=CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
)
for u in urls:
client._breaker_for(u).record_failure()
assert client._breaker_for(u).is_open
picks = [client._pick_url() for _ in range(6)]
assert set(picks) == set(urls)
def test_from_config_wires_retry_and_breaker():
"""Retry/breaker knobs set in DoclingServeConfig reach the client via
get_converter / get_chunker (DoclingServeClient.from_config)."""
from haiku.rag.chunkers.docling_serve import DoclingServeChunker
from haiku.rag.config import AppConfig
from haiku.rag.converters.docling_serve import DoclingServeConverter
config = AppConfig()
ds = config.providers.docling_serve
ds.base_url = "http://cfg-n:5001"
ds.max_attempts = 7
ds.circuit_breaker = CircuitBreakerConfig(failure_threshold=9, cooldown_s=90.0)
for component in (DoclingServeConverter(config), DoclingServeChunker(config)):
client = component.client
assert client._max_attempts == 7
assert client._breaker_config.failure_threshold == 9
assert client._breaker_config.cooldown_s == 90.0