diff --git a/docs/source/development/plugin.rst b/docs/source/development/plugin.rst index 06a6b3254..abce2df78 100644 --- a/docs/source/development/plugin.rst +++ b/docs/source/development/plugin.rst @@ -225,7 +225,7 @@ The ``get_builder`` function will return a ``builder`` for the ``EOSWorkChain``, def get_builder(codes, structure, parameters, **kwargs): protocol = parameters["workchain"].pop('protocol', "fast") - pw_code = codes.get("pw") + pw_code = codes.get("pw")['code'] overrides = { "pw": parameters["advanced"], } @@ -334,9 +334,9 @@ Here is the example of the built-in `pdos` plugins with codes `dos.x` and `projw .. code-block:: python - from aiidalab_widgets_base import ComputationalResourcesWidget + from aiidalab_qe.common.widgets import QEAppComputationalResourcesWidget - dos_code = ComputationalResourcesWidget( + dos_code = QEAppComputationalResourcesWidget( description="dos.x", default_calc_job_plugin="quantumespresso.dos", ) diff --git a/src/aiidalab_qe/app/parameters/qeapp.yaml b/src/aiidalab_qe/app/parameters/qeapp.yaml index e553edf28..c00c0f1c1 100644 --- a/src/aiidalab_qe/app/parameters/qeapp.yaml +++ b/src/aiidalab_qe/app/parameters/qeapp.yaml @@ -23,11 +23,17 @@ advanced: tot_charge: 0 vdw_corr: none -## Codes +## Computational resources codes: - dos: dos-7.2@localhost - projwfc: projwfc-7.2@localhost - pw: pw-7.2@localhost - pp: pp-7.2@localhost - xspectra: xspectra-7.2@localhost - hp: hp-7.2@localhost + dos: + code: dos-7.2@localhost + projwfc: + code: projwfc-7.2@localhost + pw: + code: pw-7.2@localhost + pp: + code: pp-7.2@localhost + xspectra: + code: xspectra-7.2@localhost + hp: + code: hp-7.2@localhost diff --git a/src/aiidalab_qe/app/submission/__init__.py b/src/aiidalab_qe/app/submission/__init__.py index 9f38741d2..eec0f36b3 100644 --- a/src/aiidalab_qe/app/submission/__init__.py +++ b/src/aiidalab_qe/app/submission/__init__.py @@ -6,24 +6,24 @@ 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 ( + PwCodeResourceSetupWidget, + QEAppComputationalResourcesWidget, +) from aiidalab_qe.workflows import QeAppWorkChain -from .resource import ParallelizationSettings, ResourceSelectionWidget - class SubmitQeAppWorkChainStep(ipw.VBox, WizardAppWidgetStep): """Step for submission of a bands workchain.""" @@ -69,17 +69,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 = PwCodeResourceSetupWidget( 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} @@ -94,8 +88,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"]) # set process label and description self.process_label = ipw.Text( description="Label:", layout=ipw.Layout(width="auto", indent="0px") @@ -138,8 +130,6 @@ 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, @@ -150,6 +140,8 @@ def __init__(self, qe_auto_setup=True, **kwargs): 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): @@ -183,6 +175,16 @@ def _identify_submission_blockers(self): if not self.sssp_installation_status.installed: yield "The SSSP library is not installed." + # check if the QEAppComputationalResourcesWidget is used + for name, code in self.codes.items(): + # skip if the code is not displayed, convenient for the plugin developer + if code.layout.display == "none": + continue + if not isinstance(code, QEAppComputationalResourcesWidget): + yield ( + f"Error: hi, plugin developer, please use the QEAppComputationalResourcesWidget from aiidalab_qe.common.widgets for code {name}." + ) + def _update_state(self, _=None): # If the previous step has failed, this should fail as well. if self.previous_step_state is self.State.FAIL: @@ -215,14 +217,12 @@ 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(): if not DEFAULT_PARAMETERS["codes"].get(name): continue try: - code_widget.refresh() - code_widget.value = orm.load_code( - DEFAULT_PARAMETERS["codes"][name] - ).uuid + code.code_selection.refresh() + code.value = orm.load_code(DEFAULT_PARAMETERS["codes"][name]).uuid except NotExistent: pass @@ -241,55 +241,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: @@ -349,10 +300,14 @@ 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() + if code.layout.display != "none" + } 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 @@ -365,12 +320,20 @@ def _get_code_uuid(code): with self.hold_trait_notifications(): for name, code in self.codes.items(): + if name not in code_data: + continue # check if the code is installed and usable # note: if code is imported from another user, it is not usable and thus will not be # treated as an option in the ComputationalResourcesWidget. - code_options = [o[1] for o in self.pw_code.code_select_dropdown.options] - if _get_code_uuid(codes.get(name)) in code_options: - code.value = _get_code_uuid(codes.get(name)) + code_options = [ + o[1] for o in code.code_selection.code_select_dropdown.options + ] + if _get_code_uuid(code_data.get(name)["code"]) in code_options: + # get code uuid from code label in case of using DEFAULT_PARAMETERS + 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.""" @@ -432,46 +395,44 @@ 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")["ntasks_per_node"], + "num_cores_per_mpiproc": codes.get("pw")["cpus_per_task"], + } + 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"]) # label and description are not stored in the parameters, but in the process directly if self.process: @@ -482,7 +443,6 @@ 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): @@ -491,4 +451,3 @@ def reset(self): self.process = None self.input_structure = None self.set_selected_codes(DEFAULT_PARAMETERS["codes"]) - self.set_resource_defaults() diff --git a/src/aiidalab_qe/app/submission/resource.py b/src/aiidalab_qe/app/submission/resource.py deleted file mode 100644 index a12bc173c..000000000 --- a/src/aiidalab_qe/app/submission/resource.py +++ /dev/null @@ -1,86 +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( - """
-

Resources

-
""" - ) - prompt = ipw.HTML( - """
-

- Specify the resources to use for the pw.x calculation. -

""" - ) - - 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( - """
-

Parallelization

-
""" - ) - prompt = ipw.HTML( - """
-

- Specify the number of k-points pools for the calculations. -

""" - ) - - 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 diff --git a/src/aiidalab_qe/common/widgets.py b/src/aiidalab_qe/common/widgets.py index 91b1b515a..2d9b08f62 100644 --- a/src/aiidalab_qe/common/widgets.py +++ b/src/aiidalab_qe/common/widgets.py @@ -17,7 +17,8 @@ import traitlets from aiida.orm import CalcJobNode from aiida.orm import Data as orm_Data -from aiida.orm import load_node +from aiida.orm import load_code, load_node +from aiidalab_widgets_base import ComputationalResourcesWidget from aiidalab_widgets_base.utils import ( StatusHTML, list_to_string_range, @@ -620,6 +621,259 @@ def _select_periodicity(self, _=None): self.structure = deepcopy(new_structure) +class QEAppComputationalResourcesWidget(ipw.VBox): + value = traitlets.Unicode(allow_none=True) + nodes = traitlets.Int(default_value=1) + cpus = traitlets.Int(default_value=1) + + def __init__(self, **kwargs): + """Widget to setup the compute resources, which include the code, + the number of nodes and the number of cpus. + """ + self.code_selection = ComputationalResourcesWidget(**kwargs) + self.code_selection.layout.width = "80%" + + self.num_nodes = ipw.BoundedIntText( + value=1, step=1, min=1, max=1000, description="Nodes", width="10%" + ) + self.num_cpus = ipw.BoundedIntText( + value=1, step=1, min=1, description="CPUs", width="10%" + ) + self.btn_setup_resource_detail = ipw.ToggleButton(description="More") + self.btn_setup_resource_detail.observe(self._setup_resource_detail, "value") + self._setup_resource_detail_output = ipw.Output(layout={"width": "500px"}) + + # combine code, nodes and cpus + children = [ + ipw.HBox( + children=[ + self.code_selection, + self.num_nodes, + self.num_cpus, + self.btn_setup_resource_detail, + ] + ), + self._setup_resource_detail_output, + ] + super().__init__(children=children, **kwargs) + + self.resource_detail = ResourceDetailSettings() + traitlets.dlink( + (self.num_cpus, "value"), (self.resource_detail.ntasks_per_node, "value") + ) + traitlets.link((self.code_selection, "value"), (self, "value")) + + @traitlets.observe("value") + def _update_resources(self, change): + if change["new"]: + self.set_resource_defaults(load_code(change["new"]).computer) + + def set_resource_defaults(self, computer=None): + import os + + if computer is None or computer.hostname == "localhost": + self.num_nodes.disabled = True + self.num_nodes.value = 1 + self.num_cpus.max = os.cpu_count() + self.num_cpus.value = 1 + self.num_cpus.description = "CPUs" + else: + default_mpiprocs = computer.get_default_mpiprocs_per_machine() + self.num_nodes.disabled = False + self.num_cpus.max = default_mpiprocs + self.num_cpus.value = default_mpiprocs + self.num_cpus.description = "CPUs" + + @property + def parameters(self): + return self.get_parameters() + + def get_parameters(self): + """Return the parameters.""" + parameters = { + "code": self.code_selection.value, + "nodes": self.num_nodes.value, + "cpus": self.num_cpus.value, + } + parameters.update(self.resource_detail.parameters) + return parameters + + @parameters.setter + def parameters(self, parameters): + self.set_parameters(parameters) + + def set_parameters(self, parameters): + """Set the parameters.""" + self.code_selection.value = parameters["code"] + if "nodes" in parameters: + self.num_nodes.value = parameters["nodes"] + if "cpus" in parameters: + self.num_cpus.value = parameters["cpus"] + if "ntasks_per_node" in parameters: + self.resource_detail.ntasks_per_node.value = parameters["ntasks_per_node"] + if "cpus_per_task" in parameters: + self.resource_detail.cpus_per_task.value = parameters["cpus_per_task"] + + def _setup_resource_detail(self, _=None): + with self._setup_resource_detail_output: + clear_output() + if self.btn_setup_resource_detail.value: + self._setup_resource_detail_output.layout = { + "width": "500px", + "border": "1px solid gray", + } + + children = [ + self.resource_detail, + ] + display(*children) + else: + self._setup_resource_detail_output.layout = { + "width": "500px", + "border": "none", + } + + +class ResourceDetailSettings(ipw.VBox): + """Widget for setting the Resource detail.""" + + prompt = ipw.HTML( + """
+

+ Specify the parameters for the scheduler (only for advanced user). +

""" + ) + + def __init__(self, **kwargs): + self.ntasks_per_node = ipw.BoundedIntText( + value=1, + step=1, + min=1, + max=1000, + description="ntasks-per-node", + style={"description_width": "100px"}, + ) + self.cpus_per_task = ipw.BoundedIntText( + value=1, + step=1, + min=1, + description="cpus-per-task", + style={"description_width": "100px"}, + ) + super().__init__( + children=[self.prompt, self.ntasks_per_node, self.cpus_per_task], **kwargs + ) + + @property + def parameters(self): + return self.get_parameters() + + def get_parameters(self): + """Return the parameters.""" + return { + "ntasks_per_node": self.ntasks_per_node.value, + "cpus_per_task": self.cpus_per_task.value, + } + + @parameters.setter + def parameters(self, parameters): + self.ntasks_per_node.value = parameters.get("ntasks_per_node", 1) + self.cpus_per_task.value = parameters.get("cpus_per_task", 1) + + def reset(self): + """Reset the settings.""" + self.ntasks_per_node.value = 1 + self.cpus_per_task.value = 1 + + +class ParallelizationSettings(ipw.VBox): + """Widget for setting the parallelization settings.""" + + prompt = ipw.HTML( + """
+

+ Specify the number of k-points pools for the pw.x calculations (only for advanced user). +

""" + ) + + def __init__(self, **kwargs): + extra = { + "style": {"description_width": "150px"}, + "layout": {"min_width": "180px"}, + } + self.npool = ipw.BoundedIntText( + value=1, step=1, min=1, max=128, description="Number of k-pools", **extra + ) + self.override = ipw.Checkbox( + escription="", + indent=False, + value=False, + layout=ipw.Layout(max_width="20px"), + ) + self.override.observe(self.set_visibility, "value") + super().__init__( + children=[ + ipw.HBox( + children=[self.override, self.prompt, self.npool], + layout=ipw.Layout(justify_content="flex-start"), + ), + ] + ) + # set the default visibility of the widget + self.npool.layout.display = "none" + + def set_visibility(self, change): + if change["new"]: + self.npool.layout.display = "block" + else: + self.npool.layout.display = "none" + + def reset(self): + """Reset the parallelization settings.""" + self.npool.value = 1 + + +class PwCodeResourceSetupWidget(QEAppComputationalResourcesWidget): + """ComputationalResources Widget for the pw.x calculation.""" + + nodes = traitlets.Int(default_value=1) + + def __init__(self, **kwargs): + # By definition, npool must be a divisor of the total number of k-points + # thus we can not set a default value here, or from the computer. + self.parallelization = ParallelizationSettings() + super().__init__(**kwargs) + # add nodes and cpus into the children of the widget + self.children += (self.parallelization,) + + def get_parallelization(self): + """Return the parallelization settings.""" + parallelization = ( + {"npool": self.parallelization.npool.value} + if self.parallelization.override.value + else {} + ) + return parallelization + + def set_parallelization(self, parallelization): + """Set the parallelization settings.""" + if "npool" in parallelization: + self.parallelization.override.value = True + self.parallelization.npool.value = parallelization["npool"] + + def get_parameters(self): + """Return the parameters.""" + parameters = super().get_parameters() + parameters.update({"parallelization": self.get_parallelization()}) + return parameters + + def set_parameters(self, parameters): + """Set the parameters.""" + super().set_parameters(parameters) + if "parallelization" in parameters: + self.set_parallelization(parameters["parallelization"]) + + class HubbardWidget(ipw.VBox): """Widget for setting up Hubbard parameters.""" diff --git a/src/aiidalab_qe/plugins/bands/workchain.py b/src/aiidalab_qe/plugins/bands/workchain.py index 56594b2b7..efa234028 100644 --- a/src/aiidalab_qe/plugins/bands/workchain.py +++ b/src/aiidalab_qe/plugins/bands/workchain.py @@ -1,4 +1,5 @@ import numpy as np +from aiida import orm from aiida.plugins import DataFactory, WorkflowFactory from aiida_quantumespresso.common.types import ElectronicType, SpinType @@ -171,11 +172,26 @@ def generate_kpath_2d(structure, kpoints_distance, kpath_2d): return kpoints +def update_resources(builder, codes): + builder.scf.pw.metadata.options.resources = { + "num_machines": codes.get("pw")["nodes"], + "num_mpiprocs_per_machine": codes.get("pw")["ntasks_per_node"], + "num_cores_per_mpiproc": codes.get("pw")["cpus_per_task"], + } + builder.scf.pw.parallelization = orm.Dict(dict=codes["pw"]["parallelization"]) + builder.bands.pw.metadata.options.resources = { + "num_machines": codes.get("pw")["nodes"], + "num_mpiprocs_per_machine": codes.get("pw")["ntasks_per_node"], + "num_cores_per_mpiproc": codes.get("pw")["cpus_per_task"], + } + builder.bands.pw.parallelization = orm.Dict(dict=codes["pw"]["parallelization"]) + + def get_builder(codes, structure, parameters, **kwargs): """Get a builder for the PwBandsWorkChain.""" from copy import deepcopy - pw_code = codes.get("pw") + pw_code = codes.get("pw")["code"] protocol = parameters["workchain"]["protocol"] scf_overrides = deepcopy(parameters["advanced"]) relax_overrides = { @@ -217,6 +233,8 @@ def get_builder(codes, structure, parameters, **kwargs): bands.pop("relax") bands.pop("structure", None) bands.pop("clean_workdir", None) + # update resources + update_resources(bands, codes) if scf_overrides["pw"]["parameters"]["SYSTEM"].get("tot_magnetization") is not None: bands.scf["pw"]["parameters"]["SYSTEM"].pop("starting_magnetization", None) diff --git a/src/aiidalab_qe/plugins/pdos/__init__.py b/src/aiidalab_qe/plugins/pdos/__init__.py index e84e4f31f..8bc3fbd54 100644 --- a/src/aiidalab_qe/plugins/pdos/__init__.py +++ b/src/aiidalab_qe/plugins/pdos/__init__.py @@ -1,6 +1,5 @@ -from aiidalab_widgets_base import ComputationalResourcesWidget - from aiidalab_qe.common.panel import OutlinePanel +from aiidalab_qe.common.widgets import QEAppComputationalResourcesWidget from .result import Result from .setting import Setting @@ -12,12 +11,12 @@ class PdosOutline(OutlinePanel): help = """""" -dos_code = ComputationalResourcesWidget( +dos_code = QEAppComputationalResourcesWidget( description="dos.x", default_calc_job_plugin="quantumespresso.dos", ) -projwfc_code = ComputationalResourcesWidget( +projwfc_code = QEAppComputationalResourcesWidget( description="projwfc.x", default_calc_job_plugin="quantumespresso.projwfc", ) diff --git a/src/aiidalab_qe/plugins/pdos/workchain.py b/src/aiidalab_qe/plugins/pdos/workchain.py index 2d5fd3e40..18ddb5a49 100644 --- a/src/aiidalab_qe/plugins/pdos/workchain.py +++ b/src/aiidalab_qe/plugins/pdos/workchain.py @@ -1,3 +1,4 @@ +from aiida import orm from aiida.plugins import WorkflowFactory from aiida_quantumespresso.common.types import ElectronicType, SpinType @@ -31,12 +32,40 @@ def check_codes(pw_code, dos_code, projwfc_code): ) +def update_resources(builder, codes): + builder.scf.pw.metadata.options.resources = { + "num_machines": codes.get("pw")["nodes"], + "num_mpiprocs_per_machine": codes.get("pw")["ntasks_per_node"], + "num_cores_per_mpiproc": codes.get("pw")["cpus_per_task"], + } + builder.scf.pw.parallelization = orm.Dict(dict=codes["pw"]["parallelization"]) + builder.nscf.pw.metadata.options.resources = { + "num_machines": codes.get("pw")["nodes"], + "num_mpiprocs_per_machine": codes.get("pw")["ntasks_per_node"], + "num_cores_per_mpiproc": codes.get("pw")["cpus_per_task"], + } + builder.nscf.pw.parallelization = orm.Dict(dict=codes["pw"]["parallelization"]) + builder.dos.metadata.options.resources = { + "num_machines": codes.get("dos")["nodes"], + "num_mpiprocs_per_machine": codes.get("dos")["ntasks_per_node"], + "num_cores_per_mpiproc": codes.get("dos")["cpus_per_task"], + } + builder.projwfc.metadata.options.resources = { + "num_machines": codes.get("projwfc")["nodes"], + "num_mpiprocs_per_machine": codes.get("projwfc")["ntasks_per_node"], + "num_cores_per_mpiproc": codes.get("projwfc")["cpus_per_task"], + } + # disable the parallelization setting for projwfc + # npool = codes["pw"]["parallelization"]["npool"] + # builder.projwfc.settings = orm.Dict(dict={"cmdline": ["-nk", str(npool)]}) + + def get_builder(codes, structure, parameters, **kwargs): from copy import deepcopy - pw_code = codes.get("pw") - dos_code = codes.get("dos") - projwfc_code = codes.get("projwfc") + pw_code = codes.get("pw")["code"] + dos_code = codes.get("dos")["code"] + projwfc_code = codes.get("projwfc")["code"] check_codes(pw_code, dos_code, projwfc_code) protocol = parameters["workchain"]["protocol"] @@ -66,6 +95,8 @@ def get_builder(codes, structure, parameters, **kwargs): # pop the inputs that are exclueded from the expose_inputs pdos.pop("structure", None) pdos.pop("clean_workdir", None) + # update resources + update_resources(pdos, codes) if ( scf_overrides["pw"]["parameters"]["SYSTEM"].get("tot_magnetization") diff --git a/src/aiidalab_qe/plugins/xas/__init__.py b/src/aiidalab_qe/plugins/xas/__init__.py index 1ca6f4473..3d6debb89 100644 --- a/src/aiidalab_qe/plugins/xas/__init__.py +++ b/src/aiidalab_qe/plugins/xas/__init__.py @@ -1,7 +1,7 @@ from importlib import resources import yaml -from aiidalab_widgets_base import ComputationalResourcesWidget +from aiidalab_qe.common.widgets import QEAppComputationalResourcesWidget from aiidalab_qe.common.panel import OutlinePanel from aiidalab_qe.plugins import xas as xas_folder @@ -18,7 +18,7 @@ class XasOutline(OutlinePanel): help = """""" -xs_code = ComputationalResourcesWidget( +xs_code = QEAppComputationalResourcesWidget( description="xspectra.x", default_calc_job_plugin="quantumespresso.xspectra" ) diff --git a/src/aiidalab_qe/workflows/__init__.py b/src/aiidalab_qe/workflows/__init__.py index 2ae51269d..e26a24c7f 100644 --- a/src/aiidalab_qe/workflows/__init__.py +++ b/src/aiidalab_qe/workflows/__init__.py @@ -112,11 +112,10 @@ def get_builder_from_protocol( parameters = parameters or {} properties = parameters["workchain"].pop("properties", []) codes = parameters.pop("codes", {}) - codes = { - key: orm.load_node(value) - for key, value in codes.items() - if value is not None - } + # load codes from uuid + for _, value in codes.items(): + if value["code"] is not None: + value["code"] = orm.load_node(value["code"]) # update pseudos for kind, uuid in parameters["advanced"]["pw"]["pseudos"].items(): parameters["advanced"]["pw"]["pseudos"][kind] = orm.load_node(uuid) @@ -147,7 +146,7 @@ def get_builder_from_protocol( } protocol = parameters["workchain"]["protocol"] relax_builder = PwRelaxWorkChain.get_builder_from_protocol( - code=codes.get("pw"), + code=codes.get("pw")["code"], structure=structure, protocol=protocol, relax_type=RelaxType(parameters["workchain"]["relax_type"]), diff --git a/tests/conftest.py b/tests/conftest.py index 768415c43..100ebdcfd 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -327,9 +327,9 @@ def app(pw_code, dos_code, projwfc_code): app.submit_step.sssp_installation_status.installed = True # set up codes - app.submit_step.pw_code.refresh() - app.submit_step.codes["dos"].refresh() - app.submit_step.codes["projwfc"].refresh() + app.submit_step.pw_code.code_selection.refresh() + app.submit_step.codes["dos"].code_selection.refresh() + app.submit_step.codes["projwfc"].code_selection.refresh() app.submit_step.pw_code.value = pw_code.uuid app.submit_step.codes["dos"].value = dos_code.uuid @@ -387,7 +387,7 @@ def _submit_app_generator( # submit_step = app.submit_step submit_step.input_structure = generate_structure_data() - submit_step.resources_config.num_cpus.value = 2 + submit_step.pw_code.num_cpus.value = 2 return app @@ -675,7 +675,7 @@ def _generate_qeapp_workchain( s2.confirm() # step 3 setup code and resources s3: SubmitQeAppWorkChainStep = app.submit_step - s3.resources_config.num_cpus.value = 4 + s3.pw_code.num_cpus.value = 4 builder = s3._create_builder() inputs = builder._inputs() inputs["relax"]["base_final_scf"] = deepcopy(inputs["relax"]["base"]) diff --git a/tests/test_app.py b/tests/test_app.py index 10e804f44..a3c80f166 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -44,7 +44,7 @@ def test_reload_and_reset(submit_app_generator, generate_qeapp_workchain): ) == 0 ) - assert app.submit_step.resources_config.num_cpus.value == 1 + assert app.submit_step.pw_code.num_cpus.value == 4 def test_select_new_structure(app_to_submit, generate_structure_data): diff --git a/tests/test_codes.py b/tests/test_codes.py index 37f46dd4c..f351df525 100644 --- a/tests/test_codes.py +++ b/tests/test_codes.py @@ -59,3 +59,22 @@ def test_identify_submission_blockers(app): submit.codes["dos"].value = dos_value blockers = list(submit._identify_submission_blockers()) assert len(blockers) == 0 + + +def test_qeapp_computational_resources_widget(): + """Test QEAppComputationalResourcesWidget.""" + from aiidalab_qe.app.submission import SubmitQeAppWorkChainStep + + new_submit_step = SubmitQeAppWorkChainStep(qe_auto_setup=False) + assert new_submit_step.codes["pw"].parallelization.npool.layout.display == "none" + new_submit_step.codes["pw"].parallelization.override.value = True + new_submit_step.codes["pw"].parallelization.npool.value = 2 + assert new_submit_step.codes["pw"].parallelization.npool.layout.display == "block" + assert new_submit_step.codes["pw"].parameters == { + "code": None, + "cpus": 1, + "cpus_per_task": 1, + "nodes": 1, + "ntasks_per_node": 1, + "parallelization": {"npool": 2}, + } diff --git a/tests/test_submit_qe_workchain/test_create_builder_default.yml b/tests/test_submit_qe_workchain/test_create_builder_default.yml index 2f9f23575..5a4add9d7 100644 --- a/tests/test_submit_qe_workchain/test_create_builder_default.yml +++ b/tests/test_submit_qe_workchain/test_create_builder_default.yml @@ -16,13 +16,24 @@ advanced: bands: kpath_2d: hexagonal codes: - xspectra: null + dos: + cpus: 1 + cpus_per_task: 1 + nodes: 1 + ntasks_per_node: 1 + projwfc: + cpus: 1 + cpus_per_task: 1 + nodes: 1 + ntasks_per_node: 1 + pw: + cpus: 2 + cpus_per_task: 1 + nodes: 1 + ntasks_per_node: 2 + parallelization: {} pdos: nscf_kpoints_distance: 0.1 -resources: - npools: 1 - num_machines: 1 - num_mpiprocs_per_machine: 2 workchain: electronic_type: metal properties: