diff --git a/src/modelscope_hub/_legacy_api.py b/src/modelscope_hub/_legacy_api.py index 2e70d52..aa3d217 100644 --- a/src/modelscope_hub/_legacy_api.py +++ b/src/modelscope_hub/_legacy_api.py @@ -14,6 +14,7 @@ from __future__ import annotations +import time import uuid from concurrent.futures import ThreadPoolExecutor, as_completed from typing import IO, Any, BinaryIO @@ -24,18 +25,28 @@ from .constants import ( API_CONNECT_TIMEOUT, + API_CONNECTION_POOL_MAXSIZE, API_MAX_RETRIES, API_TIMEOUT, LEGACY_API_PREFIX, REPO_FILES_TRUNCATION_LIMIT, REPO_TREE_MAX_REQUESTS, + REPO_TREE_PAGE_MAX_ATTEMPTS, + REPO_TREE_PAGE_RETRY_MAX_DELAY_SECONDS, REPO_TREE_WALK_WORKERS, UPLOAD_BLOB_CONNECT_TIMEOUT_SECONDS, UPLOAD_BLOB_READ_TIMEOUT_SECONDS, UPLOAD_HTTP_RETRY_ALLOWED_METHODS, RepoType, ) -from .errors import InvalidParameter, NetworkError, RequestTimeoutError, ServerError, raise_for_status +from .errors import ( + InvalidParameter, + NetworkError, + PermissionDeniedError, + RequestTimeoutError, + ServerError, + raise_for_status, +) from .utils.logger import get_logger logger = get_logger("legacy_api") @@ -100,6 +111,10 @@ def __init__( self._endpoint = endpoint.rstrip("/") self._timeout: int | tuple[int, int] = (API_CONNECT_TIMEOUT, timeout) self._session_authenticated = False + # Set once any file-tree read succeeds. From then on a 403 on a tree + # request cannot be an authorization result, so it is retried instead of + # aborting an enumeration that spans hundreds of requests. + self._tree_reads_ok = False self._session = requests.Session() if user_agent: @@ -110,7 +125,11 @@ def __init__( status_forcelist=[429, 500, 502, 503, 504], allowed_methods=UPLOAD_HTTP_RETRY_ALLOWED_METHODS, ) - adapter = HTTPAdapter(max_retries=retry) + adapter = HTTPAdapter( + max_retries=retry, + pool_connections=API_CONNECTION_POOL_MAXSIZE, + pool_maxsize=API_CONNECTION_POOL_MAXSIZE, + ) self._session.mount("https://", adapter) self._session.mount("http://", adapter) @@ -394,7 +413,12 @@ def _list_files_page( params["Root"] = root suffix = "repo/tree" if _is_dataset(repo_type) else "repo/files" - resp = self._request("GET", f"{segment}/{repo_id}/{suffix}", params=params) + resp = self._request_repo_tree( + f"{segment}/{repo_id}/{suffix}", + params, + repo_id=repo_id, + authorized=self._tree_reads_ok, + ) data = self._json_data(resp) if isinstance(data, list): return data @@ -567,6 +591,13 @@ def list_dataset_files_paginated( Datasets can have millions of files, so this method pages through ``GET /api/v1/datasets/{repo_id}/repo/tree`` with ``PageNumber``/``PageSize`` params. + + A page occasionally answers ``403 无权访问该数据集`` on a repository the + caller demonstrably can read. Once any page has succeeded the credential + is proven, so a later-page denial is a server-side hiccup rather than an + authorization result, and it is retried instead of discarding every page + collected so far -- a large listing spans hundreds of pages, which makes + hitting it near-certain. """ all_files: list[dict] = [] page_number = 1 @@ -579,10 +610,11 @@ def list_dataset_files_paginated( } if root_path and root_path != "/": params["Root"] = root_path - resp = self._request( - "GET", + resp = self._request_repo_tree( f"datasets/{repo_id}/repo/tree", - params=params, + params, + repo_id=repo_id, + authorized=self._tree_reads_ok, ) data = self._json_data(resp) if isinstance(data, list): @@ -598,6 +630,54 @@ def list_dataset_files_paginated( page_number += 1 return all_files + def _request_repo_tree( + self, + path: str, + params: dict[str, Any], + *, + repo_id: str, + authorized: bool, + ) -> Any: + """Fetch one file-tree listing, retrying a spurious denial. + + The server intermittently answers a tree request with ``403 无权访问该数据 + 集`` on a repository the caller has just read successfully -- observed on + both ``PageNumber``-paginated and ``Root``-scoped listings. Enumerating a + large repository takes hundreds of such requests, so at that scale a + single-request failure rate is effectively a guaranteed whole-listing + failure, and it discards every entry gathered so far. + + ``authorized`` means some tree read already succeeded on this client, so + the credential is proven and a denial cannot be an authorization result. + Until then a 403 is taken at face value, keeping a real permission error + fast and honest. + """ + last_error: PermissionDeniedError | None = None + for attempt in range(REPO_TREE_PAGE_MAX_ATTEMPTS): + try: + resp = self._request("GET", path, params=params) + except PermissionDeniedError as error: + if not authorized: + raise + last_error = error + if attempt < REPO_TREE_PAGE_MAX_ATTEMPTS - 1: + wait = min(2**attempt, REPO_TREE_PAGE_RETRY_MAX_DELAY_SECONDS) + logger.warning( + "Repo %s: tree listing (%s) denied on an already-authorized repo, retrying in %ds ...", + repo_id, + params.get("Root") or params.get("PageNumber") or "/", + wait, + ) + time.sleep(wait) + continue + self._tree_reads_ok = True + return resp + raise NetworkError( + f"Repo {repo_id}: tree listing ({params.get('Root') or params.get('PageNumber') or '/'}) kept " + f"returning a denial after {REPO_TREE_PAGE_MAX_ATTEMPTS} attempts on an already-authorized " + f"repo: {last_error}" + ) from last_error + # ------------------------------------------------------------------ # Revisions # ------------------------------------------------------------------ diff --git a/src/modelscope_hub/_openapi.py b/src/modelscope_hub/_openapi.py index 93a71b1..ebdd71f 100644 --- a/src/modelscope_hub/_openapi.py +++ b/src/modelscope_hub/_openapi.py @@ -30,9 +30,17 @@ from urllib.parse import urljoin, urlsplit import requests +from requests.adapters import HTTPAdapter from .config import HubConfig, get_default_config -from .constants import API_CONNECT_TIMEOUT, API_MAX_RETRIES, API_TIMEOUT, OPENAPI_PREFIX, TokenScope +from .constants import ( + API_CONNECT_TIMEOUT, + API_CONNECTION_POOL_MAXSIZE, + API_MAX_RETRIES, + API_TIMEOUT, + OPENAPI_PREFIX, + TokenScope, +) from .errors import ( APIError, AuthenticationError, @@ -200,6 +208,16 @@ def __init__( ) -> None: self._config = config or get_default_config() self._session = session or requests.Session() + if session is None: + # A bare Session caps the pool at urllib3's default of 10, which is + # below the concurrency bulk transfers use; the excess connections + # are discarded and pay for a new TLS handshake on next use. + adapter = HTTPAdapter( + pool_connections=API_CONNECTION_POOL_MAXSIZE, + pool_maxsize=API_CONNECTION_POOL_MAXSIZE, + ) + self._session.mount("https://", adapter) + self._session.mount("http://", adapter) self._timeout: float | tuple[float, float] = ( float(timeout) if timeout is not None else (float(API_CONNECT_TIMEOUT), float(API_TIMEOUT)) ) diff --git a/src/modelscope_hub/_upload.py b/src/modelscope_hub/_upload.py index cb6cd18..0b4f0de 100644 --- a/src/modelscope_hub/_upload.py +++ b/src/modelscope_hub/_upload.py @@ -20,6 +20,7 @@ import io import json import os +import posixpath import tempfile import threading import time @@ -30,6 +31,8 @@ from tqdm.auto import tqdm from .constants import ( + API_CONNECTION_POOL_MAXSIZE, + COMMIT_MAX_ACTIONS_PER_REQUEST, DATASET_LFS_SUFFIX, DEFAULT_IGNORE_PATTERNS, MODEL_LFS_SUFFIX, @@ -44,8 +47,12 @@ UPLOAD_COMMIT_BATCH_MAX_OPERATIONS, UPLOAD_COMMIT_MAX_ATTEMPTS, UPLOAD_COMMIT_MAX_CONSECUTIVE_FAILED_BATCHES, + UPLOAD_COMMIT_MAX_INLINE_BYTES, + UPLOAD_COMMIT_MAX_PER_HOUR, + UPLOAD_COMMIT_MAX_RETRY_AFTER_SECONDS, UPLOAD_COMMIT_RETRY_TOTAL_WAIT_SECONDS, UPLOAD_FAILED_FILE_MAX_RETRY_ROUNDS, + UPLOAD_INLINE_METADATA_PATHS, UPLOAD_LEGACY_PROGRESS_FILE, UPLOAD_LFS_FORCE_THRESHOLD_BYTES, UPLOAD_MAX_CONCURRENT_WORKERS, @@ -53,6 +60,7 @@ UPLOAD_MAX_FILE_SIZE_BYTES, UPLOAD_MAX_FILES_PER_DIRECTORY, UPLOAD_NORMAL_FILES_TOTAL_SIZE_BYTES, + UPLOAD_PROGRESS_MIN_INTERVAL_SECONDS, UPLOAD_RECOVERY_BACKOFF_MAX_EXPONENT, UPLOAD_RECOVERY_ENABLED, UPLOAD_RECOVERY_MAX_DELAY_SECONDS, @@ -64,6 +72,7 @@ HubError, InvalidParameter, NetworkError, + RateLimitError, StorageError, ) from .utils.file_utils import compute_hash @@ -82,6 +91,25 @@ _TRACKER_VERSION = 3 +class _DuplicateBlob: + """Marker: an earlier file in this run uploads this exact content. + + Identical content hashes to one oid, and the batch pre-sign step hands every + occurrence the same upload URL -- so without this marker each occurrence + would PUT the same bytes again. The server only reports "already stored" + once a blob has landed, which cannot help when all the pre-signing happens + before any upload starts. + """ + + __slots__ = () + + def __repr__(self) -> str: # pragma: no cover - diagnostic only + return "" + + +DUPLICATE_BLOB = _DuplicateBlob() + + # ==================================================================== # Helpers # ==================================================================== @@ -119,8 +147,21 @@ def verify_complete(self) -> None: ) +def _is_inline_metadata(path: str | Path) -> bool: + """Return whether *path* must stay inline in the commit rather than go to LFS. + + Checked before the size and suffix rules so that lowering the LFS threshold + can never turn a repository file the Hub parses server-side into a pointer. + """ + if not isinstance(path, (str, Path)): + return False + return Path(path).name.upper() in UPLOAD_INLINE_METADATA_PATHS + + def _is_lfs(path: str | Path, size: int, repo_type: str) -> bool: """Determine if a file should use LFS upload mode (suffix + size threshold).""" + if _is_inline_metadata(path): + return False if size > UPLOAD_LFS_FORCE_THRESHOLD_BYTES: return True suffix = Path(path).suffix.lower() if isinstance(path, (str, Path)) else "" @@ -136,15 +177,88 @@ def _upload_mode(path: str | Path, size: int, repo_type: str) -> str: return "lfs" if _is_lfs(path, size, repo_type) else "normal" -def _calculate_adaptive_batch_size(total_files: int) -> int: - """Calculate optimal commit batch size based on total file count.""" +def _calculate_adaptive_batch_size(total_files: int, max_operations: int) -> int: + """Commit batch size from the file count alone, capped by *max_operations*. + + Fewer, fuller commits are strictly better: the Hub throttles commits per + repository, so commit count -- not commit size -- is the scarce resource. + The only reason to stop growing a batch is the operation cap, or the inlined + content limit that :func:`_plan_commit_batches` applies on top of this. + + The cap is itself clamped to :data:`COMMIT_MAX_ACTIONS_PER_REQUEST`, which + the server rejects outright rather than truncating. + """ if total_files <= 0: return 1 - if total_files <= 100: - return total_files - if total_files <= 10_000: - return max(64, min(256, total_files // 80)) - return 512 + ceiling = max(1, COMMIT_MAX_ACTIONS_PER_REQUEST) + cap = max_operations if max_operations > 0 else total_files + return max(1, min(cap, ceiling, total_files)) + + +def _plan_commit_batches( + files: list[tuple[str, str]], + repo_type: str, + *, + max_operations: int, + max_inline_bytes: int, + sizes: dict[str, int] | None = None, +) -> list[int]: + """Split *files* into commit batches, returning each batch's file count. + + Two limits close a batch, whichever is reached first: ``max_operations`` + files, or ``max_inline_bytes`` of content that will ride *inside* the commit + request. Only non-LFS files contribute to the byte total -- an LFS file adds + a fixed-size pointer -- and base64 expansion is accounted for, because the + request carries the encoded form. A batch always holds at least one file, so + a single oversized inline file still makes progress instead of deadlocking. + """ + if not files: + return [] + cap = max_operations if max_operations > 0 else len(files) + cap = max(1, min(cap, max(1, COMMIT_MAX_ACTIONS_PER_REQUEST))) + + batches: list[int] = [] + count = 0 + inline_bytes = 0 + for path_in_repo, file_path in files: + size = sizes.get(file_path, 0) if sizes is not None else _safe_size(file_path) + encoded = 0 if _is_lfs(path_in_repo, size, repo_type) else (size + 2) // 3 * 4 + if count > 0 and (count >= cap or (max_inline_bytes > 0 and inline_bytes + encoded > max_inline_bytes)): + batches.append(count) + count = 0 + inline_bytes = 0 + count += 1 + inline_bytes += encoded + if count: + batches.append(count) + return batches + + +def _safe_size(file_path: str) -> int: + try: + return os.stat(file_path).st_size + except OSError: + return 0 + + +def _normalize_path_in_repo(path_in_repo: str | None) -> str: + """Collapse a repo destination prefix to a clean, root-relative form. + + ``"."``, ``"./"``, ``""`` and ``"/"`` all denote the repository root and + must yield no prefix. Left literal, a value like ``"."`` becomes a ``"./"`` + prefix on every file and rides into each commit action's ``path``, which the + Hub rejects wholesale as an invalid commit action (E3021). Separators are + normalized and ``.``/``..`` segments resolved; a path that escapes the root + is refused rather than silently rewritten. + """ + if not path_in_repo: + return "" + cleaned = posixpath.normpath(path_in_repo.strip().replace("\\", "/")).strip("/") + if cleaned in ("", "."): + return "" + if cleaned == ".." or cleaned.startswith("../"): + raise InvalidParameter(f"path_in_repo must stay within the repository root, got {path_in_repo!r}") + return cleaned def _compute_file_hash( @@ -482,18 +596,38 @@ def clear(self) -> None: class BatchTracker: - """Thread-safe tracker for pre-assigned upload batches.""" + """Thread-safe tracker for pre-assigned upload batches. + + Batch sizes are supplied as a plan rather than a single number so that a + batch can be closed on inlined-content volume as well as on file count. + """ + + def __init__(self, total_files: int, batch_sizes: list[int] | int) -> None: + if isinstance(batch_sizes, int): + step = max(1, batch_sizes) + sizes = [min(step, total_files - start) for start in range(0, total_files, step)] + else: + sizes = [size for size in batch_sizes if size > 0] + assigned = sum(sizes) + if assigned < total_files: + # Never drop files: a short plan gets the remainder as a final batch. + sizes.append(total_files - assigned) + self._batch_sizes = sizes + self._num_batches = len(sizes) + + # file index -> batch index, so a completed upload can find its batch + # without assuming batches are uniform. + self._owner: list[int] = [] + self._batch_start: list[int] = [] + offset = 0 + for batch_idx, size in enumerate(sizes): + self._batch_start.append(offset) + self._owner.extend([batch_idx] * size) + offset += size - def __init__(self, total_files: int, batch_size: int) -> None: - self._batch_size = batch_size - self._num_batches = (total_files - 1) // batch_size + 1 if total_files > 0 else 0 self._batch_results: list[list[dict]] = [[] for _ in range(self._num_batches)] self._batch_failures: list[list[tuple]] = [[] for _ in range(self._num_batches)] - self._batch_expected: list[int] = [] - for i in range(self._num_batches): - start = i * batch_size - end = min(start + batch_size, total_files) - self._batch_expected.append(end - start) + self._batch_expected: list[int] = list(sizes) self._batch_events: list[threading.Event] = [threading.Event() for _ in range(self._num_batches)] self._lock = threading.Lock() @@ -501,8 +635,13 @@ def __init__(self, total_files: int, batch_size: int) -> None: def num_batches(self) -> int: return self._num_batches + def batch_range(self, batch_idx: int) -> tuple[int, int]: + """Return the ``[start, end)`` file-index range owned by *batch_idx*.""" + start = self._batch_start[batch_idx] + return start, start + self._batch_sizes[batch_idx] + def batch_index(self, file_index: int) -> int: - return file_index // self._batch_size + return self._owner[file_index] def record_success(self, file_index: int, result: dict) -> None: idx = self.batch_index(file_index) @@ -538,6 +677,49 @@ def _is_batch_complete(self, batch_idx: int) -> bool: return count >= self._batch_expected[batch_idx] +class _CommitRateGovernor: + """Sliding-window limiter that keeps commits under a per-hour budget. + + Reacting to a throttle costs a wasted round trip, and when the server holds + the connection open instead of answering, a full read timeout. Pacing ahead + of the limit avoids both. Disabled when the budget is not positive, so it is + inert for interactive uploads that never approach the ceiling. + """ + + _WINDOW_SECONDS = 3600.0 + + def __init__(self, max_per_hour: int) -> None: + self._max_per_hour = max_per_hour + self._timestamps: list[float] = [] + self._lock = threading.Lock() + + @property + def enabled(self) -> bool: + return self._max_per_hour > 0 + + def acquire(self) -> float: + """Block until a commit slot is free; return the seconds spent waiting.""" + if not self.enabled: + return 0.0 + waited = 0.0 + while True: + with self._lock: + now = time.monotonic() + cutoff = now - self._WINDOW_SECONDS + self._timestamps = [ts for ts in self._timestamps if ts > cutoff] + if len(self._timestamps) < self._max_per_hour: + self._timestamps.append(now) + return waited + sleep_for = self._timestamps[0] - cutoff + logger.info( + "Commit budget reached (%d/hour), pausing %.0fs before the next commit ...", + self._max_per_hour, + sleep_for, + ) + time.sleep(max(sleep_for, 0.1)) + waited += max(sleep_for, 0.1) + + # ==================================================================== # Upload Manager # ==================================================================== @@ -558,6 +740,12 @@ def __init__( self._config = config self._openapi = openapi_client self._create_repo_fn = create_repo_fn + # The commit budget is a server-side property of the repository, not of + # one call, so the governor is shared by every commit this manager makes + # -- batch commits, recovery rounds, sync deletes and single files alike. + # Pacing only the happy path would leave recovery free to hammer a + # server that is already throttling. + self._commit_governor = _CommitRateGovernor(UPLOAD_COMMIT_MAX_PER_HOUR) # ------------------------------------------------------------------ # Public: upload_file @@ -579,6 +767,8 @@ def upload_file( if path_or_fileobj is None: raise InvalidParameter("Path or file object cannot be None!") + path_in_repo = _normalize_path_in_repo(path_in_repo) + if isinstance(path_or_fileobj, (str, Path)): path_or_fileobj = os.path.abspath(os.path.expanduser(str(path_or_fileobj))) path_in_repo = path_in_repo or os.path.basename(path_or_fileobj) @@ -627,7 +817,10 @@ def upload_file( ) print(f"Committing file to {repo_id} ...", flush=True) - return self._client.create_commit( + # Same commit path as folder uploads: a single file gets the transient + # retry and the Retry-After handling too. Committing directly meant a + # throttled or briefly unavailable server failed the call outright. + return self._commit_with_retry( repo_id=repo_id, repo_type=repo_type, operations=[operation], @@ -650,25 +843,55 @@ def delete_files( """Delete repository files through a commit operation. The direct repository DELETE endpoints reject API-token authentication. - Commit ``delete`` actions use the same supported write path as uploads - and are applied atomically in a single commit. + Commit ``delete`` actions use the same supported write path as uploads. + + A single commit is capped by the server at + :data:`COMMIT_MAX_ACTIONS_PER_REQUEST` actions, so a larger request is + split across sequential commits. Deletions within one commit are atomic; + across a split they are not, and a failure part-way leaves the earlier + commits applied -- which is reported rather than hidden, so the caller can + retry with the remaining paths. """ paths = list(dict.fromkeys(path for path in file_paths if path)) if not paths: - raise InvalidParameter( - "file_paths must contain at least one non-empty path.") + raise InvalidParameter("file_paths must contain at least one non-empty path.") + + chunk_size = max(1, COMMIT_MAX_ACTIONS_PER_REQUEST) + chunks = [paths[i : i + chunk_size] for i in range(0, len(paths), chunk_size)] + if len(chunks) > 1: + logger.info( + "Deleting %d file(s) in %d commit(s) (server caps one commit at %d actions).", + len(paths), + len(chunks), + chunk_size, + ) + + deleted: list[str] = [] + for index, chunk in enumerate(chunks): + message = commit_message if len(chunks) == 1 else f"{commit_message} ({index + 1}/{len(chunks)})" + try: + self._commit_with_retry( + repo_id=repo_id, + repo_type=repo_type, + operations=self._build_delete_operations(chunk), + commit_message=message, + revision=revision, + ) + except Exception: + if deleted: + logger.error( + "Delete commit %d/%d failed after %d file(s) were already removed.", + index + 1, + len(chunks), + len(deleted), + ) + raise + deleted.extend(chunk) - self._commit_with_retry( - repo_id=repo_id, - repo_type=repo_type, - operations=self._build_delete_operations(paths), - commit_message=commit_message, - revision=revision, - ) return { - "deleted_files": paths, + "deleted_files": deleted, "failed_files": [], - "total_files": len(paths), + "total_files": len(deleted), } # ------------------------------------------------------------------ @@ -690,6 +913,8 @@ def upload_folder( use_cache: bool | None = None, disable_tqdm: bool = False, sync_remote_repo: bool = False, + tracker_path: str | Path | None = None, + progress_callback: Any = None, ) -> dict | list[dict] | None: """Upload a folder with resumable support, adaptive batching, and retry.""" start_time = time.time() @@ -727,12 +952,14 @@ def upload_folder( # Collect files logger.info("Preparing files to upload ...") + file_sizes: dict[str, int] = {} sorted_files = self._prepare_upload_folder( folder_path=folder_path, path_in_repo=path_in_repo, repo_type=repo_type, allow_patterns=allow_patterns, ignore_patterns=ignore_patterns, + sizes_out=file_sizes, ) # For sync mode: collect ALL local files (unfiltered) to avoid @@ -760,27 +987,40 @@ def upload_folder( # Sort for deterministic batch assignment sorted_files = sorted(sorted_files, key=lambda x: x[0]) - # Calculate batch size - if UPLOAD_ADAPTIVE_BATCHING_ENABLED: - commit_batch_size = _calculate_adaptive_batch_size(len(sorted_files)) - logger.info( - "Adaptive batch size: %d (for %d files)", - commit_batch_size, - len(sorted_files), - ) - else: - commit_batch_size = ( - UPLOAD_COMMIT_BATCH_MAX_OPERATIONS if UPLOAD_COMMIT_BATCH_MAX_OPERATIONS > 0 else len(sorted_files) - ) + # Plan commit batches. The operation cap bounds the file count; the + # inline-content cap bounds how many bytes a commit body carries, which + # only non-LFS files add to. + max_operations = ( + _calculate_adaptive_batch_size(len(sorted_files), UPLOAD_COMMIT_BATCH_MAX_OPERATIONS) + if UPLOAD_ADAPTIVE_BATCHING_ENABLED + else (UPLOAD_COMMIT_BATCH_MAX_OPERATIONS if UPLOAD_COMMIT_BATCH_MAX_OPERATIONS > 0 else len(sorted_files)) + ) + batch_plan = _plan_commit_batches( + sorted_files, + repo_type, + max_operations=max_operations, + max_inline_bytes=UPLOAD_COMMIT_MAX_INLINE_BYTES, + sizes=file_sizes, + ) + logger.info( + "Commit plan: %d batch(es) for %d file(s) (max %d ops, max %d inline bytes per commit).", + len(batch_plan), + len(sorted_files), + max_operations, + UPLOAD_COMMIT_MAX_INLINE_BYTES, + ) - # Initialize tracker + # Initialize tracker. The cache normally lives in the uploaded folder, + # but a caller that stages files into a throwaway tree (a link tree, a + # per-chunk directory) must be able to keep it outside, or every run + # rediscovers hashes and re-commits what was already committed. folder_path_resolved = Path(folder_path).resolve() if use_cache: - cache_path = folder_path_resolved / UPLOAD_CACHE_FILE + cache_path = Path(tracker_path).expanduser() if tracker_path else folder_path_resolved / UPLOAD_CACHE_FILE tracker: UploadTracker | NullTracker = UploadTracker(cache_path, repo_id=repo_id) else: tracker = NullTracker() - batch_tracker = BatchTracker(len(sorted_files), commit_batch_size) + batch_tracker = BatchTracker(len(sorted_files), batch_plan) # Skip individually committed files files_to_upload: list[tuple[int, tuple[str, str]]] = [] @@ -800,38 +1040,99 @@ def upload_folder( ) files_to_upload.append((file_idx, (file_path_in_repo, file_path))) - # Batch pre-validation for LFS files with cached hashes + # Batch pre-validation for every LFS candidate. + # + # Without a cached hash the per-file upload path would ask the git-lfs + # batch endpoint for its own pre-signed URL, one round trip per file. + # Hashing up front lets all candidates be pre-signed in groups instead, + # which is the difference between one request per file and one per group + # once small files are routed to LFS. pre_validated_map: dict[str, str | None] = {} - lfs_hash_info_map: dict[int, tuple[dict, os.stat_result]] = {} + lfs_hash_info_map: dict[int, dict] = {} for file_idx, (file_path_in_repo, file_path) in files_to_upload: + size = file_sizes.get(file_path, 0) + if _upload_mode(file_path_in_repo, size, repo_type) != "lfs": + continue try: st = os.stat(file_path) - cached = tracker.get_hash(file_path_in_repo, st.st_mtime, st.st_size) - if cached is not None: - if ( - _upload_mode( - file_path_in_repo, - cached["file_size"], - repo_type, - ) - == "lfs" - ): - lfs_hash_info_map[file_idx] = (cached, st) - continue except OSError: - pass + continue + cached = tracker.get_hash(file_path_in_repo, st.st_mtime, st.st_size) + if cached is None: + continue + if _upload_mode(file_path_in_repo, cached["file_size"], repo_type) == "lfs": + lfs_hash_info_map[file_idx] = cached + + uncached_lfs = [ + (file_idx, file_info) + for file_idx, file_info in files_to_upload + if file_idx not in lfs_hash_info_map + and _upload_mode(file_info[0], file_sizes.get(file_info[1], 0), repo_type) == "lfs" + ] + if uncached_lfs: + lfs_hash_info_map.update(self._hash_files_parallel(uncached_lfs, tracker, max_workers)) if lfs_hash_info_map: - objects = [{"oid": info["file_hash"], "size": info["file_size"]} for info, _ in lfs_hash_info_map.values()] - validated = self._validate_blobs_batch(repo_id=repo_id, repo_type=repo_type, objects=objects) + objects = [{"oid": info["file_hash"], "size": info["file_size"]} for info in lfs_hash_info_map.values()] + # Identical content shares one oid, so the request set is keyed by + # digest rather than by file. + unique_objects = list({obj["oid"]: obj for obj in objects}.values()) + validated = self._validate_blobs_batch( + repo_id=repo_id, + repo_type=repo_type, + objects=unique_objects, + max_workers=max_workers, + ) pre_validated_map = validated reused = sum(1 for v in validated.values() if v is None) + unresolved = len(unique_objects) - len(validated) logger.info( - "Pre-validated %d cached LFS hash(es): %d globally existing, %d need upload.", + "Pre-validated %d/%d distinct blob(s) for %d file(s) in %d request(s): " + "%d already stored, %d to upload%s.", + len(validated), + len(unique_objects), len(objects), + -(-len(unique_objects) // max(1, UPLOAD_BLOB_VALIDATION_BATCH_MAX_OBJECTS)), reused, - len(objects) - reused, + len(validated) - reused, + f", {unresolved} unresolved (will negotiate per file)" if unresolved else "", + ) + + # Elect one owner per distinct blob. + # + # `files_to_upload` is in ascending file index, and a batch owns a + # contiguous ascending index range, so the first file holding a given oid + # always lands in a batch no later than any of its duplicates. Batches are + # committed in order and each waits for its own files, so by the time a + # duplicate's batch commits, its owner has already finished -- which is + # what lets the duplicates skip the transfer with no locking and no risk + # of a worker pool deadlocking on itself. + blob_owner: dict[str, int] = {} + blob_ready: dict[str, bool] = {} + blob_state_lock = threading.Lock() + deduped_files = 0 + deduped_bytes = 0 + for file_idx, _file_info in files_to_upload: + info = lfs_hash_info_map.get(file_idx) + if info is None: + continue + oid = info["file_hash"] + if pre_validated_map.get(oid, "") is None: + # Already stored server-side; nobody needs to transfer it. + blob_ready[oid] = True + continue + if oid not in blob_owner: + blob_owner[oid] = file_idx + else: + deduped_files += 1 + deduped_bytes += info["file_size"] + if deduped_files: + logger.info( + "Deduplicated %d file(s) sharing content with an earlier file: %d byte(s) that " + "would otherwise be uploaded twice.", + deduped_files, + deduped_bytes, ) skipped_count = len(skipped_indices) @@ -846,15 +1147,18 @@ def upload_folder( ) logger.info( - "Uploading %d file(s) in %d batch(es) of size %d (pipeline mode).", + "Uploading %d file(s) in %d batch(es) (pipeline mode).", len(files_to_upload), batch_tracker.num_batches, - commit_batch_size, ) # Pipeline: upload workers def _upload_worker(file_idx: int, file_info: tuple, pre_validated: Any = None) -> None: path_in_repo_w, file_path_w = file_info + owned_oid: str | None = None + info = lfs_hash_info_map.get(file_idx) + if info is not None and blob_owner.get(info["file_hash"]) == file_idx: + owned_oid = info["file_hash"] try: logger.debug("Uploading: %s ...", path_in_repo_w) result = self._upload_single_file( @@ -864,12 +1168,22 @@ def _upload_worker(file_idx: int, file_info: tuple, pre_validated: Any = None) - repo_type=repo_type, tracker=tracker, pre_validated=pre_validated, + hash_info=lfs_hash_info_map.get(file_idx), disable_tqdm=disable_tqdm, ) logger.debug("Uploaded: %s", path_in_repo_w) + # Publish the blob outcome before the batch is marked complete: + # the consumer reads it as soon as the batch event fires. + if owned_oid is not None: + with blob_state_lock: + blob_ready[owned_oid] = True batch_tracker.record_success(file_idx, result) + _report_wire(result) except Exception as e: logger.error("Upload failed: %s - %s", path_in_repo_w, e) + if owned_oid is not None: + with blob_state_lock: + blob_ready[owned_oid] = False batch_tracker.record_failure(file_idx, file_info, e) # Pipeline: consume batches in order @@ -877,16 +1191,151 @@ def _upload_worker(file_idx: int, file_info: tuple, pre_validated: Any = None) - all_results: list[dict] = [] total_failed_files: list[tuple] = [] num_batches = batch_tracker.num_batches + committed_files = 0 + committed_bytes = 0 + total_bytes = sum(file_sizes.values()) + # Wire-level counters, advanced as each file's transfer finishes rather + # than when its commit lands. + wire_lock = threading.Lock() + wire_state = {"bytes": 0, "files": 0, "reported_bytes": 0, "last_emit": 0.0} + + def _emit(payload: dict) -> None: + if progress_callback is None: + return + try: + progress_callback(payload) + except Exception as cb_error: # noqa: BLE001 - a reporter must not fail the upload + logger.warning("Progress callback raised %s, continuing upload.", cb_error) + + def _report( + event: str, + batch_idx: int, + files: int, + num_bytes: int, + inline_bytes: int = 0, + error: str | None = None, + ) -> None: + """Emit a batch-level progress event. + + A folder upload is otherwise silent between batches, so a long run is + indistinguishable from a hung one. Byte counts are included because a + consumer that only learns file counts cannot compute a throughput rate + or an ETA, which is most of what progress is for. ``inline_bytes`` is + the part of the batch that travels inside this commit rather than + having already gone to object storage, so a consumer can attribute + wire traffic to the right moment without double counting. + """ + with wire_lock: + wire_bytes, wire_files = wire_state["bytes"], wire_state["files"] + _emit( + { + "event": event, + "repo_id": repo_id, + "batch_index": batch_idx, + "num_batches": num_batches, + "batch_files": files, + "batch_bytes": num_bytes, + "batch_inline_bytes": inline_bytes, + "committed_files": committed_files, + "committed_bytes": committed_bytes, + "uploaded_bytes": wire_bytes, + "uploaded_files": wire_files, + "total_files": len(sorted_files), + "total_bytes": total_bytes, + "skipped_files": skipped_count, + "elapsed": time.time() - start_time, + "error": error, + } + ) + + def _report_wire(result: dict) -> None: + """Account one finished file transfer, emitting at a throttled rate. + + Commits land in lumps tens of seconds apart, so a consumer fed only by + commit events sees a rate that alternates between a spike and zero and + cannot tell a slow batch from a hung one. Blob uploads finish + continuously, which is the signal a rate should be built from. Only + bytes that really went to object storage count: a deduplicated blob + transfers nothing. + """ + if progress_callback is None: + return + moved = result["file_size_on_disk"] if result.get("is_blob_uploaded") else 0 + now = time.monotonic() + with wire_lock: + wire_state["bytes"] += moved + wire_state["files"] += 1 + due = now - wire_state["last_emit"] >= UPLOAD_PROGRESS_MIN_INTERVAL_SECONDS + if not due: + return + wire_state["last_emit"] = now + delta = wire_state["bytes"] - wire_state["reported_bytes"] + wire_state["reported_bytes"] = wire_state["bytes"] + snapshot = (wire_state["bytes"], wire_state["files"]) + _emit( + { + "event": "upload_progress", + "repo_id": repo_id, + "uploaded_bytes": snapshot[0], + "uploaded_bytes_delta": delta, + "uploaded_files": snapshot[1], + "committed_files": committed_files, + "committed_bytes": committed_bytes, + "total_files": len(sorted_files), + "total_bytes": total_bytes, + "skipped_files": skipped_count, + "elapsed": time.time() - start_time, + "error": None, + } + ) + + def _flush_wire() -> None: + """Emit whatever wire bytes the throttle has not reported yet.""" + if progress_callback is None: + return + with wire_lock: + delta = wire_state["bytes"] - wire_state["reported_bytes"] + if delta <= 0: + return + wire_state["reported_bytes"] = wire_state["bytes"] + snapshot = (wire_state["bytes"], wire_state["files"]) + _emit( + { + "event": "upload_progress", + "repo_id": repo_id, + "uploaded_bytes": snapshot[0], + "uploaded_bytes_delta": delta, + "uploaded_files": snapshot[1], + "committed_files": committed_files, + "committed_bytes": committed_bytes, + "total_files": len(sorted_files), + "total_bytes": total_bytes, + "skipped_files": skipped_count, + "elapsed": time.time() - start_time, + "error": None, + } + ) try: with ThreadPoolExecutor(max_workers=max_workers) as executor: for file_idx, file_info in files_to_upload: - pv: str | bool | None = None + pv: str | bool | _DuplicateBlob | None = None if file_idx in lfs_hash_info_map: - cached_hash = lfs_hash_info_map[file_idx][0]["file_hash"] - pv = pre_validated_map.get(cached_hash) - if pv is None: - pv = True + cached_hash = lfs_hash_info_map[file_idx]["file_hash"] + # "Answered, and the server did not ask for an upload" + # means the blob already exists and is reused. "Never + # answered" means the group failed, and the file has to + # negotiate its own URL -- treating that as reuse would + # skip the transfer and commit a pointer to a blob that + # was never stored. + if cached_hash in pre_validated_map: + url = pre_validated_map[cached_hash] + if url is None: + pv = True + elif blob_owner.get(cached_hash) == file_idx: + pv = url + else: + pv = DUPLICATE_BLOB executor.submit(_upload_worker, file_idx, file_info, pv) consecutive_failures = 0 @@ -896,8 +1345,7 @@ def _upload_worker(file_idx: int, file_info: tuple, pre_validated: Any = None) - total=num_batches, disable=disable_tqdm, ): - batch_start = batch_idx * commit_batch_size - batch_end = min(batch_start + commit_batch_size, len(sorted_files)) + batch_start, batch_end = batch_tracker.batch_range(batch_idx) if all(i in skipped_indices for i in range(batch_start, batch_end)): logger.info( "Batch %d/%d fully committed, skipping.", @@ -907,12 +1355,53 @@ def _upload_worker(file_idx: int, file_info: tuple, pre_validated: Any = None) - continue results, failures = batch_tracker.wait_for_batch(batch_idx) + # Every file of this batch has finished its transfer by now, + # so publish the bytes the throttle may still be holding + # before the commit event reports the batch as done. + _flush_wire() if failures: total_failed_files.extend(failures) for item, err in failures: logger.error(" Failed: %s - %s", item[0], err) + # A file that skipped its transfer because a duplicate owned + # it must not be committed if that owner's upload failed: + # the commit would reference a blob that was never stored. + # Its owner is in this batch or an earlier one, both already + # resolved, so the outcome is known here. + orphaned: list[dict] = [] + if blob_owner: + with blob_state_lock: + ready_snapshot = dict(blob_ready) + committable = [] + for item_r in results: + oid_r = item_r["file_hash_info"]["file_hash"] + if ( + item_r.get("upload_mode") == "lfs" + and oid_r in blob_owner + and not ready_snapshot.get(oid_r, False) + ): + orphaned.append(item_r) + continue + committable.append(item_r) + if orphaned: + logger.warning( + "Batch %d/%d: %d file(s) deferred, the upload of the content they " + "share failed; they will be retried on their own.", + batch_idx + 1, + num_batches, + len(orphaned), + ) + total_failed_files.extend( + ( + (item_r["file_path_in_repo"], item_r["file_path"]), + StorageError("shared blob upload failed"), + ) + for item_r in orphaned + ) + results = committable + self._track_uploaded_batch(tracker, results) operations = self._build_batch_operations(results, repo_type) @@ -922,6 +1411,7 @@ def _upload_worker(file_idx: int, file_info: tuple, pre_validated: Any = None) - batch_idx + 1, num_batches, ) + _report("batch_failed", batch_idx, len(failures), 0, error="all files failed to upload") continue batch_commit_message = f"{commit_message} (batch {batch_idx + 1}/{num_batches})" @@ -943,6 +1433,13 @@ def _upload_worker(file_idx: int, file_info: tuple, pre_validated: Any = None) - ) self._track_committed_batch(tracker, results) consecutive_failures = 0 + batch_bytes = sum(r["file_size_on_disk"] for r in results) + batch_inline_bytes = sum( + r["file_size_on_disk"] for r in results if r.get("upload_mode") != "lfs" + ) + committed_files += len(results) + committed_bytes += batch_bytes + _report("batch_committed", batch_idx, len(results), batch_bytes, batch_inline_bytes) except Exception as e: logger.error( "Batch %d/%d commit failed: %s", @@ -950,6 +1447,13 @@ def _upload_worker(file_idx: int, file_info: tuple, pre_validated: Any = None) - num_batches, e, ) + _report( + "batch_failed", + batch_idx, + len(results), + sum(r["file_size_on_disk"] for r in results), + error=str(e), + ) category = classify_error(e) if not _ErrorCategory.is_retryable(category): for r in results: @@ -993,6 +1497,39 @@ def _upload_worker(file_idx: int, file_info: tuple, pre_validated: Any = None) - tracker.save() # ReAct progressive retry fallback + # + # Recovery has to report progress too. It is exactly the moment an + # operator is watching, and a run whose recovered volume never reaches + # the metrics under-reports by however much it rescued: an 8 GiB run that + # lost one 512-file batch to a rejected commit finished with 40000 files + # on the Hub but 91 MB missing from done_bytes. + def _report_recovery(results: list[dict], label: str) -> None: + nonlocal committed_files, committed_bytes + recovered_bytes = sum(r["file_size_on_disk"] for r in results) + inline_bytes = sum(r["file_size_on_disk"] for r in results if r.get("upload_mode") != "lfs") + committed_files += len(results) + committed_bytes += recovered_bytes + _flush_wire() + _emit( + { + "event": "recovery_committed", + "repo_id": repo_id, + "stage": label, + "batch_index": -1, + "num_batches": num_batches, + "batch_files": len(results), + "batch_bytes": recovered_bytes, + "batch_inline_bytes": inline_bytes, + "committed_files": committed_files, + "committed_bytes": committed_bytes, + "total_files": len(sorted_files), + "total_bytes": total_bytes, + "skipped_files": skipped_count, + "elapsed": time.time() - start_time, + "error": None, + } + ) + if total_failed_files and UPLOAD_RECOVERY_ENABLED: total_failed_files, react_commits, react_results = self._retry_failed_files_react( failed_files=total_failed_files, @@ -1003,6 +1540,8 @@ def _upload_worker(file_idx: int, file_info: tuple, pre_validated: Any = None) - revision=revision, max_workers=max_workers, disable_tqdm=disable_tqdm, + on_uploaded=_report_wire, + on_committed=_report_recovery, ) commit_infos.extend(react_commits) all_results.extend(react_results) @@ -1018,6 +1557,8 @@ def _upload_worker(file_idx: int, file_info: tuple, pre_validated: Any = None) - commit_infos=commit_infos, all_results=all_results, disable_tqdm=disable_tqdm, + on_uploaded=_report_wire, + on_committed=_report_recovery, ) tracker.save() @@ -1170,6 +1711,7 @@ def _upload_single_file( repo_type: str, tracker: UploadTracker | NullTracker | None = None, pre_validated: Any = None, + hash_info: dict | None = None, disable_tqdm: bool = False, ) -> dict: if tracker is None: @@ -1178,7 +1720,13 @@ def _upload_single_file( file_stat = None is_real_path = isinstance(file_path, (str, os.PathLike)) - if is_real_path: + if hash_info is not None: + # Already computed during batch pre-validation; re-reading the file + # to hash it again would double the disk cost of every LFS file. + hash_info_d = dict(hash_info) + hash_info_d["file_path_or_obj"] = file_path + + if hash_info_d is None and is_real_path: try: file_stat = os.stat(file_path) cached = tracker.get_hash(file_path_in_repo, file_stat.st_mtime, file_stat.st_size) @@ -1319,6 +1867,15 @@ def _upload_blob( res_d["is_reused"] = True return res_d + if pre_validated is DUPLICATE_BLOB: + # An earlier file in this run owns the transfer for this content. + # Batch ordering guarantees it has finished before any commit that + # references this file, so nothing has to be waited on here. + logger.debug("Blob %s is uploaded by an earlier duplicate, skipping transfer.", sha256[:8]) + res_d["is_uploaded"] = True + res_d["is_reused"] = True + return res_d + if isinstance(pre_validated, str): upload_url: str = pre_validated else: @@ -1366,24 +1923,97 @@ def _upload_blob( # ------------------------------------------------------------------ # Internal: batch blob validation # ------------------------------------------------------------------ + def _hash_files_parallel( + self, + files: list[tuple[int, tuple[str, str]]], + tracker: UploadTracker | NullTracker, + max_workers: int, + ) -> dict[int, dict]: + """Hash *files* concurrently and record the results in *tracker*. + + Hashing ahead of the upload pipeline is what makes group pre-signing + possible: the git-lfs batch endpoint is keyed by ``sha256``, so without + the digests up front each file has to negotiate its own upload URL. + """ + hashed: dict[int, dict] = {} + + def _hash_one(entry: tuple[int, tuple[str, str]]) -> tuple[int, dict] | None: + file_idx, (path_in_repo, file_path) = entry + try: + st = os.stat(file_path) + info = _compute_file_hash(file_path_or_obj=file_path) + except OSError as error: + # Leave it to the upload worker, which reports per-file failures. + logger.debug("Cannot pre-hash %s: %s", path_in_repo, error) + return None + tracker.put_hash(path_in_repo, st.st_mtime, st.st_size, info) + return file_idx, info + + with ThreadPoolExecutor(max_workers=max(1, max_workers)) as executor: + for outcome in executor.map(_hash_one, files): + if outcome is not None: + hashed[outcome[0]] = outcome[1] + + logger.info("Hashed %d LFS candidate(s) for batch pre-validation.", len(hashed)) + return hashed + def _validate_blobs_batch( self, repo_id: str, repo_type: str, objects: list[dict], + max_workers: int = 1, ) -> dict[str, str | None]: - result: dict[str, str | None] = {} - batch_size = UPLOAD_BLOB_VALIDATION_BATCH_MAX_OBJECTS + """Pre-sign every object, in parallel groups, tolerating group failures. + + The groups are independent requests, so running them serially made the + pre-sign phase scale linearly with the file count -- measurably so: 150 + groups took 25s of pure round-trip latency before a single byte moved. + + The batch endpoint answers only about objects that *need* uploading; an + object it does not mention already exists server-side. That is normalised + here into an explicit ``oid -> None`` entry, so the returned map covers + every object of every group that answered. + + A group that fails is not fatal: its oids are simply absent from the map + and the per-file upload path negotiates their URL itself. Keeping + "answered: already exists" and "never answered" distinguishable is what + makes that safe -- conflating them would skip the transfer and commit a + pointer to a blob that was never stored. + """ + batch_size = max(1, UPLOAD_BLOB_VALIDATION_BATCH_MAX_OBJECTS) + chunks = [objects[i : i + batch_size] for i in range(0, len(objects), batch_size)] + if not chunks: + return {} - for i in range(0, len(objects), batch_size): - chunk = objects[i : i + batch_size] - validated = self._client.validate_blobs( - repo_id=repo_id, - repo_type=repo_type, - objects=chunk, - ) - result.update(validated) + def validate(chunk: list[dict]) -> dict[str, str | None]: + try: + validated = self._client.validate_blobs( + repo_id=repo_id, + repo_type=repo_type, + objects=chunk, + ) + except Exception as error: # noqa: BLE001 - degrades to per-file negotiation + logger.warning( + "Blob pre-validation failed for %d object(s) (%s); those files will negotiate " + "their upload URL individually.", + len(chunk), + error, + ) + return {} + answered: dict[str, str | None] = {obj["oid"]: None for obj in chunk} + answered.update(validated) + return answered + result: dict[str, str | None] = {} + if len(chunks) == 1: + return validate(chunks[0]) + # Bounded by the HTTP pool: more in-flight requests than pooled + # connections just trades latency for TLS handshakes. + workers = max(1, min(max_workers, API_CONNECTION_POOL_MAXSIZE, len(chunks))) + with ThreadPoolExecutor(max_workers=workers) as executor: + for validated in executor.map(validate, chunks): + result.update(validated) return result # ------------------------------------------------------------------ @@ -1401,7 +2031,9 @@ def _commit_with_retry( ) -> dict: last_error: Exception | None = None start_time = time.monotonic() + throttled_wait = 0.0 for attempt in range(max_attempts): + self._commit_governor.acquire() try: return self._client.create_commit( repo_id=repo_id, @@ -1422,8 +2054,32 @@ def _commit_with_retry( except Exception as e: last_error = e + # A throttled commit carries the wait the server wants; honoring it + # beats guessing, and it is budgeted apart from the transient-error + # allowance because its duration is known and can legitimately + # exceed it. + retry_after = getattr(last_error, "retry_after", None) if isinstance(last_error, RateLimitError) else None + if retry_after is not None: + wait = float(retry_after) + if wait > UPLOAD_COMMIT_MAX_RETRY_AFTER_SECONDS: + logger.error( + "Commit throttled with Retry-After=%.0fs, above the %ds ceiling; aborting retries.", + wait, + UPLOAD_COMMIT_MAX_RETRY_AFTER_SECONDS, + ) + break + throttled_wait += wait + logger.warning( + "Commit attempt %d/%d throttled, honoring Retry-After=%.0fs ...", + attempt + 1, + max_attempts, + wait, + ) + time.sleep(wait) + continue + wait = min(2**attempt, 60) - elapsed = time.monotonic() - start_time + elapsed = time.monotonic() - start_time - throttled_wait if elapsed + wait > UPLOAD_COMMIT_RETRY_TOTAL_WAIT_SECONDS: logger.error( "Commit total wait time would exceed %ds (already %.1fs elapsed), aborting retries.", @@ -1536,6 +2192,7 @@ def _prepare_upload_folder( repo_type: str = "model", allow_patterns: list[str] | None = None, ignore_patterns: list[str] | None = None, + sizes_out: dict[str, int] | None = None, ) -> list[tuple[str, str]]: folder = Path(folder_path).expanduser().resolve() if not folder.is_dir(): @@ -1558,7 +2215,9 @@ def _prepare_upload_folder( f"max allowed per directory: {UPLOAD_MAX_FILES_PER_DIRECTORY}" ) - # File size checks + # File size checks. Sizes are handed back through ``sizes_out`` because + # batch planning needs them next; re-stating a large tree costs one + # syscall per file for no new information. total_size = 0 normal_size = 0 for path in all_files: @@ -1571,6 +2230,8 @@ def _prepare_upload_folder( total_size += fsize if not _is_lfs(str(path), fsize, repo_type): normal_size += fsize + if sizes_out is not None: + sizes_out[str(path)] = fsize if normal_size > UPLOAD_NORMAL_FILES_TOTAL_SIZE_BYTES: logger.warning( @@ -1588,7 +2249,13 @@ def _prepare_upload_folder( ignore_patterns=ignore_patterns, ) - prefix = f"{path_in_repo.strip('/')}/" if path_in_repo else "" + # ``path_in_repo`` is a destination prefix, and "." / "./" / "" / "/" + # all mean the repo root. Collapsing them is not cosmetic: a literal + # value like "." otherwise rides into every commit action's ``path`` as + # a "./" prefix, which the Hub rejects wholesale with E3021 "invalid + # commit action". + norm_prefix = _normalize_path_in_repo(path_in_repo) + prefix = f"{norm_prefix}/" if norm_prefix else "" prepared = [(prefix + relpath, relpath_to_abspath[relpath]) for relpath in filtered_keys] logger.info("Prepared %d files for upload.", len(prepared)) @@ -1607,6 +2274,8 @@ def _retry_failed_files_react( revision: str, max_workers: int, disable_tqdm: bool = False, + on_uploaded: Any = None, + on_committed: Any = None, ) -> tuple[list[tuple], list[dict], list[dict]]: commit_infos: list[dict] = [] all_successes: list[dict] = [] @@ -1697,6 +2366,8 @@ def _retry_failed_files_react( try: result = future.result() round_successes.append(result) + if on_uploaded is not None: + on_uploaded(result) except Exception as e: round_failures.append(((path_in_repo_r, file_path_r), e)) else: @@ -1724,6 +2395,8 @@ def _retry_failed_files_react( disable_tqdm=disable_tqdm, ) round_successes.append(result) + if on_uploaded is not None: + on_uploaded(result) except Exception as e: logger.error( "[ReAct] %s: failed %s - %s", @@ -1753,6 +2426,8 @@ def _retry_failed_files_react( ) commit_infos.append(commit_info) self._track_committed_batch(tracker, batch) + if on_committed is not None: + on_committed(batch, round_name) logger.info( "[ReAct] %s: committed %d file(s).", round_name, @@ -1852,6 +2527,8 @@ def _retry_failed_simple( commit_infos: list[dict], all_results: list[dict], disable_tqdm: bool = False, + on_uploaded: Any = None, + on_committed: Any = None, ) -> list[tuple]: total_failed_files = list(failed_files) for retry_round in range(UPLOAD_FAILED_FILE_MAX_RETRY_ROUNDS): @@ -1876,6 +2553,8 @@ def _retry_failed_simple( disable_tqdm=disable_tqdm, ) retry_successes.append(result) + if on_uploaded is not None: + on_uploaded(result) except Exception as e: logger.error(" Retry failed: %s - %s", path_in_repo_r, e) retry_failures.append(((path_in_repo_r, file_path_r), e)) @@ -1894,6 +2573,8 @@ def _retry_failed_simple( commit_infos.append(commit_info) all_results.extend(retry_successes) self._track_committed_batch(tracker, retry_successes) + if on_committed is not None: + on_committed(retry_successes, f"retry round {retry_round + 1}") logger.info( " Retry round %d: committed %d file(s).", retry_round + 1, diff --git a/src/modelscope_hub/agent/_api.py b/src/modelscope_hub/agent/_api.py index 531bf41..86c2d0f 100644 --- a/src/modelscope_hub/agent/_api.py +++ b/src/modelscope_hub/agent/_api.py @@ -241,7 +241,7 @@ def _decode_dolphin_list_response(response: object, *, list_url: str) -> object: from an owner with no repositories. """ try: - payload = response.json() # type: ignore[union-attr] + payload = response.json() # type: ignore[attr-defined] except (AttributeError, ValueError) as exc: raise APIError( "Agent repository list endpoint returned a non-JSON response.", diff --git a/src/modelscope_hub/api.py b/src/modelscope_hub/api.py index 0244ea5..2718b29 100644 --- a/src/modelscope_hub/api.py +++ b/src/modelscope_hub/api.py @@ -25,7 +25,7 @@ import fnmatch import time -from collections.abc import Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping from pathlib import Path from typing import Any, BinaryIO, TypeAlias from urllib.parse import urlparse @@ -1368,6 +1368,8 @@ def upload_folder( use_cache: bool | None = None, disable_tqdm: bool = False, sync_remote_repo: bool = False, + tracker_path: str | Path | None = None, + progress_callback: Callable[[dict], None] | None = None, ) -> dict | list[dict] | None: """Upload an entire folder to a repository with resumable support. @@ -1406,6 +1408,18 @@ def upload_folder( sync_remote_repo : bool, optional If True, delete remote files that are not present locally after a successful upload (sync semantics). Default False. + tracker_path : str or Path, optional + Where to keep the resumable-upload cache. Defaults to + ``.ms_upload_cache`` inside ``folder_path``. Point this outside the + uploaded tree when ``folder_path`` is a staging directory that gets + discarded between runs, so resume state survives. + progress_callback : callable, optional + Called with one dict per batch outcome, carrying ``event`` + (``"batch_committed"`` or ``"batch_failed"``), ``batch_index``, + ``num_batches``, ``batch_files``, ``batch_bytes``, + ``committed_files``, ``committed_bytes``, ``total_files``, + ``total_bytes``, ``skipped_files``, ``elapsed`` and ``error``. + Exceptions raised by the callback are logged and swallowed. Returns ------- @@ -1441,6 +1455,8 @@ def upload_folder( use_cache=use_cache, disable_tqdm=disable_tqdm, sync_remote_repo=sync_remote_repo, + tracker_path=tracker_path, + progress_callback=progress_callback, ) def download_file( diff --git a/src/modelscope_hub/constants.py b/src/modelscope_hub/constants.py index 1a9f970..4d53fb6 100644 --- a/src/modelscope_hub/constants.py +++ b/src/modelscope_hub/constants.py @@ -227,6 +227,7 @@ def _warn_deprecated_env( name: str, *, expects_mb: bool = False, + expects_bytes: bool = False, stacklevel: int = 3, ) -> None: """Warn that a legacy environment variable remains temporarily supported.""" @@ -235,9 +236,19 @@ def _warn_deprecated_env( ) if expects_mb: message += f" {name!r} expects a value in MB." + if expects_bytes: + message += f" {name!r} expects bytes, or a value with a unit suffix such as '32KiB'." warnings.warn(message, FutureWarning, stacklevel=stacklevel) +def _format_bytes(num_bytes: int) -> str: + """Render a byte count using the largest binary unit that stays exact.""" + for unit, scale in (("GiB", 1024**3), ("MiB", 1024**2), ("KiB", 1024)): + if num_bytes and num_bytes % scale == 0: + return f"{num_bytes // scale}{unit}" + return f"{num_bytes}" + + def _env(name: str, *deprecated_names: str) -> str | None: """Read an env var, falling back to deprecated names with a warning.""" value = os.environ.get(name) @@ -251,12 +262,71 @@ def _env(name: str, *deprecated_names: str) -> str | None: return None +def _warn_invalid_env(name: str, raw: str, reason: str, fallback: object) -> None: + """Warn that an environment value was rejected, naming the value used instead. + + Silently falling back to the default made misconfiguration invisible: a + pipeline could export a tuning value, observe none of its effect, and have + no signal to look at. Every rejected value now says so. + """ + warnings.warn( + f"Environment variable {name}={raw!r} is invalid ({reason}); using {fallback!r} instead.", + UserWarning, + stacklevel=3, + ) + + +_BYTE_UNITS: dict[str, int] = { + "": 1, + "B": 1, + "K": 1024, + "KB": 1000, + "KIB": 1024, + "M": 1024**2, + "MB": 1000**2, + "MIB": 1024**2, + "G": 1024**3, + "GB": 1000**3, + "GIB": 1024**3, +} + + +def _parse_byte_size(raw: str, *, bare_unit: int = 1) -> int: + """Parse a size string into bytes, accepting an optional unit suffix. + + ``bare_unit`` scales a value given without a suffix, which is what lets a + deprecated ``*_MB`` alias keep its megabyte meaning while the canonical + name treats a bare number as bytes. Raises :class:`ValueError` so callers + decide between warning and propagating. + """ + text = raw.strip() + if not text: + raise ValueError("empty value") + digits = text + suffix = "" + while digits and not (digits[-1].isdigit() or digits[-1] == "."): + suffix = digits[-1] + suffix + digits = digits[:-1] + digits = digits.strip() + suffix = suffix.strip().upper() + if not digits: + raise ValueError("no numeric part") + if suffix not in _BYTE_UNITS: + raise ValueError(f"unknown size unit {suffix!r}") + number = float(digits) + if number != int(number): + raise ValueError("fractional byte counts are not supported") + unit = bare_unit if suffix == "" else _BYTE_UNITS[suffix] + return int(number) * unit + + def _env_int( name: str, default: int, description: str = "", category: str = "", *deprecated_names: str, + allow_zero: bool = False, ) -> int: """Read a positive integer from the environment and register it.""" all_deprecated = deprecated_names or _DEPRECATED_LOOKUP.get(name, ()) @@ -268,8 +338,13 @@ def _env_int( try: value = int(raw) except ValueError: + _warn_invalid_env(name, raw, "not an integer", default) + return default + if value < 0 or (value == 0 and not allow_zero): + reason = "must not be negative" if allow_zero else "must be a positive integer" + _warn_invalid_env(name, raw, reason, default) return default - return value if value > 0 else default + return value def _env_int_mb( @@ -294,8 +369,12 @@ def _env_int_mb( try: value = int(raw) except ValueError: + _warn_invalid_env(name, raw, "not an integer", f"{default_mb} MB") + return default_mb * 1024 * 1024 + if value <= 0: + _warn_invalid_env(name, raw, "must be a positive integer", f"{default_mb} MB") return default_mb * 1024 * 1024 - return value * 1024 * 1024 if value > 0 else default_mb * 1024 * 1024 + return value * 1024 * 1024 # Fall back to deprecated names (value already in bytes) for old in all_deprecated: raw = os.environ.get(old) @@ -304,8 +383,12 @@ def _env_int_mb( try: value = int(raw) except ValueError: + _warn_invalid_env(old, raw, "not an integer", f"{default_mb} MB") return default_mb * 1024 * 1024 - return value if value > 0 else default_mb * 1024 * 1024 + if value <= 0: + _warn_invalid_env(old, raw, "must be a positive integer", f"{default_mb} MB") + return default_mb * 1024 * 1024 + return value return default_mb * 1024 * 1024 @@ -328,8 +411,12 @@ def _env_int_mb_with_deprecated_units( try: value = int(raw) except ValueError: + _warn_invalid_env(name, raw, "not an integer", f"{default_mb} MB") + return default_bytes + if value <= 0: + _warn_invalid_env(name, raw, "must be a positive integer", f"{default_mb} MB") return default_bytes - return value * 1024 * 1024 if value > 0 else default_bytes + return value * 1024 * 1024 for old in deprecated_mb_names: raw = os.environ.get(old) @@ -338,8 +425,12 @@ def _env_int_mb_with_deprecated_units( try: value = int(raw) except ValueError: + _warn_invalid_env(old, raw, "not an integer", f"{default_mb} MB") return default_bytes - return value * 1024 * 1024 if value > 0 else default_bytes + if value <= 0: + _warn_invalid_env(old, raw, "must be a positive integer", f"{default_mb} MB") + return default_bytes + return value * 1024 * 1024 for old in deprecated_byte_names: raw = os.environ.get(old) @@ -348,8 +439,68 @@ def _env_int_mb_with_deprecated_units( try: value = int(raw) except ValueError: + _warn_invalid_env(old, raw, "not an integer", f"{default_mb} MB") + return default_bytes + if value <= 0: + _warn_invalid_env(old, raw, "must be a positive integer", f"{default_mb} MB") return default_bytes - return value if value > 0 else default_bytes + return value + + return default_bytes + + +def _env_bytes( + name: str, + default_bytes: int, + description: str, + category: str, + *, + deprecated_mb_names: tuple[str, ...] = (), + deprecated_byte_names: tuple[str, ...] = (), + allow_zero: bool = False, +) -> int: + """Read a byte-size setting whose canonical name accepts a unit suffix. + + The canonical name treats a bare number as **bytes** and understands the + suffixes ``B``, ``K``/``KiB``, ``KB``, ``M``/``MiB``, ``MB``, ``G``/``GiB`` + and ``GB`` (binary for the ``iB``/bare-letter forms, decimal for ``KB``/ + ``MB``/``GB``). A megabyte-only knob cannot express thresholds below 1 MB, + which is exactly the range that matters when deciding whether a small file + rides inline in a commit or goes to object storage. + + Names in ``deprecated_mb_names`` keep their megabyte meaning for a bare + number; names in ``deprecated_byte_names`` keep their byte meaning. + """ + deprecated_names = deprecated_mb_names + deprecated_byte_names + _env_register(name, _format_bytes(default_bytes), description, category, deprecated_names=deprecated_names) + + def _accept(source: str, raw: str, bare_unit: int) -> int: + try: + value = _parse_byte_size(raw, bare_unit=bare_unit) + except ValueError as exc: + _warn_invalid_env(source, raw, str(exc), _format_bytes(default_bytes)) + return default_bytes + if value < 0 or (value == 0 and not allow_zero): + reason = "must not be negative" if allow_zero else "must be a positive size" + _warn_invalid_env(source, raw, reason, _format_bytes(default_bytes)) + return default_bytes + return value + + raw = os.environ.get(name) + if raw is not None and raw.strip(): + return _accept(name, raw, 1) + + for old in deprecated_mb_names: + raw = os.environ.get(old) + if raw is not None and raw.strip(): + _warn_deprecated_env(old, name, expects_bytes=True, stacklevel=2) + return _accept(old, raw, 1024 * 1024) + + for old in deprecated_byte_names: + raw = os.environ.get(old) + if raw is not None and raw.strip(): + _warn_deprecated_env(old, name, expects_bytes=True, stacklevel=2) + return _accept(old, raw, 1) return default_bytes @@ -439,6 +590,42 @@ def _env_register( "API_MAX_RETRIES", ) +API_CONNECTION_POOL_MAXSIZE: int = _env_int( + "MODELSCOPE_API_CONNECTION_POOL_MAXSIZE", + 32, + "Per-host HTTP connection pool size", + "Network", +) +"""Connections kept alive per host, and the concurrency the pool can serve. + +urllib3 defaults this to 10. A folder upload runs ``max_workers`` requests at +once -- commonly 16 or more for bulk transfers -- so the default silently +discards the excess connections ("Connection pool is full, discarding +connection") and every discarded one costs a fresh TLS handshake on its next +use. This must be at least as large as the worker count to avoid that churn. +""" + +REPO_TREE_PAGE_MAX_ATTEMPTS: int = _env_int( + "MODELSCOPE_REPO_TREE_PAGE_MAX_ATTEMPTS", + 4, + "Attempts for one repo-tree page denied on an already-authorized listing", + "Network", +) +"""Retries for a spurious ``403`` on a single page of a paginated tree listing. + +A large dataset listing spans hundreds of pages and an occasional page answers +``403 无权访问该数据集`` on a repository the caller has just read successfully. +Treating that as an authorization result throws away every page already +collected, so it is retried once the credential has been proven by an earlier +page. +""" +REPO_TREE_PAGE_RETRY_MAX_DELAY_SECONDS: int = _env_int( + "MODELSCOPE_REPO_TREE_PAGE_RETRY_MAX_DELAY_SECONDS", + 8, + "Maximum backoff between repo-tree page retries (seconds)", + "Network", +) + REPO_FILES_TRUNCATION_LIMIT: int = 3000 """Server-side hard cap on a single ``repo/files`` listing. @@ -649,6 +836,78 @@ def _env_register( "Upload", "UPLOAD_ADAPTIVE_BATCH_SIZE", ) +UPLOAD_COMMIT_MAX_INLINE_BYTES: int = _env_bytes( + "MODELSCOPE_UPLOAD_COMMIT_MAX_INLINE_BYTES", + 8 * 1024 * 1024, + "Maximum inlined (non-LFS) content carried by one commit request", + "Upload", +) +"""Byte ceiling on the base64 content a single commit may carry. + +Non-LFS files travel *inside* the commit request body, so a batch sized purely +by file count can produce a request tens of megabytes large -- big enough for +the server to time out mid-read, which surfaces on the client as an unrelated +write timeout. Batching therefore closes a batch on whichever limit is reached +first, this one or :data:`UPLOAD_COMMIT_BATCH_MAX_OPERATIONS`. LFS files +contribute only a pointer, so they do not count against it. +""" +UPLOAD_COMMIT_MAX_PER_HOUR: int = _env_int( + "MODELSCOPE_UPLOAD_COMMIT_MAX_PER_HOUR", + 0, + "Client-side commit rate ceiling per hour (0 disables the governor)", + "Upload", + allow_zero=True, +) +"""Opt-in client-side commit budget, disabled by default. + +The Hub throttles commits per repository. Reacting to a throttle costs a failed +round trip and, when the server holds the connection instead of answering, a +full read timeout. A bulk pipeline that knows its budget can set this to spread +commits out and never trip the limit; interactive uploads stay unthrottled. +""" +UPLOAD_COMMIT_MAX_RETRY_AFTER_SECONDS: int = _env_int( + "MODELSCOPE_UPLOAD_COMMIT_MAX_RETRY_AFTER_SECONDS", + 1800, + "Longest server-provided Retry-After a commit will honor (seconds)", + "Upload", +) +"""Upper bound on an honored ``Retry-After`` for a throttled commit. + +A rate limit is a bounded wait the server declares, unlike a transient failure +of unknown duration, so it is budgeted separately from +:data:`UPLOAD_COMMIT_RETRY_TOTAL_WAIT_SECONDS` rather than exhausting it. +""" +COMMIT_MAX_ACTIONS_PER_REQUEST: int = _env_int( + "MODELSCOPE_COMMIT_MAX_ACTIONS_PER_REQUEST", + 2000, + "Server-enforced maximum actions in one commit request", + "Upload", +) +"""Hard ceiling the server puts on a single commit request. + +Exceeding it is rejected outright with ``HTTP 422``:: + + commit request exceeds actions limit: 3300 > 2000; split the commit into + smaller batches + +This is a server contract, not a tuning preference, so every commit path clamps +to it -- uploads, and deletes, which otherwise put every path in one request. +""" + +UPLOAD_PROGRESS_MIN_INTERVAL_SECONDS: int = _env_int( + "MODELSCOPE_UPLOAD_PROGRESS_MIN_INTERVAL_SECONDS", + 1, + "Minimum gap between wire-level upload progress events (seconds)", + "Upload", +) +"""Throttle for per-file upload progress events. + +Commits land in lumps tens of seconds apart, so a rate built only from commit +events alternates between a spike and zero and cannot distinguish a slow batch +from a hung one. Blob uploads finish continuously and are the honest source for +a rate -- but there is one per file, so the events are coalesced to this interval +rather than fanning out tens of thousands of callback invocations. +""" UPLOAD_COMMIT_MAX_ATTEMPTS: int = _env_int( "MODELSCOPE_UPLOAD_COMMIT_MAX_ATTEMPTS", 5, @@ -751,13 +1010,42 @@ def get_upload_ignore_file_pattern() -> str | None: UPLOAD_LEGACY_PROGRESS_FILE: str = ".ms_upload_progress" # Upload: limits -UPLOAD_LFS_FORCE_THRESHOLD_BYTES: int = _env_int_mb_with_deprecated_units( - "MODELSCOPE_UPLOAD_LFS_FORCE_THRESHOLD_MB", - 1, - "File-size threshold that forces LFS mode (MB)", +UPLOAD_LFS_FORCE_THRESHOLD_BYTES: int = _env_bytes( + "MODELSCOPE_UPLOAD_LFS_FORCE_THRESHOLD", + 1024 * 1024, + "File size above which LFS mode is forced (bytes; accepts a unit suffix, 0 forces LFS for every file)", "Upload", + deprecated_mb_names=("MODELSCOPE_UPLOAD_LFS_FORCE_THRESHOLD_MB",), deprecated_byte_names=("UPLOAD_LFS_ENFORCE_THRESHOLD", "UPLOAD_SIZE_THRESHOLD_TO_ENFORCE_LFS"), + allow_zero=True, ) +"""Size above which a file is uploaded as LFS regardless of its suffix. + +Files at or below the threshold are committed inline as base64, which puts their +bytes in the commit request body and caps how many of them one commit can carry. +Lowering the threshold moves that content onto the pre-signed object-storage +path instead, leaving the commit with only ``sha256`` plus ``size``. + +``0`` forces LFS for every non-empty file. It stays safe because +:data:`UPLOAD_INLINE_METADATA_PATHS` is consulted first, so the repository files +the Hub itself parses are never turned into LFS pointers. +""" +UPLOAD_INLINE_METADATA_PATHS: frozenset[str] = _env_csv_frozenset( + "MODELSCOPE_UPLOAD_INLINE_METADATA_PATHS", + "README.md,.gitattributes,.gitignore,configuration.json,configuration.yaml,configuration.yml," + "dataset_infos.json,config.json,.msc,.mdl", + "Repository file names always committed inline, never as LFS", + "Upload", +) +"""Repository-relative file names that must stay inline in the commit. + +The Hub parses these server-side -- the dataset/model card front matter, the +configuration files, the git attribute rules. Stored as an LFS pointer, the +server would read the 130-byte pointer text instead of the real content and the +card or configuration would silently render empty. Matching is on the file name +(case-insensitive), so the rule holds at any depth in the tree, and it is +checked before the size and suffix rules. +""" UPLOAD_MAX_FILE_SIZE_BYTES: int = _env_int_mb_with_deprecated_units( "MODELSCOPE_UPLOAD_MAX_FILE_SIZE_MB", 100 * 1024, @@ -955,9 +1243,11 @@ def get_upload_ignore_file_pattern() -> str | None: __all__ = [ "API_CONNECT_TIMEOUT", + "API_CONNECTION_POOL_MAXSIZE", "API_MAX_RETRIES", "API_TIMEOUT", "CATEGORY_ORDER", + "COMMIT_MAX_ACTIONS_PER_REQUEST", "CONFIG_DIR_NAME", "DATASET_LFS_SUFFIX", "DEFAULT_CACHE_DIR_NAME", @@ -996,6 +1286,8 @@ def get_upload_ignore_file_pattern() -> str | None: "MODEL_ID_SEPARATOR", "MODEL_LFS_SUFFIX", "OPENAPI_PREFIX", + "REPO_TREE_PAGE_MAX_ATTEMPTS", + "REPO_TREE_PAGE_RETRY_MAX_DELAY_SECONDS", "REPO_TYPE_DATASET", "REPO_TYPE_MODEL", "REPO_TYPE_STUDIO", @@ -1029,12 +1321,16 @@ def get_upload_ignore_file_pattern() -> str | None: "UPLOAD_COMMIT_BATCH_SIZE", "UPLOAD_COMMIT_MAX_ATTEMPTS", "UPLOAD_COMMIT_MAX_CONSECUTIVE_FAILED_BATCHES", + "UPLOAD_COMMIT_MAX_INLINE_BYTES", + "UPLOAD_COMMIT_MAX_PER_HOUR", "UPLOAD_COMMIT_MAX_RETRIES", + "UPLOAD_COMMIT_MAX_RETRY_AFTER_SECONDS", "UPLOAD_COMMIT_MAX_TOTAL_WAIT", "UPLOAD_COMMIT_RETRY_TOTAL_WAIT_SECONDS", "UPLOAD_FAILED_FILE_MAX_RETRIES", "UPLOAD_FAILED_FILE_MAX_RETRY_ROUNDS", "UPLOAD_HTTP_RETRY_ALLOWED_METHODS", + "UPLOAD_INLINE_METADATA_PATHS", "UPLOAD_LEGACY_PROGRESS_FILE", "UPLOAD_LFS_ENFORCE_THRESHOLD", "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", @@ -1047,6 +1343,7 @@ def get_upload_ignore_file_pattern() -> str | None: "UPLOAD_MAX_FILES_PER_DIRECTORY", "UPLOAD_NORMAL_FILE_SIZE_TOTAL_LIMIT", "UPLOAD_NORMAL_FILES_TOTAL_SIZE_BYTES", + "UPLOAD_PROGRESS_MIN_INTERVAL_SECONDS", "UPLOAD_REACT_BACKOFF_MAX_EXPONENT", "UPLOAD_REACT_ENABLED", "UPLOAD_REACT_MAX_DELAY", diff --git a/src/modelscope_hub/version.py b/src/modelscope_hub/version.py index 06d4093..62ee45e 100644 --- a/src/modelscope_hub/version.py +++ b/src/modelscope_hub/version.py @@ -1,3 +1,3 @@ """Version information for modelscope_hub.""" -__version__ = "0.4.3+main" +__version__ = "0.4.4+main" diff --git a/tests/test_commit_action_limit.py b/tests/test_commit_action_limit.py new file mode 100644 index 0000000..9bbbf94 --- /dev/null +++ b/tests/test_commit_action_limit.py @@ -0,0 +1,110 @@ +"""Commit requests must respect the server's hard action ceiling. + +The server rejects an oversized commit outright rather than truncating it:: + + HTTP 422 commit request exceeds actions limit: 3300 > 2000; + split the commit into smaller batches + +Observed while removing 3300 verification files from a test dataset. Uploads were +accidentally safe because the default operation cap is well below the ceiling, +but an explicitly configured larger cap, or any delete of more than 2000 paths, +would hit it. +""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +import modelscope_hub._upload as upload_module +from modelscope_hub._upload import ( + UploadManager, + _calculate_adaptive_batch_size, + _plan_commit_batches, +) +from modelscope_hub.errors import APIError + + +def _make_manager() -> tuple[UploadManager, MagicMock]: + client = MagicMock() + client.create_commit.return_value = {"ok": True} + client.validate_blobs.side_effect = lambda **kw: {o["oid"]: f"https://upload/{o['oid']}" for o in kw["objects"]} + client.upload_blob.side_effect = lambda **kw: None + return UploadManager(client, MagicMock()), client + + +def test_batch_size_is_clamped_to_the_server_ceiling(monkeypatch) -> None: + monkeypatch.setattr(upload_module, "COMMIT_MAX_ACTIONS_PER_REQUEST", 2000) + + # An explicit cap above the ceiling must not be honored: the server would + # reject the request instead of trimming it. + assert _calculate_adaptive_batch_size(50_000, 4096) == 2000 + assert _calculate_adaptive_batch_size(50_000, 512) == 512 + # "No cap" must not mean "one commit for everything" either. + assert _calculate_adaptive_batch_size(50_000, 0) == 2000 + + +def test_batch_plan_never_exceeds_the_server_ceiling(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr(upload_module, "COMMIT_MAX_ACTIONS_PER_REQUEST", 10) + files = [(f"f{i}.bin", str(tmp_path / f"f{i}.bin")) for i in range(35)] + sizes = dict.fromkeys((path for _, path in files), 1) + + plan = _plan_commit_batches( + files, + "dataset", + max_operations=1000, + max_inline_bytes=0, + sizes=sizes, + ) + + assert sum(plan) == 35 + assert max(plan) <= 10 + + +def test_delete_files_splits_at_the_server_ceiling(monkeypatch) -> None: + monkeypatch.setattr(upload_module, "COMMIT_MAX_ACTIONS_PER_REQUEST", 100) + manager, client = _make_manager() + paths = [f"junk/file-{i}.bin" for i in range(250)] + + result = manager.delete_files(repo_id="owner/repo", repo_type="dataset", file_paths=paths) + + assert client.create_commit.call_count == 3 + sent = [len(call.kwargs["operations"]) for call in client.create_commit.call_args_list] + assert sent == [100, 100, 50] + assert result["total_files"] == 250 + assert result["deleted_files"] == paths + + +def test_delete_files_keeps_one_commit_when_it_fits(monkeypatch) -> None: + monkeypatch.setattr(upload_module, "COMMIT_MAX_ACTIONS_PER_REQUEST", 2000) + manager, client = _make_manager() + + manager.delete_files(repo_id="owner/repo", repo_type="dataset", file_paths=["a.bin", "b.bin"]) + + client.create_commit.assert_called_once() + assert len(client.create_commit.call_args.kwargs["operations"]) == 2 + + +def test_delete_files_reports_what_a_partial_split_removed(monkeypatch) -> None: + # Across a split the deletion is no longer atomic. Failing loudly beats + # returning a success that claims paths which are still present. + monkeypatch.setattr(upload_module, "COMMIT_MAX_ACTIONS_PER_REQUEST", 2) + manager, client = _make_manager() + client.create_commit.side_effect = [{"ok": True}, APIError("commit rejected", status_code=422)] + + with pytest.raises(APIError): + manager.delete_files( + repo_id="owner/repo", + repo_type="dataset", + file_paths=["a", "b", "c", "d"], + ) + + assert client.create_commit.call_count == 2 + + +def test_server_ceiling_is_env_tunable() -> None: + from modelscope_hub.constants import COMMIT_MAX_ACTIONS_PER_REQUEST + + assert COMMIT_MAX_ACTIONS_PER_REQUEST == 2000 diff --git a/tests/test_repo_tree_page_retry.py b/tests/test_repo_tree_page_retry.py new file mode 100644 index 0000000..c81ba9f --- /dev/null +++ b/tests/test_repo_tree_page_retry.py @@ -0,0 +1,143 @@ +"""File-tree listing resilience against spurious denials. + +The server intermittently answers a tree request with ``403 无权访问该数据集`` on a +repository the caller has just read successfully. Observed live on a 25k-file +dataset (paginated listing: page 7 denied, then page 145 denied, then success) +and again on a 10k-file dataset (``Root``-scoped listing denied outright). +Enumerating a large repo takes hundreds of such requests, so a per-request +failure rate becomes a near-certain whole-listing failure. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +import modelscope_hub._legacy_api as legacy_module +from modelscope_hub._legacy_api import LegacyClient +from modelscope_hub.errors import NetworkError, PermissionDeniedError + + +def _client() -> LegacyClient: + return LegacyClient(endpoint="https://modelscope.cn", token="ms-test") + + +def _page(count: int, offset: int = 0) -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.headers = {} + response.json.return_value = {"Data": {"Files": [{"Path": f"f{offset + i}"} for i in range(count)]}} + return response + + +def _denied() -> PermissionDeniedError: + return PermissionDeniedError("无权访问该数据集", status_code=403) + + +def test_transient_denial_on_a_later_page_is_retried(monkeypatch) -> None: + client = _client() + monkeypatch.setattr(legacy_module.time, "sleep", lambda _s: None) + calls: list[int] = [] + + def request(method: str, path: str, **kwargs): + page = kwargs["params"]["PageNumber"] + calls.append(page) + if page == 2 and calls.count(2) == 1: + raise _denied() + return _page(200, offset=page * 200) if page < 3 else _page(7, offset=600) + + with patch.object(client, "_request", side_effect=request): + files = client.list_dataset_files_paginated("owner/ds", page_size=200) + + # Page 2 was retried rather than aborting and discarding page 1. + assert calls == [1, 2, 2, 3] + assert len(files) == 407 + + +def test_denial_before_any_successful_tree_read_fails_fast() -> None: + # Nothing has proven the credential yet, so this really is an authorization + # answer and must not be retried into a slow, misleading failure. + client = _client() + calls: list[int] = [] + + def request(method: str, path: str, **kwargs): + calls.append(kwargs["params"]["PageNumber"]) + raise _denied() + + with patch.object(client, "_request", side_effect=request): + with pytest.raises(PermissionDeniedError): + client.list_dataset_files_paginated("owner/ds", page_size=200) + + assert calls == [1] + + +def test_root_scoped_listing_also_retries_a_spurious_denial(monkeypatch) -> None: + # The per-directory walk that works around the 3000-entry cap issues the same + # kind of request, and was seen to be denied the same way. + client = _client() + monkeypatch.setattr(legacy_module.time, "sleep", lambda _s: None) + client._tree_reads_ok = True # an earlier read already proved the credential + calls: list[str | None] = [] + + def request(method: str, path: str, **kwargs): + root = kwargs["params"].get("Root") + calls.append(root) + if calls.count(root) == 1: + raise _denied() + return _page(3) + + with patch.object(client, "_request", side_effect=request): + entries = client._list_files_page("owner/ds", "dataset", "master", recursive=False, root="level1_000") + + assert calls == ["level1_000", "level1_000"] + assert len(entries) == 3 + + +def test_a_successful_tree_read_arms_the_retry_for_later_requests(monkeypatch) -> None: + client = _client() + monkeypatch.setattr(legacy_module.time, "sleep", lambda _s: None) + assert client._tree_reads_ok is False + outcomes = [_page(2), _denied(), _page(2)] + + def request(method: str, path: str, **kwargs): + result = outcomes.pop(0) + if isinstance(result, Exception): + raise result + return result + + with patch.object(client, "_request", side_effect=request): + client._list_files_page("owner/ds", "dataset", "master", recursive=False, root="a") + assert client._tree_reads_ok is True + # The denial that follows is now retried instead of raised. + entries = client._list_files_page("owner/ds", "dataset", "master", recursive=False, root="b") + + assert len(entries) == 2 + assert not outcomes + + +def test_a_persistently_denied_listing_is_reported_as_a_transport_failure(monkeypatch) -> None: + client = _client() + monkeypatch.setattr(legacy_module.time, "sleep", lambda _s: None) + monkeypatch.setattr(legacy_module, "REPO_TREE_PAGE_MAX_ATTEMPTS", 3) + + def request(method: str, path: str, **kwargs): + if kwargs["params"]["PageNumber"] == 1: + return _page(200) + raise _denied() + + with patch.object(client, "_request", side_effect=request): + # Surfaced as a transport failure, not as "permission denied": the caller + # can read the repo, so reporting a denial would misdirect the diagnosis. + with pytest.raises(NetworkError, match="already-authorized"): + client.list_dataset_files_paginated("owner/ds", page_size=200) + + +def test_retry_budget_is_env_tunable() -> None: + from modelscope_hub.constants import ( + REPO_TREE_PAGE_MAX_ATTEMPTS, + REPO_TREE_PAGE_RETRY_MAX_DELAY_SECONDS, + ) + + assert REPO_TREE_PAGE_MAX_ATTEMPTS >= 2 + assert REPO_TREE_PAGE_RETRY_MAX_DELAY_SECONDS >= 1 diff --git a/tests/test_upload_batching.py b/tests/test_upload_batching.py new file mode 100644 index 0000000..157a458 --- /dev/null +++ b/tests/test_upload_batching.py @@ -0,0 +1,885 @@ +"""Commit-batch planning, LFS routing and commit throttling behaviour.""" + +from __future__ import annotations + +import json +import threading +import time +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +import modelscope_hub._upload as upload_module +from modelscope_hub._upload import ( + BatchTracker, + UploadManager, + _calculate_adaptive_batch_size, + _is_inline_metadata, + _normalize_path_in_repo, + _plan_commit_batches, + _upload_mode, +) +from modelscope_hub.errors import InvalidParameter, NetworkError, RateLimitError + + +def _make_manager() -> tuple[UploadManager, MagicMock]: + client = MagicMock() + client.create_commit.return_value = {"ok": True} + + def validate_blobs(*, repo_id: str, repo_type: str, objects: list[dict]) -> dict[str, str]: + return {obj["oid"]: f"https://upload/{obj['oid']}" for obj in objects} + + def upload_blob(*, upload_url: str, data, size: int) -> None: + while data.read(1024 * 1024): + pass + + client.validate_blobs.side_effect = validate_blobs + client.upload_blob.side_effect = upload_blob + return UploadManager(client, MagicMock()), client + + +# ---------------------------------------------------------------- batch sizing + + +@pytest.mark.parametrize( + ("total_files", "cap", "expected"), + [ + (0, 256, 1), + (1, 256, 1), + (100, 256, 100), + (2_000, 256, 256), + (60_000, 256, 256), + (60_000, 512, 512), + # "No cap" still respects the server's action ceiling, not the file count. + (5_000, 0, 2_000), + (1_500, 0, 1_500), + # An explicit cap above the ceiling cannot be honored either. + (60_000, 4_096, 2_000), + ], +) +def test_adaptive_batch_size_fills_up_to_the_cap(total_files: int, cap: int, expected: int) -> None: + assert _calculate_adaptive_batch_size(total_files, cap) == expected + + +def test_adaptive_batch_size_is_monotonic_and_never_exceeds_cap() -> None: + sizes = [_calculate_adaptive_batch_size(n, 256) for n in range(1, 3_000)] + assert sizes == sorted(sizes) + assert max(sizes) == 256 + + +def test_batch_plan_closes_on_inline_bytes_before_operation_cap(tmp_path: Path) -> None: + # 8 inline files of 100 KiB: the operation cap would take all 8 in one + # commit, but their base64 form exceeds a 512 KiB inline budget. + files = [] + sizes = {} + for index in range(8): + path = tmp_path / f"part-{index}.txt" + path.write_bytes(b"x" * (100 * 1024)) + files.append((f"part-{index}.txt", str(path))) + sizes[str(path)] = 100 * 1024 + + plan = _plan_commit_batches( + files, + "dataset", + max_operations=256, + max_inline_bytes=512 * 1024, + sizes=sizes, + ) + + assert sum(plan) == 8 + assert max(plan) < 8 + encoded_per_file = (100 * 1024 + 2) // 3 * 4 + assert all(count * encoded_per_file <= 512 * 1024 for count in plan) + + +def test_batch_plan_ignores_lfs_bytes_in_the_inline_budget(tmp_path: Path) -> None: + # .parquet is an LFS suffix for datasets, so these contribute a pointer + # rather than inline content and must not split the batch. + files = [] + sizes = {} + for index in range(8): + path = tmp_path / f"shard-{index}.parquet" + files.append((f"shard-{index}.parquet", str(path))) + sizes[str(path)] = 100 * 1024 + + plan = _plan_commit_batches( + files, + "dataset", + max_operations=256, + max_inline_bytes=512 * 1024, + sizes=sizes, + ) + + assert plan == [8] + + +def test_batch_plan_keeps_one_oversized_inline_file_per_batch(tmp_path: Path) -> None: + # Below the 1 MiB LFS threshold, so these really do travel inline; each one + # alone blows the inline budget, which must not stall the plan. + files = [] + sizes = {} + for index in range(3): + path = tmp_path / f"big-{index}.txt" + files.append((f"big-{index}.txt", str(path))) + sizes[str(path)] = 512 * 1024 + + plan = _plan_commit_batches( + files, + "dataset", + max_operations=256, + max_inline_bytes=1024, + sizes=sizes, + ) + + assert plan == [1, 1, 1] + + +def test_batch_tracker_handles_uneven_batches() -> None: + tracker = BatchTracker(7, [3, 1, 3]) + + assert tracker.num_batches == 3 + assert tracker.batch_range(0) == (0, 3) + assert tracker.batch_range(1) == (3, 4) + assert tracker.batch_range(2) == (4, 7) + assert [tracker.batch_index(i) for i in range(7)] == [0, 0, 0, 1, 2, 2, 2] + + +def test_batch_tracker_absorbs_a_short_plan() -> None: + tracker = BatchTracker(5, [2]) + + assert tracker.num_batches == 2 + assert tracker.batch_range(1) == (2, 5) + + +def test_batch_tracker_accepts_a_uniform_size() -> None: + tracker = BatchTracker(5, 2) + + assert tracker.num_batches == 3 + assert [tracker.batch_index(i) for i in range(5)] == [0, 0, 1, 1, 2] + + +# ------------------------------------------------------- inline metadata gate + + +@pytest.mark.parametrize( + "path", + ["README.md", "nested/README.md", ".gitattributes", "configuration.json", "sub/dir/config.json"], +) +def test_inline_metadata_is_recognized_at_any_depth(path: str) -> None: + assert _is_inline_metadata(path) + + +def test_metadata_stays_inline_even_below_a_zero_threshold(monkeypatch) -> None: + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 0) + + assert _upload_mode("README.md", 500 * 1024, "dataset") == "normal" + assert _upload_mode("configuration.json", 10, "model") == "normal" + # Everything that is not repository metadata does move to LFS. + assert _upload_mode("data/sample.txt", 1, "dataset") == "lfs" + + +def test_small_files_move_to_lfs_when_the_threshold_is_lowered(monkeypatch) -> None: + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 32 * 1024) + + assert _upload_mode("data/sample.txt", 115 * 1024, "dataset") == "lfs" + # A small card or config still rides inline at this threshold. + assert _upload_mode("data/notes.txt", 2 * 1024, "dataset") == "normal" + + +def test_upload_folder_routes_small_files_to_blob_storage(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 32 * 1024) + manager, client = _make_manager() + (tmp_path / "README.md").write_bytes(b"# card") + for index in range(4): + (tmp_path / f"sample-{index}.txt").write_bytes(bytes([index]) * (64 * 1024)) + + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + max_workers=2, + use_cache=False, + disable_tqdm=True, + ) + + operations = client.create_commit.call_args.kwargs["operations"] + modes = {op["path"]: op["type"] for op in operations} + assert modes == { + "README.md": "normal", + "sample-0.txt": "lfs", + "sample-1.txt": "lfs", + "sample-2.txt": "lfs", + "sample-3.txt": "lfs", + } + # The four data files were pre-signed together, not one request each. + assert client.validate_blobs.call_count == 1 + assert len(client.validate_blobs.call_args.kwargs["objects"]) == 4 + assert client.upload_blob.call_count == 4 + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("", ""), + (".", ""), + ("./", ""), + ("/", ""), + (" . ", ""), + ("sub", "sub"), + ("./sub", "sub"), + ("sub/", "sub"), + ("/sub/", "sub"), + ("a/./b", "a/b"), + ("a/../b", "b"), + ("data\\shard", "data/shard"), + ], +) +def test_normalize_path_in_repo_collapses_root_aliases(raw: str, expected: str) -> None: + assert _normalize_path_in_repo(raw) == expected + + +@pytest.mark.parametrize("escaping", ["..", "../x", "a/../../b"]) +def test_normalize_path_in_repo_refuses_escaping_root(escaping: str) -> None: + with pytest.raises(InvalidParameter): + _normalize_path_in_repo(escaping) + + +def test_upload_folder_dot_path_in_repo_commits_root_relative_paths(tmp_path: Path, monkeypatch) -> None: + # `ms upload REPO LOCAL .` passes path_in_repo=".". Left literal it became a + # "./" prefix on every commit action path, which the Hub rejects wholesale + # as an invalid commit action (E3021). "." must map to the repo root. + manager, client = _make_manager() + (tmp_path / "README.md").write_bytes(b"# card") + (tmp_path / "data").mkdir() + (tmp_path / "data" / "train.txt").write_bytes(b"rows") + + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + path_in_repo=".", + max_workers=2, + use_cache=False, + disable_tqdm=True, + ) + + operations = client.create_commit.call_args.kwargs["operations"] + assert sorted(op["path"] for op in operations) == ["README.md", "data/train.txt"] + assert not any(op["path"].startswith("./") for op in operations) + + +def test_upload_folder_batch_presigns_in_configured_group_size(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 0) + monkeypatch.setattr(upload_module, "UPLOAD_BLOB_VALIDATION_BATCH_MAX_OBJECTS", 3) + manager, client = _make_manager() + for index in range(7): + (tmp_path / f"sample-{index}.txt").write_bytes(bytes([index]) * 128) + + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + max_workers=2, + use_cache=False, + disable_tqdm=True, + ) + + assert client.validate_blobs.call_count == 3 + assert client.upload_blob.call_count == 7 + + +def test_presign_groups_run_concurrently(tmp_path: Path, monkeypatch) -> None: + # The groups are independent requests; serialising them made the pre-sign + # phase scale with the file count in pure round-trip latency. + monkeypatch.setattr(upload_module, "UPLOAD_BLOB_VALIDATION_BATCH_MAX_OBJECTS", 1) + manager, client = _make_manager() + peak = 0 + inflight = 0 + lock = threading.Lock() + + def validate_blobs(*, repo_id: str, repo_type: str, objects: list[dict]) -> dict[str, str]: + nonlocal peak, inflight + with lock: + inflight += 1 + peak = max(peak, inflight) + time.sleep(0.05) + with lock: + inflight -= 1 + return {obj["oid"]: f"https://upload/{obj['oid']}" for obj in objects} + + client.validate_blobs.side_effect = validate_blobs + objects = [{"oid": f"{index:064x}", "size": 1} for index in range(8)] + + result = manager._validate_blobs_batch(repo_id="owner/repo", repo_type="dataset", objects=objects, max_workers=8) + + assert len(result) == 8 + assert peak > 1, "pre-sign groups must overlap" + + +def test_presign_concurrency_stays_within_the_connection_pool(tmp_path: Path, monkeypatch) -> None: + # More in-flight requests than pooled connections just trades latency for + # discarded connections and fresh TLS handshakes. + monkeypatch.setattr(upload_module, "UPLOAD_BLOB_VALIDATION_BATCH_MAX_OBJECTS", 1) + monkeypatch.setattr(upload_module, "API_CONNECTION_POOL_MAXSIZE", 2) + manager, client = _make_manager() + peak = 0 + inflight = 0 + lock = threading.Lock() + + def validate_blobs(*, repo_id: str, repo_type: str, objects: list[dict]) -> dict[str, str]: + nonlocal peak, inflight + with lock: + inflight += 1 + peak = max(peak, inflight) + time.sleep(0.02) + with lock: + inflight -= 1 + return {obj["oid"]: "https://upload/x" for obj in objects} + + client.validate_blobs.side_effect = validate_blobs + objects = [{"oid": f"{index:064x}", "size": 1} for index in range(12)] + + manager._validate_blobs_batch(repo_id="owner/repo", repo_type="dataset", objects=objects, max_workers=32) + + assert peak <= 2 + + +def test_presign_normalizes_unmentioned_objects_as_existing(monkeypatch) -> None: + # The batch endpoint answers only about objects that need uploading and + # returns empty arrays for ones that already exist. Leaving those out of the + # map would send every file of a re-run back to per-file negotiation. + monkeypatch.setattr(upload_module, "UPLOAD_BLOB_VALIDATION_BATCH_MAX_OBJECTS", 2) + manager, client = _make_manager() + objects = [{"oid": f"{index:064x}", "size": 1} for index in range(4)] + # Server mentions only the last object; the rest already exist. + wanted = objects[-1]["oid"] + client.validate_blobs.side_effect = lambda **kw: ( + {wanted: "https://upload/x"} if any(o["oid"] == wanted for o in kw["objects"]) else {} + ) + + result = manager._validate_blobs_batch(repo_id="owner/repo", repo_type="dataset", objects=objects, max_workers=4) + + assert set(result) == {o["oid"] for o in objects} + assert result[wanted] == "https://upload/x" + assert all(result[o["oid"]] is None for o in objects[:-1]) + + +def test_a_failed_presign_group_falls_back_to_per_file_negotiation(tmp_path: Path, monkeypatch) -> None: + # Losing the optimisation must never escalate into losing the upload -- and + # critically, an oid missing from the map must not read as "already exists", + # which would skip the transfer and commit a pointer to a blob that was + # never stored. + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 0) + monkeypatch.setattr(upload_module, "UPLOAD_BLOB_VALIDATION_BATCH_MAX_OBJECTS", 2) + manager, client = _make_manager() + for index in range(6): + (tmp_path / f"blob-{index}.dat").write_bytes(bytes([index]) * 512) + + seen: list[int] = [] + + def validate_blobs(*, repo_id: str, repo_type: str, objects: list[dict]) -> dict[str, str]: + seen.append(len(objects)) + if len(objects) > 1 and len(seen) == 1: + raise NetworkError("pre-sign group failed") + return {obj["oid"]: f"https://upload/{obj['oid']}" for obj in objects} + + client.validate_blobs.side_effect = validate_blobs + + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + max_workers=1, + use_cache=False, + disable_tqdm=True, + ) + + # Every file still reached object storage, none was silently treated as reused. + assert client.upload_blob.call_count == 6 + operations = client.create_commit.call_args.kwargs["operations"] + assert len(operations) == 6 + assert all(op["type"] == "lfs" and op["sha256"] for op in operations) + + +@pytest.mark.parametrize("use_cache", [True, False]) +def test_pre_hashed_files_are_not_hashed_twice(tmp_path: Path, monkeypatch, use_cache: bool) -> None: + # Pre-hashing feeds group pre-signing; without threading the digest through + # to the upload worker it would read and hash every file a second time. + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 0) + manager, _ = _make_manager() + for index in range(5): + (tmp_path / f"sample-{index}.txt").write_bytes(bytes([index]) * 256) + + calls: list[object] = [] + original = upload_module._compute_file_hash + + def counting_hash(*args, **kwargs): + calls.append(kwargs.get("file_path_or_obj", args[0] if args else None)) + return original(*args, **kwargs) + + monkeypatch.setattr(upload_module, "_compute_file_hash", counting_hash) + + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + max_workers=2, + use_cache=use_cache, + disable_tqdm=True, + ) + + assert len(calls) == 5 + assert len(set(calls)) == 5 + + +# ------------------------------------------------------------- tracker path + + +def test_tracker_path_keeps_resume_state_outside_the_uploaded_tree(tmp_path: Path) -> None: + manager, _ = _make_manager() + staging = tmp_path / "staging" + staging.mkdir() + (staging / "README.md").write_bytes(b"hello") + cache_path = tmp_path / "state" / "upload.json" + cache_path.parent.mkdir() + + manager.upload_folder( + repo_id="owner/repo", + repo_type="model", + folder_path=staging, + max_workers=1, + use_cache=True, + disable_tqdm=True, + tracker_path=cache_path, + ) + + assert not (staging / ".ms_upload_cache").exists() + assert cache_path.exists() + assert json.loads(cache_path.read_text(encoding="utf-8"))["repo_id"] == "owner/repo" + + +def test_external_tracker_lets_a_second_run_skip_committed_files(tmp_path: Path) -> None: + manager, client = _make_manager() + staging = tmp_path / "staging" + staging.mkdir() + (staging / "README.md").write_bytes(b"hello") + cache_path = tmp_path / "upload.json" + + for _ in range(2): + manager.upload_folder( + repo_id="owner/repo", + repo_type="model", + folder_path=staging, + max_workers=1, + use_cache=True, + disable_tqdm=True, + tracker_path=cache_path, + ) + + # The second run recognises the committed file and issues no new commit. + assert client.create_commit.call_count == 1 + + +# ------------------------------------------------------------ progress events + + +def test_progress_callback_reports_every_batch(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr(upload_module, "UPLOAD_COMMIT_BATCH_MAX_OPERATIONS", 2) + monkeypatch.setattr(upload_module, "UPLOAD_ADAPTIVE_BATCHING_ENABLED", False) + manager, _ = _make_manager() + for index in range(4): + (tmp_path / f"file-{index}.txt").write_bytes(bytes([index]) * (index + 1)) + events: list[dict] = [] + + manager.upload_folder( + repo_id="owner/repo", + repo_type="model", + folder_path=tmp_path, + max_workers=1, + use_cache=False, + disable_tqdm=True, + progress_callback=events.append, + ) + + commits = [e for e in events if e["event"] == "batch_committed"] + assert [e["batch_index"] for e in commits] == [0, 1] + assert commits[-1]["committed_files"] == 4 + assert commits[-1]["num_batches"] == 2 + assert all(e["total_files"] == 4 for e in commits) + + +def test_progress_events_carry_byte_counts(tmp_path: Path, monkeypatch) -> None: + # A consumer that only learns file counts cannot derive a throughput rate or + # an ETA, which is most of what a progress feed is for. + monkeypatch.setattr(upload_module, "UPLOAD_COMMIT_BATCH_MAX_OPERATIONS", 2) + monkeypatch.setattr(upload_module, "UPLOAD_ADAPTIVE_BATCHING_ENABLED", False) + manager, _ = _make_manager() + expected_total = 0 + for index in range(4): + payload = bytes([index]) * (1024 * (index + 1)) + (tmp_path / f"file-{index}.txt").write_bytes(payload) + expected_total += len(payload) + events: list[dict] = [] + + manager.upload_folder( + repo_id="owner/repo", + repo_type="model", + folder_path=tmp_path, + max_workers=1, + use_cache=False, + disable_tqdm=True, + progress_callback=events.append, + ) + + commits = [e for e in events if e["event"] == "batch_committed"] + assert all(e["total_bytes"] == expected_total for e in commits) + assert sum(e["batch_bytes"] for e in commits) == expected_total + assert commits[-1]["committed_bytes"] == expected_total + # Cumulative counters advance monotonically alongside the per-batch ones. + assert commits[0]["committed_bytes"] == commits[0]["batch_bytes"] + # These files ride inline, so all of their volume is attributed to the commit + # rather than to an earlier object-storage transfer. + assert sum(e["batch_inline_bytes"] for e in commits) == expected_total + + +def test_wire_progress_tracks_object_storage_transfers(tmp_path: Path, monkeypatch) -> None: + # Commits land in lumps, so a rate built from them alone reads as a spike + # followed by zero. Blob uploads finish continuously and are the honest + # source; only bytes that really moved are counted. + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 0) + monkeypatch.setattr(upload_module, "UPLOAD_PROGRESS_MIN_INTERVAL_SECONDS", 0) + manager, _ = _make_manager() + expected_total = 0 + for index in range(6): + payload = bytes([index]) * (2048 * (index + 1)) + (tmp_path / f"blob-{index}.dat").write_bytes(payload) + expected_total += len(payload) + events: list[dict] = [] + + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + max_workers=2, + use_cache=False, + disable_tqdm=True, + progress_callback=events.append, + ) + + wire = [e for e in events if e["event"] == "upload_progress"] + assert wire, "per-file transfers must report progress before the commit lands" + # Cumulative and never decreasing. + assert [e["uploaded_bytes"] for e in wire] == sorted(e["uploaded_bytes"] for e in wire) + assert wire[-1]["uploaded_bytes"] == expected_total + assert wire[-1]["uploaded_files"] == 6 + # Deltas partition the total exactly, so a consumer can sum them into a rate. + assert sum(e["uploaded_bytes_delta"] for e in wire) == expected_total + # Nothing rode inline, so the commit attributes no wire traffic to itself. + commits = [e for e in events if e["event"] == "batch_committed"] + assert sum(e["batch_inline_bytes"] for e in commits) == 0 + + +def test_reused_blobs_report_no_wire_bytes(tmp_path: Path, monkeypatch) -> None: + # A deduplicated blob transfers nothing; counting its size would inflate the + # reported throughput above what the link actually carried. + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 0) + monkeypatch.setattr(upload_module, "UPLOAD_PROGRESS_MIN_INTERVAL_SECONDS", 0) + manager, client = _make_manager() + client.validate_blobs.side_effect = lambda **kwargs: {obj["oid"]: None for obj in kwargs["objects"]} + for index in range(4): + (tmp_path / f"blob-{index}.dat").write_bytes(bytes([index]) * 4096) + events: list[dict] = [] + + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + max_workers=2, + use_cache=False, + disable_tqdm=True, + progress_callback=events.append, + ) + + client.upload_blob.assert_not_called() + wire = [e for e in events if e["event"] == "upload_progress"] + assert all(e["uploaded_bytes"] == 0 for e in wire) + assert wire[-1]["uploaded_files"] == 4 + # The files are still committed, and the commit still reports their volume. + commits = [e for e in events if e["event"] == "batch_committed"] + assert commits[-1]["committed_bytes"] == 4 * 4096 + + +def test_recovered_files_are_reported_so_totals_stay_complete(tmp_path: Path, monkeypatch) -> None: + # Recovery is exactly when an operator is watching. An 8 GiB run that lost one + # 512-file batch to a rejected commit put all 40000 files on the Hub but left + # 91 MB missing from done_bytes, because the recovery path emitted nothing. + monkeypatch.setattr(upload_module, "UPLOAD_COMMIT_BATCH_MAX_OPERATIONS", 2) + monkeypatch.setattr(upload_module, "UPLOAD_ADAPTIVE_BATCHING_ENABLED", False) + monkeypatch.setattr(upload_module.time, "sleep", lambda _s: None) + manager, client = _make_manager() + expected_total = 0 + for index in range(4): + payload = bytes([index]) * (1024 * (index + 1)) + (tmp_path / f"file-{index}.txt").write_bytes(payload) + expected_total += len(payload) + events: list[dict] = [] + + # Exhaust every in-commit attempt for the first batch so it really falls + # through to the recovery path instead of being absorbed earlier. The attempt + # budget is a default argument, so it is spelled out rather than patched. + attempts = upload_module.UPLOAD_COMMIT_MAX_ATTEMPTS + calls = {"n": 0} + + def create_commit(**kwargs): + calls["n"] += 1 + if calls["n"] <= attempts: + raise NetworkError("commit rejected") + return {"ok": True} + + client.create_commit.side_effect = create_commit + + manager.upload_folder( + repo_id="owner/repo", + repo_type="model", + folder_path=tmp_path, + max_workers=1, + use_cache=False, + disable_tqdm=True, + progress_callback=events.append, + ) + + recovery = [e for e in events if e["event"] == "recovery_committed"] + assert recovery, "a recovered commit must report its progress" + accounted = [e for e in events if e["event"] in ("batch_committed", "recovery_committed")] + # Every byte is accounted for exactly once across the two commit paths. + assert sum(e["batch_bytes"] for e in accounted) == expected_total + assert accounted[-1]["committed_bytes"] == expected_total + assert accounted[-1]["committed_files"] == 4 + + +def test_progress_callback_failure_does_not_abort_the_upload(tmp_path: Path) -> None: + manager, client = _make_manager() + (tmp_path / "README.md").write_bytes(b"hello") + + def explode(_event: dict) -> None: + raise RuntimeError("reporter is broken") + + manager.upload_folder( + repo_id="owner/repo", + repo_type="model", + folder_path=tmp_path, + max_workers=1, + use_cache=False, + disable_tqdm=True, + progress_callback=explode, + ) + + client.create_commit.assert_called_once() + + +# ------------------------------------------------------------ commit throttle + + +def test_commit_honors_server_retry_after(monkeypatch) -> None: + manager, client = _make_manager() + slept: list[float] = [] + monkeypatch.setattr(upload_module.time, "sleep", slept.append) + client.create_commit.side_effect = [ + RateLimitError("commit budget exhausted", retry_after=42), + {"ok": True}, + ] + + result = manager._commit_with_retry( + repo_id="owner/repo", + repo_type="model", + operations=[{"action": "create", "path": "a.txt"}], + commit_message="retry", + ) + + assert result == {"ok": True} + assert slept == [42.0] + + +def test_commit_retry_after_above_the_ceiling_aborts(monkeypatch) -> None: + manager, client = _make_manager() + slept: list[float] = [] + monkeypatch.setattr(upload_module.time, "sleep", slept.append) + monkeypatch.setattr(upload_module, "UPLOAD_COMMIT_MAX_RETRY_AFTER_SECONDS", 60) + client.create_commit.side_effect = RateLimitError("commit budget exhausted", retry_after=3600) + + with pytest.raises(RateLimitError): + manager._commit_with_retry( + repo_id="owner/repo", + repo_type="model", + operations=[{"action": "create", "path": "a.txt"}], + commit_message="retry", + ) + + assert slept == [] + + +def test_throttled_wait_does_not_consume_the_transient_error_budget(monkeypatch) -> None: + # A long, server-declared wait must not exhaust the allowance reserved for + # failures of unknown duration, or a throttle would look like a hard error. + manager, client = _make_manager() + monkeypatch.setattr(upload_module, "UPLOAD_COMMIT_RETRY_TOTAL_WAIT_SECONDS", 10) + monkeypatch.setattr(upload_module.time, "sleep", lambda _s: None) + client.create_commit.side_effect = [ + RateLimitError("throttled", retry_after=600), + RateLimitError("throttled", retry_after=600), + {"ok": True}, + ] + + assert manager._commit_with_retry( + repo_id="owner/repo", + repo_type="model", + operations=[{"action": "create", "path": "a.txt"}], + commit_message="retry", + max_attempts=4, + ) == {"ok": True} + + +def test_commit_rate_governor_paces_once_the_budget_is_spent(monkeypatch) -> None: + governor = upload_module._CommitRateGovernor(2) + clock = {"now": 0.0} + slept: list[float] = [] + monkeypatch.setattr(upload_module.time, "monotonic", lambda: clock["now"]) + + def fake_sleep(seconds: float) -> None: + slept.append(seconds) + clock["now"] += seconds + + monkeypatch.setattr(upload_module.time, "sleep", fake_sleep) + + assert governor.acquire() == 0.0 + assert governor.acquire() == 0.0 + waited = governor.acquire() + + assert slept, "the third commit within the window must wait" + assert waited == pytest.approx(3600.0, abs=1.0) + + +def test_commit_rate_governor_is_inert_when_disabled() -> None: + governor = upload_module._CommitRateGovernor(0) + + assert not governor.enabled + assert [governor.acquire() for _ in range(50)] == [0.0] * 50 + + +def test_every_commit_path_goes_through_the_shared_governor(tmp_path: Path, monkeypatch) -> None: + # The budget belongs to the repository, not to one call. Pacing only the + # happy path would leave recovery rounds free to hammer a throttled server. + monkeypatch.setattr(upload_module, "UPLOAD_COMMIT_MAX_PER_HOUR", 500) + monkeypatch.setattr(upload_module, "UPLOAD_COMMIT_BATCH_MAX_OPERATIONS", 2) + monkeypatch.setattr(upload_module, "UPLOAD_ADAPTIVE_BATCHING_ENABLED", False) + manager, client = _make_manager() + assert manager._commit_governor.enabled + + acquired: list[int] = [] + original_acquire = manager._commit_governor.acquire + monkeypatch.setattr( + manager._commit_governor, + "acquire", + lambda: (acquired.append(1), original_acquire())[1], + ) + + for index in range(4): + (tmp_path / f"file-{index}.txt").write_bytes(bytes([index])) + + manager.upload_folder( + repo_id="owner/repo", + repo_type="model", + folder_path=tmp_path, + max_workers=1, + use_cache=False, + disable_tqdm=True, + ) + manager.upload_file( + repo_id="owner/repo", + repo_type="model", + path_or_fileobj=b"single", + path_in_repo="extra.txt", + disable_tqdm=True, + ) + + # Two batch commits plus the single-file commit. + assert client.create_commit.call_count == 3 + assert len(acquired) == 3 + + +def test_upload_file_retries_a_transient_commit_failure(monkeypatch) -> None: + manager, client = _make_manager() + monkeypatch.setattr(upload_module.time, "sleep", lambda _s: None) + client.create_commit.side_effect = [ + NetworkError("commit temporarily unavailable"), + {"ok": True}, + ] + + result = manager.upload_file( + repo_id="owner/repo", + repo_type="model", + path_or_fileobj=b"hello", + path_in_repo="README.md", + disable_tqdm=True, + ) + + assert result == {"ok": True} + assert client.create_commit.call_count == 2 + + +def test_upload_file_honors_retry_after(monkeypatch) -> None: + manager, client = _make_manager() + slept: list[float] = [] + monkeypatch.setattr(upload_module.time, "sleep", slept.append) + client.create_commit.side_effect = [ + RateLimitError("commit budget exhausted", retry_after=7), + {"ok": True}, + ] + + manager.upload_file( + repo_id="owner/repo", + repo_type="model", + path_or_fileobj=b"hello", + path_in_repo="README.md", + disable_tqdm=True, + ) + + assert slept == [7.0] + + +def test_upload_file_still_fails_fast_on_a_permanent_error() -> None: + manager, client = _make_manager() + client.create_commit.side_effect = InvalidParameter("path is not allowed") + + with pytest.raises(InvalidParameter): + manager.upload_file( + repo_id="owner/repo", + repo_type="model", + path_or_fileobj=b"hello", + path_in_repo="README.md", + disable_tqdm=True, + ) + + assert client.create_commit.call_count == 1 + + +def test_prepare_upload_folder_reports_sizes_for_reuse(tmp_path: Path) -> None: + # Planning needs these sizes immediately afterwards; re-stating the tree + # costs one syscall per file and yields nothing new. + manager, _ = _make_manager() + (tmp_path / "a.txt").write_bytes(b"x" * 10) + nested = tmp_path / "sub" + nested.mkdir() + (nested / "b.txt").write_bytes(b"y" * 20) + + sizes: dict[str, int] = {} + prepared = manager._prepare_upload_folder( + folder_path=tmp_path, + path_in_repo="", + repo_type="model", + sizes_out=sizes, + ) + + assert sizes == {str(tmp_path / "a.txt"): 10, str(nested / "b.txt"): 20} + assert {path for _, path in prepared} == set(sizes) diff --git a/tests/test_upload_blob_dedup.py b/tests/test_upload_blob_dedup.py new file mode 100644 index 0000000..13c79fe --- /dev/null +++ b/tests/test_upload_blob_dedup.py @@ -0,0 +1,256 @@ +"""Identical content inside one upload must be transferred once. + +The batch pre-sign step asks the server about every distinct oid *before* any +upload starts, so the server cannot answer "already stored" for a duplicate -- +it is not stored yet. Verified against the Hub: pre-signing the same fresh oid +twice yields two upload URLs, and only after the blob lands does the endpoint +report it as existing. Without an owner election, every occurrence therefore +PUTs the same bytes. +""" + +from __future__ import annotations + +import threading +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +import modelscope_hub._upload as upload_module +from modelscope_hub._upload import DUPLICATE_BLOB, UploadManager +from modelscope_hub.errors import NetworkError, StorageError + + +def _make_manager() -> tuple[UploadManager, MagicMock, list[str]]: + """Manager whose blob PUTs are recorded by oid.""" + client = MagicMock() + client.create_commit.return_value = {"ok": True} + put_oids: list[str] = [] + lock = threading.Lock() + + def validate_blobs(*, repo_id: str, repo_type: str, objects: list[dict]) -> dict[str, str]: + return {obj["oid"]: f"https://upload/{obj['oid']}" for obj in objects} + + def upload_blob(*, upload_url: str, data, size: int) -> None: + with lock: + put_oids.append(upload_url.rsplit("/", 1)[-1]) + while data.read(1024 * 1024): + pass + + client.validate_blobs.side_effect = validate_blobs + client.upload_blob.side_effect = upload_blob + return UploadManager(client, MagicMock()), client, put_oids + + +def _write_duplicates(root: Path, groups: dict[str, int], size: int = 4096) -> int: + """Write ``{content_tag: copies}``; returns the number of distinct contents.""" + index = 0 + for tag, copies in groups.items(): + payload = tag.encode() * (size // len(tag.encode())) + for _ in range(copies): + (root / f"file-{index:04d}.dat").write_bytes(payload) + index += 1 + return len(groups) + + +def test_duplicate_content_is_uploaded_once(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 0) + manager, client, put_oids = _make_manager() + distinct = _write_duplicates(tmp_path, {"aa": 5, "bb": 3, "cc": 1}) + total_files = 9 + + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + max_workers=4, + use_cache=False, + disable_tqdm=True, + ) + + assert len(put_oids) == distinct, f"expected one PUT per distinct content, got {len(put_oids)}" + assert len(set(put_oids)) == distinct + # All nine files are still committed, each as an LFS pointer. + operations = client.create_commit.call_args.kwargs["operations"] + assert len(operations) == total_files + assert all(op["type"] == "lfs" and op["sha256"] for op in operations) + # The pointers cover exactly the distinct contents that were uploaded. + assert {op["sha256"] for op in operations} == set(put_oids) + + +def test_duplicates_report_no_wire_bytes(tmp_path: Path, monkeypatch) -> None: + # Only the owner moves bytes, so throughput must not count the copies. + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 0) + monkeypatch.setattr(upload_module, "UPLOAD_PROGRESS_MIN_INTERVAL_SECONDS", 0) + manager, _, put_oids = _make_manager() + _write_duplicates(tmp_path, {"aa": 4, "bb": 4}, size=8192) + events: list[dict] = [] + + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + max_workers=4, + use_cache=False, + disable_tqdm=True, + progress_callback=events.append, + ) + + wire = [e for e in events if e["event"] == "upload_progress"] + assert len(put_oids) == 2 + # Two owners at 8192 bytes each; the six copies contribute nothing. + assert wire[-1]["uploaded_bytes"] == 2 * 8192 + assert wire[-1]["uploaded_files"] == 8 + + +def test_owner_is_the_earliest_file_so_no_commit_precedes_its_blob(tmp_path: Path, monkeypatch) -> None: + # Correctness rests on the owner never landing in a later batch than a copy: + # batches commit in order, so a copy in batch j needs its owner in batch <= j. + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 0) + monkeypatch.setattr(upload_module, "UPLOAD_COMMIT_BATCH_MAX_OPERATIONS", 2) + monkeypatch.setattr(upload_module, "UPLOAD_ADAPTIVE_BATCHING_ENABLED", False) + manager, client, put_oids = _make_manager() + # Same content in the first and last file: owner in batch 0, copy in batch 2. + payload = b"x" * 4096 + for name in ("a.dat", "b.dat", "c.dat", "d.dat", "e.dat"): + (tmp_path / name).write_bytes(payload if name in ("a.dat", "e.dat") else name.encode() * 512) + + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + max_workers=1, + use_cache=False, + disable_tqdm=True, + ) + + # Four distinct contents for five files. + assert len(put_oids) == 4 + committed = [op["path"] for call in client.create_commit.call_args_list for op in call.kwargs["operations"]] + assert committed == ["a.dat", "b.dat", "c.dat", "d.dat", "e.dat"] + + +def test_a_copy_is_not_committed_when_its_owner_fails(tmp_path: Path, monkeypatch) -> None: + # The whole point of the owner election is that copies skip the transfer. + # If the owner fails, committing a copy would publish an LFS pointer to a + # blob that was never stored. + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 0) + monkeypatch.setattr(upload_module, "UPLOAD_RECOVERY_ENABLED", False) + monkeypatch.setattr(upload_module, "UPLOAD_FAILED_FILE_MAX_RETRY_ROUNDS", 0) + manager, client, _ = _make_manager() + payload = b"shared" * 700 + for name in ("dup-0.dat", "dup-1.dat", "dup-2.dat"): + (tmp_path / name).write_bytes(payload) + (tmp_path / "solo.dat").write_bytes(b"solo" * 900) + + def failing_upload(*, upload_url: str, data, size: int) -> None: + if size == len(payload): + raise NetworkError("owner transfer failed") + while data.read(1024 * 1024): + pass + + client.upload_blob.side_effect = failing_upload + monkeypatch.setattr(upload_module.time, "sleep", lambda _s: None) + + # The owner and both copies count as failures, so the upload reports itself + # as failed rather than quietly publishing broken pointers. + with pytest.raises(StorageError, match="3 file"): + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + max_workers=1, + use_cache=False, + disable_tqdm=True, + ) + + committed = {op["path"] for call in client.create_commit.call_args_list for op in call.kwargs["operations"]} + # Only the independent file is committed; no copy of the failed content is. + assert committed == {"solo.dat"} + assert not any(path.startswith("dup-") for path in committed) + + +def test_globally_existing_content_needs_no_owner(tmp_path: Path, monkeypatch) -> None: + # When the server already holds the blob, nobody transfers it and the copies + # are not deferred either. + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 0) + manager, client, put_oids = _make_manager() + client.validate_blobs.side_effect = lambda **kw: {obj["oid"]: None for obj in kw["objects"]} + _write_duplicates(tmp_path, {"aa": 3, "bb": 2}) + + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + max_workers=2, + use_cache=False, + disable_tqdm=True, + ) + + assert put_oids == [] + operations = client.create_commit.call_args.kwargs["operations"] + assert len(operations) == 5 + + +def test_inline_files_are_not_deduplicated(tmp_path: Path, monkeypatch) -> None: + # Inline content travels in the commit body, not as a blob, so there is no + # transfer to skip and every occurrence must carry its own bytes. + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 1024 * 1024) + manager, client, put_oids = _make_manager() + for name in ("a.txt", "b.txt", "c.txt"): + (tmp_path / name).write_bytes(b"identical") + + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + max_workers=2, + use_cache=False, + disable_tqdm=True, + ) + + assert put_oids == [] + operations = client.create_commit.call_args.kwargs["operations"] + assert len(operations) == 3 + assert all(op["type"] == "normal" and op["content"] for op in operations) + + +def test_duplicate_marker_skips_the_transfer_without_claiming_global_reuse() -> None: + manager, client, put_oids = _make_manager() + + result = manager._upload_blob( + repo_id="owner/repo", + repo_type="dataset", + sha256="a" * 64, + size=123, + data=b"x" * 123, + disable_tqdm=True, + pre_validated=DUPLICATE_BLOB, + ) + + assert put_oids == [] + client.validate_blobs.assert_not_called() + assert result["is_uploaded"] is True + assert result["is_reused"] is True + # Nothing went over the wire, so throughput accounting must see zero. + assert result["is_blob_uploaded"] is False + + +@pytest.mark.parametrize("workers", [1, 4, 16]) +def test_dedup_holds_at_any_concurrency(tmp_path: Path, monkeypatch, workers: int) -> None: + # Election happens before any worker starts, so the outcome must not depend + # on scheduling -- and no worker may ever block on another. + monkeypatch.setattr(upload_module, "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", 0) + manager, _, put_oids = _make_manager() + distinct = _write_duplicates(tmp_path, {"aa": 8, "bb": 8, "cc": 8, "dd": 1}) + + manager.upload_folder( + repo_id="owner/repo", + repo_type="dataset", + folder_path=tmp_path, + max_workers=workers, + use_cache=False, + disable_tqdm=True, + ) + + assert len(put_oids) == distinct diff --git a/tests/test_upload_config.py b/tests/test_upload_config.py index 6176727..acaf9ab 100644 --- a/tests/test_upload_config.py +++ b/tests/test_upload_config.py @@ -280,6 +280,146 @@ def test_dead_lfs_threshold_is_not_registered() -> None: assert "UPLOAD_LFS_THRESHOLD" not in names +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("32KiB", 32 * 1024), + ("512K", 512 * 1024), + ("65536", 65536), + ("1MiB", 1024 * 1024), + ("2MB", 2 * 1000 * 1000), + ("0", 0), + ], +) +def test_lfs_threshold_accepts_byte_sizes_and_unit_suffixes(raw: str, expected: int) -> None: + # A megabyte-only knob cannot express the sub-MB range that decides whether + # small files ride inline in a commit or go to object storage. + result = _run_constants( + "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", + env={"MODELSCOPE_UPLOAD_LFS_FORCE_THRESHOLD": raw}, + ) + assert json.loads(result.stdout)["UPLOAD_LFS_FORCE_THRESHOLD_BYTES"] == expected + assert "FutureWarning" not in result.stderr + + +def test_deprecated_lfs_threshold_mb_alias_keeps_megabyte_units() -> None: + result = _run_constants( + "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", + env={"MODELSCOPE_UPLOAD_LFS_FORCE_THRESHOLD_MB": "4"}, + ) + assert json.loads(result.stdout)["UPLOAD_LFS_FORCE_THRESHOLD_BYTES"] == 4 * 1024 * 1024 + assert "MODELSCOPE_UPLOAD_LFS_FORCE_THRESHOLD_MB" in result.stderr + assert "expects bytes" in result.stderr + + +def test_canonical_lfs_threshold_wins_over_deprecated_aliases() -> None: + result = _run_constants( + "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", + env={ + "MODELSCOPE_UPLOAD_LFS_FORCE_THRESHOLD": "32KiB", + "MODELSCOPE_UPLOAD_LFS_FORCE_THRESHOLD_MB": "4", + "UPLOAD_LFS_ENFORCE_THRESHOLD": "999", + }, + ) + assert json.loads(result.stdout)["UPLOAD_LFS_FORCE_THRESHOLD_BYTES"] == 32 * 1024 + assert "FutureWarning" not in result.stderr + + +@pytest.mark.parametrize( + ("env", "constant", "expected", "needle"), + [ + ( + {"MODELSCOPE_UPLOAD_LFS_FORCE_THRESHOLD": "32ZB"}, + "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", + 1024 * 1024, + "unknown size unit", + ), + ( + {"MODELSCOPE_UPLOAD_LFS_FORCE_THRESHOLD": "-1"}, + "UPLOAD_LFS_FORCE_THRESHOLD_BYTES", + 1024 * 1024, + "must not be negative", + ), + ( + {"MODELSCOPE_UPLOAD_COMMIT_BATCH_MAX_OPERATIONS": "lots"}, + "UPLOAD_COMMIT_BATCH_MAX_OPERATIONS", + 256, + "not an integer", + ), + ( + {"MODELSCOPE_UPLOAD_COMMIT_BATCH_MAX_OPERATIONS": "0"}, + "UPLOAD_COMMIT_BATCH_MAX_OPERATIONS", + 256, + "must be a positive integer", + ), + ( + {"MODELSCOPE_UPLOAD_COMMIT_MAX_INLINE_BYTES": "8bogus"}, + "UPLOAD_COMMIT_MAX_INLINE_BYTES", + 8 * 1024 * 1024, + "unknown size unit", + ), + ], +) +def test_invalid_upload_env_values_warn_instead_of_silently_reverting( + env: dict[str, str], + constant: str, + expected: int, + needle: str, +) -> None: + # Silently falling back made misconfiguration invisible: the value had no + # effect and nothing said so. + result = _run_constants(constant, env=env) + assert json.loads(result.stdout)[constant] == expected + assert needle in result.stderr + assert "is invalid" in result.stderr + + +def test_commit_rate_governor_accepts_explicit_zero_without_warning() -> None: + result = _run_constants( + "UPLOAD_COMMIT_MAX_PER_HOUR", + env={"MODELSCOPE_UPLOAD_COMMIT_MAX_PER_HOUR": "0"}, + ) + assert json.loads(result.stdout)["UPLOAD_COMMIT_MAX_PER_HOUR"] == 0 + assert "is invalid" not in result.stderr + + +def test_new_upload_knobs_are_registered_under_upload() -> None: + result = _run_constants("ENV_REGISTRY") + registry = {item["name"]: item for item in json.loads(result.stdout)["ENV_REGISTRY"]} + for name in ( + "MODELSCOPE_UPLOAD_LFS_FORCE_THRESHOLD", + "MODELSCOPE_UPLOAD_COMMIT_MAX_INLINE_BYTES", + "MODELSCOPE_UPLOAD_COMMIT_MAX_PER_HOUR", + "MODELSCOPE_UPLOAD_COMMIT_MAX_RETRY_AFTER_SECONDS", + "MODELSCOPE_UPLOAD_INLINE_METADATA_PATHS", + ): + assert registry[name]["category"] == "Upload" + assert registry["MODELSCOPE_UPLOAD_LFS_FORCE_THRESHOLD"]["default"] == "1MiB" + assert "MODELSCOPE_UPLOAD_LFS_FORCE_THRESHOLD_MB" not in registry + + +def test_http_sessions_size_the_pool_for_concurrent_workers() -> None: + # urllib3 defaults the pool to 10. A folder upload runs max_workers requests + # at once, so the default discarded the excess connections and paid for a + # fresh TLS handshake on their next use. + from modelscope_hub.constants import API_CONNECTION_POOL_MAXSIZE + + api = HubApi(config=HubConfig(token="ms-test")) + for session in (api.legacy._session, api.openapi._session): + adapter = session.get_adapter("https://modelscope.cn") + pool_kw = adapter.poolmanager.connection_pool_kw + assert pool_kw["maxsize"] == API_CONNECTION_POOL_MAXSIZE + assert API_CONNECTION_POOL_MAXSIZE >= 16 + + +def test_connection_pool_size_is_env_tunable() -> None: + result = _run_constants( + "API_CONNECTION_POOL_MAXSIZE", + env={"MODELSCOPE_API_CONNECTION_POOL_MAXSIZE": "64"}, + ) + assert json.loads(result.stdout)["API_CONNECTION_POOL_MAXSIZE"] == 64 + + def _response(data: dict | None = None) -> MagicMock: response = MagicMock() response.status_code = 200 diff --git a/tests/test_upload_lfs_gate.py b/tests/test_upload_lfs_gate.py index 6532b8c..c01c16a 100644 --- a/tests/test_upload_lfs_gate.py +++ b/tests/test_upload_lfs_gate.py @@ -80,17 +80,14 @@ def test_delete_files_rejects_empty_paths() -> None: manager, client = _make_manager() with pytest.raises(InvalidParameter, match="at least one"): - manager.delete_files( - repo_id="owner/repo", repo_type="model", file_paths=["", ""]) + manager.delete_files(repo_id="owner/repo", repo_type="model", file_paths=["", ""]) client.create_commit.assert_not_called() @pytest.mark.parametrize("repo_type", ["model", "dataset"]) @pytest.mark.parametrize("file_paths", [["", "old.bin"], "old.bin"]) -def test_hub_api_delete_files_delegates_to_upload_manager( - repo_type: str, file_paths: list[str] | str -) -> None: +def test_hub_api_delete_files_delegates_to_upload_manager(repo_type: str, file_paths: list[str] | str) -> None: api = HubApi(token="test-token") api._uploader = MagicMock() api._uploader.delete_files.return_value = {"deleted_files": ["old.bin"]} @@ -128,7 +125,8 @@ def test_hub_api_delete_patterns_resolve_remote_paths(repo_type: str) -> None: SimpleNamespace(path="nested/metadata.json", type="blob"), SimpleNamespace(path="weights.bin", type="blob"), SimpleNamespace(path="nested", type="tree"), - ]) + ] + ) result = api.delete_files( "owner/repo", @@ -138,8 +136,7 @@ def test_hub_api_delete_patterns_resolve_remote_paths(repo_type: str) -> None: revision="main", ) - api.list_repo_files.assert_called_once_with( - "owner/repo", repo_type, revision="main", recursive=True) + api.list_repo_files.assert_called_once_with("owner/repo", repo_type, revision="main", recursive=True) api._uploader.delete_files.assert_called_once_with( repo_id="owner/repo", repo_type=repo_type, @@ -153,11 +150,9 @@ def test_hub_api_delete_patterns_resolve_remote_paths(repo_type: str) -> None: def test_hub_api_delete_patterns_with_no_match_is_noop() -> None: api = HubApi(token="test-token") api._uploader = MagicMock() - api.list_repo_files = MagicMock( - return_value=[SimpleNamespace(path="weights.bin", type="blob")]) + api.list_repo_files = MagicMock(return_value=[SimpleNamespace(path="weights.bin", type="blob")]) - result = api.delete_files( - "owner/repo", "model", delete_patterns="*.json") + result = api.delete_files("owner/repo", "model", delete_patterns="*.json") api._uploader.delete_files.assert_not_called() assert result == {