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)',
    }

    def initSetUp(self):
        super().initSetUp()
        self.num_subjobs_sb = QtWidgets.QSpinBox()

        self._label = QtWidgets.QLabel('Host: ')
        self.units_lbl = QtWidgets.QLabel()
        self.host_combo = mapperwidgets.MappableComboBox()
        self.host_combo.setStyleSheet(self.host_combo.styleSheet() +
                                      self.VALIDATION_STYLE_SHEET)
        self._total_lbl = QtWidgets.QLabel('Total:')
        self._group_box = QtWidgets.QGroupBox(parent=self)
        self._subjob_widgets = (self.units_lbl, self.num_subjobs_sb,
                                self._total_lbl)

    def initLayOut(self):
        super().initLayOut()
        self._group_box.setLayout(QtWidgets.QHBoxLayout())
        self._group_box.layout().addWidget(self._label)
        self._group_box.layout().addWidget(self.host_combo)
        self._group_box.layout().addWidget(self._total_lbl)
        self._group_box.layout().addWidget(self.num_subjobs_sb)
        self._group_box.layout().addWidget(self.units_lbl)
        self.main_layout.addWidget(self._group_box)

    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 getNumSubjobsSpinbox(self):
        return self.num_subjobs_sb

    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 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,
        Mode.APPENDNEW: config_dialog.DISP_APPENDNEW
    }

    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.jobname_wdgt = JobnameWidget(parent=self)
        self.job_widgets = (self.host_selector, 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_selector)

    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