hdf5: Allow saving custom attributes

Signed-off-by: yuk <yuk@m-labs.hk>

artiq/language/environment.py

@@ -383,6 +383,21 @@ def append_to_dataset(self, key, value):
         efficiently as incremental modifications in broadcast mode."""
         self.__dataset_mgr.append_to(key, value)
 
+    @rpc(flags={"async"})
+    def set_dataset_metadata(self, key, metadata):
+        """Attach metadata to the dataset itself, or to a key in the dataset.
+
+        The metadata is saved as HDF5 attributes if there was a call to
+        ``set_dataset(..., archive=True)`` with the same key.
+
+        :param key: If ``None``, attach the metadata to the dataset itself
+            (to the ``datasets`` group when saving to HDF5).
+            Otherwise, attach the metadata to the specified key.
+        :param metadata: Dict of key-value pair.
+            Must be compatible with HDF5 attributes.
+        """
+        self.__dataset_mgr.set_metadata(key, metadata)
+
     def get_dataset(self, key, default=NoDefault, archive=True):
         """Returns the contents of a dataset.
 

artiq/master/worker_db.py

@@ -110,6 +110,7 @@ class DatasetManager:
     def __init__(self, ddb):
         self._broadcaster = Notifier(dict())
         self.local = dict()
+        self.hdf5_attributes = dict()
         self.archive = dict()
 
         self.ddb = ddb
@@ -163,6 +164,12 @@ def get(self, key, archive=False):
             self.archive[key] = data
         return data
 
+    def set_metadata(self, key, metadata):
+        if key:
+            self.hdf5_attributes["datasets/" + key] = metadata
+        else:
+            self.hdf5_attributes["datasets"] = metadata
+
     def write_hdf5(self, f):
         datasets_group = f.create_group("datasets")
         for k, v in self.local.items():
@@ -172,6 +179,11 @@ def write_hdf5(self, f):
         for k, v in self.archive.items():
             _write(archive_group, k, v)
 
+    def write_hdf5_attributes(self, f):
+        for k, attrs in self.hdf5_attributes.items():
+            if k in f:
+                _write_attributes(f, k, attrs)
+
 
 def _write(group, k, v):
     # Add context to exception message when the user writes a dataset that is
@@ -181,3 +193,14 @@ def _write(group, k, v):
     except TypeError as e:
         raise TypeError("Error writing dataset '{}' of type '{}': {}".format(
             k, type(v), e))
+
+
+def _write_attributes(f, k, attrs):
+    # Add context to exception message when the user writes a attribute that is
+    # not representable in HDF5.
+    try:
+        for attr_k, attr_v in attrs.items():
+            f[k].attrs[attr_k] = attr_v
+    except TypeError as e:
+        raise TypeError("Error writing attribute '{}' of type '{}': {}".format(
+            attr_k, type(attr_v), e))

artiq/master/worker_impl.py

@@ -289,6 +289,7 @@ def write_results():
         filename = "{:09}-{}.h5".format(rid, exp.__name__)
         with h5py.File(filename, "w") as f:
             dataset_mgr.write_hdf5(f)
+            dataset_mgr.write_hdf5_attributes(f)
             f["artiq_version"] = artiq_version
             f["rid"] = rid
             f["start_time"] = start_time