From eb8710d83af1a30a719f5ddefa4856fb6ced58ba Mon Sep 17 00:00:00 2001 From: bryan davis Date: Fri, 26 Jun 2026 11:24:13 -0500 Subject: [PATCH] requested updates --- .../haiku/rag/chunkers/docling_serve.py | 10 +- haiku_rag_slim/haiku/rag/config/models.py | 26 +++ .../haiku/rag/converters/docling_serve.py | 10 +- .../haiku/rag/providers/docling_serve.py | 154 ++++++++++++++---- tests/test_docling_serve_client.py | 139 ++++++++++++++++ 5 files changed, 306 insertions(+), 33 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py index 7f46abf3..2d7f594c 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py @@ -56,9 +56,15 @@ class DoclingServeChunker(DocumentChunker): def __init__(self, config: AppConfig = Config): self.config = config + ds = config.providers.docling_serve self.client = DoclingServeClient( - base_urls=config.providers.docling_serve.base_urls, - api_key=config.providers.docling_serve.api_key, + base_urls=ds.base_urls, + api_key=ds.api_key, + max_attempts=ds.max_attempts, + retry_base_delay=ds.retry_base_delay, + retry_max_delay=ds.retry_max_delay, + breaker_failure_threshold=ds.breaker_failure_threshold, + breaker_cooldown_s=ds.breaker_cooldown_s, ) self.chunker_type = config.processing.chunker_type diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 08622db7..5fd218e1 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -227,6 +227,32 @@ class DoclingServeConfig(BaseModel): base_url: str | list[str] = "http://localhost:5001" api_key: str = "" + max_attempts: int = Field( + default=3, + ge=1, + description="Total attempts per request before giving up; each retry " + "fails over to another instance (instances crash on memory leaks).", + ) + retry_base_delay: float = Field( + default=0.5, + ge=0, + description="Base seconds for the exponential backoff between retries.", + ) + retry_max_delay: float = Field( + default=8.0, ge=0, description="Cap on the retry backoff, in seconds." + ) + breaker_failure_threshold: int = Field( + default=3, + ge=1, + description="Consecutive failures before an instance's circuit breaker " + "opens and routing skips it.", + ) + breaker_cooldown_s: float = Field( + default=30.0, + ge=0, + description="How long an open breaker skips an instance before allowing " + "a probe.", + ) @property def base_urls(self) -> list[str]: diff --git a/haiku_rag_slim/haiku/rag/converters/docling_serve.py b/haiku_rag_slim/haiku/rag/converters/docling_serve.py index 08f0f42b..40951327 100644 --- a/haiku_rag_slim/haiku/rag/converters/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/converters/docling_serve.py @@ -59,9 +59,15 @@ class DoclingServeConverter(DocumentConverter): config: Application configuration containing docling-serve settings. """ self.config = config + ds = config.providers.docling_serve self.client = DoclingServeClient( - base_urls=config.providers.docling_serve.base_urls, - api_key=config.providers.docling_serve.api_key, + base_urls=ds.base_urls, + api_key=ds.api_key, + max_attempts=ds.max_attempts, + retry_base_delay=ds.retry_base_delay, + retry_max_delay=ds.retry_max_delay, + breaker_failure_threshold=ds.breaker_failure_threshold, + breaker_cooldown_s=ds.breaker_cooldown_s, ) @property diff --git a/haiku_rag_slim/haiku/rag/providers/docling_serve.py b/haiku_rag_slim/haiku/rag/providers/docling_serve.py index 4b587eff..d4b7fd2c 100644 --- a/haiku_rag_slim/haiku/rag/providers/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/providers/docling_serve.py @@ -3,6 +3,7 @@ import asyncio import itertools import logging +import time from collections.abc import Awaitable, Callable from typing import Any, TypeVar @@ -12,31 +13,75 @@ logger = logging.getLogger(__name__) _T = TypeVar("_T") -# Statuses worth retrying on another instance: 408/429 are transient overload, +# Statuses signalling an unhealthy instance: 408/429 are transient overload, # 5xx covers a crashed or restarting docling-serve worker (the OOM-leak failure # mode). Other 4xx are the caller's fault and won't succeed elsewhere. _RETRYABLE_STATUS = frozenset({408, 429}) def _is_retryable(exc: BaseException) -> bool: - """Whether a failed docling-serve request should be retried on another - instance. True for transport-level failures (connection reset / timeout — - an instance that crashed or went unresponsive) and server-side 5xx/overload; - False for other 4xx and task-level ``ValueError``\\ s, which won't succeed on - a retry (a deterministically bad document is handled upstream, not here).""" + """Whether a failure reflects an unhealthy docling-serve instance. + + The single predicate behind both behaviours: such a failure is retried on + another instance *and* counted against the instance's circuit breaker. True + for transport-level failures (connection reset / timeout — an instance that + crashed or went unresponsive) and server-side 5xx/overload; False for other + 4xx and task-level ``ValueError``\\ s, which won't succeed on a retry and + aren't the instance's fault (a deterministically bad document is handled + upstream, not here).""" if isinstance(exc, httpx.HTTPStatusError): status = exc.response.status_code return status in _RETRYABLE_STATUS or status >= 500 return isinstance(exc, httpx.TransportError) -# Process-global registry of round-robin rotators over docling-serve -# instances, keyed by the sorted tuple of base URLs. Clients are constructed -# per-job by `get_converter` / `get_chunker`, so the rotation cursor has to -# 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. +class _InstanceBreaker: + """Per-instance circuit breaker (closed / open with cooldown). + + Self-contained so the provider layer doesn't depend on the ingester + package. ``is_open`` auto-probes after the cooldown elapses: a single + request is allowed through, and its success closes the breaker while another + failure re-opens it. + """ + + def __init__( + self, + failure_threshold: int, + cooldown_s: float, + now_fn: Callable[[], float], + ): + self._threshold = failure_threshold + self._cooldown_s = cooldown_s + self._now = now_fn + self._consecutive_failures = 0 + self._opened_at: float | None = None + + @property + def is_open(self) -> bool: + if self._opened_at is None: + return False + # Cooldown elapsed → let the next request probe the instance. + return self._now() - self._opened_at < self._cooldown_s + + def record_success(self) -> None: + self._consecutive_failures = 0 + self._opened_at = None + + def record_failure(self) -> None: + self._consecutive_failures += 1 + if self._consecutive_failures >= self._threshold: + self._opened_at = self._now() + + +# Process-global registries over docling-serve instances. Clients are +# constructed per-job by `get_converter` / `get_chunker`, so this shared state +# has to live OUTSIDE the instance to persist across jobs: the round-robin +# cursor (keyed by the sorted tuple of base URLs) advances across jobs, and the +# per-URL breakers keep a crashed instance skipped across subsequent jobs until +# its cooldown elapses. The first client to touch a URL fixes that breaker's +# threshold/cooldown/clock; in practice all clients share one config. _instance_rotators: dict[tuple[str, ...], "itertools.cycle[str]"] = {} +_instance_breakers: dict[str, "_InstanceBreaker"] = {} class DoclingServeClient: @@ -58,6 +103,9 @@ class DoclingServeClient: max_attempts: int = 3, retry_base_delay: float = 0.5, retry_max_delay: float = 8.0, + breaker_failure_threshold: int = 3, + breaker_cooldown_s: float = 30.0, + now_fn: Callable[[], float] = time.monotonic, ): urls = [base_urls] if isinstance(base_urls, str) else list(base_urls) if not urls: @@ -67,12 +115,17 @@ class DoclingServeClient: self.timeout = timeout # transport is for testing — production callers leave it None. self._transport = transport - # Bounded retry with failover: docling-serve instances crash (memory - # leaks), so a request that hits a dying instance is retried, preferring - # an instance that hasn't already failed this request. + # Bounded retry with failover + per-instance circuit breaking: + # docling-serve instances crash (memory leaks), so a request that hits a + # dying instance is retried on another instance, and repeated failures + # trip that instance's breaker so later requests skip it while it + # recovers. self._max_attempts = max(1, max_attempts) self._retry_base_delay = retry_base_delay self._retry_max_delay = retry_max_delay + self._breaker_failure_threshold = breaker_failure_threshold + self._breaker_cooldown_s = breaker_cooldown_s + 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)) @@ -89,25 +142,59 @@ class DoclingServeClient: callers should let `_pick_url` round-robin per request.""" return self.base_urls[0] + def _breaker_for(self, url: str) -> _InstanceBreaker: + breaker = _instance_breakers.get(url) + if breaker is None: + breaker = _InstanceBreaker( + self._breaker_failure_threshold, + self._breaker_cooldown_s, + self._now, + ) + _instance_breakers[url] = breaker + return breaker + def _pick_url(self, exclude: frozenset[str] = frozenset()) -> str: - """Next instance in the round-robin, preferring one not in ``exclude`` - (instances that already failed this request). Falls back to a - possibly-excluded instance when every instance is excluded — a - single-instance fleet, or one where all instances failed, still gets - retried after backoff in case the instance has since restarted.""" - url = next(self._instance_rotator) - if url not in exclude: - return url - for _ in range(len(self.base_urls) - 1): + """Next instance in the round-robin, skipping instances that already + failed this request (``exclude``) or whose circuit breaker is open + (recently crashed / overloaded). + + Prefers a not-excluded, breaker-closed instance. Falls back to a + not-excluded breaker-open one (worth a probe over re-hitting one that + already failed this request), and finally to any instance — so a + single-instance fleet, or one where everything is excluded/open, still + gets an attempt rather than nothing to pick.""" + not_excluded: str | None = None + for _ in range(len(self.base_urls)): url = next(self._instance_rotator) - if url not in exclude: + if url in exclude: + continue + if not_excluded is None: + not_excluded = url + if not self._breaker_for(url).is_open: return url - return url + if not_excluded is not None: + return not_excluded + 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) + def _record_failure(self, base_url: str) -> None: + """Count an instance-health failure against ``base_url``'s breaker, + logging when the breaker transitions to open.""" + breaker = self._breaker_for(base_url) + was_open = breaker.is_open + breaker.record_failure() + if not was_open and breaker.is_open: + logger.warning( + "docling-serve instance %s breaker opened after %d consecutive " + "failure(s); skipping it for %.0fs", + base_url, + self._breaker_failure_threshold, + self._breaker_cooldown_s, + ) + def _get_headers(self) -> dict[str, str]: """Get headers for API requests.""" headers: dict[str, str] = {} @@ -126,8 +213,11 @@ class DoclingServeClient: instance (task IDs are instance-local, so a trio can't be split across instances). On a retryable failure — a transport error (crashed / unresponsive instance) or 5xx/overload — the next attempt prefers an - instance that hasn't already failed this request. Non-retryable errors - (4xx, task-level ``ValueError``) propagate immediately. + instance that hasn't already failed this request, and the failure is + counted against that instance's circuit breaker so later requests skip + it while it recovers; a successful attempt closes the breaker. + Non-retryable errors (4xx, task-level ``ValueError``) propagate + immediately and leave the breaker untouched. """ tried: set[str] = set() last_exc: BaseException | None = None @@ -135,10 +225,13 @@ class DoclingServeClient: base_url = self._pick_url(exclude=frozenset(tried)) try: async with self._httpx_client() as client: - return await attempt(client, base_url) + result = await attempt(client, base_url) except Exception as exc: if not _is_retryable(exc): raise + # Retryable == an unhealthy instance: count it against the + # breaker so later requests route around it. + self._record_failure(base_url) last_exc = exc tried.add(base_url) remaining = self._max_attempts - attempt_no - 1 @@ -153,6 +246,9 @@ class DoclingServeClient: remaining, ) await asyncio.sleep(self._retry_delay(attempt_no)) + else: + self._breaker_for(base_url).record_success() + return result # Unreachable: max_attempts >= 1, and the final iteration either returns # on success or re-raises on failure (remaining <= 0). raise last_exc or RuntimeError( # pragma: no cover diff --git a/tests/test_docling_serve_client.py b/tests/test_docling_serve_client.py index a93226cf..801fdd40 100644 --- a/tests/test_docling_serve_client.py +++ b/tests/test_docling_serve_client.py @@ -321,6 +321,145 @@ def test_pick_url_all_excluded_falls_back(): assert picked == "http://solo:5001" +async def _convert(client: DoclingServeClient) -> dict: + 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_breaker_opens_after_repeated_failures_then_skips_instance(): + """Repeated failures on an instance (each failing over to a healthy one) + trip its breaker; once open, _pick_url skips it entirely.""" + down = {"crash-h"} + transport, seen = _failover_transport(down, "t", {"ok": True}) + crash = "http://crash-h:5001" + client = DoclingServeClient( + base_urls=[crash, "http://live-h:5001"], + transport=transport, + breaker_failure_threshold=2, + retry_base_delay=0.0, + ) + + # Each call fails on crash-h (recorded) then fails over to live-h. + assert await _convert(client) == {"ok": True} + assert await _convert(client) == {"ok": True} # 2nd failure opens crash-h + assert client._breaker_for(crash).is_open + + seen.clear() + assert await _convert(client) == {"ok": True} + assert "crash-h" not in seen # breaker open → never attempted + assert set(seen) == {"live-h"} + + +@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 it.""" + clock = [1000.0] + down = {"flip-j"} + transport, seen = _failover_transport(down, "t", {"ok": True}) + flip = "http://flip-j:5001" + client = DoclingServeClient( + base_urls=[flip, "http://spare-j:5001"], + transport=transport, + breaker_failure_threshold=2, + breaker_cooldown_s=30.0, + retry_base_delay=0.0, + now_fn=lambda: clock[0], + ) + + await _convert(client) + await _convert(client) # opens flip-j's breaker + assert client._breaker_for(flip).is_open + + down.clear() # instance recovers + assert client._breaker_for(flip).is_open # still skipped within cooldown + clock[0] += 31.0 + assert not client._breaker_for(flip).is_open # cooldown elapsed → probe ok + + seen.clear() + await _convert(client) + await _convert(client) + assert "flip-j" in seen # traffic returned + assert not client._breaker_for(flip).is_open # success closed it + + +@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.""" + seen: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request.url.host) + return httpx.Response(400, json={"detail": "bad request"}) + + bad = "http://bad-k:5001" + client = DoclingServeClient( + base_urls=[bad], + transport=httpx.MockTransport(handler), + breaker_failure_threshold=1, + retry_base_delay=0.0, + ) + + for _ in range(3): + with pytest.raises(httpx.HTTPStatusError): + await _convert(client) + + assert not client._breaker_for(bad).is_open + + +@pytest.mark.asyncio +async def test_pick_url_falls_back_when_all_breakers_open(): + """When every instance's breaker is open (here a single-instance fleet that + just failed), _pick_url returns one anyway so the request can still probe + it rather than having nothing to pick.""" + lone = "http://lone-m:5001" + transport, _ = _failover_transport({"lone-m"}, "t", {"ok": True}) + client = DoclingServeClient( + base_urls=[lone], + transport=transport, + breaker_failure_threshold=1, + max_attempts=1, + retry_base_delay=0.0, + ) + + with pytest.raises(httpx.ConnectError): + await _convert(client) + assert client._breaker_for(lone).is_open + + # All (one) instances open → fall back to it. + assert client._pick_url() == lone + + +def test_config_knobs_reach_the_client(): + """Retry/breaker knobs set in DoclingServeConfig must reach the client via + get_converter / get_chunker (previously hardcoded constructor defaults).""" + 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.retry_base_delay = 1.25 + ds.retry_max_delay = 20.0 + ds.breaker_failure_threshold = 9 + ds.breaker_cooldown_s = 90.0 + + for component in (DoclingServeConverter(config), DoclingServeChunker(config)): + client = component.client + assert client._max_attempts == 7 + assert client._retry_base_delay == 1.25 + assert client._retry_max_delay == 20.0 + assert client._breaker_failure_threshold == 9 + assert client._breaker_cooldown_s == 90.0 + + @pytest.mark.asyncio async def test_zip_endpoint_uses_round_robin_too(): transport, seen = _scripted_transport(