From 9dc3a4dc5085cb5eb92d18ca1eafd5eef3dc405e Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Fri, 14 Jun 2024 13:13:01 -0700 Subject: [PATCH] Add python serialization API's for ivf-pq and ivf_flat (#186) Authors: - Ben Frederickson (https://github.com/benfred) Approvers: - Dante Gama Dessavre (https://github.com/dantegd) URL: https://github.com/rapidsai/cuvs/pull/186 --- cpp/include/cuvs/neighbors/ivf_flat.h | 44 +++++++++ cpp/include/cuvs/neighbors/ivf_flat.hpp | 9 ++ cpp/include/cuvs/neighbors/ivf_pq.h | 40 ++++++++ cpp/include/cuvs/neighbors/ivf_pq.hpp | 14 ++- .../neighbors/detail/cagra/cagra_build.cpp | 68 +++++++------- cpp/src/neighbors/ivf_flat_c.cpp | 67 +++++++++++++- cpp/src/neighbors/ivf_flat_index.cpp | 6 ++ cpp/src/neighbors/ivf_pq_c.cpp | 92 ++++++++++--------- cpp/src/neighbors/ivf_pq_index.cpp | 18 ++++ cpp/test/neighbors/ann_ivf_flat.cuh | 2 +- cpp/test/neighbors/ann_ivf_pq.cuh | 2 +- python/cuvs/cuvs/common/exceptions.pyx | 2 +- .../cuvs/cuvs/neighbors/ivf_flat/__init__.py | 20 +++- .../cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd | 8 ++ .../cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx | 71 ++++++++++++++ python/cuvs/cuvs/neighbors/ivf_pq/__init__.py | 12 ++- python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd | 8 ++ python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx | 71 ++++++++++++++ python/cuvs/cuvs/test/test_cagra.py | 50 ---------- python/cuvs/cuvs/test/test_serialization.py | 76 +++++++++++++++ 20 files changed, 538 insertions(+), 142 deletions(-) create mode 100644 python/cuvs/cuvs/test/test_serialization.py diff --git a/cpp/include/cuvs/neighbors/ivf_flat.h b/cpp/include/cuvs/neighbors/ivf_flat.h index 2e8c551e4..5e91a4f61 100644 --- a/cpp/include/cuvs/neighbors/ivf_flat.h +++ b/cpp/include/cuvs/neighbors/ivf_flat.h @@ -278,6 +278,50 @@ cuvsError_t cuvsIvfFlatSearch(cuvsResources_t res, * @} */ +/** + * @defgroup ivf_flat_c_serialize IVF-Flat C-API serialize functions + * @{ + */ +/** + * Save the index to file. + * + * Experimental, both the API and the serialization format are subject to change. + * + * @code{.cpp} + * #include + * + * // Create cuvsResources_t + * cuvsResources_t res; + * cuvsError_t res_create_status = cuvsResourcesCreate(&res); + * + * // create an index with `cuvsIvfFlatBuild` + * cuvsIvfFlatSerialize(res, "/path/to/index", index, true); + * @endcode + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] filename the file name for saving the index + * @param[in] index IVF-Flat index + */ +cuvsError_t cuvsIvfFlatSerialize(cuvsResources_t res, + const char* filename, + cuvsIvfFlatIndex_t index); + +/** + * Load index from file. + * + * Experimental, both the API and the serialization format are subject to change. + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] filename the name of the file that stores the index + * @param[out] index IVF-Flat index loaded disk + */ +cuvsError_t cuvsIvfFlatDeserialize(cuvsResources_t res, + const char* filename, + cuvsIvfFlatIndex_t index); +/** + * @} + */ + #ifdef __cplusplus } #endif diff --git a/cpp/include/cuvs/neighbors/ivf_flat.hpp b/cpp/include/cuvs/neighbors/ivf_flat.hpp index b5f3b54fa..ee6102f52 100644 --- a/cpp/include/cuvs/neighbors/ivf_flat.hpp +++ b/cpp/include/cuvs/neighbors/ivf_flat.hpp @@ -138,6 +138,15 @@ struct index : cuvs::neighbors::index { index& operator=(const index&) = delete; index& operator=(index&&) = default; ~index() = default; + + /** + * @brief Construct an empty index. + * + * Constructs an empty index. This index will either need to be trained with `build` + * or loaded from a saved copy with `deserialize` + */ + index(raft::resources const& res); + /** Construct an empty index. It needs to be trained and then populated. */ index(raft::resources const& res, const index_params& params, uint32_t dim); /** Construct an empty index. It needs to be trained and then populated. */ diff --git a/cpp/include/cuvs/neighbors/ivf_pq.h b/cpp/include/cuvs/neighbors/ivf_pq.h index 9424729ff..f94448a93 100644 --- a/cpp/include/cuvs/neighbors/ivf_pq.h +++ b/cpp/include/cuvs/neighbors/ivf_pq.h @@ -351,6 +351,46 @@ cuvsError_t cuvsIvfPqSearch(cuvsResources_t res, * @} */ +/** + * @defgroup ivf_pq_c_serialize IVF-PQ C-API serialize functions + * @{ + */ +/** + * Save the index to file. + * + * Experimental, both the API and the serialization format are subject to change. + * + * @code{.cpp} + * #include + * + * // Create cuvsResources_t + * cuvsResources_t res; + * cuvsError_t res_create_status = cuvsResourcesCreate(&res); + * + * // create an index with `cuvsIvfPqBuild` + * cuvsIvfPqSerialize(res, "/path/to/index", index, true); + * @endcode + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] filename the file name for saving the index + * @param[in] index IVF-PQ index + */ +cuvsError_t cuvsIvfPqSerialize(cuvsResources_t res, const char* filename, cuvsIvfPqIndex_t index); + +/** + * Load index from file. + * + * Experimental, both the API and the serialization format are subject to change. + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] filename the name of the file that stores the index + * @param[out] index IVF-PQ index loaded disk + */ +cuvsError_t cuvsIvfPqDeserialize(cuvsResources_t res, const char* filename, cuvsIvfPqIndex_t index); +/** + * @} + */ + #ifdef __cplusplus } #endif diff --git a/cpp/include/cuvs/neighbors/ivf_pq.hpp b/cpp/include/cuvs/neighbors/ivf_pq.hpp index ea9e8b0ae..f38b6cbc4 100644 --- a/cpp/include/cuvs/neighbors/ivf_pq.hpp +++ b/cpp/include/cuvs/neighbors/ivf_pq.hpp @@ -279,6 +279,14 @@ struct index : cuvs::neighbors::index { auto operator=(index&&) -> index& = default; ~index() = default; + /** + * @brief Construct an empty index. + * + * Constructs an empty index. This index will either need to be trained with `build` + * or loaded from a saved copy with `deserialize` + */ + index(raft::resources const& handle); + /** Construct an empty index. It needs to be trained and then populated. */ index(raft::resources const& handle, cuvs::distance::DistanceType metric, @@ -1366,7 +1374,7 @@ void serialize(raft::resources const& handle, * * using IdxT = int64_t; // type of the index * // create an empty index - * cuvs::neighbors::ivf_pq::index index(handl, index_params, dim); + * cuvs::neighbors::ivf_pq::index index(handle); * * cuvs::neighbors::ivf_pq::deserialize(handle, is, index); * @endcode @@ -1390,8 +1398,8 @@ void deserialize(raft::resources const& handle, * // create a string with a filepath * std::string filename("/path/to/index"); * using IdxT = int64_t; // type of the index - * // create an empty index with - * ivf_pq::index index(handle, index_params, dim); + * // create an empty index + * ivf_pq::index index(handle); * * cuvs::neighbors::ivf_pq::deserialize(handle, filename, &index); * @endcode diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cpp b/cpp/src/neighbors/detail/cagra/cagra_build.cpp index 7ea45d063..574a02097 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cpp +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cpp @@ -1,35 +1,35 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include -#include -#include -#include - -namespace cuvs::neighbors::cagra::graph_build_params { -ivf_pq_params::ivf_pq_params(raft::matrix_extent dataset_extents, - cuvs::distance::DistanceType metric) -{ - build_params = cuvs::neighbors::ivf_pq::index_params::from_dataset(dataset_extents, metric); - - search_params = cuvs::neighbors::ivf_pq::search_params{}; - search_params.n_probes = std::max(10, build_params.n_lists * 0.01); - search_params.lut_dtype = CUDA_R_16F; - search_params.internal_distance_dtype = CUDA_R_16F; - - refinement_rate = 2; -} +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra::graph_build_params { +ivf_pq_params::ivf_pq_params(raft::matrix_extent dataset_extents, + cuvs::distance::DistanceType metric) +{ + build_params = cuvs::neighbors::ivf_pq::index_params::from_dataset(dataset_extents, metric); + + search_params = cuvs::neighbors::ivf_pq::search_params{}; + search_params.n_probes = std::max(10, build_params.n_lists * 0.01); + search_params.lut_dtype = CUDA_R_16F; + search_params.internal_distance_dtype = CUDA_R_16F; + + refinement_rate = 2; +} } // namespace cuvs::neighbors::cagra::graph_build_params \ No newline at end of file diff --git a/cpp/src/neighbors/ivf_flat_c.cpp b/cpp/src/neighbors/ivf_flat_c.cpp index c79ca8df9..ae06a55b6 100644 --- a/cpp/src/neighbors/ivf_flat_c.cpp +++ b/cpp/src/neighbors/ivf_flat_c.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include #include @@ -83,6 +84,22 @@ void _search(cuvsResources_t res, *res_ptr, search_params, *index_ptr, queries_mds, neighbors_mds, distances_mds); } +template +void _serialize(cuvsResources_t res, const char* filename, cuvsIvfFlatIndex index) +{ + auto res_ptr = reinterpret_cast(res); + auto index_ptr = reinterpret_cast*>(index.addr); + cuvs::neighbors::ivf_flat::serialize(*res_ptr, std::string(filename), *index_ptr); +} + +template +void* _deserialize(cuvsResources_t res, const char* filename) +{ + auto res_ptr = reinterpret_cast(res); + auto index = new cuvs::neighbors::ivf_flat::index(*res_ptr); + cuvs::neighbors::ivf_flat::deserialize(*res_ptr, std::string(filename), index); + return index; +} } // namespace extern "C" cuvsError_t cuvsIvfFlatIndexCreate(cuvsIvfFlatIndex_t* index) @@ -120,18 +137,16 @@ extern "C" cuvsError_t cuvsIvfFlatBuild(cuvsResources_t res, return cuvs::core::translate_exceptions([=] { auto dataset = dataset_tensor->dl_tensor; + index->dtype = dataset.dtype; if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { index->addr = reinterpret_cast(_build(res, *params, dataset_tensor)); - index->dtype.code = kDLFloat; } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) { index->addr = reinterpret_cast(_build(res, *params, dataset_tensor)); - index->dtype.code = kDLInt; } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) { index->addr = reinterpret_cast(_build(res, *params, dataset_tensor)); - index->dtype.code = kDLUInt; } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, @@ -213,3 +228,49 @@ extern "C" cuvsError_t cuvsIvfFlatSearchParamsDestroy(cuvsIvfFlatSearchParams_t { return cuvs::core::translate_exceptions([=] { delete params; }); } + +extern "C" cuvsError_t cuvsIvfFlatDeserialize(cuvsResources_t res, + const char* filename, + cuvsIvfFlatIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + // read the numpy dtype from the beginning of the file + std::ifstream is(filename, std::ios::in | std::ios::binary); + if (!is) { RAFT_FAIL("Cannot open file %s", filename); } + char dtype_string[4]; + is.read(dtype_string, 4); + auto dtype = raft::detail::numpy_serializer::parse_descr(std::string(dtype_string, 4)); + + index->dtype.bits = dtype.itemsize * 8; + if (dtype.kind == 'f' && dtype.itemsize == 4) { + index->addr = reinterpret_cast(_deserialize(res, filename)); + index->dtype.code = kDLFloat; + } else if (dtype.kind == 'i' && dtype.itemsize == 1) { + index->addr = reinterpret_cast(_deserialize(res, filename)); + index->dtype.code = kDLInt; + } else if (dtype.kind == 'u' && dtype.itemsize == 1) { + index->addr = reinterpret_cast(_deserialize(res, filename)); + index->dtype.code = kDLUInt; + } else { + RAFT_FAIL( + "Unsupported dtype in file %s itemsize %i kind %i", filename, dtype.itemsize, dtype.kind); + } + }); +} + +extern "C" cuvsError_t cuvsIvfFlatSerialize(cuvsResources_t res, + const char* filename, + cuvsIvfFlatIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + _serialize(res, filename, *index); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + _serialize(res, filename, *index); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + _serialize(res, filename, *index); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); +} diff --git a/cpp/src/neighbors/ivf_flat_index.cpp b/cpp/src/neighbors/ivf_flat_index.cpp index 61d720665..b249a9c29 100644 --- a/cpp/src/neighbors/ivf_flat_index.cpp +++ b/cpp/src/neighbors/ivf_flat_index.cpp @@ -18,6 +18,12 @@ namespace cuvs::neighbors::ivf_flat { +template +index::index(raft::resources const& res) + : index(res, cuvs::distance::DistanceType::L2Expanded, 0, false, false, 0) +{ +} + template index::index(raft::resources const& res, const index_params& params, uint32_t dim) : index(res, diff --git a/cpp/src/neighbors/ivf_pq_c.cpp b/cpp/src/neighbors/ivf_pq_c.cpp index 14a879123..3e835523a 100644 --- a/cpp/src/neighbors/ivf_pq_c.cpp +++ b/cpp/src/neighbors/ivf_pq_c.cpp @@ -20,8 +20,10 @@ #include #include #include +#include #include +#include #include #include #include @@ -88,30 +90,38 @@ void _search(cuvsResources_t res, *res_ptr, search_params, *index_ptr, queries_mds, neighbors_mds, distances_mds); } +template +void _serialize(cuvsResources_t res, const char* filename, cuvsIvfPqIndex index) +{ + auto res_ptr = reinterpret_cast(res); + auto index_ptr = reinterpret_cast*>(index.addr); + cuvs::neighbors::ivf_pq::serialize(*res_ptr, std::string(filename), *index_ptr); +} + +template +void* _deserialize(cuvsResources_t res, const char* filename) +{ + auto res_ptr = reinterpret_cast(res); + auto index = new cuvs::neighbors::ivf_pq::index(*res_ptr); + cuvs::neighbors::ivf_pq::deserialize(*res_ptr, std::string(filename), index); + return index; +} } // namespace extern "C" cuvsError_t cuvsIvfPqIndexCreate(cuvsIvfPqIndex_t* index) { - try { - *index = new cuvsIvfPqIndex{}; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + return cuvs::core::translate_exceptions([=] { *index = new cuvsIvfPqIndex{}; }); } extern "C" cuvsError_t cuvsIvfPqIndexDestroy(cuvsIvfPqIndex_t index_c_ptr) { - try { + return cuvs::core::translate_exceptions([=] { auto index = *index_c_ptr; auto index_ptr = reinterpret_cast*>(index.addr); delete index_ptr; delete index_c_ptr; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + }); } extern "C" cuvsError_t cuvsIvfPqBuild(cuvsResources_t res, @@ -119,7 +129,7 @@ extern "C" cuvsError_t cuvsIvfPqBuild(cuvsResources_t res, DLManagedTensor* dataset_tensor, cuvsIvfPqIndex_t index) { - try { + return cuvs::core::translate_exceptions([=] { auto dataset = dataset_tensor->dl_tensor; if ((dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) || @@ -127,15 +137,13 @@ extern "C" cuvsError_t cuvsIvfPqBuild(cuvsResources_t res, (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8)) { index->addr = reinterpret_cast(_build(res, *params, dataset_tensor)); index->dtype.code = dataset.dtype.code; + index->dtype.bits = dataset.dtype.bits; } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, dataset.dtype.bits); } - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + }); } extern "C" cuvsError_t cuvsIvfPqSearch(cuvsResources_t res, @@ -145,7 +153,7 @@ extern "C" cuvsError_t cuvsIvfPqSearch(cuvsResources_t res, DLManagedTensor* neighbors_tensor, DLManagedTensor* distances_tensor) { - try { + return cuvs::core::translate_exceptions([=] { auto queries = queries_tensor->dl_tensor; auto neighbors = neighbors_tensor->dl_tensor; auto distances = distances_tensor->dl_tensor; @@ -163,7 +171,6 @@ extern "C" cuvsError_t cuvsIvfPqSearch(cuvsResources_t res, "distances should be of type float32"); auto index = *index_c_ptr; - RAFT_EXPECTS(queries.dtype.code == index.dtype.code, "type mismatch between index and queries"); if ((queries.dtype.code == kDLFloat && queries.dtype.bits == 32) || (queries.dtype.code == kDLInt && queries.dtype.bits == 8) || @@ -174,16 +181,12 @@ extern "C" cuvsError_t cuvsIvfPqSearch(cuvsResources_t res, queries.dtype.code, queries.dtype.bits); } - - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + }); } extern "C" cuvsError_t cuvsIvfPqIndexParamsCreate(cuvsIvfPqIndexParams_t* params) { - try { + return cuvs::core::translate_exceptions([=] { *params = new cuvsIvfPqIndexParams{.metric = L2Expanded, .metric_arg = 2.0f, .add_data_on_build = true, @@ -195,41 +198,40 @@ extern "C" cuvsError_t cuvsIvfPqIndexParamsCreate(cuvsIvfPqIndexParams_t* params .codebook_kind = codebook_gen::PER_SUBSPACE, .force_random_rotation = false, .conservative_memory_allocation = false}; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + }); } extern "C" cuvsError_t cuvsIvfPqIndexParamsDestroy(cuvsIvfPqIndexParams_t params) { - try { - delete params; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + return cuvs::core::translate_exceptions([=] { delete params; }); } extern "C" cuvsError_t cuvsIvfPqSearchParamsCreate(cuvsIvfPqSearchParams_t* params) { - try { + return cuvs::core::translate_exceptions([=] { *params = new cuvsIvfPqSearchParams{.n_probes = 20, .lut_dtype = CUDA_R_32F, .internal_distance_dtype = CUDA_R_32F, .preferred_shmem_carveout = 1.0}; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + }); } extern "C" cuvsError_t cuvsIvfPqSearchParamsDestroy(cuvsIvfPqSearchParams_t params) { - try { - delete params; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + return cuvs::core::translate_exceptions([=] { delete params; }); +} + +extern "C" cuvsError_t cuvsIvfPqDeserialize(cuvsResources_t res, + const char* filename, + cuvsIvfPqIndex_t index) +{ + return cuvs::core::translate_exceptions( + [=] { index->addr = reinterpret_cast(_deserialize(res, filename)); }); +} + +extern "C" cuvsError_t cuvsIvfPqSerialize(cuvsResources_t res, + const char* filename, + cuvsIvfPqIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { _serialize(res, filename, *index); }); } diff --git a/cpp/src/neighbors/ivf_pq_index.cpp b/cpp/src/neighbors/ivf_pq_index.cpp index ee6c5dfb4..8f4e5b331 100644 --- a/cpp/src/neighbors/ivf_pq_index.cpp +++ b/cpp/src/neighbors/ivf_pq_index.cpp @@ -33,6 +33,23 @@ index_params index_params::from_dataset(raft::matrix_extent dataset, return params; } +template +index::index(raft::resources const& handle) + // this constructor is just for a temporary index, for use in the deserialization + // api. all the parameters here will get replaced with loaded values - that aren't + // necessarily known ahead of time before deserialization. + // TODO: do we even need a handle here - could just construct one? + : index(handle, + cuvs::distance::DistanceType::L2Expanded, + codebook_gen::PER_SUBSPACE, + 0, + 0, + 8, + 0, + true) +{ +} + template index::index(raft::resources const& handle, const index_params& params, uint32_t dim) : index(handle, @@ -76,6 +93,7 @@ index::index(raft::resources const& handle, check_consistency(); accum_sorted_sizes_(n_lists) = 0; } + template IdxT index::size() const noexcept { diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh index 49cb3ec2a..5dd6ebb2b 100644 --- a/cpp/test/neighbors/ann_ivf_flat.cuh +++ b/cpp/test/neighbors/ann_ivf_flat.cuh @@ -185,7 +185,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { distances_ivfflat_dev.data(), ps.num_queries, ps.k); const std::string filename = "ivf_flat_index"; cuvs::neighbors::ivf_flat::serialize(handle_, filename, index_2); - cuvs::neighbors::ivf_flat::index index_loaded(handle_, index_params, ps.dim); + cuvs::neighbors::ivf_flat::index index_loaded(handle_); cuvs::neighbors::ivf_flat::deserialize(handle_, filename, &index_loaded); ASSERT_EQ(index_2.size(), index_loaded.size()); diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index 662b701af..fa26d7ef9 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -242,7 +242,7 @@ class ivf_pq_test : public ::testing::TestWithParam { { std::string filename = "ivf_pq_index"; cuvs::neighbors::ivf_pq::serialize(handle_, filename, build_only()); - cuvs::neighbors::ivf_pq::index index(handle_, ps.index_params, ps.dim); + cuvs::neighbors::ivf_pq::index index(handle_); cuvs::neighbors::ivf_pq::deserialize(handle_, filename, &index); return index; } diff --git a/python/cuvs/cuvs/common/exceptions.pyx b/python/cuvs/cuvs/common/exceptions.pyx index 603e0711b..30134594e 100644 --- a/python/cuvs/cuvs/common/exceptions.pyx +++ b/python/cuvs/cuvs/common/exceptions.pyx @@ -28,7 +28,7 @@ def get_last_error_text(): if c_err is NULL: return cdef bytes err = c_err - return err.decode("utf8") + return err.decode("utf8", "ignore") def check_cuvs(status: cuvsError_t): diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/__init__.py b/python/cuvs/cuvs/neighbors/ivf_flat/__init__.py index a9c3e9b2d..d6a4afdde 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/__init__.py +++ b/python/cuvs/cuvs/neighbors/ivf_flat/__init__.py @@ -13,6 +13,22 @@ # limitations under the License. -from .ivf_flat import Index, IndexParams, SearchParams, build, search +from .ivf_flat import ( + Index, + IndexParams, + SearchParams, + build, + load, + save, + search, +) -__all__ = ["Index", "IndexParams", "SearchParams", "build", "search"] +__all__ = [ + "Index", + "IndexParams", + "SearchParams", + "build", + "load", + "save", + "search", +] diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd index 8859d5d96..65fcad5cf 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd +++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd @@ -72,3 +72,11 @@ cdef extern from "cuvs/neighbors/ivf_flat.h" nogil: DLManagedTensor* queries, DLManagedTensor* neighbors, DLManagedTensor* distances) except + + + cuvsError_t cuvsIvfFlatSerialize(cuvsResources_t res, + const char * filename, + cuvsIvfFlatIndex_t index) except + + + cuvsError_t cuvsIvfFlatDeserialize(cuvsResources_t res, + const char * filename, + cuvsIvfFlatIndex_t index) except + diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx index daa723099..03d254995 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx @@ -23,6 +23,7 @@ from cuvs.common.resources import auto_sync_resources from cython.operator cimport dereference as deref from libcpp cimport bool, cast +from libcpp.string cimport string from cuvs.common cimport cydlpack from cuvs.distance_type cimport cuvsDistanceType @@ -358,3 +359,73 @@ def search(SearchParams search_params, )) return (distances, neighbors) + + +@auto_sync_resources +def save(filename, Index index, bool include_dataset=True, resources=None): + """ + Saves the index to a file. + + Saving / loading the index is experimental. The serialization format is + subject to change. + + Parameters + ---------- + filename : string + Name of the file. + index : Index + Trained IVF-Flat index. + {resources_docstring} + + Examples + -------- + >>> import cupy as cp + >>> from cuvs.neighbors import ivf_flat + >>> n_samples = 50000 + >>> n_features = 50 + >>> dataset = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + >>> # Build index + >>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset) + >>> # Serialize and deserialize the ivf_flat index built + >>> ivf_flat.save("my_index.bin", index) + >>> index_loaded = ivf_flat.load("my_index.bin") + """ + cdef string c_filename = filename.encode('utf-8') + cdef cuvsResources_t res = resources.get_c_obj() + check_cuvs(cuvsIvfFlatSerialize(res, + c_filename.c_str(), + index.index)) + + +@auto_sync_resources +def load(filename, resources=None): + """ + Loads index from file. + + Saving / loading the index is experimental. The serialization format is + subject to change, therefore loading an index saved with a previous + version of cuvs is not guaranteed to work. + + Parameters + ---------- + filename : string + Name of the file. + {resources_docstring} + + Returns + ------- + index : Index + + """ + cdef Index idx = Index() + cdef cuvsResources_t res = resources.get_c_obj() + cdef string c_filename = filename.encode('utf-8') + + check_cuvs(cuvsIvfFlatDeserialize( + res, + c_filename.c_str(), + idx.index + )) + idx.trained = True + return idx diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/__init__.py b/python/cuvs/cuvs/neighbors/ivf_pq/__init__.py index 3c939aa27..bcbcff7e8 100644 --- a/python/cuvs/cuvs/neighbors/ivf_pq/__init__.py +++ b/python/cuvs/cuvs/neighbors/ivf_pq/__init__.py @@ -13,6 +13,14 @@ # limitations under the License. -from .ivf_pq import Index, IndexParams, SearchParams, build, search +from .ivf_pq import Index, IndexParams, SearchParams, build, load, save, search -__all__ = ["Index", "IndexParams", "SearchParams", "build", "search"] +__all__ = [ + "Index", + "IndexParams", + "SearchParams", + "build", + "load", + "save", + "search", +] diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd index f036c2c61..17d2e4030 100644 --- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd +++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd @@ -91,3 +91,11 @@ cdef extern from "cuvs/neighbors/ivf_pq.h" nogil: DLManagedTensor* queries, DLManagedTensor* neighbors, DLManagedTensor* distances) except + + + cuvsError_t cuvsIvfPqSerialize(cuvsResources_t res, + const char * filename, + cuvsIvfPqIndex_t index) except + + + cuvsError_t cuvsIvfPqDeserialize(cuvsResources_t res, + const char * filename, + cuvsIvfPqIndex_t index) except + diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx index 81baed1f0..0dc10e663 100644 --- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx @@ -23,6 +23,7 @@ from cuvs.common.resources import auto_sync_resources from cython.operator cimport dereference as deref from libcpp cimport bool, cast +from libcpp.string cimport string from cuvs.common cimport cydlpack from cuvs.distance_type cimport cuvsDistanceType @@ -432,3 +433,73 @@ def search(SearchParams search_params, )) return (distances, neighbors) + + +@auto_sync_resources +def save(filename, Index index, bool include_dataset=True, resources=None): + """ + Saves the index to a file. + + Saving / loading the index is experimental. The serialization format is + subject to change. + + Parameters + ---------- + filename : string + Name of the file. + index : Index + Trained IVF-PQ index. + {resources_docstring} + + Examples + -------- + >>> import cupy as cp + >>> from cuvs.neighbors import ivf_pq + >>> n_samples = 50000 + >>> n_features = 50 + >>> dataset = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + >>> # Build index + >>> index = ivf_pq.build(ivf_pq.IndexParams(), dataset) + >>> # Serialize and deserialize the ivf_pq index built + >>> ivf_pq.save("my_index.bin", index) + >>> index_loaded = ivf_pq.load("my_index.bin") + """ + cdef string c_filename = filename.encode('utf-8') + cdef cuvsResources_t res = resources.get_c_obj() + check_cuvs(cuvsIvfPqSerialize(res, + c_filename.c_str(), + index.index)) + + +@auto_sync_resources +def load(filename, resources=None): + """ + Loads index from file. + + Saving / loading the index is experimental. The serialization format is + subject to change, therefore loading an index saved with a previous + version of cuvs is not guaranteed to work. + + Parameters + ---------- + filename : string + Name of the file. + {resources_docstring} + + Returns + ------- + index : Index + + """ + cdef Index idx = Index() + cdef cuvsResources_t res = resources.get_c_obj() + cdef string c_filename = filename.encode('utf-8') + + check_cuvs(cuvsIvfPqDeserialize( + res, + c_filename.c_str(), + idx.index + )) + idx.trained = True + return idx diff --git a/python/cuvs/cuvs/test/test_cagra.py b/python/cuvs/cuvs/test/test_cagra.py index 38d0d1865..92b88f013 100644 --- a/python/cuvs/cuvs/test/test_cagra.py +++ b/python/cuvs/cuvs/test/test_cagra.py @@ -183,53 +183,3 @@ def test_cagra_vpq_compression(): run_cagra_build_search_test( n_cols=dim, compression=cagra.CompressionParams(pq_dim=dim / pq_len) ) - - -@pytest.mark.parametrize("dtype", [np.float32, np.int8, np.ubyte]) -# TODO: expose update_dataset -# @pytest.mark.parametrize("include_dataset", [True, False]) -@pytest.mark.parametrize("include_dataset", [True]) -def test_save_load(dtype, include_dataset): - n_rows = 10000 - n_cols = 50 - n_queries = 1000 - - dataset = generate_data((n_rows, n_cols), dtype) - dataset_device = device_ndarray(dataset) - - build_params = cagra.IndexParams() - index = cagra.build(build_params, dataset_device) - - assert index.trained - filename = "my_index.bin" - cagra.save(filename, index, include_dataset=include_dataset) - loaded_index = cagra.load(filename) - - # if we didn't save the dataset with the index, we need to update the - # index with an already loaded copy - if not include_dataset: - loaded_index.update_dataset(dataset) - - queries = generate_data((n_queries, n_cols), dtype) - - queries_device = device_ndarray(queries) - search_params = cagra.SearchParams() - k = 10 - - distance_dev, neighbors_dev = cagra.search( - search_params, index, queries_device, k - ) - - neighbors = neighbors_dev.copy_to_host() - dist = distance_dev.copy_to_host() - del index - - distance_dev, neighbors_dev = cagra.search( - search_params, loaded_index, queries_device, k - ) - - neighbors2 = neighbors_dev.copy_to_host() - dist2 = distance_dev.copy_to_host() - - assert np.all(neighbors == neighbors2) - assert np.allclose(dist, dist2, rtol=1e-6) diff --git a/python/cuvs/cuvs/test/test_serialization.py b/python/cuvs/cuvs/test/test_serialization.py new file mode 100644 index 000000000..4ffccf121 --- /dev/null +++ b/python/cuvs/cuvs/test/test_serialization.py @@ -0,0 +1,76 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import numpy as np +import pytest +from pylibraft.common import device_ndarray + +from cuvs.neighbors import cagra, ivf_flat, ivf_pq +from cuvs.test.ann_utils import generate_data + + +@pytest.mark.parametrize("dtype", [np.float32, np.int8, np.ubyte]) +def test_save_load_ivf_flat(dtype): + run_save_load(ivf_flat, dtype) + + +@pytest.mark.parametrize("dtype", [np.float32, np.int8, np.ubyte]) +def test_save_load_cagra(dtype): + run_save_load(cagra, dtype) + + +def test_save_load_ivf_pq(): + run_save_load(ivf_pq, np.float32) + + +def run_save_load(ann_module, dtype): + n_rows = 10000 + n_cols = 50 + n_queries = 1000 + + dataset = generate_data((n_rows, n_cols), dtype) + dataset_device = device_ndarray(dataset) + + build_params = ann_module.IndexParams() + index = ann_module.build(build_params, dataset_device) + + assert index.trained + filename = "my_index.bin" + ann_module.save(filename, index) + loaded_index = ann_module.load(filename) + + queries = generate_data((n_queries, n_cols), dtype) + + queries_device = device_ndarray(queries) + search_params = ann_module.SearchParams() + k = 10 + + distance_dev, neighbors_dev = ann_module.search( + search_params, index, queries_device, k + ) + + neighbors = neighbors_dev.copy_to_host() + dist = distance_dev.copy_to_host() + del index + + distance_dev, neighbors_dev = ann_module.search( + search_params, loaded_index, queries_device, k + ) + + neighbors2 = neighbors_dev.copy_to_host() + dist2 = distance_dev.copy_to_host() + + assert np.all(neighbors == neighbors2) + assert np.allclose(dist, dist2, rtol=1e-6)