[Download] Tolerate missing HEAD Content-Length (#4805)

* [Download] Tolerate missing HEAD Content-Length

* [Download] Allow missing HEAD size when Xet is unavailable

---------

Co-authored-by: Celina Hanouti <hanouticelina@gmail.com>

Author: Wauplin

src/huggingface_hub/cli/download.py

@@ -208,15 +208,18 @@ def _print_result(result: str | DryRunFileInfo | list[DryRunFileInfo]) -> None:
         if isinstance(result, DryRunFileInfo):
             result = [result]
         will_download = [r for r in result if r.will_download]
+        total_size = (
+            None if any(r.file_size is None for r in will_download) else sum(r.file_size or 0 for r in will_download)
+        )
         out.text(
             f"[dry-run] Will download {len(will_download)} files"
             f" (out of {len(result)})"
-            f" totalling {_format_size(sum(r.file_size for r in will_download))}."
+            f" totalling {_format_size(total_size) if total_size is not None else 'an unknown size'}."
         )
         items = [
             {
                 "file": info.filename,
-                "size": _format_size(info.file_size) if info.will_download else "-",
+                "size": _format_size(info.file_size) if info.will_download and info.file_size is not None else "-",
             }
             for info in sorted(result, key=lambda x: x.filename)
         ]

src/huggingface_hub/file_download.py

@@ -181,8 +181,9 @@ class DryRunFileInfo:
     Args:
         commit_hash (`str`):
             The commit_hash related to the file.
-        file_size (`int`):
-            Size of the file. In case of an LFS file, contains the size of the actual LFS file, not the pointer.
+        file_size (`int`, *optional*):
+            Size of the file, if known. In case of an LFS file, contains the size of the actual LFS file, not the
+            pointer.
         filename (`str`):
             Name of the file in the repo.
         is_cached (`bool`):
@@ -193,7 +194,7 @@ class DryRunFileInfo:
     """
 
     commit_hash: str
-    file_size: int
+    file_size: int | None
     filename: str
     local_path: str
     is_cached: bool
@@ -407,7 +408,9 @@ def http_get(
                 resume_size = 0
 
             total: int | None = _get_file_length_from_http_response(response)
-            if total is None:
+            if expected_size is None:
+                expected_size = total
+            elif total is None:
                 # Hub serves compressible text files (e.g. vocab.json) with `Content-Encoding: gzip` and
                 # `Transfer-Encoding: chunked`, so the response carries no `Content-Length`. Fall back to the caller's
                 # `expected_size` (always known from the metadata HEAD on the hf_hub path) so the progress bar, and any
@@ -1178,11 +1181,10 @@ def _hf_hub_download_to_cache_dir(
         if head_call_error is not None:
             _raise_on_head_call_error(head_call_error, force_download, local_files_only)
 
-    # From now on, etag, commit_hash, url and size are not None.
+    # From now on, etag, commit_hash and url are not None.
     assert etag is not None, "etag must have been retrieved from server"
     assert commit_hash is not None, "commit_hash must have been retrieved from server"
     assert url_to_download is not None, "file location must have been retrieved from server"
-    assert expected_size is not None, "expected_size must have been retrieved from server"
     blob_path = os.path.join(storage_folder, "blobs", etag)
     pointer_path = _get_pointer_path(storage_folder, commit_hash, relative_filename)
 
@@ -1384,11 +1386,10 @@ def _hf_hub_download_to_local_dir(
         if head_call_error is not None:
             _raise_on_head_call_error(head_call_error, force_download, local_files_only)
 
-    # From now on, etag, commit_hash, url and size are not None.
+    # From now on, etag, commit_hash and url are not None.
     assert etag is not None, "etag must have been retrieved from server"
     assert commit_hash is not None, "commit_hash must have been retrieved from server"
     assert url_to_download is not None, "file location must have been retrieved from server"
-    assert expected_size is not None, "expected_size must have been retrieved from server"
 
     # Local file exists => check if it's up-to-date
     if not force_download and paths.file_path.is_file():
@@ -1675,7 +1676,7 @@ def _get_metadata_or_catch_error(
     |
     # Or the metadata is returned as
     # `(url_to_download, etag, commit_hash, expected_size, xet_file_data, None)`
-    tuple[str, str, str, int, XetFileData | None, None]
+    tuple[str, str, str, int | None, XetFileData | None, None]
 ):
     """Get metadata for a file on the Hub, safely handling network issues.
 
@@ -1767,12 +1768,11 @@ def _get_metadata_or_catch_error(
                     "Distant resource does not have an ETag, we won't be able to reliably ensure reproducibility."
                 )
 
-            # Size must exist
+            # Xet downloads require a known size, but regular HTTP downloads can recover it from the GET response.
             expected_size = metadata.size
-            if expected_size is None:
-                raise FileMetadataError("Distant resource does not have a Content-Length.")
-
             xet_file_data = metadata.xet_file_data
+            if expected_size is None and xet_file_data is not None and is_xet_available():
+                raise FileMetadataError("Distant resource does not have a Content-Length.")
 
             # In case of a redirect, save an extra redirect on the request.get call,
             # and ensure we download the exact atomic version even if it changed

tests/test_file_download.py

@@ -26,7 +26,7 @@
 
 from huggingface_hub import HfApi, constants
 from huggingface_hub._local_folder import write_download_metadata
-from huggingface_hub.errors import EntryNotFoundError, GatedRepoError, LocalEntryNotFoundError
+from huggingface_hub.errors import EntryNotFoundError, FileMetadataError, GatedRepoError, LocalEntryNotFoundError
 from huggingface_hub.file_download import (
     _CACHED_NO_EXIST,
     HfFileMetadata,
@@ -70,6 +70,52 @@
 DATASET_SAMPLE_PY_FILE = "custom_squad.py"
 
 
+@pytest.mark.parametrize("use_local_dir", [False, True])
+@pytest.mark.parametrize("xet_mode", ["no_metadata", "disabled", "not_installed", "enabled"])
+def test_download_without_head_content_length(tmp_path: Path, use_local_dir: bool, xet_mode: str) -> None:
+    content = b"content"
+
+    def _mock_head(*, url: str, **kwargs) -> httpx.Response:
+        headers = {constants.HUGGINGFACE_HEADER_X_REPO_COMMIT: "a" * 40, "ETag": '"etag"'}
+        if xet_mode != "no_metadata":
+            headers[constants.HUGGINGFACE_HEADER_X_XET_HASH] = "b" * 64
+            headers[constants.HUGGINGFACE_HEADER_X_XET_REFRESH_ROUTE] = "https://huggingface.co/xet-refresh"
+        return httpx.Response(
+            200,
+            headers=headers,
+            request=httpx.Request("HEAD", url),
+        )
+
+    @contextmanager
+    def _mock_get(*args, **kwargs):
+        yield httpx.Response(
+            200,
+            headers={"Content-Length": str(len(content))},
+            content=content,
+            request=httpx.Request("GET", "https://huggingface.co/user/repo/resolve/main/file.txt"),
+        )
+
+    download_kwargs = {"cache_dir": tmp_path / "cache"}
+    if use_local_dir:
+        download_kwargs["local_dir"] = tmp_path / "local"
+
+    with (
+        patch("huggingface_hub.file_download._httpx_follow_hub_redirects_with_backoff", side_effect=_mock_head),
+        patch("huggingface_hub.file_download.http_stream_backoff", side_effect=_mock_get) as mock_get,
+        patch("huggingface_hub.constants.HF_HUB_DISABLE_XET", xet_mode == "disabled"),
+        patch("huggingface_hub.utils._runtime.is_package_available", return_value=xet_mode != "not_installed"),
+    ):
+        if xet_mode == "enabled":
+            with pytest.raises(LocalEntryNotFoundError) as exc:
+                hf_hub_download("user/repo", "file.txt", **download_kwargs)
+            assert isinstance(exc.value.__cause__, FileMetadataError)
+            mock_get.assert_not_called()
+            return
+        path = hf_hub_download("user/repo", "file.txt", **download_kwargs)
+
+    assert Path(path).read_bytes() == content
+
+
 class TestDiskUsageWarning:
     @pytest.fixture(scope="class", autouse=True)
     def setup(self, request):
@@ -1049,6 +1095,12 @@ def test_get_pointer_path_but_invalid_relative_filename(self) -> None:
 
 
 class TestHttpGet:
+    def test_http_get_validates_content_length_when_expected_size_is_missing(self):
+        with pytest.raises(OSError, match="file should be of size 100 but has size 50"):
+            self._http_get_with_mocked_responses(
+                [self._mock_response(headers={"Content-Length": "100"}, iter_bytes=iter([b"A" * 50]))]
+            )
+
     def test_http_get_with_ssl_and_timeout_error(self, caplog):
         def _iter_content_1() -> Iterable[bytes]:
             yield b"0" * 10