Merge pull request #129 from bluescarni/pr/numpy_mem_opt

Make the custom numpy memory handler an opt-in feature

.github/workflows/gh_actions_ci.yml

@@ -102,7 +102,7 @@ jobs:
       - uses: conda-incubator/setup-miniconda@v2
         with:
           auto-update-conda: true
-          python-version: 3.9
+          python-version: "3.10"
           channels: conda-forge
           channel-priority: strict
       - name: Build
@@ -121,8 +121,9 @@ jobs:
           cd build
           cmake .. -G "Visual Studio 16 2019" -A x64 -DCMAKE_PREFIX_PATH=C:\Miniconda\envs\test\Library -DCMAKE_INSTALL_PREFIX=C:\Miniconda\envs\test\Library -DBoost_NO_BOOST_CMAKE=ON -DHEYOKA_PY_ENABLE_IPO=yes
           cmake --build . --config Release --target install
-          cd c:\
-          python -c "import numba; from heyoka import test; test.run_test_suite()"
+          cd ..
+          cd tools
+          python ci_test_runner.py --with-numba
   binder_cache:
     runs-on: ubuntu-latest
     if: github.ref == 'refs/heads/main'

doc/changelog.rst

@@ -40,6 +40,9 @@ New
 Changes
 ~~~~~~~
 
+- The custom NumPy memory manager that prevents memory leaks
+  with ``real`` arrays is now disabled by default
+  (`#129 <https://github.com/bluescarni/heyoka.py/pull/129>`__).
 - The step callbacks are now deep-copied in multithreaded
   :ref:`ensemble propagations <ensemble_prop>`
   rather then being shared among threads. The aim of this change

heyoka/expose_real.cpp

@@ -2243,8 +2243,8 @@ void pyreal_ensure_array(py::array &arr, mpfr_prec_t prec)
 
 void expose_real(py::module_ &m)
 {
-    // Install the custom NumPy memory management functions.
-    install_custom_numpy_mem_handler();
+    // Setup the custom NumPy memory management functions.
+    setup_custom_numpy_mem_handler(m);
 
     // Fill out the entries of py_real_type.
     py_real_type.tp_base = &PyGenericArrType_Type;

heyoka/numpy_memory.cpp

@@ -11,10 +11,12 @@
 #include <cstdlib>
 #include <cstring>
 #include <functional>
+#include <iostream>
 #include <map>
 #include <memory>
 #include <mutex>
 #include <new>
+#include <optional>
 #include <tuple>
 #include <typeindex>
 #include <utility>
@@ -469,51 +471,83 @@ PyDataMem_Handler npy_custom_mem_handler
        1,
        {nullptr, numpy_custom_malloc, numpy_custom_calloc, numpy_custom_realloc, numpy_custom_free}};
 
-// Flag to signal if the default NumPy memory handler
-// has been overriden by npy_custom_mem_handler.
+// When our custom NumPy memory handler is installed, the original
+// memory handler is saved in this variable.
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-bool numpy_mh_overridden = false;
+std::optional<py::object> numpy_orig_mem_handler;
 
 } // namespace
 
 } // namespace detail
 
