Compact output in snapshot_download and hf download (#3523)

* Multi-threaded snapshot download

* Compact output in snapshot_download

* comment

Author: Wauplin

src/huggingface_hub/_snapshot_download.py

@@ -17,7 +17,7 @@
 )
 from .file_download import REGEX_COMMIT_HASH, DryRunFileInfo, hf_hub_download, repo_folder_name
 from .hf_api import DatasetInfo, HfApi, ModelInfo, RepoFile, SpaceInfo
-from .utils import OfflineModeIsEnabled, filter_repo_objects, logging, validate_hf_hub_args
+from .utils import OfflineModeIsEnabled, filter_repo_objects, is_tqdm_disabled, logging, validate_hf_hub_args
 from .utils import tqdm as hf_tqdm
 
 
@@ -379,12 +379,56 @@ def snapshot_download(
 
     results: List[Union[str, DryRunFileInfo]] = []
 
+    # User can use its own tqdm class or the default one from `huggingface_hub.utils`
+    tqdm_class = tqdm_class or hf_tqdm
+
+    # Create a progress bar for the bytes downloaded
+    # This progress bar is shared across threads/files and gets updated each time we fetch
+    # metadata for a file.
+    bytes_progress = tqdm_class(
+        desc="Downloading (incomplete total...)",
+        disable=is_tqdm_disabled(log_level=logger.getEffectiveLevel()),
+        total=0,
+        initial=0,
+        unit="B",
+        unit_scale=True,
+        name="huggingface_hub.snapshot_download",
+    )
+
+    class _AggregatedTqdm:
+        """Fake tqdm object to aggregate progress into the parent `bytes_progress` bar.
+
+        In practice the `_AggregatedTqdm` object won't be displayed, it's just used to update
+        the `bytes_progress` bar from each thread/file download.
+        """
+
+        def __init__(self, *args, **kwargs):
+            # Adjust the total of the parent progress bar
+            total = kwargs.pop("total", None)
+            if total is not None:
+                bytes_progress.total += total
+                bytes_progress.refresh()
+
+            # Adjust initial of the parent progress bar
+            initial = kwargs.pop("initial", 0)
+            if initial:
+                bytes_progress.update(initial)
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            pass
+
+        def update(self, n: Optional[Union[int, float]] = 1) -> None:
+            bytes_progress.update(n)
+
     # we pass the commit_hash to hf_hub_download
     # so no network call happens if we already
     # have the file locally.
     def _inner_hf_hub_download(repo_file: str) -> None:
         results.append(
-            hf_hub_download(  # type: ignore[no-matching-overload] # ty not happy, don't know why :/
+            hf_hub_download(  # type: ignore
                 repo_id,
                 filename=repo_file,
                 repo_type=repo_type,
@@ -399,6 +443,7 @@ def _inner_hf_hub_download(repo_file: str) -> None:
                 force_download=force_download,
                 token=token,
                 headers=headers,
+                tqdm_class=_AggregatedTqdm,  # type: ignore
                 dry_run=dry_run,
             )
         )
@@ -408,10 +453,11 @@ def _inner_hf_hub_download(repo_file: str) -> None:
         filtered_repo_files,
         desc=tqdm_desc,
         max_workers=max_workers,
-        # User can use its own tqdm class or the default one from `huggingface_hub.utils`
-        tqdm_class=tqdm_class or hf_tqdm,
+        tqdm_class=tqdm_class,
     )
 
+    bytes_progress.set_description("Download complete")
+
     if dry_run:
         assert all(isinstance(r, DryRunFileInfo) for r in results)
         return results  # type: ignore

src/huggingface_hub/file_download.py

@@ -13,6 +13,7 @@
 from urllib.parse import quote, urlparse
 
 import httpx
+from tqdm.auto import tqdm as base_tqdm
 
 from . import constants
 from ._local_folder import get_local_download_paths, read_download_metadata, write_download_metadata
@@ -348,6 +349,7 @@ def http_get(
     headers: Optional[dict[str, Any]] = None,
     expected_size: Optional[int] = None,
     displayed_filename: Optional[str] = None,
+    tqdm_class: Optional[type[base_tqdm]] = None,
     _nb_retries: int = 5,
     _tqdm_bar: Optional[tqdm] = None,
 ) -> None:
@@ -425,6 +427,7 @@ def http_get(
             total=total,
             initial=resume_size,
             name="huggingface_hub.http_get",
+            tqdm_class=tqdm_class,
             _tqdm_bar=_tqdm_bar,
         )
 
@@ -453,6 +456,7 @@ def http_get(
                     resume_size=new_resume_size,
                     headers=initial_headers,
                     expected_size=expected_size,
+                    tqdm_class=tqdm_class,
                     _nb_retries=_nb_retries - 1,
                     _tqdm_bar=_tqdm_bar,
                 )
@@ -472,6 +476,7 @@ def xet_get(
     headers: dict[str, str],
     expected_size: Optional[int] = None,
     displayed_filename: Optional[str] = None,
+    tqdm_class: Optional[type[base_tqdm]] = None,
     _tqdm_bar: Optional[tqdm] = None,
 ) -> None:
     """
@@ -554,6 +559,7 @@ def token_refresher() -> tuple[str, int]:
         total=expected_size,
         initial=0,
         name="huggingface_hub.xet_get",
+        tqdm_class=tqdm_class,
         _tqdm_bar=_tqdm_bar,
     )
 
@@ -685,10 +691,10 @@ def _create_symlink(src: str, dst: str, new_blob: bool = False) -> None:
 
     # Symlinks are not supported => let's move or copy the file.
     if new_blob:
-        logger.info(f"Symlink not supported. Moving file from {abs_src} to {abs_dst}")
+        logger.debug(f"Symlink not supported. Moving file from {abs_src} to {abs_dst}")
         shutil.move(abs_src, abs_dst, copy_function=_copy_no_matter_what)
     else:
-        logger.info(f"Symlink not supported. Copying file from {abs_src} to {abs_dst}")
+        logger.debug(f"Symlink not supported. Copying file from {abs_src} to {abs_dst}")
         shutil.copyfile(abs_src, abs_dst)
 
 
@@ -763,6 +769,7 @@ def hf_hub_download(
     local_files_only: bool = False,
     headers: Optional[dict[str, str]] = None,
     endpoint: Optional[str] = None,
+    tqdm_class: Optional[type[base_tqdm]] = None,
     dry_run: Literal[False] = False,
 ) -> str: ...
 
@@ -786,6 +793,7 @@ def hf_hub_download(
     local_files_only: bool = False,
     headers: Optional[dict[str, str]] = None,
     endpoint: Optional[str] = None,
+    tqdm_class: Optional[type[base_tqdm]] = None,
     dry_run: Literal[True] = True,
 ) -> DryRunFileInfo: ...
 
@@ -809,6 +817,7 @@ def hf_hub_download(
     local_files_only: bool = False,
     headers: Optional[dict[str, str]] = None,
     endpoint: Optional[str] = None,
+    tqdm_class: Optional[type[base_tqdm]] = None,
     dry_run: bool = False,
 ) -> Union[str, DryRunFileInfo]: ...
 
@@ -832,6 +841,7 @@ def hf_hub_download(
     local_files_only: bool = False,
     headers: Optional[dict[str, str]] = None,
     endpoint: Optional[str] = None,
+    tqdm_class: Optional[type[base_tqdm]] = None,
     dry_run: bool = False,
 ) -> Union[str, DryRunFileInfo]:
     """Download a given file if it's not already present in the local cache.
@@ -908,6 +918,11 @@ def hf_hub_download(
             local cached file if it exists.
         headers (`dict`, *optional*):
             Additional headers to be sent with the request.
+        tqdm_class (`tqdm`, *optional*):
+            If provided, overwrites the default behavior for the progress bar. Passed
+            argument must inherit from `tqdm.auto.tqdm` or at least mimic its behavior.
+            Defaults to the custom HF progress bar that can be disabled by setting
+            `HF_HUB_DISABLE_PROGRESS_BARS` environment variable.
         dry_run (`bool`, *optional*, defaults to `False`):
             If `True`, perform a dry run without actually downloading the file. Returns a
             [`DryRunFileInfo`] object containing information about what would be downloaded.
@@ -985,6 +1000,7 @@ def hf_hub_download(
             cache_dir=cache_dir,
             force_download=force_download,
             local_files_only=local_files_only,
+            tqdm_class=tqdm_class,
             dry_run=dry_run,
         )
     else:
@@ -1004,6 +1020,7 @@ def hf_hub_download(
             # Additional options
             local_files_only=local_files_only,
             force_download=force_download,
+            tqdm_class=tqdm_class,
             dry_run=dry_run,
         )
 
@@ -1025,6 +1042,7 @@ def _hf_hub_download_to_cache_dir(
     # Additional options
     local_files_only: bool,
     force_download: bool,
+    tqdm_class: Optional[type[base_tqdm]],
     dry_run: bool,
 ) -> Union[str, DryRunFileInfo]:
     """Download a given file to a cache folder, if not already present.
@@ -1189,6 +1207,7 @@ def _hf_hub_download_to_cache_dir(
             force_download=force_download,
             etag=etag,
             xet_file_data=xet_file_data,
+            tqdm_class=tqdm_class,
         )
         if not os.path.exists(pointer_path):
             _create_symlink(blob_path, pointer_path, new_blob=True)
@@ -1214,6 +1233,7 @@ def _hf_hub_download_to_local_dir(
     cache_dir: str,
     force_download: bool,
     local_files_only: bool,
+    tqdm_class: Optional[type[base_tqdm]],
     dry_run: bool,
 ) -> Union[str, DryRunFileInfo]:
     """Download a given file to a local folder, if not already present.
@@ -1377,6 +1397,7 @@ def _hf_hub_download_to_local_dir(
             force_download=force_download,
             etag=etag,
             xet_file_data=xet_file_data,
+            tqdm_class=tqdm_class,
         )
 
     write_download_metadata(local_dir=local_dir, filename=filename, commit_hash=commit_hash, etag=etag)
@@ -1727,6 +1748,7 @@ def _download_to_tmp_and_move(
     force_download: bool,
     etag: Optional[str],
     xet_file_data: Optional[XetFileData],
+    tqdm_class: Optional[type[base_tqdm]] = None,
 ) -> None:
     """Download content from a URL to a destination path.
 
@@ -1749,15 +1771,15 @@ def _download_to_tmp_and_move(
         # By default, we will try to resume the download if possible.
         # However, if the user has set `force_download=True`, then we should
         # not resume the download => delete the incomplete file.
-        logger.info(f"Removing incomplete file '{incomplete_path}' (force_download=True)")
+        logger.debug(f"Removing incomplete file '{incomplete_path}' (force_download=True)")
         incomplete_path.unlink(missing_ok=True)
 
     with incomplete_path.open("ab") as f:
         resume_size = f.tell()
         message = f"Downloading '{filename}' to '{incomplete_path}'"
         if resume_size > 0 and expected_size is not None:
             message += f" (resume from {resume_size}/{expected_size})"
-        logger.info(message)
+        logger.debug(message)
 
         if expected_size is not None:  # might be None if HTTP header not set correctly
             # Check disk space in both tmp and destination path
@@ -1772,6 +1794,7 @@ def _download_to_tmp_and_move(
                 headers=headers,
                 expected_size=expected_size,
                 displayed_filename=filename,
+                tqdm_class=tqdm_class,
             )
         else:
             if xet_file_data is not None and not constants.HF_HUB_DISABLE_XET:
@@ -1787,9 +1810,10 @@ def _download_to_tmp_and_move(
                 resume_size=resume_size,
                 headers=headers,
                 expected_size=expected_size,
+                tqdm_class=tqdm_class,
             )
 
-    logger.info(f"Download complete. Moving file to {destination_path}")
+    logger.debug(f"Download complete. Moving file to {destination_path}")
     _chmod_and_move(incomplete_path, destination_path)
 
 

src/huggingface_hub/utils/__init__.py

@@ -119,4 +119,11 @@
     parse_xet_file_data_from_response,
     refresh_xet_connection_info,
 )
-from .tqdm import are_progress_bars_disabled, disable_progress_bars, enable_progress_bars, tqdm, tqdm_stream_file
+from .tqdm import (
+    are_progress_bars_disabled,
+    disable_progress_bars,
+    enable_progress_bars,
+    is_tqdm_disabled,
+    tqdm,
+    tqdm_stream_file,
+)

src/huggingface_hub/utils/tqdm.py

@@ -288,6 +288,7 @@ def _get_progress_bar_context(
     unit: str = "B",
     unit_scale: bool = True,
     name: Optional[str] = None,
+    tqdm_class: Optional[type[old_tqdm]] = None,
     _tqdm_bar: Optional[tqdm] = None,
 ) -> ContextManager[tqdm]:
     if _tqdm_bar is not None:
@@ -296,7 +297,7 @@ def _get_progress_bar_context(
         #   Makes it easier to use the same code path for both cases but in the later
         #   case, the progress bar is not closed when exiting the context manager.
 
-    return tqdm(
+    return (tqdm_class or tqdm)(  # type: ignore[return-value]
         unit=unit,
         unit_scale=unit_scale,
         total=total,