Beispiel #1
0
class SkywalkerGui(Display):
    """
    Display class to define all the logic for the skywalker alignment gui.
    Refers to widgets in the .ui file.

    Parameters
    ----------
    live : bool, optional
        Whether to launch application with live or simulated devices

    cfg : str, optional
        Configuration directory to use if not the default

    dark : bool, optional
        Choice to launch the application with a dark stylesheet

    parent : QWidget
        Parent Widget of application
    """
    def __init__(self, parent=None, live=False, cfg=None, dark=True):
        super().__init__(parent=parent)
        ui = self.ui

        #Change the stylesheet
        if dark:
            try:
                import qdarkstyle
            except ImportError:
                logger.error("Can not use dark theme, "
                             "qdarkstyle package not available")
            else:
                self.setStyleSheet(qdarkstyle.load_stylesheet_pyqt5())

        # Configure debug file after all the qt logs
        logging.basicConfig(level=logging.DEBUG,
                            format=('%(asctime)s '
                                    '%(name)-12s '
                                    '%(levelname)-8s '
                                    '%(message)s'),
                            datefmt='%m-%d %H:%M:%S',
                            filename='./skywalker_debug.log',
                            filemode='a')

        # Set self.sim, self.loader, self.nominal_config
        self.sim = not live
        self.config_folder = cfg
        self.init_config()

        # Load things
        self.config_cache = {}
        self.cache_config()

        # Load system and alignments into the combo box objects
        ui.image_title_combo.clear()
        ui.procedure_combo.clear()
        ui.procedure_combo.addItem('None')
        self.all_imager_names = [
            entry['imager'] for entry in self.loader.live_systems.values()
        ]
        for imager_name in self.all_imager_names:
            ui.image_title_combo.addItem(imager_name)
        for align in self.alignments.keys():
            ui.procedure_combo.addItem(align)

        # Pick out some initial parameters from system and alignment dicts
        first_system_key = list(self.alignments.values())[0][0][0]
        first_set = self.loader.get_subsystem(first_system_key)
        first_imager = first_set.get('imager', None)
        first_slit = first_set.get('slits', None)
        first_rotation = first_set.get('rotation', 0)

        # self.procedure and self.image_obj keep track of the gui state
        self.procedure = 'None'
        self.image_obj = first_imager

        # Initialize slit readback
        self.slit_group = ObjWidgetGroup([
            ui.slit_x_width, ui.slit_y_width, ui.slit_x_setpoint,
            ui.slit_y_setpoint, ui.slit_circle
        ], [
            'xwidth.readback', 'ywidth.readback', 'xwidth.setpoint',
            'ywidth.setpoint', 'xwidth.done'
        ],
                                         first_slit,
                                         label=ui.readback_slits_title)

        # Initialize mirror control
        self.mirror_groups = []
        mirror_labels = self.get_widget_set('mirror_name')
        mirror_rbvs = self.get_widget_set('mirror_readback')
        mirror_vals = self.get_widget_set('mirror_setpos')
        mirror_circles = self.get_widget_set('mirror_circle')
        mirror_nominals = self.get_widget_set('move_nominal')
        for label, rbv, val, circle, nom, mirror in zip(
                mirror_labels, mirror_rbvs, mirror_vals, mirror_circles,
                mirror_nominals, self.mirrors_padded()):
            mirror_group = ObjWidgetGroup([rbv, val, circle, nom], [
                'pitch.user_readback', 'pitch.user_setpoint',
                'pitch.motor_done_move'
            ],
                                          mirror,
                                          label=label)
            if mirror is None:
                mirror_group.hide()
            self.mirror_groups.append(mirror_group)

        # Initialize the goal entry fields
        self.goals_groups = []
        goal_labels = self.get_widget_set('goal_name')
        goal_edits = self.get_widget_set('goal_value')
        slit_checks = self.get_widget_set('slit_check')
        for label, edit, check, img, slit in zip(goal_labels, goal_edits,
                                                 slit_checks,
                                                 self.imagers_padded(),
                                                 self.slits_padded()):
            if img is None:
                name = None
            else:
                name = img.name
            validator = QDoubleValidator(0, 5000, 3)
            goal_group = ValueWidgetGroup(edit,
                                          label,
                                          checkbox=check,
                                          name=name,
                                          cache=self.config_cache,
                                          validator=validator)
            if img is None:
                goal_group.hide()
            elif slit is None:
                goal_group.checkbox.setEnabled(False)
            self.goals_groups.append(goal_group)

        # Initialize image and centroids. Needs goals defined first.
        self.image_group = ImgObjWidget(ui.image, first_imager,
                                        ui.beam_x_value, ui.beam_y_value,
                                        ui.beam_x_delta, ui.beam_y_delta,
                                        ui.image_state, ui.image_state_select,
                                        ui.readback_imager_title, self,
                                        first_rotation)
        ui.image.setColorMapToPreset('jet')

        # Initialize the settings window.
        first_step = Setting('first_step', 6.0)
        tolerance = Setting('tolerance', 5.0)
        averages = Setting('averages', 100)
        timeout = Setting('timeout', 600.0)
        tol_scaling = Setting('tol_scaling', 8.0)
        min_beam = Setting('min_beam', 1.0, required=False)
        min_rate = Setting('min_rate', 1.0, required=False)
        slit_width = Setting('slit_width', 0.2)
        samples = Setting('samples', 100)
        close_fee_att = Setting('close_fee_att', True)
        self.settings = SettingsGroup(
            parent=self,
            collumns=[['alignment'], ['slits', 'suspenders', 'setup']],
            alignment=[first_step, tolerance, averages, timeout, tol_scaling],
            suspenders=[min_beam, min_rate],
            slits=[slit_width, samples],
            setup=[close_fee_att])
        self.settings_cache = {}
        self.load_settings()
        self.restore_settings()
        self.cache_settings()  # Required in case nothing is loaded

        # Create the RunEngine that will be used in the alignments.
        # This gives us the ability to pause, etc.
        self.RE = RunEngine({})
        install_qt_kicker()

        # Some hax to keep the state string updated
        # There is probably a better way to do this
        # This might break on some package update
        self.RE.state  # Yes this matters
        old_set = RunEngine.state._memory[self.RE].set_

        def new_set(state):  # NOQA
            old_set(state)
            txt = " Status: " + state.capitalize()
            self.ui.status_label.setText(txt)

        RunEngine.state._memory[self.RE].set_ = new_set

        # Connect relevant signals and slots
        procedure_changed = ui.procedure_combo.currentIndexChanged[str]
        procedure_changed.connect(self.on_procedure_combo_changed)

        imager_changed = ui.image_title_combo.currentIndexChanged[str]
        imager_changed.connect(self.on_image_combo_changed)

        for goal_value in self.get_widget_set('goal_value'):
            goal_changed = goal_value.editingFinished
            goal_changed.connect(self.on_goal_changed)

        start_pressed = ui.start_button.clicked
        start_pressed.connect(self.on_start_button)

        pause_pressed = ui.pause_button.clicked
        pause_pressed.connect(self.on_pause_button)

        abort_pressed = ui.abort_button.clicked
        abort_pressed.connect(self.on_abort_button)

        slits_pressed = ui.slit_run_button.clicked
        slits_pressed.connect(self.on_slits_button)

        save_mirrors_pressed = ui.save_mirrors_button.clicked
        save_mirrors_pressed.connect(self.on_save_mirrors_button)

        save_goals_pressed = ui.save_goals_button.clicked
        save_goals_pressed.connect(self.on_save_goals_button)

        settings_pressed = ui.settings_button.clicked
        settings_pressed.connect(self.on_settings_button)

        for i, nominal_button in enumerate(mirror_nominals):
            nominal_pressed = nominal_button.clicked
            nominal_pressed.connect(partial(self.on_move_nominal_button, i))

        self.cam_lock = RLock()

        # Store some info about our screen size.
        QApp = QCoreApplication.instance()
        desktop = QApp.desktop()
        geometry = desktop.screenGeometry()
        self.screen_size = (geometry.width(), geometry.height())
        window_qsize = self.window().size()
        self.preferred_size = (window_qsize.width(), window_qsize.height())

        # Setup the post-init hook
        post_init = PostInit(self)
        self.installEventFilter(post_init)
        post_init.post_init.connect(self.on_post_init)

        # Setup the on-screen logger
        console = self.setup_gui_logger()

        # Stop the run if we get closed
        close_dict = dict(RE=self.RE, console=console)
        self.destroyed.connect(partial(SkywalkerGui.on_close, close_dict))

        # Put out the initialization message.
        init_base = 'Skywalker GUI initialized in '
        if self.sim:
            init_str = init_base + 'sim mode.'
        else:
            init_str = init_base + 'live mode.'
        logger.info(init_str)

    def init_config(self):
        if self.config_folder is None:
            this_dir = path.dirname(__file__)
            config_rel = path.join(this_dir, '..', 'config')
            self.config_folder = path.abspath(config_rel)
        self.nominal_config = self.get_cfg_path('nominal')
        self.happi_config = self.get_cfg_path('metadata')
        self.system_config = self.get_cfg_path('system')
        self.alignment_config = self.get_cfg_path('alignments')

        # Load files needed during __init__
        self.load_system()
        self.load_alignments()

    def get_cfg_path(self, name):
        if self.sim:
            name = 'sim_' + name
        return path.join(self.config_folder, name + '.json')

    def load_system(self):
        if self.sim:
            self.loader = SimConfigReader()
        else:
            self.loader = ConfigReader(self.happi_config, self.system_config)

    def load_alignments(self):
        if self.sim:
            self.alignments = sim_alignments
        else:
            with open(self.alignment_config, 'r') as f:
                d = json.load(f)
                self.alignments = d

    @pyqtSlot()
    def on_post_init(self):
        x = min(self.preferred_size[0], self.screen_size[0])
        y = min(self.preferred_size[1], self.screen_size[1])
        self.window().resize(x, y)

    # Close handler needs to be a static class method because it is run after
    # the object instance is already completely gone
    @staticmethod
    def on_close(close_dict):
        RE = close_dict['RE']
        console = close_dict['console']
        console.close()
        if RE.state != 'idle':
            RE.abort()

    def setup_gui_logger(self):
        """
        Initializes the text stream at the bottom of the gui. This text stream
        is actually just the log messages from Python!
        """
        console = GuiHandler(self.ui.log_text)
        console.setLevel(logging.INFO)
        formatter = logging.Formatter(fmt='%(asctime)s %(message)s',
                                      datefmt='%m-%d %H:%M:%S')
        console.setFormatter(formatter)
        logging.getLogger('').addHandler(console)
        return console

    @pyqtSlot(str)
    def on_image_combo_changed(self, imager_name):
        """
        Slot for the combo box above the image feed. This swaps out the imager,
        centroid, and slit readbacks.

        Parameters
        ----------
        imager_name: str
            name of the imager to activate
        """
        try:
            logger.info('Selecting imager %s', imager_name)
            systems = self.loader.get_systems_with(imager_name)
            if len(systems) == 0:
                logger.error('Invalid imager name.')
                return
            # Assume that imagers have exactly one slit and one rotation
            # Therefore, we can pick an arbitrary system entry that includes
            # the imager
            objs = self.loader.get_subsystem(systems[0])
            # This may have entries or may be missing entries if there was a
            # problem.
            try:
                image_obj = objs['imager']
                rotation = objs.get('rotation', 0)
                self.image_obj = image_obj
                self.image_group.change_obj(image_obj, rotation=rotation)
            except KeyError:
                logger.error('Failed to connect to imager')
            # Slits wasn't a mandatory field.
            slits_obj = objs.get('slits')
            if slits_obj is not None:
                self.slit_group.change_obj(slits_obj)
        except:
            logger.exception('Error on selecting imager')

    @pyqtSlot(str)
    def on_procedure_combo_changed(self, procedure_name):
        """
        Slot for the main procedure combo box. This swaps out the mirror and
        goals sections to match the chosen procedure, and determines what
        happens when we press go.

        Parameters
        ----------
        procedure_name: str
            name of the procedure to activate
        """
        try:
            logger.info('Selecting procedure %s', procedure_name)
            self.procedure = procedure_name
            if procedure_name == 'None':
                return
            else:
                self.load_active_system()
            for obj, widgets in zip(self.mirrors_padded(), self.mirror_groups):
                if obj is None:
                    widgets.hide()
                    widgets.change_obj(None)
                else:
                    widgets.change_obj(obj)
                    widgets.show()
            for obj, widgets in zip(self.imagers_padded(), self.goals_groups):
                widgets.save_value()
                widgets.clear()
            for obj, slit, widgets in zip(self.imagers_padded(),
                                          self.slits_padded(),
                                          self.goals_groups):
                if obj is None:
                    widgets.hide()
                else:
                    widgets.setup(name=obj.name)
                    if slit is None:
                        widgets.checkbox.setEnabled(False)
                    else:
                        widgets.checkbox.setEnabled(True)
                    widgets.show()
        except:
            logger.exception('Error on selecting procedure')

    @pyqtSlot()
    def on_goal_changed(self):
        """
        Slot for when the user picks a new goal. Updates the goal delta so it
        reflects the new chosen value.
        """
        try:
            self.image_group.update_deltas()
        except:
            logger.exception('Error on changing goal')

    @pyqtSlot()
    def on_start_button(self):
        """
        Slot for the start button. This begins from an idle state or resumes
        from a paused state.
        """
        try:
            if self.RE.state == 'idle':
                # Check for valid procedure
                if self.procedure == 'None':
                    logger.info("Please select a procedure.")
                    return

                # Check for valid goals
                active_size = len(self.active_system())
                raw_goals = []
                for i, goal in enumerate(self.goals()):
                    if i >= active_size:
                        break
                    elif goal is None:
                        msg = 'Please fill all goal fields before alignment.'
                        logger.info(msg)
                        return
                    raw_goals.append(goal)

                logger.info("Starting %s procedure with goals %s",
                            self.procedure, raw_goals)
                self.install_pick_cam()
                self.auto_switch_cam = True
                alignment = self.alignments[self.procedure]
                for key_set in alignment:
                    yags = [self.loader[key]['imager'] for key in key_set]
                    mots = [self.loader[key]['mirror'] for key in key_set]
                    rots = [
                        self.loader[key].get('rotation') for key in key_set
                    ]

                    # Make sure nominal positions are correct
                    for mot in mots:
                        try:
                            mot.nominal_position = self.config_cache[mot.name]
                        except KeyError:
                            pass

                    mot_rbv = 'pitch'
                    # We need to select det_rbv and interpret goals based on
                    # the camera rotation, converting things to the unrotated
                    # coordinates.
                    det_rbv = []
                    goals = []
                    for rot, yag, goal in zip(rots, yags, raw_goals):
                        rot_info = ad_stats_x_axis_rot(yag, rot)
                        det_rbv.append(rot_info['key'])
                        modifier = rot_info['mod_x']
                        if modifier is not None:
                            goal = modifier - goal
                        goals.append(goal)
                    first_steps = self.settings_cache['first_step']
                    tolerances = self.settings_cache['tolerance']
                    average = self.settings_cache['averages']
                    timeout = self.settings_cache['timeout']
                    tol_scaling = self.settings_cache['tol_scaling']

                    extra_stage = []
                    close_fee_att = self.settings_cache['close_fee_att']
                    if close_fee_att and not self.sim:
                        extra_stage.append(self.fee_att())

                    # Temporary fix: undo skywalker's goal mangling.
                    # TODO remove goal mangling from skywalker.
                    goals = [480 - g for g in goals]
                    plan = skywalker(yags,
                                     mots,
                                     det_rbv,
                                     mot_rbv,
                                     goals,
                                     first_steps=first_steps,
                                     tolerances=tolerances,
                                     averages=average,
                                     timeout=timeout,
                                     sim=self.sim,
                                     use_filters=not self.sim,
                                     tol_scaling=tol_scaling,
                                     extra_stage=extra_stage)
                    self.initialize_RE()
                    self.RE(plan)
            elif self.RE.state == 'paused':
                logger.info("Resuming procedure.")
                self.install_pick_cam()
                self.auto_switch_cam = True
                self.RE.resume()
        except:
            logger.exception('Error in running procedure')
        finally:
            self.auto_switch_cam = False

    @pyqtSlot()
    def on_pause_button(self):
        """
        Slot for the pause button. This brings us from the running state to the
        paused state.
        """
        self.auto_switch_cam = False
        if self.RE.state == 'running':
            logger.info("Pausing procedure.")
            try:
                self.RE.request_pause()
            except:
                logger.exception("Error on pause.")

    @pyqtSlot()
    def on_abort_button(self):
        """
        Slot for the abort button. This brings us from any state to the idle
        state.
        """
        self.auto_switch_cam = False
        if self.RE.state != 'idle':
            logger.info("Aborting procedure.")
            try:
                self.RE.abort()
            except:
                logger.exception("Error on abort.")

    @pyqtSlot()
    def on_slits_button(self):
        """
        Slot for the slits procedure. This checks the slit fiducialization.
        """
        try:
            logger.info('Starting slit check process.')
            image_to_check = []
            slits_to_check = []

            # First, check the slit checkboxes.
            for img_obj, slit_obj, goal_group in zip(self.imagers_padded(),
                                                     self.slits_padded(),
                                                     self.goals_groups):
                if slit_obj is not None and goal_group.is_checked:
                    image_to_check.append(img_obj)
                    slits_to_check.append(slit_obj)
            if not slits_to_check:
                logger.info('No valid slits selected!')
                return
            logger.info('Checking the following slits: %s',
                        [slit.name for slit in slits_to_check])

            self.install_pick_cam()
            self.auto_switch_cam = True

            slit_width = self.settings_cache['slit_width']
            samples = self.settings_cache['samples']

            def plan(img,
                     slit,
                     rot,
                     output_obj,
                     slit_width=slit_width,
                     samples=samples):
                rot_info = ad_stats_x_axis_rot(img, rot)
                det_rbv = rot_info['key']
                fidu = slit_scan_fiducialize(slit,
                                             img,
                                             centroid=det_rbv,
                                             x_width=slit_width,
                                             samples=samples)
                output = yield from fidu
                modifier = rot_info['mod_x']
                if modifier is not None:
                    output = modifier - output
                output_obj[img.name] = output

            self.initialize_RE()
            results = {}
            for img, slit in zip(image_to_check, slits_to_check):
                systems = self.loader.get_systems_with(img.name)
                objs = self.loader.get_subsystem(systems[0])
                rotation = objs.get('rotation', 0)
                this_plan = plan(img, slit, rotation, results)
                wrapped = run_wrapper(this_plan)
                wrapped = stage_wrapper(wrapped, [img, slit])
                self.RE(wrapped)

            logger.info('Slit scan found the following goals: %s', results)
            if self.ui.slit_fill_check.isChecked():
                logger.info('Filling goal fields automatically.')
                for img, fld in zip(self.imagers_padded(), self.goals_groups):
                    if img is not None:
                        try:
                            fld.value = round(results[img.name], 1)
                        except KeyError:
                            pass
        except:
            logger.exception('Error on slits button')
        finally:
            self.auto_switch_cam = False

    @pyqtSlot()
    def on_save_mirrors_button(self):
        try:
            if self.nominal_config is None:
                logger.info('No config file chosen.')
            else:
                logger.info('Saving mirror positions.')
                self.save_active_mirrors()
                self.cache_config()
        except:
            logger.exception('Error on saving mirrors')

    @pyqtSlot()
    def on_save_goals_button(self):
        try:
            logger.info('Saving goals.')
            self.save_active_goals()
            self.cache_config()
        except:
            logger.exception('Error on saving goals')

    @pyqtSlot()
    def on_settings_button(self):
        try:
            pos = self.ui.mapToGlobal(self.settings_button.pos())
            dialog_return = self.settings.dialog_at(pos)
            if dialog_return == QDialog.Accepted:
                self.cache_settings()
                self.save_settings()
                logger.info('Settings saved.')
            elif dialog_return == QDialog.Rejected:
                self.restore_settings()
                logger.info('Changes to settings cancelled.')
        except:
            logger.exception('Error on opening settings')

    @pyqtSlot(int)
    def on_move_nominal_button(self, index):
        try:
            nominal_positions = self.read_config() or {}
            try:
                mirror = self.mirrors()[index]
            except IndexError:
                logger.exception('Mirror index out of range')
                return
            try:
                pos = nominal_positions[mirror.name]
            except KeyError:
                logger.info('No mirror position saved')
                return
            logger.info('Moving %s to %s', mirror.name, pos)
            mirror.move(pos)
        except Exception:
            logger.exception('Misc error on move nominal button')

    def initialize_RE(self):
        """
        Set up the RunEngine for the current cached settings.
        """
        self.RE.clear_suspenders()
        min_beam = self.settings_cache['min_beam']
        min_rate = self.settings_cache['min_rate']
        if min_beam is not None:
            self.RE.install_suspender(
                BeamEnergySuspendFloor(min_beam, sleep=5, averages=100))
        if min_rate is not None:
            self.RE.install_suspender(BeamRateSuspendFloor(min_rate, sleep=5))

    def fee_att(self):
        try:
            att = self._fee_att
        except AttributeError:
            att = FeeAtt()
            self._fee_att = att
        return att

    def cache_settings(self):
        """
        Pull settings from the settings object to the local cache.
        """
        self.settings_cache = self.settings.values

    def restore_settings(self):
        """
        Push settings from the local cache into the settings object.
        """
        self.settings.values = self.settings_cache

    def save_settings(self):
        """
        Write settings from the local cache to disk.
        """
        pass

    def load_settings(self):
        """
        Load settings from disk to the local cache.
        """
        pass

    def install_pick_cam(self):
        """
        For every camera that we've successfully loaded, subscribe the pick_cam
        method if we haven't done so already.
        """
        try:
            installed = self.installed
        except AttributeError:
            installed = set()
            self.installed = installed
        for system in self.loader.cache.values():
            imager = system['imager']
            if imager not in installed:
                imager.subscribe(self.pick_cam,
                                 event_type=imager.SUB_STATE,
                                 run=False)
                installed.add(imager)

    def pick_cam(self, *args, **kwargs):
        """
        Callback to switch the active imager as the procedures progress.
        """
        if self.auto_switch_cam:
            with self.cam_lock:
                chosen_imager = None
                for img in self.imagers():
                    pos = img.position
                    if pos == "Unknown":
                        return
                    elif pos == "IN":
                        chosen_imager = img
                        break
                combo = self.ui.image_title_combo
                if chosen_imager is not None:
                    name = chosen_imager.name
                    if name != combo.currentText():
                        logger.info('Automatically switching cam to %s', name)
                        index = self.all_imager_names.index(name)
                        combo.setCurrentIndex(index)

    def read_config(self):
        if self.nominal_config is not None:
            try:
                with open(self.nominal_config, 'r') as f:
                    d = json.load(f)
            except:
                return None
            return d
        return None

    def save_config(self, d):
        if self.nominal_config is not None:
            with open(self.nominal_config, 'w') as f:
                json.dump(d, f)

    def cache_config(self):
        d = self.read_config()
        if d is not None:
            self.config_cache.update(d)

    def save_goal(self, goal_group):
        if goal_group.value is None:
            logger.info('No value to save for this goal.')
            return
        d = self.read_config() or {}
        d[goal_group.text()] = goal_group.value
        self.save_config(d)

    def save_active_goals(self):
        text = []
        values = []
        for i, goal_group in enumerate(self.goals_groups):
            if i >= len(self.active_system()):
                break
            val = goal_group.value
            if val is not None:
                values.append(val)
                text.append(goal_group.text())
        d = self.read_config() or {}
        for t, v in zip(text, values):
            d[t] = v
        self.save_config(d)

    def save_mirror(self, mirror_group):
        d = self.read_config() or {}
        mirror = mirror_group.obj
        d[mirror.name] = mirror.position
        self.save_config(d)

    def save_active_mirrors(self):
        saves = {}
        averages = 1000
        all_mirrors = self.mirrors()
        for mirror in all_mirrors:
            saves[mirror.name] = 0
        for i in range(averages):
            for mirror in all_mirrors:
                saves[mirror.name] += mirror.position / averages
        logger.info('Saving positions: %s', saves)
        read = self.read_config() or {}
        read.update(saves)
        self.save_config(read)

    def active_system(self):
        """
        List of system keys that are part of the active procedure.
        """
        active_system = []
        if self.procedure != 'None':
            for part in self.alignments[self.procedure]:
                active_system.extend(part)
        return active_system

    def load_active_system(self):
        for system in self.active_system():
            self.loader.get_subsystem(system)

    def _objs(self, key):
        objs = []
        for act in self.active_system():
            subsystem = self.loader[act]
            if subsystem is None:
                objs.append(None)
            else:
                objs.append(subsystem[key])
        return objs

    def mirrors(self):
        """
        List of active mirror objects.
        """
        return self._objs('mirror')

    def imagers(self):
        """
        List of active imager objects.
        """
        return self._objs('imager')

    def slits(self):
        """
        List of active slits objects.
        """
        return self._objs('slits')

    def goals(self):
        """
        List of goals in the user entry boxes, or None for empty or invalid
        goals.
        """
        return [goal.value for goal in self.goals_groups]

    def goal(self):
        """
        The goal associated with the visible imager, or None if the visible
        imager is not part of the active procedure.
        """
        index = self.procedure_index()
        if index is None:
            return None
        else:
            return self.goals()[index]

    def procedure_index(self):
        """
        Goal index of the active imager, or None if the visible imager is not
        part of the active procedure.
        """
        try:
            return self.imagers_padded().index(self.image_obj)
        except ValueError:
            return None

    def none_pad(self, obj_list):
        """
        Helper function to extend a list with 'None' objects until it's the
        length of MAX_MIRRORS.
        """
        padded = []
        padded.extend(obj_list)
        while len(padded) < MAX_MIRRORS:
            padded.append(None)
        return padded

    def mirrors_padded(self):
        return self.none_pad(self.mirrors())

    def imagers_padded(self):
        return self.none_pad(self.imagers())

    def slits_padded(self):
        return self.none_pad(self.slits())

    def get_widget_set(self, name, num=MAX_MIRRORS):
        """
        Widgets that come in sets of count MAX_MIRRORS are named carefully so
        we can use this macro to grab related widgets.

        Parameters
        ----------
        name: str
            Base name of widget set e.g. 'name'

        num: int, optional
            Number of widgets to return

        Returns
        -------
        widget_set: list
            List of widgets e.g. 'name_1', 'name_2', 'name_3'...
        """
        widgets = []
        for n in range(1, num + 1):
            widget = getattr(self.ui, name + "_" + str(n))
            widgets.append(widget)
        return widgets

    def ui_filename(self):
        return 'gui.ui'

    def ui_filepath(self):
        return path.join(path.dirname(path.realpath(__file__)),
                         self.ui_filename())
