diff --git a/include/knowhere/index/index_node.h b/include/knowhere/index/index_node.h index 118429f3f..2d4d2728e 100644 --- a/include/knowhere/index/index_node.h +++ b/include/knowhere/index/index_node.h @@ -12,7 +12,11 @@ #ifndef INDEX_NODE_H #define INDEX_NODE_H +#define KNOWHERE_SEARCH_CONFIG_CACHE_VERSION 1 + +#include #include +#include #include #include #include @@ -618,6 +622,22 @@ class IndexNode : public Object { SearchEmbList(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, milvus::OpContext* op_context = nullptr) const; + public: + virtual bool + SupportsSearchConfigCache() const { + return false; + } + + virtual expected + SearchWithPreparedConfig(const DataSetPtr, std::shared_ptr, const BitsetView&, + milvus::OpContext* = nullptr) const { + return expected::Err(Status::not_implemented, "prepared search config is not supported"); + } + + expected> + GetOrCreateSearchConfig(const Json& json) const; + + protected: static EmbListMetaHeader ParseEmbListMetaHeader(const uint8_t* data, int64_t size); @@ -637,6 +657,17 @@ class IndexNode : public Object { std::shared_ptr pool, milvus::OpContext* op_context = nullptr) const; Version version_; + + private: + struct SearchConfigCacheEntry { + Json json; + std::shared_ptr config; + }; + + mutable std::shared_ptr search_config_cache_; + mutable std::mutex search_config_cache_mutex_; + + protected: std::shared_ptr emb_list_offset_; // emb_list group offset structure (shared with strategy) std::string el_metric_type_; EmbListStrategyPtr emb_list_strategy_; // emb_list encoding strategy (tokenann/muvera) diff --git a/src/index/index.cc b/src/index/index.cc index 70f615494..d17616140 100644 --- a/src/index/index.cc +++ b/src/index/index.cc @@ -133,12 +133,23 @@ inline expected Index::Search(const DataSetPtr dataset, const Json& json, const BitsetView& bitset_, milvus::OpContext* op_context) const noexcept { return GuardedCall([&]() -> expected { - auto cfg = this->node->CreateConfig(); + std::unique_ptr owned_cfg; + std::shared_ptr prepared_cfg; std::string msg; - const Status load_status = LoadConfig(cfg.get(), json, knowhere::SEARCH, "Search", &msg); - if (load_status != Status::success) { - return expected::Err(load_status, msg); + if (this->node->SupportsSearchConfigCache()) { + auto result = this->node->GetOrCreateSearchConfig(json); + if (!result.has_value()) { + return expected::Err(result.error(), result.what()); + } + prepared_cfg = std::move(result.value()); + } else { + owned_cfg = this->node->CreateConfig(); + const Status load_status = LoadConfig(owned_cfg.get(), json, knowhere::SEARCH, "Search", &msg); + if (load_status != Status::success) { + return expected::Err(load_status, msg); + } } + const Config* cfg = prepared_cfg != nullptr ? prepared_cfg.get() : owned_cfg.get(); // when index is immutable, bitset size should always equal to data count in index // when index is mutable, it could happen that data count larger than bitset size, see // https://github.com/zilliztech/knowhere/issues/70 @@ -177,14 +188,18 @@ Index::Search(const DataSetPtr dataset, const Json& json, const BitsetView& b // LCOV_EXCL_STOP TimeRecorder rc("Search"); - auto k = cfg->k.value(); - auto res = this->node->SearchEmbListIfNeed(dataset, std::move(cfg), bitset, op_context); + auto k = b_cfg.k.value(); + auto res = prepared_cfg != nullptr + ? this->node->SearchWithPreparedConfig(dataset, std::move(prepared_cfg), bitset, op_context) + : this->node->SearchEmbListIfNeed(dataset, std::move(owned_cfg), bitset, op_context); auto time = rc.ElapseFromBegin("done"); time *= 0.001; // convert to ms this->node->GetSearchLatencyMetric().Observe(time); knowhere_search_topk.Observe(k); #else - auto res = this->node->SearchEmbListIfNeed(dataset, std::move(cfg), bitset, op_context); + auto res = prepared_cfg != nullptr + ? this->node->SearchWithPreparedConfig(dataset, std::move(prepared_cfg), bitset, op_context) + : this->node->SearchEmbListIfNeed(dataset, std::move(owned_cfg), bitset, op_context); #endif return res; }); diff --git a/src/index/index_node.cc b/src/index/index_node.cc index a5d41e3f8..d5514c7a1 100644 --- a/src/index/index_node.cc +++ b/src/index/index_node.cc @@ -32,6 +32,40 @@ namespace knowhere { +expected> +IndexNode::GetOrCreateSearchConfig(const Json& json) const { + auto cached = std::atomic_load_explicit(&search_config_cache_, std::memory_order_acquire); + if (cached != nullptr && cached->json == json) { + return cached->config; + } + + std::scoped_lock lock(search_config_cache_mutex_); + cached = std::atomic_load_explicit(&search_config_cache_, std::memory_order_relaxed); + if (cached != nullptr && cached->json == json) { + return cached->config; + } + + auto cfg = CreateConfig(); + Json normalized_json(json); + std::string msg; + auto status = Config::FormatAndCheck(*cfg, normalized_json, &msg); + LOG_KNOWHERE_DEBUG_ << "Search config dump: " << normalized_json.dump(); + if (status != Status::success) { + return expected>::Err(status, msg); + } + cfg->CaptureRawJson(normalized_json); + status = Config::Load(*cfg, normalized_json, knowhere::SEARCH, &msg); + if (status != Status::success) { + return expected>::Err(status, msg); + } + + std::shared_ptr prepared_config(std::move(cfg)); + auto entry = + std::make_shared(SearchConfigCacheEntry{.json = json, .config = prepared_config}); + std::atomic_store_explicit(&search_config_cache_, std::move(entry), std::memory_order_release); + return prepared_config; +} + // NOLINTBEGIN(google-default-arguments) expected IndexNode::RangeSearch(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, diff --git a/tests/ut/test_index_node.cc b/tests/ut/test_index_node.cc index 563bf67d7..4fcda9982 100644 --- a/tests/ut/test_index_node.cc +++ b/tests/ut/test_index_node.cc @@ -9,6 +9,8 @@ // 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 "catch2/catch_approx.hpp" @@ -136,6 +138,49 @@ class BaseFlatIndexNode : public IndexNode { } }; +template +class CachedSearchConfigIndexNode : public BaseFlatIndexNode { + public: + CachedSearchConfigIndexNode(const int32_t& version, const Object& object) + : BaseFlatIndexNode(version, object) { + } + + bool + SupportsSearchConfigCache() const override { + return true; + } + + expected + SearchWithPreparedConfig(const DataSetPtr, std::shared_ptr cfg, const BitsetView&, + milvus::OpContext*) const override { + const Config* expected = nullptr; + first_config_.compare_exchange_strong(expected, cfg.get()); + reused_same_config_.store(reused_same_config_.load() && first_config_.load() == cfg.get()); + return std::make_shared(); + } + + std::unique_ptr + CreateConfig() const override { + create_config_calls_.fetch_add(1); + return std::make_unique(); + } + + int + CreateConfigCalls() const { + return create_config_calls_.load(); + } + + bool + ReusedSameConfig() const { + return reused_same_config_.load(); + } + + private: + mutable std::atomic create_config_calls_{0}; + mutable std::atomic first_config_{nullptr}; + mutable std::atomic reused_same_config_{true}; +}; + TEST_CASE("Test index node") { auto version = GenTestVersionList(); DataSetPtr ds = std::make_shared(); @@ -208,3 +253,44 @@ TEST_CASE("Test index node") { } #pragma GCC diagnostic pop } + +TEST_CASE("Search reuses an immutable prepared config", "[search_config_cache]") { + KNOWHERE_SIMPLE_REGISTER_GLOBAL(SEARCH_CONFIG_CACHE, CachedSearchConfigIndexNode, fp32, knowhere::feature::FLOAT32); + const auto version = GenTestVersionList(); + auto dataset = std::make_shared(); + const Json base_search_config = {{meta::METRIC_TYPE, metric::L2}, {meta::TOPK, 10}}; + + SECTION("same config reuses the prepared object") { + auto index = IndexFactory::Instance().Create("SEARCH_CONFIG_CACHE", version).value(); + auto* node = dynamic_cast*>(index.Node()); + REQUIRE(node != nullptr); + + REQUIRE(index.Search(dataset, base_search_config, nullptr).has_value()); + REQUIRE(index.Search(dataset, base_search_config, nullptr).has_value()); + REQUIRE(node->CreateConfigCalls() == 1); + REQUIRE(node->ReusedSameConfig()); + + auto changed_search_config = base_search_config; + changed_search_config[meta::TOPK] = 20; + REQUIRE(index.Search(dataset, changed_search_config, nullptr).has_value()); + REQUIRE(node->CreateConfigCalls() == 2); + } + + SECTION("concurrent searches prepare the config once") { + auto index = IndexFactory::Instance().Create("SEARCH_CONFIG_CACHE", version).value(); + auto* node = dynamic_cast*>(index.Node()); + REQUIRE(node != nullptr); + + std::vector>> searches; + for (int i = 0; i < 32; ++i) { + searches.emplace_back( + std::async(std::launch::async, [&] { return index.Search(dataset, base_search_config, nullptr); })); + } + for (auto& search : searches) { + REQUIRE(search.get().has_value()); + } + + REQUIRE(node->CreateConfigCalls() == 1); + REQUIRE(node->ReusedSameConfig()); + } +}