Repository: a2i2/surround

Issue:
Throwing an exception doesn't stop the pipeline

Discussion:

Patch:
diff --git a/surround/assembler.py b/surround/assembler.py
index c028198..b3bb053 100644
--- a/surround/assembler.py
+++ b/surround/assembler.py
@@ -105,6 +105,8 @@ class Assembler(ABC):
 
         :param batch_mode: Whether batch mode should be used
         :type batch_mode: bool
+        :returns: whether the initialisation was successful
+        :rtype: bool
         """
 
         self.batch_mode = batch_mode
@@ -123,6 +125,9 @@ class Assembler(ABC):
                 self.finaliser.initialise(self.config)
         except Exception:
             LOGGER.exception("Failed initiating Assembler")
+            return False
+
+        return True
 
     def run(self, state=None, is_training=False):
         """
@@ -152,7 +157,12 @@ class Assembler(ABC):
         self.state = state
         try:
             self.validator.validate(self.state, self.config)
-            self.__run_pipeline(is_training)
+
+            if self.state.errors:
+                LOGGER.error("Error while validating")
+                LOGGER.error(self.state.errors)
+            else:
+                self.__run_pipeline(is_training)
         except Exception:
             LOGGER.exception("Failed running Assembler")
         finally:
@@ -169,15 +179,24 @@ class Assembler(ABC):
         """
 
         if self.pre_filters:
-            self.__execute_filters(self.pre_filters, self.state)
+            if not self.__execute_filters(self.pre_filters, self.state):
+                LOGGER.error("Failed running Assembler")
+                return
 
         if is_training:
             self.__execute_fit(self.state)
         else:
             self.__execute_main(self.state)
 
+        if self.state.errors:
+            LOGGER.error("Errors during executing the estimator")
+            LOGGER.error(self.state.errors)
+            return
+
         if self.post_filters:
-            self.__execute_filters(self.post_filters, self.state)
+            if not self.__execute_filters(self.post_filters, self.state):
+                LOGGER.error("Failed running Assembler")
+                return
 
         if (is_training or self.batch_mode) and self.visualiser:
             self.visualiser.visualise(self.state, self.config)
@@ -191,6 +210,8 @@ class Assembler(ABC):
         :type filters: list of :class:`surround.stage.Filter`
         :param state: the data being filtered
         :type state: :class:`surround.State`
+        :returns: whether execution succeeded
+        :rtype: bool
         """
 
         state.freeze()
@@ -202,14 +223,17 @@ class Assembler(ABC):
                 if state.errors:
                     LOGGER.error("Error during processing")
                     LOGGER.error(state.errors)
-                    break
+                    return False
+
             execution_time = datetime.now() - start_time
             state.execution_time = str(execution_time)
             LOGGER.info("Filter took %s secs", execution_time)
         except Exception:
             LOGGER.exception("Failed processing Surround Filters")
+            return False
 
         state.thaw()
+        return True
 
     def __execute_filter(self, stage, stage_data):
         """
diff --git a/templates/new/batch_runner.py.txt b/templates/new/batch_runner.py.txt
index c5a38ca..31b1a31 100644
--- a/templates/new/batch_runner.py.txt
+++ b/templates/new/batch_runner.py.txt
@@ -6,16 +6,18 @@ logging.basicConfig(level=logging.INFO)
 
 class BatchRunner(Runner):
     def run(self, is_training=False):
-        self.assembler.init_assembler(True)
-        data = AssemblerState()
+        if self.assembler.init_assembler(True):
+            data = AssemblerState()
 
-        # Load data to be processed
-        raw_data = "TODO: Load raw data here"
+            # Load data to be processed
+            raw_data = "TODO: Load raw data here"
 
-        # Setup input data
-        data.input_data = raw_data
+            # Setup input data
+            data.input_data = raw_data
 
-        # Run assembler
-        self.assembler.run(data, is_training)
+            # Run assembler
+            self.assembler.run(data, is_training)
 
-        logging.info("Batch Runner: %s", data.output_data)
+            logging.info("Batch Runner: %s", data.output_data)
+        else:
+            logging.exception("Failed to run the Assembler")


Test patch:
diff --git a/surround/tests/surround_test.py b/surround/tests/surround_test.py
index 1924842..026b34f 100644
--- a/surround/tests/surround_test.py
+++ b/surround/tests/surround_test.py
@@ -8,6 +8,11 @@ test_text = "hello"
 
 class HelloStage(Estimator):
     def estimate(self, state, config):
+        if state.estimator_add_error:
+            state.errors.append("Error!!")
+        elif state.estimator_throw:
+            raise Exception("Error!!")
+
         state.text = test_text
         if "helloStage" in config:
             state.config_value = config["helloStage"]["suffix"]
@@ -16,13 +21,20 @@ class HelloStage(Estimator):
         print("No training implemented")
 
 
+# pylint: disable=too-many-instance-attributes
 class AssemblerState(State):
-    text = None
-    config_value = None
-    stage1 = None
-    stage2 = None
-    final_ran = False
-
+    def __init__(self):
+        self.errors = []
+        self.text = None
+        self.config_value = None
+        self.stage1 = None
+        self.final_ran = False
+        self.use_errors_instead = False
+        self.post_filter_ran = False
+        self.post_filter_throw = False
+        self.validator_add_error = False
+        self.estimator_throw = False
+        self.estimator_add_error = False
 
 class InputValidator(Validator):
     def validate(self, state, config):
@@ -35,9 +47,22 @@ class InputValidator(Validator):
         if state.stage1:
             raise ValueError("'stage1' is not None")
 
-        if state.stage2:
-            raise ValueError("'stage2' is not None")
+        if state.validator_add_error:
+            state.errors.append("Error!!")
+
+class BadFilter(Filter):
+    def initialise(self, config):
+        raise Exception("This will fail always")
 
+    def operate(self, state, config):
+        if state.use_errors_instead:
+            state.errors.append("This will fail always")
+        else:
+            raise Exception("This will fail always")
+
+class PostFilter(Filter):
+    def operate(self, state, config):
+        state.post_filter_ran = True
 
 class TestFinalStage(Filter):
     def operate(self, state, config):
@@ -94,3 +119,88 @@ class TestSurround(unittest.TestCase):
 
         # Finalizer should still be executed
         self.assertTrue(data.final_ran)
+
+    def test_assembler_init_pass(self):
+        assembler = Assembler("Pass test", InputValidator(), HelloStage())
+        self.assertTrue(assembler.init_assembler())
+
+    def test_assembler_init_fail(self):
+        assembler = Assembler("Fail test", InputValidator())
+        assembler.set_estimator(HelloStage(), [BadFilter()])
+        self.assertFalse(assembler.init_assembler())
+
+    def test_pipeline_stop_on_exception_estimator(self):
+        data = AssemblerState()
+        data.estimator_throw = True
+
+        assembler = Assembler("Fail test", InputValidator())
+        assembler.set_estimator(HelloStage(), [], [PostFilter()])
+
+        # This should fail to execute PostFilter
+        assembler.run(data)
+
+        self.assertFalse(data.post_filter_ran)
+
+    def test_pipeline_stop_on_error_estimator(self):
+        data = AssemblerState()
+        data.estimator_add_error = True
+
+        assembler = Assembler("Fail test", InputValidator())
+        assembler.set_estimator(HelloStage(), [], [PostFilter()])
+
+        # This should fail to execute PostFilter
+        assembler.run(data)
+
+        self.assertFalse(data.post_filter_ran)
+
+    def test_pipeline_stop_on_exception_filter(self):
+        data = AssemblerState()
+
+        assembler = Assembler("Fail test", InputValidator())
+        assembler.set_estimator(HelloStage(), [BadFilter()], [PostFilter()])
+
+        # This should fail to execute HelloStage
+        assembler.run(data)
+
+        self.assertIsNone(data.text)
+        self.assertFalse(data.post_filter_ran)
+
+    def test_pipeline_stop_on_error_filter(self):
+        # Use the errors list in state instead
+        data = AssemblerState()
+        data.use_errors_instead = True
+
+        assembler = Assembler("Fail test", InputValidator())
+        assembler.set_estimator(HelloStage(), [BadFilter()], [PostFilter()])
+
+        # This should fail to execute HelloStage & PostFilter
+        assembler.run(data)
+
+        self.assertIsNone(data.text)
+        self.assertFalse(data.post_filter_ran)
+
+    def test_pipeline_stop_on_exception_validator(self):
+        data = AssemblerState()
+        data.stage1 = "Now it will fail in the validator"
+
+        assembler = Assembler("Fail test", InputValidator())
+        assembler.set_estimator(HelloStage(), [], [PostFilter()])
+
+        # This should fail to execute HelloStage & PostFilter
+        assembler.run(data)
+
+        self.assertIsNone(data.text)
+        self.assertFalse(data.post_filter_ran)
+
+    def test_pipeline_stop_on_error_validator(self):
+        data = AssemblerState()
+        data.validator_add_error = True
+
+        assembler = Assembler("Fail test", InputValidator())
+        assembler.set_estimator(HelloStage(), [], [PostFilter()])
+
+        # This should fail to execute HelloStage & PostFilter
+        assembler.run(data)
+
+        self.assertIsNone(data.text)
+        self.assertFalse(data.post_filter_ran)