Xet Upload with byte array (#3035)

* Using upload_bytes from hf_xet

* Factoring out the file key

* Using the new datastructures

* Adding bytes based test

* Adding tests for byte array based upload and testing xet backwards compatibility

* Test cleanup

* Lint

* Quality check

* Updating hf-xet version in setup.py

* Update src/huggingface_hub/_commit_api.py

Co-authored-by: Lucain <[email protected]>

* Update src/huggingface_hub/_commit_api.py

Co-authored-by: Lucain <[email protected]>

* Update src/huggingface_hub/hf_api.py

Co-authored-by: Lucain <[email protected]>

* Update tests/test_xet_upload.py

Co-authored-by: Lucain <[email protected]>

* Update src/huggingface_hub/hf_api.py

Co-authored-by: Lucain <[email protected]>

* PR comments

* Fixing compilation

* Quality

* Fixing spelling

---------

Co-authored-by: Lucain <[email protected]>

Author: bpronan

setup.py

@@ -14,7 +14,7 @@ def get_version() -> str:
 install_requires = [
     "filelock",
     "fsspec>=2023.5.0",
-    "hf-xet>=1.0.2,<2.0.0; platform_machine=='x86_64' or platform_machine=='amd64' or platform_machine=='arm64' or platform_machine=='aarch64'",
+    "hf-xet>=1.1.0,<2.0.0; platform_machine=='x86_64' or platform_machine=='amd64' or platform_machine=='arm64' or platform_machine=='aarch64'",
     "packaging>=20.9",
     "pyyaml>=5.1",
     "requests",
@@ -56,7 +56,7 @@ def get_version() -> str:
     "keras<3.0",
 ]
 
-extras["hf_xet"] = ["hf_xet>=1.0.2,<2.0.0"]
+extras["hf_xet"] = ["hf_xet>=1.1.0,<2.0.0"]
 
 extras["testing"] = (
     extras["cli"]

src/huggingface_hub/_commit_api.py

@@ -530,7 +530,7 @@ def _upload_xet_files(
     if len(additions) == 0:
         return
     # at this point, we know that hf_xet is installed
-    from hf_xet import upload_files
+    from hf_xet import upload_bytes, upload_files
 
     try:
         xet_connection_info = fetch_xet_connection_info_from_repo_info(
@@ -571,8 +571,10 @@ def token_refresher() -> Tuple[str, int]:
     num_chunks_num_digits = int(math.log10(num_chunks)) + 1
     for i, chunk in enumerate(chunk_iterable(additions, chunk_size=UPLOAD_BATCH_MAX_NUM_FILES)):
         _chunk = [op for op in chunk]
-        paths = [str(op.path_or_fileobj) for op in _chunk]
-        expected_size = sum([os.path.getsize(path) for path in paths])
+
+        bytes_ops = [op for op in _chunk if isinstance(op.path_or_fileobj, bytes)]
+        paths_ops = [op for op in _chunk if isinstance(op.path_or_fileobj, (str, Path))]
+        expected_size = sum(op.upload_info.size for op in bytes_ops + paths_ops)
 
         if num_chunks > 1:
             description = f"Uploading Batch [{str(i + 1).zfill(num_chunks_num_digits)}/{num_chunks}]..."
@@ -592,7 +594,24 @@ def token_refresher() -> Tuple[str, int]:
             def update_progress(increment: int):
                 progress.update(increment)
 
-            upload_files(paths, xet_endpoint, access_token_info, token_refresher, update_progress, repo_type)
+            if len(paths_ops) > 0:
+                upload_files(
+                    [str(op.path_or_fileobj) for op in paths_ops],
+                    xet_endpoint,
+                    access_token_info,
+                    token_refresher,
+                    update_progress,
+                    repo_type,
+                )
+            if len(bytes_ops) > 0:
+                upload_bytes(
+                    [op.path_or_fileobj for op in bytes_ops],
+                    xet_endpoint,
+                    access_token_info,
+                    token_refresher,
+                    update_progress,
+                    repo_type,
+                )
     return
 
 

src/huggingface_hub/file_download.py

@@ -582,7 +582,7 @@ def xet_get(
 
     """
     try:
-        from hf_xet import PyPointerFile, download_files  # type: ignore[no-redef]
+        from hf_xet import PyXetDownloadInfo, download_files  # type: ignore[no-redef]
     except ImportError:
         raise ValueError(
             "To use optimized download using Xet storage, you need to install the hf_xet package. "
@@ -597,8 +597,10 @@ def token_refresher() -> Tuple[str, int]:
             raise ValueError("Failed to refresh token using xet metadata.")
         return connection_info.access_token, connection_info.expiration_unix_epoch
 
-    pointer_files = [
-        PyPointerFile(path=str(incomplete_path.absolute()), hash=xet_file_data.file_hash, filesize=expected_size)
+    xet_download_info = [
+        PyXetDownloadInfo(
+            destination_path=str(incomplete_path.absolute()), hash=xet_file_data.file_hash, file_size=expected_size
+        )
     ]
 
     if not displayed_filename:
@@ -623,7 +625,7 @@ def progress_updater(progress_bytes: float):
             progress.update(progress_bytes)
 
         download_files(
-            pointer_files,
+            xet_download_info,
             endpoint=connection_info.endpoint,
             token_info=(connection_info.access_token, connection_info.expiration_unix_epoch),
             token_refresher=token_refresher,

src/huggingface_hub/hf_api.py

@@ -4475,18 +4475,17 @@ def preupload_lfs_files(
             expand="xetEnabled",
             token=token,
         ).xet_enabled
-        has_binary_data = any(
-            isinstance(addition.path_or_fileobj, (bytes, io.BufferedIOBase))
-            for addition in new_lfs_additions_to_upload
+        has_buffered_io_data = any(
+            isinstance(addition.path_or_fileobj, io.BufferedIOBase) for addition in new_lfs_additions_to_upload
         )
-        if xet_enabled and not has_binary_data and is_xet_available():
+        if xet_enabled and not has_buffered_io_data and is_xet_available():
             logger.info("Uploading files using Xet Storage..")
             _upload_xet_files(**upload_kwargs, create_pr=create_pr)  # type: ignore [arg-type]
         else:
             if xet_enabled and is_xet_available():
-                if has_binary_data:
+                if has_buffered_io_data:
                     logger.warning(
-                        "Uploading files as bytes or binary IO objects is not supported by Xet Storage. "
+                        "Uploading files as a binary IO buffer is not supported by Xet Storage. "
                         "Falling back to HTTP upload."
                     )
             _upload_lfs_files(**upload_kwargs, num_threads=num_threads)  # type: ignore [arg-type]

tests/test_xet_download.py

@@ -1,6 +1,7 @@
 import os
 from contextlib import contextmanager
 from pathlib import Path
+from typing import Tuple
 from unittest.mock import DEFAULT, Mock, patch
 
 from huggingface_hub import snapshot_download
@@ -293,3 +294,42 @@ def test_snapshot_download_cache_reuse(self, tmp_path):
 
             # Verify xet_get was not called (files were cached)
             mock_xet_get.assert_not_called()
+
+    def test_download_backward_compatibility(self, tmp_path):
+        """Test that xet download works with the old pointer file protocol.
+
+        Until the next major version of hf-xet is released, we need to support the old
+        pointer file based download to support old huggingface_hub versions.
+        """
+
+        file_path = os.path.join(tmp_path, DUMMY_XET_FILE)
+
+        file_metadata = get_hf_file_metadata(
+            hf_hub_url(
+                repo_id=DUMMY_XET_MODEL_ID,
+                filename=DUMMY_XET_FILE,
+            )
+        )
+
+        xet_file_data = file_metadata.xet_file_data
+
+        # Mock the response to not include xet metadata
+        from hf_xet import PyPointerFile, download_files
+
+        connection_info = refresh_xet_connection_info(file_data=xet_file_data, headers={})
+
+        def token_refresher() -> Tuple[str, int]:
+            connection_info = refresh_xet_connection_info(file_data=xet_file_data, headers={})
+            return connection_info.access_token, connection_info.expiration_unix_epoch
+
+        pointer_files = [PyPointerFile(path=file_path, hash=xet_file_data.file_hash, filesize=file_metadata.size)]
+
+        download_files(
+            pointer_files,
+            endpoint=connection_info.endpoint,
+            token_info=(connection_info.access_token, connection_info.expiration_unix_epoch),
+            token_refresher=token_refresher,
+            progress_updater=None,
+        )
+
+        assert os.path.exists(file_path)

tests/test_xet_upload.py

@@ -136,6 +136,21 @@ def test_upload_file_with_bytesio(self, api, tmp_path, repo_url):
             downloaded_content = f.read()
             assert downloaded_content == self.bin_content
 
+    def test_upload_file_with_byte_array(self, api, tmp_path, repo_url):
+        repo_id = repo_url.repo_id
+        content = self.bin_content
+        with assert_upload_mode("xet"):
+            api.upload_file(
+                path_or_fileobj=content,
+                path_in_repo="bytearray_file.bin",
+                repo_id=repo_id,
+            )
+        # Download and verify content
+        downloaded_file = hf_hub_download(repo_id=repo_id, filename="bytearray_file.bin", cache_dir=tmp_path)
+        with open(downloaded_file, "rb") as f:
+            downloaded_content = f.read()
+            assert downloaded_content == self.bin_content
+
     def test_fallback_to_lfs_when_xet_not_available(self, api, repo_url):
         repo_id = repo_url.repo_id
         with patch("huggingface_hub.hf_api.is_xet_available", return_value=False):
@@ -284,7 +299,7 @@ def test_hf_xet_with_token_refresher(self, api, tmp_path, repo_url):
 
         This test ensures that the downloaded file is the same as the uploaded file.
         """
-        from hf_xet import PyPointerFile, download_files
+        from hf_xet import PyXetDownloadInfo, download_files
 
         filename_in_repo = "binary_file.bin"
         repo_id = repo_url.repo_id
@@ -327,13 +342,15 @@ def token_refresher() -> Tuple[str, int]:
         mock_token_refresher = MagicMock(side_effect=token_refresher)
 
         incomplete_path = Path(tmp_path) / "file.bin.incomplete"
-        py_file = [
-            PyPointerFile(path=str(incomplete_path.absolute()), hash=xet_filedata.file_hash, filesize=expected_size)
+        file_info = [
+            PyXetDownloadInfo(
+                destination_path=str(incomplete_path.absolute()), hash=xet_filedata.file_hash, file_size=expected_size
+            )
         ]
 
         # Call the download_files function with the token refresher, set expiration to 0 forcing a refresh
         download_files(
-            py_file,
+            file_info,
             endpoint=xet_connection_info.endpoint,
             token_info=(xet_connection_info.access_token, 0),
             token_refresher=mock_token_refresher,