Support TQDM_MINITERS env variable (#19381)

Co-authored-by: awaelchli <[email protected]>

Author: wouterzwerink

src/lightning/pytorch/CHANGELOG.md

@@ -27,6 +27,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added a utility function and CLI to consolidate FSDP sharded checkpoints into a single file ([#19213](https://github.com/Lightning-AI/lightning/pull/19213))
 
 
+- The TQDM progress bar now respects the env variable `TQDM_MINITERS` for setting the refresh rate ([#19381](https://github.com/Lightning-AI/lightning/pull/19381))
+
+
 ### Changed
 
 - `seed_everything()` without passing in a seed no longer randomly selects a seed, and now defaults to `0` ([#18846](https://github.com/Lightning-AI/lightning/pull/18846))

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

@@ -431,6 +431,9 @@ def _resolve_refresh_rate(refresh_rate: int) -> int:
             # smaller refresh rate on colab causes crashes, choose a higher value
             rank_zero_debug("Using a higher refresh rate on Colab. Setting it to `20`")
             return 20
+        # Support TQDM_MINITERS environment variable, which sets the minimum refresh rate
+        if "TQDM_MINITERS" in os.environ:
+            return max(int(os.environ["TQDM_MINITERS"]), refresh_rate)
         return refresh_rate
 
 

tests/tests_pytorch/callbacks/progress/test_tqdm_progress_bar.py

@@ -334,6 +334,19 @@ def test_tqdm_progress_bar_value_on_colab(tmp_path):
     assert trainer.progress_bar_callback.refresh_rate == 19
 
 
+@pytest.mark.parametrize(("refresh_rate", "env_value", "expected"), [
+    (0, 1, 1),
+    (1, 0, 1),
+    (1, 1, 1),
+    (2, 1, 2),
+    (1, 2, 2),
+])
+def test_tqdm_progress_bar_refresh_rate_via_env_variable(refresh_rate, env_value, expected):
+    with mock.patch.dict(os.environ, {"TQDM_MINITERS": str(env_value)}):
+        bar = TQDMProgressBar(refresh_rate=refresh_rate)
+    assert bar.refresh_rate == expected
+
+
 @pytest.mark.parametrize(
     ("train_batches", "val_batches", "refresh_rate", "train_updates", "val_updates"),
     [