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

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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>
Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
This commit is contained in:
Aakash Kattelu 2026-09-01 09:17:34 -04:00 committed by GitHub
parent 03253d7a08
commit c300236c11
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 310 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

114
src/deriver/scope_backfill.py
View File

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

212
tests/deriver/test_scope_backfill.py
View File

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