feat(ingest): split OCR into tier2 in-cluster (GPU, gateway-only) + tier3 upstream
Insert a configurable in-cluster OCR rung into the escalation ladder (Deck #353): a tier2-eligible doc is OCR'd on the on-demand burst GPU before falling through to paid upstream OCR. The in-cluster backend is reached ONLY via the embedding gateway (model prefix routes to the GPU over the tailnet) and is a config value (default surya/surya-ocr-2, swappable to e.g. lightonocr) — never hard-coded. Ladder: fast -> structured -> ocr-incluster -> ocr-upstream (queues ingest-ocr-incluster / ingest-ocr-upstream). - escalation.py: 4-tier ladder; in-cluster flag folded into the dead-letter signature. - ocr.py: OcrProcessor(name, tier, model_setting, gateway_only); build_ocr_backend( ..., model=, gateway_only=) — gateway_only forces the gateway backend (never the direct Mistral fallback), disabling the tier with a warning if no gateway URL. - registry.py: per-rung enable map; scanned docs target minimum="ocr-incluster"; inline path runs the cheapest available OCR rung. - procrastinate.py: two OCR queues; legacy ingest-ocr kept as a drain target. - config.py: DOCUMENT_OCR_INCLUSTER_ENABLED (off) + DOCUMENT_OCR_INCLUSTER_MODEL. - __init__.py: register the two OCR instances; vector/processor.py: pages_ocr metered for the upstream (paid) rung only; cli.py: new --tier choices + legacy drain. - metrics.py: zero the legacy ingest-ocr queue gauge during rollout. - tests: migrated to the split ladder + new tests (gateway-only forcing, per-tier model incl. lightonocr override, no-hard-coded-surya guard). 1792 pass; ruff + ty green. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Chris Coutinho
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
060084029f
commit
c21804fbbc
@@ -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",
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
|
||||
|
||||
@@ -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,52 @@ 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
|
||||
model = model or settings.document_ocr_model
|
||||
if provider == "none":
|
||||
return None
|
||||
|
||||
if gateway_only:
|
||||
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 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,
|
||||
)
|
||||
|
||||
@@ -260,7 +289,24 @@ def build_ocr_backend(settings: Settings) -> _OcrBackend | None:
|
||||
class OcrProcessor(DocumentProcessor):
|
||||
"""Tier-3 OCR processor (gateway or direct Mistral backend)."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
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.
|
||||
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 +329,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 +373,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:
|
||||
@@ -393,7 +443,10 @@ class OcrProcessor(DocumentProcessor):
|
||||
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())
|
||||
self._batch_client = build_gateway_batch_client(
|
||||
get_settings(),
|
||||
model=getattr(get_settings(), self._model_setting),
|
||||
)
|
||||
self._batch_client_resolved = True
|
||||
return self._batch_client
|
||||
|
||||
|
||||
@@ -349,10 +349,18 @@ 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
|
||||
if ocr is not None and ocr_tier is not None:
|
||||
reason = (
|
||||
"corrupt_glyphs"
|
||||
if classification.recommended_tier == "structured"
|
||||
@@ -360,11 +368,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 +388,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"),
|
||||
)
|
||||
@@ -503,11 +513,14 @@ class ProcessorRegistry:
|
||||
"""
|
||||
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
|
||||
|
||||
@@ -645,7 +658,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
|
||||
|
||||
Reference in New Issue
Block a user