From 79a3b9d416c7755adb00e045fe63a57f39026d16 Mon Sep 17 00:00:00 2001 From: Ulysse Pence Date: Wed, 26 Aug 2026 19:50:17 -0200 Subject: [PATCH] feat(deriver): Moves scheduling events from deriver to API process --- src/backlog.py | 210 ++++++++++++++++ src/config.py | 2 + src/crud/__init__.py | 7 +- src/crud/deriver.py | 121 +++++++++- src/deriver/queue_manager.py | 77 +++--- src/main.py | 16 ++ src/reconciler/scheduler.py | 2 +- src/schemas/__init__.py | 2 + src/schemas/internal.py | 8 + src/telemetry/prometheus/metrics.py | 91 ++++++- tests/crud/test_deriver_backlog.py | 290 +++++++++++++++++++++++ tests/telemetry/test_metric_zero_init.py | 22 +- 12 files changed, 788 insertions(+), 60 deletions(-) create mode 100644 src/backlog.py create mode 100644 tests/crud/test_deriver_backlog.py diff --git a/src/backlog.py b/src/backlog.py new file mode 100644 index 00000000..45a731d2 --- /dev/null +++ b/src/backlog.py @@ -0,0 +1,210 @@ +"""Publishes the deriver backlog as Prometheus gauges from the API process.""" + +import asyncio +import contextlib +import logging +from datetime import datetime, timedelta, timezone +from typing import Any, cast + +import sentry_sdk +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from src import crud, models +from src.config import settings +from src.dependencies import tracked_db +from src.telemetry.prometheus import metrics as prometheus_metrics +from src.utils.work_unit import construct_work_unit_key + +logger = logging.getLogger(__name__) + +DREAM_POLL_INTERVAL_SECONDS = 300 + + +class BacklogMetricsPoller: + """Refreshes the deriver-backlog gauges on a timer.""" + + def __init__(self) -> None: + self._task: asyncio.Task[None] | None = None + self._shutdown_event: asyncio.Event = asyncio.Event() + self._next_dream_poll: datetime | None = None + + async def start(self) -> None: + if self._task is not None: + logger.warning("BacklogMetricsPoller already running") + return + + interval = settings.DERIVER.BACKLOG_METRICS_POLL_INTERVAL_SECONDS + self._task = asyncio.create_task(self._loop()) + logger.info("BacklogMetricsPoller started (interval=%ss)", interval) + + async def shutdown(self) -> None: + if self._task is None: + return + + logger.info("Shutting down BacklogMetricsPoller...") + self._shutdown_event.set() + + try: + await asyncio.wait_for(self._task, timeout=5.0) + except asyncio.TimeoutError: + logger.warning("BacklogMetricsPoller shutdown timed out, cancelling task") + self._task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._task + + self._task = None + logger.info("BacklogMetricsPoller stopped") + + async def _loop(self) -> None: + interval = settings.DERIVER.BACKLOG_METRICS_POLL_INTERVAL_SECONDS + + while not self._shutdown_event.is_set(): + try: + await self._refresh() + except Exception as e: + logger.error("BacklogMetricsPoller refresh failed: %s", e) + if settings.SENTRY.ENABLED: + sentry_sdk.capture_exception(e) + + try: + await asyncio.wait_for( + self._shutdown_event.wait(), timeout=float(interval) + ) + except asyncio.TimeoutError: + continue + + async def _refresh(self) -> None: + async with tracked_db("backlog_metrics", read_only=True) as db: + backlog = await crud.get_deriver_backlog(db) + prometheus_metrics.prometheus_metrics.set_deriver_backlog( + eligible_work_units=backlog.eligible_work_units, + pending_items=backlog.pending_items, + oldest_pending_age_seconds=backlog.oldest_pending_age_seconds, + ) + + prometheus_metrics.prometheus_metrics.set_seconds_since_last_vector_sync( + seconds=await _seconds_since_last_vector_sync(db) + ) + + if self._dream_poll_due(): + prometheus_metrics.prometheus_metrics.set_dreams_pending( + count=await count_pending_dreams(db) + ) + + def _dream_poll_due(self) -> bool: + now = datetime.now(timezone.utc) + if self._next_dream_poll is not None and now < self._next_dream_poll: + return False + self._next_dream_poll = now + timedelta(seconds=DREAM_POLL_INTERVAL_SECONDS) + return True + + +async def _seconds_since_last_vector_sync(db: AsyncSession) -> float: + """-1 when nothing has ever synced""" + newest = await db.scalar(select(func.max(models.MessageEmbedding.last_sync_at))) + if newest is None: + return -1.0 + return (datetime.now(timezone.utc) - newest).total_seconds() + + +async def count_pending_dreams(db: AsyncSession) -> int: + if not settings.DREAM.ENABLED: + return 0 + + explicit_counts = ( + select( + models.Document.workspace_name, + models.Document.observer, + models.Document.observed, + func.count(models.Document.id).label("explicit_count"), + ) + .where(models.Document.level == "explicit") + .group_by( + models.Document.workspace_name, + models.Document.observer, + models.Document.observed, + ) + .subquery() + ) + + rows = ( + await db.execute( + select( + models.Collection.workspace_name, + models.Collection.observer, + models.Collection.observed, + models.Collection.internal_metadata, + func.coalesce(explicit_counts.c.explicit_count, 0), + ).outerjoin( + explicit_counts, + (models.Collection.workspace_name == explicit_counts.c.workspace_name) + & (models.Collection.observer == explicit_counts.c.observer) + & (models.Collection.observed == explicit_counts.c.observed), + ) + ) + ).all() + + now = datetime.now(timezone.utc) + due: list[str] = [] + + for row in rows: + workspace_name = cast(str, row[0]) + observer = cast(str, row[1]) + observed = cast(str, row[2]) + internal_metadata = cast("dict[str, Any] | None", row[3]) + explicit_count = cast(int, row[4]) + + dream_metadata: dict[str, Any] = (internal_metadata or {}).get("dream", {}) + since_last_dream = explicit_count - int( + dream_metadata.get("last_dream_document_count", 0) + ) + if since_last_dream < settings.DREAM.DOCUMENT_THRESHOLD: + continue + + last_dream_at = cast("str | None", dream_metadata.get("last_dream_at")) + if last_dream_at and _within_min_hours_gate(last_dream_at, now): + continue + + for dream_type in settings.DREAM.ENABLED_TYPES: + due.append( + construct_work_unit_key( + workspace_name, + { + "task_type": "dream", + "observer": observer, + "observed": observed, + "dream_type": dream_type, + }, + ) + ) + + if not due: + return 0 + + already_queued = set( + ( + await db.execute( + select(models.QueueItem.work_unit_key).where( + models.QueueItem.task_type == "dream", + ~models.QueueItem.processed, + models.QueueItem.work_unit_key.in_(due), + ) + ) + ) + .scalars() + .all() + ) + + return len([key for key in due if key not in already_queued]) + + +def _within_min_hours_gate(last_dream_at: str, now: datetime) -> bool: + """True when the last dream is too recent for another one.""" + try: + last_dream_time = datetime.fromisoformat(last_dream_at) + except (ValueError, TypeError): + return False + + hours_since = (now - last_dream_time).total_seconds() / 3600 + return hours_since < settings.DREAM.MIN_HOURS_BETWEEN_DREAMS diff --git a/src/config.py b/src/config.py index 993f9bfb..0115b598 100644 --- a/src/config.py +++ b/src/config.py @@ -972,6 +972,8 @@ class DeriverSettings(HonchoSettings): # When enabled, bypasses the batch token threshold and processes work immediately FLUSH_ENABLED: bool = False + BACKLOG_METRICS_POLL_INTERVAL_SECONDS: Annotated[int, Field(default=30, ge=1)] = 30 + @model_validator(mode="before") @classmethod def _merge_model_config_defaults(cls, data: Any) -> Any: diff --git a/src/crud/__init__.py b/src/crud/__init__.py index 0e920717..1129648c 100644 --- a/src/crud/__init__.py +++ b/src/crud/__init__.py @@ -3,7 +3,11 @@ from .collection import ( get_or_create_collection, update_collection_internal_metadata, ) -from .deriver import get_deriver_status, get_queue_status +from .deriver import ( + get_deriver_backlog, + get_deriver_status, + get_queue_status, +) from .document import ( CreateDocumentsResult, create_documents, @@ -105,6 +109,7 @@ __all__ = [ "get_or_create_collection", "update_collection_internal_metadata", # Deriver + "get_deriver_backlog", "get_deriver_status", "get_queue_status", # Document diff --git a/src/crud/deriver.py b/src/crud/deriver.py index 0852a479..ca6d4d19 100644 --- a/src/crud/deriver.py +++ b/src/crud/deriver.py @@ -1,15 +1,19 @@ from collections.abc import Sequence +from datetime import timedelta from logging import getLogger from typing import Any -from sqlalchemy import Select, case, func, or_, select +from sqlalchemy import ColumnElement, Select, case, func, or_, select from sqlalchemy.engine import Row from sqlalchemy.ext.asyncio import AsyncSession from src import models, schemas +from src.config import settings logger = getLogger(__name__) +REPRESENTATION_WORK_UNIT_PREFIX = "representation:" + async def get_queue_status( db: AsyncSession, @@ -75,6 +79,121 @@ async def get_deriver_status( ) +def representation_batch_threshold_clause( + *, + work_unit_key: ColumnElement[str], + total_tokens: ColumnElement[Any], + oldest_created_at: ColumnElement[Any], +) -> ColumnElement[bool] | None: + if settings.DERIVER.FLUSH_ENABLED: + return None + + target_tokens = settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS + if target_tokens <= 0: + return None + + threshold: ColumnElement[bool] = func.coalesce(total_tokens, 0) >= target_tokens + + max_age_seconds = settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS + if max_age_seconds > 0: + threshold = or_( + threshold, + oldest_created_at <= func.now() - timedelta(seconds=max_age_seconds), + ) + + return or_( + ~work_unit_key.startswith(REPRESENTATION_WORK_UNIT_PREFIX), + threshold, + ) + + +def unclaimed_work_unit_clause( + work_unit_key: ColumnElement[str], + *, + tolerate_stale_claims: bool = False, +) -> ColumnElement[bool]: + claim = select(models.ActiveQueueSession.id).where( + models.ActiveQueueSession.work_unit_key == work_unit_key + ) + + if tolerate_stale_claims: + claim = claim.where( + models.ActiveQueueSession.last_updated + >= func.now() + - timedelta(minutes=settings.DERIVER.STALE_SESSION_TIMEOUT_MINUTES) + ) + + return ~claim.exists() + + +async def get_deriver_backlog(db: AsyncSession) -> schemas.DeriverBacklog: + """Count outstanding deriver work across the whole database. + + Read-only, and deliberately unfiltered by workspace: the caller wants the + service-wide total, not one tenant's share of it. + """ + token_stats = ( + select( + models.QueueItem.work_unit_key, + func.sum(models.Message.token_count).label("total_tokens"), + func.min(models.QueueItem.created_at).label("oldest_created_at"), + ) + .join(models.Message, models.QueueItem.message_id == models.Message.id) + .where(~models.QueueItem.processed) + .where( + models.QueueItem.work_unit_key.startswith(REPRESENTATION_WORK_UNIT_PREFIX) + ) + .group_by(models.QueueItem.work_unit_key) + .subquery() + ) + + work_units = ( + select(models.QueueItem.work_unit_key) + .where(~models.QueueItem.processed) + .group_by(models.QueueItem.work_unit_key) + .subquery() + ) + + eligible = ( + select(func.count()) + .select_from(work_units) + .outerjoin( + token_stats, + work_units.c.work_unit_key == token_stats.c.work_unit_key, + ) + .where( + unclaimed_work_unit_clause( + work_units.c.work_unit_key, tolerate_stale_claims=True + ) + ) + ) + + threshold_clause = representation_batch_threshold_clause( + work_unit_key=work_units.c.work_unit_key, + total_tokens=token_stats.c.total_tokens, + oldest_created_at=token_stats.c.oldest_created_at, + ) + if threshold_clause is not None: + eligible = eligible.where(threshold_clause) + + pending = select( + func.count(models.QueueItem.id), + func.coalesce( + func.extract("epoch", func.now() - func.min(models.QueueItem.created_at)), + 0, + ), + ).where(~models.QueueItem.processed) + + eligible_count = (await db.execute(eligible)).scalar_one() + pending_count, oldest_age = (await db.execute(pending)).one() + + return schemas.DeriverBacklog( + eligible_work_units=int(eligible_count), + pending_items=int(pending_count), + oldest_pending_age_seconds=float(oldest_age), + ) + + # Task types surfaced by the queue status endpoint. _TRACKED_TASK_TYPES = ("representation", "summary", "dream") diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 1c19c131..d40bed3d 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -15,13 +15,13 @@ from dotenv import load_dotenv from nanoid import generate as generate_nanoid from sentry_sdk.integrations.asyncio import AsyncioIntegration from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration -from sqlalchemy import and_, delete, or_, select, update +from sqlalchemy import and_, delete, select, update from sqlalchemy.dialects.postgresql import insert from sqlalchemy.engine import CursorResult from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql import func -from src import models +from src import crud, models from src.cache.client import close_cache, init_cache from src.config import settings from src.dependencies import tracked_db @@ -31,15 +31,11 @@ from src.deriver.consumer import ( ) from src.dreamer.dream_scheduler import ( DreamScheduler, + check_and_schedule_dream, get_dream_scheduler, set_dream_scheduler, ) from src.models import QueueItem -from src.reconciler import ( - ReconcilerScheduler, - get_reconciler_scheduler, - set_reconciler_scheduler, -) from src.schemas import ResolvedConfiguration from src.telemetry import prometheus_metrics from src.telemetry.sentry import initialize_sentry @@ -157,14 +153,6 @@ class QueueManager: else: self.dream_scheduler = existing_scheduler - # Get or create the singleton reconciler scheduler - existing_reconciler = get_reconciler_scheduler() - if existing_reconciler is None: - self.reconciler_scheduler: ReconcilerScheduler = ReconcilerScheduler() - set_reconciler_scheduler(self.reconciler_scheduler) - else: - self.reconciler_scheduler = existing_reconciler - # Initialize Sentry if enabled, using settings if settings.SENTRY.ENABLED: initialize_sentry( @@ -209,11 +197,10 @@ class QueueManager: ) logger.debug("Signal handlers registered") - # Start the reconciler scheduler try: - await self.reconciler_scheduler.start() + await self._reschedule_pending_dreams() except Exception: - logger.exception("Failed to start reconciler scheduler") + logger.exception("Failed to reschedule pending dreams at startup") # Run the polling loop directly in this task logger.debug("Starting polling loop directly") @@ -223,6 +210,20 @@ class QueueManager: finally: await self.cleanup() + async def _reschedule_pending_dreams(self) -> None: + if not settings.DREAM.ENABLED: + return + + async with tracked_db("reschedule_pending_dreams") as db: + collections = (await db.execute(select(models.Collection))).scalars().all() + rescheduled = 0 + for collection in collections: + if await check_and_schedule_dream(db, collection): + rescheduled += 1 + + if rescheduled: + logger.info("Rescheduled %d dream(s) at startup", rescheduled) + async def shutdown(self, sig: signal.Signals) -> None: """Handle graceful shutdown""" logger.info(f"Received exit signal {sig.name}...") @@ -231,9 +232,6 @@ class QueueManager: # Cancel all pending dreams await self.dream_scheduler.shutdown() - # Stop the reconciler scheduler - await self.reconciler_scheduler.shutdown() - if self.active_tasks: logger.info( f"Waiting for {len(self.active_tasks)} active tasks to complete..." @@ -345,7 +343,7 @@ class QueueManager: ) async with tracked_db("get_available_work_units") as db: - representation_prefix = "representation:" + representation_prefix = crud.deriver.REPRESENTATION_WORK_UNIT_PREFIX token_stats_subq = ( select( models.QueueItem.work_unit_key, @@ -383,12 +381,9 @@ class QueueManager: work_units_subq.c.work_unit_key == token_stats_subq.c.work_unit_key, ) .where( - ~select(models.ActiveQueueSession.id) - .where( - models.ActiveQueueSession.work_unit_key - == work_units_subq.c.work_unit_key + crud.deriver.unclaimed_work_unit_clause( + work_units_subq.c.work_unit_key ) - .exists() ) .order_by( work_units_subq.c.oldest_created_at.asc(), @@ -397,27 +392,13 @@ class QueueManager: .limit(limit) ) - # Apply batch threshold filter (skip if FLUSH_ENABLED is True) - if not settings.DERIVER.FLUSH_ENABLED and work_unit_target_tokens > 0: - max_age_seconds = settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS - threshold_clause = ( - func.coalesce(token_stats_subq.c.total_tokens, 0) - >= work_unit_target_tokens - ) - if max_age_seconds > 0: - threshold_clause = or_( - threshold_clause, - token_stats_subq.c.oldest_created_at - <= func.now() - timedelta(seconds=max_age_seconds), - ) - query = query.where( - or_( - ~work_units_subq.c.work_unit_key.startswith( - representation_prefix - ), - threshold_clause, - ) - ) + threshold_clause = crud.deriver.representation_batch_threshold_clause( + work_unit_key=work_units_subq.c.work_unit_key, + total_tokens=token_stats_subq.c.total_tokens, + oldest_created_at=token_stats_subq.c.oldest_created_at, + ) + if threshold_clause is not None: + query = query.where(threshold_clause) result = await db.execute(query) available_rows = result.all() diff --git a/src/main.py b/src/main.py index a1ec9765..6c072064 100644 --- a/src/main.py +++ b/src/main.py @@ -15,6 +15,7 @@ from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration from sentry_sdk.integrations.starlette import StarletteIntegration from src._version import HONCHO_VERSION +from src.backlog import BacklogMetricsPoller from src.cache.client import close_cache, init_cache from src.config import settings from src.db import ( @@ -24,6 +25,7 @@ from src.db import ( request_context, ) from src.exceptions import HonchoException +from src.reconciler import ReconcilerScheduler, set_reconciler_scheduler from src.routers import ( conclusions, keys, @@ -135,12 +137,26 @@ async def lifespan(_: FastAPI): "Error initializing cache in api process; proceeding without cache: %s", e ) + reconciler_scheduler = ReconcilerScheduler() + set_reconciler_scheduler(reconciler_scheduler) + backlog_metrics_poller = BacklogMetricsPoller() + try: + await reconciler_scheduler.start() + except Exception as e: + logger.error("Failed to start reconciler scheduler: %s", e) + try: + await backlog_metrics_poller.start() + except Exception as e: + logger.error("Failed to start backlog metrics poller: %s", e) + try: yield finally: # Import here to avoid circular import at module load time from src.vector_store import close_external_vector_store + await backlog_metrics_poller.shutdown() + await reconciler_scheduler.shutdown() await close_external_vector_store() await close_cache() await engine.dispose() diff --git a/src/reconciler/scheduler.py b/src/reconciler/scheduler.py index 08941970..21f37b51 100644 --- a/src/reconciler/scheduler.py +++ b/src/reconciler/scheduler.py @@ -111,7 +111,7 @@ class ReconcilerScheduler: logger.warning("ReconcilerScheduler already running") return - self._shutdown_event.clear() + self._shutdown_event = asyncio.Event() # Initialize next run times to first interval now = datetime.now(timezone.utc) for task_name, task in RECONCILER_TASKS.items(): diff --git a/src/schemas/__init__.py b/src/schemas/__init__.py index 9f414583..598a4c3e 100644 --- a/src/schemas/__init__.py +++ b/src/schemas/__init__.py @@ -78,6 +78,7 @@ from src.schemas.configuration import ( WorkspaceConfiguration, ) from src.schemas.internal import ( + DeriverBacklog, DocumentBase, DocumentCreate, DocumentMetadata, @@ -163,6 +164,7 @@ __all__ = [ "WorkspaceMessageSearchOptions", "WorkspaceUpdate", # internal + "DeriverBacklog", "DocumentBase", "DocumentCreate", "DocumentMetadata", diff --git a/src/schemas/internal.py b/src/schemas/internal.py index e014431f..729d1f90 100644 --- a/src/schemas/internal.py +++ b/src/schemas/internal.py @@ -144,6 +144,14 @@ class QueueCounts(BaseModel): sessions: dict[str, SessionCounts] +class DeriverBacklog(BaseModel): + """Database-wide view of outstanding deriver work.""" + + eligible_work_units: int + pending_items: int + oldest_pending_age_seconds: float + + class QueueStatusRow(BaseModel): """Represents a row from the queue status SQL query result.""" diff --git a/src/telemetry/prometheus/metrics.py b/src/telemetry/prometheus/metrics.py index cead893a..cc6fc0b6 100644 --- a/src/telemetry/prometheus/metrics.py +++ b/src/telemetry/prometheus/metrics.py @@ -187,18 +187,58 @@ telemetry_buffer_size_gauge = NamespacedGauge( # region ai # Distinct from embed_now_tasks_in_flight (in-flight fast-path work in the API # process) — this is the durable, DB-wide backlog the reconciler drains. Every -# deriver replica refreshes it on its own timer from +# API replica refreshes it on its own timer from # ReconcilerScheduler._scheduler_loop, so replicas disagree by at most one interval. # Service-wide, not per-process — hence the help string's "never sum()". # endregion message_embeddings_pending_gauge = NamespacedGauge( "message_embeddings_pending", "MessageEmbedding rows awaiting embedding (sync_state='pending'). " - + "Service-wide DB count, reported independently by every replica — " + + "Service-wide DB count, reported independently by every API replica — " + "aggregate with max() or avg(), never sum()", ["namespace"], ) +deriver_queue_work_units_eligible_gauge = NamespacedGauge( + "deriver_queue_work_units_eligible", + "Work units a deriver could claim right now. " + + "Service-wide DB count, reported independently by every API replica — " + + "aggregate with max() or avg(), never sum()", + ["namespace"], +) + +deriver_queue_items_pending_gauge = NamespacedGauge( + "deriver_queue_items_pending", + "Unprocessed queue rows, whether or not they are claimable yet. " + + "Service-wide DB count, reported independently by every API replica — " + + "aggregate with max() or avg(), never sum()", + ["namespace"], +) + +deriver_queue_oldest_pending_age_seconds_gauge = NamespacedGauge( + "deriver_queue_oldest_pending_age_seconds", + "Age of the oldest unprocessed queue row, 0 when the queue is empty. " + + "Service-wide DB value, reported independently by every API replica — " + + "aggregate with max() or avg(), never sum()", + ["namespace"], +) + +seconds_since_last_vector_sync_gauge = NamespacedGauge( + "seconds_since_last_vector_sync", + "Seconds since the newest MessageEmbedding.last_sync_at, -1 when none exists. " + + "Service-wide DB value, reported independently by every API replica — " + + "aggregate with max() or avg(), never sum()", + ["namespace"], +) + +dreams_pending_gauge = NamespacedGauge( + "dreams_pending", + "Collections past the dream document threshold and the min-hours gate with " + + "no dream already queued. Service-wide DB count, reported independently by " + + "every API replica — aggregate with max() or avg(), never sum()", + ["namespace"], +) + # DB connection-pool health. The in-flight gauge counts statements actually # executing on the wire, so checked_out minus in_flight reveals connections held # but parked (the "idle in transaction during an external call" antipattern). @@ -467,10 +507,11 @@ class PrometheusMetrics: # replica count. Scale-preserving aggregations (``max()``, ``avg()``, # quantiles) ARE correct, but only while the witnesses disagree by a bounded # amount — which requires every instance to refresh on its own timer (see - # ``message_embeddings_pending``, refreshed per replica from - # ``ReconcilerScheduler._scheduler_loop``). A bucket-3 metric that cannot - # meet that bar does not belong in the app at all — it belongs in an exporter - # that yields exactly one series. + # ``message_embeddings_pending``, refreshed per API replica from + # ``ReconcilerScheduler._scheduler_loop``, and the deriver-queue gauges, + # refreshed per API replica from ``BacklogMetricsPoller``). A bucket-3 + # metric that cannot meet that bar does not belong in the app at all — it + # belongs in an exporter that yields exactly one series. # # Prometheus stamps ``instance``/``job`` at scrape time, which is why buckets 1 # and 2 need no special handling. ``telemetry_events_dropped`` is handled @@ -507,6 +548,14 @@ class PrometheusMetrics: # ai: embed_now fast path runs as an API-process background task self._touch(embed_now_tasks_shed_counter) self.set_embed_now_tasks_in_flight(0) + self.set_message_embeddings_pending(count=0) + self.set_deriver_backlog( + eligible_work_units=0, + pending_items=0, + oldest_pending_age_seconds=0, + ) + self.set_seconds_since_last_vector_sync(seconds=-1) + self.set_dreams_pending(count=0) elif instance_type == "deriver": # deriver tokens: only the valid (token_type, component) tuples per @@ -533,8 +582,6 @@ class PrometheusMetrics: specialist_name=specialist.name, token_type=token_type.value, ) - # ai: init at 0 so the gauge is visible before its first per-replica refresh - self.set_message_embeddings_pending(count=0) def set_telemetry_buffer_size(self, *, size: int) -> None: try: @@ -548,6 +595,34 @@ class PrometheusMetrics: except Exception as e: self._handle_metric_error("set_message_embeddings_pending", e) + def set_deriver_backlog( + self, + *, + eligible_work_units: int, + pending_items: int, + oldest_pending_age_seconds: float, + ) -> None: + try: + deriver_queue_work_units_eligible_gauge.labels().set(eligible_work_units) + deriver_queue_items_pending_gauge.labels().set(pending_items) + deriver_queue_oldest_pending_age_seconds_gauge.labels().set( + oldest_pending_age_seconds + ) + except Exception as e: + self._handle_metric_error("set_deriver_backlog", e) + + def set_seconds_since_last_vector_sync(self, *, seconds: float) -> None: + try: + seconds_since_last_vector_sync_gauge.labels().set(seconds) + except Exception as e: + self._handle_metric_error("set_seconds_since_last_vector_sync", e) + + def set_dreams_pending(self, *, count: int) -> None: + try: + dreams_pending_gauge.labels().set(count) + except Exception as e: + self._handle_metric_error("set_dreams_pending", e) + prometheus_metrics = PrometheusMetrics() diff --git a/tests/crud/test_deriver_backlog.py b/tests/crud/test_deriver_backlog.py new file mode 100644 index 00000000..cc9b6de9 --- /dev/null +++ b/tests/crud/test_deriver_backlog.py @@ -0,0 +1,290 @@ +import datetime + +import pytest +from nanoid import generate as generate_nanoid +from sqlalchemy.ext.asyncio import AsyncSession + +from src import crud, models +from src.config import settings + +pytestmark = pytest.mark.asyncio + + +async def _make_session( + db: AsyncSession, workspace: models.Workspace +) -> models.Session: + session = models.Session(name=str(generate_nanoid()), workspace_name=workspace.name) + db.add(session) + await db.flush() + return session + + +async def _add_representation_item( + db: AsyncSession, + workspace: models.Workspace, + peer: models.Peer, + session: models.Session, + *, + work_unit_key: str, + token_count: int, + age_seconds: int = 0, + seq: int = 1, +) -> models.QueueItem: + message = models.Message( + session_name=session.name, + content="x", + token_count=token_count, + seq_in_session=seq, + peer_name=peer.name, + workspace_name=workspace.name, + ) + db.add(message) + await db.flush() + + item = models.QueueItem( + session_id=session.id, + work_unit_key=work_unit_key, + task_type="representation", + payload={}, + processed=False, + workspace_name=workspace.name, + message_id=message.id, + created_at=datetime.datetime.now(datetime.timezone.utc) + - datetime.timedelta(seconds=age_seconds), + ) + db.add(item) + await db.flush() + return item + + +class TestDeriverBacklog: + async def test_empty_queue_reports_zero( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + backlog = await crud.get_deriver_backlog(db_session) + + assert backlog.eligible_work_units == 0 + assert backlog.pending_items == 0 + assert backlog.oldest_pending_age_seconds == 0.0 + + async def test_sub_threshold_batch_is_pending_but_not_eligible( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """A small, fresh batch is real work that a deriver would not yet claim. + + This is the case that makes eligible and pending different numbers, and + the reason pending exists as its own gauge. + """ + workspace, peer = sample_data + session = await _make_session(db_session, workspace) + + await _add_representation_item( + db_session, + workspace, + peer, + session, + work_unit_key="representation:small", + token_count=1, + ) + await db_session.commit() + + backlog = await crud.get_deriver_backlog(db_session) + + assert backlog.pending_items == 1 + assert backlog.eligible_work_units == 0 + + async def test_token_threshold_makes_batch_eligible( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + workspace, peer = sample_data + session = await _make_session(db_session, workspace) + + await _add_representation_item( + db_session, + workspace, + peer, + session, + work_unit_key="representation:big", + token_count=settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS, + ) + await db_session.commit() + + backlog = await crud.get_deriver_backlog(db_session) + + assert backlog.eligible_work_units == 1 + + async def test_age_flush_makes_sub_threshold_batch_eligible( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + workspace, peer = sample_data + session = await _make_session(db_session, workspace) + + await _add_representation_item( + db_session, + workspace, + peer, + session, + work_unit_key="representation:old", + token_count=1, + age_seconds=settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS + 60, + ) + await db_session.commit() + + backlog = await crud.get_deriver_backlog(db_session) + + assert backlog.eligible_work_units == 1 + assert backlog.oldest_pending_age_seconds >= ( + settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS + ) + + async def test_non_representation_work_is_eligible_immediately( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + workspace, _peer = sample_data + + db_session.add( + models.QueueItem( + work_unit_key="reconciler:sync_vectors", + task_type="reconciler", + payload={}, + processed=False, + workspace_name=workspace.name, + ) + ) + await db_session.commit() + + backlog = await crud.get_deriver_backlog(db_session) + + assert backlog.eligible_work_units == 1 + + async def test_live_claim_hides_work_from_the_count( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + workspace, peer = sample_data + session = await _make_session(db_session, workspace) + + await _add_representation_item( + db_session, + workspace, + peer, + session, + work_unit_key="representation:claimed", + token_count=settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS, + ) + db_session.add( + models.ActiveQueueSession(work_unit_key="representation:claimed") + ) + await db_session.commit() + + backlog = await crud.get_deriver_backlog(db_session) + + assert backlog.eligible_work_units == 0 + + async def test_stale_claim_does_not_hide_work( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + workspace, peer = sample_data + session = await _make_session(db_session, workspace) + + await _add_representation_item( + db_session, + workspace, + peer, + session, + work_unit_key="representation:abandoned", + token_count=settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS, + ) + stale_cutoff = datetime.datetime.now( + datetime.timezone.utc + ) - datetime.timedelta( + minutes=settings.DERIVER.STALE_SESSION_TIMEOUT_MINUTES + 1 + ) + db_session.add( + models.ActiveQueueSession( + work_unit_key="representation:abandoned", + last_updated=stale_cutoff, + ) + ) + await db_session.commit() + + backlog = await crud.get_deriver_backlog(db_session) + + assert backlog.eligible_work_units == 1 + + async def test_processed_items_are_not_counted( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + workspace, peer = sample_data + session = await _make_session(db_session, workspace) + + item = await _add_representation_item( + db_session, + workspace, + peer, + session, + work_unit_key="representation:done", + token_count=settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS, + ) + item.processed = True + await db_session.commit() + + backlog = await crud.get_deriver_backlog(db_session) + + assert backlog.pending_items == 0 + assert backlog.eligible_work_units == 0 + assert backlog.oldest_pending_age_seconds == 0.0 + + +class TestBacklogAgreesWithDeriver: + @pytest.mark.parametrize( + "token_count,age_seconds", + [ + (1, 0), + (settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS, 0), + (1, settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS + 60), + ], + ids=["sub-threshold", "token-threshold", "age-flush"], + ) + async def test_eligible_count_matches_what_the_deriver_claims( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + token_count: int, + age_seconds: int, + ): + from src.deriver.queue_manager import QueueManager + + workspace, peer = sample_data + session = await _make_session(db_session, workspace) + + await _add_representation_item( + db_session, + workspace, + peer, + session, + work_unit_key="representation:agreement", + token_count=token_count, + age_seconds=age_seconds, + ) + await db_session.commit() + + expected = (await crud.get_deriver_backlog(db_session)).eligible_work_units + claimed = await QueueManager().get_and_claim_work_units() + + assert len(claimed) == expected diff --git a/tests/telemetry/test_metric_zero_init.py b/tests/telemetry/test_metric_zero_init.py index e69f50df..6e95cb8d 100644 --- a/tests/telemetry/test_metric_zero_init.py +++ b/tests/telemetry/test_metric_zero_init.py @@ -131,6 +131,16 @@ def test_deriver_token_combos_are_valid_and_complete(): ) not in ingestion +_API_BUCKET_3_GAUGES = ( + "message_embeddings_pending", + "deriver_queue_work_units_eligible", + "deriver_queue_items_pending", + "deriver_queue_oldest_pending_age_seconds", + "seconds_since_last_vector_sync", + "dreams_pending", +) + + # --------------------------------------------------------------------------- # API-process zero-init # --------------------------------------------------------------------------- @@ -161,6 +171,15 @@ def test_api_init_materializes_dialectic_and_embed(): ) assert sample("embed_now_tasks_shed_total") is not None assert sample("embed_now_tasks_in_flight") == 0.0 # gauge, explicit .set(0) + for gauge in _API_BUCKET_3_GAUGES: + assert sample(gauge) is not None, f"{gauge} was not zero-initialized" + assert sample("message_embeddings_pending") == 0.0 + assert sample("deriver_queue_work_units_eligible") == 0.0 + assert sample("deriver_queue_items_pending") == 0.0 + assert sample("deriver_queue_oldest_pending_age_seconds") == 0.0 + assert sample("dreams_pending") == 0.0 + # -1, not 0: "never synced" must not read as "just synced" to a threshold + assert sample("seconds_since_last_vector_sync") == -1.0 @pytest.mark.usefixtures("metrics_enabled") @@ -212,7 +231,6 @@ def test_deriver_init_materializes_token_and_backlog(): ) is not None ), f"specialist {specialist_name!r} was not zero-initialized" - assert sample("message_embeddings_pending") == 0.0 # gauge zero-init @pytest.mark.usefixtures("metrics_enabled") @@ -310,6 +328,8 @@ def test_deriver_init_does_not_touch_api_counters(): # the API-process embed_now counters are equally off-limits assert sample("embed_now_tasks_shed_total") is None assert sample("embed_now_tasks_in_flight") is None + for gauge in _API_BUCKET_3_GAUGES: + assert sample(gauge) is None, f"{gauge} must be API-only" # ---------------------------------------------------------------------------