diff --git a/src/aiidalab_qe/app/configuration/__init__.py b/src/aiidalab_qe/app/configuration/__init__.py index 93f24f74f..c71bc1757 100644 --- a/src/aiidalab_qe/app/configuration/__init__.py +++ b/src/aiidalab_qe/app/configuration/__init__.py @@ -63,7 +63,7 @@ def __init__(self, model: ConfigurationStepModel, **kwargs): self._model.observe( self._on_input_structure_change, - "input_structure", + "structure_uuid", ) self.settings = { @@ -254,7 +254,7 @@ def _fetch_plugin_calculation_settings(self): if identifier == "bands": ipw.dlink( - (self._model, "input_structure"), + (self._model, "structure_uuid"), (outline.include, "disabled"), lambda _: not self._model.has_pbc, ) diff --git a/src/aiidalab_qe/app/configuration/advanced/advanced.py b/src/aiidalab_qe/app/configuration/advanced/advanced.py index 6a3ac96a1..05088ca8b 100644 --- a/src/aiidalab_qe/app/configuration/advanced/advanced.py +++ b/src/aiidalab_qe/app/configuration/advanced/advanced.py @@ -79,7 +79,7 @@ def __init__(self, model: AdvancedConfigurationSettingsModel, **kwargs): self._model.observe( self._on_input_structure_change, - "input_structure", + "structure_uuid", ) self._model.observe( self._on_spin_type_change, @@ -150,7 +150,7 @@ def _update_tabs(self): self.advanced_tabs.selected_index = 0 def _on_advanced_tab_change(self, change): - tab: ConfigurationSettingsPanel = self.advanced_tabs.children[change["new"]] + tab: ConfigurationSettingsPanel = self.advanced_tabs.children[change["new"]] # type: ignore tab.render() def _on_reset_to_defaults_button_click(self, _): diff --git a/src/aiidalab_qe/app/configuration/advanced/convergence/convergence.py b/src/aiidalab_qe/app/configuration/advanced/convergence/convergence.py index 393d55b4f..910ee6799 100644 --- a/src/aiidalab_qe/app/configuration/advanced/convergence/convergence.py +++ b/src/aiidalab_qe/app/configuration/advanced/convergence/convergence.py @@ -13,8 +13,8 @@ def __init__(self, model: ConvergenceConfigurationSettingsModel, **kwargs): super().__init__(model, **kwargs) self._model.observe( - self._on_structure_change, - "input_structure", + self._on_input_structure_change, + "structure_uuid", ) self._model.observe( self._on_protocol_change, @@ -107,7 +107,7 @@ def render(self): (self.kpoints_distance, "value"), ) ipw.dlink( - (self._model, "input_structure"), + (self._model, "structure_uuid"), (self.kpoints_distance, "disabled"), lambda _: not self._model.has_pbc, ) @@ -239,7 +239,7 @@ def render(self): self.rendered = True - def _on_structure_change(self, _): + def _on_input_structure_change(self, _): self.refresh(specific="structure") def _on_protocol_change(self, _): diff --git a/src/aiidalab_qe/app/configuration/advanced/convergence/model.py b/src/aiidalab_qe/app/configuration/advanced/convergence/model.py index 64906612b..0a041887d 100644 --- a/src/aiidalab_qe/app/configuration/advanced/convergence/model.py +++ b/src/aiidalab_qe/app/configuration/advanced/convergence/model.py @@ -20,7 +20,7 @@ class ConvergenceConfigurationSettingsModel( identifier = "convergence" dependencies = [ - "input_structure", + "structure_uuid", "protocol", ] diff --git a/src/aiidalab_qe/app/configuration/advanced/hubbard/hubbard.py b/src/aiidalab_qe/app/configuration/advanced/hubbard/hubbard.py index ce8f3624d..e12519481 100644 --- a/src/aiidalab_qe/app/configuration/advanced/hubbard/hubbard.py +++ b/src/aiidalab_qe/app/configuration/advanced/hubbard/hubbard.py @@ -14,7 +14,7 @@ def __init__(self, model: HubbardConfigurationSettingsModel, **kwargs): self._model.observe( self._on_input_structure_change, - "input_structure", + "structure_uuid", ) self._model.observe( self._on_hubbard_activation, diff --git a/src/aiidalab_qe/app/configuration/advanced/hubbard/model.py b/src/aiidalab_qe/app/configuration/advanced/hubbard/model.py index 205993075..ed6682bda 100644 --- a/src/aiidalab_qe/app/configuration/advanced/hubbard/model.py +++ b/src/aiidalab_qe/app/configuration/advanced/hubbard/model.py @@ -19,7 +19,7 @@ class HubbardConfigurationSettingsModel( identifier = "hubbard" dependencies = [ - "input_structure", + "structure_uuid", ] is_active = tl.Bool(False) @@ -95,7 +95,7 @@ def set_active_eigenvalues(self, eigenvalues: list): self.has_eigenvalues = True def get_parameters_from_hubbard_structure(self): - hubbard_parameters = self.input_structure.hubbard.dict()["parameters"] + hubbard_parameters = self.input_structure.hubbard.model_dump()["parameters"] sites = self.input_structure.sites return { f"{sites[hp['atom_index']].kind_name} - {hp['atom_manifold']}": hp["value"] diff --git a/src/aiidalab_qe/app/configuration/advanced/magnetization/magnetization.py b/src/aiidalab_qe/app/configuration/advanced/magnetization/magnetization.py index 4e6d80940..a96825e70 100644 --- a/src/aiidalab_qe/app/configuration/advanced/magnetization/magnetization.py +++ b/src/aiidalab_qe/app/configuration/advanced/magnetization/magnetization.py @@ -27,7 +27,7 @@ def __init__(self, model: MagnetizationConfigurationSettingsModel, **kwargs): self._model.observe( self._on_input_structure_change, - "input_structure", + "structure_uuid", ) self._model.observe( self._on_electronic_type_change, @@ -151,7 +151,7 @@ def _build_kinds_widget(self): kind_names = ( self._model.input_structure.get_kind_names() - if self._model.input_structure + if self._model.has_structure else [] ) diff --git a/src/aiidalab_qe/app/configuration/advanced/magnetization/model.py b/src/aiidalab_qe/app/configuration/advanced/magnetization/model.py index bd76cc7ec..88241ff19 100644 --- a/src/aiidalab_qe/app/configuration/advanced/magnetization/model.py +++ b/src/aiidalab_qe/app/configuration/advanced/magnetization/model.py @@ -18,7 +18,7 @@ class MagnetizationConfigurationSettingsModel( identifier = "magnetization" dependencies = [ - "input_structure", + "structure_uuid", "electronic_type", "spin_type", "pseudos.dictionary", diff --git a/src/aiidalab_qe/app/configuration/advanced/model.py b/src/aiidalab_qe/app/configuration/advanced/model.py index 1af5ac9d7..e484d502d 100644 --- a/src/aiidalab_qe/app/configuration/advanced/model.py +++ b/src/aiidalab_qe/app/configuration/advanced/model.py @@ -26,7 +26,7 @@ class AdvancedConfigurationSettingsModel( identifier = "advanced" dependencies = [ - "input_structure", + "structure_uuid", "workchain.protocol", "workchain.spin_type", "workchain.electronic_type", @@ -233,7 +233,7 @@ def _set_pw_parameters(self, pw_parameters): control_params: dict = pw_parameters.get("CONTROL", {}) electron_params: dict = pw_parameters.get("ELECTRONS", {}) - num_atoms = len(self.input_structure.sites) if self.input_structure else 1 + num_atoms = len(self.input_structure.sites) if self.has_structure else 1 general = t.cast( GeneralConfigurationSettingsModel, diff --git a/src/aiidalab_qe/app/configuration/advanced/pseudos/model.py b/src/aiidalab_qe/app/configuration/advanced/pseudos/model.py index a16bfc2ac..a8239f573 100644 --- a/src/aiidalab_qe/app/configuration/advanced/pseudos/model.py +++ b/src/aiidalab_qe/app/configuration/advanced/pseudos/model.py @@ -31,7 +31,7 @@ class PseudosConfigurationSettingsModel( identifier = "pseudos" dependencies = [ - "input_structure", + "structure_uuid", "protocol", "spin_orbit", ] @@ -249,7 +249,7 @@ def update_cutoffs(self): if self.locked or not self.dictionary: return - kinds = self.input_structure.kinds if self.input_structure else [] + kinds = self.input_structure.kinds if self.has_structure else [] self.status_message = "" if self.family: diff --git a/src/aiidalab_qe/app/configuration/advanced/pseudos/pseudos.py b/src/aiidalab_qe/app/configuration/advanced/pseudos/pseudos.py index 37749326b..e48f90908 100644 --- a/src/aiidalab_qe/app/configuration/advanced/pseudos/pseudos.py +++ b/src/aiidalab_qe/app/configuration/advanced/pseudos/pseudos.py @@ -22,7 +22,7 @@ def __init__(self, model: PseudosConfigurationSettingsModel, **kwargs): self._model.observe( self._on_input_structure_change, - "input_structure", + "structure_uuid", ) self._model.observe( self._on_protocol_change, diff --git a/src/aiidalab_qe/app/configuration/basic/basic.py b/src/aiidalab_qe/app/configuration/basic/basic.py index e1d24c269..792257132 100644 --- a/src/aiidalab_qe/app/configuration/basic/basic.py +++ b/src/aiidalab_qe/app/configuration/basic/basic.py @@ -17,7 +17,7 @@ def __init__(self, model: BasicConfigurationSettingsModel, **kwargs): super().__init__(model, **kwargs) self._model.observe( self._on_input_structure_change, - "input_structure", + "structure_uuid", ) def render(self): diff --git a/src/aiidalab_qe/app/configuration/basic/model.py b/src/aiidalab_qe/app/configuration/basic/model.py index 0d8a91392..06a2b51b4 100644 --- a/src/aiidalab_qe/app/configuration/basic/model.py +++ b/src/aiidalab_qe/app/configuration/basic/model.py @@ -20,7 +20,7 @@ class BasicConfigurationSettingsModel( identifier = "workchain" dependencies = [ - "input_structure", + "structure_uuid", ] protocol_options = tl.List( diff --git a/src/aiidalab_qe/app/result/__init__.py b/src/aiidalab_qe/app/result/__init__.py index 480944510..96faef671 100644 --- a/src/aiidalab_qe/app/result/__init__.py +++ b/src/aiidalab_qe/app/result/__init__.py @@ -121,7 +121,7 @@ def _post_render(self): self.toggle_controls.value = ( "Results" - if (process := self._model.fetch_process_node()) and process.is_finished_ok + if self._model.has_process and self._model.process.is_finished_ok else "Status" ) @@ -152,10 +152,9 @@ def _on_state_change(self, change): def _on_previous_step_state_change(self, _): if self.previous_step_state is WizardAppWidgetStep.State.SUCCESS: - process_node = self._model.fetch_process_node() message = ( "Loading results" - if process_node and process_node.is_finished + if self._model.has_process and self._model.process.is_finished else "Submitting calculation" ) self.children = [LoadingWidget(message)] @@ -203,11 +202,10 @@ def _toggle_view(self, panel: ResultsComponent): def _update_kill_button_layout(self): if not self.rendered: return - process_node = self._model.fetch_process_node() if ( - not process_node - or process_node.is_finished - or process_node.is_excepted + not self._model.has_process + or self._model.process.is_finished + or self._model.process.is_excepted or self.state in ( self.State.SUCCESS, @@ -221,8 +219,7 @@ def _update_kill_button_layout(self): def _update_clean_scratch_button_layout(self): if not self.rendered: return - process_node = self._model.fetch_process_node() - if process_node and process_node.is_terminated: + if self._model.has_process and self._model.process.is_terminated: self.clean_scratch_button.layout.display = "block" else: self.clean_scratch_button.layout.display = "none" @@ -231,12 +228,12 @@ def _update_status(self): self._model.monitor_counter += 1 def _update_state(self): - if not (process_node := self._model.fetch_process_node()): + if not self._model.has_process: self.state = self.State.INIT self._update_controls() return - if process_state := process_node.process_state: + if process_state := self._model.process.process_state: status = self._get_process_status(process_state.value) else: status = "Unknown" @@ -254,9 +251,9 @@ def _update_state(self): ProcessState.KILLED, ): self.state = self.State.FAIL - elif process_node.is_failed: + elif self._model.process.is_failed: self.state = self.State.FAIL - elif process_node.is_finished_ok: + elif self._model.process.is_finished_ok: self.state = self.State.SUCCESS self._model.process_info = self.STATUS_TEMPLATE.format(status) diff --git a/src/aiidalab_qe/app/result/components/summary/model.py b/src/aiidalab_qe/app/result/components/summary/model.py index f83097047..10b9b41ce 100644 --- a/src/aiidalab_qe/app/result/components/summary/model.py +++ b/src/aiidalab_qe/app/result/components/summary/model.py @@ -124,16 +124,15 @@ def generate_report_text(self, report_dict): def generate_failure_report(self): """Generate a html for reporting the failure of the `QeAppWorkChain`.""" - process_node = self.fetch_process_node() - if not (process_node and process_node.exit_status): + if not (self.has_process and self.process.exit_status): return - final_calcjob = self._get_final_calcjob(process_node) + final_calcjob = self._get_final_calcjob(self.process) env = Environment() template = files(templates).joinpath("workflow_failure.jinja").read_text() style = files(styles).joinpath("style.css").read_text() self.failed_calculation_report = env.from_string(template).render( style=style, - process_report=get_workchain_report(process_node, "REPORT"), + process_report=get_workchain_report(self.process, "REPORT"), calcjob_exit_message=final_calcjob.exit_message, ) self.has_failure_report = True @@ -151,9 +150,7 @@ def _generate_report_parameters(self): """ from aiida.orm.utils.serialize import deserialize_unsafe - qeapp_wc = self.fetch_process_node() - - ui_parameters = qeapp_wc.base.extras.get("ui_parameters", {}) + ui_parameters = self.process.base.extras.get("ui_parameters", {}) if isinstance(ui_parameters, str): ui_parameters = deserialize_unsafe(ui_parameters) # Construct the report parameters needed for the report @@ -161,19 +158,20 @@ def _generate_report_parameters(self): if "workchain" not in ui_parameters: return {} - inputs = qeapp_wc.inputs + inputs = self.inputs + assert inputs.structure, "BUG! Missing structure input" # shouldn't happen! structure: orm.StructureData = inputs.structure basic = ui_parameters["workchain"] advanced = ui_parameters["advanced"] - ctime = qeapp_wc.ctime - mtime = qeapp_wc.mtime + ctime = self.process.ctime + mtime = self.process.mtime report = { "workflow_properties": { - "pk": qeapp_wc.pk, - "uuid": str(qeapp_wc.uuid), - "label": qeapp_wc.label, - "description": qeapp_wc.description, + "pk": self.process.pk, + "uuid": str(self.process.uuid), + "label": self.process.label, + "description": self.process.description, "creation_time": f"{format_time(ctime)} ({relative_time(ctime)})", "modification_time": f"{format_time(mtime)} ({relative_time(mtime)})", }, diff --git a/src/aiidalab_qe/app/result/components/summary/summary.py b/src/aiidalab_qe/app/result/components/summary/summary.py index 71028c0d1..28563320c 100644 --- a/src/aiidalab_qe/app/result/components/summary/summary.py +++ b/src/aiidalab_qe/app/result/components/summary/summary.py @@ -81,9 +81,8 @@ def _render_summary(self): self.has_settings_report = True def _render_download_widget(self): - process_node = self._model.fetch_process_node() - if process_node and process_node.is_terminated: - output_download_widget = WorkChainOutputs(node=process_node) + if self._model.has_process and self._model.process.is_terminated: + output_download_widget = WorkChainOutputs(node=self._model.process) output_download_widget.layout.width = "100%" self.output_download_container.children = [ self.output_download_help, diff --git a/src/aiidalab_qe/app/result/model.py b/src/aiidalab_qe/app/result/model.py index c6b6c5564..bc19a66e0 100644 --- a/src/aiidalab_qe/app/result/model.py +++ b/src/aiidalab_qe/app/result/model.py @@ -26,13 +26,13 @@ def update(self): self._update_process_remote_folder_state() def kill_process(self): - if process_node := self.fetch_process_node(): - control.kill_processes([process_node]) + if self.has_process: + control.kill_processes([self.process]) def clean_remote_data(self): - if not (process_node := self.fetch_process_node()): + if not self.has_process: return - for called_descendant in process_node.called_descendants: + for called_descendant in self.process.called_descendants: if isinstance(called_descendant, orm.CalcJobNode): with contextlib.suppress(Exception): called_descendant.outputs.remote_folder._clean() @@ -43,11 +43,10 @@ def reset(self): self.process_info = "" def _update_process_remote_folder_state(self): - process_node = self.fetch_process_node() - if not (process_node and process_node.called_descendants): + if not (self.has_process and self.process.called_descendants): return cleaned = [] - for called_descendant in process_node.called_descendants: + for called_descendant in self.process.called_descendants: if isinstance(called_descendant, orm.CalcJobNode): with contextlib.suppress(Exception): cleaned.append(called_descendant.outputs.remote_folder.is_empty) diff --git a/src/aiidalab_qe/app/structure/__init__.py b/src/aiidalab_qe/app/structure/__init__.py index 0c3750d9a..784533002 100644 --- a/src/aiidalab_qe/app/structure/__init__.py +++ b/src/aiidalab_qe/app/structure/__init__.py @@ -17,6 +17,7 @@ ShakeNBreakEditor, ) from aiidalab_qe.common.infobox import InAppGuide +from aiidalab_qe.common.setup_pseudos import PseudosInstallWidget from aiidalab_qe.common.widgets import CategorizedStructureExamplesWidget from aiidalab_qe.common.wizard import QeConfirmableWizardStep from aiidalab_widgets_base import ( @@ -66,7 +67,7 @@ def __init__(self, model: StructureStepModel, auto_setup=True, **kwargs): ) self._model.observe( self._on_input_structure_change, - "input_structure", + "structure_uuid", ) self._install_sssp(auto_setup) @@ -133,7 +134,8 @@ def _render(self): ipw.dlink( (self.manager, "structure_node"), - (self._model, "input_structure"), + (self._model, "structure_uuid"), + lambda node: node.uuid if node else None, ) ipw.link( (self._model, "manager_output"), @@ -198,8 +200,6 @@ def _on_input_structure_change(self, _): self._update_state() def _install_sssp(self, auto_setup): - from aiidalab_qe.common.setup_pseudos import PseudosInstallWidget - self.sssp_installation = PseudosInstallWidget(auto_start=False) ipw.dlink( (self.sssp_installation, "busy"), @@ -224,7 +224,7 @@ def _toggle_sssp_installation_widget(self): def _update_state(self): if self._model.confirmed: self.state = self.State.SUCCESS - elif self._model.input_structure is None: + elif not self._model.has_structure: self.state = self.State.READY else: self.state = self.State.CONFIGURED diff --git a/src/aiidalab_qe/app/structure/model.py b/src/aiidalab_qe/app/structure/model.py index 25b264d19..b2b846a7d 100644 --- a/src/aiidalab_qe/app/structure/model.py +++ b/src/aiidalab_qe/app/structure/model.py @@ -31,7 +31,7 @@ def update_widget_text(self): self.structure_name = str(self.input_structure.get_formula()) def reset(self): - self.input_structure = None + self.structure_uuid = None self.structure_name = "" self.manager_output = "" diff --git a/src/aiidalab_qe/app/submission/__init__.py b/src/aiidalab_qe/app/submission/__init__.py index f58627a0a..c3e1a4b1a 100644 --- a/src/aiidalab_qe/app/submission/__init__.py +++ b/src/aiidalab_qe/app/submission/__init__.py @@ -73,6 +73,10 @@ def __init__(self, model: SubmissionStepModel, auto_setup=True, **kwargs): self._on_input_parameters_change, "input_parameters", ) + self._model.observe( + self._on_process_change, + "process_uuid", + ) self.settings = { "global": self.global_resources, @@ -210,6 +214,9 @@ def _on_qe_installed(self, _): if self._model.qe_installed: self._model.update() + def _on_process_change(self, _): + self._model.update_process_metadata() + def _set_up_qe(self, auto_setup): self.qe_setup = QESetupWidget(auto_start=False) ipw.dlink( diff --git a/src/aiidalab_qe/app/submission/global_settings/model.py b/src/aiidalab_qe/app/submission/global_settings/model.py index 3c5f82f1b..d7dae2212 100644 --- a/src/aiidalab_qe/app/submission/global_settings/model.py +++ b/src/aiidalab_qe/app/submission/global_settings/model.py @@ -20,7 +20,7 @@ class GlobalResourceSettingsModel( identifier = "global" dependencies = [ - "input_structure", + "structure_uuid", "input_parameters", ] diff --git a/src/aiidalab_qe/app/submission/global_settings/setting.py b/src/aiidalab_qe/app/submission/global_settings/setting.py index 8955e97ba..05b32e4f8 100644 --- a/src/aiidalab_qe/app/submission/global_settings/setting.py +++ b/src/aiidalab_qe/app/submission/global_settings/setting.py @@ -20,7 +20,7 @@ def __init__(self, model: GlobalResourceSettingsModel, **kwargs): self._model.observe( self._on_input_structure_change, - "input_structure", + "structure_uuid", ) self._model.observe( self._on_input_parameters_change, diff --git a/src/aiidalab_qe/app/submission/model.py b/src/aiidalab_qe/app/submission/model.py index b70a50447..f10c38b9c 100644 --- a/src/aiidalab_qe/app/submission/model.py +++ b/src/aiidalab_qe/app/submission/model.py @@ -8,7 +8,7 @@ from aiida.engine import ProcessBuilderNamespace, submit from aiida.orm.utils.serialize import serialize from aiidalab_qe.app.parameters import DEFAULT_PARAMETERS -from aiidalab_qe.common.mixins import HasInputStructure, HasModels +from aiidalab_qe.common.mixins import HasInputStructure, HasModels, HasProcess from aiidalab_qe.common.panel import PluginResourceSettingsModel, ResourceSettingsModel from aiidalab_qe.common.wizard import QeConfirmableWizardStepModel from aiidalab_qe.utils import shallow_copy_nested_dict @@ -21,12 +21,12 @@ class SubmissionStepModel( QeConfirmableWizardStepModel, HasModels[ResourceSettingsModel], HasInputStructure, + HasProcess, ): identifier = "submission" input_parameters = tl.Dict() - process_node = tl.Instance(orm.WorkChainNode, allow_none=True) process_label = tl.Unicode("") process_description = tl.Unicode("") @@ -62,7 +62,7 @@ def __init__(self, *args, **kwargs): def confirm(self): super().confirm() - if not self.process_node: + if not self.has_process: self._submit() def update(self): @@ -70,7 +70,7 @@ def update(self): model.update() def update_process_label(self): - if not self.input_structure: + if not self.has_structure: self.process_label = "" return structure_label = ( @@ -165,16 +165,17 @@ def set_model_state(self, state: dict): model.set_model_state(codes[identifier]) model.locked = True - if self.process_node: - self.process_label = self.process_node.label - self.process_description = self.process_node.description - self.locked = True + def update_process_metadata(self): + if not self.has_process: + return + self.process_label = self.process.label + self.process_description = self.process.description + self.locked = True def reset(self): with self.hold_trait_notifications(): - self.input_structure = None self.input_parameters = {} - self.process_node = None + self.process_uuid = None for identifier, model in self.get_models(): if identifier not in self._default_models: model.include = False @@ -197,13 +198,10 @@ def _submit(self): "structure", self.input_structure.get_formula(), ) - self.process_node = process_node - - self._update_url() + self.process_uuid = process_node.uuid - def _update_url(self): - pk = self.process_node.pk - display(Javascript(f"window.history.pushState(null, '', '?pk={pk}');")) + pk = process_node.pk + display(Javascript(f"window.history.pushState(null, '', '?pk={pk}');")) def _link_model(self, model: ResourceSettingsModel): for dependency in model.dependencies: diff --git a/src/aiidalab_qe/app/wizard_app.py b/src/aiidalab_qe/app/wizard_app.py index 106355464..7a56bf93c 100644 --- a/src/aiidalab_qe/app/wizard_app.py +++ b/src/aiidalab_qe/app/wizard_app.py @@ -145,23 +145,22 @@ def _render_step(self, step_index): def _update_configuration_step(self): if self.structure_model.confirmed: - self.configure_model.input_structure = self.structure_model.input_structure + self.configure_model.structure_uuid = self.structure_model.structure_uuid else: - self.configure_model.input_structure = None + self.configure_model.structure_uuid = None def _update_submission_step(self): if self.configure_model.confirmed: - self.submit_model.input_structure = self.structure_model.input_structure + self.submit_model.structure_uuid = self.structure_model.structure_uuid self.submit_model.input_parameters = self.configure_model.get_model_state() else: - self.submit_model.input_structure = None + self.submit_model.structure_uuid = None self.submit_model.input_parameters = {} def _update_results_step(self): ipw.dlink( - (self.submit_model, "process_node"), + (self.submit_model, "process_uuid"), (self.results_model, "process_uuid"), - lambda node: node.uuid if node is not None else None, ) def _lock_app(self): @@ -179,14 +178,14 @@ def _update_from_process(self, pk): else: self._show_process_loading_message() process_node = load_node(pk) - self.structure_model.input_structure = process_node.inputs.structure + self.structure_model.structure_uuid = process_node.inputs.structure.uuid self.structure_model.confirm() parameters = process_node.base.extras.get("ui_parameters", {}) if parameters and isinstance(parameters, str): parameters = deserialize_unsafe(parameters) self.configure_model.set_model_state(parameters) self.configure_model.confirm() - self.submit_model.process_node = process_node + self.submit_model.process_uuid = process_node.uuid self.submit_model.set_model_state(parameters) self.submit_model.confirm() self._wizard_app_widget.selected_index = 3 diff --git a/src/aiidalab_qe/common/mixins.py b/src/aiidalab_qe/common/mixins.py index 124666ddc..74c1a9da9 100644 --- a/src/aiidalab_qe/common/mixins.py +++ b/src/aiidalab_qe/common/mixins.py @@ -9,15 +9,20 @@ from aiida_quantumespresso.data.hubbard_structure import HubbardStructureData from aiidalab_qe.common.mvc import Model +StructureType = t.Union[orm.StructureData, HubbardStructureData] + class HasInputStructure(tl.HasTraits): - input_structure = tl.Union( - [ - tl.Instance(orm.StructureData), - tl.Instance(HubbardStructureData), - ], - allow_none=True, - ) + structure_uuid = tl.Unicode(None, allow_none=True) + + @property + def input_structure(self) -> StructureType | None: + if not self.structure_uuid: + return None + try: + return t.cast(StructureType, orm.load_node(self.structure_uuid)) + except NotExistent: + return None @property def has_structure(self): @@ -29,7 +34,7 @@ def has_pbc(self): @property def has_tags(self): - return any( + return self.has_structure and any( not kind_name.isalpha() for kind_name in self.input_structure.get_kind_names() ) @@ -99,36 +104,36 @@ class HasProcess(tl.HasTraits): process_uuid = tl.Unicode(None, allow_none=True) monitor_counter = tl.Int(0) # used for continuous updates + @property + def process(self) -> orm.WorkChainNode | None: + if not self.process_uuid: + return None + try: + return t.cast(orm.WorkChainNode, orm.load_node(self.process_uuid)) + except NotExistent: + return None + @property def has_process(self): - return self.fetch_process_node() is not None + return self.process is not None @property - def inputs(self): - process_node = self.fetch_process_node() - return process_node.inputs if process_node else [] + def inputs(self) -> orm.NodeLinksManager | list: + return self.process.inputs if self.has_process else [] @property - def properties(self): - process_node = self.fetch_process_node() + def properties(self) -> list: # read the attributes directly instead of using the `get_list` method # to avoid error in case of the orm.List object being converted to a orm.Data object return ( - process_node.inputs.properties.base.attributes.get("list") - if process_node + self.inputs.properties.base.attributes.get("list") + if self.has_process else [] ) @property def outputs(self): - process_node = self.fetch_process_node() - return process_node.outputs if process_node else [] - - def fetch_process_node(self) -> orm.ProcessNode | None: - try: - return orm.load_node(self.process_uuid) if self.process_uuid else None # type: ignore - except NotExistent: - return None + return self.process.outputs if self.has_process else [] class Confirmable(tl.HasTraits): diff --git a/src/aiidalab_qe/common/panel.py b/src/aiidalab_qe/common/panel.py index 7ed6f862b..a32839422 100644 --- a/src/aiidalab_qe/common/panel.py +++ b/src/aiidalab_qe/common/panel.py @@ -526,22 +526,20 @@ def fetch_child_process_node(self, which="this") -> orm.ProcessNode | None: uuid = getattr(self, f"_{which}_process_uuid") label = getattr(self, f"_{which}_process_label") if not uuid: - root = self.fetch_process_node() + root = self.process child = next((c for c in root.called if c.process_label == label), None) uuid = child.uuid if child else None return orm.load_node(uuid) if uuid else None # type: ignore def save_state(self): """Saves the current state of the model to the AiiDA database.""" - node = self.fetch_process_node() - results = node.base.extras.get("results", {}) + results = self.process.base.extras.get("results", {}) results[self.identifier] = self.get_model_state() - node.base.extras.set("results", results) + self.process.base.extras.set("results", results) def load_state(self): """Loads the state of the model from the AiiDA database.""" - node = self.fetch_process_node() - results = node.base.extras.get("results", {}) + results = self.process.base.extras.get("results", {}) if self.identifier in results: self.set_model_state(results[self.identifier]) diff --git a/src/aiidalab_qe/common/process/tree.py b/src/aiidalab_qe/common/process/tree.py index 98632c678..4ed0b9eef 100644 --- a/src/aiidalab_qe/common/process/tree.py +++ b/src/aiidalab_qe/common/process/tree.py @@ -1,6 +1,5 @@ from __future__ import annotations -import threading import typing as t from copy import deepcopy @@ -8,6 +7,7 @@ import traitlets as tl from aiida import orm +from aiida.common.exceptions import NotExistent from aiida.common.links import LinkType from aiida.engine import ProcessState from aiida.tools.graph.graph_traversers import traverse_graph @@ -102,9 +102,10 @@ def _render(self): ) self.collapse_button.on_click(self._collapse_all) - root = self._model.fetch_process_node() - - self.trunk = WorkChainTreeNode(node=root, on_inspect=self._on_inspect) + self.trunk = WorkChainTreeNode( + node=self._model.process, + on_inspect=self._on_inspect, + ) self.trunk.add_class("tree-trunk") self.trunk.initialize() self.trunk.expand() @@ -157,18 +158,19 @@ def __init__( on_inspect: t.Callable[[str], None] | None = None, **kwargs, ): - self.uuid = node.uuid + self.process_uuid = node.uuid self.level = level self.on_inspect = on_inspect super().__init__(**kwargs) - self._node: dict[int, ProcessNodeType] = {} # thread_id: node @property - def node(self) -> ProcessNodeType: - if (tid := threading.get_ident()) not in self._node: - node = orm.load_node(self.uuid) - self._node[tid] = node - return self._node[tid] + def process(self) -> ProcessNodeType | None: + if not self.process_uuid: + return None + try: + return t.cast(ProcessNodeType, orm.load_node(self.process_uuid)) + except NotExistent: + return None def initialize(self): self._build_header() @@ -193,11 +195,11 @@ def _get_emoji(self, state): return STATE_ICONS.get(state, "❓") def _get_state(self): - if not hasattr(self.node, "process_state"): + if not hasattr(self.process, "process_state"): return "queued" - if self.node.is_failed: + if self.process.is_failed: return "failed" - state = self.node.process_state + state = self.process.process_state return ( "running" if state is ProcessState.WAITING @@ -207,7 +209,7 @@ def _get_state(self): ) def _get_human_readable_title(self): - node = self.node + node = self.process if not (label := node.process_label): return "Unknown" if label in ("PwBaseWorkChain", "PwCalculation"): @@ -249,12 +251,12 @@ def metadata_inputs(self): # summary to extract pw parameters. #1163 removes this dependency, thus allowing # the popping of the "relax" port if not running relaxation. However, this is # kept for backwards compatibility, for jobs ran before #1163. - inputs = deepcopy(self.node.get_metadata_inputs()) or {} - if "properties" in self.node.inputs: + inputs = deepcopy(self.process.get_metadata_inputs()) or {} + if "properties" in self.process.inputs: inputs = { key: value for key, value in inputs.items() - if key in self.node.inputs.properties.base.attributes.get("list") + if key in self.process.inputs.properties.base.attributes.get("list") } return inputs @@ -332,9 +334,9 @@ def _add_branches(self): self._add_branches_recursive() self._adding_branches = False - def _add_branches_recursive(self, node: orm.ProcessNode | None = None): - node = node or self.node - for child in sorted(node.called, key=lambda child: child.ctime): + def _add_branches_recursive(self, process: orm.ProcessNode | None = None): + process = process or self.process + for child in sorted(process.called, key=lambda child: child.ctime): if isinstance(child, orm.CalcFunctionNode): continue if child.pk in self.pks: @@ -373,7 +375,7 @@ def _get_tally(self): def _count_finished(self): traverser = traverse_graph( - starting_pks=[self.node.pk], + starting_pks=[self.process.pk], links_forward=[ LinkType.CALL_WORK, LinkType.CALL_CALC, @@ -450,4 +452,4 @@ def _build_header(self): def _on_label_click(self, _): if self.on_inspect is None: return - self.on_inspect(self.uuid) + self.on_inspect(self.process_uuid) diff --git a/src/aiidalab_qe/plugins/electronic_structure/result/result.py b/src/aiidalab_qe/plugins/electronic_structure/result/result.py index 8eb10e322..c51456b1d 100644 --- a/src/aiidalab_qe/plugins/electronic_structure/result/result.py +++ b/src/aiidalab_qe/plugins/electronic_structure/result/result.py @@ -100,7 +100,7 @@ def _render_bands_pdos_widget(self, node_identifiers): identifier: self._model.fetch_child_process_node(identifier) for identifier in node_identifiers }, - "root": self._model.fetch_process_node(), + "root": self._model.process, } model = BandsPdosModel.from_nodes(**nodes) widget = BandsPdosWidget(model=model) diff --git a/src/aiidalab_qe/plugins/pdos/model.py b/src/aiidalab_qe/plugins/pdos/model.py index 9f5b08ea0..5de6bdb9d 100644 --- a/src/aiidalab_qe/plugins/pdos/model.py +++ b/src/aiidalab_qe/plugins/pdos/model.py @@ -14,7 +14,7 @@ class PdosConfigurationSettingsModel(PanelModel, HasInputStructure): identifier = "pdos" dependencies = [ - "input_structure", + "structure_uuid", "workchain.protocol", ] diff --git a/src/aiidalab_qe/plugins/pdos/setting.py b/src/aiidalab_qe/plugins/pdos/setting.py index ce3e1e243..c05000541 100644 --- a/src/aiidalab_qe/plugins/pdos/setting.py +++ b/src/aiidalab_qe/plugins/pdos/setting.py @@ -18,7 +18,7 @@ def __init__(self, model: PdosConfigurationSettingsModel, **kwargs): self._model.observe( self._on_input_structure_change, - "input_structure", + "structure_uuid", ) self._model.observe( self._on_protocol_change, @@ -44,7 +44,7 @@ def render(self): (self.nscf_kpoints_distance, "value"), ) ipw.dlink( - (self._model, "input_structure"), + (self._model, "structure_uuid"), (self.nscf_kpoints_distance, "disabled"), lambda _: not self._model.has_pbc, ) @@ -65,7 +65,7 @@ def render(self): (self.use_pdos_degauss, "value"), ) ipw.dlink( - (self._model, "input_structure"), + (self._model, "structure_uuid"), (self.use_pdos_degauss, "disabled"), lambda _: not self._model.has_pbc, ) diff --git a/tests/configuration/test_advanced.py b/tests/configuration/test_advanced.py index 0d888e95b..b42db6900 100644 --- a/tests/configuration/test_advanced.py +++ b/tests/configuration/test_advanced.py @@ -59,7 +59,7 @@ def test_advanced_convergence_settings(generate_structure_data): _ = ConvergenceConfigurationSettingsPanel(model=model) # Test structure-dependent convergence change - model.input_structure = generate_structure_data("silica") + model.structure_uuid = generate_structure_data("silica").uuid assert "num_atoms = 6" in model.help_message @@ -106,7 +106,7 @@ def test_advanced_kpoints_mesh(generate_structure_data): _ = ConvergenceConfigurationSettingsPanel(model=model) structure = generate_structure_data(name="silicon") - model.input_structure = structure + model.structure_uuid = structure.uuid assert model.mesh_grid == "Mesh [14, 14, 14]" @@ -122,7 +122,7 @@ def test_advanced_molecule_settings(generate_structure_data): # Create molecule structure = generate_structure_data(name="H2O", pbc=(False, False, False)) - model.input_structure = structure + model.structure_uuid = structure.uuid # Confirm the value of kpoints_distance is fixed assert model.kpoints_distance == 100.0 @@ -175,7 +175,7 @@ def test_advanced_hubbard_settings(generate_structure_data): hubbard.render() structure = generate_structure_data(name="LiCoO2") - model.input_structure = structure + model.structure_uuid = structure.uuid # Activate Hubbard U widget model.is_active = True @@ -244,7 +244,7 @@ def test_advanced_magnetic_settings(generate_structure_data): pseudo_family = get_pseudo_family_by_label("SSSP/1.3/PBE/efficiency") structure = generate_structure_data(name="LiCoO2") - model.input_structure = structure + model.structure_uuid = structure.uuid model.spin_type = "collinear" model.dictionary = { kind.name: pseudo_family.get_pseudo(kind.symbol).uuid @@ -282,8 +282,9 @@ def test_advanced_magnetic_settings(generate_structure_data): symbols="O", name="O2", ) + structure.store() - model.input_structure = structure + model.structure_uuid = structure.uuid model.dictionary = { kind.name: pseudo_family.get_pseudo(kind.symbol).uuid for kind in structure.kinds diff --git a/tests/conftest.py b/tests/conftest.py index 931bad0a8..a93a9f649 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -82,7 +82,7 @@ def fixture_localhost(aiida_localhost): def generate_structure_data(): """generate a `StructureData` object.""" - def _generate_structure_data(name="silicon", pbc=(True, True, True)): + def _generate_structure_data(name="silicon", pbc=(True, True, True), store=True): if name == "silicon": structure = orm.StructureData( cell=[ @@ -157,6 +157,9 @@ def _generate_structure_data(name="silicon", pbc=(True, True, True)): structure.pbc = pbc + if store: + structure.store() + return structure return _generate_structure_data @@ -420,7 +423,7 @@ def _submit_app_generator( initial_magnetic_moments=0.0, electron_maxstep=80, ): - app.structure_model.input_structure = generate_structure_data() + app.structure_model.structure_uuid = generate_structure_data().uuid app.structure_model.confirm() parameters = { @@ -472,7 +475,7 @@ def _submit_app_generator( @pytest.fixture def app_to_submit(app: WizardApp, generate_structure_data): # Step 1: select structure from example - app.structure_model.input_structure = generate_structure_data() + app.structure_model.structure_uuid = generate_structure_data().uuid app.structure_model.confirm() # Step 2: configure calculation # TODO do we need to include bands and pdos here? @@ -711,14 +714,8 @@ def _generate_qeapp_workchain( structure = generate_structure_data() else: structure.store() - # TODO is this necessary? - # aiida_database_wrapper = app.structure_step.manager.children[0].children[2] # type: ignore - # aiida_database_wrapper.render() - # aiida_database = aiida_database_wrapper.children[0] # type: ignore - # aiida_database.search() - # aiida_database.results.value = structure - - app.structure_model.input_structure = structure + + app.structure_model.structure_uuid = structure.uuid app.structure_model.confirm() # step 2 configure diff --git a/tests/test_app.py b/tests/test_app.py index 843bf37ba..434b0619f 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -28,8 +28,8 @@ def test_selecting_new_structure_unconfirms_model(generate_structure_data): from aiidalab_qe.app.structure.model import StructureStepModel model = StructureStepModel() - model.input_structure = generate_structure_data() + model.structure_uuid = generate_structure_data().uuid assert model.has_structure model.confirm() - model.input_structure = generate_structure_data() + model.structure_uuid = generate_structure_data().uuid assert not model.confirmed diff --git a/tests/test_pseudo.py b/tests/test_pseudo.py index 4d18ae1df..47f42e315 100644 --- a/tests/test_pseudo.py +++ b/tests/test_pseudo.py @@ -154,7 +154,7 @@ def test_pseudos_settings(generate_structure_data, generate_upf_data): pseudos = PseudosConfigurationSettingsPanel(model=model) silicon = generate_structure_data("silicon") - model.input_structure = silicon + model.structure_uuid = silicon.uuid # Test the default family assert model.family == f"SSSP/{SSSP_VERSION}/PBEsol/efficiency" @@ -184,7 +184,7 @@ def test_pseudos_settings(generate_structure_data, generate_upf_data): # Test that changing the structure triggers a reset silica = generate_structure_data("silica") - model.input_structure = silica + model.structure_uuid = silica.uuid assert model.family == f"SSSP/{SSSP_VERSION}/PBEsol/efficiency" assert "Si" in model.dictionary.keys() assert "O" in model.dictionary.keys() @@ -355,7 +355,7 @@ def test_missing_pseudos(generate_structure_data): """Test that the model handles missing pseudos correctly.""" model = PseudosConfigurationSettingsModel() _ = PseudosConfigurationSettingsPanel(model=model) - model.input_structure = generate_structure_data("CeO") + model.structure_uuid = generate_structure_data("CeO").uuid model.functional = "PBEsol" model.library = "PseudoDojo standard (SR)" assert model.family == "PseudoDojo/0.4/PBEsol/SR/standard/upf" @@ -369,7 +369,7 @@ def test_functional_mismatch_blocker(generate_structure_data): """Test blocker for inconsistent functional across selected pseudopotentials.""" model = PseudosConfigurationSettingsModel() _ = PseudosConfigurationSettingsPanel(model=model) - model.input_structure = generate_structure_data("silica") + model.structure_uuid = generate_structure_data("silica").uuid model.functionals = ["PBE", "PBEsol"] assert len(model.blockers) == 1 assert "must have the same exchange-correlation" in model.blockers[0] @@ -381,7 +381,7 @@ def test_relativistic_mismatch_blocker(generate_structure_data): """ model = PseudosConfigurationSettingsModel() _ = PseudosConfigurationSettingsPanel(model=model) - model.input_structure = generate_structure_data("silica") + model.structure_uuid = generate_structure_data("silica").uuid model.spin_orbit = "soc" model.family = "SSSP/1.3/PBEsol/efficiency" assert len(model.blockers) == 1 diff --git a/tests/test_result.py b/tests/test_result.py index 5757a882a..8f7a09b52 100644 --- a/tests/test_result.py +++ b/tests/test_result.py @@ -150,12 +150,11 @@ def test_table_data(model): wc = generate_qeapp_workchain(relax_type="none") model.process_uuid = wc.node.uuid - node = model.fetch_process_node() assert "Si2" in model.header assert "Initial" in model.sub_header assert "properties" in model.source # inputs - assert model.structure.pk == node.inputs.structure.pk - assert str(node.inputs.structure.pk) in model.info + assert model.structure.pk == model.inputs.structure.pk + assert str(model.inputs.structure.pk) in model.info test_table_data(model) panel.render() @@ -163,7 +162,6 @@ def test_table_data(model): wc = generate_qeapp_workchain(relax_type="positions_cell") model.process_uuid = wc.node.uuid - node = model.fetch_process_node() assert "Initial" in model.sub_header assert panel.view_toggle_button.layout.display == "block" assert panel.view_toggle_button.description == "View relaxed" @@ -171,6 +169,6 @@ def test_table_data(model): assert panel.view_toggle_button.description == "View initial" assert "Relaxed" in model.sub_header assert "properties" not in model.source # outputs - assert model.structure.pk == node.outputs.structure.pk - assert str(node.outputs.structure.pk) in model.info + assert model.structure.pk == model.outputs.structure.pk + assert str(model.outputs.structure.pk) in model.info test_table_data(model) diff --git a/tests/test_status.py b/tests/test_status.py index 613356597..b45cd5d79 100644 --- a/tests/test_status.py +++ b/tests/test_status.py @@ -237,7 +237,7 @@ def update_monitor(): def delayed_add_branches_recursive( node: orm.ProcessNode | None = None, ): - node = node or self.tree.trunk.node + node = node or self.tree.trunk.process for child in sorted(node.called, key=lambda child: child.ctime): if isinstance(child, orm.CalcFunctionNode): continue @@ -306,8 +306,8 @@ def test_calcjob_node_link(self): trunk = self.tree.trunk trunk.expand(recursive=True) self.calcjob_node.label.click() - assert self.panel.process_tree.value == self.calcjob_node.uuid - assert self.panel.process_tree._tree.nodes == (self.calcjob_node.node,) + assert self.panel.process_tree.value == self.calcjob_node.process_uuid + assert self.panel.process_tree._tree.nodes == (self.calcjob_node.process,) # TODO understand why the following does not trigger automatically as in the app # TODO understand why the following triggers a thread # self.panel.process_tree.set_trait("selected_nodes", [self.calcjob_node.node]) diff --git a/tests/test_submit_qe_workchain.py b/tests/test_submit_qe_workchain.py index f6ccd1082..26592b0ab 100644 --- a/tests/test_submit_qe_workchain.py +++ b/tests/test_submit_qe_workchain.py @@ -167,7 +167,7 @@ def test_warning_messages( # now we use a large structure, so we should have the Warning-1 (and 2 if not on localhost) structure = generate_structure_data("H2O-larger") - submit_model.input_structure = structure + submit_model.structure_uuid = structure.uuid pw_code.num_cpus = 1 global_model.check_resources() num_sites = len(structure.sites)