From 6bcc2f6357eccc60a8096df67e751385dad62ad3 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 25 May 2026 11:53:01 +0300 Subject: [PATCH] PDF split-convert-merge for memory-bound large PDFs --- docs/configuration/processing.md | 46 ++++++++ docs/remote-processing.md | 24 ++++ haiku_rag_slim/haiku/rag/client/processing.py | 19 ++- haiku_rag_slim/haiku/rag/config/models.py | 10 ++ .../haiku/rag/converters/pdf_split.py | 101 ++++++++++++++++ haiku_rag_slim/pyproject.toml | 1 + tests/test_pdf_split.py | 111 ++++++++++++++++++ uv.lock | 2 + 8 files changed, 311 insertions(+), 3 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/converters/pdf_split.py create mode 100644 tests/test_pdf_split.py diff --git a/docs/configuration/processing.md b/docs/configuration/processing.md index 576bce5a..871a8323 100644 --- a/docs/configuration/processing.md +++ b/docs/configuration/processing.md @@ -111,6 +111,52 @@ you can eyeball the ratio. Conversion options work identically for both local and remote processing. +### Large PDFs and docling memory + +Docling's parser is memory-hungry and has confirmed leaks in current versions +([docling #2209](https://github.com/docling-project/docling/issues/2209), +[#1343](https://github.com/docling-project/docling/issues/1343), +[#2954](https://github.com/docling-project/docling/issues/2954); +[docling-serve #366](https://github.com/docling-project/docling-serve/issues/366), +[#474](https://github.com/docling-project/docling-serve/issues/474)). +Single-pass conversion of 400-page PDFs can OOM a workstation in local mode, +and long-running docling-serve containers see RSS grow monotonically. + +Mitigation in haiku.rag — set `processing.split_pages`: + +```yaml +processing: + split_pages: 10 # 0 disables (default) +``` + +When `split_pages > 0`, PDFs are split at the byte level into N-page slices +(using pypdfium2, already bundled), each slice converted independently, then +merged back via `DoclingDocument.concatenate` — preserving page numbers and +re-indexing `self_ref` values across slices. Peak memory per conversion is +bounded by one slice's working set rather than the whole document; in +docling-serve mode each slice is also an independent task that lets the +server release task-local state between requests. + +Recommendation: `10` is a sensible starting point for any consistently-large +PDF workload. Smaller slices reduce peak memory but multiply task overhead +(per-slice docling startup + HTTP round-trips for docling-serve). Cross-page +references (named destinations, multi-page link annotations) are dropped at +the split — accepted loss; haiku.rag doesn't surface them downstream. + +**Operational note for long-running ingest**: even with `split_pages`, +docling's per-process leak rate is non-zero. For deployments running +continuously: + +- *docling-serve mode*: set `mem_limit` on the container in Compose + (or `resources.limits.memory` in Kubernetes) plus `restart: unless-stopped` + so the kernel OOM-kills and the runtime restarts. Run multiple + docling-serve replicas behind the round-robin `base_url` list above so a + restart of one doesn't stop ingest. +- *docling-local mode*: the leak is inside the `haiku-ingester` process + itself. Apply the same `mem_limit` + restart policy to the ingester + container. Restarts are graceful — in-flight jobs land in the queue's + reaper window and resume on next start. + **Note:** When using `chunker: docling-serve`, OCR options (`do_ocr`, `force_ocr`, `ocr_engine`, `ocr_lang`) from `conversion_options` are passed to the chunking API. This is useful when running docling-serve in a read-only container where OCR model downloads fail. Set `do_ocr: false` to disable OCR entirely. ### Conversion Options diff --git a/docs/remote-processing.md b/docs/remote-processing.md index be8a89b2..1e84414d 100644 --- a/docs/remote-processing.md +++ b/docs/remote-processing.md @@ -215,6 +215,30 @@ providers: haiku-rag add-src document.pdf ``` +## Operational notes + +Long-running docling-serve containers see CPU memory grow monotonically +([docling-serve #366](https://github.com/docling-project/docling-serve/issues/366), +[#474](https://github.com/docling-project/docling-serve/issues/474)). The +underlying parser leaks are in core docling +([#2209](https://github.com/docling-project/docling/issues/2209), +[#1343](https://github.com/docling-project/docling/issues/1343)) and affect +docling-local too. + +Recommended deployment shape: + +- Set `mem_limit` on the docling-serve container (or `resources.limits.memory` + in Kubernetes) at a value comfortably above your largest expected job. +- Combine with `restart: unless-stopped` so the runtime restarts when the + kernel OOM-kills. +- Run multiple docling-serve replicas behind haiku.rag's round-robin + `providers.docling_serve.base_url` list (see + [Document Processing](configuration/processing.md)). A restart of one + replica doesn't stop ingest. +- In haiku.rag, set `processing.split_pages` for large-PDF workloads so each + slice is an independent docling-serve task and the per-task working set + stays bounded. + ## Resources - [docling-serve GitHub](https://github.com/docling-project/docling-serve) diff --git a/haiku_rag_slim/haiku/rag/client/processing.py b/haiku_rag_slim/haiku/rag/client/processing.py index 64133bda..0b8583ca 100644 --- a/haiku_rag_slim/haiku/rag/client/processing.py +++ b/haiku_rag_slim/haiku/rag/client/processing.py @@ -84,6 +84,19 @@ async def convert( """ converter = get_converter(config) + async def _convert_file( + file_path: Path, effective_uri: str | None + ) -> "DoclingDocument": + """Dispatch through split-and-merge for large PDFs when configured, + otherwise call the converter directly.""" + if file_path.suffix.lower() == ".pdf" and config.processing.split_pages > 0: + from haiku.rag.converters.pdf_split import convert_pdf_with_splitting + + return await convert_pdf_with_splitting( + converter, file_path, effective_uri, config.processing.split_pages + ) + return await converter.convert_file(file_path, source_uri=effective_uri) + # Path object - convert file directly if isinstance(source, Path): if not source.exists(): @@ -91,7 +104,7 @@ async def convert( if source.suffix.lower() not in converter.supported_extensions: raise ValueError(f"Unsupported file extension: {source.suffix}") effective_uri = source_uri or source.absolute().as_uri() - doc = await converter.convert_file(source, source_uri=effective_uri) + doc = await _convert_file(source, effective_uri) _warn_if_descriptions_missing(config, doc, str(source)) return doc @@ -123,7 +136,7 @@ async def convert( try: effective_uri = source_uri or source - doc = await converter.convert_file(temp_path, source_uri=effective_uri) + doc = await _convert_file(temp_path, effective_uri) _warn_if_descriptions_missing(config, doc, source) return doc finally: @@ -137,7 +150,7 @@ async def convert( if file_path.suffix.lower() not in converter.supported_extensions: raise ValueError(f"Unsupported file extension: {file_path.suffix}") effective_uri = source_uri or file_path.absolute().as_uri() - doc = await converter.convert_file(file_path, source_uri=effective_uri) + doc = await _convert_file(file_path, effective_uri) _warn_if_descriptions_missing(config, doc, str(file_path)) return doc diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 631305b8..0816bfed 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -157,6 +157,16 @@ class ProcessingConfig(BaseModel): chunking_merge_peers: bool = True chunking_use_markdown_tables: bool = False conversion_options: ConversionOptions = Field(default_factory=ConversionOptions) + split_pages: int = Field( + default=0, + ge=0, + description=( + "If >0, PDFs are split into N-page slices, each converted " + "independently, and merged via DoclingDocument.concatenate. 0 " + "disables splitting (single-pass conversion). Recommended: 10 " + "for memory-bound or large (>100 page) PDFs." + ), + ) pictures: PicturesMode = "image" """How embedded pictures are handled at ingest. diff --git a/haiku_rag_slim/haiku/rag/converters/pdf_split.py b/haiku_rag_slim/haiku/rag/converters/pdf_split.py new file mode 100644 index 00000000..4c415c0a --- /dev/null +++ b/haiku_rag_slim/haiku/rag/converters/pdf_split.py @@ -0,0 +1,101 @@ +"""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 +from collections.abc import Iterator +from pathlib import Path +from typing import TYPE_CHECKING + +import pypdfium2 as pdfium + +if TYPE_CHECKING: + from docling_core.types.doc.document import DoclingDocument + + from haiku.rag.converters.base import DocumentConverter + + +def iter_pdf_slices(path: Path, slice_size: int) -> Iterator[tuple[int, int, bytes]]: + """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. + """ + if slice_size <= 0: + raise ValueError(f"slice_size must be >= 1, got {slice_size}") + src = pdfium.PdfDocument(str(path)) + try: + total = len(src) + for start in range(0, total, slice_size): + 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( + converter: "DocumentConverter", + path: Path, + source_uri: str | None, + slice_size: int, +) -> "DoclingDocument": + """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. + """ + from docling_core.types.doc.document import DoclingDocument + + slices = await asyncio.to_thread(lambda: list(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: + 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) + + return DoclingDocument.concatenate(converted) diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 345b7e78..5ebd0dd4 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -32,6 +32,7 @@ dependencies = [ "pydantic>=2.12.5", "pydantic-ai-slim[openai,fastmcp,logfire,ag-ui]>=1.100.0", "pydantic-monty>=0.0.17", + "pypdfium2>=5.0", "python-dotenv>=1.2.2", "pyyaml>=6.0.3", "rich>=14.3.3", diff --git a/tests/test_pdf_split.py b/tests/test_pdf_split.py new file mode 100644 index 00000000..2042b39e --- /dev/null +++ b/tests/test_pdf_split.py @@ -0,0 +1,111 @@ +"""Unit tests for haiku.rag.converters.pdf_split. + +The integration test that pins split-and-merge against a real-PDF baseline +lives in tests/test_converters.py — that path requires docling installed and +is gated accordingly. These tests cover the byte-level split mechanism and +the docling-core concatenate contract in isolation. +""" + +import io +from pathlib import Path + +import pypdfium2 as pdfium +import pytest + +from haiku.rag.converters.pdf_split import iter_pdf_slices + + +def _make_pdf(page_count: int, tmp_path: Path) -> Path: + """Synthesize a minimal valid multi-page PDF with pypdfium2. Each page is + an A4-sized blank. Returns the path.""" + doc = pdfium.PdfDocument.new() + try: + for _ in range(page_count): + doc.new_page(width=595.0, height=842.0) # A4 in points + out = tmp_path / "synth.pdf" + with open(out, "wb") as f: + doc.save(f) + return out + finally: + doc.close() + + +def test_iter_pdf_slices_partitions_pages(tmp_path): + src = _make_pdf(7, tmp_path) + slices = list(iter_pdf_slices(src, slice_size=3)) + + assert [(s, e) for s, e, _ in slices] == [(1, 3), (4, 6), (7, 7)] + + page_counts = [] + for _, _, pdf_bytes in slices: + d = pdfium.PdfDocument(io.BytesIO(pdf_bytes)) + try: + page_counts.append(len(d)) + finally: + d.close() + assert page_counts == [3, 3, 1] + assert sum(page_counts) == 7 + + +def test_iter_pdf_slices_single_slice_when_doc_fits(tmp_path): + src = _make_pdf(4, tmp_path) + slices = list(iter_pdf_slices(src, slice_size=10)) + assert len(slices) == 1 + start, end, pdf_bytes = slices[0] + assert (start, end) == (1, 4) + d = pdfium.PdfDocument(io.BytesIO(pdf_bytes)) + try: + assert len(d) == 4 + finally: + d.close() + + +def test_iter_pdf_slices_rejects_zero_slice_size(tmp_path): + src = _make_pdf(2, tmp_path) + with pytest.raises(ValueError, match="slice_size must be >= 1"): + list(iter_pdf_slices(src, slice_size=0)) + + +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 + and self_refs across both stay unique.""" + pytest.importorskip("docling_core") + from docling_core.types.doc.base import BoundingBox, CoordOrigin + from docling_core.types.doc.document import ( + DoclingDocument, + PageItem, + ProvenanceItem, + Size, + ) + from docling_core.types.doc.labels import DocItemLabel + + def _make_one_page_doc(name: str, text: str) -> DoclingDocument: + d = DoclingDocument(name=name) + d.pages[1] = PageItem(page_no=1, size=Size(width=595.0, height=842.0)) + d.add_text( + label=DocItemLabel.TEXT, + text=text, + prov=ProvenanceItem( + page_no=1, + bbox=BoundingBox( + l=0.0, t=0.0, r=100.0, b=20.0, coord_origin=CoordOrigin.TOPLEFT + ), + charspan=(0, len(text)), + ), + ) + return d + + a = _make_one_page_doc("a", "alpha") + b = _make_one_page_doc("b", "beta") + + merged = DoclingDocument.concatenate([a, b]) + + assert len(merged.texts) == 2 + refs = [t.self_ref for t in merged.texts] + assert len(set(refs)) == 2, f"self_refs collided: {refs}" + + page_nos = sorted({p.page_no for t in merged.texts for p in t.prov}) + assert page_nos == [1, 2], f"expected b's page 1 to shift to page 2, got {page_nos}" + + assert sorted(merged.pages.keys()) == [1, 2] diff --git a/uv.lock b/uv.lock index 16a9aa4c..b002a066 100644 --- a/uv.lock +++ b/uv.lock @@ -1571,6 +1571,7 @@ dependencies = [ { name = "pydantic" }, { name = "pydantic-ai-slim", extra = ["ag-ui", "fastmcp", "logfire", "openai"] }, { name = "pydantic-monty" }, + { name = "pypdfium2" }, { name = "python-dotenv" }, { name = "pyyaml" }, { name = "rich" }, @@ -1667,6 +1668,7 @@ requires-dist = [ { name = "pydantic-ai-slim", extras = ["vertexai"], marker = "extra == 'vertexai'" }, { name = "pydantic-ai-slim", extras = ["voyageai"], marker = "extra == 'voyageai'" }, { name = "pydantic-monty", specifier = ">=0.0.17" }, + { name = "pypdfium2", specifier = ">=5.0" }, { name = "python-dotenv", specifier = ">=1.2.2" }, { name = "pyyaml", specifier = ">=6.0.3" }, { name = "rich", specifier = ">=14.3.3" },