Repository: huggingface/huggingface_hub
Pull request #2217 (2024-04-10 18:54:45)

[wip] Implement hierarchical progress bar control in huggingface_hub
As discussed in #1595, i worked on enabling/disabling specific progress bars. 

Changes:
- Added private utility functions `_set_progress_bar_state` and `_get_progress_bar_state` to the `tqdm` module for hierarchical state management via the `progress_bar_states` dict.
- Updated `disable_progress_bars`, `enable_progress_bars`, and `are_progress_bars_disabled` to utilize this hierarchical approach.

Marked as [WIP] since this is an initial draft for discussion. I'm unsure about the implementation of the hierarchical logic and look forward to your feedback, @Wauplin.

HuggingFaceDocBuilderDev: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/huggingface_hub/pr_2217). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.
lappemic: Additionally to the simplified logic i tried to improve the docstrings according to the example in #2209 and added an example for the hierarchical control. I hope this helps the clearity.

Again, thank you very much @Wauplin for the guidance on the logic implementation, i really appreciate you taking time for this!
lappemic: > Hi @lappemic thanks a lot for the changes! I haven't tested it myself yet but it looks much cleaner like this :)
Thanks @Wauplin, this is good to hear. 👍 

> Best way to be sure will be to implement tests! 
Absoultely. I haven't thought about it then. Implemented them now. I hope i got all of the corner cases. It feels like an extensive test at least. (Not really experienced with testing, so happy for every input as well)

> And finally, to make this PR complete you will have to update the documentation.
I updated the documantation. About the wording and example style i was not sure. Just took the one you used in the original issue to showcase it. Let me know if i can adjust something there.

Patch:
diff --git a/docs/source/en/package_reference/utilities.md b/docs/source/en/package_reference/utilities.md
index 313598b2..5ea9d56b 100644
--- a/docs/source/en/package_reference/utilities.md
+++ b/docs/source/en/package_reference/utilities.md
@@ -63,7 +63,6 @@ By default, progress bars are enabled. You can disable them globally by setting
 environment variable. You can also enable/disable them using [`~utils.enable_progress_bars`] and
 [`~utils.disable_progress_bars`]. If set, the environment variable has priority on the helpers.
 
-
 ```py
 >>> from huggingface_hub import snapshot_download
 >>> from huggingface_hub.utils import are_progress_bars_disabled, disable_progress_bars, enable_progress_bars
@@ -81,6 +80,34 @@ True
 >>> enable_progress_bars()
 ```
 
+### Group-specific control of progress bars
+
+You can also enable or disable progress bars for specific groups. This allows you to manage progress bar visibility more granularly within different parts of your application or library. When a progress bar is disabled for a group, all subgroups under it are also affected unless explicitly overridden.
+
+```py
+# Disable progress bars for a specific group
+>>> disable_progress_bars("peft.foo")
+>>> assert not are_progress_bars_disabled("peft")
+>>> assert not are_progress_bars_disabled("peft.something")
+>>> assert are_progress_bars_disabled("peft.foo")
+>>> assert are_progress_bars_disabled("peft.foo.bar")
+
+# Re-enable progress bars for a subgroup
+>>> enable_progress_bars("peft.foo.bar")
+>>> assert are_progress_bars_disabled("peft.foo")
+>>> assert not are_progress_bars_disabled("peft.foo.bar")
+
+# Use groups with tqdm
+# No progress bar for `name="peft.foo"`
+>>> for _ in tqdm(range(5), name="peft.foo"):
+...     pass
+
+# Progress bar will be shown for `name="peft.foo.bar"`
+>>> for _ in tqdm(range(5), name="peft.foo.bar"):
+...     pass
+100%|███████████████████████████████████████| 5/5 [00:00<00:00, 117817.53it/s]
+```
+
 ### are_progress_bars_disabled
 
 [[autodoc]] huggingface_hub.utils.are_progress_bars_disabled
