cuvs.neighbors.mg.cagra

cuvs.neighbors.mg.cagra#

8 min read time

Applies to Linux

Submodules#

cuvs.neighbors.mg.cagra.cagra

Classes#

`Index`	Multi-GPU CAGRA index object. Stores the trained multi-GPU CAGRA index
`IndexParams`	IndexParams(distribution_mode=u'sharded', , *kwargs)
`SearchParams`	SearchParams(search_mode=u'load_balancer', , merge_mode=u'merge_on_root_rank', n_rows_per_batch=1000, *kwargs)

Functions#

`build`(*args[, resources])	build(IndexParams index_params, dataset, resources=None)
`distribute`(*args[, resources])	distribute(filename, resources=None)
`extend`(*args[, resources])	extend(Index index, new_vectors, new_indices=None, resources=None)
`load`(*args[, resources])	load(filename, resources=None)
`save`(*args[, resources])	save(Index index, filename, resources=None)
`search`(*args[, resources])	search(SearchParams search_params, Index index, queries, k, neighbors=None, distances=None, resources=None)

Package Contents#

class cuvs.neighbors.mg.cagra.Index#

Multi-GPU CAGRA index object. Stores the trained multi-GPU CAGRA index state which can be used to perform nearest neighbors searches across multiple GPUs.

static __reduce__(*args, **kwargs)#: Index.__reduce_cython__(self)

static __repr__(*args, **kwargs)#

static __setstate__(*args, **kwargs)#: Index.__setstate_cython__(self, __pyx_state)

class cuvs.neighbors.mg.cagra.IndexParams#

Bases: cuvs.neighbors.cagra.cagra.IndexParams

IndexParams(distribution_mode=u’sharded’, *, **kwargs)

Parameters to build multi-GPU CAGRA index for efficient search. Extends single-GPU IndexParams with multi-GPU specific parameters.

distribution_modestr, default = “sharded”
Distribution mode for multi-GPU setup. Valid values: [“replicated”, “sharded”]

**kwargs : Additional parameters passed to single-GPU IndexParams

CAGRA currently only supports “sqeuclidean” and “inner_product” metrics.

static __reduce__(*args, **kwargs)#: IndexParams.__reduce_cython__(self)

static __setstate__(*args, **kwargs)#: IndexParams.__setstate_cython__(self, __pyx_state)

get_handle()#: IndexParams.get_handle(self)

class cuvs.neighbors.mg.cagra.SearchParams#

Bases: cuvs.neighbors.cagra.cagra.SearchParams

SearchParams(search_mode=u’load_balancer’, *, merge_mode=u’merge_on_root_rank’, n_rows_per_batch=1000, **kwargs)

Parameters to search multi-GPU CAGRA index.

static __reduce__(*args, **kwargs)#: SearchParams.__reduce_cython__(self)

static __setstate__(*args, **kwargs)#: SearchParams.__setstate_cython__(self, __pyx_state)

get_handle()#: SearchParams.get_handle(self)

cuvs.neighbors.mg.cagra.build(*args, resources=None, **kwargs)#

build(IndexParams index_params, dataset, resources=None)

Build the multi-GPU CAGRA index from the dataset for efficient search.

index_params : cuvs.neighbors.cagra.IndexParams dataset : Array interface compliant matrix shape (n_samples, dim)

Supported dtype [float32, float16, int8, uint8] IMPORTANT: For multi-GPU CAGRA, the dataset MUST be in host memory (CPU). If using CuPy/device arrays, transfer to host with array.get() or cp.asnumpy(array).

resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.
If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.

index: py:class:cuvs.neighbors.cagra.Index
>>> import numpy as np
>>> from cuvs.neighbors.mg import cagra
>>> n_samples = 50000
>>> n_features = 50
>>> n_queries = 1000
>>> k = 10
>>> # For multi-GPU CAGRA, use host (NumPy) arrays
>>> dataset = np.random.random_sample((n_samples, n_features)).astype(
...     np.float32)
>>> build_params = cagra.IndexParams(metric="sqeuclidean")
>>> index = cagra.build(build_params, dataset)
>>> distances, neighbors = cagra.search(cagra.SearchParams(),
...                                         index, dataset, k)
>>> # Results are already in host memory (NumPy arrays)

cuvs.neighbors.mg.cagra.distribute(*args, resources=None, **kwargs)#

distribute(filename, resources=None)

Distribute a single-GPU CAGRA index across multiple GPUs from a file.

filenamestr
The filename to distribute the index from.

resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.
If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.

indexIndex
The distributed index.
>>> from cuvs.neighbors.mg import cagra
>>> index = cagra.distribute("single_gpu_index.bin")

cuvs.neighbors.mg.cagra.extend(*args, resources=None, **kwargs)#

extend(Index index, new_vectors, new_indices=None, resources=None)

