Add ability for TQDMProgressBar to retain prior epoch training bars (L…

…ightning-AI#19322)

docs/source-pytorch/common/progress_bar.rst

@@ -36,6 +36,14 @@ You can update ``refresh_rate`` (rate (number of batches) at which the progress
 
     trainer = Trainer(callbacks=[TQDMProgressBar(refresh_rate=10)])
 
+By default the training progress bar is reset (overwritten) at each new epoch.
+If you wish for a new progress bar to be displayed at the end of every epoch, set
+:paramref:`TQDMProgressBar.leave <lightning.pytorch.callbacks.TQDMProgressBar.leave>` to ``True``.
+
+.. code-block:: python
+
+    trainer = Trainer(callbacks=[TQDMProgressBar(leave=True)])
+
 If you want to customize the default :class:`~lightning.pytorch.callbacks.TQDMProgressBar` used by Lightning, you can override
 specific methods of the callback class and pass your custom implementation to the :class:`~lightning.pytorch.trainer.trainer.Trainer`.
 

src/lightning/pytorch/callbacks/progress/tqdm_progress.py

@@ -99,12 +99,13 @@ class TQDMProgressBar(ProgressBar):
             together. This corresponds to
             :paramref:`~lightning.pytorch.trainer.trainer.Trainer.process_position` in the
             :class:`~lightning.pytorch.trainer.trainer.Trainer`.
+        leave: Leaves the finished progress bar in the terminal at the end of the epoch. Default: False
 
     """
 
     BAR_FORMAT = "{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_noinv_fmt}{postfix}]"
 
-    def __init__(self, refresh_rate: int = 1, process_position: int = 0):
+    def __init__(self, refresh_rate: int = 1, process_position: int = 0, leave: bool = False):
         super().__init__()
         self._refresh_rate = self._resolve_refresh_rate(refresh_rate)
         self._process_position = process_position
@@ -113,6 +114,7 @@ def __init__(self, refresh_rate: int = 1, process_position: int = 0):
         self._val_progress_bar: Optional[_tqdm] = None
         self._test_progress_bar: Optional[_tqdm] = None
         self._predict_progress_bar: Optional[_tqdm] = None
+        self._leave = leave
 
     def __getstate__(self) -> Dict:
         # can't pickle the tqdm objects
@@ -262,6 +264,8 @@ def on_train_start(self, *_: Any) -> None:
 
     @override
     def on_train_epoch_start(self, trainer: "pl.Trainer", *_: Any) -> None:
+        if self._leave:
+            self.train_progress_bar = self.init_train_tqdm()
         self.train_progress_bar.reset(convert_inf(self.total_train_batches))
         self.train_progress_bar.initial = 0
         self.train_progress_bar.set_description(f"Epoch {trainer.current_epoch}")
@@ -282,7 +286,8 @@ def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModu
 
     @override
     def on_train_end(self, *_: Any) -> None:
-        self.train_progress_bar.close()
+        if not self._leave:
+            self.train_progress_bar.close()
 
     @override
     def on_validation_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None: