Improve simplified process tree (#1174)

This PR does the following:
- Refactor tree classes to aiidalab_qe/common/process module
- Add missing "not rendered" guards for monitor-triggered events
- Improve tree styles, particularly height (max 500px)
- Improve performance of counting finished jobs by replacing recursion with single DB query
- Use a flag to signal when adding branches to avoid race conditions with monitoring thread

src/aiidalab_qe/app/result/components/status/status.py

@@ -3,11 +3,14 @@
 import ipywidgets as ipw
 
 from aiidalab_qe.app.result.components import ResultsComponent
+from aiidalab_qe.common.process.tree import (
+    SimplifiedProcessTree,
+    SimplifiedProcessTreeModel,
+)
 from aiidalab_widgets_base import ProcessNodesTreeWidget
 from aiidalab_widgets_base.viewers import AiidaNodeViewWidget
 
 from .model import WorkChainStatusModel
-from .tree import SimplifiedProcessTree, SimplifiedProcessTreeModel
 
 
 class WorkChainStatusPanel(ResultsComponent[WorkChainStatusModel]):
@@ -115,7 +118,8 @@ def _post_render(self):
         self._select_tree_root()
 
     def _on_monitor_counter_change(self, _):
-        self.process_tree.update()
+        if self.rendered:
+            self.process_tree.update()
 
     def _on_accordion_change(self, change):
         if change["new"] == 0:

src/aiidalab_qe/app/result/components/summary/summary.py

@@ -18,6 +18,8 @@ def _on_process_change(self, _):
             self._render_summary()
 
     def _on_monitor_counter_change(self, _):
+        if not self.rendered:
+            return
         if not self.has_download_widget:
             self._render_download_widget()
         if not self._model.has_failure_report:

src/aiidalab_qe/app/static/styles/status.css

@@ -1,18 +1,20 @@
 .simplified-view {
   flex-direction: row;
+  height: 500px;
 }
 
 .simplified-view > div:first-child {
-  width: 50%;
+  width: 45%;
 }
 
 .simplified-view > div:last-child {
-  width: 70%;
+  width: 55%;
 }
 
 @media (max-width: 1199px) {
   .simplified-view {
     flex-direction: column;
+    height: auto;
   }
 
   .simplified-view > div:first-child,
@@ -25,12 +27,14 @@
   min-height: 28px;
 }
 
+.simplified-view .aiida-node-view-widget {
+  flex: 1;
+}
+
 .filename-display {
   min-height: 32px;
 }
 
-/* Process tree */
-
 .rolling-output {
   flex: 1;
   min-height: 300px;
@@ -45,8 +49,20 @@
   position: relative !important;
 }
 
+/* Process tree */
+
 .simplified-process-tree {
   height: 100%;
+  max-height: 500px;
+}
+
+.simplified-process-tree .tree-container {
+  flex: 1;
+  overflow: auto;
+}
+
+.simplified-process-tree .tree-container .widget-box {
+  overflow: unset;
 }
 
 .simplified-process-tree .tree-node-header {
@@ -115,11 +131,10 @@
 
 .simplified-process-tree .tree-node-branches {
   height: 0;
-  overflow: hidden;
-  visibility: hidden;
+  display: none;
 }
 
 .simplified-process-tree .tree-node-branches.open {
   height: auto;
-  visibility: visible;
+  display: block;
 }

src/aiidalab_qe/common/process/__init__.py

@@ -0,0 +1,6 @@
+from .process import QeAppWorkChainSelector, WorkChainSelector
+
+__all__ = [
+    "QeAppWorkChainSelector",
+    "WorkChainSelector",
+]

src/aiidalab_qe/app/result/components/status/tree.py → src/aiidalab_qe/common/process/tree.py

@@ -7,7 +7,9 @@
 import traitlets as tl
 
 from aiida import orm
+from aiida.common.links import LinkType
 from aiida.engine import ProcessState
+from aiida.tools.graph.graph_traversers import traverse_graph
 from aiidalab_qe.common.mixins import HasProcess
 from aiidalab_qe.common.mvc import Model
 from aiidalab_qe.common.widgets import LoadingWidget
@@ -50,7 +52,8 @@ def _on_process_change(self, _):
         self._update()
 
     def _on_monitor_counter_change(self, _):
-        self._update()
+        if self.rendered:
+            self._update()
 
     def _on_inspect(self, uuid: str):
         self._model.clicked = None  # ensure event is triggered when label is reclicked
@@ -68,16 +71,23 @@ 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.layout.flex = "1"
+        self.trunk.add_class("tree-trunk")
         self.trunk.initialize()
         self.trunk.expand()
-        self.rendered = True
         self._update()
+
+        self.tree_container = ipw.VBox(
+            children=[self.trunk],
+        )
+        self.tree_container.add_class("tree-container")
+
         self.children = [
             self.collapse_button,
-            self.trunk,
+            self.tree_container,
             ipw.HBox(
                 children=[
                     ipw.HTML(
@@ -97,9 +107,10 @@ def _render(self):
             ),
         ]
 
+        self.rendered = True
+
     def _update(self):
-        if self.rendered:
-            self.trunk.update()
+        self.trunk.update()
 
     def _collapse_all(self, _):
         self.trunk.collapse(recursive=True)
@@ -207,6 +218,9 @@ def _get_human_readable_title(self):
 
 
 class ProcessTreeBranches(ipw.VBox):
+    def clear(self):
+        self.children = []
+
     def __len__(self):
         return len(self.children)
 
@@ -224,6 +238,8 @@ def __iadd__(self, other: ProcessTreeNode):
 
 
 class WorkChainTreeNode(ProcessTreeNode[orm.WorkChainNode]):
+    _adding_branches = False
+
     @property
     def metadata_inputs(self):
         # BACKWARDS COMPATIBILITY: originally this was added because "relax" was in the
@@ -269,6 +285,10 @@ def update(self):
         for branch in self.branches:
             branch.update()
 
+    def clear(self):
+        self.branches.clear()
+        self.pks.clear()
+
     def expand(self, recursive=False):
         if self.collapsed:
             self.toggle.click()
@@ -303,57 +323,72 @@ def _build_header(self):
             self.tally,
         ]
 
-    def _add_branches(self, node: orm.ProcessNode | None = None):
+    def _add_branches(self):
+        if self._adding_branches:
+            return
+        self._adding_branches = True
+        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):
-            if child.pk in self.pks or isinstance(child, orm.CalcFunctionNode):
+            if isinstance(child, orm.CalcFunctionNode):
+                continue
+            if child.pk in self.pks:
                 continue
-            TreeNodeClass = (
-                WorkChainTreeNode
-                if isinstance(child, orm.WorkChainNode)
-                else CalcJobTreeNode
-            )
-            branch = TreeNodeClass(
-                node=child,
-                level=self.level + 1,
-                on_inspect=self.on_inspect,
-            )
             if child.process_label in (
                 "BandsWorkChain",
                 "ProjwfcBaseWorkChain",
             ):
-                self._add_branches(child)
+                self._add_branches_recursive(child)
             else:
-                branch.initialize()
-                self.branches += branch
-                self.pks.add(child.pk)
+                self._add_branch(child)
+
+    def _add_branch(self, child: orm.ProcessNode):
+        TreeNodeClass = (
+            WorkChainTreeNode
+            if isinstance(child, orm.WorkChainNode)
+            else CalcJobTreeNode
+        )
+        branch = TreeNodeClass(
+            node=child,
+            level=self.level + 1,
+            on_inspect=self.on_inspect,
+        )
+        branch.initialize()
+        self.branches += branch
+        self.pks.add(child.pk)
 
     def _get_tally(self):
         total = self.expected_jobs["count"]
         dynamic = self.expected_jobs["dynamic"]
-        finished = self._count_finished(self.node)
+        finished = self._count_finished()
         tally = f"{finished}/{total}"
         tally += "*" if dynamic else ""
         tally += " job" if total == 1 else " jobs"
         return tally
 
-    def _get_current_total(self, node: orm.ProcessNode):
-        total = 0
-        for child in node.called:
-            if isinstance(child, orm.WorkChainNode):
-                total += self._get_current_total(child)
-            elif isinstance(child, orm.CalcJobNode):
-                total += 1
-        return total
-
-    def _count_finished(self, node: orm.ProcessNode):
-        count = 0
-        for child in node.called:
-            if isinstance(child, orm.WorkChainNode):
-                count += self._count_finished(child)
-            elif isinstance(child, orm.CalcJobNode) and child.is_finished_ok:
-                count += 1
-        return count
+    def _count_finished(self):
+        traverser = traverse_graph(
+            starting_pks=[self.node.pk],
+            links_forward=[
+                LinkType.CALL_WORK,
+                LinkType.CALL_CALC,
+            ],
+        )
+        return (
+            orm.QueryBuilder()
+            .append(
+                orm.CalcJobNode,
+                filters=(
+                    (orm.CalcJobNode.fields.pk.in_(traverser.get("nodes")))
+                    & (orm.CalcJobNode.fields.process_state == "finished")
+                    & (orm.CalcJobNode.fields.exit_status == 0)
+                ),
+            )
+            .count()
+        )
 
     def _get_expected(self, inputs: dict[str, dict]) -> dict:
         expected = {}

tests/test_status.py

@@ -1,3 +1,5 @@
+import threading
+
 import ipywidgets as ipw
 import pytest
 
@@ -8,7 +10,7 @@
     WorkChainStatusModel,
     WorkChainStatusPanel,
 )
-from aiidalab_qe.app.result.components.status.tree import (
+from aiidalab_qe.common.process.tree import (
     CalcJobTreeNode,
     ProcessTreeNode,
     SimplifiedProcessTree,
@@ -216,6 +218,44 @@ def test_calcjob_node(self):
         assert isinstance(self.calcjob_node.label, ipw.Button)
         assert self.calcjob_node.label.description == human_label
 
+    @pytest.mark.parametrize("without_flag", [True, False])
+    def test_tree_monitor_updates(self, without_flag):
+        self.tree.trunk.collapse()
+        self.tree.trunk.clear()
+        assert not self.tree.trunk.branches
+
+        def update_monitor():
+            for _ in range(10):
+                self.model.monitor_counter += 1
+
+        if without_flag:
+            # Skip `_adding_branches` flag check
+            add_branches = self.tree.trunk._add_branches  # store original method
+            self.tree.trunk._add_branches = self.tree.trunk._add_branches_recursive
+
+        # Start monitoring thread to update the tree, which will also
+        # attempt to add branches to expanded parents
+        monitor_thread = threading.Thread(target=update_monitor)
+        monitor_thread.start()
+
+        # Expand the trunk in the main thread
+        self.tree.trunk.toggle.click()
+
+        # Wait for the background thread to finish before asserting
+        monitor_thread.join()
+
+        if without_flag:
+            # Without the `_adding_branches` flag, the monitor thread will
+            # attempt to add branches at the same time as the main thread,
+            # which will result in duplicate branches
+            assert len(self.tree.trunk.branches) > 1
+            self.tree.trunk._add_branches = add_branches  # restore original method
+        else:
+            # With the flag in place, the monitor thread bails from adding
+            # branches when the parent branch is presently doing so, leading
+            # to only the adding of the single intended branch
+            assert len(self.tree.trunk.branches) == 1
+
 
 class TestWorkChainStatusPanel(TreeTestingMixin):
     model: WorkChainStatusModel