round-robin docling-serve across multiple base_urls

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

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.

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

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)

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,
         )
 

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)
 

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:

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"]