Skip to content

Commit 0e83df5

Browse files
authored
Merge pull request #985 from dwhswenson/simstore-updates
Integration tests for SimStore StorableFunctions
2 parents d32c198 + 12b0158 commit 0e83df5

7 files changed

Lines changed: 402 additions & 22 deletions

File tree

‎openpathsampling/experimental/simstore/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
)
55
from .class_info import ClassInfoContainer, ClassInfo
66
from .sql_backend import SQLStorageBackend
7+
from .memory_backend import MemoryStorageBackend
78
from . import dict_serialization_helpers
89
from .storable_functions import (
910
Processor, StorableFunctionConfig, StorableFunctionResults,

‎openpathsampling/experimental/simstore/memory_backend.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,9 +79,16 @@ def __init__(self, mode='w'):
7979
self.uuid_table = {} # maps UUID to table name
8080
self._table_to_class = {}
8181

82+
@property
83+
def identifier(self):
84+
return hex(id(self))
85+
8286
def close(self):
8387
pass # no such thing as closing here, but may be needed for API
8488

89+
def register_type(self, type_str, backend_type):
90+
pass # no need to do anything here?
91+
8592
def register_schema(self, schema, table_to_class, metadata=None):
8693
for table_name, schema_entries in schema.items():
8794
logger.debug("Registering table: %s" % table_name)

‎openpathsampling/experimental/simstore/sql_backend.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -413,10 +413,10 @@ def load_storable_function_results(self, table_name, uuids):
413413
table = self.metadata.tables[table_name]
414414
results = []
415415
for uuid_block in tools.block(uuids, self.max_query_size):
416-
# logger
417-
uuid_sel = table.select(
418-
sql.exists().where(table.c.uuid.in_(uuid_block))
419-
)
416+
# uuid_sel = table.select(
417+
# sql.exists().where(table.c.uuid.in_(uuid_block))
418+
# )
419+
uuid_sel = table.select().where(table.c.uuid.in_(uuid_block))
420420
with self.engine.connect() as conn:
421421
res = conn.execute(uuid_sel)
422422
res = res.fetchall()

‎openpathsampling/experimental/simstore/storable_functions.py‎

Lines changed: 54 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,21 @@ def from_dict(cls, dct):
222222
obj.result_dict = dct['result_dict']
223223

224224

225+
def _storage_exception_msg(func, storage, exc):
226+
func_id = str(get_uuid(func))
227+
if func.is_named:
228+
func_id += " (" + func.name + ")"
229+
try:
230+
storage_id = " in " + str(storage.identifier) + "."
231+
except:
232+
storage_id = "."
233+
234+
msg = ("A problem occured while loading disk-cached results. "
235+
"Function " + func_id + storage_id + "\n" + str(type(exc))
236+
+ str(exc))
237+
return msg
238+
239+
225240
class StorableFunction(StorableNamedObject):
226241
"""Function wrapper, providing result caching and storage to disk.
227242
@@ -254,18 +269,18 @@ def __init__(self, func, func_config=None, store_source=None, **kwargs):
254269

255270
self.local_cache = None # set correctly by self.mode setter
256271
self._disk_cache = True
257-
self._handler = None
272+
self._handlers = set([])
258273
self.mode = 'analysis'
259274

260-
def set_handler(self, storage, override=False):
261-
if not override and self._handler is not None:
262-
raise RuntimeError("Handler for this StorableFunction has "
263-
+ "already been set:" + str(self._handler))
264-
self._handler = storage
275+
def add_handler(self, storage, override=False):
276+
self._handlers.add(storage)
277+
278+
def remove_handler(self, handler):
279+
self._handlers.discard(handler)
265280

266281
@property
267282
def has_handler(self):
268-
return self._handler is not None
283+
return len(self._handlers) > 0
269284

270285
@property
271286
def disk_cache(self):
@@ -314,11 +329,14 @@ def from_dict(cls, dct):
314329

315330
def preload_cache(self, storage=None):
316331
if storage is None:
317-
storage = self._handler.storage
332+
storages = [h.storage for h in self._handlers]
333+
else:
334+
storages = [storage]
318335

319336
uuid = get_uuid(self)
320-
cache_values = storage.backend.load_storable_function_table(uuid)
321-
self.local_cache.cache_results(cache_values)
337+
for storage in storages:
338+
cache = storage.backend.load_storable_function_table(uuid)
339+
self.local_cache.cache_results(cache)
322340

323341
def is_scalar(self, item):
324342
"""Determine whether the input needs to be wrapped in a list.
@@ -360,10 +378,27 @@ def _get_cached(self, uuid_items):
360378
return self.local_cache.get_results_as_dict(uuid_items)
361379

362380
def _get_storage(self, uuid_items):
363-
if not self._handler:
381+
if not self.has_handler:
364382
return {}, uuid_items
365383

366-
return self._handler.get_function_results(get_uuid(self), uuid_items)
384+
missing = uuid_items
385+
found = {}
386+
my_uuid = get_uuid(self)
387+
for handler in self._handlers:
388+
try:
389+
uuid_map, missing = handler.get_function_results(
390+
my_uuid, missing
391+
)
392+
except Exception as e:
393+
# for any error, we just warn -- can be correctly calculated
394+
# in eval
395+
msg = _storage_exception_msg(self, handler.storage, e)
396+
warnings.warn(msg)
397+
uuid_map = {}
398+
399+
found.update(uuid_map)
400+
401+
return found, missing
367402

368403
def __call__(self, items):
369404
# important: implementation is that we always try to take an
@@ -475,9 +510,13 @@ def register_function(self, func, example_result=None):
475510
if not is_registered:
476511
self.all_functions[func_uuid].append(func)
477512

478-
# set handler; only if the table exists
479-
if not func.has_handler and (add_table or not needs_table):
480-
func.set_handler(self)
513+
if add_table or not needs_table:
514+
func.add_handler(self)
515+
516+
def close(self):
517+
for uuid, copies in self.all_functions.items():
518+
for func in copies:
519+
func.remove_handler(self)
481520

482521
def clear_non_canonical(self):
483522
self.all_functions = collections.defaultdict(list)

‎openpathsampling/experimental/simstore/storage.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,8 @@ def initialize_with_mode(self, mode):
113113
self._load_missing_info_tables(table_to_class)
114114
sim_objs = {get_uuid(obj): obj
115115
for obj in self.simulation_objects}
116+
sim_objs.update({get_uuid(obj): obj
117+
for obj in self.storable_functions})
116118
self._simulation_objects.update(sim_objs)
117119
self._update_pseudo_tables(sim_objs)
118120

@@ -136,13 +138,18 @@ def _load_missing_info_tables(self, table_to_class):
136138
raise RuntimeError("Unable to register existing database "
137139
+ "tables: " + str(missing_info_tables))
138140

141+
def register_from_tables(self, table_names, classes):
142+
# override in subclass to handle special lookups
143+
pass
144+
139145
def stash(self, objects):
140146
objects = tools.listify(objects)
141147
self._stashed.extend(objects)
142148

143149
def close(self):
144150
# TODO: should sync on close
145151
self.backend.close()
152+
self._sf_handler.close()
146153
for fallback in self.fallbacks:
147154
fallback.close()
148155

@@ -253,7 +260,6 @@ def save(self, obj_list, use_cache=True):
253260
if not input_uuids:
254261
return # exit early if everything is already in storage
255262

256-
257263
# check default table for things to register; register them
258264
# TODO: move to function: self.register_missing(by_table)
259265
# TODO: convert to while?
@@ -285,6 +291,7 @@ def save(self, obj_list, use_cache=True):
285291
old_missing = missing
286292

287293
# TODO: move to function self.store_sfr_results(by_table)
294+
self.save_function_results() # always for canonical
288295
has_sfr = (self.class_info.sfr_info is not None
289296
and self.class_info.sfr_info.table in by_table)
290297
if has_sfr:
@@ -441,7 +448,10 @@ def sync_all(self):
441448
def _cache_simulation_objects(self):
442449
# load up all the simulation objects
443450
try:
444-
backend_iter = self.backend.table_iterator('simulation_objects')
451+
backend_iter = itertools.chain(
452+
self.backend.table_iterator('simulation_objects'),
453+
self.backend.table_iterator('storable_functions')
454+
)
445455
sim_obj_uuids = [row.uuid for row in backend_iter]
446456
except KeyError:
447457
# TODO: this should probably be a custom error; don't rely on

‎openpathsampling/experimental/simstore/test_storable_function.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -386,7 +386,7 @@ def test_register_function(self, has_table, with_result):
386386
assert self.sf_handler.all_functions[uuid] == [self.func]
387387
if not unable_to_register:
388388
assert self.func.has_handler
389-
assert self.func._handler == self.sf_handler
389+
assert self.func._handlers == {self.sf_handler}
390390
assert self.sf_handler.functions == [self.func]
391391

392392
# make a copy of the func

0 commit comments

Comments
 (0)