From fe954a090e0baccd26afb9fd9e1a06b4b2544d0d Mon Sep 17 00:00:00 2001 From: bryan davis Date: Thu, 18 Jun 2026 16:35:48 -0500 Subject: [PATCH] move cpu bound actions off of main loop --- .../haiku/rag/chunkers/docling_serve.py | 11 +++-- .../haiku/rag/converters/docling_serve.py | 6 ++- .../haiku/rag/converters/pdf_split.py | 5 ++- .../haiku/rag/ingester/sources/fs.py | 23 +++++++--- tests/ingester/test_fs_source.py | 29 +++++++++++++ tests/test_chunker.py | 38 +++++++++++++++++ tests/test_converters.py | 34 +++++++++++++++ tests/test_pdf_split.py | 42 +++++++++++++++++++ 8 files changed, 177 insertions(+), 11 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py index 3cf4dc03..c0380134 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py @@ -1,3 +1,4 @@ +import asyncio import re from io import BytesIO from typing import TYPE_CHECKING @@ -99,9 +100,13 @@ class DoclingServeChunker(DocumentChunker): else: endpoint = "/v1/chunk/hybrid/file/async" - # Export document to JSON - doc_json = document.model_dump_json() - doc_bytes = doc_json.encode("utf-8") + # Export document to JSON off the event loop. model_dump_json over a + # document carrying inlined base64 page/picture images is CPU-heavy and + # proportional to document size; running it inline would block every + # other worker's coroutine for the duration of the serialization. + doc_bytes = await asyncio.to_thread( + lambda: document.model_dump_json().encode("utf-8") + ) # Prepare multipart request with DoclingDocument JSON files = {"files": ("document.json", BytesIO(doc_bytes), "application/json")} diff --git a/haiku_rag_slim/haiku/rag/converters/docling_serve.py b/haiku_rag_slim/haiku/rag/converters/docling_serve.py index b07366f4..0a0e2668 100644 --- a/haiku_rag_slim/haiku/rag/converters/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/converters/docling_serve.py @@ -225,7 +225,11 @@ class DoclingServeConverter(DocumentConverter): data=data, name=name, ) - return self._parse_zip_to_docling(zip_bytes, name) + # Parse off the event loop: the zip decompress, per-image base64 + # re-encoding, and DoclingDocument.model_validate are all synchronous + # and CPU-heavy (full-resolution page rasters when generate_page_images + # is on), so running inline would stall every other worker's coroutine. + return await asyncio.to_thread(self._parse_zip_to_docling, zip_bytes, name) async def convert_file( self, path: Path, source_uri: str | None = None diff --git a/haiku_rag_slim/haiku/rag/converters/pdf_split.py b/haiku_rag_slim/haiku/rag/converters/pdf_split.py index 4057b60e..02d6a04b 100644 --- a/haiku_rag_slim/haiku/rag/converters/pdf_split.py +++ b/haiku_rag_slim/haiku/rag/converters/pdf_split.py @@ -136,4 +136,7 @@ async def convert_pdf_with_splitting( # Off the event loop because the close path acquires the lock. await asyncio.to_thread(it.close) - return DoclingDocument.concatenate(converted) + # Merge off the event loop: concatenating slice documents that carry + # inlined base64 page/picture images is CPU-heavy and proportional to the + # total document size, so running it inline would block other coroutines. + return await asyncio.to_thread(DoclingDocument.concatenate, converted) diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py index 1cf6a6c4..57b2b23d 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py @@ -1,3 +1,4 @@ +import asyncio import hashlib import mimetypes import os @@ -87,23 +88,33 @@ class FSSource: return None return str(path.stat().st_mtime_ns) + def _read_body(self, path: Path, uri: str) -> tuple[bytes, str, str]: + """Size-check, read, and hash the file. Runs in a worker thread (see + ``fetch``) because the read and the md5 are both proportional to file + size and would otherwise block the event loop for the whole read.""" + check_file_size(path.stat().st_size, self._max_file_size, uri) + body = path.read_bytes() + content_hash = hashlib.md5(body, usedforsecurity=False).hexdigest() + # mtime_ns rather than st_mtime: nanosecond integer avoids float + # precision collisions on rapid edits. + revision = str(path.stat().st_mtime_ns) + return body, content_hash, revision + async def fetch(self, uri: str) -> FetchResult: path = self._resolve_within_root(uri) if path is None: raise UnsupportedSourceError(f"Path escapes FS root ({self.root}): {uri}") - check_file_size(path.stat().st_size, self._max_file_size, uri) - body = path.read_bytes() + body, content_hash, revision = await asyncio.to_thread( + self._read_body, path, uri + ) content_type, _ = mimetypes.guess_type(path.name) if content_type is None: content_type = "application/octet-stream" - # mtime_ns rather than st_mtime: nanosecond integer avoids float - # precision collisions on rapid edits. - revision = str(path.stat().st_mtime_ns) return FetchResult( uri=path.as_uri(), body=body, content_type=content_type, - content_hash=hashlib.md5(body, usedforsecurity=False).hexdigest(), + content_hash=content_hash, revision=revision, disk_path=path, ) diff --git a/tests/ingester/test_fs_source.py b/tests/ingester/test_fs_source.py index 10f57829..73c7cf34 100644 --- a/tests/ingester/test_fs_source.py +++ b/tests/ingester/test_fs_source.py @@ -300,3 +300,32 @@ async def test_fs_source_fetch_no_limit_when_max_size_is_none(fs_root: Path): src = FSSource(root=fs_root, max_file_size=None) result = await src.fetch((fs_root / "a.md").as_uri()) assert result.body == b"alpha" + + +@pytest.mark.asyncio +async def test_fs_source_fetch_reads_off_event_loop_thread(fs_root: Path): + """The file read and md5 are both proportional to file size and must run + off the event-loop thread, or a large file would freeze every other + worker's coroutine for the duration of the read. Capture the thread the + read+hash runs on and assert it is not the main thread.""" + import threading + + src = FSSource(root=fs_root) + target = fs_root / "a.md" + + called_from: list[threading.Thread] = [] + original = src._read_body + + def spy(path, uri): + called_from.append(threading.current_thread()) + return original(path, uri) + + src._read_body = spy # type: ignore[method-assign] + + result = await src.fetch(target.as_uri()) + assert result.body == b"alpha" + assert called_from, "_read_body was never called" + assert called_from[0] is not threading.main_thread(), ( + "FSSource._read_body ran on the event-loop thread; the read+hash must " + "be dispatched via asyncio.to_thread" + ) diff --git a/tests/test_chunker.py b/tests/test_chunker.py index 1ad0b770..9b897fc2 100644 --- a/tests/test_chunker.py +++ b/tests/test_chunker.py @@ -609,6 +609,44 @@ This is content. assert meta1.headings == ["Chapter 1", "Section 1.1"] assert meta1.page_numbers == [1, 2] + @pytest.mark.asyncio + @patch("haiku.rag.providers.docling_serve.httpx.AsyncClient") + async def test_chunk_serializes_document_off_event_loop_thread( + self, mock_client_class, chunker + ): + """model_dump_json over a document carrying inlined base64 page/picture + images is CPU-heavy and proportional to document size; it must run off + the event-loop thread or it stalls every other worker's coroutine. + + A minimal fake document records the thread its model_dump_json runs on; + the API response carries no doc_items so the document is touched only + for serialization.""" + import threading + + result_data = {"chunks": [{"text": "Chunk 1", "chunk_index": 0}]} + submit_resp, poll_resp, result_resp = create_async_workflow_mocks(result_data) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=submit_resp) + mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp]) + mock_client_class.return_value.__aenter__.return_value = mock_client + + called_from: list[threading.Thread] = [] + + class FakeDoc: + def model_dump_json(self): + called_from.append(threading.current_thread()) + return "{}" + + chunks = await chunker.chunk(FakeDoc()) + + assert len(chunks) == 1 + assert called_from, "model_dump_json was never called" + assert called_from[0] is not threading.main_thread(), ( + "DoclingDocument.model_dump_json ran on the event-loop thread; it " + "must be dispatched via asyncio.to_thread" + ) + @pytest.mark.vcr() @pytest.mark.asyncio diff --git a/tests/test_converters.py b/tests/test_converters.py index 8081edd3..916307e0 100644 --- a/tests/test_converters.py +++ b/tests/test_converters.py @@ -94,6 +94,40 @@ def create_async_workflow_zip_mocks( return submit_response, poll_response, result_response +@pytest.mark.asyncio +async def test_parse_zip_runs_off_event_loop_thread(): + """_parse_zip_to_docling does zip decompress, per-image base64 re-encoding, + and DoclingDocument.model_validate — all synchronous and CPU-heavy (full- + resolution page rasters when generate_page_images is on). It must run off + the event-loop thread, or it stalls every other worker's coroutine. Capture + the thread it runs on and assert it is not the main thread.""" + import threading + + config = AppConfig() + config.processing.converter = "docling-serve" + converter = get_converter(config) + assert isinstance(converter, DoclingServeConverter) + + converter.client.submit_and_poll_zip = AsyncMock(return_value=b"zip-bytes") + + called_from: list[threading.Thread] = [] + + def spy(zip_bytes, name): + called_from.append(threading.current_thread()) + return Mock() + + converter._parse_zip_to_docling = spy # type: ignore[method-assign] + + files = {"files": ("doc.pdf", b"pdf", "application/octet-stream")} + await converter._make_request(files, "doc.pdf") + + assert called_from, "_parse_zip_to_docling was never called" + assert called_from[0] is not threading.main_thread(), ( + "_parse_zip_to_docling ran on the event-loop thread; it must be " + "dispatched via asyncio.to_thread" + ) + + class TestTextFileHandler: """Tests for TextFileHandler utility class.""" diff --git a/tests/test_pdf_split.py b/tests/test_pdf_split.py index accb1fd9..7c253f9a 100644 --- a/tests/test_pdf_split.py +++ b/tests/test_pdf_split.py @@ -157,6 +157,48 @@ async def test_convert_aborts_and_cleans_up_on_mid_stream_slice_failure( assert len(calls) == 2 +@pytest.mark.asyncio +async def test_concatenate_runs_off_event_loop_thread(tmp_path, monkeypatch): + """DoclingDocument.concatenate merges slice documents that carry inlined + base64 page/picture images — CPU-heavy and proportional to total document + size. It must run off the event-loop thread so it doesn't stall other + workers' coroutines. Capture the thread it runs on and assert it is not the + main thread.""" + import threading + + from docling_core.types.doc.document import DoclingDocument + + src = _make_pdf(4, tmp_path) + + class _Converter: + async def convert_file(self, path: Path, *, source_uri): + return DoclingDocument(name="slice") + + called_from: list[threading.Thread] = [] + + def spy(docs): + called_from.append(threading.current_thread()) + # Return a slice doc rather than exercising the real concatenate — + # this test only asserts the dispatch thread, not merge correctness + # (covered by test_concatenate_shifts_page_nos_and_unique_self_refs). + return docs[0] + + monkeypatch.setattr(DoclingDocument, "concatenate", staticmethod(spy)) + + await convert_pdf_with_splitting( + _Converter(), # ty: ignore[invalid-argument-type] + src, + source_uri=None, + slice_size=2, + ) + + assert called_from, "concatenate was never called" + assert called_from[0] is not threading.main_thread(), ( + "DoclingDocument.concatenate ran on the event-loop thread; it must be " + "dispatched via asyncio.to_thread" + ) + + def test_concatenate_shifts_page_nos_and_unique_self_refs(): """Pins the docling-core contract we rely on: when two docs (each with items on page 1) are concatenated, the second doc's items move to page 2