fix: give vector sync a substantial retry budget (#604)

This commit is contained in:
Rajat Ahuja 2026-04-28 16:01:33 -04:00 committed by GitHub
parent b778d82319
commit 03a2374ea1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 135 additions and 16 deletions

View File

@ -11,8 +11,10 @@ import time
from dataclasses import dataclass
from typing import cast
from sqlalchemy import and_, delete, select, update
from sqlalchemy import and_, delete, or_, select, update
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm.attributes import InstrumentedAttribute
from sqlalchemy.sql import ColumnElement
from sqlalchemy.sql.functions import func
from src import models
@ -27,7 +29,20 @@ logger = logging.getLogger(__name__)
# Constants
RECONCILIATION_BATCH_SIZE = 50
RECONCILIATION_TIME_BUDGET_SECONDS = 240 # Leave headroom for other maintenance work
MAX_SYNC_ATTEMPTS = 5 # After this many failures, mark as failed
MAX_SYNC_ATTEMPTS = 20 # After this many failures, mark as failed
# Flat wait between sync attempts. With MAX_SYNC_ATTEMPTS=20 this gives ~3 hours
# of outage headroom before a row is marked failed.
SYNC_BACKOFF = datetime.timedelta(minutes=10)
def _backoff_eligible(
last_sync_at: InstrumentedAttribute[datetime.datetime | None],
) -> ColumnElement[bool]:
"""Rows are eligible for sync if never attempted or past the backoff window."""
return or_(
last_sync_at.is_(None),
last_sync_at < func.now() - SYNC_BACKOFF,
)
@dataclass
@ -73,6 +88,7 @@ async def _get_documents_needing_sync(
and_(
models.Document.deleted_at.is_(None),
models.Document.sync_state == "pending", # Only pending items
_backoff_eligible(models.Document.last_sync_at),
)
)
.order_by(models.Document.last_sync_at.asc().nullsfirst())
@ -101,7 +117,12 @@ async def _get_message_embeddings_needing_sync(
"""
stmt = (
select(models.MessageEmbedding)
.where(models.MessageEmbedding.sync_state == "pending")
.where(
and_(
models.MessageEmbedding.sync_state == "pending",
_backoff_eligible(models.MessageEmbedding.last_sync_at),
)
)
.order_by(models.MessageEmbedding.last_sync_at.asc().nullsfirst())
.limit(batch_size)
.with_for_update(skip_locked=True)
@ -494,19 +515,7 @@ async def _reconcile_message_embeddings_batch(
if not embs:
return False
try:
synced, failed = await _sync_message_embeddings(
db, embs, external_vector_store
)
except Exception:
logger.exception(
"Message embedding reconciliation failed for %s embeddings",
len(embs),
)
await _bump_message_embedding_sync_attempts(db, embs)
synced = 0
failed = len(embs)
synced, failed = await _sync_message_embeddings(db, embs, external_vector_store)
metrics.message_embeddings_synced += synced
metrics.message_embeddings_failed += failed
await db.commit()

View File

@ -799,6 +799,7 @@ def mock_tracked_db(request: pytest.FixtureRequest):
patch("src.utils.search.tracked_db", mock_tracked_db_context),
patch("src.crud.document.tracked_db", mock_tracked_db_context),
patch("src.crud.message.tracked_db", mock_tracked_db_context),
patch("src.reconciler.sync_vectors.tracked_db", mock_tracked_db_context),
patch("src.dialectic.core.tracked_db", mock_tracked_db_context),
patch("src.dreamer.specialists.tracked_db", mock_tracked_db_context),
patch("src.dreamer.surprisal.tracked_db", mock_tracked_db_context),

View File

@ -20,6 +20,7 @@ from src.reconciler.sync_vectors import (
ReconciliationMetrics,
_get_documents_needing_sync, # pyright: ignore[reportPrivateUsage]
_get_message_embeddings_needing_sync, # pyright: ignore[reportPrivateUsage]
_reconcile_message_embeddings_batch, # pyright: ignore[reportPrivateUsage]
_sync_documents, # pyright: ignore[reportPrivateUsage]
_sync_message_embeddings, # pyright: ignore[reportPrivateUsage]
run_vector_reconciliation_cycle,
@ -377,6 +378,60 @@ class TestBatchProcessing:
# Verify batch size respected
assert len(batch) == 100
async def test_documents_respect_retry_backoff(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
) -> None:
"""Pending documents should only be fetched once their backoff has elapsed."""
workspace, peer1 = sample_data
collection = models.Collection(
workspace_name=workspace.name,
observer=peer1.name,
observed=peer1.name,
)
db_session.add(collection)
await db_session.commit()
session = models.Session(
name=str(generate_nanoid()), workspace_name=workspace.name
)
db_session.add(session)
await db_session.commit()
now = datetime.datetime.now(datetime.timezone.utc)
ineligible_doc = models.Document(
content="too soon",
workspace_name=workspace.name,
observer=peer1.name,
observed=peer1.name,
session_name=session.name,
sync_state="pending",
sync_attempts=1,
last_sync_at=now - datetime.timedelta(minutes=9, seconds=59),
embedding=[1.0] * 1536,
)
eligible_doc = models.Document(
content="ready",
workspace_name=workspace.name,
observer=peer1.name,
observed=peer1.name,
session_name=session.name,
sync_state="pending",
sync_attempts=1,
last_sync_at=now - datetime.timedelta(minutes=10, seconds=1),
embedding=[2.0] * 1536,
)
db_session.add_all([ineligible_doc, eligible_doc])
await db_session.commit()
pending = await _get_documents_needing_sync(db_session)
pending_ids = {doc.id for doc in pending}
assert eligible_doc.id in pending_ids
assert ineligible_doc.id not in pending_ids
@pytest.mark.asyncio
class TestReEmbedding:
@ -705,6 +760,33 @@ class TestMessageEmbeddings:
pending = await _get_message_embeddings_needing_sync(db_session)
assert any(emb.id == pending_emb.id for emb in pending)
async def test_message_embeddings_respect_retry_backoff(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
) -> None:
"""Pending embeddings should only be fetched once their backoff has elapsed."""
workspace, peer = sample_data
ineligible_emb = await self._create_pending_message_embedding(
db_session, workspace, peer
)
eligible_emb = await self._create_pending_message_embedding(
db_session, workspace, peer
)
now = datetime.datetime.now(datetime.timezone.utc)
ineligible_emb.sync_attempts = 1
ineligible_emb.last_sync_at = now - datetime.timedelta(minutes=9, seconds=59)
eligible_emb.sync_attempts = 1
eligible_emb.last_sync_at = now - datetime.timedelta(minutes=10, seconds=1)
await db_session.commit()
pending = await _get_message_embeddings_needing_sync(db_session)
pending_ids = {emb.id for emb in pending}
assert eligible_emb.id in pending_ids
assert ineligible_emb.id not in pending_ids
async def test_missing_embeddings_reembedded_and_synced(
self,
db_session: AsyncSession,
@ -764,6 +846,33 @@ class TestMessageEmbeddings:
assert pending_emb.sync_state in {"pending", "failed"}
assert pending_emb.sync_attempts == 1
async def test_unexpected_batch_exception_does_not_bump_unattempted_rows(
self,
db_session: AsyncSession,
sample_data: tuple[models.Workspace, models.Peer],
mock_vector_store: VectorStore,
) -> None:
"""Unexpected wrapper-level failures should not penalize the whole batch."""
workspace, peer = sample_data
pending_emb = await self._create_pending_message_embedding(
db_session, workspace, peer
)
metrics = ReconciliationMetrics()
with (
patch(
"src.reconciler.sync_vectors._sync_message_embeddings",
side_effect=RuntimeError("unexpected"),
),
pytest.raises(RuntimeError, match="unexpected"),
):
await _reconcile_message_embeddings_batch(mock_vector_store, metrics)
await db_session.refresh(pending_emb)
assert pending_emb.sync_state == "pending"
assert pending_emb.sync_attempts == 0
assert pending_emb.last_sync_at is None
@pytest.mark.asyncio
class TestEndToEndReconciliation: