feat: Support instance caching for docling (#11030)

* feat: Support instance caching for docling

Co-Authored-By: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 2/3)

* [autofix.ci] apply automated fixes (attempt 3/3)

* Update test_docling_utils.py

Co-Authored-By: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 2/3)

* [autofix.ci] apply automated fixes (attempt 3/3)

* Update component_index.json

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Eric Hare
2025-12-23 12:25:04 -08:00
committed by GitHub
parent 315f1ab0b6
commit bd29a42b60
4 changed files with 404 additions and 103 deletions

File diff suppressed because one or more lines are too long

View File

@ -3,7 +3,7 @@ import signal
import sys
import traceback
from contextlib import suppress
from typing import TYPE_CHECKING
from functools import lru_cache
from docling_core.types.doc import DoclingDocument
from pydantic import BaseModel, SecretStr, TypeAdapter
@ -12,9 +12,6 @@ from lfx.log.logger import logger
from lfx.schema.data import Data
from lfx.schema.dataframe import DataFrame
if TYPE_CHECKING:
from langchain_core.language_models.chat_models import BaseChatModel
class DoclingDependencyError(Exception):
"""Custom exception for missing Docling dependencies."""
@ -152,6 +149,81 @@ def _deserialize_pydantic_model(data: dict):
return adapter.validate_python(data["config"])
# Global cache for DocumentConverter instances
# This cache persists across multiple runs and thread invocations
@lru_cache(maxsize=4)
def _get_cached_converter(
pipeline: str,
ocr_engine: str,
*,
do_picture_classification: bool,
pic_desc_config_hash: str | None,
):
"""Create and cache a DocumentConverter instance based on configuration.
This function uses LRU caching to maintain DocumentConverter instances in memory,
eliminating the 15-20 minute model loading time on subsequent runs.
Args:
pipeline: The pipeline type ("standard" or "vlm")
ocr_engine: The OCR engine to use
do_picture_classification: Whether to enable picture classification
pic_desc_config_hash: Hash of the picture description config (for cache key)
Returns:
A cached or newly created DocumentConverter instance
"""
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import OcrOptions, PdfPipelineOptions, VlmPipelineOptions
from docling.document_converter import DocumentConverter, FormatOption, PdfFormatOption
from docling.models.factories import get_ocr_factory
from docling.pipeline.vlm_pipeline import VlmPipeline
logger.info(f"Creating DocumentConverter for pipeline={pipeline}, ocr_engine={ocr_engine}")
# Configure the standard PDF pipeline
def _get_standard_opts() -> PdfPipelineOptions:
pipeline_options = PdfPipelineOptions()
pipeline_options.do_ocr = ocr_engine not in {"", "None"}
if pipeline_options.do_ocr:
ocr_factory = get_ocr_factory(
allow_external_plugins=False,
)
ocr_options: OcrOptions = ocr_factory.create_options(
kind=ocr_engine,
)
pipeline_options.ocr_options = ocr_options
pipeline_options.do_picture_classification = do_picture_classification
# Note: pic_desc_config_hash is for cache key only
# Actual picture description is handled separately (non-cached path)
_ = pic_desc_config_hash # Mark as intentionally unused
return pipeline_options
# Configure the VLM pipeline
def _get_vlm_opts() -> VlmPipelineOptions:
return VlmPipelineOptions()
if pipeline == "standard":
pdf_format_option = PdfFormatOption(
pipeline_options=_get_standard_opts(),
)
elif pipeline == "vlm":
pdf_format_option = PdfFormatOption(pipeline_cls=VlmPipeline, pipeline_options=_get_vlm_opts())
else:
msg = f"Unknown pipeline: {pipeline!r}"
raise ValueError(msg)
format_options: dict[InputFormat, FormatOption] = {
InputFormat.PDF: pdf_format_option,
InputFormat.IMAGE: pdf_format_option,
}
return DocumentConverter(format_options=format_options)
def docling_worker(
*,
file_paths: list[str],
@ -162,7 +234,12 @@ def docling_worker(
pic_desc_config: dict | None,
pic_desc_prompt: str,
):
"""Worker function for processing files with Docling in a separate process."""
"""Worker function for processing files with Docling using threading.
This function now uses a globally cached DocumentConverter instance,
significantly reducing processing time on subsequent runs from 15-20 minutes
to just seconds.
"""
# Signal handling for graceful shutdown
shutdown_requested = False
@ -205,12 +282,12 @@ def docling_worker(
check_shutdown()
try:
from docling.datamodel.base_models import ConversionStatus, InputFormat
from docling.datamodel.pipeline_options import OcrOptions, PdfPipelineOptions, VlmPipelineOptions
from docling.document_converter import DocumentConverter, FormatOption, PdfFormatOption
from docling.models.factories import get_ocr_factory
from docling.pipeline.vlm_pipeline import VlmPipeline
from langchain_docling.picture_description import PictureDescriptionLangChainOptions
from docling.datamodel.base_models import ConversionStatus, InputFormat # noqa: F401
from docling.datamodel.pipeline_options import OcrOptions, PdfPipelineOptions, VlmPipelineOptions # noqa: F401
from docling.document_converter import DocumentConverter, FormatOption, PdfFormatOption # noqa: F401
from docling.models.factories import get_ocr_factory # noqa: F401
from docling.pipeline.vlm_pipeline import VlmPipeline # noqa: F401
from langchain_docling.picture_description import PictureDescriptionLangChainOptions # noqa: F401
# Check for shutdown after imports
check_shutdown()
@ -233,27 +310,34 @@ def docling_worker(
queue.put({"error": "Worker interrupted during imports", "shutdown": True})
return
# Configure the standard PDF pipeline
def _get_standard_opts() -> PdfPipelineOptions:
# Use cached converter instead of creating new one each time
# This is the key optimization that eliminates 15-20 minute model load times
def _get_converter() -> DocumentConverter:
check_shutdown() # Check before heavy operations
pipeline_options = PdfPipelineOptions()
pipeline_options.do_ocr = ocr_engine not in {"", "None"}
if pipeline_options.do_ocr:
ocr_factory = get_ocr_factory(
allow_external_plugins=False,
)
ocr_options: OcrOptions = ocr_factory.create_options(
kind=ocr_engine,
)
pipeline_options.ocr_options = ocr_options
pipeline_options.do_picture_classification = do_picture_classification
# For now, we don't support pic_desc_config caching due to serialization complexity
# This is a known limitation that can be addressed in a future enhancement
if pic_desc_config:
pic_desc_llm: BaseChatModel = _deserialize_pydantic_model(pic_desc_config)
logger.warning(
"Picture description with LLM is not yet supported with cached converters. "
"Using non-cached converter for this request."
)
# Fall back to creating a new converter (old behavior)
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions
from docling.document_converter import DocumentConverter, FormatOption, PdfFormatOption
from docling.models.factories import get_ocr_factory
from langchain_docling.picture_description import PictureDescriptionLangChainOptions
pipeline_options = PdfPipelineOptions()
pipeline_options.do_ocr = ocr_engine not in {"", "None"}
if pipeline_options.do_ocr:
ocr_factory = get_ocr_factory(allow_external_plugins=False)
ocr_options = ocr_factory.create_options(kind=ocr_engine)
pipeline_options.ocr_options = ocr_options
pipeline_options.do_picture_classification = do_picture_classification
pic_desc_llm = _deserialize_pydantic_model(pic_desc_config)
logger.info("Docling enabling the picture description stage.")
pipeline_options.do_picture_description = True
pipeline_options.allow_external_plugins = True
@ -261,33 +345,24 @@ def docling_worker(
llm=pic_desc_llm,
prompt=pic_desc_prompt,
)
return pipeline_options
# Configure the VLM pipeline
def _get_vlm_opts() -> VlmPipelineOptions:
check_shutdown() # Check before heavy operations
return VlmPipelineOptions()
pdf_format_option = PdfFormatOption(pipeline_options=pipeline_options)
format_options: dict[InputFormat, FormatOption] = {
InputFormat.PDF: pdf_format_option,
InputFormat.IMAGE: pdf_format_option,
}
return DocumentConverter(format_options=format_options)
# Configure the main format options and create the DocumentConverter()
def _get_converter() -> DocumentConverter:
check_shutdown() # Check before heavy operations
if pipeline == "standard":
pdf_format_option = PdfFormatOption(
pipeline_options=_get_standard_opts(),
)
elif pipeline == "vlm":
pdf_format_option = PdfFormatOption(pipeline_cls=VlmPipeline, pipeline_options=_get_vlm_opts())
else:
msg = f"Unknown pipeline: {pipeline!r}"
raise ValueError(msg)
format_options: dict[InputFormat, FormatOption] = {
InputFormat.PDF: pdf_format_option,
InputFormat.IMAGE: pdf_format_option,
}
return DocumentConverter(format_options=format_options)
# Use cached converter - this is where the magic happens!
# First run: creates and caches converter (15-20 min)
# Subsequent runs: reuses cached converter (seconds)
pic_desc_config_hash = None # Will be None since we checked above
return _get_cached_converter(
pipeline=pipeline,
ocr_engine=ocr_engine,
do_picture_classification=do_picture_classification,
pic_desc_config_hash=pic_desc_config_hash,
)
try:
# Check for shutdown before creating converter (can be slow)

View File

@ -1,6 +1,6 @@
import queue
import threading
import time
from multiprocessing import Queue, get_context
from queue import Empty
from lfx.base.data import BaseFileComponent
from lfx.base.data.docling_utils import _serialize_pydantic_model, docling_worker
@ -92,60 +92,57 @@ class DoclingInlineComponent(BaseFileComponent):
*BaseFileComponent.get_base_outputs(),
]
def _wait_for_result_with_process_monitoring(self, queue: Queue, proc, timeout: int = 300):
"""Wait for result from queue while monitoring process health.
def _wait_for_result_with_thread_monitoring(
self, result_queue: queue.Queue, thread: threading.Thread, timeout: int = 300
):
"""Wait for result from queue while monitoring thread health.
Handles cases where process crashes without sending result.
Handles cases where thread crashes without sending result.
"""
start_time = time.time()
while time.time() - start_time < timeout:
# Check if process is still alive
if not proc.is_alive():
# Process died, try to get any result it might have sent
# Check if thread is still alive
if not thread.is_alive():
# Thread finished, try to get any result it might have sent
try:
result = queue.get_nowait()
except Empty:
# Process died without sending result
msg = f"Worker process crashed unexpectedly without producing result. Exit code: {proc.exitcode}"
result = result_queue.get_nowait()
except queue.Empty:
# Thread finished without sending result
msg = "Worker thread crashed unexpectedly without producing result."
raise RuntimeError(msg) from None
else:
self.log("Process completed and result retrieved")
self.log("Thread completed and result retrieved")
return result
# Poll the queue instead of blocking
try:
result = queue.get(timeout=1)
except Empty:
result = result_queue.get(timeout=1)
except queue.Empty:
# No result yet, continue monitoring
continue
else:
self.log("Result received from worker process")
self.log("Result received from worker thread")
return result
# Overall timeout reached
msg = f"Process timed out after {timeout} seconds"
msg = f"Thread timed out after {timeout} seconds"
raise TimeoutError(msg)
def _terminate_process_gracefully(self, proc, timeout_terminate: int = 10, timeout_kill: int = 5):
"""Terminate process gracefully with escalating signals.
def _stop_thread_gracefully(self, thread: threading.Thread, timeout: int = 10):
"""Wait for thread to complete gracefully.
First tries SIGTERM, then SIGKILL if needed.
Note: Python threads cannot be forcefully killed, so we just wait.
The thread should respond to shutdown signals via the queue.
"""
if not proc.is_alive():
if not thread.is_alive():
return
self.log("Attempting graceful process termination with SIGTERM")
proc.terminate() # Send SIGTERM
proc.join(timeout=timeout_terminate)
self.log("Waiting for thread to complete gracefully")
thread.join(timeout=timeout)
if proc.is_alive():
self.log("Process didn't respond to SIGTERM, using SIGKILL")
proc.kill() # Send SIGKILL
proc.join(timeout=timeout_kill)
if proc.is_alive():
self.log("Warning: Process still alive after SIGKILL")
if thread.is_alive():
self.log("Warning: Thread still alive after timeout")
def process_files(self, file_list: list[BaseFileComponent.BaseFile]) -> list[BaseFileComponent.BaseFile]:
try:
@ -167,44 +164,37 @@ class DoclingInlineComponent(BaseFileComponent):
if self.pic_desc_llm is not None:
pic_desc_config = _serialize_pydantic_model(self.pic_desc_llm)
ctx = get_context("spawn")
queue: Queue = ctx.Queue()
proc = ctx.Process(
# Use threading instead of multiprocessing for memory sharing
# This enables the global DocumentConverter cache to work across runs
result_queue: queue.Queue = queue.Queue()
thread = threading.Thread(
target=docling_worker,
kwargs={
"file_paths": file_paths,
"queue": queue,
"queue": result_queue,
"pipeline": self.pipeline,
"ocr_engine": self.ocr_engine,
"do_picture_classification": self.do_picture_classification,
"pic_desc_config": pic_desc_config,
"pic_desc_prompt": self.pic_desc_prompt,
},
daemon=False, # Allow thread to complete even if main thread exits
)
result = None
proc.start()
thread.start()
try:
result = self._wait_for_result_with_process_monitoring(queue, proc, timeout=300)
result = self._wait_for_result_with_thread_monitoring(result_queue, thread, timeout=300)
except KeyboardInterrupt:
self.log("Docling process cancelled by user")
self.log("Docling thread cancelled by user")
result = []
except Exception as e:
self.log(f"Error during processing: {e}")
raise
finally:
# Improved cleanup with graceful termination
try:
self._terminate_process_gracefully(proc)
finally:
# Always close and cleanup queue resources
try:
queue.close()
queue.join_thread()
except Exception as e: # noqa: BLE001
# Ignore cleanup errors, but log them
self.log(f"Warning: Error during queue cleanup - {e}")
# Wait for thread to complete gracefully
self._stop_thread_gracefully(thread)
# Enhanced error checking with dependency-specific handling
if isinstance(result, dict) and "error" in result:

View File

@ -1,5 +1,8 @@
"""Tests for docling_utils module."""
import time
from unittest.mock import MagicMock, patch
import pytest
try:
@ -134,3 +137,236 @@ class TestExtractDoclingDocuments:
"""Test extracting from None raises error."""
with pytest.raises(TypeError, match="No data inputs provided"):
extract_docling_documents(None, "doc")
class TestDocumentConverterCaching:
"""Test DocumentConverter caching functionality."""
def test_cached_converter_function_exists(self):
"""Test that _get_cached_converter function exists and is properly decorated."""
from lfx.base.data.docling_utils import _get_cached_converter
# Verify function exists
assert callable(_get_cached_converter)
# Verify it has cache_info method (indicates lru_cache decorator)
assert hasattr(_get_cached_converter, "cache_info")
assert callable(_get_cached_converter.cache_info)
def test_cached_converter_cache_key(self):
"""Test that cache uses correct parameters as key."""
from lfx.base.data.docling_utils import _get_cached_converter
# Clear cache before test
_get_cached_converter.cache_clear()
# Mock the DocumentConverter creation to avoid heavy imports
# Patch at import source since DocumentConverter is imported inside _get_cached_converter
with patch("docling.document_converter.DocumentConverter") as mock_converter:
mock_instance1 = MagicMock()
mock_instance2 = MagicMock()
mock_converter.side_effect = [mock_instance1, mock_instance2]
# First call with specific parameters
result1 = _get_cached_converter(
pipeline="standard",
ocr_engine="None",
do_picture_classification=False,
pic_desc_config_hash=None,
)
# Second call with same parameters should return cached result
result2 = _get_cached_converter(
pipeline="standard",
ocr_engine="None",
do_picture_classification=False,
pic_desc_config_hash=None,
)
# Third call with different parameters should create new instance
result3 = _get_cached_converter(
pipeline="vlm",
ocr_engine="None",
do_picture_classification=False,
pic_desc_config_hash=None,
)
# Verify caching behavior
assert result1 is result2, "Same parameters should return cached instance"
assert result1 is not result3, "Different parameters should return new instance"
# Verify DocumentConverter was only called twice (not three times)
assert mock_converter.call_count == 2
# Verify cache statistics
cache_info = _get_cached_converter.cache_info()
assert cache_info.hits >= 1, "Should have at least one cache hit"
assert cache_info.misses == 2, "Should have exactly two cache misses"
def test_cached_converter_lru_eviction(self):
"""Test that LRU cache properly evicts old entries when maxsize is reached."""
from lfx.base.data.docling_utils import _get_cached_converter
# Clear cache before test
_get_cached_converter.cache_clear()
# Patch at import source since DocumentConverter is imported inside _get_cached_converter
with patch("docling.document_converter.DocumentConverter") as mock_converter:
mock_instances = [MagicMock() for _ in range(5)]
mock_converter.side_effect = mock_instances
# Create 5 different cache entries (maxsize=4, so one should be evicted)
configs = [
("standard", "None", False, None),
("standard", "easyocr", False, None),
("vlm", "None", False, None),
("standard", "None", True, None),
("vlm", "easyocr", False, None),
]
for pipeline, ocr, pic_class, pic_hash in configs:
_get_cached_converter(
pipeline=pipeline,
ocr_engine=ocr,
do_picture_classification=pic_class,
pic_desc_config_hash=pic_hash,
)
# Cache size should be at most 4 (maxsize)
cache_info = _get_cached_converter.cache_info()
assert cache_info.currsize <= 4, "Cache size should not exceed maxsize"
def test_cached_converter_performance_improvement(self):
"""Test that caching provides performance improvement."""
from lfx.base.data.docling_utils import _get_cached_converter
# Clear cache before test
_get_cached_converter.cache_clear()
# Patch at import source since DocumentConverter is imported inside _get_cached_converter
with patch("docling.document_converter.DocumentConverter") as mock_converter:
# Simulate slow converter creation
def slow_creation(*args, **kwargs): # noqa: ARG001
time.sleep(0.05) # 50ms delay
return MagicMock()
mock_converter.side_effect = slow_creation
# First call (cache miss - should be slow)
start_time = time.time()
_get_cached_converter(
pipeline="standard",
ocr_engine="None",
do_picture_classification=False,
pic_desc_config_hash=None,
)
first_call_duration = time.time() - start_time
# Second call (cache hit - should be fast)
start_time = time.time()
_get_cached_converter(
pipeline="standard",
ocr_engine="None",
do_picture_classification=False,
pic_desc_config_hash=None,
)
second_call_duration = time.time() - start_time
# Cache hit should be significantly faster (at least 10x)
assert second_call_duration < first_call_duration / 10, (
f"Cache hit should be much faster: first={first_call_duration:.4f}s, second={second_call_duration:.4f}s"
)
def test_cache_clear(self):
"""Test that cache can be cleared."""
from lfx.base.data.docling_utils import _get_cached_converter
# Clear cache
_get_cached_converter.cache_clear()
# Patch at import source since DocumentConverter is imported inside _get_cached_converter
with patch("docling.document_converter.DocumentConverter"):
# Add something to cache
_get_cached_converter(
pipeline="standard",
ocr_engine="None",
do_picture_classification=False,
pic_desc_config_hash=None,
)
# Verify cache has content
cache_info = _get_cached_converter.cache_info()
assert cache_info.currsize > 0
# Clear cache
_get_cached_converter.cache_clear()
# Verify cache is empty
cache_info = _get_cached_converter.cache_info()
assert cache_info.currsize == 0
assert cache_info.hits == 0
assert cache_info.misses == 0
def test_different_ocr_engines_create_different_caches(self):
"""Test that different OCR engines result in different cached converters."""
from lfx.base.data.docling_utils import _get_cached_converter
_get_cached_converter.cache_clear()
# Patch at import source since DocumentConverter is imported inside _get_cached_converter
with patch("docling.document_converter.DocumentConverter") as mock_converter:
mock_instance1 = MagicMock()
mock_instance2 = MagicMock()
mock_converter.side_effect = [mock_instance1, mock_instance2]
# Create converter with no OCR
result1 = _get_cached_converter(
pipeline="standard",
ocr_engine="None",
do_picture_classification=False,
pic_desc_config_hash=None,
)
# Create converter with EasyOCR
result2 = _get_cached_converter(
pipeline="standard",
ocr_engine="easyocr",
do_picture_classification=False,
pic_desc_config_hash=None,
)
# Should be different instances
assert result1 is not result2
assert mock_converter.call_count == 2
def test_different_pipelines_create_different_caches(self):
"""Test that different pipelines result in different cached converters."""
from lfx.base.data.docling_utils import _get_cached_converter
_get_cached_converter.cache_clear()
# Patch at import source since DocumentConverter is imported inside _get_cached_converter
with patch("docling.document_converter.DocumentConverter") as mock_converter:
mock_instance1 = MagicMock()
mock_instance2 = MagicMock()
mock_converter.side_effect = [mock_instance1, mock_instance2]
# Create converter with standard pipeline
result1 = _get_cached_converter(
pipeline="standard",
ocr_engine="None",
do_picture_classification=False,
pic_desc_config_hash=None,
)
# Create converter with VLM pipeline
result2 = _get_cached_converter(
pipeline="vlm",
ocr_engine="None",
do_picture_classification=False,
pic_desc_config_hash=None,
)
# Should be different instances
assert result1 is not result2
assert mock_converter.call_count == 2