RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat(deriver): Moves scheduling events from deriver to API process

This commit is contained in:
Ulysse Pence 2026-08-26 19:50:17 -02:00
parent 168185ae2b
commit 79a3b9d416
12 changed files with 788 additions and 60 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

210
src/backlog.py Normal file
View File

@ -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

2
src/config.py
View File

@ -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:

7
src/crud/__init__.py
View File

@ -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

121
src/crud/deriver.py
View File

@ -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")

77
src/deriver/queue_manager.py
View File

@ -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()

16
src/main.py
View File

@ -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()

2
src/reconciler/scheduler.py
View File

@ -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():

2
src/schemas/__init__.py
View File

@ -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",

8
src/schemas/internal.py
View File

@ -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."""

91
src/telemetry/prometheus/metrics.py
View File

@ -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()

290
tests/crud/test_deriver_backlog.py Normal file
View File

@ -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

22
tests/telemetry/test_metric_zero_init.py
View File

@ -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"
# ---------------------------------------------------------------------------