Allow to specify number of workers in verdi daemon start

Now you can start more than just one worker when starting the daemon.

aiida/backends/tests/__init__.py

@@ -44,6 +44,7 @@
         'cmdline.commands.comment': ['aiida.backends.tests.cmdline.commands.test_comment'],
         'cmdline.commands.computer': ['aiida.backends.tests.cmdline.commands.test_computer'],
         'cmdline.commands.config': ['aiida.backends.tests.cmdline.commands.test_config'],
+        'cmdline.commands.daemon': ['aiida.backends.tests.cmdline.commands.test_daemon'],
         'cmdline.commands.data': ['aiida.backends.tests.cmdline.commands.test_data'],
         'cmdline.commands.database': ['aiida.backends.tests.cmdline.commands.test_database'],
         'cmdline.commands.export': ['aiida.backends.tests.cmdline.commands.test_export'],

aiida/backends/tests/cmdline/commands/test_daemon.py

@@ -0,0 +1,64 @@
+# -*- coding: utf-8 -*-
+###########################################################################
+# 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               #
+###########################################################################
+"""Tests for `verdi daemon`."""
+from __future__ import division
+from __future__ import print_function
+from __future__ import absolute_import
+
+from click.testing import CliRunner
+
+from aiida.backends.testbase import AiidaTestCase
+from aiida.cmdline.commands import cmd_daemon
+from aiida.engine.daemon.client import DaemonClient
+from aiida.manage.configuration import get_config
+
+
+class TestVerdiDaemon(AiidaTestCase):
+    """Tests for `verdi daemon` commands."""
+
+    def setUp(self):
+        super(TestVerdiDaemon, self).setUp()
+        self.daemon_client = DaemonClient(get_config().current_profile)
+        self.cli_runner = CliRunner()
+
+    def test_daemon_start(self):
+        """Test `verdi daemon start`."""
+        try:
+            result = self.cli_runner.invoke(cmd_daemon.start, [])
+            self.assertClickResultNoException(result)
+
+            daemon_response = self.daemon_client.get_daemon_info()
+            worker_response = self.daemon_client.get_worker_info()
+
+            self.assertIn('status', daemon_response, daemon_response)
+            self.assertEqual(daemon_response['status'], 'ok', daemon_response)
+
+            self.assertIn('info', worker_response, worker_response)
+            self.assertEqual(len(worker_response['info']), 1, worker_response)
+        finally:
+            self.daemon_client.stop_daemon(wait=True)
+
+    def test_daemon_start_number(self):
+        """Test `verdi daemon start` with a specific number of workers."""
+        try:
+            number = 4
+            result = self.cli_runner.invoke(cmd_daemon.start, [str(number)])
+            self.assertClickResultNoException(result)
+
+            daemon_response = self.daemon_client.get_daemon_info()
+            worker_response = self.daemon_client.get_worker_info()
+
+            self.assertIn('status', daemon_response, daemon_response)
+            self.assertEqual(daemon_response['status'], 'ok', daemon_response)
+
+            self.assertIn('info', worker_response, worker_response)
+            self.assertEqual(len(worker_response['info']), number, worker_response)
+        finally:
+            self.daemon_client.stop_daemon(wait=True)

aiida/cmdline/commands/cmd_daemon.py

@@ -27,19 +27,29 @@
 from aiida.manage.configuration import get_config
 
 
+def validate_positive_non_zero_integer(ctx, param, value):  # pylint: disable=unused-argument,invalid-name
+    """Validate that `value` is a positive integer."""
+    if not isinstance(value, int):
+        raise click.BadParameter('{} is not an integer'.format(value))
+
+    if value <= 0:
+        raise click.BadParameter('{} is not a positive non-zero integer'.format(value))
+
+    return value
+
+
 @verdi.group('daemon')
 def verdi_daemon():
     """Inspect and manage the daemon."""
 
 
 @verdi_daemon.command()
 @click.option('--foreground', is_flag=True, help='Run in foreground.')
+@click.argument('number', default=1, type=int, callback=validate_positive_non_zero_integer)
 @decorators.with_dbenv()
 @decorators.check_circus_zmq_version
-def start(foreground):
-    """
-    Start the daemon
-    """
+def start(foreground, number):
+    """Start the daemon with NUMBER workers [default=1]."""
     from aiida.engine.daemon.client import get_daemon_client
 
     client = get_daemon_client()
