From 1d95f58171c6c84cce56ccd7079a5467df9e4596 Mon Sep 17 00:00:00 2001
From: suguanYang <wangbinqi77@gmail.com>
Date: Tue, 8 Sep 2026 12:06:59 +0800
Subject: [PATCH] perf: isolate scanned PDF OCR

---
 .../document_agent/tools/ocr_pages.py         | 52 +++++++++++++------
 .../tests/contract/test_ocr_pages_contract.py | 28 ++++++----
 2 files changed, 53 insertions(+), 27 deletions(-)

diff --git a/apps/worker/app/services/document_agent/tools/ocr_pages.py b/apps/worker/app/services/document_agent/tools/ocr_pages.py
index 01ff59d1..80565118 100644
--- a/apps/worker/app/services/document_agent/tools/ocr_pages.py
+++ b/apps/worker/app/services/document_agent/tools/ocr_pages.py
@@ -7,6 +7,10 @@
 
 from app.services.document_agent.manifest import ToolContext, ToolResult
 from app.services.document_agent.pdf_text import PageTextBands
+from app.services.document_parser.formats.pdf.pymupdf_subprocess import (
+    run_in_child_process,
+    worker,
+)
 from app.services.document_agent.registry import has_page_features, register_tool
 from app.services.document_agent.visual import render_pages
 
@@ -32,6 +36,28 @@ def _line_score(item: Any) -> float:
     return 0.0
 
 
+@worker
+def _run_ocr_worker(queue: Any, page_paths: dict[int, str]) -> None:
+    """Run the local OCR model outside the heartbeat-bearing worker process."""
+    from rapidocr_onnxruntime import RapidOCR
+
+    engine = RapidOCR(intra_op_num_threads=2, inter_op_num_threads=1)
+    page_lines: dict[int, list[dict[str, Any]]] = {}
+    for page, image_path in page_paths.items():
+        lines: list[dict[str, Any]] = []
+        result, _elapse = engine(image_path)
+        for item in result or []:
+            lines.append(
+                {
+                    "box": _line_box(item),
+                    "text": _line_text(item),
+                    "score": _line_score(item),
+                }
+            )
+        page_lines[page] = lines
+    queue.put({"ok": True, "page_lines": page_lines})
+
+
 @register_tool(
     name="ocr.pages",
     description="Run RapidOCR on specified pages and return positioned text lines.",
@@ -70,26 +96,18 @@ def ocr_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
         if item.get("page") is not None and item.get("png_path")
     }
 
-    from rapidocr_onnxruntime import RapidOCR
-
-    engine = RapidOCR()
+    page_paths = {
+        page: image_path for page, image_path in png_by_page.items() if page in pages
+    }
+    result = run_in_child_process(_run_ocr_worker, page_paths, timeout=300)
+    page_lines = {
+        int(page): list(lines)
+        for page, lines in (result.get("page_lines") or {}).items()
+    }
     page_texts: dict[int, str] = {}
     page_bands: dict[int, PageTextBands] = {}
-    page_lines: dict[int, list[dict[str, Any]]] = {}
     for page in pages:
-        image_path = png_by_page.get(page)
-        lines: list[dict[str, Any]] = []
-        if image_path:
-            result, _elapse = engine(image_path)
-            for item in result or []:
-                text = _line_text(item)
-                lines.append(
-                    {
-                        "box": _line_box(item),
-                        "text": text,
-                        "score": _line_score(item),
-                    }
-                )
+        lines = page_lines.get(page, [])
         page_lines[page] = lines
         content = "\n".join(line["text"] for line in lines if line["text"])
         page_texts[page] = content
diff --git a/apps/worker/tests/contract/test_ocr_pages_contract.py b/apps/worker/tests/contract/test_ocr_pages_contract.py
index efb28a8c..1fbcde9c 100644
--- a/apps/worker/tests/contract/test_ocr_pages_contract.py
+++ b/apps/worker/tests/contract/test_ocr_pages_contract.py
@@ -3,8 +3,6 @@
 from __future__ import annotations
 
 import os
-import sys
-from types import ModuleType
 from unittest.mock import patch
 
 os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test")
@@ -55,19 +53,29 @@ def test_ocr_pages_requires_pages() -> None:
 def test_ocr_pages_writes_joined_text_to_blackboard() -> None:
     ctx = _ctx()
 
-    class FakeEngine:
-        def __call__(self, _image_path: str):
-            return [[[[0, 0], [1, 0], [1, 1], [0, 1]], "Hello", 0.9]], 0.01
-
-    fake_mod = ModuleType("rapidocr_onnxruntime")
-    fake_mod.RapidOCR = lambda: FakeEngine()  # type: ignore[attr-defined]
-
     def fake_render(*_args, **_kwargs):
         return [{"page": 1, "png_path": "/tmp/ocr_page_1.png"}]
 
+    def fake_child(worker_fn, page_paths, *, timeout):
+        assert worker_fn.__name__ == "_run_ocr_worker"
+        assert page_paths == {1: "/tmp/ocr_page_1.png"}
+        assert timeout == 300
+        return {
+            "ok": True,
+            "page_lines": {
+                1: [
+                    {
+                        "box": [[0, 0], [1, 0], [1, 1], [0, 1]],
+                        "text": "Hello",
+                        "score": 0.9,
+                    }
+                ]
+            },
+        }
+
     with (
         patch.dict(ocr_pages.__globals__, {"render_pages": fake_render}),
-        patch.dict(sys.modules, {"rapidocr_onnxruntime": fake_mod}),
+        patch.dict(ocr_pages.__globals__, {"run_in_child_process": fake_child}),
     ):
         result = ocr_pages(ctx, {"pages": [1]})