Enable progress bar under pmap

[andrewdipper]

Enables proper progress bar behavior of progress_bar_scan under pmap. However the progress bars are not tied to physical devices since io_callback won't expose the device id. The bars behave as if they are physically tied but are always sorted by most to least complete (they're actually randomly shuffled but one bar will always be first and another last).

I haven't visually tested this under a multi-gpu setup as I don't have an identical pair of gpus - just using JAX_PLATFORMS='cpu'

#655 - Related: changes progress bar to tqdm and requires the user to label devices - a different approach

A few important guidelines and requirements before we can merge your PR:

• We should be able to understand what the PR does from its title only;
• There is a high-level description of the changes;
• There are links to all the relevant issues, discussions and PRs;
• The branch is rebased on the latest main commit;
• Commit messages follow these guidelines;
• The code respects the current naming conventions;
• Docstrings follow the numpy style guide
• pre-commit is installed and configured on your machine, and you ran it before opening the PR;
• There are tests covering the changes;
• The doc is up-to-date;

[junpenglao]

Thanks! I like this better than #655 because it does not require explicit num_chain. @zaxtax thoughts?

[zaxtax]

I don't like how the bars get shuffled. I worry that it would be hard to follow the progress. Maybe the approaches can be combined somehow

…

On Sun, 4 Aug 2024, 06:45 Junpeng Lao, ***@***.***> wrote: Thanks! I like this better than #655 <#655> because it does not require explicit num_chain. @zaxtax <https://github.com/zaxtax> thoughts? — Reply to this email directly, view it on GitHub <#712 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUNIJYC34Z7BSLGMM5LZPWWUZAVCNFSM6AAAAABL5V26F6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDENRXGMYTMMRQHE> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[andrewdipper]

The initial random shuffle is a bit wonky. I believe it comes from fastprogress - the bars should always be in sorted order if we used the change from fastprogress to tqdm. That's based on similar tests with numpyro's progressbar and tqdm - but I'll verify it. Would that be better?

[junpenglao]

The initial random shuffle is a bit wonky. I believe it comes from fastprogress - the bars should always be in sorted order if we used the change from fastprogress to tqdm. That's based on similar tests with numpyro's progressbar and tqdm - but I'll verify it. Would that be better?

Yes - if that works, how about I merge this PR and then @zaxtax you can update #655?

[andrewdipper]

I had it wrong - the randomization was a race condition. Added a lock and now the bars are always ordered by most to least complete

[zaxtax]

But do the bars move around if one chain overtakes another?

…

On Mon, 5 Aug 2024, 20:44 andrewdipper, ***@***.***> wrote: I had it wrong - the randomization was a race condition. Added a lock and now the bars are always ordered by most to least complete — Reply to this email directly, view it on GitHub <#712 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCULA6YA2N4AI3UANSQDZP7BYVAVCNFSM6AAAAABL5V26F6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDENRZGY4DKOJSGI> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[andrewdipper]

No, the order will always be in decreasing order of chain progress. As such the chain to progress bar mapping will change if one overtakes another. We'd need the equivalent of a chain/device id to avoid this - but from what I've seen so far the new callbacks don't support identifying the underlying device.

[junpenglao]

No, the order will always be in decreasing order of chain progress. As such the chain to progress bar mapping will change if one overtakes another. We'd need the equivalent of a chain/device id to avoid this - but from what I've seen so far the new callbacks don't support identifying the underlying device.

So IIUC, we cannot have identifier like Chain 0, Chain 1 ...? I am not too bothered by the sorting of the chains, but would be great if there is some identifier, so users are not surprised by the behavior.

[zaxtax]

In my PR, I handle this by making a chain id that's carried along with the sample iteration. It's a bit inelegant but straightforward and doesn't require special machinery from the io callback code

…

On Tue, 6 Aug 2024, 09:03 Junpeng Lao, ***@***.***> wrote: No, the order will always be in decreasing order of chain progress. As such the chain to progress bar mapping will change if one overtakes another. We'd need the equivalent of a chain/device id to avoid this - but from what I've seen so far the new callbacks don't support identifying the underlying device. So IIUC, we cannot have identifier like Chain 0, Chain 1 ...? I am not too bothered by the sorting of the chains, but would be great if there is some identifier, so users are not surprised by the behavior. — Reply to this email directly, view it on GitHub <#712 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUIH2POMBI7KYHPLZCDZQBYK5AVCNFSM6AAAAABL5V26F6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDENZQGUZTEMJWGU> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[junpenglao]

In my PR, I handle this by making a chain id that's carried along with the sample iteration. It's a bit inelegant but straightforward and doesn't require special machinery from the io callback code

Yeah but then it requires an explicit user input of chain_id - I would rather not pushing this handling to user, hence my preference to the current solution here.

[zaxtax]

We could make a helper that hides that plumbing?

…

On Tue, 6 Aug 2024, 11:22 Junpeng Lao, ***@***.***> wrote: In my PR, I handle this by making a chain id that's carried along with the sample iteration. It's a bit inelegant but straightforward and doesn't require special machinery from the io callback code Yeah but then it requires an explicit user input of chain_id - I would rather not pushing this handling to user, hence my preference to the current solution here. — Reply to this email directly, view it on GitHub <#712 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUJE5YQU5PFB37G4LWDZQCIWDAVCNFSM6AAAAABL5V26F6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDENZQHAYDMOJZHA> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[andrewdipper]

So IIUC, we cannot have identifier like Chain 0, Chain 1 ...? I am not too bothered by the sorting of the chains, but would be great if there is some identifier, so users are not surprised by the behavior.

Correct. Progress is being used as a surrogate id.

We could make a helper that hides that plumbing?

If I'm not mistaken I think this helper has to reside outside the pmap and wouldn't be completely transparent. But #655 is more straightforward from a handling perspective.

It might be possible to use the io_callback return value to set a chain id to be used in future io_callbacks. And as a fallback hashing the initial (or whole sequence) of (input, carry) could resolve into a unique id. I'm not sure how friendly jax is to those solutions but I'll take a look.

[andrewdipper]

The latest commit uses the return value from the io_callback to uniquely identify the chains. The map from chains to progress bars won't change.

However, this changes the carry used in progress_bar_scan which is a problem if it is used externally. That'd have to be a slow change. But it does avoid forcing the user to do the device identification with part of the input. A wrapper over scan itself might clean things up a bit.

[zaxtax]

@junpenglao I'm now ok with the PR as is and definitely think we should explore a wrapper over scan in the future.

[junpenglao]

Great, thank you @andrewdipper for the contribution!

[zaxtax]

Oh should we update any examples to show how the new functionality works @andrewdipper ? Or does this not change the API?

[andrewdipper]

I didn't see any examples using progress_bar_scan and only the two uses within blackjax so we should be good.

[zaxtax]

No it looks like the examples in the sampling book don't https://blackjax-devs.github.io/sampling-book/algorithms/cyclical_sgld.html

…

On Thu, 8 Aug 2024, 17:34 andrewdipper, ***@***.***> wrote: I didn't see any examples using progress_bar_scan and only the two uses within blackjax so we should be good. — Reply to this email directly, view it on GitHub <#712 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCULMAXONNYETUY42UDTZQOFZVAVCNFSM6AAAAABL5V26F6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDENZWGEZDEMBXGI> . You are receiving this because you were mentioned.Message ID: ***@***.***>