Bound pdf_split memory + release lock between slices

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)
+# Multiple workers calling iter_pdf_slices concurrently race on that state
-# race on that state and corrupt it — first error surfaces as e.g.
+# and corrupt it — first error surfaces as e.g. "Failed to import pages",
-# "Failed to import pages", and after that every subsequent PDF load fails
+# and after that every subsequent PDF load fails with "Failed to load
-# with "Failed to load document" until the process restarts. Hold this lock
+# document" until the process restarts. Hold this lock around every pdfium
-# around every pdfium call so only one worker's slicing runs at a time. The
+# call so only one worker's slicing runs at a time. Crucially we release it
-# actual heavy work (docling-serve convert) happens AFTER iter_pdf_slices
+# between slices so other workers' slicing can interleave; the heavy work
-# returns and continues to parallelize across workers.
+# (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)
+    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()
+            yield (start + 1, end, slice_bytes)
     finally:
+        with _PDFIUM_LOCK:
             src.close()
 
 
@@ -93,26 +79,33 @@ 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``
+    Slices are produced lazily — only one slice's bytes live in memory at a
-    so the ingester pipeline's classifier maps it to ``TransientError`` and
+    time, so peak working set stays bounded regardless of page count. A
-    the queue retries the entire document. Per-slice retry would risk
+    single slice failure aborts the whole job (raised as ``ValueError`` so
-    interleaved partial state with subsequent runs and is not worth the
+    the ingester pipeline classifies it as ``TransientError`` and the queue
-    complexity.
+    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:
+    try:
-        with tempfile.NamedTemporaryFile(mode="wb", suffix=".pdf", delete=False) as tmp:
+        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:
-            # 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,
@@ -130,5 +123,10 @@ async def convert_pdf_with_splitting(
                 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)

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