Make rebuild --embed-only idempotent across crashes
This commit is contained in:
Yiorgis Gozadinos
parent
a035b0f9e4
commit
849eca94d0
5 changed files with 654 additions and 88 deletions
|
|
@ -3,7 +3,8 @@
|
|||
|
||||
### Fixed
|
||||
|
||||
- **`rebuild --embed-only` no longer buffers the entire corpus in memory.** The previous implementation accumulated every chunk's id, content, content_fts, metadata, and new embedding vector in a single Python list before flushing — on a multi-million-chunk corpus with a 3072-dim model this trivially exceeded 40 GB of RAM. The rebuild now stream-copies non-vector columns into a `chunks_rebuild_staging` table (1000 rows / page), recreates the chunks table fresh to honour vector-dim changes, then streams from staging one document at a time, embedding in batches of `embeddings.batch_size` and flushing to the new chunks table every 50 documents. Peak memory is bounded to one batch regardless of corpus size; LanceDB OSS does not support `rename_table`, so the staging copy is the only safe way to preserve chunk identity while the live table is recreated. A leftover staging table from an interrupted previous rebuild is dropped at the top of `rebuild_database` — note that an interruption mid-rebuild may still leave the chunks table in a partial state until a follow-up adds idempotent recovery from staging.
|
||||
- **`rebuild --embed-only` no longer buffers the entire corpus in memory.** The previous implementation accumulated every chunk's id, content, content_fts, metadata, and new embedding vector in a single Python list before flushing — on a multi-million-chunk corpus with a 3072-dim model this trivially exceeded 40 GB of RAM. The rebuild now stream-copies non-vector columns into a `chunks_rebuild_staging` table (1000 rows / page), recreates the chunks table fresh to honour vector-dim changes, then streams from staging one document at a time, embedding in batches of `embeddings.batch_size` and flushing to the new chunks table every 50 documents. Peak memory is bounded to one batch regardless of corpus size; LanceDB OSS does not support `rename_table`, so the staging copy is the only safe way to preserve chunk identity while the live table is recreated.
|
||||
- **`rebuild --embed-only` is now idempotent across crashes.** A second table, `chunks_rebuild_marker`, is written immediately after phase 1 (staging copy) finishes. Its presence flips the next rebuild into resume mode: phase 1 is skipped, the live chunks table is recreated, and phase 2 (re-embed) runs from the existing staging snapshot. Cleanup drops the marker before the staging table, so an interruption between the two drops leaves a markerless staging that the next run discards harmlessly. A staging table without a marker is treated as a partial phase 1 and dropped (the live chunks table is still authoritative). Running a non-embed-only mode (FULL / RECHUNK / DESCRIPTIONS / TITLE_ONLY) after a crashed embed-only correctly discards the staging recovery state. Phase 1's pagination was switched from `offset/limit` to `to_batches`, removing the latent offset-drift risk and the O(N²) cost at high offsets.
|
||||
|
||||
## [0.46.0] - 2026-05-13
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from lancedb.pydantic import LanceModel
|
|||
|
||||
from haiku.rag.client.documents import check_source_accessible
|
||||
from haiku.rag.converters import get_converter
|
||||
from haiku.rag.store.engine import ChunkRecordBase
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
from haiku.rag.store.models.document import Document
|
||||
from haiku.rag.store.models.document_item import extract_items
|
||||
|
|
@ -20,7 +21,8 @@ logger = logging.getLogger(__name__)
|
|||
|
||||
_REBUILD_BATCH_SIZE = 50
|
||||
_STAGING_TABLE_NAME = "chunks_rebuild_staging"
|
||||
_STAGING_COPY_PAGE_SIZE = 1000
|
||||
_STAGING_MARKER_TABLE_NAME = "chunks_rebuild_marker"
|
||||
_STAGING_COPY_BATCH_SIZE = 1000
|
||||
|
||||
|
||||
class _StagingChunkRecord(LanceModel):
|
||||
|
|
@ -39,6 +41,19 @@ class _StagingChunkRecord(LanceModel):
|
|||
order: int
|
||||
|
||||
|
||||
class _StagingMarkerRecord(LanceModel):
|
||||
"""Sentinel marking the staging table as complete.
|
||||
|
||||
The marker table is created only after ``_populate_staging_table`` writes
|
||||
every chunk into staging. Its presence at the top of a rebuild means
|
||||
phase 2 (the embed loop) was interrupted by an earlier crash, so staging
|
||||
is the authoritative source for the original chunk identities and we
|
||||
must resume from it instead of rerunning phase 1.
|
||||
"""
|
||||
|
||||
id: str
|
||||
|
||||
|
||||
async def rebuild_database(
|
||||
client: "HaikuRAG", mode: "RebuildMode | None" = None
|
||||
) -> AsyncGenerator[str, None]:
|
||||
|
|
@ -51,9 +66,11 @@ async def rebuild_database(
|
|||
if mode is None:
|
||||
mode = RebuildMode.FULL
|
||||
|
||||
# If a previous embed-only rebuild was interrupted, a leftover staging
|
||||
# table may remain. Drop it so the new rebuild starts from a clean slate.
|
||||
await _drop_leftover_staging_table(client)
|
||||
# Resolve any leftover staging/marker tables from a previously
|
||||
# interrupted rebuild. Returns True only when phase 1 was already
|
||||
# complete and the current mode is EMBED_ONLY, in which case we resume
|
||||
# phase 2 from the existing staging table instead of recopying.
|
||||
resume_from_staging = await _resolve_rebuild_recovery(client, mode)
|
||||
|
||||
# Wait for any already-scheduled background vacuum before the destructive
|
||||
# table operations at the top of RECHUNK / FULL. Rebuild drops and
|
||||
|
|
@ -78,7 +95,9 @@ async def rebuild_database(
|
|||
async for doc_id in _rebuild_title_only(client, documents):
|
||||
yield doc_id
|
||||
elif mode == RebuildMode.EMBED_ONLY:
|
||||
async for doc_id in _rebuild_embed_only(client, documents):
|
||||
async for doc_id in _rebuild_embed_only(
|
||||
client, documents, resume_from_staging=resume_from_staging
|
||||
):
|
||||
yield doc_id
|
||||
elif mode == RebuildMode.RECHUNK:
|
||||
await client.chunk_repository.delete_all()
|
||||
|
|
@ -145,29 +164,73 @@ async def _rebuild_title_only(
|
|||
yield doc.id
|
||||
|
||||
|
||||
async def _drop_leftover_staging_table(client: "HaikuRAG") -> None:
|
||||
"""Drop a stale staging table left over from an interrupted rebuild.
|
||||
async def _resolve_rebuild_recovery(client: "HaikuRAG", mode: "RebuildMode") -> bool:
|
||||
"""Resolve any partially-completed rebuild state from a previous crash.
|
||||
|
||||
Same crash-recovery semantics as the other rebuild modes (RECHUNK/FULL
|
||||
also leave the database non-atomic if interrupted): the user is expected
|
||||
to re-run rebuild from whatever state the database is in.
|
||||
Returns ``True`` if ``_rebuild_embed_only`` should resume from the
|
||||
existing staging table (phase 1 was already complete). In all other
|
||||
cases stale recovery tables are dropped and the rebuild starts fresh.
|
||||
|
||||
State at entry → action
|
||||
--------------------------------
|
||||
no staging, no marker → return False (normal start)
|
||||
staging only → drop staging (phase 1 was interrupted; ``chunks`` is intact)
|
||||
marker only → drop marker (corrupted state)
|
||||
staging + marker, embed → return True (resume phase 2 from staging)
|
||||
staging + marker, other → drop both (staging is for embed-only; user picked a different mode)
|
||||
"""
|
||||
tables = (await client.store.db.list_tables()).tables
|
||||
if _STAGING_TABLE_NAME not in tables:
|
||||
return
|
||||
from haiku.rag.client import RebuildMode
|
||||
|
||||
db = client.store.db
|
||||
tables = (await db.list_tables()).tables
|
||||
has_staging = _STAGING_TABLE_NAME in tables
|
||||
has_marker = _STAGING_MARKER_TABLE_NAME in tables
|
||||
|
||||
if not has_staging and not has_marker:
|
||||
return False
|
||||
|
||||
if has_marker and not has_staging:
|
||||
logger.warning(
|
||||
"Found '%s' without staging table; dropping orphaned marker.",
|
||||
_STAGING_MARKER_TABLE_NAME,
|
||||
)
|
||||
await db.drop_table(_STAGING_MARKER_TABLE_NAME)
|
||||
return False
|
||||
|
||||
if not has_marker:
|
||||
logger.warning(
|
||||
"Dropping incomplete '%s' from an interrupted phase 1.",
|
||||
_STAGING_TABLE_NAME,
|
||||
)
|
||||
await db.drop_table(_STAGING_TABLE_NAME)
|
||||
return False
|
||||
|
||||
# has_staging and has_marker
|
||||
if mode == RebuildMode.EMBED_ONLY:
|
||||
logger.warning(
|
||||
"Resuming interrupted embed-only rebuild: phase 2 will run from "
|
||||
"existing '%s'.",
|
||||
_STAGING_TABLE_NAME,
|
||||
)
|
||||
return True
|
||||
|
||||
logger.warning(
|
||||
"Dropping leftover '%s' table from a previous interrupted rebuild.",
|
||||
_STAGING_TABLE_NAME,
|
||||
"Dropping staging tables from a prior embed-only rebuild — current "
|
||||
"mode (%s) does not consume them.",
|
||||
mode.name,
|
||||
)
|
||||
await client.store.db.drop_table(_STAGING_TABLE_NAME)
|
||||
await db.drop_table(_STAGING_MARKER_TABLE_NAME)
|
||||
await db.drop_table(_STAGING_TABLE_NAME)
|
||||
return False
|
||||
|
||||
|
||||
async def _populate_staging_table(client: "HaikuRAG") -> None:
|
||||
"""Stream-copy the non-vector columns of the chunks table into staging.
|
||||
"""Stream the non-vector columns of the chunks table into staging.
|
||||
|
||||
Reads ``_STAGING_COPY_PAGE_SIZE`` rows per page so peak memory stays
|
||||
bounded regardless of corpus size. The vector column is omitted — the
|
||||
point of embed-only rebuild is to regenerate it.
|
||||
Uses ``to_batches`` for a single streaming read (no offset/limit
|
||||
pagination drift), so peak memory stays bounded regardless of corpus
|
||||
size. The vector column is omitted — the point of embed-only rebuild is
|
||||
to regenerate it.
|
||||
"""
|
||||
db = client.store.db
|
||||
tables = (await db.list_tables()).tables
|
||||
|
|
@ -178,17 +241,15 @@ async def _populate_staging_table(client: "HaikuRAG") -> None:
|
|||
if "chunks" not in tables:
|
||||
return
|
||||
|
||||
offset = 0
|
||||
while True:
|
||||
rows = (
|
||||
await client.store.chunks_table.query()
|
||||
.select(["id", "document_id", "content", "metadata", "order"])
|
||||
.offset(offset)
|
||||
.limit(_STAGING_COPY_PAGE_SIZE)
|
||||
.to_arrow()
|
||||
).to_pylist()
|
||||
stream = (
|
||||
await client.store.chunks_table.query()
|
||||
.select(["id", "document_id", "content", "metadata", "order"])
|
||||
.to_batches(max_batch_length=_STAGING_COPY_BATCH_SIZE)
|
||||
)
|
||||
async for batch in stream:
|
||||
rows = batch.to_pylist()
|
||||
if not rows:
|
||||
break
|
||||
continue
|
||||
records = [
|
||||
_StagingChunkRecord(
|
||||
id=r["id"],
|
||||
|
|
@ -200,7 +261,37 @@ async def _populate_staging_table(client: "HaikuRAG") -> None:
|
|||
for r in rows
|
||||
]
|
||||
await staging.add(records)
|
||||
offset += _STAGING_COPY_PAGE_SIZE
|
||||
|
||||
|
||||
async def _mark_phase1_complete(client: "HaikuRAG") -> None:
|
||||
"""Create the marker table that designates staging as authoritative.
|
||||
|
||||
Called after ``_populate_staging_table`` finishes. On crash recovery the
|
||||
marker's presence flips ``_rebuild_embed_only`` into resume mode.
|
||||
"""
|
||||
db = client.store.db
|
||||
if _STAGING_MARKER_TABLE_NAME in (await db.list_tables()).tables:
|
||||
return
|
||||
marker = await db.create_table(
|
||||
_STAGING_MARKER_TABLE_NAME, schema=_StagingMarkerRecord
|
||||
)
|
||||
await marker.add([_StagingMarkerRecord(id="phase1_complete")])
|
||||
|
||||
|
||||
async def _drop_staging_tables(client: "HaikuRAG") -> None:
|
||||
"""Drop the marker first, then the staging table.
|
||||
|
||||
Ordering matters: if a crash interrupts cleanup between the two drops,
|
||||
the next rebuild sees ``staging`` without ``marker`` and treats it as a
|
||||
partial phase 1 → drops staging harmlessly. The reverse order would
|
||||
leak a marker pointing at nothing.
|
||||
"""
|
||||
db = client.store.db
|
||||
tables = (await db.list_tables()).tables
|
||||
if _STAGING_MARKER_TABLE_NAME in tables:
|
||||
await db.drop_table(_STAGING_MARKER_TABLE_NAME)
|
||||
if _STAGING_TABLE_NAME in tables:
|
||||
await db.drop_table(_STAGING_TABLE_NAME)
|
||||
|
||||
|
||||
async def _read_chunks_from_staging(staging_table, document_id: str) -> list[Chunk]:
|
||||
|
|
@ -232,82 +323,101 @@ async def _read_chunks_from_staging(staging_table, document_id: str) -> list[Chu
|
|||
|
||||
|
||||
async def _rebuild_embed_only(
|
||||
client: "HaikuRAG", documents: list[Document]
|
||||
client: "HaikuRAG",
|
||||
documents: list[Document],
|
||||
*,
|
||||
resume_from_staging: bool = False,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Re-embed all chunks without changing chunk boundaries.
|
||||
|
||||
Two-phase pattern that keeps peak memory bounded regardless of corpus
|
||||
size:
|
||||
size and is idempotent across crashes:
|
||||
|
||||
1. Stream-copy the chunks table's non-vector columns into a staging
|
||||
table. LanceDB OSS doesn't support ``rename_table``, so this is the
|
||||
only safe way to preserve the original chunk identity while we drop
|
||||
and recreate the live chunks table with a (potentially) new vector
|
||||
dim.
|
||||
2. Drop and recreate ``chunks`` with the current schema, then stream
|
||||
table, then write a marker row that designates staging as complete.
|
||||
LanceDB OSS does not support ``rename_table``, so the staging copy
|
||||
is the only safe way to preserve chunk identity while the live
|
||||
``chunks`` table is dropped and recreated.
|
||||
2. Drop-and-recreate ``chunks`` with the current schema, then stream
|
||||
from staging one document at a time, re-embed in batches of
|
||||
``embeddings.batch_size``, and flush to the new chunks table every
|
||||
``_REBUILD_BATCH_SIZE`` documents.
|
||||
|
||||
Cleanup runs only on success: a crash anywhere in phase 2 leaves both
|
||||
staging and marker in place so the next rebuild can re-enter phase 2
|
||||
via ``resume_from_staging=True``. The order of the success cleanup —
|
||||
drop marker before staging — keeps an interruption between the two
|
||||
drops recoverable: the next rebuild sees staging without marker and
|
||||
treats it as a partial phase 1, which is harmless because phase 2 has
|
||||
already finished writing the new chunks table.
|
||||
"""
|
||||
from haiku.rag.embeddings import contextualize
|
||||
|
||||
db = client.store.db
|
||||
batch_size = client._config.embeddings.batch_size
|
||||
|
||||
# Phase 1: copy chunks into staging. After this we can safely destroy
|
||||
# the live chunks table without losing chunk identity / content.
|
||||
await _populate_staging_table(client)
|
||||
if not resume_from_staging:
|
||||
# Phase 1: copy chunks into staging, then mark it complete. After the
|
||||
# marker exists, a crash will resume phase 2 from staging.
|
||||
await _populate_staging_table(client)
|
||||
await _mark_phase1_complete(client)
|
||||
|
||||
# Recreate the chunks table fresh (handles vector-dim changes).
|
||||
# Recreate the chunks table fresh (idempotent; handles vector-dim
|
||||
# changes and discards any partial new chunks from a prior crashed
|
||||
# phase 2).
|
||||
await client.store.recreate_embeddings_table()
|
||||
|
||||
staging_table = await db.open_table(_STAGING_TABLE_NAME)
|
||||
|
||||
pending_records: list = []
|
||||
pending_records: list[ChunkRecordBase] = []
|
||||
yielded_docs: set[str] = set()
|
||||
|
||||
try:
|
||||
for doc in documents:
|
||||
assert doc.id is not None
|
||||
chunks = await _read_chunks_from_staging(staging_table, doc.id)
|
||||
if not chunks:
|
||||
continue
|
||||
for doc in documents:
|
||||
assert doc.id is not None
|
||||
chunks = await _read_chunks_from_staging(staging_table, doc.id)
|
||||
if not chunks:
|
||||
continue
|
||||
|
||||
texts = contextualize(chunks)
|
||||
embeddings: list[list[float]] = []
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch_embeddings = (
|
||||
await client.chunk_repository.embedder.embed_documents(
|
||||
texts[i : i + batch_size]
|
||||
)
|
||||
texts = contextualize(chunks)
|
||||
embeddings: list[list[float]] = []
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch_embeddings = await client.chunk_repository.embedder.embed_documents(
|
||||
texts[i : i + batch_size]
|
||||
)
|
||||
embeddings.extend(batch_embeddings)
|
||||
|
||||
for chunk, content_fts, embedding in zip(chunks, texts, embeddings):
|
||||
pending_records.append(
|
||||
client.store.ChunkRecord(
|
||||
id=chunk.id,
|
||||
document_id=chunk.document_id,
|
||||
content=chunk.content,
|
||||
content_fts=content_fts,
|
||||
metadata=json.dumps(chunk.metadata),
|
||||
order=chunk.order,
|
||||
vector=embedding,
|
||||
)
|
||||
embeddings.extend(batch_embeddings)
|
||||
)
|
||||
|
||||
for chunk, content_fts, embedding in zip(chunks, texts, embeddings):
|
||||
pending_records.append(
|
||||
client.store.ChunkRecord(
|
||||
id=chunk.id,
|
||||
document_id=chunk.document_id,
|
||||
content=chunk.content,
|
||||
content_fts=content_fts,
|
||||
metadata=json.dumps(chunk.metadata),
|
||||
order=chunk.order,
|
||||
vector=embedding,
|
||||
)
|
||||
)
|
||||
yielded_docs.add(doc.id)
|
||||
# Yield per-doc for progress reporting; the actual write batches up
|
||||
# to _REBUILD_BATCH_SIZE docs. If the process is interrupted between
|
||||
# yield and the next flush, the next rebuild resumes phase 2 from
|
||||
# the staging table and redoes the batch (see _rebuild_rechunk for
|
||||
# the original comment on the yield/flush gap).
|
||||
yield doc.id
|
||||
|
||||
yielded_docs.add(doc.id)
|
||||
yield doc.id
|
||||
|
||||
if len(yielded_docs) % _REBUILD_BATCH_SIZE == 0 and pending_records:
|
||||
await client.store.chunks_table.add(pending_records)
|
||||
pending_records = []
|
||||
|
||||
if pending_records:
|
||||
if len(yielded_docs) % _REBUILD_BATCH_SIZE == 0 and pending_records:
|
||||
await client.store.chunks_table.add(pending_records)
|
||||
finally:
|
||||
if _STAGING_TABLE_NAME in (await db.list_tables()).tables:
|
||||
await db.drop_table(_STAGING_TABLE_NAME)
|
||||
pending_records = []
|
||||
|
||||
if pending_records:
|
||||
await client.store.chunks_table.add(pending_records)
|
||||
|
||||
# Phase 2 finished. Drop the recovery state — marker first so a crash
|
||||
# between the two drops leaves only staging behind, which the next
|
||||
# rebuild discards harmlessly.
|
||||
await _drop_staging_tables(client)
|
||||
|
||||
# Yield docs with no chunks
|
||||
for doc in documents:
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -1,3 +1,4 @@
|
|||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import TypedDict
|
||||
|
|
@ -133,11 +134,11 @@ async def test_rebuild_embed_only_multi_doc_streams_via_staging(
|
|||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_path):
|
||||
"""A staging table left behind by a previous interrupted rebuild gets dropped.
|
||||
"""Staging table without marker is treated as partial phase 1 and dropped.
|
||||
|
||||
Simulates a crash mid-rebuild by manually creating a ``chunks_rebuild_staging``
|
||||
table before a fresh rebuild starts. ``rebuild_database`` should detect
|
||||
and drop it before doing anything else.
|
||||
Simulates a phase-1 interruption by creating only the staging table (no
|
||||
marker). On the next rebuild ``_resolve_rebuild_recovery`` should drop
|
||||
the partial staging — the live chunks table is still authoritative.
|
||||
"""
|
||||
from haiku.rag.client.rebuild import _StagingChunkRecord
|
||||
|
||||
|
|
@ -145,11 +146,13 @@ async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_
|
|||
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
|
||||
assert doc.id is not None
|
||||
|
||||
# Simulate a stale staging table from an interrupted rebuild.
|
||||
# Simulate a partial phase 1 (staging exists, marker absent).
|
||||
await client.store.db.create_table(
|
||||
"chunks_rebuild_staging", schema=_StagingChunkRecord
|
||||
)
|
||||
assert "chunks_rebuild_staging" in (await client.store.db.list_tables()).tables
|
||||
tables = (await client.store.db.list_tables()).tables
|
||||
assert "chunks_rebuild_staging" in tables
|
||||
assert "chunks_rebuild_marker" not in tables
|
||||
|
||||
processed_ids = [
|
||||
doc_id
|
||||
|
|
@ -157,9 +160,145 @@ async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_
|
|||
]
|
||||
assert doc.id in processed_ids
|
||||
|
||||
# Stale staging gone, fresh staging from this rebuild also gone.
|
||||
tables = (await client.store.db.list_tables()).tables
|
||||
assert "chunks_rebuild_staging" not in tables
|
||||
assert "chunks_rebuild_marker" not in tables
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_resumes_phase2_from_staging_after_crash(
|
||||
qa_corpus: Dataset, temp_db_path
|
||||
):
|
||||
"""Marker + staging present → phase 2 resumes from staging instead of
|
||||
redoing phase 1.
|
||||
|
||||
Simulates a phase-2 crash: pre-populate staging with the original chunks,
|
||||
create the marker, then drop the live chunks table entirely (the worst
|
||||
case — crash right after ``recreate_embeddings_table`` succeeded but
|
||||
before any phase-2 batch flushed). The rebuild must reconstruct the
|
||||
chunks table from staging without losing data.
|
||||
"""
|
||||
from haiku.rag.client.rebuild import (
|
||||
_StagingChunkRecord,
|
||||
_StagingMarkerRecord,
|
||||
)
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
|
||||
assert doc.id is not None
|
||||
original_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
||||
assert original_chunks
|
||||
original_ids = {c.id for c in original_chunks}
|
||||
original_contents = {c.id: c.content for c in original_chunks}
|
||||
|
||||
# Snapshot chunk data into staging (simulating phase 1's output).
|
||||
staging = await client.store.db.create_table(
|
||||
"chunks_rebuild_staging", schema=_StagingChunkRecord
|
||||
)
|
||||
await staging.add(
|
||||
[
|
||||
_StagingChunkRecord(
|
||||
id=c.id or "",
|
||||
document_id=c.document_id or "",
|
||||
content=c.content,
|
||||
metadata=json.dumps(c.metadata),
|
||||
order=c.order,
|
||||
)
|
||||
for c in original_chunks
|
||||
]
|
||||
)
|
||||
|
||||
# Mark phase 1 complete (simulating the marker write that happens
|
||||
# just before phase 2 starts).
|
||||
marker = await client.store.db.create_table(
|
||||
"chunks_rebuild_marker", schema=_StagingMarkerRecord
|
||||
)
|
||||
await marker.add([_StagingMarkerRecord(id="phase1_complete")])
|
||||
|
||||
# Wipe the live chunks table to simulate a worst-case phase-2 crash
|
||||
# after recreate_embeddings_table but before any chunks were
|
||||
# written.
|
||||
await client.store.db.drop_table("chunks")
|
||||
|
||||
# Recovery: rebuild_database should detect marker+staging and have
|
||||
# _rebuild_embed_only skip phase 1.
|
||||
processed_ids = [
|
||||
doc_id
|
||||
async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY)
|
||||
]
|
||||
assert doc.id in processed_ids
|
||||
|
||||
recovered = await client.chunk_repository.get_by_document_id(doc.id)
|
||||
assert {c.id for c in recovered} == original_ids
|
||||
for chunk in recovered:
|
||||
assert chunk.content == original_contents[chunk.id]
|
||||
|
||||
tables = (await client.store.db.list_tables()).tables
|
||||
assert "chunks_rebuild_staging" not in tables
|
||||
assert "chunks_rebuild_marker" not in tables
|
||||
|
||||
|
||||
async def test_rebuild_drops_orphan_marker(temp_db_path):
|
||||
"""Marker without staging is treated as corrupted and dropped.
|
||||
|
||||
No embeddings are needed: ``_resolve_rebuild_recovery`` decides on
|
||||
tables before any embed call, and the empty database has no documents
|
||||
to embed.
|
||||
"""
|
||||
from haiku.rag.client.rebuild import _StagingMarkerRecord
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
marker = await client.store.db.create_table(
|
||||
"chunks_rebuild_marker", schema=_StagingMarkerRecord
|
||||
)
|
||||
await marker.add([_StagingMarkerRecord(id="phase1_complete")])
|
||||
|
||||
_ = [
|
||||
doc_id
|
||||
async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY)
|
||||
]
|
||||
|
||||
tables = (await client.store.db.list_tables()).tables
|
||||
assert "chunks_rebuild_marker" not in tables
|
||||
assert "chunks_rebuild_staging" not in tables
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_non_embed_mode_drops_staging_recovery_state(
|
||||
qa_corpus: Dataset, temp_db_path
|
||||
):
|
||||
"""Staging + marker from a prior embed-only crash → dropped on RECHUNK.
|
||||
|
||||
If a user runs a different rebuild mode after a crashed embed-only, the
|
||||
staging tables are stale: the new mode recreates chunks from a
|
||||
different source (e.g. the stored docling blob), so the staging copy is
|
||||
not useful.
|
||||
"""
|
||||
from haiku.rag.client.rebuild import (
|
||||
_StagingChunkRecord,
|
||||
_StagingMarkerRecord,
|
||||
)
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
|
||||
assert doc.id is not None
|
||||
|
||||
await client.store.db.create_table(
|
||||
"chunks_rebuild_staging", schema=_StagingChunkRecord
|
||||
)
|
||||
marker = await client.store.db.create_table(
|
||||
"chunks_rebuild_marker", schema=_StagingMarkerRecord
|
||||
)
|
||||
await marker.add([_StagingMarkerRecord(id="phase1_complete")])
|
||||
|
||||
processed_ids = [
|
||||
doc_id async for doc_id in client.rebuild_database(mode=RebuildMode.RECHUNK)
|
||||
]
|
||||
assert doc.id in processed_ids
|
||||
|
||||
tables = (await client.store.db.list_tables()).tables
|
||||
assert "chunks_rebuild_staging" not in tables
|
||||
assert "chunks_rebuild_marker" not in tables
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
|
|
|
|||
Loading…
Reference in a new issue