Fix Notebook Progress Bars (#1313)

#1264 broke the progress bars in notebooks. It screwed up the formatting and caused an io.UnsupportedOperation error in Colab when calling sys.stderr.fileno().

This PR fixes these issues.

Closes #1312
Closes https://mosaicml.atlassian.net/browse/CO-770

Author: ravi-mosaicml

composer/loggers/progress_bar_logger.py

@@ -15,7 +15,7 @@
 from composer.core.time import Timestamp, TimeUnit
 from composer.loggers.logger import Logger, LogLevel, format_log_data_value
 from composer.loggers.logger_destination import LoggerDestination
-from composer.utils import dist
+from composer.utils import dist, is_notebook
 
 __all__ = ['ProgressBarLogger']
 
@@ -27,7 +27,7 @@ class _ProgressBar:
     def __init__(
         self,
         total: Optional[int],
-        position: int,
+        position: Optional[int],
         bar_format: str,
         file: TextIO,
         metrics: Dict[str, Any],
@@ -40,7 +40,7 @@ def __init__(
         self.position = position
         self.timestamp_key = timestamp_key
         self.file = file
-        is_atty = os.isatty(self.file.fileno())
+        is_atty = is_notebook() or os.isatty(self.file.fileno())
         self.pbar = tqdm.auto.tqdm(
             total=total,
             position=position,
@@ -51,7 +51,8 @@ def __init__(
             # We set `leave=False` so TQDM does not jump around, but we emulate `leave=True` behavior when closing
             # by printing a dummy newline and refreshing to force tqdm to print to a stale line
             # But on k8s, we need `leave=True`, as it would otherwise overwrite the pbar in place
-            leave=not is_atty,
+            # If in a notebook, then always set leave=True, as otherwise jupyter would remote the progress bars
+            leave=True if is_notebook() else not is_atty,
             postfix=metrics,
             unit=unit,
         )
@@ -67,18 +68,22 @@ def update(self, n=1):
     def update_to_timestamp(self, timestamp: Timestamp):
         n = int(getattr(timestamp, self.timestamp_key))
         n = n - self.pbar.n
-        self.pbar.update(int(n))
+        self.update(int(n))
 
     def close(self):
-        if self.position != 0:
-            # Force a (potentially hidden) progress bar to re-render itself
-            # Don't render the dummy pbar (at position 0), since that will clear a real pbar (at position 1)
+        if is_notebook():
+            # If in a notebook, always refresh before closing, so the
+            # finished progress is displayed
             self.pbar.refresh()
-        # Create a newline that will not be erased by leave=False. This allows for the finished pbar to be cached in the terminal
-        # This emulates `leave=True` without progress bar jumping
-        print('', file=self.file, flush=True)
-
-        self.pbar.close()
+        else:
+            if self.position != 0:
+                # Force a (potentially hidden) progress bar to re-render itself
+                # Don't render the dummy pbar (at position 0), since that will clear a real pbar (at position 1)
+                self.pbar.refresh()
+            # Create a newline that will not be erased by leave=False. This allows for the finished pbar to be cached in the terminal
+            # This emulates `leave=True` without progress bar jumping
+            print('', file=self.file, flush=True)
+            self.pbar.close()
 
     def state_dict(self) -> Dict[str, Any]:
         pbar_state = self.pbar.format_dict
@@ -226,7 +231,8 @@ def _build_pbar(self, state: State, is_train: bool) -> _ProgressBar:
             with the time (in units of ``max_duration.unit``) at which evaluation runs.
         """
         # Always using position=1 to avoid jumping progress bars
-        position = 1
+        # In jupyter notebooks, no need for the dummy pbar, so use the default position
+        position = None if is_notebook() else 1
         desc = f'{state.dataloader_label:15}'
         max_duration_unit = None if state.max_duration is None else state.max_duration.unit
 
@@ -267,23 +273,27 @@ def _build_pbar(self, state: State, is_train: bool) -> _ProgressBar:
             total=total,
             position=position,
             keys_to_log=_IS_TRAIN_TO_KEYS_TO_LOG[is_train],
-            bar_format=desc + ' {l_bar}{bar:25}{r_bar}{bar:-1b}',
+            # In a notebook, the `bar_format` should not include the {bar}, as otherwise
+            # it would appear twice.
+            bar_format=desc + ' {l_bar}' + ('' if is_notebook() else '{bar:25}') + '{r_bar}{bar:-1b}',
             unit=unit.value.lower(),
             metrics={},
             timestamp_key=timestamp_key,
         )
 
     def init(self, state: State, logger: Logger) -> None:
         del state, logger  # unused
-        self.dummy_pbar = _ProgressBar(
-            file=self.stream,
-            position=0,
-            total=1,
-            metrics={},
-            keys_to_log=[],
-            bar_format='{bar:-1b}',
-            timestamp_key='',
-        )
+        if not is_notebook():
+            # Notebooks don't need the dummy progress bar; otherwise, it would be visible.
+            self.dummy_pbar = _ProgressBar(
+                file=self.stream,
+                position=0,
+                total=1,
+                metrics={},
+                keys_to_log=[],
+                bar_format='{bar:-1b}',
+                timestamp_key='',
+            )
 
     def epoch_start(self, state: State, logger: Logger) -> None:
         if self.show_pbar and not self.train_pbar:

composer/utils/__init__.py

@@ -11,7 +11,7 @@
 from composer.utils.import_helpers import MissingConditionalImportError, import_object
 from composer.utils.inference import export_for_inference
 from composer.utils.iter_helpers import IteratorFileStream, ensure_tuple, map_collection
-from composer.utils.misc import is_model_deepspeed
+from composer.utils.misc import is_model_deepspeed, is_notebook
 from composer.utils.object_store import (LibcloudObjectStore, ObjectStore, ObjectStoreTransientError, S3ObjectStore,
                                          SFTPObjectStore)
 from composer.utils.retrying import retry
@@ -31,6 +31,7 @@
     'MissingConditionalImportError',
     'import_object',
     'is_model_deepspeed',
+    'is_notebook',
     'StringEnum',
     'load_checkpoint',
     'save_checkpoint',

composer/utils/collect_env.py

@@ -49,6 +49,8 @@
 import importlib_metadata
 import psutil
 
+from composer.utils.misc import is_notebook
+
 __all__ = ['configure_excepthook', 'disable_env_report', 'enable_env_report', 'print_env']
 
 # Check if PyTorch is installed
@@ -70,14 +72,11 @@
     COMPOSER_AVAILABLE = False
 
 # Check if we're running in a notebook
-try:
-    __IPYTHON__  #type: ignore
+IPYTHON_AVAILABLE = is_notebook()
+if IPYTHON_AVAILABLE:
     from composer.utils.import_helpers import import_object
     get_ipython = import_object('IPython:get_ipython')
     nb = get_ipython()
-    IPYTHON_AVAILABLE = True
-except (NameError,):
-    IPYTHON_AVAILABLE = False
 
 # Place to keep track of the original excepthook
 _orig_excepthook = None

composer/utils/misc.py

@@ -5,7 +5,7 @@
 
 import torch
 
-__all__ = ['is_model_deepspeed']
+__all__ = ['is_model_deepspeed', 'is_notebook']
 
 
 def is_model_deepspeed(model: torch.nn.Module) -> bool:
@@ -16,3 +16,12 @@ def is_model_deepspeed(model: torch.nn.Module) -> bool:
         return False
     else:
         return isinstance(model, deepspeed.DeepSpeedEngine)
+
+
+def is_notebook():
+    """Whether Composer is running in a IPython/Jupyter Notebook."""
+    try:
+        __IPYTHON__  #type: ignore
+        return True
+    except NameError:
+        return False