Enhanced computational resource widget with resource setup

src/aiidalab_qe/app/parameters/qeapp.yaml

@@ -22,8 +22,11 @@ advanced:
         accuracy: efficiency
     tot_charge: 0
 
-## Codes
+## Computational resources
 codes:
-    dos: dos-7.2@localhost
-    projwfc: projwfc-7.2@localhost
-    pw: pw-7.2@localhost
+    dos:
+        code: dos-7.2@localhost
+    projwfc:
+        code: projwfc-7.2@localhost
+    pw:
+        code: pw-7.2@localhost

src/aiidalab_qe/app/submission/__init__.py

@@ -5,24 +5,21 @@
 """
 from __future__ import annotations
 
-import os
-
 import ipywidgets as ipw
 import traitlets as tl
 from aiida import orm
 from aiida.common import NotExistent
 from aiida.engine import ProcessBuilderNamespace, submit
-from aiidalab_widgets_base import ComputationalResourcesWidget, WizardAppWidgetStep
+from aiidalab_widgets_base import WizardAppWidgetStep
 from IPython.display import display
 
 from aiidalab_qe.app.parameters import DEFAULT_PARAMETERS
 from aiidalab_qe.app.utils import get_entry_items
 from aiidalab_qe.common.setup_codes import QESetupWidget
 from aiidalab_qe.common.setup_pseudos import PseudosInstallWidget
+from aiidalab_qe.common.widgets import PWscfWidget
 from aiidalab_qe.workflows import QeAppWorkChain
 
-from .resource import ParallelizationSettings, ResourceSelectionWidget
-
 
 class SubmitQeAppWorkChainStep(ipw.VBox, WizardAppWidgetStep):
     """Step for submission of a bands workchain."""
@@ -61,17 +58,11 @@ def __init__(self, qe_auto_setup=True, **kwargs):
         self.message_area = ipw.Output()
         self._submission_blocker_messages = ipw.HTML()
 
-        self.pw_code = ComputationalResourcesWidget(
+        self.pw_code = PWscfWidget(
             description="pw.x:", default_calc_job_plugin="quantumespresso.pw"
         )
 
-        self.resources_config = ResourceSelectionWidget()
-        self.parallelization = ParallelizationSettings()
-
-        self.set_resource_defaults()
-
         self.pw_code.observe(self._update_state, "value")
-        self.pw_code.observe(self._update_resources, "value")
 
         # add plugin's entry points
         self.codes = {"pw": self.pw_code}
@@ -86,8 +77,6 @@ def __init__(self, qe_auto_setup=True, **kwargs):
                 self.codes[name] = code
                 code.observe(self._update_state, "value")
                 self.code_children.append(self.codes[name])
-        # set default codes
-        self.set_selected_codes(DEFAULT_PARAMETERS["codes"])
         #
         self.submit_button = ipw.Button(
             description="Submit",
@@ -123,15 +112,15 @@ def __init__(self, qe_auto_setup=True, **kwargs):
         super().__init__(
             children=[
                 *self.code_children,
-                self.resources_config,
-                self.parallelization,
                 self.message_area,
                 self.sssp_installation_status,
                 self.qe_setup_status,
                 self._submission_blocker_messages,
                 self.submit_button,
             ]
         )
+        # set default codes
+        self.set_selected_codes(DEFAULT_PARAMETERS["codes"])
 
     @tl.observe("internal_submission_blockers", "external_submission_blockers")
     def _observe_submission_blockers(self, _change):
@@ -197,12 +186,10 @@ def _toggle_install_widgets(self, change):
 
     def _auto_select_code(self, change):
         if change["new"] and not change["old"]:
-            for name, code_widget in self.codes.items():
+            for name, code in self.codes.items():
                 try:
-                    code_widget.refresh()
-                    code_widget.value = orm.load_code(
-                        DEFAULT_PARAMETERS["codes"][name]
-                    ).uuid
+                    code.refresh()
+                    code.value = orm.load_code(DEFAULT_PARAMETERS["codes"][name]).uuid
                 except NotExistent:
                     pass
 
@@ -221,55 +208,6 @@ def _show_alert_message(self, message, alert_class="info"):
                 )
             )
 
-    def _update_resources(self, change):
-        if change["new"] and (
-            change["old"] is None
-            or orm.load_code(change["new"]).computer.pk
-            != orm.load_code(change["old"]).computer.pk
-        ):
-            self.set_resource_defaults(orm.load_code(change["new"]).computer)
-
-    def get_resources(self):
-        resources = {
-            "num_machines": self.resources_config.num_nodes.value,
-            "num_mpiprocs_per_machine": self.resources_config.num_cpus.value,
-            "npools": self.parallelization.npools.value,
-        }
-        return resources
-
-    def set_resources(self, resources):
-        self.resources_config.num_nodes.value = resources["num_machines"]
-        self.resources_config.num_cpus.value = resources["num_mpiprocs_per_machine"]
-        self.parallelization.npools.value = resources["npools"]
-
-    def set_resource_defaults(self, computer=None):
-        if computer is None or computer.hostname == "localhost":
-            self.resources_config.num_nodes.disabled = True
-            self.resources_config.num_nodes.value = 1
-            self.resources_config.num_cpus.max = os.cpu_count()
-            self.resources_config.num_cpus.value = 1
-            self.resources_config.num_cpus.description = "CPUs"
-            self.parallelization.npools.value = 1
-        else:
-            default_mpiprocs = computer.get_default_mpiprocs_per_machine()
-            self.resources_config.num_nodes.disabled = False
-            self.resources_config.num_cpus.max = default_mpiprocs
-            self.resources_config.num_cpus.value = default_mpiprocs
-            self.resources_config.num_cpus.description = "CPUs/node"
-            self.parallelization.npools.value = self._get_default_parallelization()
-
-        self._check_resources()
-
-    def _get_default_parallelization(self):
-        """A _very_ rudimentary approach for obtaining a minimal npools setting."""
-        num_mpiprocs = (
-            self.resources_config.num_nodes.value * self.resources_config.num_cpus.value
-        )
-
-        for i in range(1, num_mpiprocs + 1):
-            if num_mpiprocs % i == 0 and num_mpiprocs // i < self.MAX_MPI_PER_POOL:
-                return i
-
     def _check_resources(self):
         """Check whether the currently selected resources will be sufficient and warn if not."""
         if not self.pw_code.value:
@@ -328,10 +266,10 @@ def get_selected_codes(self):
 
         return: A dict with the code names as keys and the code UUIDs as values.
         """
