Avoid storing process nodes - use UUID

Author: edan-bainglass

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)

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,29 +150,28 @@ 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
         # drop support for old ui parameters
         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)})",
             },

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,

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)

src/aiidalab_qe/app/submission/model.py

@@ -5,10 +5,11 @@
 from IPython.display import Javascript, display
 
 from aiida import orm
+from aiida.common.exceptions import NotExistent
 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 +22,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("")
 
@@ -60,7 +61,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):
@@ -163,15 +164,22 @@ 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._update_process_metadata()
+
+    def _update_process_metadata(self):
+        if not self.has_process:
+            return
+        try:
+            self.process_label = self.process.label
+            self.process_description = self.process.description
             self.locked = True
+        except NotExistent:
+            return
 
     def reset(self):
         with self.hold_trait_notifications():
             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
@@ -194,13 +202,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:

src/aiidalab_qe/app/wizard_app.py

@@ -159,9 +159,8 @@ def _update_submission_step(self):
 
     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):
@@ -186,7 +185,7 @@ def _update_from_process(self, pk):
                 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

src/aiidalab_qe/common/mixins.py

@@ -106,36 +106,37 @@ class HasProcess(tl.HasTraits):
     process_uuid = tl.Unicode(None, allow_none=True)
     monitor_counter = tl.Int(0)  # used for continuous updates
 
+    @property
+    @cache_per_thread(invalidator="process_uuid")
+    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):

src/aiidalab_qe/common/panel.py

@@ -529,22 +529,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])
 

src/aiidalab_qe/common/process/tree.py

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

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)