Lock pypdfium2 under concurrent workers; per-slice convert spans

This commit is contained in:
Yiorgis Gozadinos 2026-05-25 15:30:34 +03:00
parent 567a0daf67
commit 14beba6786
No known key found for this signature in database
4 changed files with 66 additions and 25 deletions

View file

@ -3,7 +3,8 @@
### Added ### Added
- New `haiku-ingester` service for continuous document ingestion: persistent SQLite job queue, async worker pool with retries and a dead-letter queue, FS/HTTP/S3/WebDAV source adapters with per-source circuit breakers, and a FastAPI control plane (`/health`, `/jobs`, `/sources`, `/dlq`). Configured under `ingester:` in `haiku.rag.yaml`. Shipped behind the `[ingester]` extra, with Logfire spans (`ingester.poller.sweep` → `ingester.job``document.{fetch,convert,chunk,embed,store}`) for traceable ingestion. See [docs/ingester.md](docs/ingester.md). - New `haiku-ingester` service for continuous document ingestion: persistent SQLite job queue, async worker pool with retries and a dead-letter queue, FS/HTTP/S3/WebDAV source adapters with per-source circuit breakers, and a FastAPI control plane (`/health`, `/jobs`, `/sources`, `/dlq`). Configured under `ingester:` in `haiku.rag.yaml`. Shipped behind the `[ingester]` extra, with Logfire spans (`ingester.poller.sweep` → `ingester.job``document.{fetch,convert,chunk,embed,store}`, plus `document.convert_slice` per slice when splitting) for traceable ingestion. See [docs/ingester.md](docs/ingester.md).
- `processing.split_pages` (default `0`): split large PDFs into N-page slices, convert each independently through docling-local or docling-serve, and merge with `DoclingDocument.concatenate()`. Bounds peak working set on memory-hungry docs and lets multiple docling-serve replicas parallelize per-document. `0` disables (single-pass conversion).
### Removed ### Removed

View file

@ -72,6 +72,10 @@ services:
condition: service_healthy condition: service_healthy
docling-serve-2: docling-serve-2:
condition: service_healthy condition: service_healthy
# Compose default is 10s — too short for the ingester's own 60s
# shutdown_grace_s. Give it 180s so SIGTERM → drain → exit completes
# before the SIGKILL fallback fires.
stop_grace_period: 180s
healthcheck: healthcheck:
# No curl in the slim image; use Python's stdlib instead. # No curl in the slim image; use Python's stdlib instead.
test: test:
@ -115,4 +119,5 @@ services:
depends_on: depends_on:
haiku-ingester: haiku-ingester:
condition: service_healthy condition: service_healthy
stop_grace_period: 30s
restart: unless-stopped restart: unless-stopped

View file

@ -23,12 +23,26 @@ them.
import asyncio import asyncio
import io import io
import tempfile import tempfile
import threading
from collections.abc import Iterator from collections.abc import Iterator
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
import logfire
import pypdfium2 as pdfium import pypdfium2 as pdfium
from haiku.rag.client.exceptions import UnsupportedSourceError
# pypdfium2 wraps libpdfium, which has global C state and is not thread-safe.
# Multiple workers calling iter_pdf_slices concurrently (via asyncio.to_thread)
# race on that state and corrupt it — first error surfaces as e.g.
# "Failed to import pages", and after that every subsequent PDF load fails
# with "Failed to load document" until the process restarts. Hold this lock
# around every pdfium call so only one worker's slicing runs at a time. The
# actual heavy work (docling-serve convert) happens AFTER iter_pdf_slices
# returns and continues to parallelize across workers.
_PDFIUM_LOCK = threading.Lock()
if TYPE_CHECKING: if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
@ -43,21 +57,31 @@ def iter_pdf_slices(path: Path, slice_size: int) -> Iterator[tuple[int, int, byt
""" """
if slice_size <= 0: if slice_size <= 0:
raise ValueError(f"slice_size must be >= 1, got {slice_size}") raise ValueError(f"slice_size must be >= 1, got {slice_size}")
src = pdfium.PdfDocument(str(path)) with _PDFIUM_LOCK:
try: try:
total = len(src) src = pdfium.PdfDocument(str(path))
for start in range(0, total, slice_size): except pdfium.PdfiumError as exc:
end = min(start + slice_size, total) # Malformed / unreadable PDF — retrying won't help. Raise the
dst = pdfium.PdfDocument.new() # typed error so the ingester pipeline classifies it as
try: # PermanentError and dead-letters the job instead of cycling it
dst.import_pages(src, list(range(start, end))) # through attempts.
buf = io.BytesIO() raise UnsupportedSourceError(
dst.save(buf) f"pypdfium2 cannot open PDF {path}: {exc}"
yield (start + 1, end, buf.getvalue()) ) from exc
finally: try:
dst.close() total = len(src)
finally: for start in range(0, total, slice_size):
src.close() end = min(start + slice_size, total)
dst = pdfium.PdfDocument.new()
try:
dst.import_pages(src, list(range(start, end)))
buf = io.BytesIO()
dst.save(buf)
yield (start + 1, end, buf.getvalue())
finally:
dst.close()
finally:
src.close()
async def convert_pdf_with_splitting( async def convert_pdf_with_splitting(
@ -86,14 +110,23 @@ async def convert_pdf_with_splitting(
tmp.flush() tmp.flush()
tmp_path = Path(tmp.name) tmp_path = Path(tmp.name)
try: try:
try: # One span per slice so each docling-serve round-trip (or local
slice_doc = await converter.convert_file( # docling invocation) is visible in Logfire; the outer
tmp_path, source_uri=source_uri # `document.convert` span aggregates the lot.
) with logfire.span(
except Exception as exc: "document.convert_slice",
raise ValueError( uri=source_uri,
f"Failed to convert slice pages {start}-{end} of {path}: {exc}" start_page=start,
) from exc end_page=end,
):
try:
slice_doc = await converter.convert_file(
tmp_path, source_uri=source_uri
)
except Exception as exc:
raise ValueError(
f"Failed to convert slice pages {start}-{end} of {path}: {exc}"
) from exc
converted.append(slice_doc) converted.append(slice_doc)
finally: finally:
tmp_path.unlink(missing_ok=True) tmp_path.unlink(missing_ok=True)

View file

@ -52,5 +52,7 @@ def build_app(
app.include_router(dlq.router, dependencies=auth_dep) app.include_router(dlq.router, dependencies=auth_dep)
# Every request becomes a span when logfire is configured; no-op otherwise. # Every request becomes a span when logfire is configured; no-op otherwise.
logfire.instrument_fastapi(app) # /health is the docker healthcheck endpoint — polled every few seconds
# by Compose, would otherwise drown the trace stream in idle GETs.
logfire.instrument_fastapi(app, excluded_urls=r"^.*/health$")
return app return app