feat(ocr): opt-in batch OCR mode via the gateway's async batch routes
Add DOCUMENT_OCR_MODE=sync|batch (default sync). In batch mode the tier-3 OCR
processor submits documents to the embedding gateway's async Batch OCR routes
(POST /v1/ocr/batch + GET /v1/ocr/batch/{job_id}, astrolabe-cloud-website#372)
for ~50% cheaper large-corpus backfill. The direct Mistral OCR path is left
untouched. Tracked on Deck #332.
Batch jobs run minutes-hours, so the OCR tier cannot block (the procrastinate
worker reclaims jobs in `doing` after INGEST_STALLED_JOB_SECONDS). Instead it
submits, records the gateway job id in a new per-tenant `batch_ocr_jobs` table
(procrastinate args are immutable across retries), and raises a BatchPending
signal that TieredEscalationStrategy turns into a same-queue deferred re-poll —
releasing the worker slot between polls. On completion the per-page markdown is
indexed like the sync path; a failure or a job past
DOCUMENT_OCR_BATCH_MAX_WAIT_SECONDS marks the document parse-failed.
Batch is opt-in and gateway-only: with the direct mistral backend, no gateway
URL, or the inline/memory pipeline (which can't defer), it falls back to sync.
One batch job per document (coalescing N docs/job is a follow-up).
- embedding/gateway_batch_client.py: submit/poll client (reuses GatewayTokenProvider).
- vector/batch_ocr_store.py + migration 008: job tracking (portable SQLite+PG).
- document_processors/escalation.py: BatchPending control-flow signal.
- document_processors/ocr.py: batch state machine + sync fallback.
- vector/processor.py: thread doc identity to the OCR tier; raise BatchPending
from the pending sentinel; propagate it as control flow (not a failure).
- vector/queue/procrastinate.py: BatchPending -> same-queue retry_in, exempt
from the transient cap (bounded by the processor's deadline).
- config + docs; tests across client/store/processor/strategy/parse-tier.
1653 unit tests 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
07ee91399b
commit
3b7e8d779b
@@ -0,0 +1,121 @@
|
||||
"""Tracking store for in-flight async batch OCR jobs (Deck #332).
|
||||
|
||||
When ``DOCUMENT_OCR_MODE=batch`` the OCR tier submits a document to the gateway's
|
||||
async batch route and then re-polls across procrastinate retries. procrastinate
|
||||
job args are immutable, so the gateway ``job_id`` (and submit time, for the poll
|
||||
deadline) live in the ``batch_ocr_jobs`` app-DB table, keyed on the document +
|
||||
its content version (``etag``).
|
||||
|
||||
Engine reuse mirrors :class:`~nextcloud_mcp_server.usage.store.UsageEventStore`:
|
||||
rather than open its own engine this store borrows the process-wide
|
||||
:class:`RefreshTokenStorage` singleton (``get_shared_storage()``) — same app DB,
|
||||
dialect handling, ``?``-placeholder shim, and the guarantee that Alembic
|
||||
migrations (incl. ``batch_ocr_jobs``) already ran.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
import anyio
|
||||
|
||||
from nextcloud_mcp_server.auth.storage import RefreshTokenStorage, get_shared_storage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BatchOcrJob:
|
||||
"""A tracked in-flight batch OCR job."""
|
||||
|
||||
job_id: str
|
||||
status: str
|
||||
submitted_at: int
|
||||
|
||||
|
||||
class BatchOcrJobStore:
|
||||
"""CRUD for the ``batch_ocr_jobs`` table (one row per in-flight job)."""
|
||||
|
||||
_shared_instance: "BatchOcrJobStore | None" = None
|
||||
_shared_lock: anyio.Lock = anyio.Lock()
|
||||
|
||||
def __init__(self, storage: RefreshTokenStorage) -> None:
|
||||
self._storage = storage
|
||||
|
||||
@classmethod
|
||||
async def shared(cls) -> "BatchOcrJobStore":
|
||||
"""Process-wide store backed by the storage singleton (mirrors
|
||||
``UsageEventStore.shared``). Tests should construct
|
||||
``BatchOcrJobStore(storage)`` directly — the cache is a process global
|
||||
with no teardown hook."""
|
||||
async with cls._shared_lock:
|
||||
if cls._shared_instance is None:
|
||||
cls._shared_instance = cls(await get_shared_storage())
|
||||
return cls._shared_instance
|
||||
|
||||
async def get(
|
||||
self, *, user_id: str, doc_id: str, doc_type: str, etag: str
|
||||
) -> BatchOcrJob | None:
|
||||
"""The in-flight job for this document+version, or ``None``."""
|
||||
async with self._storage.acquire() as db:
|
||||
async with db.execute(
|
||||
"SELECT job_id, status, submitted_at FROM batch_ocr_jobs "
|
||||
"WHERE user_id = ? AND doc_id = ? AND doc_type = ? AND etag = ?",
|
||||
(user_id, doc_id, doc_type, etag),
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return BatchOcrJob(job_id=row[0], status=row[1], submitted_at=int(row[2]))
|
||||
|
||||
async def insert_pending(
|
||||
self,
|
||||
*,
|
||||
user_id: str,
|
||||
doc_id: str,
|
||||
doc_type: str,
|
||||
etag: str,
|
||||
job_id: str,
|
||||
submitted_at: int | None = None,
|
||||
) -> None:
|
||||
"""Record a freshly-submitted job. ``ON CONFLICT DO NOTHING`` makes a
|
||||
racing double-submit harmless (the first row wins; the loser's job id is
|
||||
abandoned and reaped by the gateway-side file purge)."""
|
||||
now = submitted_at if submitted_at is not None else int(time.time())
|
||||
async with self._storage.acquire() as db:
|
||||
await db.execute(
|
||||
"INSERT INTO batch_ocr_jobs "
|
||||
"(user_id, doc_id, doc_type, etag, job_id, status, submitted_at, updated_at) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?, ?, ?) "
|
||||
"ON CONFLICT (user_id, doc_id, doc_type, etag) DO NOTHING",
|
||||
(user_id, doc_id, doc_type, etag, job_id, "pending", now, now),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def delete(
|
||||
self, *, user_id: str, doc_id: str, doc_type: str, etag: str
|
||||
) -> None:
|
||||
"""Drop the row once the job is terminal (succeeded or failed)."""
|
||||
async with self._storage.acquire() as db:
|
||||
await db.execute(
|
||||
"DELETE FROM batch_ocr_jobs "
|
||||
"WHERE user_id = ? AND doc_id = ? AND doc_type = ? AND etag = ?",
|
||||
(user_id, doc_id, doc_type, etag),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def delete_stale_for_doc(
|
||||
self, *, user_id: str, doc_id: str, doc_type: str, keep_etag: str
|
||||
) -> None:
|
||||
"""Remove superseded-version rows for a document (any etag other than the
|
||||
current one) before a resubmit, so a re-edited file doesn't leave its
|
||||
old in-flight job tracked forever."""
|
||||
async with self._storage.acquire() as db:
|
||||
await db.execute(
|
||||
"DELETE FROM batch_ocr_jobs "
|
||||
"WHERE user_id = ? AND doc_id = ? AND doc_type = ? AND etag != ?",
|
||||
(user_id, doc_id, doc_type, keep_etag),
|
||||
)
|
||||
await db.commit()
|
||||
@@ -126,6 +126,7 @@ async def _parse_pdf_tier(
|
||||
filename: str | None,
|
||||
tier: str,
|
||||
settings: Any,
|
||||
options: dict[str, Any] | None = None,
|
||||
) -> "ProcessingResult":
|
||||
"""Run a single extraction tier and apply the post-parse escalation gate.
|
||||
|
||||
@@ -141,18 +142,31 @@ async def _parse_pdf_tier(
|
||||
the "OCR is an enhancement, never worse than off" invariant: a tenant who has
|
||||
not enabled a higher tier (or has no processor for it) simply indexes the
|
||||
cheap tier's output.
|
||||
|
||||
Batch OCR (Deck #332): when the OCR tier's batch job is still in flight the
|
||||
processor returns a *pending sentinel* result; we translate it here into a
|
||||
``BatchPending`` raise (same decision point as ``EscalateError``) so the
|
||||
retry strategy re-runs this tier after a delay instead of indexing empty text.
|
||||
"""
|
||||
# Lazy import: keep the document stack (pymupdf/_isolation) off the module
|
||||
# load path; this runs only on the per-tier worker, which needs it anyway.
|
||||
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
|
||||
BatchPending,
|
||||
EscalateError,
|
||||
)
|
||||
from nextcloud_mcp_server.document_processors.ocr import ( # noqa: PLC0415
|
||||
OCR_BATCH_PENDING_KEY,
|
||||
OCR_BATCH_RETRY_IN_KEY,
|
||||
)
|
||||
|
||||
# options / progress_callback are not threaded here -- the indexing caller
|
||||
# passes neither today, and the inline path (registry.process) omits them
|
||||
# too. Forward them if a tier processor ever needs per-call tuning (e.g. OCR
|
||||
# DPI); keeping the two paths symmetric until then.
|
||||
result = await registry.process_tier(content, content_type, filename, tier)
|
||||
# ``options`` threads per-document identity (user_id/doc_id/doc_type/etag) to
|
||||
# the OCR tier so batch mode can key its job-tracking table (Deck #332). Other
|
||||
# tiers ignore it. The inline path (registry.process) passes None.
|
||||
result = await registry.process_tier(
|
||||
content, content_type, filename, tier, options=options
|
||||
)
|
||||
if result.metadata.get(OCR_BATCH_PENDING_KEY):
|
||||
raise BatchPending(retry_in=int(result.metadata[OCR_BATCH_RETRY_IN_KEY]))
|
||||
if result.success:
|
||||
decision = registry.evaluate_escalation(
|
||||
result, content, tier, settings, filename=filename
|
||||
@@ -525,19 +539,21 @@ async def process_document(
|
||||
re-picked on the next scan; the procrastinate path (max_retries=1) caps it at
|
||||
one outer attempt (~30s) and defers. Don't stack a third retry layer here.
|
||||
"""
|
||||
# EscalateError is a control-flow signal that arises ONLY on the per-tier
|
||||
# external path (tier set). Bind the class lazily here, and only when a tier
|
||||
# is set, so the document stack is never imported at *module load* (the #877
|
||||
# invariant) nor on the delete / text-doc call paths (file processing already
|
||||
# imports it via get_registry regardless). When tier is None it can't be
|
||||
# raised, so the guards below stay inert.
|
||||
escalate_error_cls: type[BaseException] | None = None
|
||||
# EscalateError and BatchPending are control-flow signals that arise ONLY on
|
||||
# the per-tier external path (tier set). Bind them lazily here, and only when a
|
||||
# tier is set, so the document stack is never imported at *module load* (the
|
||||
# #877 invariant) nor on the delete / text-doc call paths (file processing
|
||||
# already imports them via get_registry regardless). When tier is None neither
|
||||
# can be raised, so the guards below stay inert. Bound as a tuple so the guards
|
||||
# treat both identically: propagate untouched, never record an error/drop.
|
||||
control_flow_excs: tuple[type[BaseException], ...] = ()
|
||||
if tier is not None:
|
||||
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
|
||||
BatchPending,
|
||||
EscalateError,
|
||||
)
|
||||
|
||||
escalate_error_cls = EscalateError
|
||||
control_flow_excs = (EscalateError, BatchPending)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
@@ -632,13 +648,11 @@ async def process_document(
|
||||
return # Success
|
||||
|
||||
except Exception as e:
|
||||
# An escalation signal is control flow, not a failure:
|
||||
# propagate it untouched so the procrastinate retry strategy
|
||||
# can hop the job to the next tier's queue. Never retry it
|
||||
# A control-flow signal (escalation hop, or batch-OCR re-poll
|
||||
# deferral) is not a failure: propagate it untouched so the
|
||||
# procrastinate retry strategy handles it. Never retry it
|
||||
# in-process and never count it as a drop.
|
||||
if escalate_error_cls is not None and isinstance(
|
||||
e, escalate_error_cls
|
||||
):
|
||||
if isinstance(e, control_flow_excs):
|
||||
raise
|
||||
if attempt < max_retries - 1:
|
||||
logger.warning(
|
||||
@@ -691,10 +705,11 @@ async def process_document(
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
# An escalation signal must reach the procrastinate retry strategy
|
||||
# un-recorded -- it is neither a processing success nor an error
|
||||
# (the hop is its own event, counted via record_document_escalation).
|
||||
if escalate_error_cls is not None and isinstance(e, escalate_error_cls):
|
||||
# A control-flow signal must reach the procrastinate retry strategy
|
||||
# un-recorded -- it is neither a processing success nor an error (an
|
||||
# escalation hop is counted via record_document_escalation; a batch
|
||||
# re-poll deferral is not an event at all).
|
||||
if isinstance(e, control_flow_excs):
|
||||
raise
|
||||
# Single processing-error call site: catches exhausted-retry
|
||||
# re-raises, delete failures, and setup errors (get_qdrant_client /
|
||||
@@ -956,6 +971,7 @@ async def _index_document(
|
||||
get_registry,
|
||||
)
|
||||
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
|
||||
BatchPending,
|
||||
EscalateError,
|
||||
)
|
||||
|
||||
@@ -968,6 +984,15 @@ async def _index_document(
|
||||
# and the in-process/memory pool (tier is None) -- runs the inline
|
||||
# tiered pipeline (fast -> OCR escalation in one call).
|
||||
if tier is not None and _is_pdf(content_type):
|
||||
# Thread per-document identity to the OCR tier so batch mode
|
||||
# (Deck #332) can key its job-tracking table; other tiers
|
||||
# ignore it.
|
||||
ocr_options = {
|
||||
"user_id": doc_task.user_id,
|
||||
"doc_id": doc_task.doc_id,
|
||||
"doc_type": doc_task.doc_type,
|
||||
"etag": doc_task.etag or "",
|
||||
}
|
||||
result = await _parse_pdf_tier(
|
||||
registry,
|
||||
content_bytes,
|
||||
@@ -975,6 +1000,7 @@ async def _index_document(
|
||||
file_path,
|
||||
tier,
|
||||
settings,
|
||||
options=ocr_options,
|
||||
)
|
||||
else:
|
||||
result = await registry.process(
|
||||
@@ -1047,10 +1073,12 @@ async def _index_document(
|
||||
)
|
||||
else:
|
||||
logger.debug("No page_boundaries in metadata for %s", file_path)
|
||||
except EscalateError:
|
||||
# Control-flow signal (per-tier path): re-raise untouched so the
|
||||
# queue hops the job to the next tier. NOT a "failed to process"
|
||||
# error -- don't log it as one.
|
||||
except (EscalateError, BatchPending):
|
||||
# Control-flow signals (per-tier path): re-raise untouched.
|
||||
# EscalateError hops the job to the next tier; BatchPending defers
|
||||
# a re-poll on the same tier (batch OCR still in flight, Deck
|
||||
# #332). Neither is a "failed to process" error -- don't log them
|
||||
# as one.
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error("Failed to process file %s: %s", file_path, e)
|
||||
|
||||
@@ -328,13 +328,26 @@ class TieredEscalationStrategy(BaseRetryStrategy):
|
||||
def get_retry_decision(
|
||||
self, *, exception: BaseException, job: Job
|
||||
) -> RetryDecision | None:
|
||||
# Lazy import: EscalateError lives in the document stack, which the API
|
||||
# pod (it also builds this App to defer) must not load. get_retry_decision
|
||||
# runs only in the worker, where the stack is already imported.
|
||||
from ...document_processors.escalation import EscalateError # noqa: PLC0415
|
||||
# Lazy import: these live in the document stack, which the API pod (it
|
||||
# also builds this App to defer) must not load. get_retry_decision runs
|
||||
# only in the worker, where the stack is already imported.
|
||||
from ...document_processors.escalation import ( # noqa: PLC0415
|
||||
BatchPending,
|
||||
EscalateError,
|
||||
)
|
||||
|
||||
exc = _first_leaf(exception)
|
||||
|
||||
if isinstance(exc, BatchPending):
|
||||
# Batch OCR job still in flight (Deck #332): defer a re-poll on the
|
||||
# SAME queue after retry_in seconds. Deliberately exempt from the
|
||||
# transient cap below — a batch job can take minutes-hours, so the
|
||||
# poll count is unbounded here; the OCR processor's own deadline
|
||||
# (DOCUMENT_OCR_BATCH_MAX_WAIT_SECONDS) is what terminates a stuck job.
|
||||
# Releasing the worker between polls keeps the job out of `doing`, so
|
||||
# it's never stall-reclaimed.
|
||||
return RetryDecision(retry_in={"seconds": exc.retry_in})
|
||||
|
||||
if isinstance(exc, EscalateError):
|
||||
queue = TIER_QUEUES.get(exc.to_tier)
|
||||
if queue is None:
|
||||
|
||||
Reference in New Issue
Block a user