fix(RichProgressBar): Convert tensor metrics to float

The RichProgressBar was failing in multi-GPU environments because it could not handle tensor metrics from different devices. This commit overrides the get_metrics method to convert all tensor metrics to floats before they are rendered, preventing errors. An accompanying test is added to verify the fix.

Author: littlebullGit

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

@@ -17,6 +17,7 @@
 from datetime import timedelta
 from typing import Any, Optional, Union, cast
 
+import torch
 from lightning_utilities.core.imports import RequirementCache
 from typing_extensions import override
 
@@ -612,6 +613,17 @@ def _reset_progress_bar_ids(self) -> None:
         self.test_progress_bar_id = None
         self.predict_progress_bar_id = None
 
+    @override
+    def get_metrics(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
+    ) -> dict[str, Union[int, str, float, dict[str, float]]]:
+        items = super().get_metrics(trainer, pl_module)
+        # convert all metrics to float before sending to rich
+        for k, v in items.items():
+            if isinstance(v, torch.Tensor):
+                items[k] = v.item()
+        return items
+
     def _update_metrics(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
         metrics = self.get_metrics(trainer, pl_module)
         if self._metric_component:

tests/tests_pytorch/callbacks/progress/test_tqdm_progress_bar.py

@@ -109,98 +109,98 @@ def test_tqdm_progress_bar_misconfiguration():
         Trainer(callbacks=TQDMProgressBar(), enable_progress_bar=False)
 
 
+@patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False)
 @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."""
-    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)]
+
+    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
 
 
@@ -414,24 +414,30 @@ def test_test_progress_bar_update_amount(tmp_path, test_batches: int, refresh_ra
     assert progress_bar.test_progress_bar.n_values == updates
 
 
+@patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", False)
 def test_tensor_to_float_conversion(tmp_path):
     """Check tensor gets converted to float."""
-    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
-        )
+
+    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
+    )
+
+    with mock.patch.object(sys.stdout, "write") as mock_write:
         trainer.fit(TestModel())
+    bar_updates = "".join(call.args[0] for call in mock_write.call_args_list)
+    assert "a=0.123" in bar_updates
+    assert "b=1.000" in bar_updates
+    assert "c=2.000" in bar_updates
 
-        torch.testing.assert_close(trainer.progress_bar_metrics["a"], 0.123)
-        assert trainer.progress_bar_metrics["b"] == 1.0
+    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/trainer/connectors/test_rich_integration.py

@@ -14,8 +14,12 @@
 
 from unittest.mock import patch
 
+import pytest
+import torch
+
 from lightning.pytorch import Trainer
 from lightning.pytorch.callbacks import ModelSummary, ProgressBar, RichModelSummary, RichProgressBar, TQDMProgressBar
+from lightning.pytorch.demos.boring_classes import BoringModel
 
 
 class TestRichIntegration:
@@ -133,3 +137,33 @@ def test_model_summary_disabled_with_rich(self, tmp_path):
             default_root_dir=tmp_path, enable_model_summary=False, logger=False, enable_checkpointing=False
         )
         assert not any(isinstance(cb, ModelSummary) for cb in trainer.callbacks)
+
+    @patch("lightning.pytorch.trainer.connectors.callback_connector._RICH_AVAILABLE", True)
+    def test_rich_progress_bar_tensor_metric(self, tmp_path):
+        """Test that tensor metrics are converted to float for RichProgressBar."""
+
+        class MyModel(BoringModel):
+            def training_step(self, batch, batch_idx):
+                self.log("my_tensor_metric", torch.tensor(1.23), prog_bar=True)
+                return super().training_step(batch, batch_idx)
+
+        model = MyModel()
+        trainer = Trainer(
+            default_root_dir=tmp_path,
+            limit_train_batches=1,
+            limit_val_batches=0,
+            max_epochs=1,
+            logger=False,
+            enable_checkpointing=False,
+        )
+
+        with patch("lightning.pytorch.callbacks.progress.rich_progress.MetricsTextColumn.update") as mock_update:
+            trainer.fit(model)
+
+        assert mock_update.call_count > 0
+        # The metrics are updated multiple times, check the last call
+        last_call_metrics = mock_update.call_args[0][0]
+        assert "my_tensor_metric" in last_call_metrics
+        metric_val = last_call_metrics["my_tensor_metric"]
+        assert isinstance(metric_val, float)
+        assert metric_val == pytest.approx(1.23)