danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add suggested changes

Browse source
This commit is contained in:
bryan davis 2026-06-26 11:52:55 -05:00
parent d8a55dbf7b
commit c3b4fc6ca7
No known key found for this signature in database
GPG key ID: D11B4A4C0C731E5E
14 changed files with 120 additions and 91 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
haiku_rag_slim/haiku/rag/chunkers/docling_serve.py
View file

@ -56,9 +56,11 @@ 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,
breaker_config=ds.circuit_breaker,
)
self.chunker_type = config.processing.chunker_type

6
haiku_rag_slim/haiku/rag/ingester/pollers/circuit_breaker.py → haiku_rag_slim/haiku/rag/circuit_breaker.py
View file

@ -5,7 +5,11 @@ from haiku.rag.config import CircuitBreakerConfig
class CircuitBreaker:
"""Three-state breaker over discover() failures.
"""Three-state circuit breaker (closed / open with cooldown).
Shared by the ingester's per-source discover() breaker and the
docling-serve per-instance breaker — lives here, below both
``providers`` and ``ingester``, so neither has to depend on the other.
- closed: failures are counted; threshold flips to open.
- open: probes are blocked until cooldown elapses, then a single probe

35
haiku_rag_slim/haiku/rag/config/models.py
View file

@ -218,6 +218,20 @@ class OllamaConfig(BaseModel):
)
class CircuitBreakerConfig(BaseModel):
"""Three-state breaker (closed -> open -> probe) over consecutive failures.
Used by the ingester's per-source discover() breaker and the docling-serve
per-instance breaker."""
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
@ -227,6 +241,14 @@ class DoclingServeConfig(BaseModel):
base_url: str | list[str] = "http://localhost:5001"
api_key: str = ""
circuit_breaker: CircuitBreakerConfig = Field(
default_factory=lambda: CircuitBreakerConfig(
failure_threshold=3, cooldown_s=30.0
),
description="Per-instance circuit breaker: skip an instance after this "
"many consecutive failures for this cooldown. Defaults are tuned faster "
"than the discover breaker since instances restart quickly.",
)
@property
def base_urls(self) -> list[str]:
@ -295,19 +317,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):
worker_count: int = Field(
default=4,

6
haiku_rag_slim/haiku/rag/converters/docling_serve.py
View file

@ -59,9 +59,11 @@ 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,
breaker_config=ds.circuit_breaker,
)
@property

2
haiku_rag_slim/haiku/rag/ingester/pollers/__init__.py
View file

@ -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

2
haiku_rag_slim/haiku/rag/ingester/pollers/base.py
View file

@ -5,7 +5,7 @@ from datetime import UTC, datetime
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.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 (

2
haiku_rag_slim/haiku/rag/ingester/pollers/fs.py
View file

@ -12,7 +12,7 @@ from haiku.rag.telemetry import logfire
if TYPE_CHECKING:
from haiku.rag.config import FSSourceConfig
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
from haiku.rag.circuit_breaker import CircuitBreaker
from haiku.rag.ingester.sources.fs import FSSource
logger = logging.getLogger(__name__)

2
haiku_rag_slim/haiku/rag/ingester/pollers/manager.py
View file

@ -7,7 +7,7 @@ from typing import TYPE_CHECKING
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.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

2
haiku_rag_slim/haiku/rag/ingester/pollers/periodic.py
View file

@ -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):

2
haiku_rag_slim/haiku/rag/ingester/workers/pool.py
View file

@ -5,7 +5,7 @@ from typing import TYPE_CHECKING
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.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

82
haiku_rag_slim/haiku/rag/providers/docling_serve.py
View file

@ -9,6 +9,9 @@ from typing import Any
import httpx
from haiku.rag.circuit_breaker import CircuitBreaker
from haiku.rag.config import CircuitBreakerConfig
logger = logging.getLogger(__name__)
# HTTP statuses that count against an instance's health: 408/429 are transient
@ -29,44 +32,6 @@ def _is_instance_failure(exc: BaseException) -> bool:
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
@ -80,7 +45,7 @@ _instance_rotators: dict[tuple[str, ...], "itertools.cycle[str]"] = {}
# 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"] = {}
_instance_breakers: dict[str, "CircuitBreaker"] = {}
class DoclingServeClient:
@ -99,8 +64,7 @@ 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,
breaker_config: CircuitBreakerConfig | None = None,
now_fn: Callable[[], float] = time.monotonic,
):
urls = [base_urls] if isinstance(base_urls, str) else list(base_urls)
@ -111,8 +75,12 @@ 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
# Per-instance breaker config; defaults are tuned faster than the
# ingester's discover breaker since docling-serve instances restart
# quickly. Overridden from DoclingServeConfig.circuit_breaker.
self._breaker_config = breaker_config or CircuitBreakerConfig(
failure_threshold=3, cooldown_s=30.0
)
self._now = now_fn
# setdefault is atomic under the GIL — concurrent constructors with
# the same instance set will end up sharing one rotator.
@ -130,31 +98,27 @@ class DoclingServeClient:
callers should let `_pick_url` round-robin per request."""
return self.base_urls[0]
def _breaker_for(self, url: str) -> _InstanceBreaker:
def _breaker_for(self, url: str) -> CircuitBreaker:
breaker = _instance_breakers.get(url)
if breaker is None:
breaker = _InstanceBreaker(
self._breaker_failure_threshold,
self._breaker_cooldown_s,
self._now,
)
breaker = CircuitBreaker(self._breaker_config, now_fn=self._now)
_instance_breakers[url] = breaker
return breaker
def _pick_url(self) -> str:
"""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
is open (an instance that recently crashed / overloaded). When every
breaker is open — a single-instance fleet, or a fully-down one — probe
one anyway rather than failing with nothing to pick."""
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
# All breakers open. Take one more step (the loop consumed a full
# rotation, landing the cursor back at the start) so consecutive
# all-open calls rotate across instances instead of pinning the first
# one — an all-429 overload shouldn't pile every retry on one node.
return next(self._instance_rotator)
def _record_outcome(self, base_url: str, exc: BaseException) -> None:
"""Record a failed request against the instance's breaker, but only when
@ -170,8 +134,8 @@ class DoclingServeClient:
"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,
self._breaker_config.failure_threshold,
self._breaker_config.cooldown_s,
)
def _get_headers(self) -> dict[str, str]:

