From b4f66cf515677987077834c8150eb73daee415f7 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Thu, 11 Sep 2025 17:27:23 +0200 Subject: [PATCH 1/3] Multi-code export --- docs/source/reference/command_line.rst | 4 +- src/aiida/cmdline/commands/cmd_code.py | 118 +++++++++++++++++++------ 2 files changed, 91 insertions(+), 31 deletions(-) diff --git a/docs/source/reference/command_line.rst b/docs/source/reference/command_line.rst index 6bc2272694..27409608f5 100644 --- a/docs/source/reference/command_line.rst +++ b/docs/source/reference/command_line.rst @@ -75,7 +75,7 @@ Below is a list with all available subcommands. create Create a new code. delete Delete a code. duplicate Duplicate a code allowing to change some parameters. - export Export code to a yaml file. + export Export code(s) to yaml file(s). hide Hide one or more codes from `verdi code list`. list List the available codes. relabel Relabel a code. @@ -453,7 +453,7 @@ Below is a list with all available subcommands. --broker-host HOSTNAME Hostname for the message broker. [default: 127.0.0.1] --broker-port INTEGER Port for the message broker. [default: 5672] --broker-virtual-host TEXT Name of the virtual host for the message broker without - leading forward slash. [default: ""] + leading forward slash. --repository DIRECTORY Absolute path to the file repository. --test-profile Designate the profile to be used for running the test suite only. diff --git a/src/aiida/cmdline/commands/cmd_code.py b/src/aiida/cmdline/commands/cmd_code.py index a4c19d51a1..31844e278f 100644 --- a/src/aiida/cmdline/commands/cmd_code.py +++ b/src/aiida/cmdline/commands/cmd_code.py @@ -250,43 +250,103 @@ def show(code): @verdi_code.command() -@arguments.CODE() -@arguments.OUTPUT_FILE(type=click.Path(exists=False, path_type=pathlib.Path), required=False) +@click.argument('codes_and_output', nargs=-1, required=True) +@click.option('-a', '--all', 'export_all', is_flag=True, help='Export all codes.') @options.OVERWRITE() @options.SORT() @with_dbenv() -def export(code, output_file, overwrite, sort): - """Export code to a yaml file. If no output file is given, default name is created based on the code label.""" +def export(codes_and_output, export_all, overwrite, sort): + """Export code(s) to yaml file(s). + + Usage: + - verdi code export CODE_ID [CODE_ID ...] [OUTPUT_FILE] + - verdi code export --all + + If no output file is given, default names are created based on the code labels. + Custom output filename can only be provided when exporting a single code. + """ + from aiida import orm + from aiida.common import NotExistent + other_args = {'sort': sort} fileformat = 'yaml' - if output_file is None: - output_file = pathlib.Path(f'{code.full_label}.{fileformat}') - - try: - # In principle, output file validation is also done in the `data_export` function. However, the - # `validate_output_filename` function is applied here, as well, as it is also used in the `Computer` export, and - # as `Computer` is not derived from `Data`, it cannot be exported by `data_export`, so - # `validate_output_filename` cannot be removed in favor of the validation done in `data_export`. - validate_output_filename( - output_file=output_file, - overwrite=overwrite, - ) - except (FileExistsError, IsADirectoryError) as exception: - raise click.BadParameter(str(exception), param_hint='OUTPUT_FILE') from exception + # Handle --all option + if export_all: + if codes_and_output: + echo.echo_critical('Cannot specify both --all and individual codes.') + + # Get all non-hidden codes + query = orm.QueryBuilder() + query.append(orm.Code, filters={f'extras.{orm.Code.HIDDEN_KEY}': {'!==': True}}) + codes = [code for [code] in query.all()] + if not codes: + echo.echo_report('No codes found in the database.') + return + output_file = None + else: + # Parse codes and potential output file + if not codes_and_output: + echo.echo_critical('Must specify either --all or individual code identifiers.') + + codes = [] + output_file = None + + # Try to parse all arguments as codes first + for i, arg in enumerate(codes_and_output): + try: + code = orm.load_code(arg) + codes.append(code) + except NotExistent: + # This argument is not a valid code identifier + if i == len(codes_and_output) - 1: + # Last argument and not a code - treat as output file + output_file = pathlib.Path(arg) + break + else: + # Not last argument and not a code - error + echo.echo_critical(f'Invalid code identifier: {arg}') + + if not codes: + echo.echo_critical('No valid code identifiers provided.') + + # Validate output file usage + if output_file and len(codes) > 1: + msg = 'Custom output filename can only be provided if a single code is being exported.' + raise click.BadParameter(msg) + + # Export each code + for code in codes: + # Determine the output file for this specific code + if output_file is None: + current_output_file = pathlib.Path(f'{code.full_label}.{fileformat}') + else: + current_output_file = output_file - try: - data_export( - node=code, - output_fname=output_file, - fileformat=fileformat, - other_args=other_args, - overwrite=overwrite, - ) - except Exception as exception: - echo.echo_critical(f'Error in the `data_export` function: {exception}') + try: + # In principle, output file validation is also done in the `data_export` function. However, the + # `validate_output_filename` function is applied here, as well, as it is also used in the `Computer` export, + # and as `Computer` is not derived from `Data`, it cannot be exported by `data_export`, so + # `validate_output_filename` cannot be removed in favor of the validation done in `data_export`. + validate_output_filename( + output_file=current_output_file, + overwrite=overwrite, + ) + except (FileExistsError, IsADirectoryError) as exception: + raise click.BadParameter(str(exception), param_hint='OUTPUT_FILE') from exception - echo.echo_success(f'Code<{code.pk}> {code.label} exported to file `{output_file}`.') + try: + data_export( + node=code, + output_fname=current_output_file, + fileformat=fileformat, + other_args=other_args, + overwrite=overwrite, + ) + except Exception as exception: + echo.echo_critical(f'Error in the `data_export` function: {exception}') + + echo.echo_success(f'Code<{code.pk}> {code.label} exported to file `{current_output_file}`.') @verdi_code.command() From 2d1ad3a3cab4f978b4eaffba5a63b7d5dc59241d Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Thu, 11 Sep 2025 17:31:53 +0200 Subject: [PATCH 2/3] Multi-computer export --- src/aiida/cmdline/commands/cmd_computer.py | 255 ++++++++++++++++----- 1 file changed, 193 insertions(+), 62 deletions(-) diff --git a/src/aiida/cmdline/commands/cmd_computer.py b/src/aiida/cmdline/commands/cmd_computer.py index 7c071b72e1..eddd207640 100644 --- a/src/aiida/cmdline/commands/cmd_computer.py +++ b/src/aiida/cmdline/commands/cmd_computer.py @@ -745,91 +745,222 @@ def computer_export(): @computer_export.command('setup') -@arguments.COMPUTER() -@arguments.OUTPUT_FILE(type=click.Path(exists=False, path_type=pathlib.Path), required=False) +@click.argument('computers_and_output', nargs=-1, required=False) +@click.option('-a', '--all', 'export_all', is_flag=True, help='Export all computers.') @options.OVERWRITE() @options.SORT() @with_dbenv() -def computer_export_setup(computer, output_file, overwrite, sort): - """Export computer setup to a YAML file.""" - import yaml +def computer_export_setup(computers_and_output, export_all, overwrite, sort): + """Export computer setup(s) to YAML file(s). - 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(), - } + Usage: + - verdi computer export setup COMPUTER_ID [COMPUTER_ID ...] [OUTPUT_FILE] + - verdi computer export setup --all - if output_file is None: - output_file = pathlib.Path(f'{computer.label}-setup.yaml') - try: - validate_output_filename(output_file=output_file, overwrite=overwrite) - except (FileExistsError, IsADirectoryError) as exception: - raise click.BadParameter(str(exception), param_hint='OUTPUT_FILE') from exception + If no output file is given, default names are created based on the computer labels. + Custom output filename can only be provided when exporting a single computer. + """ + import yaml - 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}).' - ) + from aiida import orm + from aiida.common import NotExistent + from aiida.orm import load_computer + + # Handle --all option + if export_all: + if computers_and_output: + echo.echo_critical('Cannot specify both --all and individual computers.') + + # Get all computers + query = orm.QueryBuilder() + query.append(orm.Computer) + computers = [computer for [computer] in query.all()] + if not computers: + echo.echo_report('No computers found in the database.') + return + output_file = None else: - echo.echo_success(f"Computer<{computer.pk}> {computer.label} setup exported to file '{output_file}'.") + # Parse computers and potential output file + if not computers_and_output: + echo.echo_critical('Must specify either --all or individual computer identifiers.') + + computers = [] + output_file = None + + # Try to parse all arguments as computers first + for i, arg in enumerate(computers_and_output): + try: + computer = load_computer(arg) + computers.append(computer) + except NotExistent: + # This argument is not a valid computer identifier + if i == len(computers_and_output) - 1: + # Last argument and not a computer - treat as output file + output_file = pathlib.Path(arg) + break + else: + # Not last argument and not a computer - error + echo.echo_critical(f'Invalid computer identifier: {arg}') + + if not computers: + echo.echo_critical('No valid computer identifiers provided.') + + # Validate output file usage + if output_file and len(computers) > 1: + msg = 'Custom output filename can only be provided if a single computer is being exported.' + raise click.BadParameter(msg) + + # Export each computer + for computer in computers: + 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(), + } + + # Determine the output file for this specific computer + if output_file is None: + current_output_file = pathlib.Path(f'{computer.label}-setup.yaml') + else: + current_output_file = output_file + + try: + validate_output_filename(output_file=current_output_file, overwrite=overwrite) + except (FileExistsError, IsADirectoryError) as exception: + raise click.BadParameter(str(exception), param_hint='OUTPUT_FILE') from exception + + try: + current_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 '{current_output_file}'." + ) @computer_export.command('config') -@arguments.COMPUTER() -@arguments.OUTPUT_FILE(type=click.Path(exists=False, path_type=pathlib.Path), required=False) +@click.argument('computers_and_output', nargs=-1, required=False) +@click.option('-a', '--all', 'export_all', is_flag=True, help='Export all computers.') @options.USER( help='Email address of the AiiDA user from whom to export this computer (if different from default user).' ) @options.OVERWRITE() @options.SORT() @with_dbenv() -def computer_export_config(computer, output_file, user, overwrite, sort): - """Export computer transport configuration for a user to a YAML file.""" +def computer_export_config(computers_and_output, export_all, user, overwrite, sort): + """Export computer transport configuration(s) for a user to YAML file(s). + + Usage: + - verdi computer export config COMPUTER_ID [COMPUTER_ID ...] [OUTPUT_FILE] + - verdi computer export config --all + + If no output file is given, default names are created based on the computer labels. + Custom output filename can only be provided when exporting a single computer. + """ 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.' - ) + from aiida import orm + from aiida.common import NotExistent + from aiida.orm import load_computer + + # Handle --all option + if export_all: + if computers_and_output: + echo.echo_critical('Cannot specify both --all and individual computers.') + + # Get all configured computers + query = orm.QueryBuilder() + query.append(orm.Computer) + all_computers = [computer for [computer] in query.all()] + computers = [comp for comp in all_computers if comp.is_configured] + if not computers: + echo.echo_report('No configured computers found in the database.') + return + output_file = None else: + # Parse computers and potential output file + if not computers_and_output: + echo.echo_critical('Must specify either --all or individual computer identifiers.') + + computers = [] + output_file = None + + # Try to parse all arguments as computers first + for i, arg in enumerate(computers_and_output): + try: + computer = load_computer(arg) + computers.append(computer) + except NotExistent: + # This argument is not a valid computer identifier + if i == len(computers_and_output) - 1: + # Last argument and not a computer - treat as output file + output_file = pathlib.Path(arg) + break + else: + # Not last argument and not a computer - error + echo.echo_critical(f'Invalid computer identifier: {arg}') + + if not computers: + echo.echo_critical('No valid computer identifiers provided.') + + # Validate output file usage + if output_file and len(computers) > 1: + msg = 'Custom output filename can only be provided if a single computer is being exported.' + raise click.BadParameter(msg) + + # Export each computer + for computer in computers: + if not computer.is_configured: + echo.echo_warning( + f'Computer<{computer.pk}> {computer.label} configuration cannot be exported,' + ' because computer has not been configured yet. Skipping.' + ) + continue + + # Determine the output file for this specific computer if output_file is None: - output_file = pathlib.Path(f'{computer.label}-config.yaml') + current_output_file = pathlib.Path(f'{computer.label}-config.yaml') + else: + current_output_file = output_file + try: - validate_output_filename(output_file=output_file, overwrite=overwrite) + validate_output_filename(output_file=current_output_file, overwrite=overwrite) except (FileExistsError, IsADirectoryError) as exception: raise click.BadParameter(str(exception), param_hint='OUTPUT_FILE') from exception - try: - computer_configuration = computer.get_configuration(user) - output_file.write_text(yaml.dump(computer_configuration, sort_keys=sort), 'utf-8') - - except Exception as exception: - 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}: {exception!s}.' # noqa: E501 - ) + try: + computer_configuration = computer.get_configuration(user) + current_output_file.write_text(yaml.dump(computer_configuration, sort_keys=sort), 'utf-8') + + except Exception as exception: + error_traceback = traceback.format_exc() + echo.CMDLINE_LOGGER.debug(error_traceback) + if user is None: + echo.echo_critical( + 'Unexpected error while exporting configuration for ' + f'Computer<{computer.pk}> {computer.label}: {exception!s}.' + ) + else: + echo.echo_critical( + f'Unexpected error while exporting configuration for Computer<{computer.pk}> {computer.label}' + f' and User<{user.pk}> {user.email}: {exception!s}.' + ) else: - echo.echo_critical( - f'Unexpected error while exporting configuration for Computer<{computer.pk}> {computer.label}' - f' and User<{user.pk}> {user.email}: {exception!s}.' + echo.echo_success( + f'Computer<{computer.pk}> {computer.label} configuration exported to file `{current_output_file}`.' ) - else: - echo.echo_success(f'Computer<{computer.pk}> {computer.label} configuration exported to file `{output_file}`.') From 50e976c61f1b79c0d029617db2952176dc0dbf72 Mon Sep 17 00:00:00 2001 From: Julian Geiger Date: Fri, 12 Sep 2025 10:26:34 +0200 Subject: [PATCH 3/3] verdi computer setup-many --- docs/source/reference/command_line.rst | 23 +- src/aiida/cmdline/commands/cmd_computer.py | 288 +++++++-------------- src/aiida/cmdline/params/options/main.py | 8 + 3 files changed, 115 insertions(+), 204 deletions(-) diff --git a/docs/source/reference/command_line.rst b/docs/source/reference/command_line.rst index 27409608f5..ffe619e7b6 100644 --- a/docs/source/reference/command_line.rst +++ b/docs/source/reference/command_line.rst @@ -100,17 +100,18 @@ Below is a list with all available subcommands. --help Show this message and exit. Commands: - 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. - show Show detailed information for a computer. - test Test the connection to a computer. + 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. + setup-many Create multiple computers from YAML configuration files. + show Show detailed information for a computer. + test Test the connection to a computer. .. _reference:command-line:verdi-config: diff --git a/src/aiida/cmdline/commands/cmd_computer.py b/src/aiida/cmdline/commands/cmd_computer.py index eddd207640..960d9e1493 100644 --- a/src/aiida/cmdline/commands/cmd_computer.py +++ b/src/aiida/cmdline/commands/cmd_computer.py @@ -322,6 +322,39 @@ def computer_setup(ctx, non_interactive, **kwargs): echo.echo_report(f' verdi -p {profile.name} computer configure {computer.transport_type} {computer.label}') +@verdi_computer.command('setup-many') +@click.argument('config_files', nargs=-1, required=True, type=click.Path(exists=True, path_type=pathlib.Path)) +@with_dbenv() +def computer_setup_many(config_files): + """Create multiple computers from YAML configuration files.""" + import yaml + + from aiida.common.exceptions import IntegrityError + from aiida.orm.utils.builders.computer import ComputerBuilder + + for config_path in config_files: + try: + with open(config_path, 'r', encoding='utf-8') as f: + config_data = yaml.safe_load(f) + + computer_builder = ComputerBuilder(**config_data) + computer = computer_builder.new() + computer.store() + + echo.echo_success(f'Computer<{computer.pk}> {computer.label} created') + except IntegrityError as e: + if 'UNIQUE constraint failed: db_dbcomputer.label' in str(e): + msg = ( + f'Error processing {config_path}: Computer with label "{config_data.get("label", "unknown")}"' + 'already exists' + ) + echo.echo_error(msg) + else: + echo.echo_error(f'Error processing {config_path}: Database integrity error - {e}') + except Exception as e: + echo.echo_error(f'Error processing {config_path}: {e}') + + @verdi_computer.command('duplicate') @arguments.COMPUTER(callback=set_computer_builder) @options_computer.LABEL(contextual_default=partial(get_parameter_default, 'label')) @@ -745,222 +778,91 @@ def computer_export(): @computer_export.command('setup') -@click.argument('computers_and_output', nargs=-1, required=False) -@click.option('-a', '--all', 'export_all', is_flag=True, help='Export all computers.') +@arguments.COMPUTER() +@arguments.OUTPUT_FILE(type=click.Path(exists=False, path_type=pathlib.Path), required=False) @options.OVERWRITE() @options.SORT() @with_dbenv() -def computer_export_setup(computers_and_output, export_all, overwrite, sort): - """Export computer setup(s) to YAML file(s). - - Usage: - - verdi computer export setup COMPUTER_ID [COMPUTER_ID ...] [OUTPUT_FILE] - - verdi computer export setup --all - - If no output file is given, default names are created based on the computer labels. - Custom output filename can only be provided when exporting a single computer. - """ +def computer_export_setup(computer, output_file, overwrite, sort): + """Export computer setup to a YAML file.""" import yaml - from aiida import orm - from aiida.common import NotExistent - from aiida.orm import load_computer - - # Handle --all option - if export_all: - if computers_and_output: - echo.echo_critical('Cannot specify both --all and individual computers.') - - # Get all computers - query = orm.QueryBuilder() - query.append(orm.Computer) - computers = [computer for [computer] in query.all()] - if not computers: - echo.echo_report('No computers found in the database.') - return - output_file = None - else: - # Parse computers and potential output file - if not computers_and_output: - echo.echo_critical('Must specify either --all or individual computer identifiers.') - - computers = [] - output_file = None - - # Try to parse all arguments as computers first - for i, arg in enumerate(computers_and_output): - try: - computer = load_computer(arg) - computers.append(computer) - except NotExistent: - # This argument is not a valid computer identifier - if i == len(computers_and_output) - 1: - # Last argument and not a computer - treat as output file - output_file = pathlib.Path(arg) - break - else: - # Not last argument and not a computer - error - echo.echo_critical(f'Invalid computer identifier: {arg}') - - if not computers: - echo.echo_critical('No valid computer identifiers provided.') - - # Validate output file usage - if output_file and len(computers) > 1: - msg = 'Custom output filename can only be provided if a single computer is being exported.' - raise click.BadParameter(msg) - - # Export each computer - for computer in computers: - 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(), - } - - # Determine the output file for this specific computer - if output_file is None: - current_output_file = pathlib.Path(f'{computer.label}-setup.yaml') - else: - current_output_file = output_file + 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: - validate_output_filename(output_file=current_output_file, overwrite=overwrite) - except (FileExistsError, IsADirectoryError) as exception: - raise click.BadParameter(str(exception), param_hint='OUTPUT_FILE') from exception + if output_file is None: + output_file = pathlib.Path(f'{computer.label}-setup.yaml') + try: + validate_output_filename(output_file=output_file, overwrite=overwrite) + except (FileExistsError, IsADirectoryError) as exception: + raise click.BadParameter(str(exception), param_hint='OUTPUT_FILE') from exception - try: - current_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 '{current_output_file}'." - ) + 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') -@click.argument('computers_and_output', nargs=-1, required=False) -@click.option('-a', '--all', 'export_all', is_flag=True, help='Export all computers.') +@arguments.COMPUTER() +@arguments.OUTPUT_FILE(type=click.Path(exists=False, path_type=pathlib.Path), required=False) @options.USER( help='Email address of the AiiDA user from whom to export this computer (if different from default user).' ) @options.OVERWRITE() @options.SORT() @with_dbenv() -def computer_export_config(computers_and_output, export_all, user, overwrite, sort): - """Export computer transport configuration(s) for a user to YAML file(s). - - Usage: - - verdi computer export config COMPUTER_ID [COMPUTER_ID ...] [OUTPUT_FILE] - - verdi computer export config --all - - If no output file is given, default names are created based on the computer labels. - Custom output filename can only be provided when exporting a single computer. - """ +def computer_export_config(computer, output_file, user, overwrite, sort): + """Export computer transport configuration for a user to a YAML file.""" import yaml - from aiida import orm - from aiida.common import NotExistent - from aiida.orm import load_computer - - # Handle --all option - if export_all: - if computers_and_output: - echo.echo_critical('Cannot specify both --all and individual computers.') - - # Get all configured computers - query = orm.QueryBuilder() - query.append(orm.Computer) - all_computers = [computer for [computer] in query.all()] - computers = [comp for comp in all_computers if comp.is_configured] - if not computers: - echo.echo_report('No configured computers found in the database.') - return - output_file = None + 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.' + ) else: - # Parse computers and potential output file - if not computers_and_output: - echo.echo_critical('Must specify either --all or individual computer identifiers.') - - computers = [] - output_file = None - - # Try to parse all arguments as computers first - for i, arg in enumerate(computers_and_output): - try: - computer = load_computer(arg) - computers.append(computer) - except NotExistent: - # This argument is not a valid computer identifier - if i == len(computers_and_output) - 1: - # Last argument and not a computer - treat as output file - output_file = pathlib.Path(arg) - break - else: - # Not last argument and not a computer - error - echo.echo_critical(f'Invalid computer identifier: {arg}') - - if not computers: - echo.echo_critical('No valid computer identifiers provided.') - - # Validate output file usage - if output_file and len(computers) > 1: - msg = 'Custom output filename can only be provided if a single computer is being exported.' - raise click.BadParameter(msg) - - # Export each computer - for computer in computers: - if not computer.is_configured: - echo.echo_warning( - f'Computer<{computer.pk}> {computer.label} configuration cannot be exported,' - ' because computer has not been configured yet. Skipping.' - ) - continue - - # Determine the output file for this specific computer if output_file is None: - current_output_file = pathlib.Path(f'{computer.label}-config.yaml') - else: - current_output_file = output_file - + output_file = pathlib.Path(f'{computer.label}-config.yaml') try: - validate_output_filename(output_file=current_output_file, overwrite=overwrite) + validate_output_filename(output_file=output_file, overwrite=overwrite) except (FileExistsError, IsADirectoryError) as exception: raise click.BadParameter(str(exception), param_hint='OUTPUT_FILE') from exception - try: - computer_configuration = computer.get_configuration(user) - current_output_file.write_text(yaml.dump(computer_configuration, sort_keys=sort), 'utf-8') - - except Exception as exception: - error_traceback = traceback.format_exc() - echo.CMDLINE_LOGGER.debug(error_traceback) - if user is None: - echo.echo_critical( - 'Unexpected error while exporting configuration for ' - f'Computer<{computer.pk}> {computer.label}: {exception!s}.' - ) - else: - echo.echo_critical( - f'Unexpected error while exporting configuration for Computer<{computer.pk}> {computer.label}' - f' and User<{user.pk}> {user.email}: {exception!s}.' - ) + try: + computer_configuration = computer.get_configuration(user) + output_file.write_text(yaml.dump(computer_configuration, sort_keys=sort), 'utf-8') + + except Exception as exception: + 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}: {exception!s}.' # noqa: E501 + ) else: - echo.echo_success( - f'Computer<{computer.pk}> {computer.label} configuration exported to file `{current_output_file}`.' + echo.echo_critical( + f'Unexpected error while exporting configuration for Computer<{computer.pk}> {computer.label}' + f' and User<{user.pk}> {user.email}: {exception!s}.' ) + else: + echo.echo_success(f'Computer<{computer.pk}> {computer.label} configuration exported to file `{output_file}`.') diff --git a/src/aiida/cmdline/params/options/main.py b/src/aiida/cmdline/params/options/main.py index 4171d99213..02570c9693 100644 --- a/src/aiida/cmdline/params/options/main.py +++ b/src/aiida/cmdline/params/options/main.py @@ -47,6 +47,7 @@ 'COMPUTER', 'COMPUTERS', 'CONFIG_FILE', + 'CONFIG_FILES', 'DATA', 'DATUM', 'DB_BACKEND', @@ -735,6 +736,13 @@ def set_log_level(ctx: click.Context, _param: click.Parameter, value: t.Any) -> help='Load option values from configuration file in yaml format (local path or URL).', ) +CONFIG_FILES = click.option( + '--config-files', + multiple=True, + type=click.Path(exists=True, path_type=pathlib.Path), + help='Load computer setup from multiple configuration files in YAML format. Can be specified multiple times.', +) + IDENTIFIER = OverridableOption( '-i', '--identifier',