import asyncio import logging from quamash import QtGui, QtCore from pyqtgraph import dockarea from pyqtgraph import LayoutWidget from artiq.protocols.sync_struct import Subscriber from artiq.protocols import pyon from artiq.gui.tools import DictSyncModel from artiq.gui.scan import ScanController from artiq.gui.shortcuts import ShortcutManager class _ExplistModel(DictSyncModel): def __init__(self, explorer, parent, init): self.explorer = explorer DictSyncModel.__init__(self, ["Experiment"], parent, init) 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 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, schedule_ctl, repository_ctl): dockarea.Dock.__init__(self, "Explorer", size=(1500, 500)) self.main_window = main_window self.status_bar = status_bar 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.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 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 sub_connect(self, host, port): self.explist_subscriber = Subscriber("explist", self.init_explist_model) await self.explist_subscriber.connect(host, port) async def sub_close(self): await self.explist_subscriber.close() def init_explist_model(self, init): self.explist_model = _ExplistModel(self, self.el, init) self.el.setModel(self.explist_model) return self.explist_model 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] # attribute get_current_schedule must be set externally after # instance creation current_schedule = self.get_current_schedule() 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)