artiq/artiq/browser/experiments.py

422 lines
15 KiB
Python

import asyncio
import logging
import os
from functools import partial
from collections import OrderedDict
from PyQt5 import QtCore, QtGui, QtWidgets
import h5py
from sipyco import pyon
from artiq import __artiq_dir__ as artiq_dir
from artiq.gui.tools import (LayoutWidget, log_level_to_name, get_open_file_name)
from artiq.gui.entries import procdesc_to_entry, EntryTreeWidget
from artiq.master.worker import Worker, log_worker_exception
logger = logging.getLogger(__name__)
class _ArgumentEditor(EntryTreeWidget):
def __init__(self, dock):
EntryTreeWidget.__init__(self)
self._dock = dock
if not self._dock.arguments:
self.insertTopLevelItem(0, QtWidgets.QTreeWidgetItem(["No arguments"]))
for name, argument in self._dock.arguments.items():
self.set_argument(name, argument)
self.quickStyleClicked.connect(self._dock._run_clicked)
recompute_arguments = QtWidgets.QPushButton("Recompute all arguments")
recompute_arguments.setIcon(
QtWidgets.QApplication.style().standardIcon(
QtWidgets.QStyle.SP_BrowserReload))
recompute_arguments.clicked.connect(self._recompute_arguments_clicked)
load = QtWidgets.QPushButton("Set arguments from HDF5")
load.setToolTip("Set arguments from currently selected HDF5 file")
load.setIcon(QtWidgets.QApplication.style().standardIcon(
QtWidgets.QStyle.SP_DialogApplyButton))
load.clicked.connect(self._load_clicked)
buttons = LayoutWidget()
buttons.addWidget(recompute_arguments, 1, 1)
buttons.addWidget(load, 1, 2)
for i, s in enumerate((1, 0, 0, 1)):
buttons.layout.setColumnStretch(i, s)
self.setItemWidget(self.bottom_item, 1, buttons)
def _load_clicked(self):
asyncio.ensure_future(self._dock.load_hdf5_task())
def _recompute_arguments_clicked(self):
asyncio.ensure_future(self._dock._recompute_arguments())
def reset_entry(self, key):
asyncio.ensure_future(self._recompute_argument(key))
async def _recompute_argument(self, name):
try:
arginfo = await self._dock.compute_arginfo()
except:
logger.error("Could not recompute argument '%s' of '%s'",
name, self._dock.expurl, exc_info=True)
return
argument = self._dock.arguments[name]
procdesc = arginfo[name][0]
state = procdesc_to_entry(procdesc).default_state(procdesc)
argument["desc"] = procdesc
argument["state"] = state
self.update_argument(name, argument)
log_levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]
class _ExperimentDock(QtWidgets.QMdiSubWindow):
sigClosed = QtCore.pyqtSignal()
def __init__(self, area, expurl, arguments):
QtWidgets.QMdiSubWindow.__init__(self)
qfm = QtGui.QFontMetrics(self.font())
self.resize(100*qfm.averageCharWidth(), 30*qfm.lineSpacing())
self.setWindowTitle(expurl)
self.setWindowIcon(QtWidgets.QApplication.style().standardIcon(
QtWidgets.QStyle.SP_FileDialogContentsView))
self.setAcceptDrops(True)
self.layout = QtWidgets.QGridLayout()
top_widget = QtWidgets.QWidget()
top_widget.setLayout(self.layout)
self.setWidget(top_widget)
self.layout.setSpacing(5)
self.layout.setContentsMargins(5, 5, 5, 5)
self._area = area
self._run_task = None
self.expurl = expurl
self.arguments = arguments
self.options = {"log_level": logging.WARNING}
self.argeditor = _ArgumentEditor(self)
self.layout.addWidget(self.argeditor, 0, 0, 1, 5)
self.layout.setRowStretch(0, 1)
log_level = QtWidgets.QComboBox()
log_level.addItems(log_levels)
log_level.setCurrentIndex(1)
log_level.setToolTip("Minimum level for log entry production")
log_level_label = QtWidgets.QLabel("Logging level:")
log_level_label.setToolTip("Minimum level for log message production")
self.layout.addWidget(log_level_label, 3, 0)
self.layout.addWidget(log_level, 3, 1)
log_level.setCurrentIndex(log_levels.index(
log_level_to_name(self.options["log_level"])))
def update_log_level(index):
self.options["log_level"] = getattr(logging,
log_level.currentText())
log_level.currentIndexChanged.connect(update_log_level)
self.log_level = log_level
run = QtWidgets.QPushButton("Analyze")
run.setIcon(QtWidgets.QApplication.style().standardIcon(
QtWidgets.QStyle.SP_DialogOkButton))
run.setToolTip("Run analysis stage (Ctrl+Return)")
run.setShortcut("CTRL+RETURN")
run.setSizePolicy(QtWidgets.QSizePolicy.Expanding,
QtWidgets.QSizePolicy.Expanding)
self.layout.addWidget(run, 2, 4)
run.clicked.connect(self._run_clicked)
self._run = run
terminate = QtWidgets.QPushButton("Terminate")
terminate.setIcon(QtWidgets.QApplication.style().standardIcon(
QtWidgets.QStyle.SP_DialogCancelButton))
terminate.setToolTip("Terminate analysis (Ctrl+Backspace)")
terminate.setShortcut("CTRL+BACKSPACE")
terminate.setSizePolicy(QtWidgets.QSizePolicy.Expanding,
QtWidgets.QSizePolicy.Expanding)
self.layout.addWidget(terminate, 3, 4)
terminate.clicked.connect(self._terminate_clicked)
terminate.setEnabled(False)
self._terminate = terminate
def dragEnterEvent(self, ev):
if ev.mimeData().hasFormat("text/uri-list"):
ev.acceptProposedAction()
def dropEvent(self, ev):
for uri in ev.mimeData().urls():
if uri.scheme() == "file":
filename = QtCore.QDir.toNativeSeparators(uri.toLocalFile())
logger.debug("Loading HDF5 arguments from %s", filename)
asyncio.ensure_future(self.load_hdf5_task(filename))
break
async def compute_arginfo(self):
return await self._area.compute_arginfo(self.expurl)
async def _recompute_arguments(self, overrides={}):
try:
arginfo = await self.compute_arginfo()
except:
logger.error("Could not recompute arguments of '%s'",
self.expurl, exc_info=True)
return
for k, v in overrides.items():
# Some values (e.g. scans) may have multiple defaults in a list
if isinstance(arginfo[k][0].get("default"), list):
arginfo[k][0]["default"].insert(0, v)
else:
arginfo[k][0]["default"] = v
self.arguments = self._area.initialize_submission_arguments(arginfo)
state = self.argeditor.save_state()
self.argeditor.deleteLater()
self.argeditor = _ArgumentEditor(self)
self.layout.addWidget(self.argeditor, 0, 0, 1, 5)
self.argeditor.restore_state(state)
async def load_hdf5_task(self, filename=None):
if filename is None:
if self._area.dataset is None:
return
filename = self._area.dataset
try:
with h5py.File(filename, "r") as f:
expid = f["expid"][()]
expid = pyon.decode(expid)
arguments = expid["arguments"]
except:
logger.error("Could not retrieve expid from HDF5 file",
exc_info=True)
return
try:
self.log_level.setCurrentIndex(log_levels.index(
log_level_to_name(expid["log_level"])))
except:
logger.error("Could not set submission options from HDF5 expid",
exc_info=True)
return
await self._recompute_arguments(arguments)
def _run_clicked(self):
class_name, file = self.expurl.split("@", maxsplit=1)
expid = {
"repo_rev": "N/A",
"file": file,
"class_name": class_name,
"log_level": self.options["log_level"],
"arguments": {
name: procdesc_to_entry(argument["desc"]).state_to_value(
argument["state"])
for name, argument in self.arguments.items()},
}
self._run_task = asyncio.ensure_future(self._get_run_task(expid))
self._run.setEnabled(False)
self._terminate.setEnabled(True)
def done(fut):
logger.debug("Analysis done")
self._run_task = None
self._run.setEnabled(True)
self._terminate.setEnabled(False)
self._run_task.add_done_callback(done)
async def _get_run_task(self, expid):
logger.info("Running '%s'...", self.expurl)
worker = Worker(self._area.worker_handlers)
try:
await worker.build(rid=None, pipeline_name="browser",
wd=os.path.abspath("."),
expid=expid, priority=0)
await worker.analyze()
except:
logger.error("Failed to run '%s'", self.expurl)
log_worker_exception()
else:
logger.info("Finished running '%s'", self.expurl)
finally:
await worker.close()
def _terminate_clicked(self):
try:
self._run_task.cancel()
except:
logger.error("Unexpected failure terminating '%s'",
self.expurl, exc_info=True)
def closeEvent(self, event):
self.sigClosed.emit()
QtWidgets.QMdiSubWindow.closeEvent(self, event)
def save_state(self):
return {
"argeditor": self.argeditor.save_state(),
"geometry": bytes(self.saveGeometry()),
"options": self.options,
}
def restore_state(self, state):
self.argeditor.restore_state(state["argeditor"])
self.restoreGeometry(QtCore.QByteArray(state["geometry"]))
self.options = state["options"]
class LocalDatasetDB:
def __init__(self, dataset_sub):
self.dataset_sub = dataset_sub
dataset_sub.add_setmodel_callback(self.init)
def init(self, data):
self._data = data
def get(self, key):
return self._data.backing_store[key][1]
def update(self, mod):
self.dataset_sub.update(mod)
class ExperimentsArea(QtWidgets.QMdiArea):
def __init__(self, root, dataset_sub):
QtWidgets.QMdiArea.__init__(self)
self.pixmap = QtGui.QPixmap(os.path.join(
artiq_dir, "gui", "logo_ver.svg"))
self.current_dir = root
self.dataset = None
self.open_experiments = []
self._ddb = LocalDatasetDB(dataset_sub)
self.worker_handlers = {
"get_device_db": lambda: {},
"get_device": lambda key, resolve_alias=False: {"type": "dummy"},
"get_dataset": self._ddb.get,
"update_dataset": self._ddb.update,
}
def dataset_changed(self, path):
self.dataset = path
def dataset_activated(self, path):
sub = self.currentSubWindow()
if sub is None:
return
asyncio.ensure_future(sub.load_hdf5_task(path))
def mousePressEvent(self, ev):
if ev.button() == QtCore.Qt.LeftButton:
self.select_experiment()
def paintEvent(self, event):
QtWidgets.QMdiArea.paintEvent(self, event)
painter = QtGui.QPainter(self.viewport())
x = (self.width() - self.pixmap.width())//2
y = (self.height() - self.pixmap.height())//2
painter.setOpacity(0.5)
painter.drawPixmap(x, y, self.pixmap)
def save_state(self):
return {"experiments": [{
"expurl": dock.expurl,
"arguments": dock.arguments,
"dock": dock.save_state(),
} for dock in self.open_experiments]}
def restore_state(self, state):
if self.open_experiments:
raise NotImplementedError
for ex_state in state["experiments"]:
dock = self.open_experiment(ex_state["expurl"],
ex_state["arguments"])
dock.restore_state(ex_state["dock"])
def select_experiment(self):
asyncio.ensure_future(self._select_experiment_task())
async def _select_experiment_task(self):
try:
file = await get_open_file_name(
self, "Open experiment", self.current_dir,
"Experiments (*.py);;All files (*.*)")
except asyncio.CancelledError:
return
self.current_dir = os.path.dirname(file)
logger.debug("Opening experiment %s", file)
try:
description = await self.examine(file)
except:
logger.error("Could not examine experiment '%s'",
file, exc_info=True)
return
for class_name, class_desc in description.items():
expurl = "{}@{}".format(class_name, file)
arguments = self.initialize_submission_arguments(
class_desc["arginfo"])
self.open_experiment(expurl, arguments)
def initialize_submission_arguments(self, arginfo):
arguments = OrderedDict()
for name, (procdesc, group, tooltip) in arginfo.items():
if procdesc["ty"] == "EnumerationValue" and procdesc["quickstyle"]:
procdesc["quickstyle"] = False
state = procdesc_to_entry(procdesc).default_state(procdesc)
arguments[name] = {
"desc": procdesc,
"group": group,
"tooltip": tooltip,
"state": state # mutated by entries
}
return arguments
async def examine(self, file):
worker = Worker(self.worker_handlers)
try:
return await worker.examine("examine", file)
finally:
await worker.close()
async def compute_arginfo(self, expurl):
class_name, file = expurl.split("@", maxsplit=1)
try:
desc = await self.examine(file)
except:
logger.error("Could not examine experiment '%s'",
file, exc_info=True)
return
return desc[class_name]["arginfo"]
def open_experiment(self, expurl, arguments):
try:
dock = _ExperimentDock(self, expurl, arguments)
except:
logger.warning("Failed to create experiment dock for %s, "
"retrying with arguments reset", expurl,
exc_info=True)
dock = _ExperimentDock(self, expurl, {})
asyncio.ensure_future(dock._recompute_arguments())
dock.setAttribute(QtCore.Qt.WA_DeleteOnClose)
self.addSubWindow(dock)
dock.show()
dock.sigClosed.connect(partial(self.on_dock_closed, dock))
self.open_experiments.append(dock)
return dock
def set_argument_value(self, expurl, name, value):
logger.warning("Unable to set argument '%s', dropping change. "
"'set_argument_value' not supported in browser.", name)
def on_dock_closed(self, dock):
self.open_experiments.remove(dock)