Improve MPI lifecycle management in PyOmega_h

pytest/conftest.py

@@ -1,21 +1,39 @@
 """
 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 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 sys
+import atexit
 import os
+import sys
 
+import pytest
+import PyOmega_h as omega_h
 
-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.
+_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.
     """
-    # Force immediate exit without Python cleanup to avoid MPI finalization conflicts
-    os._exit(exitstatus)
+    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/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."""

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."""

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

src/PyOmega_h.cpp

@@ -1,7 +1,20 @@
 #include <PyOmega_h.hpp>
 
+#include <stdexcept>
+
+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);

src/PyOmega_h.hpp

@@ -18,7 +18,8 @@ namespace py = pybind11;
 
 namespace Omega_h {
 class Library;
-extern Library* pybind11_global_library;
+void initialize();
+void finalize();
 void pybind11_defines(py::module& m);
 void pybind11_array(py::module& m);
 void pybind11_comm(py::module& m);

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

src/PyOmega_h_comm.cpp

@@ -14,24 +14,27 @@ void pybind11_comm(py::module& m) {
   py::class_<Omega_h::Comm, std::shared_ptr<Omega_h::Comm>>(m, "Comm")
   // Constructors
 #ifdef OMEGA_H_USE_MPI
-    .def(py::init([](Omega_h::Library* library, CommHandle impl_handle) {
+    .def(py::init([](std::shared_ptr<Omega_h::Library> 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<Omega_h::Library> library,
+                     CommHandle impl_handle,
                      py::array_t<const Omega_h::I32> srcs,
                      py::array_t<const Omega_h::I32> dsts) {
            MPI_Comm impl;
            std::memcpy(&impl, &impl_handle, sizeof(MPI_Comm));
            auto srcs_view = numpy_to_omega_h_read<Omega_h::I32>(srcs);
            auto dsts_view = numpy_to_omega_h_read<Omega_h::I32>(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<Omega_h::Library*, bool, bool>(), py::arg("library"),
-         py::arg("is_graph"), py::arg("sends_to_self"))
+    .def(py::init([](std::shared_ptr<Omega_h::Library> 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,

src/PyOmega_h_file.cpp

@@ -18,26 +18,27 @@ 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<Omega_h::Library> 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",
     [](const std::string& filepath, std::shared_ptr<Omega_h::Comm> comm) {
       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<Omega_h::Library> 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",

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_<Omega_h::Library>(
+  // Bind Omega_h::Library with shared_ptr holder so that keep_alive
+  // policies on dependent objects (Mesh, Comm) can extend its lifetime.
+  py::class_<Omega_h::Library, std::shared_ptr<Omega_h::Library>>(
     m, "OmegaHLibrary")
-    .def(py::init<>(), "Default constructor")
-    .def("world", &Omega_h::Library::world, "Get the world communicator");
+    .def(py::init([]() { return std::make_shared<Omega_h::Library>(); }),
+         "Default constructor")
+    .def("world",
+         [](std::shared_ptr<Omega_h::Library> self) { return self->world(); },
+         py::keep_alive<0, 1>(),
+         "Get the world communicator (keeps library alive)");
 }
 
 }  // namespace Omega_h

src/PyOmega_h_mesh.cpp

@@ -22,19 +22,27 @@ void pybind11_mesh(py::module& m) {
   void (Mesh::*balance)(bool) = &Mesh::balance;
   py::class_<Omega_h::Mesh, std::shared_ptr<Omega_h::Mesh>>(m, "OmegaHMesh")
     .def(py::init<>(), "Default constructor")
-    .def(py::init<Omega_h::Library*>(), 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<Omega_h::Library> lib) {
+           return std::make_shared<Omega_h::Mesh>(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<Omega_h::Library> 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")
 
     .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")
@@ -458,8 +466,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<Omega_h::Library> lib) {
+      return std::make_shared<Omega_h::Mesh>(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",