diff --git a/docs/docs/en/src/advanced/memory.md b/docs/docs/en/src/advanced/memory.md index 7e377a775a..06079233e6 100644 --- a/docs/docs/en/src/advanced/memory.md +++ b/docs/docs/en/src/advanced/memory.md @@ -97,6 +97,10 @@ Properties: allocator overhead, scratch buffers, and any data held outside the index (e.g. user-owned input vectors). For SINDI in particular, call `GetMemoryUsage()` **after** the build completes to get a representative value. +- For an HGraph loaded through `ReaderIO`, resident memory owned by a retained `Reader` is included + when the reader overrides `Reader::GetMemoryUsage()`. The default implementation returns zero; + custom readers with an in-memory page cache should override it. The underlying file or remote + data source is not part of the reported memory. See `examples/cpp/319_feature_get_memory_usage.cpp` for a runnable example, including a helper that compares the interface value with the process resident size. @@ -118,8 +122,9 @@ understanding *where* the memory is going inside an index. Currently only HGraph provides a meaningful implementation, returning components such as `basic_flatten_codes`, `bottom_graph`, `route_graph`, `neighbors_mutex`, `pool`, -`label_table`, `high_precise_codes`, `extra_infos`, and `raw_vector`. SINDI returns an empty -map. Other index types throw an exception by default. +`label_table`, `high_precise_codes`, `extra_infos`, `raw_vector`, `reader`, and +`precise_reader`. Reader keys are present only when the corresponding retained reader exists. +SINDI returns an empty map. Other index types throw an exception by default. ### Capability Flags @@ -140,5 +145,6 @@ full control over parallelism and resource ownership. See ## Notes - A custom allocator must be thread-safe. +- A custom `Reader::GetMemoryUsage()` override must be thread-safe. - The allocator's lifetime must outlive any index and result object referencing it. - If nothing is configured, VSAG falls back to a default `malloc`-based allocator. diff --git a/docs/docs/zh/src/advanced/memory.md b/docs/docs/zh/src/advanced/memory.md index 57b1453c13..eaeca723f5 100644 --- a/docs/docs/zh/src/advanced/memory.md +++ b/docs/docs/zh/src/advanced/memory.md @@ -88,6 +88,9 @@ uint64_t bytes = index->GetMemoryUsage(); RSS 还包含 allocator 的开销、临时 scratch buffer、以及索引外部持有的数据(例如用户自有的输入 向量缓冲)。SINDI 索引尤其建议在构建完成**之后**调用 `GetMemoryUsage()` 才能拿到具有代表性的 数值。 +- HGraph 通过 `ReaderIO` 加载时,如果被索引持有的 `Reader` 重写了 + `Reader::GetMemoryUsage()`,其常驻内存会计入统计。该方法默认返回 0;带内存页缓存的自定义 + Reader 应重写此方法。底层文件或远端数据源本身不计入内存统计。 可运行示例:`examples/cpp/319_feature_get_memory_usage.cpp`,其中包含一个辅助函数将接口值与进程 驻留内存进行对照。 @@ -107,7 +110,8 @@ for (const auto& [component, bytes] : detail) { 目前仅 HGraph 提供了有效实现,返回的组件包括 `basic_flatten_codes`、`bottom_graph`、 `route_graph`、`neighbors_mutex`、`pool`、`label_table`、`high_precise_codes`、 -`extra_infos` 和 `raw_vector`。SINDI 返回空 map,其他索引类型默认抛出异常。 +`extra_infos`、`raw_vector`、`reader` 和 `precise_reader`。仅在索引持有对应 Reader 时才会 +返回 Reader 相关项。SINDI 返回空 map,其他索引类型默认抛出异常。 ### 能力标志 @@ -127,5 +131,6 @@ for (const auto& [component, bytes] : detail) { ## 注意事项 - 自定义 Allocator 必须是线程安全的。 +- 自定义 `Reader::GetMemoryUsage()` 实现必须是线程安全的。 - `Allocator` 生命周期必须覆盖所有引用它的索引与结果对象。 - 若未显式指定,VSAG 会创建一个默认的基于 `malloc` 的 allocator。 diff --git a/include/vsag/readerset.h b/include/vsag/readerset.h index ee944240ae..c0e26821f7 100644 --- a/include/vsag/readerset.h +++ b/include/vsag/readerset.h @@ -131,6 +131,20 @@ class Reader { */ [[nodiscard]] virtual uint64_t Size() const = 0; + + /** + * @brief Returns memory owned by the reader, excluding the underlying data source. + * + * Reader implementations with resident caches should override this method so indexes that + * retain the reader can include that memory in their usage statistics. Overrides must be + * thread-safe because the method may be called concurrently with reads. + * + * @return Resident memory usage in bytes. The default implementation returns zero. + */ + [[nodiscard]] virtual uint64_t + GetMemoryUsage() const { + return 0; + } }; /** diff --git a/src/algorithm/hgraph/hgraph.cpp b/src/algorithm/hgraph/hgraph.cpp index 29d02335d0..32dccd5df6 100644 --- a/src/algorithm/hgraph/hgraph.cpp +++ b/src/algorithm/hgraph/hgraph.cpp @@ -538,6 +538,7 @@ HGraph::SetImmutable() { this->parallel_searcher_->SetMutexArray(empty_mutex); this->neighbors_mutex_ = empty_mutex; this->immutable_.store(true, std::memory_order_release); + this->cal_memory_usage(); } void @@ -549,6 +550,8 @@ HGraph::SetIO(const std::shared_ptr reader) { } basic_flatten_codes_->InitIO(reader_param); bottom_graph_->InitIO(reader_param); + this->reader_ = reader; + this->precise_reader_.reset(); } void @@ -556,6 +559,7 @@ HGraph::SetPreciseCodesIO(const std::shared_ptr& reader) { auto reader_param = std::make_shared(); reader_param->reader = reader; high_precise_codes_->InitIO(reader_param); + this->precise_reader_ = reader; } const static uint64_t QUERY_SAMPLE_SIZE = 10; @@ -745,11 +749,25 @@ HGraph::GetAttributeSetByInnerId(InnerIdType inner_id, AttributeSet* attr) const this->attr_filter_index_->GetAttribute(0, inner_id, attr); } +uint64_t +HGraph::GetMemoryUsage() const { + auto memory = InnerIndexInterface::GetMemoryUsage(); + if (this->pool_ != nullptr) { + memory += this->pool_->GetMemoryUsage(); + } + if (this->reader_ != nullptr) { + memory += this->reader_->GetMemoryUsage(); + } + if (this->precise_reader_ != nullptr && this->precise_reader_.get() != this->reader_.get()) { + memory += this->precise_reader_->GetMemoryUsage(); + } + return memory; +} + void HGraph::cal_memory_usage() { auto memory = sizeof(HGraph); memory += this->neighbors_mutex_->GetMemoryUsage(); - memory += this->pool_->GetMemoryUsage(); memory += this->label_table_->GetMemoryUsage(); memory += this->basic_flatten_codes_->GetMemoryUsage(); if (this->code_slot_map_ != nullptr) { @@ -770,7 +788,6 @@ HGraph::cal_memory_usage() { if (this->create_new_raw_vector_ and this->raw_vector_ != nullptr) { memory += raw_vector_->GetMemoryUsage(); } - std::unique_lock lock(this->memory_usage_mutex_); this->current_memory_usage_.store(memory); } diff --git a/src/algorithm/hgraph/hgraph.h b/src/algorithm/hgraph/hgraph.h index 427ebef1c4..540a3eb174 100644 --- a/src/algorithm/hgraph/hgraph.h +++ b/src/algorithm/hgraph/hgraph.h @@ -109,6 +109,9 @@ class HGraph : public InnerIndexInterface { uint64_t EstimateMemory(uint64_t num_elements) const override; + [[nodiscard]] uint64_t + GetMemoryUsage() const override; + void GetAttributeSetByInnerId(InnerIdType inner_id, AttributeSet* attr) const override; @@ -792,6 +795,9 @@ class HGraph : public InnerIndexInterface { std::shared_ptr> optimizer_; // search parameter optimizer + ReaderPtr reader_{nullptr}; // shared reader used by base, graph, and precise codes + ReaderPtr precise_reader_{nullptr}; // optional dedicated reader for precise codes + bool create_new_raw_vector_{false}; // whether a separate raw vector exists FlattenInterfacePtr raw_vector_{nullptr}; // raw float vectors (for distance calc) diff --git a/src/algorithm/hgraph/hgraph_serialize.cpp b/src/algorithm/hgraph/hgraph_serialize.cpp index 576756a0e3..9f658095ca 100644 --- a/src/algorithm/hgraph/hgraph_serialize.cpp +++ b/src/algorithm/hgraph/hgraph_serialize.cpp @@ -1073,6 +1073,12 @@ HGraph::GetMemoryUsageDetail() const { if (this->create_new_raw_vector_ && this->raw_vector_ != nullptr) { memory_usage["raw_vector"] = this->raw_vector_->GetMemoryUsage(); } + if (this->reader_ != nullptr) { + memory_usage["reader"] = this->reader_->GetMemoryUsage(); + } + if (this->precise_reader_ != nullptr && this->precise_reader_.get() != this->reader_.get()) { + memory_usage["precise_reader"] = this->precise_reader_->GetMemoryUsage(); + } return memory_usage; } diff --git a/tests/fixtures/framework/test_reader.cpp b/tests/fixtures/framework/test_reader.cpp index cc30739b13..ab0e582d03 100644 --- a/tests/fixtures/framework/test_reader.cpp +++ b/tests/fixtures/framework/test_reader.cpp @@ -18,7 +18,8 @@ namespace fixtures { -TestReader::TestReader(vsag::Binary binary) : binary_(std::move(binary)) { +TestReader::TestReader(vsag::Binary binary, uint64_t memory_usage) + : binary_(std::move(binary)), memory_usage_(memory_usage) { } void @@ -37,4 +38,9 @@ TestReader::Size() const { return binary_.size; } +uint64_t +TestReader::GetMemoryUsage() const { + return memory_usage_; +} + } // namespace fixtures diff --git a/tests/fixtures/framework/test_reader.h b/tests/fixtures/framework/test_reader.h index 150af953af..f8bebada8d 100644 --- a/tests/fixtures/framework/test_reader.h +++ b/tests/fixtures/framework/test_reader.h @@ -35,7 +35,7 @@ class TestReader : public vsag::Reader { * @brief Constructs a TestReader from a Binary object. * @param binary The binary data to wrap for reading. */ - explicit TestReader(vsag::Binary binary); + explicit TestReader(vsag::Binary binary, uint64_t memory_usage = 0); /** * @brief Reads data from the binary at the specified offset. @@ -63,8 +63,12 @@ class TestReader : public vsag::Reader { [[nodiscard]] uint64_t Size() const override; + [[nodiscard]] uint64_t + GetMemoryUsage() const override; + private: vsag::Binary binary_; // The wrapped binary data. + uint64_t memory_usage_{0}; }; } // namespace fixtures diff --git a/tests/test_hgraph.cpp b/tests/test_hgraph.cpp index 421bd4f789..8886b54f35 100644 --- a/tests/test_hgraph.cpp +++ b/tests/test_hgraph.cpp @@ -2312,6 +2312,67 @@ TestHGraphReaderIO(const fixtures::HGraphTestIndexPtr& test_index, HGRAPH_PR_DAILY_CASE("HGraph Reader IO", "[ft][serialize][hgraph]", TestHGraphReaderIO) +TEST_CASE("HGraph Reader Memory Usage", "[ft][memory][hgraph]") { + constexpr uint64_t reader_memory = 64ULL * 1024ULL * 1024ULL; + constexpr int64_t dim = 16; + constexpr int64_t count = 200; + const auto build_param = R"({ + "dtype": "float32", + "metric_type": "l2", + "dim": 16, + "index_param": { + "base_quantization_type": "sq8", + "precise_quantization_type": "fp32", + "use_reorder": true, + "max_degree": 16, + "ef_construction": 100 + } + })"; + const auto reader_param = R"({ + "dtype": "float32", + "metric_type": "l2", + "dim": 16, + "index_param": { + "base_quantization_type": "sq8", + "precise_quantization_type": "fp32", + "use_reorder": true, + "max_degree": 16, + "ef_construction": 100, + "base_io_type": "reader_io", + "precise_io_type": "reader_io", + "graph_io_type": "reader_io" + } + })"; + + auto dataset = fixtures::HGraphTestIndex::pool.GetDatasetAndCreate(dim, count, "l2"); + auto source = fixtures::TestIndex::TestFactory("hgraph", build_param, true); + fixtures::TestIndex::TestBuildIndex(source, dataset, true); + auto binary_set = source->Serialize().value(); + + auto restored = fixtures::TestIndex::TestFactory("hgraph", reader_param, true); + vsag::ReaderSet readers; + readers.Set("hgraph", + std::make_shared(binary_set.Get("hgraph"), reader_memory)); + REQUIRE(restored->Deserialize(readers).has_value()); + + auto detail = restored->GetMemoryUsageDetail(); + REQUIRE(detail.at("reader") == reader_memory); + auto mutable_memory = restored->GetMemoryUsage(); + REQUIRE(restored->SetImmutable().has_value()); + auto immutable_memory = restored->GetMemoryUsage(); + REQUIRE(immutable_memory < mutable_memory); + + detail = restored->GetMemoryUsageDetail(); + REQUIRE(detail.at("neighbors_mutex") == 0); + REQUIRE(detail.at("reader") == reader_memory); + uint64_t detail_sum = 0; + for (const auto& [_, memory] : detail) { + detail_sum += memory; + } + REQUIRE(immutable_memory >= detail_sum); + REQUIRE(immutable_memory - detail_sum < 4096); +} + static void TestHGraphClone(const fixtures::HGraphTestIndexPtr& test_index, const fixtures::HGraphResourcePtr& resource) {