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 command verdi presto #6351

Merged
merged 1 commit into from
May 3, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 38 additions & 0 deletions docs/source/reference/command_line.rst
Original file line number Diff line number Diff line change
Expand Up @@ -313,6 +313,44 @@ Below is a list with all available subcommands.
list Display a list of all available plugins.


.. _reference:command-line:verdi-presto:

``verdi presto``
----------------

.. code:: console

Usage: [OPTIONS]

Set up a new profile in a jiffy.

This command aims to make setting up a new profile as easy as possible. It intentionally
provides only a limited amount of options to customize the profile and by default does
not require any options to be specified at all. For full control, please use `verdi
profile setup`.

After running `verdi presto` you can immediately start using AiiDA without additional
setup. The created profile uses the `core.sqlite_dos` storage plugin which does not
require any services, such as PostgreSQL. The broker service RabbitMQ is also optional.
The command tries to connect to it using default settings and configures it for the
profile if found. Otherwise, the profile is created without a broker, in which case some
functionality will be unavailable, most notably running the daemon and submitting
processes to said daemon.

The command performs the following actions:

* Create a new profile that is set as the new default
* Create a default user for the profile (email can be configured through the `--email` option)
* Set up the localhost as a `Computer` and configure it
* Set a number of configuration options with sensible defaults

Options:
--profile-name TEXT Name of the profile. By default, a unique name starting with
`presto` is automatically generated. [default: (dynamic)]
--email TEXT Email of the default user. [default: aiida@localhost]
--help Show this message and exit.


.. _reference:command-line:verdi-process:

``verdi process``
Expand Down
21 changes: 21 additions & 0 deletions src/aiida/brokers/rabbitmq/defaults.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
"""Defaults related to RabbitMQ."""

from __future__ import annotations

import typing as t

from aiida.common.extendeddicts import AttributeDict

__all__ = ('BROKER_DEFAULTS',)
Expand All @@ -19,3 +23,20 @@
'heartbeat': 600,
}
)


def detect_rabbitmq_config() -> dict[str, t.Any] | None:
"""Try to connect to a RabbitMQ server with the default connection parameters.

:returns: The connection parameters if the RabbitMQ server was successfully connected to, or ``None`` otherwise.
"""
from kiwipy.rmq.threadcomms import connect

connection_params = dict(BROKER_DEFAULTS)

try:
connect(connection_params=connection_params)
except ConnectionError:
return None

return connection_params
1 change: 1 addition & 0 deletions src/aiida/cmdline/commands/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
cmd_help,
cmd_node,
cmd_plugin,
cmd_presto,
cmd_process,
cmd_profile,
cmd_rabbitmq,
Expand Down
133 changes: 133 additions & 0 deletions src/aiida/cmdline/commands/cmd_presto.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,133 @@
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved. #
# This file is part of the AiiDA code. #
# #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
# For further information on the license, see the LICENSE.txt file #
# For further information please visit http://www.aiida.net #
###########################################################################
"""``verdi presto`` command."""

from __future__ import annotations

import pathlib
import re
import typing as t

import click

from aiida.cmdline.commands.cmd_verdi import verdi
from aiida.cmdline.utils import echo
from aiida.manage.configuration import get_config_option

DEFAULT_PROFILE_NAME_PREFIX: str = 'presto'


def get_default_presto_profile_name():
from aiida.manage import get_config

profile_names = get_config().profile_names
indices = []

for profile_name in profile_names:
if match := re.search(r'presto[-]?(\d+)?', profile_name):
indices.append(match.group(1) or '0')

if not indices:
return DEFAULT_PROFILE_NAME_PREFIX

last_index = int(sorted(indices)[-1])

return f'{DEFAULT_PROFILE_NAME_PREFIX}-{last_index + 1}'


@verdi.command('presto')
@click.option(
'--profile-name',
default=lambda: get_default_presto_profile_name(),
show_default=True,
help=f'Name of the profile. By default, a unique name starting with `{DEFAULT_PROFILE_NAME_PREFIX}` is '
'automatically generated.',
)
@click.option(
'--email',
default=get_config_option('autofill.user.email') or 'aiida@localhost',
show_default=True,
help='Email of the default user.',
)
@click.pass_context
def verdi_presto(ctx, profile_name, email):
"""Set up a new profile in a jiffy.

This command aims to make setting up a new profile as easy as possible. It intentionally provides only a limited
amount of options to customize the profile and by default does not require any options to be specified at all. For
full control, please use `verdi profile setup`.

After running `verdi presto` you can immediately start using AiiDA without additional setup. The created profile
uses the `core.sqlite_dos` storage plugin which does not require any services, such as PostgreSQL. The broker
service RabbitMQ is also optional. The command tries to connect to it using default settings and configures it for
the profile if found. Otherwise, the profile is created without a broker, in which case some functionality will be
unavailable, most notably running the daemon and submitting processes to said daemon.

The command performs the following actions:

\b
* Create a new profile that is set as the new default
* Create a default user for the profile (email can be configured through the `--email` option)
* Set up the localhost as a `Computer` and configure it
* Set a number of configuration options with sensible defaults

"""
from aiida.brokers.rabbitmq.defaults import detect_rabbitmq_config
from aiida.common import exceptions
from aiida.manage.configuration import create_profile, load_profile
from aiida.orm import Computer

storage_config: dict[str, t.Any] = {}
storage_backend = 'core.sqlite_dos'

broker_config = detect_rabbitmq_config()
broker_backend = 'core.rabbitmq' if broker_config is not None else None

