Source code for schrodinger.ui.qt.tasks.configwidgets

from schrodinger import get_maestro
from schrodinger.models import advanced_mappers
from schrodinger.models import mappers
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtWidgets
from schrodinger.tasks import jobtasks
from schrodinger.tasks.jobtasks import AllowedHostTypes
from schrodinger.ui.qt import basewidgets
from schrodinger.ui.qt import config_dialog
from schrodinger.ui.qt import mapperwidgets
from schrodinger.ui.qt import utils
from schrodinger.ui.qt.tasks import config_dialog_ui
from schrodinger.utils.qt_utils import suppress_signals

maestro = get_maestro()


class ConfigWidgetMixin(advanced_mappers.ModelBaseClassMapperMixin):
    model_base_class = jobtasks.JobConfig

    @property
    def config(self):
        return self.model

    def _setupModelClass(self, model):
        super()._setupModelClass(model)
        self._onConfigClassSet()

    def _onConfigClassSet(self):
        pass

    def makeInitialModel(self):
        # This is hard to figure out
        return None


class HostSelector(mappers.MapperMixin, basewidgets.BaseWidget):
    """
    :cvar PLACEHOLDER_TEXTS: Placeholders to put in the host combo when no hosts
        are available based on the currently allowed host types.
    """

    model_class = jobtasks.HostSettings
    VALID_STATE_PROP = 'valid_state'
    VALIDATION_STYLE_SHEET = f"""
        QComboBox[{VALID_STATE_PROP}="true"] {{font-style: normal}}
        QComboBox[{VALID_STATE_PROP}="false"] {{font-style: italic}}
        """
    PLACEHOLDER_TEXTS = {
        AllowedHostTypes.CPU_ONLY: '(no CPUs found)',
        AllowedHostTypes.GPU_ONLY: '(no GPUs found)',
        AllowedHostTypes.CPU_AND_GPU: '(no hosts found)',
    }

    # The following attributes used to be accessible but are now accessed
    # with methods. We raise errors if you try to access them to catch any
    # cases where we might still rely on them.
    @property
    def host_combo(self):
        raise AttributeError(
            "host_combo is not accessible as an attribute. "
            "To get the host combobox use getHostCombo() instead.")

    @host_combo.setter
    def host_combo(self, _):
        raise AttributeError(
            "host_combo is not accessible as an attribute. "
            "To get the host combobox use getHostCombo() instead.")

    @property
    def units_lbl(self):
        raise AttributeError(
            "units_lbl is not accessible as an attribute. "
            "To get the host combobox use getSubjobsUnitLabel() instead.")

    @units_lbl.setter
    def units_lbl(self, _):
        raise AttributeError(
            "units_lbl is not accessible as an attribute. "
            "To get the host combobox use getSubjobsUnitLabel() instead.")

    @property
    def num_subjobs_sb(self):
        raise AttributeError(
            "num_subjobs_sb is not accessible as an attribute. "
            "To get the host combobox use getNumSubjobsSpinbox() instead.")

    @num_subjobs_sb.setter
    def num_subjobs_sb(self, _):
        raise AttributeError(
            "num_subjobs_sb is not accessible as an attribute. "
            "To get the host combobox use getNumSubjobsSpinbox() instead.")

    def initSetUp(self):
        super().initSetUp()

        self._label = QtWidgets.QLabel('Host: ')
        self._host_combo = mapperwidgets.MappableComboBox()
        self._host_combo.setSizePolicy(QtWidgets.QSizePolicy.MinimumExpanding,
                                       QtWidgets.QSizePolicy.Fixed)
        self._host_combo.setStyleSheet(self._host_combo.styleSheet() +
                                       self.VALIDATION_STYLE_SHEET)
        self._total_lbl = QtWidgets.QLabel('Total:')
        self._num_subjobs_sb = QtWidgets.QSpinBox()
        self._units_lbl = QtWidgets.QLabel()
        self._subjob_widgets = (self._units_lbl, self._num_subjobs_sb,
                                self._total_lbl)

    def initLayOut(self):
        super().initLayOut()
        layout = QtWidgets.QHBoxLayout()
        layout.addWidget(self.getLabel())
        layout.addWidget(self.getHostCombo())
        layout.addWidget(self.getSubjobsLabel())
        layout.addWidget(self.getNumSubjobsSpinbox())
        layout.addWidget(self.getSubjobsUnitLabel())
        layout.setContentsMargins(0, 0, 0, 0)
        self.host_layout = layout
        self.main_layout.addLayout(self.host_layout)

    def initSetDefaults(self):
        super().initSetDefaults()
        self._setValidState(True)

    def getLabelText(self):
        return self._label.text()

    def setLabelText(self, new_text: str):
        return self._label.setText(new_text)

    def getLabel(self) -> QtWidgets.QLabel:
        """
        :return: Leading host label.
        """
        return self._label

    def getHostCombo(self) -> mapperwidgets.MappableComboBox:
        """
        :return: Host combo box.
        """
        return self._host_combo

    def getSubjobsLabel(self) -> QtWidgets.QLabel:
        """
        :return: Sub jobs label.
        """
        return self._total_lbl

    def getNumSubjobsSpinbox(self) -> QtWidgets.QSpinBox:
        """
        :return: Sub jobs spinbox.
        """
        return self._num_subjobs_sb

    def getSubjobsUnitLabel(self) -> QtWidgets.QLabel:
        """
        :return: Sub jobs unit label.
        """
        return self._units_lbl

    def _updateUnitsLabel(self):
        host = self.model.host
        if host is None:
            return
        if not host.num_gpus:
            self._units_lbl.setText('processors')
        else:
            self._units_lbl.setText('GPUs')

    def _updateNumLimits(self):
        host = self.model.host
        if host is None:
            return
        self._num_subjobs_sb.setMinimum(1)
        if not host.num_gpus:
            self._num_subjobs_sb.setMaximum(host.processors)
        else:
            self._num_subjobs_sb.setMaximum(host.num_gpus)

    def _updateHostOptions(self):
        self._host_combo.clear()
        if self.model.allowed_host_types is None:
            self.setVisible(False)
            return
        self.setVisible(True)
        hosts = jobtasks.get_hosts()
        allowed_types = self.model.allowed_host_types
        any_host_added = False
        for host in hosts:
            if allowed_types is AllowedHostTypes.CPU_ONLY and host.num_gpus:
                continue
            if allowed_types is AllowedHostTypes.GPU_ONLY and not host.num_gpus:
                continue
            text = self._makeHostText(host)
            self._host_combo.addItem(text, host)
            any_host_added = True
        if not any_host_added:
            self._host_combo.addItem(self.getPlaceholderText(), None)

    def _makeHostText(self, host):
        text = host.name
        if self.model.num_subjobs is not None:
            if self.model.allowed_host_types is AllowedHostTypes.CPU_ONLY:
                text += f' ({host.processors})'
            else:
                text += f' ({host.processors}, {host.num_gpus})'
        return text

    def defineMappings(self):
        M = self.model_class
        subjobs_target = mappers.TargetSpec(self._num_subjobs_sb,
                                            setter=self._numSubjobSetter,
                                            datatype=None)
        return [
                (self._updateHostOptions, M.allowed_host_types),
                (self._host_combo, M.host),
                (self._onHostChanged, M.host),
                (subjobs_target, M.num_subjobs)] # yapf:disable

    def _numSubjobSetter(self, new_val):
        if new_val is None:
            return
        else:
            self._num_subjobs_sb.setValue(new_val)

    def _setNumSubjobVisible(self, visibility):
        for wdg in self._subjob_widgets:
            wdg.setVisible(visibility)

    def _onHostChanged(self):
        self._updateSubjobSpinbox()
        valid = True
        if self._host_combo.count() == 1 and self._host_combo.currentText(
        ) == self.getPlaceholderText():
            valid = False
        self._setValidState(valid)

    def _updateSubjobSpinbox(self):
        if self.model.num_subjobs is None:
            self._setNumSubjobVisible(False)
            return
        self._setNumSubjobVisible(True)
        self._updateNumLimits()
        self._updateUnitsLabel()

    def _setValidState(self, valid: bool):
        """
        Set whether this widget is in the valid or invalid state. The invalid
        state is simply when no GPU subhosts are available and thus the combo
        box would otherwise be empty. Special styling is applied. The widget is
        in the valid state at all other times.

        :param valid: Whether or not the widget state is valid.
        """
        self._host_combo.setProperty(self.VALID_STATE_PROP, valid)
        utils.update_widget_style(self._host_combo)
        self._host_combo.setEnabled(valid)

    def getPlaceholderText(self) -> str:
        """
        Return the placeholder text for the currently allowed host types.
        """
        return self.PLACEHOLDER_TEXTS.get(self.model.allowed_host_types, '')


