From 03a2374ea15334eacf5632a0fbbeafab76abcad5 Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Tue, 28 Apr 2026 16:01:33 -0400 Subject: [PATCH] fix: give vector sync a substantial retry budget (#604) --- src/reconciler/sync_vectors.py | 41 +++++--- tests/conftest.py | 1 + tests/deriver/test_vector_reconciliation.py | 109 ++++++++++++++++++++ 3 files changed, 135 insertions(+), 16 deletions(-) diff --git a/src/reconciler/sync_vectors.py b/src/reconciler/sync_vectors.py index 4a17e40e..34a2fab4 100644 --- a/src/reconciler/sync_vectors.py +++ b/src/reconciler/sync_vectors.py @@ -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() diff --git a/tests/conftest.py b/tests/conftest.py index 3c9b8e63..80389d34 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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), diff --git a/tests/deriver/test_vector_reconciliation.py b/tests/deriver/test_vector_reconciliation.py index cc637597..20958e00 100644 --- a/tests/deriver/test_vector_reconciliation.py +++ b/tests/deriver/test_vector_reconciliation.py @@ -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: