Add marimo progress bar for pymc sampler #7883
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #7883 +/- ##
==========================================
- Coverage 92.94% 92.83% -0.12%
==========================================
Files 116 116
Lines 18845 18900 +55
==========================================
+ Hits 17516 17545 +29
- Misses 1329 1355 +26
🚀 New features to boost your workflow:
|
|
Don't love this. Is there any plan for rich to work with marimo? |
rich shut it down |
|
Also we use |
| def compute_draw_speed(elapsed, draws): | ||
| speed = draws / max(elapsed, 1e-6) | ||
|
|
||
| if speed > 1 or speed == 0: | ||
| unit = "draws/s" | ||
| else: | ||
| unit = "s/draws" | ||
| speed = 1 / speed | ||
|
|
||
| return speed, unit | ||
|
|
||
|
|
||
| def create_rich_progress_bar(full_stats, step_columns, progressbar, progressbar_theme): |
There was a problem hiding this comment.
Any reason why these aren't methods anymore? It's not like they are used outside of ProgressBarManager (compute_draw_speed was static method anyway)
There was a problem hiding this comment.
ProgressBarManager seemed to be doing too much. The second one was specific to the rich implementation
There was a problem hiding this comment.
This still seems like a function that would ever only be called by ProgressBarManager, the inputs are way too specific
There was a problem hiding this comment.
And the object that is created is also called very specifically by ProgressBarManager
There was a problem hiding this comment.
ideally it would be passed in as many of the init kwargs of the ProgressBarManager are specific to this type of ProgressBar
There was a problem hiding this comment.
What would be passed in? I was just arguing against this function being moved out of the class, because it seems tightly coupled to it anyway
CustomProgress isn't changed and would still work. The ProgressBarManager now excepts other ProgressBar protocols |
| def update(self, task_id, **kwargs): | ||
| """Update the task with the given ID with the provided keyword arguments.""" | ||
| if self.bar.progress.current >= cast(int, self.bar.progress.total): | ||
| return | ||
|
|
||
| self.divergences[task_id.chain_idx] = kwargs.get("divergences", 0) | ||
|
|
||
| total_divergences = sum(self.divergences.values()) | ||
|
|
||
| update_kwargs = {} | ||
| if total_divergences > 0: | ||
| word = "draws" if total_divergences > 1 else "draw" | ||
| update_kwargs["subtitle"] = f"{total_divergences} diverging {word}" | ||
|
|
||
| self.bar.progress.update(**update_kwargs) |
There was a problem hiding this comment.
We changed the logic to not be obsessed with NUTS.
There is now a failing variable that the step methods report back, and that's passed to update to indicate sampling is failing. Also each step decides what is a meaningful stat to report.
This reverts back to over-specializing for nuts
There was a problem hiding this comment.
Does failing report the number of divergences? I can remove this logic and just not have a subtitle which would be specific to NUTS
There was a problem hiding this comment.
Failing is just binary, we use for coloring. The relevant summary stats are reported as part of all_step_stats, we don't treat them specially.
Line 517 in 886b584
Other methods don't use ProgressBarManager, but CustomProgress, which is a rich object |
|
Such as Lines 686 to 688 in 011fb35 |
That is using rich (inherited) I was focusing on the NUTS sampler progress bar specifically but can broaden to the where |
|
Furthermore, is there a way to have multiple progress_bars, color them manually? Fine if not, but keeping things similar between backends may help with code complexity |
You can, but it takes but a bit of real estate so I opted for the single one to be simple. I'd think that'd be determined by the |
Yeah and that's a valid strategy. I'm raising it in case that also solving the CustomProgress fixes more cases / provides a solution that is different than the one you took here specifically for the ProgressBarManager |
Description
The rich progress bars don't work well in marimo notebooks. Here is a basic implementation using their progress bar component to
track the
pymcsampler.An implementation note: tasks being attached to the progress bar and progress bar manager is a bit confusing, but I tried to keep as lean as possible to get to work.
Screen.Recording.2025-07-30.at.10.56.29.PM.mov
Related Issue
Checklist
Type of change
📚 Documentation preview 📚: https://pymc--7883.org.readthedocs.build/en/7883/