Python: Fix keep_alive specifications, add tests for keep_alive (#1851)

* Fix keepalives

* Add simple GC test

* Add keepalive for snapshots api

* Add more extensive keepalive test

* Slightly API-breaking.. need to del everything

* tmp, check sth

* tmp check ci

* Revert "tmp check ci"

This reverts commit 1cf6973.

* Revert "tmp, check sth"

This reverts commit 93ed467.

* Fix typing issues with load/store_chunk in Python

Author: franzpoeschel

src/binding/python/Iteration.cpp

@@ -101,15 +101,15 @@ void init_Iteration(py::module &m)
                 py::return_value_policy::copy,
                 // garbage collection: return value must be freed before
                 // Iteration
-                py::keep_alive<1, 0>()))
+                py::keep_alive<0, 1>()))
         .def_property_readonly(
             "particles",
             py::cpp_function(
                 [](Iteration &i) { return i.particles; },
                 py::return_value_policy::copy,
                 // garbage collection: return value must be freed before
                 // Iteration
-                py::keep_alive<1, 0>()));
+                py::keep_alive<0, 1>()));
 
     add_pickle(
         cl, [](openPMD::Series series, std::vector<std::string> const &group) {

src/binding/python/ParticleSpecies.cpp

@@ -54,7 +54,7 @@ void init_ParticleSpecies(py::module &m)
                 [](ParticleSpecies &ps) { return ps.particlePatches; },
                 py::return_value_policy::copy,
                 // garbage collection: return value must be freed before Series
-                py::keep_alive<1, 0>()));
+                py::keep_alive<0, 1>()));
     add_pickle(
         cl, [](openPMD::Series series, std::vector<std::string> const &group) {
             uint64_t const n_it = std::stoull(group.at(1));

src/binding/python/Series.cpp

@@ -301,7 +301,8 @@ not possible once it has been closed.
                     throw std::runtime_error("Unreachable");
                 },
                 // copy + keepalive
-                py::return_value_policy::copy)
+                py::return_value_policy::copy,
+                py::keep_alive<0, 1>())
             .def(
                 "current_iteration",
                 [](Snapshots &s) -> std::optional<IndexedIteration> {
@@ -315,6 +316,7 @@ not possible once it has been closed.
                         return std::nullopt;
                     }
                 },
+                py::keep_alive<0, 1>(),
                 "Return the iteration that is currently being written to, if "
                 "it "
                 "exists.");
@@ -503,7 +505,7 @@ this method.
                 [](Series &s) { return s.iterations; },
                 py::return_value_policy::copy,
                 // garbage collection: return value must be freed before Series
-                py::keep_alive<1, 0>()))
+                py::keep_alive<0, 1>()))
         .def(
             "read_iterations",
             [](Series &s) {
@@ -645,5 +647,18 @@ users to overwrite default options, while keeping any other ones.
             py::arg("comm"),
             docs_merge_json)
 #endif
-        ;
+        .def("__del__", [](Series &s) {
+            try
+            {
+                s.close();
+            }
+            catch (std::exception const &e)
+            {
+                std::cerr << "Error during close: " << e.what() << std::endl;
+            }
+            catch (...)
+            {
+                std::cerr << "Unknown error during close." << std::endl;
+            }
+        });
 }

test/python/unittest/API/APITest.py

@@ -2026,6 +2026,9 @@ def makeAvailableChunksRoundTrip(self, ext):
         # Cleaner: write.close()
         # But let's keep this instance to test that that workflow stays
         # functional.
