diff --git a/CMakeLists.txt b/CMakeLists.txt index 21ac656880..66f9a53a6e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -605,6 +605,7 @@ if(openPMD_HAVE_PYTHON) src/binding/python/ParticleSpecies.cpp src/binding/python/PatchRecord.cpp src/binding/python/PatchRecordComponent.cpp + src/binding/python/Pickle.cpp src/binding/python/Record.cpp src/binding/python/RecordComponent.cpp src/binding/python/MeshRecordComponent.cpp diff --git a/include/openPMD/Iteration.hpp b/include/openPMD/Iteration.hpp index 0892627f2d..caef902ac1 100644 --- a/include/openPMD/Iteration.hpp +++ b/include/openPMD/Iteration.hpp @@ -174,8 +174,8 @@ class Iteration friend class Container; friend class Series; friend class internal::AttributableData; - template - friend T &internal::makeOwning(T &self, Series); + template + friend T &internal::makeOwning(T &self, Series_type); friend class Writable; friend class StatefulIterator; friend class StatefulSnapshotsContainer; diff --git a/include/openPMD/ParticleSpecies.hpp b/include/openPMD/ParticleSpecies.hpp index 1ec1ff8d9c..df6024db31 100644 --- a/include/openPMD/ParticleSpecies.hpp +++ b/include/openPMD/ParticleSpecies.hpp @@ -40,8 +40,8 @@ class ParticleSpecies friend class Container; friend class Container; friend class Iteration; - template - friend T &internal::makeOwning(T &self, Series); + template + friend T &internal::makeOwning(T &self, Series_type); friend class internal::ScientificDefaults; friend class Attributable; diff --git a/include/openPMD/RecordComponent.hpp b/include/openPMD/RecordComponent.hpp index 3def700f71..e8052d2f3e 100644 --- a/include/openPMD/RecordComponent.hpp +++ b/include/openPMD/RecordComponent.hpp @@ -130,8 +130,8 @@ class RecordComponent friend class DynamicMemoryView; friend class internal::RecordComponentData; friend class MeshRecordComponent; - template - friend T &internal::makeOwning(T &self, Series); + template + friend T &internal::makeOwning(T &self, Series_type); friend class internal::ScientificDefaults; friend class Attributable; diff --git a/include/openPMD/Series.hpp b/include/openPMD/Series.hpp index 93dfe333b4..000b53a05c 100644 --- a/include/openPMD/Series.hpp +++ b/include/openPMD/Series.hpp @@ -294,6 +294,8 @@ class Series : public Attributable friend class internal::SeriesData; friend class internal::AttributableData; friend class StatefulSnapshotsContainer; + template + friend T &internal::makeOwning(T &self, Series_type); public: explicit Series(); @@ -778,6 +780,8 @@ class Series : public Attributable */ void close(); + [[nodiscard]] bool closed() const; + void visitHierarchy(HierarchyVisitor &v, bool recursive) override; /** @@ -804,6 +808,11 @@ OPENPMD_private using Data_t = internal::SeriesData; std::shared_ptr m_series = nullptr; + inline std::shared_ptr getShared() + { + return m_series; + } + inline Data_t &get() { if (m_series) diff --git a/include/openPMD/backend/Attributable.hpp b/include/openPMD/backend/Attributable.hpp index f05cc8d15b..0ed82cdbcd 100644 --- a/include/openPMD/backend/Attributable.hpp +++ b/include/openPMD/backend/Attributable.hpp @@ -207,8 +207,8 @@ namespace internal * Instantiations for T exist for types RecordComponent, * MeshRecordComponent, Mesh, Record, ParticleSpecies, Iteration. */ - template - T &makeOwning(T &self, Series); + template + T &makeOwning(T &self, Series_type); } // namespace internal namespace debug @@ -241,8 +241,8 @@ class Attributable friend class Writable; friend class internal::RecordComponentData; friend void debug::printDirty(Series const &); - template - friend T &internal::makeOwning(T &self, Series); + template + friend T &internal::makeOwning(T &self, Series_type); friend class StatefulSnapshotsContainer; friend class internal::AttributableData; friend class Snapshots; @@ -457,6 +457,13 @@ class Attributable [[nodiscard]] OpenpmdStandard openPMDStandard() const; + /** Returns the persistent immutable memory ID of the underlying Series. + * + * Useful when trying to determine which API handles refer to the same IO + * instance. + */ + [[nodiscard]] uintptr_t memoryID() const; + // clang-format off OPENPMD_protected // clang-format on diff --git a/include/openPMD/backend/BaseRecord.hpp b/include/openPMD/backend/BaseRecord.hpp index 295b12e909..3acdb04b86 100644 --- a/include/openPMD/backend/BaseRecord.hpp +++ b/include/openPMD/backend/BaseRecord.hpp @@ -199,8 +199,8 @@ class BaseRecord friend class internal::BaseRecordData; template friend class internal::ScalarIterator; - template - friend T &internal::makeOwning(T &self, Series); + template + friend T &internal::makeOwning(T &self, Series_type); friend class internal::ScientificDefaults; using Data_t = diff --git a/include/openPMD/binding/python/Pickle.hpp b/include/openPMD/binding/python/Pickle.hpp index 37eceb22ef..e27ea4ff99 100644 --- a/include/openPMD/binding/python/Pickle.hpp +++ b/include/openPMD/binding/python/Pickle.hpp @@ -27,13 +27,55 @@ #include "Common.hpp" +#include #include +#include +#include #include +#include #include #include namespace openPMD { +/* + * unpickled_series, as in "plural series"; this is a cache structure for series + * objects that have been unpickled. This cache structure fixes the issue + * described in https://github.com/openPMD/openPMD-api/issues/1919. + * Idea: One single Series instance may have multiple handles referencing it. + * When pickling and unpickling these references, the underlying Series must be + * restored once only, in order to keep the reference structure. Otherwise + * something like `data = E_x[:]; series.flush();` will not work, because `E_x` + * no longer references the same Series instance as `series`. + * + * For this, the pickle structure contains as first entry the internal + * (immutable) SharedAttributable pointer address of the `Series` object + * referenced by any handle. When unpickling, this is used to restore shared + * handles in accordance with their original reference structure. The pointers + * themselves are not restored (this would not be possible), but they are used + * as equivalence classes. + */ +struct unpickled_series +{ + // Cache restored object by original Series ID (i.e. internal immutable + // pointer address). IDs are not restored equivalently, but this does not + // matter. They are necessary only for figuring out which handles point to + // the same objects. + // The cached Series objects are stored as weak_ptr, since they are memory + // managed by the Python side. The C++ side just needs to check if the + // weak_ptr is still valid when handing out a new reference. If not, reopen. + std::map> m_series_by_former_id; + std::shared_mutex m_mutex; + + auto get(uintptr_t id, std::string const &filename) + -> std::shared_ptr; +}; + +/* + * Cache the Series per thread. + */ +extern unpickled_series cache; + /** Helper to Pickle Attributable Classes * * @tparam T_Args the types in pybind11::class_ - the first type will be pickled @@ -48,7 +90,7 @@ add_pickle(pybind11::class_ &cl, T_SeriesAccessor &&seriesAccessor) { // helper: get first class in py::class_ - that's the type we pickle using PickledClass = - typename std::tuple_element<0, std::tuple >::type; + typename std::tuple_element<0, std::tuple>::type; cl.def( py::pickle( @@ -56,54 +98,27 @@ add_pickle(pybind11::class_ &cl, T_SeriesAccessor &&seriesAccessor) [](const PickledClass &a) { // Return a tuple that fully encodes the state of the object Attributable::MyPath const myPath = a.myPath(); - return py::make_tuple(myPath.filePath(), myPath.group); + // retrieve Series even though retrieveSeries is protected... + return py::make_tuple( + a.memoryID(), myPath.filePath(), myPath.group); }, // __setstate__ [&seriesAccessor](py::tuple const &t) { - // our tuple has exactly two elements: filePath & group - if (t.size() != 2) + // Our tuple has exactly three elements: Series ID, filePath & + // group. + // Check the documentation of unpickled_series above for + // the reasoning behind Series ID. + if (t.size() != 3) throw std::runtime_error("Invalid state!"); - std::string const filename = t[0].cast(); + auto id = t[0].cast(); + std::string const filename = t[1].cast(); std::vector const group = - t[1].cast >(); - - /* - * Cache the Series per thread. - */ - thread_local std::optional series; - bool re_initialize = [&]() { - try - { - return !series.has_value() || - !series->operator bool() || - auxiliary::replace_all( - series->myPath().filePath(), "\\", "/") != - auxiliary::replace_all(filename, "\\", "/"); - } - /* - * Better safe than sorry, if anything goes wrong because - * the Series is in a weird state, just reinitialize it. - */ - catch (...) - { - return true; - } - }(); - if (re_initialize) - { - /* - * Do NOT close the old Series, it might still be active in - * terms of handed-out handles. - */ - series = std::make_optional( - filename, - Access::READ_ONLY, - "defer_iteration_parsing = true"); - } + t[2].cast>(); - return seriesAccessor(*series, group); + auto series = cache.get(id, filename); + return seriesAccessor(std::move(series), group); })); } } // namespace openPMD diff --git a/src/Series.cpp b/src/Series.cpp index 855178c18d..758ab67524 100644 --- a/src/Series.cpp +++ b/src/Series.cpp @@ -3564,6 +3564,20 @@ void Series::close() m_attri.reset(); } +bool Series::closed() const +{ + if (!operator bool()) + { + return true; + } + auto &w = writable(); + if (!w.IOHandler) + { + throw error::Internal("Series went into illegal state"); + } + return !w.IOHandler->has_value(); +} + void Series::visitHierarchy(HierarchyVisitor &v, bool recursive) { if (recursive) diff --git a/src/backend/Attributable.cpp b/src/backend/Attributable.cpp index d19fa31a00..33c3c41952 100644 --- a/src/backend/Attributable.cpp +++ b/src/backend/Attributable.cpp @@ -343,6 +343,11 @@ OpenpmdStandard Attributable::openPMDStandard() const return IOHandler()->m_standard; } +uintptr_t Attributable::memoryID() const +{ + return reinterpret_cast(&retrieveSeries().Attributable::get()); +} + template void Attributable::seriesFlush_impl(internal::FlushParams const &flushParams) { @@ -625,8 +630,8 @@ void Attributable::linkHierarchy(Writable &w) namespace internal { - template - T &makeOwning(T &self, Series s) + template + T &makeOwning(T &self, Series_type s) { /* * `self` is a handle object such as RecordComponent or Mesh (see @@ -665,11 +670,15 @@ namespace internal return self; } - template RecordComponent &makeOwning(RecordComponent &, Series); - template MeshRecordComponent &makeOwning(MeshRecordComponent &, Series); - template Mesh &makeOwning(Mesh &, Series); - template Record &makeOwning(Record &, Series); - template ParticleSpecies &makeOwning(ParticleSpecies &, Series); - template Iteration &makeOwning(Iteration &, Series); + template Series &makeOwning(Series &, std::shared_ptr); + template RecordComponent & + makeOwning(RecordComponent &, std::shared_ptr); + template MeshRecordComponent & + makeOwning(MeshRecordComponent &, std::shared_ptr); + template Mesh &makeOwning(Mesh &, std::shared_ptr); + template Record &makeOwning(Record &, std::shared_ptr); + template ParticleSpecies & + makeOwning(ParticleSpecies &, std::shared_ptr); + template Iteration &makeOwning(Iteration &, std::shared_ptr); } // namespace internal } // namespace openPMD diff --git a/src/binding/python/Attributable.cpp b/src/binding/python/Attributable.cpp index f96d66cc6b..047317cf91 100644 --- a/src/binding/python/Attributable.cpp +++ b/src/binding/python/Attributable.cpp @@ -657,7 +657,10 @@ void init_Attributable(py::module &m) .def( "populate_missing_metadata", &Attributable::populateMissingMetadata, - py::arg("recursive")); + py::arg("recursive")) + .def_property_readonly("memory_id", [](Attributable const &attr) { + return attr.memoryID(); + }); py::bind_vector(m, "Attribute_Keys"); } diff --git a/src/binding/python/Iteration.cpp b/src/binding/python/Iteration.cpp index f4fd2cc84b..496389ef2f 100644 --- a/src/binding/python/Iteration.cpp +++ b/src/binding/python/Iteration.cpp @@ -117,9 +117,11 @@ void init_Iteration(py::module &m) py::keep_alive<0, 1>())); add_pickle( - cl, [](openPMD::Series series, std::vector const &group) { + cl, + [](std::shared_ptr series, + std::vector const &group) { uint64_t const n_it = std::stoull(group.at(1)); - auto res = series.iterations[n_it]; + auto res = series->iterations[n_it]; return internal::makeOwning(res, std::move(series)); }); diff --git a/src/binding/python/Mesh.cpp b/src/binding/python/Mesh.cpp index 39e20488c2..272f2daa67 100644 --- a/src/binding/python/Mesh.cpp +++ b/src/binding/python/Mesh.cpp @@ -196,9 +196,11 @@ Ref.: https://github.com/openPMD/openPMD-standard/pull/193)"[1]) "set_grid_unit_SI", py::overload_cast(&Mesh::setGridUnitSI)); add_pickle( - cl, [](openPMD::Series series, std::vector const &group) { + cl, + [](std::shared_ptr series, + std::vector const &group) { uint64_t const n_it = std::stoull(group.at(1)); - auto res = series.iterations[n_it].open().meshes[group.at(3)]; + auto res = series->iterations[n_it].open().meshes[group.at(3)]; return internal::makeOwning(res, std::move(series)); }); diff --git a/src/binding/python/MeshRecordComponent.cpp b/src/binding/python/MeshRecordComponent.cpp index 3af8fe6c4d..cae244c177 100644 --- a/src/binding/python/MeshRecordComponent.cpp +++ b/src/binding/python/MeshRecordComponent.cpp @@ -82,10 +82,12 @@ void init_MeshRecordComponent(py::module &m) "Relative position of the component on an element " "(node/cell/voxel) of the mesh"); add_pickle( - cl, [](openPMD::Series series, std::vector const &group) { + cl, + [](std::shared_ptr series, + std::vector const &group) { uint64_t const n_it = std::stoull(group.at(1)); auto res = - series.iterations[n_it] + series->iterations[n_it] .open() .meshes[group.at(3)] [group.size() < 5 ? MeshRecordComponent::SCALAR diff --git a/src/binding/python/ParticleSpecies.cpp b/src/binding/python/ParticleSpecies.cpp index a097c37e84..48ec21bf0c 100644 --- a/src/binding/python/ParticleSpecies.cpp +++ b/src/binding/python/ParticleSpecies.cpp @@ -56,10 +56,12 @@ void init_ParticleSpecies(py::module &m) // garbage collection: return value must be freed before Series py::keep_alive<0, 1>())); add_pickle( - cl, [](openPMD::Series series, std::vector const &group) { + cl, + [](std::shared_ptr series, + std::vector const &group) { uint64_t const n_it = std::stoull(group.at(1)); ParticleSpecies res = - series.iterations[n_it].open().particles[group.at(3)]; + series->iterations[n_it].open().particles[group.at(3)]; return internal::makeOwning(res, std::move(series)); }); diff --git a/src/binding/python/Pickle.cpp b/src/binding/python/Pickle.cpp new file mode 100644 index 0000000000..6a40de4234 --- /dev/null +++ b/src/binding/python/Pickle.cpp @@ -0,0 +1,123 @@ +/* Copyright 2026 Franz Poeschel * + * + * This file is part of openPMD-api. + * + * openPMD-api is free software: you can redistribute it and/or modify + * it under the terms of of either the GNU General Public License or + * the GNU Lesser General Public License as published by + * the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * openPMD-api is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU General Public License and the GNU Lesser General Public License + * for more details. + * + * You should have received a copy of the GNU General Public License + * and the GNU Lesser General Public License along with openPMD-api. + * If not, see . + */ + +#include "openPMD/binding/python/Pickle.hpp" +#include "openPMD/binding/python/Common.hpp" + +#include + +namespace openPMD +{ +unpickled_series cache; + +auto unpickled_series::get(uintptr_t id, std::string const &filename) + -> std::shared_ptr +{ + // Check if a Series object with the given id is already in cache, still + // valid and points to the specified filename. + auto check_for_cached_series = + [&]() -> std::optional> { + std::shared_lock lock(m_mutex); + auto it = m_series_by_former_id.find(id); + if (it == m_series_by_former_id.end()) + { + return std::nullopt; + } + + auto candidate = it->second.lock(); + if (!candidate) + { + return std::nullopt; + } + + if (candidate->closed()) + { + return std::nullopt; + } + + if (auxiliary::replace_all(candidate->myPath().filePath(), "\\", "/") != + auxiliary::replace_all(filename, "\\", "/")) + { + return std::nullopt; + } + + return candidate; + }; + // There is a chance that the cached Series state is weird from previous + // usage, so catch any error and reinitialize in doubt. + auto maybe_series = [&]() -> std::optional> { + try + { + return check_for_cached_series(); + } + catch (...) + { + /* + * Better safe than sorry, if anything goes wrong because + * the Series is in a weird state, just reinitialize it. + */ + return std::nullopt; + } + }(); + + if (maybe_series) + { + return std::move(*maybe_series); + } + + // Else reinitialize. + { + std::unique_lock lock(m_mutex); + + // use the chance to do some cleanup + std::deque delete_me; + for (auto it = m_series_by_former_id.begin(); + it != m_series_by_former_id.end(); + ++it) + { + if (auto locked = it->second.lock(); !locked || locked->closed()) + { + delete_me.push_back(it); + } + } + + for (auto it : delete_me) + { + // References and iterators to the erased elements are + // invalidated. Other references and iterators are not affected. + m_series_by_former_id.erase(it); + } + + auto res = std::shared_ptr{ + new Series( + filename, Access::READ_ONLY, "defer_iteration_parsing = true"), + [this, id](Series const *s) { + { + std::unique_lock lock_lambda(this->m_mutex); + this->m_series_by_former_id.erase(id); + } + delete s; + }}; + m_series_by_former_id[id] = res; + return res; + } +} +} // namespace openPMD diff --git a/src/binding/python/Record.cpp b/src/binding/python/Record.cpp index dff5b50901..2e902f2f26 100644 --- a/src/binding/python/Record.cpp +++ b/src/binding/python/Record.cpp @@ -85,10 +85,12 @@ void init_Record(py::module &m) .def("set_time_offset", &Record::setTimeOffset) .def("set_time_offset", &Record::setTimeOffset); add_pickle( - cl, [](openPMD::Series series, std::vector const &group) { + cl, + [](std::shared_ptr series, + std::vector const &group) { uint64_t const n_it = std::stoull(group.at(1)); - auto res = series.iterations[n_it].open().particles[group.at(3)] - [group.at(4)]; + auto res = series->iterations[n_it].open().particles[group.at(3)] + [group.at(4)]; return internal::makeOwning(res, std::move(series)); }); diff --git a/src/binding/python/RecordComponent.cpp b/src/binding/python/RecordComponent.cpp index 4e2a8674c0..d6f8cf075a 100644 --- a/src/binding/python/RecordComponent.cpp +++ b/src/binding/python/RecordComponent.cpp @@ -1159,9 +1159,11 @@ void init_RecordComponent(py::module &m) .def("set_unit_SI", &RecordComponent::setUnitSI) // deprecated ; add_pickle( - cl, [](openPMD::Series series, std::vector const &group) { + cl, + [](std::shared_ptr series, + std::vector const &group) { uint64_t const n_it = std::stoull(group.at(1)); - auto res = series.iterations[n_it] + auto res = series->iterations[n_it] .open() .particles[group.at(3)][group.at(4)] [group.size() < 6 ? RecordComponent::SCALAR diff --git a/src/binding/python/Series.cpp b/src/binding/python/Series.cpp index 7054d3755b..3b80d5659b 100644 --- a/src/binding/python/Series.cpp +++ b/src/binding/python/Series.cpp @@ -568,8 +568,15 @@ Look for the WriteIterations class for further documentation. "TODO FILL IN DOCUMENTATION"); add_pickle( - cl, [](openPMD::Series series, std::vector const &) { - return series; + cl, + [](std::shared_ptr series, + std::vector const &) { + // Need to work on a copy since makeOwning will change the internal + // change pointer to capture also the cached Series. For this, the + // Series must be a different object than the cached Series, + // otherwise the cat will bite its own tail here. + Series copy = *series; + return openPMD::internal::makeOwning(copy, std::move(series)); }); constexpr char const *docs_merge_json = &R"END( diff --git a/test/SerialIOTest.cpp b/test/SerialIOTest.cpp index 49aab1db18..8637389d0f 100644 --- a/test/SerialIOTest.cpp +++ b/test/SerialIOTest.cpp @@ -171,7 +171,10 @@ void char_roundtrip(std::string const &extension) ::detail::writeChar(write, "char"); ::detail::writeChar(write, "uchar"); ::detail::writeChar(write, "schar"); + auto copy = write; write.close(); + REQUIRE(copy.closed()); + REQUIRE(write.closed()); Series read("../samples/char_rountrip." + extension, Access::READ_ONLY); ::detail::readChar(read, "char"); diff --git a/test/python/unittest/API/APITest.py b/test/python/unittest/API/APITest.py index 124cf2882d..390816021b 100644 --- a/test/python/unittest/API/APITest.py +++ b/test/python/unittest/API/APITest.py @@ -2523,6 +2523,268 @@ def get_component_only(): read.flush() np.testing.assert_array_equal(loaded, np.array([50, 20], dtype=np.uint64)) + class ReadMesh: + def __init__(self, series): + self.mesh_name = "rho" + self.series = series + self.it = self.series.iterations[400] + self.mesh = self.it.meshes[self.mesh_name] + + def __call__(self, use_stored_mesh=False) -> np.ndarray: + if use_stored_mesh: + dens = self.mesh[:] + else: + self.series.iterations[400].open() + mesh = self.series.iterations[400].meshes[self.mesh_name] + dens = mesh[:] + self.series.flush() + return dens + + class ReadMomentum: + """Helper for multiprocessing pickle test - must be a class to be picklable""" + + def __init__(self, pickled_series): + self.pickled_series = pickled_series + + def __call__(self, use_stored=False) -> np.ndarray: + import pickle + + series = pickle.loads(self.pickled_series) + it = series.iterations[400] + it.open() + electrons = it.particles["electrons"] + momentum = electrons["momentum"] + data = momentum["y"][()] + series.flush() + return data + + def testPickleSeriesIdentity(self): + # This tests the bug reported in + # https://github.com/openPMD/openPMD-api/issues/1919 + # The code is adapted from the reproducer in there. + try: + import multiprocessing + except (ImportError, ModuleNotFoundError): + return + + try: + series = io.Series("../samples/git-sample/data%T.h5", io.Access.read_only) + except io.ReadError: + return + + reader = self.ReadMesh(series) + assert np.array_equal( + reader(use_stored_mesh=False), reader(use_stored_mesh=True) + ) + + params = [[False, False], [True, True], [True, False]] + try: + for param in params: + with multiprocessing.Pool(processes=1) as pool: + results = pool.map(reader, param) + self.assertTrue(np.array_equal(results[0], results[1])) + except ModuleNotFoundError: + # happens on pyodide run. ignore. + pass + + def testPickleMultipleSeriesMultipleReferences(self): + # Test that the unpickle cache correctly handles multiple Series objects, + # each with multiple handles referencing it. + try: + import pickle + import multiprocessing + except ImportError: + return + + try: + series1 = io.Series("../samples/git-sample/data%T.h5", io.Access.read_only) + series2 = io.Series("../samples/git-sample/data%T.h5", io.Access.read_only) + except io.ReadError: + return + + # Create multiple references to each series + # Series 1 references (using particle records with components) + s1_it = series1.iterations[400] + s1_electrons = s1_it.particles["electrons"] + s1_momentum = s1_electrons["momentum"] + s1_mom_x = s1_momentum["x"] + + # Series 2 references + s2_it = series2.iterations[400] + s2_electrons = s2_it.particles["electrons"] + s2_momentum = s2_electrons["momentum"] + s2_mom_x = s2_momentum["x"] + + # Verify all references work and return same data + data_s1_mom = s1_momentum["y"][()] + series1.flush() + data_s1_mom_x = s1_mom_x[()] + series1.flush() + data_s2_mom = s2_momentum["y"][()] + series2.flush() + data_s2_mom_x = s2_mom_x[()] + series2.flush() + + np.testing.assert_array_equal(data_s1_mom, data_s2_mom) + np.testing.assert_array_equal(data_s1_mom_x, data_s2_mom_x) + + # Pickle all references + pickled_s1 = pickle.dumps(series1) + pickled_s1_it = pickle.dumps(s1_it) + pickled_s1_electrons = pickle.dumps(s1_electrons) + pickled_s1_momentum = pickle.dumps(s1_momentum) + pickled_s1_mom_x = pickle.dumps(s1_mom_x) + + pickled_s2 = pickle.dumps(series2) + pickled_s2_it = pickle.dumps(s2_it) + pickled_s2_electrons = pickle.dumps(s2_electrons) + pickled_s2_momentum = pickle.dumps(s2_momentum) + pickled_s2_mom_x = pickle.dumps(s2_mom_x) + + # Explicitly delete all objects to simulate GC + del s1_it, s1_electrons, s1_momentum, s1_mom_x, series1 + del s2_it, s2_electrons, s2_momentum, s2_mom_x, series2 + + # Unpickle - this should restore the cache properly + series1 = pickle.loads(pickled_s1) + s1_it = pickle.loads(pickled_s1_it) + s1_electrons = pickle.loads(pickled_s1_electrons) + s1_momentum = pickle.loads(pickled_s1_momentum) + s1_mom_x = pickle.loads(pickled_s1_mom_x) + + del s1_it, s1_electrons + + series2 = pickle.loads(pickled_s2) + s2_it = pickle.loads(pickled_s2_it) + s2_electrons = pickle.loads(pickled_s2_electrons) + s2_momentum = pickle.loads(pickled_s2_momentum) + s2_mom_x = pickle.loads(pickled_s2_mom_x) + + del s2_it, s2_electrons + + # Verify all unpickled references still work correctly + data_s1_mom_unpickled = s1_momentum["y"][()] + series1.flush() + data_s1_mom_x_unpickled = s1_mom_x[()] + series1.flush() + data_s2_mom_unpickled = s2_momentum["y"][()] + series2.flush() + data_s2_mom_x_unpickled = s2_mom_x[()] + series2.flush() + + np.testing.assert_array_equal(data_s1_mom, data_s1_mom_unpickled) + np.testing.assert_array_equal(data_s1_mom_x, data_s1_mom_x_unpickled) + np.testing.assert_array_equal(data_s2_mom, data_s2_mom_unpickled) + np.testing.assert_array_equal(data_s2_mom_x, data_s2_mom_x_unpickled) + + # Test multiprocessing with multiple series + pickled_for_mp1 = pickle.dumps(series1) + pickled_for_mp2 = pickle.dumps(series2) + + reader1 = self.ReadMomentum(pickled_for_mp1) + reader2 = self.ReadMomentum(pickled_for_mp2) + + params = [ + [False], + [True], + ] + + with multiprocessing.Pool(processes=2) as pool: + results1 = pool.map(reader1, params) + results2 = pool.map(reader2, params) + + # All results should be equal (both series point to same file) + for i, (r1, r2) in enumerate(zip(results1, results2)): + np.testing.assert_array_equal(r1, r2) + + def workerTestPickleCleanup(self, do_close): + # Test that Series.close() triggers cleanup properly + # (cleanup happens on next unpickle of a previously not opened Series) + try: + import pickle + except ImportError: + return + + try: + series = io.Series("../samples/git-sample/data%T.h5", io.Access.read_only) + series_2 = io.Series("../samples/git-sample/data%T.h5", io.Access.read_only) + except io.ReadError: + return + + # Create multiple references + it = series.iterations[400] + electrons = it.particles["electrons"] + momentum = electrons["momentum"] + mom_x = momentum["x"] + + # Pickle everything + pickled_series = pickle.dumps(series) + pickled_momentum = pickle.dumps(momentum) + pickled_mom_x = pickle.dumps(mom_x) + + # Create multiple references + it_2 = series.iterations[400] + electrons_2 = it.particles["electrons"] + momentum_2 = electrons["momentum"] + mom_x_2 = momentum["x"] + + # Pickle everything + pickled_series_2 = pickle.dumps(series) + pickled_momentum_2 = pickle.dumps(momentum) + pickled_mom_x_2 = pickle.dumps(mom_x) + + # Close the series explicitly + series.close() + del series, it, electrons, momentum, mom_x + + series_2.close() + del series_2, it_2, electrons_2, momentum_2, mom_x_2 + + # Unpickle + series = pickle.loads(pickled_series) + momentum = pickle.loads(pickled_momentum) + mom_x = pickle.loads(pickled_mom_x) + + # Verify data is still accessible + data_momentum = momentum["y"][()] + series.flush() + data_mom_x = mom_x[()] + series.flush() + + # Basic sanity check that we got data + self.assertIsNotNone(data_momentum) + self.assertIsNotNone(data_mom_x) + + # Remove information unpickled so far from cache again, either through API call or through GC + if do_close: + # print("EXPLICITLY CLOSING") + series.close() + else: + # print("EXPLICITLY DELETING") + del series, momentum, mom_x + # print("DONE") + + # Unpickle the second Series + series_2 = pickle.loads(pickled_series_2) + momentum_2 = pickle.loads(pickled_momentum_2) + mom_x_2 = pickle.loads(pickled_mom_x_2) + + # Verify data is still accessible + data_momentum_2 = momentum_2["y"][()] + series_2.flush() + data_mom_x_2 = mom_x_2[()] + series_2.flush() + + # Basic sanity check that we got data + self.assertIsNotNone(data_momentum_2) + self.assertIsNotNone(data_mom_x_2) + + def testPickleCleanupWithClose(self): + self.workerTestPickleCleanup(do_close=True) + + def testPickleCleanupWithGC(self): + self.workerTestPickleCleanup(do_close=False) + if __name__ == "__main__": unittest.main()