483 lines
17 KiB
Python
483 lines
17 KiB
Python
"""DoclingServeClient round-robin distribution tests.
|
|
|
|
Each test here uses a unique base-URL set so the process-global cycle
|
|
map gives it a fresh itertools.cycle. Don't reuse URL strings across
|
|
tests — cycles persist for the lifetime of the process and would
|
|
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
|
|
|
|
|
|
def _scripted_transport(responses_by_path: dict[str, httpx.Response]):
|
|
"""MockTransport routing on (host, path) — lets us assert which URL got hit."""
|
|
seen_hosts: list[str] = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
seen_hosts.append(request.url.host)
|
|
key = request.url.path
|
|
if key in responses_by_path:
|
|
return responses_by_path[key]
|
|
return httpx.Response(404)
|
|
|
|
return httpx.MockTransport(handler), seen_hosts
|
|
|
|
|
|
def _success_routes(task_id: str, result: dict) -> dict[str, httpx.Response]:
|
|
return {
|
|
"/v1/convert/file/async": httpx.Response(200, json={"task_id": task_id}),
|
|
f"/v1/status/poll/{task_id}": httpx.Response(
|
|
200, json={"task_status": "success"}
|
|
),
|
|
f"/v1/result/{task_id}": httpx.Response(200, json=result),
|
|
}
|
|
|
|
|
|
def test_single_url_input_normalises_to_list():
|
|
client = DoclingServeClient(base_urls="http://only:5001")
|
|
assert client.base_urls == ["http://only:5001"]
|
|
|
|
|
|
def test_empty_list_raises():
|
|
with pytest.raises(ValueError, match="at least one"):
|
|
DoclingServeClient(base_urls=[])
|
|
|
|
|
|
def test_trailing_slashes_stripped():
|
|
client = DoclingServeClient(base_urls=["http://a:5001/", "http://b:5001//"])
|
|
assert client.base_urls == ["http://a:5001", "http://b:5001"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_round_robin_across_three_urls():
|
|
transport, seen = _scripted_transport(_success_routes("t", {"ok": True}))
|
|
client = DoclingServeClient(
|
|
base_urls=["http://a:5001", "http://b:5001", "http://c:5001"],
|
|
transport=transport,
|
|
)
|
|
|
|
for _ in range(6):
|
|
await client.submit_and_poll(
|
|
endpoint="/v1/convert/file/async",
|
|
files={"file": ("x.md", b"x", "text/markdown")},
|
|
data={},
|
|
)
|
|
|
|
# 3 round-trips per call (POST + GET poll + GET result) * 6 calls = 18 requests.
|
|
# Each call must stay on one host; calls rotate through a, b, c, a, b, c.
|
|
per_call = [seen[i : i + 3] for i in range(0, 18, 3)]
|
|
assert all(len(set(triple)) == 1 for triple in per_call), (
|
|
"submit/poll/result split across hosts — task_id wouldn't resolve"
|
|
)
|
|
hosts_picked = [triple[0] for triple in per_call]
|
|
assert hosts_picked == ["a", "b", "c", "a", "b", "c"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_task_lifecycle_pinned_to_same_url():
|
|
"""A single submit/poll/result trio must all hit the same instance —
|
|
task IDs are local to the instance that issued them."""
|
|
transport, seen = _scripted_transport(_success_routes("task-42", {"r": 1}))
|
|
client = DoclingServeClient(
|
|
base_urls=["http://primary:5001", "http://secondary:5001"],
|
|
transport=transport,
|
|
)
|
|
|
|
await client.submit_and_poll(
|
|
endpoint="/v1/convert/file/async",
|
|
files={"file": ("x.md", b"x", "text/markdown")},
|
|
data={},
|
|
)
|
|
|
|
# All three requests must be on the same host.
|
|
assert len(set(seen)) == 1
|
|
|
|
|
|
def test_round_robin_shared_across_fresh_clients():
|
|
"""get_converter / get_chunker build a NEW DoclingServeClient per job.
|
|
The cycle has to live outside the instance so successive jobs (each
|
|
with its own freshly-constructed client) actually rotate."""
|
|
urls = ["http://x:5001", "http://y:5001", "http://z:5001"]
|
|
c1 = DoclingServeClient(base_urls=urls)
|
|
c2 = DoclingServeClient(base_urls=urls)
|
|
c3 = DoclingServeClient(base_urls=urls)
|
|
c4 = DoclingServeClient(base_urls=urls)
|
|
|
|
picks = [c1._pick_url(), c2._pick_url(), c3._pick_url(), c4._pick_url()]
|
|
assert picks == [urls[0], urls[1], urls[2], urls[0]]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_zip_endpoint_uses_round_robin_too():
|
|
transport, seen = _scripted_transport(
|
|
{
|
|
"/v1/convert/file/async": httpx.Response(200, json={"task_id": "t"}),
|
|
"/v1/status/poll/t": httpx.Response(200, json={"task_status": "success"}),
|
|
"/v1/result/t": httpx.Response(200, content=b"zip-bytes"),
|
|
}
|
|
)
|
|
client = DoclingServeClient(
|
|
base_urls=["http://a:5001", "http://b:5001"], transport=transport
|
|
)
|
|
await client.submit_and_poll_zip(
|
|
endpoint="/v1/convert/file/async",
|
|
files={"file": ("x.md", b"x", "text/markdown")},
|
|
data={},
|
|
)
|
|
await client.submit_and_poll_zip(
|
|
endpoint="/v1/convert/file/async",
|
|
files={"file": ("x.md", b"x", "text/markdown")},
|
|
data={},
|
|
)
|
|
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"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_request_span_records_instance_per_attempt(monkeypatch):
|
|
"""Each attempt opens a docling_serve.request span tagged with the
|
|
instance URL, so failover is traceable in Logfire."""
|
|
from contextlib import nullcontext
|
|
|
|
from haiku.rag.providers import docling_serve as ds_module
|
|
|
|
spans: list[dict] = []
|
|
|
|
def _fake_span(span_name, /, **attrs):
|
|
spans.append({"span_name": span_name, **attrs})
|
|
return nullcontext()
|
|
|
|
monkeypatch.setattr(ds_module.logfire, "span", _fake_span)
|
|
|
|
transport, _ = _failover_transport({"down-s"}, "t", {"ok": True})
|
|
client = DoclingServeClient(
|
|
base_urls=["http://down-s:5001", "http://up-s:5001"],
|
|
transport=transport,
|
|
retry_base_delay=0.0,
|
|
)
|
|
|
|
await _poll(client)
|
|
|
|
requests = [s for s in spans if s["span_name"] == "docling_serve.request"]
|
|
assert [s["url"] for s in requests] == ["http://down-s:5001", "http://up-s:5001"]
|
|
assert [s["attempt"] for s in requests] == [0, 1]
|
|
|
|
|
|
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
|