From 98c8d0b4f9f298498c921192f3c9a3c9995e8ad5 Mon Sep 17 00:00:00 2001 From: bryan davis Date: Thu, 25 Jun 2026 13:03:44 -0500 Subject: [PATCH] change breakers to be per-instance --- .../haiku/rag/providers/docling_serve.py | 168 ++++++++++++++++-- tests/test_docling_serve_client.py | 123 +++++++++++++ 2 files changed, 272 insertions(+), 19 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/providers/docling_serve.py b/haiku_rag_slim/haiku/rag/providers/docling_serve.py index 2d32e98c..40154ad0 100644 --- a/haiku_rag_slim/haiku/rag/providers/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/providers/docling_serve.py @@ -2,10 +2,71 @@ import asyncio import itertools +import logging +import time +from collections.abc import Callable from typing import Any import httpx +logger = logging.getLogger(__name__) + +# HTTP statuses that count against an instance's health: 408/429 are transient +# overload, 5xx is a crashed/restarting worker (the OOM-leak failure mode). +_UNHEALTHY_STATUS = frozenset({408, 429}) + + +def _is_instance_failure(exc: BaseException) -> bool: + """Whether a failure reflects an unhealthy docling-serve instance and should + count against its circuit breaker: transport errors (connection reset / + timeout — a crashed or unresponsive instance) and 5xx/overload. Other 4xx + are the caller's fault and the task-level ``ValueError`` from + ``_submit_and_wait`` is a document problem — neither reflects instance + health, so they don't trip the breaker.""" + if isinstance(exc, httpx.HTTPStatusError): + status = exc.response.status_code + return status in _UNHEALTHY_STATUS or status >= 500 + return isinstance(exc, httpx.TransportError) + + +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 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 @@ -14,6 +75,13 @@ import httpx # different sets get independent rotators. _instance_rotators: dict[tuple[str, ...], "itertools.cycle[str]"] = {} +# Process-global per-instance breakers, keyed by base URL. Like the rotators, +# these must outlive the per-job client so an instance that crashed stays +# 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 (the docling-serve provider config). +_instance_breakers: dict[str, "_InstanceBreaker"] = {} + class DoclingServeClient: """Client for docling-serve async workflow. @@ -31,6 +99,9 @@ class DoclingServeClient: api_key: str | None = None, timeout: float = 300, transport: httpx.AsyncBaseTransport | None = None, + 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: @@ -40,6 +111,9 @@ class DoclingServeClient: self.timeout = timeout # transport is for testing — production callers leave it None. self._transport = transport + 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)) @@ -56,8 +130,49 @@ 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) -> str: - return next(self._instance_rotator) + """Next instance in the round-robin, skipping any whose circuit breaker + is open (an instance that recently crashed / overloaded). Falls back to + the first instance seen when every breaker is open — a single-instance + fleet, or a fully-down one, still gets a probe rather than no attempt.""" + fallback: str | None = None + for _ in range(len(self.base_urls)): + url = next(self._instance_rotator) + if fallback is None: + fallback = url + if not self._breaker_for(url).is_open: + return url + assert fallback is not None # base_urls is non-empty (checked in __init__) + return fallback + + def _record_outcome(self, base_url: str, exc: BaseException) -> None: + """Record a failed request against the instance's breaker, but only when + the failure reflects instance health (transport / 5xx). 4xx and + task-level errors are left alone — they aren't the instance's fault.""" + if not _is_instance_failure(exc): + return + 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.""" @@ -122,18 +237,26 @@ class DoclingServeClient: 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. + failure or returns no task_id. Instance-health failures (transport / + 5xx) trip the picked instance's circuit breaker so later requests skip + it while it recovers. """ headers = self._get_headers() base_url = self._pick_url() - async with self._httpx_client() as client: - task_id = await self._submit_and_wait( - client, base_url, endpoint, files, data, headers, name - ) - result_url = f"{base_url}/v1/result/{task_id}" - result_response = await client.get(result_url, headers=headers) - result_response.raise_for_status() - return result_response.json() + try: + async with self._httpx_client() as client: + task_id = await self._submit_and_wait( + client, base_url, endpoint, files, data, headers, name + ) + result_url = f"{base_url}/v1/result/{task_id}" + result_response = await client.get(result_url, headers=headers) + result_response.raise_for_status() + result = result_response.json() + except Exception as exc: + self._record_outcome(base_url, exc) + raise + self._breaker_for(base_url).record_success() + return result async def submit_and_poll_zip( self, @@ -147,15 +270,22 @@ 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 circuit-breaker bookkeeping); only the + result-fetching step differs. """ headers = self._get_headers() base_url = self._pick_url() - async with self._httpx_client() as client: - task_id = await self._submit_and_wait( - client, base_url, endpoint, files, data, headers, name - ) - result_url = f"{base_url}/v1/result/{task_id}" - result_response = await client.get(result_url, headers=headers) - result_response.raise_for_status() - return result_response.content + try: + async with self._httpx_client() as client: + task_id = await self._submit_and_wait( + client, base_url, endpoint, files, data, headers, name + ) + result_url = f"{base_url}/v1/result/{task_id}" + result_response = await client.get(result_url, headers=headers) + result_response.raise_for_status() + result = result_response.content + except Exception as exc: + self._record_outcome(base_url, exc) + raise + self._breaker_for(base_url).record_success() + return result diff --git a/tests/test_docling_serve_client.py b/tests/test_docling_serve_client.py index b19360d4..41676185 100644 --- a/tests/test_docling_serve_client.py +++ b/tests/test_docling_serve_client.py @@ -110,6 +110,129 @@ def test_round_robin_shared_across_fresh_clients(): assert picks == [urls[0], urls[1], urls[2], urls[0]] +def _health_transport( + down: set[str], task_id: str, result: dict +) -> tuple[httpx.MockTransport, list[str]]: + """MockTransport where requests to a host in the mutable ``down`` set raise + ConnectError (a crashed instance); other hosts serve a normal trio. The + caller can mutate ``down`` 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: + 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) -> None: + """Issue one submit_and_poll, swallowing the transport error raised when the + picked instance is down (this baseline has no in-request failover yet).""" + try: + await client.submit_and_poll( + endpoint="/v1/convert/file/async", + files={"file": ("x.md", b"x", "text/markdown")}, + data={}, + ) + except httpx.TransportError: + pass + + +@pytest.mark.asyncio +async def test_open_breaker_skips_crashed_instance(): + """Once an instance has failed enough to open its breaker, _pick_url skips + it — later requests route straight to a healthy instance without even + attempting the dead one.""" + down = {"crash-x"} + transport, seen = _health_transport(down, "t", {"ok": True}) + client = DoclingServeClient( + base_urls=["http://crash-x:5001", "http://live-x:5001"], + transport=transport, + breaker_failure_threshold=2, + ) + + await _poll(client) # picks crash-x → fail (failures=1) + await _poll(client) # picks live-x → ok + await _poll(client) # picks crash-x → fail (failures=2 → breaker opens) + + assert client._breaker_for("http://crash-x:5001").is_open + + seen.clear() + await _poll(client) + 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"} + transport, seen = _health_transport(down, "t", {"ok": True}) + flip = "http://flip-y:5001" + client = DoclingServeClient( + base_urls=[flip, "http://spare-y:5001"], + transport=transport, + breaker_failure_threshold=2, + breaker_cooldown_s=30.0, + now_fn=lambda: clock[0], + ) + + await _poll(client) # flip-y → fail (1) + await _poll(client) # spare-y → ok + await _poll(client) # flip-y → fail (2 → open) + 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 again. + clock[0] += 31.0 + assert not client._breaker_for(flip).is_open + + # Traffic can return to flip-y and success closes the breaker for good. + seen.clear() + for _ in range(2): + 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 with 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-z:5001" + client = DoclingServeClient( + base_urls=[bad], + transport=httpx.MockTransport(handler), + breaker_failure_threshold=1, + ) + + for _ in range(3): + with pytest.raises(httpx.HTTPStatusError): + await client.submit_and_poll( + endpoint="/v1/convert/file/async", + files={"file": ("x.md", b"x", "text/markdown")}, + data={}, + ) + + assert not client._breaker_for(bad).is_open + + @pytest.mark.asyncio async def test_zip_endpoint_uses_round_robin_too(): transport, seen = _scripted_transport(