Cannot pickle model when format is "python"

Description

Instances of BaseModel and derived classes cannot be pickled when attribute convert_to_format is "python". This hampers use of multiprocessing, since it relies on pickling objects for passing them between processes, see #1261 .

To reproduce

import pickle
import pybamm

model = pybamm.lithium_ion.SPMe()
geometry = model.default_geometry
param = model.default_parameter_values
param.process_model(model)
param.process_geometry(geometry)
mesh = pybamm.Mesh(geometry, model.default_submesh_types, model.default_var_pts)
# discretise model
disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
disc.process_model(model)

model.convert_to_format = "python"
model.default_solver.set_up(model)

try:
    with open("model.pickle", mode="wb") as f:
	pickle.dump(model, f)
except pickle.PicklingError as e:
    print(e)

Can't pickle <function evaluate at 0x7f1c4c5d3280>: attribute lookup evaluate on pybamm.expression_tree.operations.evaluate failed

Details

Function base_solver.set_up does several things but in particular is creates the SolverCallable objects the solver relies on to evaluate the RHS, initial conditions and jacobian.

# base_solver.py

# Add the solver attributes
model.init_eval = init_eval # SolverCallable instance
model.rhs_eval = rhs_eval
model.algebraic_eval = algebraic_eval
model.jac_algebraic_eval = jac_algebraic
model.terminate_events_eval = terminate_events_eval
model.discontinuity_events_eval = discontinuity_events_eval

# Calculate initial conditions
model.y0 = init_eval(inputs)

When model.convert_to_format == "python", SolverCallable.__call__ is a wrapper around EvaluatorPython.evaluate() which evaluates the Python code generated in EvaluatorPython.__init__ to translate the expression tree. When an EvaluatorPython instance is created, a global function evaluate is defined by compiling its definition on the fly. This function is then pointed to by EvaluatorPython._evaluate.

According to the docs for Pickle, only functions defined at the top level of a module can be pickled and I think this is the reason why evaluate cannot be pickled.

Note that functions (built-in and user-defined) are pickled by “fully qualified” name reference, not by value. This means that only the function name is pickled, along with the name of the module the function is defined in.

Note that evaluate is a global function and not a local function in EvaluatorPython.__init__, i.e print(self._evaluate) evaluates to <function evaluate at 0x7fa4fee843a0> and not <function EvaluatorPython.__init__.<locals>.evaluate 0x7fa4fee843a0>. However since evaluate is compiled on the fly, I don’t think Pickle can find the name of the module it is defined in…

Minimal example:

import pickle

python_str = """def dummy_function():
    return "Hello"
self._method = dummy_function
"""

class myclass():
    def __init__(self):
	compiled_function = compile(python_str, "", "exec")
	exec(compiled_function)

inst = myclass()
with open("data.pickle", mode="w+b") as f:
    try:
	pickle.dump(inst, f)
    except pickle.PicklingError as e:
	print(e)

Can't pickle <function dummy_function at 0x7fc39202ec10>: attribute lookup dummy_function on __main__ failed

Models can be pickled fine when mode.convert_to_format == "casadi", since in this case SolverCallable.function calls an instance of Casadi.Function, which can pickled.

[martinjrobins]

couple of points on this:

1. should the base solver be instead storing these “compiled” models in the solver itself? The user probably isn’t expecting the solver to modify their model when you call set_up.
2. there is no particular reason that the compiled evaluate needs to be pickled, can we just tell pickle not to include it? The EvaluatorPython class can instead store a string of the function to be compiled (I think it might already do this), and if the compiled version doesn’t exist when EvaluatorPython.evaluate is called, compile it then.

[tinosulzer]

Fixed by #1298