-// Helper to install the custom memory handling functions.
-// This can be called multiple times, all invocations past the
-// first one are no-ops.
-void install_custom_numpy_mem_handler()
+// Helper to setup the custom memory handling functions.
+// NOTE: these functions are NOT thread safe, this needs
+// to be highlighted in the docs.
+void setup_custom_numpy_mem_handler(py::module_ &m)
 {
-    if (detail::numpy_mh_overridden) {
-        // Don't do anything if we have overridden
-        // the memory management functions already.
-        return;
-    }
+    m.def("install_custom_numpy_mem_handler", []() {
+        if (detail::numpy_orig_mem_handler) {
+            // Don't do anything if we have already overridden
+            // the memory management functions.
+            return;
+        }
 
-    // NOTE: in principle here we could fetch the original memory handling
-    // capsule (which is also called "mem_handler"), and re-use the original
-    // memory functions in our implementations, instead of calling malloc/calloc/etc.
-    // This would make our custom implementations "good citizens", in the sense
-    // that we would respect existing custom memory allocating routines instead of
-    // outright overriding and ignoring them. Probably this is not an immediate concern
-    // as the memory management API is rather new, but it is something we should
-    // keep in mind moving forward.
-    auto *new_mem_handler = PyCapsule_New(&detail::npy_custom_mem_handler, "mem_handler", nullptr);
-    if (new_mem_handler == nullptr) {
-        // NOTE: if PyCapsule_New() fails, it already sets the error flag.
-        throw pybind11::error_already_set();
-    }
+        // NOTE: in principle here we could fetch the original memory handling
+        // capsule (which is also called "mem_handler"), and re-use the original
+        // memory functions in our implementations, instead of calling malloc/calloc/etc.
+        // This would make our custom implementations "good citizens", in the sense
+        // that we would respect existing custom memory allocating routines instead of
+        // outright overriding and ignoring them. Probably this is not an immediate concern
+        // as the memory management API is rather new, but it is something we should
+        // keep in mind moving forward.
+        auto *new_mem_handler = PyCapsule_New(&detail::npy_custom_mem_handler, "mem_handler", nullptr);
+        if (new_mem_handler == nullptr) {
+            // NOTE: if PyCapsule_New() fails, it already sets the error flag.
+            throw pybind11::error_already_set();
+        }
 
-    auto *old = PyDataMem_SetHandler(new_mem_handler);
-    Py_DECREF(new_mem_handler);
-    if (old == nullptr) {
-        // NOTE: if PyDataMem_SetHandler() fails, it already sets the error flag.
-        throw pybind11::error_already_set();
-    }
+        auto *old = PyDataMem_SetHandler(new_mem_handler);
+        Py_DECREF(new_mem_handler);
+        if (old == nullptr) {
+            // NOTE: if PyDataMem_SetHandler() fails, it already sets the error flag.
+            throw pybind11::error_already_set();
+        }
 
-    Py_DECREF(old);
+        // Store the original memory handler.
+        // NOTE: we use reinterpret_steal because PyDataMem_SetHandler() returned a new reference.
+        detail::numpy_orig_mem_handler.emplace(py::reinterpret_steal<py::object>(py::handle(old)));
+    });
+
+    m.def("remove_custom_numpy_mem_handler", []() {
+        if (!detail::numpy_orig_mem_handler) {
+            // Don't do anything if we have not overridden
+            // the memory management functions yet.
+            return;
+        }
+
+        // Try to restore the original memory handler.
+        auto *tmp = PyDataMem_SetHandler(detail::numpy_orig_mem_handler->ptr());
+        if (tmp == nullptr) {
+            // NOTE: if PyDataMem_SetHandler() fails, it already sets the error flag.
+            throw pybind11::error_already_set();
+        }
+
+        // The original memory handler was restored successfully. As cleanup actions we need to:
+        // - decrease the refcount of tmp, the handler that we just replaced, and
+        // - destroy the content of numpy_orig_mem_handler (which includes decreasing the refcount).
+        Py_DECREF(tmp);
+        detail::numpy_orig_mem_handler.reset();
+    });
 
-    // Set the flag.
-    detail::numpy_mh_overridden = true;
+    // NOTE: as usual, ensure that Pythonic global variables
+    // are cleaned up before shutting down the interpreter.
+    auto atexit = py::module_::import("atexit");
+    atexit.attr("register")(py::cpp_function([]() {
+#if !defined(NDEBUG)
+        std::cout << "Cleaning up the custom NumPy memory management data" << std::endl;
+#endif
+        detail::numpy_orig_mem_handler.reset();
+    }));
 }
 
 } // namespace heyoka_py

heyoka/numpy_memory.hpp

@@ -18,6 +18,8 @@
 #include <typeinfo>
 #include <utility>
 
+#include <pybind11/pybind11.h>
+
 #include "common_utils.hpp"
 
 #if defined(__GNUC__)
@@ -107,7 +109,7 @@ struct numpy_mem_metadata {
 
 std::pair<unsigned char *, numpy_mem_metadata *> get_memory_metadata(const void *) noexcept;
 
-void install_custom_numpy_mem_handler();
+void setup_custom_numpy_mem_handler(pybind11::module_ &);
 
 } // namespace heyoka_py
 

heyoka/test.py

@@ -2854,25 +2854,31 @@ def run_test_suite():
 
     retval = 0
 
