xiaoxue_admin
/
mempalace
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(chroma): route stale hnsw divergence to sqlite fallback

This commit is contained in:
fatkobra 2026-06-18 10:37:58 +00:00
parent 9f434e0bfd
commit d704433701
4 changed files with 176 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

92
mempalace/backends/chroma.py
View File

@ -9,6 +9,7 @@ import os
import pickle
import re
import sqlite3
import time
from collections import defaultdict
from numbers import Integral
from pathlib import Path
@ -606,6 +607,7 @@ def _hnsw_element_count(palace_path: str, segment_id: str) -> Optional[int]:
# sync_threshold) from expected steady-state lag.
_HNSW_DIVERGENCE_FALLBACK_FLOOR = 2000
_HNSW_DIVERGENCE_FRACTION = 0.10
_HNSW_PERSISTENT_DIVERGENCE_GRACE_SECONDS = 300.0
def _read_sync_threshold(palace_path: str, collection_name: str) -> int:
@ -649,6 +651,45 @@ def _read_sync_threshold(palace_path: str, collection_name: str) -> int:
return 1000
def _collection_has_sync_threshold_metadata(palace_path: str, collection_name: str) -> bool:
"""Return True when the collection explicitly stores hnsw:sync_threshold."""
db_path = os.path.join(palace_path, "chroma.sqlite3")
if not os.path.isfile(db_path):
return False
try:
conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
try:
row = conn.execute(
"""
SELECT 1
FROM collection_metadata cm
JOIN collections c ON cm.collection_id = c.id
WHERE c.name = ?
AND cm.key = 'hnsw:sync_threshold'
LIMIT 1
""",
(collection_name,),
).fetchone()
return row is not None
finally:
conn.close()
except Exception:
logger.debug("_collection_has_sync_threshold_metadata failed", exc_info=True)
return False
def _hnsw_metadata_age_seconds(palace_path: str, segment_id: str) -> Optional[float]:
"""Return index_metadata.pickle age in seconds, or None when unreadable."""
pickle_path = os.path.join(palace_path, segment_id, "index_metadata.pickle")
try:
return max(0.0, time.time() - os.path.getmtime(pickle_path))
except OSError:
return None
def hnsw_capacity_status(palace_path: str, collection_name: str = "mempalace_drawers") -> dict:
"""Compare sqlite embedding count against HNSW element count.
@ -693,11 +734,15 @@ def hnsw_capacity_status(palace_path: str, collection_name: str = "mempalace_dra
hnsw_count = _hnsw_element_count(palace_path, seg_id)
out["hnsw_count"] = hnsw_count
sync_threshold = _read_sync_threshold(palace_path, collection_name)
# Two synchronization windows worth — see comment above
# _HNSW_DIVERGENCE_FALLBACK_FLOOR for the rationale.
divergence_floor = max(_HNSW_DIVERGENCE_FALLBACK_FLOOR, 2 * sync_threshold)
has_explicit_sync_threshold = _collection_has_sync_threshold_metadata(
palace_path,
collection_name,
)
metadata_age_seconds = (
_hnsw_metadata_age_seconds(palace_path, seg_id) if hnsw_count is not None else None
)
out["hnsw_metadata_age_seconds"] = metadata_age_seconds
if hnsw_count is None:
# No pickle yet, so this probe cannot measure HNSW capacity.
@ -715,21 +760,52 @@ def hnsw_capacity_status(palace_path: str, collection_name: str = "mempalace_dra
divergence = sqlite_count - hnsw_count
out["divergence"] = divergence
threshold = max(divergence_floor, int(sqlite_count * _HNSW_DIVERGENCE_FRACTION))
if divergence > threshold:
# Newer palaces explicitly store mempalace's low sync threshold
# (currently 2), so a gap of dozens of rows is far beyond ordinary
# flush lag. Older palaces may lack the metadata row; keep the
# historical floor for fresh lag there, but do not let a stale pickle
# sit below the floor forever (#1816).
if has_explicit_sync_threshold:
threshold = max(0, 2 * sync_threshold)
else:
divergence_floor = max(_HNSW_DIVERGENCE_FALLBACK_FLOOR, 2 * sync_threshold)
threshold = max(
divergence_floor,
int(sqlite_count * _HNSW_DIVERGENCE_FRACTION),
)
out["threshold"] = threshold
stale_below_threshold = (
divergence > 0
and metadata_age_seconds is not None
and metadata_age_seconds >= _HNSW_PERSISTENT_DIVERGENCE_GRACE_SECONDS
)
if divergence > threshold or stale_below_threshold:
out["status"] = "diverged"
out["diverged"] = True
pct = 100.0 * divergence / max(sqlite_count, 1)
if divergence > threshold:
reason = f"exceeds threshold {threshold:,}"
else:
age = metadata_age_seconds or 0.0
reason = f"persisted below the old flush-lag floor for {age:.0f}s"
out["message"] = (
f"HNSW index holds {hnsw_count:,} elements but sqlite has "
f"{sqlite_count:,} embeddings — {divergence:,} drawers ({pct:.0f}%) "
"are invisible to vector search. Run `mempalace repair` to rebuild."
f"{sqlite_count:,} embeddings - {divergence:,} drawers "
f"({pct:.0f}%) are missing from the flushed HNSW index "
f"({reason}). Vector reads are disabled until "
"`mempalace repair` rebuilds it."
)
else:
out["status"] = "ok"
out["message"] = (
f"HNSW {hnsw_count:,} / sqlite {sqlite_count:,} (within flush-lag tolerance)"
)
if divergence < 0:
out["message"] += " (HNSW has extra flushed elements; treating as safe)"
except Exception:
logger.debug("hnsw_capacity_status failed", exc_info=True)
out["message"] = "HNSW capacity probe raised; skipping"

32
mempalace/repair.py
View File

@ -1059,6 +1059,37 @@ def extract_via_sqlite(palace_path: str, collection_name: str) -> Iterator[tuple
conn.close()
def _preserve_knowledge_graph_sqlite(source_palace: str, dest_palace: str) -> list[str]:
"""Copy KG SQLite sidecars when rebuilding a palace from chroma.sqlite3.
rebuild_from_sqlite reconstructs Chroma collections into a fresh
destination directory. The knowledge graph is a separate SQLite database,
so it must be copied explicitly or the repair succeeds while silently
dropping KG state (#1816).
"""
copied: list[str] = []
for suffix in ("", "-wal", "-shm"):
filename = f"knowledge_graph.sqlite3{suffix}"
src = os.path.join(source_palace, filename)
dst = os.path.join(dest_palace, filename)
if not os.path.isfile(src):
continue
if os.path.abspath(src) == os.path.abspath(dst):
continue
os.makedirs(dest_palace, exist_ok=True)
shutil.copy2(src, dst)
copied.append(filename)
if copied:
print(" Preserved knowledge graph: " + ", ".join(copied))
return copied
def rebuild_from_sqlite(
source_palace: str,
dest_palace: str,
@ -1205,6 +1236,7 @@ def rebuild_from_sqlite(
)
os.makedirs(dest_palace, exist_ok=True)
_preserve_knowledge_graph_sqlite(source_palace, dest_palace)
# Backend lifetime is wrapped in try/finally so the dest palace's
# PersistentClient handle (opened lazily inside ``create_collection``

37
tests/test_hnsw_capacity.py
View File

@ -11,6 +11,7 @@ from __future__ import annotations
import os
import pickle
import sqlite3
import time
import pytest
@ -640,3 +641,39 @@ def test_tool_status_via_sqlite_returns_breakdown(palace_with_drawers, monkeypat
# ops×2 (incident + repair runbook), design×1 (metaphor).
assert out["wings"].get("ops") == 2
assert out["wings"].get("design") == 1
def test_capacity_status_flags_small_gap_with_explicit_low_sync_threshold(tmp_path):
"""New palaces use a low explicit sync threshold, so 57 missing rows is unsafe."""
seg = "seg-1816-explicit-low-sync"
_seed_chroma_db(str(tmp_path), sqlite_count=1768, segment_id=seg, sync_threshold=2)
_write_pickle(str(tmp_path), seg, hnsw_count=1711)
info = hnsw_capacity_status(str(tmp_path), COLLECTION)
assert info["divergence"] == 57
assert info["threshold"] == 4
assert info["status"] == "diverged"
assert info["diverged"] is True
assert "repair" in info["message"].lower()
def test_capacity_status_flags_stale_below_floor_divergence(tmp_path):
"""A persistent below-floor sqlite>HNSW gap must not be treated as fresh lag."""
from mempalace.backends import chroma
seg = "seg-1816-stale-below-floor"
_seed_chroma_db(str(tmp_path), sqlite_count=1768, segment_id=seg)
_write_pickle(str(tmp_path), seg, hnsw_count=1711)
pickle_path = tmp_path / seg / "index_metadata.pickle"
old = time.time() - chroma._HNSW_PERSISTENT_DIVERGENCE_GRACE_SECONDS - 10
os.utime(pickle_path, (old, old))
info = hnsw_capacity_status(str(tmp_path), COLLECTION)
assert info["divergence"] == 57
assert info["threshold"] >= 2000
assert info["status"] == "diverged"
assert info["diverged"] is True
assert "persisted below" in info["message"]

23
tests/test_repair.py
View File

@ -1995,3 +1995,26 @@ def test_rebuild_index_calls_vacuum(mock_backend_cls, mock_shutil, tmp_path):
args, kwargs = mock_vacuum.call_args
assert args[0] == str(tmp_path)
assert "progress" in kwargs
def test_rebuild_from_sqlite_preserves_knowledge_graph_sidecar(tmp_path):
"""The from-sqlite repair path must not drop the KG SQLite sidecar."""
src = tmp_path / "source"
dest = tmp_path / "dest"
src.mkdir()
dest.mkdir()
(src / "knowledge_graph.sqlite3").write_text("kg-db", encoding="utf-8")
(src / "knowledge_graph.sqlite3-wal").write_text("kg-wal", encoding="utf-8")
(src / "knowledge_graph.sqlite3-shm").write_text("kg-shm", encoding="utf-8")
copied = repair._preserve_knowledge_graph_sqlite(str(src), str(dest))
assert copied == [
"knowledge_graph.sqlite3",
"knowledge_graph.sqlite3-wal",
"knowledge_graph.sqlite3-shm",
]
assert (dest / "knowledge_graph.sqlite3").read_text(encoding="utf-8") == "kg-db"
assert (dest / "knowledge_graph.sqlite3-wal").read_text(encoding="utf-8") == "kg-wal"
assert (dest / "knowledge_graph.sqlite3-shm").read_text(encoding="utf-8") == "kg-shm"