Migrate to support AiiDA 2.x (#271)

It is now migrating to AiiDA 2.x with also its new features supported such as `InstalledCode`.
The CI downstream test is added. This will fix issue #205. 
To conform with the thread-local/traitlets issue, we change also all traitlets related to threads in the code base. 
We are now getting rid of all the widgets using AiiDA node as traitlets except `NodeViewWidget` and `MinimalStructureViewer`, which will be further tested.

.github/workflows/di.yml

@@ -7,35 +7,53 @@ on:
     workflow_dispatch:
     schedule:
         - cron: 41 3 * * *
+    pull_request:
 
 jobs:
 
     test-app:
 
-        runs-on: ubuntu-latest
-        timeout-minutes: 10
-
         strategy:
             matrix:
-                tag: [latest, stable]
-                browser: [chrome, firefox]
+                tag: [latest]
+                browser: [Chrome, Firefox]
+                python-version: ['3.8', '3.10']
+                firefox: ['84.0']
             fail-fast: false
 
-        steps:
+        runs-on: ubuntu-latest
+        timeout-minutes: 60
 
+        steps:
             - name: Check out app
               uses: actions/checkout@v2
 
-            - name: Test app
-              uses: aiidalab/aiidalab-test-app-action@v2
+            - name: Cache Python dependencies
+              uses: actions/cache@v1
               with:
-                  image: aiidalab/aiidalab-docker-stack:${{ matrix.tag }}
-                  browser: ${{ matrix.browser }}
-                  name: quantum-espresso
+                  path: ~/.cache/pip
+                  key: pip-${{ matrix.python-version }}-tests-${{ hashFiles('**/setup.json') }}
+                  restore-keys: pip-${{ matrix.python-version }}-tests
+
+            - name: Set up Python ${{ matrix.python-version }}
+              uses: actions/setup-python@v2
+              with:
+                  python-version: ${{ matrix.python-version }}
+
+            - name: Install dependencies for test
+              run: |
+                  pip install -U -r requirements_test.txt
+
+            - name: Set jupyter token env
+              run: echo "JUPYTER_TOKEN=$(openssl rand -hex 32)" >> $GITHUB_ENV
 
-            - name: Upload screenshots as artifacts
-              uses: actions/upload-artifact@v2
-              if: ${{ always() }}
+            - uses: browser-actions/setup-firefox@latest
               with:
-                  name: Screenshots-${{ matrix.tag }}-${{ matrix.browser }}
-                  path: screenshots/
+                  firefox-version: ${{ matrix.firefox }}
+              run: which firefox
+
+            - name: Run pytest
+              run: |
+                  pytest --driver ${{ matrix.browser }}
+              env:
+                  TAG: ${{ matrix.tag }}

.pre-commit-config.yaml

@@ -44,3 +44,8 @@ repos:
       rev: 0.19.2
       hooks:
           - id: check-github-workflows
+
+    - repo: https://github.com/kynan/nbstripout
+      rev: 0.5.0
+      hooks:
+          - id: nbstripout

aiidalab_qe/process.py

@@ -4,8 +4,8 @@
 
 import ipywidgets as ipw
 import traitlets
-from aiida.cmdline.utils.query.calculation import CalculationQueryBuilder
 from aiida.orm import load_node
+from aiida.tools.query.calculation import CalculationQueryBuilder
 
 
 class WorkChainSelector(ipw.HBox):

aiidalab_qe/report.py

@@ -1,3 +1,4 @@
+from aiida.orm import WorkChainNode
 from aiida.plugins import WorkflowFactory
 
 PwBaseWorkChain = WorkflowFactory("quantumespresso.pw.base")
@@ -25,8 +26,8 @@
 }
 
 
-def _generate_report_dict(qeapp_wc):
-    builder_parameters = qeapp_wc.get_extra("builder_parameters", {})
+def _generate_report_dict(qeapp_wc: WorkChainNode):
+    builder_parameters = qeapp_wc.base.extras.get("builder_parameters", {})
 
     # Properties
     run_relax = builder_parameters.get("relax_type") != "none"

