NeoAnd/Docling
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: enable locks for threadsafe pdfium (#1052)

Browse Source 
* enable locks for threadsafe pdfium

Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>

* fix deadlock in pypdfium2 backend

Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>

---------

Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>
This commit is contained in:
Michele Dolfi 2025-03-02 20:06:44 +01:00 committed by GitHub
parent e25d557c06
commit 8dc0562542
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 98 additions and 71 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

68
docling/backend/docling_parse_v2_backend.py
View File

@ -12,6 +12,7 @@ from pypdfium2 import PdfPage
from docling.backend.pdf_backend import PdfDocumentBackend, PdfPageBackend from docling.backend.pdf_backend import PdfDocumentBackend, PdfPageBackend
from docling.datamodel.base_models import Cell, Size from docling.datamodel.base_models import Cell, Size
from docling.utils.locks import pypdfium2_lock
if TYPE_CHECKING: if TYPE_CHECKING:
from docling.datamodel.document import InputDocument from docling.datamodel.document import InputDocument
@ -182,20 +183,24 @@ class DoclingParseV2PageBackend(PdfPageBackend):
padbox.r = page_size.width - padbox.r padbox.r = page_size.width - padbox.r
padbox.t = page_size.height - padbox.t padbox.t = page_size.height - padbox.t
image = ( with pypdfium2_lock:
self._ppage.render( image = (
scale=scale * 1.5, self._ppage.render(
rotation=0, # no additional rotation scale=scale * 1.5,
crop=padbox.as_tuple(), rotation=0, # no additional rotation
) crop=padbox.as_tuple(),
.to_pil() )
.resize(size=(round(cropbox.width * scale), round(cropbox.height * scale))) .to_pil()
) # We resize the image from 1.5x the given scale to make it sharper. .resize(
size=(round(cropbox.width * scale), round(cropbox.height * scale))
)
) # We resize the image from 1.5x the given scale to make it sharper.
return image return image
def get_size(self) -> Size: def get_size(self) -> Size:
return Size(width=self._ppage.get_width(), height=self._ppage.get_height()) with pypdfium2_lock:
return Size(width=self._ppage.get_width(), height=self._ppage.get_height())
def unload(self): def unload(self):
self._ppage = None self._ppage = None
@ -206,23 +211,24 @@ class DoclingParseV2DocumentBackend(PdfDocumentBackend):
def __init__(self, in_doc: "InputDocument", path_or_stream: Union[BytesIO, Path]): def __init__(self, in_doc: "InputDocument", path_or_stream: Union[BytesIO, Path]):
super().__init__(in_doc, path_or_stream) super().__init__(in_doc, path_or_stream)
self._pdoc = pdfium.PdfDocument(self.path_or_stream) with pypdfium2_lock:
self.parser = pdf_parser_v2("fatal") self._pdoc = pdfium.PdfDocument(self.path_or_stream)
self.parser = pdf_parser_v2("fatal")
success = False success = False
if isinstance(self.path_or_stream, BytesIO): if isinstance(self.path_or_stream, BytesIO):
success = self.parser.load_document_from_bytesio( success = self.parser.load_document_from_bytesio(
self.document_hash, self.path_or_stream self.document_hash, self.path_or_stream
) )
elif isinstance(self.path_or_stream, Path): elif isinstance(self.path_or_stream, Path):
success = self.parser.load_document( success = self.parser.load_document(
self.document_hash, str(self.path_or_stream) self.document_hash, str(self.path_or_stream)
) )
if not success: if not success:
raise RuntimeError( raise RuntimeError(
f"docling-parse v2 could not load document {self.document_hash}." f"docling-parse v2 could not load document {self.document_hash}."
) )
def page_count(self) -> int: def page_count(self) -> int:
# return len(self._pdoc) # To be replaced with docling-parse API # return len(self._pdoc) # To be replaced with docling-parse API
@ -236,9 +242,10 @@ class DoclingParseV2DocumentBackend(PdfDocumentBackend):
return len_2 return len_2
def load_page(self, page_no: int) -> DoclingParseV2PageBackend: def load_page(self, page_no: int) -> DoclingParseV2PageBackend:
return DoclingParseV2PageBackend( with pypdfium2_lock:
self.parser, self.document_hash, page_no, self._pdoc[page_no] return DoclingParseV2PageBackend(
) self.parser, self.document_hash, page_no, self._pdoc[page_no]
)
def is_valid(self) -> bool: def is_valid(self) -> bool:
return self.page_count() > 0 return self.page_count() > 0
@ -246,5 +253,6 @@ class DoclingParseV2DocumentBackend(PdfDocumentBackend):
def unload(self): def unload(self):
super().unload() super().unload()
self.parser.unload_document(self.document_hash) self.parser.unload_document(self.document_hash)
self._pdoc.close() with pypdfium2_lock:
self._pdoc = None self._pdoc.close()
self._pdoc = None

