Log stats on array openings and query submissions. (#52)

teo-tsirpanis · web-flow · commit 2ed2004b7bb4 · 2023-06-27T23:10:00.000+03:00
* Add wrappers over opening arrays and submitting queries that optionally log stats.

* Use the wrapper functions throughout the codebase.

* Support collecting stats in the Python API and add functions to enable it.

* Log the name of the calling function alongside the stats.

* Support customizing the file the stats are written and default it to stdout.

* Log Core stats in JSON and add specialized dump and reset Python API functions.

* Support logging stats in the CLI programs.

* Use `raw_dump` to get the stats.

`dump` and `raw_dump` do the same, but only the latter is documented to return JSON.
diff --git a/apis/python/CMakeLists.txt b/apis/python/CMakeLists.txt
@@ -49,7 +49,7 @@ target_link_libraries(${VSPY_TARGET_NAME}
     kmeans_lib
     TileDB::tiledb_shared)
 
-target_compile_definitions(${VSPY_TARGET_NAME} PRIVATE VERSION_INFO=${PROJECT_VERSION})
+target_compile_definitions(${VSPY_TARGET_NAME} PRIVATE TILEDBVS_ENABLE_STATS VERSION_INFO=${PROJECT_VERSION})
 
 if (APPLE)
   set_target_properties(${VSPY_TARGET_NAME} PROPERTIES INSTALL_RPATH "@loader_path/lib")
diff --git a/apis/python/src/tiledb/vector_search/module.cc b/apis/python/src/tiledb/vector_search/module.cc
@@ -14,6 +14,9 @@ using Ctx = tiledb::Context;
 bool global_debug = true;
 double global_time_of_interest;
 
+bool enable_stats = false;
+std::vector<json> core_stats;
+
 PYBIND11_MAKE_OPAQUE(std::vector<uint32_t>);
 PYBIND11_MAKE_OPAQUE(std::vector<uint64_t>);
 
@@ -249,6 +252,24 @@ PYBIND11_MODULE(_tiledbvspy, m) {
         return validate_top_k(top_k, ground_truth);
       });
 
+  m.def("stats_enable", []() {
+    enable_stats = true;
+    tiledb::Stats::enable();
+  });
+
+  m.def("stats_disable", []() {
+    enable_stats = false;
+    tiledb::Stats::disable();
+  });
+
+  m.def("stats_reset", []() {
+    core_stats.clear();
+  });
+
+  m.def("stats_dump", []() {
+    return json{core_stats}.dump();
+  });
+
   declare_kmeans_query<uint8_t>(m, "u8");
   declare_kmeans_query<float>(m, "f32");
 
diff --git a/src/include/detail/linalg/tdb_helpers.h b/src/include/detail/linalg/tdb_helpers.h
@@ -0,0 +1,81 @@
+/**
+ * @file   tiledb_helpers.h
+ *
+ * @section LICENSE
+ *
+ * The MIT License
+ *
+ * @copyright Copyright (c) 2023 TileDB, Inc.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ *
+ * @section DESCRIPTION
+ *
+ * Helper functions for certain TileDB operations.
+ *
+ */
+
+#ifndef TILEDB_HELPERS_H
+#define TILEDB_HELPERS_H
+
+#include "stats.h"
+#include <tiledb/tiledb>
+
+namespace tiledb_helpers {
+
+/**
+ * @brief Opens a TileDB array and displays stats to stderr.
+ *
+ * Stats are only collected if the TILEDBVS_ENABLE_STATS symbol is
+ * defined, and a variable named enable_stats is set to true.
+ * The stats are written to a FILE* specified by the variable named stats_file.
+ *
+ * @param function_name The name of the function calling this. You can use the tdb_func__ macro.
+ * @param ctx The TileDB context to use.
+ * @param uri The URI of the array to open.
+ * @param query_type The mode to open the array.
+ */
+inline tiledb::Array open_array(const std::string &function_name,
+                                const tiledb::Context &ctx,
+                                const std::string &uri,
+                                tiledb_query_type_t query_type) {
+  StatsCollectionScope stats_scope(uri, function_name, "open_array");
+  return tiledb::Array(ctx, uri, query_type);
+}
+
+/**
+ * @brief Submits a TileDB query and displays stats to stderr.
+ *
+ * Stats are only collected if the TILEDBVS_ENABLE_STATS symbol is
+ * defined, and a variable named enable_stats is set to true.
+ * The stats are written to a FILE* specified by the variable named stats_file.
+ *
+ * @param function_name The name of the function calling this. You can use the tdb_func__ macro.
+ * @param query The query to submit.
+ */
+inline void submit_query(const std::string &function_name,
+                         const std::string &uri,
+                         tiledb::Query &query) {
+  StatsCollectionScope stats_scope(uri, function_name, "submit_query");
+  query.submit();
+}
+
+} // namespace tiledb_helpers
+
+#endif
diff --git a/src/include/detail/linalg/tdb_io.h b/src/include/detail/linalg/tdb_io.h
@@ -85,7 +85,7 @@ void write_matrix(
       0, (int)A.num_rows() - 1, 0, (int)A.num_cols() - 1};
 
   // Open array for writing
