diff --git a/include/svs/concurrent/README.md b/include/svs/concurrent/README.md index 48726788..a8ba73bf 100644 --- a/include/svs/concurrent/README.md +++ b/include/svs/concurrent/README.md @@ -226,3 +226,36 @@ and save/load all work with no extra code. `resize()` publishes the new size with a release store, and `size()` reads it with an acquire load. A search that overlaps an `add_points` therefore sees either the old size or the new one, never a broken value. + +The graph uses the same storage: `graphs::SimpleBlockedGraph` keeps its +adjacency lists in `SegmentedBlockedData`. `Alloc` defaults to +`HugepageAllocator`; pass a `SegmentedBlocked` to `auto_dynamic_build` +to choose another one. + +## 4. Type-erased wrapper + +`svs::ConcurrentDynamicVamana` +([../orchestrators/concurrent_dynamic_vamana.h](../orchestrators/concurrent_dynamic_vamana.h)) +wraps this index behind the `svs::DynamicVamana` API. It derives from +`svs::DynamicVamana` and only replaces the factory functions: + +- `build(parameters, data, ids, distance, threadpool, graph_allocator)` +- `assemble(config_dir, graph_loader, data_loader, distance, threadpool)` +- `assemble(stream, distance, threadpool, data_args...)` + +The dataset must use a `SegmentedBlocked` allocator; this is checked at compile +time. `distance` may be a functor or a `svs::DistanceType`. `batch_iterator()` +returns this index's own `BatchIterator`, so it takes the locks described in 1a. + +The rules from section 1 apply unchanged: `save()`, `set_threadpool()` and the +build-parameter setters still need exclusive access. + +```cpp +using Data = svs::index::vamana::concurrent::SegmentedBlockedData; +auto index = svs::ConcurrentDynamicVamana::build( + parameters, Data::load(path), ids, svs::DistanceType::L2, num_threads +); +// Safe from several threads at once: +index.add_points(points, new_ids); +auto results = index.search(queries, 10); +``` diff --git a/include/svs/concurrent/blocked_data.h b/include/svs/concurrent/blocked_data.h index 72a07a4a..de825fac 100644 --- a/include/svs/concurrent/blocked_data.h +++ b/include/svs/concurrent/blocked_data.h @@ -73,9 +73,13 @@ template class SegmentedBlocked : public svs::data::Blocked friend class SegmentedBlocked; template SegmentedBlocked(const SegmentedBlocked& other) - : parent_type{other.parameters(), other.get_allocator()} {} + : parent_type{other.parameters(), allocator_type(other.get_allocator())} {} }; +template inline constexpr bool is_segmented_blocked_v = false; +template +inline constexpr bool is_segmented_blocked_v> = true; + } // namespace svs::index::vamana::concurrent namespace svs::data { diff --git a/include/svs/concurrent/dynamic_index.h b/include/svs/concurrent/dynamic_index.h index 69aee51f..2c852b33 100644 --- a/include/svs/concurrent/dynamic_index.h +++ b/include/svs/concurrent/dynamic_index.h @@ -24,6 +24,8 @@ #include #include #include +#include +#include // Include the flat index to spin-up exhaustive searches on demand. #include "svs/index/flat/flat.h" @@ -270,10 +272,12 @@ class MutableVamanaIndex { svs::logging::logger_ptr logger_; // Methods - public: // Constructors + struct DeferReverseEdgesTag {}; + // Private constructor which defers the reverse edges building in the graph. template MutableVamanaIndex( + DeferReverseEdgesTag, Graph graph, Data data, Idx entry_point, @@ -294,53 +298,85 @@ class MutableVamanaIndex { , distance_{std::move(distance_function)} , threadpool_{threads::as_threadpool(std::move(threadpool_proto))} , search_parameters_{vamana::construct_default_search_parameters(data_)} - , construction_window_size_{2 * graph.max_degree()} + , construction_window_size_{2 * graph_.max_degree()} // Ctor accept logger in parameter , logger_{std::move(logger)} { + if (graph_.n_nodes() != data_.size()) { + throw ANNEXCEPTION("Graph node count does not match data size"); + } + if (graph_.n_nodes() != external_ids.size()) { + throw ANNEXCEPTION("Graph node count does not match external IDs size"); + } translator_.insert(external_ids, threads::UnitRange(0, external_ids.size())); + } + + public: + template + MutableVamanaIndex( + Graph graph, + Data data, + Idx entry_point, + Dist distance_function, + const ExternalIds& external_ids, + ThreadPoolProto threadpool_proto, + // Optional logger parameter + svs::logging::logger_ptr logger = svs::logging::get() + ) + : MutableVamanaIndex{ + DeferReverseEdgesTag{}, + std::move(graph), + std::move(data), + entry_point, + std::move(distance_function), + external_ids, + std::move(threadpool_proto), + std::move(logger)} { graph_.enable_reverse_edges(); graph_.rebuild_reverse_edges(threadpool_); } /// - /// Build a graph from scratch. + /// Build the graph into ``graph`` starting from ``entry_point``. /// template MutableVamanaIndex( const VamanaBuildParameters& parameters, + Graph graph, Data data, - const ExternalIds& external_ids, Dist distance_function, + const ExternalIds& external_ids, ThreadPoolProto threadpool_proto, + // Optional logger parameter svs::logging::logger_ptr logger = svs::logging::get() ) - : graph_(Graph{data.size(), parameters.graph_max_degree}) - , data_(std::move(data)) - , entry_point_{NO_ENTRY} - , status_(data_.size(), SlotMetadata::Valid) - , first_empty_{std::make_unique>(data_.size())} - , first_reusable_{std::make_unique>(data_.size())} - , translator_() - , num_valid_{std::make_unique>(data_.size())} - , distance_(std::move(distance_function)) - , threadpool_(threads::as_threadpool(std::move(threadpool_proto))) - , search_parameters_(vamana::construct_default_search_parameters(data_)) - , build_parameters_(parameters) - , logger_{std::move(logger)} { - // Verify and set defaults directly on the input parameters - verify_and_set_default_index_parameters(build_parameters_, distance_function); + : MutableVamanaIndex{ + DeferReverseEdgesTag{}, + std::move(graph), + std::move(data), + NO_ENTRY, + std::move(distance_function), + external_ids, + std::move(threadpool_proto), + std::move(logger)} { + if (graph_.n_nodes() != data_.size()) { + throw ANNEXCEPTION("Wrong sizes!"); + } + build_parameters_ = parameters; + // Verify and set defaults before using the parameters to set other member + // variables. + verify_and_set_default_index_parameters(build_parameters_, distance_); - // Set graph again as verify function might change graph_max_degree parameter - graph_ = Graph{data_.size(), build_parameters_.graph_max_degree}; + if (build_parameters_.graph_max_degree != graph_.max_degree()) { + throw ANNEXCEPTION( + "Graph max degree does not match the build parameters' max degree!" + ); + } + alpha_ = build_parameters_.alpha; construction_window_size_ = build_parameters_.window_size; max_candidates_ = build_parameters_.max_candidate_pool_size; prune_to_ = build_parameters_.prune_to; - alpha_ = build_parameters_.alpha; use_full_search_history_ = build_parameters_.use_full_search_history; - // Setup the initial translation of external to internal ids. - translator_.insert(external_ids, threads::UnitRange(0, external_ids.size())); - // An empty index has no medoid and no graph-construction work. if (data_.size() == 0) { graph_.enable_reverse_edges(); @@ -371,21 +407,41 @@ class MutableVamanaIndex { graph_.rebuild_reverse_edges(threadpool_); } - /// @brief Post re-load constructor. + /// + /// Build a graph from scratch. + /// + template + MutableVamanaIndex( + const VamanaBuildParameters& parameters, + Data data, + const ExternalIds& external_ids, + Dist distance_function, + ThreadPoolProto threadpool_proto, + svs::logging::logger_ptr logger = svs::logging::get() + ) + : MutableVamanaIndex{ + parameters, + Graph{data.size(), parameters.graph_max_degree}, + std::move(data), + std::move(distance_function), + external_ids, + std::move(threadpool_proto), + std::move(logger)} {} + + /// @brief Post re-load / copy constructor; the state may contain holes. /// /// Preconditions /// - /// * data.size() == graph.n_nodes(): The graph and the data have the same number of - /// entries. - /// * The data and graph were saved with no "holes". In otherwords, the index was - /// consolidated and compacted prior to saving. - /// * The span of internal ID's in translator covers exactly ``[0, data.size())``. + /// * data.size() == graph.n_nodes() == status.size(). + /// * ``status`` contains no ``Pending`` slots. + /// * ``translator`` maps exactly the ``Valid`` slots of ``status``. template MutableVamanaIndex( const VamanaIndexParameters& config, data_type data, graph_type graph, const Dist& distance_function, + const std::vector& status, IDTranslator translator, Pool threadpool, svs::logging::logger_ptr logger = svs::logging::get() @@ -393,11 +449,11 @@ class MutableVamanaIndex { : graph_{std::move(graph)} , data_{std::move(data)} , entry_point_{data_.size() == 0 ? NO_ENTRY : lib::narrow(config.entry_point)} - , status_{data_.size(), SlotMetadata::Valid} + , status_{data_.size(), SlotMetadata::Empty} // filled later in the body , first_empty_{std::make_unique>(data_.size())} , first_reusable_{std::make_unique>(data_.size())} , translator_{std::move(translator)} - , num_valid_{std::make_unique>(data_.size())} + , num_valid_{std::make_unique>(0)} , distance_{distance_function} , threadpool_{std::move(threadpool)} , search_parameters_{config.search_parameters} @@ -407,10 +463,96 @@ class MutableVamanaIndex { , alpha_{config.build_parameters.alpha} , use_full_search_history_{config.build_parameters.use_full_search_history} , logger_{std::move(logger)} { + const size_t n = data_.size(); + if (graph_.n_nodes() != n) { + throw ANNEXCEPTION("Graph node count does not match data size"); + } + if (status.size() != n) { + throw ANNEXCEPTION("Status size does not match data size"); + } + + size_t num_valid = 0; + size_t first_reusable = n; + for (size_t i = 0; i < n; ++i) { + const auto s = status[i]; + switch (s) { + case SlotMetadata::Valid: { + ++num_valid; + if (!translator_.has_internal(i)) { + throw ANNEXCEPTION("Translator is missing internal id {}", i); + } + break; + } + case SlotMetadata::Deleted: { + break; + } + case SlotMetadata::Empty: { + first_reusable = std::min(first_reusable, i); + // Reused slots are cleared on insertion; clear now so stale edges do + // not leak into the reverse-edge index. + graph_.clear_node(lib::narrow_cast(i)); + break; + } + case SlotMetadata::Pending: { + throw ANNEXCEPTION("Cannot assemble an index with pending slot {}", i); + } + } + status_[i] = s; + } + + if (translator_.size() != num_valid) { + throw ANNEXCEPTION( + "Translator has {} IDs but status has {} valid slots", + translator_.size(), + num_valid + ); + } + for (auto pair : translator_) { + if (pair.second >= n || status[pair.second] != SlotMetadata::Valid) { + throw ANNEXCEPTION( + "Translator maps external id {} to non-valid slot {}", + pair.first, + pair.second + ); + } + } + + num_valid_->store(num_valid); + first_reusable_->store(first_reusable); + if (n != 0) { + const auto ep = entry_point_[0]; + if (ep >= n || status[ep] == SlotMetadata::Empty) { + throw ANNEXCEPTION("Entry point {} is not an occupied slot", ep); + } + } + graph_.enable_reverse_edges(); graph_.rebuild_reverse_edges(threadpool_); } + /// @brief Post-reload backward compatible .ctor + /// @remarks just calls the main post-reload constructor with full-valid status. + template + MutableVamanaIndex( + const VamanaIndexParameters& config, + data_type data, + graph_type graph, + const Dist& distance_function, + IDTranslator translator, + Pool threadpool, + svs::logging::logger_ptr logger = svs::logging::get() + ) + : MutableVamanaIndex{ + config, + std::move(data), + std::move(graph), + distance_function, + // braced initializer is left-to-right - translator is not yet moved. + std::vector(translator.size(), SlotMetadata::Valid), + std::move(translator), + std::move(threadpool), + std::move(logger)} {} + ///// Scratchspace scratchspace_type scratchspace(const search_parameters_type& sp) const { return scratchspace_type{ @@ -1758,6 +1900,18 @@ class MutableVamanaIndex { const Data& view_data() const { return data_; } const Graph& view_graph() const { return graph_; } + // Not synchronized: callers must exclude concurrent mutation. + const IDTranslator& view_translator() const { return translator_; } + + // Not synchronized: callers must exclude concurrent mutation. + std::vector status_snapshot() const { + auto result = std::vector(status_.size()); + for (size_t i = 0, imax = result.size(); i < imax; ++i) { + result[i] = status_[i]; + } + return result; + } + /// /// @brief Verify the invariants of this data structure. /// @@ -2042,6 +2196,7 @@ MutableVamanaIndex( size_t, svs::logging::logger_ptr ) -> MutableVamanaIndex, Data, Dist>; + namespace detail { struct VamanaStateLoader { @@ -2060,30 +2215,107 @@ struct VamanaStateLoader { if (debug_load_from_static) { return VamanaStateLoader{ lib::load(table), - IDTranslator::Identity(assume_datasize)}; + IDTranslator::Identity(assume_datasize), + std::vector(assume_datasize, SlotMetadata::Valid)}; } return VamanaStateLoader{ SVS_LOAD_MEMBER_AT_(table, parameters), svs::lib::load_at(table, "translation"), + std::vector(assume_datasize, SlotMetadata::Valid), }; } ///// Members VamanaIndexParameters parameters_; IDTranslator translator_; + std::vector status_; }; +template +VamanaStateLoader auto_load_state( + ConfigProto&& config_proto, + [[maybe_unused]] bool debug_load_from_static, + [[maybe_unused]] size_t assume_datasize +) { + if constexpr (std::is_convertible_v) { + return lib::load_from_disk( + std::filesystem::path(std::forward(config_proto)), + debug_load_from_static, + assume_datasize + ); + } else { + return svs::detail::dispatch_load(std::forward(config_proto)); + } +} + } // namespace detail +/// +/// @brief Build an index whose data and graph both use grow-stable storage. +/// +/// The loaded dataset must use a ``SegmentedBlocked`` allocator, otherwise concurrent +/// ``add_points`` could relocate storage under lock-free readers. +/// +template < + typename DataProto, + typename Distance, + typename ExternalIdsProto, + typename ThreadPoolProto, + typename GraphAllocator = SegmentedBlocked>> +auto auto_dynamic_build( + const VamanaBuildParameters& parameters, + DataProto&& data_proto, + ExternalIdsProto&& external_ids_proto, + Distance distance, + ThreadPoolProto threadpool_proto, + const GraphAllocator& graph_allocator = {}, + svs::logging::logger_ptr logger = svs::logging::get() +) { + static_assert( + is_segmented_blocked_v, + "The concurrent index graph requires a SegmentedBlocked allocator." + ); + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + auto data = svs::detail::dispatch_load(SVS_FWD(data_proto), threadpool); + auto external_ids = svs::detail::dispatch_load(SVS_FWD(external_ids_proto), threadpool); + + using data_type = decltype(data); + static_assert( + is_segmented_blocked_v, + "The concurrent index requires a dataset with a SegmentedBlocked allocator." + ); + using graph_type = + graphs::SimpleBlockedGraph; + using index_type = MutableVamanaIndex; + + auto verified_parameters = parameters; + verify_and_set_default_index_parameters(verified_parameters, distance); + + auto graph = + graph_type{data.size(), verified_parameters.graph_max_degree, graph_allocator}; + return index_type( + verified_parameters, + std::move(graph), + std::move(data), + std::move(distance), + external_ids, + std::move(threadpool), + std::move(logger) + ); +} + // Assembly template < + typename ConfigProto, typename GraphLoader, typename DataLoader, typename Distance, typename ThreadPoolProto> +// Streams derived from std::istream would otherwise bind here over the stream overload. + requires(!std::is_base_of_v>) auto auto_dynamic_assemble( - const std::filesystem::path& config_path, + ConfigProto&& config_proto, GraphLoader&& graph_loader, DataLoader&& data_loader, Distance distance, @@ -2099,9 +2331,17 @@ auto auto_dynamic_assemble( // Load the dataset auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); auto data = svs::detail::dispatch_load(SVS_FWD(data_loader), threadpool); + static_assert( + is_segmented_blocked_v, + "The concurrent index requires a dataset with a SegmentedBlocked allocator." + ); // Load the graph. auto graph = svs::detail::dispatch_load(SVS_FWD(graph_loader), threadpool); + static_assert( + is_segmented_blocked_v, + "The concurrent index requires a graph with a SegmentedBlocked allocator." + ); // Make sure the data and the graph have the same size. auto datasize = data.size(); @@ -2134,31 +2374,17 @@ auto auto_dynamic_assemble( // } // }}; // auto [parameters, translator] = lib::load_from_disk(reloader, config_path); - auto [parameters, translator] = lib::load_from_disk( - config_path, debug_load_from_static, datasize + auto [parameters, translator, status] = detail::auto_load_state( + std::forward(config_proto), debug_load_from_static, datasize ); - // Make sure that the translator covers all the IDs in the graph and data. - auto translator_size = translator.size(); - if (translator_size != datasize) { - throw ANNEXCEPTION( - "Translator has {} IDs but should have {}", translator_size, datasize - ); - } - - for (size_t i = 0; i < datasize; ++i) { - if (!translator.has_internal(i)) { - throw ANNEXCEPTION("Translator is missing internal id {}", i); - } - } - - // At this point, we should be completely validated. - // Construct the index! + // Copied state may contain holes; the constructor validates status against translator. return MutableVamanaIndex{ parameters, std::move(data), std::move(graph), std::move(distance), + status, std::move(translator), std::move(threadpool), std::move(logger)}; @@ -2215,6 +2441,7 @@ auto auto_dynamic_assemble( std::move(data), std::move(graph), std::move(distance), + std::vector(datasize, SlotMetadata::Valid), std::move(translator), std::move(threadpool), std::move(logger)}; diff --git a/include/svs/concurrent/graph.h b/include/svs/concurrent/graph.h index 35e6b32b..da28a457 100644 --- a/include/svs/concurrent/graph.h +++ b/include/svs/concurrent/graph.h @@ -48,6 +48,20 @@ template struct unwrap_blocked_allocator { using type = typename A::allocator_type; }; +// True if `A` can allocate every element type `ReverseEdges` needs. Type-erased +// allocators such as `AllocatorHandle` only rebind to arithmetic types. +template +inline constexpr bool reverse_edges_rebindable_v = [] { + using traits = std::allocator_traits; + using idx_alloc = typename traits::template rebind_alloc; + using list_alloc = typename traits::template rebind_alloc>; + using lock_alloc = typename traits::template rebind_alloc; + return std::is_constructible_v && + std::is_constructible_v && + std::is_constructible_v && + std::equality_comparable; +}(); + // // We rely on an implicit layout for the graphs where length is stored inline with the // adjacency list like: @@ -85,9 +99,13 @@ template class SimpleGrap /// Base allocator underlying the adjacency storage; the reverse-edge index (when /// enabled) is allocated through it so its bytes are accounted the same way as the - /// graph's own storage. - using reverse_edge_allocator_type = + /// graph's own storage. Falls back to `std::allocator` if it cannot be rebound. + using graph_base_allocator_type = typename unwrap_blocked_allocator::type; + using reverse_edge_allocator_type = std::conditional_t< + reverse_edges_rebindable_v, + graph_base_allocator_type, + std::allocator>; using reverse_edges_type = ReverseEdges; /// @@ -625,7 +643,11 @@ template class SimpleGrap // Recover the graph's base allocator instance for the reverse-edge index, unwrapping // the Blocked<> layer (Blocked derives from its base allocator) when present. reverse_edge_allocator_type reverse_edge_allocator_() const { - if constexpr (data::is_blocked_v) { + if constexpr (!std::is_same_v< + reverse_edge_allocator_type, + graph_base_allocator_type>) { + return reverse_edge_allocator_type{}; + } else if constexpr (data::is_blocked_v) { return reverse_edge_allocator_type( static_cast( data_.get_allocator() @@ -722,19 +744,21 @@ bool operator==(const SimpleGraph& x, const SimpleGraph& y) { return graphs_equal(x, y); } -template -class SimpleBlockedGraph : public SimpleGraphBase< - Idx, - SegmentedBlockedData>> { +template > +class SimpleBlockedGraph + : public SimpleGraphBase> { public: - using parent_type = - SimpleGraphBase>>; + using parent_type = SimpleGraphBase>; using data_type = typename parent_type::data_type; + using allocator_type = typename data_type::allocator_type; // Constructors SimpleBlockedGraph(size_t num_nodes, size_t max_degree) : parent_type{num_nodes, max_degree} {} + SimpleBlockedGraph(size_t num_nodes, size_t max_degree, const allocator_type& allocator) + : parent_type{num_nodes, max_degree, allocator} {} + explicit SimpleBlockedGraph(data_type data) : parent_type{std::move(data)} {} @@ -742,29 +766,35 @@ class SimpleBlockedGraph : public SimpleGraphBase< : parent_type(std::move(parent)) {} ///// Loading - static constexpr SimpleBlockedGraph load(const lib::LoadTable& table) { + static constexpr SimpleBlockedGraph + load(const lib::LoadTable& table, const allocator_type& allocator = {}) { auto lazy = lib::Lazy([](data_type data) { return SimpleBlockedGraph(std::move(data)); }); - return parent_type::load(table, lazy); + return parent_type::load(table, lazy, allocator); } - static constexpr SimpleBlockedGraph - load(const lib::ContextFreeLoadTable& table, std::istream& is) { + static constexpr SimpleBlockedGraph load( + const lib::ContextFreeLoadTable& table, + std::istream& is, + const allocator_type& allocator = {} + ) { auto lazy = lib::Lazy([](data_type data) { return SimpleBlockedGraph(std::move(data)); }); - return parent_type::load(table, lazy, is); + return parent_type::load(table, lazy, is, allocator); } - static constexpr SimpleBlockedGraph load(const std::filesystem::path& path) { + static constexpr SimpleBlockedGraph + load(const std::filesystem::path& path, const allocator_type& allocator = {}) { if (data::detail::is_likely_reload(path)) { - return lib::load_from_disk(path); + return lib::load_from_disk(path, allocator); } else { - return SimpleBlockedGraph(data_type::load(path)); + return SimpleBlockedGraph(data_type::load(path, allocator)); } } - static constexpr SimpleBlockedGraph load(std::istream& is) { - return lib::load_from_stream(is); + static constexpr SimpleBlockedGraph + load(std::istream& is, const allocator_type& allocator = {}) { + return lib::load_from_stream(is, allocator); } }; diff --git a/include/svs/concurrent/multi.h b/include/svs/concurrent/multi.h index f8f41506..682d64dc 100644 --- a/include/svs/concurrent/multi.h +++ b/include/svs/concurrent/multi.h @@ -315,6 +315,7 @@ class MultiMutableVamanaIndex { std::move(data), std::move(graph), distance_, + std::vector(adds.size(), SlotMetadata::Valid), remapped_translator, std::move(threadpool), std::move(logger) @@ -358,6 +359,7 @@ class MultiMutableVamanaIndex { std::move(data), std::move(graph), distance_, + std::vector(adds.size(), SlotMetadata::Valid), remapped_translator, std::move(threadpool), std::move(logger) diff --git a/include/svs/core/allocator.h b/include/svs/core/allocator.h index 4a6f0670..b1bd82d9 100644 --- a/include/svs/core/allocator.h +++ b/include/svs/core/allocator.h @@ -233,7 +233,8 @@ template class HugepageAllocator { HugepageAllocator(const HugepageAllocator& other) : force_{other.force_} {} - template bool operator==(const HugepageAllocator& SVS_UNUSED(other)) { + template + bool operator==(const HugepageAllocator& SVS_UNUSED(other)) const { return true; } diff --git a/include/svs/orchestrators/concurrent_dynamic_vamana.h b/include/svs/orchestrators/concurrent_dynamic_vamana.h new file mode 100644 index 00000000..7167f302 --- /dev/null +++ b/include/svs/orchestrators/concurrent_dynamic_vamana.h @@ -0,0 +1,236 @@ +/* + * Copyright 2026 Intel 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. + */ + +#pragma once + +#include "svs/concurrent/blocked_data.h" +#include "svs/concurrent/dynamic_index.h" +#include "svs/concurrent/graph.h" +#include "svs/concurrent/iterator.h" + +#include "svs/orchestrators/dynamic_vamana.h" + +// stdlib +#include +#include +#include +#include + +namespace svs { + +/// +/// @brief Type-erased wrapper for ``svs::index::vamana::concurrent::MutableVamanaIndex``. +/// +/// Exposes the same API as ``svs::DynamicVamana``; only the factory functions differ. +/// The wrapped index may be searched and mutated from several threads at once: search, +/// batch iteration, ``add_points``, ``delete_points``, ``consolidate``, ``compact``, ID +/// inspection, ``get_distance`` and ``reconstruct_at`` are safe to call concurrently. +/// ``save``, ``set_threadpool`` and the build-parameter setters are not; call them only +/// when no other thread uses the index. See ``include/svs/concurrent/README.md``. +/// +/// The dataset must use a ``svs::index::vamana::concurrent::SegmentedBlocked`` allocator so +/// that growing it never relocates storage under concurrent readers. +/// +class ConcurrentDynamicVamana : public DynamicVamana { + public: + using base_type = DynamicVamana; + using default_graph_allocator_type = + index::vamana::concurrent::SegmentedBlocked>; + using default_graph_type = + index::vamana::concurrent::graphs::SimpleBlockedGraph; + + template + explicit ConcurrentDynamicVamana(AssembleTag tag, QueryTypes types, Impl impl) + : base_type{tag, types, std::move(impl)} {} + + /// + /// @brief Construct a ConcurrentDynamicVamana index from a data loader or dataset. + /// + /// @tparam QueryTypes The set of query element types supported by the resulting index. + /// + /// @param parameters Build parameters controlling graph construction. + /// @param data_loader Loader (or dataset) producing ``SegmentedBlocked`` storage. + /// @param ids External IDs to assign to each row; must be unique. + /// @param distance Distance functor or ``svs::DistanceType`` enum. + /// @param threadpool_proto Thread pool or number of threads to use. + /// @param graph_allocator ``SegmentedBlocked`` allocator used for the graph. + /// + template < + manager::QueryTypeDefinition QueryTypes, + typename DataLoader, + typename Distance, + typename ThreadPoolProto, + typename GraphAllocator = default_graph_allocator_type> + static ConcurrentDynamicVamana build( + const index::vamana::VamanaBuildParameters& parameters, + DataLoader&& data_loader, + std::span ids, + Distance distance, + ThreadPoolProto threadpool_proto, + const GraphAllocator& graph_allocator = {} + ) { + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + auto data = + svs::detail::dispatch_load(std::forward(data_loader), threadpool); + auto make = [&](auto distance_function) { + return ConcurrentDynamicVamana( + AssembleTag{}, + manager::as_typelist(), + index::vamana::concurrent::auto_dynamic_build( + parameters, + std::move(data), + ids, + std::move(distance_function), + std::move(threadpool), + graph_allocator + ) + ); + }; + if constexpr (std::is_same_v, DistanceType>) { + return DistanceDispatcher(distance)(make); + } else { + return make(std::move(distance)); + } + } + + /// + /// @brief Reload a ConcurrentDynamicVamana index from separate config, graph and data. + /// + /// @param config_proto Directory holding the saved index configuration, or an + /// already-loaded ``index::vamana::concurrent::detail::VamanaStateLoader``. + /// @param graph_loader Loader (or graph) producing a concurrent ``SimpleBlockedGraph``. + /// @param data_loader Loader (or dataset) producing ``SegmentedBlocked`` storage. + /// @param distance Distance functor or ``svs::DistanceType`` enum. + /// @param threadpool_proto Thread pool or number of threads to use. + /// @param debug_load_from_static Load a static index config with identity IDs. + /// + template < + manager::QueryTypeDefinition QueryTypes, + typename ConfigProto, + typename GraphLoader, + typename DataLoader, + typename Distance, + typename ThreadPoolProto> + static ConcurrentDynamicVamana assemble( + ConfigProto&& config_proto, + GraphLoader&& graph_loader, + DataLoader&& data_loader, + Distance distance, + ThreadPoolProto threadpool_proto, + bool debug_load_from_static = false + ) { + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + auto make = [&](auto distance_function) { + return ConcurrentDynamicVamana( + AssembleTag{}, + manager::as_typelist(), + index::vamana::concurrent::auto_dynamic_assemble( + std::forward(config_proto), + std::forward(graph_loader), + std::forward(data_loader), + std::move(distance_function), + std::move(threadpool), + debug_load_from_static + ) + ); + }; + if constexpr (std::is_same_v, DistanceType>) { + return DistanceDispatcher(distance)(make); + } else { + return make(std::move(distance)); + } + } + + /// + /// @brief Reload a ConcurrentDynamicVamana index saved with ``save(std::ostream&)``. + /// + /// Accepts both the native stream format and the directory-archive format. + /// + /// @tparam Data The ``SegmentedBlocked`` dataset type to load. + /// + template < + manager::QueryTypeDefinition QueryTypes, + typename Data, + typename Distance, + typename ThreadPoolProto, + typename... DataLoaderArgs> + static ConcurrentDynamicVamana assemble( + std::istream& stream, + Distance distance, + ThreadPoolProto threadpool_proto, + DataLoaderArgs&&... data_args + ) { + static_assert( + index::vamana::concurrent::is_segmented_blocked_v< + typename Data::allocator_type>, + "The concurrent index requires a dataset with a SegmentedBlocked allocator." + ); + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + auto deserializer = svs::lib::detail::Deserializer::build(stream); + if (deserializer.is_native()) { + auto make = [&](auto distance_function) { + return ConcurrentDynamicVamana( + AssembleTag{}, + manager::as_typelist(), + index::vamana::concurrent::auto_dynamic_assemble( + stream, + [&]() -> default_graph_type { + return default_graph_type::load(stream); + }, + [&]() -> Data { + return lib::load_from_stream( + stream, SVS_FWD(data_args)... + ); + }, + std::move(distance_function), + std::move(threadpool) + ) + ); + }; + if constexpr (std::is_same_v, DistanceType>) { + return DistanceDispatcher(distance)(make); + } else { + return make(std::move(distance)); + } + } + + namespace fs = std::filesystem; + lib::UniqueTempDirectory tempdir{"svs_concurrent_vamana_load"}; + lib::DirectoryArchiver::unpack(stream, tempdir, deserializer.magic()); + + const auto config_path = tempdir.get() / "config"; + const auto graph_path = tempdir.get() / "graph"; + const auto data_path = tempdir.get() / "data"; + for (const auto& path : {config_path, graph_path, data_path}) { + if (!fs::is_directory(path)) { + throw ANNEXCEPTION( + "Invalid Vamana index archive: missing {} directory!", + path.filename().string() + ); + } + } + + return assemble( + config_path, + SVS_LAZY(default_graph_type::load(graph_path)), + lib::load_from_disk(data_path, SVS_FWD(data_args)...), + std::move(distance), + std::move(threadpool) + ); + } +}; + +} // namespace svs diff --git a/include/svs/orchestrators/vamana_iterator.h b/include/svs/orchestrators/vamana_iterator.h index 4f227130..76ca8c9f 100644 --- a/include/svs/orchestrators/vamana_iterator.h +++ b/include/svs/orchestrators/vamana_iterator.h @@ -21,6 +21,8 @@ // stl #include +#include +#include namespace svs { @@ -46,16 +48,17 @@ class VamanaIterator { }; template struct Implementation : Interface { - // For the type-erased implementation - require the schedule to be type-erased as - // well. - using type = svs::index::vamana::BatchIterator; + // Each index type selects its own batch iterator. + using type = decltype(std::declval().make_batch_iterator( + std::declval>(), size_t{} + )); Implementation( const Index& index, std::span query, size_t extra_search_buffer_capacity = svs::UNSIGNED_INTEGER_PLACEHOLDER ) - : impl_{index, query, extra_search_buffer_capacity} {} + : impl_{index.make_batch_iterator(query, extra_search_buffer_capacity)} {} svs::index::vamana::VamanaSearchParameters parameters_for_current_iteration() const override { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2cfeb998..98a1cb29 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -176,6 +176,7 @@ SET(INTEGRATION_TESTS # Higher level constructs ${TEST_DIR}/svs/orchestrators/vamana.cpp ${TEST_DIR}/svs/orchestrators/dynamic_vamana.cpp + ${TEST_DIR}/svs/orchestrators/concurrent_dynamic_vamana.cpp # Integration Tests ${TEST_DIR}/integration/exhaustive.cpp ${TEST_DIR}/integration/vamana/index_search.cpp diff --git a/tests/svs/orchestrators/concurrent_dynamic_vamana.cpp b/tests/svs/orchestrators/concurrent_dynamic_vamana.cpp new file mode 100644 index 00000000..6e75631e --- /dev/null +++ b/tests/svs/orchestrators/concurrent_dynamic_vamana.cpp @@ -0,0 +1,402 @@ +/* + * Copyright 2026 Intel 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. + */ + +// Orchestrator under test +#include "svs/orchestrators/concurrent_dynamic_vamana.h" + +// svs +#include "svs/core/data/simple.h" +#include "svs/core/distance.h" +#include "svs/core/recall.h" +#include "svs/lib/file.h" + +// tests +#include "tests/utils/test_dataset.h" +#include "tests/utils/utils.h" +#include "tests/utils/vamana_reference.h" + +// catch2 +#include "catch2/catch_approx.hpp" +#include "catch2/catch_test_macros.hpp" + +// stl +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +namespace cc = svs::index::vamana::concurrent; +using ConcurrentData = cc::SegmentedBlockedData; + +const size_t num_threads = 2; + +// Allocator class which records allocated bytes +template class RecordingAllocator { + public: + using value_type = T; + + RecordingAllocator() = default; + + T* allocate(size_t n) { + *allocated_bytes += n * sizeof(T); + return static_cast(::operator new(n * sizeof(T))); + } + + void deallocate(T* p, size_t n) { + *allocated_bytes -= n * sizeof(T); + ::operator delete(p); + } + + template bool operator==(const RecordingAllocator& other) const { + return allocated_bytes == other.allocated_bytes; + } + template bool operator!=(const RecordingAllocator& other) const { + return !(*this == other); + } + + template + RecordingAllocator(const RecordingAllocator& other) + : allocated_bytes(other.allocated_bytes) {} + + size_t& allocated() { return *allocated_bytes; } + + std::shared_ptr allocated_bytes = std::make_shared(0); +}; + +ConcurrentData load_data() { return ConcurrentData::load(test_dataset::data_svs_file()); } + +std::vector iota_ids(size_t n, size_t start = 0) { + auto ids = std::vector(n); + std::iota(ids.begin(), ids.end(), start); + return ids; +} + +svs::index::vamana::VamanaBuildParameters small_build_parameters() { + return svs::index::vamana::VamanaBuildParameters{1.2, 32, 64, 128, 28, true}; +} + +// Copy rows [begin, end) of `data` into a plain dataset suitable for `add_points`. +svs::data::SimpleData rows(const ConcurrentData& data, size_t begin, size_t end) { + auto out = svs::data::SimpleData(end - begin, data.dimensions()); + for (size_t i = begin; i < end; ++i) { + out.set_datum(i - begin, data.get_datum(i)); + } + return out; +} + +svs::ConcurrentDynamicVamana build_index(const ConcurrentData& data, size_t n) { + auto subset = ConcurrentData(n, data.dimensions()); + for (size_t i = 0; i < n; ++i) { + subset.set_datum(i, data.get_datum(i)); + } + auto ids = iota_ids(n); + return svs::ConcurrentDynamicVamana::build( + small_build_parameters(), std::move(subset), ids, svs::DistanceType::L2, num_threads + ); +} + +template +svs::ConcurrentDynamicVamana +test_build(Distance distance, const GraphAllocator&... graph_allocator) { + auto expected_result = test_dataset::vamana::expected_build_results( + distance, svsbenchmark::Uncompressed(svs::DataType::float32) + ); + auto build_params = expected_result.build_parameters_.value(); + auto queries = svs::data::SimpleData::load(test_dataset::query_file()); + auto groundtruth = test_dataset::load_groundtruth(distance); + + auto data = load_data(); + const size_t n = data.size(); + auto ids = iota_ids(n); + + auto index = svs::ConcurrentDynamicVamana::build( + build_params, std::move(data), ids, distance, num_threads, graph_allocator... + ); + + CATCH_REQUIRE(index.size() == n); + CATCH_REQUIRE(index.get_alpha() == Catch::Approx(build_params.alpha)); + CATCH_REQUIRE(index.get_graph_max_degree() == build_params.graph_max_degree); + CATCH_REQUIRE(index.get_num_threads() == num_threads); + CATCH_REQUIRE(index.has_id(0)); + CATCH_REQUIRE(index.has_id(n - 1)); + + const double epsilon = 0.01; + for (const auto& expected : expected_result.config_and_recall_) { + auto these_queries = test_dataset::get_test_set(queries, expected.num_queries_); + auto these_groundtruth = + test_dataset::get_test_set(groundtruth, expected.num_queries_); + index.set_search_parameters(expected.search_parameters_); + auto results = index.search(these_queries, expected.num_neighbors_); + double recall = svs::k_recall_at_n( + these_groundtruth, results, expected.num_neighbors_, expected.recall_k_ + ); + CATCH_REQUIRE(recall > expected.recall_ - epsilon); + CATCH_REQUIRE(recall < expected.recall_ + epsilon); + } + return index; +} + +void require_same_results( + svs::ConcurrentDynamicVamana& expected, svs::DynamicVamana& actual, size_t k +) { + auto queries = test_dataset::queries(); + auto a = expected.search(queries, k); + auto b = actual.search(queries, k); + for (size_t q = 0; q < queries.size(); ++q) { + for (size_t i = 0; i < k; ++i) { + CATCH_REQUIRE(a.index(q, i) == b.index(q, i)); + } + } +} + +} // namespace + +CATCH_TEST_CASE( + "ConcurrentDynamicVamana Build", "[managers][concurrent_dynamic_vamana][build]" +) { + for (auto distance_enum : test_dataset::vamana::available_build_distances()) { + CATCH_SECTION(std::string("Functor ") + std::string(svs::name(distance_enum))) { + svs::DistanceDispatcher dispatcher(distance_enum); + dispatcher([&](auto distance) { test_build(distance); }); + } + } +} + +CATCH_TEST_CASE( + "ConcurrentDynamicVamana Build with Graph Allocator", + "[managers][concurrent_dynamic_vamana][build]" +) { + using GraphAllocator = cc::SegmentedBlocked>; + auto blocking = svs::data::BlockingParameters{}; + blocking.blocksize_elements = svs::lib::PowerOfTwo(7); + const size_t blocksize = blocking.blocksize_elements->value(); + auto recorder = RecordingAllocator{}; + + auto index = + test_build(svs::distance::DistanceL2{}, GraphAllocator{blocking, recorder}); + const size_t n = index.size(); + + // Graph capacity is a whole number of the custom blocks, not the 1 GiB default. + const size_t node_bytes = (index.get_graph_max_degree() + 1) * sizeof(uint32_t); + const size_t capacity = (n + blocksize - 1) / blocksize * blocksize; + auto breakdown = index.get_memory_breakdown(); + CATCH_REQUIRE(breakdown.graph_bytes == capacity * node_bytes); + // The reverse-edge index is allocated through the graph allocator as well. + const size_t allocated = recorder.allocated(); + CATCH_REQUIRE(allocated > breakdown.graph_bytes); + + // Growing past the current capacity appends exactly one more custom-sized block. + const size_t num_new = capacity - n + 1; + auto new_points = rows(load_data(), 0, num_new); + index.add_points(new_points.cview(), iota_ids(num_new, n)); + CATCH_REQUIRE(index.size() == n + num_new); + breakdown = index.get_memory_breakdown(); + CATCH_REQUIRE(breakdown.graph_bytes == (capacity + blocksize) * node_bytes); + CATCH_REQUIRE(recorder.allocated() >= allocated + blocksize * node_bytes); +} + +CATCH_TEST_CASE( + "ConcurrentDynamicVamana Mutation", "[managers][concurrent_dynamic_vamana]" +) { + auto data = load_data(); + const size_t n = data.size(); + const size_t half = n / 2; + auto index = build_index(data, half); + CATCH_REQUIRE(index.size() == half); + + const size_t usage_before = index.get_memory_breakdown().total(); + auto rest = rows(data, half, n); + index.add_points(rest.cview(), iota_ids(n - half, half)); + CATCH_REQUIRE(index.size() == n); + CATCH_REQUIRE(index.get_memory_breakdown().total() > usage_before); + + auto to_delete = iota_ids(half / 2); + index.delete_points(to_delete); + CATCH_REQUIRE(index.size() == n - to_delete.size()); + CATCH_REQUIRE_FALSE(index.has_id(0)); + CATCH_REQUIRE(index.has_id(n - 1)); + + index.consolidate().compact(); + CATCH_REQUIRE(index.size() == n - to_delete.size()); + + auto all = index.all_ids(); + CATCH_REQUIRE(all.size() == index.size()); + auto unique = std::unordered_set(all.begin(), all.end()); + CATCH_REQUIRE(unique.size() == all.size()); + CATCH_REQUIRE(unique.count(0) == 0); + + // Distance and reconstruction for a surviving id. + const size_t id = n - 1; + auto datum = data.get_datum(id); + auto query = std::vector(datum.begin(), datum.end()); + CATCH_REQUIRE(index.get_distance(id, query) == Catch::Approx(0.0).margin(1e-3)); + + auto reconstructed = svs::data::SimpleData(1, data.dimensions()); + auto reconstruct_ids = std::vector{id}; + index.reconstruct_at(reconstructed.view(), reconstruct_ids); + for (size_t j = 0; j < data.dimensions(); ++j) { + CATCH_REQUIRE(reconstructed.get_datum(0)[j] == datum[j]); + } + + CATCH_REQUIRE( + index.experimental_backend_string().find("concurrent") != std::string::npos + ); +} + +CATCH_TEST_CASE( + "ConcurrentDynamicVamana Batch Iterator", "[managers][concurrent_dynamic_vamana]" +) { + auto data = load_data(); + auto index = build_index(data, data.size()); + auto queries = test_dataset::queries(); + auto query = std::span(queries.get_datum(0)); + + auto iterator = index.batch_iterator(query); + auto seen = std::unordered_set(); + const size_t batch_size = 10; + for (size_t batch = 0; batch < 5; ++batch) { + iterator.next(batch_size); + CATCH_REQUIRE(iterator.size() == batch_size); + for (const auto& neighbor : iterator.results()) { + CATCH_REQUIRE(index.has_id(neighbor.id())); + CATCH_REQUIRE(seen.insert(neighbor.id()).second); + } + } + + // The first batch matches a regular search. + auto expected = index.search(queries, batch_size); + auto restarted = index.batch_iterator(query); + restarted.next(batch_size); + auto first = std::unordered_set(); + for (const auto& neighbor : restarted.results()) { + first.insert(neighbor.id()); + } + size_t overlap = 0; + for (size_t i = 0; i < batch_size; ++i) { + overlap += first.count(expected.index(0, i)); + } + CATCH_REQUIRE(overlap >= batch_size - 1); +} + +CATCH_TEST_CASE( + "ConcurrentDynamicVamana Save and Load", "[managers][concurrent_dynamic_vamana]" +) { + auto data = load_data(); + auto index = build_index(data, data.size()); + const size_t k = 10; + + CATCH_SECTION("Directories") { + svs_test::prepare_temp_directory(); + auto dir = svs_test::temp_directory(); + index.save(dir / "config", dir / "graph", dir / "data"); + svs::DynamicVamana loaded = svs::ConcurrentDynamicVamana::assemble( + dir / "config", + SVS_LAZY(cc::graphs::SimpleBlockedGraph::load(dir / "graph")), + SVS_LAZY(ConcurrentData::load(dir / "data")), + svs::DistanceType::L2, + num_threads + ); + CATCH_REQUIRE(loaded.size() == index.size()); + require_same_results(index, loaded, k); + } + + CATCH_SECTION("Native stream") { + std::stringstream stream; + index.save(stream); + svs::DynamicVamana loaded = + svs::ConcurrentDynamicVamana::assemble( + stream, svs::distance::DistanceL2(), num_threads + ); + CATCH_REQUIRE(loaded.size() == index.size()); + require_same_results(index, loaded, k); + } + + CATCH_SECTION("Directory archive stream") { + std::stringstream stream; + { + svs::lib::UniqueTempDirectory tempdir{"svs_concurrent_orchestrator_save"}; + index.save( + tempdir.get() / "config", tempdir.get() / "graph", tempdir.get() / "data" + ); + svs::lib::DirectoryArchiver::pack(tempdir, stream); + } + svs::DynamicVamana loaded = + svs::ConcurrentDynamicVamana::assemble( + stream, svs::DistanceType::L2, num_threads + ); + CATCH_REQUIRE(loaded.size() == index.size()); + require_same_results(index, loaded, k); + } +} + +CATCH_TEST_CASE( + "ConcurrentDynamicVamana Concurrent Search and Add", + "[managers][concurrent_dynamic_vamana]" +) { + auto data = load_data(); + const size_t n = data.size(); + const size_t initial = n / 2; + auto index = build_index(data, initial); + auto queries = test_dataset::queries(); + + constexpr size_t num_writers = 2; + constexpr size_t num_readers = 2; + const size_t per_writer = (n - initial) / num_writers; + + std::atomic failures{0}; + std::vector threads; + for (size_t w = 0; w < num_writers; ++w) { + threads.emplace_back([&, w]() { + const size_t begin = initial + w * per_writer; + for (size_t i = begin; i < begin + per_writer; i += 8) { + const size_t end = std::min(i + 8, begin + per_writer); + auto points = rows(data, i, end); + try { + index.add_points(points.cview(), iota_ids(end - i, i)); + } catch (...) { failures.fetch_add(1); } + } + }); + } + for (size_t r = 0; r < num_readers; ++r) { + threads.emplace_back([&]() { + for (size_t iter = 0; iter < 20; ++iter) { + try { + auto results = index.search(queries, 10); + for (size_t q = 0; q < results.n_queries(); ++q) { + if (results.index(q, 0) >= n) { + failures.fetch_add(1); + } + } + } catch (...) { failures.fetch_add(1); } + } + }); + } + for (auto& t : threads) { + t.join(); + } + + CATCH_REQUIRE(failures.load() == 0); + CATCH_REQUIRE(index.size() == initial + num_writers * per_writer); + for (size_t i = initial; i < initial + num_writers * per_writer; ++i) { + CATCH_REQUIRE(index.has_id(i)); + } +}