diff --git a/src/huggingface_hub/_commit_api.py b/src/huggingface_hub/_commit_api.py
index d4368bab..b363e4e4 100644
--- a/src/huggingface_hub/_commit_api.py
+++ b/src/huggingface_hub/_commit_api.py
@@ -406,7 +406,7 @@ def _upload_lfs_files(
 
     if HF_HUB_ENABLE_HF_TRANSFER:
         logger.debug(f"Uploading {len(filtered_actions)} LFS files to the Hub using `hf_transfer`.")
-        for action in hf_tqdm(filtered_actions):
+        for action in hf_tqdm(filtered_actions, name="huggingface_hub.lfs_upload"):
             _wrapped_lfs_upload(action)
     elif len(filtered_actions) == 1:
         logger.debug("Uploading 1 LFS file to the Hub")
diff --git a/src/huggingface_hub/file_download.py b/src/huggingface_hub/file_download.py
index 05503c27..23d1d106 100644
--- a/src/huggingface_hub/file_download.py
+++ b/src/huggingface_hub/file_download.py
@@ -498,6 +498,7 @@ def http_get(
             disable=True if (logger.getEffectiveLevel() == logging.NOTSET) else None,
             # ^ set `disable=None` rather than `disable=False` by default to disable progress bar when no TTY attached
             # see https://github.com/huggingface/huggingface_hub/pull/2000
+            name="huggingface_hub.http_get",
         )
 
     if hf_transfer and total is not None and total > 5 * DOWNLOAD_CHUNK_SIZE:
diff --git a/src/huggingface_hub/lfs.py b/src/huggingface_hub/lfs.py
index d6fb9e21..f4322e2a 100644
--- a/src/huggingface_hub/lfs.py
+++ b/src/huggingface_hub/lfs.py
@@ -430,7 +430,15 @@ def _upload_parts_hf_transfer(
     # see https://github.com/huggingface/huggingface_hub/pull/2000
     disable = True if (logger.getEffectiveLevel() == logging.NOTSET) else None
 
-    with tqdm(unit="B", unit_scale=True, total=total, initial=0, desc=desc, disable=disable) as progress:
+    with tqdm(
+        unit="B",
+        unit_scale=True,
+        total=total,
+        initial=0,
+        desc=desc,
+        disable=disable,
+        name="huggingface_hub.lfs_upload",
+    ) as progress:
         try:
             output = multipart_upload(
                 file_path=operation.path_or_fileobj,
diff --git a/src/huggingface_hub/repository.py b/src/huggingface_hub/repository.py
index d7aff2eb..3f06fcaf 100644
--- a/src/huggingface_hub/repository.py
+++ b/src/huggingface_hub/repository.py
@@ -395,6 +395,7 @@ def _lfs_log_progress():
                         unit="B",
                         unit_scale=True,
                         unit_divisor=1024,
+                        name="huggingface_hub.lfs_upload",
                     ),
                     "past_bytes": int(current_bytes),
                 }
diff --git a/src/huggingface_hub/utils/tqdm.py b/src/huggingface_hub/utils/tqdm.py
index da1e421a..13bf7a68 100644
--- a/src/huggingface_hub/utils/tqdm.py
+++ b/src/huggingface_hub/utils/tqdm.py
@@ -26,32 +26,58 @@ NOTE: Environment variable `HF_HUB_DISABLE_PROGRESS_BARS` has the priority.
 
 Example:
     ```py
-    from huggingface_hub.utils import (
-        are_progress_bars_disabled,
-        disable_progress_bars,
-        enable_progress_bars,
-        tqdm,
-    )
+    >>> from huggingface_hub.utils import are_progress_bars_disabled, disable_progress_bars, enable_progress_bars, tqdm
 
     # Disable progress bars globally
-    disable_progress_bars()
+    >>> disable_progress_bars()
 
     # Use as normal `tqdm`
-    for _ in tqdm(range(5)):
-       do_something()
+    >>> for _ in tqdm(range(5)):
+    ...    pass
 
     # Still not showing progress bars, as `disable=False` is overwritten to `True`.
-    for _ in tqdm(range(5), disable=False):
-       do_something()
+    >>> for _ in tqdm(range(5), disable=False):
+    ...    pass
 
-    are_progress_bars_disabled() # True
+    >>> are_progress_bars_disabled()
+    True
 
     # Re-enable progress bars globally
-    enable_progress_bars()
+    >>> enable_progress_bars()
 
     # Progress bar will be shown !
-    for _ in tqdm(range(5)):
-       do_something()
+    >>> for _ in tqdm(range(5)):
+    ...   pass
+    100%|███████████████████████████████████████| 5/5 [00:00<00:00, 117817.53it/s]
+    ```
+
+Group-based control:
+    ```python
+    # Disable progress bars for a specific group
+    >>> disable_progress_bars("peft.foo")
+
+    # Check state of different groups
+    >>> assert not are_progress_bars_disabled("peft"))
+    >>> assert not are_progress_bars_disabled("peft.something")
+    >>> assert are_progress_bars_disabled("peft.foo"))
+    >>> assert are_progress_bars_disabled("peft.foo.bar"))
+
+    # Enable progress bars for a subgroup
+    >>> enable_progress_bars("peft.foo.bar")
+
+    # Check if enabling a subgroup affects the parent group
+    >>> assert are_progress_bars_disabled("peft.foo"))
+    >>> assert not are_progress_bars_disabled("peft.foo.bar"))
+
+    # No progress bar for `name="peft.foo"`
+    >>> for _ in tqdm(range(5), name="peft.foo"):
+    ...     pass
+
+    # Progress bar will be shown for `name="peft.foo.bar"`
+    >>> for _ in tqdm(range(5), name="peft.foo.bar"):
+    ...     pass
+    100%|███████████████████████████████████████| 5/5 [00:00<00:00, 117817.53it/s]
+
     ```
 """
 
@@ -59,65 +85,116 @@ import io
 import warnings
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Iterator, Optional, Union
+from typing import Dict, Iterator, Optional, Union
 
 from tqdm.auto import tqdm as old_tqdm
 
 from ..constants import HF_HUB_DISABLE_PROGRESS_BARS
 
 
-# `HF_HUB_DISABLE_PROGRESS_BARS` is `Optional[bool]` while `_hf_hub_progress_bars_disabled`
-# is a `bool`. If `HF_HUB_DISABLE_PROGRESS_BARS` is set to True or False, it has priority.
-# If `HF_HUB_DISABLE_PROGRESS_BARS` is None, it means the user have not set the
-# environment variable and is free to enable/disable progress bars programmatically.
-# TL;DR: env variable has priority over code.
+# The `HF_HUB_DISABLE_PROGRESS_BARS` environment variable can be True, False, or not set (None),
+# allowing for control over progress bar visibility. When set, this variable takes precedence
+# over programmatic settings, dictating whether progress bars should be shown or hidden globally.
+# Essentially, the environment variable's setting overrides any code-based configurations.
 #
-# By default, progress bars are enabled.
-_hf_hub_progress_bars_disabled: bool = HF_HUB_DISABLE_PROGRESS_BARS or False
+# If `HF_HUB_DISABLE_PROGRESS_BARS` is not defined (None), it implies that users can manage
+# progress bar visibility through code. By default, progress bars are turned on.
+
 
+progress_bar_states: Dict[str, bool] = {}
 
-def disable_progress_bars() -> None:
+
+def disable_progress_bars(name: Optional[str] = None) -> None:
     """
-    Disable globally progress bars used in `huggingface_hub` except if `HF_HUB_DISABLE_PROGRESS_BARS` environment
-    variable has been set.
+    Disable progress bars either globally or for a specified group.
+
+    This function updates the state of progress bars based on a group name.
+    If no group name is provided, all progress bars are disabled. The operation
+    respects the `HF_HUB_DISABLE_PROGRESS_BARS` environment variable's setting.
+
+    Args:
+        name (`str`, *optional*):
+            The name of the group for which to disable the progress bars. If None,
+            progress bars are disabled globally.
 
-    Use [`~utils.enable_progress_bars`] to re-enable them.
+    Raises:
+        Warning: If the environment variable precludes changes.
     """
     if HF_HUB_DISABLE_PROGRESS_BARS is False:
         warnings.warn(
-            "Cannot disable progress bars: environment variable `HF_HUB_DISABLE_PROGRESS_BARS=0` is set and has"
-            " priority."
+            "Cannot disable progress bars: environment variable `HF_HUB_DISABLE_PROGRESS_BARS=0` is set and has priority."
         )
         return
-    global _hf_hub_progress_bars_disabled
-    _hf_hub_progress_bars_disabled = True
 
+    if name is None:
+        progress_bar_states.clear()
+        progress_bar_states["_global"] = False
+    else:
+        keys_to_remove = [key for key in progress_bar_states if key.startswith(f"{name}.")]
+        for key in keys_to_remove:
+            del progress_bar_states[key]
+        progress_bar_states[name] = False
 
-def enable_progress_bars() -> None:
+
+def enable_progress_bars(name: Optional[str] = None) -> None:
     """
-    Enable globally progress bars used in `huggingface_hub` except if `HF_HUB_DISABLE_PROGRESS_BARS` environment
-    variable has been set.
+    Enable progress bars either globally or for a specified group.
+
+    This function sets the progress bars to enabled for the specified group or globally
+    if no group is specified. The operation is subject to the `HF_HUB_DISABLE_PROGRESS_BARS`
+    environment setting.
+
+    Args:
+        name (`str`, *optional*):
+            The name of the group for which to enable the progress bars. If None,
+            progress bars are enabled globally.
 
-    Use [`~utils.disable_progress_bars`] to disable them.
+    Raises:
+        Warning: If the environment variable precludes changes.
     """
     if HF_HUB_DISABLE_PROGRESS_BARS is True:
         warnings.warn(
-            "Cannot enable progress bars: environment variable `HF_HUB_DISABLE_PROGRESS_BARS=1` is set and has"
-            " priority."
+            "Cannot enable progress bars: environment variable `HF_HUB_DISABLE_PROGRESS_BARS=1` is set and has priority."
         )
         return
-    global _hf_hub_progress_bars_disabled
-    _hf_hub_progress_bars_disabled = False
 
+    if name is None:
+        progress_bar_states.clear()
+        progress_bar_states["_global"] = True
+    else:
+        keys_to_remove = [key for key in progress_bar_states if key.startswith(f"{name}.")]
+        for key in keys_to_remove:
+            del progress_bar_states[key]
+        progress_bar_states[name] = True
 
-def are_progress_bars_disabled() -> bool:
-    """Return whether progress bars are globally disabled or not.
 
-    Progress bars used in `huggingface_hub` can be enable or disabled globally using [`~utils.enable_progress_bars`]
-    and [`~utils.disable_progress_bars`] or by setting `HF_HUB_DISABLE_PROGRESS_BARS` as environment variable.
+def are_progress_bars_disabled(name: Optional[str] = None) -> bool:
     """
-    global _hf_hub_progress_bars_disabled
-    return _hf_hub_progress_bars_disabled
+    Check if progress bars are disabled globally or for a specific group.
+
+    This function returns whether progress bars are disabled for a given group or globally.
+    It checks the `HF_HUB_DISABLE_PROGRESS_BARS` environment variable first, then the programmatic
+    settings.
+
+    Args:
+        name (`str`, *optional*):
+            The group name to check; if None, checks the global setting.
+
+    Returns:
+        `bool`: True if progress bars are disabled, False otherwise.
+    """
+    if HF_HUB_DISABLE_PROGRESS_BARS is True:
+        return True
+
+    if name is None:
+        return not progress_bar_states.get("_global", True)
+
+    while name:
+        if name in progress_bar_states:
+            return not progress_bar_states[name]
+        name = ".".join(name.split(".")[:-1])
+
+    return not progress_bar_states.get("_global", True)
 
 
 class tqdm(old_tqdm):
@@ -128,7 +205,8 @@ class tqdm(old_tqdm):
     """
 
     def __init__(self, *args, **kwargs):
-        if are_progress_bars_disabled():
+        name = kwargs.pop("name", None)  # do not pass `name` to `tqdm`
+        if are_progress_bars_disabled(name):
             kwargs["disable"] = True
         super().__init__(*args, **kwargs)
 


Test patch:
diff --git a/tests/test_utils_tqdm.py b/tests/test_utils_tqdm.py
index bbf3dc71..0e936bc8 100644
--- a/tests/test_utils_tqdm.py
+++ b/tests/test_utils_tqdm.py
@@ -16,7 +16,7 @@ from huggingface_hub.utils import (
 )
 
 
-class TestTqdmUtils(unittest.TestCase):
+class CapsysBaseTest(unittest.TestCase):
     @pytest.fixture(autouse=True)
     def capsys(self, capsys: CaptureFixture) -> None:
         """Workaround to make capsys work in unittest framework.
@@ -28,6 +28,8 @@ class TestTqdmUtils(unittest.TestCase):
         """
         self.capsys = capsys
 
+
+class TestTqdmUtils(CapsysBaseTest):
     def setUp(self) -> None:
         """Get verbosity to set it back after the tests."""
         self._previous_are_progress_bars_disabled = are_progress_bars_disabled()
@@ -44,10 +46,10 @@ class TestTqdmUtils(unittest.TestCase):
     def test_tqdm_helpers(self) -> None:
         """Test helpers to enable/disable progress bars."""
         disable_progress_bars()
-        self.assertTrue(are_progress_bars_disabled())
+        assert are_progress_bars_disabled()
 
         enable_progress_bars()
-        self.assertFalse(are_progress_bars_disabled())
+        assert not are_progress_bars_disabled()
 
     @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", True)
     def test_cannot_enable_tqdm_when_env_variable_is_set(self) -> None:
@@ -56,11 +58,11 @@ class TestTqdmUtils(unittest.TestCase):
         `HF_HUB_DISABLE_PROGRESS_BARS` is set.
         """
         disable_progress_bars()
-        self.assertTrue(are_progress_bars_disabled())
+        assert are_progress_bars_disabled()
 
         with self.assertWarns(UserWarning):
             enable_progress_bars()
-        self.assertTrue(are_progress_bars_disabled())  # Still disabled !
+        assert are_progress_bars_disabled()  # Still disabled
 
     @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", False)
     def test_cannot_disable_tqdm_when_env_variable_is_set(self) -> None:
@@ -69,11 +71,11 @@ class TestTqdmUtils(unittest.TestCase):
         `HF_HUB_DISABLE_PROGRESS_BARS` is set.
         """
         enable_progress_bars()
-        self.assertFalse(are_progress_bars_disabled())
+        assert not are_progress_bars_disabled()
 
         with self.assertWarns(UserWarning):
             disable_progress_bars()
-        self.assertFalse(are_progress_bars_disabled())  # Still enabled !
+        assert not are_progress_bars_disabled()  # Still enabled
 
     @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", None)
     def test_tqdm_disabled(self) -> None:
@@ -137,3 +139,97 @@ class TestTqdmUtils(unittest.TestCase):
             self.assertIn("config.json: 100%", captured.err)  # log file name
             self.assertIn("|█████████", captured.err)  # tqdm bar
             self.assertIn("1.00k/1.00k", captured.err)  # size in B
+
+
+class TestTqdmGroup(CapsysBaseTest):
+    def setUp(self):
+        """Set up the initial condition for each test."""
+        super().setUp()
+        enable_progress_bars()  # Ensure all are enabled before each test
+
+    def tearDown(self):
+        """Clean up after each test."""
+        super().tearDown()
+        enable_progress_bars()
+
+    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", None)
+    def test_disable_specific_group(self):
+        """Test disabling a specific group only affects that group and its subgroups."""
+        disable_progress_bars("peft.foo")
+        assert not are_progress_bars_disabled("peft")
+        assert not are_progress_bars_disabled("peft.something")
+        assert are_progress_bars_disabled("peft.foo")
+        assert are_progress_bars_disabled("peft.foo.bar")
+
+    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", None)
+    def test_enable_specific_subgroup(self):
+        """Test that enabling a subgroup does not affect the disabled state of its parent."""
+        disable_progress_bars("peft.foo")
+        enable_progress_bars("peft.foo.bar")
+        assert are_progress_bars_disabled("peft.foo")
+        assert not are_progress_bars_disabled("peft.foo.bar")
+
+    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", True)
+    def test_disable_override_by_environment_variable(self):
+        """Ensure progress bars are disabled regardless of local settings when environment variable is set."""
+        with self.assertWarns(UserWarning):
+            enable_progress_bars()
+        assert are_progress_bars_disabled("peft")
+        assert are_progress_bars_disabled("peft.foo")
+
+    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", False)
+    def test_enable_override_by_environment_variable(self):
+        """Ensure progress bars are enabled regardless of local settings when environment variable is set."""
+        with self.assertWarns(UserWarning):
+            disable_progress_bars("peft.foo")
+        assert not are_progress_bars_disabled("peft.foo")
+
+    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", None)
+    def test_partial_group_name_not_affected(self):
+        """Ensure groups with similar names but not exactly matching are not affected."""
+        disable_progress_bars("peft.foo")
+        assert not are_progress_bars_disabled("peft.footprint")
+
+    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", None)
+    def test_nested_subgroup_behavior(self):
+        """Test enabling and disabling nested subgroups."""
+        disable_progress_bars("peft")
+        enable_progress_bars("peft.foo")
+        disable_progress_bars("peft.foo.bar")
+        assert are_progress_bars_disabled("peft")
+        assert not are_progress_bars_disabled("peft.foo")
+        assert are_progress_bars_disabled("peft.foo.bar")
+
+    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", None)
+    def test_empty_group_is_root(self):
+        """Test the behavior with invalid or empty group names."""
+        disable_progress_bars("")
+        assert not are_progress_bars_disabled("peft")
+
+        enable_progress_bars("123.invalid.name")
+        assert not are_progress_bars_disabled("123.invalid.name")
+
+    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", None)
+    def test_multiple_level_toggling(self):
+        """Test multiple levels of enabling and disabling."""
+        disable_progress_bars("peft")
+        enable_progress_bars("peft.foo")
+        disable_progress_bars("peft.foo.bar.something")
+        assert are_progress_bars_disabled("peft")
+        assert not are_progress_bars_disabled("peft.foo")
+        assert are_progress_bars_disabled("peft.foo.bar.something")
+
+    def test_progress_bar_respects_group(self) -> None:
+        disable_progress_bars("foo.bar")
+        for _ in tqdm(range(10), name="foo.bar.something"):
+            pass
+        captured = self.capsys.readouterr()
+        assert captured.out == ""
+        assert captured.err == ""
+
+        enable_progress_bars("foo.bar.something")
+        for _ in tqdm(range(10), name="foo.bar.something"):
+            pass
+        captured = self.capsys.readouterr()
+        assert captured.out == ""
+        assert "10/10" in captured.err