xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(diskann, hnsw sparse builder): lost-wakeup race in builder progress loops (#530)

This commit is contained in:
Jalin Wang 2026-06-26 16:35:24 +08:00 committed by GitHub
parent 0a008ebfbf
commit cdef8686c7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 115 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
src/core/algorithm/cluster/multi_chunk_cluster.cc
View File

@ -201,16 +201,18 @@ int MultiChunkClusterAlgorithm::cluster(IndexThreads::Pointer threads,
threads->count(), &cents, &finished));
}
while (!task_group->is_finished()) {
{
std::unique_lock<std::mutex> lk(mutex_);
cond_.wait_until(lk, std::chrono::system_clock::now() +
std::chrono::seconds(check_interval_secs_));
if (error_.load(std::memory_order_acquire)) {
LOG_ERROR("Failed to cluster while waiting finish");
return errcode_;
while (finished.load() < chunk_count_) {
cond_.wait_until(lk, std::chrono::system_clock::now() +
std::chrono::seconds(check_interval_secs_));
if (error_.load(std::memory_order_acquire)) {
LOG_ERROR("Failed to cluster while waiting finish");
return errcode_;
}
LOG_INFO("Finish Chunk Count %zu, Finished Percent %.3f%%",
finished.load(), finished.load() * 100.0f / chunk_count_);
}
LOG_INFO("Finish Chunk Count %zu, Finished Percent %.3f%%", finished.load(),
finished.load() * 100.0f / chunk_count_);
}
if (error_.load(std::memory_order_acquire)) {
@ -284,16 +286,18 @@ int MultiChunkClusterAlgorithm::label(IndexThreads::Pointer threads,
threads->count(), cents, out, &finished));
}
while (!task_group->is_finished()) {
{
std::unique_lock<std::mutex> lk(mutex_);
cond_.wait_until(lk, std::chrono::system_clock::now() +
std::chrono::seconds(check_interval_secs_));
if (error_.load(std::memory_order_acquire)) {
LOG_ERROR("Failed to cluster while waiting finish");
return errcode_;
while (finished.load() < features_count) {
cond_.wait_until(lk, std::chrono::system_clock::now() +
std::chrono::seconds(check_interval_secs_));
if (error_.load(std::memory_order_acquire)) {
LOG_ERROR("Failed to cluster while waiting finish");
return errcode_;
}
LOG_INFO("Finish label cnt %zu, finished percent %.3f%%", finished.load(),
finished.load() * 100.0f / features_count);
}
LOG_INFO("Finish label cnt %zu, finished percent %.3f%%", finished.load(),
finished.load() * 100.0f / features_count);
}
if (error_.load(std::memory_order_acquire)) {

16
src/core/algorithm/cluster/multi_chunk_cluster.h
View File

@ -231,6 +231,10 @@ void MultiChunkNumericalAlgorithm<T>::do_cluster(
chunk, algorithm.centroids().count(), features_->count(), cost);
(*finished)++;
{
std::lock_guard<std::mutex> lk(mutex_);
cond_.notify_one();
}
}
return;
@ -267,6 +271,10 @@ void MultiChunkNumericalAlgorithm<T>::do_label(
}
(*finished)++;
{
std::lock_guard<std::mutex> lk(mutex_);
cond_.notify_one();
}
}
}
@ -372,6 +380,10 @@ void MultiChunkNumericalInnerProductAlgorithm<T>::do_cluster(
chunk, algorithm.centroids().count(), features_->count(), cost);
(*finished)++;
{
std::lock_guard<std::mutex> lk(mutex_);
cond_.notify_one();
}
}
}
@ -406,6 +418,10 @@ void MultiChunkNumericalInnerProductAlgorithm<T>::do_label(
}
(*finished)++;
{
std::lock_guard<std::mutex> lk(mutex_);
cond_.notify_one();
}
}
}

40
src/core/algorithm/diskann/diskann_builder.cc
View File

