Fix memory corruption issue in ThreadBatchDataLoader

- Return a shared pointer of CuImageIterator, taking ownership of CuImageIterator instance, instead of returning a copy of CuImageIterator - Do not create new instance for CuImageIterator when __iter__ is called. Instead, return self. - Wait until all threads are finished in ThreadBatchDataLoader destructor Signed-off-by: Gigon Bae <[email protected]>
rapidsai · Oct 27, 2023 · 4384f57 · 4384f57
1 parent 5f78da9
commit 4384f57
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/cpp/src/loader/thread_batch_data_loader.cpp b/cpp/src/loader/thread_batch_data_loader.cpp
@@ -1,6 +1,6 @@
 /*
  * Apache License, Version 2.0
- * Copyright 2021 NVIDIA Corporation
+ * Copyright 2021-2023 NVIDIA Corporation
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -86,6 +86,9 @@ ThreadBatchDataLoader::ThreadBatchDataLoader(LoadFunc load_func,
 
 ThreadBatchDataLoader::~ThreadBatchDataLoader()
 {
+    // Wait until all tasks are done.
+    while (wait_batch() > 0);
+
     cucim::io::DeviceType device_type = out_device_.type();
     for (auto& raster_ptr : raster_data_)
     {

diff --git a/python/pybind11/cucim_py.cpp b/python/pybind11/cucim_py.cpp
@@ -200,7 +200,7 @@ PYBIND11_MODULE(_cucim, m)
             },
             py::call_guard<py::gil_scoped_release>());
 
-    py::class_<CuImageIterator<CuImage>>(m, "CuImageIterator") //
+    py::class_<CuImageIterator<CuImage>, std::shared_ptr<CuImageIterator<CuImage>>>(m, "CuImageIterator") //
         .def(py::init<std::shared_ptr<CuImage>, bool>(), doc::CuImageIterator::doc_CuImageIterator,
              py::arg("cuimg"), //
              py::arg("ending") = false, py::call_guard<py::gil_scoped_release>())
@@ -212,8 +212,8 @@ PYBIND11_MODULE(_cucim, m)
             py::call_guard<py::gil_scoped_release>())
         .def(
             "__iter__", //
-            [](CuImageIterator<CuImage>& it) { //
-                return CuImageIterator<CuImage>(it); //
+            [](CuImageIterator<CuImage>* it) { //
+                return it; //
             }, //
             py::call_guard<py::gil_scoped_release>())
         .def("__next__", &py_cuimage_iterator_next, py::call_guard<py::gil_scoped_release>())
@@ -517,11 +517,11 @@ py::object py_read_region(const CuImage& cuimg,
     auto loader = region_ptr->loader();
     if (batch_size > 1 || (loader && loader->size() > 1))
     {
-        auto iter_ptr = region_ptr->begin();
+        auto iter_ptr = std::make_shared<CuImageIterator<CuImage>>(region_ptr->shared_from_this());
 
         py::gil_scoped_acquire scope_guard;
 
-        py::object iter = py::cast(iter_ptr);
+        py::object iter = py::cast(iter_ptr, py::return_value_policy::take_ownership);
 
         return iter;
     }