class TestTickingEngine(AiidaTestCase):
def setUp(self):
super(TestTickingEngine, self).setUp()
self.assertEquals(len(util.ProcessStack.stack()), 0)
self.ticking_engine = TickingEngine()
self.executor = ThreadPoolExecutor(max_workers=1)
def tearDown(self):
super(TestTickingEngine, self).tearDown()
self.ticking_engine.shutdown()
self.assertEquals(len(util.ProcessStack.stack()), 0)
def test_get_process(self):
# Test cancelling a future before the process runs
future = self.ticking_engine.submit(DummyProcess, inputs={'a': Int(5)})
def test_submit(self):
fut = self.ticking_engine.submit(DummyProcess, inputs={'a': Int(5)})
self._tick_till_finished()
res = fut.result()
self.assertTrue(res['ran'].value)
def test_cancel(self):
# Test cancelling a future before the process runs
future = self.ticking_engine.submit(DummyProcess, inputs={'a': Int(5)})
self.assertTrue(future.running())
future.cancel()
self.assertTrue(future.cancelled())
def _tick_till_finished(self):
self.executor.submit(self._keep_ticking())
def _keep_ticking(self):
while self.ticking_engine.tick():
pass
def _run_with_checkpoints(self, wf_class, inputs=None):
finished_steps = {}
te = TickingEngine()
fut = te.submit(wf_class, inputs)
while not fut.done():
pid = fut.pid
te.tick()
finished_steps.update(wf_class.finished_steps)
te.shutdown()
return finished_steps
def test_simple_run(self):
"""
Run the workchain which should hit the exception and therefore end
up in the FAILED state
"""
engine = TickingEngine()
future = engine.submit(TestWorkChainAbortChildren.MainWorkChain)
while not future.done():
engine.tick()
self.assertEquals(future.process.calc.has_finished_ok(), False)
self.assertEquals(future.process.calc.has_failed(), True)
self.assertEquals(future.process.calc.has_aborted(), False)
engine.shutdown()
def _run_with_checkpoints(self, wf_class, inputs=None):
finished_steps = {}
te = TickingEngine()
fut = te.submit(wf_class, inputs)
while not fut.done():
pid = fut.pid
te.tick()
finished_steps.update(wf_class.finished_steps)
# if not fut.done():
# te.stop(pid)
# fut = te.run_from(storage.load_checkpoint(pid))
te.shutdown()
return finished_steps
def test_simple_kill_through_process(self):
"""
Run the workchain for one step and then kill it by calling kill
on the workchain itself. This should have the workchain end up
in the ABORTED state.
"""
engine = TickingEngine()
future = engine.submit(TestWorkChainAbort.AbortableWorkChain)
while not future.done():
engine.tick()
future.process.abort()
self.assertEquals(future.process.calc.has_finished_ok(), False)
self.assertEquals(future.process.calc.has_failed(), False)
self.assertEquals(future.process.calc.has_aborted(), True)
engine.shutdown()
def test_simple_kill_through_node(self):
"""
Run the workchain for one step and then kill it by calling kill
on the underlying WorkCalculation node. This should have the
workchain end up in the ABORTED state.
"""
engine = TickingEngine()
future = engine.submit(TestWorkChainAbortChildren.MainWorkChain, {'kill': Bool(True)})
while not future.done():
engine.tick()
child = future.process.calc.get_outputs(link_type=LinkType.CALL)[0]
self.assertEquals(child.has_finished_ok(), False)
self.assertEquals(child.has_failed(), False)
self.assertEquals(child.has_aborted(), True)
self.assertEquals(future.process.calc.has_finished_ok(), False)
self.assertEquals(future.process.calc.has_failed(), False)
self.assertEquals(future.process.calc.has_aborted(), True)
engine.shutdown()