Add SafeSolver that uses pebble

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+- [#948](https://github.com/pybop-team/PyBOP/pull/948) - Adds `SafeSolver`, which uses pebble multiprocessing with a timeout option.
 - [#918](https://github.com/pybop-team/PyBOP/pull/918) - Adds a plot for predictions sampled from a posterior distribution (`pybop.plot.predictive`).
 - [#940](https://github.com/pybop-team/PyBOP/pull/940) - Adds support for Python 3.14 (EP-BOLFI optimiser and PyProBE still restricted to Python 3.12 or below).
 

pybop/pybamm/__init__.py

@@ -2,4 +2,4 @@
 from .eis_simulator import EISSimulator
 from .parameter_utils import set_formation_concentrations, cell_mass, cell_volume
 from .design_variables import add_variable_to_model
-from .utils import RecommendedSolver, SymbolReplacer
+from .utils import RecommendedSolver, SafeSolver, SymbolReplacer

pybop/pybamm/utils.py

@@ -1,7 +1,9 @@
 import multiprocessing as mp
 import platform
+from concurrent.futures import TimeoutError as CFTimeoutError
 
 import pybamm
+from pebble import ProcessPool
 
 
 class RecommendedSolver(pybamm.IDAKLUSolver):
@@ -22,6 +24,88 @@ def __init__(self, output_variables: list[str] | None = None):
         )
 
 
+class SafeSolver(pybamm.CasadiSolver):
+    """
+    A version of PyBaMM's CasadiSolver with a timeout option.
+
+    Additional parameters
+    ---------------------
+    timeout : float, optional
+        If timeout is a positive number, simulations are terminated after timeout
+        seconds if not completed successfully within this time. Default is None.
+    """
+
+    def __init__(self, timeout: float | None = None, **kwargs):
+        super().__init__(**kwargs)
+        self.timeout = timeout
+
+    def _integrate(
+        self,
+        model: pybamm.BaseModel,
+        t_eval,
+        inputs_list: list[dict] | None = None,
+        t_interp=None,
+        nproc=1,
+    ):
+        """
+        Solve a DAE model defined by residuals with initial conditions y0.
+
+        Parameters
+        ----------
+        model : :class:`pybamm.BaseModel`
+            The model whose solution to calculate.
+        t_eval : numeric type
+            The times at which to compute the solution
+        inputs_list : list of dict, optional
+            Any input parameters to pass to the model when solving
+        """
+
+        inputs_list = inputs_list or [{}]
+
+        ninputs = len(inputs_list)
+        if ninputs == 1 and self.timeout is None:
+            new_solution = self._integrate_single(
+                model,
+                t_eval,
+                inputs_list[0],
+                model.y0_list[0],
+            )
+            new_solutions = [new_solution]
+        else:
+            with ProcessPool(
+                context=mp.get_context(self._mp_context),
+                max_workers=nproc or mp.cpu_count(),
+            ) as p:
+                model_list = [model] * ninputs
+                t_eval_list = [t_eval] * ninputs
+                y0_list = model.y0_list
+
+                futures = p.map(
+                    self._integrate_single,
+                    model_list,
+                    t_eval_list,
+                    inputs_list,
+                    y0_list,
+                    timeout=self.timeout,
+                )
+                iterator = futures.result()
+
+                new_solutions = []
+                while True:
+                    try:
+                        new_solutions.append(next(iterator))
+                    except StopIteration:
+                        break
+                    except (TimeoutError, CFTimeoutError) as e:
+                        raise pybamm.SolverError(
+                            f"Timeout after {e.args[1]:.1f} seconds."
+                        ) from e
+                    except Exception as e:
+                        raise pybamm.SolverError(str(e)) from None
+
+        return new_solutions
+
+
 class SymbolReplacer:
     """
     Helper class to replace all instances of one or more symbols in an expression tree

pyproject.toml

@@ -33,6 +33,7 @@ dependencies = [
     "numpy>=1.26",
     "scipy>=1.12",
     "pints>=0.6.0",
+    "pebble",
 ]
 
 [project.optional-dependencies]

tests/unit/test_simulator.py

@@ -78,3 +78,29 @@ def test_set_output_variables(self):
             ValueError, match="Not a variable is not a variable in the model."
         ):
             simulator.set_output_variables(["Not a variable"])
+
+    def test_safe_solver(self):
+        model = pybamm.lithium_ion.DFN()
+
+        # Test SafeSolver on short simulations
+        experiment = pybamm.Experiment(["Discharge at C/10 for 10 seconds"])
+        simulator = pybop.pybamm.Simulator(
+            model, protocol=experiment, solver=pybop.pybamm.SafeSolver()
+        )
+        solution = simulator.solve()
+        assert isinstance(solution, pybamm.Solution)
+
+        simulator = pybop.pybamm.Simulator(
+            model, protocol=experiment, solver=pybop.pybamm.SafeSolver(timeout=10)
+        )
+        solution = simulator.solve()
+        assert isinstance(solution, pybamm.Solution)
+
+        # Test SafeSolver timeout on long simulation
+        experiment = pybamm.Experiment(["Discharge at C/10 for 10 hours"])
+        simulator = pybop.pybamm.Simulator(
+            model, protocol=experiment, solver=pybop.pybamm.SafeSolver(timeout=0.1)
+        )
+        simulator.debug_mode = True
+        with pytest.raises(pybamm.SolverError, match="Timeout"):
+            simulator.solve()