Skip to content

[MOD-14958] Cover quant training caps and concurrent transition queries #1052

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dor-forer wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[MOD-14958] Cover quant training caps and concurrent transition queries #1052

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
135 changes: 0 additions & 135 deletions tests/unit/test_hnsw_sq8.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1164,141 +1164,6 @@ TEST(SQ8TieredHNSWTest, BatchIteratorCreatedBeforeNormalizationSeesMigratedLabel
}
}

TEST(SQ8TieredHNSWTest, ConcurrentQueriesDuringNormalizationTransition) {
constexpr size_t dim = 4;
constexpr size_t normalization_set_size = 2;
HNSWParams hnsw_params = {.type = VecSimType_FLOAT32,
.dim = dim,
.metric = VecSimMetric_L2,
.quantType = VecSimQuant_SQ8};
VecSimParams primary_index_params = CreateParams(hnsw_params);
tieredIndexMock mock_thread_pool;
TieredIndexParams tiered_params = {
.jobQueue = &mock_thread_pool.jobQ,
.jobQueueCtx = mock_thread_pool.ctx,
.submitCb = tieredIndexMock::submit_callback,
.primaryIndexParams = &primary_index_params,
.specificParams = {TieredHNSWParams{.QuantNormalizationSetSize = normalization_set_size}}};
VecSimParams params = CreateParams(tiered_params);
auto *index = VecSimIndex_New(&params);
ASSERT_NE(index, nullptr);
mock_thread_pool.ctx->index_strong_ref.reset(index);

auto *tiered_index = dynamic_cast<TieredHNSWIndex<float, float> *>(index);
ASSERT_NE(tiered_index, nullptr);

float first_vector[dim] = {1.0f, 1.0f, 1.0f, 1.0f};
float second_vector[dim] = {2.0f, 2.0f, 2.0f, 2.0f};
float query[dim] = {1.0f, 1.0f, 1.0f, 1.0f};
ASSERT_EQ(VecSimIndex_AddVector(index, first_vector, 0), 1);

std::mutex transition_mutex;
std::condition_variable transition_cv;
bool finalization_entered = false;
bool allow_finalization = false;
tiered_index->setBeforeQuantizationFinalizationHook([&] {
std::unique_lock lock(transition_mutex);
finalization_entered = true;
transition_cv.notify_all();
transition_cv.wait(lock, [&] { return allow_finalization; });
});

std::thread writer([&] { EXPECT_EQ(VecSimIndex_AddVector(index, second_vector, 1), 1); });
{
std::unique_lock lock(transition_mutex);
ASSERT_TRUE(transition_cv.wait_for(lock, std::chrono::seconds(10),
[&] { return finalization_entered; }));
}

std::atomic_bool keep_reading = true;
std::atomic_int failures = 0;
std::condition_variable readers_cv;
std::mutex readers_mutex;
bool reader_completed_iteration = false;
std::thread reader([&] {
bool announced = false;
while (keep_reading.load()) {
auto *reply = VecSimIndex_TopKQuery(index, query, 2, nullptr, BY_SCORE);
if (!reply || reply->code != VecSim_QueryReply_OK ||
VecSimQueryReply_Len(reply) != normalization_set_size) {
failures.fetch_add(1);
}
if (reply) {
VecSimQueryReply_Free(reply);
}

reply = VecSimIndex_RangeQuery(index, query, 100.0, nullptr, BY_SCORE);
if (!reply || reply->code != VecSim_QueryReply_OK ||
VecSimQueryReply_Len(reply) != normalization_set_size) {
failures.fetch_add(1);
}
if (reply) {
VecSimQueryReply_Free(reply);
}

auto *iterator = VecSimBatchIterator_New(index, query, nullptr);
if (!iterator) {
failures.fetch_add(1);
} else {
if (VecSimBatchIterator_HasNext(iterator)) {
reply = VecSimBatchIterator_Next(iterator, 2, BY_SCORE);
if (!reply || reply->code != VecSim_QueryReply_OK) {
failures.fetch_add(1);
}
if (reply) {
VecSimQueryReply_Free(reply);
}
}
VecSimBatchIterator_Free(iterator);
}

(void)VecSimIndex_IndexSize(index);
(void)VecSimIndex_DebugInfo(index);
(void)VecSimIndex_StatsInfo(index);
(void)VecSimIndex_PreferAdHocSearch(index, normalization_set_size, 1, true);

if (!announced) {
{
std::lock_guard lock(readers_mutex);
reader_completed_iteration = true;
}
readers_cv.notify_one();
announced = true;
}
}
});

{
std::unique_lock lock(readers_mutex);
EXPECT_TRUE(readers_cv.wait_for(lock, std::chrono::seconds(10),
[&] { return reader_completed_iteration; }));
}

mock_thread_pool.init_threads();

{
std::lock_guard lock(transition_mutex);
allow_finalization = true;
}
transition_cv.notify_all();
writer.join();
mock_thread_pool.thread_pool_wait();

keep_reading.store(false);
reader.join();
EXPECT_EQ(failures.load(), 0);

auto *reply = VecSimIndex_TopKQuery(index, query, 2, nullptr, BY_SCORE);
ASSERT_NE(reply, nullptr);
EXPECT_EQ(reply->code, VecSim_QueryReply_OK);
EXPECT_EQ(VecSimQueryReply_Len(reply), normalization_set_size);
VecSimQueryReply_Free(reply);

// Keep the allocator alive while reset_ctx releases the index's final reference.
auto allocator = index->getAllocator();
mock_thread_pool.reset_ctx();
}