class HostLayoutWidget(basewidgets.BaseWidget):
    """
    A display widget for showing one or more HostSelectors. When multiple host
    selectors are added, the host selectors are shown inside of a
    group box. To use, instantiate and add host selectors using the
    `addHostSelector` method.
    """

    def initSetUp(self):
        super().initSetUp()
        self._group_box = QtWidgets.QGroupBox('Hosts:')
        self._selectors = []

    def initLayOut(self):
        super().initLayOut()
        self._layout = layout = QtWidgets.QGridLayout(self._group_box)
        layout.setContentsMargins(0, 0, 0, 0)

    def addHostSelector(self, selector):
        self._selectors.append(selector)
        if len(self._selectors) == 1:
            # Show only the single selector in the main layout
            self.main_layout.addWidget(selector)
            return

        if len(self._selectors) == 2:
            # Adding a second selector, switch to displaying a HostLayoutWidget
            # instead of individual selectors
            self.main_layout.removeWidget(self._selectors[0])
            self.main_layout.addWidget(self._group_box)
            self._addSelectorToRow(self._selectors[0], 0)

        self._addSelectorToRow(selector, len(self._selectors) - 1)

    def _addSelectorToRow(self, selector, row_idx):
        layout = self._group_box.layout()
        layout.addWidget(selector.getLabel(), row_idx, 0)
        layout.addWidget(selector.getHostCombo(), row_idx, 1)
        layout.addWidget(selector.getSubjobsLabel(), row_idx, 2)
        layout.addWidget(selector.getNumSubjobsSpinbox(), row_idx, 3)
        layout.addWidget(selector.getSubjobsUnitLabel(), row_idx, 4)

    def setHostSelectorOrder(self, selectors):
        if any(current_selector not in selectors
               for current_selector in self._selectors):
            raise ValueError("All selectors currently in the HostLayoutWidget "
                             "must be included in the new ordering.")
        elif any(selector not in self._selectors for selector in selectors):
            raise ValueError("New selector ordering must only include "
                             "selectors already added to the HostLayoutWidget")

        # Remove all host selector widgets from the layout.
        layout = self._group_box.layout()
        for selector in self._selectors:
            layout.removeWidget(selector.getLabel())
            layout.removeWidget(selector.getHostCombo())
            layout.removeWidget(selector.getSubjobsLabel())
            layout.removeWidget(selector.getNumSubjobsSpinbox())
            layout.removeWidget(selector.getSubjobsUnitLabel())

        # Set self._selectors to the new ordering just to be consistent
        self._selectors = selectors

        # Add selectors back in
        for row_idx, selector in enumerate(self._selectors):
            self._addSelectorToRow(selector, row_idx)


