Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 26 additions & 12 deletions mempalace/backends/chroma.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,13 @@ def _hnsw_payload_appears_sane(seg_dir: str) -> bool:
"hnsw:sync_threshold": 50_000,
}

# Missing index_metadata.pickle is normal only while a segment is still fresh
# or effectively empty. Once data_level0.bin has non-trivial payload, a
# missing metadata pickle means the segment was interrupted after writing HNSW
# data but before writing its metadata. Letting Chroma open that shape can
# segfault or hang in native HNSW code.
_HNSW_MISSING_METADATA_DATA_FLOOR = 1024


def _validate_where(where: Optional[dict]) -> None:
"""Scan a where-clause for unknown operators and raise ``UnsupportedFilterError``.
Expand Down Expand Up @@ -132,16 +139,13 @@ def _segment_appears_healthy(seg_dir: str) -> bool:
parsing it. ChromaDB writes that file after a successful HNSW flush;
a complete write starts with byte ``0x80`` and ends with byte
``0x2e`` (the protocol/terminator byte sequence chromadb serializes
with). If both bytes are present and the file is non-trivially sized,
chromadb will load the segment cleanly even when its on-disk mtime
trails ``chroma.sqlite3`` — which is the *steady state* under
chromadb 1.5.x's async batched flush, not corruption.
with).

A missing metadata file is treated as "fresh / never-flushed" and
considered healthy. Renaming an empty dir orphans nothing, and a
real corruption case manifests as a present-but-malformed file or a
chromadb load error caught downstream by palace-daemon's
``_auto_repair`` retry path.
Missing metadata is healthy only while the segment still looks fresh or
empty. If ``data_level0.bin`` already has non-trivial payload but
``index_metadata.pickle`` is missing, the segment is partially flushed:
Chroma wrote vector data without the metadata it needs to reopen the
HNSW reader safely.

Deliberately format-sniffs only; never deserializes. Deserialization
can execute arbitrary code, and the byte-sniff is sufficient to
Expand All @@ -152,16 +156,26 @@ def _segment_appears_healthy(seg_dir: str) -> bool:
chromadb writes today; if a future chromadb version emits protocol
0/1 segments, this check would start returning False on healthy
files and quarantine_stale_hnsw would conservatively rename them
out of the way (lazy rebuild on next open recovers).
out of the way.
"""
if not _hnsw_payload_appears_sane(seg_dir):
return False

meta_path = os.path.join(seg_dir, "index_metadata.pickle")
if not os.path.isfile(meta_path):
# No metadata file yet — segment hasn't flushed (fresh / empty).
# Renaming would orphan nothing; consider healthy.
data_path = os.path.join(seg_dir, "data_level0.bin")
try:
if (
os.path.isfile(data_path)
and os.path.getsize(data_path) > _HNSW_MISSING_METADATA_DATA_FLOOR
):
return False
except OSError:
return False

# No metadata and no meaningful vector payload yet: fresh/empty segment.
return True

try:
size = os.path.getsize(meta_path)
# A real chromadb metadata file is at least tens of bytes; a
Expand Down
54 changes: 51 additions & 3 deletions tests/test_backends.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,10 @@
from mempalace.backends.chroma import (
ChromaBackend,
ChromaCollection,
_HNSW_MISSING_METADATA_DATA_FLOOR,
_fix_blob_seq_ids,
_pin_hnsw_threads,
_segment_appears_healthy,
quarantine_invalid_hnsw_metadata,
quarantine_stale_hnsw,
)
Expand Down Expand Up @@ -685,21 +687,67 @@ def test_quarantine_stale_hnsw_leaves_healthy_segment_with_drift_alone(tmp_path)
assert seg.exists()


def test_quarantine_stale_hnsw_leaves_segment_without_metadata_alone(tmp_path):
"""Segment with no metadata file is treated as fresh / never-flushed
and not quarantined — renaming an empty dir orphans nothing."""
def test_quarantine_stale_hnsw_leaves_empty_segment_without_metadata_alone(tmp_path):
"""Missing metadata is okay only when the segment has no meaningful data yet."""

now = 1_700_000_000.0
palace, seg = _make_palace_with_segment(
tmp_path,
hnsw_mtime=now - 7200,
sqlite_mtime=now,
meta_bytes=None,
)

moved = quarantine_stale_hnsw(str(palace), stale_seconds=3600.0)

assert moved == []
assert seg.exists()


def test_segment_without_metadata_but_with_nontrivial_data_is_unhealthy(tmp_path):
"""Data without index_metadata.pickle is a partial flush, not a fresh segment."""

seg = tmp_path / "abcd-1234-5678"
seg.mkdir()
(seg / "data_level0.bin").write_bytes(b"\0" * (_HNSW_MISSING_METADATA_DATA_FLOOR + 1))

assert not _segment_appears_healthy(str(seg))


def test_segment_without_metadata_and_tiny_data_is_still_treated_as_fresh(tmp_path):
"""Tiny data payloads can occur before metadata has flushed; leave them alone."""

seg = tmp_path / "abcd-1234-5678"
seg.mkdir()
(seg / "data_level0.bin").write_bytes(b"\0" * _HNSW_MISSING_METADATA_DATA_FLOOR)

assert _segment_appears_healthy(str(seg))


def test_quarantine_stale_hnsw_renames_missing_metadata_with_nontrivial_data(tmp_path):
"""Regression for #1274: missing pickle + non-trivial data must quarantine."""

now = 1_700_000_000.0
palace, seg = _make_palace_with_segment(
tmp_path,
hnsw_mtime=now - 7200,
sqlite_mtime=now,
meta_bytes=None,
)
(seg / "data_level0.bin").write_bytes(b"\0" * (_HNSW_MISSING_METADATA_DATA_FLOOR + 1))
os.utime(seg / "data_level0.bin", (now - 7200, now - 7200))

moved = quarantine_stale_hnsw(str(palace), stale_seconds=3600.0)

assert len(moved) == 1
assert ".drift-" in moved[0]
assert not seg.exists()

drift_dirs = [p for p in palace.iterdir() if ".drift-" in p.name]
assert len(drift_dirs) == 1
assert (drift_dirs[0] / "data_level0.bin").exists()


def test_quarantine_stale_hnsw_renames_truncated_metadata(tmp_path):
"""Segment with a truncated (under-floor-size) metadata file is
quarantined — shape of a partial-flush during process kill."""
Expand Down
Loading