change breakers to be per-instance

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,10 +237,13 @@ 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()
+        try:
             async with self._httpx_client() as client:
                 task_id = await self._submit_and_wait(
                     client, base_url, endpoint, files, data, headers, name
@@ -133,7 +251,12 @@ class DoclingServeClient:
                 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()
+                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,10 +270,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 circuit-breaker bookkeeping); only the
+        result-fetching step differs.
         """
         headers = self._get_headers()
         base_url = self._pick_url()
+        try:
             async with self._httpx_client() as client:
                 task_id = await self._submit_and_wait(
                     client, base_url, endpoint, files, data, headers, name
@@ -158,4 +283,9 @@ class DoclingServeClient:
                 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
+                result = result_response.content
+        except Exception as exc:
+            self._record_outcome(base_url, exc)
+            raise
+        self._breaker_for(base_url).record_success()
+        return result

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(