Source code for aiidalab_widgets_base.computational_resources
from __future__ import annotations
import copy
import enum
import re
import subprocess
import threading
from collections import namedtuple
from pathlib import Path
from uuid import UUID
import ipywidgets as ipw
import jinja2
import pexpect
import shortuuid
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 MessageLevel, StatusHTML, wrap_message
STYLE = {"description_width": "180px"}
LAYOUT = {"width": "400px"}
[docs]
class ComputationalResourcesWidget(ipw.VBox):
"""Code selection widget.
Attributes:
value(code UUID): Trait that points to the selected UUID of the code instance.
It can be set either to an AiiDA code UUID or to a code label.
It is linked to the `value` trait of the `self.code_select_dropdown` widget.
codes(Dict): Trait that contains a dictionary (label => Code UUID) for all
codes found in the AiiDA database for the selected plugin. It is linked
to the 'options' trait of the `self.code_select_dropdown` widget.
allow_hidden_codes(Bool): Trait that defines whether to show hidden codes or not.
allow_disabled_computers(Bool): Trait that defines whether to show codes on disabled
computers.
"""
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"
[docs]
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=clear_after)
self.code_select_dropdown = ipw.Dropdown(
description=description,
disabled=True,
value=None,
style={"description_width": "initial"},
)
# 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],
)
tl.directional_link(
(self.code_select_dropdown, "options"),
(self, "codes"),
transform=lambda x: {c[0]: c[1] for c in x},
)
tl.link((self.code_select_dropdown, "value"), (self, "value"))
self.observe(
self.refresh, names=["allow_disabled_computers", "allow_hidden_codes"]
)
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": self._output_width})
self._default_user_email = orm.User.collection.get_default().email
children = [
ipw.HBox([self.code_select_dropdown, self.btn_setup_new_code]),
self._setup_new_code_output,
self.output,
]
super().__init__(children=children, **kwargs)
# 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.resource_setup.observe(self.refresh, "success")
tl.dlink(
(self.resource_setup, "message"),
(self.setup_message, "message"),
)
self.refresh()
[docs]
def _get_codes(self):
"""Query the list of available codes."""
user = orm.User.collection.get(email=self._default_user_email)
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=filters,
)
.all()
if c[0].computer.is_user_configured(user)
and (self.allow_hidden_codes or not c[0].is_hidden)
and (self.allow_disabled_computers or c[0].computer.is_user_enabled(user))
]
[docs]
def refresh(self, _=None):
"""Refresh available codes.
The job of this function is to look in AiiDA database, find available codes and
put them in the code_select_dropdown widget."""
self.output.value = ""
with self.hold_trait_notifications():
self.code_select_dropdown.options = self._get_codes()
if not self.code_select_dropdown.options:
self.output.value = f"No codes found for default calcjob plugin {self.default_calc_job_plugin!r}."
self.code_select_dropdown.disabled = True
else:
self.code_select_dropdown.disabled = False
self.code_select_dropdown.value = None
@tl.validate("value")
def _validate_value(self, change):
"""Check if the code is valid in DB"""
code_uuid = change["value"]
self.output.value = ""
# If code None, set value to None.
if code_uuid is None:
return None
try:
_ = UUID(code_uuid, version=4)
except ValueError:
self.output.value = f"""'<span style="color:red">{code_uuid}</span>'
is not a valid UUID."""
else:
return code_uuid
[docs]
def _setup_new_code(self, _=None):
with self._setup_new_code_output:
clear_output()
if self.btn_setup_new_code.value:
self._setup_new_code_output.layout = {
"width": self._output_width,
"border": "1px solid gray",
}
children = [
self.resource_setup,
self.setup_message,
]
display(*children)
else:
self._setup_new_code_output.layout = {
"width": self._output_width,
"border": "none",
}
[docs]
class SshConnectionState(enum.Enum):
waiting_for_input = -1
enter_password = 0
success = 1
keys_already_present = 2
do_you_want_to_continue = 3
no_keys = 4
unknown_hostname = 5
connection_refused = 6
end_of_file = 7
[docs]
class SshComputerSetup(ipw.VBox):
"""Setup a passwordless access to a computer."""
ssh_config = tl.Dict()
ssh_connection_state = tl.UseEnum(
SshConnectionState, allow_none=True, default_value=None
)
SSH_POSSIBLE_RESPONSES = [
# order matters! the index will return by pexpect and compare
# with SshConnectionState
"[Pp]assword:", # 0
"Now try logging into", # 1
"All keys were skipped because they already exist on the remote system", # 2
"Are you sure you want to continue connecting (yes/no)?", # 3
"ERROR: No identities found", # 4
"Could not resolve hostname", # 5
"Connection refused", # 6
pexpect.EOF, # 7
]
message = tl.Unicode()
password_message = tl.Unicode("The passwordless enabling log.")
[docs]
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
self._ssh_connection_message = None
self._password_message = ipw.HTML()
tl.dlink(
(self, "password_message"),
(self._password_message, "value"),
transform=lambda x: f"<b>SSH log: {x}</b>",
)
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", "display": "none"},
)
self._continue_with_password_button.on_click(self._send_password)
self.password_box = ipw.VBox(
[
ipw.HBox([self._ssh_password, self._continue_with_password_button]),
self._password_message,
]
)
# Username.
self.username = ipw.Text(description="Username:", layout=LAYOUT, style=STYLE)
# Port.
self.port = ipw.IntText(
description="Port:",
value=22,
layout=LAYOUT,
style=STYLE,
)
# Hostname.
self.hostname = ipw.Text(
description="Computer hostname:",
layout=LAYOUT,
style=STYLE,
)
# ProxyJump.
self.proxy_jump = ipw.Text(
description="ProxyJump:",
layout=LAYOUT,
style=STYLE,
)
# ProxyJump.
self.proxy_command = ipw.Text(
description="ProxyCommand:",
layout=LAYOUT,
style=STYLE,
)
self._inp_private_key = ipw.FileUpload(
accept="",
layout=LAYOUT,
description="Private key",
multiple=False,
)
self._verification_mode = ipw.Dropdown(
options=[
("Password", "password"),
("Use custom private key", "private_key"),
("Download public key", "public_key"),
("Multiple factor authentication", "mfa"),
],
layout=LAYOUT,
style=STYLE,
value="password",
description="Verification mode:",
disabled=False,
)
self._verification_mode.observe(
self._on_verification_mode_change, names="value"
)
self._verification_mode_output = ipw.Output()
self._continue_button = ipw.ToggleButton(
description="Continue", layout={"width": "100px"}, value=False
)
# Setup ssh button and output.
btn_setup_ssh = ipw.Button(description="Setup ssh")
btn_setup_ssh.on_click(self._on_setup_ssh_button_pressed)
children = [
self.hostname,
self.port,
self.username,
self.proxy_jump,
self.proxy_command,
self._verification_mode,
self._verification_mode_output,
btn_setup_ssh,
]
super().__init__(children, **kwargs)
[docs]
def _ssh_keygen(self):
"""Generate ssh key pair."""
self.message = wrap_message(
"Generating SSH key pair.",
MessageLevel.SUCCESS,
)
fpath = self._ssh_folder / "id_rsa"
keygen_cmd = [
"ssh-keygen",
"-f",
fpath,
"-t",
"rsa",
"-b",
"4096",
"-m",
"PEM",
"-N",
"",
]
if not fpath.exists():
subprocess.run(keygen_cmd, capture_output=True, check=True)
[docs]
def _can_login(self):
"""Check if it is possible to login into the remote host."""
# With BatchMode on, no password prompt or other interaction is attempted,
# so a connect that requires a password will fail.
ret = subprocess.call(
[
"ssh",
self.hostname.value,
"-o",
"BatchMode=yes",
"-o",
"ConnectTimeout=5",
"true",
]
)
return ret == 0
[docs]
def _is_in_config(self):
"""Check if the SSH config file contains host information."""
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
# even if it is not in the config file.
# We require at least the user to be specified.
if "user" not in sshcfg:
return False
return True
[docs]
def _write_ssh_config(self, private_key_abs_fname=None):
"""Put host information into the config file."""
config_path = self._ssh_folder / "config"
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")
if self.proxy_jump.value != "":
file.write(
f" ProxyJump {self.proxy_jump.value.format(username=self.username.value)}\n"
)
if self.proxy_command.value != "":
file.write(
f" ProxyCommand {self.proxy_command.value.format(username=self.username.value)}\n"
)
if private_key_abs_fname:
file.write(f" IdentityFile {private_key_abs_fname}\n")
file.write(" ServerAliveInterval 5\n")
[docs]
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
message = "Please specify the computer name (for SSH)"
raise ValueError(message)
if self.username.value == "": # check username
message = "Please specify your SSH username."
raise ValueError(message)
[docs]
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()
[docs]
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.
# 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_fname)
[docs]
def _ssh_copy_id(self):
"""Run the ssh-copy-id command and follow it until it is completed."""
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}"
)
while True:
try:
idx = self._ssh_connection_process.expect(
self.SSH_POSSIBLE_RESPONSES,
timeout=timeout,
)
self.ssh_connection_state = SshConnectionState(idx)
except pexpect.TIMEOUT:
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.
if self.ssh_connection_state in (
SshConnectionState.success,
SshConnectionState.keys_already_present,
SshConnectionState.no_keys,
SshConnectionState.unknown_hostname,
SshConnectionState.connection_refused,
SshConnectionState.end_of_file,
):
break
self._ssh_connection_message = None
self._ssh_connection_process = None
[docs]
def _send_password(self, _=None):
self._continue_with_password_button.disabled = True
self._ssh_connection_process.sendline(self._ssh_password.value)
@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:
return
if self.ssh_connection_state is SshConnectionState.success:
self.password_message = (
"The passwordless connection has been set up successfully."
)
return
if self.ssh_connection_state is SshConnectionState.keys_already_present:
self.password_message = "The passwordless connection has already been setup. Nothing to be done."
return
if self.ssh_connection_state is SshConnectionState.no_keys:
self.password_message = (
" Failed\nLooks like the key pair is not present in ~/.ssh folder."
)
return
if self.ssh_connection_state is SshConnectionState.unknown_hostname:
self.password_message = "Failed\nUnknown hostname."
return
if self.ssh_connection_state is SshConnectionState.connection_refused:
self.password_message = "Failed\nConnection refused."
return
if self.ssh_connection_state is SshConnectionState.end_of_file:
self.password_message = (
"Didn't manage to connect. Please check your username/password."
)
return
if self.ssh_connection_state is SshConnectionState.enter_password:
self._handle_ssh_password()
elif self.ssh_connection_state is SshConnectionState.do_you_want_to_continue:
self._ssh_connection_process.sendline("yes")
[docs]
def _handle_ssh_password(self):
"""Send a password to a remote computer."""
message = (
self._ssh_connection_process.before.splitlines()[-1]
+ self._ssh_connection_process.after
)
if self._ssh_connection_message == message:
self._ssh_connection_process.sendline(self._ssh_password.value)
else:
self.password_message = (
f"Please enter {self.username.value}@{self.hostname.value}'s password:"
if message == b"Password:"
else f"Please enter {message.decode('utf-8')}"
)
self._ssh_password.disabled = False
self._continue_with_password_button.disabled = False
self._ssh_connection_message = message
# 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 <i class="fa fa-spinner" aria-hidden="true"></i> (timeout 10s)'
self._send_password()
[docs]
def _on_verification_mode_change(self, change):
"""which verification mode is chosen."""
with self._verification_mode_output:
clear_output()
if self._verification_mode.value == "private_key":
display(self._inp_private_key)
elif self._verification_mode.value == "public_key":
public_key = self._ssh_folder / "id_rsa.pub"
if public_key.exists():
display(
ipw.HTML(
f"""<pre style="background-color: #253239; color: #cdd3df; line-height: normal; custom=test">{public_key.read_text()}</pre>""",
layout={"width": "100%"},
)
)
@property
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.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
[docs]
@staticmethod
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)
[docs]
def _reset(self):
self.hostname.value = ""
self.port.value = 22
self.username.value = ""
self.proxy_jump.value = ""
self.proxy_command.value = ""
@tl.observe("ssh_config")
def _observe_ssh_config(self, change):
"""Pre-filling the input fields."""
self._reset()
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"]
[docs]
class AiidaComputerSetup(ipw.VBox):
"""Inform AiiDA about a computer."""
computer_setup = tl.Dict(allow_none=True)
computer_configure = tl.Dict(allow_none=True)
message = tl.Unicode()
[docs]
def __init__(self, **kwargs):
self._on_setup_computer_success = []
self._default_user_email = orm.User.collection.get_default().email
# List of widgets to be displayed.
self.label = ipw.Text(
value="",
placeholder="Will only be used within AiiDA",
description="Computer name:",
layout=LAYOUT,
style=STYLE,
)
# Hostname.
self.hostname = ipw.Text(description="Hostname:", layout=LAYOUT, style=STYLE)
# Computer description.
self.description = ipw.Text(
value="",
placeholder="No description (yet)",
description="Computer description:",
layout=LAYOUT,
style=STYLE,
)
# Directory where to run the simulations.
self.work_dir = ipw.Text(
value="",
placeholder="/home/{username}/aiida_run",
description="AiiDA working directory:",
layout=LAYOUT,
style=STYLE,
)
# Mpirun command.
self.mpirun_command = ipw.Text(
value="mpirun -n {tot_num_mpiprocs}",
description="Mpirun command:",
layout=LAYOUT,
style=STYLE,
)
# Number of CPUs per node.
self.mpiprocs_per_machine = ipw.IntText(
value=1,
step=1,
description="#CPU(s) per node:",
layout=LAYOUT,
style=STYLE,
)
# Memory per node.
self.default_memory_per_machine_widget = ipw.Text(
value="",
placeholder="not specified",
description="Memory per node:",
layout=LAYOUT,
style=STYLE,
)
memory_wrong_syntax = ipw.HTML(
value="""<i class="fa fa-times" style="color:red;font-size:2em;" ></i> wrong syntax""",
)
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.display = "none"
if not self.default_memory_per_machine_widget.value:
self.default_memory_per_machine = None
return
try:
self.default_memory_per_machine = (
int(parse_size(change["new"], binary=True) / 1024) or None
)
memory_wrong_syntax.layout.display = "none"
except InvalidSize:
memory_wrong_syntax.layout.display = "block"
self.default_memory_per_machine = None
self.default_memory_per_machine_widget.observe(
observe_memory_per_machine, names="value"
)
# Transport type.
self.transport = ipw.Dropdown(
value="core.local",
options=plugins.entry_point.get_entry_point_names("aiida.transports"),
description="Transport type:",
style=STYLE,
)
# Safe interval.
self.safe_interval = ipw.FloatText(
value=30.0,
description="Min. connection interval (sec):",
layout=LAYOUT,
style=STYLE,
)
# Scheduler.
self.scheduler = ipw.Dropdown(
value="core.slurm",
options=plugins.entry_point.get_entry_point_names("aiida.schedulers"),
description="Scheduler:",
style=STYLE,
)
self.shebang = ipw.Text(
value="#!/usr/bin/env bash",
description="Shebang:",
layout=LAYOUT,
style=STYLE,
)
# Use double quotes to escape.
self.use_double_quotes = ipw.Checkbox(
value=False,
description="Use double quotes to escape environment variable of job script.",
)
# Use login shell.
self.use_login_shell = ipw.Checkbox(value=True, description="Use login shell")
# Prepend text.
self.prepend_text = ipw.Textarea(
placeholder="Text to prepend to each command execution",
description="Prepend text:",
layout=LAYOUT,
)
# Append text.
self.append_text = ipw.Textarea(
placeholder="Text to append to each command execution",
description="Append text:",
layout=LAYOUT,
)
# Buttons and outputs.
self.setup_button = ipw.Button(description="Setup computer")
self.setup_button.on_click(self.on_setup_computer)
test_button = ipw.Button(description="Test computer")
test_button.on_click(self.test)
self._test_out = ipw.HTML(layout=LAYOUT)
# Organize the widgets
children = [
self.label,
self.hostname,
self.description,
self.work_dir,
self.mpirun_command,
self.mpiprocs_per_machine,
ipw.HBox([self.default_memory_per_machine_widget, memory_wrong_syntax]),
self.transport,
self.safe_interval,
self.scheduler,
self.shebang,
self.use_login_shell,
self.use_double_quotes,
self.prepend_text,
self.append_text,
self.setup_button,
test_button,
self._test_out,
]
super().__init__(children, **kwargs)
[docs]
def _configure_computer(self, computer: orm.Computer, transport: str):
# Use default AiiDA user
user = orm.User.collection.get(email=self._default_user_email)
if transport == "core.ssh":
self._configure_computer_ssh(computer, user)
elif transport == "core.local":
self._configure_computer_local(computer, user)
else:
msg = f"invalid transport type '{transport}'"
raise common.ValidationError(msg)
[docs]
def _configure_computer_ssh(self, computer: orm.Computer, user: orm.User):
"""Configure the computer with SSH transport
There are three sources of authparams information ordered by priority increase:
- the default values
- the SSH config file
- the computer_configure dictionary
At the moment, there is no overlap between the SSH config file and the computer_configure dictionary.
The proxyjump and proxycommend can be read from both the SSH config file and the computer_configure dictionary,
since the SSH config file is generated from the computer_configure dictionary in `SshComputerSetup._write_ssh_config`.
"""
sshcfg = aiida_ssh_plugin.parse_sshconfig(self.hostname.value)
authparams = {
"port": int(sshcfg.get("port", 22)),
"look_for_keys": True,
"key_filename": str(
Path(sshcfg.get("identityfile", ["~/.ssh/id_rsa"])[0]).expanduser()
),
"timeout": 60,
"allow_agent": True,
"proxy_jump": "",
"proxy_command": "",
"compress": True,
"gss_auth": False,
"gss_kex": False,
"gss_deleg_creds": False,
"gss_host": self.hostname.value,
"load_system_host_keys": True,
"key_policy": "WarningPolicy",
"use_login_shell": self.use_login_shell.value,
"safe_interval": self.safe_interval.value,
}
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 "proxy_jump" in self.computer_configure:
authparams["proxy_jump"] = self.computer_configure["proxy_jump"]
if "proxy_command" in self.computer_configure:
authparams["proxy_command"] = self.computer_configure["proxy_command"]
if "key_filename" in self.computer_configure:
authparams["key_filename"] = str(
Path(self.computer_configure["key_filename"]).expanduser()
)
if "key_policy" in self.computer_configure:
authparams["key_policy"] = self.computer_configure["key_policy"]
computer.configure(user=user, **authparams)
return True
[docs]
def _configure_computer_local(self, computer: orm.Computer, user: orm.User):
"""Configure the computer with local transport"""
authparams = {
"use_login_shell": self.use_login_shell.value,
"safe_interval": self.safe_interval.value,
}
computer.configure(user=user, **authparams)
return True
[docs]
def _run_callbacks_if_computer_exists(self, label):
"""Run things on an existing computer"""
if self._computer_exists(label):
for function in self._on_setup_computer_success:
function()
return True
return False
[docs]
def _computer_exists(self, label):
try:
orm.load_computer(label=label)
except common.NotExistent:
return False
return True
[docs]
def _validate_computer_settings(self):
if self.label.value == "": # check computer label
self.message = wrap_message(
"Please specify the computer name (for AiiDA)",
MessageLevel.WARNING,
)
return False
if self.work_dir.value == "":
self.message = wrap_message(
"Please specify working directory",
MessageLevel.WARNING,
)
return False
if self.hostname.value == "":
self.message = wrap_message(
"Please specify hostname",
MessageLevel.WARNING,
)
return False
return True
[docs]
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 = wrap_message(
f"A computer {self.label.value} already exists. Skipping creation.",
MessageLevel.INFO,
)
return True
items_to_configure = [
"label",
"hostname",
"description",
"work_dir",
"mpirun_command",
"mpiprocs_per_machine",
"transport",
"use_double_quotes",
"scheduler",
"prepend_text",
"append_text",
"shebang",
]
kwargs = {key: getattr(self, key).value for key in items_to_configure}
computer_builder = ComputerBuilder(
default_memory_per_machine=self.default_memory_per_machine, **kwargs
)
try:
computer = computer_builder.new()
self._configure_computer(computer, self.transport.value)
computer.store()
except (
ComputerBuilder.ComputerValidationError,
common.exceptions.ValidationError,
RuntimeError,
) as 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 = wrap_message(
f"Computer<{computer.pk}> {computer.label} created",
MessageLevel.SUCCESS,
)
return True
self.message = wrap_message(
f"Failed to create computer {computer.label}",
MessageLevel.ERROR,
)
return False
[docs]
def on_setup_computer_success(self, function):
self._on_setup_computer_success.append(function)
[docs]
def test(self, _=None):
if self.label.value == "":
self._test_out.value = "Please specify the computer name (for AiiDA)."
return False
elif not self._computer_exists(self.label.value):
self._test_out.value = (
f"A computer called <b>{self.label.value}</b> does not exist."
)
return False
self._test_out.value = '<i class="fa fa-spinner fa-pulse"></i>'
process_result = subprocess.run(
["verdi", "computer", "test", "--print-traceback", self.label.value],
capture_output=True,
check=False,
)
if process_result.returncode == 0:
self._test_out.value = process_result.stdout.decode("utf-8").replace(
"\n", "<br>"
)
return True
else:
self._test_out.value = process_result.stderr.decode("utf-8").replace(
"\n", "<br>"
)
return False
[docs]
def _reset(self):
self.label.value = ""
self.hostname.value = ""
self.description.value = ""
self.work_dir.value = ""
self.mpirun_command.value = "mpirun -n {tot_num_mpiprocs}"
self.default_memory_per_machine_widget.value = ""
self.mpiprocs_per_machine.value = 1
self.transport.value = "core.ssh"
self.safe_interval.value = 30.0
self.scheduler.value = "core.slurm"
self.shebang.value = "#!/usr/bin/env bash"
self.use_login_shell.value = True
self.use_double_quotes.value = False
self.prepend_text.value = ""
self.append_text.value = ""
@tl.observe("computer_setup")
def _observe_computer_setup(self, _=None):
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()
@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))
[docs]
class AiidaCodeSetup(ipw.VBox):
"""Class that allows to setup AiiDA code"""
code_setup = tl.Dict(allow_none=True)
message = tl.Unicode()
[docs]
def __init__(self, **kwargs):
self._on_setup_code_success = []
# Code label.
self.label = ipw.Text(
description="AiiDA code label:",
layout=LAYOUT,
style=STYLE,
)
# Computer on which the code is installed. The value of this widget is
# the UUID of the selected computer.
self.computer = ComputerDropdownWidget()
# Code plugin.
self.default_calc_job_plugin = ipw.Dropdown(
options=sorted(
(ep.name, ep.name)
for ep in plugins.entry_point.get_entry_points("aiida.calculations")
),
description="Code plugin:",
layout=LAYOUT,
style=STYLE,
)
# Code description.
self.description = ipw.Text(
placeholder="No description (yet)",
description="Code description:",
layout=LAYOUT,
style=STYLE,
)
self.filepath_executable = ipw.Text(
placeholder="/path/to/executable",
description="Absolute path to executable:",
layout=LAYOUT,
style=STYLE,
)
# Use double quotes to escape.
self.use_double_quotes = ipw.Checkbox(
value=False,
description="Use double quotes to escape environment variable of job script.",
)
self.prepend_text = ipw.Textarea(
placeholder="Text to prepend to each command execution",
description="Prepend text:",
layout=LAYOUT,
)
self.append_text = ipw.Textarea(
placeholder="Text to append to each command execution",
description="Append text:",
layout=LAYOUT,
)
btn_setup_code = ipw.Button(description="Setup code")
btn_setup_code.on_click(self.on_setup_code)
self.setup_code_out = ipw.Output()
children = [
self.label,
self.computer,
self.default_calc_job_plugin,
self.description,
self.filepath_executable,
self.use_double_quotes,
self.prepend_text,
self.append_text,
btn_setup_code,
self.setup_code_out,
]
super().__init__(children, **kwargs)
@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
[docs]
def on_setup_code(self, _=None):
"""Setup an AiiDA code."""
with self.setup_code_out:
clear_output()
if not self.computer.value:
self.message = wrap_message(
"Please select an existing computer.",
MessageLevel.WARNING,
)
return False
items_to_configure = [
"label",
"description",
"default_calc_job_plugin",
"filepath_executable",
"use_double_quotes",
"prepend_text",
"append_text",
]
kwargs = {key: getattr(self, key).value for key in items_to_configure}
# set computer from its widget value the UUID of the computer.
computer = orm.load_computer(self.computer.value)
# Checking if the code with this name already exists
qb = orm.QueryBuilder()
qb.append(orm.Computer, filters={"uuid": computer.uuid}, tag="computer")
qb.append(
orm.InstalledCode,
with_computer="computer",
filters={"label": kwargs["label"]},
)
if qb.count() > 0:
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 = wrap_message(
f"Invalid input for code creation: <code>{exception}</code>",
MessageLevel.ERROR,
)
return False
try:
code.store()
code.is_hidden = False
except common.exceptions.ValidationError as exception:
self.message = wrap_message(
f"Unable to store the code: <code>{exception}</code>",
MessageLevel.ERROR,
)
return False
for function in self._on_setup_code_success:
function()
self.message = wrap_message(
f"Code<{code.pk}> {code.full_label} created",
MessageLevel.SUCCESS,
)
return True
[docs]
def _reset(self):
self.label.value = ""
self.computer.value = ""
self.description.value = ""
self.filepath_executable.value = ""
self.use_double_quotes.value = False
self.prepend_text.value = ""
self.append_text.value = ""
@tl.observe("code_setup")
def _observe_code_setup(self, _=None):
# Setup.
self.computer.refresh()
if not self.code_setup:
self._reset()
for key, value in self.code_setup.items():
if hasattr(self, key):
if key == "default_calc_job_plugin":
try:
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
# ComputerDropdownWidget it will refresh after the computer set up.
# if the computer is set pass the UUID to ComputerDropdownWdiget.
try:
computer = orm.load_computer(value)
except common.NotExistent:
getattr(self, key).value = None
else:
getattr(self, key).value = computer.uuid
else:
getattr(self, key).value = value
[docs]
class ComputerDropdownWidget(ipw.VBox):
"""Widget to select a configured computer.
Attributes:
value(computer UUID): Trait that points to the selected Computer instance. It can be set to an AiiDA Computer UUID. It is linked to the 'value' trait of `self._dropdown` widget.
computers(Dict): Trait that contains a dictionary (label => Computer UUID) for all computers found in the AiiDA database. It is linked to the 'options' trait of `self._dropdown` widget.
allow_select_disabled(Bool): Trait that defines whether to show disabled computers.
"""
value = tl.Unicode(allow_none=True)
computers = tl.Dict(allow_none=True)
allow_select_disabled = tl.Bool(False)
[docs]
def __init__(self, description="Select computer:", **kwargs):
"""Dropdown for configured AiiDA Computers.
description (str): Text to display before dropdown.
"""
self._default_user_email = orm.User.collection.get_default().email
self.output = ipw.HTML()
self._dropdown = ipw.Dropdown(
value=None,
description=description,
style=STYLE,
layout=LAYOUT,
disabled=True,
)
tl.directional_link(
(self, "computers"),
(self._dropdown, "options"),
transform=lambda x: [(key, x[key]) for key in x],
)
tl.directional_link(
(self._dropdown, "options"),
(self, "computers"),
transform=lambda x: {c[0]: c[1] for c in x},
)
tl.link((self._dropdown, "value"), (self, "value"))
self.observe(self.refresh, names="allow_select_disabled")
children = [
ipw.HBox(
[
self._dropdown,
]
),
self.output,
]
self.refresh()
super().__init__(children=children, **kwargs)
[docs]
def _get_computers(self) -> list:
"""Get the list of available computers."""
# Getting the current user.
user = orm.User.collection.get(email=self._default_user_email)
return [
(c[0].label, c[0].uuid)
for c in orm.QueryBuilder().append(orm.Computer).all()
if c[0].is_user_configured(user)
and (self.allow_select_disabled or c[0].is_user_enabled(user))
]
[docs]
def refresh(self, _=None):
"""Refresh the list of configured computers."""
self.output.value = ""
with self.hold_trait_notifications():
self._dropdown.options = self._get_computers()
if not self.computers:
self.output.value = "No computers found."
self._dropdown.disabled = True
else:
self._dropdown.disabled = False
self._dropdown.value = None
@tl.validate("value")
def _validate_value(self, change):
"""Select computer by computer UUID."""
computer_uuid = change["value"]
self.output.value = ""
if not computer_uuid:
return None
try:
_ = UUID(computer_uuid, version=4)
except ValueError:
self.output.value = f"""'<span style="color:red">{computer_uuid}</span>'
is not a valid UUID."""
else:
return computer_uuid
[docs]
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"""'<span style="color:red">{label}</span>'
can not find this computer."""
else:
self.value = computer_uuid
TemplateVariableLine = namedtuple("TemplateVariableLine", ["key", "str", "vars"])
TemplateVariable = namedtuple("TemplateVariable", ["widget", "lines"])
[docs]
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)
[docs]
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(
"""<div>Please fill the template variables below.</div>"""
)
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,
]
)
[docs]
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()
[docs]
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"""<div>{tooltip}</div>"""
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"
[docs]
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
[docs]
def _on_template_variable_filled(self, _):
"""Callback when a template variable is filled."""
self.fill()
[docs]
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"
[docs]
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(
"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(
"""<div><b>Tick checkbox to setup resource step by step.</b></div>"""
)
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 = """<div>
<b>Go through the steps to setup SSH connection to remote machine, computer, and code into database. </br>
The SSH connection step can be skipped and setup afterwards.</br>
</b>
</div>"""
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(
"""<div>Please select the computer/code from a database to pre-fill the fields below.</div>
"""
),
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()
[docs]
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,
]
[docs]
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
[docs]
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()
[docs]
def _on_reset(self, _=None):
"""Reset the database and the widget."""
with self.hold_trait_notifications():
self.reset()
self.comp_resources_database.reset()
[docs]
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"<b>{v}</b>" 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()
[docs]
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)
[docs]
def _on_setup_code_success(self):
"""Callback that is called when the code is successfully set up."""
self.success = True
[docs]
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