Enable pickling of models in "python" format #1298

tlestang · 2020-12-14T15:27:30Z

Description

This defines the two methods __getstate__ and __setstate__ that can be used to control what is pickled when pickling a class
instance, as well as what happens when the object in unpickled.

In this case this is useful to avoid pickling the _evaluate attribute, which is a reference to a function that cannot be pickled.
This attribute is restored when unpicking by re-compiling the definition of the function evaluate.

Note that this does not make it possible to pickle model in "jax" format.

Fixes #1283 #734

Type of change

Please add a line in the relevant section of CHANGELOG.md to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

New feature (non-breaking change which adds functionality)
Optimization (back-end change that speeds up the code)
Bug fix (non-breaking change which fixes an issue)

Key checklist:

No style issues: $ flake8
All tests pass: $ python run-tests.py --unit
The documentation builds: $ cd docs and then $ make clean; make html

You can run all three at once, using $ python run-tests.py --quick.

Further checks:

Code is commented, particularly in hard-to-understand areas
Tests added that prove fix is effective or that feature works

… to pickle

… python and casadi mode

…n jax format.

So that `inputs_list` is defined

codecov · 2020-12-14T16:29:32Z

Codecov Report

Merging #1298 (e5ae0a1) into issue-849-parallel-processing (35491ba) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@                      Coverage Diff                       @@
##           issue-849-parallel-processing    #1298   +/-   ##
==============================================================
  Coverage                          98.09%   98.09%           
==============================================================
  Files                                270      270           
  Lines                              15176    15189   +13     
==============================================================
+ Hits                               14887    14900   +13     
  Misses                               289      289

Impacted Files	Coverage Δ
pybamm/expression_tree/operations/evaluate.py	`98.15% <100.00%> (+0.06%)`	⬆️
pybamm/solvers/base_solver.py	`99.07% <100.00%> (+<0.01%)`	⬆️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 35491ba...e5ae0a1. Read the comment docs.

valentinsulzer

Thanks, neat fix :) I know you didn't introduce this but I just noticed lots of the comments in test_scipy_solver contradict the associated code ...

valentinsulzer · 2020-12-14T16:46:06Z

tests/unit/test_solvers/test_scipy_solver.py

+            # Create model
+            model = pybamm.BaseModel()
+            # Covert to casadi instead of python to avoid pickling of
+            # "EvaluatorPython" objects.


comment out of date

Thanks, good catch

valentinsulzer · 2020-12-14T16:46:22Z

tests/unit/test_solvers/test_scipy_solver.py

+            model.rhs = {var: -pybamm.InputParameter("rate") * var}
+            model.initial_conditions = {var: 1}
+            # No need to set parameters; can use base discretisation (no spatial
+            # operators)


comment out of date

valentinsulzer · 2020-12-14T16:48:27Z

tests/unit/test_solvers/test_scipy_solver.py

@@ -258,7 +298,7 @@ def test_model_solver_multiple_inputs_discontinuity_error(self):
        domain = ["negative electrode", "separator", "positive electrode"]
        var = pybamm.Variable("var", domain=domain)
        model.rhs = {var: -pybamm.InputParameter("rate") * var}
-        model.initial_conditions = {var: 1}
+        model.initial_conditions = {var: 2 * pybamm.InputParameter("rate")}
        # No need to set parameters; can use base discretisation (no spatial
        # operators)


comment out of date

valentinsulzer · 2020-12-14T16:48:43Z

tests/unit/test_solvers/test_scipy_solver.py

        ):
            solver.solve(model, t_eval, inputs=inputs_list, nproc=2)

-    def test_model_solver_multiple_inputs_initial_conditions_error(self):
+    def test_model_solver_multiple_inputs_jax_format_error(self):
        # Create model
        model = pybamm.BaseModel()
        # Covert to casadi instead of python to avoid pickling of
        # "EvaluatorPython" objects.


comment out of date

pybamm/expression_tree/operations/evaluate.py

martinjrobins

looks great @tlestang , happy to merge. Had a couple of points below.

martinjrobins · 2020-12-16T11:54:15Z

pybamm/expression_tree/operations/evaluate.py

@@ -507,6 +508,23 @@ def evaluate(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None):
        else:
            return result

+    def __getstate__(self):


This looks great @tlestang, do you think we can use the same trick for the Jax evaluator?

I naively tried the same trick with the Jax evaluator but couldn't get it to work as smoothly.
Something in jax.jit can't be pickle, but I don't know what.

AttributeError: Can't pickle local object 'jit.<locals>.f_jitted'

I tried not pickling attributes jit_evaluate and jac_evaluate in the same way, i.e.

def __getstate__(self): state = self.__dict__.copy() del state["_evaluate_jax"] del state["_jit_evaluate"] del state["_jac_evaluate"] return state

but still get the AttributeError..

martinjrobins · 2020-12-16T12:38:06Z

pybamm/solvers/base_solver.py

@@ -589,6 +589,13 @@ def solve(
            for inputs in inputs_list
        ]

+        # Cannot use multiprocessing with model in "jax" format


Note that Jax is designed to support parallel solution of models with different inputs, so this should be doable. But it can be left for a PR later on

Okay good, would be better to rely on Jax's native behavior anyway :)

tlestang · 2020-12-16T15:25:13Z

Thanks @tinosulzer and @martinjrobins for the review. I will update the comments but before we merge that PR I'd like to investigate why the content of __getstate__ isn't tested..

…red by Pool

tlestang · 2021-01-06T09:15:35Z

This brought coverage to 100%. Now waiting to merge #1319 to fix CI on MacOS

tlestang added 3 commits December 14, 2020 11:20

Add methods getstate and setstate to avoid including method _evaluate…

dab096d

… to pickle

Test solving multiple inputs using multiprocessing with model in both…

35b26dd

… python and casadi mode

Raise SolverError if attempt to solve for list of inputs with model i…

35db29d

…n jax format.

tlestang requested review from martinjrobins and valentinsulzer December 14, 2020 15:27

relocate check for jax format

d3e132a

So that `inputs_list` is defined

tlestang mentioned this pull request Dec 14, 2020

Solving for an ensemble of input parameters in parallel #1261

Merged

8 tasks

valentinsulzer approved these changes Dec 14, 2020

View reviewed changes

martinjrobins approved these changes Dec 16, 2020

View reviewed changes

tlestang added 2 commits January 5, 2021 20:30

Add call to close() and join() for coverage to pick up processes fork…

43da47a

…red by Pool

Enable multiprocessing mode for coverage

7729e87

Remove out of date comments

e5ae0a1

tlestang merged commit b273887 into issue-849-parallel-processing Jan 6, 2021

valentinsulzer mentioned this pull request Jan 10, 2021

Cannot pickle model when format is "python" #1283

Closed

valentinsulzer deleted the issue-1283-pickle-python-format branch September 13, 2021 15:50

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enable pickling of models in "python" format #1298

Enable pickling of models in "python" format #1298

tlestang commented Dec 14, 2020

codecov bot commented Dec 14, 2020 •

edited

Loading

valentinsulzer left a comment

valentinsulzer Dec 14, 2020

tlestang Dec 14, 2020

valentinsulzer Dec 14, 2020

valentinsulzer Dec 14, 2020

valentinsulzer Dec 14, 2020

martinjrobins left a comment

martinjrobins Dec 16, 2020

tlestang Dec 16, 2020

martinjrobins Dec 16, 2020

tlestang Dec 16, 2020

tlestang commented Dec 16, 2020

tlestang commented Jan 6, 2021

Enable pickling of models in "python" format #1298

Enable pickling of models in "python" format #1298

Conversation

tlestang commented Dec 14, 2020

Description

Type of change

Key checklist:

Further checks:

codecov bot commented Dec 14, 2020 • edited Loading

Codecov Report

valentinsulzer left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

martinjrobins left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tlestang commented Dec 16, 2020

tlestang commented Jan 6, 2021

codecov bot commented Dec 14, 2020 •

edited

Loading