aiidalab_qe/setup_codes.py

@@ -71,19 +71,20 @@ def _setup_code(code_name, computer_name="localhost"):
             [
                 "verdi",
                 "code",
-                "setup",
+                "create",
+                "core.code.installed",
                 "--non-interactive",
                 "--label",
                 f"{code_name}-{QE_VERSION}",
                 "--description",
                 f"{code_name}.x ({QE_VERSION}) setup by AiiDAlab.",
-                "--input-plugin",
+                "--default-calc-job-plugin",
                 f"quantumespresso.{code_name}",
                 "--computer",
                 computer_name,
                 "--prepend-text",
-                f"conda activate {CONDA_ENV_PREFIX}\nexport OMP_NUM_THREADS=1",
-                "--remote-abs-path",
+                f'eval "$(conda shell.posix hook)"\nconda activate {CONDA_ENV_PREFIX}\nexport OMP_NUM_THREADS=1',
+                "--filepath-executable",
                 CONDA_ENV_PREFIX.joinpath("bin", f"{code_name}.x"),
             ],
             check=True,

aiidalab_qe/steps.py

@@ -11,7 +11,7 @@
 import traitlets
 from aiida.common import NotExistent
 from aiida.engine import ProcessBuilderNamespace, ProcessState, submit
-from aiida.orm import ProcessNode, WorkChainNode, load_code
+from aiida.orm import WorkChainNode, load_code, load_node
 from aiida.plugins import DataFactory
 from aiida_quantumespresso.common.types import ElectronicType, RelaxType, SpinType
 from aiida_quantumespresso.workflows.pw.base import PwBaseWorkChain
@@ -34,10 +34,10 @@
 )
 from aiidalab_qe_workchain import QeAppWorkChain
 
-StructureData = DataFactory("structure")
-Float = DataFactory("float")
-Dict = DataFactory("dict")
-Str = DataFactory("str")
+StructureData = DataFactory("core.structure")
+Float = DataFactory("core.float")
+Dict = DataFactory("core.dict")
+Str = DataFactory("core.str")
 
 
 class WorkChainSettings(ipw.VBox):
@@ -601,9 +601,9 @@ def _identify_submission_blockers(self):
             and len(
                 set(
                     (
-                        self.pw_code.value.computer.pk,
-                        self.dos_code.value.computer.pk,
-                        self.projwfc_code.value.computer.pk,
+                        load_code(self.pw_code.value).computer.pk,
+                        load_code(self.dos_code.value).computer.pk,
+                        load_code(self.projwfc_code.value).computer.pk,
                     )
                 )
             )
@@ -654,7 +654,7 @@ def _auto_select_code(self, change):
                 try:
                     code_widget = getattr(self, code)
                     code_widget.refresh()
-                    code_widget.value = load_code(DEFAULT_PARAMETERS[code])
+                    code_widget.value = load_code(DEFAULT_PARAMETERS[code]).uuid
                 except NotExistent:
                     pass
 
@@ -676,9 +676,10 @@ def _show_alert_message(self, message, alert_class="info"):
     def _update_resources(self, change):
         if change["new"] and (
             change["old"] is None
-            or change["new"].computer.pk != change["old"].computer.pk
+            or load_code(change["new"]).computer.pk
+            != load_code(change["old"]).computer.pk
         ):
-            self.set_resource_defaults(change["new"].computer)
+            self.set_resource_defaults(load_code(change["new"]).computer)
 
     def set_resource_defaults(self, computer=None):
 
@@ -715,7 +716,7 @@ def _check_resources(self):
             return  # No code selected, nothing to do.
 
         num_cpus = self.resources_config.num_cpus.value
