dwhswenson · Mar 1, 2021 · Feb 25, 2021 · Feb 26, 2021 · Feb 26, 2021 · Feb 26, 2021
diff --git a/openpathsampling/experimental/simstore/__init__.py b/openpathsampling/experimental/simstore/__init__.py
@@ -4,6 +4,7 @@
 )
 from .class_info import ClassInfoContainer, ClassInfo
 from .sql_backend import SQLStorageBackend
+from .memory_backend import MemoryStorageBackend
 from . import dict_serialization_helpers
 from .storable_functions import (
     Processor, StorableFunctionConfig, StorableFunctionResults,

diff --git a/openpathsampling/experimental/simstore/memory_backend.py b/openpathsampling/experimental/simstore/memory_backend.py
@@ -79,9 +79,16 @@ def __init__(self, mode='w'):
         self.uuid_table = {}  # maps UUID to table name
         self._table_to_class = {}
 
+    @property
+    def identifier(self):
+        return hex(id(self))
+
     def close(self):
         pass  # no such thing as closing here, but may be needed for API
 
+    def register_type(self, type_str, backend_type):
+        pass  # no need to do anything here?
+
     def register_schema(self, schema, table_to_class, metadata=None):
         for table_name, schema_entries in schema.items():
             logger.debug("Registering table: %s" % table_name)

diff --git a/openpathsampling/experimental/simstore/sql_backend.py b/openpathsampling/experimental/simstore/sql_backend.py
@@ -413,10 +413,10 @@ def load_storable_function_results(self, table_name, uuids):
         table = self.metadata.tables[table_name]
         results = []
         for uuid_block in tools.block(uuids, self.max_query_size):
-            # logger
-            uuid_sel = table.select(
-                sql.exists().where(table.c.uuid.in_(uuid_block))
-            )
+            # uuid_sel = table.select(
+                # sql.exists().where(table.c.uuid.in_(uuid_block))
+            # )
+            uuid_sel = table.select().where(table.c.uuid.in_(uuid_block))
             with self.engine.connect() as conn:
                 res = conn.execute(uuid_sel)
                 res = res.fetchall()

diff --git a/openpathsampling/experimental/simstore/storable_functions.py b/openpathsampling/experimental/simstore/storable_functions.py
@@ -222,6 +222,21 @@ def from_dict(cls, dct):
         obj.result_dict = dct['result_dict']
 
 
+def _storage_exception_msg(func, storage, exc):
+    func_id = str(get_uuid(func))
+    if func.is_named:
+        func_id += " (" + func.name + ")"
+    try:
+        storage_id = " in " + str(storage.identifier) + "."
+    except:
+        storage_id = "."
+
+    msg = ("A problem occured while loading disk-cached results. "
+           "Function " + func_id + storage_id + "\n" + str(type(exc))
+           + str(exc))
+    return msg
+
+
 class StorableFunction(StorableNamedObject):
     """Function wrapper, providing result caching and storage to disk.
 
@@ -254,18 +269,18 @@ def __init__(self, func, func_config=None, store_source=None, **kwargs):
 
         self.local_cache = None  # set correctly by self.mode setter
         self._disk_cache = True
-        self._handler = None
+        self._handlers = set([])
         self.mode = 'analysis'
 
-    def set_handler(self, storage, override=False):
-        if not override and self._handler is not None:
-            raise RuntimeError("Handler for this StorableFunction has "
-                               + "already been set:" + str(self._handler))
-        self._handler = storage
+    def add_handler(self, storage, override=False):
+        self._handlers.add(storage)
+
+    def remove_handler(self, handler):
+        self._handlers.discard(handler)
 
     @property
     def has_handler(self):
-        return self._handler is not None
+        return len(self._handlers) > 0
 
     @property
     def disk_cache(self):
@@ -314,11 +329,14 @@ def from_dict(cls, dct):
 
     def preload_cache(self, storage=None):
         if storage is None:
-            storage = self._handler.storage
+            storages = [h.storage for h in self._handlers]
+        else:
+            storages = [storage]
 
         uuid = get_uuid(self)
-        cache_values = storage.backend.load_storable_function_table(uuid)
-        self.local_cache.cache_results(cache_values)
+        for storage in storages:
+            cache = storage.backend.load_storable_function_table(uuid)
+            self.local_cache.cache_results(cache)
 
     def is_scalar(self, item):
         """Determine whether the input needs to be wrapped in a list.
@@ -360,10 +378,27 @@ def _get_cached(self, uuid_items):
         return self.local_cache.get_results_as_dict(uuid_items)
 
     def _get_storage(self, uuid_items):
-        if not self._handler:
+        if not self.has_handler:
             return {}, uuid_items
 
-        return self._handler.get_function_results(get_uuid(self), uuid_items)
+        missing = uuid_items
+        found = {}
+        my_uuid = get_uuid(self)
+        for handler in self._handlers:
+            try:
+                uuid_map, missing = handler.get_function_results(
+                    my_uuid, missing
+                )
+            except Exception as e:
+                # for any error, we just warn -- can be correctly calculated
+                # in eval
+                msg = _storage_exception_msg(self, handler.storage, e)
+                warnings.warn(msg)
+                uuid_map = {}
+
+            found.update(uuid_map)
+
+        return found, missing
 
     def __call__(self, items):
         # important: implementation is that we always try to take an
@@ -475,9 +510,13 @@ def register_function(self, func, example_result=None):
         if not is_registered:
             self.all_functions[func_uuid].append(func)
 
-        # set handler; only if the table exists
-        if not func.has_handler and (add_table or not needs_table):
-            func.set_handler(self)
+        if add_table or not needs_table:
+            func.add_handler(self)
+
+    def close(self):
+        for uuid, copies in self.all_functions.items():
+            for func in copies:
+                func.remove_handler(self)
 
     def clear_non_canonical(self):
         self.all_functions = collections.defaultdict(list)

diff --git a/openpathsampling/experimental/simstore/storage.py b/openpathsampling/experimental/simstore/storage.py
@@ -113,6 +113,8 @@ def initialize_with_mode(self, mode):
             self._load_missing_info_tables(table_to_class)
             sim_objs = {get_uuid(obj): obj
                         for obj in self.simulation_objects}
+            sim_objs.update({get_uuid(obj): obj
+                             for obj in self.storable_functions})
             self._simulation_objects.update(sim_objs)
             self._update_pseudo_tables(sim_objs)
 
@@ -136,13 +138,18 @@ def _load_missing_info_tables(self, table_to_class):
             raise RuntimeError("Unable to register existing database "
                                + "tables: " + str(missing_info_tables))
 
+    def register_from_tables(self, table_names, classes):
+        # override in subclass to handle special lookups
+        pass
+
     def stash(self, objects):
         objects = tools.listify(objects)
         self._stashed.extend(objects)
 
     def close(self):
         # TODO: should sync on close
         self.backend.close()
+        self._sf_handler.close()
         for fallback in self.fallbacks:
             fallback.close()
 
@@ -253,7 +260,6 @@ def save(self, obj_list, use_cache=True):
         if not input_uuids:
             return  # exit early if everything is already in storage
 
-
         # check default table for things to register; register them
         # TODO: move to function: self.register_missing(by_table)
         # TODO: convert to while?
@@ -285,6 +291,7 @@ def save(self, obj_list, use_cache=True):
             old_missing = missing
 
         # TODO: move to function self.store_sfr_results(by_table)
+        self.save_function_results()  # always for canonical
         has_sfr = (self.class_info.sfr_info is not None
                    and self.class_info.sfr_info.table in by_table)
         if has_sfr:
@@ -441,7 +448,10 @@ def sync_all(self):
     def _cache_simulation_objects(self):
         # load up all the simulation objects
         try:
-            backend_iter = self.backend.table_iterator('simulation_objects')
+            backend_iter = itertools.chain(
+                self.backend.table_iterator('simulation_objects'),
+                self.backend.table_iterator('storable_functions')
+            )
             sim_obj_uuids = [row.uuid for row in backend_iter]
         except KeyError:
             # TODO: this should probably be a custom error; don't rely on

diff --git a/openpathsampling/experimental/simstore/test_storable_function.py b/openpathsampling/experimental/simstore/test_storable_function.py
@@ -386,7 +386,7 @@ def test_register_function(self, has_table, with_result):
         assert self.sf_handler.all_functions[uuid] == [self.func]
         if not unable_to_register:
             assert self.func.has_handler
-            assert self.func._handler == self.sf_handler
+            assert self.func._handlers == {self.sf_handler}
             assert self.sf_handler.functions == [self.func]
 
         # make a copy of the func