Spaces:
Running on L40S
Running on L40S
File size: 2,097 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""
Progress bar wrapper that gets automatically disabled when in a Timer region, or any other context
where we'd want to disable progress bars, including when TQDM is not present, or when user sets
DISABLE_TQDM=1. We can eventually add a simple ascii progress bar as fallback for missing
dependencies.
"""
import os
from cosmos_framework.utils import distributed
from cosmos_framework.utils.timer import in_timer_region
try:
import tqdm as _tqdm # noqa: F401
HAS_TQDM = True
except ImportError:
HAS_TQDM = False
except Exception as e:
HAS_TQDM = True
def _tqdm_wrapper(*args, **kwargs):
if HAS_TQDM:
import tqdm
return tqdm.tqdm(*args, **kwargs)
raise ImportError("TQDM is not installed. Please install it and try again.")
def progress_bar(fn, desc=None, total=None, force_display: bool = False):
"""
Progress bars a great, but they're not for everybody, certainly not for everywhere.
They must be guarded against:
* We're benchmarking performance (with Timer)
* If tqdm / other progress bars aren't available, skip instead of failing.
* If multi-process / GPU, only one (usually rank 0) must display it, just like prints.
* If the user just doesn't want progress bars (toggle via environment variables.
This function consideres all of those cases
"""
disable_tqdm = os.environ.get("DISABLE_TQDM", "0") == "1"
is_in_timer_region = in_timer_region()
is_rank0 = True
# Wide-scope try/except on determining rank, in case distributed context is uninitialized in a
# single-process program. If exception occurs, it's better to just assume single-process.
try:
is_rank0 = distributed.get_rank() == 0
except Exception as e:
pass
if not force_display and (not is_rank0 or is_in_timer_region or disable_tqdm):
return fn
return _tqdm_wrapper(fn, desc=desc, total=total)
__all__ = ["progress_bar"]
|