if broker_config is None:
echo.echo_report('RabbitMQ server not found: configuring the profile without a broker.')
else:
echo.echo_report('RabbitMQ server detected: configuring the profile with a broker.')

try:
profile = create_profile(
ctx.obj.config,
name=profile_name,
email=email,
storage_backend=storage_backend,
storage_config=storage_config,
broker_backend=broker_backend,
broker_config=broker_config,
)
except (ValueError, TypeError, exceptions.EntryPointError, exceptions.StorageMigrationError) as exception:
echo.echo_critical(str(exception))

echo.echo_success(f'Created new profile `{profile.name}`.')

ctx.obj.config.set_option('runner.poll.interval', 1, scope=profile.name)
ctx.obj.config.set_default_profile(profile.name, overwrite=True)
ctx.obj.config.store()

load_profile(profile.name, allow_switch=True)
echo.echo_info(f'Loaded newly created profile `{profile.name}`.')

filepath_scratch = pathlib.Path(ctx.obj.config.dirpath) / 'scratch' / profile.name

computer = Computer(
label='localhost',
hostname='localhost',
description='Localhost automatically created by `verdi presto`',
transport_type='core.local',
scheduler_type='core.direct',
workdir=str(filepath_scratch),
).store()
computer.configure(safe_interval=0)
computer.set_minimum_job_poll_interval(1)
computer.set_default_mpiprocs_per_machine(1)

echo.echo_success('Configured the localhost as a computer.')
5 changes: 4 additions & 1 deletion src/aiida/cmdline/commands/cmd_profile.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,10 @@ def profile_list():
echo.echo_report(f'configuration folder: {config.dirpath}')

if not config.profiles:
echo.echo_warning('no profiles configured: run `verdi setup` to create one')
echo.echo_warning(
'no profiles configured: Run `verdi presto` to automatically setup a profile using all defaults or use '
'`verdi profile setup` for more control.'
)
else:
sort = lambda profile: profile.name # noqa: E731
highlight = lambda profile: profile.name == config.default_profile_name # noqa: E731
Expand Down
5 changes: 4 additions & 1 deletion src/aiida/cmdline/commands/cmd_status.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,10 @@ def verdi_status(print_traceback, no_rmq):

if profile is None:
print_status(ServiceStatus.WARNING, 'profile', 'no profile configured yet')
echo.echo_report('Configure a profile by running `verdi quicksetup` or `verdi setup`.')
echo.echo_report(
'Run `verdi presto` to automatically setup a profile using all defaults or use `verdi profile setup` '
'for more control.'
)
return

print_status(ServiceStatus.UP, 'profile', profile.name)
Expand Down
5 changes: 2 additions & 3 deletions src/aiida/manage/configuration/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,8 +188,7 @@ def profile_context(profile: 'Profile' | str | None = None, allow_switch=False)

manager = get_manager()
current_profile = manager.get_profile()
manager.load_profile(profile, allow_switch)
yield profile
yield manager.load_profile(profile, allow_switch)
if current_profile is None:
manager.unload_profile()
else:
Expand Down Expand Up @@ -234,7 +233,7 @@ def create_default_user(
if user:
manager.set_default_user_email(profile, user.email)

return
return user


def create_profile(
Expand Down
2 changes: 1 addition & 1 deletion src/aiida/storage/sqlite_dos/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ class Model(BaseModel):
filepath: str = Field(
title='Directory of the backend',
description='Filepath of the directory in which to store data for this backend.',
default_factory=lambda: AIIDA_CONFIG_FOLDER / 'repository' / f'sqlite_dos_{uuid4().hex}',
default_factory=lambda: str(AIIDA_CONFIG_FOLDER / 'repository' / f'sqlite_dos_{uuid4().hex}'),
)

@field_validator('filepath')
Expand Down
50 changes: 50 additions & 0 deletions tests/cmdline/commands/test_presto.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
"""Tests for ``verdi presto``."""

import pytest
from aiida.cmdline.commands.cmd_presto import get_default_presto_profile_name, verdi_presto
from aiida.manage.configuration import profile_context
from aiida.manage.configuration.config import Config
from aiida.orm import Computer


@pytest.mark.parametrize(
'profile_names, expected',
(
([], 'presto'),
(['main', 'sqlite'], 'presto'),
(['presto'], 'presto-1'),
(['presto', 'presto-5', 'presto-2'], 'presto-6'),
(['presto', 'main', 'presto-2', 'sqlite'], 'presto-3'),
),
)
def test_get_default_presto_profile_name(monkeypatch, profile_names, expected):
"""Test the dynamic default profile function."""

def get_profile_names(self):
return profile_names

monkeypatch.setattr(Config, 'profile_names', property(get_profile_names))
assert get_default_presto_profile_name() == expected


@pytest.mark.usefixtures('empty_config')
@pytest.mark.parametrize('with_broker', (True, False))
def test_presto(run_cli_command, with_broker, monkeypatch):
"""Test the ``verdi presto``."""
from aiida.brokers.rabbitmq import defaults

if not with_broker:
# Patch the RabbitMQ detection function to pretend it could not find the service
monkeypatch.setattr(defaults, 'detect_rabbitmq_config', lambda: None)

result = run_cli_command(verdi_presto)
assert 'Created new profile `presto`.' in result.output

with profile_context('presto', allow_switch=True) as profile:
assert profile.name == 'presto'
localhost = Computer.collection.get(label='localhost')
assert localhost.is_configured
if with_broker:
assert profile.process_control_backend == 'core.rabbitmq'
else:
assert profile.process_control_backend is None
Loading