diff --git a/aiida/cmdline/commands/work.py b/aiida/cmdline/commands/work.py index c4c079cf28..7e05084f52 100644 --- a/aiida/cmdline/commands/work.py +++ b/aiida/cmdline/commands/work.py @@ -34,6 +34,7 @@ def __init__(self): 'report': (self.cli, self.complete_none), 'tree': (self.cli, self.complete_none), 'checkpoint': (self.cli, self.complete_none), + 'kill': (self.cli, self.complete_none), } def cli(self, *args): @@ -264,3 +265,31 @@ def _build_query(order_by=None, limit=None, past_days=None): qb.limit(limit) return qb.iterdict() + +@work.command('kill', context_settings=CONTEXT_SETTINGS) +@click.argument('pks', nargs=-1, type=int) +def kill(pks): + from aiida import try_load_dbenv + try_load_dbenv() + from aiida.orm import load_node + from aiida.orm.calculation.work import WorkCalculation + + nodes = [load_node(pk) for pk in pks] + workchain_nodes = [n for n in nodes if isinstance(n, WorkCalculation)] + running_workchain_nodes = [n for n in nodes if not n.has_finished()] + + num_workchains = len(running_workchain_nodes) + if num_workchains > 0: + answer = click.prompt( + 'Are you sure you want to kill {} workflows and all their children? [y/n]'.format( + num_workchains + ) + ).lower() + if answer == 'y': + click.echo('Killing workflows.') + for n in running_workchain_nodes: + n.kill() + else: + click.echo('Abort!') + else: + click.echo('No pks of valid running workchains given.') diff --git a/aiida/orm/implementation/general/calculation/work.py b/aiida/orm/implementation/general/calculation/work.py index 973933c2d7..e7d0685653 100644 --- a/aiida/orm/implementation/general/calculation/work.py +++ b/aiida/orm/implementation/general/calculation/work.py @@ -10,6 +10,7 @@ from aiida.orm.implementation.calculation import Calculation from aiida.common.lang import override +from aiida.common.links import LinkType class WorkCalculation(Calculation): """ @@ -18,6 +19,7 @@ class WorkCalculation(Calculation): """ FINISHED_KEY = '_finished' FAILED_KEY = '_failed' + ABORTED_KEY = '_aborted' @override def has_finished(self): @@ -43,4 +45,17 @@ def has_failed(self): :return: True if the calculation has failed, False otherwise. :rtype: bool """ - return self.get_attr(self.FAILED_KEY, False) is not False \ No newline at end of file + return ( + self.get_attr(self.FAILED_KEY, False) or + self.get_attr(self.ABORTED_KEY, False) + ) + + def kill(self): + """ + Kill a WorkCalculation and all its children. + """ + if not self.is_sealed: + self._set_attr(self.ABORTED_KEY, True) + self.seal() + for child in self.get_outputs(link_type=LinkType.CALL): + child.kill() diff --git a/aiida/work/process.py b/aiida/work/process.py index 388d136a72..3171a056d6 100644 --- a/aiida/work/process.py +++ b/aiida/work/process.py @@ -511,6 +511,7 @@ def on_monitored_process_failed(self, pid): except ValueError: pass else: + calc_node._set_attr(calc_node.FAILED_KEY, True) calc_node.seal() aiida.work.util.ProcessStack.pop(pid=pid) diff --git a/aiida/work/workchain.py b/aiida/work/workchain.py index 7a792a0316..e3bc3e9048 100644 --- a/aiida/work/workchain.py +++ b/aiida/work/workchain.py @@ -221,8 +221,17 @@ def _run(self, **kwargs): self._stepper = self.spec().get_outline().create_stepper(self) return self._do_step() + @property + def _aborted(self): + return self._aborted_attr or self.calc.get_attr(self.calc.ABORTED_KEY, False) + + @_aborted.setter + def _aborted(self, value): + self._aborted_attr = value + def _do_step(self, wait_on=None): if self._aborted: + self.calc.kill() return for interstep in self._intersteps: