From defcc60683a3f3fb9cb0360178fcbc357c9aed0c Mon Sep 17 00:00:00 2001 From: Zachary Bennett Date: Tue, 14 Apr 2026 08:00:57 -0500 Subject: [PATCH] feat(rust)!: tie CAGRA and brute-force index lifetime to dataset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Simplify the dataset-lifetime-safety work per review feedback on https://github.com/rapidsai/cuvs/pull/1869: - IVF-Flat and IVF-PQ always copy their input in the C++ library, so their Rust bindings do not need a lifetime parameter and have been left on their upstream signatures. - CAGRA's `make_aligned_dataset()` stores a non-owning view when the input is device-accessible, row-major and 16-byte aligned, so the Rust `Index` now carries a lifetime tied to the `ManagedTensor` it was built from. The borrow checker then prevents use-after-free regardless of which runtime path the C++ side takes. - Brute Force always stores a non-owning view and gets the same lifetime treatment. - `DatasetOwnership` and `build_owned` are dropped — the `ManagedTensor` type itself doesn't express ownership in the type system, so the `Owned` variant was fragile. `Index<'a>` now holds a `PhantomData<&'a ()>` directly, which matches the semantics of the underlying C++ library exactly. - `cagra::Index::deserialize` now returns `Index<'static>`; the deserialized index owns its data and is not tied to any input `ManagedTensor`. `serialize` and `serialize_to_hnswlib` live on `impl<'a> Index<'a>` so they work against both borrowed and deserialized indices. - `pub fn new()` is gone from both `cagra::Index` and `brute_force::Index`. An empty index has no real dataset lifetime and the builders call the private `create_handle` helper directly. - Test suites for both index types have been reorganised so the dataset + self-neighbor search boilerplate is factored into shared helpers, including a negative test that confirms a path containing an interior NUL byte surfaces as `Error::InvalidArgument`. Breaking change: `cagra::Index::build` and `brute_force::Index::build` now take `&ManagedTensor` instead of `impl Into`, and the public `new()` constructors have been removed. IVF-Flat and IVF-PQ are unchanged. Signed-off-by: Zach Bennett --- rust/cuvs/src/brute_force.rs | 183 +++++++++++++++++++++-------------- rust/cuvs/src/cagra/index.rs | 172 ++++++++++++++++++++------------ rust/cuvs/src/cagra/mod.rs | 9 +- 3 files changed, 225 insertions(+), 139 deletions(-) diff --git a/rust/cuvs/src/brute_force.rs b/rust/cuvs/src/brute_force.rs index 1440bb3205..54beac16f4 100644 --- a/rust/cuvs/src/brute_force.rs +++ b/rust/cuvs/src/brute_force.rs @@ -5,18 +5,53 @@ //! Brute Force KNN use std::io::{stderr, Write}; +use std::marker::PhantomData; use crate::distance_type::DistanceType; use crate::dlpack::ManagedTensor; use crate::error::{check_cuvs, Result}; use crate::resources::Resources; -/// Brute Force KNN Index +/// Brute Force KNN Index. +/// +/// The brute force C++ implementation always stores a non-owning view of the +/// dataset, so the index cannot outlive the `ManagedTensor` it was built +/// from. The lifetime parameter `'a` ties the index to that tensor, so the +/// borrow checker catches use-after-free at compile time. +/// +/// # Example +/// +/// ```no_run +/// # use cuvs::{ManagedTensor, Resources}; +/// # use cuvs::brute_force::Index; +/// # use cuvs::distance_type::DistanceType; +/// let res = Resources::new().unwrap(); +/// let arr = ndarray::Array::::zeros((64, 8)); +/// let tensor = ManagedTensor::from(&arr).to_device(&res).unwrap(); +/// let index = Index::build(&res, DistanceType::L2Expanded, None, &tensor).unwrap(); +/// // `tensor` must outlive `index`. +/// ``` #[derive(Debug)] -pub struct Index(ffi::cuvsBruteForceIndex_t); +pub struct Index<'a> { + inner: ffi::cuvsBruteForceIndex_t, + _dataset: PhantomData<&'a ()>, +} -impl Index { - /// Builds a new Brute Force KNN Index from the dataset for efficient search. +impl<'a> Index<'a> { + /// Creates a new empty FFI handle. + fn create_handle() -> Result { + unsafe { + let mut index = std::mem::MaybeUninit::::uninit(); + check_cuvs(ffi::cuvsBruteForceIndexCreate(index.as_mut_ptr()))?; + Ok(index.assume_init()) + } + } + + /// Builds a new Brute Force KNN Index from `dataset`. + /// + /// The compiler enforces that `dataset` outlives the returned index, + /// because the C++ brute force implementation stores a non-owning view + /// of the input and would otherwise dangle. /// /// # Arguments /// @@ -24,36 +59,29 @@ impl Index { /// * `metric` - DistanceType to use for building the index /// * `metric_arg` - Optional value of `p` for Minkowski distances /// * `dataset` - A row-major matrix on either the host or device to index - pub fn build>( + pub fn build( res: &Resources, metric: DistanceType, metric_arg: Option, - dataset: T, - ) -> Result { - let dataset: ManagedTensor = dataset.into(); - let index = Index::new()?; + dataset: &'a ManagedTensor, + ) -> Result> { + let inner = Self::create_handle()?; unsafe { check_cuvs(ffi::cuvsBruteForceBuild( res.0, dataset.as_ptr(), metric, metric_arg.unwrap_or(2.0), - index.0, + inner, ))?; } - Ok(index) + Ok(Index { + inner, + _dataset: PhantomData, + }) } - /// Creates a new empty index - pub fn new() -> Result { - unsafe { - let mut index = std::mem::MaybeUninit::::uninit(); - check_cuvs(ffi::cuvsBruteForceIndexCreate(index.as_mut_ptr()))?; - Ok(Index(index.assume_init())) - } - } - - /// Perform a Nearest Neighbors search on the Index + /// Perform a Nearest Neighbors search on the Index. /// /// # Arguments /// @@ -76,7 +104,7 @@ impl Index { check_cuvs(ffi::cuvsBruteForceSearch( res.0, - self.0, + self.inner, queries.as_ptr(), neighbors.as_ptr(), distances.as_ptr(), @@ -86,9 +114,9 @@ impl Index { } } -impl Drop for Index { +impl Drop for Index<'_> { fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsBruteForceIndexDestroy(self.0) }) { + if let Err(e) = check_cuvs(unsafe { ffi::cuvsBruteForceIndexDestroy(self.inner) }) { write!(stderr(), "failed to call bruteForceIndexDestroy {:?}", e) .expect("failed to write to stderr"); } @@ -99,66 +127,60 @@ impl Drop for Index { mod tests { use super::*; use mark_flaky_tests::flaky; - use ndarray::s; + use ndarray::{s, Array2}; use ndarray_rand::rand_distr::Uniform; use ndarray_rand::RandomExt; - fn test_bfknn(metric: DistanceType) { - let res = Resources::new().unwrap(); - - // Create a new random dataset to index - let n_datapoints = 16; - let n_features = 8; - let dataset_host = - ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); - - let dataset = ManagedTensor::from(&dataset_host).to_device(&res).unwrap(); - - println!("dataset {:#?}", dataset_host); - - // build the brute force index - let index = - Index::build(&res, metric, None, dataset).expect("failed to create brute force index"); - - res.sync_stream().unwrap(); + /// Build a small random host dataset for the brute-force tests. + fn make_dataset(n_datapoints: usize, n_features: usize) -> Array2 { + ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)) + } - // use the first 4 points from the dataset as queries : will test that we get them back - // as their own nearest neighbor - let n_queries = 4; + /// Search `index` for the first `n_queries` rows of `dataset_host` and + /// assert each query is its own nearest neighbor. + fn search_and_verify_self_neighbors( + res: &Resources, + index: &Index<'_>, + dataset_host: &Array2, + n_queries: usize, + k: usize, + ) { let queries = dataset_host.slice(s![0..n_queries, ..]); + let queries = ManagedTensor::from(&queries).to_device(res).unwrap(); - let k = 4; - - println!("queries! {:#?}", queries); - let queries = ManagedTensor::from(&queries).to_device(&res).unwrap(); let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); - let neighbors = ManagedTensor::from(&neighbors_host) - .to_device(&res) - .unwrap(); + let neighbors = ManagedTensor::from(&neighbors_host).to_device(res).unwrap(); let mut distances_host = ndarray::Array::::zeros((n_queries, k)); - let distances = ManagedTensor::from(&distances_host) - .to_device(&res) - .unwrap(); + let distances = ManagedTensor::from(&distances_host).to_device(res).unwrap(); index - .search(&res, &queries, &neighbors, &distances) - .unwrap(); + .search(res, &queries, &neighbors, &distances) + .expect("search failed"); - // Copy back to host memory - distances.to_host(&res, &mut distances_host).unwrap(); - neighbors.to_host(&res, &mut neighbors_host).unwrap(); + distances.to_host(res, &mut distances_host).unwrap(); + neighbors.to_host(res, &mut neighbors_host).unwrap(); res.sync_stream().unwrap(); - println!("distances {:#?}", distances_host); - println!("neighbors {:#?}", neighbors_host); + for i in 0..n_queries { + assert_eq!( + neighbors_host[[i, 0]], + i as i64, + "query {i} should be its own nearest neighbor" + ); + } + } - // nearest neighbors should be themselves, since queries are from the - // dataset - assert_eq!(neighbors_host[[0, 0]], 0); - assert_eq!(neighbors_host[[1, 0]], 1); - assert_eq!(neighbors_host[[2, 0]], 2); - assert_eq!(neighbors_host[[3, 0]], 3); + fn test_bfknn(metric: DistanceType) { + let res = Resources::new().unwrap(); + let dataset_host = make_dataset(16, 8); + let dataset = ManagedTensor::from(&dataset_host).to_device(&res).unwrap(); + + let index = + Index::build(&res, metric, None, &dataset).expect("failed to build brute force index"); + res.sync_stream().unwrap(); + + search_and_verify_self_neighbors(&res, &index, &dataset_host, 4, 4); } /* @@ -173,10 +195,21 @@ mod tests { test_bfknn(DistanceType::L2Expanded); } - // NOTE: brute_force multiple-search test is omitted here because the C++ - // brute_force::index stores a non-owning view into the dataset. Building - // from device data via `build()` drops the ManagedTensor after the call, - // leaving a dangling pointer. A follow-up PR will add dataset lifetime - // enforcement (DatasetOwnership<'a>) to make this safe. - // See: https://github.com/rapidsai/cuvs/issues/1838 + /// Validates that a borrowed-dataset brute-force index can be searched + /// multiple times while the dataset is kept alive by the caller. + #[test] + fn test_brute_force_multiple_searches() { + let res = Resources::new().unwrap(); + let dataset_host = make_dataset(64, 8); + let dataset_device = ManagedTensor::from(&dataset_host).to_device(&res).unwrap(); + + let index = Index::build(&res, DistanceType::L2Expanded, None, &dataset_device) + .expect("failed to build brute force index"); + res.sync_stream().unwrap(); + + for _ in 0..3 { + search_and_verify_self_neighbors(&res, &index, &dataset_host, 4, 4); + } + // `dataset_device` is still alive here — the borrow checker enforces it. + } } diff --git a/rust/cuvs/src/cagra/index.rs b/rust/cuvs/src/cagra/index.rs index 85570d4e8e..d2a28a5d03 100644 --- a/rust/cuvs/src/cagra/index.rs +++ b/rust/cuvs/src/cagra/index.rs @@ -5,6 +5,7 @@ use std::ffi::CString; use std::io::{stderr, Write}; +use std::marker::PhantomData; use std::path::Path; use crate::cagra::{IndexParams, SearchParams}; @@ -12,9 +13,37 @@ use crate::dlpack::ManagedTensor; use crate::error::{check_cuvs, Error, Result}; use crate::resources::Resources; -/// CAGRA ANN Index +/// CAGRA ANN Index. +/// +/// CAGRA's behavior is runtime-dependent: the underlying C++ helper +/// `make_aligned_dataset()` stores a non-owning view of the dataset when the +/// input is device-accessible, row-major and 16-byte aligned, and copies +/// otherwise. Because the non-owning path can't be ruled out at compile time, +/// the returned index carries the lifetime of the `ManagedTensor` used to +/// build it, so the borrow checker prevents use-after-free regardless of which +/// path the C++ side takes. +/// +/// Indices produced by [`Index::deserialize`] own their data and carry a +/// `'static` lifetime, since the serialized file is fully read into memory +/// managed by the C++ library. +/// +/// # Example +/// +/// ```no_run +/// # use cuvs::{ManagedTensor, Resources}; +/// # use cuvs::cagra::{Index, IndexParams}; +/// let res = Resources::new().unwrap(); +/// let arr = ndarray::Array::::zeros((256, 16)); +/// let params = IndexParams::new().unwrap(); +/// let tensor = ManagedTensor::from(&arr); +/// let index = Index::build(&res, ¶ms, &tensor).unwrap(); +/// // `arr` and `tensor` must outlive `index`. +/// ``` #[derive(Debug)] -pub struct Index(ffi::cuvsCagraIndex_t); +pub struct Index<'a> { + inner: ffi::cuvsCagraIndex_t, + _dataset: PhantomData<&'a ()>, +} /// Convert a filesystem path into a `CString` suitable for the cuVS C API, /// returning `Error::InvalidArgument` instead of panicking for paths that are @@ -27,42 +56,48 @@ fn path_to_cstring(path: &Path) -> Result { .map_err(|e| Error::InvalidArgument(format!("path contains an interior NUL byte: {e}"))) } -impl Index { - /// Builds a new Index from the dataset for efficient search. +impl<'a> Index<'a> { + /// Creates a new empty FFI index handle. + fn create_handle() -> Result { + unsafe { + let mut index = std::mem::MaybeUninit::::uninit(); + check_cuvs(ffi::cuvsCagraIndexCreate(index.as_mut_ptr()))?; + Ok(index.assume_init()) + } + } + + /// Builds a new CAGRA Index from `dataset`. + /// + /// The compiler enforces that `dataset` outlives the returned index, so + /// the C++ library's internal view (when it takes the non-owning path) + /// can never dangle. /// /// # Arguments /// /// * `res` - Resources to use /// * `params` - Parameters for building the index /// * `dataset` - A row-major matrix on either the host or device to index - pub fn build>( + pub fn build( res: &Resources, params: &IndexParams, - dataset: T, - ) -> Result { - let dataset: ManagedTensor = dataset.into(); - let index = Index::new()?; + dataset: &'a ManagedTensor, + ) -> Result> { + let inner = Self::create_handle()?; unsafe { check_cuvs(ffi::cuvsCagraBuild( res.0, params.0, dataset.as_ptr(), - index.0, + inner, ))?; } - Ok(index) + Ok(Index { + inner, + _dataset: PhantomData, + }) } - /// Creates a new empty index - pub fn new() -> Result { - unsafe { - let mut index = std::mem::MaybeUninit::::uninit(); - check_cuvs(ffi::cuvsCagraIndexCreate(index.as_mut_ptr()))?; - Ok(Index(index.assume_init())) - } - } - - /// Perform a Approximate Nearest Neighbors search on the Index + /// Perform an Approximate Nearest Neighbors search on the Index. /// /// # Arguments /// @@ -88,7 +123,7 @@ impl Index { check_cuvs(ffi::cuvsCagraSearch( res.0, params.0, - self.0, + self.inner, queries.as_ptr(), neighbors.as_ptr(), distances.as_ptr(), @@ -117,7 +152,7 @@ impl Index { check_cuvs(ffi::cuvsCagraSerialize( res.0, c_filename.as_ptr(), - self.0, + self.inner, include_dataset, )) } @@ -140,36 +175,40 @@ impl Index { check_cuvs(ffi::cuvsCagraSerializeToHnswlib( res.0, c_filename.as_ptr(), - self.0, + self.inner, )) } } +} +impl Index<'static> { /// Load a CAGRA index from file. /// + /// The returned index owns the data read from disk, so it has a `'static` + /// lifetime and isn't tied to any input `ManagedTensor`. + /// /// Experimental, both the API and the serialization format are subject to change. /// /// # Arguments /// /// * `res` - Resources to use /// * `filename` - The path of the file that stores the index - pub fn deserialize>(res: &Resources, filename: P) -> Result { + pub fn deserialize>(res: &Resources, filename: P) -> Result> { let c_filename = path_to_cstring(filename.as_ref())?; - let index = Index::new()?; + let inner = Self::create_handle()?; unsafe { - check_cuvs(ffi::cuvsCagraDeserialize( - res.0, - c_filename.as_ptr(), - index.0, - ))?; + check_cuvs(ffi::cuvsCagraDeserialize(res.0, c_filename.as_ptr(), inner))?; } - Ok(index) + Ok(Index { + inner, + _dataset: PhantomData, + }) } } -impl Drop for Index { +impl Drop for Index<'_> { fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsCagraIndexDestroy(self.0) }) { + if let Err(e) = check_cuvs(unsafe { ffi::cuvsCagraIndexDestroy(self.inner) }) { write!(stderr(), "failed to call cagraIndexDestroy {:?}", e) .expect("failed to write to stderr"); } @@ -179,31 +218,25 @@ impl Drop for Index { #[cfg(test)] mod tests { use super::*; - use ndarray::s; + use ndarray::{s, Array2}; use ndarray_rand::rand_distr::Uniform; use ndarray_rand::RandomExt; const N_DATAPOINTS: usize = 256; const N_FEATURES: usize = 16; - /// Build a small random dataset and a CAGRA index over it. - fn build_test_index( - res: &Resources, - build_params: &IndexParams, - ) -> (ndarray::Array2, Index) { - let dataset = - ndarray::Array::::random((N_DATAPOINTS, N_FEATURES), Uniform::new(0., 1.0)); - let index = Index::build(res, build_params, &dataset).expect("failed to build cagra index"); - (dataset, index) + /// Build a random host dataset for the CAGRA tests. + fn make_dataset() -> Array2 { + ndarray::Array::::random((N_DATAPOINTS, N_FEATURES), Uniform::new(0., 1.0)) } - /// Search the first `n_queries` rows of `dataset` against `index` and - /// assert each query finds itself as the top-1 neighbor. CAGRA search - /// requires queries and outputs to live in device memory. + /// Run a self-neighbor search against `index` using the first `n_queries` + /// rows of `dataset` and assert each query finds itself as its top-1 + /// neighbor. CAGRA search requires queries and outputs on device memory. fn search_and_verify_self_neighbors( res: &Resources, - index: &Index, - dataset: &ndarray::Array2, + index: &Index<'_>, + dataset: &Array2, n_queries: usize, k: usize, ) { @@ -235,7 +268,10 @@ mod tests { fn test_cagra(build_params: IndexParams) { let res = Resources::new().unwrap(); - let (dataset, index) = build_test_index(&res, &build_params); + let dataset = make_dataset(); + let tensor = ManagedTensor::from(&dataset); + let index = + Index::build(&res, &build_params, &tensor).expect("failed to build cagra index"); search_and_verify_self_neighbors(&res, &index, &dataset, 4, 10); } @@ -254,13 +290,16 @@ mod tests { test_cagra(build_params); } - /// Test that an index can be searched multiple times without rebuilding. - /// This validates that `search()` takes `&self` instead of `self`. + /// Validates that an index built against a borrowed `ManagedTensor` can be + /// searched multiple times (covers the `search(&self, …)` contract). #[test] fn test_cagra_multiple_searches() { let res = Resources::new().unwrap(); let build_params = IndexParams::new().unwrap(); - let (dataset, index) = build_test_index(&res, &build_params); + let dataset = make_dataset(); + let tensor = ManagedTensor::from(&dataset); + let index = + Index::build(&res, &build_params, &tensor).expect("failed to build cagra index"); for _ in 0..3 { search_and_verify_self_neighbors(&res, &index, &dataset, 4, 5); @@ -271,7 +310,10 @@ mod tests { fn test_cagra_serialize_deserialize() { let res = Resources::new().unwrap(); let build_params = IndexParams::new().unwrap(); - let (dataset, index) = build_test_index(&res, &build_params); + let dataset = make_dataset(); + let tensor = ManagedTensor::from(&dataset); + let index = + Index::build(&res, &build_params, &tensor).expect("failed to build cagra index"); let filepath = std::env::temp_dir().join("test_cagra_index.bin"); index @@ -284,11 +326,10 @@ mod tests { "serialized index file should not be empty" ); - let loaded_index = + // `deserialize` returns `Index<'static>` — the file owns its data, so + // it can safely outlive any input `ManagedTensor`. + let loaded_index: Index<'static> = Index::deserialize(&res, &filepath).expect("failed to deserialize cagra index"); - - // The deserialized index should still find each query as its own - // nearest neighbor. search_and_verify_self_neighbors(&res, &loaded_index, &dataset, 4, 10); let _ = std::fs::remove_file(&filepath); @@ -298,7 +339,10 @@ mod tests { fn test_cagra_serialize_without_dataset() { let res = Resources::new().unwrap(); let build_params = IndexParams::new().unwrap(); - let (_dataset, index) = build_test_index(&res, &build_params); + let dataset = make_dataset(); + let tensor = ManagedTensor::from(&dataset); + let index = + Index::build(&res, &build_params, &tensor).expect("failed to build cagra index"); let filepath = std::env::temp_dir().join("test_cagra_index_no_dataset.bin"); index @@ -314,7 +358,10 @@ mod tests { fn test_cagra_serialize_to_hnswlib() { let res = Resources::new().unwrap(); let build_params = IndexParams::new().unwrap(); - let (_dataset, index) = build_test_index(&res, &build_params); + let dataset = make_dataset(); + let tensor = ManagedTensor::from(&dataset); + let index = + Index::build(&res, &build_params, &tensor).expect("failed to build cagra index"); let filepath = std::env::temp_dir().join("test_cagra_index_hnsw.bin"); index @@ -339,7 +386,10 @@ mod tests { fn test_cagra_serialize_rejects_interior_nul() { let res = Resources::new().unwrap(); let build_params = IndexParams::new().unwrap(); - let (_dataset, index) = build_test_index(&res, &build_params); + let dataset = make_dataset(); + let tensor = ManagedTensor::from(&dataset); + let index = + Index::build(&res, &build_params, &tensor).expect("failed to build cagra index"); // `PathBuf::from` on Unix preserves arbitrary bytes, so we can embed a // NUL byte in the path and confirm the helper rejects it. diff --git a/rust/cuvs/src/cagra/mod.rs b/rust/cuvs/src/cagra/mod.rs index 9043b17386..76884af8d9 100644 --- a/rust/cuvs/src/cagra/mod.rs +++ b/rust/cuvs/src/cagra/mod.rs @@ -27,7 +27,8 @@ //! //! // build the cagra index //! let build_params = IndexParams::new()?; -//! let index = Index::build(&res, &build_params, &dataset)?; +//! let tensor = ManagedTensor::from(&dataset); +//! let index = Index::build(&res, &build_params, &tensor)?; //! println!( //! "Indexed {}x{} datapoints into cagra index", //! n_datapoints, n_features @@ -76,12 +77,14 @@ //! //! // Build an index (using some dataset) //! let build_params = IndexParams::new()?; -//! // let index = Index::build(&res, &build_params, &dataset)?; +//! // let tensor = ManagedTensor::from(&dataset); +//! // let index = Index::build(&res, &build_params, &tensor)?; //! //! // Save the index to disk (including the dataset) //! // index.serialize(&res, "/path/to/index.bin", true)?; //! -//! // Later, load the index from disk +//! // Later, load the index from disk. `deserialize` returns +//! // `Index<'static>` because the loaded index owns its data. //! let loaded_index = Index::deserialize(&res, "/path/to/index.bin")?; //! //! // The loaded index can be used for search just like the original