-        on_localhost = self.pw_code.value.computer.get_hostname() == "localhost"
+        on_localhost = load_node(self.pw_code.value).computer.hostname == "localhost"
         if self.pw_code.value and on_localhost and num_cpus > 1:
             self._show_alert_message(
                 "The selected code would be executed on the local host, but "
@@ -754,7 +755,9 @@ def _observe_process(self, change):
             process_node = change["new"]
             if process_node is not None:
                 self.input_structure = process_node.inputs.structure
-                builder_parameters = process_node.get_extra("builder_parameters", None)
+                builder_parameters = process_node.base.extras.get(
+                    "builder_parameters", None
+                )
                 if builder_parameters is not None:
                     self.set_selected_codes(builder_parameters)
             self._update_state()
@@ -775,9 +778,9 @@ def get_input_parameters(self):
             run_pdos=self.workchain_settings.pdos_run.value,
             protocol=self.workchain_settings.workchain_protocol.value,
             # Codes
-            pw_code=self.pw_code.value.uuid,
-            dos_code=self.dos_code.value.uuid,
-            projwfc_code=self.projwfc_code.value.uuid,
+            pw_code=self.pw_code.value,
+            dos_code=self.dos_code.value,
+            projwfc_code=self.projwfc_code.value,
             # Advanced settings
             pseudo_family=self.pseudo_family_selector.value,
         )
@@ -795,18 +798,18 @@ def set_selected_codes(self, parameters):
         """Set the inputs in the GUI based on a set of parameters."""
 
         # Codes
-        def _load_code(code):
+        def _get_code_uuid(code):
             if code is not None:
                 try:
-                    return load_code(code)
+                    return load_code(code).uuid
                 except NotExistent:
                     return None
 
         with self.hold_trait_notifications():
             # Codes
-            self.pw_code.value = _load_code(parameters["pw_code"])
-            self.dos_code.value = _load_code(parameters["dos_code"])
-            self.projwfc_code.value = _load_code(parameters["projwfc_code"])
+            self.pw_code.value = _get_code_uuid(parameters["pw_code"])
+            self.dos_code.value = _get_code_uuid(parameters["dos_code"])
+            self.projwfc_code.value = _get_code_uuid(parameters["projwfc_code"])
 
     def set_pdos_status(self):
         if self.workchain_settings.pdos_run.value:
@@ -887,7 +890,7 @@ def update_builder(buildy, resources, npools):
         with self.hold_trait_notifications():
             self.process = submit(builder)
             # Set the builder parameters on the work chain
-            self.process.set_extra("builder_parameters", parameters)
+            self.process.base.extras.set("builder_parameters", parameters)
 
     def reset(self):
         with self.hold_trait_notifications():
@@ -897,11 +900,14 @@ def reset(self):
 
 class ViewQeAppWorkChainStatusAndResultsStep(ipw.VBox, WizardAppWidgetStep):
 
-    process = traitlets.Instance(ProcessNode, allow_none=True)
+    process = traitlets.Unicode(allow_none=True)
 
     def __init__(self, **kwargs):
         self.process_tree = ProcessNodesTreeWidget()
-        ipw.dlink((self, "process"), (self.process_tree, "process"))
+        ipw.dlink(
+            (self, "process"),
+            (self.process_tree, "value"),
+        )
 
         self.node_view = NodeViewWidget(layout={"width": "auto", "height": "auto"})
         ipw.dlink(
@@ -919,7 +925,7 @@ def __init__(self, **kwargs):
                 self._update_state,
             ],
         )
-        ipw.dlink((self, "process"), (self.process_monitor, "process"))
+        ipw.dlink((self, "process"), (self.process_monitor, "value"))
 
         super().__init__([self.process_status], **kwargs)
 
@@ -934,7 +940,8 @@ def _update_state(self):
         if self.process is None:
             self.state = self.State.INIT
         else:
-            process_state = self.process.process_state
+            process = load_node(self.process)
+            process_state = process.process_state
             if process_state in (
                 ProcessState.CREATED,
                 ProcessState.RUNNING,
@@ -943,10 +950,10 @@ def _update_state(self):
                 self.state = self.State.ACTIVE
             elif (
                 process_state in (ProcessState.EXCEPTED, ProcessState.KILLED)
-                or self.process.is_failed
+                or process.is_failed
             ):
                 self.state = self.State.FAIL
-            elif self.process.is_finished_ok:
+            elif process.is_finished_ok:
                 self.state = self.State.SUCCESS
 
     @traitlets.observe("process")

aiidalab_qe/widgets.py

@@ -14,7 +14,7 @@
 
 import ipywidgets as ipw
 import traitlets
-from aiida.orm import CalcJobNode, Node
+from aiida.orm import CalcJobNode, Node, load_node
 from aiidalab_widgets_base import register_viewer_widget, viewer
 from IPython.display import HTML, Javascript, clear_output, display
 
@@ -242,7 +242,7 @@ def _observe_value(self, change):
 
 class CalcJobOutputFollower(traitlets.HasTraits):
 
-    calcjob = traitlets.Instance(CalcJobNode, allow_none=True)
+    calcjob_uuid = traitlets.Unicode(allow_none=True)
     filename = traitlets.Unicode(allow_none=True)
     output = traitlets.List(trait=traitlets.Unicode)
     lineno = traitlets.Int()
@@ -258,14 +258,11 @@ def __init__(self, **kwargs):
 
         super().__init__(**kwargs)
 
-    @traitlets.observe("calcjob")
+    @traitlets.observe("calcjob_uuid")
     def _observe_calcjob(self, change):
-        try:
-            if change["old"].pk == change["new"].pk:
-                # Old and new process are identical.
-                return
-        except AttributeError:
-            pass
+        calcjob_uuid = change["new"]
+        if change["old"] == calcjob_uuid:
+            return
 
         with self._lock:
             # Stop following
@@ -283,30 +280,30 @@ def _observe_calcjob(self, change):
             # (Re/)start following
             if change["new"]:
                 self._follow_output_thread = Thread(
-                    target=self._follow_output, args=(change["new"],)
+                    target=self._follow_output, args=(calcjob_uuid,)
                 )
                 self._follow_output_thread.start()
 
-    def _follow_output(self, calcjob):
+    def _follow_output(self, calcjob_uuid):
         """Monitor calcjob and orchestrate pushing and pulling of output."""
-        self._pull_thread = Thread(target=self._pull_output, args=(calcjob,))
+        self._pull_thread = Thread(target=self._pull_output)
         self._pull_thread.start()
-        self._push_thread = Thread(target=self._push_output, args=(calcjob,))
+        self._push_thread = Thread(target=self._push_output, args=(calcjob_uuid,))
         self._push_thread.start()
 
     def _fetch_output(self, calcjob):
         assert isinstance(calcjob, CalcJobNode)
         if "retrieved" in calcjob.outputs:
             try:
-                self.filename = calcjob.attributes["output_filename"]
+                self.filename = calcjob.base.attributes.get("output_filename")
                 with calcjob.outputs.retrieved.open(self.filename) as f:
                     return f.read().splitlines()
             except OSError:
                 return list()
 
         elif "remote_folder" in calcjob.outputs:
             try:
-                fn_out = calcjob.attributes["output_filename"]
+                fn_out = calcjob.base.attributes.get("output_filename")
                 self.filename = fn_out
                 with NamedTemporaryFile() as tmpfile:
                     calcjob.outputs.remote_folder.getfile(fn_out, tmpfile.name)
@@ -318,9 +315,10 @@ def _fetch_output(self, calcjob):
 
     _EOF = None
 
-    def _push_output(self, calcjob, delay=0.2):
+    def _push_output(self, calcjob_uuid, delay=0.2):
         """Push new log lines onto the queue."""
         lineno = 0
+        calcjob = load_node(calcjob_uuid)
         while True:
             try:
                 lines = self._fetch_output(calcjob)
@@ -335,7 +333,7 @@ def _push_output(self, calcjob, delay=0.2):
                     self._output_queue.put(self._EOF)
                     break  # noqa: B012
 
-    def _pull_output(self, calcjob):
+    def _pull_output(self):
         """Pull new log lines from the queue and update traitlets."""
         while True:
             item = self._output_queue.get()