TEST(SQ8TieredHNSWTest, WarnsForDimensionsBelow64) {
auto previous_log_callback = VecSimIndexInterface::logCallback;
struct LogCallbackRestorer {
Expand Down
145 changes: 145 additions & 0 deletions tests/unit/test_hnsw_tiered.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,9 @@
#include "mock_thread_pool.h"

#include <thread>
#include <atomic>
#include <condition_variable>
#include <set>

// Runs the test for all combination of data type(float/double) - label type (single/multi)

Expand Down Expand Up @@ -5313,6 +5316,148 @@ TYPED_TEST(HNSWTieredIndexTestSQ8, AccumulationPhaseInitialization) {
ASSERT_EQ(tiered_index->indexSize(), 0);
}

TYPED_TEST(HNSWTieredIndexTestSQ8, TrainingThresholdIsCapped) {
constexpr size_t cap = MAX_QUANT_NORMALIZATION_SET_SIZE;
for (auto metric : {VecSimMetric_L2, VecSimMetric_IP, VecSimMetric_Cosine}) {
for (size_t requested : {cap - 1, cap, cap + 1, 2 * cap}) {
SCOPED_TRACE(::testing::Message() << "metric=" << metric << " threshold=" << requested);
tieredIndexMock pool;
auto *index = this->CreateSQ8TieredIndex(pool, 64, metric, requested);
if (!index) {
pool.reset_ctx();
FAIL() << "SQ8 index construction failed";
}
ASSERT_TRUE(this->getIsInAccumulationPhase(index));
EXPECT_EQ(this->getNormalizationSetSize(index), requested < cap ? requested : cap);
}
}
}

TYPED_TEST(HNSWTieredIndexTestSQ8, ConcurrentQueriesDuringNormalizationTransition) {
constexpr size_t dim = 65;
constexpr size_t count = 8;
const size_t labels = TypeParam::isMulti() ? count / 2 : count;
for (auto metric : {VecSimMetric_L2, VecSimMetric_IP, VecSimMetric_Cosine}) {
for (auto mode : {VecSim_WriteAsync, VecSim_WriteInPlace}) {
SCOPED_TRACE(::testing::Message() << "metric=" << metric << " mode=" << mode);
VecSimIndexInterface::asyncWriteMode = mode;
tieredIndexMock pool(2);
auto *index = this->CreateSQ8TieredIndex(pool, dim, metric, count);
if (!index) {
pool.reset_ctx();
FAIL() << "SQ8 index construction failed";
}
std::vector<std::vector<TEST_DATA_T>> vectors(count, std::vector<TEST_DATA_T>(dim));
for (size_t i = 0; i < count; ++i) {
this->GenerateVectorData(vectors[i].data(), dim, 10.0f * (i + 1));
}
for (size_t i = 0; i < count - 1; ++i) {
ASSERT_EQ(VecSimIndex_AddVector(index, vectors[i].data(),
TypeParam::isMulti() ? i / 2 : i),
1);
}

std::mutex mutex;
std::condition_variable cv;
unsigned phase = 0, completed_phase = 0;
const auto pause_for_queries = [&](unsigned next_phase) {
std::unique_lock lock(mutex);
phase = next_phase;
cv.notify_all();
EXPECT_TRUE(cv.wait_for(lock, std::chrono::seconds(10),
[&] { return completed_phase >= next_phase; }));
};
index->setBeforeQuantizationFinalizationHook([&] { pause_for_queries(1); });
index->setAfterQuantizationFinalizationHook([&] { pause_for_queries(2); });
std::atomic_bool first_backend_insert{true};
index->setAfterBackendInsertBeforeFlatRemovalHook([&] {
if (first_backend_insert.exchange(false)) {
pause_for_queries(3);
}
});

const auto check_reply = [&](VecSimQueryReply *reply, std::set<size_t> &seen) {
EXPECT_NE(reply, nullptr);
if (!reply) {
return;
}
EXPECT_EQ(reply->code, VecSim_QueryReply_OK);
for (const auto &result : reply->results) {
const auto label = VecSimQueryResult_GetId(&result);
EXPECT_LT(label, labels);
EXPECT_TRUE(seen.insert(label).second) << "duplicate label " << label;
EXPECT_TRUE(std::isfinite(VecSimQueryResult_GetScore(&result)));
}
VecSimQueryReply_Free(reply);
};
const auto check_queries = [&] {
std::set<size_t> seen;
check_reply(
VecSimIndex_TopKQuery(index, vectors[0].data(), labels, nullptr, BY_SCORE),
seen);
EXPECT_EQ(seen.size(), labels);
seen.clear();
check_reply(
VecSimIndex_RangeQuery(index, vectors[0].data(), 100.0, nullptr, BY_SCORE),
seen);
EXPECT_EQ(seen.size(), labels);
seen.clear();
auto *iterator = VecSimBatchIterator_New(index, vectors[0].data(), nullptr);
ASSERT_NE(iterator, nullptr);
size_t batches = 0;
while (VecSimBatchIterator_HasNext(iterator) && batches++ < labels + 1) {
check_reply(VecSimBatchIterator_Next(iterator, 2, BY_SCORE), seen);
}
EXPECT_FALSE(VecSimBatchIterator_HasNext(iterator));
EXPECT_EQ(seen.size(), labels);
VecSimBatchIterator_Free(iterator);
(void)VecSimIndex_IndexSize(index);
(void)VecSimIndex_DebugInfo(index);
(void)VecSimIndex_StatsInfo(index);
(void)VecSimIndex_PreferAdHocSearch(index, count, 1, true);
};

// Each hook waits for a complete query pass begun in that phase. The reader also
// keeps running between hooks, while the writer and migration workers advance.
std::jthread reader([&](std::stop_token stop) {
while (!stop.stop_requested()) {
unsigned observed_phase;
{
std::unique_lock lock(mutex);
if (!cv.wait_for(lock, std::chrono::seconds(10),
[&] { return phase > 0; })) {
ADD_FAILURE() << "Writer did not reach quantization finalization";
return;
}
observed_phase = phase;
}
check_queries();
{
std::lock_guard lock(mutex);
completed_phase = observed_phase;
}
cv.notify_all();
}
});
if (mode == VecSim_WriteAsync) {
pool.init_threads();
}
std::jthread writer([&] {
EXPECT_EQ(VecSimIndex_AddVector(index, vectors.back().data(), labels - 1), 1);
});
writer.join();
pool.thread_pool_wait();
reader.request_stop();
reader.join();
EXPECT_EQ(completed_phase, 3u);
EXPECT_FALSE(this->getIsInAccumulationPhase(index));
EXPECT_EQ(this->getFrontendIndex(index)->indexSize(), 0);
EXPECT_EQ(this->getBackendIndex(index)->indexSize(), count);
check_queries();
}
}
}

TYPED_TEST(HNSWTieredIndexTestSQ8, InitialMemoryAndSuppliedMeanWithOptionalTraining) {
constexpr size_t dim = 17;
const std::vector<float> supplied_mean(dim, -0.25f);
Expand Down
Loading