Fix bug in process function interrupt handlers causing verdi daemon stop to timeout

aiida/engine/daemon/runner.py

@@ -30,7 +30,9 @@ def start_daemon():
     configure_logging(daemon=True, daemon_log_file=daemon_client.daemon_log_file)
 
     try:
-        runner = get_manager().create_daemon_runner()
+        manager = get_manager()
+        runner = manager.create_daemon_runner()
+        manager.set_runner(runner)
     except Exception as exception:
         LOGGER.exception('daemon runner failed to start')
         raise

aiida/engine/processes/functions.py

@@ -119,15 +119,16 @@ def run_get_node(*args, **kwargs):
             Run the FunctionProcess with the supplied inputs in a local runner.
 
             The function will have to create a new runner for the FunctionProcess instead of using the global runner,
-            because otherwise if this workfunction were to call another one from within its scope, that would use
+            because otherwise if this process function were to call another one from within its scope, that would use
             the same runner and it would be blocking the event loop from continuing.
 
             :param args: input arguments to construct the FunctionProcess
             :param kwargs: input keyword arguments to construct the FunctionProcess
             :return: tuple of the outputs of the process and the process node pk
             :rtype: (dict, int)
             """
-            runner = get_manager().create_runner(with_persistence=False)
+            manager = get_manager()
+            runner = manager.create_runner(with_persistence=False)
             inputs = process_class.create_inputs(*args, **kwargs)
 
             # Remove all the known inputs from the kwargs
@@ -140,18 +141,32 @@ def run_get_node(*args, **kwargs):
 
             process = process_class(inputs=inputs, runner=runner)
 
-            def kill_process(_num, _frame):
-                """Send the kill signal to the process in the current scope."""
-                LOGGER.critical('runner received interrupt, killing process %s', process.pid)
-                process.kill(msg='Process was killed because the runner received an interrupt')
+            # Only add handlers for interrupt signal to kill the process if we are in a local and not a daemon runner.
+            # Without this check, running process functions in a daemon worker would be killed if the daemon is shutdown
+            current_runner = manager.get_runner()
+            original_handler = None
+            kill_signal = signal.SIGINT
 
-            signal.signal(signal.SIGINT, kill_process)
-            signal.signal(signal.SIGTERM, kill_process)
+            if not current_runner.is_daemon_runner:
 
-            result = process.execute()
+                def kill_process(_num, _frame):
+                    """Send the kill signal to the process in the current scope."""
+                    from tornado import gen
+                    LOGGER.critical('runner received interrupt, killing process %s', process.pid)
+                    result = process.kill(msg='Process was killed because the runner received an interrupt')
+                    raise gen.Return(result)
 
-            # Close the runner properly
-            runner.close()
+                # Store the current handler on the signal such that it can be restored after process has terminated
+                original_handler = signal.getsignal(kill_signal)
+                signal.signal(kill_signal, kill_process)
+
+            try:
+                result = process.execute()
+            finally:
+                # If the `original_handler` is set, that means the `kill_process` was bound, which needs to be reset
+                if original_handler:
+                    signal.signal(signal.SIGINT, original_handler)
+                runner.close()
 
             store_provenance = inputs.get('metadata', {}).get('store_provenance', True)
             if not store_provenance:

aiida/engine/processes/process.py

@@ -293,7 +293,7 @@ def kill(self, msg=None):
 
             if killing:
                 # We are waiting for things to be killed, so return the 'gathered' future
-                result = plumpy.gather(result)
+                result = plumpy.gather(killing)
 
         return result
 

aiida/engine/runners.py

@@ -116,6 +116,15 @@ def job_manager(self):
     def controller(self):
         return self._controller
 
+    @property
+    def is_daemon_runner(self):
+        """Return whether the runner is a daemon runner, which means it submits processes over RabbitMQ.
+
+        :return: True if the runner is a daemon runner
+        :rtype: bool
+        """
+        return self._rmq_submit
+
     def is_closed(self):
         return self._closed