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)