diff --git a/conf/config.yaml b/conf/config.yaml index 8179cb9be..a0c2ec1a4 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -232,6 +232,13 @@ loader: docling_max_task_retry: 3 docling_retry_base_delay: 2.0 + # Env: PARSE_TIMEOUT + # Outer wall-clock bound (seconds) for a single file's parse stage. marker and + # docling self-limit at marker_timeout/docling_timeout; pymupdf has no internal + # timeout, so this prevents a wedged pymupdf parse from stalling indexing — the + # bad file fails and is reported instead. + parse_timeout: 3600 + # Env: TRANSCRIBER_BASE_URL, TRANSCRIBER_API_KEY, TRANSCRIBER_MODEL, # TRANSCRIBER_TIMEOUT, TRANSCRIBER_MAX_CONCURRENT_CHUNKS, USE_WHISPER_LANG_DETECTOR, # TRANSCRIBER_DIRECT_UPLOAD_SUFFIXES (pipe-separated, e.g. ".mp3|.mp4|.wav") diff --git a/openrag/api/routers/admin/presets.py b/openrag/api/routers/admin/presets.py index b279147cb..5893536d1 100644 --- a/openrag/api/routers/admin/presets.py +++ b/openrag/api/routers/admin/presets.py @@ -1,7 +1,5 @@ """Admin routes for the Phase 14 pipeline preset registry.""" -from typing import get_args - from api.dependencies.auth import require_admin from api.schemas.admin.preset_schemas import ( CreatePresetRequest, @@ -11,7 +9,7 @@ UpdatePresetRequest, ) from core.chunking import chunking_registry -from core.config.indexation_pipeline import IndexationPipelineConfig +from core.config.indexation_pipeline import PARSING_STRATEGIES from core.rerankers.registry import reranker_registry from core.retrieval import retriever_registry from di.providers import get_preset_service @@ -21,9 +19,11 @@ _DEFAULT_RERANKER_PROVIDERS = ["infinity", "openai"] -# Derived from the validated Literal so the exposed options can never drift from -# what IndexationPipelineConfig actually accepts. -_PARSING_STRATEGIES = list(get_args(IndexationPipelineConfig.model_fields["parsing_strategy"].annotation)) +# The selectable PDF backends. Shares the IndexationPipelineConfig constant so +# the exposed options can never drift from what the model accepts. ``None`` +# (inherit the global PDFLoader) is the field default, not an explicit choice, +# so it is not surfaced here. +_PARSING_STRATEGIES = list(PARSING_STRATEGIES) def _registered_or_default(registered: list[str], defaults: list[str]) -> list[str]: diff --git a/openrag/core/config/indexation.py b/openrag/core/config/indexation.py index 0469ca5ea..152145f14 100644 --- a/openrag/core/config/indexation.py +++ b/openrag/core/config/indexation.py @@ -178,6 +178,14 @@ class LoaderConfig(ConfigMixin): docling_timeout: int = 3600 docling_max_task_retry: int = 3 docling_retry_base_delay: float = 2.0 + # Outer wall-clock bound (seconds) for a single file's parse stage. Unlike + # marker/docling — which self-limit at marker_timeout/docling_timeout — the + # pymupdf backend has no internal timeout, so without this a wedged pymupdf + # parse would stall indexing indefinitely. Bounds any backend (and any future + # unbounded one); a slow/wedged parse fails *that* file instead of hanging. + # Must be > 0: it feeds asyncio.wait_for, so 0/negative would fail every + # parse immediately rather than disable the bound. + parse_timeout: int = Field(default=3600, gt=0) transcriber: TranscriberConfig = Field(default_factory=TranscriberConfig) openai: OpenAILoaderConfig = Field(default_factory=OpenAILoaderConfig) # Max depth of nested .eml-in-.eml attachments the EmlLoader will descend diff --git a/openrag/core/config/indexation_pipeline.py b/openrag/core/config/indexation_pipeline.py index 793137c5c..ad0e98079 100644 --- a/openrag/core/config/indexation_pipeline.py +++ b/openrag/core/config/indexation_pipeline.py @@ -12,6 +12,13 @@ from core.config.chunking import ChunkerConfig from pydantic import BaseModel, ConfigDict, Field +# PDF parsing backends a preset may explicitly select. ``None`` (the default) +# means "inherit the deployment's global PDFLoader" — a preset that doesn't care +# about PDF parsing follows the operator's global ``file_loaders.pdf`` choice +# instead of silently forcing one backend (and, with it, lazily spinning up that +# backend's Ray pool). See pipeline_builder._select_parser. +PARSING_STRATEGIES: tuple[str, ...] = ("pymupdf", "marker", "docling") + class IndexationPipelineConfig(BaseModel): """Indexation pipeline settings for one partition preset.""" @@ -19,7 +26,8 @@ class IndexationPipelineConfig(BaseModel): model_config = ConfigDict(extra="ignore") chunking: ChunkerConfig = Field(default_factory=ChunkerConfig) - parsing_strategy: Literal["pymupdf", "marker", "docling"] = "marker" + # None => inherit the global PDFLoader (see PARSING_STRATEGIES above). + parsing_strategy: Literal["pymupdf", "marker", "docling"] | None = None # VLM / image captioning vlm: str | None = None # endpoint name; None = use global default @@ -50,4 +58,4 @@ class IndexationPipelineConfig(BaseModel): topic_tagging_llm: str | None = None -__all__ = ["IndexationPipelineConfig"] +__all__ = ["IndexationPipelineConfig", "PARSING_STRATEGIES"] diff --git a/openrag/core/config/loader.py b/openrag/core/config/loader.py index c5d8beb61..a800afae2 100644 --- a/openrag/core/config/loader.py +++ b/openrag/core/config/loader.py @@ -106,6 +106,7 @@ ("DOCLING_NUM_GPUS", "loader.docling_num_gpus", float), ("DOCLING_POOL_SIZE", "loader.docling_pool_size", int), ("DOCLING_MAX_TASKS_PER_WORKER", "loader.docling_max_tasks_per_worker", int), + ("PARSE_TIMEOUT", "loader.parse_timeout", int), ("WHISPER_MODEL", "loader.local_whisper.model", str), ("WHISPER_N_WORKERS", "loader.local_whisper.whisper_n_workers", int), ("WHISPER_NUM_GPUS", "loader.local_whisper.whisper_num_gpus", float), diff --git a/openrag/core/indexing/parsers/pdf/pymupdf.py b/openrag/core/indexing/parsers/pdf/pymupdf.py index f42b33755..503518aba 100644 --- a/openrag/core/indexing/parsers/pdf/pymupdf.py +++ b/openrag/core/indexing/parsers/pdf/pymupdf.py @@ -31,7 +31,6 @@ import pymupdf4llm from ....models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock -from ...image_preprocessor import extract_data_uri_image_blocks from ..document_parser import DocumentParser from ..registry import parser_registry @@ -48,29 +47,25 @@ def _extract_text(raw: bytes) -> tuple[list[str], list[ImageBlock]]: def _extract_markdown(raw: bytes) -> tuple[list[str], list[ImageBlock]]: - """Return Markdown per page + ``ImageBlock``s built from embedded data URIs. - - ``embed_images=True`` makes ``pymupdf4llm`` write images as base64 - data URIs in-line. We decode each ref into an ``ImageBlock`` and - leave the ref in the page text untouched so the caption stage can - substitute later via ``ImageBlock.metadata['markdown_ref']``. + """Return structured Markdown per page (no images). + + pymupdf is the lightweight, no-VLM backend. ``pymupdf4llm`` preserves + document structure (headings, lists, tables) — which the markdown-aware + chunker needs to cut on real boundaries instead of mid-sentence — while + ``embed_images=False`` keeps base64 image data out of the text. That keeps + chunks small (no Milvus gRPC overflow) and skips image rendering entirely + (fast). Image-aware parsing is marker/docling's job, so no ``ImageBlock``s + are produced here. """ with pymupdf.open(stream=raw, filetype="pdf") as doc: chunks = pymupdf4llm.to_markdown( doc, page_chunks=True, - embed_images=True, + embed_images=False, write_images=False, - dpi=300, ) - pages: list[str] = [] - images: list[ImageBlock] = [] - for i, chunk in enumerate(chunks, start=1): - text = (chunk.get("text") or "").strip() - pages.append(text) - if text: - images.extend(extract_data_uri_image_blocks(text, page_number=i)) - return pages, images + pages = [(chunk.get("text") or "").strip() for chunk in chunks] + return pages, [] @parser_registry.register("pymupdf") diff --git a/openrag/services/orchestrators/preset_service.py b/openrag/services/orchestrators/preset_service.py index 095f1cc17..f33825e53 100644 --- a/openrag/services/orchestrators/preset_service.py +++ b/openrag/services/orchestrators/preset_service.py @@ -28,15 +28,21 @@ _VALID_PRESET_TYPES = frozenset({"indexation", "retrieval"}) _DEFAULT_SEEDS: dict[str, dict[str, dict[str, Any]]] = { - # ``enable_contextualization`` is intentionally omitted from the ``default`` - # indexation preset below — it is injected at seed time from the global - # ``chunker.contextual_retrieval`` (``CONTEXTUAL_RETRIEVAL``) toggle in - # _finalize_seed, so the env flag still drives contextualization on fresh - # deployments. Named presets (legal/finance) keep their explicit choice. + # ``enable_contextualization`` and ``parsing_strategy`` are intentionally + # omitted from the ``default`` indexation preset below. The former is injected + # at seed time from the global ``chunker.contextual_retrieval`` + # (``CONTEXTUAL_RETRIEVAL``) toggle in _finalize_seed; the latter stays unset + # so the default preset inherits the global ``file_loaders.pdf`` (``PDFLoader``) + # backend at parse time. Named presets (legal/finance) keep their explicit choice. "indexation": { "default": { "chunking": {"name": "recursive_splitter", "chunk_size": 512, "chunk_overlap_rate": 0.2}, - "parsing_strategy": "marker", + # ``parsing_strategy`` is intentionally omitted so the default preset + # inherits the deployment's global PDFLoader (``file_loaders.pdf``) + # rather than forcing one PDF backend on every partition. A hardcoded + # value here would override the operator's global choice and lazily + # spin up that backend's Ray pool — e.g. forcing marker on a + # pymupdf-configured deployment. Named presets below opt in explicitly. "enable_image_captioning": True, "enable_entity_extraction": True, "enable_topic_tagging": True, diff --git a/openrag/services/workers/bootstrap.py b/openrag/services/workers/bootstrap.py index 529c524b2..7e9cc513a 100644 --- a/openrag/services/workers/bootstrap.py +++ b/openrag/services/workers/bootstrap.py @@ -51,10 +51,10 @@ def wrapper(name, cls, namespace="openrag", remote_args=(), **options): def get_or_create_actor(name, cls, namespace="openrag", remote_args=(), **options): # ``get_if_exists`` makes get-or-create atomic inside Ray: concurrent callers # (e.g. the multiple Ray Serve replicas that all run worker bootstrap in their - # lifespan at once) either all attach to the same actor or exactly one creates - # it and the rest attach. The previous get_actor()-then-create pattern had a - # TOCTOU race — two replicas both saw "not found" and then collided on - # creation, crashing a replica's startup with "Failed to look up actor". + # lifespan at once, or a per-preset backend selecting the same pool) either all + # attach to the same actor or exactly one creates it and the rest attach. The + # previous get_actor()-then-create pattern had a TOCTOU race — two callers both + # saw "not found" and collided on creation, crashing with "Failed to look up actor". return cls.options(name=name, namespace=namespace, get_if_exists=True, **options).remote(*remote_args) diff --git a/openrag/services/workers/indexer_pool.py b/openrag/services/workers/indexer_pool.py index f82643612..2361bef8c 100644 --- a/openrag/services/workers/indexer_pool.py +++ b/openrag/services/workers/indexer_pool.py @@ -35,6 +35,7 @@ def __init__(self) -> None: cfg = load_config() parser = build_parser_dispatcher(cfg) + parser_factory = _build_parser_factory(parser) vlm = build_caption_vlm(cfg) chunker = _build_chunker(cfg) embedder_factory = _build_embedder_factory(cfg) @@ -62,7 +63,9 @@ def __init__(self) -> None: vector_store=self._vector_store, vlm=vlm, image_captioning=cfg.loader.image_captioning, + timeouts=_build_pipeline_timeouts(cfg), chunker_factory=_build_chunker_from_config, + parser_factory=parser_factory, embedder_factory=embedder_factory, vlm_factory=vlm_factory, contextualizer_factory=contextualizer_factory, @@ -332,6 +335,39 @@ def _build_chunker_from_config(chunker_config: Any) -> Any: return _build_chunker(SimpleNamespace(chunker=chunker_config)) +def _build_parser_factory(parser: Any) -> Any: + """Factory honoring a preset's ``parsing_strategy`` for PDFs. + + Without this, the pipeline falls back to the single global-config dispatcher + and every PDF uses the global default loader — silently ignoring a preset's + pymupdf/docling choice. Each per-strategy wrapper reuses ``parser``'s shared + backend cache, so selecting a strategy never builds a duplicate + marker/docling Ray pool. + """ + cache: dict[str, Any] = {} + + def factory(strategy: str = "marker") -> Any: + wrapper = cache.get(strategy) + if wrapper is None: + wrapper = parser.for_pdf_strategy(strategy) + cache[strategy] = wrapper + return wrapper + + return factory + + +def _build_pipeline_timeouts(cfg: Settings) -> Any: + """Per-stage timeouts for the indexing pipeline. + + Bounds the parse stage at ``loader.parse_timeout`` so a wedged parse (notably + pymupdf, which has no internal timeout) fails that file instead of stalling + indexing. Other stages stay unbounded here (their backends self-limit). See #571. + """ + from services.workers.pipeline_builder import PipelineTimeouts + + return PipelineTimeouts(parse=cfg.loader.parse_timeout) + + def _build_embedder_factory(cfg: Settings) -> Any: from core.embeddings import embedder_registry diff --git a/openrag/services/workers/parsers/docling_workers.py b/openrag/services/workers/parsers/docling_workers.py index 1716516b4..9edddff93 100644 --- a/openrag/services/workers/parsers/docling_workers.py +++ b/openrag/services/workers/parsers/docling_workers.py @@ -120,7 +120,12 @@ class DoclingLoader(BasePooledParser): def __init__(self) -> None: self.config = load_config() - self.worker: DoclingPool = ray.get_actor("DoclingPool", namespace="openrag") + # Create the pool lazily if bootstrap didn't: bootstrap only pre-warms the + # globally-configured PDF backend, but a per-preset parsing_strategy can + # select docling even when the global default is marker (see #569/#575). + from services.workers.bootstrap import get_or_create_actor + + self.worker: DoclingPool = get_or_create_actor("DoclingPool", DoclingPool, lifetime="detached") def supported_types(self) -> list[str]: return [DocumentType.PDF.value] diff --git a/openrag/services/workers/parsers/marker_workers.py b/openrag/services/workers/parsers/marker_workers.py index b08ca8a48..085800450 100644 --- a/openrag/services/workers/parsers/marker_workers.py +++ b/openrag/services/workers/parsers/marker_workers.py @@ -361,7 +361,12 @@ class MarkerLoader(BasePooledParser): def __init__(self) -> None: self.config = load_config() - self.worker = ray.get_actor("MarkerPool", namespace="openrag") + # Lazily create the pool if bootstrap didn't (it only pre-warms the + # globally-configured PDF backend; a preset can select marker even when + # the global default is docling — see #569/#575). + from services.workers.bootstrap import get_or_create_actor + + self.worker = get_or_create_actor("MarkerPool", MarkerPool, lifetime="detached") def supported_types(self) -> list[str]: return [DocumentType.PDF.value] diff --git a/openrag/services/workers/parsers/parser_dispatcher.py b/openrag/services/workers/parsers/parser_dispatcher.py index 713c4bd63..92e194d93 100644 --- a/openrag/services/workers/parsers/parser_dispatcher.py +++ b/openrag/services/workers/parsers/parser_dispatcher.py @@ -38,6 +38,10 @@ "OpenAIAudioLoader": "audio_client", } +# Backend names a preset ``parsing_strategy`` may select for PDFs (the values of +# _PDF_BACKENDS). Used to validate a strategy before routing a PDF to it. +_PDF_BACKEND_NAMES: frozenset[str] = frozenset(_PDF_BACKENDS.values()) + # Attachment ext (lowercased, no dot) → DocumentType, used to wire the EML # parser's per-attachment sub-parsers. _MAX_EML_ATTACHMENT_DEPTH = 3 @@ -84,6 +88,21 @@ async def parse(self, document: Document) -> ProcessedDocument: parser = self._get(backend) return await parser.parse(document) + def for_pdf_strategy(self, strategy: str) -> DocumentParser: + """Return a parser that forces ``strategy`` (a PDF backend name such as + ``"pymupdf"`` or ``"docling"``) for PDF documents while dispatching every + other content type exactly as this dispatcher does. + + Backends come from this dispatcher's shared cache, so honoring a + per-preset ``parsing_strategy`` never spins up a duplicate marker/docling + pool. + """ + if strategy not in _PDF_BACKEND_NAMES: + raise ValueError( + f"Unsupported PDF parsing strategy {strategy!r}; expected one of {sorted(_PDF_BACKEND_NAMES)}" + ) + return _PdfStrategyParser(self, strategy) + # ----- backend resolution ----- def _resolve_backend(self, content_type: DocumentType, ext: str) -> str: @@ -202,10 +221,32 @@ async def language_detector(path): # noqa: E731 - small adapter to the (path) - return _create("core.indexing.parsers.audio.client_based", "audio_client", client=client) +class _PdfStrategyParser(DocumentParser): + """Force a specific PDF backend (a preset's ``parsing_strategy``) for PDF + documents, delegating every other content type to the shared dispatcher so + the per-preset choice never duplicates a backend or its pool.""" + + def __init__(self, dispatcher: ParserDispatcher, pdf_backend: str) -> None: + self._dispatcher = dispatcher + self._pdf_backend = pdf_backend + + def supported_types(self) -> list[str]: + return self._dispatcher.supported_types() + + async def parse(self, document: Document) -> ProcessedDocument: + if document.content_type is DocumentType.PDF: + return await self._dispatcher._get(self._pdf_backend).parse(document) + return await self._dispatcher.parse(document) + + # Only backends that need a non-conventional module path (``pymupdf`` lives # under ``pdf/``) or injected dependencies (pooled/client/eml). Everything else # is built by convention in ``ParserDispatcher._build``. _BUILDERS: dict[str, Any] = { + # pymupdf is the lightweight, fast PDF backend. It builds in markdown mode + # (the default): pymupdf4llm preserves structure (headings/tables) for the + # markdown-aware chunker, with embed_images=False so no base64 bloats chunks + # and no image rendering happens. Images/captioning are marker/docling's job. "pymupdf": lambda d: _create("core.indexing.parsers.pdf.pymupdf", "pymupdf"), "eml": lambda d: d._build_eml(), "marker": lambda d: d._build_marker(), diff --git a/openrag/services/workers/pipeline_builder.py b/openrag/services/workers/pipeline_builder.py index 817b3b5e5..b90c4c1f9 100644 --- a/openrag/services/workers/pipeline_builder.py +++ b/openrag/services/workers/pipeline_builder.py @@ -127,7 +127,11 @@ def _effective_indexation_config(self, row: MutableMapping[str, Any]) -> Indexat raise TypeError("indexation_config must be an IndexationPipelineConfig or dict") def _select_parser(self, config: IndexationPipelineConfig | None) -> DocumentParser: - if config is not None and self.parser_factory is not None: + # parsing_strategy is None => the preset doesn't override PDF parsing, so + # defer to the global dispatcher (self.parser), which routes PDFs to the + # deployment's configured file_loaders.pdf. Only an explicit strategy + # goes through the factory (and lazily builds that backend's pool). + if config is not None and self.parser_factory is not None and config.parsing_strategy is not None: return self.parser_factory(config.parsing_strategy) return self.parser diff --git a/openrag/services/workers/stages/parse.py b/openrag/services/workers/stages/parse.py index 29f225aed..22a43979a 100644 --- a/openrag/services/workers/stages/parse.py +++ b/openrag/services/workers/stages/parse.py @@ -39,4 +39,14 @@ async def _parse_with_timeout( document: Document, timeout: float | None, ) -> ProcessedDocument: - return await run_with_optional_timeout(lambda: parser.parse(document), timeout) + try: + return await run_with_optional_timeout(lambda: parser.parse(document), timeout) + except TimeoutError as exc: + if timeout is None: + # No outer bound was applied, so this TimeoutError is internal to the + # parser — surface it as-is (and avoid formatting {timeout:g} on None). + raise + # asyncio.wait_for raises a bare TimeoutError whose str() is empty; give + # the failed-file report (row["error"]) a message that names the file and + # the bound that was exceeded. + raise TimeoutError(f"parse timed out after {timeout:g}s for {document.filename!r}") from exc diff --git a/tests/unit/services/workers/parsers/test_parser_dispatcher.py b/tests/unit/services/workers/parsers/test_parser_dispatcher.py index 802bcfcf0..acc405ace 100644 --- a/tests/unit/services/workers/parsers/test_parser_dispatcher.py +++ b/tests/unit/services/workers/parsers/test_parser_dispatcher.py @@ -115,6 +115,63 @@ async def test_parse_dispatches_to_cached_backend() -> None: assert result.text_blocks[0].text == "ok" +@pytest.mark.asyncio +async def test_for_pdf_strategy_overrides_pdf_backend_and_shares_cache() -> None: + """A preset's ``parsing_strategy`` must override the global PDF backend for + PDFs while non-PDF types still dispatch normally — reusing the dispatcher's + cached backends (no duplicate pools).""" + disp = ParserDispatcher(_config(pdf="MarkerLoader")) # global default = marker + marker, pymupdf, text = _FakeParser(), _FakeParser(), _FakeParser() + disp._by_name.update({"marker": marker, "pymupdf": pymupdf, "text": text}) + + pdf_parser = disp.for_pdf_strategy("pymupdf") + + pdf = Document(filename="a.pdf", content_type=DocumentType.PDF, raw_bytes=b"%PDF-1.4") + await pdf_parser.parse(pdf) + assert pymupdf.seen is pdf # routed to the preset strategy, not the global marker + assert marker.seen is None + + txt = Document(filename="a.txt", content_type=DocumentType.TEXT, raw_bytes=b"hi") + await pdf_parser.parse(txt) + assert text.seen is txt # non-PDF content still dispatches by content type + + +def test_for_pdf_strategy_rejects_unknown_strategy() -> None: + with pytest.raises(ValueError, match="Unsupported PDF parsing strategy"): + ParserDispatcher(_config()).for_pdf_strategy("nope") + + +def test_pymupdf_backend_builds_in_markdown_mode_without_images() -> None: + """The lightweight pymupdf backend builds in markdown mode — structured text + for the markdown-aware chunker — but with embed_images=False so it never + inlines base64 images into chunk text (which bloats chunks / breaks Milvus + inserts). Images are marker/docling's job.""" + parser = ParserDispatcher(_config(pdf="PyMuPDFLoader"))._get("pymupdf") + assert getattr(parser, "_mode", None) == "markdown" + + # The "without images" contract: build a PDF that actually contains an image + # and confirm the markdown extractor produces NO ImageBlocks and inlines no + # base64 data URIs. Catches a regression that re-enables embed_images. + import io + + import pymupdf + from core.indexing.parsers.pdf.pymupdf import _extract_markdown + from PIL import Image + + buf = io.BytesIO() + Image.new("RGB", (8, 8), "red").save(buf, format="PNG") + doc = pymupdf.open() + page = doc.new_page() + page.insert_text((72, 72), "Hello world.") + page.insert_image(pymupdf.Rect(0, 0, 8, 8), stream=buf.getvalue()) + raw = doc.tobytes() + doc.close() + + pages, images = _extract_markdown(raw) + assert images == [] # pymupdf must not extract/inline images + assert not any("data:image" in p for p in pages) + + def test_build_caption_vlm_requires_endpoint() -> None: # No VLM endpoint configured -> unavailable, regardless of the captioning flag. assert build_caption_vlm(_config(image_captioning=True, vlm_base_url="")) is None diff --git a/tests/unit/services/workers/parsers/test_pool_loaders.py b/tests/unit/services/workers/parsers/test_pool_loaders.py new file mode 100644 index 000000000..c5dc42817 --- /dev/null +++ b/tests/unit/services/workers/parsers/test_pool_loaders.py @@ -0,0 +1,44 @@ +"""The pooled PDF loaders must create their Ray pool lazily via +``get_or_create_actor`` rather than assuming bootstrap already created it. + +Bootstrap only pre-warms the *globally* configured PDF backend, but a per-preset +``parsing_strategy`` can select any backend at runtime (#569). A get-only +``ray.get_actor`` then fails with "Failed to look up actor 'DoclingPool'" +(#575); lazy creation makes whichever backend a preset picks work on first use. +""" + +from __future__ import annotations + +import pytest + + +@pytest.mark.parametrize( + ("module_name", "loader_name", "pool_name", "pool_attr"), + [ + ("services.workers.parsers.docling_workers", "DoclingLoader", "DoclingPool", "DoclingPool"), + ("services.workers.parsers.marker_workers", "MarkerLoader", "MarkerPool", "MarkerPool"), + ], +) +def test_pool_loader_lazily_creates_its_pool(monkeypatch, module_name, loader_name, pool_name, pool_attr): + import importlib + + import services.workers.bootstrap as bootstrap + + module = importlib.import_module(module_name) + Loader = getattr(module, loader_name) + PoolCls = getattr(module, pool_attr) + + calls: list[tuple] = [] + + def fake_get_or_create_actor(name, cls, **options): + calls.append((name, cls, options)) + return "pool-handle" + + monkeypatch.setattr(bootstrap, "get_or_create_actor", fake_get_or_create_actor) + # A get-only ray.get_actor must NOT be used anymore — fail loudly if it is. + monkeypatch.setattr(module.ray, "get_actor", lambda *a, **k: pytest.fail("loader used get-only ray.get_actor")) + + loader = Loader() + + assert loader.worker == "pool-handle" + assert calls == [(pool_name, PoolCls, {"lifetime": "detached"})] diff --git a/tests/unit/services/workers/stages/test_parse.py b/tests/unit/services/workers/stages/test_parse.py index 85b77e000..48c607804 100644 --- a/tests/unit/services/workers/stages/test_parse.py +++ b/tests/unit/services/workers/stages/test_parse.py @@ -61,6 +61,31 @@ async def test_parse_stage_marks_error_and_scrubs_credentials_when_parser_fails( assert "api_key" not in row +@pytest.mark.asyncio +async def test_parse_stage_times_out_a_wedged_parse_and_reports_the_file(): + """A slow/wedged parse must fail *that* file (with a descriptive, non-empty + error naming the file) instead of hanging the pipeline indefinitely (#571).""" + import asyncio + + class _WedgedParser(DocumentParser): + async def parse(self, document: Document) -> ProcessedDocument: + await asyncio.sleep(10) # longer than the stage timeout + raise AssertionError("should have timed out") + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + document = Document(id="doc-1", filename="wedged.pdf", content_type=DocumentType.TEXT, text="x") + row = {"document": document, "token": "secret"} + + with pytest.raises(TimeoutError, match=r"parse timed out after 0.05s for 'wedged.pdf'"): + await parse_stage(row, _WedgedParser(), timeout=0.05) + + assert row["stage"] == "parse_failed" + assert row["error"] == "parse timed out after 0.05s for 'wedged.pdf'" + assert "token" not in row + + @pytest.mark.asyncio async def test_parse_stage_requires_document_in_row(): row = {"api_key": "secret"} @@ -71,3 +96,17 @@ async def test_parse_stage_requires_document_in_row(): assert row["stage"] == "parse_failed" assert row["error"] == "parse_stage row must contain a Document under 'document'" assert "api_key" not in row + + +@pytest.mark.asyncio +async def test_parse_stage_surfaces_internal_timeout_when_no_outer_bound(): + """With timeout=None, a TimeoutError raised inside the parser is internal — + it must surface as-is (not be relabeled or crash formatting None).""" + document = Document(id="doc-1", filename="x.pdf", content_type=DocumentType.TEXT, text="hi") + row = {"document": document} + + with pytest.raises(TimeoutError, match="internal parser timeout"): + await parse_stage(row, FakeParser(error=TimeoutError("internal parser timeout")), timeout=None) + + assert row["stage"] == "parse_failed" + assert row["error"] == "internal parser timeout" diff --git a/tests/unit/services/workers/test_indexer_pool.py b/tests/unit/services/workers/test_indexer_pool.py index d888108d8..c99ec872c 100644 --- a/tests/unit/services/workers/test_indexer_pool.py +++ b/tests/unit/services/workers/test_indexer_pool.py @@ -20,6 +20,18 @@ class _NonCallableChunker: chunk = None +def test_build_pipeline_timeouts_bounds_parse_from_config() -> None: + """The pipeline must bound the parse stage at loader.parse_timeout so a wedged + parse fails that file instead of stalling indexing (#571).""" + from services.workers.indexer_pool import _build_pipeline_timeouts + + cfg = SimpleNamespace(loader=SimpleNamespace(parse_timeout=42)) + + timeouts = _build_pipeline_timeouts(cfg) + + assert timeouts.parse == 42 + + def test_build_chunker_returns_native_chunker(monkeypatch: pytest.MonkeyPatch) -> None: import core.chunking.factory as factory from services.workers.indexer_pool import _build_chunker @@ -156,6 +168,28 @@ def test_build_contextualizer_factory_returns_factory_for_later_hydration(tmp_pa factory("default") +def test_build_parser_factory_delegates_to_strategy_and_caches() -> None: + # The parser factory must route a preset's parsing_strategy through the + # dispatcher's for_pdf_strategy (so pymupdf/docling are honored, not the + # global default) and cache per strategy so no backend/pool is duplicated. + from services.workers.indexer_pool import _build_parser_factory + + calls: list[str] = [] + + class _FakeDispatcher: + def for_pdf_strategy(self, strategy: str): + calls.append(strategy) + return SimpleNamespace(strategy=strategy) + + factory = _build_parser_factory(_FakeDispatcher()) + + first = factory("pymupdf") + assert first.strategy == "pymupdf" + assert factory("pymupdf") is first # cached: built once per strategy + assert factory("docling").strategy == "docling" + assert calls == ["pymupdf", "docling"] # no rebuild for the repeated strategy + + def test_contextualizer_factory_reads_live_registry(tmp_path) -> None: # The factory holds a live reference to cfg.models.llm: a name added to the # registry AFTER the factory is built (mimicking the indexer's lazy DB @@ -757,7 +791,7 @@ def model_copy(self, *, update): batch_size=32, embed_concurrency=2, ), - loader=SimpleNamespace(image_captioning=True), + loader=SimpleNamespace(image_captioning=True, parse_timeout=3600), vectordb=SimpleNamespace(collection_name="vdb_test"), rdb=RDBConfig(), )