Merge wait and timeout to one CLI option for verdi process {kill|play|pause} (aiidateam#6902)

The original design of `wait` and `timeout` was to distinguish between actions
that immediately return and actions that scheduled. This mechanism was however
never used and resulted in an misinterpretation in the force-kill PR aiidateam#6793
introducing a bug fixed in PR $6870. The mechanism of `wait` and `timeout` was
also never correctly implemented. In this PR we rectify the logic and simplify
it by handling immediate and scheduled actions the same way.

My interpretation is that the original logic (seen in commit 8388018) is that it
unwraps the future once, if it is a another future then it is a scheduled action
otherwise an immediate one. Then 1b6ecb8 in plumpy introduced a double wrapping
for `play`, `pause` and `kill` of the return value which I think is correct as
these are scheduled actions. The problem appears in cd0d15c where on the
aiida-core side in the first round unwrapping is always done till the actual
(nonfuture) result by using `unwrap_kiwi_future`. This makes the usage of
`wait` completely obsolete as we always get the final nonfuture result in the
first step. Also the `timeout` was not passed to the `result` which made it a
blocking action for scheduled tasked (the `timeout` is only applied on the
first layer of future in `futures.as_completed` which is meaningless since the
first layer is always a future and returns immediately).

This is fixed by simplifying the two step procedure, unwrap once and if future
unwrap again, to one step: Unwrap until nonfuture result gained with timeout.
One can specify a `timeout == 0` to express that the action should not wait for
a response while one can specify `timeout==float('inf')` (default value) to wait
until a response has been received without a timeout.

Author: agoscinski

src/aiida/cmdline/commands/cmd_process.py

@@ -25,6 +25,16 @@
     verdi daemon start
 """
 
+ACTION_TIMEOUT = OverridableOption(
+    '-t',
+    '--timeout',
+    type=click.FloatRange(0, float('inf')),
+    default=float('inf'),
+    show_default=True,
+    help='Time in seconds to wait for a response before timing out. '
+    'If timeout is 0 the command does not wait for response.',
+)
+
 
 def valid_projections():
     """Return list of valid projections for the ``--project`` option of ``verdi process list``.
@@ -320,15 +330,7 @@ def process_status(call_link_label, most_recent_node, max_depth, processes):
 @verdi_process.command('kill')
 @arguments.PROCESSES()
 @options.ALL(help='Kill all processes if no specific processes are specified.')
-@OverridableOption(
-    '-t',
-    '--timeout',
-    type=click.FLOAT,
-    default=5.0,
-    show_default=True,
-    help='Time in seconds to wait for a response of the kill task before timing out.',
-)()
-@options.WAIT()
+@ACTION_TIMEOUT()
 @OverridableOption(
     '-F',
     '--force',
@@ -338,7 +340,7 @@ def process_status(call_link_label, most_recent_node, max_depth, processes):
     'Note: This may lead to orphaned jobs on your HPC and should be used with caution.',
 )()
 @decorators.with_dbenv()
-def process_kill(processes, all_entries, timeout, wait, force):
+def process_kill(processes, all_entries, timeout, force):
     """Kill running processes.
 
     Kill one or multiple running processes."""
@@ -368,7 +370,6 @@ def process_kill(processes, all_entries, timeout, wait, force):
                 force=force,
                 all_entries=all_entries,
                 timeout=timeout,
-                wait=wait,
             )
         except control.ProcessTimeoutException as exception:
             echo.echo_critical(f'{exception}\n{REPAIR_INSTRUCTIONS}')
@@ -380,10 +381,9 @@ def process_kill(processes, all_entries, timeout, wait, force):
 @verdi_process.command('pause')
 @arguments.PROCESSES()
 @options.ALL(help='Pause all active processes if no specific processes are specified.')
-@options.TIMEOUT()
-@options.WAIT()
+@ACTION_TIMEOUT()
 @decorators.with_dbenv()
-def process_pause(processes, all_entries, timeout, wait):
+def process_pause(processes, all_entries, timeout):
     """Pause running processes.
 
     Pause one or multiple running processes."""