+        # Need to delete everything as garbage collection will keep `write`
+        # alive as long as E_x is around.
+        del E_x
         del write
 
         read = io.Series(
@@ -2340,6 +2343,170 @@ def testScalarHdf5Fields(self):
         self.assertEqual(loaded_from_scalar, np.array([45]))
         series_read_again.close()
 
+    def testKeepaliveComponentExtraction(self):
+        """Test that keepalive specifications
+        guard root objects from garbage collection."""
+        self.testKeepaliveMeshComponent()
+        self.testKeepaliveParticlePosition()
+        self.testKeepaliveParticlePatches()
+
+    def testKeepaliveMeshComponent(self):
+        """Test keepalive for mesh component extraction."""
+        for ext in tested_file_extensions:
+            self.backend_keepalive_mesh_component(ext)
+
+    def testKeepaliveParticlePosition(self):
+        """Test keepalive for particle position component extraction."""
+        for ext in tested_file_extensions:
+            self.backend_keepalive_particle_position(ext)
+
+    def testKeepaliveParticlePatches(self):
+        """Test keepalive for particle patches component extraction."""
+        for ext in tested_file_extensions:
+            self.backend_keepalive_particle_patches(ext)
+
+    def backend_keepalive_mesh_component(self, file_ending):
+        """Helper function that tests keepalive
+        for mesh component extraction."""
+        import gc
+
+        filename = "unittest_py_keepalive_mesh." + file_ending
+        path = filename
+
+        def get_component_only():
+            series = io.Series(path, io.Access.create_linear)
+            backend = series.backend
+            iteration = series.snapshots()[0]
+            mesh = iteration.meshes["E"]
+            component = mesh["x"]
+
+            mesh.axis_labels = ["x", "y"]
+            component.reset_dataset(io.Dataset(np.dtype("float"), [10, 10]))
+
+            del iteration
+            del mesh
+            del series
+            gc.collect()
+
+            return component, backend
+
+        component, backend = get_component_only()
+        gc.collect()
+
+        component[:, :] = np.reshape(
+            np.arange(100, dtype=np.dtype("float")),
+            [10, 10]
+        )
+
+        component.series_flush()
+        if backend == "ADIOS2":
+            del component
+            gc.collect()
+
+        read = io.Series(path, io.Access.read_only)
+        loaded = read.snapshots()[0].meshes["E"]["x"][:]
+        read.flush()
+        np.testing.assert_array_equal(
+            loaded,
+            np.reshape(np.arange(100, dtype=np.dtype("float")), [10, 10])
+        )
+
+    def backend_keepalive_particle_position(self, file_ending):
+        """Helper function that tests keepalive
+        for particle position component extraction."""
+        import gc
+
+        filename = "unittest_py_keepalive_particle." + file_ending
+        path = filename
+        num_particles = 100
+
+        def get_component_only():
+            series = io.Series(path, io.Access.create_linear)
+            backend = series.backend
+            iteration = series.snapshots()[0]
+            particles = iteration.particles["electrons"]
+            position = particles["position"]["x"]
+
+            position.reset_dataset(
+                io.Dataset(np.dtype("float"), [num_particles]))
+
+            del iteration
+            del particles
+            del series
+            gc.collect()
+
+            return position, backend
+
+        position, backend = get_component_only()
+        gc.collect()
+
+        position[:] = np.arange(num_particles, dtype=np.dtype("float"))
+
+        position.series_flush()
+        if backend == "ADIOS2":
+            del position
+            gc.collect()
+
+        read = io.Series(path, io.Access.read_only)
+        loaded = read.snapshots()[0] \
+            .particles["electrons"]["position"]["x"][:]
+        read.flush()
+        np.testing.assert_array_equal(
+            loaded,
+            np.arange(num_particles, dtype=np.dtype("float"))
+        )
+
+    def backend_keepalive_particle_patches(self, file_ending):
+        """Helper function that tests keepalive
+        for particle patches extraction."""
+        import gc
+
+        filename = "unittest_py_keepalive_patches." + file_ending
+        path = filename
+
+        def get_component_only():
+            series = io.Series(path, io.Access.create_linear)
+            backend = series.backend
+            iteration = series.snapshots()[0]
+            particles = iteration.particles["electrons"]
+
+            dset = io.Dataset(np.dtype(np.float32), [30])
+            position_x = particles["position"]["x"]
+            position_x.reset_dataset(dset)
+            position_x[:] = np.arange(30, dtype=np.float32)
+
+            dset = io.Dataset(np.dtype("uint64"), [2])
+            num_particles_comp = particles.particle_patches["numParticles"]
+            num_particles_comp.reset_dataset(dset)
+            num_particles_comp.store(0, np.uint64(10))
+            num_particles_comp.store(1, np.uint64(20))
+
+            del iteration
+            del particles
+            del series
+            gc.collect()
+
+            return num_particles_comp, backend
+
+        component, backend = get_component_only()
+        gc.collect()
+
+        component.store(0, np.uint64(50))
+
+        component.series_flush()
+        if backend == "ADIOS2":
+            del component
+            gc.collect()
+
+        read = io.Series(path, io.Access.read_only)
+        loaded = read.snapshots()[0] \
+            .particles["electrons"].particle_patches["numParticles"].load()
+        read.flush()
+        np.testing.assert_array_equal(
+            loaded[0],
+            np.uint64(50)
+        )
+
 
 if __name__ == '__main__':
     unittest.main()