Feature/5275 clean progress bar print (#5470)

* Trainer.test should return only test metrics (#5214)

* resolve bug

* merge tests

* Fix metric state reset (#5273)

* Fix metric state reset

* Fix test

* Improve formatting

Co-authored-by: Ananya Harsh Jha <[email protected]>

* print() method added to ProgressBar

* printing alongside progress bar added to LightningModule.print()

* LightningModule.print() method documentation updated

* ProgressBarBase.print() stub added

* stub

* add progress bar tests

* fix isort

* Progress Callback fixes

* test_metric.py duplicate DummyList removed

* PEP and isort fixes

* CHANGELOG updated

* test_progress_bar_print win linesep fix

* test_progress_bar.py remove whitespaces

* Update CHANGELOG.md

Co-authored-by: chaton <[email protected]>
Co-authored-by: Tadej Svetina <[email protected]>
Co-authored-by: Ananya Harsh Jha <[email protected]>
Co-authored-by: Alexander Snorkin <[email protected]>
Co-authored-by: rohitgr7 <[email protected]>
Co-authored-by: Adrian Wälchli <[email protected]>
Co-authored-by: Carlos Mocholí <[email protected]>

Author: 8 people

CHANGELOG.md

@@ -89,6 +89,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added DeepSpeed integration ([#5954](https://github.com/PyTorchLightning/pytorch-lightning/pull/5954),
     [#6042](https://github.com/PyTorchLightning/pytorch-lightning/pull/6042))
 
+- Added a way to print to terminal without breaking up the progress bar ([#5470](https://github.com/PyTorchLightning/pytorch-lightning/pull/5470))
+
+
 ### Changed
 
 - Changed `stat_scores` metric now calculates stat scores over all classes and gains new parameters, in line with the new `StatScores` metric ([#4839](https://github.com/PyTorchLightning/pytorch-lightning/pull/4839))

pytorch_lightning/callbacks/progress.py

@@ -19,6 +19,8 @@
 
 """
 import importlib
+import io
+import os
 import sys
 
 # check if ipywidgets is installed before importing tqdm.auto
@@ -187,6 +189,12 @@ def enable(self):
         """
         raise NotImplementedError
 
+    def print(self, *args, **kwargs):
+        """
+        You should provide a way to print without breaking the progress bar.
+        """
+        print(*args, **kwargs)
+
     def on_init_end(self, trainer):
         self._trainer = trainer
 
@@ -451,6 +459,22 @@ def on_predict_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, da
     def on_predict_end(self, trainer, pl_module):
         self.predict_progress_bar.close()
 
+    def print(
+        self, *args, sep: str = ' ', end: str = os.linesep, file: Optional[io.TextIOBase] = None, nolock: bool = False
+    ):
+        active_progress_bar = None
+
+        if not self.main_progress_bar.disable:
+            active_progress_bar = self.main_progress_bar
+        elif not self.val_progress_bar.disable:
+            active_progress_bar = self.val_progress_bar
+        elif not self.test_progress_bar.disable:
+            active_progress_bar = self.test_progress_bar
+
+        if active_progress_bar is not None:
+            s = sep.join(map(str, args))
+            active_progress_bar.write(s, end=end, file=file, nolock=nolock)
+
     def _should_update(self, current, total):
         return self.is_enabled and (current % self.refresh_rate == 0 or current == total)
 

pytorch_lightning/core/lightning.py

@@ -189,8 +189,8 @@ def print(self, *args, **kwargs) -> None:
         Prints only from process 0. Use this in any distributed mode to log only once.
 
         Args:
-            *args: The thing to print. Will be passed to Python's built-in print function.
-            **kwargs: Will be passed to Python's built-in print function.
+            *args: The thing to print. The same as for Python's built-in print function.
+            **kwargs: The same as for Python's built-in print function.
 
         Example::
 
@@ -199,7 +199,11 @@ def forward(self, x):
 
         """
         if self.trainer.is_global_zero:
-            print(*args, **kwargs)
+            progress_bar = self.trainer.progress_bar_callback
+            if progress_bar is not None and progress_bar.is_enabled:
+                progress_bar.print(*args, **kwargs)
+            else:
+                print(*args, **kwargs)
 
     def log(
         self,

tests/callbacks/test_progress_bar.py

@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+import sys
 from unittest import mock
-from unittest.mock import call, Mock
+from unittest.mock import ANY, call, Mock
 
 import pytest
 import torch
@@ -381,3 +382,69 @@ def training_step(self, batch, batch_idx):
 def test_tqdm_format_num(input_num, expected):
     """ Check that the specialized tqdm.format_num appends 0 to floats and strings """
     assert tqdm.format_num(input_num) == expected
+
+
+class PrintModel(BoringModel):
+
+    def training_step(self, *args, **kwargs):
+        self.print("training_step", end="")
+        return super().training_step(*args, **kwargs)
+
+    def validation_step(self, *args, **kwargs):
+        self.print("validation_step", file=sys.stderr)
+        return super().validation_step(*args, **kwargs)
+
+    def test_step(self, *args, **kwargs):
+        self.print("test_step")
+        return super().test_step(*args, **kwargs)
+
+
+@mock.patch("pytorch_lightning.callbacks.progress.tqdm.write")
+def test_progress_bar_print(tqdm_write, tmpdir):
+    """ Test that printing in the LightningModule redirects arguments to the progress bar. """
+    model = PrintModel()
+    bar = ProgressBar()
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        num_sanity_val_steps=0,
+        limit_train_batches=1,
+        limit_val_batches=1,
+        limit_test_batches=1,
+        max_steps=1,
+        callbacks=[bar],
+    )
+    trainer.fit(model)
+    trainer.test(model)
+    assert tqdm_write.call_count == 3
+    assert tqdm_write.call_args_list == [
+        call("training_step", end="", file=None, nolock=False),
+        call("validation_step", end=os.linesep, file=sys.stderr, nolock=False),
+        call("test_step", end=os.linesep, file=None, nolock=False),
+    ]
+
+
+@mock.patch('builtins.print')
+@mock.patch("pytorch_lightning.callbacks.progress.tqdm.write")
+def test_progress_bar_print_disabled(tqdm_write, mock_print, tmpdir):
+    """ Test that printing in LightningModule goes through built-in print functin when progress bar is disabled. """
+    model = PrintModel()
+    bar = ProgressBar()
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        num_sanity_val_steps=0,
+        limit_train_batches=1,
+        limit_val_batches=1,
+        limit_test_batches=1,
+        max_steps=1,
+        callbacks=[bar],
+    )
+    bar.disable()
+    trainer.fit(model)
+    trainer.test(model)
+
+    mock_print.assert_has_calls([
+        call("training_step", end=""),
+        call("validation_step", file=ANY),
+        call("test_step"),
+    ])
+    tqdm_write.assert_not_called()