|
| 1 | +""" |
| 2 | +Page-chunked, multi-process PDF conversion. |
| 3 | +
|
| 4 | +markitdown's ``PdfConverter`` holds the parsed pdfplumber object graph for the |
| 5 | +*whole* document at once (and, for prose PDFs, then re-parses everything with |
| 6 | +pdfminer), so both peak memory and wall-clock grow with document size in a |
| 7 | +single process. |
| 8 | +
|
| 9 | +Splitting the PDF into page ranges and converting each range in its own process |
| 10 | +bounds both: a worker only ever holds N pages of parsed objects, and it returns |
| 11 | +that memory to the OS when it exits. |
| 12 | +
|
| 13 | +Both pdfplumber and pdfminer.six are pure Python, so this has to use processes — |
| 14 | +threads would serialise on the GIL and buy nothing. |
| 15 | +""" |
| 16 | +import logging |
| 17 | +import os |
| 18 | +import shutil |
| 19 | +import tempfile |
| 20 | +from concurrent.futures import ProcessPoolExecutor |
| 21 | +from typing import Optional |
| 22 | + |
| 23 | +logger = logging.getLogger(__name__) |
| 24 | + |
| 25 | +# Pages per chunk. Smaller chunks lower peak memory per worker but add |
| 26 | +# per-chunk parser startup overhead. |
| 27 | +PAGES_PER_CHUNK = int(os.environ.get("PDF_PAGES_PER_CHUNK", "20")) |
| 28 | +# Worker processes. Each holds ~140-190 MiB just for the markitdown import, so |
| 29 | +# this must be sized against the container memory limit, not just its CPU. |
| 30 | +MAX_WORKERS = int(os.environ.get("PDF_MAX_WORKERS", "2")) |
| 31 | +# Below this page count the process startup cost outweighs the parallelism. |
| 32 | +MIN_PAGES_FOR_PARALLEL = int(os.environ.get("PDF_MIN_PAGES_FOR_PARALLEL", "40")) |
| 33 | + |
| 34 | + |
| 35 | +def page_count(path: str) -> Optional[int]: |
| 36 | + """ |
| 37 | + Return the number of pages in ``path``, or None if it is not a readable PDF. |
| 38 | +
|
| 39 | + Used both to decide whether chunking is worthwhile and as the PDF sniff, so |
| 40 | + a corrupt or non-PDF file transparently falls back to the normal path. |
| 41 | + """ |
| 42 | + try: |
| 43 | + from pypdf import PdfReader |
| 44 | + |
| 45 | + return len(PdfReader(path).pages) |
| 46 | + except Exception: |
| 47 | + return None |
| 48 | + |
| 49 | + |
| 50 | +def split_pdf(path: str, pages_per_chunk: int, out_dir: str) -> list[str]: |
| 51 | + """ |
| 52 | + Split ``path`` into single-file PDFs of ``pages_per_chunk`` pages each. |
| 53 | +
|
| 54 | + Chunks are written to ``out_dir`` and passed to workers by path rather than |
| 55 | + by value — pickling the bytes through the process pool would duplicate the |
| 56 | + document in memory, which is the thing this module exists to avoid. |
| 57 | + """ |
| 58 | + from pypdf import PdfReader, PdfWriter |
| 59 | + |
| 60 | + if pages_per_chunk < 1: |
| 61 | + raise ValueError("pages_per_chunk must be >= 1") |
| 62 | + |
| 63 | + reader = PdfReader(path) |
| 64 | + total = len(reader.pages) |
| 65 | + chunk_paths: list[str] = [] |
| 66 | + |
| 67 | + for index, start in enumerate(range(0, total, pages_per_chunk)): |
| 68 | + writer = PdfWriter() |
| 69 | + for page in reader.pages[start:start + pages_per_chunk]: |
| 70 | + writer.add_page(page) |
| 71 | + chunk_path = os.path.join(out_dir, f"chunk_{index:05d}.pdf") |
| 72 | + with open(chunk_path, "wb") as fh: |
| 73 | + writer.write(fh) |
| 74 | + chunk_paths.append(chunk_path) |
| 75 | + |
| 76 | + return chunk_paths |
| 77 | + |
| 78 | + |
| 79 | +def convert_chunk(chunk_path: str) -> str: |
| 80 | + """ |
| 81 | + Convert one chunk file to markdown. Runs in a worker process. |
| 82 | +
|
| 83 | + Must stay importable at module level so the pool can pickle it. markitdown |
| 84 | + is imported lazily here so the parent does not pay for it when the chunked |
| 85 | + path is never taken. |
| 86 | + """ |
| 87 | + from markitdown._stream_info import StreamInfo |
| 88 | + from markitdown.converters._pdf_converter import PdfConverter |
| 89 | + |
| 90 | + stream_info = StreamInfo(extension=".pdf", mimetype="application/pdf") |
| 91 | + with open(chunk_path, "rb") as fh: |
| 92 | + return PdfConverter().convert(fh, stream_info).markdown |
| 93 | + |
| 94 | + |
| 95 | +def convert_pdf_chunked( |
| 96 | + path: str, |
| 97 | + pages_per_chunk: int = PAGES_PER_CHUNK, |
| 98 | + max_workers: int = MAX_WORKERS, |
| 99 | +) -> str: |
| 100 | + """ |
| 101 | + Convert a PDF by splitting it into page ranges converted in parallel. |
| 102 | +
|
| 103 | + Chunk results are concatenated in page order, so output ordering matches a |
| 104 | + whole-document conversion. |
| 105 | + """ |
| 106 | + work_dir = tempfile.mkdtemp(prefix="pdf_chunks_") |
| 107 | + try: |
| 108 | + chunk_paths = split_pdf(path, pages_per_chunk, work_dir) |
| 109 | + logger.info("split %s into %d chunks of <=%d pages", |
| 110 | + path, len(chunk_paths), pages_per_chunk) |
| 111 | + |
| 112 | + if len(chunk_paths) == 1: |
| 113 | + return convert_chunk(chunk_paths[0]) |
| 114 | + |
| 115 | + # "spawn" rather than the Linux default "fork": this runs inside a |
| 116 | + # FastAPI threadpool worker, and forking a process that holds threads |
| 117 | + # can deadlock if a lock is held at fork time. |
| 118 | + import multiprocessing |
| 119 | + |
| 120 | + context = multiprocessing.get_context("spawn") |
| 121 | + workers = max(1, min(max_workers, len(chunk_paths))) |
| 122 | + with ProcessPoolExecutor(max_workers=workers, mp_context=context) as pool: |
| 123 | + # map() preserves input order, so pages stay sequential. |
| 124 | + parts = list(pool.map(convert_chunk, chunk_paths)) |
| 125 | + |
| 126 | + return "\n\n".join(part for part in parts if part and part.strip()) |
| 127 | + finally: |
| 128 | + shutil.rmtree(work_dir, ignore_errors=True) |
| 129 | + |
| 130 | + |
| 131 | +def should_chunk(path: str, min_pages: int = MIN_PAGES_FOR_PARALLEL) -> bool: |
| 132 | + """True if ``path`` is a PDF large enough to be worth chunking.""" |
| 133 | + pages = page_count(path) |
| 134 | + return pages is not None and pages >= min_pages |
0 commit comments