class IncorporationSelector(ConfigWidgetMixin, basewidgets.BaseWidget):
    Mode = jobtasks.IncorporationMode
    MODE_TEXT = {
        Mode.APPEND: config_dialog.DISP_APPEND,
        Mode.APPENDINPLACE: config_dialog.DISP_APPENDINPLACE,
        Mode.REPLACE: config_dialog.DISP_REPLACE,
        Mode.IGNORE: config_dialog.DISP_IGNORE,
    }

    def initSetUp(self):
        super().initSetUp()
        self.incorp_combo = QtWidgets.QComboBox()

    def initLayOut(self):
        super().initLayOut()
        self.main_layout.addWidget(self.incorp_combo)

    def _onConfigClassSet(self):
        super()._onConfigClassSet()
        with suppress_signals(self.incorp_combo):
            self._setUpIncorpCombo()

    def _setUpIncorpCombo(self):
        Config = self.model_class
        self.incorp_combo.clear()
        if Config.incorporation is None:
            return
        allowed_modes = Config.incorporation.allowed_modes
        for mode in allowed_modes:
            self.incorp_combo.addItem(self.MODE_TEXT[mode], mode)

    def defineMappings(self):
        M = self.model_class
        if M.incorporation is None:
            return []
        target = mappers.TargetSpec(
            getter=self._currentIncorp,
            setter=self._setCurrentIncorp,
            signal=self.incorp_combo.currentIndexChanged)
        return [(target, M.incorporation)]

    def _currentIncorp(self):
        return self.incorp_combo.currentData()

    def _setCurrentIncorp(self, mode):
        for index in range(self.incorp_combo.count()):
            item_mode = self.incorp_combo.itemData(index)
            if mode == item_mode:
                break
        else:
            index = 0
        self.incorp_combo.setCurrentIndex(index)


