feat: add mempalace repair --mode hnsw --segment <uuid> (#1046)

[sha2fiddy]

What this does

Adds mempalace repair --mode hnsw --segment <uuid> for in-place rescue of a corrupt HNSW segment. Default --mode legacy is unchanged.

mempalace repair --mode hnsw --segment <uuid>
                 [--max-elements N] [--backup/--no-backup]
                 [--purge-queue] [--quarantine-orphans] [--dry-run]

Why

mempalace migrate on a chromadb 1.5.x palace (2026-04-19) produced a 5 TB sparse link_lists.bin and a truncated index_metadata.pickle. repair rebuild reproduced it (122 GB). Upstream classifies this "by design" (chroma-core/chroma#2594), so the workaround sits in mempalace.

How it works

1. Parse the hnswlib header from header.bin (or first 100 bytes of data_level0.bin).
2. Extract vectors directly from data_level0.bin via struct (no chromadb in the read path).
3. Reconcile labels against index_metadata.pickle; dedupe / drop zero-labels.
4. Detect space (l2 / ip / cosine) from collection_metadata; fall back to l2.
5. Build a fresh persistent index via chroma-hnswlib (is_persistent_index=True).
6. Self-query verify on a small sample. A mismatch raises RebuildVerificationError and aborts before the swap.
7. Atomic swap via os.replace. Optional timestamped backup. Rollback via os.rename.
8. Optional cleanup: purge stuck embeddings_queue rows, quarantine orphans to quarantined_orphans.json.

Dependencies

numpy and chroma-hnswlib are imported lazily. Missing deps print an install hint and abort cleanly. No pyproject.toml changes (per CONTRIBUTING.md).

Note: chroma-hnswlib, not upstream hnswlib. The persistent-mode kwargs (is_persistent_index=, persistence_location=) are fork additions. Both packages install under the same hnswlib Python name.

Tests

20 new tests in tests/test_repair.py (synthetic-segment fixture reproduces the legacy-format incident on disk), 5 in tests/test_cli.py. Full non-benchmark suite: 1092/1092. Ruff clean. repair.py coverage: 90%.

Known coverage gap

ChromaDB 1.5.7+ with rust bindings doesn't write index_metadata.pickle and leaves header.bin cur_count=0. The rescue path's input files don't exist on a 1.5.x-born palace. The synthetic fixture covers the 0.5.x / 0.6.x-origin case (which is the actual incident).

Follow-up commits (2026-04-26)

Two compat fixes, both turned up running this PR's code against a ~500k-drawer live palace:

dc9dc92 — Pickle-format compat. _reconcile_with_pickle and the total_elements_added write at the end of rebuild_hnsw_segment now handle index_metadata.pickle whether it's the chromadb 0.6.x attribute-style object shape (meta.label_to_id) or the 1.5.x dict shape (meta["label_to_id"]). Without this, repair raises AttributeError on a palace whose pickle was originally written by 0.6.x.

448c83d — Verify loosened to top-10 with wider ef. _self_query_verify previously required top-1. Two false positives on healthy indexes: byte-identical near-duplicate vectors (the same code snippet appearing across multiple transcripts) can legitimately rank #1, and hnswlib's default ef≈10 is too narrow for ~500k-element indexes with M=16. Now requires "self appears in top-10" with index.set_ef(max(200, k*4)), and clamps k to healthy_n so small palaces still pass.

Related

• Closes migrate and repair rebuild trigger unbounded link_lists.bin growth on large palaces #1046.
• fix: narrow _fix_blob_seq_ids + add repair --mode max-seq-id #1135 adds --mode max-seq-id to the same dispatcher; whichever lands second takes a small cli.py rebase.
• mempalace mine --mode convos segfaults after chromadb upgrade (stale HNSW index) #1132, mempalace migrate/status crash with SIGSEGV on chromadb version mismatch — palace unrecoverable via official tooling #1218: bug reports --mode hnsw unblocks.

Checklist

• pytest tests/ -v --ignore=tests/benchmarks passes (1092 / 1092)
• ruff check . clean
• ruff format --check . clean
• No new dependencies added to pyproject.toml
• --help text documents chroma-hnswlib install hint
• Conventional-commits messages
• Branch off develop, PR targets develop

[sha2fiddy]

Pushed two follow-up commits with details in the PR description: pickle-format compat (dc9dc92) and verify loosening (448c83d). Both turned up running this PR's code against a ~500k-drawer live palace.

[bleathem]

Hit this on a 730k-drawer palace last week (3.3.2, chromadb 1.5.8, macOS ARM). mempalace repair --yes exited 0 saying "Repair complete. 730531 drawers rebuilt" but had created a brand-new collection with an 857GB sparse link_lists.bin instead of rebuilding the existing one. The bug isn't migrate/repair specific either — once I rolled back, every subsequent add_drawer kept extending link_lists, hit 73TB apparent before I disabled the auto-save hooks.

The Rust HNSW count() segfaults on the corrupt segment so quarantine_stale_hnsw can't fire even after #1108. Recovery for me was manually quarantining the bins and accepting BM25-only — the surgical extract from data_level0.bin you have here would've saved me a lot of hand-rolling.

[sha2fiddy]

Rebased on upstream/develop @fdfaf01. How the conflicts shook out:

• --mode accepts {legacy, hnsw, max-seq-id}. Both new modes are peer subcommands of mempalace repair. Shared flags (--segment, --backup, --dry-run, --yes) get mode-specific help text; mode-only flags (--from-sidecar, --max-elements, --purge-queue, --quarantine-orphans) are wired through.
• mempalace repair-status prints a fix-it hint when hnsw_capacity_status (HNSW max_elements frozen at 16K while collection grows to 100K+ entries — MCP server segfaults on every tool call #1222) reports divergence. It surfaces both options: mempalace repair --mode hnsw --segment <uuid> (faster, no re-embed) with the diverged segment UUID inline, and mempalace repair (full rebuild). Covered by a new test in tests/test_repair.py.
• rebuild_hnsw_segment reuses HNSW max_elements frozen at 16K while collection grows to 100K+ entries — MCP server segfaults on every tool call #1222's hnsw_capacity_status probe via the existing status() entry point instead of duplicating the sqlite/pickle read. The two paths do different jobs: status() is the read-only health check, rebuild_hnsw_segment is the in-place fix.
• _vector_disabled (HNSW max_elements frozen at 16K while collection grows to 100K+ entries — MCP server segfaults on every tool call #1222) lives in the running MCP process, not on disk. It clears on the next _refresh_vector_disabled_flag() call (startup or reconnect), so after a successful rebuild the CLI now prints a hint pointing at mempalace_reconnect. No need to restart the server.
• The __main__ block in mempalace/repair.py got status and max-seq-id subparsers so the standalone module entry point matches the CLI surface.

Targeted tests 170/170 (test_repair, test_cli, test_backends, test_hnsw_capacity). Full non-benchmark suite 1480/1480. Ruff clean.

[bleathem]

Tested this branch against the 730k-drawer palace I described earlier. Two things to share — one bug in the size check, then a successful end-to-end run.

Size-check too strict

_extract_vectors in repair.py rejects the file:

ValueError: data_level0.bin is 1784301444 bytes, expected >= 1970976000

The check is len(data) >= max_elements * stride, but the read loop only walks cur_count * stride. Our data_level0.bin was sized exactly to cur_count * stride (1,064,619 × 1,676 = 1,784,301,444), not to the pre-allocated max_elements. So files captured mid-write — or wherever the corruption left things — fail to extract even though every byte the loop wants is present.

Locally patched the comparison to n * stride and everything below worked. One-line change in _extract_vectors.

Run results (with that one-line patch)

Header:  dim=384, cur_count=1,064,619, max_elements=1,176,000, size_per_element=1676

Report
  raw labels           1,064,619
  after dedup/zeros    1,063,946
  healthy (in pickle)    729,121
  orphan HNSW labels     334,825
  stale pickle ids           608
  new max_elements       947,857

Built:   729,121 vectors in 312.7s
Peak RSS ~4 GB

The 335k orphan-HNSW labels lines up with what I described before — HNSW had inflated to 1M entries during the writer corruption, well beyond the ~730k drawers that actually exist in sqlite. Quarantine wrote them out cleanly.

After rebuild:

• link_lists.bin = 5.9 MB (vs 73 TB sparse before this PR's repair)
• Collection.count() returns 712,288 — no segfault
• Direct coll.query(...) returns relevant hits with sensible L2 distances on the test palace

mempalace search CLI itself returns "Error executing plan: Internal error: Error finding id", but that's coming from the BM25/search-wrapper layer, not the HNSW segment this PR rebuilt — coll.query against the rebuilt collection works fine.

Backup at <seg>.hnsw-backup-<ts> and quarantined_orphans.json both got written as advertised.

This PR + that one size-check fix is the cleanest path off the corrupted state I've seen for this failure mode. Thanks for sticking with it.

[sha2fiddy]

Thanks @bleathem — pushed the fix as f739378 with a regression test. Targeted suite 58/58, ruff clean.

[igorls]

Thanks for this contribution, and apologies for the slow turnaround.

develop has moved a fair way since this was opened and the branch no longer merges cleanly. If you're still interested in landing it, could you rebase onto current develop? Once it merges cleanly and CI is green I'll get it reviewed for the 3.8.0 cycle.

If you'd rather not pick it back up, no problem at all — just say so and I'll close it out, and thanks either way for taking the time to send it.