Merge 0b8109a into 37e6fc3

IDSIA · Jul 19, 2019 · 87d9d3e · 87d9d3e
2 parents 37e6fc3 + 0b8109a
commit 87d9d3e
Showing 1 changed file with 21 additions and 42 deletions.
diff --git a/sacred/observers/tinydb_hashfs.py b/sacred/observers/tinydb_hashfs.py
@@ -49,6 +49,25 @@ def __deepcopy__(self, memo):
         return BufferedReaderWrapper(f)
 
 
+def get_db_file_manager(root_dir):
+    fs = HashFS(os.path.join(root_dir, 'hashfs'), depth=3,
+                width=2, algorithm='md5')
+
+    # Setup Serialisation object for non list/dict objects
+    serialization_store = SerializationMiddleware()
+    serialization_store.register_serializer(DateTimeSerializer(), 'TinyDate')
+    serialization_store.register_serializer(FileSerializer(fs), 'TinyFile')
+
+    if opt.has_numpy:
+        serialization_store.register_serializer(NdArraySerializer(), 'TinyArray')
+    if opt.has_pandas:
+        serialization_store.register_serializer(DataFrameSerializer(), 'TinyDataFrame')
+        serialization_store.register_serializer(SeriesSerializer(), 'TinySeries')
+
+    db = TinyDB(os.path.join(root_dir, 'metadata.json'), storage=serialization_store)
+    return db, fs
+
+
 class TinyDbObserver(RunObserver):
 
     VERSION = "TinyDbObserver-{}".format(__version__)
@@ -60,28 +79,7 @@ def create(path='./runs_db', overwrite=None):
         if not os.path.exists(root_dir):
             os.makedirs(root_dir)
 
-        fs = HashFS(os.path.join(root_dir, 'hashfs'), depth=3,
-                    width=2, algorithm='md5')
-
-        # Setup Serialisation object for non list/dict objects
-        serialization_store = SerializationMiddleware()
-        serialization_store.register_serializer(DateTimeSerializer(),
-                                                'TinyDate')
-        serialization_store.register_serializer(FileSerializer(fs),
-                                                'TinyFile')
-
-        if opt.has_numpy:
-            serialization_store.register_serializer(NdArraySerializer(),
-                                                    'TinyArray')
-        if opt.has_pandas:
-            serialization_store.register_serializer(DataFrameSerializer(),
-                                                    'TinyDataFrame')
-            serialization_store.register_serializer(SeriesSerializer(),
-                                                    'TinySeries')
-
-        db = TinyDB(os.path.join(root_dir, 'metadata.json'),
-                    storage=serialization_store)
-
+        db, fs = get_db_file_manager(root_dir)
         return TinyDbObserver(db, fs, overwrite=overwrite, root=root_dir)
 
     def __init__(self, db, fs, overwrite=None, root=None):
@@ -237,26 +235,7 @@ def __init__(self, path):
         if not os.path.exists(root_dir):
             raise IOError('Path does not exist: %s' % path)
 
-        fs = HashFS(os.path.join(root_dir, 'hashfs'), depth=3,
-                    width=2, algorithm='md5')
-
-        # Setup Serialisation for non list/dict objects
-        serialization_store = SerializationMiddleware()
-        serialization_store.register_serializer(DateTimeSerializer(),
-                                                'TinyDate')
-        serialization_store.register_serializer(FileSerializer(fs),
-                                                'TinyFile')
-        if opt.has_numpy:
-            serialization_store.register_serializer(NdArraySerializer(),
-                                                    'TinyArray')
-        if opt.has_pandas:
-            serialization_store.register_serializer(DataFrameSerializer(),
-                                                    'TinyDataFrame')
-            serialization_store.register_serializer(SeriesSerializer(),
-                                                    'TinySeries')
-
-        db = TinyDB(os.path.join(root_dir, 'metadata.json'),
-                    storage=serialization_store)
+        db, fs = get_db_file_manager(root_dir)
 
         self.db = db
         self.runs = db.table('runs')