Additional benchmarks - memory and search (#387)

Author: alonre24

src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h

@@ -50,3 +50,4 @@ INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_preferAdHocOptimization_Test)
 INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_runGCAPI_Test)
 
 INDEX_TEST_FRIEND_CLASS(BM_VecSimBasics)
+INDEX_TEST_FRIEND_CLASS(BM_VecSimCommon)

src/VecSim/index_factories/tiered_factory.h

@@ -34,7 +34,8 @@ VecSimIndex *NewIndex(const TieredIndexParams *params, HNSWIndex<DataType, DistT
                           .multi = params->primaryIndexParams->algoParams.hnswParams.multi,
                           .blockSize = params->primaryIndexParams->algoParams.hnswParams.blockSize};
 
-    AbstractIndexInitParams abstractInitParams = {.allocator = hnsw_index->getAllocator(),
+    std::shared_ptr<VecSimAllocator> flat_allocator = VecSimAllocator::newVecsimAllocator();
+    AbstractIndexInitParams abstractInitParams = {.allocator = flat_allocator,
                                                   .dim = bf_params.dim,
                                                   .vecType = bf_params.type,
                                                   .metric = bf_params.metric,

src/VecSim/vec_sim_tiered_index.h

@@ -93,7 +93,8 @@ class VecSimTieredIndex : public VecSimIndexInterface {
     static VecSimWriteMode getWriteMode() { return VecSimIndexInterface::asyncWriteMode; }
 
 #ifdef BUILD_TESTS
-    inline VecSimIndexAbstract<DistType> *getFlatbufferIndex() { return this->frontendIndex; }
+    inline VecSimIndexAbstract<DistType> *getFlatBufferIndex() { return this->frontendIndex; }
+    inline size_t getFlatBufferLimit() { return this->flatBufferLimit; }
 #endif
 
 private:

src/python_bindings/CMakeLists.txt

@@ -21,9 +21,9 @@ endif()
 
 add_subdirectory(${root}/src/VecSim VectorSimilarity)
 
-include_directories(${root}/src)
+include_directories(${root}/src ${root}/tests/utils)
 
-pybind11_add_module(VecSim bindings.cpp)
+pybind11_add_module(VecSim ../../tests/utils/mock_thread_pool.cpp bindings.cpp)
 
 target_link_libraries(VecSim VectorSimilarity)
 

src/python_bindings/bindings.cpp

@@ -16,10 +16,9 @@
 #include <thread>
 #include <VecSim/algorithms/hnsw/hnsw_single.h>
 #include <VecSim/algorithms/brute_force/brute_force_single.h>
-#include "tiered_index_mock.h"
+#include "mock_thread_pool.h"
 
 namespace py = pybind11;
-using namespace tiered_index_mock;
 
 // Helper function that iterates query results and wrap them in python numpy object -
 // a tuple of two 2D arrays: (labels, distances)
@@ -366,92 +365,64 @@ class PyHNSWLibIndex : public PyVecSimIndex {
 };
 
 class PyTieredIndex : public PyVecSimIndex {
-private:
+protected:
+    tieredIndexMock mock_thread_pool;
+
     VecSimIndexAbstract<float> *getFlatBuffer() {
         return reinterpret_cast<VecSimTieredIndex<float, float> *>(this->index.get())
-            ->getFlatbufferIndex();
+            ->getFlatBufferIndex();
     }
 
-protected:
-    JobQueue jobQueue;       // External queue that holds the jobs.
-    IndexExtCtx jobQueueCtx; // External context to be sent to the submit callback.
-    SubmitCB submitCb;       // A callback that submits an array of jobs into a given jobQueue.
-    size_t flatBufferLimit; // Maximum size allowed for the flat buffer. If flat buffer is full, use
-                            // in-place insertion.
-    bool run_thread;
-    std::bitset<MAX_POOL_SIZE> executions_status;
-
-    TieredIndexParams TieredIndexParams_Init() {
-        TieredIndexParams ret = {
-            .jobQueue = &this->jobQueue,
-            .jobQueueCtx = &this->jobQueueCtx,
-            .submitCb = this->submitCb,
-            .flatBufferLimit = this->flatBufferLimit,
+    TieredIndexParams getTieredIndexParams(size_t buffer_limit) {
+        // Create TieredIndexParams using the mock thread pool.
+        return TieredIndexParams{
+            .jobQueue = &(this->mock_thread_pool.jobQ),
+            .jobQueueCtx = this->mock_thread_pool.ctx,
+            .submitCb = tieredIndexMock::submit_callback,
+            .flatBufferLimit = buffer_limit,
         };
-
-        return ret;
     }
 
 public:
-    explicit PyTieredIndex(size_t BufferLimit = 3000000)
-        : submitCb(submit_callback), flatBufferLimit(BufferLimit), run_thread(true) {
-
-        for (size_t i = 0; i < THREAD_POOL_SIZE; i++) {
-            ThreadParams params(run_thread, executions_status, i, jobQueue);
-            thread_pool.emplace_back(thread_main_loop, params);
-        }
-    }
-
-    virtual ~PyTieredIndex() = 0;
+    explicit PyTieredIndex() { mock_thread_pool.init_threads(); }
 
     void WaitForIndex(size_t waiting_duration = 10) {
-        bool keep_wating = true;
-        while (keep_wating) {
-            std::this_thread::sleep_for(std::chrono::milliseconds(waiting_duration));
-            std::unique_lock<std::mutex> lock(queue_guard);
-            if (jobQueue.empty()) {
-                while (true) {
-                    if (executions_status.count() == 0) {
-                        keep_wating = false;
-                        break;
-                    }
-                    std::this_thread::sleep_for(std::chrono::milliseconds(waiting_duration));
-                }
-            }
-        }
+        mock_thread_pool.thread_pool_wait(waiting_duration);
     }
 
     size_t getFlatIndexSize() { return getFlatBuffer()->indexLabelCount(); }
 
-    static size_t GetThreadsNum() { return THREAD_POOL_SIZE; }
+    size_t getThreadsNum() { return mock_thread_pool.thread_pool_size; }
 
-    size_t getBufferLimit() { return flatBufferLimit; }
+    size_t getBufferLimit() {
+        return reinterpret_cast<VecSimTieredIndex<float, float> *>(this->index.get())
+            ->getFlatBufferLimit();
+    }
 };
 
-PyTieredIndex::~PyTieredIndex() { thread_pool_terminate(jobQueue, run_thread); }
 class PyTiered_HNSWIndex : public PyTieredIndex {
 public:
     explicit PyTiered_HNSWIndex(const HNSWParams &hnsw_params,
-                                const TieredHNSWParams &tiered_hnsw_params) {
+                                const TieredHNSWParams &tiered_hnsw_params, size_t buffer_limit) {
 
         // Create primaryIndexParams and specific params for hnsw tiered index.
         VecSimParams primary_index_params = {.algo = VecSimAlgo_HNSWLIB,
                                              .algoParams = {.hnswParams = HNSWParams{hnsw_params}}};
 
-        // create TieredIndexParams
-        TieredIndexParams tiered_params = TieredIndexParams_Init();
-
+        auto tiered_params = this->getTieredIndexParams(buffer_limit);
         tiered_params.primaryIndexParams = &primary_index_params;
         tiered_params.specificParams.tieredHnswParams = tiered_hnsw_params;
 
-        // create VecSimParams for TieredIndexParams
+        // Create VecSimParams for TieredIndexParams
         VecSimParams params = {.algo = VecSimAlgo_TIERED,
                                .algoParams = {.tieredParams = TieredIndexParams{tiered_params}}};
 
         this->index = std::shared_ptr<VecSimIndex>(VecSimIndex_New(&params), VecSimIndex_Free);
+
         // Set the created tiered index in the index external context.
-        this->jobQueueCtx.index_strong_ref = this->index;
+        this->mock_thread_pool.ctx->index_strong_ref = this->index;
     }
+
     size_t HNSWLabelCount() {
         return this->index->info().tieredInfo.backendCommonInfo.indexLabelCount;
     }
@@ -568,17 +539,17 @@ PYBIND11_MODULE(VecSim, m) {
              py::arg("radius"), py::arg("query_param") = nullptr, py::arg("num_threads") = -1);
 
     py::class_<PyTieredIndex, PyVecSimIndex>(m, "TieredIndex")
-        .def("wait_for_index", &PyTiered_HNSWIndex::WaitForIndex, py::arg("waiting_duration") = 10)
-        .def("get_curr_bf_size", &PyTiered_HNSWIndex::getFlatIndexSize)
-        .def("get_buffer_limit", &PyTiered_HNSWIndex::getBufferLimit)
-        .def_static("get_threads_num", &PyTieredIndex::GetThreadsNum);
+        .def("wait_for_index", &PyTieredIndex::WaitForIndex, py::arg("waiting_duration") = 10)
+        .def("get_curr_bf_size", &PyTieredIndex::getFlatIndexSize)
+        .def("get_buffer_limit", &PyTieredIndex::getBufferLimit)
+        .def("get_threads_num", &PyTieredIndex::getThreadsNum);
 
     py::class_<PyTiered_HNSWIndex, PyTieredIndex>(m, "Tiered_HNSWIndex")
-        .def(
-            py::init([](const HNSWParams &hnsw_params, const TieredHNSWParams &tiered_hnsw_params) {
-                return new PyTiered_HNSWIndex(hnsw_params, tiered_hnsw_params);
-            }),
-            py::arg("hnsw_params"), py::arg("tiered_hnsw_params"))
+        .def(py::init([](const HNSWParams &hnsw_params, const TieredHNSWParams &tiered_hnsw_params,
+                         size_t flat_buffer_size = DEFAULT_BLOCK_SIZE) {
+                 return new PyTiered_HNSWIndex(hnsw_params, tiered_hnsw_params, flat_buffer_size);
+             }),
+             py::arg("hnsw_params"), py::arg("tiered_hnsw_params"), py::arg("flat_buffer_size"))
         .def("hnsw_label_count", &PyTiered_HNSWIndex::HNSWLabelCount);
 
     py::class_<PyBFIndex, PyVecSimIndex>(m, "BFIndex")

src/python_bindings/tiered_index_mock.h

@@ -24,19 +24,17 @@ class BM_BatchIterator : public BM_VecSimIndex<index_type_t> {
     static void HNSW_VariableBatchSize(benchmark::State &st);
     static void HNSW_BatchesToAdhocBF(benchmark::State &st);
 
-    static void RunBatchedSearch_HNSW(benchmark::State &st, size_t &correct, size_t iter,
+    static void RunBatchedSearch_HNSW(benchmark::State &st, std::atomic_int &correct, size_t iter,
                                       size_t num_batches, size_t batch_size, size_t &total_res_num,
                                       size_t batch_increase_factor, size_t index_memory,
                                       double &memory_delta);
 };
 
 template <typename index_type_t>
-void BM_BatchIterator<index_type_t>::RunBatchedSearch_HNSW(benchmark::State &st, size_t &correct,
-                                                           size_t iter, size_t num_batches,
-                                                           size_t batch_size, size_t &total_res_num,
-                                                           size_t batch_increase_factor,
-                                                           size_t index_memory,
-                                                           double &memory_delta) {
+void BM_BatchIterator<index_type_t>::RunBatchedSearch_HNSW(
+    benchmark::State &st, std::atomic_int &correct, size_t iter, size_t num_batches,
+    size_t batch_size, size_t &total_res_num, size_t batch_increase_factor, size_t index_memory,
+    double &memory_delta) {
     VecSimBatchIterator *batchIterator = VecSimBatchIterator_New(
         INDICES.at(VecSimAlgo_HNSWLIB), QUERIES[iter % N_QUERIES].data(), nullptr);
     VecSimQueryResult_List accumulated_results[num_batches];
@@ -158,7 +156,7 @@ void BM_BatchIterator<index_type_t>::HNSW_FixedBatchSize(benchmark::State &st) {
     size_t num_batches = st.range(1);
     size_t total_res_num = num_batches * batch_size;
     size_t iter = 0;
-    size_t correct = 0;
+    std::atomic_int correct = 0;
     size_t index_memory = VecSimIndex_Info(INDICES[VecSimAlgo_HNSWLIB]).commonInfo.memory;
     double memory_delta = 0.0;
 
@@ -177,7 +175,7 @@ void BM_BatchIterator<index_type_t>::HNSW_VariableBatchSize(benchmark::State &st
     size_t num_batches = st.range(1);
     size_t total_res_num;
     size_t iter = 0;
-    size_t correct = 0;
+    std::atomic_int correct = 0;
     size_t index_memory = VecSimIndex_Info(INDICES[VecSimAlgo_HNSWLIB]).commonInfo.memory;
     double memory_delta = 0.0;
 
@@ -196,7 +194,7 @@ void BM_BatchIterator<index_type_t>::HNSW_BatchesToAdhocBF(benchmark::State &st)
     size_t num_batches = st.range(1);
     size_t total_res_num;
     size_t iter = 0;
-    size_t correct = 0;
+    std::atomic_int correct = 0;
     size_t index_memory = VecSimIndex_Info(INDICES[VecSimAlgo_HNSWLIB]).commonInfo.memory;
     double memory_delta = 0.0;