import asyncio import logging import os from functools import partial from collections import OrderedDict from PyQt5 import QtCore, QtGui, QtWidgets import h5py 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 argty_to_entry from artiq.protocols import pyon from artiq.master.worker import Worker, log_worker_exception logger = logging.getLogger(__name__) class _WheelFilter(QtCore.QObject): def eventFilter(self, obj, event): if (event.type() == QtCore.QEvent.Wheel and event.modifiers() != QtCore.Qt.NoModifier): event.ignore() return True return False class _ArgumentEditor(QtWidgets.QTreeWidget): def __init__(self, dock): QtWidgets.QTreeWidget.__init__(self) self.setColumnCount(3) self.header().setStretchLastSection(False) try: set_resize_mode = self.header().setSectionResizeMode except AttributeError: set_resize_mode = self.header().setResizeMode set_resize_mode(0, QtWidgets.QHeaderView.ResizeToContents) set_resize_mode(1, QtWidgets.QHeaderView.Stretch) set_resize_mode(2, QtWidgets.QHeaderView.ResizeToContents) self.header().setVisible(False) self.setSelectionMode(self.NoSelection) self.setHorizontalScrollMode(self.ScrollPerPixel) self.setVerticalScrollMode(self.ScrollPerPixel) self.viewport().installEventFilter(_WheelFilter(self.viewport())) self._groups = dict() self._arg_to_entry_widgetitem = dict() self._dock = dock if not self._dock.arguments: self.addTopLevelItem(QtWidgets.QTreeWidgetItem(["No arguments"])) for name, argument in self._dock.arguments.items(): try: entry = argty_to_entry[argument["desc"]["ty"]](argument) except: print(name, argument) widget_item = QtWidgets.QTreeWidgetItem([name]) self._arg_to_entry_widgetitem[name] = entry, widget_item if argument["group"] is None: self.addTopLevelItem(widget_item) else: self._get_group(argument["group"]).addChild(widget_item) self.setItemWidget(widget_item, 1, entry) recompute_argument = QtWidgets.QToolButton() recompute_argument.setToolTip("Re-run the experiment's build " "method and take the default value") recompute_argument.setIcon( QtWidgets.QApplication.style().standardIcon( QtWidgets.QStyle.SP_BrowserReload)) recompute_argument.clicked.connect( partial(self._recompute_argument_clicked, name)) fix_layout = LayoutWidget() fix_layout.addWidget(recompute_argument) self.setItemWidget(widget_item, 2, fix_layout) widget_item = QtWidgets.QTreeWidgetItem() self.addTopLevelItem(widget_item) 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) buttons = LayoutWidget() buttons.addWidget(recompute_arguments, 1, 1) for i, s in enumerate((1, 0, 0, 1)): buttons.layout.setColumnStretch(i, s) self.setItemWidget(widget_item, 1, buttons) def _get_group(self, name): if name in self._groups: return self._groups[name] group = QtWidgets.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 _recompute_arguments_clicked(self): asyncio.ensure_future(self._dock._recompute_arguments()) def _recompute_argument_clicked(self, name): asyncio.ensure_future(self._recompute_argument(name)) 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 = argty_to_entry[procdesc["ty"]].default_state(procdesc) argument["desc"] = procdesc argument["state"] = state old_entry, widget_item = self._arg_to_entry_widgetitem[name] old_entry.deleteLater() entry = argty_to_entry[procdesc["ty"]](argument) self._arg_to_entry_widgetitem[name] = entry, widget_item self.setItemWidget(widget_item, 1, entry) def save_state(self): expanded = [] for k, v in self._groups.items(): if v.isExpanded(): expanded.append(k) return {"expanded": expanded} def restore_state(self, state): for e in state["expanded"]: try: self._groups[e].setExpanded(True) except KeyError: pass log_levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] class _ExperimentDock(QtWidgets.QMdiSubWindow): sigClosed = QtCore.pyqtSignal() def __init__(self, area, expurl, arguments): QtWidgets.QMdiSubWindow.__init__(self) 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 load = QtWidgets.QPushButton("Set arguments") load.setToolTip("Set arguments from currently selected HDF5 " "file (Ctrl+Space)") load.setIcon(QtWidgets.QApplication.style().standardIcon( QtWidgets.QStyle.SP_DialogApplyButton)) load.setShortcut("CTRL+SPACE") load.clicked.connect(self._load_clicked) self.layout.addWidget(load, 1, 4) 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": logger.debug("Loading HDF5 arguments from %s", uri.path()) asyncio.ensure_future(self.load_hdf5_task(uri.path())) 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) self.argeditor.deleteLater() self.argeditor = _ArgumentEditor(self) self.layout.addWidget(self.argeditor, 0, 0, 1, 5) async def load_hdf5_task(self, filename): 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 _load_clicked(self): if self._area.dataset is None: return asyncio.ensure_future(self.load_hdf5_task(self._area.dataset)) 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: argty_to_entry[argument["desc"]["ty"]].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, datasets_sub): self.datasets_sub = datasets_sub datasets_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.datasets_sub.update(mod) class ExperimentsArea(QtWidgets.QMdiArea): def __init__(self, root, datasets_sub): QtWidgets.QMdiArea.__init__(self) self.pixmap = QtGui.QPixmap(os.path.join( artiq_dir, "gui", "logo20.svg")) self.current_dir = root self.dataset = None self.open_experiments = [] self._ddb = LocalDatasetDB(datasets_sub) self.worker_handlers = { "get_device_db": lambda: None, "get_device": lambda k: None, "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) in arginfo.items(): state = argty_to_entry[procdesc["ty"]].default_state(procdesc) arguments[name] = { "desc": procdesc, "group": group, "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()) self.addSubWindow(dock) dock.show() dock.sigClosed.connect(partial(self.on_dock_closed, dock)) self.open_experiments.append(dock) return dock def on_dock_closed(self, dock): self.open_experiments.remove(dock)