test: disable rich progress bar in callback tests to ensure tqdm usage

Author: littlebullGit

tests/tests_pytorch/callbacks/progress/test_tqdm_progress_bar.py

@@ -18,7 +18,7 @@
 from collections import defaultdict
 from typing import Union
 from unittest import mock
-from unittest.mock import ANY, Mock, PropertyMock, call
+from unittest.mock import ANY, Mock, PropertyMock, call, patch
 
 import pytest
 import torch
@@ -112,120 +112,122 @@ def test_tqdm_progress_bar_misconfiguration():
 @pytest.mark.parametrize("num_dl", [1, 2])
 def test_tqdm_progress_bar_totals(tmp_path, num_dl):
     """Test that the progress finishes with the correct total steps processed."""
-
-    class CustomModel(BoringModel):
-        def _get_dataloaders(self):
-            dls = [DataLoader(RandomDataset(32, 64)), DataLoader(RandomDataset(32, 64))]
-            return dls[0] if num_dl == 1 else dls
-
-        def val_dataloader(self):
-            return self._get_dataloaders()
-
-        def test_dataloader(self):
-            return self._get_dataloaders()
-
-        def predict_dataloader(self):
-            return self._get_dataloaders()
-
-        def validation_step(self, batch, batch_idx, dataloader_idx=0):
-            return
-
-        def test_step(self, batch, batch_idx, dataloader_idx=0):
-            return
-
-        def predict_step(self, batch, batch_idx, dataloader_idx=0):
-            return
-
-    model = CustomModel()
-
-    # check the sanity dataloaders
-    num_sanity_val_steps = 4
-    trainer = Trainer(
-        default_root_dir=tmp_path, max_epochs=1, limit_train_batches=0, num_sanity_val_steps=num_sanity_val_steps
-    )
-    pbar = trainer.progress_bar_callback
-    with mock.patch("lightning.pytorch.callbacks.progress.tqdm_progress.Tqdm", MockTqdm):
-        trainer.fit(model)
-
-    expected_sanity_steps = [num_sanity_val_steps] * num_dl
-    assert not pbar.val_progress_bar.leave
-    assert trainer.num_sanity_val_batches == expected_sanity_steps
-    assert pbar.val_progress_bar.total_values == expected_sanity_steps
-    assert pbar.val_progress_bar.n_values == list(range(num_sanity_val_steps + 1)) * num_dl
-    assert pbar.val_progress_bar.descriptions == [f"Sanity Checking DataLoader {i}: " for i in range(num_dl)]
-
-    # fit
-    trainer = Trainer(default_root_dir=tmp_path, max_epochs=1)
-    pbar = trainer.progress_bar_callback
-    with mock.patch("lightning.pytorch.callbacks.progress.tqdm_progress.Tqdm", MockTqdm):
-        trainer.fit(model)
-
-    n = trainer.num_training_batches
-    m = trainer.num_val_batches
-    assert len(trainer.train_dataloader) == n
-    # train progress bar should have reached the end
-    assert pbar.train_progress_bar.total == n
-    assert pbar.train_progress_bar.n == n
-    assert pbar.train_progress_bar.leave
-
-    # check val progress bar total
-    assert pbar.val_progress_bar.total_values == m
-    assert pbar.val_progress_bar.n_values == list(range(m[0] + 1)) * num_dl
-    assert pbar.val_progress_bar.descriptions == [f"Validation DataLoader {i}: " for i in range(num_dl)]
-    assert not pbar.val_progress_bar.leave
-
-    # validate
-    with mock.patch("lightning.pytorch.callbacks.progress.tqdm_progress.Tqdm", MockTqdm):
-        trainer.validate(model)
-    assert trainer.num_val_batches == m
-    assert pbar.val_progress_bar.total_values == m
-    assert pbar.val_progress_bar.n_values == list(range(m[0] + 1)) * num_dl
-    assert pbar.val_progress_bar.descriptions == [f"Validation DataLoader {i}: " for i in range(num_dl)]
-
-    # test
-    with mock.patch("lightning.pytorch.callbacks.progress.tqdm_progress.Tqdm", MockTqdm):
-        trainer.test(model)
-    assert pbar.test_progress_bar.leave
-    k = trainer.num_test_batches
-    assert pbar.test_progress_bar.total_values == k
-    assert pbar.test_progress_bar.n_values == list(range(k[0] + 1)) * num_dl
-    assert pbar.test_progress_bar.descriptions == [f"Testing DataLoader {i}: " for i in range(num_dl)]
-    assert pbar.test_progress_bar.leave
-
-    # predict
-    with mock.patch("lightning.pytorch.callbacks.progress.tqdm_progress.Tqdm", MockTqdm):
-        trainer.predict(model)
-    assert pbar.predict_progress_bar.leave
-    k = trainer.num_predict_batches
-    assert pbar.predict_progress_bar.total_values == k
-    assert pbar.predict_progress_bar.n_values == list(range(k[0] + 1)) * num_dl
-    assert pbar.predict_progress_bar.descriptions == [f"Predicting DataLoader {i}: " for i in range(num_dl)]
+    with patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False):
+
+        class CustomModel(BoringModel):
+            def _get_dataloaders(self):
+                dls = [DataLoader(RandomDataset(32, 64)), DataLoader(RandomDataset(32, 64))]
+                return dls[0] if num_dl == 1 else dls
+
+            def val_dataloader(self):
+                return self._get_dataloaders()
+
+            def test_dataloader(self):
+                return self._get_dataloaders()
+
+            def predict_dataloader(self):
+                return self._get_dataloaders()
+
+            def validation_step(self, batch, batch_idx, dataloader_idx=0):
+                return
+
+            def test_step(self, batch, batch_idx, dataloader_idx=0):
+                return
+
+            def predict_step(self, batch, batch_idx, dataloader_idx=0):
+                return
+
+        model = CustomModel()
+
+        # check the sanity dataloaders
+        num_sanity_val_steps = 4
+        trainer = Trainer(
+            default_root_dir=tmp_path, max_epochs=1, limit_train_batches=0, num_sanity_val_steps=num_sanity_val_steps
+        )
+        pbar = trainer.progress_bar_callback
+        with mock.patch("lightning.pytorch.callbacks.progress.tqdm_progress.Tqdm", MockTqdm):
+            trainer.fit(model)
+
+        expected_sanity_steps = [num_sanity_val_steps] * num_dl
+        assert not pbar.val_progress_bar.leave
+        assert trainer.num_sanity_val_batches == expected_sanity_steps
+        assert pbar.val_progress_bar.total_values == expected_sanity_steps
+        assert pbar.val_progress_bar.n_values == list(range(num_sanity_val_steps + 1)) * num_dl
+        assert pbar.val_progress_bar.descriptions == [f"Sanity Checking DataLoader {i}: " for i in range(num_dl)]
+
+        # fit
+        trainer = Trainer(default_root_dir=tmp_path, max_epochs=1)
+        pbar = trainer.progress_bar_callback
+        with mock.patch("lightning.pytorch.callbacks.progress.tqdm_progress.Tqdm", MockTqdm):
+            trainer.fit(model)
+
+        n = trainer.num_training_batches
+        m = trainer.num_val_batches
+        assert len(trainer.train_dataloader) == n
+        # train progress bar should have reached the end
+        assert pbar.train_progress_bar.total == n
+        assert pbar.train_progress_bar.n == n
+        assert pbar.train_progress_bar.leave
+
+        # check val progress bar total
+        assert pbar.val_progress_bar.total_values == m
+        assert pbar.val_progress_bar.n_values == list(range(m[0] + 1)) * num_dl
+        assert pbar.val_progress_bar.descriptions == [f"Validation DataLoader {i}: " for i in range(num_dl)]
+        assert not pbar.val_progress_bar.leave
+
+        # validate
+        with mock.patch("lightning.pytorch.callbacks.progress.tqdm_progress.Tqdm", MockTqdm):
+            trainer.validate(model)
+        assert trainer.num_val_batches == m
+        assert pbar.val_progress_bar.total_values == m
+        assert pbar.val_progress_bar.n_values == list(range(m[0] + 1)) * num_dl
+        assert pbar.val_progress_bar.descriptions == [f"Validation DataLoader {i}: " for i in range(num_dl)]
+
+        # test
+        with mock.patch("lightning.pytorch.callbacks.progress.tqdm_progress.Tqdm", MockTqdm):
+            trainer.test(model)
+        assert pbar.test_progress_bar.leave
+        k = trainer.num_test_batches
+        assert pbar.test_progress_bar.total_values == k
+        assert pbar.test_progress_bar.n_values == list(range(k[0] + 1)) * num_dl
+        assert pbar.test_progress_bar.descriptions == [f"Testing DataLoader {i}: " for i in range(num_dl)]
+        assert pbar.test_progress_bar.leave
+
+        # predict
+        with mock.patch("lightning.pytorch.callbacks.progress.tqdm_progress.Tqdm", MockTqdm):
+            trainer.predict(model)
+        assert pbar.predict_progress_bar.leave
+        k = trainer.num_predict_batches
+        assert pbar.predict_progress_bar.total_values == k
+        assert pbar.predict_progress_bar.n_values == list(range(k[0] + 1)) * num_dl
+        assert pbar.predict_progress_bar.descriptions == [f"Predicting DataLoader {i}: " for i in range(num_dl)]
     assert pbar.predict_progress_bar.leave
 
 
 def test_tqdm_progress_bar_fast_dev_run(tmp_path):