Beispiel #2
0
class BlueskyDispatcher:

    def __init__(self, port=8765, handle_exceptions=False):
        self.version = VERSION

        self._RE = RunEngine(context_managers=[])
        # passing empty array as context managers so that the run engine will
        # not try to register signal handlers or complain that it's not the
        # main thread

        # subscribe our own callback in order to extract the UID of a run and
        # store it in our self._run_uid:
        self._RE.subscribe(self._extract_uid_callback, name="start")

        self._websocket_port = port

        self._handle_exceptions = handle_exceptions
        # determines whether or not this dispatcher should handle exceptions
        # raised by user supplied scan functions by simply logging the
        # exception and also updating any currently connected websocket
        # subscriber clients. Or if it should re raise the exceptions to the
        # user's calling code.

        self._selected_scan = None
        # acts as a pointer to whichever scan function is to be executed.

        self._run_uid = None
        # during execution of a bluesky plan: hold the UID of the run
        # this is intended to aid any clients know where to retrieve their
        # data at a later point if needed, for transparency.

        self._run_start_time = None
        # during execution of a bluesky plan: store the time
        # this is intended to aid any clients that wish to determine how
        # long a run has been going for.
        # created with time_stamp = datetime.datetime.now(datetime.timezone.utc)
        # stringified with: time_str = time_stamp.__str__()
        # and decoded back to an object with:
        # new_time_stamp = datetime.datetime.strptime(time_str, '%Y-%m-%d %H:%M:%S.%f%z')

        self._supplied_params = None
        # when the selected_scan is called, these will be supplied as long as
        # they are not None

        self._scan_functions = {}
        # The dict of bluesky plan scan functions registered against this
        # dispatcher.
        # key = name of the scan
        # value = a python function matching the required signature

        self._callback_count = 0

        self._callbacks_to_subscribe = []
        # The list of callback functions to be subscribed to the run engine
        # instance (no matter the scan/plan being run),
        # Example value:
        # [
        #   {func: callback_function_1, name: "all"},
        #   {func: callback_function_2, name: "stop"}
        # ]
        # Why?: (the need to maintain a list is to ensure that if in future
        # the run engine instance is discarded and recreated anew after every
        # 'run' that the subsequent instance can be set up with the same
        # subscriptions.)

        self._busy = threading.Lock()
        # Used by the RunEngine to signal that it's currently busy. Read by the
        # websocket threads to decide to refuse or accept a request to start a
        # scan.

        self._start_scanning = threading.Event()
        # Used by the websocket thread to signal the main thread that it should
        # start a scan. Which scan to be run should be selected before
        # set()'ing this event.

        self._re_state_changes_q = Queue()
        # Used to capture RunEngine state changes so no events are dropped,
        # even if they occur in rapid succession.
        # The addition of a new item in this queue is what 'wakes up' the
        # coroutine_to_bridge_thread_events_to_asyncio

        self._re_state_changes_data = None
        # Used to hold the information provided by the RunEngine whenever its
        # state changes OR by this dispatcher if an exception was raised when
        # running a plan.
        #   eg.:
        #   {'exception': False,
        #    'new_state': 'idle',
        #    'old_state': 'running'}
        #   OR
        #   {'exception': True,
        #    'message': 'cannot divide by zero'}
        # Consumed by websocket connections serving 'subscribers'.

        self._websocket_state_update = asyncio.Event()
        # the event that 'wakes up' sleeping websocket subscriber connections
        # so they can consume above re_state_changes_data OR
        # exception data to feed to their respective subscribers

        # register our state_hook function on the RunEngine, it will add updates
        # as items on the _re_state_changes_q Queue where the
        # coroutine_to_bridge_thread_events_to_asyncio function can work through
        # the items to one-by-one notify waiting websocket coroutines through
        # the above asyncio.Event()
        self._RE.state_hook = self._state_hook
        logger.info(f'BlueskyDispatcher({VERSION}) initialised.')

    def add_scan(self, scan_function, scan_name):
        if scan_name in self._scan_functions:
            raise ScanNameError(f'The supplied scan name "{scan_name}" is '
                                f'already in use, If you wish to overwrite the '
                                f'previously supplied scan function remove the '
                                f'existing one first with '
                                f'remove_scan("{scan_name}")')
        if self._good_function_signature(scan_function):
            self._scan_functions[scan_name] = scan_function

    def subscribe_callback_function(self, func, name="all"):
        """ subscribe function to mirror the run engine's subscribe function
        which uses the same signature, The purpose of this function is to
        abstract away the access to the run engine, so that the
        dispatcher may have the freedom to recreate the run engine instance
        at any time. The reason for storing the provided callback function in
        the list called _callbacks_to_subscribe is so that in the event that the
        runengine IS recreated, the same callbacks can be applied to it.

        The reason we don't return to the user the actual_token value,
        (but instead one we control) is so that we don't have to depend on the
        returned value from the actual runengine instance being deterministic
        every time this dispatcher may decide to recreate the instance and
        re-subscribe all the callbacks.

        "actual_token" key is to store the value as returned by the LAST and
        most RECENT time THAT callback function was subscribed to THIS
        self._RE instance (so we don't depend on the run engine's
        implementation being deterministic about the assignment of return
        tokenint values)

        "token_int" key is the value we returned to the user and so need to
        remember for when they wish to unsubscribe
        """
        actual_token = self._RE.subscribe(func, name)
        token_to_return = self._get_next_callback_token()
        self._callbacks_to_subscribe.append(
            {
                "func": func,
                "name": name,
                "actual_token": actual_token,
                "user_held_token": token_to_return
            })
        return token_to_return

    def unsubscribe_callback_function(self, user_held_token):
        """ unsubscribes the callback function from the run engine denoted by
        the user_held_token value"""
        # figure out the actual token from the user held token:
        actual_token = None
        for cb in self._callbacks_to_subscribe:
            if cb["user_held_token"] == user_held_token:
                actual_token = cb["actual_token"]
                break
        # unsubscribe the function from the current RE instance:
        if actual_token is None:
            raise LookupError(f"No callback function previously registered "
                              f"with token: {user_held_token}")
        else:
            self._RE.unsubscribe(actual_token)
            # and prevent it being reapplied in the event of RE recreation,
            # amend our _callbacks_to_subscribe list:
            self._callbacks_to_subscribe = [
                entry for entry in self._callbacks_to_subscribe if
                entry["user_held_token"] != user_held_token]

    def remove_scan(self, scan_name):
        if scan_name not in self._scan_functions:
            raise ScanNameError(f'The supplied scan name "{scan_name}" is not '
                                f'in the list of scan functions')
        del self._scan_functions[scan_name]

    def start(self):
        # This code runs in the main thread.
        logger.info('BlueskyDispatcher starting…')
        ws_thread_ready = threading.Event()
        #   .   .   .   .   .   .   .   .   .   .   .   .   .   .   .   .   .
        ws_thread = threading.Thread(target=self._websocket_thread,
                                     kwargs={'ready_signal': ws_thread_ready},
                                     daemon=True,
                                     name="websocket_thread")
        # Regarding above Thread instantiation as a daemon thread:
        # running the thread as a daemon means that the main thread will
        # effectively forget about it once its started (which is how I've
        # been mentally thinking about it anyway) the implications of this
        # are then:
        # a)    the websocket thread will be automatically killed when the main
        #       thread is, this means you need to keep main thread alive for as
        #       long as you wish the websocket thread to be (which we are)
        # b)    You may run into issues if the websocket thread needed to shut
        #       down gracefully, for example it had opened files for writing,
        #       since the main thread will have lost control over it.
        # c)    Most importantly, I don't have to press ctrl-C twice to kill
        #       this script when I fail to have the main thread perform reaping
        #       of the websocket thread.
        #   .   .   .   .   .   .   .   .   .   .   .   .   .   .   .   .   .
        ws_thread.start()
        # wait for the websocket thread to get up and running and ready before
        # continuing
        ws_thread_ready.wait()
        # now go ahead and perform the scan

        # Todo: From here on assumptions are made that the other thread
        #  correctly sets things up for us - check we're not exposing ourselves
        #  to race conditions or deadlocks.

        while True:
            self._start_scanning.wait()

            # we don't acquire the self._busy lock here but instead rely on the
            # websocket thread coroutine that called _start_scanning.set() to
            # acquire it on behalf of this code, because the coroutine that
            # called self._start_scanning.set() is guaranteed to run
            # sequentially to a competing coroutine that might try to do the
            # same thing with a different plan, due to the nature of the
            # asyncio event loop. However if left up to this code to acquire the
            # lock, this code runs in a separate thread to the coroutines and
            # so the risk is that multiple coroutines (running in their own
            # thread) both attempt to set the parameters for the next scan and
            # respond to their respective websocket clients with success
            # messages, yet this thread had not executed anything yet and so you
            # have multiple websocket coroutines attempt to initiate plans
            # before this code (which is in a separate thread to the websockets
            # thread) had a chance to acquire the busy lock.

            self._start_scanning.clear()
            scan_func = self._scan_functions[self._selected_scan]
            try:
                self._run_start_time = datetime.datetime.now(
                        datetime.timezone.utc)
                if self._supplied_params is not None:
                    scan_func(
                        self._RE,
                        self._state_hook,  # Todo: NO state_hook (just testing!)
                        **self._supplied_params)
                else:
                    # this avoids executing self._selected_scan(…, **None) which
                    # raises "TypeError: arg after ** must be a mapping"
                    scan_func(self._RE, self._state_hook)

            except Exception as e:
                name_and_shame = (f'supplied plan with name '
                                  f'"{self._selected_scan}" raised '
                                  f'exception')
                logger.exception(name_and_shame)
                # send message to any connected websocket client there was
                # an issue:
                self._re_state_changes_q.put({
                    'exception': True,
                    'message': name_and_shame + '\n' + repr(e)})
            # IMPORTANT!
            # We rely on the supplied_params being set by the code that
            # .set() the start_scanning Event BEFORE it .set() the
            # start_scanning Event, (we appropriately reset it afterwards)
            self._reset_dispatcher_state_after_run()
            # I'm not sure how to detect that if _supplied_params is not
            # None, that it then pertains to the function we are about to
            # call and not a different previous function that was called.
            # rely on defaults if no params:

            # if we reset self._selected_scan here I think we might risk having
            # this thread be paused while a websocket thread takes a client that
            # requests to start a scan, and it tries to respond that we are
            # still busy with a current scan, and try to get which scan based on
            # self._selected_scan, but we had just reset it to None
            logger.debug("releasing busy lock")
            self._busy.release()
            # if we reset self._selected_scan here AFTER we release the _busy
            # lock, I fear that we could have a scenario where a websocket takes
            # a client that requests to start a scan, and since we have released
            # busy lock, the websocket client assumes it's ok to get a scan
            # started and so it would start by setting up some things including
            # setting the self._selected_scan ... If that happens between the
            # above self._busy.release() and our call here to
            # self._selected_scan = None then we risk resetting the value of
            # _selected_scan when a new scan request was 'coming through the
            # pipeline'. We now risk this execution looping back around to find
            # that we should go through this while loops body again, but when
            # we go to refer to the _selected_scan it will have been reset by
            # us.
            #
            # We don't HAVE to reset the self._selected_scan value if we rely
            # on the websocket thread to appropriately set it before invoking us
            # again, and resetting it here would not remove the reliance on the
            # websocket thread appropriately setting it anyway
            #
            # also we don't have to worry about competing websocket
            # threads from stepping over each other to modify the values of
            # _selected_scan or _supplied_params because they are not running
            # in parallel thanks to asyncio.

    ####    Start of Private methods:

    def _reset_dispatcher_state_after_run(self):
        self._selected_scan = None
        self._run_start_time = None
        self._run_uid = None
        # this is where we would also recycle our runEngine instance if we
        # wanted to ensure a defined start state for the next run (The reason
        # you may end up with an undefined state prior to a subsequent run,
        # would be if the previous run left the runEngine instance altered in
        # some way (subscribing a callback is one way). )
        # a call to recycle the RE instance should be followed by a call to
        # _reapply_callback_functions()

    def _get_next_callback_token(self):
        """ function to get the next usable token int value for returning to
        users that call the subscribe_callback_function method to ensure
        always a unique token value, but without using the value returned by
        the runEngine so as to abstract the implementation of the assignment
        of tokens within the run engine, thereby allowing this dispatcher to
        have the freedom to destroy and recreate the runEngine instance. """
        tokenint_to_return = self._callback_count
        # the token is basically just a value that increments whenever a
        # new function is subscribed.
        self._callback_count += 1
        return tokenint_to_return

    def _extract_uid_callback(self, name, doc):
        """callback to register against the run engine in order to extract
        the run UID (for benefit of any subsequent subscribers)
        (only registered against "start" bluesky documents so we ignore
        "name" param)"""
        self._run_uid = doc["uid"]

        # This callback is being called by the run engine after it's already
        # started, so the state hook has already detected a change in the run
        # engine state and started the "push notification out to websocket
        # clients" process so by now the clients will already have gotten a
        # message describing basically all the below but where the run_uid
        # was not yet set (since we only just set it here now) so this
        # basically sends the same document but this time the run_uid will be
        # populated (sending the seemingly same information for the other
        # fields again because, a) it can't hurt and b) the client need not
        # write logic to deal with the case when certain attributes are omitted.
        self._re_state_changes_q.put(
            {'exception': False,
             'state': self._RE.state,
             'run_uid': self._run_uid,
             'plan_name': self._selected_scan,
             'run_start_time': self._run_start_time})

    def _reapply_callback_functions(self):
        """to be used in the event that the self._RE instance gets recreated,
        so as to ensure the new instance is subscribed to all the same
        callback functions as the previous instance was"""

        # firstly reapply our own callback to extract the UID:
        self._RE.subscribe(self._extract_uid_callback, name="start")

        for cb in self._callbacks_to_subscribe:
            new_actual_token = self._RE.subscribe(cb["func"], cb["name"])
            # and update the stored actual token (leaving the user held token
            # unchanged):
            cb["actual_token"] = new_actual_token

    def _good_function_signature(self, function):
        """checks that the supplied function has a signature that matches our
        requirements (first two args are positional, remaining args are named
        and optional"""
        params = inspect.signature(function).parameters
        # params is now an ordered dict
        count = 0
        for i, (key, val) in enumerate(params.items()):
            count += 1
            if i < 2:
                assert val.default is val.empty, "Don't provide default " \
                                                 "values for the first two " \
                                                 "args, the dispatcher will " \
                                                 "supply those."
            else:
                assert val.default is not val.empty, "You need to provide " \
                                                     "sane defaults for your " \
                                                     "plan's parameters"
        assert count > 1, "You need to supply at least 2 positional args for " \
                          "RunEngine and StateHook"
        return True

    def _state_hook(self, new_state, old_state):
        """this function will be called from the main thread, by the
        RunEngine code whenever its state changes.
        (This code will NOT execute concurrently with the run engine code)
        (This code WILL execute concurrently with the websocket/s code)
        """
        # prepare start time as string form, not object. it's for sending as
        # json.
        start_time_str = self._run_start_time
        if start_time_str is not None:
            start_time_str = start_time_str.__str__()
        # https://docs.python.org/3.6/library/queue.html
        self._re_state_changes_q.put(
            {'new_state': new_state,
             'old_state': old_state,
             'exception': False,
             'run_uid': self._run_uid,
             'plan_name': self._selected_scan,
             'run_start_time': start_time_str})

    async def _bridge_queue_events_to_coroutines(self,
                                                  loop,
                                                  queue,
                                                  asyncio_event):
        """This function serves to bridge the events between the main threading
        driven thread and the secondary asyncio driven thread by using an
        executor to get around the blocking aspects of waiting on a threading
        Queue primitive and then subsequently echo that item using an
        asyncio.Event() for any awaiting open websocket connection coroutines.

        Using a queue ensures every single RunEngine state update event is
        accounted for and an appropriate message dispatched to any waiting
        subscriber clients.

        It is expected that the queue will remain relatively empty and
        only during very brief moments where the RunEngine changes state in
        quick succession would there be more than one event item in the queue

        The state_hook function, called by the RunEngine whenever its state
        changes, is what adds items to the queue
        """
        while True:
            state_update = await loop.run_in_executor(
                None,
                functools.partial(queue.get, block=True))
            # (this is the only consumer of the queue)
            logger.debug(f'return value from the queue: {repr(state_update)}')
            self._re_state_changes_data = state_update

            asyncio_event.set()
            # at this point any websocket connection coroutines that were
            # shelved out of the event loop awaiting this asyncio_event will be
            # scheduled back into the event loop. and continue execution after
            # their `await event.wait()` line.

            asyncio_event.clear()
            # it's ok to immediately reset the Event, any coroutines that
            # were waiting on it will still be run once this coroutine yields
            # control back to the event loop via an await or completion, it
            # just means that if they hit another line (or the same line) that
            # makes them await it they will again be shelved out of the event
            # loop until it is once again set().

            await self._yield_control_back_to_event_loop()
            # deliberately pass control back to the event loop so that the
            # coroutines that are now back in the loop can run before we reset
            # the global variable they will be referring to.

            self._re_state_changes_data = None

            queue.task_done()

    async def _yield_control_back_to_event_loop(self):
        await asyncio.sleep(0)
        # https://github.com/python/asyncio/issues/284

    def _websocket_thread(self, ready_signal=None):
        # runs in the other (asyncio driven) thread (not the main thread)
        logger.debug("websocket thread: establishing new event loop")
        loop_for_this_thread = asyncio.new_event_loop()
        asyncio.set_event_loop(loop_for_this_thread)

        if ready_signal is not None:
            logger.debug("websocket thread: signalling that I'm ready")
            ready_signal.set()

        logger.debug("websocket thread: starting websocket control server")
        loop_for_this_thread.run_until_complete(
            websockets.serve(
                self._websocket_server,
                "0.0.0.0",
                self._websocket_port))

        logger.debug("websocket thread: Setting the websocket state update "
                     "Event() object")
        self._websocket_state_update = asyncio.Event(loop=loop_for_this_thread)
        logger.debug("websocket thread: starting events bridging function")
        asyncio.ensure_future(
            self._bridge_queue_events_to_coroutines(
                loop_for_this_thread,
                self._re_state_changes_q,
                self._websocket_state_update), loop=loop_for_this_thread)

        logger.debug("websocket thread: ready for connections!")

        loop_for_this_thread.run_forever()
        loop_for_this_thread.close()

    async def _websocket_server(self, websocket, path):
        # runs in the other (asyncio driven) thread (not the main thread)
        logger.info("websocket server: client connected to me!")
        logger.debug(f'websocket server: runengine state: {self._RE.state}')
        while True:
            # This loop is only effectively infinite if their message
            # contains a 'keep_open' key (see bottom of loop)
            json_msg = await websocket.recv()
            try:
                obj = json.loads(json_msg)
                """
                    valid example obj:
                    
                    start a new plan running:
                    {
                        'type': 'start',
                        'plan': 'simulated',  # <-- this is the scan_name
                        'keep_open': True     # <-- this is optional
                    }

                    ( the presence of a keep_open key is enough, 
                    the value doesn't matter )

                    pause a running plan:
                    {'type': 'pause'}
                    
                    interrogate the current state of the RunEngine:
                    {'type': 'state'}

                    subscribe to a stream of RunEngine state update events
                    (the onus is on the client to close the connection):
                    {'type': 'subscribe'}
                """
            except json.JSONDecodeError:
                logger.warning(f'websocket server: the following received '
                               f'message could not be properly decoded as '
                               f'json: {json_msg}')
                await websocket.send(json.dumps({
                    'success': False,
                    'status': "Your message could not be properly decoded, "
                              "is it valid json?"}))
                await websocket.close()
                return None
            # now at this point we can be relatively
            # confident obj is a proper object
            logger.debug(f'received websocket json message: {repr(obj)}')
            if obj['type'] == 'halt':
                # "stop a running plan,
                # don't wait for it to clean up, mark as aborted"
                try:
                    self._RE.halt()
                    # halt() takes no args but may raise
                    # runtime error/transition error, or returns task.result()
                    await websocket.send(json.dumps({
                        'success': True,
                        'status': "halt requested"}))
                except Exception as err:
                    err_repr = repr(err)
                    await websocket.send(json.dumps({
                        'success': False,
                        'status': "exception was raised",
                        'exception': err_repr}))

            elif obj['type'] == 'abort':
                # "stop a running plan, wait for it to cleanup, mark as aborted"
                try:
                    self._RE.abort(reason='Todo: get reason from websocket '
                                          'client')
                    # Todo: get reason from websocket client
                    await websocket.send(json.dumps({
                        'success': True,
                        'status': "abort requested"}))
                except Exception as err:
                    err_repr = repr(err)
                    await websocket.send(json.dumps({
                        'success': False,
                        'status': "exception was raised",
                        'exception': err_repr}))

            elif obj['type'] == 'stop':
                # "stop a running plan, wait for it to clean up, mark as
                # successful"
                try:
                    self._RE.stop()
                    # .stop() takes no args, returns tuple(self._run_start_uids)
                    await websocket.send(json.dumps({
                        'success': True,
                        'status': "stop requested"}))
                except Exception as err:
                    err_repr = repr(err)
                    await websocket.send(json.dumps({
                        'success': False,
                        'status': "exception was raised",
                        'exception': err_repr}))

            elif obj['type'] == 'pause':
                # "if checkpoints have been programmed into the plan, pause on
                # the next checkpoint
                # if no checkpoints, the effect becomes the same as abort
                # this is analogous to pressing Ctrl-C during execution of a
                # plan manually on the terminal"
                try:
                    self._RE.request_pause()
                    await websocket.send(json.dumps({
                        'success': True,
                        'status': "pause requested"}))
                except Exception as err:
                    err_repr = repr(err)
                    await websocket.send(json.dumps({
                        'success': False,
                        'status': "exception was raised",
                        'exception': err_repr}))

            elif obj['type'] == 'resume':
                # resume a paused plan from the last checkpoint
                try:
                    if self._RE.state == 'paused':
                        self._RE.resume()  # takes no args
                        await websocket.send(json.dumps({
                            'success': True,
                            'status': "resume requested"}))
                    else:
                        await websocket.send(json.dumps({
                            'success': False,
                            'status': "cannot resume a runEngine that "
                                      "isn't paused"}))
                except Exception as err:
                    err_repr = repr(err)
                    await websocket.send(json.dumps({
                        'success': False,
                        'status': "exception was raised",
                        'exception': err_repr}))

            elif obj['type'] == 'start':
                if self._busy.locked():
                    # scan is currently underway, don't attempt to start another
                    await websocket.send(json.dumps({
                        'success': False,
                        'status': f'Currently busy with '
                                  f'"{self._selected_scan}" scan'}))
                elif 'plan' not in obj:
                    # if they don't specify which plan then fail:
                    await websocket.send(json.dumps({
                        'success': False,
                        'status': f'You need to provide a "plan" key '
                                  f'specifying which scan you want to start'}))
                elif obj['plan'] not in self._scan_functions:
                    # if we don't have the plan they specified then fail:
                    requested_scan = obj['plan']
                    await websocket.send(json.dumps({
                        'success': False,
                        'status': f'I don\'t recognise any scan function '
                                  f'by the name "{requested_scan}"'}))
                else:
                    response = {  # Get a default response prepared
                        'success': True,
                        'status': "Signalling main thread to start a new scan"}
                    # pull out any supplied parameters:
                    if 'params' in obj:
                        self._supplied_params = obj['params']
                        response['params'] = obj['params']  # update response

                    self._selected_scan = obj['plan']

                    # Pre-emptively set a lock here on behalf of the 'main'
                    # thread to ensure that once we .set() the
                    # self._start_scanning event object the busy flag is
                    # already appropriately set, doing it here ensures that
                    # once we respond to our respective websocket client and
                    # relinquish control back to the event loop, that there
                    # is no chance that a competing websocket coroutine can
                    # come in and attempt to do the same thing before the
                    # 'main' thread had even had a chance to get started yet
                    # (because remember it's a separate thread so the order
                    # of execution between this thread and that is
                    # non-deterministic), the competing websocket coroutine will
                    # already see that the busy flag is set and so not
                    # attempt to modify the class instance variables (such as
                    # self._selected_scan and self._supplied_params,
                    # etc.) while the 'main' thread is starting to get
                    # started on executing the plan.
                    logger.debug("locking busy lock")
                    self._busy.acquire()

                    # initiate a scan by setting an event object that the main
                    # thread is waiting on before it proceeds to run whichever
                    # function is referenced now by self._selected_scan and
                    # the corresponding value in self.scan_functions
                    self._start_scanning.set()
                    await websocket.send(json.dumps(response))

            elif obj['type'] == 'state':
                await websocket.send(json.dumps({
                    'type': "status",
                    'about': "current state of the Bluesky RunEngine",
                    'state': self._RE.state
                }))

            elif obj['type'] == "subscribe":
                # if client wants to subscribe to runengine state updates,
                # send them an initial message containing the current state
                # as derived from the self instance attributes.
                # quick preprocessing of the start time so we don't call
                # None.__str__(), otherwise other end will get a string "None"
                start_time = None
                if self._run_start_time is not None:
                    start_time = self._run_start_time.__str__()
                await websocket.send(json.dumps({
                    'bluesky_dispatcher_version': VERSION,
                    'type': "status",
                    'about': "run engine state updates",
                    'state': self._RE.state,  # a string
                    'plan_name': self._selected_scan,  # a string
                    'run_uid': self._run_uid,  # a string
                    'run_start_time': start_time  # a string
                }))
                while True:
                    # then wait until the state is updated:
                    await self._websocket_state_update.wait()
                    assert self._re_state_changes_data is not None, (
                        "Alex's mental model of asyncio event loop behaviour "
                        "is wrong - prerequisites not met in websocket "
                        "connection that just received asyncio.Event signal "
                        "regarding update to RunEngine state")
                    # then send our client an update:
                    if self._re_state_changes_data['exception']:
                        await websocket.send(json.dumps({
                            'type': "exception",
                            'about': "exception raised during execution of "
                                     "supplied bluesky plan",
                            'message': self._re_state_changes_data['message']
                        }))
                    else:
                        # NOTE: we derive our values from those
                        # stored in the state_changes_data dict
                        # because this is not going to change between
                        # swappings of coroutines, (if we derived the
                        # current state from the actual RE instance and
                        # self. properties then we risk missing fully formed
                        # updates because of race conditions
                        await websocket.send(json.dumps({
                            'type': "status",
                            'about': "run engine state updates",
                            'state': self._re_state_changes_data['new_state'],
                            'old_state': self._re_state_changes_data['old_state'],
                            'plan_name': self._re_state_changes_data[
                                'plan_name'],  # a string
                            'run_uid': self._re_state_changes_data['run_uid'],
                            # a string
                            'run_start_time': self._re_state_changes_data[
                                'run_start_time']  # a string
                        }))
            # client's message doesn't match any of the above, therefore we
            # don't know what they want from us:
            else:
                await websocket.send(json.dumps({
                    'success': False,
                    'status': "Unsupported message type"}))
                # At this point we don't want to keep looping here if their
                # original request is not serviceable.
                await websocket.close()
                break

            if 'keep_open' not in obj:
                await websocket.close()
                break