@ -267,17 +267,21 @@ int DiskAnnBuilder::build_internal(IndexThreads::Pointer threads) {
threads->count(), &finished));
}
while (!task_group->is_finished()) {
{
std::unique_lock<std::mutex> lk(mutex_);
cond_.wait_until(lk, std::chrono::system_clock::now() +
std::chrono::seconds(check_interval_secs_));
if (error_.load(std::memory_order_acquire)) {
LOG_ERROR("Failed to build index while waiting finish");
return errcode_;
while (finished.load() < entity_.doc_cnt()) {
cond_.wait_until(lk, std::chrono::system_clock::now() +
std::chrono::seconds(check_interval_secs_));
if (error_.load(std::memory_order_acquire)) {
LOG_ERROR("Failed to build index while waiting finish");
return errcode_;
}
LOG_INFO("Built cnt %zu, finished percent %.3f%%",
(size_t)finished.load(),
finished.load() * 100.0f / entity_.doc_cnt());
}
LOG_INFO("Built cnt %zu, finished percent %.3f%%", (size_t)finished.load(),
finished.load() * 100.0f / entity_.doc_cnt());
}
if (error_.load(std::memory_order_acquire)) {
LOG_ERROR("Failed to build index while waiting finish");
return errcode_;
@ -300,17 +304,21 @@ int DiskAnnBuilder::prune_internal(IndexThreads::Pointer threads) {
threads->count(), &finished));
}
while (!task_group->is_finished()) {
{
std::unique_lock<std::mutex> lk(mutex_);
cond_.wait_until(lk, std::chrono::system_clock::now() +
std::chrono::seconds(check_interval_secs_));
if (error_.load(std::memory_order_acquire)) {
LOG_ERROR("Failed to purne index while waiting finish");
return errcode_;
while (finished.load() < entity_.doc_cnt()) {
cond_.wait_until(lk, std::chrono::system_clock::now() +
std::chrono::seconds(check_interval_secs_));
if (error_.load(std::memory_order_acquire)) {
LOG_ERROR("Failed to prune index while waiting finish");
return errcode_;
}
LOG_INFO("Prune cnt %zu, finished percent %.3f%%",
(size_t)finished.load(),
finished.load() * 100.0f / entity_.doc_cnt());
}
LOG_INFO("Prune cnt %zu, finished percent %.3f%%", (size_t)finished.load(),
finished.load() * 100.0f / entity_.doc_cnt());
}
if (error_.load(std::memory_order_acquire)) {
LOG_ERROR("Failed to prune index while waiting finish");
return errcode_;

19
src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.cc
View File

@ -307,17 +307,20 @@ int HnswSparseBuilder::build_graph(IndexThreads::Pointer threads,
i, threads->count(), &finished));
}
while (!task_group->is_finished()) {
{
std::unique_lock<std::mutex> lk(mutex_);
cond_.wait_until(lk, std::chrono::system_clock::now() +
std::chrono::seconds(check_interval_secs_));
if (error_.load(std::memory_order_acquire)) {
LOG_ERROR("Failed to build index while waiting finish");
return errcode_;
while (finished.load() < entity_.doc_cnt()) {
cond_.wait_until(lk, std::chrono::system_clock::now() +
std::chrono::seconds(check_interval_secs_));
if (error_.load(std::memory_order_acquire)) {
LOG_ERROR("Failed to build index while waiting finish");
return errcode_;
}
LOG_INFO("Built cnt %u, finished percent %.3f%%", finished.load(),
finished.load() * 100.0f / entity_.doc_cnt());
}
LOG_INFO("Built cnt %u, finished percent %.3f%%", finished.load(),
finished.load() * 100.0f / entity_.doc_cnt());
}
if (error_.load(std::memory_order_acquire)) {
LOG_ERROR("Failed to build index while waiting finish");
return errcode_;

44
tests/core/algorithm/diskann/diskann_builder_test.cc
View File

@ -16,6 +16,7 @@
#include <sys/stat.h>
#include <sys/types.h>
#include <fcntl.h>
#include <chrono>
#include <future>
#include <gtest/gtest.h>
#include <zvec/ailego/container/vector.h>
@ -99,6 +100,49 @@ TEST_F(DiskAnnBuilderTest, TestGeneral) {
ASSERT_GT(stats.built_costtime(), 0UL);
}
// Regression test: building a small DiskAnn index must complete quickly.
// A lost-wakeup bug in the condition-variable progress loops previously caused
// 15–30 second stalls during train/build on small datasets because
// notify_one() was either missing or racing against a wrong predicate.
TEST_F(DiskAnnBuilderTest, SmallDatasetBuildTime) {
constexpr size_t kSmallDim = 4;
constexpr size_t kSmallDocCnt = 12;
auto meta = make_shared<IndexMeta>(IndexMeta::DataType::DT_FP32, kSmallDim);
meta->set_metric("SquaredEuclidean", 0, Params());
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder");
ASSERT_NE(builder, nullptr);
auto holder = make_shared<MultiPassIndexHolder<IndexMeta::DataType::DT_FP32>>(
kSmallDim);
for (size_t i = 0; i < kSmallDocCnt; ++i) {
NumericalVector<float> vec(kSmallDim, static_cast<float>(i));
ASSERT_TRUE(holder->emplace(i, vec));
}
Params params;
params.set("zvec.diskann.builder.max_degree", 32);
params.set("zvec.diskann.builder.list_size", 50);
params.set("zvec.diskann.builder.max_pq_chunk_num", 2);
params.set("zvec.diskann.builder.threads", 4);
ASSERT_EQ(0, builder->init(*meta, params));
auto t0 = std::chrono::steady_clock::now();
ASSERT_EQ(0, builder->train(holder));
ASSERT_EQ(0, builder->build(holder));
auto t1 = std::chrono::steady_clock::now();
auto elapsed_ms =
std::chrono::duration_cast<std::chrono::milliseconds>(t1 - t0).count();
// Before the fix, this took 15–30 seconds. After the fix, it should
// complete in well under 5 seconds even on slow CI machines.
EXPECT_LT(elapsed_ms, 5000)
<< "DiskAnn build with " << kSmallDocCnt << " vectors took " << elapsed_ms
<< " ms — likely a lost-wakeup regression in progress loops.";
}
// DiskAnn is now exposed implicitly: no caller ever invokes a
// ``LoadDiskAnnPlugin`` / ``IsLibAioAvailable`` API (those were removed from
// the public surface together with ``zvec.load_diskann_plugin()`` in Python).