mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 14:17:08 +00:00
fix: stop the re-embed migration from destroying the index it rebuilds
Five defects from a review of the embeddings work, four of them silent. - Write local files atomically. `LocalStorage.save_file` streamed straight onto the destination, so an interrupted write left a truncated file. `reembed` rewrites every index it touches, and a half-written `index.faiss` loads at neither the old width nor the new one -- the source was unrecoverable, with no backup and no temp file left behind. Bytes now land beside the destination and move into place with `os.replace`. S3 was already safe (single PUT). - Read pgvector chunks a page at a time. `reembed_pgvector` materialised every `(id, text)` row for a source before embedding -- ~1.6 GB at 200k chunks and several times that for non-Latin scripts, with the `PGresult` held alongside until the cursor closed. Inside the shipped 4Gi limit, while also holding the model, that is an OOMKill -- which is exactly the SIGKILL the point above turned into a destroyed index. It now walks the source by keyset. - Bound the first wave of delegated embeds. The failure cooldown is only latched once the first `get()` returns, so every request already in flight paid the full EMBEDDINGS_DELEGATE_TIMEOUT: measured 64 threads all timing out together, and at the shipped 60s across a 96-thread WSGI pool that is an API serving nothing at all, health checks included. One caller now probes while the rest fail fast; after a single success the gate leaves the path entirely. - Ship EMBEDDINGS_NAME commented in .env-template. The comment directly above it says to leave it commented when upgrading, and the line shipped set. Any value reaching `.env` lands in `model_fields_set`, which makes `resolve_embeddings_pin` bail -- so a template-derived `.env` disabled the legacy pin outright and repointed a populated index at a different 768-dim model, where no width check fires. The pin already picks granite for a fresh install and mpnet for an existing one, so nothing needs to be set by hand. - Stamp `sources.model` on wiki sources. They were created with the column NULL and then embedded like any other source, and the boot check reads NULL as "pre-dates the column, therefore the legacy model" -- reporting a correctly embedded source as stale on every startup of every process. Stamped at creation, and again on each page re-embed so existing rows heal. The two docs that promised the FAISS index survives a failed run said so of the embed only; both now describe the write, and upgrading.mdx says to stop ingest for the duration.
This commit is contained in:
1 parent
00be2c05ad
commit
12dd7c7eab
14 files changed
+583
-135
No files matched your search
@@ -7,6 +7,46 @@ import pytest
|
||||
from application.scripts import reembed
|
||||
|
||||
|
||||
def paginating_cursor(chunk_rows, *, graph_rows=(), graph_table=("graph_nodes",)):
|
||||
"""A cursor answering the reads ``reembed_pgvector`` issues, from memory.
|
||||
|
||||
The chunk read is paginated by keyset, so a cursor that returns the same
|
||||
page for every ``fetchall`` never terminates. Returns the cursor and the
|
||||
mutable table backing it.
|
||||
|
||||
Args:
|
||||
chunk_rows: ``(id, text)`` rows for the source's chunk table.
|
||||
graph_rows: ``(id, name)`` rows for ``graph_nodes``.
|
||||
graph_table: What ``to_regclass`` reports; ``(None,)`` for absent.
|
||||
|
||||
Returns:
|
||||
``(cursor, tables)``, where ``tables["chunks"]`` and
|
||||
``tables["graph"]`` can be reassigned to change what is served.
|
||||
"""
|
||||
tables = {"chunks": list(chunk_rows), "graph": list(graph_rows)}
|
||||
pending = {"rows": []}
|
||||
cursor = MagicMock()
|
||||
|
||||
def execute(query, params=None):
|
||||
text = str(query)
|
||||
if "to_regclass" in text:
|
||||
pending["rows"] = [graph_table]
|
||||
elif "FROM graph_nodes" in text:
|
||||
pending["rows"] = list(tables["graph"])
|
||||
elif "count(*)" in text:
|
||||
pending["rows"] = [(len(tables["chunks"]),)]
|
||||
elif "id > %s" in text:
|
||||
_, after_id, limit = params
|
||||
pending["rows"] = [r for r in tables["chunks"] if r[0] > after_id][:limit]
|
||||
else:
|
||||
pending["rows"] = list(tables["chunks"])[: params[1]]
|
||||
|
||||
cursor.execute.side_effect = execute
|
||||
cursor.fetchall.side_effect = lambda: pending["rows"]
|
||||
cursor.fetchone.side_effect = lambda: pending["rows"][0] if pending["rows"] else None
|
||||
return cursor, tables
|
||||
|
||||
|
||||
class TestCLI:
|
||||
def test_defaults(self):
|
||||
args = reembed.build_parser().parse_args([])
|
||||
@@ -172,11 +212,18 @@ class TestPgvectorWithoutTheExtension:
|
||||
|
||||
@pytest.fixture
|
||||
def store(self):
|
||||
"""A store whose cursor serves the paginated reads, not a fixed page.
|
||||
|
||||
``reembed_pgvector`` walks the source by keyset, so a cursor that
|
||||
returns the same rows for every ``fetchall`` never terminates. The fake
|
||||
answers each of the three queries the function issues from one in-memory
|
||||
table, which tests mutate through ``rows``.
|
||||
"""
|
||||
cursor, table = paginating_cursor([(1, "alpha"), (2, "beta"), (3, "gamma")])
|
||||
|
||||
store = MagicMock()
|
||||
store._table_name = "documents"
|
||||
store._vector_column = "embedding"
|
||||
cursor = MagicMock()
|
||||
cursor.fetchall.return_value = [(1, "alpha"), (2, "beta"), (3, "gamma")]
|
||||
conn = MagicMock()
|
||||
conn.cursor.return_value = cursor
|
||||
store._get_connection.return_value = conn
|
||||
@@ -186,60 +233,85 @@ class TestPgvectorWithoutTheExtension:
|
||||
with patch.object(
|
||||
reembed.VectorCreator, "create_vectorstore", return_value=store
|
||||
), patch.object(reembed.settings, "GRAPHRAG_ENABLED", False):
|
||||
yield store, conn, cursor
|
||||
yield store, conn, cursor, table
|
||||
|
||||
def test_reads_and_rewrites_every_chunk(self, store):
|
||||
_, conn, cursor = store
|
||||
_, conn, cursor, _ = store
|
||||
seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
|
||||
assert (seen, written) == (3, 3)
|
||||
cursor.executemany.assert_called_once()
|
||||
conn.commit.assert_called()
|
||||
|
||||
def test_dry_run_neither_embeds_nor_writes(self, store):
|
||||
fake_store, conn, cursor = store
|
||||
def test_pages_are_bounded_by_batch_size(self, store):
|
||||
"""The read is paginated, so a huge source never lands in memory at once."""
|
||||
_, _, cursor, table = store
|
||||
table["chunks"] = [(i, f"chunk-{i}") for i in range(1, 11)]
|
||||
seen, written = reembed.reembed_pgvector("s1", batch_size=3, dry_run=False)
|
||||
assert (seen, written) == (10, 10)
|
||||
selects = [
|
||||
call for call in cursor.execute.call_args_list
|
||||
if "SELECT id, text" in str(call.args[0])
|
||||
]
|
||||
# Four pages of at most 3, then one empty page to end the walk.
|
||||
assert len(selects) == 5
|
||||
assert all(call.args[1][-1] == 3 for call in selects)
|
||||
|
||||
def test_dry_run_counts_without_reading_the_text(self, store):
|
||||
fake_store, _, cursor, _ = store
|
||||
seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=True)
|
||||
assert (seen, written) == (3, 0)
|
||||
fake_store._embedding.embed_documents.assert_not_called()
|
||||
cursor.executemany.assert_not_called()
|
||||
assert any(
|
||||
"count(*)" in str(call.args[0]) for call in cursor.execute.call_args_list
|
||||
)
|
||||
|
||||
def test_batches_commit_separately(self, store):
|
||||
_, conn, cursor = store
|
||||
_, conn, cursor, _ = store
|
||||
reembed.reembed_pgvector("s1", batch_size=2, dry_run=False)
|
||||
# 3 rows at batch 2 is two write transactions.
|
||||
assert cursor.executemany.call_count == 2
|
||||
assert conn.commit.call_count == 2
|
||||
|
||||
def test_failed_batch_rolls_back_and_raises(self, store):
|
||||
fake_store, conn, cursor = store
|
||||
_, conn, cursor, _ = store
|
||||
cursor.executemany.side_effect = RuntimeError("write failed")
|
||||
with pytest.raises(RuntimeError):
|
||||
reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
|
||||
conn.rollback.assert_called_once()
|
||||
|
||||
def test_connection_is_returned_even_on_failure(self, store):
|
||||
fake_store, _, cursor = store
|
||||
fake_store, _, cursor, _ = store
|
||||
cursor.executemany.side_effect = RuntimeError("boom")
|
||||
with pytest.raises(RuntimeError):
|
||||
reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
|
||||
fake_store.close.assert_called_once()
|
||||
|
||||
def test_null_text_does_not_crash_the_embed_call(self, store):
|
||||
fake_store, _, cursor = store
|
||||
cursor.fetchall.return_value = [(1, None), (2, "beta")]
|
||||
fake_store, _, _, table = store
|
||||
table["chunks"] = [(1, None), (2, "beta")]
|
||||
seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
|
||||
assert (seen, written) == (2, 2)
|
||||
assert fake_store._embedding.embed_documents.call_args.args[0] == ["", "beta"]
|
||||
|
||||
def test_empty_source_is_a_no_op(self, store):
|
||||
fake_store, _, cursor = store
|
||||
cursor.fetchall.return_value = []
|
||||
fake_store, _, _, table = store
|
||||
table["chunks"] = []
|
||||
assert reembed.reembed_pgvector("s1", batch_size=64, dry_run=False) == (0, 0)
|
||||
fake_store._embedding.embed_documents.assert_not_called()
|
||||
|
||||
def test_source_discovery_returns_sorted_ids(self, store):
|
||||
_, _, cursor = store
|
||||
def test_source_discovery_returns_sorted_ids(self):
|
||||
store = MagicMock()
|
||||
store._table_name = "documents"
|
||||
cursor = MagicMock()
|
||||
cursor.fetchall.return_value = [("b",), ("a",)]
|
||||
assert reembed.list_source_ids("pgvector") == ["b", "a"]
|
||||
conn = MagicMock()
|
||||
conn.cursor.return_value = cursor
|
||||
store._get_connection.return_value = conn
|
||||
with patch.object(
|
||||
reembed.VectorCreator, "create_vectorstore", return_value=store
|
||||
):
|
||||
assert reembed.list_source_ids("pgvector") == ["b", "a"]
|
||||
|
||||
|
||||
class TestFaissSourceDiscovery:
|
||||
@@ -353,12 +425,13 @@ class TestGraphNodeReembedding:
|
||||
|
||||
@staticmethod
|
||||
def _pgvector_mocks():
|
||||
cursor, _ = paginating_cursor(
|
||||
[(1, "alpha"), (2, "beta")],
|
||||
graph_rows=[("n1", "Alpha"), ("n2", "Beta")],
|
||||
)
|
||||
store = MagicMock()
|
||||
store._table_name = "documents"
|
||||
store._vector_column = "embedding"
|
||||
cursor = MagicMock()
|
||||
cursor.fetchone.return_value = ("graph_nodes",)
|
||||
cursor.fetchall.return_value = [(1, "alpha"), (2, "beta")]
|
||||
conn = MagicMock()
|
||||
conn.cursor.return_value = cursor
|
||||
store._get_connection.return_value = conn
|
||||
|
||||
Reference in new issue
Block a user