Fail docling-serve requests over to another instance with per-instance breaking
When a docling-serve instance crashes or returns 5xx, DoclingServeClient now retries the request on a different instance (up to max_attempts) and trips a per-instance circuit breaker so subsequent jobs skip a dead instance until its cooldown elapses. Reuses the shared CircuitBreaker; adds max_attempts and a nested circuit_breaker to providers.docling_serve. Co-Authored-By: bryan davis <bryan@monkeytronics.org> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yiorgis Gozadinos
parent
d781335868
commit
3c7699f156
8 changed files with 519 additions and 51 deletions
|
|
@ -4,6 +4,7 @@
|
|||
### Added
|
||||
|
||||
- `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
|
||||
|
||||
|
|
|
|||
|
|
@ -96,13 +96,21 @@ providers:
|
|||
- http://gpu-1:5001
|
||||
- http://cpu-1: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
|
||||
client processes pick independently, so the distribution evens out over many
|
||||
jobs without coordination. For tight load balancing, failover, or health
|
||||
checks, run a real load balancer (nginx `least_conn`, etc.) in front and
|
||||
configure a single URL here.
|
||||
jobs without coordination. When a listed instance crashes or returns 5xx, the
|
||||
client fails the request over to another instance (up to `max_attempts`) and
|
||||
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
|
||||
is usually the throughput ceiling — a default docling-serve instance
|
||||
|
|
|
|||
|
|
@ -56,10 +56,7 @@ class DoclingServeChunker(DocumentChunker):
|
|||
|
||||
def __init__(self, config: AppConfig = Config):
|
||||
self.config = config
|
||||
self.client = DoclingServeClient(
|
||||
base_urls=config.providers.docling_serve.base_urls,
|
||||
api_key=config.providers.docling_serve.api_key,
|
||||
)
|
||||
self.client = DoclingServeClient.from_config(config.providers.docling_serve)
|
||||
self.chunker_type = config.processing.chunker_type
|
||||
|
||||
def _build_chunking_data(self) -> dict[str, str | list[str]]:
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from haiku.rag.config.models import (
|
|||
AppConfig,
|
||||
CircuitBreakerConfig,
|
||||
ConversionOptions,
|
||||
DoclingServeConfig,
|
||||
EmbeddingModelConfig,
|
||||
EmbeddingsConfig,
|
||||
FSSourceConfig,
|
||||
|
|
@ -40,6 +41,7 @@ __all__ = [
|
|||
"AppConfig",
|
||||
"CircuitBreakerConfig",
|
||||
"ConversionOptions",
|
||||
"DoclingServeConfig",
|
||||
"EmbeddingModelConfig",
|
||||
"EmbeddingsConfig",
|
||||
"FSSourceConfig",
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""docling-serve endpoints. Accepts a single URL or a list — when a list
|
||||
is given, the client round-robins jobs across the URLs. Each job's
|
||||
submit/poll/result trio stays on the same instance (task IDs are
|
||||
instance-local). The round-robin counter is per-process; for true load
|
||||
balancing or failover, put an LB in front."""
|
||||
"""docling-serve endpoints. Accepts a single URL or a list — when a list is
|
||||
given, the client round-robins jobs across the URLs, fails a request over to
|
||||
another instance when one crashes or returns 5xx, and trips a per-instance
|
||||
circuit breaker so repeated failures route around a dead instance. Each
|
||||
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"
|
||||
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
|
||||
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)
|
||||
|
||||
|
||||
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):
|
||||
# Reject unknown keys so a renamed/removed setting (e.g. the former
|
||||
# claim_timeout_s) fails loudly instead of being silently ignored.
|
||||
|
|
|
|||
|
|
@ -57,10 +57,7 @@ class DoclingServeConverter(DocumentConverter):
|
|||
config: Application configuration containing docling-serve settings.
|
||||
"""
|
||||
self.config = config
|
||||
self.client = DoclingServeClient(
|
||||
base_urls=config.providers.docling_serve.base_urls,
|
||||
api_key=config.providers.docling_serve.api_key,
|
||||
)
|
||||
self.client = DoclingServeClient.from_config(config.providers.docling_serve)
|
||||
|
||||
@property
|
||||
def supported_extensions(self) -> list[str]:
|
||||
|
|
|
|||
|
|
@ -2,27 +2,61 @@
|
|||
|
||||
import asyncio
|
||||
import itertools
|
||||
from typing import Any
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any, TypeVar
|
||||
|
||||
import httpx
|
||||
|
||||
# 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
|
||||
# live OUTSIDE the instance to actually advance across jobs. Two clients
|
||||
# pointing at the same instance set share one rotator; clients with
|
||||
# different sets get independent rotators.
|
||||
from haiku.rag.circuit_breaker import CircuitBreaker
|
||||
from haiku.rag.config import CircuitBreakerConfig, DoclingServeConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_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_breakers: dict[str, CircuitBreaker] = {}
|
||||
|
||||
|
||||
class DoclingServeClient:
|
||||
"""Client for docling-serve async workflow.
|
||||
|
||||
Handles the submit → poll → fetch pattern used by both conversion and
|
||||
chunking. Accepts a list of base URLs and round-robins jobs across
|
||||
them via a process-wide rotator keyed by the URL set. Each job's
|
||||
submit/poll/result trio stays on the same instance — task IDs are
|
||||
instance-local, so picking a different URL mid-job would 404.
|
||||
chunking. Accepts a list of base URLs and round-robins jobs across them via
|
||||
a process-wide rotator keyed by the URL set. Each job's submit/poll/result
|
||||
trio stays on the same instance — task IDs are instance-local, so picking a
|
||||
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__(
|
||||
|
|
@ -31,6 +65,11 @@ class DoclingServeClient:
|
|||
api_key: str | None = None,
|
||||
timeout: float = 300,
|
||||
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)
|
||||
if not urls:
|
||||
|
|
@ -40,6 +79,11 @@ class DoclingServeClient:
|
|||
self.timeout = timeout
|
||||
# transport is for testing — production callers leave it None.
|
||||
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
|
||||
# the same instance set will end up sharing one rotator.
|
||||
key = tuple(sorted(self.base_urls))
|
||||
|
|
@ -47,6 +91,15 @@ class DoclingServeClient:
|
|||
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:
|
||||
return httpx.AsyncClient(timeout=self.timeout, transport=self._transport)
|
||||
|
||||
|
|
@ -56,9 +109,80 @@ class DoclingServeClient:
|
|||
callers should let `_pick_url` round-robin per request."""
|
||||
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)
|
||||
|
||||
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]:
|
||||
"""Get headers for API requests."""
|
||||
headers: dict[str, str] = {}
|
||||
|
|
@ -118,15 +242,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 +260,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 +274,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 +287,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)
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ resume mid-rotation, breaking specific-order assertions.
|
|||
import httpx
|
||||
import pytest
|
||||
|
||||
from haiku.rag.config import CircuitBreakerConfig
|
||||
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)]
|
||||
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
|
||||
|
|
|
|||
Loading…
Reference in a new issue