_upload_large_folder.py 29 KB

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