diff --git a/CHANGELOG.md b/CHANGELOG.md index 850eb52d..3da7b98f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/docs/configuration/processing.md b/docs/configuration/processing.md index 1174caaa..8e50e174 100644 --- a/docs/configuration/processing.md +++ b/docs/configuration/processing.md @@ -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 diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py index 7f46abf3..6398692b 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py @@ -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]]: diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/circuit_breaker.py b/haiku_rag_slim/haiku/rag/circuit_breaker.py similarity index 95% rename from haiku_rag_slim/haiku/rag/ingester/pollers/circuit_breaker.py rename to haiku_rag_slim/haiku/rag/circuit_breaker.py index 25c03e75..664e2826 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/circuit_breaker.py +++ b/haiku_rag_slim/haiku/rag/circuit_breaker.py @@ -5,7 +5,7 @@ from haiku.rag.config import CircuitBreakerConfig 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. - open: probes are blocked until cooldown elapses, then a single probe diff --git a/haiku_rag_slim/haiku/rag/config/__init__.py b/haiku_rag_slim/haiku/rag/config/__init__.py index 55645163..c2b79f0d 100644 --- a/haiku_rag_slim/haiku/rag/config/__init__.py +++ b/haiku_rag_slim/haiku/rag/config/__init__.py @@ -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", diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 0f2f682d..b9f2734c 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -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. diff --git a/haiku_rag_slim/haiku/rag/converters/docling_serve.py b/haiku_rag_slim/haiku/rag/converters/docling_serve.py index c728b4ae..c9a9ba2a 100644 --- a/haiku_rag_slim/haiku/rag/converters/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/converters/docling_serve.py @@ -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]: diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/__init__.py b/haiku_rag_slim/haiku/rag/ingester/pollers/__init__.py index 2c581799..19bf24b9 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/__init__.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/__init__.py @@ -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.fs import FSPoller from haiku.rag.ingester.pollers.manager import PollerManager diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/base.py b/haiku_rag_slim/haiku/rag/ingester/pollers/base.py index c3d6a859..e6700b8f 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/base.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/base.py @@ -3,9 +3,9 @@ import logging import random from datetime import UTC, datetime +from haiku.rag.circuit_breaker import CircuitBreaker from haiku.rag.config import SourceConfig 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.repository import JobRepo, SyncStateRepo from haiku.rag.ingester.sources.base import ( diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/fs.py b/haiku_rag_slim/haiku/rag/ingester/pollers/fs.py index e27e8b80..f30f5111 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/fs.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/fs.py @@ -11,8 +11,8 @@ from haiku.rag.ingester.sources.filter import FileFilter from haiku.rag.telemetry import logfire if TYPE_CHECKING: + from haiku.rag.circuit_breaker import CircuitBreaker from haiku.rag.config import FSSourceConfig - from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker from haiku.rag.ingester.sources.fs import FSSource logger = logging.getLogger(__name__) diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py index a947ddf0..7067e215 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py @@ -4,10 +4,10 @@ from collections.abc import Sequence from datetime import UTC, datetime from typing import TYPE_CHECKING +from haiku.rag.circuit_breaker import CircuitBreaker from haiku.rag.config import FSSourceConfig, SourceConfig from haiku.rag.ingester.batch import BatchManifest 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.fs import FSPoller from haiku.rag.ingester.pollers.periodic import PeriodicPoller diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/periodic.py b/haiku_rag_slim/haiku/rag/ingester/pollers/periodic.py index 45ec5fa8..8417c7da 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/periodic.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/periodic.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING from haiku.rag.ingester.pollers.base import BasePoller if TYPE_CHECKING: - from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker + from haiku.rag.circuit_breaker import CircuitBreaker class PeriodicPoller(BasePoller): diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index f270bcf5..501477bb 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -5,9 +5,9 @@ import time from typing import TYPE_CHECKING from uuid import uuid4 +from haiku.rag.circuit_breaker import CircuitBreaker from haiku.rag.config import CircuitBreakerConfig 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.repository import JobRepo, SyncStateRepo from haiku.rag.ingester.workers.pipeline import run_job diff --git a/haiku_rag_slim/haiku/rag/providers/docling_serve.py b/haiku_rag_slim/haiku/rag/providers/docling_serve.py index 2d32e98c..7a9164ec 100644 --- a/haiku_rag_slim/haiku/rag/providers/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/providers/docling_serve.py @@ -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) diff --git a/tests/ingester/test_circuit_breaker.py b/tests/ingester/test_circuit_breaker.py index 0267c68e..ad098d02 100644 --- a/tests/ingester/test_circuit_breaker.py +++ b/tests/ingester/test_circuit_breaker.py @@ -1,5 +1,5 @@ +from haiku.rag.circuit_breaker import CircuitBreaker from haiku.rag.config import CircuitBreakerConfig -from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker class _Clock: diff --git a/tests/ingester/test_pollers.py b/tests/ingester/test_pollers.py index c9daed5e..d7970d13 100644 --- a/tests/ingester/test_pollers.py +++ b/tests/ingester/test_pollers.py @@ -4,6 +4,7 @@ from pathlib import Path import pytest +from haiku.rag.circuit_breaker import CircuitBreaker from haiku.rag.config import ( CircuitBreakerConfig, FSSourceConfig, @@ -11,7 +12,6 @@ from haiku.rag.config import ( S3SourceConfig, WebDAVSourceConfig, ) -from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker from haiku.rag.ingester.pollers.manager import PollerManager from haiku.rag.ingester.pollers.periodic import PeriodicPoller from haiku.rag.ingester.queue.models import JobOp, JobStatus diff --git a/tests/test_docling_serve_client.py b/tests/test_docling_serve_client.py index b19360d4..8bda6d61 100644 --- a/tests/test_docling_serve_client.py +++ b/tests/test_docling_serve_client.py @@ -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