Lock pypdfium2 under concurrent workers; per-slice convert spans
This commit is contained in:
parent
567a0daf67
commit
14beba6786
4 changed files with 66 additions and 25 deletions
|
|
@ -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
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue