Add built-in support and API for exit codes in WorkChains (#1640)

Currently it is possible to return non-zero integers from a WorkChain
outline step to abort its execution and assign the exit status to the
node, however, there is no official API yet to make this easier.

Here we define the concept of an ExitCode, a named tuple that takes
an integer exit status and a descriptive message. Through the
ProcessSpec, a WorkChain developer can add exit codes that correspond
to errors that may crop up during the execution of the workchain. For
example the following spec definition:

	@classmethod
	def define(cls, spec):
   		super(CifCleanWorkChain, cls).define(spec)
    	spec.exit_code(418, 'ERROR_I_AM_A_TEAPOT',
        	message='workchain found itself in an identity crisis')

In one of the outline steps, the exit code can be used by retrieving
it through either the integer exit status or the string label:

	return self.exit_codes('I_AM_A_TEAPOT')

or

	return self.exit_codes(418)

and return it. Returning an instance of ExitCode will trigger the exact
same mechanism as returning an integer from an outline step. The engine
will detect the ExitCode and return its exit status as the result,
triggering the abort and the exit status to be set on the Node.

Note that this addition required name changes in parts of the existing
API. For example the Calculation attribute `finish_status` was renamed
to `exit_status`. This is because it can now be accompanied by an string
`exit_message`. The `exit_status` and `exit_message` together form the
`ExitCode`.

.ci/test_daemon.py

@@ -76,8 +76,8 @@ def validate_calculations(expected_results):
     for pk, expected_dict in expected_results.iteritems():
         calc = load_node(pk)
         if not calc.is_finished_ok:
-            print 'Calculation<{}> not finished ok: process_state<{}> finish_status<{}>'.format(
-                pk, calc.process_state, calc.finish_status)
+            print 'Calculation<{}> not finished ok: process_state<{}> exit_status<{}>'.format(
+                pk, calc.process_state, calc.exit_status)
             print_logshow(pk)
             valid = False
 
@@ -118,8 +118,8 @@ def validate_workchains(expected_results):
 
         # I check only if this_valid, otherwise calc could not exist
         if this_valid and not calc.is_finished_ok:
-            print 'Calculation<{}> not finished ok: process_state<{}> finish_status<{}>'.format(
-                pk, calc.process_state, calc.finish_status)
+            print 'Calculation<{}> not finished ok: process_state<{}> exit_status<{}>'.format(
+                pk, calc.process_state, calc.exit_status)
             print_logshow(pk)
             valid = False
             this_valid = False
@@ -142,8 +142,8 @@ def validate_cached(cached_calcs):
     for calc in cached_calcs:
 
         if not calc.is_finished_ok:
-            print 'Cached calculation<{}> not finished ok: process_state<{}> finish_status<{}>'.format(
-                pk, calc.process_state, calc.finish_status)
+            print 'Cached calculation<{}> not finished ok: process_state<{}> exit_status<{}>'.format(
+                pk, calc.process_state, calc.exit_status)
             print_logshow(pk)
             valid = False
 

.pre-commit-config.yaml

@@ -22,6 +22,7 @@
         aiida/work/awaitable.py|
         aiida/work/context.py|
         aiida/work/exceptions.py|
+        aiida/work/exit_code.py|
         aiida/work/futures.py|
         aiida/work/launch.py|
         aiida/work/persistence.py|
@@ -61,6 +62,7 @@
         aiida/work/awaitable.py|
         aiida/work/context.py|
         aiida/work/exceptions.py|
+        aiida/work/exit_code.py|
         aiida/work/futures.py|
         aiida/work/launch.py|
         aiida/work/persistence.py|

aiida/backends/tests/inline_calculation.py

@@ -40,16 +40,16 @@ def test_inline_calculation_process_state(self):
         self.assertEquals(calculation.is_finished_ok, True)
         self.assertEquals(calculation.is_failed, False)
 
-    def test_finish_status(self):
+    def test_exit_status(self):
         """
         If an InlineCalculation reaches the FINISHED process state, it has to have been successful
-        which means that the finish status always has to be 0
+        which means that the exit status always has to be 0
         """
         calculation, result = self.incr_inline(inp=Int(11))
         self.assertEquals(calculation.is_finished, True)
         self.assertEquals(calculation.is_finished_ok, True)
         self.assertEquals(calculation.is_failed, False)
-        self.assertEquals(calculation.finish_status, 0)
+        self.assertEquals(calculation.exit_status, 0)
 
     def test_incr(self):
         """

aiida/backends/tests/work/__init__.py

@@ -7,4 +7,3 @@
 # For further information on the license, see the LICENSE.txt file        #
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
-

aiida/backends/tests/work/test_process_spec.py

@@ -44,3 +44,42 @@ def test_dynamic_output(self):
         self.assertFalse(self.spec.validate_outputs({'key': 5})[0])
         self.assertFalse(self.spec.validate_outputs({'key': n})[0])
         self.assertTrue(self.spec.validate_outputs({'key': d})[0])
+
+    def test_exit_code(self):
+        """
+        Test the definition of error codes through the ProcessSpec
+        """
+        label = 'SOME_EXIT_CODE'
+        status = 418
+        message = 'I am a teapot'
+
+        self.spec.exit_code(status, label, message)
+
+        self.assertEquals(self.spec.exit_codes.SOME_EXIT_CODE.status, status)
+        self.assertEquals(self.spec.exit_codes.SOME_EXIT_CODE.message, message)
+
+        self.assertEquals(self.spec.exit_codes['SOME_EXIT_CODE'].status, status)
+        self.assertEquals(self.spec.exit_codes['SOME_EXIT_CODE'].message, message)
+
+        self.assertEquals(self.spec.exit_codes[label].status, status)
+        self.assertEquals(self.spec.exit_codes[label].message, message)
+
+    def test_exit_code_invalid(self):
+        """
+        Test type validation for registering new error codes
+        """
+        status = 418
+        label = 'SOME_EXIT_CODE'
+        message = 'I am a teapot'
+
+        with self.assertRaises(TypeError):
+            self.spec.exit_code(status, 256, message)
+
+        with self.assertRaises(TypeError):
+            self.spec.exit_code('string', label, message)
+
+        with self.assertRaises(ValueError):
+            self.spec.exit_code(-256, label, message)
+
+        with self.assertRaises(TypeError):
+            self.spec.exit_code(status, label, 8)

aiida/backends/tests/work/test_workfunctions.py

@@ -14,7 +14,7 @@
 from aiida.orm.data.str import Str
 from aiida.orm import load_node
 from aiida.orm.calculation.function import FunctionCalculation
-from aiida.work import run, run_get_node, submit, workfunction, Process, Exit
+from aiida.work import run, run_get_node, submit, workfunction, Process, ExitCode
 
 DEFAULT_INT = 256
 DEFAULT_LABEL = 'Default label'
@@ -64,8 +64,8 @@ def wf_default_label_description(a=Int(DEFAULT_INT), label=DEFAULT_LABEL, descri
             return a
 
         @workfunction
-        def wf_exit(exit_code):
-            raise Exit(exit_code.value)
+        def wf_exit_code(exit_status, exit_message):
+            return ExitCode(exit_status.value, exit_message.value)
 
         @workfunction
         def wf_excepts(exception):
@@ -79,7 +79,7 @@ def wf_excepts(exception):
         self.wf_args_and_kwargs = wf_args_and_kwargs
         self.wf_args_and_default = wf_args_and_default
         self.wf_default_label_description = wf_default_label_description
-        self.wf_exit = wf_exit
+        self.wf_exit_code = wf_exit_code
         self.wf_excepts = wf_excepts
 
     def tearDown(self):
@@ -217,13 +217,13 @@ def test_wf_default_label_description(self):
         self.assertEquals(node.label, CUSTOM_LABEL)
         self.assertEquals(node.description, CUSTOM_DESCRIPTION)
 
-    def test_finish_status(self):
+    def test_exit_status(self):
         """
         If a workfunction reaches the FINISHED process state, it has to have been successful
-        which means that the finish status always has to be 0
+        which means that the exit status always has to be 0
         """
         result, node = self.wf_args_with_default.run_get_node()
-        self.assertEquals(node.finish_status, 0)
+        self.assertEquals(node.exit_status, 0)
         self.assertEquals(node.is_finished_ok, True)
         self.assertEquals(node.is_failed, False)
 
@@ -242,15 +242,17 @@ def test_launchers(self):
         with self.assertRaises(AssertionError):
             submit(self.wf_return_true)
 
-    def test_exit_exception(self):
+    def test_return_exit_code(self):
         """
-        A workfunction that raises the Exit exception should not EXCEPT but be FINISHED
+        A workfunction that returns an ExitCode namedtuple should have its exit status and message set FINISHED
         """
-        finish_status = 418
-        result, node = self.wf_exit.run_get_node(exit_code=Int(finish_status))
+        exit_status = 418
+        exit_message = 'I am a teapot'
+        result, node = self.wf_exit_code.run_get_node(exit_status=Int(exit_status), exit_message=Str(exit_message))
         self.assertTrue(node.is_finished)
         self.assertFalse(node.is_finished_ok)
-        self.assertEquals(node.finish_status, finish_status)
+        self.assertEquals(node.exit_status, exit_status)
+        self.assertEquals(node.exit_message, exit_message)
 
     def test_normal_exception(self):
         """

aiida/backends/tests/work/work_chain.py

@@ -7,11 +7,9 @@
 # For further information on the license, see the LICENSE.txt file        #
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
-
 import inspect
 import plumpy
 import plumpy.test_utils
-import unittest
 
 from aiida.backends.testbase import AiidaTestCase
 from aiida.common.links import LinkType
@@ -117,38 +115,54 @@ def _set_finished(self, function_name):
         self.finished_steps[function_name] = True
 
 
-class ReturnWorkChain(WorkChain):
-    FAILURE_STATUS = 1
+class PotentialFailureWorkChain(WorkChain):
+
+    EXIT_STATUS = 1
+    EXIT_MESSAGE = 'Well you did ask for it'
 
     @classmethod
     def define(cls, spec):
-        super(ReturnWorkChain, cls).define(spec)
+        super(PotentialFailureWorkChain, cls).define(spec)
         spec.input('success', valid_type=Bool)
+        spec.input('through_exit_code', valid_type=Bool, default=Bool(False))
+        spec.exit_code(cls.EXIT_STATUS, 'EXIT_STATUS', cls.EXIT_MESSAGE)
         spec.outline(
             cls.failure,
             cls.success
         )
 
     def failure(self):
         if self.inputs.success.value is False:
-            return self.FAILURE_STATUS
+            if self.inputs.through_exit_code.value is False:
+                return self.EXIT_STATUS
+            else:
+                return self.exit_codes.EXIT_STATUS
 
     def success(self):
         return
 
 
-class TestFinishStatus(AiidaTestCase):
+class TestExitStatus(AiidaTestCase):
 
     def test_failing_workchain(self):
-        result, node = work.run_get_node(ReturnWorkChain, success=Bool(False))
-        self.assertEquals(node.finish_status, ReturnWorkChain.FAILURE_STATUS)
+        result, node = work.run_get_node(PotentialFailureWorkChain, success=Bool(False))
+        self.assertEquals(node.exit_status, PotentialFailureWorkChain.EXIT_STATUS)
+        self.assertEquals(node.exit_message, None)
+        self.assertEquals(node.is_finished, True)
+        self.assertEquals(node.is_finished_ok, False)
+        self.assertEquals(node.is_failed, True)
+
+    def test_failing_workchain_with_message(self):
+        result, node = work.run_get_node(PotentialFailureWorkChain, success=Bool(False), through_exit_code=Bool(True))
+        self.assertEquals(node.exit_status, PotentialFailureWorkChain.EXIT_STATUS)
+        self.assertEquals(node.exit_message, PotentialFailureWorkChain.EXIT_MESSAGE)
         self.assertEquals(node.is_finished, True)
         self.assertEquals(node.is_finished_ok, False)
         self.assertEquals(node.is_failed, True)
 
     def test_successful_workchain(self):
-        result, node = work.run_get_node(ReturnWorkChain, success=Bool(True))
-        self.assertEquals(node.finish_status, 0)
+        result, node = work.run_get_node(PotentialFailureWorkChain, success=Bool(True))
+        self.assertEquals(node.exit_status, 0)
         self.assertEquals(node.is_finished, True)
         self.assertEquals(node.is_finished_ok, True)
         self.assertEquals(node.is_failed, False)
@@ -547,6 +561,41 @@ def check_input(self):
 
         run_and_check_success(TestWorkChain, namespace={'value': value})
 
+    def test_exit_codes(self):
+        status = 418
+        label = 'SOME_EXIT_CODE'
+        message = 'I am a teapot'
+
+        class ExitCodeWorkChain(WorkChain):
+
+            @classmethod
+            def define(cls, spec):
+                super(ExitCodeWorkChain, cls).define(spec)
+                spec.outline(cls.run)
+                spec.exit_code(status, label, message)
+
+            def run(self):
+                pass
+
+        wc = ExitCodeWorkChain()
+
+        # The exit code can be gotten by calling it with the status or label, as well as using attribute dereferencing
+        self.assertEquals(wc.exit_codes(status).status, status)
+        self.assertEquals(wc.exit_codes(label).status, status)
+        self.assertEquals(wc.exit_codes.SOME_EXIT_CODE.status, status)
+
+        with self.assertRaises(AttributeError):
+            wc.exit_codes.NON_EXISTENT_ERROR
+
+        self.assertEquals(ExitCodeWorkChain.exit_codes.SOME_EXIT_CODE.status, status)
+        self.assertEquals(ExitCodeWorkChain.exit_codes.SOME_EXIT_CODE.message, message)
+
+        self.assertEquals(ExitCodeWorkChain.exit_codes['SOME_EXIT_CODE'].status, status)
+        self.assertEquals(ExitCodeWorkChain.exit_codes['SOME_EXIT_CODE'].message, message)
+
+        self.assertEquals(ExitCodeWorkChain.exit_codes[label].status, status)
+        self.assertEquals(ExitCodeWorkChain.exit_codes[label].message, message)
+
     def _run_with_checkpoints(self, wf_class, inputs=None):
         if inputs is None:
             inputs = {}
@@ -1084,4 +1133,3 @@ def test_nested_expose(self):
                 'sub.sub.sub_2.b': Float(1.2), 'sub.sub.sub_2.sub_3.c': Bool(False)
             }
         )
