Skip to content

Commit

Permalink
CLI: Add the verdi computer export command (aiidateam#6389)
Browse files Browse the repository at this point in the history
This command has two subcommands, `setup` and `config`, that dump the
definition of a `Computer` and an associated `AuthInfo` to a YAML file
that can be used to recreate it using `verdi computer setup` and
`verdi computer configure` using the `--config` option, respectively.
  • Loading branch information
agoscinski authored and mikibonacci committed Sep 3, 2024
1 parent fc4da66 commit dc19430
Show file tree
Hide file tree
Showing 4 changed files with 163 additions and 2 deletions.
3 changes: 2 additions & 1 deletion docs/source/reference/command_line.rst
Original file line number Diff line number Diff line change
Expand Up @@ -100,11 +100,12 @@ Below is a list with all available subcommands.
--help Show this message and exit.
Commands:
configure Configure the Authinfo details for a computer (and user).
configure Configure the transport for a computer and user.
delete Delete a computer.
disable Disable the computer for the given user.
duplicate Duplicate a computer allowing to change some parameters.
enable Enable the computer for the given user.
export Export the setup or configuration of a computer.
list List all available computers.
relabel Relabel a computer.
setup Create a new computer.
Expand Down
1 change: 1 addition & 0 deletions src/aiida/cmdline/commands/cmd_code.py
Original file line number Diff line number Diff line change
Expand Up @@ -241,6 +241,7 @@ def show(code):
is_flag=True,
default=True,
help='Sort the keys of the output YAML.',
show_default=True,
)
@with_dbenv()
def export(code, output_file, sort):
Expand Down
93 changes: 92 additions & 1 deletion src/aiida/cmdline/commands/cmd_computer.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
###########################################################################
"""`verdi computer` command."""

import pathlib
import traceback
from copy import deepcopy
from functools import partial
from math import isclose
Expand Down Expand Up @@ -673,7 +675,7 @@ def get_command(self, ctx, name):

@verdi_computer.group('configure', cls=LazyConfigureGroup)
def computer_configure():
"""Configure the Authinfo details for a computer (and user)."""
"""Configure the transport for a computer and user."""


@computer_configure.command('show')
Expand Down Expand Up @@ -730,3 +732,92 @@ def computer_config_show(computer, user, defaults, as_option_string):
else:
table.append((f'* {name}', '-'))
echo_tabulate(table, tablefmt='plain')


@verdi_computer.group('export')
def computer_export():
"""Export the setup or configuration of a computer."""


@computer_export.command('setup')
@arguments.COMPUTER()
@arguments.OUTPUT_FILE(type=click.Path(exists=False, path_type=pathlib.Path))
@click.option(
'--sort/--no-sort',
is_flag=True,
default=True,
help='Sort the keys of the output YAML.',
show_default=True,
)
@with_dbenv()
def computer_export_setup(computer, output_file, sort):
"""Export computer setup to a YAML file."""
import yaml

computer_setup = {
'label': computer.label,
'hostname': computer.hostname,
'description': computer.description,
'transport': computer.transport_type,
'scheduler': computer.scheduler_type,
'shebang': computer.get_shebang(),
'work_dir': computer.get_workdir(),
'mpirun_command': ' '.join(computer.get_mpirun_command()),
'mpiprocs_per_machine': computer.get_default_mpiprocs_per_machine(),
'default_memory_per_machine': computer.get_default_memory_per_machine(),
'use_double_quotes': computer.get_use_double_quotes(),
'prepend_text': computer.get_prepend_text(),
'append_text': computer.get_append_text(),
}
try:
output_file.write_text(yaml.dump(computer_setup, sort_keys=sort), 'utf-8')
except Exception as e:
error_traceback = traceback.format_exc()
echo.CMDLINE_LOGGER.debug(error_traceback)
echo.echo_critical(
f'Unexpected error while exporting setup for Computer<{computer.pk}> {computer.label}:\n ({e!s}).'
)
else:
echo.echo_success(f"Computer<{computer.pk}> {computer.label} setup exported to file '{output_file}'.")


@computer_export.command('config')
@arguments.COMPUTER()
@arguments.OUTPUT_FILE(type=click.Path(exists=False, path_type=pathlib.Path))
@options.USER(
help='Email address of the AiiDA user from whom to export this computer (if different from default user).'
)
@click.option(
'--sort/--no-sort',
is_flag=True,
default=True,
help='Sort the keys of the output YAML.',
show_default=True,
)
@with_dbenv()
def computer_export_config(computer, output_file, user, sort):
"""Export computer transport configuration for a user to a YAML file."""
import yaml

