Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes14kdownloads
_upload_large_folder.py770 linesDownload Raw Back to huggingface_hub
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 
codekingpro/portable-devtools · Team Ai