class JobnameWidget(ConfigWidgetMixin, basewidgets.BaseWidget):
    nameUpdated = QtCore.pyqtSignal()

    def initSetUp(self):
        super().initSetUp()
        self.jobname_le = QtWidgets.QLineEdit()
        self.jobname_le.editingFinished.connect(self._onEditingFinished)
        self.jobname_le.textChanged.connect(self._onTextChanged)

    def initLayOut(self):
        super().initLayOut()
        self.main_layout.addWidget(self.jobname_le)

    def _onEditingFinished(self):
        if not self.jobname_le.text():
            self.nameUpdated.emit()

    def _onTextChanged(self):
        if self.jobname_le.text():
            self.nameUpdated.emit()

    def setText(self, text):
        if text:
            self.jobname_le.setText(text)

    def defineMappings(self):
        M = self.model_class
        if M.jobname is None:
            return []
        target = mappers.TargetSpec(getter=self.jobname_le.text,
                                    setter=self.setText,
                                    signal=self.nameUpdated)
        return [(target, M.jobname)]


class ConfigDialog(ConfigWidgetMixin, basewidgets.BaseOptionsDialog):
    ui_module = config_dialog_ui
    startRequested = QtCore.pyqtSignal()

    def initSetOptions(self):
        super().initSetOptions()
        self.help_topic = 'MM_TOPIC_JOB_START_DIALOG'

    def initSetUp(self):
        super().initSetUp()
        self.incorp_selector = IncorporationSelector()
        self._setUpHostSelector()
        self._host_layout_widget = HostLayoutWidget(parent=self)
        self._host_layout_widget.addHostSelector(self.host_selector)
        self.jobname_wdgt = JobnameWidget(parent=self)
        self.job_widgets = (self._host_layout_widget, self.jobname_wdgt)
        self.reset_btn.hide()

    def _setUpHostSelector(self):
        self.host_selector = HostSelector(parent=self)

    def initLayOut(self):
        super().initLayOut()
        self.ui.incorp_layout.addWidget(self.incorp_selector)
        self.ui.jobname_layout.addWidget(self.jobname_wdgt)
        self.ui.job_group.layout().addWidget(self._host_layout_widget)

    def _onConfigClassSet(self):
        super()._onConfigClassSet()
        M = self.model_class
        visible = M.incorporation is not None and bool(maestro)
        self.ui.output_group.setVisible(visible)

    def setModel(self, *args, **kwargs):
        super().setModel(*args, **kwargs)
        self.ui.job_group.setVisible(True)
        self.jobname_wdgt.setVisible(self.model_class.jobname is not None)
        show_job_group = any([w.isVisibleTo(self) for w in self.job_widgets])
        self.ui.job_group.setVisible(show_job_group)

    @basewidgets.bottom_button('Run')
    def _runBtnSlot(self):
        self.accept()
        self.startRequested.emit()

    def defineMappings(self):
        M = self.model_class
        mappings = [(self.host_selector, M.host_settings),
                    (self.incorp_selector, M), (self.jobname_wdgt, M)]
        return mappings