From 055fd23d5df6e2700adaf41f4acc86e1a32620ac Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 22 May 2026 16:51:12 +0300 Subject: [PATCH] round-robin docling-serve across multiple base_urls --- CHANGELOG.md | 1 + docs/configuration/processing.md | 29 +++++ .../haiku/rag/chunkers/docling_serve.py | 2 +- haiku_rag_slim/haiku/rag/config/models.py | 15 ++- .../haiku/rag/converters/docling_serve.py | 2 +- haiku_rag_slim/haiku/rag/ingester/app.py | 24 +++- .../haiku/rag/providers/docling_serve.py | 59 +++++++-- tests/test_docling_serve_client.py | 116 ++++++++++++++++++ 8 files changed, 229 insertions(+), 19 deletions(-) create mode 100644 tests/test_docling_serve_client.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 55717a09..a2c28ac2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ - `document.metadata` now uses source-agnostic keys: `source_revision` (was `etag` — S3-only and never populated for FS, so periodic sweeps re-ingested every file) and `content_type` (was `contentType`, snake_case for consistency). The v0.50.0 startup migration rewrites existing documents. All four source adapters (FS, HTTP, S3, future WebDAV) now write their native revision (mtime_ns, ETag, etc.) under the same key, fixing the regression where FS sources never short-circuited on unchanged files. - Ingester pollers skip their periodic sweep when the source already has queued or claimed jobs in the queue — saves the listing round-trip (`PROPFIND` / `S3 LIST` / FS walk) when work is backed up. FS push events from `watchfiles` keep flowing during skipped sweeps. Visible in Logfire as `ingester.poller.sweep` spans with `skipped=true reason=pending_work`. - Ingester now drains in-flight jobs on `SIGINT` / `SIGTERM` up to `workers.shutdown_grace_s` (default 60s) before cancelling. Cancelled jobs stay `claimed` and the reaper resets them on next start. Bonus: the pipeline no longer wraps `KeyboardInterrupt` / `SystemExit` / `CancelledError` as `TransientError` — those now propagate as intended. +- `providers.docling_serve.base_url` now accepts a list. Jobs round-robin across the entries with each job's submit/poll/result pinned to one instance (task IDs are instance-local). The counter is per-process; for cross-process load balancing or failover, put an LB in front and pass a single URL here. - Drop `list_documents` and `get_document` from the default RAG skill's tool set; the skill now exposes only `search` and `cite`. Both tools dumped unbounded content into the agent's context (full document lists, full document bodies) and `get_document` returned no chunk_ids so its output was structurally uncitable. The analysis skill already covers these uses programmatically — `await list_documents()` and `Path('/documents/{id}/content.txt').read_text()` inside `execute_code`. The tool branches remain in `create_skill_tools` and the `skill_generator` `AVAILABLE_TOOLS` set so users can still opt in when building custom skills. ## [0.48.1] - 2026-05-21 diff --git a/docs/configuration/processing.md b/docs/configuration/processing.md index c0a15bbe..576bce5a 100644 --- a/docs/configuration/processing.md +++ b/docs/configuration/processing.md @@ -80,6 +80,35 @@ providers: api_key: "your-api-key" # Optional ``` +`base_url` also accepts a list — jobs round-robin across the entries, with +each job's submit / poll / result pinned to one instance (task IDs are +instance-local): + +```yaml +providers: + docling_serve: + base_url: + - http://gpu-1:5001 + - http://cpu-1:5001 + - http://cpu-2:5001 +``` + +The round-robin counter is per-process — multiple concurrent ingester or +client processes pick independently, so the distribution evens out over many +jobs without coordination. For tight load balancing, failover, or health +checks, run a real load balancer (nginx `least_conn`, etc.) in front and +configure a single URL here. + +**Tuning `ingester.workers.worker_count` for docling-serve users**: convert +is usually the throughput ceiling — a default docling-serve instance +processes one task at a time (configurable via `DOCLING_SERVE_ENG_LOC_NUM_WORKERS` +if you've set it). A reasonable starting point for `worker_count` is **1–2 × +the number of `docling_serve.base_url` entries**: enough to overlap fetch / +embed / store of one job with the convert of another, without piling jobs +into docling-serve's internal queue beyond what its workers can chew through. +The ingester logs the worker / source / docling-serve counts on startup so +you can eyeball the ratio. + Conversion options work identically for both local and remote processing. **Note:** When using `chunker: docling-serve`, OCR options (`do_ocr`, `force_ocr`, `ocr_engine`, `ocr_lang`) from `conversion_options` are passed to the chunking API. This is useful when running docling-serve in a read-only container where OCR model downloads fail. Set `do_ocr: false` to disable OCR entirely. diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py index 81d1758e..3cf4dc03 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py @@ -56,7 +56,7 @@ class DoclingServeChunker(DocumentChunker): def __init__(self, config: AppConfig = Config): self.config = config self.client = DoclingServeClient( - base_url=config.providers.docling_serve.base_url, + base_urls=config.providers.docling_serve.base_urls, api_key=config.providers.docling_serve.api_key, ) 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 7d4fcc47..631305b8 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -200,9 +200,22 @@ class OllamaConfig(BaseModel): class DoclingServeConfig(BaseModel): - base_url: str = "http://localhost:5001" + """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 + submit/poll/result trio stays on the same instance (task IDs are + instance-local). The round-robin counter is per-process; for true load + balancing or failover, put an LB in front.""" + + base_url: str | list[str] = "http://localhost:5001" api_key: str = "" + @property + def base_urls(self) -> list[str]: + """Always-a-list view of base_url. Empty input falls back to localhost.""" + if isinstance(self.base_url, str): + return [self.base_url] + return list(self.base_url) or ["http://localhost:5001"] + class ProvidersConfig(BaseModel): ollama: OllamaConfig = Field(default_factory=OllamaConfig) diff --git a/haiku_rag_slim/haiku/rag/converters/docling_serve.py b/haiku_rag_slim/haiku/rag/converters/docling_serve.py index 9d15accc..b07366f4 100644 --- a/haiku_rag_slim/haiku/rag/converters/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/converters/docling_serve.py @@ -60,7 +60,7 @@ class DoclingServeConverter(DocumentConverter): """ self.config = config self.client = DoclingServeClient( - base_url=config.providers.docling_serve.base_url, + base_urls=config.providers.docling_serve.base_urls, api_key=config.providers.docling_serve.api_key, ) diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index cb11de50..03600207 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -86,11 +86,27 @@ class IngesterApp: await self._pool.start() await self._pollers.start() - logger.info( - "Ingester running: %d worker(s), %d source(s)", - ingester_cfg.workers.worker_count, - len(ingester_cfg.sources), + # Log the docling-serve fleet size when relevant so the + # operator can eyeball the worker/instance ratio. The convert + # phase is usually the throughput ceiling. + proc = self._config.processing + uses_docling_serve = ( + proc.converter == "docling-serve" or proc.chunker == "docling-serve" ) + if uses_docling_serve: + logger.info( + "Ingester running: %d worker(s), %d source(s), " + "%d docling-serve instance(s)", + ingester_cfg.workers.worker_count, + len(ingester_cfg.sources), + len(self._config.providers.docling_serve.base_urls), + ) + else: + logger.info( + "Ingester running: %d worker(s), %d source(s)", + ingester_cfg.workers.worker_count, + len(ingester_cfg.sources), + ) api_task, api_server = await self._maybe_start_api(api) diff --git a/haiku_rag_slim/haiku/rag/providers/docling_serve.py b/haiku_rag_slim/haiku/rag/providers/docling_serve.py index 3345fed9..29b53879 100644 --- a/haiku_rag_slim/haiku/rag/providers/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/providers/docling_serve.py @@ -10,12 +10,44 @@ class DoclingServeClient: """Client for docling-serve async workflow. Handles the submit → poll → fetch pattern used by both conversion and chunking. + + Accepts a list of base URLs and round-robins jobs across them. Each job's + submit/poll/result trio stays on the same instance — task IDs are + instance-local, so picking a different URL mid-job would 404. The counter + is per-process; concurrent processes pick independently, so over many + jobs the distribution is even but not coordinated. """ - def __init__(self, base_url: str, api_key: str | None = None, timeout: float = 300): - self.base_url = base_url.rstrip("/") + def __init__( + self, + base_urls: str | list[str], + api_key: str | None = None, + timeout: float = 300, + transport: httpx.AsyncBaseTransport | None = None, + ): + urls = [base_urls] if isinstance(base_urls, str) else list(base_urls) + if not urls: + raise ValueError("DoclingServeClient requires at least one base_url") + self.base_urls: list[str] = [u.rstrip("/") for u in urls] self.api_key = api_key self.timeout = timeout + # transport is for testing — production callers leave it None. + self._transport = transport + self._counter = 0 + + def _httpx_client(self) -> httpx.AsyncClient: + return httpx.AsyncClient(timeout=self.timeout, transport=self._transport) + + @property + def base_url(self) -> str: + """First-URL view, mostly for log messages. Don't use for dispatch — + callers should let `_pick_url` round-robin per request.""" + return self.base_urls[0] + + def _pick_url(self) -> str: + url = self.base_urls[self._counter % len(self.base_urls)] + self._counter += 1 + return url def _get_headers(self) -> dict[str, str]: """Get headers for API requests.""" @@ -27,6 +59,7 @@ class DoclingServeClient: async def _submit_and_wait( self, client: httpx.AsyncClient, + base_url: str, endpoint: str, files: dict[str, Any], data: dict[str, Any], @@ -38,7 +71,7 @@ class DoclingServeClient: Shared by submit_and_poll (JSON results) and submit_and_poll_zip (binary zip results) — only the result-fetching step differs. """ - submit_url = f"{self.base_url}{endpoint}" + submit_url = f"{base_url}{endpoint}" response = await client.post( submit_url, files=files, @@ -52,7 +85,7 @@ class DoclingServeClient: if not task_id: raise ValueError("docling-serve did not return a task_id") - poll_url = f"{self.base_url}/v1/status/poll/{task_id}" + poll_url = f"{base_url}/v1/status/poll/{task_id}" while True: poll_response = await client.get(poll_url, headers=headers) poll_response.raise_for_status() @@ -88,20 +121,21 @@ class DoclingServeClient: ValueError: If the task fails or service is unavailable """ headers = self._get_headers() + base_url = self._pick_url() try: - async with httpx.AsyncClient(timeout=self.timeout) as client: + async with self._httpx_client() as client: task_id = await self._submit_and_wait( - client, endpoint, files, data, headers, name + client, base_url, endpoint, files, data, headers, name ) - result_url = f"{self.base_url}/v1/result/{task_id}" + 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() except httpx.ConnectError as e: raise ValueError( - f"Could not connect to docling-serve at {self.base_url}. " + f"Could not connect to docling-serve at {base_url}. " f"Ensure the service is running and accessible. Error: {e}" ) from e except httpx.TimeoutException as e: @@ -134,20 +168,21 @@ class DoclingServeClient: ``submit_and_poll``; only the result-fetching step differs. """ headers = self._get_headers() + base_url = self._pick_url() try: - async with httpx.AsyncClient(timeout=self.timeout) as client: + async with self._httpx_client() as client: task_id = await self._submit_and_wait( - client, endpoint, files, data, headers, name + client, base_url, endpoint, files, data, headers, name ) - result_url = f"{self.base_url}/v1/result/{task_id}" + 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 except httpx.ConnectError as e: raise ValueError( - f"Could not connect to docling-serve at {self.base_url}. " + f"Could not connect to docling-serve at {base_url}. " f"Ensure the service is running and accessible. Error: {e}" ) from e except httpx.TimeoutException as e: diff --git a/tests/test_docling_serve_client.py b/tests/test_docling_serve_client.py new file mode 100644 index 00000000..de40d705 --- /dev/null +++ b/tests/test_docling_serve_client.py @@ -0,0 +1,116 @@ +"""DoclingServeClient round-robin distribution tests.""" + +import httpx +import pytest + +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 + + +@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"]