@@ -404,7 +404,6 @@ def process_pause(processes, all_entries, timeout, wait):
                 msg_text='Paused through `verdi process pause`',
                 all_entries=all_entries,
                 timeout=timeout,
-                wait=wait,
             )
         except control.ProcessTimeoutException as exception:
             echo.echo_critical(f'{exception}\n{REPAIR_INSTRUCTIONS}')
@@ -416,10 +415,9 @@ def process_pause(processes, all_entries, timeout, wait):
 @verdi_process.command('play')
 @arguments.PROCESSES()
 @options.ALL(help='Play all paused processes if no specific processes are specified.')
-@options.TIMEOUT()
-@options.WAIT()
+@ACTION_TIMEOUT()
 @decorators.with_dbenv()
-def process_play(processes, all_entries, timeout, wait):
+def process_play(processes, all_entries, timeout):
     """Play (unpause) paused processes.
 
     Play (unpause) one or multiple paused processes."""
@@ -435,7 +433,7 @@ def process_play(processes, all_entries, timeout, wait):
 
     with capture_logging() as stream:
         try:
-            control.play_processes(processes, all_entries=all_entries, timeout=timeout, wait=wait)
+            control.play_processes(processes, all_entries=all_entries, timeout=timeout)
         except control.ProcessTimeoutException as exception:
             echo.echo_critical(f'{exception}\n{REPAIR_INSTRUCTIONS}')
 

src/aiida/cmdline/params/options/main.py

@@ -125,7 +125,6 @@
     'USER_LAST_NAME',
     'VERBOSITY',
     'VISUALIZATION_FORMAT',
-    'WAIT',
     'WITH_ELEMENTS',
     'WITH_ELEMENTS_EXCLUSIVE',
     'active_process_states',
@@ -690,12 +689,6 @@ def set_log_level(ctx, _param, value):
     help='Time in seconds to wait for a response before timing out.',
 )
 