Extend the multi-GPU CAGRA index with new vectors.

index : cuvs.neighbors.cagra.Index new_vectors : Array interface compliant matrix shape (n_new_vectors, dim)

Supported dtype [float32, float16, int8, uint8] IMPORTANT: For multi-GPU CAGRA, new_vectors MUST be in host memory (CPU). If using CuPy/device arrays, transfer to host with array.get() or cp.asnumpy(array).

new_indicesArray interface compliant matrix shape (n_new_vectors,),

optional

If provided, these indices will be used for the new vectors. If not provided, indices will be automatically assigned. IMPORTANT: Must be in host memory (CPU) for multi-GPU CAGRA. Expected dtype: uint32

resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.
If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.
>>> import numpy as np
>>> from cuvs.neighbors.mg import cagra
>>> n_samples = 50000
>>> n_features = 50
>>> n_new_vectors = 1000
>>> # For multi-GPU CAGRA, use host (NumPy) arrays
>>> dataset = np.random.random_sample((n_samples, n_features)).astype(
...     np.float32)
>>> new_vectors = np.random.random_sample(
...     (n_new_vectors, n_features)).astype(np.float32)
>>> new_indices = np.arange(n_samples, n_samples + n_new_vectors,
...                         dtype=np.uint32)
>>> build_params = cagra.IndexParams(metric="sqeuclidean")
>>> index = cagra.build(build_params, dataset)
>>> cagra.extend(index, new_vectors, new_indices)

cuvs.neighbors.mg.cagra.load(*args, resources=None, **kwargs)#

load(filename, resources=None)

Deserialize the multi-GPU CAGRA index from a file.

filenamestr
The filename to deserialize the index from.

resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.
If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.

indexIndex
The deserialized index.
>>> from cuvs.neighbors.mg import cagra
>>> index = cagra.load("index.bin")

cuvs.neighbors.mg.cagra.save(*args, resources=None, **kwargs)#

save(Index index, filename, resources=None)

Serialize the multi-GPU CAGRA index to a file.

index : cuvs.neighbors.cagra.Index filename : str

The filename to serialize the index to.

resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.
If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.
>>> import numpy as np
>>> from cuvs.neighbors.mg import cagra
>>> n_samples = 50000
>>> n_features = 50
>>> # For multi-GPU CAGRA, use host (NumPy) arrays
>>> dataset = np.random.random_sample((n_samples, n_features)).astype(
...     np.float32)
>>> build_params = cagra.IndexParams(metric="sqeuclidean")
>>> index = cagra.build(build_params, dataset)
>>> cagra.save(index, "index.bin")

cuvs.neighbors.mg.cagra.search(*args, resources=None, **kwargs)#

search(SearchParams search_params, Index index, queries, k, neighbors=None, distances=None, resources=None)

Search the multi-GPU CAGRA index for the k-nearest neighbors of each query.

search_params : cuvs.neighbors.cagra.SearchParams index : cuvs.neighbors.cagra.Index queries : Array interface compliant matrix shape (n_queries, dim)

Supported dtype [float32, float16, int8, uint8] IMPORTANT: For multi-GPU CAGRA, queries MUST be in host memory (CPU). If using CuPy/device arrays, transfer to host with array.get() or cp.asnumpy(array).

kint
The number of neighbors to search for each query.

neighborsArray interface compliant matrix shape (n_queries, k), optional
If provided, this array will be filled with the indices of the k-nearest neighbors. If not provided, a new host array will be allocated. IMPORTANT: Must be in host memory (CPU) for multi-GPU CAGRA. Expected dtype: int64

distancesArray interface compliant matrix shape (n_queries, k), optional
If provided, this array will be filled with the distances to the k-nearest neighbors. If not provided, a new host array will be allocated. IMPORTANT: Must be in host memory (CPU) for multi-GPU CAGRA.

resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.
If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.

distancesnumpy.ndarray
The distances to the k-nearest neighbors for each query (in host memory).

neighborsnumpy.ndarray
The indices of the k-nearest neighbors for each query (in host memory).
>>> import numpy as np
>>> from cuvs.neighbors.mg import cagra
>>> n_samples = 50000
>>> n_features = 50
>>> n_queries = 1000
>>> k = 10
>>> # For multi-GPU CAGRA, use host (NumPy) arrays
>>> dataset = np.random.random_sample((n_samples, n_features)).astype(
...     np.float32)
>>> queries = np.random.random_sample((n_queries, n_features)).astype(
...     np.float32)
>>> build_params = cagra.IndexParams(metric="sqeuclidean")
>>> index = cagra.build(build_params, dataset)
>>> distances, neighbors = cagra.search(cagra.SearchParams(),
...                                         index, queries, k)
>>> # Results are already in host memory (NumPy arrays)