Merging usage of wait and timeout for actions (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 aiidateam#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.

Related commits in aiida-core 8388018, cd0d15c and plumpy 1b6ecb8

One can specifiy 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.FLOAT,
+    default=float('inf'),
+    show_default=True,
+    help='Time in seconds to wait for a response before timing out. '
+    'If timeout <= 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.
     """
@@ -241,77 +234,53 @@ def _perform_actions(
         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.
+    timeout to put a maximum wait time on the actions. TODO fix docstring
 
     :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.
     """
-    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}')
+    if not timeout:
+        return
+
+    LOGGER.report(f"waiting for process(es) {','.join([str(proc.pk) for proc in futures.values()])}")
 
     try:
         for future, process in futures.items():
-            # unwrap is need here since LoopCommunicator will also wrap a future
+            # we unwrap to the end
             unwrapped = unwrap_kiwi_future(future)
             try:
-                result = unwrapped.result(timeout=timeout)
+                # future does not interpret float('inf') correctly by changing it to None we get the intended behavior
+                result = unwrapped.result(timeout=None if timeout == float('inf') else timeout)
             except communications.TimeoutError:
-                cancelled = unwrapped.cancel()
+                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.')
             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
+                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}>')
                 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()])}")
-
-        for future in concurrent.futures.as_completed(scheduled.keys(), timeout=timeout):
-            process = scheduled[future]
-
-            try:
-                result = future.result()
-            except Exception as exception:
-                LOGGER.error(f'failed to {infinitive} Process<{process.pk}>: {exception}')
-            else:
-                handle_result(result)
-
+                    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'