diff --git a/simdb/db.py b/simdb/db.py index cef7ad3..e55209d 100644 --- a/simdb/db.py +++ b/simdb/db.py @@ -28,6 +28,8 @@ PY2 = sys.version_info.major == 2 try: + from pymor.core.pickle import dumps, load +except ImportError: from cPickle import dumps, load except ImportError: from pickle import dumps, load diff --git a/simdb/run.py b/simdb/run.py index 34c2529..dbd469f 100644 --- a/simdb/run.py +++ b/simdb/run.py @@ -30,9 +30,11 @@ import atexit from collections import defaultdict try: - from cPickle import dump + from pymor.core.pickle import dump, PicklingError except ImportError: - from pickle import dump + from cPickle import dump, PicklingError +except ImportError: + from pickle import dump, PicklingError import datetime import os if PY2: @@ -288,7 +290,15 @@ def process_data(v): data = {k: process_data(v) for k, v in _current_dataset_data.items()} with open(os.path.join(_current_dataset, 'DATA'), 'wb') as f: - dump(data, f, protocol=-1) + try: + dump(data, f, protocol=-1) + except PicklingError: + for kk, vv in data.items(): + try: + dump({kk: vv}, f, protocol=-1) + except PicklingError as e: + print(f'could not pickle "{kk}"') + raise e def get_metadata(v): if isinstance(v, np.ndarray): @@ -405,8 +415,7 @@ def dump_failed(filename): _saved_excepthook, sys.excepthook = sys.excepthook, _excepthook -@atexit.register -def _exit_hook(): +def declare_finished(): if _run: finished = datetime.datetime.now() if _current_dataset: @@ -415,3 +424,8 @@ def _exit_hook(): yaml.dump(finished, f) with open(os.path.join(_db_path, 'RUNS', _run, 'FINISHED'), 'wt') as f: yaml.dump(finished, f) + + +@atexit.register +def _exit_hook(): + declare_finished()