homelab/stack
1
0
Fork 0
Code Issues 114 Pull Requests 1 Actions Packages Projects 1 Releases Wiki Activity

fix(llm): restamp must not rebuild the HNSW index (F6)

Browse Source 
`stack llm restamp` called `ensure_hnsw`, which ALTERs the embedding
column's type and rebuilds the HNSW index — ACCESS EXCLUSIVE for the
life of a rebuild over 1M+ rows, locking out every reader for a
metadata-only backfill that never touches embeddings.

Split `_HNSW_DDL` into the column/HNSW part (unchanged `ensure_hnsw`,
still called by `stack llm index` after a fresh run) and a new
`ensure_metadata_indexes(engine)` holding only the three GIN indexes
on codes/families/elements — no ACCESS EXCLUSIVE lock, safe on a live
table. Restamp now calls only the latter, and skips it entirely under
--dry-run.
This commit is contained in:
kert
2026-09-09 21:03:56 -04:00
parent 1c1a5fb39b
commit b1e6d6ef3a
4 changed files with 127 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/cli/llm.py
View File

@@ -114,12 +114,17 @@ def restamp(
stored document text and merges it into `cmetadata` with a jsonb `||`,
never touching the embedding — no GPU time, no re-chunking. Safe to
re-run: rows already stamped with the current values are skipped.
F6: never rebuilds the HNSW index (`ensure_hnsw` ALTERs the vector
column and takes an ACCESS EXCLUSIVE lock for the life of the
rebuild) — only the cheap metadata GIN indexes, and not even those
under `--dry-run`.
"""
import os
from llm import config as llm_config
from llm.index import _engine
from llm.migrate import ensure_hnsw, migrate
from llm.migrate import ensure_metadata_indexes, migrate
from llm.restamp import restamp as restamp_collection
from pfs.anchors import code_family_index
from pfs.families import FAMILIES, refresh_from
@@ -130,7 +135,8 @@ def restamp(
cfg = llm_config.load()
engine = _engine(cfg)
migrate(engine)
ensure_hnsw(engine, cfg.embed_dim)
if not dry_run:
ensure_metadata_indexes(engine)
if os.path.exists(cfg.duckdb_replica):
import duckdb

38
src/llm/migrate.py
View File

@@ -44,18 +44,29 @@ CREATE TABLE IF NOT EXISTS index_docket_state (
# langchain-postgres creates langchain_pg_embedding with an untyped vector
# column; HNSW needs a typed one. ALTER is a no-op when already typed.
# The ALTER and the HNSW build both take an ACCESS EXCLUSIVE lock on the
# whole table — fine right after a fresh index run's `add_embeddings`,
# but never run this from `stack llm restamp` (F6): a metadata-only
# backfill has no reason to lock out every reader for the minutes an
# HNSW rebuild over 1M+ rows takes.
_HNSW_DDL = [
"ALTER TABLE langchain_pg_embedding ALTER COLUMN embedding TYPE vector({dim})",
"CREATE INDEX IF NOT EXISTS ix_embedding_hnsw "
"ON langchain_pg_embedding USING hnsw (embedding vector_cosine_ops)",
# Code-cited sources (llm.evidence) match a code against the space-
# separated cmetadata->>'codes' inside one collection; without these
# two the query is a sequential scan of every embedding.
"CREATE INDEX IF NOT EXISTS ix_embedding_collection "
"ON langchain_pg_embedding (collection_id)",
]
# Code-cited sources (llm.evidence) match a code against the space-
# separated cmetadata->>'codes' inside one collection; without these the
# query is a sequential scan of every embedding. Plain CREATE INDEX (no
# ACCESS EXCLUSIVE, unlike ensure_hnsw's vector-column retype) — safe to
# run against a live table, so `stack llm restamp` calls only this, never
# ensure_hnsw.
_METADATA_INDEX_DDL = [
"CREATE INDEX IF NOT EXISTS ix_embedding_codes "
"ON langchain_pg_embedding USING gin "
"(string_to_array(coalesce(cmetadata->>'codes',''), ' '))",
"CREATE INDEX IF NOT EXISTS ix_embedding_collection "
"ON langchain_pg_embedding (collection_id)",
"CREATE INDEX IF NOT EXISTS ix_embedding_families "
"ON langchain_pg_embedding USING gin "
"(string_to_array(coalesce(cmetadata->>'families',''), ' '))",
@@ -77,8 +88,23 @@ def ensure_hnsw(engine: Engine, dim: int) -> None:
"""Type the embedding column and build the HNSW index.
Call after the PGVector store has created its tables (first
``add_embeddings``), not before.
``add_embeddings``), not before. ACCESS EXCLUSIVE for the life of
the rebuild — never call this from a restamp run (F6); see
``ensure_metadata_indexes``.
"""
with engine.begin() as conn:
for ddl in _HNSW_DDL:
conn.execute(text(ddl.format(dim=dim)))
def ensure_metadata_indexes(engine: Engine) -> None:
"""GIN indexes on the codes/families/elements metadata only.
No ACCESS EXCLUSIVE lock and no vector-column touch — safe to run
against a live, queryable table. ``stack llm restamp`` calls this
instead of ``ensure_hnsw`` (F6): a metadata-only backfill has no
reason to rebuild the HNSW index.
"""
with engine.begin() as conn:
for ddl in _METADATA_INDEX_DDL:
conn.execute(text(ddl))

66
tests/cli/test_llm_exercise.py
View File

@@ -232,6 +232,72 @@ class TestIndexLimit:
assert mock_index_refs.call_args.kwargs["mark_complete"] is False
class TestRestamp:
"""F6: `stack llm restamp` must never rebuild the HNSW index — only
the cheap metadata GIN indexes, and not even those under --dry-run."""
@patch("os.path.exists", return_value=False)
@patch("pfs.anchors.code_family_index", return_value={})
@patch("llm.restamp.restamp")
@patch("llm.migrate.ensure_metadata_indexes")
@patch("llm.migrate.migrate")
@patch("llm.index._engine")
@patch("llm.config.load")
def test_dry_run_issues_no_ddl(
self,
mock_load,
mock_engine,
mock_migrate,
mock_ensure_metadata,
mock_restamp,
mock_code_index,
mock_exists,
):
mock_load.return_value = MagicMock(duckdb_replica="nope.duckdb")
mock_restamp.return_value = {
"collection": "rules",
"scanned": 0,
"updated": 0,
"seconds": 0.0,
}
result = runner.invoke(app, ["restamp", "--collection", "rules", "--dry-run"])
assert result.exit_code == 0, result.output
mock_ensure_metadata.assert_not_called()
assert mock_restamp.call_args.kwargs["dry_run"] is True
@patch("os.path.exists", return_value=False)
@patch("pfs.anchors.code_family_index", return_value={})
@patch("llm.restamp.restamp")
@patch("llm.migrate.ensure_metadata_indexes")
@patch("llm.migrate.migrate")
@patch("llm.index._engine")
@patch("llm.config.load")
def test_real_run_calls_ensure_metadata_indexes_not_hnsw(
self,
mock_load,
mock_engine,
mock_migrate,
mock_ensure_metadata,
mock_restamp,
mock_code_index,
mock_exists,
):
mock_load.return_value = MagicMock(duckdb_replica="nope.duckdb")
mock_restamp.return_value = {
"collection": "rules",
"scanned": 0,
"updated": 0,
"seconds": 0.0,
}
result = runner.invoke(app, ["restamp", "--collection", "rules"])
assert result.exit_code == 0, result.output
mock_ensure_metadata.assert_called_once()
class TestServe:
@patch("uvicorn.run")
def test_serve_starts_uvicorn(self, mock_run):

26
tests/llm/test_migrate.py
View File

@@ -24,22 +24,38 @@ class TestDdl:
executed = " ".join(str(call.args[0]) for call in conn.execute.call_args_list)
assert "vector(768)" in executed
assert "USING hnsw" in executed
assert (
"ix_embedding_collection" in executed
and "langchain_pg_embedding (collection_id)" in executed
)
# F6: the metadata GIN indexes are ensure_metadata_indexes' job now,
# not ensure_hnsw's — this must never take the ACCESS EXCLUSIVE lock.
assert "ix_embedding_codes" not in executed
assert "ix_embedding_families" not in executed
assert "ix_embedding_elements" not in executed
def test_ensure_metadata_indexes_issues_exactly_the_three_gin_statements(self):
engine = MagicMock()
conn = engine.begin.return_value.__enter__.return_value
migrate.ensure_metadata_indexes(engine)
calls = [str(call.args[0]) for call in conn.execute.call_args_list]
assert len(calls) == 3
executed = " ".join(calls)
# code-cited lookups scan cmetadata->>'codes' inside one collection
assert (
"ix_embedding_codes" in executed
and "USING gin (string_to_array" in executed
)
assert (
"ix_embedding_collection" in executed
and "langchain_pg_embedding (collection_id)" in executed
)
# family and element indexes follow the same pattern
assert (
"ix_embedding_families" in executed and "cmetadata->>'families'" in executed
)
assert (
"ix_embedding_elements" in executed and "cmetadata->>'elements'" in executed
)
# No vector-column ALTER, no HNSW build — this must be safe on a
# live table.
assert "ALTER TABLE" not in executed
assert "hnsw" not in executed.lower()
def test_fingerprint_column_and_docket_state(self):
assert "fingerprint" in migrate.INDEX_STATE_DDL