diff --git a/tests/unit/VectorDBStrategy/test_persistent_hnsw_lib.py b/tests/unit/VectorDBStrategy/test_persistent_hnsw_lib.py index 05555ed..9aad334 100644 --- a/tests/unit/VectorDBStrategy/test_persistent_hnsw_lib.py +++ b/tests/unit/VectorDBStrategy/test_persistent_hnsw_lib.py @@ -1,3 +1,4 @@ +import json import os import tempfile import unittest @@ -41,22 +42,66 @@ def test_remove(self): assert len(knn) == 1 assert knn[0][1] == id2 - def test_persists_files_to_disk(self): - vector_db = PersistentHNSWLibVectorDB(persist_path=self.persist_path) + def test_mutations_below_checkpoint_interval_only_touch_wal(self): + """Below the checkpoint interval, mutations should be captured by + cheap WAL appends rather than a full index rewrite, so no manifest + or index snapshot should exist yet.""" + vector_db = PersistentHNSWLibVectorDB( + persist_path=self.persist_path, checkpoint_interval=100 + ) vector_db.add(embedding=[0.1, 0.2, 0.3]) - assert os.path.exists(self.persist_path) - assert os.path.exists(self.persist_path + ".meta.json") + assert os.path.exists(self.persist_path + ".wal") + assert os.path.getsize(self.persist_path + ".wal") > 0 + assert not os.path.exists(self.persist_path + ".manifest.json") - def test_survives_simulated_restart(self): - """Embeddings added by one instance should be visible to a fresh - instance pointed at the same path, simulating a process restart.""" - first_instance = PersistentHNSWLibVectorDB(persist_path=self.persist_path) + def test_checkpoint_publishes_manifest_and_clears_wal(self): + vector_db = PersistentHNSWLibVectorDB( + persist_path=self.persist_path, checkpoint_interval=100 + ) + vector_db.add(embedding=[0.1, 0.2, 0.3]) + vector_db.checkpoint() + + assert os.path.exists(self.persist_path + ".manifest.json") + assert os.path.exists(self.persist_path + ".g1") + assert os.path.getsize(self.persist_path + ".wal") == 0 + # No leftover temp files from the atomic rename. + assert not os.path.exists(self.persist_path + ".g1.tmp") + assert not os.path.exists(self.persist_path + ".manifest.json.tmp") + + def test_auto_checkpoints_after_interval_and_prunes_old_generation(self): + vector_db = PersistentHNSWLibVectorDB( + persist_path=self.persist_path, checkpoint_interval=2 + ) + vector_db.add(embedding=[0.1, 0.2, 0.3]) + vector_db.add(embedding=[0.2, 0.3, 0.4]) + + # Auto-checkpoint should have fired after the 2nd mutation. + assert os.path.exists(self.persist_path + ".manifest.json") + assert os.path.exists(self.persist_path + ".g1") + assert os.path.getsize(self.persist_path + ".wal") == 0 + + vector_db.add(embedding=[0.3, 0.4, 0.5]) + vector_db.add(embedding=[0.4, 0.5, 0.6]) + + # A second checkpoint should have fired and the stale generation + # file from the first checkpoint should have been pruned. + assert os.path.exists(self.persist_path + ".g2") + assert not os.path.exists(self.persist_path + ".g1") + + def test_survives_simulated_restart_before_any_checkpoint(self): + """Mutations that never reached the checkpoint interval should still + survive a restart, recovered by replaying the WAL.""" + first_instance = PersistentHNSWLibVectorDB( + persist_path=self.persist_path, checkpoint_interval=100 + ) first_instance.add(embedding=[0.1, 0.2, 0.3]) first_instance.add(embedding=[0.2, 0.3, 0.4]) del first_instance - second_instance = PersistentHNSWLibVectorDB(persist_path=self.persist_path) + second_instance = PersistentHNSWLibVectorDB( + persist_path=self.persist_path, checkpoint_interval=100 + ) knn = second_instance.get_knn(embedding=[0.1, 0.2, 0.3], k=2) assert len(knn) == 2 @@ -64,6 +109,57 @@ def test_survives_simulated_restart(self): new_id = second_instance.add(embedding=[0.3, 0.4, 0.5]) assert new_id == 2 + def test_survives_simulated_restart_after_checkpoint(self): + first_instance = PersistentHNSWLibVectorDB( + persist_path=self.persist_path, checkpoint_interval=1 + ) + first_instance.add(embedding=[0.1, 0.2, 0.3]) + first_instance.add(embedding=[0.2, 0.3, 0.4]) + del first_instance + + second_instance = PersistentHNSWLibVectorDB( + persist_path=self.persist_path, checkpoint_interval=1 + ) + knn = second_instance.get_knn(embedding=[0.1, 0.2, 0.3], k=2) + assert len(knn) == 2 + + def test_recovers_from_crash_between_checkpoint_index_write_and_manifest(self): + """If the process dies after the new generation's index file is + written but before the manifest is rewritten to point at it, the + on-disk manifest still refers to the previous, fully consistent + generation. Recovery must fall back to it (plus WAL replay) rather + than surface the orphaned, half-committed generation.""" + vector_db = PersistentHNSWLibVectorDB( + persist_path=self.persist_path, checkpoint_interval=100 + ) + vector_db.add(embedding=[0.1, 0.2, 0.3]) + vector_db.checkpoint() # generation 1 fully committed + + vector_db.add(embedding=[0.2, 0.3, 0.4]) # queued in the WAL + + def failing_write_manifest(): + raise RuntimeError("simulated crash before manifest commit") + + vector_db._write_manifest = failing_write_manifest + with self.assertRaises(RuntimeError): + vector_db.checkpoint() + + # The new generation's index snapshot made it to disk... + assert os.path.exists(self.persist_path + ".g2") + # ...but the manifest was never updated to point at it. + with open(vector_db._manifest_path, "r") as f: + manifest = json.load(f) + assert manifest["generation"] == 1 + del vector_db + + recovered = PersistentHNSWLibVectorDB( + persist_path=self.persist_path, checkpoint_interval=100 + ) + # Recovers via generation 1 + WAL replay, ignoring the orphaned + # generation 2 snapshot that was never published. + knn = recovered.get_knn(embedding=[0.1, 0.2, 0.3], k=2) + assert len(knn) == 2 + def test_reset(self): vector_db = PersistentHNSWLibVectorDB(persist_path=self.persist_path) vector_db.add(embedding=[0.1, 0.2, 0.3]) @@ -74,6 +170,15 @@ def test_reset(self): knn = vector_db.get_knn(embedding=[0.1, 0.2, 0.3], k=3) assert len(knn) == 0 + def test_reset_persists_across_restart(self): + vector_db = PersistentHNSWLibVectorDB(persist_path=self.persist_path) + vector_db.add(embedding=[0.1, 0.2, 0.3]) + vector_db.reset() + del vector_db + + reloaded = PersistentHNSWLibVectorDB(persist_path=self.persist_path) + assert reloaded.is_empty() + def test_euclidean_metric(self): vector_db = PersistentHNSWLibVectorDB( persist_path=self.persist_path, diff --git a/vcache/vcache_core/cache/embedding_store/vector_db/strategies/hnsw_lib_persistent.py b/vcache/vcache_core/cache/embedding_store/vector_db/strategies/hnsw_lib_persistent.py index 228042a..abe1003 100644 --- a/vcache/vcache_core/cache/embedding_store/vector_db/strategies/hnsw_lib_persistent.py +++ b/vcache/vcache_core/cache/embedding_store/vector_db/strategies/hnsw_lib_persistent.py @@ -19,12 +19,35 @@ class PersistentHNSWLibVectorDB(VectorDB): HNSWLib-based vector database implementation that persists its index to disk, so it survives process restarts. + Persistence uses a checkpoint + write-ahead log (WAL) scheme rather than + rewriting the full hnswlib index on every mutation: + + - Each `add`/`remove` is appended as one line to a WAL file and fsync'd. + This is cheap (O(1) per call) regardless of index size. + - Every `checkpoint_interval` mutations (or on an explicit `checkpoint()` + call), the in-memory index is serialized to a new, uniquely-named + generation file and a manifest is written pointing at it. Both the + generation file and the manifest are written to a temporary path first + and then atomically renamed into place via `os.replace`, so a crash + mid-write can never leave a partially-written file visible under its + final name. + - The manifest is the single source of truth for "what generation is + current" and carries all bookkeeping (dimension, space, id counters) + that used to live in a separate sidecar file, so the index and its + metadata can never desynchronize: either the manifest rename completed + and both are visible together, or it didn't and both are still the + previous, consistent generation. + - On startup, the last checkpointed generation is loaded and then any WAL + entries recorded after that checkpoint are replayed, recovering + mutations that happened after the last checkpoint but before a crash. + This is a standalone implementation (not a subclass of `HNSWLibVectorDB`) so it does not depend on, or risk altering, that class's internals. - hnswlib's own `save_index`/`load_index` only capture the graph and - vectors, so a small JSON sidecar file next to the index file stores the - remaining bookkeeping (dimension, space, and id counters) needed to - resume exactly where the cache left off. + + Note: this scheme makes on-disk state crash-safe for a single writer. It + does not coordinate concurrent writers across multiple processes; running + more than one process against the same `persist_path` concurrently is not + supported. """ def __init__( @@ -32,20 +55,30 @@ def __init__( persist_path: str, similarity_metric_type: SimilarityMetricType = SimilarityMetricType.COSINE, max_capacity: int = 100000, + checkpoint_interval: int = 100, ): """Initializes the persistent HNSWLib vector database. Args: - persist_path (str): Path to the file used to store the hnswlib - index. A sidecar file at `persist_path + ".meta.json"` stores - the remaining bookkeeping. If either file already exists, the + persist_path (str): Base path used to derive the on-disk files: + a manifest at `persist_path + ".manifest.json"`, a WAL at + `persist_path + ".wal"`, and index generation snapshots at + `persist_path + ".g"`. If a manifest already exists, the database is restored from disk. similarity_metric_type (SimilarityMetricType): The similarity metric to use for comparisons. max_capacity (int): The maximum number of vectors the database can store. + checkpoint_interval (int): Number of mutations to accumulate in the + WAL before automatically checkpointing (serializing the full + index and clearing the WAL). Lower values bound the amount of + WAL replay work needed after a crash at the cost of more + frequent full-index writes; higher values do the opposite. """ self.persist_path = persist_path - self._meta_path = persist_path + ".meta.json" + self._manifest_path = persist_path + ".manifest.json" + self._wal_path = persist_path + ".wal" + self.checkpoint_interval = checkpoint_interval + self.embedding_count = 0 self.__next_embedding_id = 0 self.similarity_metric_type = similarity_metric_type @@ -57,11 +90,18 @@ def __init__( self.ef = None self.index = None - if os.path.exists(self._meta_path) and os.path.exists(persist_path): + self._generation = 0 + self._pending_wal_entries = 0 + + if os.path.exists(self._manifest_path): self._load_from_disk() + elif os.path.exists(self._wal_path): + # No checkpoint has ever completed, but a WAL from a previous + # (possibly crashed) instance exists: replay it from scratch. + self._replay_wal() def add(self, embedding: List[float]) -> int: - """Adds an embedding vector to the database and persists it to disk. + """Adds an embedding vector to the database and durably logs it. Args: embedding (List[float]): The embedding vector to add. @@ -75,11 +115,11 @@ def add(self, embedding: List[float]) -> int: self.index.add_items(embedding, id) self.embedding_count += 1 self.__next_embedding_id += 1 - self._save_to_disk() + self._append_wal({"op": "add", "id": id, "embedding": list(embedding)}) return id def remove(self, embedding_id: int) -> int: - """Marks an embedding for deletion and persists the change to disk. + """Marks an embedding for deletion and durably logs the change. Note: HNSWLib does not physically remove data, but marks it as deleted. @@ -97,7 +137,7 @@ def remove(self, embedding_id: int) -> int: raise ValueError("Index is not initialized") self.index.mark_deleted(embedding_id) self.embedding_count -= 1 - self._save_to_disk() + self._append_wal({"op": "remove", "id": embedding_id}) return embedding_id def get_knn(self, embedding: List[float], k: int) -> List[tuple[float, int]]: @@ -125,12 +165,43 @@ def get_knn(self, embedding: List[float], k: int) -> List[tuple[float, int]]: return list(zip(similarity_scores, id_list)) def reset(self) -> None: - """Resets the vector database to an empty state and persists it.""" + """Resets the vector database to an empty state and checkpoints it.""" if self.dim is not None: self._init_vector_store(self.dim) self.embedding_count = 0 self.__next_embedding_id = 0 - self._save_to_disk() + self.checkpoint() + + def checkpoint(self) -> None: + """Forces a full, atomic checkpoint of the index and manifest. + + Serializes the current in-memory index to a new generation file, + atomically publishes a manifest pointing at it, and then clears the + WAL, since every mutation up to this point is now captured in the + checkpoint itself. Both writes go through a temp-file-plus-rename so + a crash mid-checkpoint leaves the previous, still-consistent + generation visible. + """ + previous_generation = self._generation + self._generation += 1 + + if self.index is not None: + index_path = self._generation_path(self._generation) + tmp_index_path = index_path + ".tmp" + self.index.save_index(tmp_index_path) + self._fsync_file(tmp_index_path) + os.replace(tmp_index_path, index_path) + + self._write_manifest() + self._truncate_wal() + self._pending_wal_entries = 0 + + if self.index is not None: + old_index_path = self._generation_path(previous_generation) + if previous_generation != self._generation and os.path.exists( + old_index_path + ): + os.remove(old_index_path) def _init_vector_store(self, embedding_dim: int): """Initializes the HNSWLib index. @@ -177,36 +248,109 @@ def size(self) -> int: """ return self.embedding_count - def _save_to_disk(self) -> None: - """Persists the hnswlib index and its bookkeeping sidecar to disk.""" - self.index.save_index(self.persist_path) - with open(self._meta_path, "w") as f: - json.dump( - { - "dim": self.dim, - "space": self.space, - "max_elements": self.max_elements, - "ef_construction": self.ef_construction, - "M": self.M, - "ef": self.ef, - "embedding_count": self.embedding_count, - "next_embedding_id": self.__next_embedding_id, - }, - f, - ) + def _generation_path(self, generation: int) -> str: + """Path of the index snapshot file for a given generation number.""" + return f"{self.persist_path}.g{generation}" + + def _append_wal(self, entry: dict) -> None: + """Durably appends one mutation record to the WAL, checkpointing if + the configured interval has been reached.""" + with open(self._wal_path, "a") as f: + f.write(json.dumps(entry) + "\n") + f.flush() + os.fsync(f.fileno()) + self._pending_wal_entries += 1 + if self._pending_wal_entries >= self.checkpoint_interval: + self.checkpoint() + + def _truncate_wal(self) -> None: + """Atomically clears the WAL, e.g. after its entries are captured in + a checkpoint.""" + tmp_wal_path = self._wal_path + ".tmp" + open(tmp_wal_path, "w").close() + os.replace(tmp_wal_path, self._wal_path) + + def _write_manifest(self) -> None: + """Atomically publishes a manifest pointing at the current + generation. This is the single commit point: once this rename + completes, the new generation (and the bookkeeping alongside it) is + the durable state; until then, the previous manifest is still + authoritative.""" + manifest = { + "generation": self._generation, + "dim": self.dim, + "space": self.space, + "max_elements": self.max_elements, + "ef_construction": self.ef_construction, + "M": self.M, + "ef": self.ef, + "embedding_count": self.embedding_count, + "next_embedding_id": self.__next_embedding_id, + } + tmp_manifest_path = self._manifest_path + ".tmp" + with open(tmp_manifest_path, "w") as f: + json.dump(manifest, f) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_manifest_path, self._manifest_path) + + def _fsync_file(self, path: str) -> None: + """Flushes a file's contents to durable storage.""" + fd = os.open(path, os.O_RDWR) + try: + os.fsync(fd) + finally: + os.close(fd) def _load_from_disk(self) -> None: - """Restores the hnswlib index and bookkeeping from disk.""" - with open(self._meta_path, "r") as f: - meta = json.load(f) - self.dim = meta["dim"] - self.space = meta["space"] - self.max_elements = meta["max_elements"] - self.ef_construction = meta["ef_construction"] - self.M = meta["M"] - self.ef = meta["ef"] - self.embedding_count = meta["embedding_count"] - self.__next_embedding_id = meta["next_embedding_id"] - self.index = hnswlib.Index(space=self.space, dim=self.dim) - self.index.load_index(self.persist_path, max_elements=self.max_elements) - self.index.set_ef(self.ef) + """Restores state from the last manifest, then replays any WAL + entries recorded after that checkpoint.""" + with open(self._manifest_path, "r") as f: + manifest = json.load(f) + self._generation = manifest["generation"] + self.dim = manifest["dim"] + self.space = manifest["space"] + self.max_elements = manifest["max_elements"] + self.ef_construction = manifest["ef_construction"] + self.M = manifest["M"] + self.ef = manifest["ef"] + self.embedding_count = manifest["embedding_count"] + self.__next_embedding_id = manifest["next_embedding_id"] + + index_path = self._generation_path(self._generation) + if self.dim is not None and os.path.exists(index_path): + self.index = hnswlib.Index(space=self.space, dim=self.dim) + self.index.load_index(index_path, max_elements=self.max_elements) + self.index.set_ef(self.ef) + + self._replay_wal() + + def _replay_wal(self) -> None: + """Re-applies mutations recorded in the WAL since the last + checkpoint, recovering work that a crash interrupted before it could + be checkpointed.""" + if not os.path.exists(self._wal_path): + return + with open(self._wal_path, "r") as f: + lines = [line for line in f.read().splitlines() if line] + if not lines: + return + + for line in lines: + entry = json.loads(line) + if entry["op"] == "add": + if self.index is None: + self._init_vector_store(len(entry["embedding"])) + self.index.add_items(entry["embedding"], entry["id"]) + self.embedding_count += 1 + self.__next_embedding_id = max( + self.__next_embedding_id, entry["id"] + 1 + ) + elif entry["op"] == "remove": + self.index.mark_deleted(entry["id"]) + self.embedding_count -= 1 + + # The replayed mutations are now reflected in-memory; fold them into + # a fresh checkpoint so a subsequent restart doesn't need to replay + # them again. + self.checkpoint()