From bd92305f7aebb6f449caeec41fe18f4992099fd5 Mon Sep 17 00:00:00 2001 From: safaa Date: Thu, 16 Jul 2026 17:34:42 +0300 Subject: [PATCH 1/5] Add persistent SQLite metadata storage and persistent HNSWLib vector DB Adds two opt-in, drop-in strategy implementations so the cache can survive process restarts: SQLiteEmbeddingMetadataStorage persists metadata to a SQLite file, and PersistentHNSWLibVectorDB persists its hnswlib index plus a small bookkeeping sidecar to disk. Existing defaults (in-memory storage, non-persistent HNSWLib) are unchanged. --- .../test_sqlite_metadata_storage.py | 92 ++++++++ .../test_persistent_hnsw_lib.py | 89 ++++++++ vcache/__init__.py | 4 + .../embedding_metadata_storage/__init__.py | 4 + .../strategies/__init__.py | 7 +- .../strategies/sqlite.py | 153 +++++++++++++ .../embedding_store/vector_db/__init__.py | 4 + .../vector_db/strategies/__init__.py | 8 +- .../strategies/hnsw_lib_persistent.py | 212 ++++++++++++++++++ 9 files changed, 571 insertions(+), 2 deletions(-) create mode 100644 tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py create mode 100644 tests/unit/VectorDBStrategy/test_persistent_hnsw_lib.py create mode 100644 vcache/vcache_core/cache/embedding_store/embedding_metadata_storage/strategies/sqlite.py create mode 100644 vcache/vcache_core/cache/embedding_store/vector_db/strategies/hnsw_lib_persistent.py diff --git a/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py b/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py new file mode 100644 index 0000000..d401915 --- /dev/null +++ b/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py @@ -0,0 +1,92 @@ +import os +import tempfile +import unittest + +from vcache.vcache_core.cache.embedding_store.embedding_metadata_storage import ( + SQLiteEmbeddingMetadataStorage, +) +from vcache.vcache_core.cache.embedding_store.embedding_metadata_storage.embedding_metadata_obj import ( + EmbeddingMetadataObj, +) + + +class TestSQLiteEmbeddingMetadataStorage(unittest.TestCase): + def setUp(self): + # ignore_cleanup_errors: on Windows, sqlite3 keeps the file handle + # open for the lifetime of the connection, which otherwise makes + # tearDown's rmtree fail with a PermissionError. + self.tmp_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) + self.db_path = os.path.join(self.tmp_dir.name, "metadata.sqlite3") + + def tearDown(self): + self.tmp_dir.cleanup() + + def test_sqlite_strategy(self): + embedding_metadata_storage = SQLiteEmbeddingMetadataStorage(self.db_path) + + initial_obj = EmbeddingMetadataObj(embedding_id=0, response="test") + embedding_id = embedding_metadata_storage.add_metadata( + embedding_id=0, metadata=initial_obj + ) + assert embedding_id == 0 + assert embedding_metadata_storage.get_metadata(embedding_id=0) == initial_obj + + updated_obj = EmbeddingMetadataObj(embedding_id=0, response="test2") + embedding_metadata_storage.update_metadata(embedding_id=0, metadata=updated_obj) + assert embedding_metadata_storage.get_metadata(embedding_id=0) == updated_obj + + embedding_metadata_storage.flush() + with self.assertRaises(ValueError): + embedding_metadata_storage.get_metadata(embedding_id=0) + + def test_get_missing_metadata_raises(self): + embedding_metadata_storage = SQLiteEmbeddingMetadataStorage(self.db_path) + with self.assertRaises(ValueError): + embedding_metadata_storage.get_metadata(embedding_id=42) + + def test_update_missing_metadata_raises(self): + embedding_metadata_storage = SQLiteEmbeddingMetadataStorage(self.db_path) + with self.assertRaises(ValueError): + embedding_metadata_storage.update_metadata( + embedding_id=42, + metadata=EmbeddingMetadataObj(embedding_id=42, response="test"), + ) + + def test_remove_metadata(self): + embedding_metadata_storage = SQLiteEmbeddingMetadataStorage(self.db_path) + embedding_metadata_storage.add_metadata( + embedding_id=1, + metadata=EmbeddingMetadataObj(embedding_id=1, response="test"), + ) + assert embedding_metadata_storage.remove_metadata(embedding_id=1) is True + assert embedding_metadata_storage.remove_metadata(embedding_id=1) is False + + def test_get_all_embedding_metadata_objects(self): + embedding_metadata_storage = SQLiteEmbeddingMetadataStorage(self.db_path) + for i in range(3): + embedding_metadata_storage.add_metadata( + embedding_id=i, + metadata=EmbeddingMetadataObj(embedding_id=i, response=f"test{i}"), + ) + all_metadata = embedding_metadata_storage.get_all_embedding_metadata_objects() + assert len(all_metadata) == 3 + assert {meta.response for meta in all_metadata} == {"test0", "test1", "test2"} + + def test_survives_simulated_restart(self): + """Data written by one instance should be visible to a fresh instance + pointed at the same file, simulating a process restart.""" + first_instance = SQLiteEmbeddingMetadataStorage(self.db_path) + first_instance.add_metadata( + embedding_id=7, + metadata=EmbeddingMetadataObj(embedding_id=7, response="persisted", id_set=3), + ) + del first_instance + + second_instance = SQLiteEmbeddingMetadataStorage(self.db_path) + restored = second_instance.get_metadata(embedding_id=7) + assert restored.response == "persisted" + assert restored.id_set == 3 + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/VectorDBStrategy/test_persistent_hnsw_lib.py b/tests/unit/VectorDBStrategy/test_persistent_hnsw_lib.py new file mode 100644 index 0000000..05555ed --- /dev/null +++ b/tests/unit/VectorDBStrategy/test_persistent_hnsw_lib.py @@ -0,0 +1,89 @@ +import os +import tempfile +import unittest + +from vcache.vcache_core.cache.embedding_store.vector_db import ( + PersistentHNSWLibVectorDB, + SimilarityMetricType, +) + + +class TestPersistentHNSWLibVectorDB(unittest.TestCase): + def setUp(self): + self.tmp_dir = tempfile.TemporaryDirectory() + self.persist_path = os.path.join(self.tmp_dir.name, "index.hnsw") + + def tearDown(self): + self.tmp_dir.cleanup() + + def test_add_and_get_knn(self): + vector_db = PersistentHNSWLibVectorDB(persist_path=self.persist_path) + + embedding = [0.1, 0.2, 0.3] + id1 = vector_db.add(embedding=embedding) + knn = vector_db.get_knn(embedding=embedding, k=1) + assert len(knn) == 1 + assert knn[0][1] == id1 + + vector_db.add(embedding=[0.2, 0.3, 0.4]) + vector_db.add(embedding=[0.3, 0.4, 0.5]) + knn = vector_db.get_knn(embedding=embedding, k=3) + assert len(knn) == 3 + + def test_remove(self): + vector_db = PersistentHNSWLibVectorDB(persist_path=self.persist_path) + + id1 = vector_db.add(embedding=[0.1, 0.2, 0.3]) + id2 = vector_db.add(embedding=[0.2, 0.3, 0.4]) + vector_db.remove(embedding_id=id1) + + knn = vector_db.get_knn(embedding=[0.1, 0.2, 0.3], k=2) + assert len(knn) == 1 + assert knn[0][1] == id2 + + def test_persists_files_to_disk(self): + vector_db = PersistentHNSWLibVectorDB(persist_path=self.persist_path) + 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") + + 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) + 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) + knn = second_instance.get_knn(embedding=[0.1, 0.2, 0.3], k=2) + assert len(knn) == 2 + + # New adds after reload should not collide with restored ids + new_id = second_instance.add(embedding=[0.3, 0.4, 0.5]) + assert new_id == 2 + + def test_reset(self): + vector_db = PersistentHNSWLibVectorDB(persist_path=self.persist_path) + vector_db.add(embedding=[0.1, 0.2, 0.3]) + vector_db.add(embedding=[0.2, 0.3, 0.4]) + + vector_db.reset() + + knn = vector_db.get_knn(embedding=[0.1, 0.2, 0.3], k=3) + assert len(knn) == 0 + + def test_euclidean_metric(self): + vector_db = PersistentHNSWLibVectorDB( + persist_path=self.persist_path, + similarity_metric_type=SimilarityMetricType.EUCLIDEAN, + ) + embedding = [0.1, 0.2, 0.3] + id1 = vector_db.add(embedding=embedding) + knn = vector_db.get_knn(embedding=embedding, k=1) + assert knn[0][1] == id1 + + +if __name__ == "__main__": + unittest.main() diff --git a/vcache/__init__.py b/vcache/__init__.py index 10e2261..5d8c5bb 100644 --- a/vcache/__init__.py +++ b/vcache/__init__.py @@ -37,6 +37,7 @@ from .vcache_core.cache.embedding_store.embedding_metadata_storage import ( InMemoryEmbeddingMetadataStorage, LangchainMetadataStorage, + SQLiteEmbeddingMetadataStorage, ) # Concrete Vector databases @@ -44,6 +45,7 @@ ChromaVectorDB, FAISSVectorDB, HNSWLibVectorDB, + PersistentHNSWLibVectorDB, SimilarityMetricType, VectorDB, ) @@ -112,6 +114,7 @@ # Concrete Vector databases "FAISSVectorDB", "HNSWLibVectorDB", + "PersistentHNSWLibVectorDB", "ChromaVectorDB", "SimilarityMetricType", # Concrete Similarity evaluators @@ -128,5 +131,6 @@ # Concrete Embedding metadata storage "InMemoryEmbeddingMetadataStorage", "LangchainMetadataStorage", + "SQLiteEmbeddingMetadataStorage", "EmbeddingMetadataObj", ] diff --git a/vcache/vcache_core/cache/embedding_store/embedding_metadata_storage/__init__.py b/vcache/vcache_core/cache/embedding_store/embedding_metadata_storage/__init__.py index 9a8a604..154ba88 100644 --- a/vcache/vcache_core/cache/embedding_store/embedding_metadata_storage/__init__.py +++ b/vcache/vcache_core/cache/embedding_store/embedding_metadata_storage/__init__.py @@ -7,9 +7,13 @@ from vcache.vcache_core.cache.embedding_store.embedding_metadata_storage.strategies.langchain import ( LangchainMetadataStorage, ) +from vcache.vcache_core.cache.embedding_store.embedding_metadata_storage.strategies.sqlite import ( + SQLiteEmbeddingMetadataStorage, +) __all__ = [ "EmbeddingMetadataStorage", "InMemoryEmbeddingMetadataStorage", "LangchainMetadataStorage", + "SQLiteEmbeddingMetadataStorage", ] diff --git a/vcache/vcache_core/cache/embedding_store/embedding_metadata_storage/strategies/__init__.py b/vcache/vcache_core/cache/embedding_store/embedding_metadata_storage/strategies/__init__.py index 493255e..2ca0271 100644 --- a/vcache/vcache_core/cache/embedding_store/embedding_metadata_storage/strategies/__init__.py +++ b/vcache/vcache_core/cache/embedding_store/embedding_metadata_storage/strategies/__init__.py @@ -1,4 +1,9 @@ from .in_memory import InMemoryEmbeddingMetadataStorage from .langchain import LangchainMetadataStorage +from .sqlite import SQLiteEmbeddingMetadataStorage -__all__ = ["InMemoryEmbeddingMetadataStorage", "LangchainMetadataStorage"] +__all__ = [ + "InMemoryEmbeddingMetadataStorage", + "LangchainMetadataStorage", + "SQLiteEmbeddingMetadataStorage", +] diff --git a/vcache/vcache_core/cache/embedding_store/embedding_metadata_storage/strategies/sqlite.py b/vcache/vcache_core/cache/embedding_store/embedding_metadata_storage/strategies/sqlite.py new file mode 100644 index 0000000..0f186bf --- /dev/null +++ b/vcache/vcache_core/cache/embedding_store/embedding_metadata_storage/strategies/sqlite.py @@ -0,0 +1,153 @@ +import pickle +import sqlite3 +import threading +from typing import List + +from vcache.vcache_core.cache.embedding_store.embedding_metadata_storage.embedding_metadata_obj import ( + EmbeddingMetadataObj, +) +from vcache.vcache_core.cache.embedding_store.embedding_metadata_storage.embedding_metadata_storage import ( + EmbeddingMetadataStorage, +) + + +class SQLiteEmbeddingMetadataStorage(EmbeddingMetadataStorage): + """ + SQLite-backed implementation of embedding metadata storage. + + Unlike `InMemoryEmbeddingMetadataStorage`, this implementation persists + metadata to disk, so it survives process restarts. Each metadata object + is stored as a pickled blob keyed by `embedding_id`, rather than mapped to + individual columns, since `EmbeddingMetadataObj` mixes datetimes, optional + floats, and tuples, and gains fields over time. + """ + + def __init__(self, db_path: str): + """ + Initialize SQLite-backed embedding metadata storage. + + Args: + db_path: Path to the SQLite database file. Created if it does + not yet exist. + """ + self.db_path = db_path + self._lock = threading.Lock() + self._connection = sqlite3.connect(db_path, check_same_thread=False) + with self._lock: + self._connection.execute( + "CREATE TABLE IF NOT EXISTS embedding_metadata (" + "embedding_id INTEGER PRIMARY KEY, data BLOB NOT NULL)" + ) + self._connection.commit() + + def add_metadata(self, embedding_id: int, metadata: EmbeddingMetadataObj) -> int: + """ + Add metadata for a specific embedding. + + Args: + embedding_id: The id of the embedding to add the metadata for. + metadata: The metadata to add to the embedding. + + Returns: + The id of the embedding. + """ + with self._lock: + self._connection.execute( + "INSERT OR REPLACE INTO embedding_metadata (embedding_id, data) " + "VALUES (?, ?)", + (embedding_id, pickle.dumps(metadata)), + ) + self._connection.commit() + return embedding_id + + def get_metadata(self, embedding_id: int) -> EmbeddingMetadataObj: + """ + Get metadata for a specific embedding. + + Args: + embedding_id: The id of the embedding to get the metadata for. + + Returns: + The metadata of the embedding. + + Raises: + ValueError: If embedding metadata is not found. + """ + with self._lock: + row = self._connection.execute( + "SELECT data FROM embedding_metadata WHERE embedding_id = ?", + (embedding_id,), + ).fetchone() + if row is None: + raise ValueError( + f"Embedding metadata for embedding id {embedding_id} not found" + ) + return pickle.loads(row[0]) + + def update_metadata( + self, embedding_id: int, metadata: EmbeddingMetadataObj + ) -> EmbeddingMetadataObj: + """ + Update metadata for a specific embedding. + + Args: + embedding_id: The id of the embedding to update the metadata for. + metadata: The metadata to update the embedding with. + + Returns: + The updated metadata of the embedding. + + Raises: + ValueError: If embedding metadata is not found. + """ + with self._lock: + cursor = self._connection.execute( + "UPDATE embedding_metadata SET data = ? WHERE embedding_id = ?", + (pickle.dumps(metadata), embedding_id), + ) + self._connection.commit() + not_found = cursor.rowcount == 0 + if not_found: + raise ValueError( + f"Embedding metadata for embedding id {embedding_id} not found" + ) + return metadata + + def remove_metadata(self, embedding_id: int) -> bool: + """ + Remove metadata for a specific embedding. + + Args: + embedding_id: The id of the embedding to remove metadata for. + + Returns: + True if metadata was removed, False if not found. + """ + with self._lock: + cursor = self._connection.execute( + "DELETE FROM embedding_metadata WHERE embedding_id = ?", + (embedding_id,), + ) + self._connection.commit() + return cursor.rowcount > 0 + + def flush(self) -> None: + """ + Flush all metadata from storage. + """ + with self._lock: + self._connection.execute("DELETE FROM embedding_metadata") + self._connection.commit() + + def get_all_embedding_metadata_objects(self) -> List[EmbeddingMetadataObj]: + """ + Get all embedding metadata objects in storage. + + Returns: + A list of all the embedding metadata objects in the storage. + """ + with self._lock: + rows = self._connection.execute( + "SELECT data FROM embedding_metadata" + ).fetchall() + return [pickle.loads(row[0]) for row in rows] diff --git a/vcache/vcache_core/cache/embedding_store/vector_db/__init__.py b/vcache/vcache_core/cache/embedding_store/vector_db/__init__.py index bb4a6c3..61697a0 100644 --- a/vcache/vcache_core/cache/embedding_store/vector_db/__init__.py +++ b/vcache/vcache_core/cache/embedding_store/vector_db/__init__.py @@ -7,6 +7,9 @@ from vcache.vcache_core.cache.embedding_store.vector_db.strategies.hnsw_lib import ( HNSWLibVectorDB, ) +from vcache.vcache_core.cache.embedding_store.vector_db.strategies.hnsw_lib_persistent import ( + PersistentHNSWLibVectorDB, +) from vcache.vcache_core.cache.embedding_store.vector_db.vector_db import ( SimilarityMetricType, VectorDB, @@ -16,6 +19,7 @@ "VectorDB", "SimilarityMetricType", "HNSWLibVectorDB", + "PersistentHNSWLibVectorDB", "FAISSVectorDB", "ChromaVectorDB", ] diff --git a/vcache/vcache_core/cache/embedding_store/vector_db/strategies/__init__.py b/vcache/vcache_core/cache/embedding_store/vector_db/strategies/__init__.py index 0a0c86d..8008a8f 100644 --- a/vcache/vcache_core/cache/embedding_store/vector_db/strategies/__init__.py +++ b/vcache/vcache_core/cache/embedding_store/vector_db/strategies/__init__.py @@ -1,5 +1,11 @@ from .chroma import ChromaVectorDB from .faiss import FAISSVectorDB from .hnsw_lib import HNSWLibVectorDB +from .hnsw_lib_persistent import PersistentHNSWLibVectorDB -__all__ = ["ChromaVectorDB", "FAISSVectorDB", "HNSWLibVectorDB"] +__all__ = [ + "ChromaVectorDB", + "FAISSVectorDB", + "HNSWLibVectorDB", + "PersistentHNSWLibVectorDB", +] 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 new file mode 100644 index 0000000..228042a --- /dev/null +++ b/vcache/vcache_core/cache/embedding_store/vector_db/strategies/hnsw_lib_persistent.py @@ -0,0 +1,212 @@ +import json +import os +from typing import List + +import hnswlib + +from vcache.vcache_core.cache.embedding_store.vector_db.vector_db import ( + SimilarityMetricType, + VectorDB, +) + +""" +Run 'sudo apt-get install build-essential' on Linux Debian/Ubuntu to install the build-essential package +""" + + +class PersistentHNSWLibVectorDB(VectorDB): + """ + HNSWLib-based vector database implementation that persists its index to + disk, so it survives process restarts. + + 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. + """ + + def __init__( + self, + persist_path: str, + similarity_metric_type: SimilarityMetricType = SimilarityMetricType.COSINE, + max_capacity: int = 100000, + ): + """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 + 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. + """ + self.persist_path = persist_path + self._meta_path = persist_path + ".meta.json" + self.embedding_count = 0 + self.__next_embedding_id = 0 + self.similarity_metric_type = similarity_metric_type + self.space = None + self.dim = None + self.max_elements = max_capacity + self.ef_construction = None + self.M = None + self.ef = None + self.index = None + + if os.path.exists(self._meta_path) and os.path.exists(persist_path): + self._load_from_disk() + + def add(self, embedding: List[float]) -> int: + """Adds an embedding vector to the database and persists it to disk. + + Args: + embedding (List[float]): The embedding vector to add. + + Returns: + int: The unique ID assigned to the added embedding. + """ + if self.index is None: + self._init_vector_store(len(embedding)) + id = self.__next_embedding_id + self.index.add_items(embedding, id) + self.embedding_count += 1 + self.__next_embedding_id += 1 + self._save_to_disk() + return id + + def remove(self, embedding_id: int) -> int: + """Marks an embedding for deletion and persists the change to disk. + + Note: + HNSWLib does not physically remove data, but marks it as deleted. + + Args: + embedding_id (int): The ID of the embedding to remove. + + Returns: + int: The ID of the removed embedding. + + Raises: + ValueError: If the index has not been initialized. + """ + if self.index is None: + raise ValueError("Index is not initialized") + self.index.mark_deleted(embedding_id) + self.embedding_count -= 1 + self._save_to_disk() + return embedding_id + + def get_knn(self, embedding: List[float], k: int) -> List[tuple[float, int]]: + """Gets k-nearest neighbors for a given embedding. + + Args: + embedding (List[float]): The query embedding vector. + k (int): The number of nearest neighbors to return. + + Returns: + List[tuple[float, int]]: A list of tuples containing similarity + scores and embedding IDs. + """ + if self.index is None: + return [] + k_ = min(k, self.embedding_count) + if k_ == 0: + return [] + ids, similarities = self.index.knn_query(embedding, k=k_) + metric_type = self.similarity_metric_type.value + similarity_scores = [ + self.transform_similarity_score(sim, metric_type) for sim in similarities[0] + ] + id_list = [int(id) for id in ids[0]] + return list(zip(similarity_scores, id_list)) + + def reset(self) -> None: + """Resets the vector database to an empty state and persists 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() + + def _init_vector_store(self, embedding_dim: int): + """Initializes the HNSWLib index. + + Args: + embedding_dim (int): The dimension of the embedding vectors. + + Raises: + ValueError: If the similarity metric type is invalid. + """ + metric_type = self.similarity_metric_type.value + match metric_type: + case "cosine": + self.space = "cosine" + case "euclidean": + self.space = "l2" + case _: + raise ValueError(f"Invalid similarity metric type: {metric_type}") + self.dim = embedding_dim + self.ef_construction = 350 + self.M = 52 + self.ef = 400 + self.index = hnswlib.Index(space=self.space, dim=self.dim) + self.index.init_index( + max_elements=self.max_elements, + ef_construction=self.ef_construction, + M=self.M, + ) + self.index.set_ef(self.ef) + + def is_empty(self) -> bool: + """Checks if the vector database is empty. + + Returns: + bool: True if the database contains no embeddings, False otherwise. + """ + return self.embedding_count == 0 + + def size(self) -> int: + """Gets the number of embeddings in the vector database. + + Returns: + int: The number of embeddings in the vector database. + """ + 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 _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) From 78192cb62129ea60344fc337819c6ea1a75704d4 Mon Sep 17 00:00:00 2001 From: safaa Date: Fri, 17 Jul 2026 10:17:49 +0300 Subject: [PATCH 2/5] Make PersistentHNSWLibVectorDB writes atomic and avoid full index rewrites Addresses maintainer feedback on the persistent-storage PR: the previous implementation called index.save_index() plus a separate meta.json write on every add/remove, which rewrote the whole index each time and could leave the index and metadata desynchronized if a crash landed between the two writes. Replaces this with a checkpoint + write-ahead-log scheme: mutations are appended to a WAL (cheap, O(1) per call) and only folded into a full index snapshot periodically. Checkpoints publish a new generation file and manifest via temp-file-plus-os.replace, so the previous generation stays valid and visible until the new one is fully committed, and a crash can never surface a partially-written index or an index/metadata mismatch. On load, the last checkpoint is restored and any WAL entries written after it are replayed. --- .../test_persistent_hnsw_lib.py | 123 ++++++++- .../strategies/hnsw_lib_persistent.py | 236 ++++++++++++++---- 2 files changed, 304 insertions(+), 55 deletions(-) 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() From 0baddb47f52eb1e58413507e3dcadb3654ed8686 Mon Sep 17 00:00:00 2001 From: safaa Date: Fri, 17 Jul 2026 20:36:35 +0300 Subject: [PATCH 3/5] Apply ruff formatting --- .../EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py b/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py index d401915..240ee30 100644 --- a/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py +++ b/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py @@ -78,7 +78,9 @@ def test_survives_simulated_restart(self): first_instance = SQLiteEmbeddingMetadataStorage(self.db_path) first_instance.add_metadata( embedding_id=7, - metadata=EmbeddingMetadataObj(embedding_id=7, response="persisted", id_set=3), + metadata=EmbeddingMetadataObj( + embedding_id=7, response="persisted", id_set=3 + ), ) del first_instance From a9412f557551ad8f982befcf8571128f993da270 Mon Sep 17 00:00:00 2001 From: safaa Date: Fri, 17 Jul 2026 20:38:41 +0300 Subject: [PATCH 4/5] Revert "Apply ruff formatting" This reverts commit 80ed092f0114ef83daae11b7238dbb46952c20eb. --- .../EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py b/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py index 240ee30..d401915 100644 --- a/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py +++ b/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py @@ -78,9 +78,7 @@ def test_survives_simulated_restart(self): first_instance = SQLiteEmbeddingMetadataStorage(self.db_path) first_instance.add_metadata( embedding_id=7, - metadata=EmbeddingMetadataObj( - embedding_id=7, response="persisted", id_set=3 - ), + metadata=EmbeddingMetadataObj(embedding_id=7, response="persisted", id_set=3), ) del first_instance From 8ef66526a8698b4d985efc2f6de78a0a3219192f Mon Sep 17 00:00:00 2001 From: safaa Date: Fri, 17 Jul 2026 21:20:37 +0300 Subject: [PATCH 5/5] Apply ruff formatting --- .../EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py b/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py index d401915..240ee30 100644 --- a/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py +++ b/tests/unit/EmbeddingMetadataStrategy/test_sqlite_metadata_storage.py @@ -78,7 +78,9 @@ def test_survives_simulated_restart(self): first_instance = SQLiteEmbeddingMetadataStorage(self.db_path) first_instance.add_metadata( embedding_id=7, - metadata=EmbeddingMetadataObj(embedding_id=7, response="persisted", id_set=3), + metadata=EmbeddingMetadataObj( + embedding_id=7, response="persisted", id_set=3 + ), ) del first_instance