-

aiida/cmdline/commands/calculation.py

@@ -175,10 +175,10 @@ def calculation_list(self, *args):
         parser.set_defaults(all_states=False)
         parser.add_argument('-S', '--process-state', choices=([e.value for e in ProcessState]),
                             help='Only include entries with this process state')
-        parser.add_argument('-f', '--finish-status', type=int,
-                            help='Only include entries with this finish status')
+        parser.add_argument('-E', '--exit-status', type=int,
+                            help='Only include entries with this exit status')
         parser.add_argument('-n', '--failed', dest='failed', action='store_true',
-                            help='Only include entries that are failed, i.e. whose finish status is non-zero')
+                            help='Only include entries that are failed, i.e. whose exit status is non-zero')
         parser.add_argument('-A', '--all-users',
                             dest='all_users', action='store_true',
                             help="Show calculations for all users, rather than only for the current user")
@@ -216,7 +216,7 @@ def calculation_list(self, *args):
             parsed_args.states = None
 
         PROCESS_STATE_KEY = 'attributes.{}'.format(C.PROCESS_STATE_KEY)
-        FINISH_STATUS_KEY = 'attributes.{}'.format(C.FINISH_STATUS_KEY)
+        EXIT_STATUS_KEY = 'attributes.{}'.format(C.EXIT_STATUS_KEY)
 
         filters = {}
 
@@ -227,12 +227,12 @@ def calculation_list(self, *args):
         if parsed_args.failed:
             parsed_args.states = None
             filters[PROCESS_STATE_KEY] = {'==': ProcessState.FINISHED.value}
-            filters[FINISH_STATUS_KEY] = {'!==': 0}
+            filters[EXIT_STATUS_KEY] = {'!==': 0}
 
-        if parsed_args.finish_status:
+        if parsed_args.exit_status:
             parsed_args.states = None
             filters[PROCESS_STATE_KEY] = {'==': ProcessState.FINISHED.value}
-            filters[FINISH_STATUS_KEY] = {'==': parsed_args.finish_status}
+            filters[EXIT_STATUS_KEY] = {'==': parsed_args.exit_status}
 
         C._list_calculations(
             states=parsed_args.states,