Fast download in hf file system (#2143)

* first draft

* Faster downloads in HfFileSystem

* Apply suggestions from code review

Co-authored-by: Quentin Lhoest <[email protected]>

* speed-up get_file and read

* add tests

* fix test

* fix revision

* fix on windows

* simplified file.read

* Update src/huggingface_hub/hf_file_system.py

Co-authored-by: Quentin Lhoest <[email protected]>

---------

Co-authored-by: Quentin Lhoest <[email protected]>

Author: Wauplin

src/huggingface_hub/file_download.py

@@ -404,6 +404,7 @@ def http_get(
     expected_size: Optional[int] = None,
     displayed_filename: Optional[str] = None,
     _nb_retries: int = 5,
+    _tqdm_bar: Optional[tqdm] = None,
 ) -> None:
     """
     Download a remote file. Do not gobble up errors, and will return errors tailored to the Hugging Face Hub.
@@ -483,84 +484,90 @@ def http_get(
     )
 
     # Stream file to buffer
-    with tqdm(
-        unit="B",
-        unit_scale=True,
-        total=total,
-        initial=resume_size,
-        desc=displayed_filename,
-        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
-    ) as progress:
-        if hf_transfer and total is not None and total > 5 * DOWNLOAD_CHUNK_SIZE:
-            supports_callback = "callback" in inspect.signature(hf_transfer.download).parameters
-            if not supports_callback:
-                warnings.warn(
-                    "You are using an outdated version of `hf_transfer`. "
-                    "Consider upgrading to latest version to enable progress bars "
-                    "using `pip install -U hf_transfer`."
-                )
-            try:
-                hf_transfer.download(
-                    url=url,
-                    filename=temp_file.name,
-                    max_files=HF_TRANSFER_CONCURRENCY,
-                    chunk_size=DOWNLOAD_CHUNK_SIZE,
-                    headers=headers,
-                    parallel_failures=3,
-                    max_retries=5,
-                    **({"callback": progress.update} if supports_callback else {}),
-                )
-            except Exception as e:
-                raise RuntimeError(
-                    "An error occurred while downloading using `hf_transfer`. Consider"
-                    " disabling HF_HUB_ENABLE_HF_TRANSFER for better error handling."
-                ) from e
-            if not supports_callback:
-                progress.update(total)
-            if expected_size is not None and expected_size != os.path.getsize(temp_file.name):
-                raise EnvironmentError(
-                    consistency_error_message.format(
-                        actual_size=os.path.getsize(temp_file.name),
-                    )
-                )
-            return
-        new_resume_size = resume_size
+    progress = _tqdm_bar
+    if progress is None:
+        progress = tqdm(
+            unit="B",
+            unit_scale=True,
+            total=total,
+            initial=resume_size,
+            desc=displayed_filename,
+            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
+        )
+
+    if hf_transfer and total is not None and total > 5 * DOWNLOAD_CHUNK_SIZE:
+        supports_callback = "callback" in inspect.signature(hf_transfer.download).parameters
+        if not supports_callback:
+            warnings.warn(
+                "You are using an outdated version of `hf_transfer`. "
+                "Consider upgrading to latest version to enable progress bars "
+                "using `pip install -U hf_transfer`."
+            )
         try:
-            for chunk in r.iter_content(chunk_size=DOWNLOAD_CHUNK_SIZE):
-                if chunk:  # filter out keep-alive new chunks
-                    progress.update(len(chunk))
-                    temp_file.write(chunk)
-                    new_resume_size += len(chunk)
-                    # Some data has been downloaded from the server so we reset the number of retries.
-                    _nb_retries = 5
-        except (requests.ConnectionError, requests.ReadTimeout) as e:
-            # If ConnectionError (SSLError) or ReadTimeout happen while streaming data from the server, it is most likely
-            # a transient error (network outage?). We log a warning message and try to resume the download a few times
-            # before giving up. Tre retry mechanism is basic but should be enough in most cases.
-            if _nb_retries <= 0:
-                logger.warning("Error while downloading from %s: %s\nMax retries exceeded.", url, str(e))
-                raise
-            logger.warning("Error while downloading from %s: %s\nTrying to resume download...", url, str(e))
-            time.sleep(1)
-            reset_sessions()  # In case of SSLError it's best to reset the shared requests.Session objects
-            return http_get(
+            hf_transfer.download(
                 url=url,
-                temp_file=temp_file,
-                proxies=proxies,
-                resume_size=new_resume_size,
-                headers=initial_headers,
-                expected_size=expected_size,
-                _nb_retries=_nb_retries - 1,
+                filename=temp_file.name,
+                max_files=HF_TRANSFER_CONCURRENCY,
+                chunk_size=DOWNLOAD_CHUNK_SIZE,
+                headers=headers,
+                parallel_failures=3,
+                max_retries=5,
+                **({"callback": progress.update} if supports_callback else {}),
             )
-
-        if expected_size is not None and expected_size != temp_file.tell():
+        except Exception as e:
+            raise RuntimeError(
+                "An error occurred while downloading using `hf_transfer`. Consider"
+                " disabling HF_HUB_ENABLE_HF_TRANSFER for better error handling."
+            ) from e
+        if not supports_callback:
+            progress.update(total)
+        if expected_size is not None and expected_size != os.path.getsize(temp_file.name):
             raise EnvironmentError(
                 consistency_error_message.format(
-                    actual_size=temp_file.tell(),
+                    actual_size=os.path.getsize(temp_file.name),
                 )
             )
+        return
+    new_resume_size = resume_size
+    try:
+        for chunk in r.iter_content(chunk_size=DOWNLOAD_CHUNK_SIZE):
+            if chunk:  # filter out keep-alive new chunks
+                progress.update(len(chunk))
+                temp_file.write(chunk)
+                new_resume_size += len(chunk)
+                # Some data has been downloaded from the server so we reset the number of retries.
+                _nb_retries = 5
+    except (requests.ConnectionError, requests.ReadTimeout) as e:
+        # If ConnectionError (SSLError) or ReadTimeout happen while streaming data from the server, it is most likely
+        # a transient error (network outage?). We log a warning message and try to resume the download a few times
+        # before giving up. Tre retry mechanism is basic but should be enough in most cases.
+        if _nb_retries <= 0:
+            logger.warning("Error while downloading from %s: %s\nMax retries exceeded.", url, str(e))
+            raise
+        logger.warning("Error while downloading from %s: %s\nTrying to resume download...", url, str(e))
+        time.sleep(1)
+        reset_sessions()  # In case of SSLError it's best to reset the shared requests.Session objects
+        return http_get(
+            url=url,
+            temp_file=temp_file,
+            proxies=proxies,
+            resume_size=new_resume_size,
+            headers=initial_headers,
+            expected_size=expected_size,
+            _nb_retries=_nb_retries - 1,
+            _tqdm_bar=_tqdm_bar,
+        )
+
+    progress.close()
+
+    if expected_size is not None and expected_size != temp_file.tell():
+        raise EnvironmentError(
+            consistency_error_message.format(
+                actual_size=temp_file.tell(),
+            )
+        )
 
 
 @validate_hf_hub_args

src/huggingface_hub/hf_file_system.py

@@ -6,15 +6,24 @@
 from dataclasses import dataclass, field
 from datetime import datetime
 from itertools import chain
+from pathlib import Path
 from typing import Any, Dict, List, NoReturn, Optional, Tuple, Union
 from urllib.parse import quote, unquote
 
 import fsspec
+from fsspec.callbacks import _DEFAULT_CALLBACK, NoOpCallback, TqdmCallback
+from fsspec.utils import isfilelike
 from requests import Response
 
 from ._commit_api import CommitOperationCopy, CommitOperationDelete
-from .constants import DEFAULT_REVISION, ENDPOINT, REPO_TYPE_MODEL, REPO_TYPES_MAPPING, REPO_TYPES_URL_PREFIXES
-from .file_download import hf_hub_url
+from .constants import (
+    DEFAULT_REVISION,
+    ENDPOINT,
+    REPO_TYPE_MODEL,
+    REPO_TYPES_MAPPING,
+    REPO_TYPES_URL_PREFIXES,
+)
+from .file_download import hf_hub_url, http_get
 from .hf_api import HfApi, LastCommitInfo, RepoFile
 from .utils import (
     EntryNotFoundError,
@@ -591,6 +600,58 @@ def url(self, path: str) -> str:
             url = url.replace("/resolve/", "/tree/", 1)
         return url
 
+    def get_file(self, rpath, lpath, callback=_DEFAULT_CALLBACK, outfile=None, **kwargs) -> None:
+        """Copy single remote file to local."""
+        revision = kwargs.get("revision")
+        unhandled_kwargs = set(kwargs.keys()) - {"revision"}
+        if not isinstance(callback, (NoOpCallback, TqdmCallback)) or len(unhandled_kwargs) > 0:
+            # for now, let's not handle custom callbacks
+            # and let's not handle custom kwargs
+            return super().get_file(rpath, lpath, callback=callback, outfile=outfile, **kwargs)
+
+        # Taken from https://github.com/fsspec/filesystem_spec/blob/47b445ae4c284a82dd15e0287b1ffc410e8fc470/fsspec/spec.py#L883
+        if isfilelike(lpath):
+            outfile = lpath
+        elif self.isdir(rpath):
+            os.makedirs(lpath, exist_ok=True)
+            return None
+
+        if isinstance(lpath, (str, Path)):  # otherwise, let's assume it's a file-like object
+            os.makedirs(os.path.dirname(lpath), exist_ok=True)
+
+        # Open file if not already open
+        close_file = False
+        if outfile is None:
+            outfile = open(lpath, "wb")
+            close_file = True
+        initial_pos = outfile.tell()
+
+        # Custom implementation of `get_file` to use `http_get`.
+        resolve_remote_path = self.resolve_path(rpath, revision=revision)
+        expected_size = self.info(rpath, revision=revision)["size"]
+        callback.set_size(expected_size)
+        try:
+            http_get(
+                url=hf_hub_url(
+                    repo_id=resolve_remote_path.repo_id,
+                    revision=resolve_remote_path.revision,
+                    filename=resolve_remote_path.path_in_repo,
+                    repo_type=resolve_remote_path.repo_type,
+                    endpoint=self.endpoint,
+                ),
+                temp_file=outfile,
+                displayed_filename=rpath,
+                expected_size=expected_size,
+                resume_size=0,
+                headers=self._api._build_hf_headers(),
+                _tqdm_bar=callback.tqdm if isinstance(callback, TqdmCallback) else None,
+            )
+            outfile.seek(initial_pos)
+        finally:
+            # Close file only if we opened it ourselves
+            if close_file:
+                outfile.close()
+
     @property
     def transaction(self):
         """A context within which files are committed together upon exit
@@ -618,6 +679,7 @@ def __init__(self, fs: HfFileSystem, path: str, revision: Optional[str] = None,
                 raise FileNotFoundError(
                     f"{e}.\nMake sure the repository and revision exist before writing data."
                 ) from e
+            raise
         super().__init__(fs, self.resolved_path.unresolve(), **kwargs)
         self.fs: HfFileSystem
 
@@ -667,6 +729,18 @@ def _upload_chunk(self, final: bool = False) -> None:
                 path=self.resolved_path.unresolve(),
             )
 
+    def read(self, length=-1):
+        """Read remote file.
+
+        If `length` is not provided or is -1, the entire file is downloaded and read. On POSIX systems and if
+        `hf_transfer` is not enabled, the file is loaded in memory directly. Otherwise, the file is downloaded to a
+        temporary file and read from there.
+        """
+        if self.mode == "rb" and (length is None or length == -1) and self.loc == 0:
+            with self.fs.open(self.path, "rb", block_size=0) as f:  # block_size=0 enables fast streaming
+                return f.read()
+        return super().read(length)
+
     def url(self) -> str:
         return self.fs.url(self.path)
 
@@ -695,7 +769,7 @@ def __init__(
                 raise FileNotFoundError(
                     f"{e}.\nMake sure the repository and revision exist before writing data."
                 ) from e
-        # avoid an unecessary .info() call to instantiate .details
+        # avoid an unnecessary .info() call to instantiate .details
         self.details = {"name": self.resolved_path.unresolve(), "size": None}
         super().__init__(
             fs, self.resolved_path.unresolve(), mode=mode, block_size=block_size, cache_type=cache_type, **kwargs