perf: lazy page rendering in convert_pdf_to_image

CHANGELOG.md

@@ -1,3 +1,8 @@
+## 1.5.5
+
+### Enhancement
+- Lazy page rendering in `convert_pdf_to_image` to reduce peak memory from O(N pages) to O(1 page)
+
 ## 1.5.4
 
 ### Enhancement

test_unstructured_inference/inference/test_layout.py

@@ -591,6 +591,271 @@ def test_exposed_pdf_image_dpi(pdf_image_dpi, expected, monkeypatch):
         assert mock_from_image.call_args[0][0].height == expected
 
 
+def test_convert_pdf_to_image_no_output_folder():
+    result = layout.convert_pdf_to_image(filename="sample-docs/loremipsum.pdf", dpi=72)
+    assert len(result) == 1
+    assert isinstance(result[0], Image.Image)
+
+
+def test_convert_pdf_to_image_output_folder_returns_images(tmp_path):
+    result = layout.convert_pdf_to_image(
+        filename="sample-docs/loremipsum.pdf",
+        dpi=72,
+        output_folder=tmp_path,
+        path_only=False,
+    )
+    assert len(result) == 1
+    assert isinstance(result[0], Image.Image)
+    saved = list(tmp_path.glob("*.png"))
+    assert len(saved) == 1
+
+
+def test_convert_pdf_to_image_path_only(tmp_path):
+    result = layout.convert_pdf_to_image(
+        filename="sample-docs/loremipsum.pdf",
+        dpi=72,
+        output_folder=tmp_path,
+        path_only=True,
+    )
+    assert len(result) == 1
+    assert all(isinstance(p, str) for p in result)
+    for p in result:
+        assert os.path.exists(p)
+        assert p.endswith(".png")
+    saved = sorted(tmp_path.glob("*.png"))
+    assert [str(s) for s in saved] == sorted(result)
+
+
+def test_convert_pdf_to_image_save_not_under_pdfium_lock(tmp_path):
+    """Verify that PIL save (disk I/O) is NOT performed while holding _pdfium_lock."""
+    original_save = Image.Image.save
+    lock_held_during_save = []
+
+    def spy_save(self, *args, **kwargs):
+        lock_held_during_save.append(layout._pdfium_lock.locked())
+        return original_save(self, *args, **kwargs)
+
+    with patch.object(Image.Image, "save", spy_save):
+        layout.convert_pdf_to_image(
+            filename="sample-docs/loremipsum.pdf",
+            dpi=72,
+            output_folder=tmp_path,
+            path_only=True,
+        )
+    assert lock_held_during_save, "save was never called"
+    assert not any(lock_held_during_save), "pil_image.save() was called while _pdfium_lock was held"
+
+
+def test_convert_pdf_to_image_concurrent_saves_not_serialized(tmp_path):
+    """Two concurrent callers must be able to overlap their disk writes.
+
+    Uses a threading.Barrier to verify both threads are inside save()
+    simultaneously. If saves are serialized under _pdfium_lock, the second
+    thread can never reach save() while the first is there, so the barrier
+    times out and the test fails.
+    """
+    import threading
+
+    original_save = Image.Image.save
+    barrier = threading.Barrier(2, timeout=5)
+    overlap_detected = threading.Event()
+
+    def barrier_save(self, *args, **kwargs):
+        try:
+            barrier.wait()
+            overlap_detected.set()
+        except threading.BrokenBarrierError:
+            pass
+        return original_save(self, *args, **kwargs)
+
+    errors: list[str] = []
+
+    def run(folder):
+        try:
+            layout.convert_pdf_to_image(
+                filename="sample-docs/loremipsum.pdf",
+                dpi=72,
+                output_folder=folder,
+                path_only=True,
+            )
+        except Exception as exc:
+            errors.append(str(exc))
+
+    dir_a = tmp_path / "a"
+    dir_b = tmp_path / "b"
+    dir_a.mkdir()
+    dir_b.mkdir()
+
+    with patch.object(Image.Image, "save", barrier_save):
+        t1 = threading.Thread(target=run, args=(dir_a,))
+        t2 = threading.Thread(target=run, args=(dir_b,))
+        t1.start()
+        t2.start()
+        t1.join(timeout=10)
+        t2.join(timeout=10)
+
+    assert not errors, f"threads raised: {errors}"
+    assert overlap_detected.is_set(), (
+        "saves were serialized under _pdfium_lock — threads could not overlap"
+    )
+    assert list(dir_a.glob("*.png")), "thread A produced no output"
+    assert list(dir_b.glob("*.png")), "thread B produced no output"
+
+
+def test_render_can_proceed_while_other_thread_saves(tmp_path):
+    """Thread B can acquire _pdfium_lock and render while thread A is in save().
+
+    Blocks thread A inside save() (outside the lock), then starts thread B.
+    If B completes entirely while A is still blocked, the lock was not held
+    during save — rendering and saving can overlap across callers.
+    """
+    import threading
+
+    original_save = Image.Image.save
+    a_in_save = threading.Event()
+    b_done = threading.Event()
+
+    dir_a = tmp_path / "a"
+    dir_b = tmp_path / "b"
+    dir_a.mkdir()
+    dir_b.mkdir()
+
+    def gated_save(self, *args, **kwargs):
+        fp = str(args[0]) if args else ""
+        if str(dir_a) in fp:
+            a_in_save.set()
+            b_done.wait(timeout=5)
+        return original_save(self, *args, **kwargs)
+
+    errors: list[str] = []
+
+    def run(folder, done_event=None):
+        try:
+            layout.convert_pdf_to_image(
+                filename="sample-docs/loremipsum.pdf",
+                dpi=72,
+                output_folder=folder,
+                path_only=True,
+            )
+        except Exception as exc:
+            errors.append(str(exc))
+        finally:
+            if done_event:
+                done_event.set()
+
+    with patch.object(Image.Image, "save", gated_save):
+        t_a = threading.Thread(target=run, args=(dir_a,))
+        t_b = threading.Thread(target=run, args=(dir_b, b_done))
+        t_a.start()
+        a_in_save.wait(timeout=5)
+        # A is now blocked in save (outside lock). B should render + save freely.
+        t_b.start()
+        t_b.join(timeout=10)
+        t_a.join(timeout=10)
+
+    assert not errors, f"threads raised: {errors}"
+    assert b_done.is_set(), "Thread B could not complete while A was saving"
+    assert list(dir_a.glob("*.png")), "thread A produced no output"
+    assert list(dir_b.glob("*.png")), "thread B produced no output"
+
+
+def test_multi_page_concurrent_output_complete(tmp_path):
+    """Two threads processing a multi-page PDF both produce correct, complete output."""
+    import threading
+
+    errors: list[str] = []
+
+    def run(folder):
+        try:
+            layout.convert_pdf_to_image(
+                filename="sample-docs/loremipsum_multipage.pdf",
+                dpi=72,
+                output_folder=folder,
+                path_only=True,
+            )
+        except Exception as exc:
+            errors.append(str(exc))
+
+    dir_a = tmp_path / "a"
+    dir_b = tmp_path / "b"
+    dir_a.mkdir()
+    dir_b.mkdir()
+
+    t1 = threading.Thread(target=run, args=(dir_a,))
+    t2 = threading.Thread(target=run, args=(dir_b,))
+    t1.start()
+    t2.start()
+    t1.join(timeout=60)
+    t2.join(timeout=60)
+
+    assert not errors, f"threads raised: {errors}"
+    a_files = sorted(dir_a.glob("*.png"))
+    b_files = sorted(dir_b.glob("*.png"))
+    assert len(a_files) == 10, f"thread A produced {len(a_files)} files, expected 10"
+    assert len(b_files) == 10, f"thread B produced {len(b_files)} files, expected 10"
+    for i in range(1, 11):
+        assert (dir_a / f"page_{i}.png").exists(), f"thread A missing page_{i}.png"
+        assert (dir_b / f"page_{i}.png").exists(), f"thread B missing page_{i}.png"
+
+
+def test_error_in_one_thread_does_not_block_other(tmp_path):
+    """If one thread fails mid-processing, the other still completes."""
+    import threading
+
+    original_save = Image.Image.save
+
+    dir_a = tmp_path / "a"
+    dir_b = tmp_path / "b"
+    dir_a.mkdir()
+    dir_b.mkdir()
+
+    def failing_save(self, *args, **kwargs):
+        fp = str(args[0]) if args else ""
+        if str(dir_a) in fp:
+            raise OSError("simulated disk failure")
+        return original_save(self, *args, **kwargs)
+
+    a_error: list[Exception] = []
+    b_result: list[str] = []
+    b_error: list[Exception] = []
+
+    def run_a():
+        try:
+            layout.convert_pdf_to_image(
+                filename="sample-docs/loremipsum.pdf",
+                dpi=72,
+                output_folder=dir_a,
+                path_only=True,
+            )
+        except Exception as exc:
+            a_error.append(exc)
+
+    def run_b():
+        try:
+            result = layout.convert_pdf_to_image(
+                filename="sample-docs/loremipsum.pdf",
+                dpi=72,
+                output_folder=dir_b,
+                path_only=True,
+            )
+            b_result.extend(result)
+        except Exception as exc:
+            b_error.append(exc)
+
+    with patch.object(Image.Image, "save", failing_save):
+        t_a = threading.Thread(target=run_a)
+        t_b = threading.Thread(target=run_b)
+        t_a.start()
+        t_b.start()
+        t_a.join(timeout=10)
+        t_b.join(timeout=10)
+
+    assert a_error, "Thread A should have failed"
+    assert not b_error, f"Thread B should have succeeded: {b_error}"
+    assert b_result, "Thread B produced no result"
+    assert list(dir_b.glob("*.png")), "Thread B produced no output files"
+
+
 @pytest.mark.parametrize(
     ("filename", "img_num", "should_complete"),
     [

unstructured_inference/__version__.py

@@ -1 +1 @@
-__version__ = "1.5.4"  # pragma: no cover
+__version__ = "1.5.5"  # pragma: no cover

unstructured_inference/inference/layout.py

@@ -426,43 +426,57 @@ def convert_pdf_to_image(
         raise ValueError("output_folder must be specified if path_only is true")
     if filename is None and file is None:
         raise ValueError("Either filename or file must be provided")
-    with _pdfium_lock, pdfium.PdfDocument(filename or file, password=password) as pdf:
+    if output_folder:
+        assert Path(output_folder).exists()
+        assert Path(output_folder).is_dir()
+
+    if dpi is None:
+        dpi = inference_config.PDF_RENDER_DPI
+    scale = dpi / 72.0
+
+    with _pdfium_lock:
+        pdf = pdfium.PdfDocument(filename or file, password=password)
+        n_pages = len(pdf)
+
+    try:
         images: dict[int, Image.Image] = {}
-        if dpi is None:
-            dpi = inference_config.PDF_RENDER_DPI
-        scale = dpi / 72.0
-        for i, page in enumerate(pdf, start=1):
-            try:
-                if first_page is not None and i < first_page:
-                    continue
-                if last_page is not None and i > last_page:
-                    break
-                bitmap = page.render(
-                    scale=scale,
-                    no_smoothtext=False,
-                    no_smoothimage=False,
-                    no_smoothpath=False,
-                    optimize_mode="print",
-                )
+        filenames: list[str] = []
+        for i in range(n_pages):
+            page_num = i + 1
+            if first_page is not None and page_num < first_page:
+                continue
+            if last_page is not None and page_num > last_page:
+                break
+
+            with _pdfium_lock:
+                page = pdf[i]
                 try:
-                    images[i] = bitmap.to_pil()
+                    bitmap = page.render(
+                        scale=scale,
+                        no_smoothtext=False,
+                        no_smoothimage=False,
+                        no_smoothpath=False,
+                        optimize_mode="print",
+                    )
+                    try:
+                        pil_image = bitmap.to_pil()
+                    finally:
+                        bitmap.close()
                 finally:
-                    bitmap.close()
-            finally:
-                page.close()
-
-    if not output_folder:
-        return list(images.values())
-    else:
-        # Save images to output_folder
-        filenames: list[str] = []
-        assert Path(output_folder).exists()
-        assert Path(output_folder).is_dir()
-        for i, image in images.items():
-            fn: str = os.path.join(str(output_folder), f"page_{i}.png")
-            image.save(fn, format="PNG", compress_level=1, optimize=False)
-            filenames.append(fn)
-        if path_only:
-            return filenames
-        images_values: list[Image.Image] = list(images.values())
-        return images_values
+                    page.close()
+
+            if output_folder:
+                fn: str = os.path.join(str(output_folder), f"page_{page_num}.png")
+                pil_image.save(fn, format="PNG", compress_level=1, optimize=False)
+                filenames.append(fn)
+                if not path_only:
+                    images[page_num] = pil_image
+            else:
+                images[page_num] = pil_image
+    finally:
+        with _pdfium_lock:
+            pdf.close()
+
+    if path_only:
+        return filenames
+    return list(images.values())