codekingpro/portable-devtools

1# Copyright 2021 The HuggingFace Inc. team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License14"""Utility helpers to handle progress bars in `huggingface_hub`.15 16Example:17    1. Use `huggingface_hub.utils.tqdm` as you would use `tqdm.tqdm` or `tqdm.auto.tqdm`.18    2. To disable progress bars, either use `disable_progress_bars()` helper or set the19       environment variable `HF_HUB_DISABLE_PROGRESS_BARS` to 1.20    3. To re-enable progress bars, use `enable_progress_bars()`.21    4. To check whether progress bars are disabled, use `are_progress_bars_disabled()`.22 23NOTE: Environment variable `HF_HUB_DISABLE_PROGRESS_BARS` has the priority.24 25Example:26    ```py27    >>> from huggingface_hub.utils import are_progress_bars_disabled, disable_progress_bars, enable_progress_bars, tqdm28 29    # Disable progress bars globally30    >>> disable_progress_bars()31 32    # Use as normal `tqdm`33    >>> for _ in tqdm(range(5)):34    ...    pass35 36    # Still not showing progress bars, as `disable=False` is overwritten to `True`.37    >>> for _ in tqdm(range(5), disable=False):38    ...    pass39 40    >>> are_progress_bars_disabled()41    True42 43    # Re-enable progress bars globally44    >>> enable_progress_bars()45 46    # Progress bar will be shown !47    >>> for _ in tqdm(range(5)):48    ...   pass49    100%|███████████████████████████████████████| 5/5 [00:00<00:00, 117817.53it/s]50    ```51 52Group-based control:53    ```python54    # Disable progress bars for a specific group55    >>> disable_progress_bars("peft.foo")56 57    # Check state of different groups58    >>> assert not are_progress_bars_disabled("peft"))59    >>> assert not are_progress_bars_disabled("peft.something")60    >>> assert are_progress_bars_disabled("peft.foo"))61    >>> assert are_progress_bars_disabled("peft.foo.bar"))62 63    # Enable progress bars for a subgroup64    >>> enable_progress_bars("peft.foo.bar")65 66    # Check if enabling a subgroup affects the parent group67    >>> assert are_progress_bars_disabled("peft.foo"))68    >>> assert not are_progress_bars_disabled("peft.foo.bar"))69 70    # No progress bar for `name="peft.foo"`71    >>> for _ in tqdm(range(5), name="peft.foo"):72    ...     pass73 74    # Progress bar will be shown for `name="peft.foo.bar"`75    >>> for _ in tqdm(range(5), name="peft.foo.bar"):76    ...     pass77    100%|███████████████████████████████████████| 5/5 [00:00<00:00, 117817.53it/s]78 79    ```80"""81 82import io83import logging84import os85import threading86import warnings87from collections.abc import Iterator88from contextlib import contextmanager, nullcontext89from pathlib import Path90from typing import ContextManager91 92from tqdm.auto import tqdm as old_tqdm93 94from ..constants import HF_HUB_DISABLE_PROGRESS_BARS95 96 97# The `HF_HUB_DISABLE_PROGRESS_BARS` environment variable can be True, False, or not set (None),98# allowing for control over progress bar visibility. When set, this variable takes precedence99# over programmatic settings, dictating whether progress bars should be shown or hidden globally.100# Essentially, the environment variable's setting overrides any code-based configurations.101#102# If `HF_HUB_DISABLE_PROGRESS_BARS` is not defined (None), it implies that users can manage103# progress bar visibility through code. By default, progress bars are turned on.104 105progress_bar_states: dict[str, bool] = {}106 107 108class disable_progress_bars:109    """110    Disable progress bars either globally or for a specified group.111 112    This function updates the state of progress bars based on a group name.113    If no group name is provided, all progress bars are disabled. The operation114    respects the `HF_HUB_DISABLE_PROGRESS_BARS` environment variable's setting.115 116    Works as both a regular call and a context manager:117        disable_progress_bars()           # disables until enable_progress_bars()118        with disable_progress_bars():     # disables for the block, re-enables on exit119            ...120 121    Args:122        name (`str`, *optional*):123            The name of the group for which to disable the progress bars. If None,124            progress bars are disabled globally.125 126    Raises:127        Warning: If the environment variable precludes changes.128    """129 130    def __init__(self, name: str | None = None) -> None:131        self.name = name132 133        if HF_HUB_DISABLE_PROGRESS_BARS is False:134            warnings.warn(135                "Cannot disable progress bars: environment variable `HF_HUB_DISABLE_PROGRESS_BARS=0` is set and has priority."136            )137            self._should_reenable = False138            return139 140        self._should_reenable = not are_progress_bars_disabled(name)141        if name is None:142            progress_bar_states.clear()143            progress_bar_states["_global"] = False144        else:145            keys_to_remove = [key for key in progress_bar_states if key.startswith(f"{name}.")]146            for key in keys_to_remove:147                del progress_bar_states[key]148            progress_bar_states[name] = False149 150    def __enter__(self) -> "disable_progress_bars":151        return self152 153    def __exit__(self, *exc) -> None:154        if self._should_reenable:155            enable_progress_bars(self.name)156 157 158def enable_progress_bars(name: str | None = None) -> None:159    """160    Enable progress bars either globally or for a specified group.161 162    This function sets the progress bars to enabled for the specified group or globally163    if no group is specified. The operation is subject to the `HF_HUB_DISABLE_PROGRESS_BARS`164    environment setting.165 166    Args:167        name (`str`, *optional*):168            The name of the group for which to enable the progress bars. If None,169            progress bars are enabled globally.170 171    Raises:172        Warning: If the environment variable precludes changes.173    """174    if HF_HUB_DISABLE_PROGRESS_BARS is True:175        warnings.warn(176            "Cannot enable progress bars: environment variable `HF_HUB_DISABLE_PROGRESS_BARS=1` is set and has priority."177        )178        return179 180    if name is None:181        progress_bar_states.clear()182        progress_bar_states["_global"] = True183    else:184        keys_to_remove = [key for key in progress_bar_states if key.startswith(f"{name}.")]185        for key in keys_to_remove:186            del progress_bar_states[key]187        progress_bar_states[name] = True188 189 190def are_progress_bars_disabled(name: str | None = None) -> bool:191    """192    Check if progress bars are disabled globally or for a specific group.193 194    This function returns whether progress bars are disabled for a given group or globally.195    It checks the `HF_HUB_DISABLE_PROGRESS_BARS` environment variable first, then the programmatic196    settings.197 198    Args:199        name (`str`, *optional*):200            The group name to check; if None, checks the global setting.201 202    Returns:203        `bool`: True if progress bars are disabled, False otherwise.204    """205    if HF_HUB_DISABLE_PROGRESS_BARS is True:206        return True207 208    if name is None:209        return not progress_bar_states.get("_global", True)210 211    while name:212        if name in progress_bar_states:213            return not progress_bar_states[name]214        name = ".".join(name.split(".")[:-1])215 216    return not progress_bar_states.get("_global", True)217 218 219def is_tqdm_disabled(log_level: int) -> bool | None:220    """221    Determine if tqdm progress bars should be disabled based on logging level and environment settings.222 223    see https://github.com/huggingface/huggingface_hub/pull/2000 and https://github.com/huggingface/huggingface_hub/pull/2698.224    """225    if log_level == logging.NOTSET:226        return True227    if os.getenv("TQDM_POSITION") == "-1":228        return False229    return None230 231 232class tqdm(old_tqdm):233    """234    Class to override `disable` argument in case progress bars are globally disabled.235 236    Taken from https://github.com/tqdm/tqdm/issues/619#issuecomment-619639324.237    """238 239    def __init__(self, *args, **kwargs):240        name = kwargs.pop("name", None)  # do not pass `name` to `tqdm`241        if are_progress_bars_disabled(name):242            kwargs["disable"] = True243        super().__init__(*args, **kwargs)244 245    def __delattr__(self, attr: str) -> None:246        """Fix for https://github.com/huggingface/huggingface_hub/issues/1603"""247        try:248            super().__delattr__(attr)249        except AttributeError:250            if attr != "_lock":251                raise252 253 254# Prevent tqdm's default multiprocessing write-lock from spawning a resource255# tracker subprocess via fork_exec(). That path fails when stderr has an invalid256# fd (e.g. Textual TUIs that return -1 from sys.stderr.fileno()). Inter-process257# bar coordination on the HF subclass is not a supported use case. See #4065.258tqdm.set_lock(threading.RLock())259 260 261class silent_tqdm:262    """Fake tqdm object that does nothing."""263 264    def __init__(self, *args, **kwargs):265        pass266 267    def __enter__(self):268        return self269 270    def __exit__(self, exc_type, exc_value, traceback):271        pass272 273    def update(self, n: int | float | None = 1) -> None:274        pass275 276 277@contextmanager278def tqdm_stream_file(path: Path | str) -> Iterator[io.BufferedReader]:279    """280    Open a file as binary and wrap the `read` method to display a progress bar when it's streamed.281 282    First implemented in `transformers` in 2019 but removed when switched to git-lfs. Used in `huggingface_hub` to show283    progress bar when uploading an LFS file to the Hub. See github.com/huggingface/transformers/pull/2078#discussion_r354739608284    for implementation details.285 286    Note: currently implementation handles only files stored on disk as it is the most common use case. Could be287          extended to stream any `BinaryIO` object but we might have to debug some corner cases.288 289    Example:290    ```py291    >>> with tqdm_stream_file("config.json") as f:292    >>>     httpx.put(url, data=f)293    config.json: 100%|█████████████████████████| 8.19k/8.19k [00:02<00:00, 3.72kB/s]294    ```295    """296    if isinstance(path, str):297        path = Path(path)298 299    with path.open("rb") as f:300        total_size = path.stat().st_size301        pbar = tqdm(302            unit="B",303            unit_scale=True,304            total=total_size,305            initial=0,306            desc=path.name,307        )308 309        f_read = f.read310 311        def _inner_read(size: int | None = -1) -> bytes:312            data = f_read(size)313            pbar.update(len(data))314            return data315 316        f.read = _inner_read  # type: ignore317 318        yield f319 320        pbar.close()321 322 323def _create_progress_bar(*, cls: type[old_tqdm], log_level: int, name: str | None = None, **kwargs) -> old_tqdm:324    """Create a progress bar.325 326    For our `tqdm` subclass (or subclasses of it): respects all disable signals327    (`HF_HUB_DISABLE_PROGRESS_BARS`, `disable_progress_bars()`, log level) and uses328    `disable=None` for TTY auto-detection (see https://github.com/huggingface/huggingface_hub/pull/2000),329    unless `TQDM_POSITION=-1` forces bars on (https://github.com/huggingface/huggingface_hub/pull/2698).330 331    For other classes: does not inject `disable` or `name`. the custom class is fully332    responsible for its own behavior. Vanilla tqdm defaults to `disable=False` (bar shows).333    Omits `name` which vanilla tqdm rejects with `TqdmKeyError`. See https://github.com/huggingface/huggingface_hub/issues/4050.334    """335    # issubclass() crashes on non-class callables (e.g. functools.partial), guard with isinstance.336    if not (isinstance(cls, type) and issubclass(cls, tqdm)):337        return cls(**kwargs)  # type: ignore[return-value]338 339    # HF subclass: keep the historical log-level / TTY behavior. Group-based340    # disabling is already handled in `tqdm.__init__`.341    disable = is_tqdm_disabled(log_level)342    return cls(disable=disable, name=name, **kwargs)  # type: ignore[return-value]343 344 345def _get_progress_bar_context(346    *,347    desc: str,348    log_level: int,349    total: int | None = None,350    initial: int = 0,351    unit: str = "B",352    unit_scale: bool = True,353    name: str | None = None,354    tqdm_class: type[old_tqdm] | None = None,355    _tqdm_bar: tqdm | None = None,356) -> ContextManager[tqdm]:357    if _tqdm_bar is not None:358        return nullcontext(_tqdm_bar)359        # ^ `contextlib.nullcontext` mimics a context manager that does nothing360        #   Makes it easier to use the same code path for both cases but in the later361        #   case, the progress bar is not closed when exiting the context manager.362 363    return _create_progress_bar(  # type: ignore364        cls=tqdm_class or tqdm,365        log_level=log_level,366        name=name,367        unit=unit,368        unit_scale=unit_scale,369        total=total,370        initial=initial,371        desc=desc,372    )373