codekingpro/portable-devtools
114k
1# Copyright 2024-present, the HuggingFace Inc. team.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14import enum15import logging16import os17import queue18import shutil19import sys20import threading21import time22import traceback23from datetime import datetime24from pathlib import Path25from threading import Lock26from typing import TYPE_CHECKING, Any27from urllib.parse import quote28 29from ._commit_api import CommitOperationAdd, UploadInfo, _fetch_upload_modes30from ._local_folder import LocalUploadFileMetadata, LocalUploadFilePaths, get_local_upload_paths, read_upload_metadata31from .constants import DEFAULT_REVISION, REPO_TYPES32from .utils import DEFAULT_IGNORE_PATTERNS, _format_size, filter_repo_objects, tqdm33from .utils._runtime import is_xet_available34from .utils.sha import sha_fileobj35 36 37if TYPE_CHECKING:38 from .hf_api import HfApi39 40logger = logging.getLogger(__name__)41 42WAITING_TIME_IF_NO_TASKS = 10 # seconds43MAX_NB_FILES_FETCH_UPLOAD_MODE = 10044COMMIT_SIZE_SCALE: list[int] = [20, 50, 75, 100, 125, 200, 250, 400, 600, 1000]45 46UPLOAD_BATCH_SIZE_XET = 256 # Max 256 files per upload batch for XET-enabled repos47UPLOAD_BATCH_SIZE_LFS = 1 # Otherwise, batches of 1 for regular LFS upload48 49# Repository limits (from https://huggingface.co/docs/hub/repositories-recommendations)50MAX_FILES_PER_REPO = 100_000 # Recommended maximum number of files per repository51MAX_FILES_PER_FOLDER = 10_000 # Recommended maximum number of files per folder52MAX_FILE_SIZE_GB = 200 # Recommended maximum for individual file size (split larger files)53RECOMMENDED_FILE_SIZE_GB = 20 # Recommended maximum for individual file size54 55 56def _validate_upload_limits(paths_list: list[LocalUploadFilePaths]) -> None:57 """58 Validate upload against repository limits and warn about potential issues.59 60 Args:61 paths_list: List of file paths to be uploaded62 63 Warns about:64 - Too many files in the repository (>100k)65 - Too many entries (files or subdirectories) in a single folder (>10k)66 - Files exceeding size limits (>20GB recommended, >200GB maximum)67 """68 logger.info("Running validation checks on files to upload...")69 70 # Check 1: Total file count71 if len(paths_list) > MAX_FILES_PER_REPO:72 logger.warning(73 f"You are about to upload {len(paths_list):,} files. "74 f"This exceeds the recommended limit of {MAX_FILES_PER_REPO:,} files per repository.\n"75 f"Consider:\n"76 f" - Splitting your data into multiple repositories\n"77 f" - Using fewer, larger files (e.g., parquet files)\n"78 f" - See: https://huggingface.co/docs/hub/repositories-recommendations"79 )80 81 # Check 2: Files and subdirectories per folder82 # Track immediate children (files and subdirs) for each folder83 from collections import defaultdict84 85 entries_per_folder: dict[str, Any] = defaultdict(lambda: {"files": 0, "subdirs": set()})86 87 for paths in paths_list:88 path = Path(paths.path_in_repo)89 parts = path.parts90 91 # Count this file in its immediate parent directory92 parent = str(path.parent) if str(path.parent) != "." else "."93 entries_per_folder[parent]["files"] += 194 95 # Track immediate subdirectories for each parent folder96 # Walk through the path components to track parent-child relationships97 for i, child in enumerate(parts[:-1]):98 parent = "." if i == 0 else "/".join(parts[:i])99 entries_per_folder[parent]["subdirs"].add(child)100 101 # Check limits for each folder102 for folder, data in entries_per_folder.items():103 file_count = data["files"]104 subdir_count = len(data["subdirs"])105 total_entries = file_count + subdir_count106 107 if total_entries > MAX_FILES_PER_FOLDER:108 folder_display = "root" if folder == "." else folder109 logger.warning(110 f"Folder '{folder_display}' contains {total_entries:,} entries "111 f"({file_count:,} files and {subdir_count:,} subdirectories). "112 f"This exceeds the recommended {MAX_FILES_PER_FOLDER:,} entries per folder.\n"113 "Consider reorganising into sub-folders."114 )115 116 # Check 3: File sizes117 large_files = []118 very_large_files = []119 120 for paths in paths_list:121 size = paths.file_path.stat().st_size122 size_gb = size / 1_000_000_000 # Use decimal GB as per Hub limits123 124 if size_gb > MAX_FILE_SIZE_GB:125 very_large_files.append((paths.path_in_repo, size_gb))126 elif size_gb > RECOMMENDED_FILE_SIZE_GB:127 large_files.append((paths.path_in_repo, size_gb))128 129 # Warn about very large files (>200GB)130 if very_large_files:131 files_str = "\n - ".join(f"{path}: {size:.1f}GB" for path, size in very_large_files[:5])132 more_str = f"\n ... and {len(very_large_files) - 5} more files" if len(very_large_files) > 5 else ""133 logger.warning(134 f"Found {len(very_large_files)} files exceeding the {MAX_FILE_SIZE_GB}GB recommended maximum:\n"135 f" - {files_str}{more_str}\n"136 f"Consider splitting these files into smaller chunks."137 )138 139 # Warn about large files (>20GB)140 if large_files:141 files_str = "\n - ".join(f"{path}: {size:.1f}GB" for path, size in large_files[:5])142 more_str = f"\n ... and {len(large_files) - 5} more files" if len(large_files) > 5 else ""143 logger.warning(144 f"Found {len(large_files)} files larger than {RECOMMENDED_FILE_SIZE_GB}GB (recommended limit):\n"145 f" - {files_str}{more_str}\n"146 f"Large files may slow down loading and processing."147 )148 149 logger.info("Validation checks complete.")150 151 152def upload_large_folder_internal(153 api: "HfApi",154 repo_id: str,155 folder_path: str | Path,156 *,157 repo_type: str, # Repo type is required!158 revision: str | None = None,159 private: bool | None = None,160 allow_patterns: list[str] | str | None = None,161 ignore_patterns: list[str] | str | None = None,162 num_workers: int | None = None,163 print_report: bool = True,164 print_report_every: int = 60,165):166 """Upload a large folder to the Hub in the most resilient way possible.167 168 See [`HfApi.upload_large_folder`] for the full documentation.169 """170 # 1. Check args and setup171 if repo_type is None:172 raise ValueError(173 "For large uploads, `repo_type` is explicitly required. Please set it to `model`, `dataset` or `space`."174 " If you are using the CLI, pass it as `--repo-type=model`."175 )176 if repo_type not in REPO_TYPES:177 raise ValueError(f"Invalid repo type, must be one of {REPO_TYPES}")178 if revision is None:179 revision = DEFAULT_REVISION180 181 folder_path = Path(folder_path).expanduser().resolve()182 if not folder_path.is_dir():183 raise ValueError(f"Provided path: '{folder_path}' is not a directory")184 185 if ignore_patterns is None:186 ignore_patterns = []187 elif isinstance(ignore_patterns, str):188 ignore_patterns = [ignore_patterns]189 ignore_patterns += DEFAULT_IGNORE_PATTERNS190 191 if num_workers is None:192 nb_cores = os.cpu_count() or 1193 num_workers = max(nb_cores // 2, 1) # Use at most half of cpu cores194 195 # 2. Create repo if missing196 repo_url = api.create_repo(repo_id=repo_id, repo_type=repo_type, private=private, exist_ok=True)197 logger.info(f"Repo created: {repo_url}")198 repo_id = repo_url.repo_id199 200 # Warn on too many commits201 try:202 commits = api.list_repo_commits(repo_id=repo_id, repo_type=repo_type, revision=revision)203 commit_count = len(commits)204 if commit_count > 500:205 logger.warning(206 f"\n{'=' * 80}\n"207 f"WARNING: This repository has {commit_count} commits.\n"208 f"Repositories with a large number of commits can experience performance issues.\n"209 f"\n"210 f"Consider squashing your commit history using `super_squash_history()`.\n"211 "To do so, you need to stop this process, run the snippet below and restart the upload command."212 f" from huggingface_hub import super_squash_history\n"213 f" super_squash_history(repo_id='{repo_id}', repo_type='{repo_type}')\n"214 f"\n"215 f"Note: This is a non-revertible operation. See the documentation for more details:\n"216 f"https://huggingface.co/docs/huggingface_hub/main/en/package_reference/hf_api#huggingface_hub.HfApi.super_squash_history\n"217 f"{'=' * 80}\n"218 )219 except Exception as e:220 # Don't fail the upload if we can't check commit count221 logger.debug(f"Could not check commit count: {e}")222 223 # 2.1 Check if xet is enabled to set batch file upload size224 upload_batch_size = UPLOAD_BATCH_SIZE_XET if is_xet_available() else UPLOAD_BATCH_SIZE_LFS225 226 # 3. List files to upload227 filtered_paths_list = filter_repo_objects(228 (path.relative_to(folder_path).as_posix() for path in folder_path.glob("**/*") if path.is_file()),229 allow_patterns=allow_patterns,230 ignore_patterns=ignore_patterns,231 )232 paths_list = [get_local_upload_paths(folder_path, relpath) for relpath in filtered_paths_list]233 logger.info(f"Found {len(paths_list)} candidate files to upload")234 235 # Validate upload against repository limits236 _validate_upload_limits(paths_list)237 238 logger.info("Starting upload...")239 240 # Read metadata for each file241 items = [242 (paths, read_upload_metadata(folder_path, paths.path_in_repo))243 for paths in tqdm(paths_list, desc="Recovering from metadata files")244 ]245 246 # 4. Start workers247 status = LargeUploadStatus(items, upload_batch_size)248 threads = [249 threading.Thread(250 target=_worker_job,251 kwargs={252 "status": status,253 "api": api,254 "repo_id": repo_id,255 "repo_type": repo_type,256 "revision": revision,257 },258 )259 for _ in range(num_workers)260 ]261 262 for thread in threads:263 thread.start()264 265 # 5. Print regular reports266 if print_report:267 print("\n\n" + status.current_report())268 last_report_ts = time.time()269 while True:270 time.sleep(1)271 if time.time() - last_report_ts >= print_report_every:272 if print_report:273 _print_overwrite(status.current_report())274 last_report_ts = time.time()275 if status.is_done():276 logger.info("Is done: exiting main loop")277 break278 279 for thread in threads:280 thread.join()281 282 logger.info(status.current_report())283 logger.info("Upload is complete!")284 285 286####################287# Logic to manage workers and synchronize tasks288####################289 290 291class WorkerJob(enum.Enum):292 SHA256 = enum.auto()293 GET_UPLOAD_MODE = enum.auto()294 PREUPLOAD_LFS = enum.auto()295 COMMIT = enum.auto()296 WAIT = enum.auto() # if no tasks are available but we don't want to exit297 298 299JOB_ITEM_T = tuple[LocalUploadFilePaths, LocalUploadFileMetadata]300 301 302class LargeUploadStatus:303 """Contains information, queues and tasks for a large upload process."""304 305 def __init__(self, items: list[JOB_ITEM_T], upload_batch_size: int = 1):306 self.items = items307 self.queue_sha256: "queue.Queue[JOB_ITEM_T]" = queue.Queue()308 self.queue_get_upload_mode: "queue.Queue[JOB_ITEM_T]" = queue.Queue()309 self.queue_preupload_lfs: "queue.Queue[JOB_ITEM_T]" = queue.Queue()310 self.queue_commit: "queue.Queue[JOB_ITEM_T]" = queue.Queue()311 self.lock = Lock()312 313 self.nb_workers_sha256: int = 0314 self.nb_workers_get_upload_mode: int = 0315 self.nb_workers_preupload_lfs: int = 0316 self.upload_batch_size: int = upload_batch_size317 self.nb_workers_commit: int = 0318 self.nb_workers_waiting: int = 0319 self.last_commit_attempt: float | None = None320 321 self._started_at = datetime.now()322 self._chunk_idx: int = 1323 self._chunk_lock: Lock = Lock()324 325 # Setup queues326 for item in self.items:327 paths, metadata = item328 if metadata.sha256 is None:329 self.queue_sha256.put(item)330 elif metadata.upload_mode is None:331 self.queue_get_upload_mode.put(item)332 elif metadata.upload_mode == "lfs" and not metadata.is_uploaded:333 self.queue_preupload_lfs.put(item)334 elif not metadata.is_committed:335 self.queue_commit.put(item)336 else:337 logger.debug(f"Skipping file {paths.path_in_repo} (already uploaded and committed)")338 339 def target_chunk(self) -> int:340 with self._chunk_lock:341 return COMMIT_SIZE_SCALE[self._chunk_idx]342 343 def update_chunk(self, success: bool, nb_items: int, duration: float) -> None:344 with self._chunk_lock:345 if not success:346 logger.warning(f"Failed to commit {nb_items} files at once. Will retry with less files in next batch.")347 self._chunk_idx -= 1348 elif nb_items >= COMMIT_SIZE_SCALE[self._chunk_idx] and duration < 40:349 logger.info(f"Successfully committed {nb_items} at once. Increasing the limit for next batch.")350 self._chunk_idx += 1351 352 self._chunk_idx = max(0, min(self._chunk_idx, len(COMMIT_SIZE_SCALE) - 1))353 354 def current_report(self) -> str:355 """Generate a report of the current status of the large upload."""356 nb_hashed = 0357 size_hashed = 0358 nb_preuploaded = 0359 nb_lfs = 0360 nb_lfs_unsure = 0361 size_preuploaded = 0362 nb_committed = 0363 size_committed = 0364 total_size = 0365 ignored_files = 0366 total_files = 0367 368 with self.lock:369 for _, metadata in self.items:370 if metadata.should_ignore:371 ignored_files += 1372 continue373 total_size += metadata.size374 total_files += 1375 if metadata.sha256 is not None:376 nb_hashed += 1377 size_hashed += metadata.size378 if metadata.upload_mode == "lfs":379 nb_lfs += 1380 if metadata.upload_mode is None:381 nb_lfs_unsure += 1382 if metadata.is_uploaded:383 nb_preuploaded += 1384 size_preuploaded += metadata.size385 if metadata.is_committed:386 nb_committed += 1387 size_committed += metadata.size388 total_size_str = _format_size(total_size)389 390 now = datetime.now()391 now_str = now.strftime("%Y-%m-%d %H:%M:%S")392 elapsed = now - self._started_at393 elapsed_str = str(elapsed).split(".")[0] # remove milliseconds394 395 message = "\n" + "-" * 10396 message += f" {now_str} ({elapsed_str}) "397 message += "-" * 10 + "\n"398 399 message += "Files: "400 message += f"hashed {nb_hashed}/{total_files} ({_format_size(size_hashed)}/{total_size_str}) | "401 message += f"pre-uploaded: {nb_preuploaded}/{nb_lfs} ({_format_size(size_preuploaded)}/{total_size_str})"402 if nb_lfs_unsure > 0:403 message += f" (+{nb_lfs_unsure} unsure)"404 message += f" | committed: {nb_committed}/{total_files} ({_format_size(size_committed)}/{total_size_str})"405 message += f" | ignored: {ignored_files}\n"406 407 message += "Workers: "408 message += f"hashing: {self.nb_workers_sha256} | "409 message += f"get upload mode: {self.nb_workers_get_upload_mode} | "410 message += f"pre-uploading: {self.nb_workers_preupload_lfs} | "411 message += f"committing: {self.nb_workers_commit} | "412 message += f"waiting: {self.nb_workers_waiting}\n"413 message += "-" * 51414 415 return message416 417 def is_done(self) -> bool:418 with self.lock:419 return all(metadata.is_committed or metadata.should_ignore for _, metadata in self.items)420 421 422def _worker_job(423 status: LargeUploadStatus,424 api: "HfApi",425 repo_id: str,426 repo_type: str,427 revision: str,428):429 """430 Main process for a worker. The worker will perform tasks based on the priority list until all files are uploaded431 and committed. If no tasks are available, the worker will wait for 10 seconds before checking again.432 433 If a task fails for any reason, the item(s) are put back in the queue for another worker to pick up.434 435 Read `upload_large_folder` docstring for more information on how tasks are prioritized.436 """437 while True:438 next_job: tuple[WorkerJob, list[JOB_ITEM_T]] | None = None439 440 # Determine next task441 next_job = _determine_next_job(status)442 if next_job is None:443 return444 job, items = next_job445 446 # Perform task447 match job:448 case WorkerJob.SHA256:449 item = items[0] # single item450 try:451 _compute_sha256(item)452 status.queue_get_upload_mode.put(item)453 except KeyboardInterrupt:454 raise455 except Exception as e:456 logger.error(f"Failed to compute sha256: {e}")457 traceback.format_exc()458 status.queue_sha256.put(item)459 460 with status.lock:461 status.nb_workers_sha256 -= 1462 463 case WorkerJob.GET_UPLOAD_MODE:464 try:465 _get_upload_mode(items, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision)466 except KeyboardInterrupt:467 raise468 except Exception as e:469 logger.error(f"Failed to get upload mode: {e}")470 traceback.format_exc()471 472 # Items are either:473 # - dropped (if should_ignore)474 # - put in LFS queue (if LFS)475 # - put in commit queue (if regular)476 # - or put back (if error occurred).477 for item in items:478 _, metadata = item479 if metadata.should_ignore:480 continue481 match metadata.upload_mode:482 case "lfs":483 status.queue_preupload_lfs.put(item)484 case "regular":485 status.queue_commit.put(item)486 case _:487 status.queue_get_upload_mode.put(item)488 489 with status.lock:490 status.nb_workers_get_upload_mode -= 1491 492 case WorkerJob.PREUPLOAD_LFS:493 try:494 _preupload_lfs(items, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision)495 for item in items:496 status.queue_commit.put(item)497 except KeyboardInterrupt:498 raise499 except Exception as e:500 logger.error(f"Failed to preupload LFS: {e}")501 traceback.format_exc()502 for item in items:503 status.queue_preupload_lfs.put(item)504 505 with status.lock:506 status.nb_workers_preupload_lfs -= 1507 508 case WorkerJob.COMMIT:509 start_ts = time.time()510 success = True511 try:512 _commit(items, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision)513 except KeyboardInterrupt:514 raise515 except Exception as e:516 logger.error(f"Failed to commit: {e}")517 traceback.format_exc()518 for item in items:519 status.queue_commit.put(item)520 success = False521 duration = time.time() - start_ts522 status.update_chunk(success, len(items), duration)523 with status.lock:524 status.last_commit_attempt = time.time()525 status.nb_workers_commit -= 1526 527 case WorkerJob.WAIT:528 time.sleep(WAITING_TIME_IF_NO_TASKS)529 with status.lock:530 status.nb_workers_waiting -= 1531 532 533def _determine_next_job(status: LargeUploadStatus) -> tuple[WorkerJob, list[JOB_ITEM_T]] | None:534 with status.lock:535 # 1. Commit if more than 5 minutes since last commit attempt (and at least 1 file)536 if (537 status.nb_workers_commit == 0538 and status.queue_commit.qsize() > 0539 and status.last_commit_attempt is not None540 and time.time() - status.last_commit_attempt > 5 * 60541 ):542 status.nb_workers_commit += 1543 logger.debug("Job: commit (more than 5 minutes since last commit attempt)")544 return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))545 546 # 2. Commit if at least 100 files are ready to commit547 elif status.nb_workers_commit == 0 and status.queue_commit.qsize() >= 150:548 status.nb_workers_commit += 1549 logger.debug("Job: commit (>100 files ready)")550 return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))551 552 # 3. Get upload mode if at least 100 files553 elif status.queue_get_upload_mode.qsize() >= MAX_NB_FILES_FETCH_UPLOAD_MODE:554 status.nb_workers_get_upload_mode += 1555 logger.debug(f"Job: get upload mode (>{MAX_NB_FILES_FETCH_UPLOAD_MODE} files ready)")556 return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, MAX_NB_FILES_FETCH_UPLOAD_MODE))557 558 # 4. Preupload LFS file if at least `status.upload_batch_size` files and no worker is preuploading LFS559 elif status.queue_preupload_lfs.qsize() >= status.upload_batch_size and status.nb_workers_preupload_lfs == 0:560 status.nb_workers_preupload_lfs += 1561 logger.debug("Job: preupload LFS (no other worker preuploading LFS)")562 return (WorkerJob.PREUPLOAD_LFS, _get_n(status.queue_preupload_lfs, status.upload_batch_size))563 564 # 5. Compute sha256 if at least 1 file and no worker is computing sha256565 elif status.queue_sha256.qsize() > 0 and status.nb_workers_sha256 == 0:566 status.nb_workers_sha256 += 1567 logger.debug("Job: sha256 (no other worker computing sha256)")568 return (WorkerJob.SHA256, _get_one(status.queue_sha256))569 570 # 6. Get upload mode if at least 1 file and no worker is getting upload mode571 elif status.queue_get_upload_mode.qsize() > 0 and status.nb_workers_get_upload_mode == 0:572 status.nb_workers_get_upload_mode += 1573 logger.debug("Job: get upload mode (no other worker getting upload mode)")574 return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, MAX_NB_FILES_FETCH_UPLOAD_MODE))575 576 # 7. Preupload LFS file if at least `status.upload_batch_size` files577 elif status.queue_preupload_lfs.qsize() >= status.upload_batch_size:578 status.nb_workers_preupload_lfs += 1579 logger.debug("Job: preupload LFS")580 return (WorkerJob.PREUPLOAD_LFS, _get_n(status.queue_preupload_lfs, status.upload_batch_size))581 582 # 8. Compute sha256 if at least 1 file583 elif status.queue_sha256.qsize() > 0:584 status.nb_workers_sha256 += 1585 logger.debug("Job: sha256")586 return (WorkerJob.SHA256, _get_one(status.queue_sha256))587 588 # 9. Get upload mode if at least 1 file589 elif status.queue_get_upload_mode.qsize() > 0:590 status.nb_workers_get_upload_mode += 1591 logger.debug("Job: get upload mode")592 return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, MAX_NB_FILES_FETCH_UPLOAD_MODE))593 594 # 10. Preupload LFS file if at least 1 file595 elif status.queue_preupload_lfs.qsize() > 0:596 status.nb_workers_preupload_lfs += 1597 logger.debug("Job: preupload LFS")598 return (WorkerJob.PREUPLOAD_LFS, _get_n(status.queue_preupload_lfs, status.upload_batch_size))599 600 # 11. Commit if at least 1 file and 1 min since last commit attempt601 elif (602 status.nb_workers_commit == 0603 and status.queue_commit.qsize() > 0604 and status.last_commit_attempt is not None605 and time.time() - status.last_commit_attempt > 1 * 60606 ):607 status.nb_workers_commit += 1608 logger.debug("Job: commit (1 min since last commit attempt)")609 return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))610 611 # 12. Commit if at least 1 file all other queues are empty and all workers are waiting612 # e.g. when it's the last commit613 elif (614 status.nb_workers_commit == 0615 and status.queue_commit.qsize() > 0616 and status.queue_sha256.qsize() == 0617 and status.queue_get_upload_mode.qsize() == 0618 and status.queue_preupload_lfs.qsize() == 0619 and status.nb_workers_sha256 == 0620 and status.nb_workers_get_upload_mode == 0621 and status.nb_workers_preupload_lfs == 0622 ):623 status.nb_workers_commit += 1624 logger.debug("Job: commit")625 return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))626 627 # 13. If all queues are empty, exit628 elif all(metadata.is_committed or metadata.should_ignore for _, metadata in status.items):629 logger.info("All files have been processed! Exiting worker.")630 return None631 632 # 14. If no task is available, wait633 else:634 status.nb_workers_waiting += 1635 logger.debug(f"No task available, waiting... ({WAITING_TIME_IF_NO_TASKS}s)")636 return (WorkerJob.WAIT, [])637 638 639####################640# Atomic jobs (sha256, get_upload_mode, preupload_lfs, commit)641####################642 643 644def _compute_sha256(item: JOB_ITEM_T) -> None:645 """Compute sha256 of a file and save it in metadata."""646 paths, metadata = item647 if metadata.sha256 is None:648 with paths.file_path.open("rb") as f:649 metadata.sha256 = sha_fileobj(f).hex()650 metadata.save(paths)651 652 653def _get_upload_mode(items: list[JOB_ITEM_T], api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None:654 """Get upload mode for each file and update metadata.655 656 Also receive info if the file should be ignored.657 """658 additions = [_build_hacky_operation(item) for item in items]659 _fetch_upload_modes(660 additions=additions,661 repo_type=repo_type,662 repo_id=repo_id,663 headers=api._build_hf_headers(),664 revision=quote(revision, safe=""),665 endpoint=api.endpoint,666 )667 for item, addition in zip(items, additions):668 paths, metadata = item669 metadata.upload_mode = addition._upload_mode670 metadata.should_ignore = addition._should_ignore671 metadata.remote_oid = addition._remote_oid672 metadata.save(paths)673 674 675def _preupload_lfs(items: list[JOB_ITEM_T], api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None:676 """Preupload LFS files and update metadata."""677 additions = [_build_hacky_operation(item) for item in items]678 api.preupload_lfs_files(679 repo_id=repo_id,680 repo_type=repo_type,681 revision=revision,682 additions=additions,683 )684 685 for paths, metadata in items:686 metadata.is_uploaded = True687 metadata.save(paths)688 689 690def _commit(items: list[JOB_ITEM_T], api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None:691 """Commit files to the repo."""692 additions = [_build_hacky_operation(item) for item in items]693 api.create_commit(694 repo_id=repo_id,695 repo_type=repo_type,696 revision=revision,697 operations=additions,698 commit_message="Add files using upload-large-folder tool",699 )700 for paths, metadata in items:701 metadata.is_committed = True702 metadata.save(paths)703 704 705####################706# Hacks with CommitOperationAdd to bypass checks/sha256 calculation707####################708 709 710class HackyCommitOperationAdd(CommitOperationAdd):711 def __post_init__(self) -> None:712 if isinstance(self.path_or_fileobj, Path):713 self.path_or_fileobj = str(self.path_or_fileobj)714 715 716def _build_hacky_operation(item: JOB_ITEM_T) -> HackyCommitOperationAdd:717 paths, metadata = item718 operation = HackyCommitOperationAdd(path_in_repo=paths.path_in_repo, path_or_fileobj=paths.file_path)719 with paths.file_path.open("rb") as file:720 sample = file.peek(512)[:512]721 if metadata.sha256 is None:722 raise ValueError("sha256 must have been computed by now!")723 operation.upload_info = UploadInfo(sha256=bytes.fromhex(metadata.sha256), size=metadata.size, sample=sample)724 operation._upload_mode = metadata.upload_mode # type: ignore725 operation._should_ignore = metadata.should_ignore726 operation._remote_oid = metadata.remote_oid727 operation._is_uploaded = metadata.is_uploaded728 if metadata.is_uploaded and metadata.upload_mode == "lfs":729 operation.path_or_fileobj = b""730 return operation731 732 733####################734# Misc helpers735####################736 737 738def _get_one(queue: "queue.Queue[JOB_ITEM_T]") -> list[JOB_ITEM_T]:739 return [queue.get()]740 741 742def _get_n(queue: "queue.Queue[JOB_ITEM_T]", n: int) -> list[JOB_ITEM_T]:743 return [queue.get() for _ in range(min(queue.qsize(), n))]744 745 746def _print_overwrite(report: str) -> None:747 """Print a report, overwriting the previous lines.748 749 Since tqdm in using `sys.stderr` to (re-)write progress bars, we need to use `sys.stdout`750 to print the report.751 752 Note: works well only if no other process is writing to `sys.stdout`!753 """754 report += "\n"755 # Get terminal width756 terminal_width = shutil.get_terminal_size().columns757 758 # Count number of lines that should be cleared759 nb_lines = sum(len(line) // terminal_width + 1 for line in report.splitlines())760 761 # Clear previous lines based on the number of lines in the report762 for _ in range(nb_lines):763 sys.stdout.write("\r\033[K") # Clear line764 sys.stdout.write("\033[F") # Move cursor up one line765 766 # Print the new report, filling remaining space with whitespace767 sys.stdout.write(report)768 sys.stdout.write(" " * (terminal_width - len(report.splitlines()[-1])))769 sys.stdout.flush()770 