if not computer.is_configured:
echo.echo_critical(
f'Computer<{computer.pk}> {computer.label} configuration cannot be exported,'
' because computer has not been configured yet.'
)
try:
computer_configuration = computer.get_configuration(user)
output_file.write_text(yaml.dump(computer_configuration, sort_keys=sort), 'utf-8')
except Exception as e:
error_traceback = traceback.format_exc()
echo.CMDLINE_LOGGER.debug(error_traceback)
if user is None:
echo.echo_critical(
f'Unexpected error while exporting configuration for Computer<{computer.pk}> {computer.label}: {e!s}.'
)
else:
echo.echo_critical(
f'Unexpected error while exporting configuration for Computer<{computer.pk}> {computer.label}'
f' and User<{user.pk}> {user.email}: {e!s}.'
)
else:
echo.echo_success(f"Computer<{computer.pk}> {computer.label} configuration exported to file '{output_file}'.")
68 changes: 68 additions & 0 deletions tests/cmdline/commands/test_computer.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,17 +14,21 @@
from collections import OrderedDict

import pytest
import yaml
from aiida import orm
from aiida.cmdline.commands.cmd_computer import (
computer_configure,
computer_delete,
computer_duplicate,
computer_export_config,
computer_export_setup,
computer_list,
computer_relabel,
computer_setup,
computer_show,
computer_test,
)
from aiida.cmdline.utils.echo import ExitCode


def generate_setup_options_dict(replace_args=None, non_interactive=True):
Expand Down Expand Up @@ -511,6 +515,70 @@ def test_show(self):
assert '--username=' in result.output
assert result_cur.output == result.output

@pytest.mark.parametrize('sort', ['--sort', '--no-sort'])
def test_computer_export_setup(self, tmp_path, sort):
"""Test if 'verdi computer export setup' command works"""
self.comp_builder.label = 'test_computer_export_setup' + sort
self.comp_builder.transport = 'core.ssh'
comp = self.comp_builder.new()
comp.store()

exported_setup_filename = tmp_path / 'computer-setup.yml'
result = self.cli_runner(computer_export_setup, [sort, comp.label, exported_setup_filename])
assert result.exit_code == 0, 'Command should have run successfull.'
assert str(exported_setup_filename) in result.output, 'Filename should be in terminal output but was not found.'
assert exported_setup_filename.exists(), f"'{exported_setup_filename}' was not created during export."
# verifying correctness by comparing internal and loaded yml object
configure_setup_data = yaml.safe_load(exported_setup_filename.read_text())
assert configure_setup_data == self.comp_builder.get_computer_spec(
comp
), 'Internal computer configuration does not agree with exported one.'

# we create a directory so we raise an error when exporting with the same name
# to test the except part of the function
already_existing_filename = tmp_path / 'tmp_dir'
already_existing_filename.mkdir()
result = self.cli_runner(computer_export_setup, [sort, comp.label, already_existing_filename], raises=True)
assert result.exit_code == ExitCode.CRITICAL

@pytest.mark.parametrize('sort', ['--sort', '--no-sort'])
def test_computer_export_config(self, tmp_path, sort):
"""Test if 'verdi computer export config' command works"""
self.comp_builder.label = 'test_computer_export_config' + sort
self.comp_builder.transport = 'core.ssh'
comp = self.comp_builder.new()
comp.store()

exported_config_filename = tmp_path / 'computer-configure.yml'
# We have not configured the computer yet so it should exit with an critical error
result = self.cli_runner(computer_export_config, [comp.label, exported_config_filename], raises=True)
assert result.exit_code == ExitCode.CRITICAL

comp.configure(safe_interval=0.0)
result = self.cli_runner(computer_export_config, [comp.label, exported_config_filename])
assert 'Success' in result.output, 'Command should have run successfull.'
assert (
str(exported_config_filename) in result.output
), 'Filename should be in terminal output but was not found.'
assert exported_config_filename.exists(), f"'{exported_config_filename}' was not created during export."
# verifying correctness by comparing internal and loaded yml object
configure_config_data = yaml.safe_load(exported_config_filename.read_text())
assert (
configure_config_data == comp.get_configuration()
), 'Internal computer configuration does not agree with exported one.'

# we create a directory so we raise an error when exporting with the same name
# to test the except part of the function
already_existing_filename = tmp_path / 'tmp_dir'
already_existing_filename.mkdir()
result = self.cli_runner(computer_export_config, [comp.label, already_existing_filename], raises=True)
assert result.exit_code == ExitCode.CRITICAL

result = self.cli_runner(
computer_export_config, ['--user', self.user.email, comp.label, already_existing_filename], raises=True
)
assert result.exit_code == ExitCode.CRITICAL


class TestVerdiComputerCommands:
"""Testing verdi computer commands.
Expand Down

0 comments on commit dc19430

Please sign in to comment.