Add initial dask_integration & parallel_committor

dwhswenson · Jul 21, 2020 · Jul 21, 2020 · Jul 21, 2020 · Jul 22, 2020 · Jul 22, 2020
commit d48acad6d2cb0b5d69c2bca7b6f02ba1dd0f5571
diff --git a/openpathsampling/experimental/storage/dask_integration.py b/openpathsampling/experimental/storage/dask_integration.py
@@ -0,0 +1,98 @@
+import distributed
+from distributed.protocol.serialize import register_serialization_family
+
+import cloudpickle as cp
+
+from openpathsampling.experimental.storage.memory_backend import MemoryStorageBackend
+from openpathsampling.experimental.storage.ops_storage import OPSStorage
+from openpathsampling.experimental.storage.serialization_helpers import (
+    has_uuid, get_uuid
+)
+
+def ops_dumps(obj):
+    header = {'serializer': 'openpathsampling'}
+    if not has_uuid(obj):
+        # quick exit
+        raise NotImplementedError()
+
+    storage = OPSStorage.from_backend(MemoryStorageBackend())
+    storage.save(obj)
+    uuid = get_uuid(obj)
+    frames = [cp.dumps({'uuid': uuid, 'backend': storage.backend})]
+    return header, frames
+
+def ops_loads(header, frames):
+    if len(frames) > 1:
+        frame = ''.join(frames)
+    else:
+        frame = frames[0]
+    dct = cp.loads(frame)
+    backend = dct['backend']
+    backend.mode = 'r'
+    storage = OPSStorage.from_backend(backend)
+    for func in storage.storable_functions:
+        func.preload_cache(storage)
+    obj = storage.load([dct['uuid']])[0]
+    return obj
+
+
+def serialize_task(task):
+    # allows serialization of: (1) UUID objects; (2) methods of UUID object;
+    # (3) bare functions
+    if has_uuid(task):
+        func = None
+        obj = task
+    elif hasattr(task, '__self__') and has_uuid(task.__self__):
+        func = task.__name__
+        obj = task.__self__
+    else:
+        func = task
+        obj = None
+    return obj, func
+
+def deserialize_task(ser_task):
+    obj, func = ser_task
+    if obj is None:
+        task = func
+    elif func is None:
+        task = obj
+    else:
+        task = getattr(obj, func)
+    return task
+
+register_serialization_family('openpathsampling', ops_dumps, ops_loads)
+
+
+def _save_results(storage, results):
+    # trick to move storage.save storage process to OPS serialization, not
+    # cloudpickle
+    storage.save(results)
+
+def _remote_task(serialized_task, *args, **kwargs):
+    task = deserialize_task(serialized_task)
+    result = task(*args, **kwargs)
+    return result
+
+class DaskDistributedScheduler(object):
+    def __init__(self, client=None):
+        if client is None:
+            client = distributed.Client(
+                serializers=['openpathsampling', 'dask', 'pickle',  ],
+                deserializers=['openpathsampling', 'dask', 'pickle', ],
+            )
+        self.client = client
+
+    def store_results(self, storage, results):
+        self.client.fire_and_forget(_save_results, storage, results)
+
+    def wrap_task(self, task):
+        def inner(*args, **kwargs):
+            ser_task = serialize_task(task)
+            result = self.client.submit(_remote_task,
+                                        ser_task,
+                                        pure=False,
+                                        *args,
+                                        **kwargs)
+            return result
+
+        return inner
diff --git a/openpathsampling/experimental/storage/parallel_committor.py b/openpathsampling/experimental/storage/parallel_committor.py
@@ -0,0 +1,51 @@
+import openpathsampling as paths
+
+
+class NewCommittor(paths.PathSimulator):
+    def __init__(self, vol, engine, storage):
+        super().__init__(storage)
+        self.vol = vol
+        self.init_ens = paths.LengthEnsemble(1) & paths.AllInXEnsemble(vol)
+        self.final_ens = paths.LengthEnsemble(5) & paths.AllInXEnsemble(vol)
+        self.engine = engine
+        self.mover = paths.ForwardExtendMover(
+            ensemble=self.init_ens,
+            target_ensemble=self.final_ens,
+            engine=engine
+        )
+
+
+    def to_dict(self):
+        return {'vol': self.vol,
+                'init_ens': self.init_ens,
+                'final_ens': self.final_ens,
+                'engine': self.engine,
+                'mover': self.mover,
+                'storage': self.storage}
+
+    @classmethod
+    def from_dict(cls, dct):
+        obj = cls(dct['vol'], dct['engine'], dct['storage'])
+        obj.init_ens = dct['init_ens']
+        obj.final_ens = dct['final_ens']
+        obj.mover = dct['mover']
+        return obj
+
+    def committor_task(self, snap, step_num):
+        samples = paths.Sample(trajectory=paths.Trajectory([snap]),
+                               ensemble=self.init_ens,
+                               replica=0)
+
+        change = self.mover.move_core([samples])
+        step = paths.MCStep(
+            simulation=self,
+            mccycle=step_num,
+            active=paths.SampleSet(change.trials[0]),
+            change=change
+        )
+        return step
+
+    def run(self, snaps, n_steps):
+        for snap in snaps:
+            for step_num in range(n_steps):
+                step = self.committor_task(snap, step_num)