diff --git a/.gitignore b/.gitignore index 66e045e5..385516e5 100644 --- a/.gitignore +++ b/.gitignore @@ -109,3 +109,6 @@ venv.bak/ .DS_Store .vscode + +# screenshots from tests +screenshots/ diff --git a/aiidalab_widgets_base/computational_resources.py b/aiidalab_widgets_base/computational_resources.py index 25418576..b7bd12e7 100644 --- a/aiidalab_widgets_base/computational_resources.py +++ b/aiidalab_widgets_base/computational_resources.py @@ -1,23 +1,29 @@ +from __future__ import annotations + +import copy import enum import os +import re import subprocess import threading -from copy import copy +from collections import namedtuple from pathlib import Path from uuid import UUID import ipywidgets as ipw +import jinja2 import pexpect import shortuuid -import traitlets +import traitlets as tl from aiida import common, orm, plugins from aiida.orm.utils.builders.computer import ComputerBuilder from aiida.transports.plugins import ssh as aiida_ssh_plugin from humanfriendly import InvalidSize, parse_size from IPython.display import clear_output, display +from jinja2 import meta as jinja2_meta from .databases import ComputationalResourcesDatabaseWidget -from .utils import StatusHTML +from .utils import MessageLevel, StatusHTML, wrap_message STYLE = {"description_width": "180px"} LAYOUT = {"width": "400px"} @@ -41,36 +47,50 @@ class ComputationalResourcesWidget(ipw.VBox): computers. """ - value = traitlets.Unicode(allow_none=True) - codes = traitlets.Dict(allow_none=True) - allow_hidden_codes = traitlets.Bool(False) - allow_disabled_computers = traitlets.Bool(False) - default_calc_job_plugin = traitlets.Unicode(allow_none=True) - - def __init__(self, description="Select code:", path_to_root="../", **kwargs): + value = tl.Unicode(allow_none=True) + codes = tl.Dict(allow_none=True) + allow_hidden_codes = tl.Bool(False) + allow_disabled_computers = tl.Bool(False) + + # code output layout width + _output_width = "460px" + + def __init__( + self, + description="Select code:", + enable_quick_setup=True, + enable_detailed_setup=True, + clear_after=None, + default_calc_job_plugin=None, + **kwargs, + ): """Dropdown for Codes for one input plugin. description (str): Description to display before the dropdown. """ + clear_after = clear_after or 15 + self.default_calc_job_plugin = default_calc_job_plugin self.output = ipw.HTML() - self.setup_message = StatusHTML(clear_after=30) + self.setup_message = StatusHTML(clear_after=clear_after) self.code_select_dropdown = ipw.Dropdown( description=description, disabled=True, value=None, style={"description_width": "initial"}, ) - traitlets.directional_link( + + # Create bi-directional link between the code_select_dropdown and the codes trait. + tl.directional_link( (self, "codes"), (self.code_select_dropdown, "options"), transform=lambda x: [(key, x[key]) for key in x], ) - traitlets.directional_link( + tl.directional_link( (self.code_select_dropdown, "options"), (self, "codes"), transform=lambda x: {c[0]: c[1] for c in x}, ) - traitlets.link((self.code_select_dropdown, "value"), (self, "value")) + tl.link((self.code_select_dropdown, "value"), (self, "value")) self.observe( self.refresh, names=["allow_disabled_computers", "allow_hidden_codes"] @@ -79,7 +99,7 @@ def __init__(self, description="Select code:", path_to_root="../", **kwargs): self.btn_setup_new_code = ipw.ToggleButton(description="Setup new code") self.btn_setup_new_code.observe(self._setup_new_code, "value") - self._setup_new_code_output = ipw.Output(layout={"width": "500px"}) + self._setup_new_code_output = ipw.Output(layout={"width": self._output_width}) children = [ ipw.HBox([self.code_select_dropdown, self.btn_setup_new_code]), @@ -88,93 +108,36 @@ def __init__(self, description="Select code:", path_to_root="../", **kwargs): ] super().__init__(children=children, **kwargs) - # Setting up codes and computers. - self.comp_resources_database = ComputationalResourcesDatabaseWidget( - default_calc_job_plugin=self.default_calc_job_plugin - ) - - self.ssh_computer_setup = SshComputerSetup() - ipw.dlink( - (self.ssh_computer_setup, "message"), - (self.setup_message, "message"), - ) - - ipw.dlink( - (self.comp_resources_database, "ssh_config"), - (self.ssh_computer_setup, "ssh_config"), + # Computer/code setup + self.resource_setup = _ResourceSetupBaseWidget( + default_calc_job_plugin=self.default_calc_job_plugin, + enable_quick_setup=enable_quick_setup, + enable_detailed_setup=enable_detailed_setup, ) - - self.aiida_computer_setup = AiidaComputerSetup() - ipw.dlink( - (self.aiida_computer_setup, "message"), + self.resource_setup.observe(self.refresh, "success") + tl.dlink( + (self.resource_setup, "message"), (self.setup_message, "message"), ) - ipw.dlink( - (self.comp_resources_database, "computer_setup"), - (self.aiida_computer_setup, "computer_setup"), - ) - - # Set up AiiDA code. - self.aiida_code_setup = AiidaCodeSetup() - ipw.dlink( - (self.aiida_code_setup, "message"), - (self.setup_message, "message"), - ) - ipw.dlink( - (self.comp_resources_database, "code_setup"), - (self.aiida_code_setup, "code_setup"), - ) - self.aiida_code_setup.on_setup_code_success(self.refresh) - - # After a successfull computer setup the codes widget should be refreshed. - # E.g. the list of available computers needs to be updated. - self.aiida_computer_setup.on_setup_computer_success( - self.aiida_code_setup.refresh - ) - - # Quick setup. - quick_setup_button = ipw.Button(description="Quick Setup") - quick_setup_button.on_click(self.quick_setup) - self.quick_setup = ipw.VBox( - children=[ - self.ssh_computer_setup.username, - quick_setup_button, - self.aiida_code_setup.setup_code_out, - ] - ) - - # Detailed setup. - self.detailed_setup = ipw.Accordion( - children=[ - self.ssh_computer_setup, - self.aiida_computer_setup, - self.aiida_code_setup, - ] - ) - self.detailed_setup.set_title(0, "Set up password-less SSH connection") - self.detailed_setup.set_title(1, "Set up a computer in AiiDA") - self.detailed_setup.set_title(2, "Set up a code in AiiDA") self.refresh() - def quick_setup(self, _=None): - """Go through all the setup steps automatically.""" - with self.hold_trait_notifications(): - self.ssh_computer_setup._on_setup_ssh_button_pressed() - if self.aiida_computer_setup.on_setup_computer(): - self.aiida_code_setup.on_setup_code() - def _get_codes(self): """Query the list of available codes.""" user = orm.User.collection.get_default() + filters = ( + {"attributes.input_plugin": self.default_calc_job_plugin} + if self.default_calc_job_plugin + else {} + ) return [ (self._full_code_label(c[0]), c[0].uuid) for c in orm.QueryBuilder() .append( orm.Code, - filters={"attributes.input_plugin": self.default_calc_job_plugin}, + filters=filters, ) .all() if c[0].computer.is_user_configured(user) @@ -202,7 +165,7 @@ def refresh(self, _=None): self.code_select_dropdown.disabled = False self.code_select_dropdown.value = None - @traitlets.validate("value") + @tl.validate("value") def _validate_value(self, change): """Check if the code is valid in DB""" code_uuid = change["value"] @@ -225,33 +188,18 @@ def _setup_new_code(self, _=None): clear_output() if self.btn_setup_new_code.value: self._setup_new_code_output.layout = { - "width": "500px", + "width": self._output_width, "border": "1px solid gray", } - if self.comp_resources_database.database: - setup_tab = ipw.Tab( - children=[self.quick_setup, self.detailed_setup] - ) - setup_tab.set_title(0, "Quick Setup") - setup_tab.set_title(1, "Detailed Setup") - children = [ - ipw.HTML( - """Please select the computer/code from a database to pre-fill the fields below.""" - ), - self.comp_resources_database, - self.ssh_computer_setup.password_box, - self.setup_message, - setup_tab, - ] - else: - # Display only Detailed Setup if DB is empty - setup_tab = ipw.Tab(children=[self.detailed_setup]) - setup_tab.set_title(0, "Detailed Setup") - children = [self.setup_message, setup_tab] + + children = [ + self.resource_setup, + self.setup_message, + ] display(*children) else: self._setup_new_code_output.layout = { - "width": "500px", + "width": self._output_width, "border": "none", } @@ -271,8 +219,8 @@ class SshConnectionState(enum.Enum): class SshComputerSetup(ipw.VBox): """Setup a passwordless access to a computer.""" - ssh_config = traitlets.Dict() - ssh_connection_state = traitlets.UseEnum( + ssh_config = tl.Dict() + ssh_connection_state = tl.UseEnum( SshConnectionState, allow_none=True, default_value=None ) SSH_POSSIBLE_RESPONSES = [ @@ -287,34 +235,56 @@ class SshComputerSetup(ipw.VBox): "Connection refused", # 6 pexpect.EOF, # 7 ] - message = traitlets.Unicode() - password_message = traitlets.Unicode("The passwordless enabling log.") + message = tl.Unicode() + password_message = tl.Unicode("The passwordless enabling log.") + + def __init__(self, ssh_folder: Path | None = None, **kwargs): + """Setup a passwordless access to a computer.""" + # ssh folder init + if ssh_folder is None: + ssh_folder = Path.home() / ".ssh" + + if not ssh_folder.exists(): + ssh_folder.mkdir() + ssh_folder.chmod(0o700) + + self._ssh_folder = ssh_folder - def __init__(self, **kwargs): self._ssh_connection_message = None self._password_message = ipw.HTML() - ipw.dlink((self, "password_message"), (self._password_message, "value")) - self._ssh_password = ipw.Password(layout={"width": "150px"}, disabled=True) + tl.dlink( + (self, "password_message"), + (self._password_message, "value"), + transform=lambda x: f"SSH log: {x}", + ) + self._ssh_password = ipw.Password( + description="password:", + disabled=False, + layout=LAYOUT, + style=STYLE, + ) + # Don't show the continue button until it ask for password the + # second time, which happened when the proxy jump is set. The + # first time it ask for password is for the jump host. self._continue_with_password_button = ipw.Button( - description="Continue", layout={"width": "100px"}, disabled=True + description="Continue", + layout={"width": "100px", "display": "none"}, ) self._continue_with_password_button.on_click(self._send_password) self.password_box = ipw.VBox( [ - self._password_message, ipw.HBox([self._ssh_password, self._continue_with_password_button]), + self._password_message, ] ) # Username. - self.username = ipw.Text( - description="SSH username:", layout=LAYOUT, style=STYLE - ) + self.username = ipw.Text(description="Username:", layout=LAYOUT, style=STYLE) # Port. self.port = ipw.IntText( - description="SSH port:", + description="Port:", value=22, layout=LAYOUT, style=STYLE, @@ -351,6 +321,7 @@ def __init__(self, **kwargs): ("Password", "password"), ("Use custom private key", "private_key"), ("Download public key", "public_key"), + ("Multiple factor authentication", "mfa"), ], layout=LAYOUT, style=STYLE, @@ -385,8 +356,11 @@ def __init__(self, **kwargs): def _ssh_keygen(self): """Generate ssh key pair.""" - self.message = "Generating SSH key pair." - fpath = Path.home() / ".ssh" / "id_rsa" + self.message = wrap_message( + "Generating SSH key pair.", + MessageLevel.SUCCESS, + ) + fpath = self._ssh_folder / "id_rsa" keygen_cmd = [ "ssh-keygen", "-f", @@ -422,8 +396,8 @@ def _can_login(self): def _is_in_config(self): """Check if the SSH config file contains host information.""" - fpath = Path.home() / ".ssh" / "config" - if not fpath.exists(): + config_path = self._ssh_folder / "config" + if not config_path.exists(): return False sshcfg = aiida_ssh_plugin.parse_sshconfig(self.hostname.value) # NOTE: parse_sshconfig returns a dict with a hostname @@ -435,15 +409,13 @@ def _is_in_config(self): def _write_ssh_config(self, private_key_abs_fname=None): """Put host information into the config file.""" - fpath = Path.home() / ".ssh" - if not fpath.exists(): - fpath.mkdir() - fpath.chmod(0o700) + config_path = self._ssh_folder / "config" - fpath = fpath / "config" - - self.message = f"Adding {self.hostname.value} section to {fpath}" - with open(fpath, "a") as file: + self.message = wrap_message( + f"Adding {self.hostname.value} section to {config_path}", + MessageLevel.SUCCESS, + ) + with open(config_path, "a") as file: file.write(f"Host {self.hostname.value}\n") file.write(f" User {self.username.value}\n") file.write(f" Port {self.port.value}\n") @@ -459,48 +431,76 @@ def _write_ssh_config(self, private_key_abs_fname=None): file.write(f" IdentityFile {private_key_abs_fname}\n") file.write(" ServerAliveInterval 5\n") - def _on_setup_ssh_button_pressed(self, _=None): + def key_pair_prepare(self): + """Prepare key pair for the ssh connection.""" # Always start by generating a key pair if they are not present. self._ssh_keygen() # If hostname & username are not provided - do not do anything. if self.hostname.value == "": # check hostname - self.message = "Please specify the computer hostname." - return False + message = "Please specify the computer name (for SSH)" + + raise ValueError(message) if self.username.value == "": # check username - self.message = "Please specify your SSH username." - return False + message = "Please specify your SSH username." - private_key_abs_fname = None - if self._verification_mode.value == "private_key": - # unwrap private key file and setting temporary private_key content - private_key_abs_fname, private_key_content = self._private_key - if private_key_abs_fname is None: # check private key file - self.message = "Please upload your private key file." - return False + raise ValueError(message) + + def thread_ssh_copy_id(self): + """Copy public key on the remote computer, on a separate thread.""" + ssh_connection_thread = threading.Thread(target=self._ssh_copy_id) + ssh_connection_thread.start() + + def _on_setup_ssh_button_pressed(self, _=None): + """Setup ssh connection.""" + if self._verification_mode.value == "password": + try: + self.key_pair_prepare() + except ValueError as exc: + self.message = wrap_message(str(exc), MessageLevel.ERROR) + return + + self.thread_ssh_copy_id() + + # For not password ssh auth (such as using private_key or 2FA), key pair is not needed (2FA) + # or the key pair is ready. + # There are other mechanism to set up the ssh connection. + # But we still need to write the ssh config to the ssh config file for such as + # proxy jump. + + private_key_fname = None + if self._verification_mode.value == "private_key": # Write private key in ~/.ssh/ and use the name of upload file, # if exist, generate random string and append to filename then override current name. - # TODO(danielhollas): I don't think this works as intended. If there is an existing private key, - # the new filename is never propagated to the caller here. - # https://github.com/aiidalab/aiidalab-widgets-base/issues/516 - self._add_private_key(private_key_abs_fname, private_key_content) + + # unwrap private key file and setting temporary private_key content + private_key_fname, private_key_content = self._private_key + if private_key_fname is None: # check private key file + message = "Please upload your private key file." + self.message = wrap_message(message, MessageLevel.ERROR) + return + + filename = Path(private_key_fname).name + + # if the private key filename is exist, generate random string and append to filename subfix + # then override current name. + if filename in [str(p.name) for p in Path(self._ssh_folder).iterdir()]: + private_key_fpath = self._ssh_folder / f"{filename}-{shortuuid.uuid()}" + + self._add_private_key(private_key_fpath, private_key_content) # TODO(danielhollas): I am not sure this is correct. What if the user wants # to overwrite the private key? Or any other config? The configuration would never be written. # And the user is not notified that we did not write anything. # https://github.com/aiidalab/aiidalab-widgets-base/issues/516 if not self._is_in_config(): - self._write_ssh_config(private_key_abs_fname=private_key_abs_fname) - - # Copy public key on the remote computer. - ssh_connection_thread = threading.Thread(target=self._ssh_copy_id) - ssh_connection_thread.start() + self._write_ssh_config(private_key_abs_fname=private_key_fname) def _ssh_copy_id(self): """Run the ssh-copy-id command and follow it until it is completed.""" - timeout = 30 + timeout = 10 self.password_message = f"Sending public key to {self.hostname.value}... " self._ssh_connection_process = pexpect.spawn( f"ssh-copy-id {self.hostname.value}" @@ -513,11 +513,7 @@ def _ssh_copy_id(self): ) self.ssh_connection_state = SshConnectionState(idx) except pexpect.TIMEOUT: - self._ssh_password.disabled = True - self._continue_with_password_button.disabled = True - self.password_message = ( - f"Exceeded {timeout} s timeout. Please start again." - ) + self.password_message = f"Exceeded {timeout} s timeout. Please check you username and password and try again." break # Terminating the process when nothing else can be done. @@ -535,11 +531,10 @@ def _ssh_copy_id(self): self._ssh_connection_process = None def _send_password(self, _=None): - self._ssh_password.disabled = True self._continue_with_password_button.disabled = True self._ssh_connection_process.sendline(self._ssh_password.value) - @traitlets.observe("ssh_connection_state") + @tl.observe("ssh_connection_state") def _observe_ssh_connnection_state(self, _=None): """Observe the ssh connection state and act according to the changes.""" if self.ssh_connection_state is SshConnectionState.waiting_for_input: @@ -591,7 +586,13 @@ def _handle_ssh_password(self): self._continue_with_password_button.disabled = False self._ssh_connection_message = message - self.ssh_connection_state = SshConnectionState.waiting_for_input + # If user did not provide a password, we wait for the input. + # Otherwise, we send the password. + if self._ssh_password.value == "": + self.ssh_connection_state = SshConnectionState.waiting_for_input + else: + self.password_message = 'Sending password (timeout 10s)' + self._send_password() def _on_verification_mode_change(self, change): """which verification mode is chosen.""" @@ -600,7 +601,7 @@ def _on_verification_mode_change(self, change): if self._verification_mode.value == "private_key": display(self._inp_private_key) elif self._verification_mode.value == "public_key": - public_key = Path.home() / ".ssh" / "id_rsa.pub" + public_key = self._ssh_folder / "id_rsa.pub" if public_key.exists(): display( ipw.HTML( @@ -610,35 +611,21 @@ def _on_verification_mode_change(self, change): ) @property - def _private_key(self): + def _private_key(self) -> tuple[str | None, bytes | None]: """Unwrap private key file and setting filename and file content.""" if self._inp_private_key.value: (fname, _value), *_ = self._inp_private_key.value.items() - content = copy(_value["content"]) + content = copy.copy(_value["content"]) self._inp_private_key.value.clear() self._inp_private_key._counter = 0 # pylint: disable=protected-access return fname, content return None, None @staticmethod - def _add_private_key(private_key_fname, private_key_content): - """ - param private_key_fname: string - param private_key_content: bytes - """ - fpath = Path.home() / ".ssh" / private_key_fname - if fpath.exists(): - # If the file already exist and has the same content, we do nothing. - if fpath.read_bytes() == private_key_content: - return fpath - # If the content is different, we make a new file with a unique name. - fpath = fpath / "_" / shortuuid.uuid() - - fpath.write_bytes(private_key_content) - - fpath.chmod(0o600) - - return fpath + def _add_private_key(private_key_fpath: Path, private_key_content: bytes): + """Write private key to the private key file in the ssh folder.""" + private_key_fpath.write_bytes(private_key_content) + private_key_fpath.chmod(0o600) def _reset(self): self.hostname.value = "" @@ -647,27 +634,30 @@ def _reset(self): self.proxy_jump.value = "" self.proxy_command.value = "" - @traitlets.observe("ssh_config") - def _observe_ssh_config(self, _=None): + @tl.observe("ssh_config") + def _observe_ssh_config(self, change): """Pre-filling the input fields.""" - if not self.ssh_config: - self._reset() + self._reset() - if "hostname" in self.ssh_config: - self.hostname.value = self.ssh_config["hostname"] - if "port" in self.ssh_config: - self.port.value = int(self.ssh_config["port"]) - if "proxy_jump" in self.ssh_config: - self.proxy_jump.value = self.ssh_config["proxy_jump"] - if "proxy_command" in self.ssh_config: - self.proxy_command.value = self.ssh_config["proxy_command"] + new_ssh_config = change["new"] + if "hostname" in new_ssh_config: + self.hostname.value = new_ssh_config["hostname"] + if "username" in new_ssh_config: + self.username.value = new_ssh_config["username"] + if "port" in new_ssh_config: + self.port.value = int(new_ssh_config["port"]) + if "proxy_jump" in new_ssh_config: + self.proxy_jump.value = new_ssh_config["proxy_jump"] + if "proxy_command" in new_ssh_config: + self.proxy_command.value = new_ssh_config["proxy_command"] class AiidaComputerSetup(ipw.VBox): """Inform AiiDA about a computer.""" - computer_setup = traitlets.Dict(allow_none=True) - message = traitlets.Unicode() + computer_setup = tl.Dict(allow_none=True) + computer_configure = tl.Dict(allow_none=True) + message = tl.Unicode() def __init__(self, **kwargs): self._on_setup_computer_success = [] @@ -729,13 +719,13 @@ def __init__(self, **kwargs): ) memory_wrong_syntax = ipw.HTML( value=""" wrong syntax""", - layout={"visibility": "hidden"}, ) + memory_wrong_syntax.layout.display = "none" self.default_memory_per_machine = None def observe_memory_per_machine(change): """Check if the string defining memory is valid.""" - memory_wrong_syntax.layout.visibility = "hidden" + memory_wrong_syntax.layout.display = "none" if not self.default_memory_per_machine_widget.value: self.default_memory_per_machine = None return @@ -743,9 +733,9 @@ def observe_memory_per_machine(change): self.default_memory_per_machine = ( int(parse_size(change["new"], binary=True) / 1024) or None ) - memory_wrong_syntax.layout.visibility = "hidden" + memory_wrong_syntax.layout.display = "none" except InvalidSize: - memory_wrong_syntax.layout.visibility = "visible" + memory_wrong_syntax.layout.display = "block" self.default_memory_per_machine = None self.default_memory_per_machine_widget.observe( @@ -871,11 +861,15 @@ def _configure_computer_ssh(self, computer: orm.Computer, user: orm.User): "use_login_shell": self.use_login_shell.value, "safe_interval": self.safe_interval.value, } - try: - authparams["username"] = sshcfg["user"] - except KeyError as exc: - message = "SSH username is not provided" - raise RuntimeError(message) from exc + if "username" in self.computer_configure: + authparams["username"] = self.computer_configure["username"] + else: + try: + # This require the Ssh connection setup is done before the computer setup + authparams["username"] = sshcfg["user"] + except KeyError as exc: + message = "SSH username is not provided" + raise RuntimeError(message) from exc if "proxycommand" in sshcfg: authparams["proxy_command"] = sshcfg["proxycommand"] @@ -911,28 +905,39 @@ def _computer_exists(self, label): def _validate_computer_settings(self): if self.label.value == "": # check computer label - self.message = "Please specify the computer name (for AiiDA)" + self.message = wrap_message( + "Please specify the computer name (for AiiDA)", + MessageLevel.WARNING, + ) return False if self.work_dir.value == "": - self.message = "Please specify working directory" + self.message = wrap_message( + "Please specify working directory", + MessageLevel.WARNING, + ) return False if self.hostname.value == "": - self.message = "Please specify hostname" + self.message = wrap_message( + "Please specify hostname", + MessageLevel.WARNING, + ) return False return True def on_setup_computer(self, _=None): """Create a new computer.""" - if not self._validate_computer_settings(): return False # If the computer already exists, we just run the registered functions and return if self._run_callbacks_if_computer_exists(self.label.value): - self.message = f"A computer called {self.label.value} already exists." + self.message = wrap_message( + f"A computer {self.label.value} already exists. Skipping creation.", + MessageLevel.INFO, + ) return True items_to_configure = [ @@ -964,15 +969,24 @@ def on_setup_computer(self, _=None): common.exceptions.ValidationError, RuntimeError, ) as err: - self.message = f"Computer setup failed! {type(err).__name__}: {err}" + self.message = wrap_message( + f"Computer setup failed! {type(err).__name__}: {err}", + MessageLevel.ERROR, + ) return False # Callbacks will not run if the computer is not stored if self._run_callbacks_if_computer_exists(self.label.value): - self.message = f"Computer<{computer.pk}> {computer.label} created" + self.message = wrap_message( + f"Computer<{computer.pk}> {computer.label} created", + MessageLevel.SUCCESS, + ) return True - self.message = f"Failed to create computer {computer.label}" + self.message = wrap_message( + f"Failed to create computer {computer.label}", + MessageLevel.ERROR, + ) return False def on_setup_computer_success(self, function): @@ -1022,33 +1036,41 @@ def _reset(self): self.prepend_text.value = "" self.append_text.value = "" - @traitlets.observe("computer_setup") + @tl.observe("computer_setup") def _observe_computer_setup(self, _=None): - # Setup. - if not self.computer_setup: + for key, value in self.computer_setup.items(): + if key == "default_memory_per_machine": + self.default_memory_per_machine_widget.value = f"{value} KB" + elif hasattr(self, key): + getattr(self, key).value = value + + @tl.observe("computer_configure") + def _observe_computer_configure(self, _=None): + for key, value in self.computer_configure.items(): + if hasattr(self, key): + getattr(self, key).value = value + + @tl.observe("computer_setup", "computer_configure") + def _observe_computer_setup_and_configure(self, _=None): + if not self.some_values_provided: self._reset() - return - if "setup" in self.computer_setup: - for key, value in self.computer_setup["setup"].items(): - if key == "default_memory_per_machine": - self.default_memory_per_machine_widget.value = f"{value} KB" - elif hasattr(self, key): - getattr(self, key).value = value - # Configure. - if "configure" in self.computer_setup: - for key, value in self.computer_setup["configure"].items(): - if hasattr(self, key): - getattr(self, key).value = value + @property + def some_values_provided(self): + return any((self.computer_setup, self.computer_configure)) + + @property + def all_values_provided(self): + return all((self.computer_setup, self.computer_configure)) class AiidaCodeSetup(ipw.VBox): """Class that allows to setup AiiDA code""" - code_setup = traitlets.Dict(allow_none=True) - message = traitlets.Unicode() + code_setup = tl.Dict(allow_none=True) + message = tl.Unicode() - def __init__(self, path_to_root="../", **kwargs): + def __init__(self, **kwargs): self._on_setup_code_success = [] # Code label. @@ -1060,9 +1082,7 @@ def __init__(self, path_to_root="../", **kwargs): # Computer on which the code is installed. The value of this widget is # the UUID of the selected computer. - self.computer = ComputerDropdownWidget( - path_to_root=path_to_root, - ) + self.computer = ComputerDropdownWidget() # Code plugin. self.default_calc_job_plugin = ipw.Dropdown( @@ -1126,7 +1146,7 @@ def __init__(self, path_to_root="../", **kwargs): ] super().__init__(children, **kwargs) - @traitlets.validate("default_calc_job_plugin") + @tl.validate("default_calc_job_plugin") def _validate_default_calc_job_plugin(self, proposal): plugin = proposal["value"] return plugin if plugin in self.default_calc_job_plugin.options else None @@ -1137,7 +1157,10 @@ def on_setup_code(self, _=None): clear_output() if not self.computer.value: - self.message = "Please select an existing computer." + self.message = wrap_message( + "Please select an existing computer.", + MessageLevel.WARNING, + ) return False items_to_configure = [ @@ -1164,28 +1187,44 @@ def on_setup_code(self, _=None): filters={"label": kwargs["label"]}, ) if qb.count() > 0: - self.message = ( - f"Code {kwargs['label']}@{computer.label} already exists." + self.message = wrap_message( + f"Code {kwargs['label']}@{computer.label} already exists, skipping creation.", + MessageLevel.INFO, ) + # TODO: (unkcpz) as callback function this return value not actually used + # meanwhile if the code already exists we still want to regard it as success (TBD). + # and the `_on_setup_code_success` callback functions will run. + # The return will break the flow and handle back the control to the caller. + # The proper way is to not return but raise an exception and handle it in the caller function + # for the afterward process. return False try: code = orm.InstalledCode(computer=computer, **kwargs) except (common.exceptions.InputValidationError, KeyError) as exception: - self.message = f"Invalid inputs: {exception}" + self.message = wrap_message( + f"Invalid input for code creation: {exception}", + MessageLevel.ERROR, + ) return False try: code.store() code.is_hidden = False except common.exceptions.ValidationError as exception: - self.message = f"Unable to store the Code: {exception}" + self.message = wrap_message( + f"Unable to store the code: {exception}", + MessageLevel.ERROR, + ) return False for function in self._on_setup_code_success: function() - self.message = f"Code<{code.pk}> {code.full_label} created" + self.message = wrap_message( + f"Code<{code.pk}> {code.full_label} created", + MessageLevel.SUCCESS, + ) return True @@ -1204,7 +1243,7 @@ def _reset(self): def refresh(self): self._observe_code_setup() - @traitlets.observe("code_setup") + @tl.observe("code_setup") def _observe_code_setup(self, _=None): # Setup. self.computer.refresh() @@ -1214,9 +1253,14 @@ def _observe_code_setup(self, _=None): if hasattr(self, key): if key == "default_calc_job_plugin": try: - getattr(self, key).label = value - except traitlets.TraitError: - self.message = f"Input plugin {value} is not installed." + self.default_calc_job_plugin.value = value + except tl.TraitError: + # If is a template then don't raise the error message. + if not re.match(r".*{{.+}}.*", value): + self.message = wrap_message( + f"Input plugin {value} is not installed.", + MessageLevel.WARNING, + ) elif key == "computer": # check if the computer is set by load the label. # if the computer not set put the value to None as placeholder for @@ -1241,16 +1285,14 @@ class ComputerDropdownWidget(ipw.VBox): allow_select_disabled(Bool): Trait that defines whether to show disabled computers. """ - value = traitlets.Unicode(allow_none=True) - computers = traitlets.Dict(allow_none=True) - allow_select_disabled = traitlets.Bool(False) + value = tl.Unicode(allow_none=True) + computers = tl.Dict(allow_none=True) + allow_select_disabled = tl.Bool(False) - def __init__(self, description="Select computer:", path_to_root="../", **kwargs): + def __init__(self, description="Select computer:", **kwargs): """Dropdown for configured AiiDA Computers. description (str): Text to display before dropdown. - - path_to_root (str): Path to the app's root folder. """ self.output = ipw.HTML() @@ -1261,17 +1303,17 @@ def __init__(self, description="Select computer:", path_to_root="../", **kwargs) layout=LAYOUT, disabled=True, ) - traitlets.directional_link( + tl.directional_link( (self, "computers"), (self._dropdown, "options"), transform=lambda x: [(key, x[key]) for key in x], ) - traitlets.directional_link( + tl.directional_link( (self._dropdown, "options"), (self, "computers"), transform=lambda x: {c[0]: c[1] for c in x}, ) - traitlets.link((self._dropdown, "value"), (self, "value")) + tl.link((self._dropdown, "value"), (self, "value")) self.observe(self.refresh, names="allow_select_disabled") @@ -1312,7 +1354,7 @@ def refresh(self, _=None): self._dropdown.value = None - @traitlets.validate("value") + @tl.validate("value") def _validate_value(self, change): """Select computer by computer UUID.""" computer_uuid = change["value"] @@ -1327,3 +1369,544 @@ def _validate_value(self, change): is not a valid UUID.""" else: return computer_uuid + + def select_by_label(self, label): + """Select computer by computer label.""" + self.output.value = "" + if not label: + return None + + try: + computer_uuid = self.computers[label] + except KeyError: + self.output.value = f"""'{label}' + can not find this computer.""" + else: + self.value = computer_uuid + + +TemplateVariableLine = namedtuple("TemplateVariableLine", ["key", "str", "vars"]) + +TemplateVariable = namedtuple("TemplateVariable", ["widget", "lines"]) + + +class TemplateVariablesWidget(ipw.VBox): + # The input template is a dictionary of keyname and template string. + templates = tl.Dict(allow_none=True) + + # The output template is a dictionary of keyname and filled string. + filled_templates = tl.Dict(allow_none=True) + + # the output metadata is a dictionary of keyname and metadata of this template. + metadata = tl.Dict(allow_none=True) + + def __init__(self): + # A placeholder for the template variables widget. + self.template_variables = ipw.VBox() + + # A dictionary of mapping variables. + # the key is the variable name, and the value is a tuple of (template value and widget). + self._template_variables = {} + self._help_text = ipw.HTML( + """
Please fill the template variables below.
""" + ) + self._help_text.layout.display = "none" + + # unfilled_variables is a list of variables that are not filled. + self.unfilled_variables = [] + + super().__init__( + children=[ + self._help_text, + self.template_variables, + ] + ) + + def reset(self): + """Reset the widget.""" + self.templates = {} + self.filled_templates = {} + self._template_variables = {} + self._help_text.layout.display = "none" + self.template_variables.children = [] + + @tl.observe("templates") + def _templates_changed(self, _=None): + """Render the template variables widget.""" + # reset traits and then render the widget. + self._template_variables = {} + self.filled_templates = {} + self._help_text.layout.display = "none" + + self._render() + + self.metadata = self.templates.get("metadata", {}) + + # Update the output filled template. + # After `self._render` all the widgets are created. + # We can use the widget value to fill the template even not all the widgets are filled. + self.fill() + + def _render(self): + """Render the template variables widget.""" + metadata = self.templates.get("metadata", {}) + tooltip = metadata.get("tooltip", None) + + if tooltip: + self._help_text.value = f"""
{tooltip}
""" + + for line_key, line_str in self.templates.items(): + env = jinja2.Environment() + parsed_content = env.parse(line_str) + + # vars is a set of variables in the template + line_vars = jinja2_meta.find_undeclared_variables(parsed_content) + + # Create a widget for each variable. + # The var is the name in a template string + var_metadata = metadata.get("template_variables", {}) + for var in line_vars: + # one var can be used in multiple templates, so we need to keep track of the mapping with the set of variables. + var_meta = var_metadata.get(var, {}) + if var in self._template_variables: + # use the same widget for the same variable. + temp_var = self._template_variables[var] + w = temp_var.widget + lines = temp_var.lines + lines.append(TemplateVariableLine(line_key, line_str, line_vars)) + template_var = TemplateVariable(w, lines) + + self._template_variables[var] = template_var + else: + # create a new widget for the variable. + description = var_meta.get("key_display", f"{var}") + widget_type = var_meta.get("type", "text") + if widget_type == "text": + w = ipw.Text( + description=f"{description}:", + value=var_meta.get("default", ""), + # delay notifying the observers until the user stops typing + continuous_update=False, + layout=LAYOUT, + style=STYLE, + ) + elif widget_type == "list": + w = ipw.Dropdown( + description=f"{description}:", + options=var_meta.get("options", ()), + value=var_meta.get("default", None), + layout=LAYOUT, + style=STYLE, + ) + else: + raise ValueError(f"Invalid widget type {widget_type}") + + # Every time the value of the widget changes, we update the filled template. + # This migth be too much to sync the final filled template every time. + w.observe(self._on_template_variable_filled, names="value") + + template_var = TemplateVariable( + w, [TemplateVariableLine(line_key, line_str, line_vars)] + ) + self._template_variables[var] = template_var + + # Render by change the VBox children of placeholder. + self.template_variables.children = [ + # widget is shared so we only need to get the first one. + template_var.widget + for template_var in self._template_variables.values() + ] + + # Show the help text if there are template variables. + if self.template_variables.children: + self._help_text.layout.display = "block" + + def fill(self): + """Use template and current widgets value to fill the template + and update the filled_template. + """ + self.unfilled_variables = [] + + # Remove the metadata from the templates because for the final + # filled template we don't pass it to setup traits. + filled_templates = copy.deepcopy(self.templates) + if "metadata" in filled_templates: + del filled_templates["metadata"] + + for template_var in self._template_variables.values(): + for line in template_var.lines: + # update the filled template. + # if the widget is not filled, use the original template var string e.g. {{ var }} + inp_dict = {} + for _var in line.vars: + if self._template_variables[_var].widget.value == "": + variable_key = self._template_variables[_var].widget.description + self.unfilled_variables.append(variable_key.strip(":")) + inp_dict[_var] = "{{ " + _var + " }}" + else: + inp_dict[_var] = self._template_variables[_var].widget.value + + # re-render the template + env = jinja2.Environment() + filled_str = env.from_string(line.str).render(**inp_dict) + + # Update the filled template. + filled_templates[line.key] = filled_str + + # assign back to trigger the trait change. + self.filled_templates = filled_templates + + def _on_template_variable_filled(self, _): + """Callback when a template variable is filled.""" + self.fill() + + +class _ResourceSetupBaseWidget(ipw.VBox): + """The widget that allows to setup a computer and code. + This is the building block of the `ComputationalResourcesDatabaseWidget` which + will be directly used by the user. + """ + + success = tl.Bool(False) + message = tl.Unicode() + + ssh_auth = None # store the ssh auth type. Can be "password" or "2FA" + + def __init__( + self, + default_calc_job_plugin=None, + enable_quick_setup=True, + enable_detailed_setup=True, + ): + if not any((enable_detailed_setup, enable_quick_setup)): + raise ValueError( # noqa + "At least one of `enable_quick_setup` and `enable_detailed_setup` should be True." + ) + + self.enable_quick_setup = enable_quick_setup + self.enable_detailed_setup = enable_detailed_setup + + self.quick_setup_button = ipw.Button( + description="Quick setup", + tooltip="Setup a computer and code in one click.", + icon="rocket", + button_style="success", + ) + self.quick_setup_button.on_click(self._on_quick_setup) + + reset_button = ipw.Button( + description="Reset", + tooltip="Reset the resource.", + icon="refresh", + button_style="primary", + ) + reset_button.on_click(self._on_reset) + + # Resource database for setup computer/code. + + self.comp_resources_database = ComputationalResourcesDatabaseWidget( + default_calc_job_plugin=default_calc_job_plugin, + show_reset_button=False, + ) + + # All templates + self.template_computer_setup = TemplateVariablesWidget() + self.template_computer_configure = TemplateVariablesWidget() + self.template_code = TemplateVariablesWidget() + + # SSH + self.ssh_computer_setup = SshComputerSetup() + tl.dlink( + (self.ssh_computer_setup, "message"), + (self, "message"), + ) + + # Computer setup. + self.aiida_computer_setup = AiidaComputerSetup() + tl.dlink( + (self.aiida_computer_setup, "message"), + (self, "message"), + ) + + # Code setup. + self.aiida_code_setup = AiidaCodeSetup() + tl.dlink( + (self.aiida_code_setup, "message"), + (self, "message"), + ) + + # Data flow. + + # Computer template -> SSH connection setup. + tl.dlink( + (self.comp_resources_database, "computer_configure"), + (self.template_computer_configure, "templates"), + ) + tl.dlink( + (self.template_computer_configure, "filled_templates"), + (self.ssh_computer_setup, "ssh_config"), + ) + + # Computer template -> Computer setup. + tl.dlink( + (self.comp_resources_database, "computer_setup"), + (self.template_computer_setup, "templates"), + ) + tl.dlink( + (self.template_computer_setup, "filled_templates"), + (self.aiida_computer_setup, "computer_setup"), + ) + + # Computer template -> Computer configure. + tl.dlink( + (self.template_computer_configure, "filled_templates"), + (self.aiida_computer_setup, "computer_configure"), + ) + self.template_computer_configure.observe( + self._on_template_computer_configure_metadata_change, "metadata" + ) + self.aiida_computer_setup.on_setup_computer_success( + self._on_setup_computer_success + ) + + # Code template -> Code setup. + tl.dlink( + (self.comp_resources_database, "code_setup"), + (self.template_code, "templates"), + ) + tl.dlink( + (self.template_code, "filled_templates"), + (self.aiida_code_setup, "code_setup"), + ) + self.aiida_code_setup.on_setup_code_success(self._on_setup_code_success) + + # The widget for the detailed setup. + description_toggle_detail_setup = ipw.HTML( + """
Tick checkbox to setup resource step by step.
""" + ) + self.toggle_detail_setup = ipw.Checkbox( + description="", + value=False, + indent=False, + # this narrow the checkbox to the minimum width. + layout=ipw.Layout(width="30px"), + ) + self.toggle_detail_setup.observe(self._on_toggle_detail_setup, names="value") + self.detailed_setup_switch_widgets = ipw.HBox( + children=[ + self.toggle_detail_setup, + description_toggle_detail_setup, + ], + layout=ipw.Layout(justify_content="flex-end"), + ) + + detailed_setup_description_text = """
+ Go through the steps to setup SSH connection to remote machine, computer, and code into database.
+ The SSH connection step can be skipped and setup afterwards.
+ +
""" + description = ipw.HTML(detailed_setup_description_text) + + detailed_setup = ipw.Tab( + children=[ + self.ssh_computer_setup, + self.aiida_computer_setup, + self.aiida_code_setup, + ], + ) + detailed_setup.set_title(0, "SSH connection") + detailed_setup.set_title(1, "Computer") + detailed_setup.set_title(2, "Code") + + self.detailed_setup_widget = ipw.VBox( + children=[ + description, + detailed_setup, + ] + ) + + super().__init__( + children=[ + ipw.HTML( + """
Please select the computer/code from a database to pre-fill the fields below.
+ """ + ), + self.comp_resources_database, + self.template_computer_setup, + self.template_code, + self.template_computer_configure, + self.ssh_computer_setup.password_box, + self.detailed_setup_switch_widgets, + ipw.HBox( + children=[ + self.quick_setup_button, + reset_button, + ], + layout=ipw.Layout(justify_content="flex-end"), + ), + self.detailed_setup_widget, + ], + ) + + # update the layout + self._update_layout() + + def _update_layout(self): + """Update the layout to hide or show the bundled quick_setup/detailed_setup.""" + # check if the password is asked for ssh connection. + if self.ssh_auth != "password": + self.ssh_computer_setup.password_box.layout.display = "none" + else: + self.ssh_computer_setup.password_box.layout.display = "block" + + # If both quick and detailed setup are enabled + # - show the switch widget + # - show the quick setup widget (database + template variables + poppud password box + quick setup button) + # - hide the detailed setup widget as default + if self.enable_detailed_setup and self.enable_quick_setup: + self.detailed_setup_widget.layout.display = "none" + return + + # If only quick setup is enabled + # - hide the switch widget + # - hide the detailed setup widget + if self.enable_quick_setup: + self.detailed_setup_switch_widgets.layout.display = "none" + self.detailed_setup_widget.layout.display = "none" + return + + # If only detailed setup is enabled + # - hide the switch widget + # - hide the quick setup widget + # - hide the database widget + # which means only the detailed setup widget is shown. + if self.enable_detailed_setup: + self.children = [ + self.detailed_setup_widget, + ] + + def _on_toggle_detail_setup(self, change): + """When the checkbox is toggled, show/hide the detailed setup.""" + if change["new"]: + self.detailed_setup_widget.layout.display = "block" + self.quick_setup_button.disabled = True + # fill the template variables with the default values or the filled values. + # If the template variables are not all filled raise a warning. + else: + self.detailed_setup_widget.layout.display = "none" + self.quick_setup_button.disabled = False + + def _on_template_computer_configure_metadata_change(self, change): + """callback when the metadata of computer_configure template is changed.""" + if change["new"] is None: + return + + # decide whether to show the ssh password box widget. + # Since for 2FA ssh credential, the password are not needed but set from + # independent mechanism. + self.ssh_auth = change["new"].get("ssh_auth", None) + if self.ssh_auth is None: + self.ssh_auth = "password" + + self._update_layout() + + def _on_reset(self, _=None): + """Reset the database and the widget.""" + with self.hold_trait_notifications(): + self.reset() + self.comp_resources_database.reset() + + def _on_quick_setup(self, _=None): + """Go through all the setup steps automatically.""" + # Raise error if the computer is not selected. + if not self.aiida_computer_setup.some_values_provided: + self.message = wrap_message( + "Please select a computer from the database.", + MessageLevel.ERROR, + ) + return + + # Check if all the template variables are filled. + # If not raise a warning and return (skip the setup). + if ( + unfilled_variables := self.template_computer_setup.unfilled_variables + + self.template_computer_configure.unfilled_variables + + self.template_code.unfilled_variables + ): + var_warn_message = ", ".join([f"{v}" for v in unfilled_variables]) + self.message = wrap_message( + f"Please fill the template variables: {var_warn_message}", + MessageLevel.WARNING, + ) + return + + # Setup the computer and code. + if self.aiida_computer_setup.on_setup_computer(): + self.aiida_code_setup.on_setup_code() + + # Prepare the ssh key pair and copy to remote computer. + # This only happend when the ssh_auth is password. + if self.ssh_auth == "password": + try: + self.ssh_computer_setup.key_pair_prepare() + except ValueError as exc: + self.message = wrap_message( + f"Key pair generation failed: {exc}", + level=MessageLevel.ERROR, + ) + + self.ssh_computer_setup.thread_ssh_copy_id() + + # For not password ssh auth, key pair is not needed. + # There are other mechanism to set up the ssh connection. + # But we still need to write the ssh config to the ssh config file for such as + # proxy jump. + if not self.ssh_computer_setup._is_in_config(): + self.ssh_computer_setup._write_ssh_config() + + def _on_setup_computer_success(self): + """Callback that is called when the computer is successfully set up.""" + # update the computer dropdown list of code setup + self.aiida_code_setup.refresh() + + # select the computer in the dropdown list for code setup + label = self.aiida_computer_setup.label.value + self.aiida_code_setup.computer.select_by_label(label) + + def _on_setup_code_success(self): + """Callback that is called when the code is successfully set up.""" + self.success = True + + def reset(self): + """Reset the widget.""" + # reset the database + self.comp_resources_database.reset() + + # reset template variables + self.template_computer_setup.reset() + self.template_computer_configure.reset() + self.template_code.reset() + + # reset sub widgets + self.aiida_code_setup._reset() + self.aiida_computer_setup._reset() + + self.ssh_computer_setup._reset() + + # essential, since if not, the same computer_configure won't trigger the `_observe_ssh_config` callback. + self.ssh_computer_setup.ssh_config = {} + + # reset traits + self.computer_setup = {} + self.computer_configure = {} + self.code_setup = {} + + self.message = "" + self.success = False + self.ssh_auth = None + + # The layout also reset + self._update_layout() + + # untick the detailed switch checkbox + self.toggle_detail_setup.value = False diff --git a/aiidalab_widgets_base/databases.py b/aiidalab_widgets_base/databases.py index 2302e737..9b7bd149 100644 --- a/aiidalab_widgets_base/databases.py +++ b/aiidalab_widgets_base/databases.py @@ -1,9 +1,13 @@ """Widgets that allow to query online databases.""" +import re + import ase import ipywidgets as ipw import requests import traitlets as tl +TEMPLATE_PATTERN = re.compile(r"\{\{.*\}\}") + class CodQueryWidget(ipw.VBox): """Query structures in Crystallography Open Database (COD) @@ -204,179 +208,211 @@ def _update_structure(self, change: dict) -> None: class ComputationalResourcesDatabaseWidget(ipw.VBox): """Extract the setup of a known computer from the AiiDA code registry.""" - default_calc_job_plugin = tl.Unicode(allow_none=True) - ssh_config = tl.Dict() + _default_database_source = ( + "https://aiidateam.github.io/aiida-resource-registry/database.json" + ) + + database_source = tl.Unicode(allow_none=True) + computer_setup = tl.Dict() + computer_configure = tl.Dict() code_setup = tl.Dict() - database = tl.Dict() - def __init__(self, **kwargs): + STYLE = {"description_width": "180px"} + LAYOUT = {"width": "400px"} + + def __init__( + self, + default_calc_job_plugin=None, + database_source=None, + show_reset_button=True, + **kwargs, + ): + if database_source is None: + database_source = self._default_database_source + + self.default_calc_job_plugin = default_calc_job_plugin + # Select domain. - self.inp_domain = ipw.Dropdown( - options=[], - description="Domain", + self.domain_selector = ipw.Dropdown( + options=(), + description="Domain:", disabled=False, + layout=self.LAYOUT, + style=self.STYLE, ) - self.inp_domain.observe(self._domain_changed, names=["value", "options"]) + self.domain_selector.observe(self._domain_changed, names=["value", "options"]) # Select computer. - self.inp_computer = ipw.Dropdown( - options=[], + self.computer_selector = ipw.Dropdown( + options=(), description="Computer:", disabled=False, + layout=self.LAYOUT, + style=self.STYLE, + ) + self.computer_selector.observe( + self._computer_changed, names=["value", "options"] ) - self.inp_computer.observe(self._computer_changed, names=["value", "options"]) # Select code. - self.inp_code = ipw.Dropdown( - options=[], + self.code_selector = ipw.Dropdown( + options=(), description="Code:", disabled=False, + layout=self.LAYOUT, + style=self.STYLE, ) - self.inp_code.observe(self._code_changed, names=["value", "options"]) + self.code_selector.observe(self._code_changed, names=["value", "options"]) - btn_reset = ipw.Button(description="Reset") - btn_reset.on_click(self._reset) + reset_button = ipw.Button(description="Reset") + reset_button.on_click(self.reset) + + if show_reset_button is False: + reset_button.layout.display = "none" super().__init__( children=[ - self.inp_domain, - self.inp_computer, - self.inp_code, - btn_reset, + self.domain_selector, + self.computer_selector, + self.code_selector, + reset_button, ], **kwargs, ) - self.update() - self._reset() - - def _reset(self, _=None): - self.inp_domain.value = None - - def clean_up_database(self, database, plugin): - for domain in list(database.keys()): - for computer in list(database[domain].keys() - {"default"}): - for code in list( - database[domain][computer].keys() - - {"computer-configure", "computer-setup"} - ): - if ( - plugin - != database[domain][computer][code]["default_calc_job_plugin"] - ): - del database[domain][computer][code] - # If no codes remained that correspond to the chosen plugin, remove the computer. - if ( - len( - database[domain][computer].keys() - - {"computer-configure", "computer-setup"} + self.database_source = database_source + self.reset() + + def reset(self, _=None): + """Reset widget and traits""" + with self.hold_trait_notifications(): + self.domain_selector.value = None + self.computer_selector.value = None + self.code_selector.value = None + + @tl.observe("database_source") + def _database_source_changed(self, _=None): + self.database = self._database_generator( + self.database_source, self.default_calc_job_plugin + ) + + # Update domain selector. + self.domain_selector.options = self.database.keys() + self.reset() + + @staticmethod + def _database_generator(database_source, default_calc_job_plugin): + """From database source JSON and default calc job plugin, generate resource database""" + try: + database = requests.get(database_source).json() + except Exception: + database = {} + + if default_calc_job_plugin is None: + return database + + # filter database by default calc job plugin + for domain, domain_value in database.copy().items(): + for computer, computer_value in domain_value.copy().items(): + if computer == "default": + # skip default computer + continue + + for code, code_value in list(computer_value["codes"].items()): + # check if code plugin matches default plugin + # if not, remove code + plugin = re.sub( + TEMPLATE_PATTERN, r".*", code_value["default_calc_job_plugin"] ) - == 0 - ): - del database[domain][computer] - # If no computers remained - remove the domain. - if len(database[domain].keys() - {"default"}) == 0: + if re.match(plugin, default_calc_job_plugin) is None: + del computer_value["codes"][code] + + if len(computer_value["codes"]) == 0: + # remove computer since no codes defined in this computer source + del domain_value[computer] + if domain_value.get("default") == computer: + # also remove default computer from domain + del domain_value["default"] + + if len(domain_value) == 0: + # remove domain since no computers with required codes defined in this domain source del database[domain] - # Making sure the 'default' key still points to an existing computer. - elif database[domain]["default"] not in database[domain]: - database[domain]["default"] = sorted( - database[domain].keys() - {"default"} - )[0] - return database + continue - def update(self, _=None): - database = requests.get( - "https://aiidateam.github.io/aiida-code-registry/database_v2_1.json" - ).json() - self.database = ( - self.clean_up_database(database, self.default_calc_job_plugin) - if self.default_calc_job_plugin - else database - ) + if domain_value["default"] not in domain_value: + # make sure default computer is still points to existing computer + domain_value["default"] = sorted(domain_value.keys() - {"default"})[0] - @tl.observe("database") - def _observer_database_change(self, _=None): - self.inp_domain.options = self.database.keys() - self._reset() + return database - def _domain_changed(self, _=None): - with self.hold_trait_notifications(): - self.inp_computer.value = None - try: - self.inp_computer.options = [ - key - for key in self.database[self.inp_domain.value].keys() - if key != "default" - ] - self.inp_computer.value = self.database[self.inp_domain.value][ - "default" - ] - except KeyError: - self.inp_computer.options = [] + def _domain_changed(self, change=None): + """callback when new domain selected""" + with self.hold_trait_notifications(): # To prevent multiple calls to callbacks + if change["new"] is None: + self.computer_selector.options = () + self.computer_selector.value = None + self.code_selector.options = () + self.code_selector.value = None return + else: + selected_domain = self.domain_selector.value - def _computer_changed(self, _=None): - """Read settings from the YAML files and populate self.database with them.""" + with self.hold_trait_notifications(): + self.computer_selector.options = tuple( + key + for key in self.database.get(selected_domain, {}).keys() + if key != "default" + ) + self.computer_selector.value = self.database.get( + selected_domain, {} + ).get("default") + def _computer_changed(self, change=None): + """callback when new computer selected""" with self.hold_trait_notifications(): - if self.inp_computer.value is None: - self.inp_code.options = [] - self.inp_code.value = None - self.ssh_config = {} - self.computer_setup = {} + if change["new"] is None: + self.code_selector.options = () + self.code_selector.value = None + self.computer_setup, self.computer_configure = {}, {} return + else: + self.code_selector.value = None + selected_computer = self.computer_selector.value + + selected_domain = self.domain_selector.value + + computer_dict = self.database.get(selected_domain, {}).get( + selected_computer, {} + ) + + self.code_selector.options = list(computer_dict.get("codes", {}).keys()) + self.code_selector.value = None - self.inp_code.options = [ - key - for key in self.database[self.inp_domain.value][ - self.inp_computer.value - ].keys() - if key not in ("computer-setup", "computer-configure") - ] - - setup = self.database[self.inp_domain.value][self.inp_computer.value][ - "computer-setup" - ] - config = self.database[self.inp_domain.value][self.inp_computer.value][ - "computer-configure" - ] - ssh_config = {"hostname": setup["hostname"]} - if "proxy_command" in config: - ssh_config["proxy_command"] = config["proxy_command"] - if "proxy_jump" in config: - ssh_config["proxy_jump"] = config["proxy_jump"] - - self.ssh_config = ssh_config # To notify the trait change - - self.computer_setup = { - "setup": setup, - "configure": config, - } - - self._code_changed() - - def _code_changed(self, _=None): + computer_setup = computer_dict.get("computer", {}).get("computer-setup", {}) + computer_configure = computer_dict.get("computer", {}).get( + "computer-configure", {} + ) + + if "hostname" in computer_setup: + computer_configure["hostname"] = computer_setup.get("hostname") + + self.computer_setup = computer_setup + self.computer_configure = computer_configure + + def _code_changed(self, change=None): """Update code settings.""" - with self.hold_trait_notifications(): - if self.inp_code.value is None: - self.code_setup = {} - return - code_setup = self.database[self.inp_domain.value][self.inp_computer.value][ - self.inp_code.value - ] - - if ( - "label" - in self.database[self.inp_domain.value][self.inp_computer.value][ - "computer-setup" - ] - ): - code_setup["computer"] = self.database[self.inp_domain.value][ - self.inp_computer.value - ]["computer-setup"]["label"] - - self.code_setup = code_setup - - @tl.default("default_calc_job_plugin") - def _default_calc_job_plugin(self): - return None + if change["new"] is None: + self.code_setup = {} + return + else: + selected_code = self.code_selector.value + + selected_domain = self.domain_selector.value + selected_computer = self.computer_selector.value + + self.code_setup = ( + self.database.get(selected_domain, {}) + .get(selected_computer, {}) + .get("codes", {}) + .get(selected_code, {}) + ) diff --git a/aiidalab_widgets_base/utils/__init__.py b/aiidalab_widgets_base/utils/__init__.py index 69d51898..4b8b0b37 100644 --- a/aiidalab_widgets_base/utils/__init__.py +++ b/aiidalab_widgets_base/utils/__init__.py @@ -1,5 +1,6 @@ """Some utility functions used acrross the repository.""" import threading +from enum import Enum from typing import Any, Tuple import ipywidgets as ipw @@ -172,6 +173,33 @@ def _observe_message(self, change): self.show_temporary_message(change["new"]) +# Define the message levels as Enum +class MessageLevel(Enum): + INFO = "info" + WARNING = "warning" + ERROR = "danger" + SUCCESS = "success" + + +def wrap_message(message, level=MessageLevel.INFO): + """Wrap message into HTML code with the given level.""" + # mapping level to fa icon + # https://fontawesome.com/v4.7.0/icons/ + mapping = { + MessageLevel.INFO: "info-circle", + MessageLevel.WARNING: "exclamation-triangle", + MessageLevel.ERROR: "exclamation-circle", + MessageLevel.SUCCESS: "check-circle", + } + + # The message is wrapped into a div with the class "alert" and the icon of the given level + return f""" + + """ + + def ase2spglib(ase_structure: Atoms) -> Tuple[Any, Any, Any]: """ Convert ase Atoms instance to spglib cell in the format defined at diff --git a/notebooks/computational_resources.ipynb b/notebooks/computational_resources.ipynb index 4278471e..9d31b5fa 100644 --- a/notebooks/computational_resources.ipynb +++ b/notebooks/computational_resources.ipynb @@ -51,7 +51,8 @@ "outputs": [], "source": [ "resources1 = awb.ComputationalResourcesWidget()\n", - "resources2 = awb.ComputationalResourcesWidget()" + "resources2 = awb.ComputationalResourcesWidget(enable_quick_setup=False)\n", + "resources3 = awb.ComputationalResourcesWidget(enable_detailed_setup=False)" ] }, { @@ -62,7 +63,8 @@ "outputs": [], "source": [ "display(resources1)\n", - "display(resources2)" + "display(resources2)\n", + "display(resources3)" ] } ], diff --git a/tests/test_computational_resources.py b/tests/test_computational_resources.py index 53651cdc..b850c302 100644 --- a/tests/test_computational_resources.py +++ b/tests/test_computational_resources.py @@ -1,19 +1,21 @@ +import re +from pathlib import Path + import pytest from aiida import orm from aiidalab_widgets_base import computational_resources +from aiidalab_widgets_base.computational_resources import ( + ComputationalResourcesWidget, + _ResourceSetupBaseWidget, +) +HTML_TAG_CLEANER = re.compile(r"<[^>]*>") -@pytest.mark.usefixtures("aiida_profile_clean") -def test_computational_resources_widget(aiida_local_code_bash): - """Test the ComputationalResourcesWidget.""" - widget = computational_resources.ComputationalResourcesWidget( - default_calc_job_plugin="bash" - ) - # Get the list of currently installed codes. - codes = widget._get_codes() - assert "bash@localhost" == codes[0][0] +def clean_html(raw_html): + """Remove html tags from a string.""" + return re.sub(HTML_TAG_CLEANER, "", raw_html) @pytest.mark.usefixtures("aiida_profile_clean") @@ -50,11 +52,28 @@ def test_ssh_computer_setup_widget(monkeypatch, tmp_path): widget._ssh_keygen() # Create non-default private key file. - fpath = widget._add_private_key("my_key_name", b"my_key_content") - assert fpath.exists() - with open(fpath) as f: + private_key_path = tmp_path / ".ssh" / "my_key_name" + widget._add_private_key(private_key_path, b"my_key_content") + assert private_key_path.exists() + with open(private_key_path) as f: assert f.read() == "my_key_content" + # set private key with same name to trigger the rename operation + widget._verification_mode.value = "private_key" + # mock _private_key to mimic the upload of the private key + monkeypatch.setattr( + "aiidalab_widgets_base.computational_resources.SshComputerSetup._private_key", + property(lambda _: ("my_key_name", b"my_key_content_new")), + ) + # check the private key is renamed, monkeypatch the shortuuid to make the test deterministic + monkeypatch.setattr("shortuuid.uuid", lambda: "00001111") + + widget._on_setup_ssh_button_pressed() + + assert "my_key_name-00001111" in [ + str(p.name) for p in Path(tmp_path / ".ssh").iterdir() + ] + # Setting the ssh_config to an empty dictionary should reset the widget. widget.ssh_config = {} assert widget.hostname.value == "" @@ -65,13 +84,10 @@ def test_ssh_computer_setup_widget(monkeypatch, tmp_path): @pytest.mark.usefixtures("aiida_profile_clean") -def test_aiida_computer_setup_widget(monkeypatch): - """Test the AiidaComputerSetup.""" - # monkeypatch the parse_sshconfig - monkeypatch.setattr( - "aiida.transports.plugins.ssh.parse_sshconfig", - lambda _: {"hostname": "daint.cscs.ch", "user": "aiida"}, - ) +def test_aiida_computer_setup_widget_default(): + """Test the AiidaComputerSetup. + The 'default' in name means the username is from computer configuration. + """ widget = computational_resources.AiidaComputerSetup() # At the beginning, the computer_name should be an empty string. @@ -79,29 +95,29 @@ def test_aiida_computer_setup_widget(monkeypatch): # Preparing the computer setup. computer_setup = { - "setup": { - "label": "daint", - "hostname": "daint.cscs.ch", - "description": "Daint supercomputer", - "work_dir": "/scratch/snx3000/{username}/aiida_run", - "mpirun_command": "srun -n {tot_num_mpiprocs}", - "default_memory_per_machine": 2000000000, - "mpiprocs_per_machine": 12, - "transport": "core.ssh", - "scheduler": "core.slurm", - "shebang": "#!/bin/bash", - "use_double_quotes": True, - "prepend_text": "#SBATCH --account=proj20", - "append_text": "", - }, - "configure": { - "proxy_jump": "ela.cscs.ch", - "safe_interval": 10, - "use_login_shell": True, - }, + "label": "daint", + "hostname": "daint.cscs.ch", + "description": "Daint supercomputer", + "work_dir": "/scratch/snx3000/{username}/aiida_run", + "mpirun_command": "srun -n {tot_num_mpiprocs}", + "default_memory_per_machine": 2000000000, + "mpiprocs_per_machine": 12, + "transport": "core.ssh", + "scheduler": "core.slurm", + "shebang": "#!/bin/bash", + "use_double_quotes": True, + "prepend_text": "#SBATCH --account=proj20", + "append_text": "", + } + computer_configure = { + "proxy_jump": "ela.cscs.ch", + "safe_interval": 10, + "use_login_shell": True, + "username": "aiida", } widget.computer_setup = computer_setup + widget.computer_configure = computer_configure assert widget.on_setup_computer() # Check that the computer is created. @@ -111,13 +127,79 @@ def test_aiida_computer_setup_widget(monkeypatch): # Reset the widget and check that a few attributes are reset. widget.computer_setup = {} + widget.computer_configure = {} assert widget.label.value == "" assert widget.hostname.value == "" assert widget.description.value == "" # Check that setup is failing if the configuration is missing. assert not widget.on_setup_computer() - assert widget.message.startswith("Please specify the computer name") + assert "Please specify the computer name" in widget.message + + +@pytest.mark.usefixtures("aiida_profile_clean") +def test_aiida_computer_setup_widget_ssh_username(monkeypatch, tmp_path): + """Test the AiidaComputerSetup. + The 'default' in name means the username is from computer configuration. + """ + # mock home directory for ssh config file + monkeypatch.setenv("HOME", str(tmp_path)) + + widget = computational_resources.AiidaComputerSetup() + + # At the beginning, the computer_name should be an empty string. + assert widget.label.value == "" + + # Preparing the computer setup. + computer_setup = { + "label": "daint", + "hostname": "daint.cscs.ch", + "description": "Daint supercomputer", + "work_dir": "/scratch/snx3000/{username}/aiida_run", + "mpirun_command": "srun -n {tot_num_mpiprocs}", + "default_memory_per_machine": 2000000000, + "mpiprocs_per_machine": 12, + "transport": "core.ssh", + "scheduler": "core.slurm", + "shebang": "#!/bin/bash", + "use_double_quotes": True, + "prepend_text": "#SBATCH --account=proj20", + "append_text": "", + } + computer_configure = { + "proxy_jump": "ela.cscs.ch", + "safe_interval": 10, + "use_login_shell": True, + } + + widget.computer_setup = computer_setup + widget.computer_configure = computer_configure + assert not widget.on_setup_computer() + assert "SSH username is not provided" in widget.message + + # monkeypatch the parse_sshconfig + # this will go through the ssh username from sshcfg + monkeypatch.setattr( + "aiida.transports.plugins.ssh.parse_sshconfig", + lambda _: {"hostname": "daint.cscs.ch", "user": "aiida"}, + ) + + assert widget.on_setup_computer() + + computer = orm.load_computer("daint") + assert computer.label == "daint" + assert computer.hostname == "daint.cscs.ch" + + # Reset the widget and check that a few attributes are reset. + widget.computer_setup = {} + widget.computer_configure = {} + assert widget.label.value == "" + assert widget.hostname.value == "" + assert widget.description.value == "" + + # Check that setup is failing if the configuration is missing. + assert not widget.on_setup_computer() + assert "Please specify the computer name" in widget.message @pytest.mark.usefixtures("aiida_profile_clean") @@ -130,28 +212,27 @@ def test_aiida_localhost_setup_widget(): # Preparing the computer setup. computer_setup = { - "setup": { - "label": "localhosttest", - "hostname": "localhost", - "description": "locahost computer", - "work_dir": "/home/jovyan/aiida_run", - "mpirun_command": "srun -n {tot_num_mpiprocs}", - "default_memory_per_machine": 2000000000, - "mpiprocs_per_machine": 2, - "transport": "core.local", - "scheduler": "core.direct", - "shebang": "#!/bin/bash", - "use_double_quotes": True, - "prepend_text": "", - "append_text": "", - }, - "configure": { - "safe_interval": 10, - "use_login_shell": True, - }, + "label": "localhosttest", + "hostname": "localhost", + "description": "locahost computer", + "work_dir": "/home/jovyan/aiida_run", + "mpirun_command": "srun -n {tot_num_mpiprocs}", + "default_memory_per_machine": 2000000000, + "mpiprocs_per_machine": 2, + "transport": "core.local", + "scheduler": "core.direct", + "shebang": "#!/bin/bash", + "use_double_quotes": True, + "prepend_text": "", + "append_text": "", + } + computer_configure = { + "safe_interval": 10, + "use_login_shell": True, } widget.computer_setup = computer_setup + widget.computer_configure = computer_configure assert widget.on_setup_computer() # Check that the computer is created. @@ -161,6 +242,7 @@ def test_aiida_localhost_setup_widget(): # Reset the widget and check that a few attributes are reset. widget.computer_setup = {} + widget.computer_configure = {} assert widget.label.value == "" assert widget.hostname.value == "" assert widget.description.value == "" @@ -222,3 +304,492 @@ def test_computer_dropdown_widget(aiida_localhost): # Trying to set the dropdown value to None widget.value = None assert widget._dropdown.value is None + + +def test_template_variables_widget(): + """Test template_variables_widget.""" + w = computational_resources.TemplateVariablesWidget() + + w.templates = { + "label": "{{ label }}", + "hostname": "daint.cscs.ch", + "description": "Piz Daint supercomputer at CSCS Lugano, Switzerland, multicore partition.", + "transport": "core.ssh", + "scheduler": "core.slurm", + "work_dir": "/scratch/snx3000/{username}/aiida_run/", + "shebang": "#!/bin/bash", + "mpirun_command": "srun -n {tot_num_mpiprocs}", + "mpiprocs_per_machine": 36, + "prepend_text": "#SBATCH --partition={{ slurm_partition }}\n#SBATCH --account={{ slurm_account }}\n#SBATCH --constraint=mc\n#SBATCH --cpus-per-task=1\n\nexport OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK\nsource $MODULESHOME/init/bash\nulimit -s unlimited", + "metadata": { + "slurm_partition": { + "type": "text", + "key_display": "Slurm partition", + }, + }, + } + + # Fill the template variables + for key, value in w._template_variables.items(): + if key == "label": + sub_widget = value.widget + sub_widget.value = "daint-mc-test" + + # check the filled value is updated in the filled template + assert w.filled_templates["label"] == "daint-mc-test" + + # Fill two template variables in one template line + for key, value in w._template_variables.items(): + if key == "slurm_partition": + sub_widget = value.widget + sub_widget.value = "normal-test" + + elif key == "slurm_account": + sub_widget = value.widget + sub_widget.value = "newuser" + + # check the filled value is updated in the filled template + assert ( + w.filled_templates["prepend_text"] + == "#SBATCH --partition=normal-test\n#SBATCH --account=newuser\n#SBATCH --constraint=mc\n#SBATCH --cpus-per-task=1\n\nexport OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK\nsource $MODULESHOME/init/bash\nulimit -s unlimited" + ) + + # Test the filled template is updated when the filled value is updated. + for key, value in w._template_variables.items(): + if key == "slurm_partition": + sub_widget = value.widget + sub_widget.value = "debug" + + assert "debug" in w.filled_templates["prepend_text"] + + +def test_template_variables_widget_metadata(): + """Test metadata support in template_variables_widget.""" + w = computational_resources.TemplateVariablesWidget() + + w.templates = { + "prepend_text": "#SBATCH --partition={{ slurm_partition }}\n#SBATCH --account={{ slurm_account }}\n#SBATCH --constraint=mc\n#SBATCH --cpus-per-task=1\n\nexport OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK\nsource $MODULESHOME/init/bash\nulimit -s unlimited", + "metadata": { + "template_variables": { + "slurm_partition": { + "type": "list", + "default": "normal", + "options": ["normal", "normal-test", "debug"], + "key_display": "Slurm partition", + }, + "slurm_account": { + "type": "text", + "key_display": "Slurm account", + }, + }, + }, + } + + # Test the default value is filled in correctly. + assert ( + w.filled_templates["prepend_text"] + == "#SBATCH --partition=normal\n#SBATCH --account={{ slurm_account }}\n#SBATCH --constraint=mc\n#SBATCH --cpus-per-task=1\n\nexport OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK\nsource $MODULESHOME/init/bash\nulimit -s unlimited" + ) + + # Fill two template variables in one template line + for key, value in w._template_variables.items(): + if key == "slurm_partition": + sub_widget = value.widget + + # Test set the widget from the metadata + assert sub_widget.description == "Slurm partition:" + assert sub_widget.options == ("normal", "normal-test", "debug") + assert sub_widget.value == "normal" + + sub_widget.value = "normal-test" + + elif key == "slurm_account": + sub_widget = value.widget + sub_widget.value = "newuser" + + # check the filled value is updated in the filled template + assert ( + w.filled_templates["prepend_text"] + == "#SBATCH --partition=normal-test\n#SBATCH --account=newuser\n#SBATCH --constraint=mc\n#SBATCH --cpus-per-task=1\n\nexport OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK\nsource $MODULESHOME/init/bash\nulimit -s unlimited" + ) + assert "metadata" not in w.filled_templates + + +def test_template_variables_widget_multi_template_variables(): + """This is the test for same key in multiple templates.""" + w = computational_resources.TemplateVariablesWidget() + + w.templates = { + "label": "{{ code_binary_name }}-7.2", + "description": "The code {{ code_binary_name }}.x of Quantum ESPRESSO compiled for daint-mc", + "default_calc_job_plugin": "quantumespresso.{{ code_binary_name }}", + "filepath_executable": "/apps/dom/UES/jenkins/7.0.UP03/21.09/dom-mc/software/QuantumESPRESSO/7.0-CrayIntel-21.09/bin/{{ code_binary_name }}.x", + "prepend_text": "module load daint-mc\nmodule load QuantumESPRESSO\n", + "append_text": "", + "metadata": { + "template_variables": { + "code_binary_name": { + "type": "list", + "options": ["pw", "ph", "pp"], + "key_display": "Code name", + }, + }, + }, + } + + # Fill the code_binary_name template variables for all + for key, value in w._template_variables.items(): + if key == "code_binary_name": + sub_widget = value.widget + assert sub_widget.description == "Code name:" + assert sub_widget.options == ("pw", "ph", "pp") + assert sub_widget.value is None + + sub_widget.value = "ph" + + # check all filed templates are updated + assert w.filled_templates["label"] == "ph-7.2" + assert ( + w.filled_templates["description"] + == "The code ph.x of Quantum ESPRESSO compiled for daint-mc" + ) + assert w.filled_templates["default_calc_job_plugin"] == "quantumespresso.ph" + + +def test_template_variables_widget_help_text_disappear_if_no_template_str(): + """This test when the template string is update to without template field, the help text should disappear.""" + w = computational_resources.TemplateVariablesWidget() + + # The initial help text should be not displayed. + assert w._help_text.layout.display == "none" + + w.templates = { + "label": "{{ code_binary_name }}-7.2", + } + + # The help text should be displayed. + assert w._help_text.layout.display == "block" + + # Fill the code_binary_name template variables for all + for key, value in w._template_variables.items(): + if key == "code_binary_name": + sub_widget = value.widget + sub_widget.value = "pw" + + w.templates = { + "label": "ph-7.2", + } + # The help text should be not displayed. + assert w._help_text.layout.display == "none" + + +@pytest.mark.usefixtures("aiida_profile_clean") +def test_resource_setup_widget_default(): + """Test the _ResourceSetupBaseWidget.""" + with pytest.raises(ValueError): + w = _ResourceSetupBaseWidget( + enable_detailed_setup=False, enable_quick_setup=False + ) + + w = _ResourceSetupBaseWidget() + + # Test message is update correctly. By click setup button without filling in any information. + w._on_quick_setup() + + # assert "Please select a computer from the database" in w.message + + # Test select a new resource setup will update the output interface (e.g. ssh_config, computer_setup, code_setup) + # and the computer/code setup widget will be updated accordingly. + w.comp_resources_database.domain_selector.value = "daint.cscs.ch" + w.comp_resources_database.computer_selector.value = "mc" + w.comp_resources_database.code_selector.value = "QE-7.2-exe-template" + + # Test before the template is filled, the warning message is displayed. + w._on_quick_setup() + assert "Please fill the template variables" in w.message + + # Fill in the computer name and trigger the setup button again, the message should be updated. + for ( + key, + mapping_variable, + ) in w.template_computer_setup._template_variables.items(): + if key == "label": + sub_widget = mapping_variable.widget + + # Test the default value is filled in correctly. + assert sub_widget.value == "daint-mc" + if key == "slurm_partition": + sub_widget = mapping_variable.widget + sub_widget.value = "debug" + if key == "slurm_account": + sub_widget = mapping_variable.widget + sub_widget.value = "newuser" + + # Fill the computer configure template variables + for ( + key, + mapping_variable, + ) in w.template_computer_configure._template_variables.items(): + if key == "username": + sub_widget = mapping_variable.widget + sub_widget.value = "aiida" + + # Fill the code name + for key, mapping_variable in w.template_code._template_variables.items(): + if key == "code_binary_name": + sub_widget = mapping_variable.widget + sub_widget.value = "ph" + + # select the other code and check the filled template is updated + sub_widget.value = "pw" + + w.ssh_computer_setup.username.value = "aiida" + + # Since cscs is 2FA, test the password box is not displayed. + assert w.ssh_computer_setup.password_box.layout.display == "none" + + w._on_quick_setup() + + assert w.success + assert orm.load_code("pw-7.2@daint-mc") + + # test select new resource reset the widget, success trait, and message trait. + w.reset() + + assert w.ssh_auth is None + assert w.aiida_computer_setup.computer_setup == {} + assert w.aiida_computer_setup.computer_configure == {} + assert w.aiida_code_setup.code_setup == {} + assert w.success is False + assert w.message == "" + assert w.template_code._help_text.layout.display == "none" + assert w.template_code._template_variables == {} + + # reselect after reset should update the output interface + w.comp_resources_database.domain_selector.value = "daint.cscs.ch" + assert w.template_computer_setup._template_variables != {} + + +@pytest.mark.usefixtures("aiida_profile_clean") +def test_resource_setup_widget_for_password_configure(monkeypatch, tmp_path): + """Test for computer configure with password as ssh auth. + The ssh auth is password, thus will generate ssh key pair and try to upload the key + """ + # monkeypatch home so the ssh key is generated in the temporary directory + monkeypatch.setenv("HOME", str(tmp_path)) + + w = _ResourceSetupBaseWidget() + + # Test select a new resource setup will update the output interface (e.g. ssh_config, computer_setup, code_setup) + # and the computer/code setup widget will be updated accordingly. + w.comp_resources_database.domain_selector.value = "merlin.psi.ch" + w.comp_resources_database.computer_selector.value = "cpu" + w.comp_resources_database.code_selector.value = "QE-7.0-exe-template" + + # Fill in the computer name and trigger the setup button again, the message should be updated. + for ( + key, + mapping_variable, + ) in w.template_computer_setup._template_variables.items(): + if key == "label": + sub_widget = mapping_variable.widget + + # Test the default value is filled in correctly. + assert sub_widget.value == "merlin-cpu" + + # Test the password box is displayed. + assert w.ssh_computer_setup.password_box.layout.display == "block" + + # Fill the computer configure template variables + for ( + key, + mapping_variable, + ) in w.template_computer_configure._template_variables.items(): + if key == "username": + sub_widget = mapping_variable.widget + sub_widget.value = "aiida" + + # Fill the code name + for key, mapping_variable in w.template_code._template_variables.items(): + if key == "code_binary_name": + sub_widget = mapping_variable.widget + sub_widget.value = "ph" + + # select the other code and check the filled template is updated + sub_widget.value = "pw" + + # The quick_setup with password auth will try connect which will timeout. + # Thus, mock the connect method to avoid the timeout. + monkeypatch.setattr( + "aiidalab_widgets_base.computational_resources.SshComputerSetup.thread_ssh_copy_id", + lambda _: None, + ) + w._on_quick_setup() + + assert w.success + # check the code is really created + assert orm.load_code("pw-7.0@merlin-cpu") + + # The key pair will be generated to the temporary directory + # Check the content of the config is correct + with open(tmp_path / ".ssh" / "config") as f: + content = f.read() + assert "User aiida" in content + assert "Host merlin-l-01.psi.ch" in content + + # After reset the password box should be hidden again. + w.reset() + assert w.ssh_computer_setup.password_box.layout.display == "none" + + +@pytest.mark.usefixtures("aiida_profile_clean") +def test_resource_setup_widget_computer_change_code_reset(): + """Test the _ResourceSetupBaseWidget that when computer template changed, the code selector widget is reset.""" + w = _ResourceSetupBaseWidget() + + # Test select a new resource setup will update the output interface (e.g. ssh_config, computer_setup, code_setup) + # and the computer/code setup widget will be updated accordingly. + w.comp_resources_database.domain_selector.value = "daint.cscs.ch" + w.comp_resources_database.computer_selector.value = "mc" + w.comp_resources_database.code_selector.value = "QE-7.2-exe-template" + + assert w.template_code._help_text.layout.display == "block" + + # Change the computer template, code template prompt box should stay. + w.comp_resources_database.computer_selector.value = "gpu" + + # check the option of resource database widget is reset + assert w.comp_resources_database.code_selector.value is None + + +def test_resource_setup_widget_detailed_setup(): + """Detail branch test of the resource setup widget""" + w = _ResourceSetupBaseWidget() + + w.comp_resources_database.domain_selector.value = "daint.cscs.ch" + w.comp_resources_database.computer_selector.value = "mc" + w.comp_resources_database.code_selector.value = "pw-7.0" + + # Test the detailed setup widget is displayed with the label updated because the + # information can get from the default of the template variables. + # Same for the slurm_partition and multithreading hint which has default from the template variables metadata. + assert w.aiida_computer_setup.label.value == "daint-mc" + assert "normal" in w.aiida_computer_setup.prepend_text.value + assert "nomultithread" in w.aiida_computer_setup.prepend_text.value + + # Check that the computer/code setup widget is updated accordingly in the detailed setup widget. + # By triggering the setup button one by one in the detailed setup widget, the message should be updated. + # check we the same aiida_computer_setup for resource and the detailed setup + assert id(w.detailed_setup_widget.children[1].children[0]) == id( + w.ssh_computer_setup + ) + assert id(w.detailed_setup_widget.children[1].children[1]) == id( + w.aiida_computer_setup + ) + assert id(w.detailed_setup_widget.children[1].children[2]) == id(w.aiida_code_setup) + computer_label = "daint-mc" + w.aiida_computer_setup.label.value = computer_label + w.aiida_computer_setup.on_setup_computer() + + assert "created" in w.message + + comp_uuid = orm.load_computer(computer_label).uuid + w.aiida_code_setup.computer._dropdown.value = comp_uuid + w.aiida_code_setup.on_setup_code() + + assert "created" in w.message + + # test select new resource reset the widget, success trait, and message trait, and the computer/code setup widget is cleared. + w.reset() + + assert w.aiida_computer_setup.computer_setup == {} + assert w.aiida_computer_setup.computer_configure == {} + assert w.aiida_code_setup.code_setup == {} + assert w.ssh_computer_setup.ssh_config == {} + assert w.success is False + assert w.message == "" + + +@pytest.mark.usefixtures("aiida_profile_clean") +def test_computer_resource_setup_widget_default(monkeypatch, tmp_path): + """A test for bundle widget `ComputationalResourcesWidget`.""" + # Test the enable toggle are passed to the sub widgets. + with pytest.raises(ValueError): + ComputationalResourcesWidget( + enable_detailed_setup=False, enable_quick_setup=False + ) + + # monkeypatch home so the ssh key is generated in the temporary directory + monkeypatch.setenv("HOME", str(tmp_path)) + + # Set with clear_after=1 to avoid the widget frozen at the end of test to wait the counting thread to finish. + w = ComputationalResourcesWidget(clear_after=1) + + # Go through a full setup process. + + # check no code is existing in AiiDA database + # by check dropdown is empty + assert w.code_select_dropdown.options == () + + w_resource = w.resource_setup + + w_resource.comp_resources_database.domain_selector.value = "merlin.psi.ch" + w_resource.comp_resources_database.computer_selector.value = "cpu" + w_resource.comp_resources_database.code_selector.value = "QE-7.0-exe-template" + + # Fill in the computer name and trigger the setup button again, the message should be updated. + for ( + key, + mapping_variable, + ) in w_resource.template_computer_setup._template_variables.items(): + if key == "label": + sub_widget = mapping_variable.widget + + # Test the default value is filled in correctly. + assert sub_widget.value == "merlin-cpu" + + # Fill the computer configure template variables + for ( + key, + mapping_variable, + ) in w_resource.template_computer_configure._template_variables.items(): + if key == "username": + sub_widget = mapping_variable.widget + sub_widget.value = "aiida" + + # Fill the code name + for ( + key, + mapping_variable, + ) in w_resource.template_code._template_variables.items(): + if key == "code_binary_name": + sub_widget = mapping_variable.widget + sub_widget.value = "ph" + + # select the other code and check the filled template is updated + sub_widget.value = "pw" + + w_resource.ssh_computer_setup.username.value = "aiida" + + # The quick_setup with password auth will try connect which will timeout. + # Thus, mock the connect method to avoid the timeout. + monkeypatch.setattr( + "aiidalab_widgets_base.computational_resources.SshComputerSetup.thread_ssh_copy_id", + lambda _: None, + ) + + w_resource._on_quick_setup() + + assert w_resource.success + # check the code is really created + assert orm.load_code("pw-7.0@merlin-cpu") + + # check the dropdown is updated + assert "pw-7.0@merlin-cpu" in [c[0] for c in w.code_select_dropdown.options] + + # The key pair will be generated to the temporary directory + # Check the content of the config is correct + with open(tmp_path / ".ssh" / "config") as f: + content = f.read() + assert "User aiida" in content + assert "Host merlin-l-01.psi.ch" in content diff --git a/tests/test_databases.py b/tests/test_databases.py index ed45b1ff..898e272c 100644 --- a/tests/test_databases.py +++ b/tests/test_databases.py @@ -1,6 +1,7 @@ import ase import aiidalab_widgets_base as awb +from aiidalab_widgets_base.databases import ComputationalResourcesDatabaseWidget def test_cod_query_widget(): @@ -37,32 +38,72 @@ def test_optimade_query_widget(): def test_computational_resources_database_widget(): """Test the structure browser widget.""" - # Initiate the widget with no arguments. - widget = awb.databases.ComputationalResourcesDatabaseWidget() - assert "merlin.psi.ch" in widget.database + widget = ComputationalResourcesDatabaseWidget() + assert "daint.cscs.ch" in widget.database # Initialize the widget with default_calc_job_plugin="cp2k" - widget = awb.databases.ComputationalResourcesDatabaseWidget( - default_calc_job_plugin="cp2k" - ) + # Note: after migrate to the new database with schema fixed, this test should go with + # the local defined database rather than relying on the remote one. + # Same for the quick setup widget. + widget = ComputationalResourcesDatabaseWidget(default_calc_job_plugin="cp2k") assert ( "merlin.psi.ch" not in widget.database ) # Merlin does not have CP2K installed. - widget.inp_domain.label = "daint.cscs.ch" - widget.inp_computer.value = "multicore" - widget.inp_code.value = "cp2k-9.1-multicore" + # Select computer/code + widget.domain_selector.value = "daint.cscs.ch" + widget.computer_selector.value = "mc" - # Check that the configuration is provided. + # check the code is not selected + assert widget.code_selector.value is None + + widget.code_selector.value = "cp2k-9.1" - assert "label" in widget.computer_setup["setup"] - assert "hostname" in widget.ssh_config + # Check that the configuration is provided. + assert "label" in widget.computer_setup + assert "hostname" in widget.computer_configure assert "filepath_executable" in widget.code_setup + # test after computer re-select to another, the code selector is reset + widget.computer_selector.value = "gpu" + assert widget.code_selector.value is None + # Simulate reset. - widget._reset() + widget.reset() assert widget.computer_setup == {} + assert widget.computer_configure == {} assert widget.code_setup == {} - assert widget.ssh_config == {} + + # after reset, the computer/code selector is reset + assert widget.computer_selector.options == () + assert widget.code_selector.options == () + assert widget.computer_selector.value is None + assert widget.code_selector.value is None + + # after reset, the domain selector value is reset, but the options are not + assert widget.domain_selector.value is None + assert len(widget.domain_selector.options) > 0 + + +def test_resource_database_widget_recognize_template_entry_points(): + """Test that the template like entry points are recognized.""" + # Initiate the widget with no arguments. + widget = ComputationalResourcesDatabaseWidget() + assert "daint.cscs.ch" in widget.database + + # Initialize the widget with default_calc_job_plugin="quantumespresso.pw" + # In merlin, there is a template entry point for Quantum ESPRESSO. + # Note: after migrate to the new database with schema fixed, this test should go with + # the local defined database rather than relying on the remote one. + # Same for the quick setup widget. + widget = ComputationalResourcesDatabaseWidget( + default_calc_job_plugin="quantumespresso.pw" + ) + assert "merlin.psi.ch" in widget.database + + widget = ComputationalResourcesDatabaseWidget( + default_calc_job_plugin="quantumespresso.ph" + ) + assert "merlin.psi.ch" in widget.database diff --git a/tests_notebooks/test_notebooks.py b/tests_notebooks/test_notebooks.py index f7579fd1..2cec668c 100644 --- a/tests_notebooks/test_notebooks.py +++ b/tests_notebooks/test_notebooks.py @@ -115,12 +115,12 @@ def test_computational_resources_code_setup( selenium_driver, aiidalab_exec, final_screenshot ): """Test the quicksetup of the code""" - # check the code pw-7.0 is not in code list + # check the code cp2k is not in code list output = aiidalab_exec("verdi code list").decode().strip() - assert "pw-7.0" not in output + assert "cp2k" not in output driver = selenium_driver("notebooks/computational_resources.ipynb") - driver.set_window_size(800, 800) + driver.set_window_size(800, 1600) # click the "Setup new code" button driver.find_element(By.XPATH, '(//button[text()="Setup new code"])[1]').click() @@ -128,50 +128,27 @@ def test_computational_resources_code_setup( # Select daint.cscs.ch domain driver.find_element(By.XPATH, '(//option[text()="daint.cscs.ch"])[1]').click() - # Select computer multicore - driver.find_element(By.XPATH, '(//option[text()="multicore"])[1]').click() + # Select computer mc + driver.find_element(By.XPATH, '(//option[text()="mc"])[1]').click() - # select code pw-7.0-multicore - driver.find_element(By.XPATH, '(//option[text()="pw-7.0-multicore"])[1]').click() + # select code + driver.find_element(By.XPATH, '(//option[text()="cp2k-9.1"])[1]').click() # fill the SSH username driver.find_element( By.XPATH, "(//label[text()='SSH username:'])[1]/following-sibling::input" ).send_keys("dummyuser") - # click the quick setup - driver.find_element(By.XPATH, '(//button[text()="Quick Setup"])[1]').click() - time.sleep(1.0) - - # check the new code pw-7.0@daint-mc is in code list - output = aiidalab_exec("verdi code list").decode().strip() - assert "pw-7.0@daint-mc" in output - - # Set the second code of the same computer - # issue https://github.com/aiidalab/aiidalab-widgets-base/issues/416 - # click the "Setup new code" button - driver.find_element(By.XPATH, '(//button[text()="Setup new code"])[2]').click() - - # Select daint.cscs.ch domain - driver.find_element(By.XPATH, '(//option[text()="daint.cscs.ch"])[2]').click() - - # Select computer multicore - driver.find_element(By.XPATH, '(//option[text()="multicore"])[2]').click() - - # select code pw-7.0-multicore - driver.find_element(By.XPATH, '(//option[text()="dos-7.0-multicore"])[2]').click() - - # fill the SSH username - # Get the element of index 3 which is the SSH username of second widget - # the one of index 2 is the SSH username in detail setup of the first widget. driver.find_element( - By.XPATH, "(//label[text()='SSH username:'])[3]/following-sibling::input" + By.XPATH, "(//label[text()='Slurm account:'])[1]/following-sibling::input" ).send_keys("dummyuser") - # click the quick setup - driver.find_element(By.XPATH, '(//button[text()="Quick Setup"])[2]').click() + # click the quick setup (contain text "Quick setup") + driver.find_element( + By.XPATH, '(//button[contains(text(),"Quick setup")])[1]' + ).click() time.sleep(1.0) # check the new code pw-7.0@daint-mc is in code list output = aiidalab_exec("verdi code list").decode().strip() - assert "dos-7.0@daint-mc" in output + assert "cp2k-9.1@daint-mc" in output