[BUG] CAGRA index fails to dealloc on last versions
Nobody has claimed this yet.
- Dominant language
- Cuda
- Stars
- 854
- Forks
- 236
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 62
Description
Describe the bug
When running simple code doing vector search with CAGRA the code succeeds finding the neighbors but raises dealloc exception when it finishes:
... TypeError: 'NoneType' object is not callable Exception ignored in: 'cuvs.neighbors.cagra.cagra.Index.__dealloc__' TypeError: 'NoneType' object is not callable
Steps/Code to reproduce bug
Here is the MWE causing the issue:
import numpy as np
import torch
from typing import Dict, Union, List, Tuple, Optional, Any
import pickle
import time
import os
from tqdm import tqdm
class CAGRAVectorSearch:
"""
GPU-accelerated vector search using NVIDIA's CAGRA algorithm from cuVS.
Provides a similar interface to HNSWVectorSearch but with GPU acceleration.
"""
def __init__(
self,
vector_dim: int = None,
cache_dir: str = None,
device_id: int = 0,
use_float16: bool = False,
):
"""
Initialize the GPU-accelerated vector search engine using CAGRA from cuVS.
Args:
vector_dim (int, optional): Dimension of the vectors. Can be inferred when adding vectors.
cache_dir (str, optional): Directory to store cache files. If None, no caching is used.
device_id (int): GPU device ID to use
use_float16 (bool): Whether to use FP16 for faster search (with slightly lower precision)
"""
# Import cuVS at initialization time to avoid dependency issues
try:
# The correct import pattern based on documentation
import cuvs.neighbors.cagra as cagra
self.cagra = cagra
self._has_cuvs = True
print("Successfully imported cuVS CAGRA")
except ImportError:
print("WARNING: cuVS not installed. Falling back to torch for GPU operations.")
print("For best performance, install cuVS: pip install cuvs-cu12 or https://github.com/rapidsai/cuvs")
self._has_cuvs = False
self.vector_dim = vector_dim
self.device_id = device_id
self.use_float16 = use_float16
self.index = None
self.vectors = None
self.labels = []
self.label_to_indices = {}
self._is_fitted = False
# Tensor storage on GPU for faster operations
self._gpu_vectors = None
self._gpu_data_type = torch.float16 if use_float16 else torch.float32
# Caching system
self.cache_dir = cache_dir
if self.cache_dir and not os.path.exists(self.cache_dir):
os.makedirs(self.cache_dir)
# Cache for label statistics (centroids, radii)
self._centroids_cache = {}
self._radii_cache = {}
self._cache_dirty = False # Flag to track if cache needs updating
def add_vectors(
self, vectors_dict: Dict[str, Union[np.ndarray, torch.Tensor]]
) -> None:
"""
Add vectors to the search index from a dictionary.
Args:
vectors_dict: Dictionary mapping labels to vectors (numpy arrays or torch tensors)
"""
all_vectors = []
all_labels = []
label_to_indices = {}
start_idx = 0
for label, vectors in vectors_dict.items():
# Convert to numpy array
if isinstance(vectors, torch.Tensor):
# Move tensor to CPU if on GPU
if vectors.is_cuda:
vectors = vectors.detach().cpu()
vectors_np = vectors.numpy()
elif isinstance(vectors, np.ndarray):
vectors_np = vectors
else:
raise TypeError(f"Unsupported vector type: {type(vectors)}")
# Infer dimension if not specified
if self.vector_dim is None and vectors_np.shape[0] > 0:
self.vector_dim = vectors_np.shape[1]
# Track indices for this label
end_idx = start_idx + len(vectors_np)
label_to_indices[label] = list(range(start_idx, end_idx))
start_idx = end_idx
all_vectors.append(vectors_np)
all_labels.extend([label] * len(vectors_np))
# Concatenate all vectors
if all_vectors:
self.vectors = np.vstack(all_vectors)
else:
self.vectors = np.array([])
self.labels = all_labels
self.label_to_indices = label_to_indices
self._is_fitted = False
# Clear GPU storage since data has changed
self._gpu_vectors = None
# Mark cache as dirty
self._cache_dirty = True
self._centroids_cache = {}
self._radii_cache = {}
def add_vector(self, vector: Union[np.ndarray, torch.Tensor], label: str) -> None:
"""
Add a single vector to the search index.
Args:
vector: Vector as numpy array or torch tensor
label: Label for the vector
"""
# Convert to numpy array
if isinstance(vector, torch.Tensor):
# Move tensor to CPU if on GPU
if vector.is_cuda:
vector = vector.detach().cpu()
vector_np = vector.numpy()
elif isinstance(vector, np.ndarray):
vector_np = vector
else:
raise TypeError(f"Unsupported vector type: {type(vector)}")
# Reshape to ensure 2D
if vector_np.ndim == 1:
vector_np = vector_np.reshape(1, -1)
# Infer dimension if not specified
if self.vector_dim is None:
self.vector_dim = vector_np.shape[1]
# Initialize vectors if empty
if self.vectors is None or len(self.vectors) == 0:
self.vectors = vector_np
else:
self.vectors = np.vstack([self.vectors, vector_np])
# Update labels and indices
new_idx = len(self.labels)
self.labels.append(label)
if label in self.label_to_indices:
self.label_to_indices[label].append(new_idx)
# Invalidate cached data for this label
if label in self._centroids_cache:
del self._centroids_cache[label]
if label in self._radii_cache:
del self._radii_cache[label]
else:
self.label_to_indices[label] = [new_idx]
# Mark as not fitted since data has changed
self._is_fitted = False
self._gpu_vectors = None
# Mark cache as dirty
self._cache_dirty = True
def fit(
self,
n_lists: int = 128,
metric: str = "sqeuclidean",
build_algo: str = "ivf_pq",
graph_degree: int = 64,
intermediate_graph_degree: int = 128,
precalculate_stats: bool = True,
) -> None:
"""
Build the CAGRA index for fast GPU search.
Args:
n_lists: Number of partitions for the index (higher = more accurate but slower build)
metric: Distance metric ('sqeuclidean', 'inner_product')
build_algo: Algorithm for building the graph ('ivf_pq', 'nn_descent')
graph_degree: Degree of the proximity graph (higher = more accurate but more memory)
intermediate_graph_degree: Degree during graph construction
precalculate_stats: Whether to precalculate label statistics after fitting
"""
if self.vectors is None or len(self.vectors) == 0:
raise ValueError("No vectors added yet. Call add_vectors() first.")
if self.vector_dim is None:
raise ValueError("Vector dimension not set and could not be inferred.")
if self._has_cuvs:
try:
import cupy as cp
print(f"Building CAGRA index with {len(self.vectors)} vectors...")
start_time = time.time()
# Convert vectors to CuPy array on GPU
try:
# Try to create a CuPy array directly
gpu_vectors = cp.array(self.vectors, dtype=np.float32)
except Exception:
# Fallback to going through PyTorch
vectors_tensor = torch.tensor(self.vectors, dtype=self._gpu_data_type)
self._gpu_vectors = vectors_tensor.cuda(self.device_id)
# Convert to CuPy array
gpu_vectors = cp.asarray(self._gpu_vectors.detach())
# Create index parameters according to documentation
index_params = self.cagra.IndexParams(
metric=metric,
intermediate_graph_degree=intermediate_graph_degree,
graph_degree=graph_degree,
build_algo=build_algo
)
# Build index
self.index = self.cagra.build(index_params, gpu_vectors)
# Store vectors on GPU for later use
self._gpu_vectors = torch.tensor(
self.vectors, dtype=self._gpu_data_type, device=f"cuda:{self.device_id}"
)
print(f"CAGRA index built in {time.time() - start_time:.2f} seconds")
except Exception as e:
print(f"Error building CAGRA index: {e}")
print("Falling back to PyTorch implementation")
self._has_cuvs = False
if not self._has_cuvs:
# Fallback to PyTorch implementation
print("Building PyTorch-based GPU index (not as fast as CAGRA)...")
start_time = time.time()
# Keep vectors on GPU for search
self._gpu_vectors = torch.tensor(
self.vectors, dtype=self._gpu_data_type, device=f"cuda:{self.device_id}"
)
# Store the metric type for later use in search
self._metric = metric
# We don't build an actual index, but we'll perform brute-force search with PyTorch
print(f"PyTorch index prepared in {time.time() - start_time:.2f} seconds")
self._is_fitted = True
# Precalculate label statistics if requested
if precalculate_stats:
self.precalculate_label_stats(metric=metric)
def search(
self,
query_vector: Union[np.ndarray, torch.Tensor],
k: int = 5,
exclude_label: Optional[str] = None,
) -> Tuple[List[str], List[float], float]:
"""
Search for the k nearest neighbors to a query vector using GPU.
Args:
query_vector: Query vector as numpy array or torch tensor
k: Number of nearest neighbors to return
exclude_label: Label to exclude from results
Returns:
Tuple containing:
- List of labels for the nearest neighbors
- List of distances to the nearest neighbors
- Mean distance to the nearest neighbors
"""
if not self._is_fitted:
self.fit()
# Convert query to appropriate format for search
if isinstance(query_vector, np.ndarray):
if query_vector.ndim == 2 and query_vector.shape[0] == 1:
query_vector = query_vector.flatten()
elif query_vector.ndim > 1:
raise ValueError(
f"Query must be a single vector, got shape {query_vector.shape}"
)
query_tensor = torch.tensor(
query_vector, dtype=self._gpu_data_type, device=f"cuda:{self.device_id}"
)
elif isinstance(query_vector, torch.Tensor):
if query_vector.dim() == 2 and query_vector.shape[0] == 1:
query_vector = query_vector.flatten()
elif query_vector.dim() > 1:
raise ValueError(
f"Query must be a single vector, got shape {query_vector.shape}"
)
query_tensor = query_vector.to(
device=f"cuda:{self.device_id}", dtype=self._gpu_data_type
)
else:
raise TypeError(f"Unsupported query type: {type(query_vector)}")
# Calculate how many extra neighbors to retrieve in case of filtering
extra = 20 if exclude_label is not None else 0
n_neighbors = min(k + extra, len(self.vectors))
# Search using appropriate method
if self._has_cuvs and self.index is not None:
try:
import cupy as cp
# Reshape the query for batch search (CAGRA expects batch)
query_cp = cp.asarray(query_tensor.reshape(1, -1).detach())
# Create search parameters
search_params = self.cagra.SearchParams(
itopk_size=128, # Increase for better accuracy
)
# Perform search according to documentation
distances, indices = self.cagra.search(
search_params,
self.index,
query_cp,
k=n_neighbors
)
# Convert to numpy arrays
distances = cp.asnumpy(distances).flatten()
indices = cp.asnumpy(indices).flatten()
except Exception as e:
print(f"Error during CAGRA search: {e}")
print("Falling back to PyTorch brute force search")
# Fallback to PyTorch
distances, indices = self._search_pytorch(query_tensor, n_neighbors)
else:
# Use PyTorch-based search
distances, indices = self._search_pytorch(query_tensor, n_neighbors)
# Convert indices to labels
neighbor_labels = [self.labels[i] for i in indices]
# Filter by label if needed
if exclude_label is not None:
filtered_results = [
(d, l) for d, l in zip(distances, neighbor_labels) if l != exclude_label
]
if not filtered_results:
return [], [], None
# Take the k nearest with different labels
filtered_results = filtered_results[:k]
final_distances = [d for d, _ in filtered_results]
final_labels = [l for _, l in filtered_results]
else:
# Just take the first k
final_distances = distances[:k].tolist()
final_labels = neighbor_labels[:k]
# Calculate mean distance
mean_distance = float(np.mean(final_distances))
return final_labels, final_distances, mean_distance
def _search_pytorch(self, query_tensor, n_neighbors):
"""
Perform brute force search using PyTorch operations.
Args:
query_tensor: PyTorch tensor on GPU
n_neighbors: Number of neighbors to retrieve
Returns:
Tuple of (distances, indices)
"""
query_norm = query_tensor.reshape(1, -1)
# Check if metric is defined
metric = getattr(self, "_metric", "sqeuclidean")
if metric in ["cosine", "inner_product"]:
# Normalize vectors for cosine similarity
query_norm = torch.nn.functional.normalize(query_norm, p=2, dim=1)
vectors_norm = torch.nn.functional.normalize(
self._gpu_vectors, p=2, dim=1
)
# Calculate similarities (dot products)
similarities = torch.mm(query_norm, vectors_norm.t())
if metric == "cosine":
# Convert to distances (1 - similarity) for cosine
distances = 1.0 - similarities
else:
# For inner product, smaller values are "farther" so negate
distances = -similarities
else: # "sqeuclidean" / "l2"
# Calculate squared euclidean distance
diff = query_norm.unsqueeze(1) - self._gpu_vectors.unsqueeze(0)
distances = torch.sum(diff ** 2, dim=2)
# Get top-k
topk_distances, topk_indices = torch.topk(
distances.flatten(), k=n_neighbors, largest=False
)
return topk_distances.cpu().numpy(), topk_indices.cpu().numpy()
def get_mean_distance(
self,
query_vector: Union[np.ndarray, torch.Tensor],
k: int = 5,
exclude_label: Optional[str] = None,
) -> float:
"""
Get the mean distance to the k nearest neighbors.
Args:
query_vector: Query vector
k: Number of neighbors
exclude_label: Label to exclude
Returns:
Mean distance to the k nearest neighbors
"""
_, _, mean_distance = self.search(query_vector, k, exclude_label)
return mean_distance
def precalculate_label_stats(
self, metric: str = "sqeuclidean", batch_size: int = 4096, verbose: bool = True
) -> None:
"""
Precalculate and cache statistics for all labels (centroids and radii).
Uses GPU for fast calculation.
Args:
metric: Distance metric to use ('sqeuclidean', 'inner_product')
batch_size: Process vectors in batches of this size
verbose: Whether to print progress information
"""
start_time = time.time()
if verbose:
print(
f"Precalculating label statistics for {len(set(self.labels))} unique labels..."
)
# Try to load from cache if available
if self.cache_dir and not self._cache_dirty:
cache_file = os.path.join(self.cache_dir, f"label_stats_cache_{metric}.pkl")
if os.path.exists(cache_file):
try:
with open(cache_file, "rb") as f:
cache_data = pickle.load(f)
self._centroids_cache = cache_data["centroids"]
self._radii_cache = cache_data["radii"]
if verbose:
print(
f"Loaded label statistics from cache ({time.time() - start_time:.2f}s)"
)
self._cache_dirty = False
return
except Exception as e:
if verbose:
print(f"Error loading cache: {e}. Recalculating...")
# Ensure vectors are on GPU
if self._gpu_vectors is None:
self._gpu_vectors = torch.tensor(
self.vectors, dtype=self._gpu_data_type, device=f"cuda:{self.device_id}"
)
# Calculate centroids first (reused for radius calculation)
if verbose:
print("Calculating centroids...")
self._centroids_cache = {}
centroids_gpu = {}
for label in tqdm(set(self.labels), desc="Computing centroids"):
indices = self.label_to_indices.get(label, [])
if not indices:
continue
# Get vectors for this label
label_vectors = self._gpu_vectors[indices]
# Calculate centroid on GPU
centroid = torch.mean(label_vectors, dim=0)
# Store in GPU cache for radius calculation
centroids_gpu[label] = centroid
# Store CPU version in main cache
self._centroids_cache[label] = centroid.cpu().numpy()
# Calculate radii
if verbose:
print("Calculating radii...")
self._radii_cache = {}
# Process each label
for idx, label in enumerate(tqdm(set(self.labels), desc="Computing radii")):
# Get vectors and centroid for this label
indices = self.label_to_indices.get(label, [])
if not indices:
self._radii_cache[label] = 0.0
continue
# Get GPU centroid
centroid = centroids_gpu[label]
# Process in batches to avoid memory issues
max_distance = 0.0
for batch_start in range(0, len(indices), batch_size):
batch_indices = indices[batch_start : batch_start + batch_size]
batch_vectors = self._gpu_vectors[batch_indices]
# Calculate distances based on metric
if metric == "inner_product":
# Inner product (negative values are farther)
similarities = torch.matmul(batch_vectors, centroid)
distances = -similarities
else: # default to sqeuclidean
# Calculate squared differences
diff = batch_vectors - centroid
distances = torch.sum(diff * diff, dim=1)
# Find maximum distance in this batch
batch_max = torch.max(distances).item()
max_distance = max(max_distance, batch_max)
# Store radius
self._radii_cache[label] = max_distance
# Cache is now up to date
self._cache_dirty = False
# Save to cache file if cache_dir is specified
if self.cache_dir:
cache_file = os.path.join(self.cache_dir, f"label_stats_cache_{metric}.pkl")
cache_data = {
"centroids": self._centroids_cache,
"radii": self._radii_cache,
}
try:
with open(cache_file, "wb") as f:
pickle.dump(cache_data, f)
if verbose:
print(f"Saved label statistics to cache: {cache_file}")
except Exception as e:
if verbose:
print(f"Error saving cache: {e}")
total_time = time.time() - start_time
if verbose:
print(f"Completed label statistics calculation in {total_time:.2f} seconds")
# Clear GPU centroids cache to free up memory
centroids_gpu.clear()
def get_centroid(self, label: str) -> np.ndarray:
"""
Get the centroid of vectors for a specific label.
Args:
label: The label to get centroid for
Returns:
Centroid vector as numpy array
"""
# Check if we have the centroid cached
if label in self._centroids_cache:
return self._centroids_cache[label]
# Calculate centroid if not cached
indices = self.label_to_indices.get(label, [])
if not indices:
raise ValueError(f"Label '{label}' not found or has no vectors")
# Calculate centroid on GPU for efficiency
if self._gpu_vectors is not None:
label_vectors = self._gpu_vectors[indices]
centroid = torch.mean(label_vectors, dim=0).cpu().numpy()
else:
# Fallback to CPU calculation
vectors = self.vectors[indices]
centroid = np.mean(vectors, axis=0)
# Cache the result
self._centroids_cache[label] = centroid
return centroid
def get_radius(self, label: str, metric: str = "sqeuclidean") -> float:
"""
Get the radius of vectors for a specific label.
Uses cached value if available, otherwise calculates it.
Args:
label: The label to get radius for
metric: Distance metric to use ('sqeuclidean', 'inner_product')
Returns:
Radius as float
"""
# Check if we have the radius cached
if label in self._radii_cache:
return self._radii_cache[label]
# Calculate radius if not cached
indices = self.label_to_indices.get(label, [])
if not indices:
return 0.0
# Get or calculate centroid
if label in self._centroids_cache:
centroid_np = self._centroids_cache[label]
else:
centroid_np = self.get_centroid(label)
# Calculate radius on GPU for efficiency
device = f"cuda:{self.device_id}"
centroid = torch.tensor(centroid_np, dtype=self._gpu_data_type, device=device)
# Ensure vectors are on GPU
if self._gpu_vectors is None:
self._gpu_vectors = torch.tensor(
self.vectors, dtype=self._gpu_data_type, device=device
)
# Get vectors for this label
label_vectors = self._gpu_vectors[indices]
# Calculate distances based on metric
if metric == "inner_product":
similarities = torch.matmul(label_vectors, centroid)
distances = -similarities
else: # default to sqeuclidean
diff = label_vectors - centroid
distances = torch.sum(diff * diff, dim=1)
# Get maximum distance
radius = torch.max(distances).item()
# Cache the result
self._radii_cache[label] = radius
return radius
def radius_by_label(
self,
label: str = None,
return_all: bool = False,
metric: str = "cosine",
force_recalculate: bool = False,
) -> Union[float, Dict[str, float]]:
"""
Get the radius of a cluster - the maximum distance from the centroid to any vector with the given label.
Uses cached values for fast retrieval.
Args:
label: The label to calculate radius for. If None and return_all is True, calculate for all labels.
return_all: If True, return radii for all labels as a dictionary.
metric: Distance metric to use ('cosine', 'l2', 'ip').
force_recalculate: If True, force recalculation even if cached values exist.
Returns:
If return_all is False: The radius for the specified label.
If return_all is True: Dictionary mapping labels to their radii.
"""
# If cache is dirty or we're forcing recalculation, update it
if self._cache_dirty or force_recalculate or not self._radii_cache:
self.precalculate_label_stats(metric=metric)
if return_all:
return self._radii_cache.copy()
if label is None:
raise ValueError("Must specify either 'label' or set 'return_all=True'")
if label not in self.label_to_indices:
raise ValueError(f"Label '{label}' not found in the database")
# Get cached radius or calculate if needed
return self.get_radius(label, metric=metric)
def save(self, filepath: str, save_cache: bool = True) -> None:
"""
Save the vector database to disk.
Args:
filepath: Path to save the database
save_cache: Whether to save cached label statistics
"""
# Store everything except GPU data
data = {
"vector_dim": self.vector_dim,
"vectors": self.vectors,
"labels": self.labels,
"label_to_indices": self.label_to_indices,
"is_fitted": self._is_fitted,
"cache_dir": self.cache_dir,
"device_id": self.device_id,
"use_float16": self.use_float16,
}
# Include cache if requested
if save_cache:
data["centroids_cache"] = self._centroids_cache
data["radii_cache"] = self._radii_cache
data["cache_dirty"] = self._cache_dirty
with open(filepath, "wb") as f:
pickle.dump(data, f)
# Save CAGRA index if available
if self._has_cuvs and self.index is not None:
try:
index_filepath = filepath + ".cagra_index"
self.cagra.save(index_filepath, self.index)
print(f"CAGRA index saved to {index_filepath}")
except Exception as e:
print(f"Error saving CAGRA index: {e}")
print("Only vector data was saved")
def load(
self, filepath: str, load_cache: bool = True, rebuild_index: bool = True
) -> None:
"""
Load a vector database from disk.
Args:
filepath: Path to the saved database
load_cache: Whether to load cached label statistics
rebuild_index: Whether to rebuild the index after loading
"""
with open(filepath, "rb") as f:
data = pickle.load(f)
self.vector_dim = data["vector_dim"]
self.vectors = data["vectors"]
self.labels = data["labels"]
self.label_to_indices = data["label_to_indices"]
self._is_fitted = data["is_fitted"]
# Load device settings
if "device_id" in data:
self.device_id = data["device_id"]
if "use_float16" in data:
self.use_float16 = data["use_float16"]
# Load cache directory
if "cache_dir" in data:
self.cache_dir = data["cache_dir"]
# Load cache if available and requested
if load_cache:
if "centroids_cache" in data:
self._centroids_cache = data["centroids_cache"]
if "radii_cache" in data:
self._radii_cache = data["radii_cache"]
if "cache_dirty" in data:
self._cache_dirty = data["cache_dirty"]
# Try to load CAGRA index
if self._has_cuvs:
try:
index_filepath = filepath + ".cagra_index"
if os.path.exists(index_filepath):
self.index = self.cagra.load(index_filepath)
self._is_fitted = True
# Convert vectors to GPU
self._gpu_vectors = torch.tensor(
self.vectors, dtype=self._gpu_data_type, device=f"cuda:{self.device_id}"
)
print(f"CAGRA index loaded from {index_filepath}")
return
except Exception as e:
print(f"Error loading CAGRA index: {e}")
print("Will rebuild index")
self.index = None
self._is_fitted = False
# Rebuild the index if requested or if loading the index failed
if rebuild_index and (not self._is_fitted or self.index is None):
self._is_fitted = False # Force rebuild
self.fit()
else:
# Ensure GPU vectors are loaded
self._gpu_vectors = None # Will be loaded on demand
def __len__(self) -> int:
"""Return the number of vectors in the database."""
return len(self.labels) if self.labels else 0
def calibrate_from_parquet(
parquet_file: str, cache_dir: str = None, use_gpu: bool = True
) -> Any:
"""
Load vectors from a parquet file and build a search index.
Args:
parquet_file: Path to the parquet file
cache_dir: Directory to store cache files
use_gpu: Whether to use GPU-accelerated search
Returns:
Search index (CAGRAVectorSearch if use_gpu=True, otherwise HNSWVectorSearch)
"""
import pandas as pd
print(f"Loading vectors from {parquet_file}...")
df = pd.read_parquet(parquet_file)
# Prepare vectors dictionary
vectors_dict = {}
for _, row in df.iterrows():
label = row["label"]
vector = row["vector"]
if label not in vectors_dict:
vectors_dict[label] = []
vectors_dict[label].append(vector)
# Convert lists to numpy arrays
for label in vectors_dict:
vectors_dict[label] = np.vstack(vectors_dict[label])
print(
f"Loaded {sum(len(v) for v in vectors_dict.values())} vectors for {len(vectors_dict)} labels"
)
# Create appropriate search index
if use_gpu:
vs = CAGRAVectorSearch(cache_dir=cache_dir)
else:
from corsound_vfm.inference.adaptive_thresholds.knn_signatures_cpu import (
HNSWVectorSearch,
)
vs = HNSWVectorSearch(cache_dir=cache_dir)
# Add vectors and build index
vs.add_vectors(vectors_dict)
vs.fit(precalculate_stats=True)
return vs
if __name__ == '__main__':
# MWE
# Create a random dataset
np.random.seed(42)
num_vectors = 100
vector_dim = 1536
num_labels = 5
vectors = np.random.rand(num_vectors, vector_dim).astype(np.float32)
labels = [f"label_{i % num_labels}" for i in range(num_vectors)]
vectors_dict = {}
for label in set(labels):
vectors_dict[label] = vectors[np.array(labels) == label]
# Save to parquet
import pandas as pd
df = pd.DataFrame({"label": labels, "vector": list(vectors)})
df.to_parquet("vectors.parquet")
# Load and build index
vs = calibrate_from_parquet("vectors.parquet", use_gpu=True)
# Perform a search
query_vector = np.random.rand(vector_dim).astype(np.float32)
k = 5
exclude_label = "label_0"
labels, distances, mean_distance = vs.search(query_vector, k, exclude_label)
print(f"Labels: {labels}")
print(f"Distances: {distances}")
print(f"Mean distance: {mean_distance}")
# Save the index
vs.save("vector_index.pkl")
# Load the index
vs_loaded = CAGRAVectorSearch()
vs_loaded.load("vector_index.pkl")
# Perform a search on the loaded index
labels, distances, mean_distance = vs_loaded.search(query_vector, k, exclude_label)
print(f"Labels (loaded): {labels}")
print(f"Distances (loaded): {distances}")
print(f"Mean distance (loaded): {mean_distance}")
Expected behavior
To dealloc successfully.
Environment details (please complete the following information):
- Environment location: Bare-metal
- Method of RAFT install: conda
Additional context
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reproducing the provided Python MWE with the CAGRA index and observe cleanup after the search completes. The issue names no repository file or test, so locate the cuvs.neighbors.cagra.Index lifecycle and its dealloc path; done means the search succeeds without a deallocation exception.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, search
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100