Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

CLI: Add the subcommand verdi computer export #6389

Merged
merged 12 commits into from
May 27, 2024
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
Loading