Repository: pfnet/pytorch-pfn-extras

Issue:
ExtensionsManager need to call of `Extension.on_error` handler on error
Snapshot extensions provides `snapshot_on_error` option, but currently this is not working as `on_error` handler is not called from the manager.


Discussion:

Resolving patch:
diff --git a/pytorch_pfn_extras/training/manager.py b/pytorch_pfn_extras/training/manager.py
--- a/pytorch_pfn_extras/training/manager.py
+++ b/pytorch_pfn_extras/training/manager.py
@@ -160,6 +160,7 @@ def __init__(
         self.reporter = reporting.Reporter()
         self._transform_model = transform_model
         self._start_extensions_called = False
+        self._run_on_error_called = False
 
         # Indicates whether models can be accessed from extensions in the
         # current iteration.
@@ -435,6 +436,14 @@ def get_extension(self, name: str) -> extension_module.Extension:
         else:
             raise ValueError('extension %s not found' % name)
 
+    def _run_on_error(self, exc: Exception) -> None:
+        if not self._run_on_error_called:
+            self._run_on_error_called = True
+            tb = exc.__traceback__
+            assert tb is not None
+            for _, entry in self.extensions:
+                entry.extension.on_error(self, exc, tb)
+
     def run_extensions(
             self, *, completed: bool = True, only_iterations: bool = True) -> None:
         if completed:
@@ -640,9 +649,11 @@ def complete_iteration(
                 yield
                 for name in step_optimizers_names:
                     self._optimizers[name].step()
-            finally:
                 self.iteration += 1
                 self.run_extensions(completed=True, only_iterations=True)
+            except Exception as e:
+                self._run_on_error(e)
+                raise
 
     @contextlib.contextmanager
     def run_iteration(
@@ -676,7 +687,6 @@ def run_iteration(
                 if notification._is_completed:
                     for name in step_optimizers_names:
                         self._optimizers[name].step()
-            finally:
                 # The iteration count is increased just before calling the
                 # extensions.
                 if notification._is_completed:
@@ -691,6 +701,9 @@ def run_iteration(
                     # needs to run regardless of the training state
                     self.execution += 1
                     self.run_extensions(completed=False, only_iterations=False)
+            except Exception as e:
+                self._run_on_error(e)
+                raise
 
         if self._internal_stop_trigger(self):
             self._finalize_extensions()


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/training_tests/test_extension.py b/tests/pytorch_pfn_extras_tests/training_tests/test_extension.py
--- a/tests/pytorch_pfn_extras_tests/training_tests/test_extension.py
+++ b/tests/pytorch_pfn_extras_tests/training_tests/test_extension.py
@@ -78,3 +78,32 @@ def test_make_extension_unexpected_kwargs():
         @ppe.training.make_extension(foo=1)
         def my_extension(_):
             pass
+
+
+def test_on_error():
+    class DummyExt(ppe.training.Extension):
+        def __init__(self):
+            self.call_cnt = 0
+
+        def __call__(self, manager):
+            pass
+
+        def on_error(self, manager, exc, tb):
+            assert isinstance(exc, RuntimeError)
+            self.call_cnt += 1
+
+    optimizers = {'main': object()}
+    manager = ppe.training.ExtensionsManager(
+        {}, optimizers, 1, iters_per_epoch=1)
+    ext = DummyExt()
+    manager.extend(ext)
+
+    with manager.run_iteration():
+        pass
+    assert ext.call_cnt == 0
+    try:
+        with manager.run_iteration():
+            raise RuntimeError
+    except RuntimeError:
+        pass
+    assert ext.call_cnt == 1