883 Fix exception issue in StatsHandler and CheckpointSaver (#890)

* [DLMED] update handlers

* [DLMED] add more doc-strings

* [DLMED] add more doc-string

* [DLMED] update according to comments

* [DLMED] update according to Eric's comments

Author: Nic-Ma

monai/handlers/checkpoint_saver.py

@@ -37,11 +37,14 @@ class CheckpointSaver:
         name: identifier of logging.logger to use, if None, defaulting to ``engine.logger``.
         file_prefix: prefix for the filenames to which objects will be saved.
         save_final: whether to save checkpoint or session at final iteration or exception.
+            If checkpoints are to be saved when an exception is raised, put this handler before
+            `StatsHandler` in the handler list, because the logic with Ignite can only trigger
+            the first attached handler for `EXCEPTION_RAISED` event.
         save_key_metric: whether to save checkpoint or session when the value of key_metric is
             higher than all the previous values during training.keep 4 decimal places of metric,
             checkpoint name is: {file_prefix}_key_metric=0.XXXX.pth.
         key_metric_name: the name of key_metric in ignite metrics dictionary.
-            if None, use `engine.state.key_metric` instead.
+            If None, use `engine.state.key_metric` instead.
         key_metric_n_saved: save top N checkpoints or sessions, sorted by the value of key
             metric in descending order.
         epoch_level: save checkpoint during training for every N epochs or every N iterations.
@@ -168,7 +171,8 @@ def completed(self, engine: Engine) -> None:
 
     def exception_raised(self, engine: Engine, e: Exception) -> None:
         """Callback for train or validation/evaluation exception raised Event.
-        Save current data as final checkpoint if configure save_final is True.
+        Save current data as final checkpoint if configure save_final is True. This callback may be skipped
+        because the logic with Ignite can only trigger the first attached handler for `EXCEPTION_RAISED` event.
 
         Args:
             engine: Ignite Engine, it can be a trainer, validator or evaluator.
@@ -179,6 +183,7 @@ def exception_raised(self, engine: Engine, e: Exception) -> None:
         assert self.logger is not None
         assert hasattr(self.logger, "info"), "Error, provided logger has not info attribute."
         self.logger.info(f"Exception_raised, saved exception checkpoint: {self._final_checkpoint.last_checkpoint}")
+        raise e
 
     def metrics_completed(self, engine: Engine) -> None:
         """Callback to compare metrics and save models in train or validation when epoch completed.

monai/handlers/stats_handler.py

@@ -70,7 +70,7 @@ def __init__(
                 tag_name: scalar_value to logger. Defaults to ``'Loss'``.
             key_var_format: a formatting string to control the output string format of key: value.
             logger_handler: add additional handler to handle the stats data: save to file, etc.
-                add existing python logging handlers: https://docs.python.org/3/library/logging.handlers.html
+                Add existing python logging handlers: https://docs.python.org/3/library/logging.handlers.html
         """
 
         self.epoch_print_logger = epoch_print_logger
@@ -133,16 +133,16 @@ def iteration_completed(self, engine: Engine) -> None:
     def exception_raised(self, engine: Engine, e: Exception) -> None:
         """
         Handler for train or validation/evaluation exception raised Event.
-        Print the exception information and traceback.
+        Print the exception information and traceback. This callback may be skipped because the logic
+        with Ignite can only trigger the first attached handler for `EXCEPTION_RAISED` event.
 
         Args:
             engine: Ignite Engine, it can be a trainer, validator or evaluator.
             e: the exception caught in Ignite during engine.run().
 
         """
         self.logger.exception(f"Exception: {e}")
-        # import traceback
-        # traceback.print_exc()
+        raise e
 
     def _default_epoch_print(self, engine: Engine) -> None:
         """

tests/test_handler_checkpoint_saver.py

@@ -11,7 +11,6 @@
 
 import logging
 import os
-import shutil
 import sys
 import tempfile
 import unittest
@@ -80,25 +79,43 @@ def _train_func(engine, batch):
         if multi_devices:
             net = torch.nn.DataParallel(net)
         optimizer = optim.SGD(net.parameters(), lr=0.02)
-        tempdir = tempfile.mkdtemp()
-        handler = CheckpointSaver(
-            tempdir,
-            {"net": net, "opt": optimizer},
-            "CheckpointSaver",
-            "test",
-            save_final,
-            save_key_metric,
-            key_metric_name,
-            key_metric_n_saved,
-            epoch_level,
-            save_interval,
-            n_saved,
-        )
-        handler.attach(engine)
-        engine.run(data, max_epochs=5)
-        for filename in filenames:
-            self.assertTrue(os.path.exists(os.path.join(tempdir, filename)))
-        shutil.rmtree(tempdir)
+        with tempfile.TemporaryDirectory() as tempdir:
+            handler = CheckpointSaver(
+                tempdir,
+                {"net": net, "opt": optimizer},
+                "CheckpointSaver",
+                "test",
+                save_final,
+                save_key_metric,
+                key_metric_name,
+                key_metric_n_saved,
+                epoch_level,
+                save_interval,
+                n_saved,
+            )
+            handler.attach(engine)
+            engine.run(data, max_epochs=5)
+            for filename in filenames:
+                self.assertTrue(os.path.exists(os.path.join(tempdir, filename)))
+
+    def test_exception(self):
+        logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+        net = torch.nn.PReLU()
+
+        # set up engine
+        def _train_func(engine, batch):
+            raise RuntimeError("test exception.")
+
+        engine = Engine(_train_func)
+
+        # set up testing handler
+        with tempfile.TemporaryDirectory() as tempdir:
+            stats_handler = CheckpointSaver(tempdir, {"net": net}, save_final=True)
+            stats_handler.attach(engine)
+
+            with self.assertRaises(RuntimeError):
+                engine.run(range(3), max_epochs=2)
+            self.assertTrue(os.path.exists(os.path.join(tempdir, "net_final_iteration=1.pth")))
 
 
 if __name__ == "__main__":

tests/test_handler_stats.py

@@ -143,6 +143,22 @@ def _train_func(engine, batch):
                         self.assertTrue(has_key_word.match(line))
         shutil.rmtree(tempdir)
 
+    def test_exception(self):
+        logging.basicConfig(level=logging.INFO)
+
+        # set up engine
+        def _train_func(engine, batch):
+            raise RuntimeError("test exception.")
+
+        engine = Engine(_train_func)
+
+        # set up testing handler
+        stats_handler = StatsHandler()
+        stats_handler.attach(engine)
+
+        with self.assertRaises(RuntimeError):
+            engine.run(range(3), max_epochs=2)
+
 
 if __name__ == "__main__":
     unittest.main()