Fix sanity check for RichProgressBar

CHANGELOG.md

@@ -231,9 +231,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed a bug that caused incorrect batch indices to be passed to the `BasePredictionWriter` hooks when using a dataloader with `num_workers > 0` ([#10870](https://github.com/PyTorchLightning/pytorch-lightning/pull/10870))
 
 
-
--
-
+- Fixed running sanity check with `RichProgressBar` ([#10913](https://github.com/PyTorchLightning/pytorch-lightning/pull/10913))
 
 
 ## [1.5.4] - 2021-11-30

pytorch_lightning/callbacks/progress/rich_progress.py

@@ -334,7 +334,8 @@ 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.refresh()
 
     def on_train_epoch_start(self, trainer, pl_module):

tests/callbacks/test_rich_progress_bar.py

@@ -201,3 +201,24 @@ def test_rich_progress_bar_refresh_rate(progress_update, tmpdir, refresh_rate, e
     trainer.fit(model)
 
     assert progress_update.call_count == expected_call_count
+
+
+@RunIf(rich=True)
+def test_rich_progress_bar_stages(tmpdir):
+    model = BoringModel()
+
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        num_sanity_val_steps=1,
+        limit_train_batches=1,
+        limit_val_batches=1,
+        limit_test_batches=1,
+        limit_predict_batches=1,
+        max_epochs=1,
+        callbacks=RichProgressBar(),
+    )
+
+    trainer.fit(model)
+    trainer.validate(model)
+    trainer.test(model)
+    trainer.predict(model)