Update type-erased code to keep Ctx alive when calling into C++ to fix Vamana ingestion crash (#347)

Author: jparismorgan

apis/python/src/tiledb/vector_search/ingestion.py

@@ -1506,7 +1506,6 @@ def ingest_flat(
             ids_array.close()
 
     def ingest_vamana(
-        ctx,
         index_group_uri: str,
         source_uri: str,
         source_type: str,
@@ -1626,14 +1625,15 @@ def ingest_vamana(
             parts_array.close()
             ids_array.close()
 
-            # Now that we've ingested the vectors and their IDs, train the index with the data.
-            from tiledb.vector_search import _tiledbvspy as vspy
+        # Now that we've ingested the vectors and their IDs, train the index with the data.
+        from tiledb.vector_search import _tiledbvspy as vspy
 
-            index = vspy.IndexVamana(ctx, index_group_uri)
-            data = vspy.FeatureVectorArray(ctx, parts_array_uri, ids_array_uri)
-            index.train(data)
-            index.add(data)
-            index.write_index(ctx, index_group_uri, index_timestamp)
+        ctx = vspy.Ctx(config)
+        index = vspy.IndexVamana(ctx, index_group_uri)
+        data = vspy.FeatureVectorArray(ctx, parts_array_uri, ids_array_uri)
+        index.train(data)
+        index.add(data)
+        index.write_index(ctx, index_group_uri, index_timestamp)
 
     def write_centroids(
         centroids: np.ndarray,
@@ -2183,12 +2183,8 @@ def create_ingestion_dag(
             )
             return d
         elif index_type == "VAMANA":
-            from tiledb.vector_search import _tiledbvspy as vspy
-
-            ctx = vspy.Ctx(config)
             ingest_node = submit(
                 ingest_vamana,
-                ctx=ctx,
                 index_group_uri=index_group_uri,
                 source_uri=source_uri,
                 source_type=source_type,

apis/python/src/tiledb/vector_search/type_erased_module.cc

@@ -122,7 +122,10 @@ void init_type_erased_module(py::module_& m) {
       ;
 #endif
   py::class_<FeatureVector>(m, "FeatureVector", py::buffer_protocol())
-      .def(py::init<const tiledb::Context&, const std::string&>())
+      .def(
+          py::init<const tiledb::Context&, const std::string&>(),
+          py::keep_alive<1, 2>()  // IndexIVFFlat should keep ctx alive.
+          )
       .def(py::init<size_t, const std::string&>())
       .def(py::init<size_t, void*, const std::string&>())
       .def("dimension", &FeatureVector::dimension)
@@ -165,11 +168,17 @@ void init_type_erased_module(py::module_& m) {
       }));
 
   py::class_<FeatureVectorArray>(m, "FeatureVectorArray", py::buffer_protocol())
-      .def(py::init<const tiledb::Context&, const std::string&>())
-      .def(py::init<
-           const tiledb::Context&,
-           const std::string&,
-           const std::string&>())
+      .def(
+          py::init<const tiledb::Context&, const std::string&>(),
+          py::keep_alive<1, 2>()  // FeatureVectorArray should keep ctx alive.
+          )
+      .def(
+          py::init<
+              const tiledb::Context&,
+              const std::string&,
+              const std::string&>(),
+          py::keep_alive<1, 2>()  // FeatureVectorArray should keep ctx alive.
+          )
       .def(py::init<size_t, size_t, const std::string&, const std::string&>())
       .def("dimension", &FeatureVectorArray::dimension)
       .def("num_vectors", &FeatureVectorArray::num_vectors)
@@ -222,7 +231,10 @@ void init_type_erased_module(py::module_& m) {
       }));
 
   py::class_<IndexFlatL2>(m, "IndexFlatL2")
-      .def(py::init<const tiledb::Context&, const std::string&>())
+      .def(
+          py::init<const tiledb::Context&, const std::string&>(),
+          py::keep_alive<1, 2>()  // IndexFlatL2 should keep ctx alive.
+          )
       .def("add", &IndexFlatL2::add)
       .def("add_with_ids", &IndexFlatL2::add_with_ids)
       .def("train", &IndexFlatL2::train)
@@ -248,7 +260,10 @@ void init_type_erased_module(py::module_& m) {
       }));
 
   py::class_<IndexVamana>(m, "IndexVamana")
-      .def(py::init<const tiledb::Context&, const std::string&>())
+      .def(
+          py::init<const tiledb::Context&, const std::string&>(),
+          py::keep_alive<1, 2>()  // IndexVamana should keep ctx alive.
+          )
       .def(
           "__init__",
           [](IndexVamana& instance, py::kwargs kwargs) {
@@ -288,6 +303,7 @@ void init_type_erased_module(py::module_& m) {
              const std::string& storage_version) {
             index.write_index(ctx, group_uri, timestamp, storage_version);
           },
+          py::keep_alive<1, 2>(),  // IndexVamana should keep ctx alive.
           py::arg("ctx"),
           py::arg("group_uri"),
           py::arg("timestamp") = py::none(),
@@ -300,7 +316,10 @@ void init_type_erased_module(py::module_& m) {
       .def("dimension", &IndexVamana::dimension);
 
   py::class_<IndexIVFFlat>(m, "IndexIVFFlat")
-      .def(py::init<const tiledb::Context&, const std::string&>())
+      .def(
+          py::init<const tiledb::Context&, const std::string&>(),
+          py::keep_alive<1, 2>()  // IndexIVFFlat should keep ctx alive.
+          )
       .def(
           "__init__",
           [](IndexIVFFlat& instance, py::kwargs kwargs) {

documentation/Building.md

@@ -130,3 +130,35 @@ pre-commit run --all-files
 # Run on single file and make formatting changes:
 pre-commit run --files path/to/file.py
 ```
+
+# Debugging C++ crashes in Python
+
+If you are seeing C++ crashes only when running Python, you can debug by running `lldb -f python -- apis/python/test/crash.py` and then `(lldb) run`. By default tiledbvspy does not have debug symbols, so you'll see a stack trace like:
+
+```bash
+(lldb) bt
+* thread #1, queue = 'com.apple.main-thread', stop reason = EXC_BAD_ACCESS (code=1, address=0x30001356a8f20)
+  * frame #0: 0x0000000124a8d7c0 _tiledbvspy.cpython-39-darwin.so`___lldb_unnamed_symbol832 + 88
+    frame #1: 0x0000000124a8d954 _tiledbvspy.cpython-39-darwin.so`___lldb_unnamed_symbol834 + 32
+    frame #2: 0x0000000124ac0c64 _tiledbvspy.cpython-39-darwin.so`___lldb_unnamed_symbol1139 + 104
+    frame #3: 0x0000000124b6ff00 _tiledbvspy.cpython-39-darwin.so`___lldb_unnamed_symbol2869 + 32
+    frame #4: 0x0000000124b20150 _tiledbvspy.cpython-39-darwin.so`___lldb_unnamed_symbol1879 + 52
+    frame #5: 0x0000000124c47354 _tiledbvspy.cpython-39-darwin.so`___lldb_unnamed_symbol3648 + 100
+    frame #6: 0x000000012073b4d0 cc.cpython-39-darwin.so`pybind11::detail::clear_instance(_object*) + 392
+    frame #7: 0x000000012073b1c8 cc.cpython-39-darwin.so`pybind11_object_dealloc + 44
+    frame #8: 0x000000010007feb0 python`frame_dealloc + 840
+    frame #9: 0x000000010005d260 python`_PyFunction_Vectorcall + 476
+    frame #10: 0x000000010016ba6c python`_PyEval_EvalFrameDefault + 27688
+    frame #11: 0x0000000100162dfc python`_PyEval_EvalCode + 700
+    frame #12: 0x00000001001d7e04 python`run_mod + 188
+    frame #13: 0x00000001001d6674 python`pyrun_file + 176
+    frame #14: 0x00000001001d60e4 python`pyrun_simple_file + 352
+    frame #15: 0x00000001001d5f44 python`PyRun_SimpleFileExFlags + 64
+    frame #16: 0x00000001001f883c python`pymain_run_file + 264
+    frame #17: 0x00000001001f7e84 python`pymain_run_python + 360
+    frame #18: 0x00000001001f7cc4 python`Py_RunMain + 40
+    frame #19: 0x0000000100006a20 python`main + 56
+    frame #20: 0x00000001832510e0 dyld`start + 2360
+```
+
+To enable symbols, you can update `src/CMakeLists.txt` and add `set(CMAKE_BUILD_TYPE "Debug")` after it adds `set(CMAKE_BUILD_TYPE ...`.