WIP: Implement option to allow force kill

Do not peform stepping when force_kill command is send.
This avoids getting stuck in the kill callback.

Author: agoscinski

src/plumpy/process_comms.py

@@ -30,6 +30,7 @@
 
 INTENT_KEY = 'intent'
 MESSAGE_KEY = 'message'
+FORCE_KILL_KEY = 'force_kill'
 
 
 class Intent:
@@ -196,17 +197,19 @@ async def play_process(self, pid: 'PID_TYPE') -> 'ProcessResult':
         result = await asyncio.wrap_future(future)
         return result
 
-    async def kill_process(self, pid: 'PID_TYPE', msg: Optional[Any] = None) -> 'ProcessResult':
+    async def kill_process(self, pid: 'PID_TYPE', msg: Optional[Any] = None, force_kill: bool = False) -> 'ProcessResult':
         """
         Kill the process
 
         :param pid: the pid of the process to kill
         :param msg: optional kill message
         :return: True if killed, False otherwise
         """
+        breakpoint()
         message = copy.copy(KILL_MSG)
         if msg is not None:
             message[MESSAGE_KEY] = msg
+        message[FORCE_KILL_KEY] = force_kill
 
         # Wait for the communication to go through
         kill_future = self._communicator.rpc_send(pid, message)
@@ -375,7 +378,7 @@ def play_all(self) -> None:
         """
         self._communicator.broadcast_send(None, subject=Intent.PLAY)
 
-    def kill_process(self, pid: 'PID_TYPE', msg: Optional[Any] = None) -> kiwipy.Future:
+    def kill_process(self, pid: 'PID_TYPE', msg: Optional[Any] = None, force_kill: bool = False) -> kiwipy.Future:
         """
         Kill the process
 
@@ -384,9 +387,11 @@ def kill_process(self, pid: 'PID_TYPE', msg: Optional[Any] = None) -> kiwipy.Fut
         :return: a response future from the process to be killed
 
         """
+        breakpoint()
         message = copy.copy(KILL_MSG)
         if msg is not None:
             message[MESSAGE_KEY] = msg
+        message[FORCE_KILL_KEY] = force_kill
 
         return self._communicator.rpc_send(pid, message)
 
@@ -405,6 +410,7 @@ def continue_process(
         nowait: bool = False,
         no_reply: bool = False
     ) -> Union[None, PID_TYPE, ProcessResult]:
+        breakpoint()
         message = create_continue_body(pid=pid, tag=tag, nowait=nowait)
         return self.task_send(message, no_reply=no_reply)
 
@@ -479,6 +485,7 @@ def task_send(self, message: Any, no_reply: bool = False) -> Optional[Any]:
         :param no_reply: if True, this call will be fire-and-forget, i.e. no return value
         :return: the response from the remote side (if no_reply=False)
         """
+        breakpoint()
         return self._communicator.task_send(message, no_reply=no_reply)
 
 

src/plumpy/processes.py

@@ -51,7 +51,29 @@
 
 __all__ = ['Process', 'ProcessSpec', 'BundleKeys', 'TransitionFailed']
 
+
+#file_handler = logging.FileHandler(filename='tmp.log')
+#stdout_handler = logging.StreamHandler(stream=sys.stdout)
+#handlers = [file_handler, stdout_handler]
+#
+#logging.basicConfig(
+#    level=logging.DEBUG, 
+#    format='[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s - %(message)s',
+#    handlers=handlers
+#)
+
+#file_handler = logging.FileHandler(filename="/Users/alexgo/code/aiida-core/plumpy2.log")
+#stdout_handler = logging.StreamHandler(stream=sys.stdout)
+#handlers = [file_handler, stdout_handler]
+#
+#logging.basicConfig(
+#    level=logging.DEBUG, 
+#    format='[%(asctime)s] {%(pathname)s:%(lineno)d} %(levelname)s - %(message)s',
+#    handlers=handlers
+#)
 _LOGGER = logging.getLogger(__name__)
+
+
 PROCESS_STACK = ContextVar('process stack', default=[])
 
 
@@ -389,8 +411,8 @@ def logger(self) -> logging.Logger:
         :return: The logger.
 
         """
-        if self._logger is not None:
-            return self._logger
+        #if self._logger is not None:
+        #    return self._logger
 
         return _LOGGER
 
@@ -908,6 +930,7 @@ def message_receive(self, _comm: kiwipy.Communicator, msg: Dict[str, Any]) -> An
         :param msg: the message
         :return: the outcome of processing the message, the return value will be sent back as a response to the sender
         """
+        breakpoint()
         self.logger.debug("Process<%s>: received RPC message with communicator '%s': %r", self.pid, _comm, msg)
 
         intent = msg[process_comms.INTENT_KEY]
@@ -917,7 +940,11 @@ def message_receive(self, _comm: kiwipy.Communicator, msg: Dict[str, Any]) -> An
         if intent == process_comms.Intent.PAUSE:
             return self._schedule_rpc(self.pause, msg=msg.get(process_comms.MESSAGE_KEY, None))
         if intent == process_comms.Intent.KILL:
-            return self._schedule_rpc(self.kill, msg=msg.get(process_comms.MESSAGE_KEY, None))
+            breakpoint()
+            # have problems to pass new argument get
+            # Error: failed to kill Process<699>: Process.kill() got an unexpected keyword argument 'force_kill'
+            #return self._schedule_rpc(self.kill, msg=msg.get(process_comms.MESSAGE_KEY, None), force_kill=msg.get(process_comms.FORCE_KILL_KEY, False))
+            return self._schedule_rpc(self.kill, msg=msg)
         if intent == process_comms.Intent.STATUS:
             status_info: Dict[str, Any] = {}
             self.get_status_info(status_info)
@@ -934,6 +961,7 @@ def broadcast_receive(self, _comm: kiwipy.Communicator, body: Any, sender: Any,
         :param _comm: the communicator that sent the message
         :param msg: the message
         """
+        breakpoint()
         # pylint: disable=unused-argument
         self.logger.debug(
             "Process<%s>: received broadcast message '%s' with communicator '%s': %r", self.pid, subject, _comm, body
@@ -945,6 +973,7 @@ def broadcast_receive(self, _comm: kiwipy.Communicator, body: Any, sender: Any,
         if subject == process_comms.Intent.PAUSE:
             return self._schedule_rpc(self.pause, msg=body)
         if subject == process_comms.Intent.KILL:
+            # TODO deal with this
             return self._schedule_rpc(self.kill, msg=body)
         return None
 
@@ -1126,11 +1155,21 @@ def fail(self, exception: Optional[BaseException], trace_back: Optional[Tracebac
         """
         self.transition_to(process_states.ProcessState.EXCEPTED, exception, trace_back)
 
-    def kill(self, msg: Union[str, None] = None) -> Union[bool, asyncio.Future]:
+    def kill(self, msg: Union[dict, None] = None, force_kill: bool = False) -> Union[bool, asyncio.Future]:
         """
         Kill the process
+
+        # PR_COMMENT have not figured out how to integrate force_kill as argument
+        #            so I just pass the dict
+
         :param msg: An optional kill message
         """
+        breakpoint()
+        if msg is None:
+            force_kill = False 
+        else:
+            force_kill = msg.get(process_comms.FORCE_KILL_KEY, False)
+
         if self.state == process_states.ProcessState.KILLED:
             # Already killed
             return True
@@ -1143,15 +1182,16 @@ def kill(self, msg: Union[str, None] = None) -> Union[bool, asyncio.Future]:
             # Already killing
             return self._killing
 
-        if self._stepping:
+        if self._stepping and not force_kill:
             # Ask the step function to pause by setting this flag and giving the
             # caller back a future
-            interrupt_exception = process_states.KillInterruption(msg)
+            interrupt_exception = process_states.KillInterruption(msg.get(process_comms.MESSAGE_KEY, None))
             self._set_interrupt_action_from_exception(interrupt_exception)
             self._killing = self._interrupt_action
             self._state.interrupt(interrupt_exception)
             return cast(futures.CancellableAction, self._interrupt_action)
 
+        breakpoint()
         self.transition_to(process_states.ProcessState.KILLED, msg)
         return True