joshspeagle / dynesty

Dynamic Nested Sampling package for computing Bayesian posteriors and evidences
https://dynesty.readthedocs.io/
MIT License
347 stars 76 forks source link

Saving sampler state when running in parallel #165

Closed ajdittmann closed 4 years ago

ajdittmann commented 4 years ago

127

If you want to save the state of the sampler, just pickle-ing the whole thing should work.

Is there a way to do this when using a multiprocessing or MPI pool? I tried deleting the pool during getstate but that didn't work.

joshspeagle commented 4 years ago

Ooo, I had never considered this. I can see why this would cause issues. Technically the pool is defined both explicitly (in self.pool) and implicitly (via self.M) (see here). Removing both during __getstate__ might fix this (via the #130 solution). I can try and patch this for the next release.

ajdittmann commented 4 years ago

Removing M as well as pool seems to do the trick! I suggest that you add this version of __getstate__ to the dynamic nested sampler as well.

joshspeagle commented 4 years ago

Will do! Thanks for catching this.

kcroker commented 2 years ago

Following guidance in this thread, in dynesty 1.2.2, if I set dynesty_sampler.M to my current pool's map function on a de-pickled NestedSampler and then run_nested(), I except on a KeyError 'pool'. This gets triggered within evaluation of the M function. If I don't attempt to tell it about the new pool after de-pickling, it runs on a single core. Below is the stack trace. Any suggestions?

    for it, results in enumerate(
  File "/home/o/miniconda3/envs/sufficient/lib/python3.9/site-packages/dynesty/sampler.py", line 740, in sample
    u, v, logl, nc = self._new_point(loglstar_new)
  File "/home/o/miniconda3/envs/sufficient/lib/python3.9/site-packages/dynesty/sampler.py", line 360, in _new_point
    u, v, logl, nc, blob = self._get_point_value(loglstar)
  File "/home/o/miniconda3/envs/sufficient/lib/python3.9/site-packages/dynesty/sampler.py", line 343, in _get_point_value
    self._fill_queue(loglstar)
  File "/home/o/miniconda3/envs/sufficient/lib/python3.9/site-packages/dynesty/sampler.py", line 332, in _fill_queue
    self.queue = list(self.M(evolve_point, args))
  File "/home/o/miniconda3/envs/sufficient/lib/python3.9/multiprocessing/pool.py", line 364, in map
    return self._map_async(func, iterable, mapstar, chunksize).get()
  File "/home/o/miniconda3/envs/sufficient/lib/python3.9/multiprocessing/pool.py", line 771, in get
    raise self._value
  File "/home/o/miniconda3/envs/sufficient/lib/python3.9/multiprocessing/pool.py", line 537, in _handle_tasks
    put(task)
  File "/home/o/miniconda3/envs/sufficient/lib/python3.9/multiprocessing/connection.py", line 211, in send
    self._send_bytes(_ForkingPickler.dumps(obj))
  File "/home/o/miniconda3/envs/sufficient/lib/python3.9/multiprocessing/reduction.py", line 51, in dumps
    cls(buf, protocol).dump(obj)
  File "/home/o/miniconda3/envs/sufficient/lib/python3.9/site-packages/dynesty/utils.py", line 159, in __getstate__
    del state['pool']
KeyError: 'pool'
segasai commented 2 years ago

https://github.com/joshspeagle/dynesty/issues/371#issuecomment-1124400165 -- this has the solution. Also you should update to 1.2.3 which fixes a related problem. If you still have a in problem, please open a new issue.

kcroker commented 2 years ago

#371 (comment) -- this has the solution. Also you should update to 1.2.3 which fixes a related problem. If you still have a in problem, please open a new issue.

Thanks! #371 fixed the issue with reloading. There was a related issue with creating the pickle, but based on discussions here I was able to make it work by setting sampler.pool = None before pickling, and then reassigning the correct pool back after the pickle. Will upgrade to 1.2.3 soon xD