Compare commits
1 Commits
fix-update
...
feat/diska
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
fcbcde1ea8 |
@@ -7,6 +7,7 @@ from pathlib import Path
|
|||||||
from typing import Any, Literal
|
from typing import Any, Literal
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
import psutil
|
||||||
from leann.interface import (
|
from leann.interface import (
|
||||||
LeannBackendBuilderInterface,
|
LeannBackendBuilderInterface,
|
||||||
LeannBackendFactoryInterface,
|
LeannBackendFactoryInterface,
|
||||||
@@ -84,6 +85,43 @@ def _write_vectors_to_bin(data: np.ndarray, file_path: Path):
|
|||||||
f.write(data.tobytes())
|
f.write(data.tobytes())
|
||||||
|
|
||||||
|
|
||||||
|
def _calculate_smart_memory_config(data: np.ndarray) -> tuple[float, float]:
|
||||||
|
"""
|
||||||
|
Calculate smart memory configuration for DiskANN based on data size and system specs.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
data: The embedding data array
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
tuple: (search_memory_maximum, build_memory_maximum) in GB
|
||||||
|
"""
|
||||||
|
num_vectors, dim = data.shape
|
||||||
|
|
||||||
|
# Calculate embedding storage size
|
||||||
|
embedding_size_bytes = num_vectors * dim * 4 # float32 = 4 bytes
|
||||||
|
embedding_size_gb = embedding_size_bytes / (1024**3)
|
||||||
|
|
||||||
|
# search_memory_maximum: 1/10 of embedding size for optimal PQ compression
|
||||||
|
# This controls Product Quantization size - smaller means more compression
|
||||||
|
search_memory_gb = max(0.1, embedding_size_gb / 10) # At least 100MB
|
||||||
|
|
||||||
|
# build_memory_maximum: Based on available system RAM for sharding control
|
||||||
|
# This controls how much memory DiskANN uses during index construction
|
||||||
|
available_memory_gb = psutil.virtual_memory().available / (1024**3)
|
||||||
|
total_memory_gb = psutil.virtual_memory().total / (1024**3)
|
||||||
|
|
||||||
|
# Use 50% of available memory, but at least 2GB and at most 75% of total
|
||||||
|
build_memory_gb = max(2.0, min(available_memory_gb * 0.5, total_memory_gb * 0.75))
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
f"Smart memory config - Data: {embedding_size_gb:.2f}GB, "
|
||||||
|
f"Search mem: {search_memory_gb:.2f}GB (PQ control), "
|
||||||
|
f"Build mem: {build_memory_gb:.2f}GB (sharding control)"
|
||||||
|
)
|
||||||
|
|
||||||
|
return search_memory_gb, build_memory_gb
|
||||||
|
|
||||||
|
|
||||||
@register_backend("diskann")
|
@register_backend("diskann")
|
||||||
class DiskannBackend(LeannBackendFactoryInterface):
|
class DiskannBackend(LeannBackendFactoryInterface):
|
||||||
@staticmethod
|
@staticmethod
|
||||||
@@ -121,6 +159,16 @@ class DiskannBuilder(LeannBackendBuilderInterface):
|
|||||||
f"Unsupported distance_metric '{build_kwargs.get('distance_metric', 'unknown')}'."
|
f"Unsupported distance_metric '{build_kwargs.get('distance_metric', 'unknown')}'."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Calculate smart memory configuration if not explicitly provided
|
||||||
|
if (
|
||||||
|
"search_memory_maximum" not in build_kwargs
|
||||||
|
or "build_memory_maximum" not in build_kwargs
|
||||||
|
):
|
||||||
|
smart_search_mem, smart_build_mem = _calculate_smart_memory_config(data)
|
||||||
|
else:
|
||||||
|
smart_search_mem = build_kwargs.get("search_memory_maximum", 4.0)
|
||||||
|
smart_build_mem = build_kwargs.get("build_memory_maximum", 8.0)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from . import _diskannpy as diskannpy # type: ignore
|
from . import _diskannpy as diskannpy # type: ignore
|
||||||
|
|
||||||
@@ -131,8 +179,8 @@ class DiskannBuilder(LeannBackendBuilderInterface):
|
|||||||
index_prefix,
|
index_prefix,
|
||||||
build_kwargs.get("complexity", 64),
|
build_kwargs.get("complexity", 64),
|
||||||
build_kwargs.get("graph_degree", 32),
|
build_kwargs.get("graph_degree", 32),
|
||||||
build_kwargs.get("search_memory_maximum", 4.0),
|
build_kwargs.get("search_memory_maximum", smart_search_mem),
|
||||||
build_kwargs.get("build_memory_maximum", 8.0),
|
build_kwargs.get("build_memory_maximum", smart_build_mem),
|
||||||
build_kwargs.get("num_threads", 8),
|
build_kwargs.get("num_threads", 8),
|
||||||
build_kwargs.get("pq_disk_bytes", 0),
|
build_kwargs.get("pq_disk_bytes", 0),
|
||||||
"",
|
"",
|
||||||
|
|||||||
Reference in New Issue
Block a user