make learners picklable

[basnijholt]

Description

I realized that in some cases it's very useful to pickle learners. For example to send it over the network when parallelizing code.

With these changes, most learners become picklable.

Checklist

• Fixed style issues using pre-commit run --all (first install using pip install pre-commit)
• pytest passed

Type of change

Check relevant option(s).

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

[basnijholt]

@akhmerov and @jbweston, I wonder if we should also pickle adaptive.__version__ and warn the user that if they are loading a learner that is pickled with another version, that there is no guarantee that stuff's correct?

let's do that in a future PR.

[codecov-io]

Codecov Report

Merging #264 into master will increase coverage by 0.66%.
The diff coverage is 94.23%.

@@            Coverage Diff             @@
##           master     #264      +/-   ##
==========================================
+ Coverage   79.52%   80.18%   +0.66%     
==========================================
  Files          32       33       +1     
  Lines        4425     4522      +97     
  Branches      815      819       +4     
==========================================
+ Hits         3519     3626     +107     
+ Misses        779      773       -6     
+ Partials      127      123       -4     

Impacted Files	Coverage Δ
adaptive/tests/test_pickling.py	88.67% <88.67%> (ø)
adaptive/learner/average_learner.py	83.72% <100.00%> (+4.97%)	⬆️
adaptive/learner/balancing_learner.py	75.00% <100.00%> (+1.22%)	⬆️
adaptive/learner/data_saver.py	90.00% <100.00%> (+1.76%)	⬆️
adaptive/learner/integrator_learner.py	91.37% <100.00%> (+3.42%)	⬆️
adaptive/learner/learner1D.py	92.81% <100.00%> (+0.48%)	⬆️
adaptive/learner/learner2D.py	79.28% <100.00%> (+1.26%)	⬆️
adaptive/learner/sequence_learner.py	85.52% <100.00%> (-0.19%)	⬇️
... and 4 more

---

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 a0b22ff...acc5400. Read the comment docs.

[akhmerov]

I think tests should be more strict with idempotency. For example the loss reported by the new learner should be exactly the same, the learner's response to .ask should be exactly the same, etc.

[basnijholt]

@akhmerov, good idea, I've added that in 54017f3.

[akhmerov]

@basnijholt any idea why tests fail?

[basnijholt]

The failure is intermittent and unrelated to these changes, also happens in master sometimes.

It seems like the distributed test keeps getting stuck on Windows and MacOS.

[akhmerov]

But it seems like serialization tests are failing?

[basnijholt]

Oh, you are right! I assumed it were the same failures I had seen before.

I think that is fixed by bc65213.

[jbweston]

Good work, and thanks for taking feedback on board.

The refactors I proposed are nice-to-haves, but I understand if we just want to get this merged.

OTOH we use an unseeded random number generator in such a way that I cannot see why the tests pass.

For me, explaining/fixing this is a prerequisite for merging.

[tlaeven]

Using a cached function doesn't work in this branch, is this intentional?

Minimal example:

import adaptive
from functools import lru_cache
adaptive.notebook_extension()

@lru_cache()
def g(x):
    return x**2

def f(x):
    return x-g(0)

learner = adaptive.Learner1D(f, [-1,1])

runner = adaptive.Runner(learner)
runner.live_info()

Raises an error:
loky.process_executor.BrokenProcessPool: A task has failed to un-serialize. Please ensure that the arguments of the function are all picklable.

Note

It does work when the cached function is imported from a module:

[basnijholt]

@tlaeven, this is actually unrelated to these changes. The code you posted never worked, and has nothing to do with Adaptive:

from functools import lru_cache
import loky
from concurrent.futures import ProcessPoolExecutor

@lru_cache()
def g(x):
    return x ** 2

def f(x):
    return x - g(0)

# Before
with ProcessPoolExecutor() as ex:
    fut = ex.submit(f, 0)
    try:
        fut.result()
    except Exception as e:
        print(f"ProcessPoolExecutor failed: {e}")

# Now
ex = loky.get_reusable_executor()
fut = ex.submit(f, 0)
try:
    fut.result()
except Exception as e:
    print(f"loky failed: {e}")

which prints:

ProcessPoolExecutor failed: A process in the process pool was terminated abruptly while the future was running or pending.
loky failed: A task has failed to un-serialize. Please ensure that the arguments of the function are all picklable.

I guess the difference is because of basnijholt/adaptive-scheduler#39.

[basnijholt]

@jbweston, I've fixed all that was brought up.

The LearnerND takes some more work, so I moved those commits to #272.

I think it's ready to merge. Would you take a final look?

[jbweston]

LGTM! Nice work.