artiq/artiq/gui/explorer.py

428 lines
15 KiB
Python

import asyncio
import logging
from quamash import QtGui, QtCore
from pyqtgraph import dockarea
from pyqtgraph import LayoutWidget
from artiq.protocols import pyon
from artiq.gui.models import DictSyncModel
from artiq.gui.scan import ScanController
from artiq.gui.shortcuts import ShortcutManager
class Model(DictSyncModel):
def __init__(self, init):
DictSyncModel.__init__(self,
["Experiment"],
init)
self.explorer = None
def sort_key(self, k, v):
return k
def convert(self, k, v, column):
return k
def __setitem__(self, k, v):
DictSyncModel.__setitem__(self, k, v)
if self.explorer is not None:
if k == self.explorer.selected_key:
self.explorer.update_selection(k, k)
class _FreeValueEntry(QtGui.QLineEdit):
def __init__(self, procdesc):
QtGui.QLineEdit.__init__(self)
if "default" in procdesc:
self.set_argument_value(procdesc["default"])
def get_argument_value(self):
return pyon.decode(self.text())
def set_argument_value(self, value):
self.setText(pyon.encode(value))
class _BooleanEntry(QtGui.QCheckBox):
def __init__(self, procdesc):
QtGui.QCheckBox.__init__(self)
if "default" in procdesc:
self.set_argument_value(procdesc["default"])
def get_argument_value(self):
return self.isChecked()
def set_argument_value(self, value):
self.setChecked(value)
class _EnumerationEntry(QtGui.QComboBox):
def __init__(self, procdesc):
QtGui.QComboBox.__init__(self)
self.choices = procdesc["choices"]
self.addItems(self.choices)
if "default" in procdesc:
self.set_argument_value(procdesc["default"])
def get_argument_value(self):
return self.choices[self.currentIndex()]
def set_argument_value(self, value):
idx = self.choices.index(value)
self.setCurrentIndex(idx)
class _NumberEntry(QtGui.QDoubleSpinBox):
def __init__(self, procdesc):
QtGui.QDoubleSpinBox.__init__(self)
self.scale = procdesc["scale"]
self.setDecimals(procdesc["ndecimals"])
self.setSingleStep(procdesc["step"]/self.scale)
if procdesc["min"] is not None:
self.setMinimum(procdesc["min"]/self.scale)
else:
self.setMinimum(float("-inf"))
if procdesc["max"] is not None:
self.setMaximum(procdesc["max"]/self.scale)
else:
self.setMaximum(float("inf"))
if procdesc["unit"]:
self.setSuffix(" " + procdesc["unit"])
if "default" in procdesc:
self.set_argument_value(procdesc["default"])
def get_argument_value(self):
return self.value()*self.scale
def set_argument_value(self, value):
self.setValue(value/self.scale)
class _StringEntry(QtGui.QLineEdit):
def __init__(self, procdesc):
QtGui.QLineEdit.__init__(self)
if "default" in procdesc:
self.set_argument_value(procdesc["default"])
def get_argument_value(self):
return self.text()
def set_argument_value(self, value):
self.setText(value)
_procty_to_entry = {
"FreeValue": _FreeValueEntry,
"BooleanValue": _BooleanEntry,
"EnumerationValue": _EnumerationEntry,
"NumberValue": _NumberEntry,
"StringValue": _StringEntry,
"Scannable": ScanController
}
class _ArgumentEditor(QtGui.QTreeWidget):
def __init__(self, main_window):
QtGui.QTreeWidget.__init__(self)
self.setColumnCount(2)
self.header().setResizeMode(QtGui.QHeaderView.ResizeToContents)
self.header().setVisible(False)
self.setSelectionMode(QtGui.QAbstractItemView.NoSelection)
self.main_window = main_window
self._groups = dict()
self.set_arguments([])
def clear(self):
QtGui.QTreeWidget.clear(self)
self._groups.clear()
def _get_group(self, name):
if name in self._groups:
return self._groups[name]
group = QtGui.QTreeWidgetItem([name, ""])
for c in 0, 1:
group.setBackground(c, QtGui.QBrush(QtGui.QColor(100, 100, 100)))
group.setForeground(c, QtGui.QBrush(QtGui.QColor(220, 220, 255)))
font = group.font(c)
font.setBold(True)
group.setFont(c, font)
self.addTopLevelItem(group)
self._groups[name] = group
return group
def set_arguments(self, arguments):
self.clear()
if not arguments:
self.addTopLevelItem(QtGui.QTreeWidgetItem(["No arguments", ""]))
self._args_to_entries = dict()
for n, (name, (procdesc, group)) in enumerate(arguments):
entry = _procty_to_entry[procdesc["ty"]](procdesc)
self._args_to_entries[name] = entry
widget_item = QtGui.QTreeWidgetItem([name, ""])
if group is None:
self.addTopLevelItem(widget_item)
else:
self._get_group(group).addChild(widget_item)
self.setItemWidget(widget_item, 1, entry)
def get_argument_values(self, show_error_message):
r = dict()
for arg, entry in self._args_to_entries.items():
try:
r[arg] = entry.get_argument_value()
except Exception as e:
if show_error_message:
msgbox = QtGui.QMessageBox(self.main_window)
msgbox.setWindowTitle("Error")
msgbox.setText("Failed to obtain value for argument '{}':\n{}"
.format(arg, str(e)))
msgbox.setStandardButtons(QtGui.QMessageBox.Ok)
msgbox.show()
return None
return r
def set_argument_values(self, arguments, ignore_errors):
for arg, value in arguments.items():
try:
entry = self._args_to_entries[arg]
entry.set_argument_value(value)
except:
if not ignore_errors:
raise
def save_state(self):
expanded = []
for k, v in self._groups.items():
if v.isExpanded():
expanded.append(k)
argument_values = self.get_argument_values(False)
return {
"expanded": expanded,
"argument_values": argument_values
}
def restore_state(self, state):
self.set_argument_values(state["argument_values"], True)
for e in state["expanded"]:
try:
self._groups[e].setExpanded(True)
except KeyError:
pass
class ExplorerDock(dockarea.Dock):
def __init__(self, main_window, status_bar,
explist_sub, schedule_sub,
schedule_ctl, repository_ctl):
dockarea.Dock.__init__(self, "Explorer", size=(1500, 500))
self.main_window = main_window
self.status_bar = status_bar
self.schedule_sub = schedule_sub
self.schedule_ctl = schedule_ctl
self.splitter = QtGui.QSplitter(QtCore.Qt.Horizontal)
self.addWidget(self.splitter)
grid = LayoutWidget()
self.splitter.addWidget(grid)
self.el = QtGui.QListView()
self.el.selectionChanged = self._selection_changed
self.selected_key = None
grid.addWidget(self.el, 0, 0, colspan=4)
self.datetime = QtGui.QDateTimeEdit()
self.datetime.setDisplayFormat("MMM d yyyy hh:mm:ss")
self.datetime.setDate(QtCore.QDate.currentDate())
self.datetime.dateTimeChanged.connect(self.enable_duedate)
self.datetime_en = QtGui.QCheckBox("Due date:")
grid.addWidget(self.datetime_en, 1, 0, colspan=2)
grid.addWidget(self.datetime, 1, 2, colspan=2)
self.pipeline = QtGui.QLineEdit()
self.pipeline.setText("main")
grid.addWidget(QtGui.QLabel("Pipeline:"), 2, 0, colspan=2)
grid.addWidget(self.pipeline, 2, 2, colspan=2)
self.priority = QtGui.QSpinBox()
self.priority.setRange(-99, 99)
grid.addWidget(QtGui.QLabel("Priority:"), 3, 0)
grid.addWidget(self.priority, 3, 1)
self.flush = QtGui.QCheckBox("Flush")
self.flush.setToolTip("Flush the pipeline before starting the experiment")
grid.addWidget(self.flush, 3, 2)
self.log_level = QtGui.QComboBox()
self.log_level.addItems(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"])
self.log_level.setCurrentIndex(1)
self.log_level.setToolTip("Minimum level for log entry production")
grid.addWidget(self.log_level, 3, 3)
submit = QtGui.QPushButton("Submit")
submit.setToolTip("Schedule the selected experiment (Ctrl+Return)")
grid.addWidget(submit, 4, 0, colspan=4)
submit.clicked.connect(self.submit_clicked)
self.argeditor = _ArgumentEditor(self.main_window)
self.splitter.addWidget(self.argeditor)
self.splitter.setSizes([grid.minimumSizeHint().width(), 1000])
self.argeditor_states = dict()
self.explist_model = Model(dict())
explist_sub.add_setmodel_callback(self.set_model)
self.shortcuts = ShortcutManager(self.main_window, self)
self.el.setContextMenuPolicy(QtCore.Qt.ActionsContextMenu)
submit_action = QtGui.QAction("Submit", self.el)
submit_action.triggered.connect(self.submit_clicked)
submit_action.setShortcut("CTRL+RETURN")
self.el.addAction(submit_action)
reqterm_action = QtGui.QAction("Request termination of instances", self.el)
reqterm_action.triggered.connect(self.request_inst_term)
reqterm_action.setShortcut("CTRL+BACKSPACE")
self.el.addAction(reqterm_action)
sep = QtGui.QAction(self.el)
sep.setSeparator(True)
self.el.addAction(sep)
edit_shortcuts_action = QtGui.QAction("Edit shortcuts", self.el)
edit_shortcuts_action.triggered.connect(self.edit_shortcuts)
self.el.addAction(edit_shortcuts_action)
scan_repository_action = QtGui.QAction("(Re)scan repository HEAD",
self.el)
def scan_repository():
asyncio.ensure_future(repository_ctl.scan_async())
self.status_bar.showMessage("Requested repository scan")
scan_repository_action.triggered.connect(scan_repository)
self.el.addAction(scan_repository_action)
def set_model(self, model):
model.explorer = self
self.explist_model = model
self.el.setModel(model)
def update_selection(self, selected, deselected):
if deselected:
self.argeditor_states[deselected] = self.argeditor.save_state()
if selected:
expinfo = self.explist_model.backing_store[selected]
self.argeditor.set_arguments(expinfo["arguments"])
if selected in self.argeditor_states:
self.argeditor.restore_state(self.argeditor_states[selected])
self.splitter.insertWidget(1, self.argeditor)
self.selected_key = selected
def _sel_to_key(self, selection):
selection = selection.indexes()
if selection:
row = selection[0].row()
return self.explist_model.row_to_key[row]
else:
return None
def _selection_changed(self, selected, deselected):
self.update_selection(self._sel_to_key(selected),
self._sel_to_key(deselected))
def save_state(self):
idx = self.el.selectedIndexes()
if idx:
row = idx[0].row()
key = self.explist_model.row_to_key[row]
self.argeditor_states[key] = self.argeditor.save_state()
return {
"argeditor": self.argeditor_states,
"shortcuts": self.shortcuts.save_state()
}
def restore_state(self, state):
try:
argeditor_states = state["argeditor"]
shortcuts_state = state["shortcuts"]
except KeyError:
return
self.argeditor_states = argeditor_states
self.shortcuts.restore_state(shortcuts_state)
def enable_duedate(self):
self.datetime_en.setChecked(True)
async def submit_task(self, pipeline_name, file, class_name, arguments,
priority, due_date, flush):
expid = {
"log_level": getattr(logging, self.log_level.currentText()),
"repo_rev": None,
"file": file,
"class_name": class_name,
"arguments": arguments,
}
rid = await self.schedule_ctl.submit(pipeline_name, expid,
priority, due_date, flush)
self.status_bar.showMessage("Submitted RID {}".format(rid))
def submit(self, pipeline, key, priority, due_date, flush):
# TODO: refactor explorer and cleanup.
# Argument editors should immediately modify the global state.
expinfo = self.explist_model.backing_store[key]
if key == self.selected_key:
arguments = self.argeditor.get_argument_values(True)
if arguments is None:
# There has been an error. Displaying the error message box
# was done by argeditor.
return
else:
try:
arguments = self.argeditor_states[key]["argument_values"]
except KeyError:
arguments = dict()
asyncio.ensure_future(self.submit_task(pipeline, expinfo["file"],
expinfo["class_name"],
arguments, priority, due_date,
flush))
def submit_clicked(self):
if self.selected_key is not None:
if self.datetime_en.isChecked():
due_date = self.datetime.dateTime().toMSecsSinceEpoch()/1000
else:
due_date = None
self.submit(self.pipeline.text(),
self.selected_key,
self.priority.value(),
due_date,
self.flush.isChecked())
async def request_term_multiple(self, rids):
for rid in rids:
try:
await self.schedule_ctl.request_termination(rid)
except:
pass
def request_inst_term(self):
if self.selected_key is not None:
expinfo = self.explist_model.backing_store[self.selected_key]
if self.schedule_sub.model is not None:
current_schedule = self.schedule_sub.model.backing_store
rids = []
for rid, desc in current_schedule.items():
expid = desc["expid"]
if ("repo_rev" in expid # only consider runs from repository
and expid["file"] == expinfo["file"]
and expid["class_name"] == expinfo["class_name"]):
rids.append(rid)
asyncio.ensure_future(self.request_term_multiple(rids))
def edit_shortcuts(self):
experiments = sorted(self.explist_model.backing_store.keys())
self.shortcuts.edit(experiments)