2
tests/ingester/test_circuit_breaker.py
View file

@ -1,5 +1,5 @@
from haiku.rag.config import CircuitBreakerConfig
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
from haiku.rag.circuit_breaker import CircuitBreaker
class _Clock:

2
tests/ingester/test_pollers.py
View file

@ -11,7 +11,7 @@ from haiku.rag.config import (
S3SourceConfig,
WebDAVSourceConfig,
)
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
from haiku.rag.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

60
tests/test_docling_serve_client.py
View file

@ -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
@ -152,7 +153,7 @@ async def test_open_breaker_skips_crashed_instance():
client = DoclingServeClient(
base_urls=["http://crash-x:5001", "http://live-x:5001"],
transport=transport,
breaker_failure_threshold=2,
breaker_config=CircuitBreakerConfig(failure_threshold=2, cooldown_s=30.0),
)
await _poll(client) # picks crash-x → fail (failures=1)
@ -179,8 +180,7 @@ async def test_breaker_recovers_after_cooldown():
client = DoclingServeClient(
base_urls=[flip, "http://spare-y:5001"],
transport=transport,
breaker_failure_threshold=2,
breaker_cooldown_s=30.0,
breaker_config=CircuitBreakerConfig(failure_threshold=2, cooldown_s=30.0),
now_fn=lambda: clock[0],
)
@ -219,7 +219,7 @@ async def test_4xx_does_not_trip_breaker():
client = DoclingServeClient(
base_urls=[bad],
transport=httpx.MockTransport(handler),
breaker_failure_threshold=1,
breaker_config=CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
)
for _ in range(3):
@ -243,7 +243,7 @@ async def test_pick_url_falls_back_when_all_breakers_open():
client = DoclingServeClient(
base_urls=[lone],
transport=transport,
breaker_failure_threshold=1,
breaker_config=CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
)
await _poll(client) # fails → breaker opens (threshold=1)
@ -253,6 +253,54 @@ async def test_pick_url_falls_back_when_all_breakers_open():
assert client._pick_url() == lone
@pytest.mark.asyncio
async def test_all_open_breakers_rotate_instead_of_pinning_one():
"""When every breaker is open, successive picks must 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",
]
transport, _ = _health_transport(
{"allopen-a", "allopen-b", "allopen-c"}, "t", {"ok": True}
)
client = DoclingServeClient(
base_urls=urls,
transport=transport,
breaker_config=CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
)
# One failure each opens all three breakers: each _poll picks the next
# still-closed instance, fails, and opens it.
for _ in range(3):
await _poll(client)
assert all(client._breaker_for(u).is_open for u in urls)
picks = [client._pick_url() for _ in range(6)]
# All-open fallback rotates a → b → c → a → b → c, not the same URL six times.
assert picks == urls + urls
def test_config_breaker_reaches_client():
"""The breaker config set on 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()
config.providers.docling_serve.base_url = "http://cfg-brk:5001"
config.providers.docling_serve.circuit_breaker = CircuitBreakerConfig(
failure_threshold=9, cooldown_s=123.0
)
for component in (DoclingServeConverter(config), DoclingServeChunker(config)):
assert component.client._breaker_config.failure_threshold == 9
assert component.client._breaker_config.cooldown_s == 123.0
@pytest.mark.asyncio
async def test_zip_records_instance_failure():
"""submit_and_poll_zip records a crashed instance against its breaker, the
@ -262,7 +310,7 @@ async def test_zip_records_instance_failure():
client = DoclingServeClient(
base_urls=[z],
transport=transport,
breaker_failure_threshold=1,
breaker_config=CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
)
with pytest.raises(httpx.ConnectError):