wip

Author: agoscinski

src/aiida/engine/processes/process.py

@@ -329,51 +329,87 @@ def load_instance_state(
 
         self.node.logger.info(f'Loaded process<{self.node.pk}> from saved state')
 
+    async def _launch_task(self, coro, *args, **kwargs):
+        """Launch a coroutine as a task, making sure to make it interruptable."""
+        import functools
+        from aiida.engine.utils import interruptable_task
+        task_fn = functools.partial(coro, *args, **kwargs)
+        try:
+            self._task = interruptable_task(task_fn)
+            result = await self._task
+            return result
+        finally:
+            self._task = None
+
     def kill(self, msg_text: str | None = None, force_kill: bool = False) -> Union[bool, plumpy.futures.Future]:
         """Kill the process and all the children calculations it called
 
         :param msg: message
         """
-        self.node.logger.info(f'Request to kill Process<{self.node.pk}>')
-
-        had_been_terminated = self.has_terminated()
-
+        from .calcjobs.tasks import task_kill_job
+
+        if self._killing and not force_kill:
+            # if already killing we have triggered the Interruption
+            coro = self._launch_task(task_kill_job, self.node, self.runner.transport)
+            task = asyncio.create_task(coro)
+            _ = self.loop.run_until_complete(task)
+        # the parent class invokes the transition to being killed
         result = super().kill(msg_text, force_kill)
 
-        # Only kill children if we could be killed ourselves
-        if result is not False and not had_been_terminated:
-            killing = []
-            for child in self.node.called:
-                if self.runner.controller is None:
-                    self.logger.info('no controller available to kill child<%s>', child.pk)
-                    continue
-                try:
-                    result = self.runner.controller.kill_process(child.pk, msg_text=f'Killed by parent<{self.node.pk}>')
-                    result = asyncio.wrap_future(result)  # type: ignore[arg-type]
-                    if asyncio.isfuture(result):
-                        killing.append(result)
-                except ConnectionClosed:
-                    self.logger.info('no connection available to kill child<%s>', child.pk)
-                except UnroutableError:
-                    self.logger.info('kill signal was unable to reach child<%s>', child.pk)
-
-            if asyncio.isfuture(result):
-                # We ourselves are waiting to be killed so add it to the list
-                killing.append(result)
-
-            if killing:
-                # We are waiting for things to be killed, so return the 'gathered' future
-                kill_future = plumpy.futures.gather(*killing)
-                result = self.loop.create_future()
-
-                def done(done_future: plumpy.futures.Future):
-                    is_all_killed = all(done_future.result())
-                    result.set_result(is_all_killed)
-
-                kill_future.add_done_callback(done)
-
         return result
 
+        # TODO might need to be merged with `task_kill_job`
+        #if self._killing is not None:
+        #    self._killing.set_result(True)
+        #else:
+        #    self.node.logger.info(f'killed CalcJob<{self.node.pk}> but async future was None')
+
+        #self.node.logger.info(f'Request to kill Process<{self.node.pk}>')
+        #from .
+        #await self._launch_task(task_kill_job, node, transport_queue)
+        #if self._killing is not None:
+        #    self._killing.set_result(True)
+        #else:
+        #    logger.info(f'killed CalcJob<{node.pk}> but async future was None')
+
+        #had_been_terminated = self.has_terminated()
+
+        #result = super().kill(msg_text, force_kill)
+
+        ## Only kill children if we could be killed ourselves
+        #if result is not False and not had_been_terminated:
+        #    killing = []
+        #    for child in self.node.called:
+        #        if self.runner.controller is None:
+        #            self.logger.info('no controller available to kill child<%s>', child.pk)
+        #            continue
+        #        try:
+        #            result = self.runner.controller.kill_process(child.pk, msg_text=f'Killed by parent<{self.node.pk}>')
+        #            result = asyncio.wrap_future(result)  # type: ignore[arg-type]
+        #            if asyncio.isfuture(result):
+        #                killing.append(result)
+        #        except ConnectionClosed:
+        #            self.logger.info('no connection available to kill child<%s>', child.pk)
+        #        except UnroutableError:
+        #            self.logger.info('kill signal was unable to reach child<%s>', child.pk)
+
+        #    if asyncio.isfuture(result):
+        #        # We ourselves are waiting to be killed so add it to the list
+        #        killing.append(result)
+
+        #    if killing:
+        #        # We are waiting for things to be killed, so return the 'gathered' future
+        #        kill_future = plumpy.futures.gather(*killing)
+        #        result = self.loop.create_future()
+
+        #        def done(done_future: plumpy.futures.Future):
+        #            is_all_killed = all(done_future.result())
+        #            result.set_result(is_all_killed)
+
+        #        kill_future.add_done_callback(done)
+
+        #return result
+
     @override
     def out(self, output_port: str, value: Any = None) -> None:
         """Attach output to output port.