_upload_large_folder.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766
  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 at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # 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 and
  13. # limitations under the License.
  14. import enum
  15. import logging
  16. import os
  17. import queue
  18. import shutil
  19. import sys
  20. import threading
  21. import time
  22. import traceback
  23. from datetime import datetime
  24. from pathlib import Path
  25. from threading import Lock
  26. from typing import TYPE_CHECKING, Any
  27. from urllib.parse import quote
  28. from ._commit_api import CommitOperationAdd, UploadInfo, _fetch_upload_modes
  29. from ._local_folder import LocalUploadFileMetadata, LocalUploadFilePaths, get_local_upload_paths, read_upload_metadata
  30. from .constants import DEFAULT_REVISION, REPO_TYPES
  31. from .utils import DEFAULT_IGNORE_PATTERNS, _format_size, filter_repo_objects, tqdm
  32. from .utils._runtime import is_xet_available
  33. from .utils.sha import sha_fileobj
  34. if TYPE_CHECKING:
  35. from .hf_api import HfApi
  36. logger = logging.getLogger(__name__)
  37. WAITING_TIME_IF_NO_TASKS = 10 # seconds
  38. MAX_NB_FILES_FETCH_UPLOAD_MODE = 100
  39. COMMIT_SIZE_SCALE: list[int] = [20, 50, 75, 100, 125, 200, 250, 400, 600, 1000]
  40. UPLOAD_BATCH_SIZE_XET = 256 # Max 256 files per upload batch for XET-enabled repos
  41. UPLOAD_BATCH_SIZE_LFS = 1 # Otherwise, batches of 1 for regular LFS upload
  42. # Repository limits (from https://huggingface.co/docs/hub/repositories-recommendations)
  43. MAX_FILES_PER_REPO = 100_000 # Recommended maximum number of files per repository
  44. MAX_FILES_PER_FOLDER = 10_000 # Recommended maximum number of files per folder
  45. MAX_FILE_SIZE_GB = 200 # Recommended maximum for individual file size (split larger files)
  46. RECOMMENDED_FILE_SIZE_GB = 20 # Recommended maximum for individual file size
  47. def _validate_upload_limits(paths_list: list[LocalUploadFilePaths]) -> None:
  48. """
  49. Validate upload against repository limits and warn about potential issues.
  50. Args:
  51. paths_list: List of file paths to be uploaded
  52. Warns about:
  53. - Too many files in the repository (>100k)
  54. - Too many entries (files or subdirectories) in a single folder (>10k)
  55. - Files exceeding size limits (>20GB recommended, >200GB maximum)
  56. """
  57. logger.info("Running validation checks on files to upload...")
  58. # Check 1: Total file count
  59. if len(paths_list) > MAX_FILES_PER_REPO:
  60. logger.warning(
  61. f"You are about to upload {len(paths_list):,} files. "
  62. f"This exceeds the recommended limit of {MAX_FILES_PER_REPO:,} files per repository.\n"
  63. f"Consider:\n"
  64. f" - Splitting your data into multiple repositories\n"
  65. f" - Using fewer, larger files (e.g., parquet files)\n"
  66. f" - See: https://huggingface.co/docs/hub/repositories-recommendations"
  67. )
  68. # Check 2: Files and subdirectories per folder
  69. # Track immediate children (files and subdirs) for each folder
  70. from collections import defaultdict
  71. entries_per_folder: dict[str, Any] = defaultdict(lambda: {"files": 0, "subdirs": set()})
  72. for paths in paths_list:
  73. path = Path(paths.path_in_repo)
  74. parts = path.parts
  75. # Count this file in its immediate parent directory
  76. parent = str(path.parent) if str(path.parent) != "." else "."
  77. entries_per_folder[parent]["files"] += 1
  78. # Track immediate subdirectories for each parent folder
  79. # Walk through the path components to track parent-child relationships
  80. for i, child in enumerate(parts[:-1]):
  81. parent = "." if i == 0 else "/".join(parts[:i])
  82. entries_per_folder[parent]["subdirs"].add(child)
  83. # Check limits for each folder
  84. for folder, data in entries_per_folder.items():
  85. file_count = data["files"]
  86. subdir_count = len(data["subdirs"])
  87. total_entries = file_count + subdir_count
  88. if total_entries > MAX_FILES_PER_FOLDER:
  89. folder_display = "root" if folder == "." else folder
  90. logger.warning(
  91. f"Folder '{folder_display}' contains {total_entries:,} entries "
  92. f"({file_count:,} files and {subdir_count:,} subdirectories). "
  93. f"This exceeds the recommended {MAX_FILES_PER_FOLDER:,} entries per folder.\n"
  94. "Consider reorganising into sub-folders."
  95. )
  96. # Check 3: File sizes
  97. large_files = []
  98. very_large_files = []
  99. for paths in paths_list:
  100. size = paths.file_path.stat().st_size
  101. size_gb = size / 1_000_000_000 # Use decimal GB as per Hub limits
  102. if size_gb > MAX_FILE_SIZE_GB:
  103. very_large_files.append((paths.path_in_repo, size_gb))
  104. elif size_gb > RECOMMENDED_FILE_SIZE_GB:
  105. large_files.append((paths.path_in_repo, size_gb))
  106. # Warn about very large files (>200GB)
  107. if very_large_files:
  108. files_str = "\n - ".join(f"{path}: {size:.1f}GB" for path, size in very_large_files[:5])
  109. more_str = f"\n ... and {len(very_large_files) - 5} more files" if len(very_large_files) > 5 else ""
  110. logger.warning(
  111. f"Found {len(very_large_files)} files exceeding the {MAX_FILE_SIZE_GB}GB recommended maximum:\n"
  112. f" - {files_str}{more_str}\n"
  113. f"Consider splitting these files into smaller chunks."
  114. )
  115. # Warn about large files (>20GB)
  116. if large_files:
  117. files_str = "\n - ".join(f"{path}: {size:.1f}GB" for path, size in large_files[:5])
  118. more_str = f"\n ... and {len(large_files) - 5} more files" if len(large_files) > 5 else ""
  119. logger.warning(
  120. f"Found {len(large_files)} files larger than {RECOMMENDED_FILE_SIZE_GB}GB (recommended limit):\n"
  121. f" - {files_str}{more_str}\n"
  122. f"Large files may slow down loading and processing."
  123. )
  124. logger.info("Validation checks complete.")
  125. def upload_large_folder_internal(
  126. api: "HfApi",
  127. repo_id: str,
  128. folder_path: str | Path,
  129. *,
  130. repo_type: str, # Repo type is required!
  131. revision: str | None = None,
  132. private: bool | None = None,
  133. allow_patterns: list[str] | str | None = None,
  134. ignore_patterns: list[str] | str | None = None,
  135. num_workers: int | None = None,
  136. print_report: bool = True,
  137. print_report_every: int = 60,
  138. ):
  139. """Upload a large folder to the Hub in the most resilient way possible.
  140. See [`HfApi.upload_large_folder`] for the full documentation.
  141. """
  142. # 1. Check args and setup
  143. if repo_type is None:
  144. raise ValueError(
  145. "For large uploads, `repo_type` is explicitly required. Please set it to `model`, `dataset` or `space`."
  146. " If you are using the CLI, pass it as `--repo-type=model`."
  147. )
  148. if repo_type not in REPO_TYPES:
  149. raise ValueError(f"Invalid repo type, must be one of {REPO_TYPES}")
  150. if revision is None:
  151. revision = DEFAULT_REVISION
  152. folder_path = Path(folder_path).expanduser().resolve()
  153. if not folder_path.is_dir():
  154. raise ValueError(f"Provided path: '{folder_path}' is not a directory")
  155. if ignore_patterns is None:
  156. ignore_patterns = []
  157. elif isinstance(ignore_patterns, str):
  158. ignore_patterns = [ignore_patterns]
  159. ignore_patterns += DEFAULT_IGNORE_PATTERNS
  160. if num_workers is None:
  161. nb_cores = os.cpu_count() or 1
  162. num_workers = max(nb_cores // 2, 1) # Use at most half of cpu cores
  163. # 2. Create repo if missing
  164. repo_url = api.create_repo(repo_id=repo_id, repo_type=repo_type, private=private, exist_ok=True)
  165. logger.info(f"Repo created: {repo_url}")
  166. repo_id = repo_url.repo_id
  167. # Warn on too many commits
  168. try:
  169. commits = api.list_repo_commits(repo_id=repo_id, repo_type=repo_type, revision=revision)
  170. commit_count = len(commits)
  171. if commit_count > 500:
  172. logger.warning(
  173. f"\n{'=' * 80}\n"
  174. f"WARNING: This repository has {commit_count} commits.\n"
  175. f"Repositories with a large number of commits can experience performance issues.\n"
  176. f"\n"
  177. f"Consider squashing your commit history using `super_squash_history()`.\n"
  178. "To do so, you need to stop this process, run the snippet below and restart the upload command."
  179. f" from huggingface_hub import super_squash_history\n"
  180. f" super_squash_history(repo_id='{repo_id}', repo_type='{repo_type}')\n"
  181. f"\n"
  182. f"Note: This is a non-revertible operation. See the documentation for more details:\n"
  183. f"https://huggingface.co/docs/huggingface_hub/main/en/package_reference/hf_api#huggingface_hub.HfApi.super_squash_history\n"
  184. f"{'=' * 80}\n"
  185. )
  186. except Exception as e:
  187. # Don't fail the upload if we can't check commit count
  188. logger.debug(f"Could not check commit count: {e}")
  189. # 2.1 Check if xet is enabled to set batch file upload size
  190. upload_batch_size = UPLOAD_BATCH_SIZE_XET if is_xet_available() else UPLOAD_BATCH_SIZE_LFS
  191. # 3. List files to upload
  192. filtered_paths_list = filter_repo_objects(
  193. (path.relative_to(folder_path).as_posix() for path in folder_path.glob("**/*") if path.is_file()),
  194. allow_patterns=allow_patterns,
  195. ignore_patterns=ignore_patterns,
  196. )
  197. paths_list = [get_local_upload_paths(folder_path, relpath) for relpath in filtered_paths_list]
  198. logger.info(f"Found {len(paths_list)} candidate files to upload")
  199. # Validate upload against repository limits
  200. _validate_upload_limits(paths_list)
  201. logger.info("Starting upload...")
  202. # Read metadata for each file
  203. items = [
  204. (paths, read_upload_metadata(folder_path, paths.path_in_repo))
  205. for paths in tqdm(paths_list, desc="Recovering from metadata files")
  206. ]
  207. # 4. Start workers
  208. status = LargeUploadStatus(items, upload_batch_size)
  209. threads = [
  210. threading.Thread(
  211. target=_worker_job,
  212. kwargs={
  213. "status": status,
  214. "api": api,
  215. "repo_id": repo_id,
  216. "repo_type": repo_type,
  217. "revision": revision,
  218. },
  219. )
  220. for _ in range(num_workers)
  221. ]
  222. for thread in threads:
  223. thread.start()
  224. # 5. Print regular reports
  225. if print_report:
  226. print("\n\n" + status.current_report())
  227. last_report_ts = time.time()
  228. while True:
  229. time.sleep(1)
  230. if time.time() - last_report_ts >= print_report_every:
  231. if print_report:
  232. _print_overwrite(status.current_report())
  233. last_report_ts = time.time()
  234. if status.is_done():
  235. logger.info("Is done: exiting main loop")
  236. break
  237. for thread in threads:
  238. thread.join()
  239. logger.info(status.current_report())
  240. logger.info("Upload is complete!")
  241. ####################
  242. # Logic to manage workers and synchronize tasks
  243. ####################
  244. class WorkerJob(enum.Enum):
  245. SHA256 = enum.auto()
  246. GET_UPLOAD_MODE = enum.auto()
  247. PREUPLOAD_LFS = enum.auto()
  248. COMMIT = enum.auto()
  249. WAIT = enum.auto() # if no tasks are available but we don't want to exit
  250. JOB_ITEM_T = tuple[LocalUploadFilePaths, LocalUploadFileMetadata]
  251. class LargeUploadStatus:
  252. """Contains information, queues and tasks for a large upload process."""
  253. def __init__(self, items: list[JOB_ITEM_T], upload_batch_size: int = 1):
  254. self.items = items
  255. self.queue_sha256: "queue.Queue[JOB_ITEM_T]" = queue.Queue()
  256. self.queue_get_upload_mode: "queue.Queue[JOB_ITEM_T]" = queue.Queue()
  257. self.queue_preupload_lfs: "queue.Queue[JOB_ITEM_T]" = queue.Queue()
  258. self.queue_commit: "queue.Queue[JOB_ITEM_T]" = queue.Queue()
  259. self.lock = Lock()
  260. self.nb_workers_sha256: int = 0
  261. self.nb_workers_get_upload_mode: int = 0
  262. self.nb_workers_preupload_lfs: int = 0
  263. self.upload_batch_size: int = upload_batch_size
  264. self.nb_workers_commit: int = 0
  265. self.nb_workers_waiting: int = 0
  266. self.last_commit_attempt: float | None = None
  267. self._started_at = datetime.now()
  268. self._chunk_idx: int = 1
  269. self._chunk_lock: Lock = Lock()
  270. # Setup queues
  271. for item in self.items:
  272. paths, metadata = item
  273. if metadata.sha256 is None:
  274. self.queue_sha256.put(item)
  275. elif metadata.upload_mode is None:
  276. self.queue_get_upload_mode.put(item)
  277. elif metadata.upload_mode == "lfs" and not metadata.is_uploaded:
  278. self.queue_preupload_lfs.put(item)
  279. elif not metadata.is_committed:
  280. self.queue_commit.put(item)
  281. else:
  282. logger.debug(f"Skipping file {paths.path_in_repo} (already uploaded and committed)")
  283. def target_chunk(self) -> int:
  284. with self._chunk_lock:
  285. return COMMIT_SIZE_SCALE[self._chunk_idx]
  286. def update_chunk(self, success: bool, nb_items: int, duration: float) -> None:
  287. with self._chunk_lock:
  288. if not success:
  289. logger.warning(f"Failed to commit {nb_items} files at once. Will retry with less files in next batch.")
  290. self._chunk_idx -= 1
  291. elif nb_items >= COMMIT_SIZE_SCALE[self._chunk_idx] and duration < 40:
  292. logger.info(f"Successfully committed {nb_items} at once. Increasing the limit for next batch.")
  293. self._chunk_idx += 1
  294. self._chunk_idx = max(0, min(self._chunk_idx, len(COMMIT_SIZE_SCALE) - 1))
  295. def current_report(self) -> str:
  296. """Generate a report of the current status of the large upload."""
  297. nb_hashed = 0
  298. size_hashed = 0
  299. nb_preuploaded = 0
  300. nb_lfs = 0
  301. nb_lfs_unsure = 0
  302. size_preuploaded = 0
  303. nb_committed = 0
  304. size_committed = 0
  305. total_size = 0
  306. ignored_files = 0
  307. total_files = 0
  308. with self.lock:
  309. for _, metadata in self.items:
  310. if metadata.should_ignore:
  311. ignored_files += 1
  312. continue
  313. total_size += metadata.size
  314. total_files += 1
  315. if metadata.sha256 is not None:
  316. nb_hashed += 1
  317. size_hashed += metadata.size
  318. if metadata.upload_mode == "lfs":
  319. nb_lfs += 1
  320. if metadata.upload_mode is None:
  321. nb_lfs_unsure += 1
  322. if metadata.is_uploaded:
  323. nb_preuploaded += 1
  324. size_preuploaded += metadata.size
  325. if metadata.is_committed:
  326. nb_committed += 1
  327. size_committed += metadata.size
  328. total_size_str = _format_size(total_size)
  329. now = datetime.now()
  330. now_str = now.strftime("%Y-%m-%d %H:%M:%S")
  331. elapsed = now - self._started_at
  332. elapsed_str = str(elapsed).split(".")[0] # remove milliseconds
  333. message = "\n" + "-" * 10
  334. message += f" {now_str} ({elapsed_str}) "
  335. message += "-" * 10 + "\n"
  336. message += "Files: "
  337. message += f"hashed {nb_hashed}/{total_files} ({_format_size(size_hashed)}/{total_size_str}) | "
  338. message += f"pre-uploaded: {nb_preuploaded}/{nb_lfs} ({_format_size(size_preuploaded)}/{total_size_str})"
  339. if nb_lfs_unsure > 0:
  340. message += f" (+{nb_lfs_unsure} unsure)"
  341. message += f" | committed: {nb_committed}/{total_files} ({_format_size(size_committed)}/{total_size_str})"
  342. message += f" | ignored: {ignored_files}\n"
  343. message += "Workers: "
  344. message += f"hashing: {self.nb_workers_sha256} | "
  345. message += f"get upload mode: {self.nb_workers_get_upload_mode} | "
  346. message += f"pre-uploading: {self.nb_workers_preupload_lfs} | "
  347. message += f"committing: {self.nb_workers_commit} | "
  348. message += f"waiting: {self.nb_workers_waiting}\n"
  349. message += "-" * 51
  350. return message
  351. def is_done(self) -> bool:
  352. with self.lock:
  353. return all(metadata.is_committed or metadata.should_ignore for _, metadata in self.items)
  354. def _worker_job(
  355. status: LargeUploadStatus,
  356. api: "HfApi",
  357. repo_id: str,
  358. repo_type: str,
  359. revision: str,
  360. ):
  361. """
  362. Main process for a worker. The worker will perform tasks based on the priority list until all files are uploaded
  363. and committed. If no tasks are available, the worker will wait for 10 seconds before checking again.
  364. If a task fails for any reason, the item(s) are put back in the queue for another worker to pick up.
  365. Read `upload_large_folder` docstring for more information on how tasks are prioritized.
  366. """
  367. while True:
  368. next_job: tuple[WorkerJob, list[JOB_ITEM_T]] | None = None
  369. # Determine next task
  370. next_job = _determine_next_job(status)
  371. if next_job is None:
  372. return
  373. job, items = next_job
  374. # Perform task
  375. match job:
  376. case WorkerJob.SHA256:
  377. item = items[0] # single item
  378. try:
  379. _compute_sha256(item)
  380. status.queue_get_upload_mode.put(item)
  381. except KeyboardInterrupt:
  382. raise
  383. except Exception as e:
  384. logger.error(f"Failed to compute sha256: {e}")
  385. traceback.format_exc()
  386. status.queue_sha256.put(item)
  387. with status.lock:
  388. status.nb_workers_sha256 -= 1
  389. case WorkerJob.GET_UPLOAD_MODE:
  390. try:
  391. _get_upload_mode(items, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision)
  392. except KeyboardInterrupt:
  393. raise
  394. except Exception as e:
  395. logger.error(f"Failed to get upload mode: {e}")
  396. traceback.format_exc()
  397. # Items are either:
  398. # - dropped (if should_ignore)
  399. # - put in LFS queue (if LFS)
  400. # - put in commit queue (if regular)
  401. # - or put back (if error occurred).
  402. for item in items:
  403. _, metadata = item
  404. if metadata.should_ignore:
  405. continue
  406. match metadata.upload_mode:
  407. case "lfs":
  408. status.queue_preupload_lfs.put(item)
  409. case "regular":
  410. status.queue_commit.put(item)
  411. case _:
  412. status.queue_get_upload_mode.put(item)
  413. with status.lock:
  414. status.nb_workers_get_upload_mode -= 1
  415. case WorkerJob.PREUPLOAD_LFS:
  416. try:
  417. _preupload_lfs(items, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision)
  418. for item in items:
  419. status.queue_commit.put(item)
  420. except KeyboardInterrupt:
  421. raise
  422. except Exception as e:
  423. logger.error(f"Failed to preupload LFS: {e}")
  424. traceback.format_exc()
  425. for item in items:
  426. status.queue_preupload_lfs.put(item)
  427. with status.lock:
  428. status.nb_workers_preupload_lfs -= 1
  429. case WorkerJob.COMMIT:
  430. start_ts = time.time()
  431. success = True
  432. try:
  433. _commit(items, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision)
  434. except KeyboardInterrupt:
  435. raise
  436. except Exception as e:
  437. logger.error(f"Failed to commit: {e}")
  438. traceback.format_exc()
  439. for item in items:
  440. status.queue_commit.put(item)
  441. success = False
  442. duration = time.time() - start_ts
  443. status.update_chunk(success, len(items), duration)
  444. with status.lock:
  445. status.last_commit_attempt = time.time()
  446. status.nb_workers_commit -= 1
  447. case WorkerJob.WAIT:
  448. time.sleep(WAITING_TIME_IF_NO_TASKS)
  449. with status.lock:
  450. status.nb_workers_waiting -= 1
  451. def _determine_next_job(status: LargeUploadStatus) -> tuple[WorkerJob, list[JOB_ITEM_T]] | None:
  452. with status.lock:
  453. # 1. Commit if more than 5 minutes since last commit attempt (and at least 1 file)
  454. if (
  455. status.nb_workers_commit == 0
  456. and status.queue_commit.qsize() > 0
  457. and status.last_commit_attempt is not None
  458. and time.time() - status.last_commit_attempt > 5 * 60
  459. ):
  460. status.nb_workers_commit += 1
  461. logger.debug("Job: commit (more than 5 minutes since last commit attempt)")
  462. return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))
  463. # 2. Commit if at least 100 files are ready to commit
  464. elif status.nb_workers_commit == 0 and status.queue_commit.qsize() >= 150:
  465. status.nb_workers_commit += 1
  466. logger.debug("Job: commit (>100 files ready)")
  467. return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))
  468. # 3. Get upload mode if at least 100 files
  469. elif status.queue_get_upload_mode.qsize() >= MAX_NB_FILES_FETCH_UPLOAD_MODE:
  470. status.nb_workers_get_upload_mode += 1
  471. logger.debug(f"Job: get upload mode (>{MAX_NB_FILES_FETCH_UPLOAD_MODE} files ready)")
  472. return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, MAX_NB_FILES_FETCH_UPLOAD_MODE))
  473. # 4. Preupload LFS file if at least `status.upload_batch_size` files and no worker is preuploading LFS
  474. elif status.queue_preupload_lfs.qsize() >= status.upload_batch_size and status.nb_workers_preupload_lfs == 0:
  475. status.nb_workers_preupload_lfs += 1
  476. logger.debug("Job: preupload LFS (no other worker preuploading LFS)")
  477. return (WorkerJob.PREUPLOAD_LFS, _get_n(status.queue_preupload_lfs, status.upload_batch_size))
  478. # 5. Compute sha256 if at least 1 file and no worker is computing sha256
  479. elif status.queue_sha256.qsize() > 0 and status.nb_workers_sha256 == 0:
  480. status.nb_workers_sha256 += 1
  481. logger.debug("Job: sha256 (no other worker computing sha256)")
  482. return (WorkerJob.SHA256, _get_one(status.queue_sha256))
  483. # 6. Get upload mode if at least 1 file and no worker is getting upload mode
  484. elif status.queue_get_upload_mode.qsize() > 0 and status.nb_workers_get_upload_mode == 0:
  485. status.nb_workers_get_upload_mode += 1
  486. logger.debug("Job: get upload mode (no other worker getting upload mode)")
  487. return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, MAX_NB_FILES_FETCH_UPLOAD_MODE))
  488. # 7. Preupload LFS file if at least `status.upload_batch_size` files
  489. elif status.queue_preupload_lfs.qsize() >= status.upload_batch_size:
  490. status.nb_workers_preupload_lfs += 1
  491. logger.debug("Job: preupload LFS")
  492. return (WorkerJob.PREUPLOAD_LFS, _get_n(status.queue_preupload_lfs, status.upload_batch_size))
  493. # 8. Compute sha256 if at least 1 file
  494. elif status.queue_sha256.qsize() > 0:
  495. status.nb_workers_sha256 += 1
  496. logger.debug("Job: sha256")
  497. return (WorkerJob.SHA256, _get_one(status.queue_sha256))
  498. # 9. Get upload mode if at least 1 file
  499. elif status.queue_get_upload_mode.qsize() > 0:
  500. status.nb_workers_get_upload_mode += 1
  501. logger.debug("Job: get upload mode")
  502. return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, MAX_NB_FILES_FETCH_UPLOAD_MODE))
  503. # 10. Preupload LFS file if at least 1 file
  504. elif status.queue_preupload_lfs.qsize() > 0:
  505. status.nb_workers_preupload_lfs += 1
  506. logger.debug("Job: preupload LFS")
  507. return (WorkerJob.PREUPLOAD_LFS, _get_n(status.queue_preupload_lfs, status.upload_batch_size))
  508. # 11. Commit if at least 1 file and 1 min since last commit attempt
  509. elif (
  510. status.nb_workers_commit == 0
  511. and status.queue_commit.qsize() > 0
  512. and status.last_commit_attempt is not None
  513. and time.time() - status.last_commit_attempt > 1 * 60
  514. ):
  515. status.nb_workers_commit += 1
  516. logger.debug("Job: commit (1 min since last commit attempt)")
  517. return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))
  518. # 12. Commit if at least 1 file all other queues are empty and all workers are waiting
  519. # e.g. when it's the last commit
  520. elif (
  521. status.nb_workers_commit == 0
  522. and status.queue_commit.qsize() > 0
  523. and status.queue_sha256.qsize() == 0
  524. and status.queue_get_upload_mode.qsize() == 0
  525. and status.queue_preupload_lfs.qsize() == 0
  526. and status.nb_workers_sha256 == 0
  527. and status.nb_workers_get_upload_mode == 0
  528. and status.nb_workers_preupload_lfs == 0
  529. ):
  530. status.nb_workers_commit += 1
  531. logger.debug("Job: commit")
  532. return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))
  533. # 13. If all queues are empty, exit
  534. elif all(metadata.is_committed or metadata.should_ignore for _, metadata in status.items):
  535. logger.info("All files have been processed! Exiting worker.")
  536. return None
  537. # 14. If no task is available, wait
  538. else:
  539. status.nb_workers_waiting += 1
  540. logger.debug(f"No task available, waiting... ({WAITING_TIME_IF_NO_TASKS}s)")
  541. return (WorkerJob.WAIT, [])
  542. ####################
  543. # Atomic jobs (sha256, get_upload_mode, preupload_lfs, commit)
  544. ####################
  545. def _compute_sha256(item: JOB_ITEM_T) -> None:
  546. """Compute sha256 of a file and save it in metadata."""
  547. paths, metadata = item
  548. if metadata.sha256 is None:
  549. with paths.file_path.open("rb") as f:
  550. metadata.sha256 = sha_fileobj(f).hex()
  551. metadata.save(paths)
  552. def _get_upload_mode(items: list[JOB_ITEM_T], api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None:
  553. """Get upload mode for each file and update metadata.
  554. Also receive info if the file should be ignored.
  555. """
  556. additions = [_build_hacky_operation(item) for item in items]
  557. _fetch_upload_modes(
  558. additions=additions,
  559. repo_type=repo_type,
  560. repo_id=repo_id,
  561. headers=api._build_hf_headers(),
  562. revision=quote(revision, safe=""),
  563. endpoint=api.endpoint,
  564. )
  565. for item, addition in zip(items, additions):
  566. paths, metadata = item
  567. metadata.upload_mode = addition._upload_mode
  568. metadata.should_ignore = addition._should_ignore
  569. metadata.remote_oid = addition._remote_oid
  570. metadata.save(paths)
  571. def _preupload_lfs(items: list[JOB_ITEM_T], api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None:
  572. """Preupload LFS files and update metadata."""
  573. additions = [_build_hacky_operation(item) for item in items]
  574. api.preupload_lfs_files(
  575. repo_id=repo_id,
  576. repo_type=repo_type,
  577. revision=revision,
  578. additions=additions,
  579. )
  580. for paths, metadata in items:
  581. metadata.is_uploaded = True
  582. metadata.save(paths)
  583. def _commit(items: list[JOB_ITEM_T], api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None:
  584. """Commit files to the repo."""
  585. additions = [_build_hacky_operation(item) for item in items]
  586. api.create_commit(
  587. repo_id=repo_id,
  588. repo_type=repo_type,
  589. revision=revision,
  590. operations=additions,
  591. commit_message="Add files using upload-large-folder tool",
  592. )
  593. for paths, metadata in items:
  594. metadata.is_committed = True
  595. metadata.save(paths)
  596. ####################
  597. # Hacks with CommitOperationAdd to bypass checks/sha256 calculation
  598. ####################
  599. class HackyCommitOperationAdd(CommitOperationAdd):
  600. def __post_init__(self) -> None:
  601. if isinstance(self.path_or_fileobj, Path):
  602. self.path_or_fileobj = str(self.path_or_fileobj)
  603. def _build_hacky_operation(item: JOB_ITEM_T) -> HackyCommitOperationAdd:
  604. paths, metadata = item
  605. operation = HackyCommitOperationAdd(path_in_repo=paths.path_in_repo, path_or_fileobj=paths.file_path)
  606. with paths.file_path.open("rb") as file:
  607. sample = file.peek(512)[:512]
  608. if metadata.sha256 is None:
  609. raise ValueError("sha256 must have been computed by now!")
  610. operation.upload_info = UploadInfo(sha256=bytes.fromhex(metadata.sha256), size=metadata.size, sample=sample)
  611. operation._upload_mode = metadata.upload_mode # type: ignore
  612. operation._should_ignore = metadata.should_ignore
  613. operation._remote_oid = metadata.remote_oid
  614. return operation
  615. ####################
  616. # Misc helpers
  617. ####################
  618. def _get_one(queue: "queue.Queue[JOB_ITEM_T]") -> list[JOB_ITEM_T]:
  619. return [queue.get()]
  620. def _get_n(queue: "queue.Queue[JOB_ITEM_T]", n: int) -> list[JOB_ITEM_T]:
  621. return [queue.get() for _ in range(min(queue.qsize(), n))]
  622. def _print_overwrite(report: str) -> None:
  623. """Print a report, overwriting the previous lines.
  624. Since tqdm in using `sys.stderr` to (re-)write progress bars, we need to use `sys.stdout`
  625. to print the report.
  626. Note: works well only if no other process is writing to `sys.stdout`!
  627. """
  628. report += "\n"
  629. # Get terminal width
  630. terminal_width = shutil.get_terminal_size().columns
  631. # Count number of lines that should be cleared
  632. nb_lines = sum(len(line) // terminal_width + 1 for line in report.splitlines())
  633. # Clear previous lines based on the number of lines in the report
  634. for _ in range(nb_lines):
  635. sys.stdout.write("\r\033[K") # Clear line
  636. sys.stdout.write("\033[F") # Move cursor up one line
  637. # Print the new report, filling remaining space with whitespace
  638. sys.stdout.write(report)
  639. sys.stdout.write(" " * (terminal_width - len(report.splitlines()[-1])))
  640. sys.stdout.flush()