-    model = BoringModel()
+    with patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False):
+        model = BoringModel()
 
-    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
+        trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
 
-    trainer.fit(model)
+        trainer.fit(model)
 
-    pbar = trainer.progress_bar_callback
+        pbar = trainer.progress_bar_callback
 
-    assert pbar.val_progress_bar.n == 1
-    assert pbar.val_progress_bar.total == 1
+        assert pbar.val_progress_bar.n == 1
+        assert pbar.val_progress_bar.total == 1
 
-    # the train progress bar should display 1 batch
-    assert pbar.train_progress_bar.total == 1
-    assert pbar.train_progress_bar.n == 1
+        # the train progress bar should display 1 batch
+        assert pbar.train_progress_bar.total == 1
+        assert pbar.train_progress_bar.n == 1
 
-    trainer.validate(model)
+        trainer.validate(model)
 
-    # the validation progress bar should display 1 batch
-    assert pbar.val_progress_bar.total == 1
-    assert pbar.val_progress_bar.n == 1
+        # the validation progress bar should display 1 batch
+        assert pbar.val_progress_bar.total == 1
+        assert pbar.val_progress_bar.n == 1
 
-    trainer.test(model)
+        trainer.test(model)
 
     # the test progress bar should display 1 batch
     assert pbar.test_progress_bar.total == 1
