danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Bound pdf_split memory + release lock between slices

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-05-25 16:24:50 +03:00
parent c2f681dd78
commit 9ca06df745
No known key found for this signature in database
2 changed files with 116 additions and 74 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

144
haiku_rag_slim/haiku/rag/converters/pdf_split.py
View file

@ -1,30 +1,8 @@
"""Split large PDFs into N-page slices, convert each, merge with
``DoclingDocument.concatenate``.
Rationale: docling's parser pipeline is memory-hungry and has confirmed leaks
(docling #2209, #1343, #2954; docling-serve #366, #474). Single-pass convert
of a 400-page PDF can OOM a workstation. Splitting bounds the peak working set
to one slice's pipeline state. Each slice round-trips through the existing
``DocumentConverter`` interface — so the docling-local and docling-serve
adapters work without modification, and each docling-serve task is independent
on the server side too.
Merge uses ``DoclingDocument.concatenate`` from docling-core ≥ 2.75. It
re-indexes ``self_ref`` values and shifts ``prov.page_no`` via an internal
``page_delta`` against ``_max_page`` (see
``docling_core/types/doc/document.py::_DocIndex.index``). No helper needed
on our side.
Cross-page references (named destinations, multi-page link annotations) are
dropped at the byte-level split — accepted loss; haiku.rag doesn't surface
them.
"""
import asyncio
import io
import tempfile
import threading
from collections.abc import Iterator
from collections.abc import Generator
from pathlib import Path
from typing import TYPE_CHECKING
@ -34,13 +12,13 @@ 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.
# Multiple workers calling iter_pdf_slices concurrently 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. Crucially we release it
# between slices so other workers' slicing can interleave; the heavy work
# (docling convert) happens between yields with the lock free.
_PDFIUM_LOCK = threading.Lock()
if TYPE_CHECKING:
@ -48,12 +26,21 @@ if TYPE_CHECKING:
from haiku.rag.converters.base import DocumentConverter
_SENTINEL: object = object()
def iter_pdf_slices(path: Path, slice_size: int) -> Iterator[tuple[int, int, bytes]]:
def iter_pdf_slices(
path: Path, slice_size: int
) -> Generator[tuple[int, int, bytes], None, None]:
"""Yield ``(start_page, end_page_inclusive, pdf_bytes)`` for each
``slice_size``-page slice of ``path``. Page numbers are 1-based to match
docling's ``prov.page_no`` convention. Each yielded byte string is a
standalone PDF that docling can convert.
The pdfium lock is held only around each pdfium call (open, slice
extraction, close) — never across a ``yield``. The source ``PdfDocument``
handle stays open across yields; per-document handles coexist safely as
long as no two pdfium calls execute concurrently.
"""
if slice_size <= 0:
raise ValueError(f"slice_size must be >= 1, got {slice_size}")
@ -61,26 +48,25 @@ def iter_pdf_slices(path: Path, slice_size: int) -> Iterator[tuple[int, int, byt
try:
src = pdfium.PdfDocument(str(path))
except pdfium.PdfiumError as exc:
# Malformed / unreadable PDF — retrying won't help. Raise the
# typed error so the ingester pipeline classifies it as
# PermanentError and dead-letters the job instead of cycling it
# through attempts.
raise UnsupportedSourceError(
f"pypdfium2 cannot open PDF {path}: {exc}"
) from exc
try:
total = len(src)
for start in range(0, total, slice_size):
end = min(start + slice_size, total)
total = len(src)
try:
for start in range(0, total, slice_size):
end = min(start + slice_size, total)
with _PDFIUM_LOCK:
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())
slice_bytes = buf.getvalue()
finally:
dst.close()
finally:
yield (start + 1, end, slice_bytes)
finally:
with _PDFIUM_LOCK:
src.close()
@ -93,42 +79,54 @@ async def convert_pdf_with_splitting(
"""Split a PDF, convert each slice through ``converter``, return the
merged ``DoclingDocument``.
A single slice failure aborts the whole job — surfaced as ``ValueError``
so the ingester pipeline's classifier maps it to ``TransientError`` and
the queue retries the entire document. Per-slice retry would risk
interleaved partial state with subsequent runs and is not worth the
complexity.
Slices are produced lazily — only one slice's bytes live in memory at a
time, so peak working set stays bounded regardless of page count. A
single slice failure aborts the whole job (raised as ``ValueError`` so
the ingester pipeline classifies it as ``TransientError`` and the queue
retries the entire document). Per-slice retry would risk interleaved
partial state with subsequent runs and is not worth the complexity.
"""
from docling_core.types.doc.document import DoclingDocument
slices = await asyncio.to_thread(lambda: list(iter_pdf_slices(path, slice_size)))
def _next(it):
return next(it, _SENTINEL)
it = iter_pdf_slices(path, slice_size)
converted: list[DoclingDocument] = []
for start, end, pdf_bytes in slices:
with tempfile.NamedTemporaryFile(mode="wb", suffix=".pdf", delete=False) as tmp:
tmp.write(pdf_bytes)
tmp.flush()
tmp_path = Path(tmp.name)
try:
# One span per slice so each docling-serve round-trip (or local
# docling invocation) is visible in Logfire; the outer
# `document.convert` span aggregates the lot.
with logfire.span(
"document.convert_slice",
uri=source_uri,
start_page=start,
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)
finally:
tmp_path.unlink(missing_ok=True)
try:
while True:
slice_item = await asyncio.to_thread(_next, it)
if slice_item is _SENTINEL:
break
start, end, pdf_bytes = slice_item
with tempfile.NamedTemporaryFile(
mode="wb", suffix=".pdf", delete=False
) as tmp:
tmp.write(pdf_bytes)
tmp.flush()
tmp_path = Path(tmp.name)
try:
with logfire.span(
"document.convert_slice",
uri=source_uri,
start_page=start,
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)
finally:
tmp_path.unlink(missing_ok=True)
finally:
# Close the generator under the pdfium lock so src.close() runs
# even when we abort mid-stream (slice failure, cancellation).
# Off the event loop because the close path acquires the lock.
await asyncio.to_thread(it.close)
return DoclingDocument.concatenate(converted)

46
tests/test_pdf_split.py
View file

@ -12,7 +12,10 @@ from pathlib import Path
import pypdfium2 as pdfium
import pytest
from haiku.rag.converters.pdf_split import iter_pdf_slices
from haiku.rag.converters.pdf_split import (
convert_pdf_with_splitting,
iter_pdf_slices,
)
def _make_pdf(page_count: int, tmp_path: Path) -> Path:
@ -66,6 +69,47 @@ def test_iter_pdf_slices_rejects_zero_slice_size(tmp_path):
list(iter_pdf_slices(src, slice_size=0))
@pytest.mark.asyncio
async def test_convert_aborts_and_cleans_up_on_mid_stream_slice_failure(
tmp_path, monkeypatch
):
"""When converting slice 2 of 3 fails, the whole call must raise
ValueError naming the failed slice's page range, every tempfile created
along the way must be deleted, and the source PDF handle must be closed.
"""
src = _make_pdf(7, tmp_path)
# Pin tempfiles to a per-test dir so we can list leaks deterministically.
monkeypatch.setattr("tempfile.tempdir", str(tmp_path))
calls: list[Path] = []
class _FlakyConverter:
async def convert_file(self, path: Path, *, source_uri):
calls.append(path)
if len(calls) == 2:
raise RuntimeError("docling exploded")
from docling_core.types.doc.document import DoclingDocument
return DoclingDocument(name="slice")
with pytest.raises(ValueError, match="pages 4-6"):
await convert_pdf_with_splitting(
_FlakyConverter(), # ty: ignore[invalid-argument-type]
src,
source_uri=None,
slice_size=3,
)
# All converter inputs were under tmp_path (the pinned tempdir)…
assert all(str(p).startswith(str(tmp_path)) for p in calls)
# …and none of them are still on disk.
leftover = [p for p in calls if p.exists()]
assert leftover == [], f"tempfiles leaked: {leftover}"
# We aborted after slice 2; slice 3 was never attempted.
assert len(calls) == 2
def test_concatenate_shifts_page_nos_and_unique_self_refs():
"""Pins the docling-core contract we rely on: when two docs (each with
items on page 1) are concatenated, the second doc's items move to page 2