-        codes = {key: code.value for key, code in self.codes.items()}
+        codes = {key: code.parameters for key, code in self.codes.items()}
         return codes
 
-    def set_selected_codes(self, codes):
+    def set_selected_codes(self, code_data):
         """Set the inputs in the GUI based on a set of codes."""
 
         # Codes
@@ -344,7 +282,13 @@ def _get_code_uuid(code):
 
         with self.hold_trait_notifications():
             for name, code in self.codes.items():
-                code.value = _get_code_uuid(codes.get(name))
+                # get code uuid from code label in case of using DEFAULT_PARAMETERS
+                if name not in code_data:
+                    continue
+                code_data.get(name)["code"] = _get_code_uuid(
+                    code_data.get(name)["code"]
+                )
+                code.parameters = code_data.get(name)
 
     def update_codes_display(self):
         """Hide code if no related property is selected."""
@@ -400,53 +344,49 @@ def _create_builder(self) -> ProcessBuilderNamespace:
         from copy import deepcopy
 
         self.ui_parameters = deepcopy(self.input_parameters)
-        self.ui_parameters["resources"] = self.get_resources()
         # add codes and resource info into ui_parameters
-        self.ui_parameters.update(self.get_submission_parameters())
+        submission_parameters = self.get_submission_parameters()
+        self.ui_parameters.update(submission_parameters)
         builder = QeAppWorkChain.get_builder_from_protocol(
             structure=self.input_structure,
             parameters=deepcopy(self.ui_parameters),
         )
 
-        self._update_builder(builder, self.MAX_MPI_PER_POOL)
+        self._update_builder(builder, submission_parameters["codes"])
 
         return builder
 
-    def _update_builder(self, buildy, max_mpi_per_pool):
-        resources = self.get_resources()
-        npools = resources.pop("npools", 1)
-        """Update the resources and parallelization of the ``QeAppWorkChain`` builder."""
-        for k, v in buildy.items():
-            if isinstance(v, (dict, ProcessBuilderNamespace)):
-                if k == "pw" and v["pseudos"]:
-                    v["parallelization"] = orm.Dict(dict={"npool": npools})
-                if k == "projwfc":
-                    v["settings"] = orm.Dict(dict={"cmdline": ["-nk", str(npools)]})
-                if k == "dos":
-                    v["metadata"]["options"]["resources"] = {
-                        "num_machines": 1,
-                        "num_mpiprocs_per_machine": min(
-                            max_mpi_per_pool,
-                            resources["num_mpiprocs_per_machine"],
-                        ),
-                    }
-                    # Continue to the next item to avoid overriding the resources in the
-                    # recursive `update_builder` call.
-                    continue
-                if k == "resources":
-                    buildy["resources"] = resources
-                else:
-                    self._update_builder(v, max_mpi_per_pool)
+    def _update_builder(self, builder, codes):
+        """Update the resources and parallelization of the ``relax`` builder."""
+        # update resources
+        builder.relax.base.pw.metadata.options.resources = {
+            "num_machines": codes.get("pw")["nodes"],
+            "num_mpiprocs_per_machine": codes.get("pw")["cpus"],
+        }
+        builder.relax.base.pw.parallelization = orm.Dict(
+            dict=codes["pw"]["parallelization"]
+        )
 
     def set_submission_parameters(self, parameters):
