From c300236c110c6e544ced07c843c5daf44fa133a5 Mon Sep 17 00:00:00 2001 From: Aakash Kattelu Date: Tue, 1 Sep 2026 09:17:34 -0400 Subject: [PATCH] fix(deriver): reduce scope backfill memory usage (#1104) * fix(deriver): chunk scope backfill so large sessions don't OOM the worker _run_backfill embedded, wrote, and synced every planned copy at once, holding one Python float list per document. A 14k-document session is ~580MB of vectors alone, and several backfills run concurrently, which OOM-killed the deriver at its 1000Mi limit and crash-looped it since the work units never completed. Phases 2-4 now run per chunk of 500 specs and drop each chunk's embeddings once synced. Co-Authored-By: Claude Fable 5 * fix(deriver): hydrate backfill embeddings per chunk Phase 1 no longer materializes every source embedding into plans. load_only skips the vector column on the plan queries, and each chunk reloads only its source embeddings before embed/write/sync. * fix(deriver): lock scope membership across backfill chunk writes SELECT ... FOR UPDATE on the active SessionPeer row so a concurrent leave cannot commit between the membership check and the copy inserts. Adds a concurrency test that asserts the leave blocks until commit. * fix: add test for memory bound --------- Co-authored-by: Claude Fable 5 Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> --- src/deriver/scope_backfill.py | 114 ++++++++++++-- tests/deriver/test_scope_backfill.py | 212 ++++++++++++++++++++++++++- 2 files changed, 310 insertions(+), 16 deletions(-) diff --git a/src/deriver/scope_backfill.py b/src/deriver/scope_backfill.py index 81017e28..9be082f9 100644 --- a/src/deriver/scope_backfill.py +++ b/src/deriver/scope_backfill.py @@ -30,12 +30,12 @@ from typing import Any from sqlalchemy import select, update from sqlalchemy.dialects.postgresql import array +from sqlalchemy.orm import load_only from sqlalchemy.sql.functions import func from src import crud, models from src.config import settings from src.crud.scope import ScopeBackfillState -from src.crud.session import is_peer_in_session from src.dependencies import tracked_db from src.embedding_client import embedding_client from src.schemas import DreamType @@ -49,6 +49,10 @@ logger = logging.getLogger(__name__) # was copied from. The presence of this key is the idempotency marker. COPIED_FROM_KEY = "copied_from" +# Specs embedded, written, and synced per pass. Bounds the live embeddings +# (~40KB each as Python floats) so a large session cannot OOM the deriver. +BACKFILL_CHUNK_SIZE = 500 + def _store_embeddings_in_postgres() -> bool: """Whether document embeddings are persisted to the postgres column. @@ -173,7 +177,23 @@ async def _run_backfill( plans: list[_CopySpec] = [] async with tracked_db("scope_backfill.plan") as db: source_result = await db.execute( - select(models.Document).where( + select(models.Document) + .options( + load_only( + models.Document.id, + models.Document.workspace_name, + models.Document.observer, + models.Document.observed, + models.Document.content, + models.Document.level, + models.Document.times_derived, + models.Document.internal_metadata, + models.Document.session_name, + models.Document.source_ids, + models.Document.deleted_at, + ) + ) + .where( models.Document.workspace_name == workspace_name, models.Document.session_name == session_name, models.Document.level == "explicit", @@ -196,7 +216,19 @@ async def _run_backfill( # by (observed, copied_from). Includes soft-deleted rows: those are # restore candidates, not blockers. copies_result = await db.execute( - select(models.Document).where( + select(models.Document) + .options( + load_only( + models.Document.id, + models.Document.workspace_name, + models.Document.observer, + models.Document.observed, + models.Document.session_name, + models.Document.internal_metadata, + models.Document.deleted_at, + ) + ) + .where( models.Document.workspace_name == workspace_name, models.Document.observer == scope_peer, models.Document.session_name == session_name, @@ -216,12 +248,13 @@ async def _run_backfill( key = (source.observed, source.id) if key in live_copies: continue + # Vectors hydrate per chunk; plans only carry ids + content. plans.append( _CopySpec( observed=source.observed, source_id=source.id, content=source.content, - embedding=_embedding_as_list(source.embedding), + embedding=None, internal_metadata=dict(source.internal_metadata), times_derived=source.times_derived, source_ids=list(source.source_ids) @@ -235,7 +268,54 @@ async def _run_backfill( if not plans: return 0, set() - # Phase 2 (no DB): fill missing embeddings. Source rows have NULL + # Phases 2-4 run per chunk so only one chunk's embeddings are alive at a + # time; each chunk's vectors are dropped once synced. + store_in_postgres = _store_embeddings_in_postgres() + touched_observed: set[str] = set() + copied = 0 + for start in range(0, len(plans), BACKFILL_CHUNK_SIZE): + chunk = plans[start : start + BACKFILL_CHUNK_SIZE] + if not await _copy_chunk( + workspace_name, scope_peer, session_name, chunk, store_in_postgres + ): + return None + copied += len(chunk) + touched_observed.update(spec.observed for spec in chunk) + for spec in chunk: + spec.embedding = None + + return copied, touched_observed + + +async def _hydrate_chunk_embeddings( + workspace_name: str, plans: list[_CopySpec] +) -> None: + """Load this chunk's source embeddings from postgres (if any).""" + source_ids = [spec.source_id for spec in plans] + async with tracked_db("scope_backfill.hydrate_embeddings") as db: + result = await db.execute( + select(models.Document.id, models.Document.embedding).where( + models.Document.workspace_name == workspace_name, + models.Document.id.in_(source_ids), + ) + ) + by_id = {row.id: _embedding_as_list(row.embedding) for row in result.all()} + for spec in plans: + spec.embedding = by_id.get(spec.source_id) + + +async def _copy_chunk( + workspace_name: str, + scope_peer: str, + session_name: str, + plans: list[_CopySpec], + store_in_postgres: bool, +) -> bool: + """Embed, write, and sync one chunk. False if the session left the scope.""" + # Phase 2a (DB): pull this chunk's embeddings only. + await _hydrate_chunk_embeddings(workspace_name, plans) + + # Phase 2b (no DB): fill missing embeddings. Source rows have NULL # embeddings on external-store deployments (and soft-deleted copies may # have lost their vectors) — re-embed via the embedding API only; no LLM. missing = [spec for spec in plans if spec.embedding is None] @@ -254,17 +334,22 @@ async def _run_backfill( spec.embedding = embedding # Phase 3 (DB): write the copies. - store_in_postgres = _store_embeddings_in_postgres() touched_observed = {spec.observed for spec in plans} new_rows: list[models.Document] = [] async with tracked_db("scope_backfill.write") as db: - # scope_backfill and scope_removal carry different work-unit keys, so - # nothing orders them: a removal enqueued right after the add (or one - # that landed while phase 2 was embedding) can sweep the scope before - # these copies exist. Re-checking membership here, in the transaction - # that inserts, keeps a removed session from being copied back in. - if not await is_peer_in_session(db, workspace_name, session_name, scope_peer): - return None + # Row-lock active membership for this txn so a concurrent leave + # (``left_at``) cannot commit between the check and the inserts. + membership = await db.scalar( + select(models.SessionPeer.peer_name) + .where(models.SessionPeer.workspace_name == workspace_name) + .where(models.SessionPeer.session_name == session_name) + .where(models.SessionPeer.peer_name == scope_peer) + .where(models.SessionPeer.left_at.is_(None)) + .with_for_update() + .limit(1) + ) + if membership is None: + return False for observed in sorted(touched_observed): await crud.get_or_create_collection( @@ -319,8 +404,7 @@ async def _run_backfill( # Phase 4: sync to the external vector store (or mark synced in pgvector # mode). Failures leave rows in sync_state='pending' for the reconciler. await _sync_copies_to_vector_store(workspace_name, scope_peer, plans, copied_ids) - - return len(plans), touched_observed + return True async def _sync_copies_to_vector_store( diff --git a/tests/deriver/test_scope_backfill.py b/tests/deriver/test_scope_backfill.py index 06464c6f..efca34f3 100644 --- a/tests/deriver/test_scope_backfill.py +++ b/tests/deriver/test_scope_backfill.py @@ -22,15 +22,19 @@ engine (see ``mock_tracked_db_context`` in conftest.py) — a different connection that cannot see another session's uncommitted writes. """ +import asyncio +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager from typing import Any import pytest from fastapi.testclient import TestClient from nanoid import generate as generate_nanoid from sqlalchemy import func, select, update -from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker from src import crud, models +from src.deriver import scope_backfill as scope_backfill_mod from src.deriver.scope_backfill import ( COPIED_FROM_KEY, process_scope_backfill, @@ -415,6 +419,128 @@ async def test_backfill_skips_a_session_that_left_the_scope( assert session_name not in peer.internal_metadata.get("backfill_status", {}) +@pytest.mark.asyncio +async def test_copy_chunk_membership_lock_blocks_leave_until_write_commits( + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + db_engine: AsyncEngine, + monkeypatch: pytest.MonkeyPatch, +): + """A concurrent leave cannot commit between membership check and inserts.""" + test_workspace, sender = sample_data + workspace_name = test_workspace.name + scope_name = str(generate_nanoid()) + scope_peer = await _create_scope_peer(db_session, workspace_name, scope_name) + session = await _create_session(db_session, workspace_name) + await _join_scope(db_session, workspace_name, session.name, scope_peer.name) + await _create_collection( + db_session, workspace_name, observer=sender.name, observed=sender.name + ) + await _create_collection( + db_session, workspace_name, observer=scope_peer.name, observed=sender.name + ) + source = await _create_document( + db_session, + workspace_name, + observer=sender.name, + observed=sender.name, + session_name=session.name, + content="locked membership fact", + ) + + factory = async_sessionmaker(bind=db_engine, expire_on_commit=False) + leave_finished = asyncio.Event() + leave_task_box: dict[str, asyncio.Task[None]] = {} + + async def concurrent_leave() -> None: + async with factory() as leave_db: + await leave_db.execute( + update(models.SessionPeer) + .where( + models.SessionPeer.workspace_name == workspace_name, + models.SessionPeer.session_name == session.name, + models.SessionPeer.peer_name == scope_peer.name, + models.SessionPeer.left_at.is_(None), + ) + .values(left_at=func.now()) + ) + await leave_db.commit() + leave_finished.set() + + original_tracked_db = scope_backfill_mod.tracked_db # pyright: ignore[reportPrivateLocalImportUsage] + + @asynccontextmanager + async def tracked_db_with_leave_race( + operation_name: str | None = None, *, read_only: bool = False + ) -> AsyncGenerator[AsyncSession]: + async with original_tracked_db(operation_name, read_only=read_only) as db: + if operation_name == "scope_backfill.write": + real_scalar = db.scalar + raced = False + + async def scalar_then_race(statement: Any, *args: Any, **kwargs: Any): + nonlocal raced + result = await real_scalar(statement, *args, **kwargs) + if not raced and result is not None: + raced = True + leave_task_box["task"] = asyncio.create_task(concurrent_leave()) + # Leave's UPDATE must block on this txn's row lock. + for _ in range(50): + await asyncio.sleep(0.01) + if leave_task_box["task"].done(): + break + assert not leave_task_box["task"].done() + return result + + db.scalar = scalar_then_race # type: ignore[method-assign] + yield db + + monkeypatch.setattr(scope_backfill_mod, "tracked_db", tracked_db_with_leave_race) + + ok = await scope_backfill_mod._copy_chunk( # pyright: ignore[reportPrivateUsage] + workspace_name, + scope_peer.name, + session.name, + [ + scope_backfill_mod._CopySpec( # pyright: ignore[reportPrivateUsage] + observed=sender.name, + source_id=source.id, + content=source.content, + embedding=None, + internal_metadata={}, + times_derived=1, + source_ids=None, + session_name=session.name, + ) + ], + store_in_postgres=True, + ) + assert ok is True + + leave_task = leave_task_box["task"] + await asyncio.wait_for(leave_task, timeout=2.0) + assert leave_finished.is_set() + + copies = await _get_docs( + db_session, + workspace_name, + observer=scope_peer.name, + observed=sender.name, + include_deleted=False, + ) + assert len(copies) == 1 + assert copies[0].internal_metadata.get(COPIED_FROM_KEY) == source.id + + membership = await db_session.scalar( + select(models.SessionPeer.left_at).where( + models.SessionPeer.workspace_name == workspace_name, + models.SessionPeer.session_name == session.name, + models.SessionPeer.peer_name == scope_peer.name, + ) + ) + assert membership is not None + + # --------------------------------------------------------------------------- # 3. Multi-peer session # --------------------------------------------------------------------------- @@ -956,3 +1082,87 @@ async def test_backfill_status_writes_preserve_the_scope_kind_flag( await db_session.commit() metadata = await assert_still_a_scope("clearing the status") assert session_name not in metadata.get("backfill_status", {}) + + +@pytest.mark.asyncio +async def test_backfill_embeds_and_writes_in_bounded_chunks( + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + monkeypatch: pytest.MonkeyPatch, +): + """Phases 2-4 run per chunk, so a large session never holds every vector.""" + from src.deriver import scope_backfill + from src.embedding_client import embedding_client + + test_workspace, sender = sample_data + workspace_name = test_workspace.name + scope_name = str(generate_nanoid()) + scope_peer = await _create_scope_peer(db_session, workspace_name, scope_name) + session = await _create_session(db_session, workspace_name) + await _join_scope(db_session, workspace_name, session.name, scope_peer.name) + await _create_collection( + db_session, workspace_name, observer=sender.name, observed=sender.name + ) + await _create_collection( + db_session, workspace_name, observer=scope_peer.name, observed=sender.name + ) + for i in range(3): + source = await _create_document( + db_session, + workspace_name, + observer=sender.name, + observed=sender.name, + session_name=session.name, + content=f"fact {i}", + ) + source.embedding = None + await db_session.commit() + + batch_sizes: list[int] = [] + seen_specs: list[scope_backfill._CopySpec] = [] # pyright: ignore[reportPrivateUsage] + peak_live_embeddings = 0 + original_embed = embedding_client.simple_batch_embed + original_copy_chunk = scope_backfill._copy_chunk # pyright: ignore[reportPrivateUsage] + + async def recording_embed(texts: list[str], **kwargs: Any) -> list[list[float]]: + batch_sizes.append(len(texts)) + return await original_embed(texts, **kwargs) + + async def counting_copy_chunk( + ws_name: str, + peer_name: str, + sess_name: str, + plans: list[scope_backfill._CopySpec], # pyright: ignore[reportPrivateUsage] + store_in_postgres: bool, + ) -> bool: + nonlocal peak_live_embeddings + seen_specs.extend(plans) + result = await original_copy_chunk( + ws_name, peer_name, sess_name, plans, store_in_postgres + ) + # Sampled after this chunk syncs but before _run_backfill drops its + # vectors, so every *earlier* chunk must already be cleared and the + # live count can never exceed one chunk. That drop is the whole + # memory bound; without it this peaks at 3 instead of 2. + peak_live_embeddings = max( + peak_live_embeddings, + sum(1 for spec in seen_specs if spec.embedding is not None), + ) + return result + + monkeypatch.setattr(scope_backfill, "BACKFILL_CHUNK_SIZE", 2) + monkeypatch.setattr(embedding_client, "simple_batch_embed", recording_embed) + monkeypatch.setattr(scope_backfill, "_copy_chunk", counting_copy_chunk) + + await process_scope_backfill( + ScopeBackfillPayload(scope_peer=scope_peer.name, session_name=session.name), + workspace_name, + ) + + assert batch_sizes == [2, 1] + assert peak_live_embeddings == 2 + copies = await _get_docs( + db_session, workspace_name, observer=scope_peer.name, observed=sender.name + ) + assert len(copies) == 3 + assert all(copy.embedding is not None for copy in copies)