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 mapperwidgets
from schrodinger.ui.qt import config_dialog
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):

    model_class = jobtasks.HostSettings

    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._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 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('', 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
        host_combo_target = mappers.TargetSpec(
            self.host_combo, slot=self._updateSubjobSpinbox)
        subjobs_target = mappers.TargetSpec(
            self.num_subjobs_sb, setter=self._numSubjobSetter, datatype=None)
        return [
                (self._updateHostOptions, M.allowed_host_types),
                (host_combo_target, 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 _updateSubjobSpinbox(self):
        if self.model.num_subjobs is None:
            self._setNumSubjobVisible(False)
            return
        self._setNumSubjobVisible(True)
        self._updateNumLimits()
        self._updateUnitsLabel()


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