Add API to attach hdf5 attributes to datasets

RELEASE_NOTES.rst

@@ -11,6 +11,8 @@ Highlights:
 * Implemented Phaser-servo. This requires recent gateware on Phaser.
 * Implemented Phaser-MIQRO support. This requires the Phaser MIQRO gateware
   variant.
+* HDF5 attributes can be attached to datasets using ``set_dataset_metadata()``.
+
 
 ARTIQ-7
 -------

artiq/examples/no_hardware/repository/hdf5_attributes.py

@@ -0,0 +1,24 @@
+import numpy as np
+
+from artiq.experiment import *
+
+
+class HDF5Attributes(EnvExperiment):
+    """Archive data to HDF5 with attributes"""
+    def run(self):
+        # Attach attributes to the HDF5 group `datasets`
+        self.set_dataset_metadata(None, {
+            "arr": np.array([1, 2, 3]),
+            "description": "demo",
+        })
+
+        dummy = np.empty(20)
+        dummy.fill(np.nan)
+        # `archive=True` is required in order to
+        # attach attributes to HDF5 datasets
+        self.set_dataset("dummy", dummy,
+                         broadcast=True, archive=True)
+        self.set_dataset_metadata("dummy", {"k1": "v1", "k2": "v2"})
+
+        # Attach metadata to an absent key is no-op
+        self.set_dataset_metadata("nothing", {"no": "op"})

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,14 @@ def get(self, key, archive=False):
             self.archive[key] = data
         return data
 
+    def set_metadata(self, key, metadata):
+        if key:
+            if key not in self.local:
+                logger.warning(f"Key '{key}' not found in dataset.")
+            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 +181,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 +195,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