-    suite = _ut.TestLoader().loadTestsFromTestCase(taylor_add_jet_test_case)
-    suite.addTest(_ut.makeSuite(_test_scalar_integrator.scalar_integrator_test_case))
-    suite.addTest(_ut.makeSuite(_test_batch_integrator.batch_integrator_test_case))
-    suite.addTest(_ut.makeSuite(_test_dtens.dtens_test_case))
-    suite.addTest(_ut.makeSuite(_test_mp.mp_test_case))
-    suite.addTest(_ut.makeSuite(_test_model.model_test_case))
-    suite.addTest(_ut.makeSuite(_test_real.real_test_case))
-    suite.addTest(_ut.makeSuite(_test_real128.real128_test_case))
-    suite.addTest(_ut.makeSuite(_test_cfunc.cfunc_test_case))
-    suite.addTest(_ut.makeSuite(_test_ensemble.ensemble_test_case))
-    suite.addTest(_ut.makeSuite(s11n_backend_test_case))
-    suite.addTest(_ut.makeSuite(recommended_simd_size_test_case))
-    suite.addTest(_ut.makeSuite(c_output_test_case))
-    suite.addTest(_ut.makeSuite(_test_expression.expression_test_case))
-    suite.addTest(_ut.makeSuite(llvm_state_test_case))
-    suite.addTest(_ut.makeSuite(event_classes_test_case))
-    suite.addTest(_ut.makeSuite(event_detection_test_case))
-    suite.addTest(_ut.makeSuite(kepE_test_case))
-    suite.addTest(_ut.makeSuite(sympy_test_case))
+    tl = _ut.TestLoader()
+
+    suite = tl.loadTestsFromTestCase(taylor_add_jet_test_case)
+    suite.addTest(
+        tl.loadTestsFromTestCase(_test_scalar_integrator.scalar_integrator_test_case)
+    )
+    suite.addTest(
+        tl.loadTestsFromTestCase(_test_batch_integrator.batch_integrator_test_case)
+    )
+    suite.addTest(tl.loadTestsFromTestCase(_test_dtens.dtens_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(_test_mp.mp_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(_test_model.model_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(_test_real.real_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(_test_real128.real128_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(_test_cfunc.cfunc_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(_test_ensemble.ensemble_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(s11n_backend_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(recommended_simd_size_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(c_output_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(_test_expression.expression_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(llvm_state_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(event_classes_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(event_detection_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(kepE_test_case))
+    suite.addTest(tl.loadTestsFromTestCase(sympy_test_case))
 
     test_result = _ut.TextTestRunner(verbosity=2).run(suite)
 

tools/ci_test_runner.py

@@ -0,0 +1,20 @@
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Test suite runner.")
+    parser.add_argument("--with-numba", action="store_true")
+
+    args = parser.parse_args()
+
+    import heyoka
+
+    if args.with_numba:
+        import numba
+
+    heyoka.test.run_test_suite()
+
+    if hasattr(heyoka, "real"):
+        heyoka.install_custom_numpy_mem_handler()
+        heyoka.test.run_test_suite()
+        heyoka.remove_custom_numpy_mem_handler()
+        heyoka.test.run_test_suite()

tools/circleci_conda_heyoka_head_310.sh

@@ -39,9 +39,9 @@ cd build
 cmake ../ -DCMAKE_INSTALL_PREFIX=$deps_dir -DCMAKE_PREFIX_PATH=$deps_dir -DCMAKE_BUILD_TYPE=Debug -DHEYOKA_PY_ENABLE_IPO=yes -DBoost_NO_BOOST_CMAKE=ON
 make -j2 VERBOSE=1 install
 
-cd
+cd ../tools
 
-python -c "from heyoka import test; test.run_test_suite()"
+python ci_test_runner.py
 
 cd $HEYOKA_PY_PROJECT_DIR
 

tools/circleci_conda_heyoka_head_38.sh

@@ -37,9 +37,9 @@ cd build
 cmake ../ -DCMAKE_INSTALL_PREFIX=$deps_dir -DCMAKE_PREFIX_PATH=$deps_dir -DCMAKE_BUILD_TYPE=Debug -DHEYOKA_PY_ENABLE_IPO=yes -DBoost_NO_BOOST_CMAKE=ON
 make -j2 VERBOSE=1 install
 
-cd
+cd ../tools
 
-python -c "from heyoka import test; test.run_test_suite()"
+python ci_test_runner.py
 
 cd $HEYOKA_PY_PROJECT_DIR
 

tools/circleci_conda_heyoka_head_release_310.sh

@@ -39,9 +39,9 @@ cd build
 cmake ../ -DCMAKE_INSTALL_PREFIX=$deps_dir -DCMAKE_PREFIX_PATH=$deps_dir -DHEYOKA_PY_ENABLE_IPO=yes -DBoost_NO_BOOST_CMAKE=ON
 make -j2 VERBOSE=1 install
 
-cd
+cd ../tools
 
-python -c "from heyoka import test; test.run_test_suite()"
+python ci_test_runner.py
 
 cd $HEYOKA_PY_PROJECT_DIR
 

tools/circleci_ubuntu_arm64.sh

@@ -36,9 +36,9 @@ cd build
 cmake ../ -DCMAKE_INSTALL_PREFIX=$deps_dir -DCMAKE_PREFIX_PATH=$deps_dir -DCMAKE_BUILD_TYPE=Debug -DHEYOKA_PY_ENABLE_IPO=yes -DBoost_NO_BOOST_CMAKE=ON
 make -j2 VERBOSE=1 install
 
-cd
+cd ../tools
 
-python -c "from heyoka import test; test.run_test_suite()"
+python ci_test_runner.py
 
 set +e
 set +x

tools/gha_conda_asan.sh

@@ -39,9 +39,9 @@ cd build
 cmake ../ -DCMAKE_INSTALL_PREFIX=$deps_dir -DCMAKE_PREFIX_PATH=$deps_dir -DCMAKE_BUILD_TYPE=Debug -DHEYOKA_PY_ENABLE_IPO=yes -DBoost_NO_BOOST_CMAKE=ON
 make -j2 VERBOSE=1 install
 
-cd
+cd ../tools
 
-ASAN_OPTIONS=detect_leaks=0 LD_PRELOAD=$CONDA_PREFIX/lib/libasan.so python -c "from heyoka import test; test.run_test_suite()"
+ASAN_OPTIONS=detect_leaks=0 LD_PRELOAD=$CONDA_PREFIX/lib/libasan.so python ci_test_runner.py
 
 set +e
 set +x

tools/gha_conda_static.sh

@@ -37,9 +37,9 @@ cd build
 cmake ../ -DCMAKE_INSTALL_PREFIX=$deps_dir -DCMAKE_PREFIX_PATH=$deps_dir -DCMAKE_BUILD_TYPE=Debug -DHEYOKA_PY_ENABLE_IPO=yes -DBoost_NO_BOOST_CMAKE=ON
 make -j2 VERBOSE=1 install
 
-cd
+cd ../tools
 
-python -c "import numba; from heyoka import test; test.run_test_suite()"
+python ci_test_runner.py --with-numba
 
 set +e
 set +x

tools/gha_manylinux.sh

@@ -103,7 +103,9 @@ auditwheel repair dist/heyoka* -w ./dist2
 unset LD_LIBRARY_PATH
 cd /
 /opt/python/${PYTHON_DIR}/bin/pip install ${GITHUB_WORKSPACE}/build/wheel/dist2/heyoka*
-/opt/python/${PYTHON_DIR}/bin/python -c "import heyoka; heyoka.test.run_test_suite();"
+cd ${GITHUB_WORKSPACE}/tools
+/opt/python/${PYTHON_DIR}/bin/python ci_test_runner.py
+cd /
 
 # Upload to PyPI.
 if [[ "${HEYOKA_PY_RELEASE_BUILD}" == "yes" ]]; then

tools/gha_osx_heyoka_head.sh

@@ -34,9 +34,9 @@ cd build
 cmake ../ -DCMAKE_INSTALL_PREFIX=$deps_dir -DCMAKE_PREFIX_PATH=$deps_dir -DCMAKE_BUILD_TYPE=Debug -DHEYOKA_PY_ENABLE_IPO=yes -DBoost_NO_BOOST_CMAKE=ON
 make -j2 VERBOSE=1 install
 
-cd
+cd ../tools
 
-python -c "from heyoka import test; test.run_test_suite()"
+python ci_test_runner.py
 
 set +e
 set +x

tools/gha_osx_heyoka_head_static.sh

@@ -34,9 +34,9 @@ cd build
 cmake ../ -DCMAKE_INSTALL_PREFIX=$deps_dir -DCMAKE_PREFIX_PATH=$deps_dir -DCMAKE_BUILD_TYPE=Debug -DHEYOKA_PY_ENABLE_IPO=yes -DBoost_NO_BOOST_CMAKE=ON
 make -j2 VERBOSE=1 install
 
-cd
+cd ../tools
 
-python -c "import numba; from heyoka import test; test.run_test_suite()"
+python ci_test_runner.py --with-numba
 
 set +e
 set +x

tools/gha_osx_heyoka_stable.sh

@@ -22,9 +22,9 @@ cd build
 cmake ../ -DCMAKE_INSTALL_PREFIX=$deps_dir -DCMAKE_PREFIX_PATH=$deps_dir -DCMAKE_BUILD_TYPE=Debug -DHEYOKA_PY_ENABLE_IPO=yes -DBoost_NO_BOOST_CMAKE=ON
 make -j2 VERBOSE=1 install
 
-cd
+cd ../tools
 
-python -c "from heyoka import test; test.run_test_suite()"
+python ci_test_runner.py
 
 set +e
 set +x

tools/travis_ubuntu_ppc64.sh

@@ -35,11 +35,11 @@ make VERBOSE=1 install
 
 echo "INSTALL DONE"
 
-cd /
+cd ../tools
 
 echo "MOVED OUT"
 
-$deps_dir/bin/python -c "from heyoka import test; test.run_test_suite()"
+$deps_dir/bin/python ci_test_runner.py
 
 echo "PYTHON RUN"