diff --git a/src/aiida/cmdline/commands/cmd_computer.py b/src/aiida/cmdline/commands/cmd_computer.py index 67123f8a24..7606daa10a 100644 --- a/src/aiida/cmdline/commands/cmd_computer.py +++ b/src/aiida/cmdline/commands/cmd_computer.py @@ -8,6 +8,7 @@ ########################################################################### """`verdi computer` command.""" +import json import pathlib import traceback from copy import deepcopy @@ -270,6 +271,19 @@ def set_computer_builder(ctx, param, value): return value +# Helper function to set template vars in context +def set_template_vars_in_context(ctx, param, value): + """Set template variables in the context for the config provider to use.""" + if value: + try: + template_var_dict = json.loads(value) + # Store template vars in context for the config provider + ctx._template_vars = template_var_dict + except json.JSONDecodeError as e: + raise click.BadParameter(f'Invalid JSON in template-vars: {e}') + return value + + @verdi_computer.command('setup') @options_computer.LABEL() @options_computer.HOSTNAME() @@ -285,22 +299,32 @@ def set_computer_builder(ctx, param, value): @options_computer.PREPEND_TEXT() @options_computer.APPEND_TEXT() @options.NON_INTERACTIVE() -@options.CONFIG_FILE() +@options.TEMPLATE_VARS() # This should come before TEMPLATE_FILE +@options.TEMPLATE_FILE() # This will process the template and set defaults @click.pass_context @with_dbenv() def computer_setup(ctx, non_interactive, **kwargs): """Create a new computer.""" from aiida.orm.utils.builders.computer import ComputerBuilder - if kwargs['label'] in get_computer_names(): + # Debug output + print(f'Debug: non_interactive = {non_interactive}') + print(f'Debug: kwargs keys = {list(kwargs.keys())}') + print(f'Debug: ctx.default_map = {ctx.default_map}') + + # Check for existing computer + if kwargs.get('label') and kwargs['label'] in get_computer_names(): echo.echo_critical( 'A computer called {c} already exists. ' 'Use "verdi computer duplicate {c}" to set up a new ' 'computer starting from the settings of {c}.'.format(c=kwargs['label']) ) - kwargs['transport'] = kwargs['transport'].name - kwargs['scheduler'] = kwargs['scheduler'].name + # Convert entry points to their names + if kwargs.get('transport'): + kwargs['transport'] = kwargs['transport'].name + if kwargs.get('scheduler'): + kwargs['scheduler'] = kwargs['scheduler'].name computer_builder = ComputerBuilder(**kwargs) try: diff --git a/src/aiida/cmdline/params/options/main.py b/src/aiida/cmdline/params/options/main.py index 364711f21a..20daf9d3af 100644 --- a/src/aiida/cmdline/params/options/main.py +++ b/src/aiida/cmdline/params/options/main.py @@ -8,11 +8,14 @@ ########################################################################### """Module with pre-defined reusable commandline options that can be used as `click` decorators.""" +import json import pathlib import click from aiida.brokers.rabbitmq.defaults import BROKER_DEFAULTS +from aiida.cmdline.utils import echo +from aiida.cmdline.utils.template_config import load_and_process_template from aiida.common.log import LOG_LEVELS, configure_logging from aiida.manage.external.postgres import DEFAULT_DBINFO @@ -113,6 +116,8 @@ 'SORT', 'START_DATE', 'SYMLINK_CALCS', + 'TEMPLATE_FILE', + 'TEMPLATE_VARS', 'TIMEOUT', 'TRAJECTORY_INDEX', 'TRANSPORT', @@ -910,3 +915,73 @@ def set_log_level(ctx, _param, value): show_default=True, help='End date for node mtime range selection for node collection dumping.', ) + +import click + +from aiida.cmdline.utils import echo + +from .overridable import OverridableOption + + +# Template processing callback +def process_template_callback(ctx, param, value): + """Process template file and update context defaults.""" + if not value: + return value + + ctx.default_map = ctx.default_map or {} + + # Get template vars from context if they were set by TEMPLATE_VARS option + template_vars = getattr(ctx, '_template_vars', None) + + # Check if we're in non-interactive mode + non_interactive = ctx.params.get('non_interactive', False) + + try: + # Load and process the template + config_data = load_and_process_template(value, interactive=not non_interactive, template_vars=template_vars) + + # Update the default map with template values + for key, template_value in config_data.items(): + if key not in ctx.default_map: + ctx.default_map[key] = template_value + + except Exception as e: + echo.echo_critical(f'Error processing template: {e}') + + return value + + +# Template vars callback +def set_template_vars_callback(ctx, param, value): + """Set template variables in the context for the template option to use.""" + if value: + try: + template_var_dict = json.loads(value) + # Store template vars in context for the template option to use + ctx._template_vars = template_var_dict + except json.JSONDecodeError as e: + raise click.BadParameter(f'Invalid JSON in template-vars: {e}') + return value + + +# Template options using simple OverridableOption with callbacks +TEMPLATE_VARS = OverridableOption( + '--template-vars', + type=click.STRING, + is_eager=True, # Process before template option + callback=set_template_vars_callback, + expose_value=False, # Don't pass to command function + help='JSON string containing template variable values for non-interactive mode. ' + 'Example: \'{"label": "my-computer", "slurm_account": "my_account"}\'', +) + +TEMPLATE_FILE = OverridableOption( + '--template', + type=click.STRING, + is_eager=True, # Process template before other options + callback=process_template_callback, + expose_value=False, # Don't pass template to the command function + help='Load computer setup from configuration file in YAML format (local path or URL). ' + 'Supports Jinja2 templates with interactive prompting.', +) diff --git a/src/aiida/cmdline/utils/template_config.py b/src/aiida/cmdline/utils/template_config.py new file mode 100644 index 0000000000..eb64b1bb78 --- /dev/null +++ b/src/aiida/cmdline/utils/template_config.py @@ -0,0 +1,150 @@ +from typing import Any, Dict, List, Optional + +import click +import requests +import yaml +from jinja2 import BaseLoader, Environment, meta + +from aiida.cmdline.utils import echo + + +class StringTemplateLoader(BaseLoader): + """Jinja2 loader that loads templates from strings.""" + + def __init__(self, template_string: str): + self.template_string = template_string + + def get_source(self, environment, template): + return self.template_string, None, lambda: True + + +def prompt_for_template_variables(template_variables: Dict[str, Any]) -> Dict[str, Any]: + """Prompt user for template variables based on metadata definitions.""" + values = {} + + echo.echo_report('Template variables detected. Please provide values:') + echo.echo('') + + for var_name, var_config in template_variables.items(): + key_display = var_config.get('key_display', var_name) + description = var_config.get('description', f'Value for {var_name}') + var_type = var_config.get('type', 'text') + default = var_config.get('default') + options = var_config.get('options', []) + + # Display help text + echo.echo(f'{click.style(key_display, fg="yellow")}') + echo.echo(f' {description}') + + if var_type == 'list' and options: + echo.echo(f' Options: {", ".join(options)}') + while True: + value = click.prompt(' Enter value', default=default, show_default=True if default else False) + if value in options: + values[var_name] = value + break + else: + echo.echo_error(f'Invalid option. Please choose from: {", ".join(options)}') + else: + value = click.prompt(' Enter value', default=default, show_default=True if default else False) + values[var_name] = value + + echo.echo('') + + return values + + +def detect_template_variables(template_content: str) -> List[str]: + """Detect Jinja2 variables in template content.""" + env = Environment(loader=StringTemplateLoader(template_content)) + ast = env.parse(template_content) + return list(meta.find_undeclared_variables(ast)) + + +def load_and_process_template( + file_path_or_url: str, interactive: bool = True, template_vars: Optional[Dict[str, Any]] = None +) -> Dict[str, Any]: + """Load and process a template configuration file.""" + + # Load content + if file_path_or_url.startswith(('http://', 'https://')): + try: + response = requests.get(file_path_or_url, timeout=10) + response.raise_for_status() + content = response.text + except requests.RequestException as e: + raise click.BadParameter(f'Failed to fetch URL {file_path_or_url}: {e}') + else: + try: + with open(file_path_or_url, 'r', encoding='utf-8') as f: + content = f.read() + except IOError as e: + raise click.BadParameter(f'Failed to read file {file_path_or_url}: {e}') + + # Parse YAML to get metadata + try: + full_config = yaml.safe_load(content) + except yaml.YAMLError as e: + raise click.BadParameter(f'Invalid YAML: {e}') + + # Extract metadata and template variables (if they exist) + metadata = full_config.pop('metadata', {}) + template_variables = metadata.get('template_variables', {}) + + # Detect variables that need values + detected_vars = detect_template_variables(content) + + # If no template variables detected, just return the config + if not detected_vars: + return full_config + + # Filter to only prompt for variables that are actually used and defined in metadata + vars_to_prompt = {var: config for var, config in template_variables.items() if var in detected_vars} + + if vars_to_prompt: + if interactive: + # Interactive prompting for template variables + template_values = prompt_for_template_variables(vars_to_prompt) + else: + # Non-interactive mode + if not template_vars: + raise click.BadParameter( + f'Template variables detected ({", ".join(detected_vars)}) but no values provided. ' + 'Use --template-vars to provide values in JSON format.' + ) + template_values = template_vars + + # Render the template with provided values + env = Environment(loader=StringTemplateLoader(content)) + template = env.from_string(content) + rendered_content = template.render(**template_values) + + # Parse the rendered YAML + try: + config = yaml.safe_load(rendered_content) + except yaml.YAMLError as e: + raise click.BadParameter(f'Invalid YAML after template rendering: {e}') + else: + # Template variables detected but none defined in metadata + # This could happen with simple Jinja variables like {{ username }} + if interactive: + echo.echo_warning(f'Template variables detected ({", ".join(detected_vars)}) but no metadata found.') + echo.echo_warning('You may need to provide values manually or the template may not render correctly.') + + if template_vars: + # Try to render with provided vars + env = Environment(loader=StringTemplateLoader(content)) + template = env.from_string(content) + rendered_content = template.render(**template_vars) + try: + config = yaml.safe_load(rendered_content) + except yaml.YAMLError as e: + raise click.BadParameter(f'Invalid YAML after template rendering: {e}') + else: + # Return original config and hope for the best + config = full_config + + # Remove metadata section if it exists + config.pop('metadata', None) + + return config