feat(ingest): per-tier escalation via procrastinate queue-hop
Split external (procrastinate) document processing into per-tier queues so a
document is attempted at most once per tier and requeued to the next tier's
queue on a low-quality parse, using procrastinate's native retry.
- escalation.py: TIER_LADDER (fast->structured->ocr) + EscalateError signal
- registry: process_tier (one tier) + evaluate_escalation post-parse gate
(reuses classify_from_text) + next_available_tier; shared _classify_result
and _oversize_result with the inline pipeline
- processor: process_document(tier=...) runs one tier and raises EscalateError
before embed (junk text never indexed); inline memory path unchanged
- queue/procrastinate: ingest-fast|structured|ocr queues; TieredEscalationStrategy
(queue-hop on EscalateError, bounded same-tier transient retry); queue-aware
task; producer defers to ingest-fast; per-queue counts + all-queue reclaim
- cli: worker --tier {fast,structured,ocr}
- billing: pages_ocr usage event + pipeline_tier metadata (paid OCR billed apart)
- observability: astrolabe_ingest_queue_depth{queue,status} gauge + per-queue
counts in nc_get_vector_sync_status / management status endpoint
- config: INGEST_ESCALATION_ENABLED (default true), INGEST_TRANSIENT_MAX_ATTEMPTS
INGEST_ESCALATION_ENABLED=false and INGEST_QUEUE=memory preserve prior behaviour.
Deck #323.
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
6fab0e2ae3
commit
9676bb3106
@@ -29,6 +29,10 @@ class IngestPending:
|
||||
# Per-status counts (todo/doing/failed/…) on the postgres backend; None on
|
||||
# the memory backend, which has no durable per-status breakdown.
|
||||
job_counts: dict[str, int] | None = None
|
||||
# Per-tier-queue breakdown ``{queue: {status: count}}`` on the postgres
|
||||
# backend (Deck #323); None on the memory backend. Feeds the per-tier status
|
||||
# surface + the astrolabe_ingest_queue_depth gauge.
|
||||
job_counts_by_queue: dict[str, dict[str, int]] | None = None
|
||||
|
||||
|
||||
async def get_ingest_pending(
|
||||
@@ -47,13 +51,27 @@ async def get_ingest_pending(
|
||||
"""
|
||||
if ingest_queue == "postgres":
|
||||
counts: dict[str, int] = {}
|
||||
if task_producer is not None and hasattr(task_producer, "job_counts"):
|
||||
by_queue: dict[str, dict[str, int]] | None = None
|
||||
# Prefer the per-queue breakdown (Deck #323) and aggregate from it, so the
|
||||
# fleet-wide totals and the per-tier view always agree. Fall back to the
|
||||
# aggregated call for any producer that predates job_counts_by_queue.
|
||||
if task_producer is not None and hasattr(task_producer, "job_counts_by_queue"):
|
||||
try:
|
||||
by_queue = await task_producer.job_counts_by_queue()
|
||||
for per_status in by_queue.values():
|
||||
for status, value in per_status.items():
|
||||
counts[status] = counts.get(status, 0) + value
|
||||
except Exception as e:
|
||||
logger.warning("Failed to read ingest job counts by queue: %s", e)
|
||||
elif task_producer is not None and hasattr(task_producer, "job_counts"):
|
||||
try:
|
||||
counts = await task_producer.job_counts()
|
||||
except Exception as e:
|
||||
logger.warning("Failed to read ingest job counts: %s", e)
|
||||
pending = counts.get("todo", 0) + counts.get("doing", 0)
|
||||
return IngestPending(pending=pending, job_counts=counts)
|
||||
return IngestPending(
|
||||
pending=pending, job_counts=counts, job_counts_by_queue=by_queue
|
||||
)
|
||||
|
||||
if document_receive_stream is None:
|
||||
return IngestPending(pending=0)
|
||||
|
||||
@@ -29,6 +29,7 @@ from qdrant_client.models import FieldCondition, Filter, MatchValue
|
||||
|
||||
from nextcloud_mcp_server.config import get_settings
|
||||
from nextcloud_mcp_server.observability.metrics import (
|
||||
update_ingest_queue_depth,
|
||||
update_vector_sync_indexed_chunks,
|
||||
update_vector_sync_indexed_documents,
|
||||
update_vector_sync_pending_documents,
|
||||
@@ -96,6 +97,8 @@ async def publish_vector_sync_metrics(
|
||||
# Keep the legacy gauge meaningful on every consumer path, not just the
|
||||
# single-user one — existing dashboards/alerts reference it.
|
||||
update_vector_sync_queue_size(pending.pending)
|
||||
# Per-tier-queue depth (Deck #323): None on the memory backend (no-op).
|
||||
update_ingest_queue_depth(pending.job_counts_by_queue)
|
||||
except Exception as exc: # noqa: BLE001 — metrics must not break ingest
|
||||
logger.warning("Failed to publish pending-documents gauge: %s", exc)
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@ Processes documents from stream: fetches content, generates embeddings, stores i
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, cast
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
@@ -14,6 +14,12 @@ from anyio.abc import TaskStatus
|
||||
from anyio.streams.memory import MemoryObjectReceiveStream
|
||||
from qdrant_client.models import PointStruct
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# Type-only: the document stack is heavy (pymupdf/_isolation) and must stay
|
||||
# off processor.py's import path (#877); the runtime import is lazy.
|
||||
from nextcloud_mcp_server.document_processors.base import ProcessingResult
|
||||
from nextcloud_mcp_server.document_processors.registry import ProcessorRegistry
|
||||
|
||||
from nextcloud_mcp_server.acl_hash import compute_acl_hash
|
||||
from nextcloud_mcp_server.client import NextcloudClient
|
||||
from nextcloud_mcp_server.config import get_settings
|
||||
@@ -21,6 +27,7 @@ from nextcloud_mcp_server.embedding import get_bm25_service, get_embedding_servi
|
||||
from nextcloud_mcp_server.models.deck import DeckCard
|
||||
from nextcloud_mcp_server.observability.metrics import (
|
||||
record_document_chunks,
|
||||
record_document_escalation,
|
||||
record_document_parse_failed,
|
||||
record_embedding,
|
||||
record_embedding_tokens,
|
||||
@@ -106,6 +113,59 @@ def _drop_reason(exc: BaseException) -> str:
|
||||
return "other"
|
||||
|
||||
|
||||
def _is_pdf(content_type: str) -> bool:
|
||||
"""Whether a MIME type is a PDF (parameter-tolerant)."""
|
||||
return content_type.split(";")[0].strip().lower() == "application/pdf"
|
||||
|
||||
|
||||
async def _parse_pdf_tier(
|
||||
registry: "ProcessorRegistry",
|
||||
content: bytes,
|
||||
content_type: str,
|
||||
filename: str | None,
|
||||
tier: str,
|
||||
settings: Any,
|
||||
) -> "ProcessingResult":
|
||||
"""Run a single extraction tier and apply the post-parse escalation gate.
|
||||
|
||||
The external per-tier ingest path (Deck #323): the procrastinate worker for
|
||||
``tier`` parses with exactly that tier, then either returns the result to
|
||||
index or raises ``EscalateError`` to hand the document to the next tier's
|
||||
queue (the queue's retry strategy turns the raise into a native queue-hop).
|
||||
The escalation metric is recorded here, at the decision point.
|
||||
|
||||
A hard parse failure (``result.success`` False) is returned as-is, not
|
||||
escalated -- a corrupt/encrypted/oversize PDF that one engine can't open
|
||||
usually defeats the others too; the caller marks it failed. This preserves
|
||||
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.
|
||||
"""
|
||||
# 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
|
||||
EscalateError,
|
||||
)
|
||||
|
||||
result = await registry.process_tier(content, content_type, filename, tier)
|
||||
if result.success:
|
||||
decision = registry.evaluate_escalation(
|
||||
result, content, tier, settings, filename=filename
|
||||
)
|
||||
if decision is not None:
|
||||
to_tier, reason = decision
|
||||
record_document_escalation(tier, to_tier, reason)
|
||||
logger.info(
|
||||
"Escalating %s %s->%s (reason=%s)",
|
||||
filename or "<bytes>",
|
||||
tier,
|
||||
to_tier,
|
||||
reason,
|
||||
)
|
||||
raise EscalateError(from_tier=tier, to_tier=to_tier, reason=reason)
|
||||
return result
|
||||
|
||||
|
||||
def assign_page_numbers(chunks, page_boundaries):
|
||||
"""Assign page numbers to chunks based on page boundaries.
|
||||
|
||||
@@ -173,6 +233,7 @@ async def record_indexing_usage(
|
||||
token_count: int,
|
||||
total_chars: int,
|
||||
page_count: int | None,
|
||||
pipeline_tier: str | None = None,
|
||||
) -> None:
|
||||
"""Record the billable usage events for one embedded document.
|
||||
|
||||
@@ -213,6 +274,11 @@ async def record_indexing_usage(
|
||||
"doc_type": doc_type,
|
||||
"user_id": user_id,
|
||||
"total_chars": total_chars,
|
||||
# Which extraction tier produced the parsed pages (Deck #323). Carried so
|
||||
# the CP rollup / a future per-tier price can attribute parsing cost to
|
||||
# the tier that incurred it (paid OCR vs CPU-cheap fast). None for text
|
||||
# doc types, which are never parsed.
|
||||
"pipeline_tier": pipeline_tier,
|
||||
}
|
||||
try:
|
||||
store = await UsageEventStore.shared()
|
||||
@@ -242,6 +308,18 @@ async def record_indexing_usage(
|
||||
metadata=metadata,
|
||||
enabled=True,
|
||||
)
|
||||
# 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":
|
||||
await store.record_usage_event(
|
||||
metric="pages_ocr",
|
||||
value=page_count,
|
||||
metadata=metadata,
|
||||
enabled=True,
|
||||
)
|
||||
except Exception:
|
||||
# Reached only when shared()/store construction itself raises
|
||||
# (record_usage_event swallows its own write failures). Metering is on,
|
||||
@@ -393,7 +471,11 @@ async def _reconcile_tag_event(
|
||||
|
||||
|
||||
async def process_document(
|
||||
doc_task: DocumentTask, nc_client: NextcloudClient, *, max_retries: int = 3
|
||||
doc_task: DocumentTask,
|
||||
nc_client: NextcloudClient,
|
||||
*,
|
||||
max_retries: int = 3,
|
||||
tier: str | None = None,
|
||||
):
|
||||
"""
|
||||
Process a single document: fetch, tokenize, embed, store in Qdrant.
|
||||
@@ -407,6 +489,11 @@ async def process_document(
|
||||
(3) suits the in-process SQLite pool, which has no durable retry. The
|
||||
procrastinate worker passes ``1`` so durable retry is owned by the
|
||||
queue (and survives worker crashes), avoiding compounding 3×N retries.
|
||||
tier: Extraction tier to run for PDFs on the external per-tier path (Deck
|
||||
#323) -- the procrastinate worker passes the tier matching its queue.
|
||||
``None`` (the default, used by the in-process/memory pool) runs the
|
||||
inline tiered pipeline (``registry.process``: fast -> OCR escalation
|
||||
in one call) and never raises ``EscalateError``.
|
||||
|
||||
Retry layering: the embedding provider adds its own transient retry (5
|
||||
attempts, 2s→60s backoff — card 309) *inside* each of these attempts. On the
|
||||
@@ -415,6 +502,19 @@ 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 there so the in-process /
|
||||
# memory path never pulls the document stack at call time (mirrors the lazy
|
||||
# get_registry import; see #877). When tier is None it can't be raised, so
|
||||
# the guards below stay inert.
|
||||
escalate_error_cls: type[BaseException] | None = None
|
||||
if tier is not None:
|
||||
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
|
||||
EscalateError,
|
||||
)
|
||||
|
||||
escalate_error_cls = EscalateError
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
logger.debug(
|
||||
@@ -484,7 +584,9 @@ async def process_document(
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
indexed = await _index_document(doc_task, nc_client, qdrant_client)
|
||||
indexed = await _index_document(
|
||||
doc_task, nc_client, qdrant_client, tier=tier
|
||||
)
|
||||
|
||||
# A permanent parse failure returns False: it was already
|
||||
# recorded (document_parse_failed_total + the registry's
|
||||
@@ -506,6 +608,14 @@ 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
|
||||
# in-process and never count it as a drop.
|
||||
if escalate_error_cls is not None and isinstance(
|
||||
e, escalate_error_cls
|
||||
):
|
||||
raise
|
||||
if attempt < max_retries - 1:
|
||||
logger.warning(
|
||||
"Retry %s/%s for %s_%s: %s",
|
||||
@@ -556,7 +666,12 @@ async def process_document(
|
||||
record_ingest_dropped(reason)
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
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):
|
||||
raise
|
||||
# Single processing-error call site: catches exhausted-retry
|
||||
# re-raises, delete failures, and setup errors (get_qdrant_client /
|
||||
# get_settings) — each counted exactly once. A failed delete is not
|
||||
@@ -571,11 +686,20 @@ async def process_document(
|
||||
|
||||
|
||||
async def _index_document(
|
||||
doc_task: DocumentTask, nc_client: NextcloudClient, qdrant_client
|
||||
doc_task: DocumentTask,
|
||||
nc_client: NextcloudClient,
|
||||
qdrant_client,
|
||||
*,
|
||||
tier: str | None = None,
|
||||
) -> bool | None:
|
||||
"""
|
||||
Index a single document (called by process_document with retry).
|
||||
|
||||
``tier`` selects the external per-tier PDF path (Deck #323): when set and the
|
||||
file is a PDF, exactly that tier is parsed and a low-quality result raises
|
||||
``EscalateError`` to hand the document to the next tier's queue. ``None``
|
||||
(default) runs the inline tiered pipeline (``registry.process``).
|
||||
|
||||
Returns ``False`` when a permanent parse failure means nothing was indexed
|
||||
(the caller must then skip the success metrics); ``None`` otherwise.
|
||||
|
||||
@@ -800,22 +924,40 @@ async def _index_document(
|
||||
"vector_sync.file_size": len(content_bytes),
|
||||
},
|
||||
):
|
||||
# The registry runs the tiered PDF pipeline (tier-0 classify ->
|
||||
# tier-1 fast -> OCR escalation) and records classification metrics.
|
||||
# Imported lazily so module import doesn't pull in the document stack
|
||||
# (document_processors -> _isolation, Unix-only ``resource``; see #877).
|
||||
# The registry runs the tiered PDF pipeline and records
|
||||
# classification metrics. Imported lazily so module import doesn't
|
||||
# pull in the document stack (document_processors -> _isolation,
|
||||
# Unix-only ``resource``; see #877).
|
||||
from nextcloud_mcp_server.document_processors import ( # noqa: PLC0415
|
||||
get_registry,
|
||||
)
|
||||
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
|
||||
EscalateError,
|
||||
)
|
||||
|
||||
registry = get_registry()
|
||||
|
||||
try:
|
||||
result = await registry.process(
|
||||
content=content_bytes,
|
||||
content_type=content_type,
|
||||
filename=file_path,
|
||||
)
|
||||
# External per-tier path (Deck #323): run only this worker's tier
|
||||
# for PDFs and let a low-quality parse raise EscalateError (a
|
||||
# queue-hop to the next tier). Everything else -- non-PDF files,
|
||||
# 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):
|
||||
result = await _parse_pdf_tier(
|
||||
registry,
|
||||
content_bytes,
|
||||
content_type,
|
||||
file_path,
|
||||
tier,
|
||||
settings,
|
||||
)
|
||||
else:
|
||||
result = await registry.process(
|
||||
content=content_bytes,
|
||||
content_type=content_type,
|
||||
filename=file_path,
|
||||
)
|
||||
|
||||
# A permanent parse failure (e.g. an isolated-worker OOM/timeout
|
||||
# on a pathological PDF) returns success=False rather than
|
||||
@@ -881,6 +1023,11 @@ 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.
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error("Failed to process file %s: %s", file_path, e)
|
||||
raise
|
||||
@@ -1037,6 +1184,14 @@ async def _index_document(
|
||||
and not isinstance(raw_page_count, bool)
|
||||
else None
|
||||
),
|
||||
# Tier that produced the parsed pages (registry stamps it on the
|
||||
# result metadata); text doc types stay "fast". Narrow defensively
|
||||
# to str|None — file_metadata is loosely typed (Any values).
|
||||
pipeline_tier=(
|
||||
pt
|
||||
if isinstance(pt := file_metadata.get("pipeline_tier"), str)
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
async def generate_sparse_embeddings():
|
||||
|
||||
@@ -5,11 +5,14 @@ This replaces NATS JetStream and the old Postgres-queue stub. The MCP server now
|
||||
owns *both* sides of ingest:
|
||||
|
||||
- **Producer** (API role / scanner) — :class:`ProcrastinateTaskProducer.send`
|
||||
*defers* one ``ingest:process_document`` job per changed document into the
|
||||
per-tenant Postgres (the same app DB; procrastinate manages its own tables).
|
||||
- **Consumer** (worker role) — ``nextcloud-mcp-server worker`` runs
|
||||
:func:`procrastinate.App.run_worker`, which drains the ``ingest`` queue and
|
||||
invokes the existing :func:`process_document` pipeline.
|
||||
*defers* one ``ingest:process_document`` job per changed document onto the
|
||||
cheapest tier's queue (``ingest-fast``) in the per-tenant Postgres (the same
|
||||
app DB; procrastinate manages its own tables).
|
||||
- **Consumer** (worker role) — ``nextcloud-mcp-server worker [--tier T]`` runs
|
||||
:func:`procrastinate.App.run_worker`, which drains its tier's queue and invokes
|
||||
the existing :func:`process_document` pipeline. A parse too poor to index hops
|
||||
the job to the next tier's queue (see :class:`TieredEscalationStrategy`), so
|
||||
cheap CPU parsing and paid OCR run on independently-scaled fleets (Deck #323).
|
||||
|
||||
Design notes:
|
||||
|
||||
@@ -35,9 +38,17 @@ from datetime import datetime, timezone
|
||||
from types import TracebackType
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from procrastinate import App, Blueprint, JobContext, PsycopgConnector, RetryStrategy
|
||||
from procrastinate import (
|
||||
App,
|
||||
BaseRetryStrategy,
|
||||
Blueprint,
|
||||
JobContext,
|
||||
PsycopgConnector,
|
||||
RetryDecision,
|
||||
)
|
||||
from procrastinate.connector import BaseConnector
|
||||
from procrastinate.exceptions import AlreadyEnqueued
|
||||
from procrastinate.jobs import Job
|
||||
|
||||
from ...config import get_procrastinate_conninfo, get_settings
|
||||
from ..scanner import DocumentTask
|
||||
@@ -47,14 +58,53 @@ if TYPE_CHECKING:
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Single queue for document ingest. KEDA scales the worker Deployment on the
|
||||
# depth of this queue (``SELECT count(*) FROM procrastinate_jobs WHERE
|
||||
# queue_name='ingest' AND status='todo'``).
|
||||
INGEST_QUEUE_NAME = "ingest"
|
||||
# One queue per extraction tier (Deck #323), aligned cheapest-first with
|
||||
# document_processors.escalation.TIER_LADDER. Each queue is drained by its own
|
||||
# worker Deployment + KEDA ScaledObject (``SELECT count(*) FROM
|
||||
# procrastinate_jobs WHERE queue_name=<queue> AND status='todo'``), so a
|
||||
# CPU-bound ``fast`` fleet, an in-cluster ``structured`` fleet, and a paid
|
||||
# network-bound ``ocr`` fleet scale (and fail) independently.
|
||||
INGEST_QUEUE_FAST = "ingest-fast"
|
||||
INGEST_QUEUE_STRUCTURED = "ingest-structured"
|
||||
INGEST_QUEUE_OCR = "ingest-ocr"
|
||||
|
||||
# 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,
|
||||
}
|
||||
_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.
|
||||
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_INGEST_QUEUE = "ingest"
|
||||
# Back-compat alias for callers that imported the old single-queue constant.
|
||||
INGEST_QUEUE_NAME = DEFAULT_INGEST_QUEUE
|
||||
|
||||
# Queues the job-count + reclaim helpers sweep (tier queues + the legacy one).
|
||||
_MANAGED_QUEUES: tuple[str, ...] = (*ALL_INGEST_QUEUES, LEGACY_INGEST_QUEUE)
|
||||
|
||||
# Blueprint namespace → registered task names are prefixed ``ingest:``.
|
||||
_NAMESPACE = "ingest"
|
||||
INGEST_TASK_NAME = f"{_NAMESPACE}:process_document"
|
||||
|
||||
|
||||
def tier_for_queue(queue: str | None) -> str:
|
||||
"""Tier a worker on ``queue`` should run. Unknown/legacy -> ``fast``.
|
||||
|
||||
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.
|
||||
"""
|
||||
return _QUEUE_TIERS.get(queue or "", "fast")
|
||||
|
||||
|
||||
# A crashed worker leaves its job in ``doing``; reclaim it once its (per-worker)
|
||||
# heartbeat is this many seconds stale. The default is sized well above the
|
||||
# longest expected ``process_document`` (PDF render + embedding) so a slow-but-
|
||||
@@ -69,6 +119,7 @@ INGEST_TASK_NAME = f"{_NAMESPACE}:process_document"
|
||||
# Blueprint cannot be added to more than one App — which the tests (in-memory +
|
||||
# real Postgres) and any re-init path require.
|
||||
async def process_document_task(
|
||||
context: JobContext,
|
||||
*,
|
||||
user_id: str,
|
||||
doc_id: str,
|
||||
@@ -80,12 +131,26 @@ async def process_document_task(
|
||||
etag: str | None = None,
|
||||
owner_id: str | None = None,
|
||||
) -> None:
|
||||
"""Worker entry: rebuild the DocumentTask, resolve creds, run the pipeline."""
|
||||
"""Worker entry: rebuild the DocumentTask, resolve creds, run the pipeline.
|
||||
|
||||
Queue-aware (Deck #323): the tier this worker runs is the tier of the job's
|
||||
current queue. A low-quality parse raises ``EscalateError``, which the
|
||||
:class:`TieredEscalationStrategy` turns into a queue-hop to the next tier.
|
||||
When per-tier escalation is disabled (``INGEST_ESCALATION_ENABLED=false``),
|
||||
``tier`` stays ``None`` and the inline pipeline runs (fast -> OCR in one
|
||||
call), reproducing the pre-#323 single-queue behaviour.
|
||||
"""
|
||||
# Local imports avoid a heavy import chain at blueprint-definition time
|
||||
# (this module is also imported by the API pod just to defer jobs).
|
||||
from ..oauth_sync import NotProvisionedError # noqa: PLC0415
|
||||
from ..processor import process_document # noqa: PLC0415
|
||||
|
||||
tier = (
|
||||
tier_for_queue(context.job.queue)
|
||||
if get_settings().ingest_escalation_enabled
|
||||
else None
|
||||
)
|
||||
|
||||
task = DocumentTask(
|
||||
user_id=user_id,
|
||||
doc_id=doc_id,
|
||||
@@ -111,7 +176,7 @@ async def process_document_task(
|
||||
|
||||
try:
|
||||
# Durable retry is procrastinate's job; disable the in-process loop.
|
||||
await process_document(task, nc_client, max_retries=1)
|
||||
await process_document(task, nc_client, max_retries=1, tier=tier)
|
||||
finally:
|
||||
await nc_client.close()
|
||||
|
||||
@@ -127,9 +192,11 @@ async def reclaim_stalled_ingest_jobs(context: JobContext, timestamp: int) -> No
|
||||
retry_at = datetime.now(tz=timezone.utc)
|
||||
stalled_after = get_settings().ingest_stalled_job_seconds
|
||||
reclaimed = 0
|
||||
for job in await manager.get_stalled_jobs(
|
||||
queue=INGEST_QUEUE_NAME, seconds_since_heartbeat=stalled_after
|
||||
):
|
||||
# queue=None sweeps every queue, so an orphaned job on any tier queue is
|
||||
# reclaimed regardless of which tier's worker happens to run this periodic.
|
||||
# retry_job_by_id_async keeps the job on its own queue, so a stalled ``ocr``
|
||||
# job re-runs on ``ingest-ocr`` (the ocr fleet), not the reclaiming worker's.
|
||||
for job in await manager.get_stalled_jobs(seconds_since_heartbeat=stalled_after):
|
||||
if job.id is None:
|
||||
continue
|
||||
await manager.retry_job_by_id_async(job_id=job.id, retry_at=retry_at)
|
||||
@@ -163,6 +230,121 @@ async def _resolve_client(user_id: str) -> NextcloudClient:
|
||||
return await get_user_client_basic_auth(user_id, host)
|
||||
|
||||
|
||||
def _first_leaf(exc: BaseException) -> BaseException:
|
||||
"""Descend nested ExceptionGroups to the first concrete leaf exception.
|
||||
|
||||
An anyio task group can wrap the real cause (and nest groups); the retry
|
||||
strategy classifies on the leaf, mirroring ``processor._drop_reason``.
|
||||
"""
|
||||
while isinstance(exc, BaseExceptionGroup) and exc.exceptions:
|
||||
exc = exc.exceptions[0]
|
||||
return exc
|
||||
|
||||
|
||||
def _is_transient_infra_error(exc: BaseException) -> bool:
|
||||
"""Whether ``exc`` is a transient infra blip worth a SAME-tier retry.
|
||||
|
||||
Mirrors the retryable subset of ``processor._drop_reason``: doc-fetch /
|
||||
embed / Qdrant timeouts, connection drops, rate limits, and 5xx. A parse
|
||||
that is merely *poor* never reaches here -- that path raises
|
||||
``EscalateError`` (handled separately) -- so this is purely about
|
||||
infrastructure that should recover on its own. Imports are lazy: this only
|
||||
runs in the worker, and the module is also imported by the API pod to defer.
|
||||
"""
|
||||
import httpx # noqa: PLC0415
|
||||
|
||||
if isinstance(exc, (httpx.TimeoutException, httpx.ConnectError)):
|
||||
return True
|
||||
try:
|
||||
import openai # noqa: PLC0415
|
||||
|
||||
if isinstance(
|
||||
exc,
|
||||
(
|
||||
openai.APITimeoutError,
|
||||
openai.APIConnectionError,
|
||||
openai.RateLimitError,
|
||||
),
|
||||
):
|
||||
return True
|
||||
if isinstance(exc, openai.APIStatusError):
|
||||
return exc.status_code >= 500
|
||||
except ImportError: # pragma: no cover -- openai is a hard dependency
|
||||
pass
|
||||
if type(exc).__module__.startswith("qdrant_client"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class TieredEscalationStrategy(BaseRetryStrategy):
|
||||
"""Native procrastinate retry that escalates across tier queues (Deck #323).
|
||||
|
||||
Three outcomes, decided from the raised exception:
|
||||
|
||||
- ``EscalateError`` -> ``RetryDecision(queue=<next tier's queue>)``: the SAME
|
||||
job hops to the next fleet's queue and is parsed once by that tier. This is
|
||||
how a document is "requeued on a failed parse" -- once per tier, with no
|
||||
same-tier parse retry.
|
||||
- a whitelisted transient infra error (doc fetch / embed / Qdrant blip) ->
|
||||
same-queue exponential backoff, while under ``max_transient_attempts``.
|
||||
- anything else, the transient cap is reached, or the target tier is unknown
|
||||
-> ``None`` (no retry); the placeholder was already marked failed by the
|
||||
pipeline, and the next scan re-picks the document.
|
||||
|
||||
Per-tier attempt accounting is intentionally approximate: a queue-hop can't
|
||||
reset ``job.attempts`` (procrastinate has no per-tier counter), so parse
|
||||
escalations *do* advance the same counter the transient cap reads. Because a
|
||||
parse escalation hops (it never retries in place) the "once per parse per
|
||||
tier" guarantee is structural; the cap is just a generous global ceiling on
|
||||
transient churn across the whole lineage, not an exact per-tier count.
|
||||
"""
|
||||
|
||||
def __init__(self, *, max_transient_attempts: int) -> None:
|
||||
self._max_transient_attempts = max_transient_attempts
|
||||
|
||||
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
|
||||
|
||||
exc = _first_leaf(exception)
|
||||
|
||||
if isinstance(exc, EscalateError):
|
||||
queue = TIER_QUEUES.get(exc.to_tier)
|
||||
if queue is None:
|
||||
# Unknown target tier: don't strand the job on a queue no worker
|
||||
# drains -- stop and let the placeholder/next scan handle it.
|
||||
logger.error(
|
||||
"ingest.escalate_unknown_tier from=%s to=%s",
|
||||
exc.from_tier,
|
||||
exc.to_tier,
|
||||
)
|
||||
return None
|
||||
logger.info(
|
||||
"ingest.escalate from=%s to=%s reason=%s queue=%s",
|
||||
exc.from_tier,
|
||||
exc.to_tier,
|
||||
exc.reason,
|
||||
queue,
|
||||
)
|
||||
# Immediate hop -- the next tier's fleet should pick it up at once.
|
||||
return RetryDecision(queue=queue, retry_in={"seconds": 0})
|
||||
|
||||
if (
|
||||
_is_transient_infra_error(exc)
|
||||
and job.attempts < self._max_transient_attempts
|
||||
):
|
||||
# 4, 8, 16, ... seconds, capped at 5 min. attempts is >=1 here (the
|
||||
# failing attempt is counted), so attempts-1 makes the first wait 4s.
|
||||
wait = min(4 * (2 ** max(0, job.attempts - 1)), 300)
|
||||
return RetryDecision(retry_in={"seconds": wait})
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _build_ingest_blueprint() -> Blueprint:
|
||||
"""Create a fresh Blueprint with the ingest tasks registered.
|
||||
|
||||
@@ -172,13 +354,22 @@ def _build_ingest_blueprint() -> Blueprint:
|
||||
bp = Blueprint()
|
||||
# Durable retry owned by the queue (survives worker crashes); the in-process
|
||||
# retry loop in process_document is disabled on this path via max_retries=1.
|
||||
bp.task(
|
||||
# The task's default queue is the cheapest tier; the producer defers there
|
||||
# explicitly and the strategy hops a job up the ladder on a poor parse.
|
||||
bp.task( # type: ignore[no-matching-overload]
|
||||
name="process_document",
|
||||
queue=INGEST_QUEUE_NAME,
|
||||
retry=RetryStrategy(max_attempts=5, exponential_wait=4),
|
||||
queue=DEFAULT_INGEST_QUEUE,
|
||||
pass_context=True,
|
||||
# procrastinate's RetryValue type only admits RetryStrategy, but a custom
|
||||
# BaseRetryStrategy subclass is the documented extension point (and is
|
||||
# accepted at runtime by get_retry_strategy). The annotation is just too
|
||||
# narrow, hence the ignore.
|
||||
retry=TieredEscalationStrategy(
|
||||
max_transient_attempts=get_settings().ingest_transient_max_attempts
|
||||
),
|
||||
)(process_document_task)
|
||||
reclaim = bp.task(
|
||||
name="reclaim_stalled_jobs", queue=INGEST_QUEUE_NAME, pass_context=True
|
||||
name="reclaim_stalled_jobs", queue=DEFAULT_INGEST_QUEUE, pass_context=True
|
||||
)(reclaim_stalled_ingest_jobs)
|
||||
bp.periodic(cron="*/5 * * * *", periodic_id="reclaim_stalled_ingest")(reclaim)
|
||||
return bp
|
||||
@@ -276,20 +467,43 @@ async def apply_ingest_queue_schema(
|
||||
_JOB_STATUSES = ("todo", "doing", "succeeded", "failed", "cancelled", "aborted")
|
||||
|
||||
|
||||
async def get_ingest_job_counts(app: App | None = None) -> dict[str, int]:
|
||||
"""Return ingest job counts by status (``todo``/``doing``/``failed``/…).
|
||||
async def get_ingest_job_counts_by_queue(
|
||||
app: App | None = None,
|
||||
) -> dict[str, dict[str, int]]:
|
||||
"""Per-queue ingest job counts by status (Deck #323).
|
||||
|
||||
Reads procrastinate's per-queue stats via the manager API (not hand-written
|
||||
SQL) so a future schema bump doesn't silently break the status surface. The
|
||||
manager flattens its per-status ``stats`` into top-level row keys, so we read
|
||||
the known status keys directly. Assumes the app's connector is already open.
|
||||
Returns ``{queue_name: {status: count}}`` for the managed ingest queues (the
|
||||
per-tier queues + the legacy single queue) that have rows. Reads
|
||||
procrastinate's per-queue stats via the manager API (not hand-written SQL) so
|
||||
a future schema bump doesn't silently break the status surface. Assumes the
|
||||
app's connector is already open. Feeds the per-tier status surface + the
|
||||
``astrolabe_ingest_queue_depth`` gauge.
|
||||
"""
|
||||
app = app or get_procrastinate_app()
|
||||
counts: dict[str, int] = {}
|
||||
for row in await app.job_manager.list_queues_async(queue=INGEST_QUEUE_NAME):
|
||||
by_queue: dict[str, dict[str, int]] = {}
|
||||
for row in await app.job_manager.list_queues_async():
|
||||
name = row.get("name")
|
||||
if name not in _MANAGED_QUEUES:
|
||||
continue
|
||||
per = by_queue.setdefault(name, {})
|
||||
for status in _JOB_STATUSES:
|
||||
if status in row:
|
||||
counts[status] = counts.get(status, 0) + int(row[status])
|
||||
per[status] = per.get(status, 0) + int(row[status])
|
||||
return by_queue
|
||||
|
||||
|
||||
async def get_ingest_job_counts(app: App | None = None) -> dict[str, int]:
|
||||
"""Aggregate ingest job counts by status across all managed queues.
|
||||
|
||||
Fleet-wide totals summed over the per-tier queues + the legacy queue, so
|
||||
``pending = todo + doing`` reflects all outstanding ingest work regardless of
|
||||
which tier a document currently sits on. Per-queue breakdown:
|
||||
:func:`get_ingest_job_counts_by_queue`.
|
||||
"""
|
||||
counts: dict[str, int] = {}
|
||||
for per in (await get_ingest_job_counts_by_queue(app)).values():
|
||||
for status, value in per.items():
|
||||
counts[status] = counts.get(status, 0) + value
|
||||
return counts
|
||||
|
||||
|
||||
@@ -336,7 +550,13 @@ class ProcrastinateTaskProducer:
|
||||
|
||||
async def send(self, task: DocumentTask, /) -> None:
|
||||
key = _doc_queueing_lock(task)
|
||||
deferrer = self._app.configure_task(INGEST_TASK_NAME, queueing_lock=key)
|
||||
# Always defer onto the cheapest tier's queue; the escalation strategy
|
||||
# hops the job up the ladder on a poor parse. queueing_lock is a global
|
||||
# partial-unique on status='todo', so a doc mid-escalation on a higher
|
||||
# tier still dedupes a fresh enqueue here -- no double-processing.
|
||||
deferrer = self._app.configure_task(
|
||||
INGEST_TASK_NAME, queue=DEFAULT_INGEST_QUEUE, queueing_lock=key
|
||||
)
|
||||
try:
|
||||
await deferrer.defer_async(**asdict(task))
|
||||
except AlreadyEnqueued:
|
||||
@@ -358,6 +578,10 @@ class ProcrastinateTaskProducer:
|
||||
"""Ingest job counts by status (for the vector-sync status surface)."""
|
||||
return await get_ingest_job_counts(self._app)
|
||||
|
||||
async def job_counts_by_queue(self) -> dict[str, dict[str, int]]:
|
||||
"""Per-tier-queue ingest job counts by status (Deck #323)."""
|
||||
return await get_ingest_job_counts_by_queue(self._app)
|
||||
|
||||
def clone(self) -> ProcrastinateTaskProducer:
|
||||
return self
|
||||
|
||||
|
||||
Reference in New Issue
Block a user