-  tiledb::Array array(ctx, uri, TILEDB_WRITE);
+  tiledb::Array array = tiledb_helpers::open_array(tdb_func__, ctx, uri, TILEDB_WRITE);
 
   tiledb::Subarray subarray(ctx, array);
   subarray.set_subarray(subarray_vals);
@@ -96,7 +96,7 @@ void write_matrix(
       .set_data_buffer(
           "values", &A(0, 0), (int)A.num_rows() * (int)A.num_cols())
       .set_subarray(subarray);
-  query.submit();
+  tiledb_helpers::submit_query(tdb_func__, uri, query);
 
   array.close();
 }
@@ -132,7 +132,7 @@ void write_vector(
   std::vector<int32_t> subarray_vals{0, (int)size(v) - 1};
 
   // Open array for writing
-  tiledb::Array array(ctx, uri, TILEDB_WRITE);
+  tiledb::Array array = tiledb_helpers::open_array(tdb_func__, ctx, uri, TILEDB_WRITE);
 
   tiledb::Subarray subarray(ctx, array);
   subarray.set_subarray(subarray_vals);
@@ -141,7 +141,7 @@ void write_vector(
   query.set_layout(TILEDB_ROW_MAJOR)
       .set_data_buffer("values", v)
       .set_subarray(subarray);
-  query.submit();
+  tiledb_helpers::submit_query(tdb_func__, uri, query);
 
   array.close();
 }
@@ -157,7 +157,7 @@ std::vector<T> read_vector(const tiledb::Context& ctx, const std::string& uri) {
     std::cerr << "# Reading std::vector: " << uri << std::endl;
   }
 
-  auto array_ = tiledb::Array{ctx, uri, TILEDB_READ};
+  tiledb::Array array_ = tiledb_helpers::open_array(tdb_func__, ctx, uri, TILEDB_READ);
   auto schema_ = array_.schema();
 
   using domain_type = int32_t;
@@ -189,7 +189,7 @@ std::vector<T> read_vector(const tiledb::Context& ctx, const std::string& uri) {
   tiledb::Query query(ctx, array_);
   query.set_subarray(subarray).set_data_buffer(
       attr_name, data_.data(), vec_rows_);
-  query.submit();
+  tiledb_helpers::submit_query(tdb_func__, uri, query);
   _memory_data.insert_entry(tdb_func__, vec_rows_ * sizeof(T));
 
   array_.close();
diff --git a/src/include/detail/linalg/tdb_matrix.h b/src/include/detail/linalg/tdb_matrix.h
@@ -72,6 +72,7 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
   log_timer constructor_timer{"tdbMatrix constructor"};
 
   std::reference_wrapper<const tiledb::Context> ctx_;
+  std::string uri_;
   tiledb::Array array_;
   tiledb::ArraySchema schema_;
   std::unique_ptr<T[]> backing_data_;
@@ -199,7 +200,8 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
       size_t col_begin,
       size_t col_end)  // noexcept
       : ctx_{ctx}
-      , array_{ctx, uri, TILEDB_READ}
+      , uri_{uri}
+      , array_{tiledb_helpers::open_array(tdb_func__, ctx, uri, TILEDB_READ)}
       , schema_{array_.schema()} {
     constructor_timer.stop();
     scoped_timer _{tdb_func__ + uri};
@@ -282,7 +284,7 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
     query.set_subarray(subarray)
         .set_layout(layout_order)
         .set_data_buffer(attr_name, data_.get(), num_rows * num_cols);
-    query.submit();
+    tiledb_helpers::submit_query(tdb_func__, uri, query);
     _memory_data.insert_entry(tdb_func__, num_rows * num_cols * sizeof(T));
 
     // assert(tiledb::Query::Status::COMPLETE == query.query_status());
@@ -395,7 +397,7 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
         query.set_subarray(subarray)
             .set_layout(layout_order)
             .set_data_buffer(attr_name, ptr, num_elements);
-        query.submit();
+        tiledb_helpers::submit_query(tdb_func__, uri, query);
         _memory_data.insert_entry(tdb_func__, num_elements * sizeof(T));
 
         // assert(tiledb::Query::Status::COMPLETE == query.query_status());
@@ -417,7 +419,7 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
        */
       auto attr_idx = 0;
 
-      auto ids_array_ = tiledb::Array{ctx_, id_uri, TILEDB_READ};
+      tiledb::Array ids_array_ = tiledb_helpers::open_array(tdb_func__, ctx_, id_uri, TILEDB_READ);
       auto ids_schema_ = ids_array_.schema();
 
       auto attr_num{ids_schema_.attribute_num()};
@@ -449,7 +451,7 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
         auto ptr = part_ids.data() + offset;
         query.set_subarray(subarray).set_data_buffer(
             attr_name, ptr, num_elements);
-        query.submit();
+        tiledb_helpers::submit_query(tdb_func__, uri, query);
         _memory_data.insert_entry(tdb_func__, num_elements * sizeof(T));
 
         if (tiledb::Query::Status::COMPLETE != query.query_status()) {
@@ -550,7 +552,7 @@ class tdbMatrix : public Matrix<T, LayoutPolicy, I> {
     query.set_subarray(subarray)
         .set_layout(layout_order)
         .set_data_buffer(attr_name, this_data, read_size);
-    query.submit();
+    tiledb_helpers::submit_query(tdb_func__, uri_, query);
     _memory_data.insert_entry(tdb_func__, read_size * sizeof(T));
 
     // assert(tiledb::Query::Status::COMPLETE == query.query_status());
diff --git a/src/include/detail/linalg/tdb_partitioned_matrix.h b/src/include/detail/linalg/tdb_partitioned_matrix.h
@@ -106,6 +106,7 @@ class tdbPartitionedMatrix : public Matrix<T, LayoutPolicy, I> {
   log_timer constructor_timer{"tdbPartitionedMatrix constructor"};
 
   std::reference_wrapper<const tiledb::Context> ctx_;
+  std::string uri_;
   tiledb::Array array_;
   tiledb::ArraySchema schema_;
   std::unique_ptr<T[]> backing_data_;
@@ -178,9 +179,10 @@ class tdbPartitionedMatrix : public Matrix<T, LayoutPolicy, I> {
       size_t nthreads)
       : constructor_timer{tdb_func__ + std::string{" constructor"}}
       , ctx_{ctx}
-      , array_{ctx_, uri, TILEDB_READ}
+      , uri_{uri}
+      , array_{tiledb_helpers::open_array(tdb_func__, ctx_, uri, TILEDB_READ)}
       , schema_{array_.schema()}
-      , ids_array_{ctx_, ids_uri, TILEDB_READ}
+      , ids_array_{tiledb_helpers::open_array(tdb_func__, ctx_, ids_uri, TILEDB_READ)}
       , ids_schema_{ids_array_.schema()}
       , indices_{std::move(in_indices)}
       , parts_{in_parts}
@@ -345,7 +347,7 @@ class tdbPartitionedMatrix : public Matrix<T, LayoutPolicy, I> {
       query.set_subarray(subarray)
           .set_layout(layout_order)
           .set_data_buffer(attr_name, ptr, col_count * dimension);
-      query.submit();
+      tiledb_helpers::submit_query(tdb_func__, uri_, query);
       _memory_data.insert_entry(tdb_func__, col_count * dimension * sizeof(T));
 
       // assert(tiledb::Query::Status::COMPLETE == query.query_status());
diff --git a/src/include/linalg.h b/src/include/linalg.h
@@ -46,6 +46,7 @@
 #include "detail/linalg/choose_blas.h"
 #include "detail/linalg/linalg_defs.h"
 #include "detail/linalg/matrix.h"
+#include "detail/linalg/tdb_helpers.h"
 #include "detail/linalg/tdb_io.h"
 #include "detail/linalg/tdb_matrix.h"
 #include "detail/linalg/tdb_partitioned_matrix.h"
diff --git a/src/include/stats.h b/src/include/stats.h
@@ -53,7 +53,47 @@
 
 using json = nlohmann::json;
 
+// Make stats support opt-in to avoid requiring to define an enable_stats variable on all projects.
+#ifdef TILEDBVS_ENABLE_STATS
+extern bool enable_stats;
+extern std::vector<json> core_stats;
+#endif
 
+class StatsCollectionScope final {
+public:
+  explicit StatsCollectionScope(const std::string &uri, const std::string& function, const std::string &operation_type) {
+#ifdef TILEDBVS_ENABLE_STATS
+    if (!enable_stats)
+      return;
+    tiledb::Stats::reset();
+    uri_ = uri;
+    function_ = function;
+    operation_type_ = operation_type;
+#else
+    std::ignore = std::make_tuple(uri, function, operation_type);
+#endif
+  }
+
+  ~StatsCollectionScope() {
+#ifdef TILEDBVS_ENABLE_STATS
+    if (!enable_stats)
+      return;
+    std::string stats_str;
+    tiledb::Stats::raw_dump(&stats_str);
+    core_stats.push_back({
+      {"uri", uri_},
+      {"function", function_},
+      {"operation_type", operation_type_},
+      {"stats", json::parse(stats_str)}
+    });
+#endif
+  }
+
+#ifdef TILEDBVS_ENABLE_STATS
+private:
+  std::string uri_, function_, operation_type_;
+#endif
+};
 
 auto dump_logs = [](std::ostream& output,
                     const std::string algorithm,
diff --git a/src/include/test/time_open.cc b/src/include/test/time_open.cc
@@ -1,13 +1,14 @@
 
 #include <string>
 #include <tiledb/tiledb>
+#include "detail/linalg/tdb_helpers.h"
 #include "utils/timer.h"
 
 void open_array(const std::string& uri) {
   scoped_timer _{"open_array " + uri};
 
   tiledb::Context ctx;
-  tiledb::Array array(ctx, uri, TILEDB_READ);
+  tiledb::Array array = tiledb_helpers::open_array(tdb_func__, ctx, uri, TILEDB_READ);
 
   scoped_timer _2{"get_schema portion"};
   tiledb::ArraySchema schema = array.schema();
diff --git a/src/include/test/unit_slicing.cc b/src/include/test/unit_slicing.cc
@@ -50,7 +50,7 @@ TEST_CASE("slice", "[linalg][ci-skip]") {
   std::vector<int> data2_(288);
   std::vector<float> value_(288);
 
-  tiledb::Array array_{ctx_, uri, TILEDB_READ};
+  tiledb::Array array_ = tiledb_helpers::open_array(tdb_func__, ctx_, uri, TILEDB_READ);
   tiledb::ArraySchema schema_{array_.schema()};
   tiledb::Query query(ctx_, array_);
 
@@ -66,7 +66,7 @@ TEST_CASE("slice", "[linalg][ci-skip]") {
       .set_data_buffer("rows", data_.data(), 288)
       .set_data_buffer("a", value_.data(), 288);
 
-  query.submit();
+  tiledb_helpers::submit_query(tdb_func__, uri, query);
 
   for (int i = 0; i < 135; i++) {
     std::cout << data_[i] << ", " << data2_[i] << ": " << value_[i]
diff --git a/src/include/utils/logging.h b/src/include/utils/logging.h
diff --git a/src/src/CMakeLists.txt b/src/src/CMakeLists.txt
diff --git a/src/src/flat_l2.cc b/src/src/flat_l2.cc
diff --git a/src/src/index.cc b/src/src/index.cc
diff --git a/src/src/ivf_flat.cc b/src/src/ivf_flat.cc
