akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add bounded memory support to HNSWVectorIndex with LRU eviction

Browse source 
HNSWVectorIndex now supports optional memory bounding:
- New max_entries parameter sets soft limit on number of entries
- When limit reached, lowest importance entries are evicted
- Eviction uses importance (ascending) then age (oldest first) ordering
- eviction_batch_size controls how many entries evicted at once

Changes:
- Add max_entries and eviction_batch_size parameters
- Add _evict_entries() method for importance-based eviction
- Update get_memory_stats() to report budget_bytes and evictions
- Update stats() to include eviction metrics
- Update save_index/load_index to persist eviction settings
- Add 5 new tests for eviction behavior
This commit is contained in:
chopratejas 2026-02-01 21:12:06 -08:00
parent f5fc371f07
commit b27c95bdac
2 changed files with 301 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

112
headroom/memory/adapters/hnsw.py
View file

@ -194,9 +194,10 @@ class HNSWVectorIndex:
- Persistence support with save_index/load_index
- Thread-safe operations with Lock
- Optional auto-save on index modifications
- **Bounded memory with LRU eviction** (when max_entries is set)
Usage:
index = HNSWVectorIndex(dimension=384)
index = HNSWVectorIndex(dimension=384, max_entries=10000)
await index.index(memory_with_embedding)
results = await index.search(VectorFilter(
query_vector=query_embedding,
@ -211,6 +212,12 @@ class HNSWVectorIndex:
better recall but use more memory. Default: 16
- ef_search: Size of dynamic candidate list during search. Higher values
give better recall but slower search. Default: 50
Memory Bounding:
- max_entries: Soft limit on number of entries. When reached, lowest
importance entries are evicted to make room. Default: None (unbounded).
- eviction_batch_size: Number of entries to evict at once when limit
is reached. Default: 100.
"""
def __init__(
@ -222,6 +229,8 @@ class HNSWVectorIndex:
ef_search: int = 50,
auto_save: bool = False,
save_path: str | Path | None = None,
max_entries: int | None = None,
eviction_batch_size: int = 100,
) -> None:
"""Initialize the HNSW vector index.
@ -238,6 +247,9 @@ class HNSWVectorIndex:
auto_save: If True and save_path is set, automatically save
index after modifications.
save_path: Path for auto-save operations. Required if auto_save=True.
max_entries: Soft limit on number of entries. When reached,
lowest importance entries are evicted. None = unbounded.
eviction_batch_size: Number of entries to evict when limit is reached.
Raises:
ValueError: If auto_save is True but save_path is not provided.
@ -263,6 +275,11 @@ class HNSWVectorIndex:
self._auto_save = auto_save
self._save_path = Path(save_path) if save_path else None
# Memory bounding
self._max_entries = max_entries
self._eviction_batch_size = eviction_batch_size
self._eviction_count = 0 # Track total evictions for stats
# Initialize HNSW index with cosine similarity
# hnswlib uses 'cosine' space which internally normalizes vectors
# Note: hnswlib is guaranteed non-None here due to _check_hnswlib_available() above
@ -302,7 +319,8 @@ class HNSWVectorIndex:
async def index(self, memory: Memory) -> None:
"""Index a memory's embedding for similarity search.
The memory must have an embedding set.
The memory must have an embedding set. If max_entries is set and
the limit is reached, low-importance entries are evicted.
Args:
memory: The memory to index.
@ -327,7 +345,13 @@ class HNSWVectorIndex:
# Update metadata
self._metadata[memory.id] = IndexedMemoryMetadata.from_memory(memory)
else:
# Resize if needed
# Evict if at capacity (before adding new entry)
if self._max_entries is not None:
current_size = len(self._memory_to_hnsw)
if current_size >= self._max_entries:
self._evict_entries(self._eviction_batch_size)
# Resize HNSW index if needed (separate from entry limit)
if self._next_hnsw_id >= self._max_elements:
self._resize_index(self._max_elements * 2)
@ -350,6 +374,55 @@ class HNSWVectorIndex:
if self._auto_save and self._save_path:
self.save_index(self._save_path)
def _evict_entries(self, count: int) -> int:
"""Evict the lowest importance entries from the index.
Must be called with lock held.
Eviction strategy: Sort by importance (ascending), then by age
(oldest first for ties). Evict the lowest scoring entries.
Args:
count: Number of entries to evict.
Returns:
Number of entries actually evicted.
"""
if not self._metadata:
return 0
# Sort entries by importance (ascending), then by created_at (oldest first)
sorted_entries = sorted(
self._metadata.items(),
key=lambda x: (x[1].importance, x[1].created_at),
)
# Evict the lowest importance entries
evicted = 0
for memory_id, _metadata in sorted_entries[:count]:
if memory_id not in self._memory_to_hnsw:
continue
hnsw_id = self._memory_to_hnsw[memory_id]
# Mark as deleted in HNSW index
self._index.mark_deleted(hnsw_id)
# Remove from mappings
del self._memory_to_hnsw[memory_id]
del self._hnsw_to_memory[hnsw_id]
# Remove metadata and embedding
if memory_id in self._metadata:
del self._metadata[memory_id]
if memory_id in self._embeddings:
del self._embeddings[memory_id]
evicted += 1
self._eviction_count += evicted
return evicted
async def index_batch(self, memories: list[Memory]) -> int:
"""Index multiple memories' embeddings.
@ -736,6 +809,9 @@ class HNSWVectorIndex:
"ef_construction": self._ef_construction,
"m": self._m,
"ef_search": self._ef_search,
"max_entries": self._max_entries,
"eviction_batch_size": self._eviction_batch_size,
"eviction_count": self._eviction_count,
"memory_to_hnsw": self._memory_to_hnsw,
"hnsw_to_memory": self._hnsw_to_memory,
"next_hnsw_id": self._next_hnsw_id,
@ -786,6 +862,11 @@ class HNSWVectorIndex:
self._m = meta_data["m"]
self._ef_search = meta_data["ef_search"]
# Restore bounding parameters (with defaults for backward compatibility)
self._max_entries = meta_data.get("max_entries")
self._eviction_batch_size = meta_data.get("eviction_batch_size", 100)
self._eviction_count = meta_data.get("eviction_count", 0)
# Create new index and load from file
self._index = hnswlib.Index(space="cosine", dim=self._dimension) # type: ignore[union-attr]
self._index.load_index(
@ -829,6 +910,7 @@ class HNSWVectorIndex:
self._next_hnsw_id = 0
self._metadata.clear()
self._embeddings.clear()
self._eviction_count = 0
if self._auto_save and self._save_path:
self.save_index(self._save_path)
@ -840,17 +922,21 @@ class HNSWVectorIndex:
Dictionary with index metrics.
"""
with self._lock:
current_size = len(self._memory_to_hnsw)
return {
"size": len(self._memory_to_hnsw),
"size": current_size,
"dimension": self._dimension,
"max_elements": self._max_elements,
"max_entries": self._max_entries,
"ef_construction": self._ef_construction,
"m": self._m,
"ef_search": self._ef_search,
"eviction_count": self._eviction_count,
"utilization": (
(len(self._memory_to_hnsw) / self._max_elements) * 100
if self._max_elements > 0
else 0.0
(current_size / self._max_elements) * 100 if self._max_elements > 0 else 0.0
),
"entry_utilization": (
(current_size / self._max_entries) * 100 if self._max_entries else None
),
}
@ -902,14 +988,22 @@ class HNSWVectorIndex:
index_size_estimate = len(self._memory_to_hnsw) * (self._dimension * 4 + self._m * 8)
size_bytes += index_size_estimate
# Calculate budget based on max_entries if set
# Budget = estimated size at max capacity
budget_bytes = None
if self._max_entries is not None:
# Estimate: each entry ~= embedding bytes + metadata overhead (~500 bytes)
per_entry_estimate = self._dimension * 4 + self._m * 8 + 500
budget_bytes = self._max_entries * per_entry_estimate
return ComponentStats(
name="vector_index",
entry_count=len(self._memory_to_hnsw),
size_bytes=size_bytes,
budget_bytes=None,
budget_bytes=budget_bytes,
hits=0,
misses=0,
evictions=0,
evictions=self._eviction_count,
)
def set_ef_search(self, ef_search: int) -> None:

198
tests/test_hnsw_only.py
View file

@ -78,3 +78,201 @@ class TestHNSWVectorIndex:
assert results[0].memory.id == memories[0].id
assert results[0].similarity > 0.99
print("[TEST] PASSED!")
@pytest.mark.asyncio
async def test_bounded_index_eviction(self, temp_db_path):
"""Test that bounded index evicts low-importance entries."""
from headroom.memory.adapters.hnsw import HNSWVectorIndex
# Create bounded index with max 5 entries
index = HNSWVectorIndex(
dimension=384,
max_entries=5,
eviction_batch_size=2,
)
np.random.seed(42)
# Add 5 memories with varying importance
memories = []
for i in range(5):
embedding = np.random.randn(384).astype(np.float32)
memory = Memory(
content=f"Content {i}",
user_id="alice",
embedding=embedding,
importance=0.1 * (i + 1), # 0.1, 0.2, 0.3, 0.4, 0.5
)
await index.index(memory)
memories.append(memory)
assert index.size == 5
# Add one more - should trigger eviction of lowest importance
new_embedding = np.random.randn(384).astype(np.float32)
new_memory = Memory(
content="New high importance",
user_id="alice",
embedding=new_embedding,
importance=0.9,
)
await index.index(new_memory)
# Should have evicted 2 entries (eviction_batch_size) then added 1
# So size should be 5 - 2 + 1 = 4
assert index.size == 4
# The lowest importance entries (0.1, 0.2) should be gone
stats = index.get_memory_stats()
assert stats.evictions == 2
# Search should not find the evicted memories
filter = VectorFilter(
query_vector=memories[0].embedding, # Lowest importance, should be evicted
top_k=10,
user_id="alice",
)
results = await index.search(filter)
# memories[0] and memories[1] should be evicted
result_ids = {r.memory.id for r in results}
assert memories[0].id not in result_ids
assert memories[1].id not in result_ids
@pytest.mark.asyncio
async def test_bounded_index_stats(self, temp_db_path):
"""Test that bounded index reports correct stats."""
from headroom.memory.adapters.hnsw import HNSWVectorIndex
index = HNSWVectorIndex(
dimension=384,
max_entries=100,
)
stats = index.get_memory_stats()
assert stats.name == "vector_index"
assert stats.entry_count == 0
assert stats.budget_bytes is not None # Should have budget when max_entries set
assert stats.evictions == 0
# Add some entries
np.random.seed(42)
for i in range(10):
embedding = np.random.randn(384).astype(np.float32)
memory = Memory(
content=f"Content {i}",
user_id="alice",
embedding=embedding,
)
await index.index(memory)
stats = index.get_memory_stats()
assert stats.entry_count == 10
assert stats.size_bytes > 0
@pytest.mark.asyncio
async def test_unbounded_index_no_eviction(self, temp_db_path):
"""Test that unbounded index doesn't evict."""
from headroom.memory.adapters.hnsw import HNSWVectorIndex
# Create unbounded index (max_entries=None)
index = HNSWVectorIndex(dimension=384)
np.random.seed(42)
# Add many memories
for i in range(20):
embedding = np.random.randn(384).astype(np.float32)
memory = Memory(
content=f"Content {i}",
user_id="alice",
embedding=embedding,
importance=0.1,
)
await index.index(memory)
# All should be present
assert index.size == 20
stats = index.get_memory_stats()
assert stats.budget_bytes is None # No budget when unbounded
assert stats.evictions == 0
@pytest.mark.asyncio
async def test_eviction_prefers_low_importance_then_old(self, temp_db_path):
"""Test eviction order: lowest importance first, then oldest."""
import time
from headroom.memory.adapters.hnsw import HNSWVectorIndex
index = HNSWVectorIndex(
dimension=384,
max_entries=3,
eviction_batch_size=1,
)
np.random.seed(42)
# Add memories with same importance but different times
memories = []
for i in range(3):
embedding = np.random.randn(384).astype(np.float32)
memory = Memory(
content=f"Content {i}",
user_id="alice",
embedding=embedding,
importance=0.5, # Same importance
)
await index.index(memory)
memories.append(memory)
time.sleep(0.01) # Small delay to ensure different created_at
# Add one more to trigger eviction
new_embedding = np.random.randn(384).astype(np.float32)
await index.index(
Memory(
content="New",
user_id="alice",
embedding=new_embedding,
importance=0.5,
)
)
# Should have evicted the oldest (first) entry
assert index.size == 3
assert memories[0].id not in index._memory_to_hnsw
@pytest.mark.asyncio
async def test_save_load_preserves_eviction_settings(self, temp_db_path):
"""Test that save/load preserves eviction settings."""
from headroom.memory.adapters.hnsw import HNSWVectorIndex
index = HNSWVectorIndex(
dimension=384,
max_entries=50,
eviction_batch_size=10,
save_path=temp_db_path,
)
np.random.seed(42)
# Add some entries
for i in range(5):
embedding = np.random.randn(384).astype(np.float32)
memory = Memory(
content=f"Content {i}",
user_id="alice",
embedding=embedding,
)
await index.index(memory)
# Save
index.save_index(temp_db_path)
# Create new index and load
index2 = HNSWVectorIndex(dimension=384)
index2.load_index(temp_db_path)
assert index2._max_entries == 50
assert index2._eviction_batch_size == 10
assert index2.size == 5