-WAIT = OverridableOption(
-    '--wait/--no-wait',
-    default=False,
-    help='Wait for the action to be completed otherwise return as soon as it is scheduled.',
-)
-
 FORMULA_MODE = OverridableOption(
     '-f',
     '--formula-mode',

src/aiida/engine/processes/control.py

@@ -104,7 +104,7 @@ def revive_processes(processes: list[ProcessNode], *, wait: bool = False) -> Non
 
 
 def play_processes(
-    processes: list[ProcessNode] | None = None, *, all_entries: bool = False, timeout: float = 5.0, wait: bool = False
+    processes: list[ProcessNode] | None = None, *, all_entries: bool = False, timeout: float = 5.0
 ) -> None:
     """Play (unpause) paused processes.
 
@@ -113,7 +113,6 @@ def play_processes(
     :param processes: List of processes to play.
     :param all_entries: Play all paused processes.
     :param timeout: Raise a ``ProcessTimeoutException`` if the process does not respond within this amount of seconds.
-    :param wait: Set to ``True`` to wait for process response, for ``False`` the action is fire-and-forget.
     :raises ``ProcessTimeoutException``: If the processes do not respond within the timeout.
     """
     if not get_daemon_client().is_daemon_running:
@@ -130,7 +129,7 @@ def play_processes(
         return
 
     controller = get_manager().get_process_controller()
-    _perform_actions(processes, controller.play_process, 'play', 'playing', timeout, wait)
+    _perform_actions(processes, controller.play_process, 'play', 'playing', timeout)
 
 
 def pause_processes(
@@ -139,7 +138,6 @@ def pause_processes(
     msg_text: str = 'Paused through `aiida.engine.processes.control.pause_processes`',
     all_entries: bool = False,
     timeout: float = 5.0,
-    wait: bool = False,
 ) -> None:
     """Pause running processes.
 
@@ -148,7 +146,6 @@ def pause_processes(
     :param processes: List of processes to play.
     :param all_entries: Pause all playing processes.
     :param timeout: Raise a ``ProcessTimeoutException`` if the process does not respond within this amount of seconds.
-    :param wait: Set to ``True`` to wait for process response, for ``False`` the action is fire-and-forget.
     :raises ``ProcessTimeoutException``: If the processes do not respond within the timeout.
     """
     if not get_daemon_client().is_daemon_running:
@@ -166,7 +163,7 @@ def pause_processes(
 
     controller = get_manager().get_process_controller()
     action = functools.partial(controller.pause_process, msg_text=msg_text)
-    _perform_actions(processes, action, 'pause', 'pausing', timeout, wait)
+    _perform_actions(processes, action, 'pause', 'pausing', timeout)
 
 
 def kill_processes(
@@ -176,7 +173,6 @@ def kill_processes(
     force: bool = False,
     all_entries: bool = False,
     timeout: float = 5.0,
-    wait: bool = False,
 ) -> None:
     """Kill running processes.
 
@@ -185,7 +181,6 @@ def kill_processes(
     :param processes: List of processes to play.
     :param all_entries: Kill all active processes.
     :param timeout: Raise a ``ProcessTimeoutException`` if the process does not respond within this amount of seconds.
-    :param wait: Set to ``True`` to wait for process response, for ``False`` the action is fire-and-forget.
     :raises ``ProcessTimeoutException``: If the processes do not respond within the timeout.
     """
     if not get_daemon_client().is_daemon_running:
@@ -203,7 +198,7 @@ def kill_processes(
 
     controller = get_manager().get_process_controller()
     action = functools.partial(controller.kill_process, msg_text=msg_text, force_kill=force)
-    _perform_actions(processes, action, 'kill', 'killing', timeout, wait)
+    _perform_actions(processes, action, 'kill', 'killing', timeout)
 
 
 def _perform_actions(
@@ -212,7 +207,6 @@ def _perform_actions(
     infinitive: str,
     present: str,
     timeout: t.Optional[float] = None,
-    wait: bool = False,
     **kwargs: t.Any,
 ) -> None:
     """Perform an action on a list of processes.
@@ -223,7 +217,6 @@ def _perform_actions(
     :param present: The present tense of the verb that represents the action.
     :param past: The past tense of the verb that represents the action.
     :param timeout: Raise a ``ProcessTimeoutException`` if the process does not respond within this amount of seconds.
-    :param wait: Set to ``True`` to wait for process response, for ``False`` the action is fire-and-forget.
     :param kwargs: Keyword arguments that will be passed to the method ``action``.
     :raises ``ProcessTimeoutException``: If the processes do not respond within the timeout.
     """
@@ -236,86 +229,71 @@ def _perform_actions(
 
         try:
             future = action(process.pk, **kwargs)
+            LOGGER.report(f'Request to {infinitive} Process<{process.pk}> sent.')
         except communications.UnroutableError:
             LOGGER.error(f'Process<{process.pk}> is unreachable.')
         else:
             futures[future] = process
 
-    _resolve_futures(futures, infinitive, present, wait, timeout)
+    _resolve_futures(futures, infinitive, present, timeout)
 
 
 def _resolve_futures(
     futures: dict[concurrent.futures.Future, ProcessNode],
     infinitive: str,
     present: str,
-    wait: bool = False,
     timeout: t.Optional[float] = None,
 ) -> None:
     """Process a mapping of futures representing an action on an active process.
 
     This function will echo the correct information strings based on the outcomes of the futures and the given verb
-    conjugations. You can optionally wait for any pending actions to be completed before the functions returns and use a
-    timeout to put a maximum wait time on the actions.
+    conjugations. The function waits for any pending actions to be completed. By specifying a timeout the function
+    aborts after the specified time and cancels pending actions.
 
     :param futures: The map of action futures and the corresponding processes.
     :param infinitive: The infinitive form of the action verb.
     :param present: The present tense form of the action verb.
-    :param wait: Set to ``True`` to wait for process response, for ``False`` the action is fire-and-forget.
-    :param timeout: Raise a ``ProcessTimeoutException`` if the process does not respond within this amount of seconds.
+    :param timeout: If None or float('inf') it waits until the actions are completed otherwise it waits for response the
+        amount in seconds.
     """
-    scheduled = {}
-
-    def handle_result(result):
-        if result is True:
-            LOGGER.report(f'request to {infinitive} Process<{process.pk}> sent')
-        elif result is False:
-            LOGGER.error(f'problem {present} Process<{process.pk}>')
-        elif isinstance(result, kiwipy.Future):
-            LOGGER.report(f'scheduled {infinitive} Process<{process.pk}>')
-            scheduled[result] = process
-        else:
-            LOGGER.error(f'got unexpected response when {present} Process<{process.pk}>: {result}')
-
-    try:
-        for future, process in futures.items():
-            # unwrap is need here since LoopCommunicator will also wrap a future
-            unwrapped = unwrap_kiwi_future(future)
-            try:
-                result = unwrapped.result(timeout=timeout)
-            except communications.TimeoutError:
-                cancelled = unwrapped.cancel()
-                if cancelled:
-                    LOGGER.error(f'call to {infinitive} Process<{process.pk}> timed out and was cancelled.')
-                else:
-                    LOGGER.error(f'call to {infinitive} Process<{process.pk}> timed out but could not be cancelled.')
-            except Exception as exception:
-                LOGGER.error(f'failed to {infinitive} Process<{process.pk}>: {exception}')
-            else:
-                if isinstance(result, kiwipy.Future):
-                    LOGGER.report(f'scheduled {infinitive} Process<{process.pk}>')
-                    scheduled[result] = process
-                else:
-                    handle_result(result)
-
-        if not wait or not scheduled:
-            return
-
-        LOGGER.report(f"waiting for process(es) {','.join([str(proc.pk) for proc in scheduled.values()])}")
+    if not timeout or not futures:
+        if futures:
+            LOGGER.report(
+                f"Request to {infinitive} process(es) {','.join([str(proc.pk) for proc in futures.values()])}"
+                ' sent. Skipping waiting for response.'
+            )
+        return
 
-        for future in concurrent.futures.as_completed(scheduled.keys(), timeout=timeout):
-            process = scheduled[future]
+    LOGGER.report(f"Waiting for process(es) {','.join([str(proc.pk) for proc in futures.values()])}.")
 
+    # Ensure that when futures are only are completed if they return an actual value (not a future)
+    unwrapped_futures = {unwrap_kiwi_future(future): process for future, process in futures.items()}
+    try:
+        # future does not interpret float('inf') correctly by changing it to None we get the intended behavior
+        for future in concurrent.futures.as_completed(
+            unwrapped_futures.keys(), timeout=None if timeout == float('inf') else timeout
+        ):
+            process = unwrapped_futures[future]
             try:
                 result = future.result()
             except Exception as exception:
-                LOGGER.error(f'failed to {infinitive} Process<{process.pk}>: {exception}')
+                LOGGER.error(f'Failed to {infinitive} Process<{process.pk}>: {exception}')
             else:
-                handle_result(result)
-
+                if result is True:
+                    LOGGER.report(f'Request to {infinitive} Process<{process.pk}> processed.')
+                elif result is False:
+                    LOGGER.error(f'Problem {present} Process<{process.pk}>')
+                else:
+                    LOGGER.error(f'Got unexpected response when {present} Process<{process.pk}>: {result}')
     except concurrent.futures.TimeoutError:
-        raise ProcessTimeoutException(
-            f'timed out trying to {infinitive} processes {futures.values()}\n'
-            'This could be because the daemon workers are too busy to respond, please try again later.\n'
-            'If the problem persists, make sure the daemon and RabbitMQ are running properly by restarting them.\n'
-            'If the processes remain unresponsive, as a last resort, try reviving them with ``revive_processes``.'
-        )
+        # We cancel the tasks that are not done
+        undone_futures = {future: process for future, process in unwrapped_futures.items() if not future.done()}
+        if not undone_futures:
+            LOGGER.error(f'Call to {infinitive} timed out but already done.')
+        for future, process in undone_futures.items():
+            if not future.done():
+                cancelled = future.cancel()
+                if cancelled:
+                    LOGGER.error(f'Call to {infinitive} Process<{process.pk}> timed out and was cancelled.')
+                else:
+                    LOGGER.error(f'Call to {infinitive} Process<{process.pk}> timed out but could not be cancelled.')