added experimental pickling support to Universe #2140

richardjgowers · 2018-11-15T00:21:54Z

This adds a basic form of pickling for a Universe, everything seems to work in the snippet below

import dask
from dask import distributed

import MDAnalysis as mda
from MDAnalysisTests.datafiles import GRO, XTC

c = distributed.Client()

u = mda.Universe(GRO, XTC)

ag = u.atoms[:5]

def get_pos(u, ag, frame):    
    u.trajectory[frame]
    
    return ag.center_of_mass()

futures = []
for i in range(10):
    futures.append(dask.delayed(get_pos)(u, ag, i))

L = c.compute(dask.delayed(futures))

ref = [get_pos(u, ag, i) for i in range(10)]

WRT what sort of schedulers this should work for, I think this is only safe for distributed or multiprocessing, but not thread (as you can't share a Universe and play with frames)

@VOD555 I think this means you can send AtomGroups directly (no more remembering .ix tricks), and you don't reread the topology file (ie faster) because that half of the Universe is pickled.

richardjgowers · 2018-11-15T00:22:40Z

package/MDAnalysis/core/universe.py

-        raise NotImplementedError
+    @classmethod
+    def _unpickle_U(cls, top, traj, anchor):
+        u = cls(top, anchor_name=anchor)


the anchor_name= here makes sure that AtomGroups that get unpickled afterwards recognise this deserialised object as the correct parent for them

richardjgowers · 2018-11-15T00:24:43Z

package/MDAnalysis/core/universe.py

-    def __setstate__(self, state):
-        raise NotImplementedError
+    def __reduce__(self):
+        return (self._unpickle_U, (self._topology, self.trajectory.filename, self.anchor_name))


the entire Topology side of things seems to pickle cleanly without any drama, so we can just do this. To serialise/deserialise the Reader side of things we will always want a unique file handle, so passing the filename around isn't too awful.

There's probably lots of state in Reader which isn't currently respected (auxs, current frame off the top of my head)

Need to be documented what does not work yet.

From the code above it seems that Aux are included but if not, then pickle should fail on anything that we cannot pickle at the moment.

Turns out a Reader just knows how to serialise itself (including file handles)

codecov · 2018-11-21T19:50:55Z

Codecov Report

Merging #2140 into develop will decrease coverage by 1.29%.
The diff coverage is 100%.

@@             Coverage Diff             @@
##           develop    #2140      +/-   ##
===========================================
- Coverage    90.71%   89.41%    -1.3%     
===========================================
  Files           15      158     +143     
  Lines         1982    19646   +17664     
  Branches         0     2761    +2761     
===========================================
+ Hits          1798    17567   +15769     
- Misses         184     1476    +1292     
- Partials         0      603     +603

Impacted Files	Coverage Δ
package/MDAnalysis/core/universe.py	`95.05% <100%> (ø)`
package/MDAnalysis/coordinates/base.py	`93.32% <100%> (ø)`
package/MDAnalysis/coordinates/GRO.py	`93.87% <0%> (ø)`
package/MDAnalysis/analysis/helanal.py	`85.19% <0%> (ø)`
package/MDAnalysis/topology/MinimalParser.py	`100% <0%> (ø)`
package/MDAnalysis/visualization/__init__.py	`100% <0%> (ø)`
package/MDAnalysis/core/qcprot.py	`100% <0%> (ø)`
...e/MDAnalysis/analysis/encore/clustering/cluster.py	`96.22% <0%> (ø)`
package/MDAnalysis/coordinates/DCD.py	`97.16% <0%> (ø)`
... and 136 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update ea05c5a...f1a7d5a. Read the comment docs.

richardjgowers · 2018-12-02T21:38:54Z

poke @VOD555 if this works then it makes pmda a hell of a lot simpler

orbeckst

Very encouraging!! Addresses/supercedes #643 (and possibly #878 ?)

More testing & documentation.

Will also eventually need tests for #1981 #1940

testsuite/MDAnalysisTests/core/test_universe.py

orbeckst · 2018-12-05T17:58:15Z

package/MDAnalysis/core/universe.py

-    def __setstate__(self, state):
-        raise NotImplementedError
+    def __reduce__(self):
+        return (self._unpickle_U, (self._topology, self.trajectory.filename, self.anchor_name))


Need to be documented what does not work yet.

From the code above it seems that Aux are included but if not, then pickle should fail on anything that we cannot pickle at the moment.

package/MDAnalysis/coordinates/base.py

orbeckst · 2018-12-05T18:03:26Z

testsuite/MDAnalysisTests/core/test_universe.py

@@ -283,10 +283,14 @@ def test_load_multiple_args(self):
        assert_equal(len(u.atoms), 3341, "Loading universe failed somehow")
        assert_equal(u.trajectory.n_frames, 2 * ref.trajectory.n_frames)

-    def test_pickle_raises_NotImplementedError(self):
+    def test_pickle(self):


Should also test that

aux

transformations

are properly pickled.

richardjgowers · 2018-12-06T14:33:08Z

@orbeckst I'm happy to push ahead and finish this once I get some confirmation it works for pmda too, else it's just a useless fix

richardjgowers · 2019-02-16T16:55:12Z

package/MDAnalysis/core/universe.py

@@ -738,7 +738,7 @@ def _gen_anchor_hash(self):
                return self._anchor_uuid
            except AttributeError:
                # store this so we can later recall it if needed
-                self._anchor_uuid = uuid.uuid4()
+                self._anchor_uuid = str(uuid.uuid4())


uuid's didn't seem to survive the serialisation well, but if I just take a string it's all good (and still uuidesque)

richardjgowers · 2019-02-16T16:55:46Z

package/MDAnalysis/core/universe.py

+        """Special method used by __reduce__ to deserialise a Universe"""
+        # top is a Topology object at this point, but Universe can handle that
+        u = cls(top)
+        u.anchor_name = anchor


this is the magic that lets future atomgroups find their 'home' universe

and broke everything

richardjgowers · 2019-02-16T17:28:43Z

Ok looks like multiprocessing works, except for ascii format trajectories. I think this is close to working and I'd like to put it in (for a project I have brewing) so @MDAnalysis/coredevs try and break this (while I fix up ascii formats)

orbeckst · 2019-02-20T00:57:24Z

@VOD555 can you have a look at this work in the context of PMDA?

orbeckst · 2019-02-20T01:02:17Z

These are surprisingly few changes.

orbeckst · 2019-02-20T01:04:22Z

As I said in the review "More testing and documentation" (I think you also need CHANGELOG etc). It should be clearly documented what works and what doesn't. If not everything works then label it as "experimental".

And it would be nice to fail with meaningful exceptions – if possible.

richardjgowers · 2019-02-20T22:40:53Z

package/MDAnalysis/coordinates/chain.py

@@ -353,6 +353,16 @@ def __init__(self, filenames, skip=1, dt=None, continuous=False, **kwargs):
        self.ts = None
        self.rewind()

+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state.pop('_ChainReader__chained_trajectories_iter', None)


So I can't really reproduce the state of this generator cleanly (or at least in a simple way). Tbh I've never used chainreader or really touched the code, so maybe I should just slap a 'NotImplementedError' on this and let someone who understands this fix it in the future.

I think what you're doing is ok — like all the others, it resets to frame 0 on unpickling.

Once pickle support is upgraded to also restore the frame index then moving to the appropriate frame should just use the usual trajectory[frame] and then let ChainReader take care of how to do it.

If you somehow managed to pickle ChainReader.readers in

mdanalysis/package/MDAnalysis/coordinates/chain.py

Lines 265 to 266 in 6d6b862

self.readers = [core.reader(filename, **kwargs)

for filename in filenames]

then it should work. Otherwise, we need a special getstate/setstate pair that blanks readers and rebuilds them from ChainReader.filenames (although we should actually store the original filenames because they can be a list of tuples — that's an addition to ChainReader itself).

richardjgowers · 2019-02-20T22:43:43Z

package/MDAnalysis/coordinates/base.py

+
+        del self._pickle_fn
+
+class _BAsciiPickle(_AsciiPickle):


This is kinda ugly, but to unpickle the file handle you need to remember the mode it was opened in. This isn't currently possible afaik

richardjgowers · 2019-02-20T22:45:43Z

I did some hacking, all readers now pickle, to the best of my knowledge..... Lots of ugly hacks everywhere though

orbeckst · 2020-02-03T20:42:55Z

testsuite/MDAnalysisTests/parallelism/test_multiprocessing.py

+    (XYZ,),
+])
+def u(request):
+    if len(request.param) == 1:


