Source code for aiidalab_widgets_base.computational_resources

import enum
import os
import subprocess
import threading
from copy import copy
from pathlib import Path
from uuid import UUID

import ipywidgets as ipw
import pexpect
import shortuuid
import traitlets
from aiida import common, orm, plugins
from aiida.common.exceptions import NotExistent
from aiida.orm.utils.builders.computer import ComputerBuilder
from aiida.transports.plugins.ssh import parse_sshconfig
from humanfriendly import InvalidSize, parse_size
from IPython.display import clear_output, display

from .databases import ComputationalResourcesDatabaseWidget
from .utils import StatusHTML

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 = 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)
[docs] def __init__(self, description="Select code:", path_to_root="../", **kwargs): """Dropdown for Codes for one input plugin. description (str): Description to display before the dropdown. """ self.output = ipw.HTML() self.setup_message = StatusHTML(clear_after=30) self.code_select_dropdown = ipw.Dropdown( description=description, disabled=True, value=None, style={"description_width": "initial"}, ) traitlets.link((self, "codes"), (self.code_select_dropdown, "options")) traitlets.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": "500px"}) children = [ ipw.HBox([self.code_select_dropdown, self.btn_setup_new_code]), self._setup_new_code_output, self.output, ] 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"), ) self.aiida_computer_setup = AiidaComputerSetup() ipw.dlink( (self.aiida_computer_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()
[docs] 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()
[docs] def _get_codes(self): """Query the list of available codes.""" user = orm.User.collection.get_default() 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}, ) .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] @staticmethod def _full_code_label(code): return f"{code.label}@{code.computer.label}"
[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}'." self.code_select_dropdown.disabled = True else: self.code_select_dropdown.disabled = False self.code_select_dropdown.value = None
@traitlets.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": "500px", "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] display(*children) else: self._setup_new_code_output.layout = { "width": "500px", "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 = traitlets.Dict() ssh_connection_state = traitlets.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 = traitlets.Unicode() password_message = traitlets.Unicode("The passwordless enabling log.")
[docs] 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) self._continue_with_password_button = ipw.Button( description="Continue", layout={"width": "100px"}, disabled=True ) 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]), ] ) # Username. self.username = ipw.Text( description="SSH username:", layout=LAYOUT, style=STYLE ) # Port. self.port = ipw.IntText( description="SSH 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, style=STYLE, description="Private key", multiple=False, ) self._verification_mode = ipw.Dropdown( options=[ ("Password", "password"), ("Use custom private key", "private_key"), ("Download public key", "public_key"), ], 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 = "Generating SSH key pair." fpath = Path.home() / ".ssh" / "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)
[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 config file contains host information.""" fpath = Path.home() / ".ssh" / "config" if not fpath.exists(): return False cfglines = open(fpath).read().split("\n") return "Host " + self.hostname.value in cfglines
[docs] def _write_ssh_config(self, private_key_abs_fname=None): """Put host information into the config file.""" fpath = Path.home() / ".ssh" / "config" self.message = f"Adding {self.hostname.value} section to {fpath}" with open(fpath, "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 _on_setup_ssh_button_pressed(self, _=None): # 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 if self.username.value == "": # check username self.message = "Please specify your SSH username." return False 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 # 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. self._add_private_key(private_key_abs_fname, private_key_content) 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()
[docs] def _ssh_copy_id(self): """Run the ssh-copy-id command and follow it until it is completed.""" timeout = 30 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._ssh_password.disabled = True self._continue_with_password_button.disabled = True self.password_message = ( f"Exceeded {timeout} s timeout. Please start 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._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") 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 self.ssh_connection_state = SshConnectionState.waiting_for_input
[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 = Path.home() / ".ssh" / "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): """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"]) 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_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 file already exist and has the same content if fpath.read_bytes() == private_key_content: return fpath.name fpath = fpath / "_" / shortuuid.uuid() fpath.write_bytes(private_key_content) fpath.chmod(0o600) return fpath
[docs] def _reset(self): self.hostname.value = "" self.port.value = 22 self.username.value = "" self.proxy_jump.value = "" self.proxy_command.value = ""
@traitlets.observe("ssh_config") def _observe_ssh_config(self, _=None): """Pre-filling the input fields.""" if not self.ssh_config: 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"]
[docs]class AiidaComputerSetup(ipw.VBox): """Inform AiiDA about a computer.""" computer_setup = traitlets.Dict(allow_none=True) message = traitlets.Unicode()
[docs] def __init__(self, **kwargs): self._on_setup_computer_success = [] # 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="/scratch/{username}/aiida_run", description="Workdir:", 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""", layout={"visibility": "hidden"}, ) 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" 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.visibility = "hidden" except InvalidSize: memory_wrong_syntax.layout.visibility = "visible" 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.Output(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): """Configure the computer""" sshcfg = parse_sshconfig(self.hostname.value) authparams = { "port": int(sshcfg.get("port", 22)), "look_for_keys": True, "key_filename": os.path.expanduser( sshcfg.get("identityfile", ["~/.ssh/id_rsa"])[0] ), "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, } try: authparams["username"] = sshcfg["user"] except KeyError as exc: message = "SSH username is not provided" self.message = message raise RuntimeError(message) from exc if "proxycommand" in sshcfg: authparams["proxy_command"] = sshcfg["proxycommand"] elif "proxyjump" in sshcfg: authparams["proxy_jump"] = sshcfg["proxyjump"] # user default AiiDA user user = orm.User.collection.get_default() computer.configure(user=user, **authparams) return True
[docs] def _run_callbacks_if_computer_exists(self, label): """Run things on an existing computer""" try: orm.Computer.objects.get(label=label) for function in self._on_setup_computer_success: function() except common.NotExistent: return False else: return True
[docs] def on_setup_computer(self, _=None): """Create a new computer.""" if self.label.value == "": # check hostname self.message = "Please specify the computer name (for AiiDA)" 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." 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) except ( ComputerBuilder.ComputerValidationError, common.exceptions.ValidationError, RuntimeError, ) as err: self.message = f"Failed to setup computer {type(err).__name__}: {err}" return False else: computer.store() # 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" return True self.message = f"Failed to create computer {computer.label}" return False
[docs] def on_setup_computer_success(self, function): self._on_setup_computer_success.append(function)
[docs] def test(self, _=None): with self._test_out: clear_output() print( subprocess.check_output( ["verdi", "computer", "test", "--print-traceback", self.label.value] ).decode("utf-8") )
[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 = ""
@traitlets.observe("computer_setup") def _observe_computer_setup(self, _=None): # Setup. if not self.computer_setup: 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
[docs]class AiidaCodeSetup(ipw.VBox): """Class that allows to setup AiiDA code""" code_setup = traitlets.Dict(allow_none=True) message = traitlets.Unicode()
[docs] def __init__(self, path_to_root="../", **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( path_to_root=path_to_root, ) # 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)
@traitlets.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 = "Please select an existing computer." 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 = ( f"Code {kwargs['label']}@{computer.label} already exists." ) return False try: code = orm.InstalledCode(computer=computer, **kwargs) except (common.exceptions.InputValidationError, KeyError) as exception: self.message = f"Invalid inputs: {exception}" return False try: code.store() code.is_hidden = False except common.exceptions.ValidationError as exception: self.message = f"Unable to store the Code: {exception}" return False for function in self._on_setup_code_success: function() self.message = f"Code<{code.pk}> {code.full_label} created" return True
[docs] def on_setup_code_success(self, function): self._on_setup_code_success.append(function)
[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 = ""
[docs] def refresh(self): self._observe_code_setup()
@traitlets.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: getattr(self, key).label = value except traitlets.TraitError: self.message = f"Input plugin {value} is not installed." 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 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 = traitlets.Unicode(allow_none=True) computers = traitlets.Dict(allow_none=True) allow_select_disabled = traitlets.Bool(False)
[docs] def __init__(self, description="Select computer:", path_to_root="../", **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() self._dropdown = ipw.Dropdown( options={}, value=None, description=description, style=STYLE, layout=LAYOUT, disabled=True, ) traitlets.link((self, "computers"), (self._dropdown, "options")) traitlets.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): """Get the list of available computers.""" # Getting the current user. user = orm.User.collection.get_default() 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
@traitlets.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