-        self.set_resources(parameters["resources"])
+        # backward compatibility for v2023.11
+        # which have a separate "resources" section for pw code
+        if "resources" in parameters:
+            parameters["codes"] = {
+                key: {"code": value} for key, value in parameters["codes"].items()
+            }
+            parameters["codes"]["pw"]["nodes"] = parameters["resources"]["num_machines"]
+            parameters["codes"]["pw"]["cpus"] = parameters["resources"][
+                "num_mpiprocs_per_machine"
+            ]
+            parameters["codes"]["pw"]["parallelization"] = {
+                "npool": parameters["resources"]["npools"]
+            }
         self.set_selected_codes(parameters["codes"])
 
     def get_submission_parameters(self):
         """Get the parameters for the submission step."""
         return {
             "codes": self.get_selected_codes(),
-            "resources": self.get_resources(),
         }
 
     def reset(self):
@@ -455,4 +395,3 @@ def reset(self):
             self.process = None
             self.input_structure = None
             self.set_selected_codes(DEFAULT_PARAMETERS["codes"])
-            self.set_resource_defaults()

src/aiidalab_qe/app/submission/resource.py

@@ -1,85 +0,0 @@
-# -*- coding: utf-8 -*-
-"""Widgets for the submission of bands work chains.
-
-Authors: AiiDAlab team
-"""
-import ipywidgets as ipw
-
-
-class ResourceSelectionWidget(ipw.VBox):
-    """Widget for the selection of compute resources."""
-
-    title = ipw.HTML(
-        """<div style="padding-top: 0px; padding-bottom: 0px">
-        <h4>Resources</h4>
-    </div>"""
-    )
-    prompt = ipw.HTML(
-        """<div style="line-height:120%; padding-top:0px">
-        <p style="padding-bottom:10px">
-        Specify the resources to use for the pw.x calculation.
-        </p></div>"""
-    )
-
-    def __init__(self, **kwargs):
-        extra = {
-            "style": {"description_width": "150px"},
-            "layout": {"min_width": "180px"},
-        }
-        self.num_nodes = ipw.BoundedIntText(
-            value=1, step=1, min=1, max=1000, description="Nodes", **extra
-        )
-        self.num_cpus = ipw.BoundedIntText(
-            value=1, step=1, min=1, description="CPUs", **extra
-        )
-
-        super().__init__(
-            children=[
-                self.title,
-                ipw.HBox(
-                    children=[self.prompt, self.num_nodes, self.num_cpus],
-                    layout=ipw.Layout(justify_content="space-between"),
-                ),
-            ]
-        )
-
-    def reset(self):
-        self.num_nodes.value = 1
-        self.num_cpus.value = 1
-
-
-class ParallelizationSettings(ipw.VBox):
-    """Widget for setting the parallelization settings."""
-
-    title = ipw.HTML(
-        """<div style="padding-top: 0px; padding-bottom: 0px">
-        <h4>Parallelization</h4>
-    </div>"""
-    )
-    prompt = ipw.HTML(
-        """<div style="line-height:120%; padding-top:0px">
-        <p style="padding-bottom:10px">
-        Specify the number of k-points pools for the calculations.
-        </p></div>"""
-    )
-
-    def __init__(self, **kwargs):
-        extra = {
-            "style": {"description_width": "150px"},
-            "layout": {"min_width": "180px"},
-        }
-        self.npools = ipw.BoundedIntText(
-            value=1, step=1, min=1, max=128, description="Number of k-pools", **extra
-        )
-        super().__init__(
-            children=[
-                self.title,
-                ipw.HBox(
-                    children=[self.prompt, self.npools],
-                    layout=ipw.Layout(justify_content="space-between"),
-                ),
-            ]
-        )
-
-    def reset(self):
-        self.npools.value = 1