98
docling/backend/pypdfium2_backend.py
View File

@ -13,6 +13,7 @@ from pypdfium2._helpers.misc import PdfiumError
from docling.backend.pdf_backend import PdfDocumentBackend, PdfPageBackend from docling.backend.pdf_backend import PdfDocumentBackend, PdfPageBackend
from docling.datamodel.base_models import Cell from docling.datamodel.base_models import Cell
from docling.utils.locks import pypdfium2_lock
if TYPE_CHECKING: if TYPE_CHECKING:
from docling.datamodel.document import InputDocument from docling.datamodel.document import InputDocument
@ -24,6 +25,7 @@ class PyPdfiumPageBackend(PdfPageBackend):
def __init__( def __init__(
self, pdfium_doc: pdfium.PdfDocument, document_hash: str, page_no: int self, pdfium_doc: pdfium.PdfDocument, document_hash: str, page_no: int
): ):
# Note: lock applied by the caller
self.valid = True # No better way to tell from pypdfium. self.valid = True # No better way to tell from pypdfium.
try: try:
self._ppage: pdfium.PdfPage = pdfium_doc[page_no] self._ppage: pdfium.PdfPage = pdfium_doc[page_no]
@ -40,51 +42,57 @@ class PyPdfiumPageBackend(PdfPageBackend):
def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]: def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]:
AREA_THRESHOLD = 0 # 32 * 32 AREA_THRESHOLD = 0 # 32 * 32
for obj in self._ppage.get_objects(filter=[pdfium_c.FPDF_PAGEOBJ_IMAGE]): page_size = self.get_size()
pos = obj.get_pos() with pypdfium2_lock:
cropbox = BoundingBox.from_tuple( for obj in self._ppage.get_objects(filter=[pdfium_c.FPDF_PAGEOBJ_IMAGE]):
pos, origin=CoordOrigin.BOTTOMLEFT pos = obj.get_pos()
).to_top_left_origin(page_height=self.get_size().height) cropbox = BoundingBox.from_tuple(
pos, origin=CoordOrigin.BOTTOMLEFT
).to_top_left_origin(page_height=page_size.height)
if cropbox.area() > AREA_THRESHOLD: if cropbox.area() > AREA_THRESHOLD:
cropbox = cropbox.scaled(scale=scale) cropbox = cropbox.scaled(scale=scale)
yield cropbox yield cropbox
def get_text_in_rect(self, bbox: BoundingBox) -> str: def get_text_in_rect(self, bbox: BoundingBox) -> str:
if not self.text_page: with pypdfium2_lock:
self.text_page = self._ppage.get_textpage() if not self.text_page:
self.text_page = self._ppage.get_textpage()
if bbox.coord_origin != CoordOrigin.BOTTOMLEFT: if bbox.coord_origin != CoordOrigin.BOTTOMLEFT:
bbox = bbox.to_bottom_left_origin(self.get_size().height) bbox = bbox.to_bottom_left_origin(self.get_size().height)
text_piece = self.text_page.get_text_bounded(*bbox.as_tuple()) with pypdfium2_lock:
text_piece = self.text_page.get_text_bounded(*bbox.as_tuple())
return text_piece return text_piece
def get_text_cells(self) -> Iterable[Cell]: def get_text_cells(self) -> Iterable[Cell]:
if not self.text_page: with pypdfium2_lock:
self.text_page = self._ppage.get_textpage() if not self.text_page:
self.text_page = self._ppage.get_textpage()
cells = [] cells = []
cell_counter = 0 cell_counter = 0
page_size = self.get_size() page_size = self.get_size()
for i in range(self.text_page.count_rects()): with pypdfium2_lock:
rect = self.text_page.get_rect(i) for i in range(self.text_page.count_rects()):
text_piece = self.text_page.get_text_bounded(*rect) rect = self.text_page.get_rect(i)
x0, y0, x1, y1 = rect text_piece = self.text_page.get_text_bounded(*rect)
cells.append( x0, y0, x1, y1 = rect
Cell( cells.append(
id=cell_counter, Cell(
text=text_piece, id=cell_counter,
bbox=BoundingBox( text=text_piece,
l=x0, b=y0, r=x1, t=y1, coord_origin=CoordOrigin.BOTTOMLEFT bbox=BoundingBox(
).to_top_left_origin(page_size.height), l=x0, b=y0, r=x1, t=y1, coord_origin=CoordOrigin.BOTTOMLEFT
).to_top_left_origin(page_size.height),
)
) )
) cell_counter += 1
cell_counter += 1
# PyPdfium2 produces very fragmented cells, with sub-word level boundaries, in many PDFs. # PyPdfium2 produces very fragmented cells, with sub-word level boundaries, in many PDFs.
# The cell merging code below is to clean this up. # The cell merging code below is to clean this up.
@ -214,20 +222,24 @@ class PyPdfiumPageBackend(PdfPageBackend):
padbox.r = page_size.width - padbox.r padbox.r = page_size.width - padbox.r
padbox.t = page_size.height - padbox.t padbox.t = page_size.height - padbox.t
image = ( with pypdfium2_lock:
self._ppage.render( image = (
scale=scale * 1.5, self._ppage.render(
rotation=0, # no additional rotation scale=scale * 1.5,
crop=padbox.as_tuple(), rotation=0, # no additional rotation
) crop=padbox.as_tuple(),
.to_pil() )
.resize(size=(round(cropbox.width * scale), round(cropbox.height * scale))) .to_pil()
) # We resize the image from 1.5x the given scale to make it sharper. .resize(
size=(round(cropbox.width * scale), round(cropbox.height * scale))
)
) # We resize the image from 1.5x the given scale to make it sharper.
return image return image
def get_size(self) -> Size: def get_size(self) -> Size:
return Size(width=self._ppage.get_width(), height=self._ppage.get_height()) with pypdfium2_lock:
return Size(width=self._ppage.get_width(), height=self._ppage.get_height())
def unload(self): def unload(self):
self._ppage = None self._ppage = None
@ -239,22 +251,26 @@ class PyPdfiumDocumentBackend(PdfDocumentBackend):
super().__init__(in_doc, path_or_stream) super().__init__(in_doc, path_or_stream)
try: try:
self._pdoc = pdfium.PdfDocument(self.path_or_stream) with pypdfium2_lock:
self._pdoc = pdfium.PdfDocument(self.path_or_stream)
except PdfiumError as e: except PdfiumError as e:
raise RuntimeError( raise RuntimeError(
f"pypdfium could not load document with hash {self.document_hash}" f"pypdfium could not load document with hash {self.document_hash}"
) from e ) from e
def page_count(self) -> int: def page_count(self) -> int:
return len(self._pdoc) with pypdfium2_lock:
return len(self._pdoc)
def load_page(self, page_no: int) -> PyPdfiumPageBackend: def load_page(self, page_no: int) -> PyPdfiumPageBackend:
return PyPdfiumPageBackend(self._pdoc, self.document_hash, page_no) with pypdfium2_lock:
return PyPdfiumPageBackend(self._pdoc, self.document_hash, page_no)
def is_valid(self) -> bool: def is_valid(self) -> bool:
return self.page_count() > 0 return self.page_count() > 0
def unload(self): def unload(self):
super().unload() super().unload()
self._pdoc.close() with pypdfium2_lock:
self._pdoc = None self._pdoc.close()
self._pdoc = None

3
docling/utils/locks.py Normal file
View File

@ -0,0 +1,3 @@
import threading
pypdfium2_lock = threading.Lock()