Skip to content

Commit

Permalink
Add missing features in resource settings tabs
Browse files Browse the repository at this point in the history
  • Loading branch information
edan-bainglass committed Nov 28, 2024
1 parent 1ae24a8 commit 63fc914
Show file tree
Hide file tree
Showing 9 changed files with 340 additions and 366 deletions.
2 changes: 1 addition & 1 deletion src/aiidalab_qe/app/configuration/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,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>
""",
Expand Down
66 changes: 31 additions & 35 deletions src/aiidalab_qe/app/submission/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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()
Expand All @@ -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()
Expand Down Expand Up @@ -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"],
Expand All @@ -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,
Expand Down
169 changes: 70 additions & 99 deletions src/aiidalab_qe/app/submission/global_settings/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,57 +5,39 @@
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.panel import DEFAULT_USER_EMAIL, 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

Expand All @@ -68,9 +50,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(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():
Expand All @@ -82,97 +69,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")
Expand Down Expand Up @@ -256,6 +196,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,
Expand Down
Loading

0 comments on commit 63fc914

Please sign in to comment.