Runtime signals (#161)

* wip implementation

* full implementation

Process runtime methods may also return signals.

* fix/update some tests

* black

* unrelated doc change

* add api docs

* fix handling CONTINUE signal

* add tests

* black

* update release notes

* update doc (general concepts) and docstrings

* doc: add user guide subsection

* typo

Author: benbovy

doc/api.rst

@@ -189,3 +189,4 @@ Model runtime monitoring
    monitoring.ProgressBar
    runtime_hook
    RuntimeHook
+   RuntimeSignal

doc/framework.rst

@@ -224,18 +224,26 @@ A model run is divided into four successive stages:
 3. finalize step
 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. 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()``,
-``.run_step()``, ``.finalize_step()`` and ``.finalize()`` for each
-stage above, respectively. Note that this is entirely optional. For
-example, time-independent processes (e.g., for setting model grids)
-usually implement stage 1 only. In a few cases, the role of a process
-may even consist of just declaring some variables that are used
-elsewhere.
+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. Stage 4 is always run even when an
+exception is raised during stage 1, 2 or 3.
+
+Each :func:`~xsimlab.process` decorated class may provide its own computation
+instructions by implementing specific "runtime" methods named ``.initialize()``,
+``.run_step()``, ``.finalize_step()`` and ``.finalize()`` for each stage above,
+respectively. Note that this is entirely optional. For example, time-independent
+processes (e.g., for setting model grids) usually implement stage 1 only. In a
+few cases, the role of a process may even consist of just declaring some
+variables that are used elsewhere.
+
+Runtime methods may be decorated by :func:`~xsimlab.runtime`. This is useful if
+one needs to access the value of some runtime-specific variables like the
+current step, time step duration, etc. from within those methods. Runtime
+methods may also return a :func:`~xsimlab.RuntimeSignal` to control the
+workflow, e.g., break the execution of the current stage.
+
+It is also possible to monitor and/or control simulations independently of any
+model, using runtime hooks. See Section :ref:`monitor`.
 
 Get / set variable values inside a process
 ------------------------------------------

doc/monitor.rst

@@ -16,25 +16,24 @@ it exemplifies how to create your own custom monitoring.
     sys.path.append('scripts')
     from advection_model import advect_model, advect_model_src
 
-The following imports are necessary for the examples below.
+Let's use the following setup for the examples below. It is based on the
+``advect_model`` created in Section :ref:`create_model`.
 
 .. ipython:: python
 
     import xsimlab as xs
 
-.. ipython:: python
-   :suppress:
-
     in_ds = xs.create_setup(
         model=advect_model,
         clocks={
-            'time': np.linspace(0., 1., 5),
+            'time': np.linspace(0., 1., 6),
         },
         input_vars={
             'grid': {'length': 1.5, 'spacing': 0.01},
             'init': {'loc': 0.3, 'scale': 0.1},
             'advect__v': 1.
         },
+        output_vars={'profile__u': 'time'}
     )
 
 
@@ -172,3 +171,41 @@ methods that may share some state:
 
    with PrintStepTime():
        in_ds.xsimlab.run(model=advect_model)
+
+
+Control simulation runtime
+--------------------------
+
+Runtime hook functions may return a :class:`~xsimlab.RuntimeSignal` so that you
+can control the simulation workflow (e.g., skip the current stage or process,
+break the simulation time steps) based on some condition or some computed value.
+
+In the example below, the simulation stops as soon as the gaussian pulse (peak
+value) has been advected past ``x = 0.4``.
+
+.. ipython::
+
+   In [2]: @xs.runtime_hook("run_step", "model", "post")
+      ...: def maybe_stop(model, context, state):
+      ...:     peak_idx = np.argmax(state[('profile', 'u')])
+      ...:     peak_x = state[('grid', 'x')][peak_idx]
+      ...:
+      ...:     if peak_x > 0.4:
+      ...:         print("Peak crossed x=0.4, stop simulation!")
+      ...:         return xs.RuntimeSignal.BREAK
+      ...:
+
+   In [3]: out_ds = in_ds.xsimlab.run(
+      ...:     model=advect_model,
+      ...:     hooks=[print_step_start, maybe_stop]
+      ...: )
+
+Even when a simulation stops early like in the example above, the resulting
+xarray Dataset still contains all time steps defined in the input Dataset.
+Output variables have fill (masked) values for the time steps that were not run,
+as shown below with the ``nan`` values for ``profile__u`` (fill values are not
+stored physically in the Zarr output store).
+
+.. ipython:: python
+
+   out_ds

doc/run_model.rst

@@ -106,20 +106,21 @@ IPython (Jupyter) magic commands
 
 Writing a new setup from scratch may be tedious, especially for big models with
 a lot of input variables. If you are using IPython (Jupyter), xarray-simlab
-provides convenient commands that can be activated with:
+provides helper commands that are available after loading the
+``xsimlab.ipython`` extension, i.e.,
 
 .. ipython:: python
 
    %load_ext xsimlab.ipython
 
 The ``%create_setup`` magic command auto-generates the
-:func:`~xsimlab.create_setup` code cell above from a given model:
+:func:`~xsimlab.create_setup` code cell above from a given model, e.g.,
 
 .. ipython:: python
 
-   %create_setup advect_model --default --comment
+   %create_setup advect_model --default --verbose
 
-The ``--default`` and ``--comment`` options respectively add default values found
+The ``--default`` and ``--verbose`` options respectively add default values found
 for input variables in the model and input variable description as line comments.
 
 Full command help:

doc/whats_new.rst

@@ -30,6 +30,8 @@ Enhancements
 - Added :func:`~xsimlab.group_dict` variable (:issue:`159`).
 - Added :func:`~xsimlab.global_ref` variable for model-wise implicit linking of
   variables in separate processes, based on global names (:issue:`160`).
+- Added :class:`~xsimlab.RuntimeSignal` for controlling simulation workflow from
+  process runtime methods and/or runtime hook functions (:issue:`161`).
 
 Bug fixes
 ~~~~~~~~~

xsimlab/__init__.py

@@ -11,6 +11,7 @@
     process,
     process_info,
     runtime,
+    RuntimeSignal,
     variable_info,
 )
 from .variable import (

xsimlab/drivers.py

@@ -5,6 +5,7 @@
 import pandas as pd
 
 from .hook import flatten_hooks, group_hooks, RuntimeHook
+from .process import RuntimeSignal
 from .stores import ZarrSimulationStore
 from .utils import get_batch_size
 
@@ -345,16 +346,27 @@ def _run(
 
             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)
+            signal = model.execute("run_step", rt_context, **execute_kwargs)
+
+            if signal == RuntimeSignal.BREAK:
+                break
 
             store.write_output_vars(batch, step, model=model)
 
-            model.execute("finalize_step", rt_context, **execute_kwargs)
+            # after writing output variables so that index positions
+            # are properly updated in store.
+            if signal == RuntimeSignal.CONTINUE:
+                continue
+
+            signal = model.execute("finalize_step", rt_context, **execute_kwargs)
+
+            if signal == RuntimeSignal.BREAK:
+                break
 
         store.write_output_vars(batch, -1, model=model)
         store.write_index_vars(model=model)
-    except Exception as error:
-        raise error
+    except Exception:
+        raise
     finally:
         model.execute("finalize", rt_context, **execute_kwargs)
 

xsimlab/hook.py

@@ -1,18 +1,17 @@
 import inspect
+from enum import Enum
 from typing import Callable, Dict, Iterable, List, Union
 
 from .process import SimulationStage
 
 
-__all__ = ("runtime_hook", "RuntimeHook")
-
-
 def runtime_hook(stage, level="model", trigger="post"):
     """Decorator that allows to call a function or a method
     at one or more specific times during a simulation.
 
     The decorated function / method must have the following signature:
-    ``func(model, context, state)`` or ``meth(self, model, context, state)``.
+    ``func(model, context, state)`` or ``meth(self, model, context, state)``. It
+    may return a :class:`RuntimeSignal` (optional).
 
     Parameters
     ----------