Merge pull request #922 from cbcoutinho/feat/tier2-incluster-ocr
feat(ingest): split OCR into tier2 in-cluster (GPU, gateway-only) + tier3 upstream
This commit is contained in:
@@ -335,7 +335,7 @@ def _init_worker_observability(settings: Settings) -> None:
|
||||
)
|
||||
@click.option(
|
||||
"--tier",
|
||||
type=click.Choice(["fast", "structured", "ocr"]),
|
||||
type=click.Choice(["fast", "structured", "ocr-incluster", "ocr-upstream"]),
|
||||
default=None,
|
||||
help=(
|
||||
"Run only this extraction tier's queue (Deck #323). Omit to drain ALL "
|
||||
@@ -354,8 +354,9 @@ def worker(concurrency: int | None, tier: str | None):
|
||||
|
||||
\b
|
||||
With --tier the worker drains only that tier's queue (``ingest-<tier>``), so
|
||||
a CPU-bound ``fast`` fleet, an in-cluster ``structured`` fleet, and a paid
|
||||
``ocr`` fleet scale independently. Without it, all tier queues are drained in
|
||||
a CPU-bound ``fast`` fleet, an in-cluster ``structured`` fleet, an on-demand
|
||||
GPU ``ocr-incluster`` fleet, and a paid ``ocr-upstream`` fleet scale
|
||||
independently. Without it, all tier queues are drained in
|
||||
one process (handy for dev / a single Deployment). A low-quality parse hops
|
||||
the job to the next tier's queue automatically (see TieredEscalationStrategy).
|
||||
|
||||
@@ -386,21 +387,31 @@ def worker(concurrency: int | None, tier: str | None):
|
||||
ALL_INGEST_QUEUES,
|
||||
INGEST_QUEUE_MAINTENANCE,
|
||||
LEGACY_INGEST_QUEUE,
|
||||
LEGACY_INGEST_QUEUE_OCR,
|
||||
TIER_QUEUES,
|
||||
apply_ingest_queue_schema,
|
||||
get_procrastinate_app,
|
||||
)
|
||||
|
||||
# Which queues this process drains. A single tier -> just its queue; no tier
|
||||
# -> every tier queue PLUS the legacy single queue, so a rolling upgrade
|
||||
# never strands jobs deferred under the pre-#323 name. Every worker also
|
||||
# -> every tier queue PLUS the legacy queues, so a rolling upgrade never
|
||||
# strands jobs deferred under the pre-#323 single queue or the pre-#353 single
|
||||
# OCR queue. The upstream OCR worker also drains the pre-split ``ingest-ocr``
|
||||
# queue (the old single OCR tier defaulted to upstream/Mistral). Every worker
|
||||
# drains the maintenance queue so the periodic stalled-job reclaim fires
|
||||
# regardless of which tier(s) are scaled up (procrastinate dedups the
|
||||
# periodic, so multiple drainers don't multiply the reclaim).
|
||||
if tier is not None:
|
||||
queues = [TIER_QUEUES[tier], INGEST_QUEUE_MAINTENANCE]
|
||||
if tier == "ocr-upstream":
|
||||
queues.insert(1, LEGACY_INGEST_QUEUE_OCR)
|
||||
else:
|
||||
queues = [*ALL_INGEST_QUEUES, LEGACY_INGEST_QUEUE, INGEST_QUEUE_MAINTENANCE]
|
||||
queues = [
|
||||
*ALL_INGEST_QUEUES,
|
||||
LEGACY_INGEST_QUEUE,
|
||||
LEGACY_INGEST_QUEUE_OCR,
|
||||
INGEST_QUEUE_MAINTENANCE,
|
||||
]
|
||||
|
||||
# This is the consumer side of the distributed (postgres) ingest backend.
|
||||
# Unlike the in-process anyio pool, the worker talks to procrastinate's App
|
||||
|
||||
@@ -162,6 +162,13 @@ _DEFAULTS: dict[str, Any] = {
|
||||
# Provider-namespaced OCR model id (gateway routes on the prefix; the direct
|
||||
# mistral backend strips it).
|
||||
"document_ocr_model": "mistral/mistral-ocr-latest",
|
||||
# In-cluster OCR tier (tier2, Deck #353): the on-demand burst GPU, reached
|
||||
# ONLY via the embedding gateway. Off + per-tenant by default; tried before the
|
||||
# paid upstream OCR rung. The model is a config value (the gateway routes on its
|
||||
# "<provider>/" prefix) — surya is the current pick but swappable (e.g.
|
||||
# "lightonocr/...") and is NEVER hard-coded, only this default.
|
||||
"document_ocr_incluster_enabled": False,
|
||||
"document_ocr_incluster_model": "surya/surya-ocr-2",
|
||||
# OCR escalation triggers (tier-0). A page is OCR-worthy when its text is
|
||||
# near-empty (< min_page_chars) OR low-quality (< min_text_quality) OR (when
|
||||
# detect_scanned) mostly a raster image; a doc escalates when the OCR-worthy
|
||||
@@ -884,6 +891,11 @@ class Settings:
|
||||
# gateway routes on the "<provider>/" prefix; the direct mistral backend
|
||||
# strips it.
|
||||
document_ocr_model: str = "mistral/mistral-ocr-latest"
|
||||
# In-cluster OCR tier (tier2, Deck #353): on-demand burst GPU reached ONLY via
|
||||
# the embedding gateway; off + per-tenant, tried before the upstream rung. The
|
||||
# model is a swappable config value (surya is the current pick, never hardcoded).
|
||||
document_ocr_incluster_enabled: bool = False
|
||||
document_ocr_incluster_model: str = "surya/surya-ocr-2"
|
||||
# OCR backend HTTP request timeout (seconds). float for parity with the
|
||||
# parse timeout / httpx.Timeout; per-tenant tunable so a gateway with a
|
||||
# shorter ceiling isn't masked by the 180s default.
|
||||
@@ -1535,6 +1547,8 @@ def get_settings() -> Settings:
|
||||
"document_ocr_enabled": "DOCUMENT_OCR_ENABLED",
|
||||
"document_ocr_provider": "DOCUMENT_OCR_PROVIDER",
|
||||
"document_ocr_model": "DOCUMENT_OCR_MODEL",
|
||||
"document_ocr_incluster_enabled": "DOCUMENT_OCR_INCLUSTER_ENABLED",
|
||||
"document_ocr_incluster_model": "DOCUMENT_OCR_INCLUSTER_MODEL",
|
||||
"document_ocr_timeout_seconds": "DOCUMENT_OCR_TIMEOUT_SECONDS",
|
||||
"document_ocr_mode": "DOCUMENT_OCR_MODE",
|
||||
"document_ocr_batch_poll_seconds": "DOCUMENT_OCR_BATCH_POLL_SECONDS",
|
||||
|
||||
@@ -8,13 +8,32 @@ from .registry import ProcessorRegistry, get_registry
|
||||
|
||||
# Register processors at module initialization. The tiered PDF pipeline selects
|
||||
# by tier (not priority): Pypdfium2FastProcessor is the ``fast`` tier,
|
||||
# PyMuPDFProcessor the ``structured`` rollback, and OcrProcessor the ``ocr``
|
||||
# escalation target (reached only when document_ocr_enabled). OcrProcessor gets
|
||||
# the lowest priority so it is never the non-tiered default for PDFs.
|
||||
# PyMuPDFProcessor the ``structured`` rollback, and TWO OcrProcessor instances are
|
||||
# the OCR rungs — ``ocr-incluster`` (the on-demand burst GPU, gateway-only, reached
|
||||
# via the embedding gateway over the tailnet; e.g. surya) tried before
|
||||
# ``ocr-upstream`` (paid Mistral). Each is reached only when its own opt-in flag is
|
||||
# set. OCR gets the lowest priorities so it's never the non-tiered default for PDFs.
|
||||
_registry = get_registry()
|
||||
_registry.register(Pypdfium2FastProcessor(), priority=20)
|
||||
_registry.register(PyMuPDFProcessor(), priority=10)
|
||||
_registry.register(OcrProcessor(), priority=1)
|
||||
_registry.register(
|
||||
OcrProcessor(
|
||||
name="ocr-incluster",
|
||||
tier="ocr-incluster",
|
||||
model_setting="document_ocr_incluster_model",
|
||||
gateway_only=True,
|
||||
),
|
||||
priority=2,
|
||||
)
|
||||
_registry.register(
|
||||
OcrProcessor(
|
||||
name="ocr-upstream",
|
||||
tier="ocr-upstream",
|
||||
model_setting="document_ocr_model",
|
||||
gateway_only=False,
|
||||
),
|
||||
priority=1,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DocumentProcessor",
|
||||
|
||||
@@ -89,7 +89,10 @@ class DocClassification:
|
||||
total_chars: int
|
||||
mean_text_quality: float
|
||||
ocr_page_fraction: float # fraction of sampled pages flagged needs_ocr
|
||||
recommended_tier: str # "fast" | "structured" | "ocr"
|
||||
# Classifier vocabulary (coarse): "fast" | "structured" | "ocr". "ocr" means
|
||||
# "needs OCR" — the registry resolves it to a concrete rung (ocr-incluster ->
|
||||
# ocr-upstream) via next_available_tier; it is NOT the TIER_LADDER tier name.
|
||||
recommended_tier: str
|
||||
mean_control_ratio: float = 0.0 # doc-level C0-control-char ratio (glyph-leak)
|
||||
flags: set[str] = field(
|
||||
default_factory=set
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
The escalation ladder is the cheapest-first ordering of extraction tiers:
|
||||
|
||||
fast -> structured -> ocr ( -> llm, reserved)
|
||||
fast -> structured -> ocr-incluster -> ocr-upstream ( -> llm, reserved)
|
||||
|
||||
It mirrors the ``tier`` vocabulary documented on
|
||||
:meth:`DocumentProcessor.tier <.base.DocumentProcessor.tier>` and the
|
||||
@@ -24,9 +24,11 @@ from __future__ import annotations
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal
|
||||
|
||||
# Cheapest-first. ``llm`` is reserved (see base.DocumentProcessor.tier) and not
|
||||
# wired yet, so it is intentionally absent from the live ladder.
|
||||
TIER_LADDER: tuple[str, ...] = ("fast", "structured", "ocr")
|
||||
# Cheapest-first. OCR is split into two rungs: ``ocr-incluster`` (the on-demand
|
||||
# burst GPU, e.g. surya, reached via the embedding gateway over the tailnet) tried
|
||||
# BEFORE ``ocr-upstream`` (paid Mistral). ``llm`` is reserved (see
|
||||
# base.DocumentProcessor.tier) and not wired yet.
|
||||
TIER_LADDER: tuple[str, ...] = ("fast", "structured", "ocr-incluster", "ocr-upstream")
|
||||
|
||||
|
||||
def escalation_tiers_signature(settings: Any) -> str:
|
||||
@@ -55,6 +57,7 @@ def escalation_tiers_signature(settings: Any) -> str:
|
||||
"""
|
||||
return (
|
||||
f"ocr={int(bool(settings.document_ocr_enabled))};"
|
||||
f"ocric={int(bool(settings.document_ocr_incluster_enabled))};"
|
||||
f"t1={settings.document_tier1_engine}"
|
||||
)
|
||||
|
||||
@@ -141,7 +144,8 @@ class BatchPending(Exception):
|
||||
``except Exception`` on the indexing path, never counted as a drop/parse
|
||||
error, and never marks the placeholder failed (the doc isn't done yet).
|
||||
Unlike ``EscalateError`` it does NOT change queue — the job stays on its own
|
||||
(``ocr``) tier queue and is simply deferred.
|
||||
(``ocr-upstream``) tier queue and is simply deferred (batch mode is the
|
||||
upstream Mistral path only; the in-cluster rung is synchronous).
|
||||
"""
|
||||
|
||||
def __init__(self, *, retry_in: int) -> None:
|
||||
|
||||
@@ -203,11 +203,16 @@ def _build_gateway_token_provider(settings: Settings) -> Any:
|
||||
)
|
||||
|
||||
|
||||
def build_gateway_batch_client(settings: Settings) -> "GatewayBatchOcrClient | None":
|
||||
def build_gateway_batch_client(
|
||||
settings: Settings, *, model: str | None = None
|
||||
) -> "GatewayBatchOcrClient | None":
|
||||
"""Build a ``GatewayBatchOcrClient`` when the gateway is the OCR backend, else
|
||||
``None`` (so batch mode falls back to sync for provider=mistral / no gateway).
|
||||
Batch OCR is gateway-only — Mistral's Batch API is reached *through* the
|
||||
gateway's batch routes, never directly from the pod."""
|
||||
gateway's batch routes, never directly from the pod.
|
||||
|
||||
``model`` overrides ``settings.document_ocr_model`` so a per-tier OCR rung
|
||||
(e.g. the in-cluster tier) submits its own provider-namespaced model id."""
|
||||
if settings.document_ocr_provider not in ("gateway", "auto"):
|
||||
return None
|
||||
if not settings.embedding_gateway_url:
|
||||
@@ -216,28 +221,66 @@ def build_gateway_batch_client(settings: Settings) -> "GatewayBatchOcrClient | N
|
||||
|
||||
return GatewayBatchOcrClient(
|
||||
settings.embedding_gateway_url,
|
||||
settings.document_ocr_model,
|
||||
model or settings.document_ocr_model,
|
||||
_build_gateway_token_provider(settings),
|
||||
)
|
||||
|
||||
|
||||
def build_ocr_backend(settings: Settings) -> _OcrBackend | None:
|
||||
"""Select an OCR backend from settings, or None when none is available."""
|
||||
def build_ocr_backend(
|
||||
settings: Settings, *, model: str | None = None, gateway_only: bool = False
|
||||
) -> _OcrBackend | None:
|
||||
"""Select an OCR backend from settings, or None when none is available.
|
||||
|
||||
``model`` overrides ``settings.document_ocr_model`` so a per-tier OCR rung
|
||||
binds its own provider-namespaced model id. ``gateway_only`` forces the
|
||||
gateway backend (never the direct Mistral fallback) — used by the **in-cluster**
|
||||
OCR tier, whose backend (e.g. surya on the burst GPU) is reachable ONLY through
|
||||
the embedding gateway over the tailnet; with no gateway URL the tier is
|
||||
disabled (warn) rather than misrouted to a direct backend that can't serve it.
|
||||
"""
|
||||
provider = settings.document_ocr_provider
|
||||
# `is not None` (not truthiness): an empty model string must NOT silently fall
|
||||
# back to the upstream default and misroute a per-tier rung.
|
||||
model = model if model is not None else settings.document_ocr_model
|
||||
|
||||
if gateway_only:
|
||||
# The in-cluster tier is gateway-only and never touches the `mistral`
|
||||
# provider, but provider=none still suppresses it. Warn so an operator who
|
||||
# set provider=none to disable Mistral doesn't silently lose the GPU tier.
|
||||
if provider == "none":
|
||||
if settings.document_ocr_incluster_enabled:
|
||||
logger.warning(
|
||||
"DOCUMENT_OCR_PROVIDER=none disables in-cluster OCR even with "
|
||||
"DOCUMENT_OCR_INCLUSTER_ENABLED=true; set it to 'gateway' or "
|
||||
"'auto' to keep the in-cluster tier"
|
||||
)
|
||||
return None
|
||||
if settings.embedding_gateway_url:
|
||||
return _GatewayOcrBackend(
|
||||
settings.embedding_gateway_url,
|
||||
model,
|
||||
_build_gateway_token_provider(settings),
|
||||
)
|
||||
logger.warning(
|
||||
"in-cluster OCR tier requires EMBEDDING_GATEWAY_URL (it routes through "
|
||||
"the gateway, never a direct backend); this OCR tier is disabled"
|
||||
)
|
||||
return None
|
||||
|
||||
if provider == "none":
|
||||
return None
|
||||
|
||||
if provider in ("gateway", "auto") and settings.embedding_gateway_url:
|
||||
return _GatewayOcrBackend(
|
||||
settings.embedding_gateway_url,
|
||||
settings.document_ocr_model,
|
||||
model,
|
||||
_build_gateway_token_provider(settings),
|
||||
)
|
||||
|
||||
if provider in ("mistral", "auto") and settings.mistral_api_key:
|
||||
return _MistralOcrBackend(
|
||||
settings.mistral_api_key,
|
||||
settings.document_ocr_model,
|
||||
model,
|
||||
settings.mistral_base_url,
|
||||
)
|
||||
|
||||
@@ -258,9 +301,37 @@ def build_ocr_backend(settings: Settings) -> _OcrBackend | None:
|
||||
|
||||
|
||||
class OcrProcessor(DocumentProcessor):
|
||||
"""Tier-3 OCR processor (gateway or direct Mistral backend)."""
|
||||
"""OCR processor for both OCR rungs — tier2 in-cluster (gateway-only GPU) and
|
||||
tier3 upstream (gateway or direct Mistral backend). One class, two registered
|
||||
instances bound to different ``(tier, model_setting, gateway_only)``."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
# The defaults describe the UPSTREAM (tier3) rung and exist for
|
||||
# test/bare-construction convenience only. Application wiring
|
||||
# (document_processors/__init__.py) ALWAYS passes every arg explicitly for
|
||||
# both rungs — don't rely on these defaults in app code.
|
||||
*,
|
||||
name: str = "ocr-upstream",
|
||||
tier: str = "ocr-upstream",
|
||||
model_setting: str = "document_ocr_model",
|
||||
gateway_only: bool = False,
|
||||
) -> None:
|
||||
# One OcrProcessor class serves BOTH OCR rungs; instances are bound to a
|
||||
# tier + the settings attribute holding their provider-namespaced model id
|
||||
# (+ whether the backend is gateway-only). The in-cluster rung is
|
||||
# gateway-only (its model, e.g. surya, is reachable solely via the
|
||||
# gateway); the upstream rung keeps the configurable gateway/mistral
|
||||
# selection. surya is NEVER hard-coded here — only a config default.
|
||||
# Fail fast on a misconfigured model_setting (a typo in a constructor call)
|
||||
# so it surfaces at startup, not as an AttributeError mid-OCR. The string
|
||||
# only ever comes from hardcoded defaults in __init__.py, never user input.
|
||||
if not hasattr(Settings, model_setting):
|
||||
raise ValueError(f"Unknown model_setting: {model_setting!r}")
|
||||
self._name = name
|
||||
self._tier = tier
|
||||
self._model_setting = model_setting
|
||||
self._gateway_only = gateway_only
|
||||
# Resolve the backend once and reuse it: rebuilding per call would create
|
||||
# a fresh GatewayTokenProvider each time (discarding its M2M-token cache
|
||||
# -> a token fetch per document) and a new Mistral SDK client per call.
|
||||
@@ -283,11 +354,11 @@ class OcrProcessor(DocumentProcessor):
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "ocr"
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def tier(self) -> str:
|
||||
return "ocr"
|
||||
return self._tier
|
||||
|
||||
@property
|
||||
def supported_mime_types(self) -> set[str]:
|
||||
@@ -327,7 +398,11 @@ class OcrProcessor(DocumentProcessor):
|
||||
self._backend_lock = anyio.Lock()
|
||||
async with self._backend_lock:
|
||||
if not self._backend_resolved: # double-checked
|
||||
self._backend = build_ocr_backend(settings)
|
||||
self._backend = build_ocr_backend(
|
||||
settings,
|
||||
model=getattr(settings, self._model_setting),
|
||||
gateway_only=self._gateway_only,
|
||||
)
|
||||
self._backend_resolved = True
|
||||
backend = self._backend
|
||||
if backend is None:
|
||||
@@ -387,13 +462,24 @@ class OcrProcessor(DocumentProcessor):
|
||||
async def _get_batch_client(self) -> "GatewayBatchOcrClient | None":
|
||||
"""Cached gateway batch client (or ``None`` when batch isn't applicable —
|
||||
provider=mistral / no gateway). Resolved once under the backend lock so the
|
||||
token provider's M2M cache survives across documents."""
|
||||
token provider's M2M cache survives across documents.
|
||||
|
||||
The in-cluster (``gateway_only``) rung never uses batch mode: it targets
|
||||
the on-demand GPU, which is synchronous/low-latency, while batch OCR is the
|
||||
upstream (Mistral) async-job path. So even with ``DOCUMENT_OCR_MODE=batch``
|
||||
set globally, the in-cluster tier stays on the synchronous backend."""
|
||||
if self._gateway_only:
|
||||
return None
|
||||
if not self._batch_client_resolved:
|
||||
if self._batch_client_lock is None:
|
||||
self._batch_client_lock = anyio.Lock()
|
||||
async with self._batch_client_lock:
|
||||
if not self._batch_client_resolved: # double-checked
|
||||
self._batch_client = build_gateway_batch_client(get_settings())
|
||||
settings = get_settings()
|
||||
self._batch_client = build_gateway_batch_client(
|
||||
settings,
|
||||
model=getattr(settings, self._model_setting),
|
||||
)
|
||||
self._batch_client_resolved = True
|
||||
return self._batch_client
|
||||
|
||||
@@ -431,6 +517,11 @@ class OcrProcessor(DocumentProcessor):
|
||||
return None
|
||||
client = await self._get_batch_client()
|
||||
if client is None:
|
||||
# The in-cluster (gateway_only) rung returns None here BY DESIGN — the
|
||||
# GPU is synchronous-only, batch is the upstream Mistral path — so don't
|
||||
# emit the "no gateway backend" warning (the gateway IS configured; that
|
||||
# warning would send an operator chasing a non-existent config problem).
|
||||
if not self._gateway_only:
|
||||
self._batch_fallback(
|
||||
"no gateway backend (provider=mistral or EMBEDDING_GATEWAY_URL unset)",
|
||||
filename,
|
||||
|
||||
@@ -349,10 +349,22 @@ class ProcessorRegistry:
|
||||
and not structured_failed
|
||||
and classification.recommended_tier in ("ocr", "structured")
|
||||
and classification.page_count > 0
|
||||
and settings.document_ocr_enabled
|
||||
and (
|
||||
settings.document_ocr_enabled or settings.document_ocr_incluster_enabled
|
||||
)
|
||||
):
|
||||
ocr = self._pdf_processor_for_tier("ocr")
|
||||
if ocr is not None:
|
||||
# Inline (memory pool) path: no queues to hop, so pick the cheapest
|
||||
# available OCR rung (in-cluster GPU before paid upstream) via the same
|
||||
# availability walk the queue path uses.
|
||||
ocr_tier = self.next_available_tier(
|
||||
from_tier, settings, minimum="ocr-incluster"
|
||||
)
|
||||
ocr = self._pdf_processor_for_tier(ocr_tier) if ocr_tier else None
|
||||
# `ocr is not None` already implies `ocr_tier is not None` at runtime,
|
||||
# but the type checker can't infer that across the conditional above,
|
||||
# so the explicit guard narrows `ocr_tier` to `str` for the
|
||||
# record_document_escalation(from_tier, ocr_tier, reason) call below.
|
||||
if ocr is not None and ocr_tier is not None:
|
||||
reason = (
|
||||
"corrupt_glyphs"
|
||||
if classification.recommended_tier == "structured"
|
||||
@@ -360,11 +372,12 @@ class ProcessorRegistry:
|
||||
if classification.total_chars == 0
|
||||
else "low_confidence"
|
||||
)
|
||||
record_document_escalation(from_tier, "ocr", reason)
|
||||
record_document_escalation(from_tier, ocr_tier, reason)
|
||||
logger.info(
|
||||
"Escalating %s %s->ocr (reason=%s)",
|
||||
"Escalating %s %s->%s (reason=%s)",
|
||||
filename or "<bytes>",
|
||||
from_tier,
|
||||
ocr_tier,
|
||||
reason,
|
||||
)
|
||||
ocr_result = await self._run_processor(
|
||||
@@ -379,13 +392,14 @@ class ProcessorRegistry:
|
||||
# OCR is an enhancement, not a gate: if it can't run (no backend
|
||||
# configured / API down) or returns nothing, keep the tier-1
|
||||
# result rather than failing the document. Otherwise an operator
|
||||
# who sets DOCUMENT_OCR_ENABLED=true without credentials would
|
||||
# make scanned docs fail entirely -- strictly worse than off.
|
||||
# who enables OCR without credentials would make scanned docs fail
|
||||
# entirely -- strictly worse than off.
|
||||
if ocr_result.success:
|
||||
return ocr_result
|
||||
logger.warning(
|
||||
"OCR escalation did not succeed for %s (%s); keeping the "
|
||||
"OCR escalation to %s did not succeed for %s (%s); keeping the "
|
||||
"tier-1 result",
|
||||
ocr_tier,
|
||||
filename or "<bytes>",
|
||||
ocr_result.metadata.get("parse_failed_reason", "error"),
|
||||
)
|
||||
@@ -448,7 +462,13 @@ class ProcessorRegistry:
|
||||
return None
|
||||
try:
|
||||
image_coverage = None
|
||||
if settings.document_ocr_enabled and settings.document_ocr_detect_scanned:
|
||||
# Scan detection feeds either OCR rung (tier2 in-cluster or tier3
|
||||
# upstream), so run it whenever EITHER is enabled — a tenant with
|
||||
# only in-cluster OCR on still needs image-coverage scan signals.
|
||||
ocr_any_enabled = (
|
||||
settings.document_ocr_enabled or settings.document_ocr_incluster_enabled
|
||||
)
|
||||
if ocr_any_enabled and settings.document_ocr_detect_scanned:
|
||||
try:
|
||||
image_coverage = image_coverage_per_page(content)
|
||||
except Exception:
|
||||
@@ -498,16 +518,19 @@ class ProcessorRegistry:
|
||||
``ignore_ocr_enabled`` drops only the OCR-enabled gate (not the registered-
|
||||
processor requirement): it answers "would this tier run if OCR were turned
|
||||
on?" — used to compute the *ideal* escalation target for the what-if-OCR
|
||||
suppressed-escalation signal. (Today only ``ocr`` has an enabled gate; a
|
||||
future per-tier gate would extend the condition below.)
|
||||
suppressed-escalation signal. (Both OCR rungs — ``ocr-incluster`` and
|
||||
``ocr-upstream`` — have their own enabled gate; non-OCR tiers have none.)
|
||||
"""
|
||||
if self._pdf_processor_for_tier(tier) is None:
|
||||
return False
|
||||
if (
|
||||
not ignore_ocr_enabled
|
||||
and tier == "ocr"
|
||||
and not settings.document_ocr_enabled
|
||||
):
|
||||
# Each OCR rung has its own opt-in flag (in-cluster vs upstream); a rung is
|
||||
# unavailable when its flag is off (unless we're computing the what-if
|
||||
# ideal target). Non-OCR tiers have no enabled gate.
|
||||
ocr_enable = {
|
||||
"ocr-incluster": settings.document_ocr_incluster_enabled,
|
||||
"ocr-upstream": settings.document_ocr_enabled,
|
||||
}
|
||||
if not ignore_ocr_enabled and tier in ocr_enable and not ocr_enable[tier]:
|
||||
return False
|
||||
return True
|
||||
|
||||
@@ -605,9 +628,11 @@ class ProcessorRegistry:
|
||||
|
||||
Target-tier routing:
|
||||
|
||||
- ``total_chars == 0`` (scanned / no text layer) -> target the ``ocr``
|
||||
tier directly. Text-extractor tiers (``structured``) cannot conjure
|
||||
text from a pure raster scan, so a structured hop would just be wasted.
|
||||
- ``total_chars == 0`` (scanned / no text layer) -> target the cheapest
|
||||
OCR rung (``ocr-incluster``) directly; ``next_available_tier`` falls
|
||||
through to ``ocr-upstream`` if in-cluster is disabled/unregistered.
|
||||
Text-extractor tiers (``structured``) cannot conjure text from a pure
|
||||
raster scan, so a structured hop would just be wasted.
|
||||
- glyph-corrupt text layer (``recommended_tier == "structured"``) -> target
|
||||
the ``structured`` tier; pymupdf re-extracts a broken-/ToUnicode layer
|
||||
correctly, so OCR is never the target for this case.
|
||||
@@ -645,7 +670,10 @@ class ProcessorRegistry:
|
||||
minimum = "structured"
|
||||
reason = "corrupt_glyphs"
|
||||
elif classification.total_chars == 0:
|
||||
minimum = "ocr"
|
||||
# Scanned / no text layer: target the cheapest OCR rung (in-cluster
|
||||
# GPU); next_available_tier then falls through to the upstream rung if
|
||||
# in-cluster is disabled/unregistered.
|
||||
minimum = "ocr-incluster"
|
||||
reason = "empty_text"
|
||||
else:
|
||||
minimum = None
|
||||
|
||||
@@ -708,9 +708,10 @@ def update_ingest_queue_depth(by_queue: dict[str, dict[str, int]] | None) -> Non
|
||||
from nextcloud_mcp_server.vector.queue.procrastinate import ( # noqa: PLC0415
|
||||
ALL_INGEST_QUEUES,
|
||||
LEGACY_INGEST_QUEUE,
|
||||
LEGACY_INGEST_QUEUE_OCR,
|
||||
)
|
||||
|
||||
for queue in (*ALL_INGEST_QUEUES, LEGACY_INGEST_QUEUE):
|
||||
for queue in (*ALL_INGEST_QUEUES, LEGACY_INGEST_QUEUE, LEGACY_INGEST_QUEUE_OCR):
|
||||
for status in _INGEST_DEPTH_STATUSES:
|
||||
ingest_queue_depth.labels(queue=queue, status=status).set(0)
|
||||
for queue, per_status in by_queue.items():
|
||||
|
||||
@@ -353,10 +353,13 @@ async def record_indexing_usage(
|
||||
)
|
||||
# Paid-OCR pages are metered as a SEPARATE line (Deck #323) so the
|
||||
# expensive tier's cost is billable independently of CPU-cheap parsing
|
||||
# -- pages_embedded counts all parsed pages, pages_ocr only the OCR
|
||||
# tier's. Gated on the tier so it's emitted exactly when the doc was
|
||||
# actually OCR'd; the same page_count guard above applies.
|
||||
if pipeline_tier == "ocr":
|
||||
# -- pages_embedded counts all parsed pages, pages_ocr only the paid
|
||||
# OCR tier's. After the OCR split (Deck #353) "paid" = the UPSTREAM
|
||||
# (Mistral) rung; the in-cluster GPU rung's cost is recovered via the
|
||||
# burst lifecycle, not per-page, so it is NOT metered here (separate
|
||||
# per-tier OCR metering is a billing follow-up). Gated on the tier so
|
||||
# it's emitted exactly when the doc hit upstream OCR.
|
||||
if pipeline_tier == "ocr-upstream":
|
||||
await store.record_usage_event(
|
||||
metric="pages_ocr",
|
||||
value=page_count,
|
||||
|
||||
@@ -66,24 +66,30 @@ logger = logging.getLogger(__name__)
|
||||
# network-bound ``ocr`` fleet scale (and fail) independently.
|
||||
INGEST_QUEUE_FAST = "ingest-fast"
|
||||
INGEST_QUEUE_STRUCTURED = "ingest-structured"
|
||||
INGEST_QUEUE_OCR = "ingest-ocr"
|
||||
# OCR split into two rungs (Deck #353): in-cluster (burst GPU via the gateway,
|
||||
# the queue the GPU autoscaler counts) tried before upstream (paid Mistral).
|
||||
INGEST_QUEUE_OCR_INCLUSTER = "ingest-ocr-incluster"
|
||||
INGEST_QUEUE_OCR_UPSTREAM = "ingest-ocr-upstream"
|
||||
|
||||
# tier -> queue. The producer always defers onto the cheapest tier's queue; a
|
||||
# low-quality parse hops the job up the ladder via the retry strategy below.
|
||||
TIER_QUEUES: dict[str, str] = {
|
||||
"fast": INGEST_QUEUE_FAST,
|
||||
"structured": INGEST_QUEUE_STRUCTURED,
|
||||
"ocr": INGEST_QUEUE_OCR,
|
||||
"ocr-incluster": INGEST_QUEUE_OCR_INCLUSTER,
|
||||
"ocr-upstream": INGEST_QUEUE_OCR_UPSTREAM,
|
||||
}
|
||||
_QUEUE_TIERS: dict[str, str] = {queue: tier for tier, queue in TIER_QUEUES.items()}
|
||||
ALL_INGEST_QUEUES: tuple[str, ...] = tuple(TIER_QUEUES.values())
|
||||
# New jobs start here; ``ocr`` is reached only by escalation.
|
||||
# New jobs start here; the OCR rungs are reached only by escalation.
|
||||
DEFAULT_INGEST_QUEUE = INGEST_QUEUE_FAST
|
||||
|
||||
# Legacy single-queue name (pre-#323). A rolling upgrade may still have jobs
|
||||
# parked on it; a worker can be told to drain it alongside the tier queues, and
|
||||
# the job-count / reclaim helpers include it so nothing is stranded.
|
||||
# Legacy single-queue name (pre-#323) and the pre-split single OCR queue
|
||||
# (pre-#353). A rolling upgrade may still have jobs parked on either; a worker
|
||||
# can drain them alongside the tier queues, and the job-count / reclaim helpers
|
||||
# include them so nothing is stranded.
|
||||
LEGACY_INGEST_QUEUE = "ingest"
|
||||
LEGACY_INGEST_QUEUE_OCR = "ingest-ocr"
|
||||
# Back-compat alias for callers that imported the old single-queue constant.
|
||||
INGEST_QUEUE_NAME = DEFAULT_INGEST_QUEUE
|
||||
|
||||
@@ -95,8 +101,12 @@ INGEST_QUEUE_NAME = DEFAULT_INGEST_QUEUE
|
||||
# queues so tier isolation (which fleet processes which docs) is preserved.
|
||||
INGEST_QUEUE_MAINTENANCE = "ingest-maintenance"
|
||||
|
||||
# Queues the job-count + reclaim helpers sweep (tier queues + the legacy one).
|
||||
_MANAGED_QUEUES: tuple[str, ...] = (*ALL_INGEST_QUEUES, LEGACY_INGEST_QUEUE)
|
||||
# Queues the job-count + reclaim helpers sweep (tier queues + the legacy ones).
|
||||
_MANAGED_QUEUES: tuple[str, ...] = (
|
||||
*ALL_INGEST_QUEUES,
|
||||
LEGACY_INGEST_QUEUE,
|
||||
LEGACY_INGEST_QUEUE_OCR,
|
||||
)
|
||||
|
||||
# Blueprint namespace → registered task names are prefixed ``ingest:``.
|
||||
_NAMESPACE = "ingest"
|
||||
@@ -109,6 +119,15 @@ def tier_for_queue(queue: str | None) -> str:
|
||||
The queue-aware task uses this to pick which single tier to parse with: the
|
||||
job's current queue *is* its tier. A job on the legacy ``ingest`` queue (or
|
||||
any unrecognised queue) defaults to the cheapest tier.
|
||||
|
||||
The pre-split legacy OCR queue ``ingest-ocr`` (``LEGACY_INGEST_QUEUE_OCR``)
|
||||
is deliberately NOT mapped here, so it also resolves to ``fast``. These are
|
||||
in-flight jobs enqueued before the tier2/tier3 split; running them at fast
|
||||
re-extracts (an empty layer for a scanned doc), which re-enters the ladder
|
||||
and naturally re-escalates to ``ocr-incluster`` (the cheap GPU rung) — one
|
||||
extra cheap hop, but it keeps stranded legacy OCR jobs OFF the paid upstream
|
||||
rung rather than mapping them straight to ``ocr-upstream``. The set is
|
||||
transient (only during a single rollout window).
|
||||
"""
|
||||
return _QUEUE_TIERS.get(queue or "", "fast")
|
||||
|
||||
|
||||
@@ -19,8 +19,14 @@ from nextcloud_mcp_server.document_processors.escalation import (
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
|
||||
def _settings(*, ocr: bool, engine: str = "pypdfium2") -> SimpleNamespace:
|
||||
return SimpleNamespace(document_ocr_enabled=ocr, document_tier1_engine=engine)
|
||||
def _settings(
|
||||
*, ocr: bool, ocr_incluster: bool = False, engine: str = "pypdfium2"
|
||||
) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
document_ocr_enabled=ocr,
|
||||
document_ocr_incluster_enabled=ocr_incluster,
|
||||
document_tier1_engine=engine,
|
||||
)
|
||||
|
||||
|
||||
def test_signature_is_stable_for_same_config() -> None:
|
||||
@@ -36,6 +42,14 @@ def test_enabling_ocr_changes_signature() -> None:
|
||||
) != escalation_tiers_signature(_settings(ocr=True))
|
||||
|
||||
|
||||
def test_enabling_ocr_incluster_changes_signature() -> None:
|
||||
# Enabling the in-cluster (tier2) rung adds an escalation tier independently of
|
||||
# the upstream rung -> previously dead-lettered scanned docs become retryable.
|
||||
assert escalation_tiers_signature(
|
||||
_settings(ocr=False, ocr_incluster=False)
|
||||
) != escalation_tiers_signature(_settings(ocr=False, ocr_incluster=True))
|
||||
|
||||
|
||||
def test_tier1_engine_change_changes_signature() -> None:
|
||||
assert escalation_tiers_signature(
|
||||
_settings(ocr=False, engine="pypdfium2")
|
||||
|
||||
@@ -18,6 +18,7 @@ def _settings(**kw) -> Any: # a Settings stand-in (only the read fields matter)
|
||||
base = dict(
|
||||
document_ocr_provider="auto",
|
||||
document_ocr_model="mistral/mistral-ocr-latest",
|
||||
document_ocr_incluster_enabled=False,
|
||||
document_ocr_timeout_seconds=180.0,
|
||||
document_ocr_mode="sync",
|
||||
document_ocr_batch_poll_seconds=120,
|
||||
@@ -56,7 +57,7 @@ def test_build_backend_none():
|
||||
|
||||
def test_build_backend_gateway():
|
||||
b = ocr.build_ocr_backend(
|
||||
_settings(document_ocr_provider="gateway", embedding_gateway_url="http://gw")
|
||||
_settings(document_ocr_provider="gateway", embedding_gateway_url="https://gw")
|
||||
)
|
||||
assert isinstance(b, ocr._GatewayOcrBackend)
|
||||
|
||||
@@ -70,7 +71,7 @@ def test_build_backend_mistral():
|
||||
|
||||
def test_build_backend_auto_prefers_gateway():
|
||||
b = ocr.build_ocr_backend(
|
||||
_settings(embedding_gateway_url="http://gw", mistral_api_key="k")
|
||||
_settings(embedding_gateway_url="https://gw", mistral_api_key="k")
|
||||
)
|
||||
assert isinstance(b, ocr._GatewayOcrBackend)
|
||||
|
||||
@@ -99,6 +100,149 @@ def test_build_gateway_batch_client_gateway_only(kw, expect_client):
|
||||
assert (client is not None) is expect_client
|
||||
|
||||
|
||||
# --- in-cluster (tier2) backend: gateway-forced + configurable model -----------
|
||||
|
||||
|
||||
def test_build_backend_gateway_only_forces_gateway():
|
||||
"""The in-cluster tier is gateway-only: even with provider=mistral + a key it
|
||||
builds the gateway backend, NEVER the direct Mistral fallback (the GPU is
|
||||
reachable solely through the gateway)."""
|
||||
b = ocr.build_ocr_backend(
|
||||
_settings(
|
||||
document_ocr_provider="mistral",
|
||||
mistral_api_key="k",
|
||||
embedding_gateway_url="https://gw",
|
||||
),
|
||||
gateway_only=True,
|
||||
)
|
||||
assert isinstance(b, ocr._GatewayOcrBackend)
|
||||
|
||||
|
||||
def test_build_backend_gateway_only_no_url_disabled():
|
||||
"""Gateway-only with no gateway URL -> disabled (None), never a direct backend."""
|
||||
assert (
|
||||
ocr.build_ocr_backend(_settings(mistral_api_key="k"), gateway_only=True) is None
|
||||
)
|
||||
|
||||
|
||||
def test_build_backend_gateway_only_provider_none_warns_when_incluster_enabled(caplog):
|
||||
"""provider=none also suppresses the gateway-only in-cluster tier (it never uses
|
||||
the mistral provider, but the global `none` gate still applies). When in-cluster
|
||||
is enabled this is almost certainly an operator mistake -> warn so it's visible."""
|
||||
with caplog.at_level(
|
||||
"WARNING", logger="nextcloud_mcp_server.document_processors.ocr"
|
||||
):
|
||||
b = ocr.build_ocr_backend(
|
||||
_settings(
|
||||
document_ocr_provider="none",
|
||||
embedding_gateway_url="https://gw",
|
||||
document_ocr_incluster_enabled=True,
|
||||
),
|
||||
gateway_only=True,
|
||||
)
|
||||
assert b is None
|
||||
assert any(
|
||||
"DOCUMENT_OCR_PROVIDER=none disables in-cluster" in r.message
|
||||
for r in caplog.records
|
||||
)
|
||||
|
||||
|
||||
def test_build_backend_gateway_only_provider_none_silent_when_incluster_disabled(
|
||||
caplog,
|
||||
):
|
||||
"""provider=none with in-cluster OFF is a deliberate disable -> no warning noise."""
|
||||
with caplog.at_level(
|
||||
"WARNING", logger="nextcloud_mcp_server.document_processors.ocr"
|
||||
):
|
||||
b = ocr.build_ocr_backend(
|
||||
_settings(document_ocr_provider="none", embedding_gateway_url="https://gw"),
|
||||
gateway_only=True,
|
||||
)
|
||||
assert b is None
|
||||
assert not any("disables in-cluster" in r.message for r in caplog.records)
|
||||
|
||||
|
||||
def test_build_backend_empty_model_does_not_fall_back(caplog):
|
||||
"""An empty model string must NOT silently fall back to the upstream default
|
||||
(an `or` would); the in-cluster rung keeps the empty id it was handed."""
|
||||
b = ocr.build_ocr_backend(
|
||||
_settings(embedding_gateway_url="https://gw"),
|
||||
model="",
|
||||
gateway_only=True,
|
||||
)
|
||||
assert isinstance(b, ocr._GatewayOcrBackend)
|
||||
assert b._model == ""
|
||||
|
||||
|
||||
def test_build_backend_model_override_is_not_hardcoded():
|
||||
"""The per-tier model is whatever config passes -- surya by default, but fully
|
||||
swappable (e.g. lightonocr) with no code change."""
|
||||
surya = ocr.build_ocr_backend(
|
||||
_settings(embedding_gateway_url="https://gw"),
|
||||
model="surya/surya-ocr-2",
|
||||
gateway_only=True,
|
||||
)
|
||||
assert isinstance(surya, ocr._GatewayOcrBackend)
|
||||
assert surya._model == "surya/surya-ocr-2"
|
||||
lit = ocr.build_ocr_backend(
|
||||
_settings(embedding_gateway_url="https://gw"),
|
||||
model="lightonocr/lightonocr-1b",
|
||||
gateway_only=True,
|
||||
)
|
||||
assert isinstance(lit, ocr._GatewayOcrBackend)
|
||||
assert lit._model == "lightonocr/lightonocr-1b" # config-driven, not hardcoded
|
||||
|
||||
|
||||
async def test_incluster_processor_resolves_its_model_gateway_only(monkeypatch):
|
||||
"""An OcrProcessor bound to the in-cluster rung builds its backend with its own
|
||||
configured model (document_ocr_incluster_model) and gateway_only=True."""
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
class _FakeBackend:
|
||||
async def ocr(self, content, mime_type):
|
||||
return "ocr text", [{"page": 1, "start_offset": 0, "end_offset": 8}]
|
||||
|
||||
def _spy(settings, *, model=None, gateway_only=False):
|
||||
captured["model"] = model
|
||||
captured["gateway_only"] = gateway_only
|
||||
return _FakeBackend()
|
||||
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", _spy)
|
||||
monkeypatch.setattr(
|
||||
ocr,
|
||||
"get_settings",
|
||||
lambda: _settings(
|
||||
document_ocr_incluster_model="surya/surya-ocr-2",
|
||||
embedding_gateway_url="https://gw",
|
||||
),
|
||||
)
|
||||
proc = ocr.OcrProcessor(
|
||||
name="ocr-incluster",
|
||||
tier="ocr-incluster",
|
||||
model_setting="document_ocr_incluster_model",
|
||||
gateway_only=True,
|
||||
)
|
||||
await proc.process(b"%PDF", "application/pdf", "x.pdf")
|
||||
assert captured == {"model": "surya/surya-ocr-2", "gateway_only": True}
|
||||
|
||||
|
||||
def test_no_surya_string_literal_in_document_processors():
|
||||
"""surya must be a CONFIG default only -- never a hard-coded behavioural literal
|
||||
in the worker (it's swappable, e.g. lightonocr). Comments may mention it; code
|
||||
string literals may not (the default lives in config.py, a different module)."""
|
||||
import pathlib # noqa: PLC0415
|
||||
|
||||
assert ocr.__file__ is not None
|
||||
pkg = pathlib.Path(ocr.__file__).parent
|
||||
offenders = [
|
||||
f"{p.name}: {ln.strip()}"
|
||||
for p in pkg.rglob("*.py") # recurse: future backends/ subdirs too
|
||||
for ln in p.read_text().splitlines()
|
||||
if '"surya' in ln.split("#", 1)[0] or "'surya" in ln.split("#", 1)[0]
|
||||
]
|
||||
assert not offenders, offenders
|
||||
|
||||
|
||||
def test_build_backend_gateway_missing_m2m_raises():
|
||||
# client_id set but token_url/secret missing -> explicit ValueError (not a
|
||||
# stripped assert), surfaced on backend resolution.
|
||||
@@ -106,17 +250,17 @@ def test_build_backend_gateway_missing_m2m_raises():
|
||||
ocr.build_ocr_backend(
|
||||
_settings(
|
||||
document_ocr_provider="gateway",
|
||||
embedding_gateway_url="http://gw",
|
||||
embedding_gateway_url="https://gw",
|
||||
embedding_gateway_client_id="cid",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_gateway_backend_url_normalization():
|
||||
b = ocr._GatewayOcrBackend("http://gw", "mistral/mistral-ocr-latest")
|
||||
assert b._url == "http://gw/v1/ocr"
|
||||
b2 = ocr._GatewayOcrBackend("http://gw/v1/", "m")
|
||||
assert b2._url == "http://gw/v1/ocr"
|
||||
b = ocr._GatewayOcrBackend("https://gw", "mistral/mistral-ocr-latest")
|
||||
assert b._url == "https://gw/v1/ocr"
|
||||
b2 = ocr._GatewayOcrBackend("https://gw/v1/", "m")
|
||||
assert b2._url == "https://gw/v1/ocr"
|
||||
|
||||
|
||||
# --- OcrProcessor ------------------------------------------------------------
|
||||
@@ -126,7 +270,7 @@ async def test_processor_unsupported_when_no_backend(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
ocr, "get_settings", lambda: _settings(document_ocr_provider="none")
|
||||
)
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: None)
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s, **kw: None)
|
||||
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
|
||||
assert r.success is False
|
||||
assert r.metadata["parse_failed_reason"] == "unsupported"
|
||||
@@ -138,12 +282,12 @@ async def test_processor_success(monkeypatch):
|
||||
return "hello world", [{"page": 1, "start_offset": 0, "end_offset": 11}]
|
||||
|
||||
monkeypatch.setattr(ocr, "get_settings", lambda: _settings())
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: _FakeBackend())
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s, **kw: _FakeBackend())
|
||||
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
|
||||
assert r.success is True
|
||||
assert r.text == "hello world"
|
||||
assert r.metadata["page_count"] == 1
|
||||
assert r.processor == "ocr"
|
||||
assert r.processor == "ocr-upstream"
|
||||
|
||||
|
||||
async def test_processor_backend_error_returns_success_false(monkeypatch):
|
||||
@@ -152,7 +296,7 @@ async def test_processor_backend_error_returns_success_false(monkeypatch):
|
||||
raise RuntimeError("api down")
|
||||
|
||||
monkeypatch.setattr(ocr, "get_settings", lambda: _settings())
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: _BoomBackend())
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s, **kw: _BoomBackend())
|
||||
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
|
||||
assert r.success is False
|
||||
assert r.metadata["parse_failed_reason"] == "error"
|
||||
@@ -168,7 +312,7 @@ async def test_processor_timeout_returns_timeout_reason(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
ocr, "get_settings", lambda: _settings(document_ocr_timeout_seconds=5.0)
|
||||
)
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: _TimeoutBackend())
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s, **kw: _TimeoutBackend())
|
||||
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
|
||||
assert r.success is False
|
||||
assert r.metadata["parse_failed_reason"] == "timeout"
|
||||
@@ -187,7 +331,9 @@ async def test_gateway_httpx_timeout_maps_to_timeout_reason(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
ocr, "get_settings", lambda: _settings(document_ocr_timeout_seconds=5.0)
|
||||
)
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: _HttpxTimeoutBackend())
|
||||
monkeypatch.setattr(
|
||||
ocr, "build_ocr_backend", lambda s, **kw: _HttpxTimeoutBackend()
|
||||
)
|
||||
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
|
||||
assert r.success is False
|
||||
assert r.metadata["parse_failed_reason"] == "timeout"
|
||||
@@ -331,7 +477,7 @@ def _wire_batch(monkeypatch, *, client, store, settings=None):
|
||||
embedding_gateway_url="https://gw",
|
||||
)
|
||||
monkeypatch.setattr(ocr, "get_settings", lambda: settings)
|
||||
monkeypatch.setattr(ocr, "build_gateway_batch_client", lambda s: client)
|
||||
monkeypatch.setattr(ocr, "build_gateway_batch_client", lambda s, **kw: client)
|
||||
|
||||
async def _shared(cls):
|
||||
return store
|
||||
@@ -339,6 +485,61 @@ def _wire_batch(monkeypatch, *, client, store, settings=None):
|
||||
monkeypatch.setattr(_bos.BatchOcrJobStore, "shared", classmethod(_shared))
|
||||
|
||||
|
||||
async def test_gateway_only_processor_never_uses_batch_mode(monkeypatch):
|
||||
"""The in-cluster (gateway_only) rung targets the synchronous GPU; batch mode
|
||||
is the upstream Mistral async-job path. _get_batch_client returns None and
|
||||
never builds a batch client even with DOCUMENT_OCR_MODE=batch set globally —
|
||||
while the upstream (gateway_only=False) processor still resolves one."""
|
||||
called = {"n": 0}
|
||||
|
||||
def _spy(settings, **kw):
|
||||
called["n"] += 1
|
||||
return _FakeBatchClient()
|
||||
|
||||
monkeypatch.setattr(
|
||||
ocr,
|
||||
"get_settings",
|
||||
lambda: _settings(
|
||||
document_ocr_mode="batch",
|
||||
document_ocr_provider="gateway",
|
||||
embedding_gateway_url="https://gw",
|
||||
document_ocr_incluster_model="surya/surya-ocr-2",
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(ocr, "build_gateway_batch_client", _spy)
|
||||
|
||||
incluster = ocr.OcrProcessor(
|
||||
name="ocr-incluster",
|
||||
tier="ocr-incluster",
|
||||
model_setting="document_ocr_incluster_model",
|
||||
gateway_only=True,
|
||||
)
|
||||
assert await incluster._get_batch_client() is None
|
||||
assert called["n"] == 0 # short-circuited before building anything
|
||||
|
||||
# _process_batch returns None for the gateway-only rung WITHOUT emitting the
|
||||
# misleading "no gateway backend" warning (the gateway IS configured; the rung
|
||||
# is simply synchronous-only). _batch_fallback_warned stays False to prove it.
|
||||
result = await incluster._process_batch(
|
||||
b"%PDF-1.7",
|
||||
"application/pdf",
|
||||
"x.pdf",
|
||||
dict(_IDENTITY),
|
||||
_settings(
|
||||
document_ocr_mode="batch",
|
||||
document_ocr_provider="gateway",
|
||||
embedding_gateway_url="https://gw",
|
||||
document_ocr_incluster_model="surya/surya-ocr-2",
|
||||
),
|
||||
)
|
||||
assert result is None
|
||||
assert incluster._batch_fallback_warned is False
|
||||
|
||||
upstream = ocr.OcrProcessor() # gateway_only=False
|
||||
assert await upstream._get_batch_client() is not None
|
||||
assert called["n"] == 1
|
||||
|
||||
|
||||
async def test_batch_first_run_submits_and_returns_pending_sentinel(monkeypatch):
|
||||
client = _FakeBatchClient()
|
||||
store = _FakeStore()
|
||||
@@ -451,8 +652,8 @@ async def test_batch_falls_back_to_sync_when_no_gateway(monkeypatch):
|
||||
|
||||
settings = _settings(document_ocr_mode="batch", document_ocr_provider="mistral")
|
||||
monkeypatch.setattr(ocr, "get_settings", lambda: settings)
|
||||
monkeypatch.setattr(ocr, "build_gateway_batch_client", lambda s: None)
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: _FakeBackend())
|
||||
monkeypatch.setattr(ocr, "build_gateway_batch_client", lambda s, **kw: None)
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s, **kw: _FakeBackend())
|
||||
|
||||
r = await ocr.OcrProcessor().process(
|
||||
b"%PDF", "application/pdf", options=dict(_IDENTITY)
|
||||
@@ -472,8 +673,8 @@ async def test_batch_falls_back_to_sync_when_no_identity(monkeypatch):
|
||||
embedding_gateway_url="https://gw",
|
||||
)
|
||||
monkeypatch.setattr(ocr, "get_settings", lambda: settings)
|
||||
monkeypatch.setattr(ocr, "build_gateway_batch_client", lambda s: client)
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: _FakeBackend())
|
||||
monkeypatch.setattr(ocr, "build_gateway_batch_client", lambda s, **kw: client)
|
||||
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s, **kw: _FakeBackend())
|
||||
|
||||
# No options -> inline path -> batch inapplicable -> sync fallback.
|
||||
r = await ocr.OcrProcessor().process(b"%PDF", "application/pdf", options=None)
|
||||
|
||||
@@ -191,7 +191,7 @@ async def test_ocr_tier_records_pages_ocr(store_spy):
|
||||
token_count=900,
|
||||
total_chars=40000,
|
||||
page_count=8,
|
||||
pipeline_tier="ocr",
|
||||
pipeline_tier="ocr-upstream",
|
||||
)
|
||||
by_metric = {
|
||||
c.kwargs["metric"]: c.kwargs["value"]
|
||||
@@ -205,7 +205,7 @@ async def test_ocr_tier_records_pages_ocr(store_spy):
|
||||
}
|
||||
# pipeline_tier is threaded into the billing metadata for CP attribution.
|
||||
for c in store_spy.record_usage_event.await_args_list:
|
||||
assert c.kwargs["metadata"]["pipeline_tier"] == "ocr"
|
||||
assert c.kwargs["metadata"]["pipeline_tier"] == "ocr-upstream"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
|
||||
@@ -74,6 +74,7 @@ class _Settings:
|
||||
engine="pypdfium2",
|
||||
classify=True,
|
||||
ocr=False,
|
||||
ocr_incluster=False,
|
||||
min_text_quality=0.5,
|
||||
page_fraction=0.5,
|
||||
min_page_chars=16,
|
||||
@@ -85,7 +86,13 @@ class _Settings:
|
||||
):
|
||||
self.document_tier1_engine = engine
|
||||
self.document_classify_enabled = classify
|
||||
# ``ocr`` enables the UPSTREAM rung (document_ocr_enabled); ``ocr_incluster``
|
||||
# the in-cluster GPU rung (tried first). The model attrs let
|
||||
# build_ocr_backend resolve a per-tier model id.
|
||||
self.document_ocr_enabled = ocr
|
||||
self.document_ocr_incluster_enabled = ocr_incluster
|
||||
self.document_ocr_model = "mistral/mistral-ocr-latest"
|
||||
self.document_ocr_incluster_model = "surya/surya-ocr-2"
|
||||
self.document_ocr_min_text_quality = min_text_quality
|
||||
self.document_ocr_page_fraction = page_fraction
|
||||
self.document_ocr_min_page_chars = min_page_chars
|
||||
@@ -195,10 +202,10 @@ async def test_ocr_escalation_on_empty_text(monkeypatch):
|
||||
monkeypatch.setattr(reg_mod, "record_document_escalation", esc)
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text=""), 20),
|
||||
(_Fake("ocr", "ocr", text="ocr text"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream", text="ocr text"), 5),
|
||||
)
|
||||
res = await r.process(b"%PDF-1.7", "application/pdf")
|
||||
assert res.processor == "ocr"
|
||||
assert res.processor == "ocr-upstream"
|
||||
esc.assert_called_once()
|
||||
|
||||
|
||||
@@ -210,7 +217,7 @@ async def test_zero_page_pdf_does_not_escalate(monkeypatch):
|
||||
monkeypatch.setattr(reg_mod, "record_document_escalation", esc)
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text="", pages=0), 20),
|
||||
(_Fake("ocr", "ocr"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
res = await r.process(b"%PDF-1.7", "application/pdf")
|
||||
assert res.processor == "fast"
|
||||
@@ -231,7 +238,7 @@ async def test_ocr_failure_falls_back_to_fast(monkeypatch):
|
||||
monkeypatch.setattr(reg_mod, "record_document_escalation", MagicMock())
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text=""), 20),
|
||||
(_Fake("ocr", "ocr", text="", success=False), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream", text="", success=False), 5),
|
||||
)
|
||||
res = await r.process(b"%PDF-1.7", "application/pdf")
|
||||
assert res.processor == "fast"
|
||||
@@ -242,7 +249,7 @@ async def test_no_ocr_escalation_when_disabled(monkeypatch):
|
||||
monkeypatch.setattr(reg_mod, "get_settings", lambda: _Settings(ocr=False))
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text=""), 20),
|
||||
(_Fake("ocr", "ocr"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
res = await r.process(b"%PDF-1.7", "application/pdf")
|
||||
# Fast tier is terminal when OCR is disabled.
|
||||
@@ -289,11 +296,11 @@ async def test_glyph_corrupt_no_structured_falls_through_to_ocr(monkeypatch):
|
||||
monkeypatch.setattr(reg_mod, "record_document_escalation", esc)
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text=_GLYPH), 20),
|
||||
(_Fake("ocr", "ocr", text="ocr recovered text"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream", text="ocr recovered text"), 5),
|
||||
) # no structured registered
|
||||
res = await r.process(b"%PDF-1.7", "application/pdf")
|
||||
assert res.processor == "ocr"
|
||||
esc.assert_called_once_with("fast", "ocr", "corrupt_glyphs")
|
||||
assert res.processor == "ocr-upstream"
|
||||
esc.assert_called_once_with("fast", "ocr-upstream", "corrupt_glyphs")
|
||||
|
||||
|
||||
async def test_inline_lowconf_tries_structured_before_ocr(monkeypatch):
|
||||
@@ -305,7 +312,7 @@ async def test_inline_lowconf_tries_structured_before_ocr(monkeypatch):
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text="x" * 40), 20), # one long token -> quality ~0
|
||||
(_Fake("structured", "structured", text="clean recovered prose text here"), 10),
|
||||
(_Fake("ocr", "ocr", text="ocr text"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream", text="ocr text"), 5),
|
||||
)
|
||||
res = await r.process(b"%PDF-1.7", "application/pdf")
|
||||
assert res.processor == "structured"
|
||||
@@ -321,11 +328,11 @@ async def test_inline_empty_skips_structured_straight_to_ocr(monkeypatch):
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text=""), 20),
|
||||
(_Fake("structured", "structured", text="should not run"), 10),
|
||||
(_Fake("ocr", "ocr", text="ocr text"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream", text="ocr text"), 5),
|
||||
)
|
||||
res = await r.process(b"%PDF-1.7", "application/pdf")
|
||||
assert res.processor == "ocr"
|
||||
esc.assert_called_once_with("fast", "ocr", "empty_text")
|
||||
assert res.processor == "ocr-upstream"
|
||||
esc.assert_called_once_with("fast", "ocr-upstream", "empty_text")
|
||||
|
||||
|
||||
async def test_inline_fast_structured_ocr_cascade(monkeypatch):
|
||||
@@ -339,13 +346,13 @@ async def test_inline_fast_structured_ocr_cascade(monkeypatch):
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text="x" * 40), 20), # quality ~0, non-empty
|
||||
(_Fake("structured", "structured", text=""), 10), # re-extract empty
|
||||
(_Fake("ocr", "ocr", text="ocr recovered text"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream", text="ocr recovered text"), 5),
|
||||
)
|
||||
res = await r.process(b"%PDF-1.7", "application/pdf")
|
||||
assert res.processor == "ocr"
|
||||
assert res.processor == "ocr-upstream"
|
||||
assert esc.call_args_list == [
|
||||
call("fast", "structured", "low_confidence"),
|
||||
call("structured", "ocr", "empty_text"),
|
||||
call("structured", "ocr-upstream", "empty_text"),
|
||||
]
|
||||
|
||||
|
||||
@@ -359,13 +366,13 @@ async def test_inline_structured_still_corrupt_escalates_to_ocr(monkeypatch):
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text=_GLYPH), 20),
|
||||
(_Fake("structured", "structured", text=_GLYPH), 10), # still corrupt
|
||||
(_Fake("ocr", "ocr", text="ocr recovered text"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream", text="ocr recovered text"), 5),
|
||||
)
|
||||
res = await r.process(b"%PDF-1.7", "application/pdf")
|
||||
assert res.processor == "ocr"
|
||||
assert res.processor == "ocr-upstream"
|
||||
assert esc.call_args_list == [
|
||||
call("fast", "structured", "corrupt_glyphs"),
|
||||
call("structured", "ocr", "corrupt_glyphs"),
|
||||
call("structured", "ocr-upstream", "corrupt_glyphs"),
|
||||
]
|
||||
|
||||
|
||||
@@ -375,7 +382,7 @@ def test_evaluate_escalation_glyph_corrupt_goes_structured(monkeypatch):
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("structured", "structured"), 10),
|
||||
(_Fake("ocr", "ocr"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
res = ProcessingResult(
|
||||
text=_GLYPH,
|
||||
@@ -399,7 +406,7 @@ def test_evaluate_escalation_glyph_corrupt_no_structured_falls_through_to_ocr(
|
||||
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("ocr", "ocr"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
) # no structured registered
|
||||
res = ProcessingResult(
|
||||
text=_GLYPH,
|
||||
@@ -412,7 +419,7 @@ def test_evaluate_escalation_glyph_corrupt_no_structured_falls_through_to_ocr(
|
||||
processor="fast",
|
||||
)
|
||||
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=True))
|
||||
assert decision == EscalationDecision("hop", "ocr", "corrupt_glyphs")
|
||||
assert decision == EscalationDecision("hop", "ocr-upstream", "corrupt_glyphs")
|
||||
|
||||
|
||||
def test_evaluate_escalation_glyph_corrupt_no_structured_ocr_disabled_suppressed(
|
||||
@@ -424,7 +431,7 @@ def test_evaluate_escalation_glyph_corrupt_no_structured_ocr_disabled_suppressed
|
||||
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("ocr", "ocr"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
) # structured not registered; ocr registered but disabled below
|
||||
res = ProcessingResult(
|
||||
text=_GLYPH,
|
||||
@@ -437,7 +444,9 @@ def test_evaluate_escalation_glyph_corrupt_no_structured_ocr_disabled_suppressed
|
||||
processor="fast",
|
||||
)
|
||||
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=False))
|
||||
assert decision == EscalationDecision("suppressed", "ocr", "corrupt_glyphs")
|
||||
assert decision == EscalationDecision(
|
||||
"suppressed", "ocr-upstream", "corrupt_glyphs"
|
||||
)
|
||||
|
||||
|
||||
# --- Per-tier external path (Deck #323) -------------------------------------
|
||||
@@ -449,7 +458,7 @@ async def test_process_tier_runs_named_tier(monkeypatch):
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("structured", "structured"), 10),
|
||||
(_Fake("ocr", "ocr"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
res = await r.process_tier(b"%PDF-1.7", "application/pdf", "f.pdf", "structured")
|
||||
assert res.processor == "structured"
|
||||
@@ -469,8 +478,10 @@ async def test_process_tier_oversize_fails_fast(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
reg_mod, "get_settings", lambda: _Settings(max_pdf_size_mb=0.001)
|
||||
)
|
||||
r = _registry((_Fake("ocr", "ocr"), 5))
|
||||
res = await r.process_tier(b"x" * 4096, "application/pdf", "big.pdf", "ocr")
|
||||
r = _registry((_Fake("ocr-upstream", "ocr-upstream"), 5))
|
||||
res = await r.process_tier(
|
||||
b"x" * 4096, "application/pdf", "big.pdf", "ocr-upstream"
|
||||
)
|
||||
assert res.success is False
|
||||
assert res.metadata["parse_failed_reason"] == "oversize"
|
||||
|
||||
@@ -479,7 +490,7 @@ def test_next_available_tier_walks_ladder():
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("structured", "structured"), 10),
|
||||
(_Fake("ocr", "ocr"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
# ocr disabled -> structured is the only target above fast.
|
||||
s = _Settings(ocr=False)
|
||||
@@ -487,21 +498,25 @@ def test_next_available_tier_walks_ladder():
|
||||
assert r.next_available_tier("structured", s) is None # ocr gated off
|
||||
# ocr enabled -> reachable; minimum skips the structured rung.
|
||||
s_ocr = _Settings(ocr=True)
|
||||
assert r.next_available_tier("structured", s_ocr) == "ocr"
|
||||
assert r.next_available_tier("fast", s_ocr, minimum="ocr") == "ocr"
|
||||
assert r.next_available_tier("structured", s_ocr) == "ocr-upstream"
|
||||
assert (
|
||||
r.next_available_tier("fast", s_ocr, minimum="ocr-upstream") == "ocr-upstream"
|
||||
)
|
||||
|
||||
|
||||
def test_next_available_tier_skips_unregistered():
|
||||
# No structured processor -> fast escalates straight to ocr.
|
||||
r = _registry((_Fake("fast", "fast"), 20), (_Fake("ocr", "ocr"), 5))
|
||||
assert r.next_available_tier("fast", _Settings(ocr=True)) == "ocr"
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20), (_Fake("ocr-upstream", "ocr-upstream"), 5)
|
||||
)
|
||||
assert r.next_available_tier("fast", _Settings(ocr=True)) == "ocr-upstream"
|
||||
|
||||
|
||||
def test_evaluate_escalation_good_text_indexes(monkeypatch):
|
||||
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text="This is clean readable prose text."), 20),
|
||||
(_Fake("ocr", "ocr"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
res = ProcessingResult(
|
||||
text="This is clean readable prose text.",
|
||||
@@ -520,7 +535,7 @@ def test_evaluate_escalation_empty_jumps_to_ocr(monkeypatch):
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("structured", "structured"), 10),
|
||||
(_Fake("ocr", "ocr"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
res = ProcessingResult(
|
||||
text="",
|
||||
@@ -531,7 +546,7 @@ def test_evaluate_escalation_empty_jumps_to_ocr(monkeypatch):
|
||||
processor="fast",
|
||||
)
|
||||
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=True))
|
||||
assert decision == EscalationDecision("hop", "ocr", "empty_text")
|
||||
assert decision == EscalationDecision("hop", "ocr-upstream", "empty_text")
|
||||
|
||||
|
||||
def test_evaluate_escalation_lowconf_goes_to_structured(monkeypatch):
|
||||
@@ -541,7 +556,7 @@ def test_evaluate_escalation_lowconf_goes_to_structured(monkeypatch):
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("structured", "structured"), 10),
|
||||
(_Fake("ocr", "ocr"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
res = ProcessingResult(
|
||||
text=junk,
|
||||
@@ -559,7 +574,9 @@ def test_evaluate_escalation_lowconf_goes_to_structured(monkeypatch):
|
||||
|
||||
def test_evaluate_escalation_failure_not_escalated(monkeypatch):
|
||||
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
|
||||
r = _registry((_Fake("fast", "fast"), 20), (_Fake("ocr", "ocr"), 5))
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20), (_Fake("ocr-upstream", "ocr-upstream"), 5)
|
||||
)
|
||||
res = ProcessingResult(
|
||||
text="",
|
||||
metadata={"parse_failed_reason": "error"},
|
||||
@@ -590,7 +607,9 @@ def test_evaluate_escalation_terminal_when_no_higher_tier(monkeypatch):
|
||||
def test_evaluate_escalation_zero_page_does_not_escalate(monkeypatch):
|
||||
"""A zero-page (empty/corrupt) PDF never escalates on the external path."""
|
||||
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
|
||||
r = _registry((_Fake("fast", "fast"), 20), (_Fake("ocr", "ocr"), 5))
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20), (_Fake("ocr-upstream", "ocr-upstream"), 5)
|
||||
)
|
||||
res = ProcessingResult(
|
||||
text="",
|
||||
metadata={"page_count": 0, "page_boundaries": []},
|
||||
@@ -604,7 +623,9 @@ def test_evaluate_escalation_lowconf_to_ocr_when_no_structured(monkeypatch):
|
||||
unregistered structured rung), not to None."""
|
||||
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
|
||||
junk = "z" * 40 # non-empty but junk -> recommended ocr, total_chars > 0
|
||||
r = _registry((_Fake("fast", "fast"), 20), (_Fake("ocr", "ocr"), 5))
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20), (_Fake("ocr-upstream", "ocr-upstream"), 5)
|
||||
)
|
||||
res = ProcessingResult(
|
||||
text=junk,
|
||||
metadata={
|
||||
@@ -616,14 +637,16 @@ def test_evaluate_escalation_lowconf_to_ocr_when_no_structured(monkeypatch):
|
||||
processor="fast",
|
||||
)
|
||||
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=True))
|
||||
assert decision == EscalationDecision("hop", "ocr", "low_confidence")
|
||||
assert decision == EscalationDecision("hop", "ocr-upstream", "low_confidence")
|
||||
|
||||
|
||||
def test_evaluate_escalation_suppressed_when_ocr_disabled(monkeypatch):
|
||||
"""OCR off: a scanned doc does NOT hop to ocr; it returns a 'suppressed'
|
||||
decision (the what-if-OCR signal) so the caller indexes at the current tier."""
|
||||
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
|
||||
r = _registry((_Fake("fast", "fast"), 20), (_Fake("ocr", "ocr"), 5))
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20), (_Fake("ocr-upstream", "ocr-upstream"), 5)
|
||||
)
|
||||
res = ProcessingResult(
|
||||
text="",
|
||||
metadata={
|
||||
@@ -633,7 +656,32 @@ def test_evaluate_escalation_suppressed_when_ocr_disabled(monkeypatch):
|
||||
processor="fast",
|
||||
)
|
||||
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=False))
|
||||
assert decision == EscalationDecision("suppressed", "ocr", "empty_text")
|
||||
assert decision == EscalationDecision("suppressed", "ocr-upstream", "empty_text")
|
||||
|
||||
|
||||
def test_evaluate_escalation_suppressed_targets_incluster_four_rung(monkeypatch):
|
||||
"""Four-rung registry, BOTH OCR flags off: the suppressed what-if-OCR signal
|
||||
names the *cheapest* ideal rung (ocr-incluster), not ocr-upstream, since the
|
||||
ideal-target walk (ignore_ocr_enabled) picks the cheapest registered OCR rung."""
|
||||
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("structured", "structured"), 10),
|
||||
(_Fake("ocr-incluster", "ocr-incluster"), 6),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
res = ProcessingResult(
|
||||
text="",
|
||||
metadata={
|
||||
"page_count": 1,
|
||||
"page_boundaries": [{"page": 1, "start_offset": 0, "end_offset": 0}],
|
||||
},
|
||||
processor="fast",
|
||||
)
|
||||
decision = r.evaluate_escalation(
|
||||
res, b"%PDF", "fast", _Settings(ocr=False, ocr_incluster=False)
|
||||
)
|
||||
assert decision == EscalationDecision("suppressed", "ocr-incluster", "empty_text")
|
||||
|
||||
|
||||
def test_evaluate_escalation_lowconf_suppressed_when_only_ocr_disabled(monkeypatch):
|
||||
@@ -641,7 +689,9 @@ def test_evaluate_escalation_lowconf_suppressed_when_only_ocr_disabled(monkeypat
|
||||
the would-be hop is suppressed (not a structured hop, which isn't registered)."""
|
||||
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
|
||||
junk = "q" * 40
|
||||
r = _registry((_Fake("fast", "fast"), 20), (_Fake("ocr", "ocr"), 5))
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20), (_Fake("ocr-upstream", "ocr-upstream"), 5)
|
||||
)
|
||||
res = ProcessingResult(
|
||||
text=junk,
|
||||
metadata={
|
||||
@@ -653,7 +703,9 @@ def test_evaluate_escalation_lowconf_suppressed_when_only_ocr_disabled(monkeypat
|
||||
processor="fast",
|
||||
)
|
||||
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=False))
|
||||
assert decision == EscalationDecision("suppressed", "ocr", "low_confidence")
|
||||
assert decision == EscalationDecision(
|
||||
"suppressed", "ocr-upstream", "low_confidence"
|
||||
)
|
||||
|
||||
|
||||
def test_evaluate_escalation_structured_hop_not_suppressed_when_ocr_off(monkeypatch):
|
||||
@@ -664,7 +716,7 @@ def test_evaluate_escalation_structured_hop_not_suppressed_when_ocr_off(monkeypa
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("structured", "structured"), 10),
|
||||
(_Fake("ocr", "ocr"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
res = ProcessingResult(
|
||||
text=junk,
|
||||
@@ -699,14 +751,15 @@ def test_evaluate_escalation_terminal_when_ocr_unregistered_and_off(monkeypatch)
|
||||
def test_evaluate_escalation_empty_suppressed_even_when_structured_registered(
|
||||
monkeypatch,
|
||||
):
|
||||
"""empty_text uses minimum='ocr', so it skips structured even when structured
|
||||
IS registered: with OCR off it suppresses to ocr, never hops to structured
|
||||
(a text extractor can't conjure text from a raster scan)."""
|
||||
"""empty_text uses minimum='ocr-incluster', so it skips structured even when
|
||||
structured IS registered: with OCR off it suppresses to the cheapest registered
|
||||
OCR rung (here ocr-upstream, the only one registered in this test), never hops
|
||||
to structured (a text extractor can't conjure text from a raster scan)."""
|
||||
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("structured", "structured"), 10), # registered but skipped for empty
|
||||
(_Fake("ocr", "ocr"), 5),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
res = ProcessingResult(
|
||||
text="",
|
||||
@@ -717,4 +770,159 @@ def test_evaluate_escalation_empty_suppressed_even_when_structured_registered(
|
||||
processor="fast",
|
||||
)
|
||||
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=False))
|
||||
assert decision == EscalationDecision("suppressed", "ocr", "empty_text")
|
||||
assert decision == EscalationDecision("suppressed", "ocr-upstream", "empty_text")
|
||||
|
||||
|
||||
# --- tier2 in-cluster OCR rung (Deck #353) -----------------------------------
|
||||
|
||||
|
||||
async def test_inline_empty_routes_to_ocr_incluster_before_upstream(monkeypatch):
|
||||
"""Both OCR rungs enabled + registered: an empty text layer hops to the
|
||||
in-cluster (tier2) rung FIRST, never straight to the paid upstream rung."""
|
||||
monkeypatch.setattr(
|
||||
reg_mod, "get_settings", lambda: _Settings(ocr=True, ocr_incluster=True)
|
||||
)
|
||||
esc = MagicMock()
|
||||
monkeypatch.setattr(reg_mod, "record_document_escalation", esc)
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text=""), 20),
|
||||
(_Fake("structured", "structured", text="should not run"), 10),
|
||||
(_Fake("ocr-incluster", "ocr-incluster", text="incluster ocr text"), 6),
|
||||
(_Fake("ocr-upstream", "ocr-upstream", text="upstream ocr text"), 5),
|
||||
)
|
||||
res = await r.process(b"%PDF-1.7", "application/pdf")
|
||||
assert res.processor == "ocr-incluster"
|
||||
esc.assert_called_once_with("fast", "ocr-incluster", "empty_text")
|
||||
|
||||
|
||||
async def test_inline_only_incluster_enabled_routes_to_incluster(monkeypatch):
|
||||
"""Tenant with ONLY in-cluster OCR on (upstream off): empty text still hops
|
||||
to the in-cluster rung (its own enable flag gates it, independent of upstream)."""
|
||||
monkeypatch.setattr(
|
||||
reg_mod, "get_settings", lambda: _Settings(ocr=False, ocr_incluster=True)
|
||||
)
|
||||
esc = MagicMock()
|
||||
monkeypatch.setattr(reg_mod, "record_document_escalation", esc)
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text=""), 20),
|
||||
(_Fake("ocr-incluster", "ocr-incluster", text="incluster ocr text"), 6),
|
||||
(_Fake("ocr-upstream", "ocr-upstream", text="upstream ocr text"), 5),
|
||||
)
|
||||
res = await r.process(b"%PDF-1.7", "application/pdf")
|
||||
assert res.processor == "ocr-incluster"
|
||||
esc.assert_called_once_with("fast", "ocr-incluster", "empty_text")
|
||||
|
||||
|
||||
async def test_inline_incluster_disabled_skips_to_upstream(monkeypatch):
|
||||
"""In-cluster registered but DISABLED, upstream enabled: the disabled tier2
|
||||
rung is skipped and the job escalates to the upstream rung."""
|
||||
monkeypatch.setattr(
|
||||
reg_mod, "get_settings", lambda: _Settings(ocr=True, ocr_incluster=False)
|
||||
)
|
||||
esc = MagicMock()
|
||||
monkeypatch.setattr(reg_mod, "record_document_escalation", esc)
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text=""), 20),
|
||||
(_Fake("ocr-incluster", "ocr-incluster", text="incluster ocr text"), 6),
|
||||
(_Fake("ocr-upstream", "ocr-upstream", text="upstream ocr text"), 5),
|
||||
)
|
||||
res = await r.process(b"%PDF-1.7", "application/pdf")
|
||||
assert res.processor == "ocr-upstream"
|
||||
esc.assert_called_once_with("fast", "ocr-upstream", "empty_text")
|
||||
|
||||
|
||||
async def test_inline_incluster_failure_falls_back_to_fast_not_upstream(monkeypatch):
|
||||
"""CURRENT behavior (pins the known follow-up gap): when the chosen in-cluster
|
||||
rung runs but FAILS (e.g. GPU 503 -> success=False), the inline path keeps the
|
||||
tier-1 result rather than cascading to the paid upstream rung. OCR is an
|
||||
enhancement, not a gate. (A future change will escalate transient GPU failures
|
||||
to ocr-upstream; this test makes that diff explicit.)"""
|
||||
monkeypatch.setattr(
|
||||
reg_mod, "get_settings", lambda: _Settings(ocr=True, ocr_incluster=True)
|
||||
)
|
||||
monkeypatch.setattr(reg_mod, "record_document_escalation", MagicMock())
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast", text=""), 20),
|
||||
(_Fake("ocr-incluster", "ocr-incluster", text="", success=False), 6),
|
||||
(_Fake("ocr-upstream", "ocr-upstream", text="upstream ocr text"), 5),
|
||||
)
|
||||
res = await r.process(b"%PDF-1.7", "application/pdf")
|
||||
assert res.processor == "fast"
|
||||
assert res.success is True
|
||||
|
||||
|
||||
def _empty_result() -> ProcessingResult:
|
||||
return ProcessingResult(
|
||||
text="",
|
||||
metadata={
|
||||
"page_count": 1,
|
||||
"page_boundaries": [{"page": 1, "start_offset": 0, "end_offset": 0}],
|
||||
},
|
||||
processor="fast",
|
||||
)
|
||||
|
||||
|
||||
def test_evaluate_escalation_empty_text_hops_to_incluster(monkeypatch):
|
||||
"""External path: empty text targets the cheapest OCR rung (in-cluster, tier2)
|
||||
when it is enabled + registered."""
|
||||
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("structured", "structured"), 10),
|
||||
(_Fake("ocr-incluster", "ocr-incluster"), 6),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
decision = r.evaluate_escalation(
|
||||
_empty_result(), b"%PDF", "fast", _Settings(ocr=True, ocr_incluster=True)
|
||||
)
|
||||
assert decision == EscalationDecision("hop", "ocr-incluster", "empty_text")
|
||||
|
||||
|
||||
def test_evaluate_escalation_incluster_disabled_hops_to_upstream(monkeypatch):
|
||||
"""External path: in-cluster disabled -> next_available_tier falls through to
|
||||
the upstream rung (a real hop, not a suppression, since upstream is on)."""
|
||||
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("ocr-incluster", "ocr-incluster"), 6),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
decision = r.evaluate_escalation(
|
||||
_empty_result(), b"%PDF", "fast", _Settings(ocr=True, ocr_incluster=False)
|
||||
)
|
||||
assert decision == EscalationDecision("hop", "ocr-upstream", "empty_text")
|
||||
|
||||
|
||||
def test_next_available_tier_walks_full_four_rung_ladder():
|
||||
"""next_available_tier walks fast -> structured -> ocr-incluster ->
|
||||
ocr-upstream when every rung is registered + enabled."""
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("structured", "structured"), 10),
|
||||
(_Fake("ocr-incluster", "ocr-incluster"), 6),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
s = _Settings(ocr=True, ocr_incluster=True)
|
||||
assert r.next_available_tier("fast", s) == "structured"
|
||||
assert r.next_available_tier("structured", s) == "ocr-incluster"
|
||||
assert r.next_available_tier("ocr-incluster", s) == "ocr-upstream"
|
||||
assert r.next_available_tier("ocr-upstream", s) is None
|
||||
# minimum pins the floor: from fast with minimum=ocr-incluster skips structured.
|
||||
assert r.next_available_tier("fast", s, minimum="ocr-incluster") == "ocr-incluster"
|
||||
|
||||
|
||||
def test_next_available_tier_incluster_disabled_skips_to_upstream():
|
||||
"""A disabled in-cluster rung is skipped; the walk lands on the upstream rung."""
|
||||
r = _registry(
|
||||
(_Fake("fast", "fast"), 20),
|
||||
(_Fake("structured", "structured"), 10),
|
||||
(_Fake("ocr-incluster", "ocr-incluster"), 6),
|
||||
(_Fake("ocr-upstream", "ocr-upstream"), 5),
|
||||
)
|
||||
s = _Settings(ocr=True, ocr_incluster=False)
|
||||
assert r.next_available_tier("structured", s) == "ocr-upstream"
|
||||
# ...but the *ideal* target ignoring the enable gate is still in-cluster.
|
||||
assert (
|
||||
r.next_available_tier("structured", s, ignore_ocr_enabled=True)
|
||||
== "ocr-incluster"
|
||||
)
|
||||
|
||||
@@ -29,6 +29,7 @@ pytestmark = pytest.mark.unit
|
||||
def _settings(*, ocr_enabled: bool) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
document_ocr_enabled=ocr_enabled,
|
||||
document_ocr_incluster_enabled=False,
|
||||
document_tier1_engine="pypdfium2",
|
||||
get_collection_name=lambda: "c",
|
||||
)
|
||||
|
||||
@@ -115,7 +115,7 @@ class TestProcessDocumentTask:
|
||||
# Calling the Task runs its wrapped function in-process. The job is on the
|
||||
# ocr queue, so the queue-aware task must derive tier="ocr".
|
||||
await pq.process_document_task(
|
||||
_ctx(pq.INGEST_QUEUE_OCR),
|
||||
_ctx(pq.INGEST_QUEUE_OCR_UPSTREAM),
|
||||
user_id="alice",
|
||||
doc_id="42",
|
||||
doc_type="note",
|
||||
@@ -131,7 +131,7 @@ class TestProcessDocumentTask:
|
||||
# Worker disables the in-process retry loop; durable retry is the queue's.
|
||||
assert captured["max_retries"] == 1
|
||||
# Tier is derived from the job's queue (escalation enabled by default).
|
||||
assert captured["tier"] == "ocr"
|
||||
assert captured["tier"] == "ocr-upstream"
|
||||
fake_client.close.assert_awaited_once()
|
||||
|
||||
async def test_pipeline_error_propagates_and_closes_client(self, monkeypatch):
|
||||
|
||||
@@ -36,19 +36,25 @@ def _job(queue: str = pq.INGEST_QUEUE_FAST, attempts: int = 1) -> Job:
|
||||
class TestLadder:
|
||||
def test_next_tier_ordering(self):
|
||||
assert next_tier("fast") == "structured"
|
||||
assert next_tier("structured") == "ocr"
|
||||
assert next_tier("ocr") is None # terminal
|
||||
assert next_tier("structured") == "ocr-incluster"
|
||||
assert next_tier("ocr-incluster") == "ocr-upstream"
|
||||
assert next_tier("ocr-upstream") is None # terminal
|
||||
assert next_tier("unknown") is None
|
||||
|
||||
def test_ladder_is_cheapest_first(self):
|
||||
assert TIER_LADDER == ("fast", "structured", "ocr")
|
||||
assert TIER_LADDER == ("fast", "structured", "ocr-incluster", "ocr-upstream")
|
||||
|
||||
def test_tier_for_queue(self):
|
||||
assert pq.tier_for_queue(pq.INGEST_QUEUE_OCR) == "ocr"
|
||||
assert pq.tier_for_queue(pq.INGEST_QUEUE_OCR_INCLUSTER) == "ocr-incluster"
|
||||
assert pq.tier_for_queue(pq.INGEST_QUEUE_OCR_UPSTREAM) == "ocr-upstream"
|
||||
assert pq.tier_for_queue(pq.INGEST_QUEUE_STRUCTURED) == "structured"
|
||||
# Legacy / unknown / None all fall back to the cheapest tier.
|
||||
assert pq.tier_for_queue(pq.LEGACY_INGEST_QUEUE) == "fast"
|
||||
assert pq.tier_for_queue(None) == "fast"
|
||||
# The pre-split legacy OCR queue also resolves to fast (NOT ocr-upstream):
|
||||
# stranded in-flight jobs re-extract empty and re-escalate via the ladder
|
||||
# to the cheap ocr-incluster rung, never straight to paid upstream.
|
||||
assert pq.tier_for_queue(pq.LEGACY_INGEST_QUEUE_OCR) == "fast"
|
||||
|
||||
|
||||
class TestTieredEscalationStrategy:
|
||||
@@ -56,10 +62,22 @@ class TestTieredEscalationStrategy:
|
||||
return pq.TieredEscalationStrategy(max_transient_attempts=max_transient)
|
||||
|
||||
def test_escalate_hops_to_target_queue(self):
|
||||
exc = EscalateError(from_tier="fast", to_tier="ocr", reason="empty_text")
|
||||
exc = EscalateError(
|
||||
from_tier="fast", to_tier="ocr-upstream", reason="empty_text"
|
||||
)
|
||||
decision = self._strategy().get_retry_decision(exception=exc, job=_job())
|
||||
assert decision is not None
|
||||
assert decision.queue == pq.INGEST_QUEUE_OCR
|
||||
assert decision.queue == pq.INGEST_QUEUE_OCR_UPSTREAM
|
||||
|
||||
def test_escalate_hops_to_incluster_queue(self):
|
||||
# tier2 in-cluster (Deck #353): structured -> ocr-incluster lands on the
|
||||
# in-cluster queue the GPU sentinel watches, not the paid upstream queue.
|
||||
exc = EscalateError(
|
||||
from_tier="structured", to_tier="ocr-incluster", reason="empty_text"
|
||||
)
|
||||
decision = self._strategy().get_retry_decision(exception=exc, job=_job())
|
||||
assert decision is not None
|
||||
assert decision.queue == pq.INGEST_QUEUE_OCR_INCLUSTER
|
||||
|
||||
def test_escalate_to_structured(self):
|
||||
exc = EscalateError(
|
||||
@@ -75,11 +93,13 @@ class TestTieredEscalationStrategy:
|
||||
assert decision is None
|
||||
|
||||
def test_escalate_unwraps_exception_group(self):
|
||||
exc = EscalateError(from_tier="fast", to_tier="ocr", reason="empty_text")
|
||||
exc = EscalateError(
|
||||
from_tier="fast", to_tier="ocr-upstream", reason="empty_text"
|
||||
)
|
||||
group = ExceptionGroup("wrapped", [exc])
|
||||
decision = self._strategy().get_retry_decision(exception=group, job=_job())
|
||||
assert decision is not None
|
||||
assert decision.queue == pq.INGEST_QUEUE_OCR
|
||||
assert decision.queue == pq.INGEST_QUEUE_OCR_UPSTREAM
|
||||
|
||||
def test_transient_retries_same_queue_under_cap(self):
|
||||
decision = self._strategy(max_transient=5).get_retry_decision(
|
||||
@@ -126,7 +146,8 @@ class TestTieredEscalationStrategy:
|
||||
# Batch OCR re-poll (Deck #332): same-queue deferral after retry_in.
|
||||
before = datetime.now(timezone.utc)
|
||||
decision = self._strategy().get_retry_decision(
|
||||
exception=BatchPending(retry_in=120), job=_job(queue=pq.INGEST_QUEUE_OCR)
|
||||
exception=BatchPending(retry_in=120),
|
||||
job=_job(queue=pq.INGEST_QUEUE_OCR_UPSTREAM),
|
||||
)
|
||||
after = datetime.now(timezone.utc)
|
||||
assert decision is not None
|
||||
|
||||
Reference in New Issue
Block a user