@@ -64,13 +74,17 @@ def start(foreground):
 
     print_client_response_status(response)
 
+    if number > 1:
+        echo.echo('Starting additional workers... ', nl=False)
+        additional_workers = number - 1
+        response = client.increase_workers(additional_workers)
+        print_client_response_status(response)
+
 
 @verdi_daemon.command()
 @click.option('--all', 'all_profiles', is_flag=True, help='Show all daemons.')
 def status(all_profiles):
-    """
-    Print the status of the current daemon or all daemons
-    """
+    """Print the status of the current daemon or all daemons."""
     from aiida.engine.daemon.client import get_daemon_client
 
     config = get_config()
@@ -92,9 +106,7 @@ def status(all_profiles):
 @click.argument('number', default=1, type=int)
 @decorators.only_if_daemon_running()
 def incr(number):
-    """
-    Add NUMBER [default=1] workers to the running daemon
-    """
+    """Add NUMBER [default=1] workers to the running daemon."""
     from aiida.engine.daemon.client import get_daemon_client
 
     client = get_daemon_client()
@@ -106,9 +118,7 @@ def incr(number):
 @click.argument('number', default=1, type=int)
 @decorators.only_if_daemon_running()
 def decr(number):
-    """
-    Remove NUMBER [default=1] workers from the running daemon
-    """
+    """Remove NUMBER [default=1] workers from the running daemon."""
     from aiida.engine.daemon.client import get_daemon_client
 
     client = get_daemon_client()
@@ -118,9 +128,7 @@ def decr(number):
 
 @verdi_daemon.command()
 def logshow():
-    """
-    Show the log of the daemon, press CTRL+C to quit
-    """
+    """Show the log of the daemon, press CTRL+C to quit."""
     from aiida.engine.daemon.client import get_daemon_client
 
     client = get_daemon_client()
@@ -137,9 +145,7 @@ def logshow():
 @click.option('--no-wait', is_flag=True, help='Do not wait for confirmation.')
 @click.option('--all', 'all_profiles', is_flag=True, help='Stop all daemons.')
 def stop(no_wait, all_profiles):
-    """
-    Stop the daemon
-    """
+    """Stop the daemon."""
     from aiida.engine.daemon.client import get_daemon_client
 
     config = get_config()
@@ -185,11 +191,11 @@ def stop(no_wait, all_profiles):
 @decorators.with_dbenv()
 @decorators.only_if_daemon_running()
 def restart(ctx, reset, no_wait):
-    """
-    Restart the daemon. By default will only reset the workers of the running daemon.
-    After the restart the same amount of workers will be running. If the --reset flag
-    is passed, however, the full circus daemon will be stopped and restarted with just
-    a single worker
+    """Restart the daemon.
+
+    By default will only reset the workers of the running daemon. After the restart the same amount of workers will be
+    running. If the `--reset` flag is passed, however, the full circus daemon will be stopped and restarted with just
+    a single worker.
     """
     from aiida.engine.daemon.client import get_daemon_client
 
@@ -218,9 +224,9 @@ def restart(ctx, reset, no_wait):
 @decorators.with_dbenv()
 @decorators.check_circus_zmq_version
 def start_circus(foreground):
-    """
-    This will actually launch the circus daemon, either daemonized in the background
-    or in the foreground, printing all logs to stdout.
+    """This will actually launch the circus daemon, either daemonized in the background or in the foreground.
+
+    If run in the foreground all logs are redirected to stdout.
 
     .. note:: this should not be called directly from the commandline!
     """

docs/source/verdi/verdi_user_guide.rst

@@ -334,13 +334,13 @@ Below is a list with all available subcommands.
       --help  Show this message and exit.
 
     Commands:
-      decr     Remove NUMBER [default=1] workers from the running daemon
-      incr     Add NUMBER [default=1] workers to the running daemon
-      logshow  Show the log of the daemon, press CTRL+C to quit
+      decr     Remove NUMBER [default=1] workers from the running daemon.
+      incr     Add NUMBER [default=1] workers to the running daemon.
+      logshow  Show the log of the daemon, press CTRL+C to quit.
       restart  Restart the daemon.
-      start    Start the daemon
-      status   Print the status of the current daemon or all daemons
-      stop     Stop the daemon
+      start    Start the daemon with NUMBER workers [default=1].
+      status   Print the status of the current daemon or all daemons.
+      stop     Stop the daemon.
 
 
 .. _verdi_data: