Pool improvements#506
Merged
Merged
Conversation
Collaborator
segasai
commented
May 5, 2026
segasai
commented
- When restoring from a checkpoint, I try to set the queue_size based on the current pool size.
- Also improve autodetection of the pool size, as the multiprocessing.Pool supports _processes keyword
- Use chunksize=1 with dynesty's pool
Coverage Report for CI Build 25501343432Coverage increased (+0.04%) to 91.962%Details
Uncovered Changes
Coverage RegressionsNo coverage regressions found. Coverage Stats
💛 - Coveralls |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR improves how dynesty interacts with multiprocessing pools by centralizing pool/queue_size parsing, enhancing queue_size handling during checkpoint restore, and adjusting mapping behavior to better handle uneven task durations.
Changes:
- Move and extend pool/
queue_sizeparsing into a shared helper that can auto-detect pool size viapool._processesorpool.size. - Update checkpoint restore to reconfigure
mapper/pooland setqueue_sizeon restored samplers. - Update tests/docs/changelog and set
chunksize=1fordynesty.pool.Pool.map.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_resume.py | Adds a resume test exercising restore behavior when pool size differs. |
| tests/test_pool.py | Expands sampler coverage and adds a queue_size (specified vs auto) test. |
| py/dynesty/utils.py | Adds shared _parse_pool_queue and updates restore logic to set queue_size on restored samplers. |
| py/dynesty/pool.py | Forces chunksize=1 for the dynesty pool wrapper’s map. |
| py/dynesty/dynesty.py | Removes local _parse_pool_queue and imports the shared implementation. |
| docs/source/quickstart.rst | Adds guidance about restoring with different CPU counts/pool sizes. |
| docs/source/faq.rst | Updates pool/queue_size guidance for variable-duration likelihoods. |
| CHANGELOG.md | Documents the pool/restore/chunksize changes. |
Comments suppressed due to low confidence (1)
tests/test_pool.py:245
- This test creates a
multiprocessingpool but doesn't callterminator(pool)afterwards. The rest of this file explicitly callsterminator(pool)even inside context managers due to the coveragepy cleanup issue noted above; skipping it here can reintroduce hanging/coverage flakiness. Consider callingterminator(pool)afterrun_nestedfor consistency.
ctx = mp.get_context('spawn')
with ctx.Pool(2) as pool:
sampler = dynesty.NestedSampler(loglike_gau,
prior_transform_gau,
ndim,
nlive=nlive,
sample='rslice',
pool=pool,
queue_size=queue_size,
rstate=rstate)
sampler.run_nested(print_progress=printing, dlogz=10)
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+2159
to
+2161
| if queue_size_new != queue_size_old and queue_size_old != 1: | ||
| warnings.warn( | ||
| f'Restoring the sampler with queue_size {queue_size_old}') |
Comment on lines
+2243
to
+2244
| "`pool.size` or pool._processes has not been provided. " | ||
| "Please `define `pool.size` or specify `queue_size` " |
Comment on lines
+657
to
+658
| You should be careful when restoring the sampler on machine with different number of CPUs when using a pool. | ||
| We will still use the original queue_size unless it was 1 before. |
| **How to decide on the number of processes in a pool and how to set queue_size** | ||
|
|
||
| Assuming that you decided on the number of live-points K that you want to use and that the likelihood evaluation is not very quick, you should use as many processes as you can up to around K. The queue_size should be equal the number of processes. If you are using the the number of processes that M is smaller than K, you may want to use :math:`M=K//2` or :math:`M=K//3` i.e integer fractions. So if you are using 1024 live-points all powers of two up to 1024 would be good choice for the number of processes. | ||
| Assuming that you decided on the number of live-points K that you want to use and that the likelihood evaluation is not very quick, you should use as many processes as you can up to around K. The queue_size should be equal the number of processes. If you are using the the number of processes that M is smaller than K, you may want to use :math:`M=K//2` or :math:`M=K//3` i.e integer fractions. So if you are using 1024 live-points all powers of two up to 1024 would be good choice for the number of processes. If you have likelihood functions that take highly variable amount of time to evaluate, you may benefit from using the queue_size that is significantly larger than the number CPUs (but queue_size should be always < K) |
Comment on lines
+10
to
+12
| - When restoring the sampler with the pool, try to use an updated value of queue_size based on pool size | ||
| - Use chunksize=1 for dynesty pool as that is better behaved for queue_size>nthreads and unequal duration of function evaluations | ||
| - When starting dynesty with multiprocessing pool, I now try to use the _processes keyword to find how many CPUs it is. That should reduce the need for manual queue_size specification |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.