Repository: IDSIA/sacred
Pull request #337 (2018-07-13 12:00:27)

Add option to add metadata to artifacts
A dict with custom metadata can now be added to artifacts. This only has
an effect for the MongoObserver. The metadata will appear in the
`.metadata` attribute of the file like objects retrieved with pymongo
from gridfs.


Patch:
diff --git a/sacred/experiment.py b/sacred/experiment.py
index 4236b2e..d3061ab 100755
--- a/sacred/experiment.py
+++ b/sacred/experiment.py
@@ -321,7 +321,7 @@ class Experiment(Ingredient):
         assert self.current_run is not None, "Can only be called during a run."
         self.current_run.add_resource(filename)
 
-    def add_artifact(self, filename, name=None):
+    def add_artifact(self, filename, name=None, metadata=None):
         """Add a file as an artifact.
 
         In Sacred terminology an artifact is a file produced by the experiment
@@ -338,10 +338,12 @@ class Experiment(Ingredient):
         name : str, optional
             optionally set the name of the artifact.
             Defaults to the relative file-path.
-
+        metadata: dict, optional
+            optionally attach metadata to the artifact.
+            This only has an effect when using the MongoObserver.
         """
         assert self.current_run is not None, "Can only be called during a run."
-        self.current_run.add_artifact(filename, name)
+        self.current_run.add_artifact(filename, name, metadata)
 
     @property
     def info(self):
diff --git a/sacred/observers/base.py b/sacred/observers/base.py
index 2a68e7b..a41a42c 100644
--- a/sacred/observers/base.py
+++ b/sacred/observers/base.py
@@ -33,5 +33,5 @@ class RunObserver(object):
     def resource_event(self, filename):
         pass
 
-    def artifact_event(self, name, filename):
+    def artifact_event(self, name, filename, metadata=None):
         pass
diff --git a/sacred/observers/file_storage.py b/sacred/observers/file_storage.py
index a491d6c..7d9430a 100644
--- a/sacred/observers/file_storage.py
+++ b/sacred/observers/file_storage.py
@@ -204,7 +204,7 @@ class FileStorageObserver(RunObserver):
         self.run_entry['resources'].append([filename, store_path])
         self.save_json(self.run_entry, 'run.json')
 
-    def artifact_event(self, name, filename):
+    def artifact_event(self, name, filename, metadata=None):
         self.save_file(filename, name)
         self.run_entry['artifacts'].append(name)
         self.save_json(self.run_entry, 'run.json')
diff --git a/sacred/observers/mongo.py b/sacred/observers/mongo.py
index a6fdc9d..ef9db79 100644
--- a/sacred/observers/mongo.py
+++ b/sacred/observers/mongo.py
@@ -184,12 +184,12 @@ class MongoObserver(RunObserver):
         self.run_entry['resources'].append((filename, md5hash))
         self.save()
 
-    def artifact_event(self, name, filename):
+    def artifact_event(self, name, filename, metadata=None):
         with open(filename, 'rb') as f:
             run_id = self.run_entry['_id']
             db_filename = 'artifact://{}/{}/{}'.format(self.runs.name, run_id,
                                                        name)
-            file_id = self.fs.put(f, filename=db_filename)
+            file_id = self.fs.put(f, filename=db_filename, metadata=metadata)
         self.run_entry['artifacts'].append({'name': name,
                                             'file_id': file_id})
         self.save()
diff --git a/sacred/observers/sql.py b/sacred/observers/sql.py
index 1ffb540..a84b47b 100644
--- a/sacred/observers/sql.py
+++ b/sacred/observers/sql.py
@@ -106,7 +106,7 @@ class SqlObserver(RunObserver):
         self.run.resources.append(res)
         self.save()
 
-    def artifact_event(self, name, filename):
+    def artifact_event(self, name, filename, metadata=None):
         a = Artifact.create(name, filename)
         self.run.artifacts.append(a)
         self.save()
diff --git a/sacred/observers/tinydb_hashfs.py b/sacred/observers/tinydb_hashfs.py
index f4896c1..1a17cf4 100644
--- a/sacred/observers/tinydb_hashfs.py
+++ b/sacred/observers/tinydb_hashfs.py
@@ -189,7 +189,7 @@ class TinyDbObserver(RunObserver):
             self.run_entry['resources'].append(resource)
             self.save()
 
-    def artifact_event(self, name, filename):
+    def artifact_event(self, name, filename, metadata=None):
 
         id_ = self.fs.put(filename).id
         handle = BufferedReaderWrapper(open(filename, 'rb'))
diff --git a/sacred/run.py b/sacred/run.py
index dfba3cc..7848d50 100755
--- a/sacred/run.py
+++ b/sacred/run.py
@@ -157,7 +157,7 @@ class Run(object):
         filename = os.path.abspath(filename)
         self._emit_resource_added(filename)
 
-    def add_artifact(self, filename, name=None):
+    def add_artifact(self, filename, name=None, metadata=None):
         """Add a file as an artifact.
 
         In Sacred terminology an artifact is a file produced by the experiment
@@ -173,10 +173,13 @@ class Run(object):
         name : str, optional
             optionally set the name of the artifact.
             Defaults to the filename.
+        metadata: dict
+            optionally attach metadata to the artifact.
+            This only has an effect when using the MongoObserver.
         """
         filename = os.path.abspath(filename)
         name = os.path.basename(filename) if name is None else name
-        self._emit_artifact_added(name, filename)
+        self._emit_artifact_added(name, filename, metadata)
 
     def __call__(self, *args):
         r"""Start this run.
@@ -374,11 +377,12 @@ class Run(object):
         for observer in self.observers:
             self._safe_call(observer, 'resource_event', filename=filename)
 
-    def _emit_artifact_added(self, name, filename):
+    def _emit_artifact_added(self, name, filename, metadata):
         for observer in self.observers:
             self._safe_call(observer, 'artifact_event',
                             name=name,
-                            filename=filename)
+                            filename=filename,
+                            metadata=metadata)
 
     def _safe_call(self, obs, method, **kwargs):
         if obs not in self._failed_observers and hasattr(obs, method):


Test patch:
diff --git a/tests/test_run.py b/tests/test_run.py
index 8903c0b..dfb2500 100755
--- a/tests/test_run.py
+++ b/tests/test_run.py
@@ -129,8 +129,9 @@ def test_run_heartbeat_event(run):
 def test_run_artifact_event(run):
     observer = run.observers[0]
     handle, f_name = tempfile.mkstemp()
-    run.add_artifact(f_name, name='foobar')
-    observer.artifact_event.assert_called_with(filename=f_name, name='foobar')
+    metadata = {'testkey': 42}
+    run.add_artifact(f_name, name='foobar', metadata=metadata)
+    observer.artifact_event.assert_called_with(filename=f_name, name='foobar', metadata=metadata)
     os.close(handle)
     os.remove(f_name)