From 8921c3fab43fb401d5cf580f0a8c8d359bcff621 Mon Sep 17 00:00:00 2001 From: Sichao25 Date: Tue, 11 Aug 2026 13:56:21 -0400 Subject: [PATCH 1/3] keep lib alive when mesh exists --- pytest/conftest.py | 21 ++++----------------- src/PyOmega_h.hpp | 1 - src/PyOmega_h_build.cpp | 7 +++---- src/PyOmega_h_comm.cpp | 23 +++++++++++++++-------- src/PyOmega_h_file.cpp | 26 ++++++++++++++------------ src/PyOmega_h_library.cpp | 30 ++++++++++++++---------------- src/PyOmega_h_mesh.cpp | 25 ++++++++++++++++++------- 7 files changed, 68 insertions(+), 65 deletions(-) diff --git a/pytest/conftest.py b/pytest/conftest.py index 7faac399..60e6513d 100644 --- a/pytest/conftest.py +++ b/pytest/conftest.py @@ -1,21 +1,8 @@ """ Pytest configuration for PyOmega_h tests. -This conftest.py handles MPI/Kokkos cleanup issues by forcing a clean exit -before Python's normal cleanup process, which can conflict with MPI finalization. +The PyOmega_h bindings now use shared_ptr with py::keep_alive +policies on Mesh and Comm objects. This ensures the Library outlives all +dependent objects, so normal Python GC cleanup works correctly without +needing os._exit(). """ - -import sys -import os - - -def pytest_sessionfinish(session, exitstatus): - """ - Hook called after all tests complete but before pytest exits. - - For MPI programs, we need to exit immediately to avoid Python's cleanup - attempting to use MPI after it has been finalized. This prevents the - "Attempting to use an MPI routine after finalizing MPICH" error. - """ - # Force immediate exit without Python cleanup to avoid MPI finalization conflicts - os._exit(exitstatus) diff --git a/src/PyOmega_h.hpp b/src/PyOmega_h.hpp index 04dd996c..edb2d65a 100644 --- a/src/PyOmega_h.hpp +++ b/src/PyOmega_h.hpp @@ -18,7 +18,6 @@ namespace py = pybind11; namespace Omega_h { class Library; -extern Library* pybind11_global_library; void pybind11_defines(py::module& m); void pybind11_array(py::module& m); void pybind11_comm(py::module& m); diff --git a/src/PyOmega_h_build.cpp b/src/PyOmega_h_build.cpp index f658e4cc..d8151875 100644 --- a/src/PyOmega_h_build.cpp +++ b/src/PyOmega_h_build.cpp @@ -5,12 +5,11 @@ namespace Omega_h { void pybind11_build(py::module& m) { m.def("build_box", &Omega_h::build_box, "Build a rectangular mesh", - // if we give this a default argument, its lifetime is that - // of a global variable, which exceeds the Library lifetime! - py::arg("comm") /*= pybind11_global_library->world()*/, + py::arg("comm"), py::arg("family") = OMEGA_H_SIMPLEX, py::arg("x") = 1.0, py::arg("y") = 1.0, py::arg("z") = 1.0, py::arg("nx") = 0, - py::arg("ny") = 0, py::arg("nz") = 0, py::arg("symmetric") = false); + py::arg("ny") = 0, py::arg("nz") = 0, py::arg("symmetric") = false, + py::keep_alive<0, 1>()); } } // namespace Omega_h diff --git a/src/PyOmega_h_comm.cpp b/src/PyOmega_h_comm.cpp index daa3d8ef..d93f6bea 100644 --- a/src/PyOmega_h_comm.cpp +++ b/src/PyOmega_h_comm.cpp @@ -14,24 +14,27 @@ void pybind11_comm(py::module& m) { py::class_>(m, "Comm") // Constructors #ifdef OMEGA_H_USE_MPI - .def(py::init([](Omega_h::Library* library, CommHandle impl_handle) { + .def(py::init([](std::shared_ptr library, + CommHandle impl_handle) { MPI_Comm impl; std::memcpy(&impl, &impl_handle, sizeof(MPI_Comm)); - return new Omega_h::Comm(library, impl); + return new Omega_h::Comm(library.get(), impl); }), - py::arg("library"), py::arg("impl_handle")) + py::arg("library"), py::arg("impl_handle"), + py::keep_alive<1, 0>()) - .def(py::init([](Omega_h::Library* library, CommHandle impl_handle, + .def(py::init([](std::shared_ptr library, + CommHandle impl_handle, py::array_t srcs, py::array_t dsts) { MPI_Comm impl; std::memcpy(&impl, &impl_handle, sizeof(MPI_Comm)); auto srcs_view = numpy_to_omega_h_read(srcs); auto dsts_view = numpy_to_omega_h_read(dsts); - return new Omega_h::Comm(library, impl, srcs_view, dsts_view); + return new Omega_h::Comm(library.get(), impl, srcs_view, dsts_view); }), py::arg("library"), py::arg("impl_handle"), py::arg("srcs"), - py::arg("dsts")) + py::arg("dsts"), py::keep_alive<1, 0>()) .def( "get_impl_handle", @@ -43,8 +46,12 @@ void pybind11_comm(py::module& m) { }, "Get the underlying MPI communicator as an opaque integer handle") #else - .def(py::init(), py::arg("library"), - py::arg("is_graph"), py::arg("sends_to_self")) + .def(py::init([](std::shared_ptr library, bool is_graph, + bool sends_to_self) { + return new Omega_h::Comm(library.get(), is_graph, sends_to_self); + }), + py::arg("library"), py::arg("is_graph"), py::arg("sends_to_self"), + py::keep_alive<1, 0>()) #endif // Methods .def("library", &Omega_h::Comm::library, py::return_value_policy::reference, diff --git a/src/PyOmega_h_file.cpp b/src/PyOmega_h_file.cpp index 85e93709..86626c9b 100644 --- a/src/PyOmega_h_file.cpp +++ b/src/PyOmega_h_file.cpp @@ -18,18 +18,19 @@ void pybind11_file(py::module& m) { return Omega_h::read_mesh_file(filepath, comm); }, py::arg("filepath"), py::arg("comm"), - "Read mesh from file (auto-detects format)", py::return_value_policy::move); + "Read mesh from file (auto-detects format)", + py::keep_alive<0, 2>(), py::return_value_policy::move); // Binary format I/O m.def( "read_mesh_binary", - [](const std::string& filepath, Omega_h::Library* lib) { - Omega_h::Mesh mesh(lib); + [](const std::string& filepath, std::shared_ptr lib) { + Omega_h::Mesh mesh(lib.get()); Omega_h::binary::read(filepath, lib->world(), &mesh); return mesh; }, py::arg("filepath"), py::arg("library"), "Read mesh from binary file", - py::return_value_policy::move); + py::keep_alive<0, 2>(), py::return_value_policy::move); m.def( "read_mesh_binary", @@ -37,7 +38,7 @@ void pybind11_file(py::module& m) { return Omega_h::binary::read(filepath, comm); }, py::arg("filepath"), py::arg("comm"), "Read mesh from binary file", - py::return_value_policy::move); + py::keep_alive<0, 2>(), py::return_value_policy::move); m.def( "write_mesh_binary", @@ -53,7 +54,7 @@ void pybind11_file(py::module& m) { return Omega_h::gmsh::read(filepath, comm); }, py::arg("filepath"), py::arg("comm"), "Read mesh from Gmsh file", - py::return_value_policy::move); + py::keep_alive<0, 2>(), py::return_value_policy::move); m.def( "write_mesh_gmsh", @@ -69,7 +70,7 @@ void pybind11_file(py::module& m) { return Omega_h::gmsh::read_parallel(filepath, comm); }, py::arg("filepath"), py::arg("comm"), "Read parallel Gmsh mesh (MSH 4.1+)", - py::return_value_policy::move); + py::keep_alive<0, 2>(), py::return_value_policy::move); m.def("write_mesh_gmsh_parallel", &Omega_h::gmsh::write_parallel, py::arg("filepath"), py::arg("mesh"), "Write parallel Gmsh mesh (MSH 4.1)"); @@ -115,7 +116,7 @@ void pybind11_file(py::module& m) { return mesh; }, py::arg("filepath"), py::arg("comm"), "Read mesh from VTU file", - py::return_value_policy::move); + py::keep_alive<0, 2>(), py::return_value_policy::move); m.def( "read_mesh_parallel_vtk", @@ -125,7 +126,7 @@ void pybind11_file(py::module& m) { return mesh; }, py::arg("pvtupath"), py::arg("comm"), "Read parallel VTK mesh", - py::return_value_policy::move); + py::keep_alive<0, 2>(), py::return_value_policy::move); #ifdef OMEGA_H_USE_SEACASEXODUS // Exodus ClassifyWith enum @@ -253,12 +254,13 @@ void pybind11_file(py::module& m) { // ADIOS2 format I/O m.def( "read_mesh_adios2", - [](const std::string& filepath, Omega_h::Library* lib, + [](const std::string& filepath, std::shared_ptr lib, const std::string& prefix) { - return Omega_h::adios::read(filepath, lib, prefix); + return Omega_h::adios::read(filepath, lib.get(), prefix); }, py::arg("filepath"), py::arg("library"), py::arg("prefix") = "", - "Read mesh from ADIOS2 file", py::return_value_policy::move); + "Read mesh from ADIOS2 file", py::keep_alive<0, 2>(), + py::return_value_policy::move); m.def( "write_mesh_adios2", diff --git a/src/PyOmega_h_library.cpp b/src/PyOmega_h_library.cpp index c2bd1867..f41fe625 100644 --- a/src/PyOmega_h_library.cpp +++ b/src/PyOmega_h_library.cpp @@ -5,26 +5,24 @@ namespace Omega_h { /* The lifetime of the Library object is quite important (it must contain the - lifetime of pretty much all other Omega_h objects), and - I'm unsure about the order in which destructors will be called. - So, for now, I'll take the approach that the Python interface will have - the Library as a hidden global variable. - This is consistent with how mpi4py seems to work. - I tried using the Python atexit mechanism, but that seems to execute prior - to final garbage collection. - - Note: We use a raw pointer and intentionally leak it to avoid CUDA/Kokkos - finalization issues during Python shutdown. The OS will clean up the memory. + lifetime of pretty much all other Omega_h objects). We use a shared_ptr + holder so that py::keep_alive policies on Mesh and Comm objects can keep + the Library alive until all dependent objects have been destroyed. + This avoids the "Attempting to use an MPI routine after finalizing MPICH" + error that occurs when Python GC destroys the Library before Meshes/Comms. */ -Library* pybind11_global_library = nullptr; - void pybind11_library(py::module& m) { - // Bind Omega_h::Library - py::class_( + // Bind Omega_h::Library with shared_ptr holder so that keep_alive + // policies on dependent objects (Mesh, Comm) can extend its lifetime. + py::class_>( m, "OmegaHLibrary") - .def(py::init<>(), "Default constructor") - .def("world", &Omega_h::Library::world, "Get the world communicator"); + .def(py::init([]() { return std::make_shared(); }), + "Default constructor") + .def("world", + [](std::shared_ptr self) { return self->world(); }, + py::keep_alive<0, 1>(), + "Get the world communicator (keeps library alive)"); } } // namespace Omega_h diff --git a/src/PyOmega_h_mesh.cpp b/src/PyOmega_h_mesh.cpp index edcf86b6..9267aed1 100644 --- a/src/PyOmega_h_mesh.cpp +++ b/src/PyOmega_h_mesh.cpp @@ -22,11 +22,18 @@ void pybind11_mesh(py::module& m) { void (Mesh::*balance)(bool) = &Mesh::balance; py::class_>(m, "OmegaHMesh") .def(py::init<>(), "Default constructor") - .def(py::init(), py::arg("library"), - "Constructor with library") - - .def("set_library", &Omega_h::Mesh::set_library, py::arg("library"), - "Set the library") + .def(py::init([](std::shared_ptr lib) { + return std::make_shared(lib.get()); + }), + py::arg("library"), py::keep_alive<1, 0>(), + "Constructor with library (keeps library alive)") + + .def("set_library", + [](Omega_h::Mesh& mesh, std::shared_ptr lib) { + mesh.set_library(lib.get()); + }, + py::arg("library"), py::keep_alive<1, 2>(), + "Set the library (keeps library alive)") .def("library", &Omega_h::Mesh::library, py::return_value_policy::reference, "Get the library") @@ -458,8 +465,12 @@ void pybind11_mesh(py::module& m) { return omega_h_read_to_numpy(sizes); }, "Get element sizes"); - m.def( - "new_empty_mesh", []() { return Mesh(pybind11_global_library); }); + m.def("new_empty_mesh", + [](std::shared_ptr lib) { + return std::make_shared(lib.get()); + }, + py::arg("library"), py::keep_alive<1, 0>(), + "Create an empty mesh associated with the given library"); // Mesh utility functions m.def( "average_field", From ca2aa7d2fbbfd22858f0cb4fa9cc301c6cbe8631 Mon Sep 17 00:00:00 2001 From: Sichao25 Date: Tue, 11 Aug 2026 17:23:04 -0400 Subject: [PATCH 2/3] optimize pytest exit --- pytest/conftest.py | 47 ++++++++++++++++++++++++++++++--- pytest/test_file_io.py | 24 ----------------- pytest/test_mesh_tags_simple.py | 19 ------------- src/PyOmega_h_mesh.cpp | 3 ++- 4 files changed, 45 insertions(+), 48 deletions(-) diff --git a/pytest/conftest.py b/pytest/conftest.py index 60e6513d..9dcf7269 100644 --- a/pytest/conftest.py +++ b/pytest/conftest.py @@ -1,8 +1,47 @@ """ Pytest configuration for PyOmega_h tests. -The PyOmega_h bindings now use shared_ptr with py::keep_alive -policies on Mesh and Comm objects. This ensures the Library outlives all -dependent objects, so normal Python GC cleanup works correctly without -needing os._exit(). +The lib fixture is defined here (conftest.py) so that exactly the same +Omega_h::Library instance exists for the entire test session. Each +test file must NOT define its own lib fixture — doing so creates duplicate +Library instances, and whichever one is destroyed first will call +MPI_Finalize, causing the other Library's Comm destructors to fail with: +"Attempting to use an MPI routine after finalizing MPICH". """ + +import atexit +import os +import sys + +import pytest +import PyOmega_h as omega_h + +_global_lib = None + + +@pytest.fixture(scope="session") +def lib(): + """Session-scoped library instance shared across all tests. + + The Library must outlive all Mesh and Comm objects. The pybind11 + bindings use shared_ptr with py::keep_alive to ensure this ordering + through normal Python GC reference counting. + """ + global _global_lib + if _global_lib is None: + _global_lib = omega_h.OmegaHLibrary() + return _global_lib + + +@pytest.fixture(scope="session") +def world(lib): + """Session-scoped communicator.""" + return lib.world() + + +def _cleanup_exit(): + """Avoid Kokkos::Cuda::finalize() and MPI_Finalize() ordering issues.""" + os._exit(0) + + +atexit.register(_cleanup_exit) diff --git a/pytest/test_file_io.py b/pytest/test_file_io.py index 08e5509c..b69fb7a4 100644 --- a/pytest/test_file_io.py +++ b/pytest/test_file_io.py @@ -14,30 +14,6 @@ import tempfile -# Global library instance - intentionally NOT cleaned up to avoid MPI finalization issues -# See PyOmega_h_library.cpp for design rationale -_global_lib = None - - -@pytest.fixture(scope="session") -def lib(): - """Session-scoped library instance shared across all tests. - - Note: Cleanup is handled by conftest.py using os._exit() to avoid - MPI/Kokkos finalization order issues. - """ - global _global_lib - if _global_lib is None: - _global_lib = omega_h.OmegaHLibrary() - return _global_lib - - -@pytest.fixture(scope="session") -def world(lib): - """Session-scoped communicator.""" - return lib.world() - - @pytest.fixture def test_mesh(world): """Create a simple 3D test mesh.""" diff --git a/pytest/test_mesh_tags_simple.py b/pytest/test_mesh_tags_simple.py index 0a1b03b4..3e511e26 100644 --- a/pytest/test_mesh_tags_simple.py +++ b/pytest/test_mesh_tags_simple.py @@ -12,25 +12,6 @@ import PyOmega_h as omega_h -# Create a global library instance to avoid Kokkos/MPI finalization issues -# The library must outlive all mesh objects and is intentionally NOT deleted -# to avoid MPI finalization order issues (see PyOmega_h_library.cpp) -_global_lib = None - - -@pytest.fixture(scope="session") -def lib(): - """Session-scoped library instance shared across all tests. - - Note: Cleanup is handled by conftest.py using os._exit() to avoid - MPI/Kokkos finalization order issues. - """ - global _global_lib - if _global_lib is None: - _global_lib = omega_h.OmegaHLibrary() - return _global_lib - - @pytest.fixture def mesh(lib): """Create a simple 2D box mesh for testing.""" diff --git a/src/PyOmega_h_mesh.cpp b/src/PyOmega_h_mesh.cpp index 9267aed1..7c81a4cd 100644 --- a/src/PyOmega_h_mesh.cpp +++ b/src/PyOmega_h_mesh.cpp @@ -41,7 +41,8 @@ void pybind11_mesh(py::module& m) { .def("set_comm", &Omega_h::Mesh::set_comm, py::arg("comm"), "Set the communicator") - .def("comm", &Omega_h::Mesh::comm, "Get the communicator") + .def("comm", &Omega_h::Mesh::comm, py::keep_alive<0, 1>(), + "Get the communicator (keeps mesh alive)") .def("set_dim", &Omega_h::Mesh::set_dim, py::arg("dim"), "Set mesh dimension") From ed0141cd1ce1d06a3a358577a8763e067b3a07b0 Mon Sep 17 00:00:00 2001 From: Sichao25 Date: Mon, 31 Aug 2026 12:06:36 -0400 Subject: [PATCH 3/3] explicitly handle python runtime initialize/finalize --- pytest/conftest.py | 8 ----- src/CMakeLists.txt | 1 + src/PyOmega_h.cpp | 13 +++++++++ src/PyOmega_h.hpp | 2 ++ src/PyOmega_h_runtime.cpp | 61 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 77 insertions(+), 8 deletions(-) create mode 100644 src/PyOmega_h_runtime.cpp diff --git a/pytest/conftest.py b/pytest/conftest.py index 9dcf7269..83aae78a 100644 --- a/pytest/conftest.py +++ b/pytest/conftest.py @@ -37,11 +37,3 @@ def lib(): def world(lib): """Session-scoped communicator.""" return lib.world() - - -def _cleanup_exit(): - """Avoid Kokkos::Cuda::finalize() and MPI_Finalize() ordering issues.""" - os._exit(0) - - -atexit.register(_cleanup_exit) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index be207aec..67f39dc3 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -868,6 +868,7 @@ install(FILES ${Omega_h_HEADERS} DESTINATION include) if (Omega_h_USE_pybind11) set(PYBIND11_SOURCES PyOmega_h.cpp + PyOmega_h_runtime.cpp PyOmega_h_defines.cpp PyOmega_h_array.cpp PyOmega_h_comm.cpp diff --git a/src/PyOmega_h.cpp b/src/PyOmega_h.cpp index 0a5da59a..da1dfef3 100644 --- a/src/PyOmega_h.cpp +++ b/src/PyOmega_h.cpp @@ -1,7 +1,20 @@ #include +#include + +namespace { +void py_atexit_finalize() { + Omega_h::finalize(); +} +} // namespace + PYBIND11_MODULE(PyOmega_h, m) { m.doc() = "Omega_h: simplex mesh adaptation"; + Omega_h::initialize(); + if (Py_AtExit(&py_atexit_finalize) != 0) { + throw std::runtime_error( + "PyOmega_h: failed to register Py_AtExit finalizer"); + } Omega_h::pybind11_defines(m); Omega_h::pybind11_array(m); Omega_h::pybind11_comm(m); diff --git a/src/PyOmega_h.hpp b/src/PyOmega_h.hpp index edb2d65a..6bbd52ab 100644 --- a/src/PyOmega_h.hpp +++ b/src/PyOmega_h.hpp @@ -18,6 +18,8 @@ namespace py = pybind11; namespace Omega_h { class Library; +void initialize(); +void finalize(); void pybind11_defines(py::module& m); void pybind11_array(py::module& m); void pybind11_comm(py::module& m); diff --git a/src/PyOmega_h_runtime.cpp b/src/PyOmega_h_runtime.cpp new file mode 100644 index 00000000..eb36a389 --- /dev/null +++ b/src/PyOmega_h_runtime.cpp @@ -0,0 +1,61 @@ +#include + +#include + +#ifdef OMEGA_H_USE_MPI +#include +#endif +#ifdef OMEGA_H_USE_KOKKOS +#include +#endif + +namespace Omega_h { + +namespace { +#ifdef OMEGA_H_USE_MPI +bool g_py_owns_mpi = false; +#endif +#ifdef OMEGA_H_USE_KOKKOS +bool g_py_owns_kokkos = false; +#endif +} // namespace + + +void initialize() { +#ifdef OMEGA_H_USE_MPI + int is_initialized = 0; + MPI_Initialized(&is_initialized); + if (!is_initialized) { + if (MPI_Init(nullptr, nullptr) != MPI_SUCCESS) { + throw std::runtime_error("PyOmega_h: MPI_Init failed"); + } + g_py_owns_mpi = true; + } +#endif +#ifdef OMEGA_H_USE_KOKKOS + if (!Kokkos::is_initialized()) { + Kokkos::initialize(); + g_py_owns_kokkos = true; + } +#endif +} + +void finalize() { +#ifdef OMEGA_H_USE_KOKKOS + if (g_py_owns_kokkos && !Kokkos::is_finalized()) { + KokkosPool::destroyGlobalPool(); + Kokkos::finalize(); + g_py_owns_kokkos = false; + } +#endif +#ifdef OMEGA_H_USE_MPI + int is_finalized = 0; + MPI_Finalized(&is_finalized); + if (g_py_owns_mpi && !is_finalized) { + MPI_Finalize(); + g_py_owns_mpi = false; + } +#endif +} + +} // namespace Omega_h