diff --git a/src/plumpy/process_comms.py b/src/plumpy/process_comms.py index c66e8431..9a8b03df 100644 --- a/src/plumpy/process_comms.py +++ b/src/plumpy/process_comms.py @@ -30,6 +30,7 @@ INTENT_KEY = 'intent' MESSAGE_KEY = 'message' +FORCE_KILL_KEY = 'force_kill' class Intent: @@ -196,7 +197,7 @@ 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 @@ -204,9 +205,11 @@ async def kill_process(self, pid: 'PID_TYPE', msg: Optional[Any] = None) -> 'Pro :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) diff --git a/src/plumpy/processes.py b/src/plumpy/processes.py index b6e14ad9..95941f70 100644 --- a/src/plumpy/processes.py +++ b/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