occheung
/
artiq
forked from M-Labs/artiq
1
0
Fork 0
Code Pull Requests Activity

worker_db: write hdf5 dataset metadata

This commit is contained in:
Simon Renblad 2023-06-30 16:54:39 +08:00 committed by Sébastien Bourdeauducq
parent bf38fc8b0f
commit bfbe13e51b
1 changed files with 17 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
artiq/master/worker_db.py
View File

@ -111,11 +111,12 @@ class DatasetManager:
self._broadcaster = Notifier(dict()) self._broadcaster = Notifier(dict())
self.local = dict() self.local = dict()
self.archive = dict() self.archive = dict()
self.metadata = dict()
self.ddb = ddb self.ddb = ddb
self._broadcaster.publish = ddb.update self._broadcaster.publish = ddb.update
def set(self, key, value, broadcast=False, persist=False, archive=True): def set(self, key, value, metadata, broadcast, persist, archive):
if persist: if persist:
broadcast = True broadcast = True
@ -123,7 +124,7 @@ class DatasetManager:
logger.warning(f"Dataset '{key}' will not be stored. Both 'broadcast' and 'archive' are set to False.") logger.warning(f"Dataset '{key}' will not be stored. Both 'broadcast' and 'archive' are set to False.")
if broadcast: if broadcast:
self._broadcaster[key] = persist, value self._broadcaster[key] = persist, value, metadata
elif key in self._broadcaster.raw_view: elif key in self._broadcaster.raw_view:
del self._broadcaster[key] del self._broadcaster[key]
@ -132,6 +133,8 @@ class DatasetManager:
elif key in self.local: elif key in self.local:
del self.local[key] del self.local[key]
self.metadata[key] = metadata
def _get_mutation_target(self, key): def _get_mutation_target(self, key):
target = self.local.get(key, None) target = self.local.get(key, None)
if key in self._broadcaster.raw_view: if key in self._broadcaster.raw_view:
@ -166,21 +169,30 @@ class DatasetManager:
self.archive[key] = data self.archive[key] = data
return data return data
def get_metadata(self, key):
if key in self.metadata:
return self.metadata[key]
return self.ddb.get_metadata(key)
def write_hdf5(self, f): def write_hdf5(self, f):
datasets_group = f.create_group("datasets") datasets_group = f.create_group("datasets")
for k, v in self.local.items(): for k, v in self.local.items():
_write(datasets_group, k, v) m = self.metadata.get(k, {})
_write(datasets_group, k, v, m)
archive_group = f.create_group("archive") archive_group = f.create_group("archive")
for k, v in self.archive.items(): for k, v in self.archive.items():
_write(archive_group, k, v) m = self.metadata.get(k, {})
_write(archive_group, k, v, m)
def _write(group, k, v): def _write(group, k, v, m):
# Add context to exception message when the user writes a dataset that is # Add context to exception message when the user writes a dataset that is
# not representable in HDF5. # not representable in HDF5.
try: try:
group[k] = v group[k] = v
for key, val in m.items():
group[k].attrs[key] = val
except TypeError as e: except TypeError as e:
raise TypeError("Error writing dataset '{}' of type '{}': {}".format( raise TypeError("Error writing dataset '{}' of type '{}': {}".format(
k, type(v), e)) k, type(v), e))