@@ -325,14 +327,15 @@ def test_tqdm_progress_bar_default_value(tmp_path):
 @mock.patch.dict(os.environ, {"COLAB_GPU": "1"})
 def test_tqdm_progress_bar_value_on_colab(tmp_path):
     """Test that Trainer will override the default in Google COLAB."""
-    trainer = Trainer(default_root_dir=tmp_path)
-    assert trainer.progress_bar_callback.refresh_rate == 20
+    with patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False):
+        trainer = Trainer(default_root_dir=tmp_path)
+        assert trainer.progress_bar_callback.refresh_rate == 20
 
-    trainer = Trainer(default_root_dir=tmp_path, callbacks=TQDMProgressBar())
-    assert trainer.progress_bar_callback.refresh_rate == 20
+        trainer = Trainer(default_root_dir=tmp_path, callbacks=TQDMProgressBar())
+        assert trainer.progress_bar_callback.refresh_rate == 20
 
-    trainer = Trainer(default_root_dir=tmp_path, callbacks=TQDMProgressBar(refresh_rate=19))
-    assert trainer.progress_bar_callback.refresh_rate == 19
+        trainer = Trainer(default_root_dir=tmp_path, callbacks=TQDMProgressBar(refresh_rate=19))
+        assert trainer.progress_bar_callback.refresh_rate == 19
 
 
 @pytest.mark.parametrize(
@@ -413,21 +416,22 @@ def test_test_progress_bar_update_amount(tmp_path, test_batches: int, refresh_ra
 
 def test_tensor_to_float_conversion(tmp_path):
     """Check tensor gets converted to float."""
-
-    class TestModel(BoringModel):
-        def training_step(self, batch, batch_idx):
-            self.log("a", torch.tensor(0.123), prog_bar=True, on_epoch=False)
-            self.log("b", torch.tensor([1]), prog_bar=True, on_epoch=False)
-            self.log("c", 2, prog_bar=True, on_epoch=False)
-            return super().training_step(batch, batch_idx)
-
-    trainer = Trainer(
-        default_root_dir=tmp_path, max_epochs=1, limit_train_batches=2, logger=False, enable_checkpointing=False
-    )
-    trainer.fit(TestModel())
-
-    torch.testing.assert_close(trainer.progress_bar_metrics["a"], 0.123)
-    assert trainer.progress_bar_metrics["b"] == 1.0
+    with patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False):
+
+        class TestModel(BoringModel):
+            def training_step(self, batch, batch_idx):
+                self.log("a", torch.tensor(0.123), prog_bar=True, on_epoch=False)
+                self.log("b", torch.tensor([1]), prog_bar=True, on_epoch=False)
+                self.log("c", 2, prog_bar=True, on_epoch=False)
+                return super().training_step(batch, batch_idx)
+
+        trainer = Trainer(
+            default_root_dir=tmp_path, max_epochs=1, limit_train_batches=2, logger=False, enable_checkpointing=False
+        )
+        trainer.fit(TestModel())
+
+        torch.testing.assert_close(trainer.progress_bar_metrics["a"], 0.123)
+        assert trainer.progress_bar_metrics["b"] == 1.0
     assert trainer.progress_bar_metrics["c"] == 2.0
     pbar = trainer.progress_bar_callback.train_progress_bar
     actual = str(pbar.postfix)

tests/tests_pytorch/callbacks/test_callbacks.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 from pathlib import Path
 from re import escape
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import pytest
 from lightning_utilities.test.warning import no_warning_call
@@ -122,17 +122,18 @@ def load_state_dict(self, state_dict) -> None:
 def test_resume_callback_state_saved_by_type_stateful(tmp_path):
     """Test that a legacy checkpoint that didn't use a state key before can still be loaded, using
     state_dict/load_state_dict."""
-    model = BoringModel()
-    callback = OldStatefulCallback(state=111)
-    trainer = Trainer(default_root_dir=tmp_path, max_steps=1, callbacks=[callback])
-    trainer.fit(model)
-    ckpt_path = Path(trainer.checkpoint_callback.best_model_path)
-    assert ckpt_path.exists()
-
-    callback = OldStatefulCallback(state=222)
-    trainer = Trainer(default_root_dir=tmp_path, max_steps=2, callbacks=[callback])
-    trainer.fit(model, ckpt_path=ckpt_path)
-    assert callback.state == 111
+    with patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False):
+        model = BoringModel()
+        callback = OldStatefulCallback(state=111)
+        trainer = Trainer(default_root_dir=tmp_path, max_steps=1, callbacks=[callback])
+        trainer.fit(model)
+        ckpt_path = Path(trainer.checkpoint_callback.best_model_path)
+        assert ckpt_path.exists()
+
+        callback = OldStatefulCallback(state=222)
+        trainer = Trainer(default_root_dir=tmp_path, max_steps=2, callbacks=[callback])
+        trainer.fit(model, ckpt_path=ckpt_path)
+        assert callback.state == 111
 
 
 def test_resume_incomplete_callbacks_list_warning(tmp_path):