Implement CustomProgress that does not output empty divs when disabled

[tomicapretto]

Description

This PR adds a new class CustomProgress that inherits from rich.progress.Progress. I overloaded the necessary methods to make it not output anything when we disable the progress bar. I feel it's a bit hacky so I would like to know what others think about this modification.

This problem is noticeable especially in VS Code. It's not such a big issue in Jupyter Notebook. Some screenshots:

VS Code progressbar=True (before and after we get the same result, I cannot get rid of those empty spaces)

VS Code progressbar=False before the changes

VS Code progressbar=False after the changes

Jupyter Notebook after the changes

Update

In the terminal, after the change

Related Issue

• Closes #7264

---

📚 Documentation preview 📚: https://pymc--7290.org.readthedocs.build/en/7290/

[tomicapretto]

Some tests are failing, it's because there's something I missed in the child class I created.

[tomicapretto]

@ricardoV94 this is ready to be reviewed

[ricardoV94]

@tomicapretto you seem to have picked commits from main, can you rebase? Hard to review otherwise

[tomicapretto]

@tomicapretto you seem to have picked commits from main, can you rebase? Hard to review otherwise

I think I did not the cleanest thing but I think it does not show anything from main now

[tomicapretto]

Converted to a draft because I found other points of improvement.

• I think the refresh argument is not needed and it causes the sampler to be slower
• There should be some kind of final update when computing predictions, see the bar is not full:

[fonnesbeck]

The refresh argument was added because (1) without it progress bars did not always end at 100% after sampling was complete (as you noted above) and (2) the divergences text did not update without it.What about adding a single refresh update at the end of sampling? Perhaps that would ensure the last update is complete.

You can use the example in #7278 to test for (2).

[tomicapretto]

@fonnesbeck I'll test some ideas / combinations and update on my findings

[drbenvincent]

FYI: This has also affected doctests in CausalPy, see https://github.com/pymc-labs/CausalPy/issues/323 Would be very pleased to see a fix :)

[tomicapretto]

Before the last commit, it wouldn't update to 100% using one chain (because it doesn't use multiprocess sampling)

Now it works well

I can confirm divergences are updated and it goes up to 100%

The progressbar for posterior predictive sampling was also broken:

Now it works fine:

[tomicapretto]

I'm now trying to see if I can get those two extra divs not added. It does not take any relevant space when using jupyter notebook, but it appears on VS Code and it's really bothering.

[tomicapretto]

Ok I have some findings. It has to do with how rich handles the output in jupyter notebooks. Console has a force_jupyter parameter which is None by default (and with notebooks it's set to True internally). See what happens when we change it.

In VS Code, it removes the extra space that gets added, but notice it changes the color and the description:

In VS Code, it doesn't do any harm, but it stillchanges the color and the description:

I'm not sure if this is something we can take care of. It's really bothering for the ones who use VS Code, but I don't see any simple solution. I'll do a bit more research but I'm not going to invest a lot of time on it.

EDIT

It's not that it does not produce the same output on both platforms. It does produce the same output and it's this:

<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"></pre>

But I guess Jupyter Notebooks must be using some CSS that causes the height of the output to be very small, while VS Code does not.

[ricardoV94]

Not to be a bummer, but should we revert to the old progessbar? @fonnesbeck @zaxtax ?

[tomicapretto]

Not to be a bummer, but should we revert to the old progessbar? @fonnesbeck @zaxtax ?

I'm not sure what was the reason to move away from the old one. But definitely the goodness from rich are not for free.

[fonnesbeck]

I can dig into it a bit later today. If it comes down to color differences, I'm fine with it. If its deeper than that then we probably should seek alternatives (including reverting).

[aseyboldt]

As I noticed in nutpie, it is also surprisingly easy to build a custom progress bar. We could also adapt something like the nutpie one into pymc if that makes sense. It won't work (without modifications) in the terminal however.

[fonnesbeck]

@aseyboldt you are using fastprogress in nutpie, correct?

[aseyboldt]

Not anymore. This included an updated progressbar: https://github.com/pymc-devs/nutpie/pull/105

It's not released yet though. The code is around here: https://github.com/pymc-devs/nutpie/blob/main/python/nutpie/sample.py#L308

[fonnesbeck]

After testing this for a bit, I'm inclined to blame VS Code here rather than rich. The Python interactive pane is very buggy, not just here but in a variety of settings. Buttons on the toolbar often do not work, the input field occasionally gets resized so you can't read the text in it, and so on. Its VS Code's responsibility to have its Jupyter interface behave like a Jupyter notebook, and it often does not.

[fonnesbeck]

In the meantime, should we at least go ahead and merge some version of this PR? It clearly fixes things for a lot of people.

[tomicapretto]

In the meantime, should we at least go ahead and merge some version of this PR? It clearly fixes things for a lot of people.

I agree. I've seen the speed being a lot worse just because of the progressbar. So this is already something.

[fonnesbeck]

The outstanding failure appears to be related to Jax versioning(?) Its complaining about using max as an argument to clip instead of a_max.

[ricardoV94]

The outstanding failure appears to be related to Jax versioning(?) Its complaining about using max as an argument to clip instead of a_max.

It's blackjax fault: https://github.com/pymc-devs/pymc/pull/7317

[fonnesbeck]

@tomicapretto can you rebase this to get the fix in #7320 ?

[tomicapretto]

@tomicapretto can you rebase this to get the fix in #7320 ?

@fonnesbeck I think rebased correctly, but I'm not 100% as I don't find myself doing it often.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 92.50%. Comparing base (0216473) to head (292dae9). Report is 138 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7290      +/-   ##
==========================================
+ Coverage   91.84%   92.50%   +0.65%     
==========================================
  Files         102      102              
  Lines       17149    17177      +28     
==========================================
+ Hits        15751    15889     +138     
+ Misses       1398     1288     -110     

Files	Coverage Δ
pymc/backends/arviz.py	96.61% <100.00%> (ø)
pymc/sampling/forward.py	95.96% <100.00%> (+0.03%)	:arrow_up:
pymc/sampling/mcmc.py	88.00% <100.00%> (+0.02%)	:arrow_up:
pymc/sampling/parallel.py	88.50% <100.00%> (ø)
pymc/sampling/population.py	73.71% <100.00%> (ø)
pymc/smc/sampling.py	99.20% <100.00%> (ø)
pymc/tuning/starting.py	91.22% <100.00%> (ø)
pymc/util.py	81.39% <100.00%> (+1.99%)	:arrow_up:
pymc/variational/inference.py	87.93% <100.00%> (ø)

... and 3 files with indirect coverage changes