From 9ca06df7450425d739069057c9e8398eeca82850 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 25 May 2026 16:24:50 +0300 Subject: [PATCH] Bound pdf_split memory + release lock between slices --- .../haiku/rag/converters/pdf_split.py | 144 +++++++++--------- tests/test_pdf_split.py | 46 +++++- 2 files changed, 116 insertions(+), 74 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/converters/pdf_split.py b/haiku_rag_slim/haiku/rag/converters/pdf_split.py index ece6d103..6139bb14 100644 --- a/haiku_rag_slim/haiku/rag/converters/pdf_split.py +++ b/haiku_rag_slim/haiku/rag/converters/pdf_split.py @@ -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) diff --git a/tests/test_pdf_split.py b/tests/test_pdf_split.py index 2042b39e..32d98567 100644 --- a/tests/test_pdf_split.py +++ b/tests/test_pdf_split.py @@ -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