Fix sanity check for RichProgressBar (#10913)

Author: kaushikb11

CHANGELOG.md

@@ -4,29 +4,17 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
-## [1.5.6] - 2021-12-14
+## [1.5.6] - 2021-12-15
 
 ### Fixed
 
 - Fixed a bug where the DeepSpeedPlugin arguments `cpu_checkpointing` and `contiguous_memory_optimization` were not being forwarded to deepspeed correctly ([#10874](https://github.com/PyTorchLightning/pytorch-lightning/issues/10874))
-
-
 - Fixed an issue with `NeptuneLogger` causing checkpoints to be uploaded with a duplicated file extension ([#11015](https://github.com/PyTorchLightning/pytorch-lightning/issues/11015))
-=======
-
-
 - Fixed support for logging within callbacks returned from `LightningModule` ([#10991](https://github.com/PyTorchLightning/pytorch-lightning/pull/10991))
-
-
+- Fixed running sanity check with `RichProgressBar` ([#10913](https://github.com/PyTorchLightning/pytorch-lightning/pull/10913))
 - Fixed support for `CombinedLoader` while checking for warning raised with eval dataloaders ([#10994](https://github.com/PyTorchLightning/pytorch-lightning/pull/10994))
 
 
--
-
-
--
-
-
 ## [1.5.5] - 2021-12-07
 
 ### Fixed

pytorch_lightning/callbacks/progress/rich_progress.py

@@ -328,7 +328,9 @@ def on_sanity_check_start(self, trainer, pl_module):
 
     def on_sanity_check_end(self, trainer, pl_module):
         super().on_sanity_check_end(trainer, pl_module)
-        self._update(self.val_sanity_progress_bar_id, visible=False)
+        if self.progress is not None:
+            self.progress.update(self.val_sanity_progress_bar_id, advance=0, visible=False)
+            self.progress.refresh()
 
     def on_train_epoch_start(self, trainer, pl_module):
         super().on_train_epoch_start(trainer, pl_module)

tests/callbacks/test_rich_progress_bar.py

@@ -180,3 +180,24 @@ def test_rich_progress_bar_leave(tmpdir, leave, reset_call_count):
         )
         trainer.fit(model)
     assert mock_progress_reset.call_count == reset_call_count
+
+
+@RunIf(rich=True)
+@pytest.mark.parametrize("limit_val_batches", (1, 5))
+def test_rich_progress_bar_num_sanity_val_steps(tmpdir, limit_val_batches: int):
+    model = BoringModel()
+
+    progress_bar = RichProgressBar()
+    num_sanity_val_steps = 3
+
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        num_sanity_val_steps=num_sanity_val_steps,
+        limit_train_batches=1,
+        limit_val_batches=limit_val_batches,
+        max_epochs=1,
+        callbacks=progress_bar,
+    )
+
+    trainer.fit(model)
+    assert progress_bar.progress.tasks[0].completed == min(num_sanity_val_steps, limit_val_batches)