I4087-multiprocessing-2

[martinjrobins]

Description

• idaklu solver can run multiple simulations in parallel using openmp
• Previously this was done using Python multiprocessing, but this has been removed and instead a list of inputs can be passed directly to the _integrate function of the idaklu solver, the solver then runs these simulations in parallel using openmp

Partially fixes #4087 (just for IDAKLU solver). This is a redo of some of #4260 (plan is to split this old PR into multiple smaller PRs, starting with the openmp parallisation of the IDAKLU solver).

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)

Key checklist:

• No style issues: $ pre-commit run (or $ nox -s pre-commit) (see CONTRIBUTING.md for how to set this up to run automatically when committing locally, in just two lines of code)
• All tests pass: $ python run-tests.py --all (or $ nox -s tests)
• The documentation builds: $ python run-tests.py --doctest (or $ nox -s doctests)

You can run integration tests, unit tests, and doctests together at once, using $ python run-tests.py --quick (or $ nox -s quick).

Further checks:

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

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.46%. Comparing base (f47c12e) to head (727efe2).
Report is 186 commits behind head on develop.

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #4449      +/-   ##
===========================================
+ Coverage    99.42%   99.46%   +0.04%     
===========================================
  Files          292      292              
  Lines        22340    22350      +10     
===========================================
+ Hits         22211    22230      +19     
+ Misses         129      120       -9     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[agriyakhetarpal]

@martinjrobins, I'm unsure how the Homebrew OpenMP detection on macOS was working before this PR (my guess is that CVODE was able to find it on its own, but maybe since we now use the header directly, we need to pass its path?). Could you also trigger the wheel build from this branch? We use a different libomp.dylib there rather than the one from Homebrew, so that the wheels don't receive a higher macOS deployment target.

[MarcBerliner]

I'm unsure how the Homebrew OpenMP detection on macOS was working before this PR (my guess is that CVODE was able to find it on its own, but maybe since we now use the header directly, we need to pass its path?). Could you also trigger the wheel build from this branch? We use a different libomp.dylib there rather than the one from Homebrew, so that the wheels don't receive a higher macOS deployment target.

I think I'm running into this issue. I'm trying to install this on macOS, and I get the error

/PyBaMM/src/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverGroup.cpp:2:10: fatal error: 'omp.h' file not found
      #include <omp.h>

[agriyakhetarpal]

@MarcBerliner, I think passing -I$(brew --prefix libomp)/include should work

[agriyakhetarpal]

CXXFLAGS="-I $(brew --prefix libomp)/include" nox -s dev works for me locally. I guess we should set OpenMP_ROOT (or CXXFLAGS) in the noxfile when running on macOS, so that we get passing builds for both CI tests and locally.

[MarcBerliner]

Thanks @martinjrobins, this will be really nice for optimization.

One high-level question: how does this handle solver convergence errors? If we batch several simulations, we don't want a single ill-conditioned set of parameters to ruin the whole batch with an exception. Do we always receive a solution for the simulations that converge?

Also, this fix to the noxfile.py solved my local build issue with macOS (thanks @agriyakhetarpal):

PYBAMM_ENV = {
    "LD_LIBRARY_PATH": f"{homedir}/.local/lib",
    "PYTHONIOENCODING": "utf-8",
    "MPLBACKEND": "Agg",
    # Expression evaluators (...EXPR_CASADI cannot be fully disabled at this time)
    "PYBAMM_IDAKLU_EXPR_CASADI": os.getenv("PYBAMM_IDAKLU_EXPR_CASADI", "ON"),
    "PYBAMM_IDAKLU_EXPR_IREE": set_iree_state(),
    "IREE_INDEX_URL": os.getenv(
        "IREE_INDEX_URL", "https://iree.dev/pip-release-links.html"
    ),
    "CXXFLAGS": f"-I {os.popen('brew --prefix libomp').read().strip()}/include",
}

[MarcBerliner]

This is outside the scope of this PR, but we should be more careful about this in pybamm. We transpose the y when we store it with the python Solution object, so it's f-contiguous there. Not a big deal, but it'd make life simpler to keep it consistent.

[martinjrobins]

Agreed, I had to write the code assuming a particular ordering hence the specification as c_style. If an f_style ordered array is passed in then it gets converted to c_style by pybind11 (according to their docs: https://pybind11.readthedocs.io/en/stable/advanced/pycpp/numpy.html). This isn't ideal since it would involve a copy, but I wasn't sure about a better way of doing it. Let me know if you have any ideas

[MarcBerliner]

Interesting, I didn't realize pybind11 made a copy here. IMO, I think we should stick with c_style and make python-side changes to y to avoid the transpose and change its ordering from y[time, var] to y[var, time]

[martinjrobins]

CXXFLAGS="-I $(brew --prefix libomp)/include" nox -s dev works for me locally. I guess we should set OpenMP_ROOT (or CXXFLAGS) in the noxfile when running on macOS, so that we get passing builds for both CI tests and locally.

Yea, I put this into the CI and it works fine (see last commit). Let me know if that is the best way to do it?

[martinjrobins]

One high-level question: how does this handle solver convergence errors?

It should be the same behaviour as currently. If the retval from IDASolve indicates it has failed the individual solver will stop early and return. All the other solves, if they converge ok, carry on and solve for the complete duration and return normally. The retval is stored in the solution so the user can tell if that individual solve has failed or not.

[agriyakhetarpal]

CXXFLAGS="-I $(brew --prefix libomp)/include" nox -s dev works for me locally. I guess we should set OpenMP_ROOT (or CXXFLAGS) in the noxfile when running on macOS, so that we get passing builds for both CI tests and locally.

Yea, I put this into the CI and it works fine (see last commit). Let me know if that is the best way to do it?

Yes. Though, this would fail if someone tries to install without our nox -s dev helper (i.e., with pip install -e .). Could we include this in our CMake subprocess call in setup.py? That should work everywhere.

I'm off to a conference this week, so I haven't re-reviewed #4352 fully yet. We'll have to think how to do it there, because setup.py would be removed with that PR.

[MarcBerliner]

It should be the same behaviour as currently. If the retval from IDASolve indicates it has failed the individual solver will stop early and return. All the other solves, if they converge ok, carry on and solve for the complete duration and return normally. The retval is stored in the solution so the user can tell if that individual solve has failed or not.

It looks like we throw an error for any simulation failure in _post_process_solution, so the user wouldn't be able to examine the retvals

PyBaMM/src/pybamm/solvers/idaklu_solver.py

Lines 830 to 835 in 5087673

	if sol.flag == 2:
	termination = "event"
	elif sol.flag >= 0:
	termination = "final time"
	else:
	raise pybamm.SolverError(f"FAILURE {self._solver_flag(sol.flag)}")

We could change this to a warning (or make error vs. warning a solver option).

[martinjrobins]

We could change this to a warning (or make error vs. warning a solver option).

hmmm, this would be a big change to the API if someone was relying on the solve throwing an error if it failed. Plus the other solvers raise errors as well (even for multiple input parameters). I agree it would sometime be better to carry on if a solve fails though, but I'd say this would be a different issue (and probably one to discuss at the developers meeting)

[martinjrobins]

Yes. Though, this would fail if someone tries to install without our nox -s dev helper (i.e., with pip install -e .). Could we include this in our CMake subprocess call in setup.py? That should work everywhere.

I'll see if I can do it in CMakeLists.txt, since we're getting rid of setup.py

[MarcBerliner]

hmmm, this would be a big change to the API if someone was relying on the solve throwing an error if it failed. Plus the other solvers raise errors as well (even for multiple input parameters). I agree it would sometime be better to carry on if a solve fails though, but I'd say this would be a different issue (and probably one to discuss at the developers meeting)

Yeah, if this behavior is consistent with the other solvers' multiprocessing, then let's table this for now and bring this up at the meeting. Otherwise LGTM (besides the build issue)

[kratman]

@martinjrobins Did you trigger the wheel build to make sure that works?