FEAT: Improve interactions of user pools

[ColmTalbot]

I noticed that it was difficult to pass a user-specified pool through run_sampler, and it isn't used at all in the post processing. This PR:

• unifies how to create a pool for Bilby with a context (a la with multiprocessing.pool()...)
• explicitly passes the user pool around including for post processing
• moves logic out of Sampler._setup_pool and Sampler._close_pool
• support MPI pools using schwimmbad. This change will effectively make parallel_bilby moot.

[mj-will]

Looks great, I look forward to trying it! I've added some initial comments but might need to have a second look.

[mj-will]

Is this safe to do? Depending on how the sampler uses the pool, are there settings where the pool the sampler has stored is not updated?

For example, if the sampler has constructed a likelihood using the pool.map from the initial input pool will this break?

[ColmTalbot]

I'm trying to work through what this would look like.

If the initial input pool is not None, all of the pool objects reference here should be the same, so I don't think that will make a difference.

One potential issue is that if a specific sampler implementation handles the pool internally by itself this will create two pools and then in the best case, we have a bunch of extra processes we don't need. I think this is what nessai does, so maybe we should game through that specific case.

[mj-will]

Coming back to this, would it make sense to have an update_pool method or similar? That could have some additional logic to prevent multiple pools etc.

[mj-will]

Here's an example of what I had in mind: https://github.com/mj-will/aspire-bilby/blob/main/src/aspire_bilby/plugin.py#L199

self.pool is used in to create the likelihood function for aspire.

Looking at this now, since it happens in run_sampler perhaps this isn't actually an issue in this specific case. What do you think?

[mj-will]

I think this would benefit from a doc-string with example usage.

[ColmTalbot]

I don't remember exactly what the behaviour for some of these questions is, so I'll go back and check, and ideally write docstrings about them.

[ColmTalbot]

I'm trying to work through what this would look like.

If the initial input pool is not None, all of the pool objects reference here should be the same, so I don't think that will make a difference.

One potential issue is that if a specific sampler implementation handles the pool internally by itself this will create two pools and then in the best case, we have a bunch of extra processes we don't need. I think this is what nessai does, so maybe we should game through that specific case.

[mj-will]

Forgot to submit my comments as a review, but see them above.

[mj-will]

I think it would be useful to allow users to specify their own (additional) initializer.

For nessai, there's another variable that needs to be global for the pool to work and that can't be intialized unless using the fork start method.

[ColmTalbot]

Does this mean that we would also need to be able to allow users to pass through arbitrary additional arguments. At that stage, is it easier to just describe how to implement an independent version?

[mj-will]

Yes, that's true.

Here's what I had to do for nessai:

• https://github.com/bilby-dev/nessai-bilby/blob/main/src/nessai_bilby/plugin.py#L29
• https://github.com/bilby-dev/nessai-bilby/blob/main/src/nessai_bilby/plugin.py#L362

It works fine but since it duplicates so much of the code, it needs to be kept up-to-date with the main code base.

[adivijaykumar]

Can we add a test for schwimmbad/MPI support?

[mj-will]

We should documentation for user pools and note that the pool will only work correctly in the initializer has been called.