Source code for sciunit.models.backends
"""Base class for simulator backends for SciUnit models."""
import os
import inspect
import tempfile
import pickle
import shelve
available_backends = {}
[docs]def register_backends(vars):
"""Register backends for use with models.
`vars` should be a dictionary of variables obtained from e.g. `locals()`,
at least some of which are Backend classes, e.g. from imports.
"""
new_backends = {x if x is None else x.replace('Backend', ''): cls
for x, cls in vars.items()
if inspect.isclass(cls) and issubclass(cls, Backend)}
available_backends.update(new_backends)
[docs]class Backend(object):
"""
Base class for simulator backends.
Should only be used with model classes derived from `RunnableModel`.
Supports caching of simulation results.
Backend classes should implement simulator-specific
details of modifying, running, and reading results from the simulation.
"""
[docs] def init_backend(self, *args, **kwargs):
"""Initialize the backend."""
self.model.attrs = {}
self.use_memory_cache = kwargs.get('use_memory_cache', True)
if self.use_memory_cache:
self.init_memory_cache()
self.use_disk_cache = kwargs.get('use_disk_cache', False)
if self.use_disk_cache:
self.init_disk_cache()
self.load_model()
self.model.unpicklable += ['_backend']
#: Name of the backend
name = None
#: The function that handles running the simulation
f = None
#: Optional list of state variables for a backend to record.
recorded_variables = None
[docs] def init_cache(self):
"""Initialize the cache."""
self.init_memory_cache()
self.init_disk_cache()
[docs] def init_memory_cache(self):
"""Initialize the in-memory version of the cache."""
self.memory_cache = {}
[docs] def init_disk_cache(self):
"""Initialize the on-disk version of the cache."""
try:
# Cleanup old disk cache files
path = self.disk_cache_location
os.remove(path)
except Exception:
pass
self.disk_cache_location = os.path.join(tempfile.mkdtemp(), 'cache')
[docs] def get_memory_cache(self, key=None):
"""Return result in memory cache for key 'key' or None if not found."""
key = self.model.hash if key is None else key
self._results = self.memory_cache.get(key)
return self._results
[docs] def get_disk_cache(self, key=None):
"""Return result in disk cache for key 'key' or None if not found."""
key = self.model.hash if key is None else key
if not getattr(self, 'disk_cache_location', False):
self.init_disk_cache()
disk_cache = shelve.open(self.disk_cache_location)
self._results = disk_cache.get(key)
disk_cache.close()
return self._results
[docs] def set_memory_cache(self, results, key=None):
"""Store result in memory cache with key matching model state."""
key = self.model.hash if key is None else key
self.memory_cache[key] = results
[docs] def set_disk_cache(self, results, key=None):
"""Store result in disk cache with key matching model state."""
if not getattr(self, 'disk_cache_location', False):
self.init_disk_cache()
disk_cache = shelve.open(self.disk_cache_location)
key = self.model.hash if key is None else key
disk_cache[key] = results
disk_cache.close()
[docs] def load_model(self):
"""Load the model into memory."""
pass
[docs] def set_attrs(self, **attrs):
"""Set model attributes on the backend."""
pass
[docs] def set_run_params(self, **run_params):
"""Set model attributes on the backend."""
pass
[docs] def backend_run(self):
"""Check for cached results; then run the model if needed."""
key = self.model.hash
if self.use_memory_cache and self.get_memory_cache(key):
return self._results
if self.use_disk_cache and self.get_disk_cache(key):
return self._results
results = self._backend_run()
if self.use_memory_cache:
self.set_memory_cache(results, key)
if self.use_disk_cache:
self.set_disk_cache(results, key)
return results
[docs] def _backend_run(self):
"""Run the model via the backend."""
raise NotImplementedError("Each backend must implement '_backend_run'")
[docs] def save_results(self, path='.'):
"""Save results on disk."""
with open(path, 'wb') as f:
pickle.dump(self.results, f)
[docs]class BackendException(Exception):
"""Generic backend exception class."""
pass
# Register the base class as a Backend just so that there is
# always something available. This Backend won't do anything
# useful other than caching.
register_backends({None: Backend})