Add missing features in resource settings tabs

src/aiidalab_qe/app/configuration/__init__.py

@@ -56,7 +56,7 @@ def __init__(self, model: ConfigurationStepModel, **kwargs):
             lambda structure: ""
             if structure
             else """
-                <div class="alert alert-info">
+                <div class="alert alert-danger">
                     <b>Please set the input structure first.</b>
                 </div>
             """,

src/aiidalab_qe/app/submission/__init__.py

@@ -39,10 +39,6 @@ def __init__(self, model: SubmissionStepModel, qe_auto_setup=True, **kwargs):
             self._on_submission,
             "confirmed",
         )
-        self._model.observe(
-            self._on_input_structure_change,
-            "input_structure",
-        )
         self._model.observe(
             self._on_input_parameters_change,
             "input_parameters",
@@ -77,22 +73,28 @@ def __init__(self, model: SubmissionStepModel, qe_auto_setup=True, **kwargs):
 
         self.rendered = False
 
-        global_code_model = GlobalResourceSettingsModel()
-        self.global_code_settings = GlobalResourceSettingsPanel(model=global_code_model)
-        self._model.add_model("global", global_code_model)
-        global_code_model.observe(
+        global_resources_model = GlobalResourceSettingsModel()
+        self.global_resources = GlobalResourceSettingsPanel(
+            model=global_resources_model
+        )
+        self._model.add_model("global", global_resources_model)
+        ipw.dlink(
+            (self._model, "plugin_overrides"),
+            (global_resources_model, "plugin_overrides"),
+        )
+        global_resources_model.observe(
             self._on_plugin_submission_blockers_change,
             ["submission_blockers"],
         )
-        global_code_model.observe(
+        global_resources_model.observe(
             self._on_plugin_submission_warning_messages_change,
             ["submission_warning_messages"],
         )
 
         self.settings = {
-            "global": self.global_code_settings,
+            "global": self.global_resources,
         }
-        self._fetch_plugin_settings()
+        self._fetch_plugin_resource_settings()
 
         self._install_sssp(qe_auto_setup)
         self._set_up_qe(qe_auto_setup)
@@ -211,14 +213,15 @@ def _on_tab_change(self, change):
         tab: ResourceSettingsPanel = self.tabs.children[tab_index]  # type: ignore
         tab.render()
 
-    def _on_input_structure_change(self, _):
-        """"""
-
     def _on_input_parameters_change(self, _):
-        self._model.update_active_models()
-        self._update_tabs()
         self._model.update_process_label()
+        self._model.update_plugin_inclusion()
+        self._model.update_plugin_overrides()
         self._model.update_submission_blockers()
+        self._update_tabs()
+
+    def _on_plugin_overrides_change(self, _):
+        self._model.update_plugin_overrides()
 
     def _on_plugin_submission_blockers_change(self, _):
         self._model.update_submission_blockers()
@@ -237,16 +240,13 @@ def _on_submission_blockers_change(self, _):
         self._model.update_submission_blocker_message()
         self._update_state()
 
-    def _on_submission_warning_change(self, _):
-        self._model.update_submission_warning_message()
-
     def _on_installation_change(self, _):
         self._model.update_submission_blockers()
 
     def _on_qe_installed(self, _):
         self._toggle_qe_installation_widget()
         if self._model.qe_installed:
-            self._model.refresh_codes()
+            self._model.update()
 
     def _on_sssp_installed(self, _):
         self._toggle_sssp_installation_widget()
@@ -325,14 +325,19 @@ def _update_state(self, _=None):
         else:
             self.state = self.state.CONFIGURED
 
-    def _fetch_plugin_settings(self):
-        eps = get_entry_items("aiidalab_qe.properties", "code")
-        for identifier, data in eps.items():
+    def _fetch_plugin_resource_settings(self):
+        entries = get_entry_items("aiidalab_qe.properties", "resources")
+        for identifier, resources in entries.items():
             for key in ("panel", "model"):
-                if key not in data:
+                if key not in resources:
                     raise ValueError(f"Entry {identifier} is missing the '{key}' key")
-            panel = data["panel"]
-            model: ResourceSettingsModel = data["model"]()
+
+            panel = resources["panel"]
+            model: ResourceSettingsModel = resources["model"]()
+            model.observe(
+                self._on_plugin_overrides_change,
+                "override",
+            )
             model.observe(
                 self._on_plugin_submission_blockers_change,
                 ["submission_blockers"],
@@ -343,15 +348,6 @@ def _fetch_plugin_settings(self):
             )
             self._model.add_model(identifier, model)
 
-            def toggle_plugin(_, model=model):
-                model.update()
-                self._update_tabs()
-
-            model.observe(
-                toggle_plugin,
-                "include",
-            )
-
             self.settings[identifier] = panel(
                 identifier=identifier,
                 model=model,

src/aiidalab_qe/app/submission/global_settings/model.py

@@ -5,55 +5,37 @@
 import traitlets as tl
 
 from aiida import orm
-from aiidalab_qe.app.parameters import DEFAULT_PARAMETERS
 from aiidalab_qe.common.code import CodeModel, PwCodeModel
 from aiidalab_qe.common.mixins import HasInputStructure
 from aiidalab_qe.common.panel import ResourceSettingsModel
 from aiidalab_qe.common.widgets import QEAppComputationalResourcesWidget
 
-DEFAULT: dict = DEFAULT_PARAMETERS  # type: ignore
-
 
 class GlobalResourceSettingsModel(
     ResourceSettingsModel,
     HasInputStructure,
 ):
     """Model for the global code setting."""
 
+    identifier = "global"
+
     dependencies = [
         "input_parameters",
         "input_structure",
     ]
 
     input_parameters = tl.Dict()
 
-    codes = tl.Dict(
-        key_trait=tl.Unicode(),  # code name
-        value_trait=tl.Instance(CodeModel),  # code metadata
-    )
-    # this is a copy of the codes trait, which is used to trigger the update of the plugin
-    global_codes = tl.Dict(
-        key_trait=tl.Unicode(),  # code name
-        value_trait=tl.Dict(),  # code metadata
-    )
-
-    plugin_mapping = tl.Dict(
-        key_trait=tl.Unicode(),  # plugin identifier
-        value_trait=tl.List(tl.Unicode()),  # list of code names
-    )
+    plugin_overrides = tl.List(tl.Unicode())
+    plugin_overrides_notification = tl.Unicode("")
 
-    submission_blockers = tl.List(tl.Unicode())
-    submission_warning_messages = tl.Unicode("")
+    plugin_mapping: dict[str, list[str]] = {}
 
     include = True
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        # Used by the code-setup thread to fetch code options
-        # This is necessary to avoid passing the User object
-        # between session in separate threads.
-        self._default_user_email = orm.User.collection.get_default().email
         self._RUN_ON_LOCALHOST_NUM_SITES_WARN_THRESHOLD = 10
         self._RUN_ON_LOCALHOST_VOLUME_WARN_THRESHOLD = 1000  # \AA^3
 
@@ -66,9 +48,14 @@ def __init__(self, *args, **kwargs):
             </div>
         """
 
-    def refresh_codes(self):
+        self.override = True
+
+    def update(self):
         for _, code_model in self.codes.items():
-            code_model.update(self._default_user_email)  # type: ignore
+            code_model.update(self.DEFAULT_USER_EMAIL)
+
+    def update_global_codes(self):
+        self.global_codes = self.get_model_state()["codes"]
 
     def update_active_codes(self):
         for name, code_model in self.codes.items():
@@ -80,97 +67,50 @@ def update_active_codes(self):
                 for code_name in code_names:
                     self.codes[code_name].activate()
 
-    def get_model_state(self):
-        codes = {name: model.get_model_state() for name, model in self.codes.items()}
-
-        return {"codes": codes}
-
-    def set_model_state(self, code_data: dict):
-        for name, code_model in self.codes.items():
-            if name in code_data and code_model.is_active:
-                code_model.set_model_state(code_data[name])
+    def update_plugin_overrides_notification(self):
+        if self.plugin_overrides:
+            formatted = "\n".join(
+                f"<li>{plugin}</li>" for plugin in self.plugin_overrides
+            )
+            self.plugin_overrides_notification = f"""
+                <div class="alert alert-info" style="margin-top: 5px; margin-bottom: 0">
+                    <strong>Currently overriding computational resources for:</strong>
+                    <ul>
+                        {formatted}
+                    </ul>
+                </div>
+            """
+        else:
+            self.plugin_overrides_notification = ""
 
-    def add_code(self, identifier: str, code: CodeModel) -> CodeModel | None:
-        """Add a code to the codes trait."""
-        code_model = None
-        default_calc_job_plugin = code.default_calc_job_plugin
+    def add_code(self, identifier: str, code_model: CodeModel) -> CodeModel | None:
+        """Registers a code with this model."""
+        base_code_model = None
+        default_calc_job_plugin = code_model.default_calc_job_plugin
         name = default_calc_job_plugin.split(".")[-1]
+
         if default_calc_job_plugin not in self.codes:
             if default_calc_job_plugin == "quantumespresso.pw":
-                code_model = PwCodeModel(
+                base_code_model = PwCodeModel(
                     name=name,
                     description=name,
                     default_calc_job_plugin=default_calc_job_plugin,
                 )
             else:
-                code_model = CodeModel(
+                base_code_model = CodeModel(
                     name=name,
                     description=name,
                     default_calc_job_plugin=default_calc_job_plugin,
                 )
-            self.codes[default_calc_job_plugin] = code_model
-        # update the plugin mapping to keep track of which codes are associated with which plugins
+            self.codes[default_calc_job_plugin] = base_code_model
+            self._link_code(base_code_model)
+
         if identifier not in self.plugin_mapping:
             self.plugin_mapping[identifier] = [default_calc_job_plugin]
         else:
             self.plugin_mapping[identifier].append(default_calc_job_plugin)
 
-        return code_model
-
-    def get_code(self, name) -> CodeModel | None:
-        if name in self.codes:  # type: ignore
-            return self.codes[name]  # type: ignore
-
-    def get_selected_codes(self) -> dict[str, dict]:
-        return {
-            name: code_model.get_model_state()
-            for name, code_model in self.codes.items()
-            if code_model.is_ready
-        }
-
-    def set_selected_codes(self, code_data=DEFAULT["codes"]):
-        for name, code_model in self.codes.items():
-            if name in code_data and code_model.is_active:
-                code_model.set_model_state(code_data[name])
-
-    def reset(self):
-        """Reset the model to its default state."""
-        for code_model in self.codes.values():
-            code_model.reset()
-
-    def _get_properties(self) -> list[str]:
-        return self.input_parameters.get("workchain", {}).get("properties", [])
-
-    def update_submission_blockers(self):
-        self.submission_blockers = list(self._check_submission_blockers())
-
-    def _check_submission_blockers(self):
-        # No pw code selected (this is ignored while the setup process is running).
-        pw_code = self.get_code("quantumespresso.pw")
-        if pw_code and not pw_code.selected and not self.installing_qe:
-            yield ("No pw code selected")
-
-        # code related to the selected property is not installed
-        properties = self._get_properties()
-        message = "Calculating the {property} property requires code {code} to be set."
-        for identifier, code_names in self.plugin_mapping.items():
-            if identifier in properties:
-                for name in code_names:
-                    code = self.get_code(name)
-                    if not code.is_ready:
-                        yield message.format(property=identifier, code=code.description)
-
-        # 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 not code.is_ready:
-                continue
-            if not issubclass(
-                code.code_widget_class, QEAppComputationalResourcesWidget
-            ):
-                yield (
-                    f"Error: hi, plugin developer, please use the QEAppComputationalResourcesWidget from aiidalab_qe.common.widgets for code {name}."
-                )
+        return base_code_model
 
     def check_resources(self):
         pw_code = self.get_code("quantumespresso.pw")
@@ -254,6 +194,37 @@ def check_resources(self):
             )
         )
 
+    def _get_properties(self) -> list[str]:
+        return self.input_parameters.get("workchain", {}).get("properties", [])
+
+    def _check_submission_blockers(self):
+        # No pw code selected
+        pw_code = self.get_code("quantumespresso.pw")
+        if pw_code and not pw_code.selected:
+            yield ("No pw code selected")
+
+        # Code related to the selected property is not installed
+        properties = self._get_properties()
+        message = "Calculating the {property} property requires code {code} to be set."
+        for identifier, code_names in self.plugin_mapping.items():
+            if identifier in properties:
+                for name in code_names:
+                    code = self.get_code(name)
+                    if not code.is_ready:
+                        yield message.format(property=identifier, code=code.description)
+
+        # 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 not code.is_ready:
+                continue
+            if not issubclass(
+                code.code_widget_class, QEAppComputationalResourcesWidget
+            ):
+                yield (
+                    f"Error: hi, plugin developer, please use the QEAppComputationalResourcesWidget from aiidalab_qe.common.widgets for code {name}."
+                )
+
     def _estimate_min_cpus(
         self,
         n,