maybe just

def u(request): return mda.Universe(*request.param)

?

Then that would also work for ChainReader shenanigans

u(TOPOL, TRJ1, TRJ2) u(TOPOL, [TRJ1, TRJ2]) u(TOPOL, [(TRJ1, "pdb"), (TRJ2, "xtc"), TRJ3])

orbeckst · 2020-05-29T18:44:12Z

@yuxuanzhuang as discussed on the GSoC call, put up your updated PR from your fork and then we can close this one as superseded. (Mention PR #2140 in the text of your PR so that they link nicely to each other.)

orbeckst · 2020-05-29T18:45:16Z

Oops – sorry, you just did it 16 minutes ago – my page hadn't been refreshed.

I am closing this PR now (superseded by PR #2704)

* Fixes #2878 * basic approach: composition instead of inheritance for pickling Universe (which was tested in PR #2704 (which was derived from PR #2140)) * Changes made in this Pull Request: - add new classes and pickle_open function to picklable_file_io.py - add new picklable `BufferedReader`, `FileIO`, and `TextIOWrapper` classes. - implement `__getstate__`, `__setstate__` to `Universe` and `BaseReader` - fix DCD, XDR pickle issue - modified gsd and ncdf to be picklabel - modified ChainReader to be picklabel - ensure chemfiles is picklable * add tests (MultiFrameReader will test for serializability) * update CHANGELOG * update docs Note: This merge squashed 120 commits. See PR #2723 for the full history with 420 comments.

* Fixes MDAnalysis#2878 * basic approach: composition instead of inheritance for pickling Universe (which was tested in PR MDAnalysis#2704 (which was derived from PR MDAnalysis#2140)) * Changes made in this Pull Request: - add new classes and pickle_open function to picklable_file_io.py - add new picklable `BufferedReader`, `FileIO`, and `TextIOWrapper` classes. - implement `__getstate__`, `__setstate__` to `Universe` and `BaseReader` - fix DCD, XDR pickle issue - modified gsd and ncdf to be picklabel - modified ChainReader to be picklabel - ensure chemfiles is picklable * add tests (MultiFrameReader will test for serializability) * update CHANGELOG * update docs Note: This merge squashed 120 commits. See PR MDAnalysis#2723 for the full history with 420 comments.

richardjgowers requested review from orbeckst and VOD555 November 15, 2018 00:21

richardjgowers commented Nov 15, 2018

View reviewed changes

orbeckst added the persistence label Dec 5, 2018

orbeckst requested changes Dec 5, 2018

View reviewed changes

richardjgowers added 4 commits February 16, 2019 10:52

added experimental pickling support to Universe

a00840e

fixed pickle test

133072f

wip of reader pickling

aa595f5

simplified serialisation support more

ba1a138

richardjgowers force-pushed the serialise branch from 0596149 to ba1a138 Compare February 16, 2019 16:52

richardjgowers commented Feb 16, 2019

View reviewed changes

richardjgowers added 3 commits February 16, 2019 11:14

make AuxReader not seralise (until tested)

caa588d

add tests for multiprocessing

e21361d

added more formats to multiprocessing tests

99bdd08

and broke everything

pickling all readers now works...

f1a7d5a

richardjgowers commented Feb 20, 2019

View reviewed changes

orbeckst reviewed Feb 3, 2020

View reviewed changes

yuxuanzhuang mentioned this pull request May 29, 2020

Updating Serialization Functionality from PR #2140 #2704

Closed

4 tasks

orbeckst closed this May 29, 2020

RMeli deleted the serialise branch December 23, 2023 21:57

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

added experimental pickling support to Universe #2140

added experimental pickling support to Universe #2140

richardjgowers commented Nov 15, 2018 •

edited

Loading

richardjgowers Nov 15, 2018

richardjgowers Nov 15, 2018

orbeckst Dec 5, 2018

richardjgowers Feb 16, 2019

codecov bot commented Nov 21, 2018 •

edited

Loading

richardjgowers commented Dec 2, 2018

orbeckst left a comment

orbeckst Dec 5, 2018

orbeckst Dec 5, 2018

richardjgowers commented Dec 6, 2018

richardjgowers Feb 16, 2019

richardjgowers Feb 16, 2019

richardjgowers commented Feb 16, 2019

orbeckst commented Feb 20, 2019

orbeckst commented Feb 20, 2019

orbeckst commented Feb 20, 2019

richardjgowers Feb 20, 2019

orbeckst Feb 3, 2020

richardjgowers Feb 20, 2019

richardjgowers commented Feb 20, 2019

orbeckst Feb 3, 2020 •

edited

Loading

orbeckst Feb 3, 2020

orbeckst commented May 29, 2020

orbeckst commented May 29, 2020

	self.readers = [core.reader(filename, **kwargs)
	for filename in filenames]

added experimental pickling support to Universe #2140

added experimental pickling support to Universe #2140

Conversation

richardjgowers commented Nov 15, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Nov 21, 2018 • edited Loading

Codecov Report

richardjgowers commented Dec 2, 2018

orbeckst left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

richardjgowers commented Dec 6, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

richardjgowers commented Feb 16, 2019

orbeckst commented Feb 20, 2019

orbeckst commented Feb 20, 2019

orbeckst commented Feb 20, 2019

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

richardjgowers commented Feb 20, 2019

orbeckst Feb 3, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

orbeckst commented May 29, 2020

orbeckst commented May 29, 2020

richardjgowers commented Nov 15, 2018 •

edited

Loading

codecov bot commented Nov 21, 2018 •

edited

Loading

orbeckst Feb 3, 2020 •

edited

Loading