Ensure finalize simulation stage is always run (#154)

* ensure finalize stage is always run

* add test

* update doc

* black

* fix PR link in release notes

Author: benbovy

doc/framework.rst

@@ -202,7 +202,8 @@ A model run is divided into four successive stages:
 4. finalization
 
 During a simulation, stages 1 and 4 are run only once while stages 2
-and 3 are repeated for a given number of (time) steps.
+and 3 are repeated for a given number of (time) steps. Stage 4 is run even if
+an exception is raised during stage 1, 2 or 3.
 
 Each process-ified class may provide its own computation instructions
 by implementing specific methods named ``.initialize()``,

doc/whats_new.rst

@@ -36,6 +36,8 @@ Bug fixes
   which inherit from process-decorated classes (:issue:`149`).
 - Fix :func:`~xarray.Dataset.xsimlab.update_clocks` when only the master clock is
   updated implicitly (:issue:`151`).
+- Ensure that the ``finalize`` simulation stage is always run, even when an
+  exception is raised during the previous stages (:issue:`154`).
 
 v0.4.1 (17 April 2020)
 ----------------------

xsimlab/drivers.py

@@ -330,30 +330,33 @@ def _run(
 
     in_vars = _get_input_vars(ds_init, model)
     model.update_state(in_vars, validate=validate_inputs, ignore_static=True)
-    model.execute("initialize", rt_context, **execute_kwargs)
 
-    for step, (_, ds_step) in enumerate(ds_gby_steps):
+    try:
+        model.execute("initialize", rt_context, **execute_kwargs)
 
-        rt_context.update(
-            step=step,
-            step_start=ds_step["_clock_start"].values,
-            step_end=ds_step["_clock_end"].values,
-            step_delta=ds_step["_clock_diff"].values,
-        )
-
-        in_vars = _get_input_vars(ds_step, model)
-        model.update_state(in_vars, validate=validate_inputs, ignore_static=False)
-        model.execute("run_step", rt_context, **execute_kwargs)
+        for step, (_, ds_step) in enumerate(ds_gby_steps):
 
-        store.write_output_vars(batch, step, model=model)
+            rt_context.update(
+                step=step,
+                step_start=ds_step["_clock_start"].values,
+                step_end=ds_step["_clock_end"].values,
+                step_delta=ds_step["_clock_diff"].values,
+            )
 
-        model.execute("finalize_step", rt_context, **execute_kwargs)
+            in_vars = _get_input_vars(ds_step, model)
+            model.update_state(in_vars, validate=validate_inputs, ignore_static=False)
+            model.execute("run_step", rt_context, **execute_kwargs)
 
-    store.write_output_vars(batch, -1, model=model)
+            store.write_output_vars(batch, step, model=model)
 
-    model.execute("finalize", rt_context, **execute_kwargs)
+            model.execute("finalize_step", rt_context, **execute_kwargs)
 
-    store.write_index_vars(model=model)
+        store.write_output_vars(batch, -1, model=model)
+        store.write_index_vars(model=model)
+    except Exception as error:
+        raise error
+    finally:
+        model.execute("finalize", rt_context, **execute_kwargs)
 
 
 class XarraySimulationDriver(BaseSimulationDriver):

xsimlab/tests/test_drivers.py

@@ -39,7 +39,7 @@ def test_runtime_context_in_model(in_dataset, model):
     class P:
         @xs.runtime(args="not_a_runtime_arg")
         def run_step(self, arg):
-            pass
+            print(arg)
 
     m = model.update_processes({"p": P})
 
@@ -93,7 +93,7 @@ def test_constructor(self, in_dataset, model):
         with pytest.raises(KeyError, match=r"Missing variables.*"):
             XarraySimulationDriver(invalid_ds, model)
 
-    def test_run_model_get_results(self, in_dataset, out_dataset, xarray_driver):
+    def test_run_model_get_results(self, out_dataset, xarray_driver):
         xarray_driver.run_model()
         out_ds_actual = xarray_driver.get_results()
 
@@ -124,3 +124,27 @@ def test_multi_index(self, in_dataset, model):
         out_dataset = driver.get_results()
 
         pd.testing.assert_index_equal(out_dataset.indexes["dummy"], midx)
+
+
+def test_finalize_always_called():
+    @xs.process
+    class P:
+        var = xs.variable(intent="out")
+
+        def initialize(self):
+            self.var = "initialized"
+            raise RuntimeError()
+
+        def finalize(self):
+            self.var = "finalized"
+
+    model = xs.Model({"p": P})
+    in_dataset = xs.create_setup(model=model, clocks={"clock": [0, 1]})
+    driver = XarraySimulationDriver(in_dataset, model)
+
+    try:
+        driver.run_model()
+    except RuntimeError:
+        pass
+
+    assert model.state[("p", "var")] == "finalized"