fix: give vector sync a substantial retry budget (#604)
This commit is contained in:
parent
b778d82319
commit
03a2374ea1
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Reference in New Issue