Skip to content

Commit

Permalink
Add verdi devel launch-multiply-add
Browse files Browse the repository at this point in the history
  • Loading branch information
GeigerJ2 committed Nov 5, 2024
1 parent 8350df0 commit b7c8286
Show file tree
Hide file tree
Showing 3 changed files with 88 additions and 1 deletion.
1 change: 1 addition & 0 deletions docs/source/reference/command_line.rst
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,7 @@ Below is a list with all available subcommands.
check-load-time Check for common indicators that slowdown `verdi`.
check-undesired-imports Check that verdi does not import python modules it shouldn't.
launch-add Launch an ``ArithmeticAddCalculation``.
launch-multiply-add Launch a ``MultipylAddWorkChain``.
rabbitmq Commands to interact with RabbitMQ.
run-sql Run a raw SQL command on the profile database (only...
validate-plugins Validate all plugins by checking they can be loaded.
Expand Down
51 changes: 51 additions & 0 deletions src/aiida/cmdline/commands/cmd_devel.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,57 @@ def devel_launch_arithmetic_add(code, daemon, sleep):
echo.echo_warning(f'ArithmeticAddCalculation<{node.pk}> did not finish successfully.')


@verdi_devel.command('launch-multiply-add')
@options.CODE(type=types.CodeParamType(entry_point='core.arithmetic.add'))
@click.option('-d', '--daemon', is_flag=True, help='Submit to the daemon instead of running blockingly.')
def devel_launch_multiply_add(code, daemon):
"""Launch a ``MultipylAddWorkChain``.
Unless specified with the option ``--code``, a suitable ``Code`` is automatically setup. By default the command
configures ``bash`` on the ``localhost``. If the localhost is not yet configured as a ``Computer``, that is also
done automatically.
"""
from shutil import which

from aiida.engine import run_get_node, submit
from aiida.orm import InstalledCode, Int, load_code
from aiida.plugins.factories import WorkflowFactory

default_calc_job_plugin = 'core.arithmetic.add'

if not code:
try:
code = load_code('bash@localhost')
except exceptions.NotExistent:
localhost = prepare_localhost()
code = InstalledCode(
label='bash',
computer=localhost,
filepath_executable=which('bash'),
default_calc_job_plugin=default_calc_job_plugin,
).store()
else:
assert code.default_calc_job_plugin == default_calc_job_plugin

multiply_add = WorkflowFactory('core.arithmetic.multiply_add')
inputs = {
'x': Int(1),
'y': Int(1),
'z': Int(1),
'code': code,
}

if daemon:
node = submit(multiply_add, inputs=inputs)
echo.echo_success(f'Submitted workflow `{node}`')
else:
_, node = run_get_node(multiply_add, inputs)
if node.is_finished_ok:
echo.echo_success(f'MultiplyAddWorkChain<{node.pk}> finished successfully.')
else:
echo.echo_warning(f'MultiplyAddWorkChain<{node.pk}> did not finish successfully.')


def prepare_localhost():
"""Prepare and return the localhost as ``Computer``.
Expand Down
37 changes: 36 additions & 1 deletion tests/cmdline/commands/test_devel.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@

import pytest
from aiida.cmdline.commands import cmd_devel
from aiida.orm import Node, ProcessNode, QueryBuilder
from aiida.orm import Node, ProcessNode, QueryBuilder, WorkChainNode


@pytest.mark.requires_psql
Expand Down Expand Up @@ -56,3 +56,38 @@ def test_launch_add_code(run_cli_command, aiida_code_installed):

node = QueryBuilder().append(ProcessNode, tag='node').order_by({'node': {'ctime': 'desc'}}).first(flat=True)
assert node.is_finished_ok


def test_launch_multiply_add(run_cli_command):
"""Test ``verdi devel launch-multiply-add``.
Start with a clean profile such that the functionality that automatically sets up the localhost computer and code is
tested properly.
"""
result = run_cli_command(cmd_devel.devel_launch_multiply_add)
assert re.search(r'Warning: No `localhost` computer exists yet: creating and configuring', result.stdout)
assert re.search(r'Success: MultiplyAddWorkChain<.*> finished successfully', result.stdout)

node = QueryBuilder().append(WorkChainNode, tag='node').order_by({'node': {'ctime': 'desc'}}).first(flat=True)
assert node.is_finished_ok


@pytest.mark.usefixtures('started_daemon_client')
def test_launch_multiply_add_daemon(run_cli_command, submit_and_await):
"""Test ``verdi devel launch-multiply-add`` with the ``--daemon`` flag."""
result = run_cli_command(cmd_devel.devel_launch_multiply_add, ['--daemon'])
assert re.search(r'Submitted workflow', result.stdout)

node = QueryBuilder().append(ProcessNode, tag='node').order_by({'node': {'ctime': 'desc'}}).first(flat=True)
submit_and_await(node)
assert node.is_finished_ok


def test_launch_add_multiply_code(run_cli_command, aiida_code_installed):
"""Test ``verdi devel launch-multiply-add`` passing an explicit ``Code``."""
code = aiida_code_installed(default_calc_job_plugin='core.arithmetic.add', filepath_executable='/bin/bash')
result = run_cli_command(cmd_devel.devel_launch_multiply_add, ['-X', str(code.pk)])
assert not re.search(r'Warning: No `localhost` computer exists yet: creating and configuring', result.stdout)

node = QueryBuilder().append(ProcessNode, tag='node').order_by({'node': {'ctime': 'desc'}}).first(flat=True)
assert node.is_finished_ok

0 comments on commit b7c8286

Please sign in to comment.