Repository: pyro-ppl/pyro

Issue:
[FR] Use normal tqdm in notebooks by default
Currently, Pyro's MCMC uses tqdm's auto_notebook to display progress bar by default. However, I find that it causes more troubles than advantages:
+ auto_notebook uses jupyter widget, which requires users install widget extension.
+ the widget is only available at running time, which leaves some weird messages when sharing results or reopening it. For example, in [this gist](https://gist.github.com/fehiepsi/fa2f7980aef1e7242670a2470694ee2e), it leaves the line

```
HBox(children=(IntProgress(value=0, description='Warmup', max=1300, style=ProgressStyle(description_width='ini…
```

I tried to replace [this line](https://github.com/uber/pyro/blob/dev/pyro/infer/mcmc/logger.py#L7) by `from tqdm import tqdm` but it only works for `num_chains=1`. For multi-chain, the issue https://github.com/tqdm/tqdm/issues/630 still remains.


Discussion:
@fehiepsi - I remember I had changed that to tqdm at one point, for exactly this reason. We don't have a good alternative at this point. The following comes to mind:

 - [ ] Resort to using tqdm and not auto_notebook when in notebook environment: This gives a really bad output and a long scroll due to the issue that you linked above, and I doubt if there is a resolution for that. I think the tqdm devs will ask us to use auto_notebook instead.
 - [ ] Resort to using tqdm but disable progress bar with multiple chains: This seems like throwing the baby out with the bath water kind of a solution. At least currently, the logging stats are available in the notebook environment for users to debug their script. The `HBox` residue on export, while not ideal, is probably no worse than no output at all.
 - [ ] Not use a progress bar at all but just use standard logging in the notebook environment (as we did in the earlier version). This works for the single chain case, but would give garbled results with multiprocessing.
 - [ ] Roll out our own version of progress bar (my attempt to do so is in the `progbar` branch), which I abandoned because while we were able to solve some of these issues, I was creating new problems that tqdm has already solved quite well.
 - [x] Only use auto_notebook for multiple chains in the notebook environment - This seems like a simple work-around, that gives us what we have right now, but preserves the output at least for the single chain case. I can make this change, if this seems acceptable.

I think the best long-term course of option for us is to figure out how to get the above line to render when the notebook is converted into HTML. It doesn't even have to work out of the box, and even if we have to change a few configurations or do it through a script to get it working for our tutorials, that's okay. The users can choose to disable the progress bar before the final export to get rid of the `HBox` line and just report the diagnostics. 

note: Above is my vote for the next step. Feel free to add your suggestion to the list.
@neerajprad About getting HBox line to be rendered in HTML, this PR of nbconvert can be a solution: https://github.com/jupyter/nbconvert/pull/900. I really love the current progressbar messages (the only minor point is when MCMC is so fast in notebook, [this issue](https://github.com/jupyter/notebook/issues/3010) will happen, but I can solve it by changing jupyter's configs)). Whether it is a normal log or a widget is not important I guess. Because there are many issues for tqdm with multiprocessing in notebooks, I think we can go with your vote and keep the current stage of progressbar for multi-chain. If we can convert `ipywidget` state to normal log at the end of sampling, then it is the best IMO.

Resolving patch:
diff --git a/pyro/infer/mcmc/logger.py b/pyro/infer/mcmc/logger.py
--- a/pyro/infer/mcmc/logger.py
+++ b/pyro/infer/mcmc/logger.py
@@ -4,7 +4,8 @@
 import sys
 from collections import OrderedDict
 
-from tqdm.auto import tqdm
+from tqdm import tqdm
+from tqdm.auto import tqdm as tqdm_nb
 
 try:
     get_ipython
@@ -38,22 +39,70 @@
 # OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 
 
-# TODO: Remove once https://github.com/tqdm/tqdm/issues/650 is resolved.
-class ProgressBar(tqdm):
-    def __init__(self, *args, **kwargs):
-        super(ProgressBar, self).__init__(*args, **kwargs)
+class ProgressBar(object):
+    """
+    Initialize progress bars using :class:`~tqdm.tqdm`.
+
+    :param int warmup_steps: Number of warmup steps.
+    :param int num_samples: Number of MCMC samples.
+    :param int min_width: Minimum column width of the bar.
+    :param int max_width: Maximum column width of the bar.
+    :param bool disable: Disable progress bar.
+    :param int num_bars: Number of progress bars to initialize.
+        If multiple bars are initialized, they need to be separately
+        updated via the ``pos`` kwarg.
+    """
+    def __init__(self, warmup_steps, num_samples, min_width=80, max_width=120,
+                 disable=False, num_bars=1):
+        total_steps = warmup_steps + num_samples
+        # Disable progress bar in "CI"
+        # (see https://github.com/travis-ci/travis-ci/issues/1337).
+        disable = disable or "CI" in os.environ or "PYTEST_XDIST_WORKER" in os.environ
+        bar_format = None
+        if not ipython_env:
+            bar_format = "{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}, {rate_fmt}{postfix}]"
+        pbar_cls = tqdm_nb if num_bars > 1 and ipython_env else tqdm
+        self.progress_bars = []
+        for i in range(num_bars):
+            description = "Warmup" if num_bars == 1 else "Warmup [{}]".format(i + 1)
+            pbar = pbar_cls(total=total_steps, desc=description, bar_format=bar_format,
+                            position=i, file=sys.stderr, disable=disable)
+            # Assume reasonable values when terminal width not available
+            if getattr(pbar, "ncols", None) is not None:
+                pbar.ncols = max(min_width, pbar.ncols)
+                pbar.ncols = min(max_width, pbar.ncols)
+            self.progress_bars.append(pbar)
+        self.disable = disable
+        self.ipython_env = ipython_env
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        self.close()
+        return False
 
     def set_description(self, *args, **kwargs):
+        pos = kwargs.pop("pos", 0)
         if not self.disable:
-            super(ProgressBar, self).set_description(*args, **kwargs)
+            self.progress_bars[pos].set_description(*args, **kwargs)
 
     def set_postfix(self, *args, **kwargs):
+        pos = kwargs.pop("pos", 0)
         if not self.disable:
-            super(ProgressBar, self).set_postfix(*args, **kwargs)
+            self.progress_bars[pos].set_postfix(*args, **kwargs)
 
     def update(self, *args, **kwargs):
+        pos = kwargs.pop("pos", 0)
         if not self.disable:
-            super(ProgressBar, self).update(*args, **kwargs)
+            self.progress_bars[pos].update(*args, **kwargs)
+
+    def close(self):
+        for pbar in self.progress_bars:
+            pbar.close()
+        # Required to not overwrite multiple progress bars on exit.
+        if not self.ipython_env and not self.disable:
+            sys.stderr.write("\n" * len(self.progress_bars))
 
 
 class QueueHandler(logging.Handler):
@@ -177,37 +226,6 @@ def filter(self, record):
         return True
 
 
-def initialize_progbar(warmup_steps, num_samples, min_width=80, max_width=120, pos=None,
-                       disable=False):
-    """
-    Initialize progress bar using :class:`~tqdm.tqdm`.
-
-    :param int warmup_steps: Number of warmup steps.
-    :param int num_samples: Number of MCMC samples.
-    :param int min_width: Minimum column width of the bar.
-    :param int max_width: Maximum column width of the bar.
-    :param int pos: Position of the bar (e.g. in the case of
-        multiple parallel samplers).
-    :param bool disable: Disable progress bar.
-    """
-    description = "Warmup" if pos is None else "Warmup [{}]".format(pos + 1)
-    total_steps = warmup_steps + num_samples
-    # Disable progress bar in "CI"
-    # (see https://github.com/travis-ci/travis-ci/issues/1337).
-    disable = disable or "CI" in os.environ or "PYTEST_XDIST_WORKER" in os.environ
-    bar_format = None
-    if not ipython_env:
-        bar_format = "{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}, {rate_fmt}{postfix}]"
-    progress_bar = ProgressBar(total=total_steps, desc=description, bar_format=bar_format,
-                               position=pos, file=sys.stderr, disable=disable)
-    progress_bar._ipython_env = ipython_env
-
-    if getattr(progress_bar, "ncols", None) is not None:
-        progress_bar.ncols = max(min_width, progress_bar.ncols)
-        progress_bar.ncols = min(max_width, progress_bar.ncols)
-    return progress_bar
-
-
 def initialize_logger(logger, logger_id, progress_bar=None, log_queue=None):
     """
     Initialize logger for the :class:`pyro.infer.mcmc` module.
diff --git a/pyro/infer/mcmc/mcmc.py b/pyro/infer/mcmc/mcmc.py
--- a/pyro/infer/mcmc/mcmc.py
+++ b/pyro/infer/mcmc/mcmc.py
@@ -3,24 +3,22 @@
 import json
 import logging
 import signal
-import sys
 import threading
 import warnings
 from collections import OrderedDict
 
 import six
-from six.moves import queue
 import torch
 import torch.multiprocessing as mp
+from six.moves import queue
 
 import pyro
+import pyro.ops.stats as stats
 from pyro.infer import TracePosterior
 from pyro.infer.abstract_infer import Marginals
-from pyro.infer.mcmc.logger import initialize_logger, initialize_progbar, DIAGNOSTIC_MSG, TqdmHandler
-import pyro.ops.stats as stats
+from pyro.infer.mcmc.logger import initialize_logger, DIAGNOSTIC_MSG, TqdmHandler, ProgressBar
 from pyro.util import optional
 
-
 MAX_SEED = 2**32 - 1
 
 
@@ -29,8 +27,7 @@ def logger_thread(log_queue, warmup_steps, num_samples, num_chains, disable_prog
     Logging thread that asynchronously consumes logging events from `log_queue`,
     and handles them appropriately.
     """
-    progress_bars = [initialize_progbar(warmup_steps, num_samples, pos=i, disable=disable_progbar)
-                     for i in range(num_chains)]
+    progress_bars = ProgressBar(warmup_steps, num_samples, disable=disable_progbar, num_bars=num_chains)
     logger = logging.getLogger(__name__)
     logger.propagate = False
     logger.addHandler(TqdmHandler())
@@ -49,18 +46,14 @@ def logger_thread(log_queue, warmup_steps, num_samples, num_chains, disable_prog
                 pbar_pos = int(logger_id.split(":")[-1])
                 num_samples[pbar_pos] += 1
                 if num_samples[pbar_pos] == warmup_steps:
-                    progress_bars[pbar_pos].set_description("Sample [{}]".format(pbar_pos + 1))
+                    progress_bars.set_description("Sample [{}]".format(pbar_pos + 1), pos=pbar_pos)
                 diagnostics = json.loads(msg, object_pairs_hook=OrderedDict)
-                progress_bars[pbar_pos].set_postfix(diagnostics)
-                progress_bars[pbar_pos].update()
+                progress_bars.set_postfix(diagnostics, pos=pbar_pos)
+                progress_bars.update(pos=pbar_pos)
             else:
                 logger.handle(record)
     finally:
-        for pbar in progress_bars:
-            pbar.close()
-            # Required to not overwrite multiple progress bars on exit.
-            if not pbar._ipython_env:
-                sys.stderr.write("\n")
+        progress_bars.close()
 
 
 class _Worker(object):
@@ -208,7 +201,7 @@ def _traces(self, *args, **kwargs):
         is_multiprocessing = log_queue is not None
         progress_bar = None
         if not is_multiprocessing:
-            progress_bar = initialize_progbar(self.warmup_steps, self.num_samples, disable=self.disable_progbar)
+            progress_bar = ProgressBar(self.warmup_steps, self.num_samples, disable=self.disable_progbar)
         self.logger = initialize_logger(self.logger, logger_id, progress_bar, log_queue)
         self.kernel.setup(self.warmup_steps, *args, **kwargs)
         trace = self.kernel.initial_trace


Test patch: