Skip to content

Commit

Permalink
Support TQDM_MINITERS env variable (#19381)
Browse files Browse the repository at this point in the history
Co-authored-by: awaelchli <[email protected]>
  • Loading branch information
wouterzwerink and awaelchli authored Feb 1, 2024
1 parent c346f4d commit 5d178d0
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 0 deletions.
3 changes: 3 additions & 0 deletions src/lightning/pytorch/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
3 changes: 3 additions & 0 deletions src/lightning/pytorch/callbacks/progress/tqdm_progress.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
13 changes: 13 additions & 0 deletions tests/tests_pytorch/callbacks/progress/test_tqdm_progress_bar.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
[
Expand Down

0 comments on commit 5d178d0

Please sign in to comment.