Beispiel #3
0
class QRunEngine(QObject):
    sigDocumentYield = Signal(str, dict)
    sigAbort = Signal()  # TODO: wireup me
    sigException = Signal(Exception)
    sigFinish = Signal()
    sigStart = Signal()
    sigPause = Signal()
    sigResume = Signal()

    def __init__(self, **kwargs):
        super(QRunEngine, self).__init__()

        self.RE = RunEngine(context_managers=[],
                            during_task=DuringTask(),
                            **kwargs)
        self.RE.subscribe(self.sigDocumentYield.emit)

        # TODO: pull from settings plugin
        from suitcase.mongo_normalized import Serializer
        #TODO create single databroker db
        #python-dotenv stores name-value pairs in .env (add to .gitginore)
        username = os.getenv("USER_MONGO")
        pw = os.getenv("PASSWD_MONGO")
        try:
            self.RE.subscribe(
                Serializer(
                    f"mongodb://{username}:{pw}@localhost:27017/mds?authsource=mds",
                    f"mongodb://{username}:{pw}@localhost:27017/fs?authsource=fs"
                ))
        except OperationFailure as err:
            msg.notifyMessage("Could not connect to local mongo database.",
                              title="xicam.Acquire Error",
                              level=msg.ERROR)
            msg.logError(err)

        self.queue = PriorityQueue()
        self.process_queue()

    @threads.method(threadkey="run_engine", showBusy=False)
    def process_queue(self):
        while True:
            try:
                priority_plan = self.queue.get(
                    block=True,
                    timeout=.1)  # timeout is arbitrary, we'll come right back
            except Empty:
                continue
            priority, (args,
                       kwargs) = priority_plan.priority, priority_plan.args

            self.sigStart.emit()
            msg.showBusy()
            try:
                self.RE(*args, **kwargs)
            except Exception as ex:
                msg.showMessage(
                    "An error occured during a Bluesky plan. See the Xi-CAM log for details."
                )
                msg.logError(ex)
                self.sigException.emit(ex)
            finally:
                msg.showReady()
            self.sigFinish.emit()

    @wraps(RunEngine.__call__)
    def __call__(self, *args, **kwargs):
        self.put(*args, **kwargs)

    @property
    def isIdle(self):
        return self.RE.state == 'idle'

    def abort(self, reason=''):
        if self.RE.state != 'idle':
            self.RE.abort(reason=reason)
            self.sigAbort.emit()

    def pause(self, defer=False):
        if self.RE.state != 'paused':
            self.RE.request_pause(defer)
            self.sigPause.emit()

    def resume(self, ):
        if self.RE.state == 'paused':
            self.threadfuture = threads.QThreadFuture(
                self.RE.resume,
                threadkey='RE',
                showBusy=True,
                finished_slot=self.sigFinish.emit)
            self.threadfuture.start()
            self.sigResume.emit()

    def put(self, *args, priority=1, **kwargs):
        # handle ParameterizedPlan's
        # plan = args[0]
        # if isinstance(args[0], ParameterizedPlan):
        #     # Ask for parameters
        #     param = plan.parameter
        #     if param:
        #         ParameterDialog(param).exec_()

        reserved = set(kwargs.keys()).union(
            ['plan_type', 'plan_args', 'scan_id', 'time', 'uid'])
        self._metadata_dialog = MetadataDialog(reserved=reserved)
        self._metadata_dialog.open()
        self._metadata_dialog.accepted.connect(
            partial(self._put, self._metadata_dialog, priority, args, kwargs))

    def _put(self, dialog: MetadataDialog, priority, args, kwargs):
        metadata = dialog.get_metadata()
        kwargs.update(metadata)
        self.queue.put(PrioritizedPlan(priority, (args, kwargs)))