upload large folder operations uses batches of files for preupload-lfs jobs for xet-enabled repositories (#3228)

* modify funcs to accept lists

* batch uploads post testing

* make style

* is_xet_enabled

* change batch size constructor param to have default size of 1

* test batch size set to greater than one

* style + quality

* apply comments from Lucain

* get n files or at least totalling 64MB

* Revert "get n files or at least totalling 64MB"

This reverts commit 04d31d9.

---------

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

Author: assafvayner

src/huggingface_hub/_upload_large_folder.py

@@ -33,6 +33,7 @@
 from .constants import DEFAULT_REVISION, REPO_TYPES
 from .utils import DEFAULT_IGNORE_PATTERNS, filter_repo_objects, tqdm
 from .utils._cache_manager import _format_size
+from .utils._runtime import is_xet_available
 from .utils.sha import sha_fileobj
 
 
@@ -45,6 +46,9 @@
 MAX_NB_FILES_FETCH_UPLOAD_MODE = 100
 COMMIT_SIZE_SCALE: List[int] = [20, 50, 75, 100, 125, 200, 250, 400, 600, 1000]
 
+UPLOAD_BATCH_SIZE_XET = 256  # Max 256 files per upload batch for XET-enabled repos
+UPLOAD_BATCH_SIZE_LFS = 1  # Otherwise, batches of 1 for regular LFS upload
+
 
 def upload_large_folder_internal(
     api: "HfApi",
@@ -93,6 +97,17 @@ def upload_large_folder_internal(
     repo_url = api.create_repo(repo_id=repo_id, repo_type=repo_type, private=private, exist_ok=True)
     logger.info(f"Repo created: {repo_url}")
     repo_id = repo_url.repo_id
+    # 2.1 Check if xet is enabled to set batch file upload size
+    is_xet_enabled = (
+        is_xet_available()
+        and api.repo_info(
+            repo_id=repo_id,
+            repo_type=repo_type,
+            revision=revision,
+            expand="xetEnabled",
+        ).xet_enabled
+    )
+    upload_batch_size = UPLOAD_BATCH_SIZE_XET if is_xet_enabled else UPLOAD_BATCH_SIZE_LFS
 
     # 3. List files to upload
     filtered_paths_list = filter_repo_objects(
@@ -110,7 +125,7 @@ def upload_large_folder_internal(
     ]
 
     # 4. Start workers
-    status = LargeUploadStatus(items)
+    status = LargeUploadStatus(items, upload_batch_size)
     threads = [
         threading.Thread(
             target=_worker_job,
@@ -168,7 +183,7 @@ class WorkerJob(enum.Enum):
 class LargeUploadStatus:
     """Contains information, queues and tasks for a large upload process."""
 
-    def __init__(self, items: List[JOB_ITEM_T]):
+    def __init__(self, items: List[JOB_ITEM_T], upload_batch_size: int = 1):
         self.items = items
         self.queue_sha256: "queue.Queue[JOB_ITEM_T]" = queue.Queue()
         self.queue_get_upload_mode: "queue.Queue[JOB_ITEM_T]" = queue.Queue()
@@ -179,6 +194,7 @@ def __init__(self, items: List[JOB_ITEM_T]):
         self.nb_workers_sha256: int = 0
         self.nb_workers_get_upload_mode: int = 0
         self.nb_workers_preupload_lfs: int = 0
+        self.upload_batch_size: int = upload_batch_size
         self.nb_workers_commit: int = 0
         self.nb_workers_waiting: int = 0
         self.last_commit_attempt: Optional[float] = None
@@ -353,16 +369,17 @@ def _worker_job(
                 status.nb_workers_get_upload_mode -= 1
 
         elif job == WorkerJob.PREUPLOAD_LFS:
-            item = items[0]  # single item
             try:
-                _preupload_lfs(item, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision)
-                status.queue_commit.put(item)
+                _preupload_lfs(items, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision)
+                for item in items:
+                    status.queue_commit.put(item)
             except KeyboardInterrupt:
                 raise
             except Exception as e:
                 logger.error(f"Failed to preupload LFS: {e}")
                 traceback.format_exc()
-                status.queue_preupload_lfs.put(item)
+                for item in items:
+                    status.queue_preupload_lfs.put(item)
 
             with status.lock:
                 status.nb_workers_preupload_lfs -= 1
@@ -417,11 +434,11 @@ def _determine_next_job(status: LargeUploadStatus) -> Optional[Tuple[WorkerJob,
             logger.debug(f"Job: get upload mode (>{MAX_NB_FILES_FETCH_UPLOAD_MODE} files ready)")
             return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, MAX_NB_FILES_FETCH_UPLOAD_MODE))
 
-        # 4. Preupload LFS file if at least 1 file and no worker is preuploading LFS
-        elif status.queue_preupload_lfs.qsize() > 0 and status.nb_workers_preupload_lfs == 0:
+        # 4. Preupload LFS file if at least `status.upload_batch_size` files and no worker is preuploading LFS
+        elif status.queue_preupload_lfs.qsize() >= status.upload_batch_size and status.nb_workers_preupload_lfs == 0:
             status.nb_workers_preupload_lfs += 1
             logger.debug("Job: preupload LFS (no other worker preuploading LFS)")
-            return (WorkerJob.PREUPLOAD_LFS, _get_one(status.queue_preupload_lfs))
+            return (WorkerJob.PREUPLOAD_LFS, _get_n(status.queue_preupload_lfs, status.upload_batch_size))
 
         # 5. Compute sha256 if at least 1 file and no worker is computing sha256
         elif status.queue_sha256.qsize() > 0 and status.nb_workers_sha256 == 0:
@@ -435,14 +452,14 @@ def _determine_next_job(status: LargeUploadStatus) -> Optional[Tuple[WorkerJob,
             logger.debug("Job: get upload mode (no other worker getting upload mode)")
             return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, MAX_NB_FILES_FETCH_UPLOAD_MODE))
 
-        # 7. Preupload LFS file if at least 1 file
+        # 7. Preupload LFS file if at least `status.upload_batch_size` files
         #    Skip if hf_transfer is enabled and there is already a worker preuploading LFS
-        elif status.queue_preupload_lfs.qsize() > 0 and (
+        elif status.queue_preupload_lfs.qsize() >= status.upload_batch_size and (
             status.nb_workers_preupload_lfs == 0 or not constants.HF_HUB_ENABLE_HF_TRANSFER
         ):
             status.nb_workers_preupload_lfs += 1
             logger.debug("Job: preupload LFS")
-            return (WorkerJob.PREUPLOAD_LFS, _get_one(status.queue_preupload_lfs))
+            return (WorkerJob.PREUPLOAD_LFS, _get_n(status.queue_preupload_lfs, status.upload_batch_size))
 
         # 8. Compute sha256 if at least 1 file
         elif status.queue_sha256.qsize() > 0:
@@ -456,7 +473,13 @@ def _determine_next_job(status: LargeUploadStatus) -> Optional[Tuple[WorkerJob,
             logger.debug("Job: get upload mode")
             return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, MAX_NB_FILES_FETCH_UPLOAD_MODE))
 
-        # 10. Commit if at least 1 file and 1 min since last commit attempt
+        # 10. Preupload LFS file if at least 1 file
+        elif status.queue_preupload_lfs.qsize() > 0:
+            status.nb_workers_preupload_lfs += 1
+            logger.debug("Job: preupload LFS")
+            return (WorkerJob.PREUPLOAD_LFS, _get_n(status.queue_preupload_lfs, status.upload_batch_size))
+
+        # 11. Commit if at least 1 file and 1 min since last commit attempt
         elif (
             status.nb_workers_commit == 0
             and status.queue_commit.qsize() > 0
@@ -467,7 +490,7 @@ def _determine_next_job(status: LargeUploadStatus) -> Optional[Tuple[WorkerJob,
             logger.debug("Job: commit (1 min since last commit attempt)")
             return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))
 
-        # 11. Commit if at least 1 file all other queues are empty and all workers are waiting
+        # 12. Commit if at least 1 file all other queues are empty and all workers are waiting
         #     e.g. when it's the last commit
         elif (
             status.nb_workers_commit == 0
@@ -483,12 +506,12 @@ def _determine_next_job(status: LargeUploadStatus) -> Optional[Tuple[WorkerJob,
             logger.debug("Job: commit")
             return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))
 
-        # 12. If all queues are empty, exit
+        # 13. If all queues are empty, exit
         elif all(metadata.is_committed or metadata.should_ignore for _, metadata in status.items):
             logger.info("All files have been processed! Exiting worker.")
             return None
 
-        # 13. If no task is available, wait
+        # 14. If no task is available, wait
         else:
             status.nb_workers_waiting += 1
             logger.debug(f"No task available, waiting... ({WAITING_TIME_IF_NO_TASKS}s)")
@@ -531,19 +554,19 @@ def _get_upload_mode(items: List[JOB_ITEM_T], api: "HfApi", repo_id: str, repo_t
         metadata.save(paths)
 
 
-def _preupload_lfs(item: JOB_ITEM_T, api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None:
-    """Preupload LFS file and update metadata."""
-    paths, metadata = item
-    addition = _build_hacky_operation(item)
+def _preupload_lfs(items: List[JOB_ITEM_T], api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None:
+    """Preupload LFS files and update metadata."""
+    additions = [_build_hacky_operation(item) for item in items]
     api.preupload_lfs_files(
         repo_id=repo_id,
         repo_type=repo_type,
         revision=revision,
-        additions=[addition],
+        additions=additions,
     )
 
-    metadata.is_uploaded = True
-    metadata.save(paths)
+    for paths, metadata in items:
+        metadata.is_uploaded = True
+        metadata.save(paths)
 
 
 def _commit(items: List[JOB_ITEM_T], api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None:

tests/test_xet_upload.py

@@ -279,6 +279,39 @@ def test_upload_large_folder(self, api, tmp_path, repo_url: RepoUrl) -> None:
             regular_file = local_dir / f"subfolder_{i}/file_regular_{i}_{j}.txt"
             assert regular_file.read_bytes() == f"content_regular_{i}_{j}".encode()
 
+    def test_upload_large_folder_batch_size_greater_than_one(self, api, tmp_path, repo_url: RepoUrl) -> None:
+        from hf_xet import upload_files as real_upload_files
+
+        N_FILES = 500
+        repo_id = repo_url.repo_id
+
+        folder = Path(tmp_path) / "large_folder"
+        folder.mkdir()
+        for i in range(N_FILES):
+            (folder / f"file_xet_{i}.bin").write_bytes(f"content_lfs_{i}".encode())
+
+        # capture the number of files passed in per call to hf_xet.upload_files
+        # to ensure that the batch size is respected.
+        num_files_per_call = []
+
+        def spy_upload_files(*args, **kwargs):
+            num_files = len(args[0])
+            num_files_per_call.append(num_files)
+            return real_upload_files(*args, **kwargs)
+
+        with assert_upload_mode("xet"):
+            with patch("hf_xet.upload_files", side_effect=spy_upload_files):
+                api.upload_large_folder(repo_id=repo_id, repo_type="model", folder_path=folder, num_workers=4)
+
+        # the batch size is set to 256 however due to speed of hashing and get_upload_mode calls it's not always guaranteed
+        # that the files will be uploaded in batches of 256. They may be uploaded in smaller batches if no other jobs
+        # are available to run; even as small as 1 file per call.
+        #
+        # However, it would be unlikely that all files are uploaded in batches of 1 if batching was correctly implemented.
+        # So we assert that not all files were uploaded in batches of 1, although it is possible even with batching.
+
+        assert any(n > 1 for n in num_files_per_call)
+
 
 @requires("hf_xet")
 class TestXetE2E(TestXetUpload):