NVIDIA / NVIDIA/cuvs

[BUG] CAGRA index fails to dealloc on last versions

Open
#907 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.