diff --git a/AGENTS.md b/AGENTS.md index 2b94ee3..995d2e3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -138,3 +138,12 @@ After changing the upload path, check whether the following must be updated as w - The upload environment-variable table and semantic documentation in `README.md`. - `__all__` exports and compatibility aliases in `src/modelscope_hub/compat/` and `constants.py`. - This file, if new constants, error codes, or scenarios were introduced. + +--- + +## 7. User-Agent (UA) Requirements + +- Every SDK HTTP request must include a `User-Agent` starting with `modelscope_hub/`. Use `utils.build_user_agent` / `append_user_agent` and enforce coverage at client send boundaries across OpenAPI, Legacy, Agent, both sessions, anonymous requests, retries, and probes. +- Custom UA information must be appended, never replacing the SDK identifier. Handle header names case-insensitively; empty values must not remove the SDK UA. Propagate CLI / compat markers (`invoked_by/cli` / `invoked_by/compat`) to derived clients and download requests. +- Prefer a persisted `session_id`; fall back to a process-stable ID when storage is not writable. Telemetry persistence failures must not block business requests. UA must not contain tokens, cookies, or other credentials. +- Changes to HTTP transport, client construction, or delegating entry points must verify these invariants against actual `PreparedRequest` objects; checking session default headers alone is insufficient. diff --git a/README.md b/README.md index 4ae6df1..12be847 100644 --- a/README.md +++ b/README.md @@ -211,6 +211,8 @@ snapshot = api.download_repo( path = api.download_file("Qwen/Qwen3-0.6B", "model", "config.json", local_files_only=True) ``` +Downloads validate available size and SHA-256 metadata before replacing files, with safe resume, bounded retries, timeouts, and cancellation. Incomplete snapshots raise `FileDownloadError` (CLI exit code `1`), preserving completed files for subsequent retries. + ### Upload ```bash @@ -905,7 +907,11 @@ Token is persisted locally after `ms-hub login` and auto-loaded in subsequent se |----------|---------|-------------| | `MODELSCOPE_API_TIMEOUT` | `60` | HTTP request timeout (seconds) | | `MODELSCOPE_API_CONNECT_TIMEOUT` | `10` | HTTP connect timeout (seconds) | -| `MODELSCOPE_API_MAX_RETRIES` | `5` | Max retry attempts for transient failures | +| `MODELSCOPE_API_MAX_RETRIES` | `5` | Regular API retries; downloads and file listings use separate retry policies | +| `MODELSCOPE_REPO_TREE_PAGE_TIMEOUT_SECONDS` | `15` | Per-page listing timeout (seconds) | +| `MODELSCOPE_REPO_TREE_PAGE_MAX_ATTEMPTS` | `4` | Maximum attempts per listing page, including the initial request | +| `MODELSCOPE_REPO_TREE_PAGE_RETRY_MAX_DELAY_SECONDS` | `8` | Maximum listing retry delay (seconds) | +| `MODELSCOPE_REPO_TREE_LIST_TIMEOUT_SECONDS` | `3600` | Total repository listing time budget (seconds) | **Download:** @@ -915,8 +921,13 @@ Token is persisted locally after `ms-hub login` and auto-loaded in subsequent se | `MODELSCOPE_DOWNLOAD_PARALLEL_THRESHOLD_MB` | `500` | Parallel download threshold (MB) | | `MODELSCOPE_DOWNLOAD_CHUNK_SIZE_MB` | `1` | Streaming chunk size (MB) | | `MODELSCOPE_DOWNLOAD_PART_SIZE_MB` | `160` | Parallel range chunk size (MB) | -| `MODELSCOPE_DOWNLOAD_MAX_RETRIES` | `5` | Per-file download retry count | -| `MODELSCOPE_DOWNLOAD_TIMEOUT` | `60` | Per-file download timeout (seconds) | +| `MODELSCOPE_DOWNLOAD_MAX_RETRIES` | `5` | Single-layer transport retries; `0` disables them; integrity retries are separately bounded | +| `MODELSCOPE_DOWNLOAD_TIMEOUT` | `60` | Idle read timeout for streams and range parts, not the whole file (seconds) | +| `MODELSCOPE_DOWNLOAD_TOTAL_TIMEOUT` | `3600` | Per-file time budget, including locks, probes, transfers, retries, and validation (seconds) | +| `MODELSCOPE_DOWNLOAD_MIN_BYTES_PER_SECOND` | `1024` | Minimum rate per connection (B/s); `0` disables low-speed detection | +| `MODELSCOPE_DOWNLOAD_LOW_SPEED_WINDOW_SECONDS` | `60` | Minimum-rate observation window (seconds) | +| `MODELSCOPE_DOWNLOAD_RETRY_MAX_DELAY_SECONDS` | `60` | Maximum download Retry-After or backoff delay (seconds) | +| `MODELSCOPE_DOWNLOAD_FILE_LOCK_TIMEOUT_SECONDS` | `600` | Maximum wait for another process's file lock (seconds) | | `MODELSCOPE_DOWNLOAD_FILE_LOCK` | `true` | File lock for multiprocess download safety | | `MODELSCOPE_DOWNLOAD_INTRA_CLOUD` | `true` | Alibaba cloud intra-cloud acceleration | | `MODELSCOPE_DOWNLOAD_INTRA_CLOUD_REGION` | (auto) | Override intra-cloud region ID | diff --git a/src/modelscope_hub/__init__.py b/src/modelscope_hub/__init__.py index 46f2d95..4db8011 100644 --- a/src/modelscope_hub/__init__.py +++ b/src/modelscope_hub/__init__.py @@ -26,7 +26,9 @@ CacheError, CacheNotFound, CorruptedCacheException, + FileDownloadError, FileIntegrityError, + FileSizeMismatchError, HubError, InvalidParameter, NetworkError, @@ -106,7 +108,9 @@ "CacheError", "CacheNotFound", "CorruptedCacheException", + "FileDownloadError", "FileIntegrityError", + "FileSizeMismatchError", "HubError", "InvalidParameter", "NetworkError", diff --git a/src/modelscope_hub/_download.py b/src/modelscope_hub/_download.py index bb5fab5..c1becfa 100644 --- a/src/modelscope_hub/_download.py +++ b/src/modelscope_hub/_download.py @@ -22,36 +22,54 @@ import errno import fnmatch import hashlib -import io import os import re import threading import time import uuid -from concurrent.futures import ThreadPoolExecutor, as_completed +from collections.abc import Callable +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, as_completed, wait from pathlib import Path from typing import TYPE_CHECKING from urllib.parse import urlparse import requests from tqdm.auto import tqdm +from urllib3.exceptions import MaxRetryError from urllib3.util.retry import Retry +from ._download_control import ( + _DownloadCancelledError, + _DownloadControl, + _DownloadDeadlineError, + _network_context, + _ReadHTTPAdapter, +) from .constants import ( DOWNLOAD_CHUNK_SIZE, + DOWNLOAD_FILE_LOCK_TIMEOUT_SECONDS, + DOWNLOAD_LOW_SPEED_WINDOW_SECONDS, + DOWNLOAD_MIN_BYTES_PER_SECOND, DOWNLOAD_PARALLEL_THRESHOLD, DOWNLOAD_PARALLELS, DOWNLOAD_PART_SIZE, + DOWNLOAD_RETRY_MAX_DELAY_SECONDS, DOWNLOAD_RETRY_TIMES, - DOWNLOAD_TIMEOUT, + DOWNLOAD_TOTAL_TIMEOUT, ENV_FILE_LOCK, ENV_INTER_CLOUD_REGIONS, ENV_INTRA_CLOUD_ACCELERATION, ENV_INTRA_CLOUD_REGION, + REPO_TREE_LIST_TIMEOUT_SECONDS, ) from .errors import ( + APIError, + AuthenticationError, CacheNotFound, + FileDownloadError, FileIntegrityError, + FileSizeMismatchError, + HubError, NetworkError, NotExistError, PermissionDeniedError, @@ -59,6 +77,7 @@ ) from .utils.file_utils import compute_hash, ensure_dir from .utils.logger import get_logger +from .utils.patterns import extract_common_prefix if TYPE_CHECKING: from ._legacy_api import LegacyClient @@ -67,6 +86,18 @@ logger = get_logger("download") DOWNLOAD_HASH_RETRY_TIMES = 3 +# Attempts for a file whose delivered size contradicts the repository listing. +# The server answered consistently, so one more try over the default route is +# enough to rule out a stale replica or CDN node before reporting the file. +DOWNLOAD_SIZE_RETRY_TIMES = 2 +# Seconds before the first integrity retry, doubled for each further one, so a +# transient CDN or replica inconsistency has time to settle. +DOWNLOAD_INTEGRITY_RETRY_BACKOFF = 1.0 + +# Response headers identifying the request and the node that served it. They +# are quoted in integrity errors so a broken object can be traced server-side. +_DIAGNOSTIC_HEADERS = ("x-oss-request-id", "x-request-id", "eagleid", "x-cache", "x-linked-etag") +_CONTENT_RANGE_RE = re.compile(r"bytes\s+(\d+)-(\d+)/(\d+|\*)", re.IGNORECASE) # --------------------------------------------------------------------------- @@ -120,7 +151,12 @@ def end(self) -> None: @contextlib.contextmanager -def _optional_file_lock(lock_path: Path | None, *, enabled: bool = True): +def _optional_file_lock( + lock_path: Path | None, + *, + enabled: bool = True, + control: _DownloadControl | None = None, +): """Acquire a file lock when *enabled*, otherwise no-op.""" if not enabled or lock_path is None: yield @@ -129,12 +165,18 @@ def _optional_file_lock(lock_path: Path | None, *, enabled: bool = True): from filelock import FileLock, SoftFileLock, Timeout lock_path.parent.mkdir(parents=True, exist_ok=True) - default_interval = 60 + default_interval = 0.2 if control is not None else 60 lock = FileLock(str(lock_path), timeout=default_interval) is_soft = False - waited = 0 + waited = 0.0 + started = time.monotonic() + next_notice = 60.0 while True: + if control is not None: + control.check() + if time.monotonic() - started >= DOWNLOAD_FILE_LOCK_TIMEOUT_SECONDS: + raise _DownloadDeadlineError(f"等待下载文件锁超过 {DOWNLOAD_FILE_LOCK_TIMEOUT_SECONDS}s：{lock_path}") try: lock.acquire(timeout=default_interval) except Timeout: @@ -156,10 +198,9 @@ def _optional_file_lock(lock_path: Path | None, *, enabled: bool = True): pass waited = 0 continue - logger.info( - "Still waiting to acquire lock on %s", - lock_path, - ) + if waited >= next_notice: + logger.info("Still waiting to acquire lock on %s", lock_path) + next_notice = waited + 60 except NotImplementedError as exc: if "use SoftFileLock instead" in str(exc): logger.warning( @@ -224,12 +265,148 @@ def _matches_patterns(path: str, patterns: list[str] | None) -> bool: return False +# --------------------------------------------------------------------------- +# Response validation +# --------------------------------------------------------------------------- +class _IncompleteBodyError(Exception): + """The body ended before the announced length arrived; safe to resume.""" + + +class _RangeNotHonouredError(Exception): + """A ranged request was answered with something other than that range.""" + + +def _listing_roots(patterns: list[str] | None) -> list[str | None]: + """Use only literal directory prefixes; keep local glob filtering as the final selection check.""" + if not patterns or any(".*" in p or "\\" in p or p.startswith("/") for p in patterns): + return [None] + roots = [extract_common_prefix([pattern]) for pattern in patterns] + if any(root is None or ".." in root.split("/") for root in roots): + return [None] + result: list[str | None] = [] + for root in sorted(set(roots), key=lambda value: (len(value or ""), value or "")): + if root is not None and not any(root == parent or root.startswith(f"{parent}/") for parent in result): + result.append(root) + return result + + +def _parse_content_range(value: str | None) -> tuple[int, int, int | None] | None: + """Parse ``Content-Range: bytes start-end/total``; *total* may be ``*``.""" + match = _CONTENT_RANGE_RE.fullmatch(value.strip()) if value else None + if match is None: + return None + total = None if match.group(3) == "*" else int(match.group(3)) + start, end = int(match.group(1)), int(match.group(2)) + if end < start or (total is not None and end >= total): + return None + return start, end, total + + +def _unsatisfied_range_total(value: str | None) -> int | None: + """Return the object size from a 416 ``Content-Range: bytes */total`` header.""" + total = (value or "").rpartition("/")[2].strip() + return int(total) if total.isdigit() else None + + +def _checkable_content_length(resp: requests.Response) -> int | None: + """Return the announced body length when it can be compared with the bytes read. + + Under a ``Content-Encoding`` the header counts encoded bytes while + ``iter_content`` yields decoded ones, so the two cannot be compared. + """ + raw = resp.headers.get("Content-Length") + if raw is None or resp.headers.get("Content-Encoding", "identity").strip().lower() not in ("", "identity"): + return None + try: + return int(raw) + except ValueError: + return None + + +def _describe_response(resp: requests.Response, request_id: str) -> str: + """Summarise a download response for diagnostics without leaking the signed URL.""" + headers = resp.headers + parts = [f"HTTP {resp.status_code}", f"Content-Length={headers.get('Content-Length', 'absent')}"] + if headers.get("Content-Range"): + parts.append(f"Content-Range={headers['Content-Range']}") + parts.append(f"host={urlparse(resp.url or '').hostname or 'unknown'}") + parts.append(f"X-Request-ID={request_id}") + parts.extend(f"{name}={headers[name]}" for name in _DIAGNOSTIC_HEADERS if headers.get(name)) + return ", ".join(parts) + + +def _incomplete_path(target: Path) -> Path: + """Return the partial-download path used while *target* is being fetched.""" + return target.with_suffix(target.suffix + ".incomplete") + + +def _wait_download_retry(retry: Retry, exc: Exception, path: str, control: _DownloadControl) -> Retry: + control.check() + if isinstance(exc, HubError) and not exc.retryable: + raise exc + try: + retry = retry.increment("GET", path, error=exc) + except MaxRetryError: + if isinstance(exc, HubError): + raise exc + raise NetworkError(f"下载 {path} 在 {DOWNLOAD_RETRY_TIMES} 次重试后失败：{type(exc).__name__}") from exc + requested = getattr(exc, "retry_after", None) + delay = min( + max(0, requested) if requested is not None else retry.get_backoff_time(), DOWNLOAD_RETRY_MAX_DELAY_SECONDS + ) + logger.warning("Download failed for %s (%s), retrying in %ss", path, type(exc).__name__, delay) + control.sleep(delay) + return retry + + +def _finish_download( + tmp: Path, + target: Path, + expected_hash: str | None, + control: _DownloadControl, + actual_hash: str | None = None, +) -> None: + if expected_hash: + if actual_hash is None: + digest = hashlib.sha256() + with tmp.open("rb") as fh: + while block := fh.read(8 * DOWNLOAD_CHUNK_SIZE): + control.check() + digest.update(block) + actual_hash = digest.hexdigest() + if actual_hash != expected_hash: + raise FileIntegrityError( + f"Hash mismatch for {target.name}: expected {expected_hash[:16]}..., got {actual_hash[:16]}..." + ) + control.check() + tmp.replace(target) + + # --------------------------------------------------------------------------- # Parallel range download # --------------------------------------------------------------------------- def _download_part_with_retry(params: tuple) -> None: - """Download a byte range with retry and resume support.""" - file_path, progress_callbacks, start, end, url, file_name, headers, cookies = params + """Download one byte range into its part file, resuming after transport errors. + + Every response must be exactly the requested range: a server that ignores + ``Range`` raises :class:`_RangeNotHonouredError` so the caller can fall back to a + single stream instead of splicing whole-file bodies into the parts. + """ + ( + file_path, + progress_callbacks, + start, + end, + url, + file_name, + headers, + cookies, + file_size, + request, + control, + part_stop, + ) = params + expected = end - start + 1 get_headers = {} if headers is None else copy.deepcopy(headers) get_headers["X-Request-ID"] = uuid.uuid4().hex retry = Retry( @@ -238,37 +415,81 @@ def _download_part_with_retry(params: tuple) -> None: allowed_methods=["GET"], ) part_file_name = f"{file_path}_{start}_{end}" + reported = 0 while True: try: - partial_length = 0 - if os.path.exists(part_file_name): - with open(part_file_name, "rb") as f: - partial_length = f.seek(0, io.SEEK_END) - for cb in progress_callbacks: - cb.update(partial_length) + control.check() + if part_stop.is_set(): + raise _DownloadCancelledError("分片组已取消") + partial_length = os.path.getsize(part_file_name) if os.path.exists(part_file_name) else 0 + if partial_length > expected: + # Left over from a different split of the file; it cannot be resumed. + os.remove(part_file_name) + partial_length = 0 + if partial_length > reported: + for cb in progress_callbacks: + cb.update(partial_length - reported) + reported = partial_length download_start = start + partial_length if download_start > end: break get_headers["Range"] = f"bytes={download_start}-{end}" - with open(part_file_name, "ab+") as f: - r = requests.get( - url, - stream=True, - headers=get_headers, - cookies=cookies, - timeout=DOWNLOAD_TIMEOUT, - ) - r.raise_for_status() - for chunk in r.iter_content(chunk_size=DOWNLOAD_CHUNK_SIZE): - if chunk: - f.write(chunk) - for cb in progress_callbacks: - cb.update(len(chunk)) + try: + r = request(get_headers) + except APIError as exc: + if exc.status_code == 416: + response = getattr(exc, "response", None) + total = ( + _unsatisfied_range_total(response.headers.get("Content-Range")) + if response is not None + else None + ) + raise FileSizeMismatchError( + f"Server reports {total if total is not None else 'fewer'} bytes for {file_name} " + f"but the repository lists {file_size} bytes (HTTP 416)" + ) from exc + raise + with control.response(r): + if part_stop.is_set(): + raise _DownloadCancelledError("分片组已取消") + content_range = _parse_content_range(r.headers.get("Content-Range")) + if content_range is not None and content_range[2] is not None and content_range[2] != file_size: + raise FileSizeMismatchError( + f"Server reports {content_range[2]} bytes for {file_name} but the repository lists " + f"{file_size} bytes ({_describe_response(r, get_headers['X-Request-ID'])})" + ) + if content_range is None or content_range[0] != download_start or content_range[1] != end: + raise _RangeNotHonouredError( + f"asked for bytes={download_start}-{end}, got " + f"{_describe_response(r, get_headers['X-Request-ID'])}" + ) + with open(part_file_name, "ab") as f: + for chunk in control.chunks(r, DOWNLOAD_CHUNK_SIZE): + if chunk: + if f.tell() + len(chunk) > expected: + raise FileSizeMismatchError(f"Server exceeded range {start}-{end} for {file_name}") + f.write(chunk) + reported += len(chunk) + for cb in progress_callbacks: + cb.update(len(chunk)) + written = os.path.getsize(part_file_name) + if written != expected: + if written > expected: + os.remove(part_file_name) + raise _IncompleteBodyError(f"part {start}-{end} holds {written} of {expected} bytes") break + except ( + _RangeNotHonouredError, + FileIntegrityError, + _DownloadCancelledError, + _DownloadDeadlineError, + NotExistError, + PermissionDeniedError, + AuthenticationError, + ): + raise except Exception as exc: - retry = retry.increment("GET", url, error=exc) - logger.warning("Downloading part %s-%s failed: %s, will retry", start, end, exc) - retry.sleep() + retry = _wait_download_retry(retry, exc, f"{file_name}:{start}-{end}", control) def _parallel_download( @@ -278,13 +499,20 @@ def _parallel_download( headers: dict[str, str] | None, cookies: dict | None, progress_callbacks: list[ProgressCallback] | None = None, + *, + request: Callable[[dict[str, str]], requests.Response], + control: _DownloadControl, + expected_sha256: str | None = None, ) -> str: """Split a large file into parts and download in parallel. - Returns the SHA256 hex digest of the merged file. + Returns the SHA256 hex digest of the merged file. When the server ignores + ranged requests, or serves an object of a different size, the parts are + removed and the error is re-raised for the caller to handle. """ callbacks = progress_callbacks or [] part_size = DOWNLOAD_PART_SIZE + part_stop = threading.Event() tasks: list[tuple] = [] file_path = str(target) target.parent.mkdir(parents=True, exist_ok=True) @@ -293,33 +521,92 @@ def _parallel_download( for idx in range(num_full_parts): start = idx * part_size end = (idx + 1) * part_size - 1 - tasks.append((file_path, callbacks, start, end, url, target.name, headers, cookies)) + tasks.append( + ( + file_path, + callbacks, + start, + end, + url, + target.name, + headers, + cookies, + file_size, + request, + control, + part_stop, + ) + ) remainder_start = num_full_parts * part_size if remainder_start < file_size: - tasks.append((file_path, callbacks, remainder_start, file_size - 1, url, target.name, headers, cookies)) + tasks.append( + ( + file_path, + callbacks, + remainder_start, + file_size - 1, + url, + target.name, + headers, + cookies, + file_size, + request, + control, + part_stop, + ) + ) + part_names = [f"{file_path}_{task[2]}_{task[3]}" for task in tasks] parallels = min(DOWNLOAD_PARALLELS, 16) - with ThreadPoolExecutor(max_workers=parallels, thread_name_prefix="download") as executor: - list(executor.map(_download_part_with_retry, tasks)) - - for cb in callbacks: - cb.end() + executor = ThreadPoolExecutor(max_workers=parallels, thread_name_prefix="download") + futures = [executor.submit(_download_part_with_retry, task) for task in tasks] + try: + for future in as_completed(futures): + future.result() + except BaseException as exc: + part_stop.set() + control.abort_responses() + for future in futures: + future.cancel() + executor.shutdown(wait=True, cancel_futures=True) + if isinstance(exc, (_RangeNotHonouredError, FileIntegrityError)): + # Parts of a whole-file body or of the wrong object cannot be resumed. + for name in part_names: + with contextlib.suppress(OSError): + os.remove(name) + raise + finally: + executor.shutdown(wait=True, cancel_futures=True) + for cb in callbacks: + cb.end() hash_sha256 = hashlib.sha256() + merged = 0 tmp_target = target.with_suffix(target.suffix + ".parallel_tmp") with open(tmp_target, "wb") as output_file: - for task in tasks: - part_file_name = f"{task[0]}_{task[2]}_{task[3]}" + for part_file_name in part_names: with open(part_file_name, "rb") as part_file: while True: chunk = part_file.read(16 * DOWNLOAD_CHUNK_SIZE) if not chunk: break + control.check() output_file.write(chunk) hash_sha256.update(chunk) + merged += len(chunk) os.remove(part_file_name) + if merged != file_size: + tmp_target.unlink(missing_ok=True) + raise FileSizeMismatchError(f"Merged {merged} of {file_size} bytes for {target.name} from ranged parts") + actual = hash_sha256.hexdigest() + if expected_sha256 and actual != expected_sha256: + tmp_target.unlink(missing_ok=True) + raise FileIntegrityError( + f"Hash mismatch for {target.name}: expected {expected_sha256[:16]}..., got {actual[:16]}..." + ) + control.check() tmp_target.replace(target) - return hash_sha256.hexdigest() + return actual # --------------------------------------------------------------------------- @@ -331,10 +618,14 @@ class DownloadManager: Dependencies are injected via constructor to keep this class testable. """ - def __init__(self, legacy_client: LegacyClient, config: HubConfig) -> None: + def __init__(self, legacy_client: LegacyClient, config: HubConfig, *, user_agent: dict | str | None = None) -> None: + from .utils import build_user_agent + self._client = legacy_client self._config = config + self._user_agent = build_user_agent(config=config, extra=user_agent) self._cached_region: str | None = None + self._region_lock = threading.Lock() # Cache only confirmed OSS-internal resolutions. A repository can contain # Git metadata and LFS objects behind different download routes, so cache # entries are scoped to repo/type/revision plus the redirect endpoint. @@ -349,27 +640,34 @@ def __init__(self, legacy_client: LegacyClient, config: HubConfig) -> None: # User-agent & headers # ------------------------------------------------------------------ def _build_user_agent(self, user_agent: dict | str | None = None) -> str: - from .utils import build_user_agent + from .utils import append_user_agent - return build_user_agent( - session_id=self._config.get_session_id(), - extra=user_agent, - ) + return append_user_agent(self._user_agent, user_agent) def _detect_region(self) -> str: - """Detect Alibaba cloud region ID for intra-cloud acceleration.""" - if self._cached_region is not None: + """Probe the region only once per DownloadManager instance.""" + with self._region_lock: + if self._cached_region is None: + self._cached_region = "" + self._cached_region = self._detect_region_uncached() return self._cached_region + def _detect_region_uncached(self) -> str: endpoint = self._client.endpoint internal_url = f"{endpoint}/api/v1/repos/internalAccelerationInfo" def _get(url: str, timeout: float): try: - resp = requests.get(url, timeout=timeout) - resp.raise_for_status() - return resp - except requests.exceptions.RequestException: + with _DownloadControl(total_timeout=timeout) as probe: + with self._client.read_context(probe): + resp = self._client._read_session.get(url, stream=True, timeout=timeout) + probe.track(resp) + with probe.response(resp): + resp.raise_for_status() + resp._content = b"".join(probe.chunks(resp, 8192)) + resp._content_consumed = True + return resp + except (requests.exceptions.RequestException, _DownloadDeadlineError): return None region_id = "" @@ -411,19 +709,27 @@ def _probe_redirect_url( headers: dict[str, str], cookies: dict | None = None, timeout: float = 5.0, + control: _DownloadControl | None = None, ) -> str: """Send a HEAD request without following redirects to get the 302 Location.""" + if control is not None: + control.check() + timeout = min(timeout, max(0.001, control.deadline - time.monotonic())) try: - r = requests.head( - url, - headers=headers, - cookies=cookies, - allow_redirects=False, - timeout=timeout, - ) - if r.status_code in (301, 302, 303, 307, 308): - return r.headers.get("Location", "") - except requests.exceptions.RequestException: + with ( + _DownloadControl( + total_timeout=timeout, + cancel_event=control.cancel_event if control is not None else None, + ) as probe, + _network_context(probe), + requests.Session() as session, + ): + session.mount("https://", _ReadHTTPAdapter(max_retries=0)) + session.mount("http://", _ReadHTTPAdapter(max_retries=0)) + with session.head(url, headers=headers, cookies=cookies, allow_redirects=False, timeout=timeout) as r: + if r.status_code in (301, 302, 303, 307, 308): + return r.headers.get("Location", "") + except (requests.exceptions.RequestException, _DownloadDeadlineError): pass return "" @@ -459,6 +765,7 @@ def _resolve_inter_region_headers( cookies: dict | None = None, peer_regions: list[str] | None = None, initial_redirect_url: str | None = None, + control: _DownloadControl | None = None, ) -> tuple[dict[str, str], str]: """Probe peer regions and return headers with the best region for OSS internal download. @@ -481,7 +788,9 @@ def _resolve_inter_region_headers( # established the storage route pass the initial redirect to avoid an # identical second HEAD request. redirect_url = ( - self._probe_redirect_url(url, headers, cookies) if initial_redirect_url is None else initial_redirect_url + self._probe_redirect_url(url, headers, cookies, control=control) + if initial_redirect_url is None + else initial_redirect_url ) if self._is_oss_internal_url(redirect_url): logger.debug("Inter-region: local region already yields OSS internal URL, skipping.") @@ -492,7 +801,7 @@ def _resolve_inter_region_headers( if peer_region == current_region: continue probe_headers = {**headers, "x-aliyun-region-id": peer_region} - redirect_url = self._probe_redirect_url(url, probe_headers, cookies) + redirect_url = self._probe_redirect_url(url, probe_headers, cookies, control=control) if self._is_oss_internal_url(redirect_url): logger.debug( 'Inter-region acceleration: using peer region "%s" (OSS internal).', @@ -507,14 +816,17 @@ def _resolve_inter_region_headers( def _build_download_headers( self, user_agent: dict | str | None = None, + *, + accelerate: bool = True, ) -> dict[str, str]: from .constants import _env, _env_bool headers: dict[str, str] = { "user-agent": self._build_user_agent(user_agent), "snapshot-identifier": uuid.uuid4().hex, + "Accept-Encoding": "identity", } - if _env_bool(ENV_INTRA_CLOUD_ACCELERATION, True): + if accelerate and _env_bool(ENV_INTRA_CLOUD_ACCELERATION, True): region = (_env(ENV_INTRA_CLOUD_REGION, "INTRA_CLOUD_ACCELERATION_REGION") or "").strip() if not region: try: @@ -542,6 +854,48 @@ def download_file( file_size: int | None = None, user_agent: dict | str | None = None, progress_callbacks: list[type[ProgressCallback]] | None = None, + _cancel_event: threading.Event | None = None, + ) -> Path: + """Download one file with cache-first lookup and a shared budget for locks, probes, and retries.""" + with _DownloadControl( + total_timeout=DOWNLOAD_TOTAL_TIMEOUT, + cancel_event=_cancel_event, + min_rate=DOWNLOAD_MIN_BYTES_PER_SECOND, + speed_window=DOWNLOAD_LOW_SPEED_WINDOW_SECONDS, + ) as control: + with self._client.read_context(control): + return self._download_file( + repo_id, + repo_type, + file_path, + revision, + cache_dir, + local_dir, + force, + expected_sha256, + local_files_only, + file_size, + user_agent, + progress_callbacks, + _control=control, + ) + + def _download_file( + self, + repo_id: str, + repo_type: str, + file_path: str, + revision: str | None = None, + cache_dir: Path | None = None, + local_dir: Path | None = None, + force: bool = False, + expected_sha256: str | None = None, + local_files_only: bool = False, + file_size: int | None = None, + user_agent: dict | str | None = None, + progress_callbacks: list[type[ProgressCallback]] | None = None, + *, + _control: _DownloadControl, ) -> Path: """Download a single file from a repository. @@ -574,7 +928,9 @@ def download_file( When ``True``, return the cached path without network access. Raises ``ValueError`` if the file is not cached. file_size: - Known file size (enables parallel range download for large files). + Known file size, e.g. from the repository listing. Received bytes + and cached files are validated against it, and it enables parallel + range download for large files. user_agent: Custom user-agent info appended to the default UA string. progress_callbacks: @@ -585,6 +941,14 @@ def download_file( ------- Path Absolute path to the downloaded (or cached) file on disk. + + Raises + ------ + FileSizeMismatchError + When the server keeps delivering a size other than *file_size*. + FileIntegrityError + When the content keeps failing the *expected_sha256* check. A + rejected download never remains at the target path. """ if local_files_only: return self._resolve_cached_file( @@ -608,47 +972,67 @@ def download_file( root = self._repo_cache_dir(repo_id, repo_type, cache_dir) target = root / "snapshots" / effective_revision / file_path - if not force and self._cache_hit(target, expected_sha256): + if not force and self._cache_hit(target, expected_sha256, file_size, control=_control): return target use_lock = _file_lock_enabled() lock_path = self._lock_path(repo_id, repo_type, cache_dir=cache_dir, file_path=file_path) if use_lock else None - with _optional_file_lock(lock_path, enabled=use_lock): + with _optional_file_lock(lock_path, enabled=use_lock, control=_control): # Re-check after acquiring lock — another process may have finished. - if not force and self._cache_hit(target, expected_sha256): + if not force and self._cache_hit(target, expected_sha256, file_size, control=_control): return target ensure_dir(target.parent) - - for attempt in range(DOWNLOAD_HASH_RETRY_TIMES): - self._download_with_resume( - repo_id, - repo_type, - file_path, - effective_revision, - target, - file_size=file_size, - user_agent=user_agent, - progress_callbacks=progress_callbacks, - ) - - if not expected_sha256: - break - + if not force: + # Whatever is at target failed validation. Drop it now so a failed + # download cannot leave it looking finished. + target.unlink(missing_ok=True) + + # A hash mismatch gets the full retry budget. A size mismatch means + # the server consistently served the wrong object, so it only gets + # one more attempt, over the default route. + failures = {"hash": 0, "size": 0} + limits = {"hash": DOWNLOAD_HASH_RETRY_TIMES, "size": DOWNLOAD_SIZE_RETRY_TIMES} + prefer_default_route = False + while True: + _control.check() try: - self.verify_file(target, expected_sha256) + self._download_with_resume( + repo_id, + repo_type, + file_path, + effective_revision, + target, + file_size=file_size, + user_agent=user_agent, + progress_callbacks=progress_callbacks, + prefer_default_route=prefer_default_route, + control=_control, + expected_sha256=expected_sha256, + ) break - except FileIntegrityError: - if attempt < DOWNLOAD_HASH_RETRY_TIMES - 1: - logger.warning( - "Hash validation failed for %s, retrying (%d/%d)", - file_path, - attempt + 1, - DOWNLOAD_HASH_RETRY_TIMES, - ) - target.unlink(missing_ok=True) - else: + except FileIntegrityError as exc: + kind = "size" if isinstance(exc, FileSizeMismatchError) else "hash" + # A size mismatch is raised before the rename, so the target is + # still the previous copy (only possible with force). + self._discard_download(target, keep_target=True) + target.with_suffix(target.suffix + ".parallel_tmp").unlink(missing_ok=True) + failures[kind] += 1 + if failures[kind] >= limits[kind]: raise + if kind == "size": + prefer_default_route = True + delay = DOWNLOAD_INTEGRITY_RETRY_BACKOFF * 2 ** (sum(failures.values()) - 1) + logger.warning( + "%s validation failed for %s, retrying in %.0fs (%d/%d): %s", + kind.capitalize(), + file_path, + delay, + failures[kind], + limits[kind], + exc.message, + ) + _control.sleep(delay) return target @@ -665,6 +1049,7 @@ def download_repo( local_files_only: bool = False, user_agent: dict | str | None = None, progress_callbacks: list[type[ProgressCallback]] | None = None, + file_list: list[dict] | None = None, ) -> Path: """Download an entire repository (snapshot download). @@ -697,6 +1082,13 @@ def download_repo( ------- Path Absolute path to the snapshot/local directory. + + Raises + ------ + FileDownloadError + After every file was attempted, when any of them failed. Completed + files stay in place; ``failed_files`` maps each failed path to its + error. """ if local_files_only: return self._resolve_cached_snapshot( @@ -732,18 +1124,42 @@ def download_repo( output_dir=output_dir, ) - if repo_type in ("dataset", "datasets"): - files = self._client.list_dataset_files_paginated( - repo_id=repo_id, - revision=effective_revision, - ) + cancel_event = threading.Event() + if file_list is not None: + files = file_list else: - files = self._client.list_repo_files( - repo_id=repo_id, - repo_type=repo_type, - revision=effective_revision, - recursive=True, - ) + files = [] + with _DownloadControl( + total_timeout=REPO_TREE_LIST_TIMEOUT_SECONDS, + cancel_event=cancel_event, + ) as listing_control: + with self._client.read_context(listing_control): + for listing_root in _listing_roots(allow_patterns): + listing_control.check() + try: + if repo_type in ("dataset", "datasets"): + files.extend( + self._client.list_dataset_files_paginated( + repo_id=repo_id, + revision=effective_revision, + root_path=listing_root or "/", + ) + ) + else: + files.extend( + self._client.list_repo_files( + repo_id=repo_id, + repo_type=repo_type, + revision=effective_revision, + recursive=True, + root=listing_root, + ) + ) + except NotExistError: + # A missing literal directory has no matches; a missing repository must still raise. + if listing_root is None: + raise + self._client.get_repo_info(repo_id, repo_type) download_items: list[tuple[str, str | None, int | None]] = [] for f in files: @@ -758,8 +1174,8 @@ def download_repo( if ignore_patterns and _matches_patterns(path, ignore_patterns): continue sha256 = f.get("Sha256") or f.get("sha256") or None - raw_size = f.get("Size") or f.get("size") - size = int(raw_size) if raw_size else None + raw_size = f.get("Size", f.get("size")) + size = int(raw_size) if raw_size is not None else None download_items.append((path, sha256, size)) if not download_items: @@ -768,38 +1184,68 @@ def download_repo( logger.info("Downloading %d files from %s@%s", len(download_items), repo_id, effective_revision) - errors: list[str] = [] - with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = { - executor.submit( - self.download_file, - repo_id=repo_id, - repo_type=repo_type, - file_path=fp, - revision=effective_revision, - cache_dir=cache_dir, - local_dir=local_dir, - expected_sha256=sha256, - file_size=size, - user_agent=user_agent, - progress_callbacks=progress_callbacks, - ): fp - for fp, sha256, size in download_items - } + failures: dict[str, BaseException] = {} + executor = ThreadPoolExecutor(max_workers=max_workers) + items = iter(download_items) + futures = {} + + def submit_next() -> None: + item = next(items, None) + if item is None: + return + fp, sha256, size = item + future = executor.submit( + self.download_file, + repo_id=repo_id, + repo_type=repo_type, + file_path=fp, + revision=effective_revision, + cache_dir=cache_dir, + local_dir=local_dir, + expected_sha256=sha256, + file_size=size, + user_agent=user_agent, + progress_callbacks=progress_callbacks, + _cancel_event=cancel_event, + ) + futures[future] = fp + try: + # Bound pending tasks instead of creating millions of Futures; stop scheduling after cancellation. + for _ in range(max_workers * 2): + submit_next() with tqdm(total=len(download_items), desc="Downloading", unit="file") as pbar: - for future in as_completed(futures): - fp = futures[future] - try: - future.result() - except Exception as exc: - errors.append(f"{fp}: {exc}") - logger.error("Failed to download %s: %s", fp, exc) - finally: - pbar.update(1) - - if errors: - logger.warning("%d file(s) failed to download", len(errors)) + while futures: + completed, _ = wait(futures, return_when=FIRST_COMPLETED) + for future in completed: + fp = futures.pop(future) + try: + future.result() + except Exception as exc: + failures[fp] = exc + logger.error("Failed to download %s: %s", fp, exc) + finally: + pbar.update(1) + submit_next() + except BaseException: + cancel_event.set() + for future in futures: + future.cancel() + raise + finally: + executor.shutdown(wait=True, cancel_futures=True) + + if failures: + # Returning normally would let callers treat a partial snapshot as + # complete. Every file was attempted first, so a re-run only fetches + # the failed ones. + raise FileDownloadError( + f"{len(failures)} of {len(download_items)} file(s) failed to download from " + f"{repo_id}@{effective_revision}", + failed_files=failures, + total=len(download_items), + output_dir=output_dir, + ) return output_dir @@ -824,17 +1270,36 @@ def _download_archive( tmp_path: Path | None = None try: - resp = self._client.download_archive( - repo_id=repo_id, - repo_type=repo_type, - revision=revision, - ) - - with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp: - tmp_path = Path(tmp.name) - for chunk in resp.iter_content(chunk_size=DOWNLOAD_CHUNK_SIZE): - if chunk: - tmp.write(chunk) + with ( + _DownloadControl( + total_timeout=DOWNLOAD_TOTAL_TIMEOUT, + min_rate=DOWNLOAD_MIN_BYTES_PER_SECOND, + speed_window=DOWNLOAD_LOW_SPEED_WINDOW_SECONDS, + ) as control, + self._client.read_context(control), + ): + retry = Retry(total=DOWNLOAD_RETRY_TIMES, backoff_factor=1, allowed_methods=["GET"]) + while True: + try: + resp = self._client.download_archive( + repo_id=repo_id, + repo_type=repo_type, + revision=revision, + ) + with control.response(resp), tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp: + tmp_path = Path(tmp.name) + for chunk in control.chunks(resp, DOWNLOAD_CHUNK_SIZE): + tmp.write(chunk) + announced = _checkable_content_length(resp) + if announced is not None and tmp_path.stat().st_size != announced: + raise _IncompleteBodyError("归档内容长度与响应头不一致") + break + except (_DownloadCancelledError, _DownloadDeadlineError): + raise + except Exception as exc: + if tmp_path is not None: + tmp_path.unlink(missing_ok=True) + retry = _wait_download_retry(retry, exc, repo_id, control) with zipfile.ZipFile(tmp_path, "r") as zf: zf.extractall(output_dir) @@ -857,12 +1322,29 @@ def _download_archive( # Internal helpers # ------------------------------------------------------------------ @staticmethod - def _cache_hit(target: Path, expected_sha256: str | None) -> bool: - """Return True if *target* exists and passes optional hash check.""" + def _cache_hit( + target: Path, + expected_sha256: str | None, + expected_size: int | None = None, + *, + control: _DownloadControl | None = None, + ) -> bool: + """Return True if *target* exists and passes the optional size and hash checks.""" if not target.exists(): return False + if expected_size is not None and target.stat().st_size != expected_size: + logger.debug("Cache stale (size mismatch): %s", target) + return False if expected_sha256: - actual = compute_hash(target, "sha256") + if control is None: + actual = compute_hash(target, "sha256") + else: + digest = hashlib.sha256() + with target.open("rb") as cached: + while block := cached.read(8 * DOWNLOAD_CHUNK_SIZE): + control.check() + digest.update(block) + actual = digest.hexdigest() if actual != expected_sha256: logger.debug("Cache stale (hash mismatch): %s", target) return False @@ -871,6 +1353,13 @@ def _cache_hit(target: Path, expected_sha256: str | None) -> bool: logger.debug("Cache hit: %s", target) return True + @staticmethod + def _discard_download(target: Path, *, keep_target: bool = False) -> None: + """Remove a rejected download so neither it nor its partial file looks reusable.""" + _incomplete_path(target).unlink(missing_ok=True) + if not keep_target: + target.unlink(missing_ok=True) + @staticmethod def _dir_has_entries(path: Path) -> bool: """Return True when a directory exists and contains at least one entry.""" @@ -1149,11 +1638,23 @@ def _download_with_resume( file_size: int | None = None, user_agent: dict | str | None = None, progress_callbacks: list[type[ProgressCallback]] | None = None, + prefer_default_route: bool = False, + control: _DownloadControl, + expected_sha256: str | None = None, ) -> Path: - """Download a file with HTTP Range resume support and retry.""" + """Download a file with HTTP Range resume support and retry. + + Before the partial file is moved into place, the bytes received are + checked against the announced ``Content-Length`` and, when *file_size* + is known, against the repository listing. A short body is resumed; a + body whose size contradicts the listing raises + :class:`~.errors.FileSizeMismatchError` and leaves no partial file. + *prefer_default_route* skips intra-cloud and peer-region routing so a + retry can avoid a replica that served bad content. + """ use_parallel = file_size is not None and file_size > DOWNLOAD_PARALLEL_THRESHOLD and DOWNLOAD_PARALLELS > 1 - download_headers = self._build_download_headers(user_agent) + download_headers = self._build_download_headers(user_agent, accelerate=not prefer_default_route) # Inter-region acceleration: first discover the file's storage route, # then cache only a confirmed internal resolution for that exact route. @@ -1161,7 +1662,7 @@ def _download_with_resume( # OSS; a metadata miss must therefore never suppress an LFS probe. # Progress bar prefix: "⚡ " = local OSS, "⇄ " = peer OSS, " " = CDN, "" = not configured source_prefix = "" - peer_regions = self._get_inter_cloud_regions() + peer_regions = [] if prefer_default_route else self._get_inter_cloud_regions() if peer_regions: source = "default" cached: tuple[str | None, str] | None = None @@ -1169,7 +1670,7 @@ def _download_with_resume( try: probe_url = self._client.get_download_url(repo_id, repo_type, file_path, revision) cookies = {"m_session_id": self._client.token} if self._client.token else None - initial_redirect_url = self._probe_redirect_url(probe_url, download_headers, cookies) + initial_redirect_url = self._probe_redirect_url(probe_url, download_headers, cookies, control=control) route_key = self._storage_route_key(initial_redirect_url) if route_key is not None: @@ -1188,7 +1689,8 @@ def _download_with_resume( if not is_probe_owner: assert probe_event is not None - probe_event.wait() + while not probe_event.wait(0.1): + control.check() continue try: @@ -1198,6 +1700,7 @@ def _download_with_resume( cookies, peer_regions=peer_regions, initial_redirect_url=initial_redirect_url, + control=control, ) resolved_region = download_headers.get("x-aliyun-region-id") if source != "default": @@ -1220,6 +1723,8 @@ def _download_with_resume( resolved_region, source = cached if resolved_region is not None: download_headers["x-aliyun-region-id"] = resolved_region + except (_DownloadCancelledError, _DownloadDeadlineError): + raise except Exception as exc: logger.warning( "Failed to resolve inter-region acceleration: %s. Falling back to default.", @@ -1233,6 +1738,17 @@ def _download_with_resume( "default": " ", }[source] + def request(headers: dict[str, str]) -> requests.Response: + control.check() + with self._client.read_context(control): + return self._client.download_stream( + repo_id, + repo_type, + file_path, + revision, + headers=headers, + ) + if use_parallel: assert file_size is not None # guaranteed by use_parallel above url = self._client.get_download_url( @@ -1249,95 +1765,206 @@ def _download_with_resume( if progress_callbacks: cb_instances = [cls(file_path, file_size) for cls in progress_callbacks] - _parallel_download( - url=url, - target=target, - file_size=file_size, - headers=download_headers, - cookies=cookies, - progress_callbacks=cb_instances, - ) - return target + try: + _parallel_download( + url=url, + target=target, + file_size=file_size, + headers=download_headers, + cookies=cookies, + progress_callbacks=cb_instances, + request=request, + control=control, + expected_sha256=expected_sha256, + ) + return target + except _RangeNotHonouredError as exc: + logger.warning("Ranged requests are not honoured for %s (%s); using a single stream", file_path, exc) # Single-stream download with retry cb_classes = list(progress_callbacks or []) cb_instances = [cls(file_path, file_size or 0) for cls in cb_classes] + reported = 0 + # The listing reports 0 both for empty files and for unknown sizes. + expected_size = file_size retry = Retry( total=DOWNLOAD_RETRY_TIMES, backoff_factor=1, allowed_methods=["GET"], ) - tmp_path = target.with_suffix(target.suffix + ".incomplete") + tmp_path = _incomplete_path(target) while True: try: - existing_size = 0 - if tmp_path.exists(): - existing_size = tmp_path.stat().st_size + control.check() + existing_size = tmp_path.stat().st_size if tmp_path.exists() else 0 + if expected_size is not None and existing_size > expected_size: + # Longer than the file itself, so it cannot be a prefix of it. + tmp_path.unlink() + existing_size = 0 + if existing_size > reported: for cb in cb_instances: - cb.update(existing_size) + cb.update(existing_size - reported) + reported = existing_size + if tmp_path.exists() and expected_size is not None and existing_size == expected_size: + # An earlier run received every byte but stopped before the + # rename. Resuming would only earn a 416 from the CDN; the + # caller's hash check decides whether the bytes are right. + _finish_download(tmp_path, target, expected_sha256, control) + for cb in cb_instances: + cb.end() + return target + request_id = uuid.uuid4().hex extra_headers: dict[str, str] = copy.deepcopy(download_headers) - extra_headers["X-Request-ID"] = uuid.uuid4().hex + extra_headers["X-Request-ID"] = request_id if existing_size > 0: extra_headers["Range"] = f"bytes={existing_size}-" logger.debug("Resuming download from byte %d", existing_size) try: - resp = self._client.download_stream( - repo_id=repo_id, - repo_type=repo_type, - file_path=file_path, - revision=revision, - headers=extra_headers, - ) + resp = request(extra_headers) except requests.Timeout as exc: raise RequestTimeoutError(f"Download timed out for {file_path}: {exc}") from exc except requests.ConnectionError as exc: raise NetworkError(f"Download connection failed for {file_path}: {exc}") from exc except requests.RequestException as exc: raise NetworkError(f"Download failed for {file_path}: {exc}") from exc + except APIError as exc: + if exc.status_code == 416 and existing_size > 0: + # The object does not extend past the partial file, so the + # partial file is stale. Start over. + tmp_path.unlink(missing_ok=True) + raise _IncompleteBodyError( + f"server refused to resume from byte {existing_size} (HTTP 416)" + ) from exc + raise + + with control.response(resp): + # The LFS CDN returns the object SHA-256, enabling validation without caller metadata. + linked_hash = resp.headers.get("X-Linked-ETag", "").strip('"') + if expected_sha256 is None and re.fullmatch(r"[0-9a-fA-F]{64}", linked_hash): + expected_sha256 = linked_hash.lower() + content_length = _checkable_content_length(resp) + content_range = _parse_content_range(resp.headers.get("Content-Range")) + if existing_size > 0 and (resp.status_code == 206 or content_range is not None): + if content_range is None or content_range[0] != existing_size: + tmp_path.unlink(missing_ok=True) + raise _IncompleteBodyError( + f"resume from byte {existing_size} was answered with " + f"{_describe_response(resp, request_id)}; restarting" + ) + mode, offset = "ab", existing_size + else: + # Anything but a matching 206 carries the whole file: files + # served without a redirect ignore Range entirely. + mode, offset = "wb", 0 + if reported: + for cb in cb_instances: + cb.update(-reported) + reported = 0 + + if resp.status_code == 206 and (content_range is None or content_range[0] != offset): + tmp_path.unlink(missing_ok=True) + raise _IncompleteBodyError("服务器返回了不匹配请求起点的分段响应") + announced: int | None = None + if content_range is not None and content_range[2] is not None: + announced = content_range[2] + elif content_length is not None: + announced = offset + content_length + if expected_size is not None and announced is not None and announced != expected_size: + tmp_path.unlink(missing_ok=True) + raise FileSizeMismatchError( + f"Server announced {announced} bytes for {file_path} but the repository lists " + f"{expected_size} bytes ({_describe_response(resp, request_id)}); the remote object " + "is empty, truncated or out of date" + ) - content_length = resp.headers.get("Content-Length") - total_size = int(content_length) if content_length else None - is_resumed = resp.status_code == 206 - - if is_resumed and total_size: - total_size += existing_size - - mode = "ab" if is_resumed else "wb" - if not is_resumed: - existing_size = 0 + received = 0 + digest = hashlib.sha256() if expected_sha256 else None + if digest is not None and offset: + with tmp_path.open("rb") as prefix: + while block := prefix.read(8 * DOWNLOAD_CHUNK_SIZE): + control.check() + digest.update(block) + with tqdm( + total=announced, + initial=offset, + unit="B", + unit_scale=True, + desc=f"{source_prefix}{Path(file_path).name}", + leave=False, + ) as pbar: + with open(tmp_path, mode) as fh: + for chunk in control.chunks(resp, DOWNLOAD_CHUNK_SIZE): + if chunk: + if expected_size is not None and offset + received + len(chunk) > expected_size: + raise FileSizeMismatchError( + f"Server exceeded {expected_size} bytes for {file_path} " + f"({_describe_response(resp, request_id)})" + ) + fh.write(chunk) + if digest is not None: + digest.update(chunk) + received += len(chunk) + reported += len(chunk) + pbar.update(len(chunk)) + for cb in cb_instances: + cb.update(len(chunk)) + + if content_length is not None and received < content_length: + # urllib3 2.x raises on a short read, but 1.x ends the body + # silently; keep the partial file and resume from it. + raise _IncompleteBodyError( + f"connection closed after {received} of {content_length} bytes " + f"({_describe_response(resp, request_id)})" + ) + if expected_size is not None and offset + received != expected_size: + tmp_path.unlink(missing_ok=True) + raise FileSizeMismatchError( + f"Server returned {offset + received} of {expected_size} bytes for {file_path} " + f"({_describe_response(resp, request_id)}); the remote object is empty, truncated " + "or out of date" + ) - with tqdm( - total=total_size, - initial=existing_size, - unit="B", - unit_scale=True, - desc=f"{source_prefix}{Path(file_path).name}", - leave=False, - ) as pbar: - with open(tmp_path, mode) as fh: - for chunk in resp.iter_content(chunk_size=DOWNLOAD_CHUNK_SIZE): - if chunk: - fh.write(chunk) - pbar.update(len(chunk)) - for cb in cb_instances: - cb.update(len(chunk)) - - tmp_path.replace(target) + _finish_download( + tmp_path, + target, + expected_sha256, + control, + actual_hash=digest.hexdigest() if digest is not None else None, + ) for cb in cb_instances: cb.end() logger.debug("Downloaded: %s", target) return target - except (NotExistError, PermissionDeniedError): + except ( + NotExistError, + PermissionDeniedError, + AuthenticationError, + FileIntegrityError, + _DownloadCancelledError, + _DownloadDeadlineError, + ): + for cb in cb_instances: + cb.end() raise except Exception as exc: - retry = retry.increment("GET", file_path, error=exc) - logger.warning("Download failed for %s: %s, will retry", file_path, exc) - retry.sleep() + if not prefer_default_route: + # Fall back to the public route for this file only, without changing other files' cached routes. + download_headers.pop("x-aliyun-region-id", None) + try: + retry = _wait_download_retry(retry, exc, file_path, control) + except BaseException: + for cb in cb_instances: + cb.end() + raise + except BaseException: + for cb in cb_instances: + cb.end() + raise def verify_file(self, file_path: Path, expected_sha256: str) -> bool: """Verify a downloaded file's SHA256 hash. diff --git a/src/modelscope_hub/_download_control.py b/src/modelscope_hub/_download_control.py new file mode 100644 index 0000000..f98b120 --- /dev/null +++ b/src/modelscope_hub/_download_control.py @@ -0,0 +1,280 @@ +"""Download cancellation, deadlines, and stream monitoring without changing uploads.""" + +from __future__ import annotations + +import contextlib +import socket +import threading +import time +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Any + +import requests +from requests.adapters import HTTPAdapter +from urllib3.exceptions import HTTPError as Urllib3HTTPError + +from .errors import NetworkError, RequestTimeoutError + + +class _DownloadCancelledError(Exception): + """The parent download was cancelled; do not retry or report a remote failure.""" + + +class _DownloadDeadlineError(RequestTimeoutError): + """The total download time budget has been exhausted.""" + + retryable = False + + +class _DownloadTooSlowError(NetworkError): + """Download progress stayed below the minimum rate for a full observation window.""" + + +@dataclass +class _WatchedResponse: + sock: Any + started: float + received: int = 0 + error: Exception | None = None + monitor_speed: bool = True + + +def _response_socket(response: requests.Response) -> Any: + # urllib3 may release connection.sock while HTTPResponse still retains + # SocketIO. Only shut down interrupted reads; leave socket timeouts and data unchanged. + raw = response.raw + connection = getattr(raw, "_connection", None) + sock = getattr(connection, "sock", None) + if sock is None: + fp = getattr(getattr(raw, "_fp", None), "fp", None) + sock = getattr(getattr(fp, "raw", None), "_sock", None) + return sock + + +class _DownloadControl: + """A per-call budget whose workers share a parent cancellation event, not global state.""" + + def __init__( + self, + *, + total_timeout: float, + cancel_event: threading.Event | None = None, + min_rate: int = 0, + speed_window: float = 60, + ) -> None: + self.cancel_event = cancel_event if cancel_event is not None else threading.Event() + self.deadline = time.monotonic() + total_timeout + self.min_rate = min_rate + self.speed_window = speed_window + self._responses: dict[requests.Response, _WatchedResponse] = {} + self._transports: dict[Any, _WatchedResponse] = {} + self._lock = threading.Lock() + self._finished = threading.Event() + self._monitor: threading.Thread | None = None + + def __enter__(self) -> _DownloadControl: + return self + + def __exit__(self, *args: Any) -> None: + self._finished.set() + if self._monitor is not None: + self._monitor.join(timeout=1) + with self._lock: + responses = list(self._responses) + self._responses.clear() + for response in responses: + response.close() + + def check(self) -> None: + if self.cancel_event.is_set(): + raise _DownloadCancelledError("下载已取消") + if time.monotonic() >= self.deadline: + raise _DownloadDeadlineError("下载总时限已耗尽；已保存可继续下载的临时文件") + + def timeouts(self, connect: float, read: float) -> tuple[float, float]: + self.check() + remaining = max(0.001, self.deadline - time.monotonic()) + return min(connect, remaining), min(read, remaining) + + def sleep(self, seconds: float) -> None: + self.check() + self.cancel_event.wait(min(max(0, seconds), max(0, self.deadline - time.monotonic()))) + self.check() + + def _start_monitor(self) -> None: + if self._monitor is None: + self._monitor = threading.Thread(target=self._watch, name="download-watchdog", daemon=True) + self._monitor.start() + + def track(self, response: requests.Response) -> None: + self.check() + with self._lock: + self._responses[response] = _WatchedResponse(_response_socket(response), time.monotonic()) + self._start_monitor() + + @contextlib.contextmanager + def transport(self, sock: Any) -> Iterator[None]: + """Interrupt sockets while reading headers so trickling headers cannot bypass deadlines.""" + self.check() + state = _WatchedResponse(sock, time.monotonic(), monitor_speed=False) + with self._lock: + self._transports[sock] = state + self._start_monitor() + try: + yield + finally: + with self._lock: + self._transports.pop(sock, None) + self.check() + if state.error is not None: + raise state.error + + def release(self, response: requests.Response) -> None: + with self._lock: + self._responses.pop(response, None) + + @contextlib.contextmanager + def response(self, response: requests.Response) -> Iterator[requests.Response]: + try: + yield response + finally: + self.release(response) + response.close() + + def _check_response(self, response: requests.Response, received: int = 0) -> None: + self.check() + with self._lock: + state = self._responses.get(response) + if state is not None: + state.received += received + if state.error is not None: + raise state.error + + def _watch(self) -> None: + # A separate watchdog interrupts reads that keep receiving tiny chunks, + # including with urllib3 1.26. Do not call response.close() here: the + # reader may hold the buffered stream lock. + while not self._finished.wait(0.1): + now = time.monotonic() + with self._lock: + for state in [*self._responses.values(), *self._transports.values()]: + if state.error is not None: + continue + if self.cancel_event.is_set(): + state.error = _DownloadCancelledError("下载已取消") + elif now >= self.deadline: + state.error = _DownloadDeadlineError("下载总时限已耗尽") + elif state.monitor_speed and self.min_rate and now - state.started >= self.speed_window: + rate = state.received / (now - state.started) + if rate < self.min_rate: + state.error = _DownloadTooSlowError( + f"下载速率 {rate:.0f} B/s 低于 {self.min_rate} B/s，观察窗口 {self.speed_window:g}s" + ) + else: + state.started, state.received = now, 0 + if state.error is not None and state.sock is not None: + with contextlib.suppress(OSError): + state.sock.shutdown(socket.SHUT_RDWR) + + def abort_responses(self) -> None: + """Abort active responses without cancelling the parent, allowing single-stream fallback.""" + with self._lock: + for state in [*self._responses.values(), *self._transports.values()]: + state.error = _DownloadCancelledError("分片组已取消") + if state.sock is not None: + with contextlib.suppress(OSError): + state.sock.shutdown(socket.SHUT_RDWR) + + def chunks(self, response: requests.Response, chunk_size: int) -> Iterator[bytes]: + """Yield received bytes promptly instead of waiting for a full chunk to report progress.""" + raw = response.raw + try: + if hasattr(raw, "stream"): + read1 = getattr(raw, "read1", None) + if read1 is None and not response.headers.get("Content-Encoding"): + read1 = getattr(getattr(raw, "_fp", None), "read1", None) + decoded = False + else: + decoded = True + while True: + self._check_response(response) + if read1 is not None: + data = read1(chunk_size, decode_content=True) if decoded else read1(chunk_size) + else: + data = raw.read(chunk_size, decode_content=True) + self._check_response(response, len(data)) + if not data: + break + yield data + else: + for chunk in response.iter_content(chunk_size=chunk_size): + self._check_response(response, len(chunk)) + if chunk: + yield chunk + except (Urllib3HTTPError, OSError) as exc: + self._check_response(response) + raise requests.ConnectionError(f"下载流读取失败：{type(exc).__name__}") from exc + finally: + self.release(response) + + +_network_state = threading.local() + + +@contextlib.contextmanager +def _network_context(control: _DownloadControl) -> Iterator[None]: + previous = getattr(_network_state, "control", None) + _network_state.control = control + try: + yield + finally: + _network_state.control = previous + + +class _WatchedConnectionMixin: + def _watched(self, method: str) -> Any: + original = getattr(super(), method) + control = getattr(_network_state, "control", None) + if control is None: + return original() + with control.transport(getattr(self, "sock", None)): + return original() + + def getresponse(self) -> Any: + return self._watched("getresponse") + + def _tunnel(self) -> Any: + # HTTP CONNECT tunnel headers can also arrive as an endless trickle. + return self._watched("_tunnel") + + +class _ReadHTTPAdapter(HTTPAdapter): + """Download/listing pools with instance-local connection classes, leaving urllib3 globals intact.""" + + @staticmethod + def _configure(manager: Any) -> None: + classes = dict(manager.pool_classes_by_scheme) + for scheme, pool in classes.items(): + if getattr(pool, "_download_controlled", False): + continue + connection = getattr(pool, "ConnectionCls", None) + if connection is None: + continue + # Preserve the original TLS, HTTP proxy, and SOCKS connection classes and arguments. + watched = type(f"_Download{connection.__name__}", (_WatchedConnectionMixin, connection), {}) + classes[scheme] = type( + f"_Download{pool.__name__}", + (pool,), + {"ConnectionCls": watched, "_download_controlled": True}, + ) + manager.pool_classes_by_scheme = classes + + def init_poolmanager(self, connections: int, maxsize: int, block: bool = False, **kwargs: Any) -> None: + super().init_poolmanager(connections, maxsize, block=block, **kwargs) + self._configure(self.poolmanager) + + def proxy_manager_for(self, proxy: str, **kwargs: Any) -> Any: + manager = super().proxy_manager_for(proxy, **kwargs) + self._configure(manager) + return manager diff --git a/src/modelscope_hub/_legacy_api.py b/src/modelscope_hub/_legacy_api.py index 211614d..b7844b5 100644 --- a/src/modelscope_hub/_legacy_api.py +++ b/src/modelscope_hub/_legacy_api.py @@ -14,25 +14,35 @@ from __future__ import annotations +import contextlib +import threading import time import uuid +from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor, as_completed from typing import IO, Any, BinaryIO from urllib.parse import quote_plus, urlparse import requests from requests.adapters import HTTPAdapter, Retry +from requests.structures import CaseInsensitiveDict +from urllib3.exceptions import InvalidHeader +from ._download_control import _DownloadControl, _DownloadDeadlineError, _network_context, _ReadHTTPAdapter +from .config import get_default_config from .constants import ( API_CONNECT_TIMEOUT, API_CONNECTION_POOL_MAXSIZE, API_MAX_RETRIES, API_TIMEOUT, + DOWNLOAD_CHUNK_SIZE, + DOWNLOAD_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_PAGE_TIMEOUT_SECONDS, REPO_TREE_WALK_WORKERS, UPLOAD_BLOB_CONNECT_TIMEOUT_SECONDS, UPLOAD_BLOB_READ_TIMEOUT_SECONDS, @@ -44,10 +54,12 @@ InvalidParameter, NetworkError, PermissionDeniedError, + RateLimitError, RequestTimeoutError, ServerError, raise_for_status, ) +from .utils import append_user_agent, build_user_agent from .utils.logger import get_logger logger = get_logger("legacy_api") @@ -118,8 +130,12 @@ def __init__( self._tree_reads_ok = False self._session = requests.Session() - if user_agent: - self._session.headers["User-Agent"] = user_agent + self._user_agent = ( + user_agent + if user_agent and user_agent.startswith("modelscope_hub/") + else build_user_agent(config=get_default_config(), extra=user_agent) + ) + self._session.headers["User-Agent"] = self._user_agent retry = Retry( total=max_retries, backoff_factor=0.5, @@ -133,6 +149,27 @@ def __init__( ) self._session.mount("https://", adapter) self._session.mount("http://", adapter) + # Callers manage download/listing retries; this pool leaves upload and regular API policies unchanged. + self._read_session = requests.Session() + self._read_session.headers.update(self._session.headers) + read_adapter = _ReadHTTPAdapter( + max_retries=0, + pool_connections=API_CONNECTION_POOL_MAXSIZE, + pool_maxsize=API_CONNECTION_POOL_MAXSIZE, + ) + self._read_session.mount("https://", read_adapter) + self._read_session.mount("http://", read_adapter) + self._read_state = threading.local() + + @contextlib.contextmanager + def read_context(self, control: _DownloadControl) -> Iterator[None]: + previous = getattr(self._read_state, "control", None) + self._read_state.control = control + try: + with _network_context(control): + yield + finally: + self._read_state.control = previous # ------------------------------------------------------------------ # Properties @@ -150,6 +187,7 @@ def token(self, value: str | None) -> None: self._token = value self._session_authenticated = False self._session.cookies.clear() + self._read_session.cookies.clear() # ------------------------------------------------------------------ # Transport @@ -165,9 +203,11 @@ def _headers(self, extra: dict[str, str] | None = None) -> dict[str, str]: } if self._token: headers["Authorization"] = f"Bearer {self._token}" + merged = CaseInsensitiveDict(headers) if extra: - headers.update(extra) - return headers + merged.update(extra) + merged["User-Agent"] = append_user_agent(self._user_agent, merged.get("User-Agent")) + return dict(merged) def _ensure_session_auth(self) -> None: """Set the ``m_session_id`` cookie on the session for legacy API auth. @@ -181,6 +221,7 @@ def _ensure_session_auth(self) -> None: return domain = urlparse(self._endpoint).hostname or "" self._session.cookies.set("m_session_id", self._token, domain=domain, path="/") + self._read_session.cookies.set("m_session_id", self._token, domain=domain, path="/") self._session_authenticated = True logger.debug("Legacy session cookie set for domain=%s", domain) @@ -194,8 +235,9 @@ def _request( data: Any | None = None, files: dict[str, Any] | None = None, headers: dict[str, str] | None = None, - timeout: int | None = None, + timeout: int | tuple[float, float] | None = None, stream: bool = False, + single_attempt: bool = False, ) -> requests.Response: """Unified request dispatcher with auth, retry, and error handling.""" self._ensure_session_auth() @@ -207,8 +249,32 @@ def _request( merged_headers.pop("Content-Type", None) logger.debug("%s %s", method, url) + control = getattr(self._read_state, "control", None) + if single_attempt and control is None: + # Bound both headers and body reads for single-page calls without an external controller. + with _DownloadControl(total_timeout=REPO_TREE_PAGE_TIMEOUT_SECONDS) as page_control: + with self.read_context(page_control): + return self._request( + method, + path, + params=params, + json_body=json_body, + data=data, + files=files, + headers=headers, + timeout=timeout, + stream=stream, + single_attempt=True, + ) + session = self._read_session if control is not None or single_attempt else self._session + request_timeout = timeout or self._timeout + if control is not None: + connect, read = ( + (API_CONNECT_TIMEOUT, request_timeout) if isinstance(request_timeout, int) else request_timeout + ) + request_timeout = control.timeouts(connect, read) try: - resp = self._session.request( + resp = session.request( method=method, url=url, params=params, @@ -216,8 +282,8 @@ def _request( data=data, files=files, headers=merged_headers, - timeout=timeout or self._timeout, - stream=stream, + timeout=request_timeout, + stream=stream or control is not None, ) except requests.exceptions.RetryError as exc: raise ServerError(f"Max retries exceeded: {exc}") from exc @@ -227,7 +293,7 @@ def _request( raise RequestTimeoutError(f"Request timed out: {exc}") from exc logger.debug("%s %s -> %s", method, url, resp.status_code) - if resp.status_code >= 400: + if resp.status_code >= 400 and control is None: logger.debug( "Request failed: %s %s params=%s status=%s body=%s", method, @@ -236,6 +302,38 @@ def _request( resp.status_code, resp.text[:500] if resp.text else "", ) + if control is not None: + try: + control.track(resp) + if not stream or resp.status_code >= 400: + # Bound error-body reads too, avoiding unbounded reads inside raise_for_status. + body = bytearray() + for chunk in control.chunks(resp, DOWNLOAD_CHUNK_SIZE): + body.extend(chunk) + if resp.status_code >= 400 and len(body) >= 65536: + break + if resp.status_code >= 400 and len(body) >= 65536: + resp.raw.close() + resp._content = bytes(body) + resp._content_consumed = True + control.release(resp) + resp.close() + try: + raise_for_status(resp) + except APIError as exc: + exc.response = resp # type: ignore[attr-defined] + retry_after = resp.headers.get("Retry-After") + if retry_after: + try: + exc.retry_after = Retry().parse_retry_after(retry_after) # type: ignore[attr-defined] + except (ValueError, InvalidHeader): + pass + raise + return resp + except BaseException: + control.release(resp) + resp.close() + raise raise_for_status(resp) return resp @@ -488,22 +586,37 @@ def fetch_level(roots: list[str | None], *, recursive: bool) -> dict[str | None, budget_spent = True requests_made += len(roots) + control = getattr(self._read_state, "control", None) + def one(subroot: str | None) -> list[dict]: - return self._list_files_page( - repo_id, - repo_type, - revision, - recursive=recursive, - root=subroot, - ) + scope = self.read_context(control) if control is not None else contextlib.nullcontext() + with scope: + if control is not None: + control.check() + return self._list_files_page( + repo_id, + repo_type, + revision, + recursive=recursive, + root=subroot, + ) if len(roots) == 1: listings = {roots[0]: one(roots[0])} else: workers = min(REPO_TREE_WALK_WORKERS, len(roots)) - with ThreadPoolExecutor(max_workers=workers) as pool: - futures = {pool.submit(one, subroot): subroot for subroot in roots} + pool = ThreadPoolExecutor(max_workers=workers) + futures = {pool.submit(one, subroot): subroot for subroot in roots} + try: listings = {futures[future]: future.result() for future in as_completed(futures)} + except BaseException: + if control is not None: + control.cancel_event.set() + for future in futures: + future.cancel() + raise + finally: + pool.shutdown(wait=True, cancel_futures=True) if requests_made >= progress_mark: logger.info( @@ -593,7 +706,7 @@ def list_dataset_files_paginated( ``GET /api/v1/datasets/{repo_id}/repo/tree`` with ``PageNumber``/``PageSize`` params. - A page occasionally answers ``403 无权访问该数据集`` on a repository the + A page occasionally answers ``403 access denied`` 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 @@ -602,6 +715,7 @@ def list_dataset_files_paginated( """ all_files: list[dict] = [] page_number = 1 + logger.info("Listing %s@%s Root=%s ...", repo_id, revision, root_path) while True: params: dict[str, Any] = { "Revision": revision, @@ -626,11 +740,37 @@ def list_dataset_files_paginated( files = [] all_files.extend(files) + if page_number % 10 == 0 or len(files) < page_size: + logger.info("Listing %s Root=%s: page %d, %d entries", repo_id, root_path, page_number, len(all_files)) if len(files) < page_size: break page_number += 1 return all_files + def _fetch_tree_page(self, path: str, params: dict[str, Any]) -> requests.Response: + parent = getattr(self._read_state, "control", None) + budget = float(REPO_TREE_PAGE_TIMEOUT_SECONDS) + if parent is not None: + parent.check() + budget = min(budget, max(0.001, parent.deadline - time.monotonic())) + try: + with _DownloadControl( + total_timeout=budget, + cancel_event=parent.cancel_event if parent is not None else None, + ) as page_control: + with self.read_context(page_control): + return self._request( + "GET", + path, + params=params, + timeout=(API_CONNECT_TIMEOUT, REPO_TREE_PAGE_TIMEOUT_SECONDS), + single_attempt=True, + ) + except _DownloadDeadlineError as exc: + if parent is not None: + parent.check() + raise RequestTimeoutError(f"列表单页请求超过 {REPO_TREE_PAGE_TIMEOUT_SECONDS}s") from exc + def _request_repo_tree( self, path: str, @@ -641,8 +781,8 @@ def _request_repo_tree( ) -> 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 + The server intermittently answers a tree request with ``403 access + denied`` 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 @@ -653,31 +793,48 @@ def _request_repo_tree( Until then a 403 is taken at face value, keeping a real permission error fast and honest. """ - last_error: PermissionDeniedError | None = None + last_error: NetworkError | ServerError | RateLimitError | 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: + resp = self._fetch_tree_page(path, params) + except (NetworkError, ServerError, RateLimitError, PermissionDeniedError) as error: + if isinstance(error, PermissionDeniedError) and not authorized: + raise + if not error.retryable and not isinstance(error, PermissionDeniedError): raise last_error = error if attempt < REPO_TREE_PAGE_MAX_ATTEMPTS - 1: - wait = min(2**attempt, REPO_TREE_PAGE_RETRY_MAX_DELAY_SECONDS) + requested_wait = getattr(error, "retry_after", None) + wait = min( + max(0, requested_wait) if requested_wait is not None else 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 %s: Root=%s page=%s listing failed (%s), retry in %ss (%s/%s)", repo_id, - params.get("Root") or params.get("PageNumber") or "/", + params.get("Root") or "/", + params.get("PageNumber", 1), + type(error).__name__, wait, + attempt + 1, + REPO_TREE_PAGE_MAX_ATTEMPTS, ) - time.sleep(wait) + control = getattr(self._read_state, "control", None) + if control is None: + time.sleep(wait) + else: + control.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 + if isinstance(last_error, PermissionDeniedError): + 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 + assert last_error is not None + raise last_error # ------------------------------------------------------------------ # Revisions @@ -890,14 +1047,16 @@ def upload_blob( if self._token: upload_headers["Authorization"] = f"Bearer {self._token}" upload_headers["Cookie"] = f"m_session_id={self._token}" + merged_headers = CaseInsensitiveDict(upload_headers) if headers: - upload_headers.update(headers) + merged_headers.update(headers) + merged_headers["User-Agent"] = append_user_agent(self._user_agent, merged_headers.get("User-Agent")) try: resp = self._session.put( upload_url, data=data, - headers=upload_headers, + headers=merged_headers, timeout=timeout or (UPLOAD_BLOB_CONNECT_TIMEOUT_SECONDS, UPLOAD_BLOB_READ_TIMEOUT_SECONDS), ) except requests.ConnectionError as exc: @@ -992,6 +1151,7 @@ def download_archive( f"{segment}/{repo_id}/archive/zip/{revision}", headers=headers, stream=True, + timeout=(API_CONNECT_TIMEOUT, DOWNLOAD_TIMEOUT), ) # ------------------------------------------------------------------ @@ -1011,10 +1171,23 @@ def download_stream( """ segment = _resolve_segment(repo_type) params = {"Revision": revision, "FilePath": file_path} + control = getattr(self._read_state, "control", None) + if control is None: + # Low-level calls do not add adapter retries; higher-level callers handle retries and integrity checks. + self._ensure_session_auth() + resp = self._read_session.get( + self.get_download_url(repo_id, repo_type, file_path, revision), + headers=self._headers(headers), + timeout=(API_CONNECT_TIMEOUT, DOWNLOAD_TIMEOUT), + stream=True, + ) + raise_for_status(resp) + return resp return self._request( "GET", f"{segment}/{repo_id}/repo", params=params, headers=headers, stream=True, + timeout=(API_CONNECT_TIMEOUT, DOWNLOAD_TIMEOUT), ) diff --git a/src/modelscope_hub/_openapi.py b/src/modelscope_hub/_openapi.py index ebdd71f..7217f09 100644 --- a/src/modelscope_hub/_openapi.py +++ b/src/modelscope_hub/_openapi.py @@ -22,22 +22,28 @@ from __future__ import annotations +import contextlib import random +import threading import time -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, Iterator, Mapping from pathlib import Path from typing import Any, BinaryIO from urllib.parse import urljoin, urlsplit import requests from requests.adapters import HTTPAdapter +from requests.structures import CaseInsensitiveDict +from ._download_control import _DownloadControl, _network_context, _ReadHTTPAdapter from .config import HubConfig, get_default_config from .constants import ( API_CONNECT_TIMEOUT, API_CONNECTION_POOL_MAXSIZE, API_MAX_RETRIES, API_TIMEOUT, + DOWNLOAD_CHUNK_SIZE, + DOWNLOAD_RETRY_MAX_DELAY_SECONDS, OPENAPI_PREFIX, TokenScope, ) @@ -64,6 +70,7 @@ UpdateSkillSettingsPayload, UpdateStudioSettingsPayload, ) +from .utils import append_user_agent, build_user_agent from .utils.logger import get_logger __all__ = ["OPERATION_REGISTRY", "OpenAPIClient"] @@ -196,6 +203,10 @@ class OpenAPIClient: ``(API_CONNECT_TIMEOUT, API_TIMEOUT)`` as ``(connect, read)`` tuple. max_retries: Maximum number of retry attempts for transient failures. + user_agent: + Additional information appended to the standard SDK UA. Custom UAs from + an injected session or request headers are appended too, preserving the + SDK identifier. Both sessions share the same base UA. """ def __init__( @@ -205,9 +216,14 @@ def __init__( session: requests.Session | None = None, timeout: float | None = None, max_retries: int | None = None, + user_agent: dict | str | None = None, ) -> None: self._config = config or get_default_config() + self._user_agent = build_user_agent(config=self._config, extra=user_agent) self._session = session or requests.Session() + session_ua = self._session.headers.get("User-Agent") + if session_ua and session_ua != requests.utils.default_user_agent(): + self._user_agent = append_user_agent(self._user_agent, session_ua) 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 @@ -225,6 +241,27 @@ def __init__( # Tri-state cache for whether this deployment serves ``GET /mcp/servers``. # ``None`` means "not probed yet"; see :meth:`list_mcp_servers`. self._mcp_list_supports_get: bool | None = None + self._read_state = threading.local() + self._read_session = requests.Session() + self._read_session.headers.update(self._session.headers) + adapter = _ReadHTTPAdapter( + max_retries=0, + pool_connections=API_CONNECTION_POOL_MAXSIZE, + pool_maxsize=API_CONNECTION_POOL_MAXSIZE, + ) + self._read_session.mount("https://", adapter) + self._read_session.mount("http://", adapter) + + @contextlib.contextmanager + def read_context(self, control: _DownloadControl) -> Iterator[None]: + """Use deadline-controlled transport for site probes, preserving regular OpenAPI policies.""" + previous = getattr(self._read_state, "control", None) + self._read_state.control = control + try: + with _network_context(control): + yield + finally: + self._read_state.control = previous # ------------------------------------------------------------------ # Lifecycle @@ -232,6 +269,7 @@ def __init__( def close(self) -> None: """Release the underlying HTTP session.""" self._session.close() + self._read_session.close() def __enter__(self) -> OpenAPIClient: return self @@ -493,27 +531,47 @@ def _send( else: merged_headers = {} request_cookies = {} + request_headers = CaseInsensitiveDict(merged_headers) if headers: - merged_headers.update(headers) + request_headers.update(headers) + # Inject UA at the sole send boundary for both sessions, anonymous requests and retries. + request_headers["User-Agent"] = append_user_agent(self._user_agent, request_headers.get("User-Agent")) method_upper = method.upper() attempts = max(1, self._max_retries) last_exc: BaseException | None = None + control = getattr(self._read_state, "control", None) for attempt in range(1, attempts + 1): _logger.debug("%s %s", method_upper, final_url) + request_timeout = timeout if timeout is not None else self._timeout + if control is not None: + connect, read = ( + (API_CONNECT_TIMEOUT, request_timeout) + if isinstance(request_timeout, (int, float)) + else request_timeout + ) + request_timeout = control.timeouts(connect, read) try: - response = self._session.request( + session = self._read_session if control is not None else self._session + stream_options: dict[str, Any] = {"stream": True} if control is not None else {} + response = session.request( method=method_upper, url=final_url, params=params, json=json_body, data=data, files=files, - headers=merged_headers, + headers=request_headers, cookies=request_cookies, - timeout=timeout if timeout is not None else self._timeout, + timeout=request_timeout, + **stream_options, ) + if control is not None: + control.track(response) + with control.response(response): + response._content = b"".join(control.chunks(response, DOWNLOAD_CHUNK_SIZE)) + response._content_consumed = True except requests.Timeout as exc: last_exc = RequestTimeoutError(f"Request timed out: {exc}") except requests.ConnectionError as exc: @@ -568,7 +626,10 @@ def _send( attempt, attempts, ) - time.sleep(backoff) + if control is None: + time.sleep(backoff) + else: + control.sleep(min(backoff, DOWNLOAD_RETRY_MAX_DELAY_SECONDS)) assert last_exc is not None # for type-checkers raise last_exc diff --git a/src/modelscope_hub/agent/_api.py b/src/modelscope_hub/agent/_api.py index 86c2d0f..8d9eb97 100644 --- a/src/modelscope_hub/agent/_api.py +++ b/src/modelscope_hub/agent/_api.py @@ -181,6 +181,9 @@ class AgentApi: API token (fallback: HubConfig default / ``ms login``). timeout : int HTTP request timeout in seconds. + user_agent : dict, str or None + Application or channel information appended to the SDK UA across + OpenAPI, legacy and blob requests. """ def __init__( @@ -190,12 +193,13 @@ def __init__( timeout: int = 60, *, config: HubConfig | None = None, + user_agent: dict | str | None = None, ): self._config = config or HubConfig(endpoint=endpoint, token=token) self.server = (self._config.endpoint or "").rstrip("/") self.token = self._config.token self.timeout = timeout - self._openapi = OpenAPIClient(config=self._config, timeout=float(timeout)) + self._openapi = OpenAPIClient(config=self._config, timeout=float(timeout), user_agent=user_agent) # ---- repository ---- diff --git a/src/modelscope_hub/api.py b/src/modelscope_hub/api.py index 2718b29..1453215 100644 --- a/src/modelscope_hub/api.py +++ b/src/modelscope_hub/api.py @@ -65,6 +65,7 @@ RepoInfo, UserInfo, ) +from .utils import build_user_agent from .utils.logger import get_logger __all__ = ["HubApi"] @@ -171,6 +172,9 @@ class HubApi: Defaults to ``https://modelscope.cn``. token : str or None, optional Override the API token. Takes precedence over ``config.token``. + user_agent : dict, str or None, optional + Application or channel information appended to the SDK UA and propagated + to OpenAPI, legacy, probe and download requests. Must not contain credentials. Examples -------- @@ -193,7 +197,9 @@ def __init__( *, endpoint: str | None = None, token: str | None = None, + user_agent: dict | str | None = None, ) -> None: + self._user_agent_extra = dict(user_agent) if isinstance(user_agent, dict) else user_agent base = config or get_default_config() if config is None and (endpoint is not None or token is not None): from dataclasses import replace @@ -224,19 +230,17 @@ def __init__( def openapi(self) -> OpenAPIClient: """Lazily-constructed OpenAPI client.""" if self._openapi is None: - self._openapi = OpenAPIClient(self._config) + self._openapi = OpenAPIClient(self._config, user_agent=self._user_agent_extra) return self._openapi @property def legacy(self) -> LegacyClient: """Lazily-constructed legacy ``/api/v1`` client.""" if self._legacy is None: - from .utils import build_user_agent - self._legacy = LegacyClient( token=self._config.token, endpoint=self._config.endpoint or DEFAULT_ENDPOINT, - user_agent=build_user_agent(self._config.get_session_id()), + user_agent=build_user_agent(config=self._config, extra=self._user_agent_extra), ) elif self._legacy.token != self._config.token: # Clears propagate as well as changes: a cached client left holding a @@ -248,7 +252,7 @@ def legacy(self) -> LegacyClient: def downloader(self) -> DownloadManager: """Lazily-constructed :class:`DownloadManager`.""" if self._downloader is None: - self._downloader = DownloadManager(self.legacy, self._config) + self._downloader = DownloadManager(self.legacy, self._config, user_agent=self._user_agent_extra) return self._downloader @property @@ -644,7 +648,9 @@ def _login_with_scoped_token(self, token: str) -> UserInfo | None: """ from .constants import API_CONNECT_TIMEOUT - probe = OpenAPIClient(self._config, timeout=API_CONNECT_TIMEOUT, max_retries=0) + probe = OpenAPIClient( + self._config, timeout=API_CONNECT_TIMEOUT, max_retries=0, user_agent=self._user_agent_extra + ) try: payload = probe.get_current_user() except Exception: # advisory only -- never mask the original failure @@ -694,7 +700,9 @@ def _explain_login_failure(self, token: str, exc: HubError) -> HubError: if not isinstance(exc, AuthenticationError): return exc peer = self._peer_site_endpoint() - if peer is None or not self._token_valid_on(token, peer): + if peer is None or not self._token_valid_on( + token, peer, user_agent=build_user_agent(config=self._config, extra=self._user_agent_extra) + ): return exc return AuthenticationError( f"{exc.message} This token is valid on {peer} instead; retry with " @@ -728,7 +736,7 @@ def site_key(url: str) -> str: return None @staticmethod - def _token_valid_on(token: str, endpoint: str) -> bool: + def _token_valid_on(token: str, endpoint: str, *, user_agent: str | None = None) -> bool: """Best-effort check of whether *token* authenticates against *endpoint*. Runs on the failure path only and is strictly advisory: any error means @@ -743,6 +751,7 @@ def _token_valid_on(token: str, endpoint: str) -> bool: endpoint=endpoint, timeout=API_CONNECT_TIMEOUT, max_retries=0, + user_agent=user_agent or build_user_agent(config=get_default_config()), ) try: probe.login(token) @@ -1248,12 +1257,29 @@ def resolve_endpoint_for_read( primary = DEFAULT_INTL_ENDPOINT if prefer_ai else DEFAULT_ENDPOINT fallback = DEFAULT_ENDPOINT if prefer_ai else DEFAULT_INTL_ENDPOINT - primary_probe = HubApi(endpoint=primary, token=effective_token) - if primary_probe.repo_exists(repo_id, repo_type): + from ._download_control import _DownloadControl + from .constants import REPO_TREE_PAGE_TIMEOUT_SECONDS + + def exists(endpoint: str) -> bool: + from dataclasses import replace + + probe = HubApi( + config=replace(self._config), + endpoint=endpoint, + token=effective_token, + user_agent=self._user_agent_extra, + ) + try: + with _DownloadControl(total_timeout=REPO_TREE_PAGE_TIMEOUT_SECONDS) as control: + with probe.openapi.read_context(control): + return probe.repo_exists(repo_id, repo_type) + finally: + probe.openapi.close() + + if exists(primary): return primary - fallback_probe = HubApi(endpoint=fallback, token=effective_token) - if fallback_probe.repo_exists(repo_id, repo_type): + if exists(fallback): logger.warning( "Repo %s not found on %s, using %s instead.", repo_id, @@ -1472,6 +1498,7 @@ def download_file( expected_sha256: str | None = None, local_files_only: bool = False, user_agent: dict | str | None = None, + file_size: int | None = None, ) -> Path: """Download a single file from a repository. @@ -1503,6 +1530,10 @@ def download_file( Raises :class:`CacheNotFound` if the file is not cached. user_agent : dict, str or None, optional Custom user-agent info appended to the default UA string. + file_size : int, optional + Expected size in bytes, e.g. from a file listing. The download is + rejected when the server delivers a different size, and files + larger than the parallel threshold use ranged parallel streams. Returns ------- @@ -1513,6 +1544,9 @@ def download_file( ------ NotExistError When the file or repository does not exist. + FileIntegrityError + When the content keeps failing ``expected_sha256`` or (as + :class:`FileSizeMismatchError`) ``file_size``. Examples -------- @@ -1542,6 +1576,7 @@ def download_file( expected_sha256=expected_sha256, local_files_only=local_files_only, user_agent=user_agent, + file_size=file_size, ) def download_repo( @@ -1558,6 +1593,7 @@ def download_repo( local_files_only: bool = False, user_agent: dict | str | None = None, progress_callbacks: list[type[ProgressCallback]] | None = None, + file_list: list[dict] | None = None, ) -> Path: """Download an entire repository snapshot. @@ -1589,12 +1625,23 @@ def download_repo( progress_callbacks : list of ProgressCallback subclasses, optional Callback *classes* (not instances); each is instantiated per file to report byte-level download progress. + file_list : list of dict, optional + Previously fetched legacy file entries (Path, Size, Sha256). When + provided, listing requests are skipped. The caller must ensure the + entries belong to the same repository and revision; filters still apply. Returns ------- Path Absolute path to the snapshot/local directory. + Raises + ------ + FileDownloadError + When any file failed after all files were attempted. Completed + files are kept, so calling again fetches only the failed ones; + ``failed_files`` maps each failed path to its error. + Examples -------- Download only the tokenizer assets of a model: @@ -1627,6 +1674,7 @@ def download_repo( local_files_only=local_files_only, user_agent=user_agent, progress_callbacks=progress_callbacks, + file_list=file_list, ) def list_repo_files( diff --git a/src/modelscope_hub/cli/agent.py b/src/modelscope_hub/cli/agent.py index 0aced6c..f8a1153 100644 --- a/src/modelscope_hub/cli/agent.py +++ b/src/modelscope_hub/cli/agent.py @@ -92,7 +92,7 @@ def _cmd_list(owner, page_number, page_size, *, endpoint, token) -> int: """List remote agent repositories.""" if not endpoint: return _fail("not logged in. Provide endpoint.") - client = AgentApi(endpoint=endpoint, token=token or "") + client = AgentApi(endpoint=endpoint, token=token or "", user_agent={"invoked_by": "cli"}) try: result = client.list_agents(owner=owner, page_number=page_number, page_size=page_size) except APIError as e: @@ -144,7 +144,7 @@ def _cmd_download(repo, local_dir, revision, *, endpoint, token, username) -> in ) group, name = _resolve_repo(repo, username or "") - client = AgentApi(endpoint=endpoint, token=token or "") + client = AgentApi(endpoint=endpoint, token=token or "", user_agent={"invoked_by": "cli"}) try: if client.repo_info(group, name) is None: return _fail(f"repository {group}/{name} not found.") @@ -209,7 +209,7 @@ def _cmd_upload(repo, local_dir, revision, dry_run, *, endpoint, token, username ) group, name = _resolve_repo(repo, username or "") - client = AgentApi(endpoint=endpoint, token=token) + client = AgentApi(endpoint=endpoint, token=token, user_agent={"invoked_by": "cli"}) try: if not client.check_repo(group, name): client.create_repo(group, name, visibility=visibility) @@ -507,7 +507,7 @@ def execute(self) -> None: from .._openapi import OpenAPIClient try: - openapi = OpenAPIClient(config=config) + openapi = OpenAPIClient(config=config, user_agent={"invoked_by": "cli"}) user_data = openapi.get_current_user() or {} username = user_data.get("username") or user_data.get("Username") or "" except Exception: diff --git a/src/modelscope_hub/cli/base.py b/src/modelscope_hub/cli/base.py index f2de5f6..cf66030 100644 --- a/src/modelscope_hub/cli/base.py +++ b/src/modelscope_hub/cli/base.py @@ -75,6 +75,7 @@ def make_api(args: Namespace) -> HubApi: return HubApi( token=getattr(args, "token", None), endpoint=getattr(args, "endpoint", None), + user_agent={"invoked_by": "cli"}, ) diff --git a/src/modelscope_hub/cli/download.py b/src/modelscope_hub/cli/download.py index 328fc17..e041662 100644 --- a/src/modelscope_hub/cli/download.py +++ b/src/modelscope_hub/cli/download.py @@ -146,7 +146,7 @@ def _make_api_with_endpoint(self) -> HubApi: if endpoint: endpoint = _resolve_cli_endpoint(endpoint) - return HubApi(token=token, endpoint=endpoint) + return HubApi(token=token, endpoint=endpoint, user_agent={"invoked_by": "cli"}) api = make_api(self.args) try: @@ -154,7 +154,7 @@ def _make_api_with_endpoint(self) -> HubApi: self.args.repo_id, repo_type=self.args.repo_type, ) - return HubApi(token=token, endpoint=resolved) + return HubApi(token=token, endpoint=resolved, user_agent={"invoked_by": "cli"}) except Exception: return api diff --git a/src/modelscope_hub/compat/file_download.py b/src/modelscope_hub/compat/file_download.py index 5d81739..1494fa1 100644 --- a/src/modelscope_hub/compat/file_download.py +++ b/src/modelscope_hub/compat/file_download.py @@ -13,6 +13,7 @@ from ..api import HubApi from ..constants import RepoType from ..errors import AuthenticationError, NotExistError, PermissionDeniedError +from ..utils import append_user_agent def _resolve_legacy_paths( @@ -54,11 +55,11 @@ def model_file_download( stacklevel=2, ) - api = HubApi(token=token, endpoint=endpoint) + api = HubApi(token=token, endpoint=endpoint, user_agent=append_user_agent("invoked_by/compat", user_agent)) if endpoint is None and not local_files_only: try: endpoint = api.resolve_endpoint_for_read(model_id, repo_type="model") - api = HubApi(token=token, endpoint=endpoint) + api = HubApi(token=token, endpoint=endpoint, user_agent=append_user_agent("invoked_by/compat", user_agent)) except Exception: pass effective_cache, effective_local = _resolve_legacy_paths( @@ -104,11 +105,11 @@ def dataset_file_download( stacklevel=2, ) - api = HubApi(token=token, endpoint=endpoint) + api = HubApi(token=token, endpoint=endpoint, user_agent=append_user_agent("invoked_by/compat", user_agent)) if endpoint is None and not local_files_only: try: endpoint = api.resolve_endpoint_for_read(dataset_id, repo_type="dataset") - api = HubApi(token=token, endpoint=endpoint) + api = HubApi(token=token, endpoint=endpoint, user_agent=append_user_agent("invoked_by/compat", user_agent)) except Exception: pass effective_cache, effective_local = _resolve_legacy_paths( diff --git a/src/modelscope_hub/compat/hub_api.py b/src/modelscope_hub/compat/hub_api.py index f011f91..5d91634 100644 --- a/src/modelscope_hub/compat/hub_api.py +++ b/src/modelscope_hub/compat/hub_api.py @@ -125,7 +125,7 @@ def __init__( if endpoint and not endpoint.startswith("http"): endpoint = f"https://{endpoint}" self._endpoint = endpoint - self._api = HubApi(endpoint=endpoint, token=token) + self._api = HubApi(endpoint=endpoint, token=token, user_agent={"invoked_by": "compat"}) # ------------------------------------------------------------------ # Authentication @@ -215,7 +215,7 @@ def delete_files( """ api = self._api if token or endpoint: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) return api.delete_files( repo_id, repo_type, @@ -242,7 +242,7 @@ def create_repo( """Create a repository (legacy signature).""" api = self._api if token or endpoint: - api = HubApi(token=token, endpoint=endpoint or self._endpoint) + api = HubApi(token=token, endpoint=endpoint or self._endpoint, user_agent={"invoked_by": "compat"}) if create_default_config: kwargs["create_default_config"] = True try: @@ -301,7 +301,7 @@ def _create_aigc_model(self, model_id: str, aigc_model: Any, kwargs: dict[str, A api = self._api if token or endpoint: - api = HubApi(token=token, endpoint=endpoint or self._endpoint) + api = HubApi(token=token, endpoint=endpoint or self._endpoint, user_agent={"invoked_by": "compat"}) owner, name = api._parse_repo_id(model_id) normalised_visibility = api._normalize_visibility(visibility) if normalised_visibility is None: @@ -376,7 +376,7 @@ def create_model_tag( api = self._api if token or endpoint: - api = HubApi(token=token, endpoint=endpoint or self._endpoint) + api = HubApi(token=token, endpoint=endpoint or self._endpoint, user_agent={"invoked_by": "compat"}) if aigc_model is None: if readme_content is not None: raise InvalidParameter("readme_content is only supported for AIGC model tags") @@ -507,6 +507,7 @@ def repo_exists( api = self._api if endpoint is not None or token is not None: api = HubApi( + user_agent={"invoked_by": "compat"}, endpoint=endpoint or self._api._config.endpoint, token=token or self._api._config.token, ) @@ -595,7 +596,9 @@ def _api_for(self, control: dict[str, Any]) -> HubApi: endpoint = control.get("endpoint") config = self._api._config if (token and token != config.token) or (endpoint and endpoint != config.endpoint): - return HubApi(token=token or config.token, endpoint=endpoint or config.endpoint) + return HubApi( + token=token or config.token, endpoint=endpoint or config.endpoint, user_agent={"invoked_by": "compat"} + ) return self._api def deploy_studio(self, studio_id: str, **kwargs: Any) -> dict: @@ -826,7 +829,7 @@ def list_models( """List models owned by a user/org.""" api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) page = api.list_repos( RepoType.MODEL, owner=owner_or_group, @@ -883,7 +886,7 @@ def list_datasets( actual_owner = owner or owner_or_group api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) return api.list_repos( RepoType.DATASET, owner=actual_owner, @@ -908,7 +911,7 @@ def create_dataset( repo_id = f"{namespace}/{dataset_name}" api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) api.create_repo( repo_id, repo_type=RepoType.DATASET, @@ -950,7 +953,7 @@ def get_dataset( ) api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) return api.get_repo(dataset_id, RepoType.DATASET, revision=revision) def delete_model( @@ -962,7 +965,7 @@ def delete_model( """Delete a model repository.""" api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) api.delete_repo(model_id, RepoType.MODEL) def delete_dataset( @@ -980,7 +983,7 @@ def delete_dataset( ) api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) api.delete_repo(dataset_id, RepoType.DATASET) def get_dataset_files( @@ -999,7 +1002,7 @@ def get_dataset_files( """Get dataset file tree via the datahub API.""" api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) if dataset_hub_id is None: if "/" in repo_id: @@ -1038,7 +1041,7 @@ def get_dataset_id_and_type( """Get the dataset hub-internal id and type.""" api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) cache_key = (namespace, dataset_name, endpoint or self._endpoint) cached = LegacyHubApi._dataset_id_type_cache.get(cache_key) @@ -1063,7 +1066,7 @@ def get_dataset_meta_file_list( """Get the meta file-list of the dataset.""" api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) params = {"Revision": revision} resp = api.legacy._request("GET", f"datasets/{dataset_id}/repo/tree", params=params) @@ -1107,7 +1110,7 @@ def get_dataset_meta_files_local_paths( api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) local_paths: dict[str, list] = defaultdict(list) dataset_formation = DatasetFormations(dataset_type) @@ -1197,7 +1200,7 @@ def get_dataset_access_config( """Get STS token config for dataset OSS access.""" api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) resp = api.legacy._request( "GET", f"datasets/{namespace}/{dataset_name}/ststoken", @@ -1217,7 +1220,7 @@ def get_dataset_access_config_session( """Get STS token config with session-based auth.""" api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) resp = api.legacy._request( "GET", f"datasets/{namespace}/{dataset_name}/ststoken", @@ -1237,7 +1240,7 @@ def get_dataset_access_config_for_unzipped( """Get STS config for unzipped dataset files.""" api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) # Get visibility data = api.legacy.get_repo_info(f"{namespace}/{dataset_name}", RepoType.DATASET) @@ -1270,7 +1273,7 @@ def dataset_download_statistics( try: api = self._api if token: - api = HubApi(endpoint=endpoint or self._endpoint, token=token) + api = HubApi(endpoint=endpoint or self._endpoint, token=token, user_agent={"invoked_by": "compat"}) api.legacy._request( "POST", f"datasets/{namespace}/{dataset_name}/download/increase", diff --git a/src/modelscope_hub/compat/snapshot_download.py b/src/modelscope_hub/compat/snapshot_download.py index 3981972..b06208d 100644 --- a/src/modelscope_hub/compat/snapshot_download.py +++ b/src/modelscope_hub/compat/snapshot_download.py @@ -16,6 +16,7 @@ from ..api import HubApi from ..constants import RepoType from ..errors import AuthenticationError, NotExistError, PermissionDeniedError +from ..utils import append_user_agent from ..utils.patterns import normalize_patterns from .file_download import _resolve_legacy_paths @@ -69,14 +70,14 @@ def snapshot_download( include = _normalize_pattern(allow_patterns) or _normalize_pattern(allow_file_pattern) exclude = _normalize_pattern(ignore_patterns) or _normalize_pattern(ignore_file_pattern) - api = HubApi(token=token, endpoint=endpoint) + api = HubApi(token=token, endpoint=endpoint, user_agent=append_user_agent("invoked_by/compat", user_agent)) if endpoint is None and not local_files_only: try: endpoint = api.resolve_endpoint_for_read( effective_id, repo_type=effective_type, ) - api = HubApi(token=token, endpoint=endpoint) + api = HubApi(token=token, endpoint=endpoint, user_agent=append_user_agent("invoked_by/compat", user_agent)) except Exception: pass effective_cache, effective_local = _resolve_legacy_paths( @@ -132,11 +133,11 @@ def dataset_snapshot_download( include = _normalize_pattern(allow_patterns) or _normalize_pattern(allow_file_pattern) exclude = _normalize_pattern(ignore_patterns) or _normalize_pattern(ignore_file_pattern) - api = HubApi(token=token, endpoint=endpoint) + api = HubApi(token=token, endpoint=endpoint, user_agent=append_user_agent("invoked_by/compat", user_agent)) if endpoint is None and not local_files_only: try: endpoint = api.resolve_endpoint_for_read(dataset_id, repo_type="dataset") - api = HubApi(token=token, endpoint=endpoint) + api = HubApi(token=token, endpoint=endpoint, user_agent=append_user_agent("invoked_by/compat", user_agent)) except Exception: pass effective_cache, effective_local = _resolve_legacy_paths( diff --git a/src/modelscope_hub/constants.py b/src/modelscope_hub/constants.py index e317f99..c4ca763 100644 --- a/src/modelscope_hub/constants.py +++ b/src/modelscope_hub/constants.py @@ -614,7 +614,7 @@ def _env_register( """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. +``403 access denied`` 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. @@ -696,19 +696,64 @@ def _env_register( DOWNLOAD_RETRY_TIMES: int = _env_int( "MODELSCOPE_DOWNLOAD_MAX_RETRIES", 5, - "Per-file download retry count", + "Per-file download retry count (one retry layer; zero disables retries)", "Download", "DOWNLOAD_RETRY_TIMES", + allow_zero=True, ) DOWNLOAD_TIMEOUT: int = _env_int( "MODELSCOPE_DOWNLOAD_TIMEOUT", 60, - "Per-file download timeout (seconds)", + "Download idle read timeout (seconds; single stream and range parts)", "Download", "DOWNLOAD_TIMEOUT", ) +DOWNLOAD_TOTAL_TIMEOUT: int = _env_int( + "MODELSCOPE_DOWNLOAD_TOTAL_TIMEOUT", + 3600, + "Total time budget per file including retries, probes and lock wait (seconds)", + "Download", +) +DOWNLOAD_MIN_BYTES_PER_SECOND: int = _env_int( + "MODELSCOPE_DOWNLOAD_MIN_BYTES_PER_SECOND", + 1024, + "Minimum transfer rate per connection (B/s; zero disables rate checking)", + "Download", + allow_zero=True, +) +DOWNLOAD_LOW_SPEED_WINDOW_SECONDS: int = _env_int( + "MODELSCOPE_DOWNLOAD_LOW_SPEED_WINDOW_SECONDS", + 60, + "Observation window for minimum download rate (seconds)", + "Download", +) +DOWNLOAD_RETRY_MAX_DELAY_SECONDS: int = _env_int( + "MODELSCOPE_DOWNLOAD_RETRY_MAX_DELAY_SECONDS", + 60, + "Maximum download Retry-After or backoff delay (seconds)", + "Download", +) +DOWNLOAD_FILE_LOCK_TIMEOUT_SECONDS: int = _env_int( + "MODELSCOPE_DOWNLOAD_FILE_LOCK_TIMEOUT_SECONDS", + 600, + "Maximum wait for another downloader's file lock (seconds)", + "Download", +) +REPO_TREE_PAGE_TIMEOUT_SECONDS: int = _env_int( + "MODELSCOPE_REPO_TREE_PAGE_TIMEOUT_SECONDS", + 15, + "Idle timeout per file-tree page (seconds; no nested HTTP retries)", + "Network", +) +REPO_TREE_LIST_TIMEOUT_SECONDS: int = _env_int( + "MODELSCOPE_REPO_TREE_LIST_TIMEOUT_SECONDS", + 3600, + "Total budget for enumerating a repository (seconds)", + "Network", +) + DOWNLOAD_PART_SIZE: int = _env_int_mb( "MODELSCOPE_DOWNLOAD_PART_SIZE_MB", 160, @@ -1302,6 +1347,13 @@ def get_upload_ignore_file_pattern() -> str | None: "DOWNLOAD_PART_SIZE", "DOWNLOAD_RETRY_TIMES", "DOWNLOAD_TIMEOUT", + "DOWNLOAD_TOTAL_TIMEOUT", + "DOWNLOAD_MIN_BYTES_PER_SECOND", + "DOWNLOAD_LOW_SPEED_WINDOW_SECONDS", + "DOWNLOAD_RETRY_MAX_DELAY_SECONDS", + "DOWNLOAD_FILE_LOCK_TIMEOUT_SECONDS", + "REPO_TREE_PAGE_TIMEOUT_SECONDS", + "REPO_TREE_LIST_TIMEOUT_SECONDS", "ENV_AGENT_PLUGIN_REPO", "ENV_FILE_LOCK", "ENV_CACHE", diff --git a/src/modelscope_hub/errors.py b/src/modelscope_hub/errors.py index c2e1623..23b821c 100644 --- a/src/modelscope_hub/errors.py +++ b/src/modelscope_hub/errors.py @@ -18,10 +18,13 @@ from __future__ import annotations import re +from collections.abc import Mapping from typing import TYPE_CHECKING, Any from urllib.parse import parse_qs, urlparse, urlunparse if TYPE_CHECKING: # pragma: no cover - type-only imports + from os import PathLike + from requests import Response # --------------------------------------------------------------------------- @@ -347,6 +350,48 @@ class StorageError(HubError): suggestion = "File upload/download failed (storage service error). Please retry later." +class FileDownloadError(StorageError): + """Raised after a snapshot download attempted every file and some failed (E1003). + + Files that downloaded successfully stay in place, so re-running the same + call fetches only the paths in :attr:`failed_files`. :attr:`retryable` is + true only when every individual failure is itself retryable. + """ + + suggestion = ( + "Some files failed to download; completed files are kept. " + "Re-run the same command to retry only the failed files." + ) + _MAX_LISTED = 10 + + def __init__( + self, + message: str, + *, + failed_files: Mapping[str, BaseException] | None = None, + total: int | None = None, + output_dir: str | PathLike[str] | None = None, + ) -> None: + super().__init__(message) + self.failed_files: dict[str, BaseException] = dict(failed_files or {}) + self.total = total + self.output_dir = output_dir + self.retryable = bool(self.failed_files) and all( + getattr(exc, "retryable", False) for exc in self.failed_files.values() + ) + + def __str__(self) -> str: + lines = [super().__str__()] + for path, exc in list(self.failed_files.items())[: self._MAX_LISTED]: + # Only the headline: API errors append multi-line request detail. + text = str(exc) + lines.append(f" - {path}: {text.splitlines()[0] if text else exc.__class__.__name__}") + hidden = len(self.failed_files) - self._MAX_LISTED + if hidden > 0: + lines.append(f" ... and {hidden} more") + return "\n".join(lines) + + class FileIntegrityError(HubError): """Raised when a downloaded or uploaded file fails integrity validation (E2020).""" @@ -355,6 +400,23 @@ class FileIntegrityError(HubError): suggestion = "File SHA256 checksum mismatch. Will retry automatically." +class FileSizeMismatchError(FileIntegrityError): + """Raised when the server delivers a different number of bytes than the repository lists (E2020). + + The transfer itself completed, so the remote object (or the replica / CDN + node serving it) is empty, truncated or out of date rather than the + connection being cut. Repeating the same request rarely helps, so it is not + marked retryable: the download makes one more attempt over the default + route and then reports the file. + """ + + retryable = False + suggestion = ( + "The server returned content that does not match the repository's file size. " + "Report the file path and request ids in this message to the repository owner or ModelScope." + ) + + # -- Cache errors (E1022) --------------------------------------------------- class CacheError(HubError): """Base for local cache filesystem or metadata corruption.""" @@ -433,11 +495,11 @@ class NotSupportedError(HubError): # --------------------------------------------------------------------------- _BUSINESS_CODE_MAP: dict[int, type[APIError]] = { # -> E3001, served with HTTP 400 by POST /api/v1/login on both sites - 10010103009: AuthenticationError, # AccessToken 无效或过期 + 10010103009: AuthenticationError, # AccessToken is invalid or expired # -> E3026 - 10020101001: AlreadyExistsError, # 国内站 - 数据集已存在 - 10010101001: AlreadyExistsError, # 国内站 - 模型已存在 - 10010202004: AlreadyExistsError, # 国际站 - 名称已被使用 + 10020101001: AlreadyExistsError, # China site: dataset already exists + 10010101001: AlreadyExistsError, # China site: model already exists + 10010202004: AlreadyExistsError, # International site: name is already in use } # The legacy commit endpoint overloads 400/409 for optimistic branch conflicts @@ -699,7 +761,9 @@ def is_repo_exists_error(exc: BaseException) -> bool: "CacheError", "CacheNotFound", "CorruptedCacheException", + "FileDownloadError", "FileIntegrityError", + "FileSizeMismatchError", "NetworkError", "RequestTimeoutError", "StorageError", diff --git a/src/modelscope_hub/utils/__init__.py b/src/modelscope_hub/utils/__init__.py index d5363ae..2117b67 100644 --- a/src/modelscope_hub/utils/__init__.py +++ b/src/modelscope_hub/utils/__init__.py @@ -2,6 +2,9 @@ from __future__ import annotations +import uuid +from typing import TYPE_CHECKING + from .file_utils import compute_hash, ensure_dir, get_cache_dir, get_file_size from .format import format_size, format_timesince, tabulate from .logger import get_logger @@ -9,7 +12,14 @@ from .patterns import extract_common_prefix, normalize_patterns from .time_utils import parse_timestamp +if TYPE_CHECKING: + from ..config import HubConfig + +# Keep retries and transports in one session when credential storage is unavailable. +_FALLBACK_SESSION_ID = uuid.uuid4().hex + __all__ = [ + "append_user_agent", "build_user_agent", "compute_hash", "encode_media_to_base64", @@ -29,6 +39,8 @@ def build_user_agent( session_id: str | None = None, extra: dict | str | None = None, + *, + config: HubConfig | None = None, ) -> str: """Build the standard ModelScope Hub SDK User-Agent string. @@ -39,12 +51,21 @@ def build_user_agent( Falls back to a random UUID if not provided. extra : dict, str or None Additional key/value pairs or free-form string appended to the UA. + config : HubConfig, optional + Resolve the persisted session_id from this config; use a process-stable + ID if storage is unavailable. An explicit session_id avoids storage access. """ import os import platform - import uuid from .. import __version__ + from ..errors import CacheError + + if config is not None and session_id is None: + try: + session_id = config.get_session_id() + except (CacheError, OSError): + session_id = _FALLBACK_SESSION_ID env = os.environ.get("MODELSCOPE_CLOUD_ENVIRONMENT", "custom") user_name = os.environ.get("MODELSCOPE_CLOUD_USERNAME", "unknown") @@ -55,8 +76,16 @@ def build_user_agent( f"session_id/{sid}; platform/{platform.platform()}; " f"processor/{platform.processor()}; env/{env}; user/{user_name}" ) + return append_user_agent(ua, extra) + + +def append_user_agent(base: str, extra: dict | str | None = None) -> str: + """Append custom UA information without removing the SDK prefix or duplicating fields.""" if isinstance(extra, dict): - ua += "; " + "; ".join(f"{k}/{v}" for k, v in extra.items()) - elif isinstance(extra, str): - ua += "; " + extra - return ua + extra = "; ".join(f"{k}/{v}" for k, v in extra.items()) + parts = base.split("; ") + if extra: + for part in extra.split("; "): + if part and part not in parts: + parts.append(part) + return "; ".join(parts) diff --git a/src/modelscope_hub/version.py b/src/modelscope_hub/version.py index a59471b..8844a57 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.6+main" +__version__ = "0.5.0+main" diff --git a/tests/cli/test_base.py b/tests/cli/test_base.py index 5bd2d3b..0ef7da3 100644 --- a/tests/cli/test_base.py +++ b/tests/cli/test_base.py @@ -97,21 +97,23 @@ def test_make_api_with_token(self): args = Namespace(token="my_token", endpoint=None) with patch("modelscope_hub.cli.base.HubApi") as mock_hub: make_api(args) - mock_hub.assert_called_once_with(token="my_token", endpoint=None) + mock_hub.assert_called_once_with(token="my_token", endpoint=None, user_agent={"invoked_by": "cli"}) def test_make_api_without_token(self): """make_api passes None when no token.""" args = Namespace(token=None, endpoint=None) with patch("modelscope_hub.cli.base.HubApi") as mock_hub: make_api(args) - mock_hub.assert_called_once_with(token=None, endpoint=None) + mock_hub.assert_called_once_with(token=None, endpoint=None, user_agent={"invoked_by": "cli"}) def test_make_api_with_endpoint(self): """make_api passes endpoint to HubApi.""" args = Namespace(token=None, endpoint="https://custom.endpoint.com") with patch("modelscope_hub.cli.base.HubApi") as mock_hub: make_api(args) - mock_hub.assert_called_once_with(token=None, endpoint="https://custom.endpoint.com") + mock_hub.assert_called_once_with( + token=None, endpoint="https://custom.endpoint.com", user_agent={"invoked_by": "cli"} + ) class TestAddRepoTypeArg: diff --git a/tests/test_agent_plugin.py b/tests/test_agent_plugin.py index f32e6c0..f07817a 100644 --- a/tests/test_agent_plugin.py +++ b/tests/test_agent_plugin.py @@ -325,12 +325,14 @@ def test_verify_manifest_refuses_keys_pointing_outside_the_package(tmp_path, key # --------------------------------------------------------------------------- # execution audit / no opt-in # --------------------------------------------------------------------------- -def test_log_execution_records_the_build_before_import(tmp_path, caplog): +def test_log_execution_records_the_build_before_import(tmp_path, caplog, monkeypatch): """There is no opt-in to wait for any more -- the allow-list is the authorisation -- but executing downloaded code still deserves an audit line naming the exact build, emitted before the import so a crash during it leaves a trace of what was being loaded.""" spec = spec_for(make_plugin(tmp_path)) + # Capture on the target logger because SDK logs do not propagate to root; monkeypatch restores its handlers. + monkeypatch.setattr(_plugin.logger, "handlers", [*_plugin.logger.handlers, caplog.handler]) with caplog.at_level("INFO", logger="modelscope_hub.agent"): _plugin.log_execution(spec) for expected in (spec.repo_id, spec.revision, spec.entry_module, "9.9.9"): diff --git a/tests/test_download_integrity.py b/tests/test_download_integrity.py new file mode 100644 index 0000000..9c7a6d7 --- /dev/null +++ b/tests/test_download_integrity.py @@ -0,0 +1,802 @@ +"""Download integrity hardening: size validation, resume, cleanup and failure reporting. + +A real HTTP server on 127.0.0.1 stands in for the hub and its LFS CDN, so the +tests exercise requests/urllib3 streaming, redirects, ``Range`` handling and +chunked bodies end to end. It mirrors what ModelScope does in production: +LFS files redirect to a CDN that honours ``Range`` (206 + ``Content-Range``, +416 past the end), while files served directly by the API ignore ``Range`` and +answer 200 with a chunked body. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import threading +import time +from collections.abc import Iterator +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from urllib.parse import parse_qs, quote, unquote, urlparse + +import pytest + +from modelscope_hub import _download +from modelscope_hub.api import HubApi +from modelscope_hub.cli.main import run_cmd +from modelscope_hub.errors import ( + FileDownloadError, + FileIntegrityError, + FileSizeMismatchError, + NetworkError, +) + +REPO = "owner/data" + + +def _sha256(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +@dataclass +class _RemoteFile: + content: bytes + # "lfs": 302 to a Range-capable CDN; "direct": 200 chunked, Range ignored. + mode: str = "lfs" + # Bytes actually served when they differ from what the listing describes. + served: bytes | None = None + # Served instead while the intra-cloud region header is present. + served_with_region: bytes | None = None + # The first body response stops after this many bytes and drops the socket. + truncate_first: int | None = None + header_delay: float = 0 + trickle_interval: float = 0 + header_trickle: bool = False + status: int = 200 + retry_after: str | None = None + + def body(self, region: str | None) -> bytes: + if region and self.served_with_region is not None: + return self.served_with_region + return self.content if self.served is None else self.served + + +@dataclass +class _Request: + kind: str + path: str + range: str | None + region: str | None + + +@dataclass +class _FakeHub: + url: str = "" + files: dict[str, _RemoteFile] = field(default_factory=dict) + requests: list[_Request] = field(default_factory=list) + truncated: set[str] = field(default_factory=set) + lock: threading.Lock = field(default_factory=threading.Lock) + + def add(self, path: str, content: bytes, **kwargs) -> bytes: + self.files[path] = _RemoteFile(content, **kwargs) + return content + + def seen(self, path: str, kind: str) -> list[_Request]: + return [r for r in self.requests if r.path == path and r.kind == kind] + + +def _handler(hub: _FakeHub) -> type[BaseHTTPRequestHandler]: + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *args) -> None: + pass + + def _record(self, kind: str, path: str) -> None: + with hub.lock: + hub.requests.append( + _Request(kind, path, self.headers.get("Range"), self.headers.get("x-aliyun-region-id")) + ) + + def _truncate_now(self, path: str) -> bool: + with hub.lock: + if hub.files[path].truncate_first is None or path in hub.truncated: + return False + hub.truncated.add(path) + return True + + def _send(self, status: int, body: bytes, headers: dict[str, str] | None = None) -> None: + self.send_response(status) + for name, value in (headers or {}).items(): + self.send_header(name, value) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self) -> None: # noqa: N802 - http.server API + url = urlparse(self.path) + if url.path.endswith("/repo/tree"): + entries = [ + {"Path": p, "Type": "blob", "Size": len(f.content), "Sha256": _sha256(f.content)} + for p, f in sorted(hub.files.items()) + if not parse_qs(url.query).get("Root") + or p.startswith(parse_qs(url.query)["Root"][0].rstrip("/") + "/") + ] + body = json.dumps({"Code": 200, "Data": {"Files": entries}}).encode() + return self._send(200, body, {"Content-Type": "application/json"}) + if url.path.endswith("/repo"): + path = parse_qs(url.query)["FilePath"][0] + self._record("api", path) + if hub.files[path].mode == "direct": + return self._serve_whole(path) + self.send_response(302) + self.send_header("Location", f"/cdn/{quote(path)}") + self.send_header("Content-Length", "0") + self.end_headers() + return None + if url.path.startswith("/cdn/"): + path = unquote(url.path[len("/cdn/") :]) + self._record("cdn", path) + return self._serve_ranged(path) + return self._send(404, b'{"Code": 404}', {"Content-Type": "application/json"}) + + def _serve_whole(self, path: str) -> None: + body = hub.files[path].body(self.headers.get("x-aliyun-region-id")) + self.send_response(200) + self.send_header("Transfer-Encoding", "chunked") + self.end_headers() + for start in range(0, len(body), 8192): + piece = body[start : start + 8192] + self.wfile.write(f"{len(piece):x}\r\n".encode() + piece + b"\r\n") + self.wfile.write(b"0\r\n\r\n") + + def _serve_ranged(self, path: str) -> None: + remote = hub.files[path] + if remote.header_delay: + time.sleep(remote.header_delay) + if remote.header_trickle: + for byte in b"HTTP/1.1 200 OK\r\nContent-Length: 100\r\n\r\n": + self.wfile.write(bytes([byte])) + self.wfile.flush() + time.sleep(0.05) + return None + if remote.status >= 400: + return self._send(remote.status, b"{}", {"Retry-After": remote.retry_after or "0"}) + body = remote.body(self.headers.get("x-aliyun-region-id")) + headers = { + "Accept-Ranges": "bytes", + "X-Cache": "HIT TCP_MEM_HIT", + "EagleId": "test-eagle", + "X-Linked-ETag": _sha256(remote.content), + } + status, payload = 200, body + requested = self.headers.get("Range") + if requested: + first, _, last = requested.removeprefix("bytes=").partition("-") + start = int(first) + if start >= len(body): + return self._send(416, b"", {"Content-Range": f"bytes */{len(body)}"}) + end = min(int(last), len(body) - 1) if last else len(body) - 1 + status, payload = 206, body[start : end + 1] + headers["Content-Range"] = f"bytes {start}-{end}/{len(body)}" + if self._truncate_now(path): + cut = hub.files[path].truncate_first or 0 + self.send_response(status) + for name, value in headers.items(): + self.send_header(name, value) + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload[:cut]) + self.close_connection = True + return None + if remote.trickle_interval: + self.send_response(status) + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + for byte in payload: + self.wfile.write(bytes([byte])) + self.wfile.flush() + time.sleep(remote.trickle_interval) + return None + return self._send(status, payload, headers) + + return Handler + + +class _QuietServer(ThreadingHTTPServer): + daemon_threads = True + + def handle_error(self, request, client_address) -> None: + # Clients abandoning a body they rejected are expected here. + pass + + +@pytest.fixture() +def hub(monkeypatch) -> Iterator[_FakeHub]: + monkeypatch.setenv("NO_PROXY", "127.0.0.1,localhost") + monkeypatch.setenv("MODELSCOPE_DOWNLOAD_INTRA_CLOUD", "false") + monkeypatch.delenv("MODELSCOPE_DOWNLOAD_INTRA_CLOUD_REGION", raising=False) + monkeypatch.delenv("MODELSCOPE_DOWNLOAD_INTER_CLOUD_REGIONS", raising=False) + monkeypatch.setattr(_download, "DOWNLOAD_INTEGRITY_RETRY_BACKOFF", 0) + fake = _FakeHub() + server = _QuietServer(("127.0.0.1", 0), _handler(fake)) + fake.url = f"http://127.0.0.1:{server.server_port}" + thread = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}, daemon=True) + thread.start() + try: + yield fake + finally: + server.shutdown() + server.server_close() + + +@pytest.fixture() +def api(hub) -> HubApi: + return HubApi(endpoint=hub.url, token="") + + +@pytest.fixture() +def dirs(tmp_path) -> tuple[Path, Path]: + return tmp_path / "out", tmp_path / "cache" + + +class TestSnapshotFailureReporting: + def test_empty_body_for_listed_file_is_reported_and_leaves_nothing(self, hub, api, dirs): + out, cache = dirs + good = hub.add("shards/a.tar", os.urandom(200_000)) + hub.add("shards/b.tar", os.urandom(200_000), served=b"") + hub.add("shards/c.tar", os.urandom(150_000)) + + with pytest.raises(FileDownloadError) as excinfo: + api.download_repo(REPO, "dataset", local_dir=out, cache_dir=cache) + + error = excinfo.value + assert list(error.failed_files) == ["shards/b.tar"] + assert error.total == 3 + assert Path(error.output_dir) == out + assert error.retryable is False + assert "1 of 3 file(s) failed" in str(error) + assert " - shards/b.tar: [E2020] Server announced 0 bytes" in str(error) + + cause = error.failed_files["shards/b.tar"] + assert isinstance(cause, FileSizeMismatchError) + assert "repository lists 200000 bytes" in str(cause) + assert "Content-Length=0" in str(cause) + assert "host=127.0.0.1" in str(cause) + assert "eagleid=test-eagle" in str(cause) + + assert (out / "shards/a.tar").read_bytes() == good + assert (out / "shards/c.tar").stat().st_size == 150_000 + assert not (out / "shards/b.tar").exists() + assert not (out / "shards/b.tar.incomplete").exists() + # One retry over the default route instead of three full re-downloads. + assert len(hub.seen("shards/b.tar", "cdn")) == 2 + + def test_complete_snapshot_returns_output_dir(self, hub, api, dirs): + out, cache = dirs + files = {f"part-{i}.bin": hub.add(f"part-{i}.bin", os.urandom(30_000 + i)) for i in range(4)} + files["README.md"] = hub.add("README.md", b"# data\n", mode="direct") + + assert api.download_repo(REPO, "dataset", local_dir=out, cache_dir=cache) == out + for path, content in files.items(): + assert (out / path).read_bytes() == content + + def test_cli_exits_non_zero_when_a_file_fails(self, hub, dirs, monkeypatch): + out, cache = dirs + monkeypatch.setenv("MODELSCOPE_API_TOKEN", "") + hub.add("a.bin", os.urandom(10_000)) + hub.add("b.bin", os.urandom(10_000), served=b"") + + argv = ["download", REPO, "--repo-type", "dataset", "--local-dir", str(out), "--cache-dir", str(cache)] + assert run_cmd(["--endpoint", hub.url, *argv]) == 1 + assert (out / "a.bin").exists() + assert not (out / "b.bin").exists() + + +class TestSizeValidation: + def test_short_chunked_body_is_rejected(self, hub, api, dirs): + out, cache = dirs + hub.add("notes.txt", b"x" * 5000, mode="direct", served=b"x" * 100) + + with pytest.raises(FileSizeMismatchError, match="returned 100 of 5000 bytes"): + api.download_file(REPO, "dataset", "notes.txt", local_dir=out, cache_dir=cache, file_size=5000) + + assert not (out / "notes.txt").exists() + assert not (out / "notes.txt.incomplete").exists() + + def test_zero_byte_leftover_is_not_a_cache_hit(self, hub, api, dirs): + out, cache = dirs + content = hub.add("z.bin", os.urandom(4096)) + (out).mkdir(parents=True) + (out / "z.bin").write_bytes(b"") + + path = api.download_file(REPO, "dataset", "z.bin", local_dir=out, cache_dir=cache, file_size=len(content)) + + assert path.read_bytes() == content + + def test_size_mismatch_retries_once_without_region_routing(self, hub, api, dirs, monkeypatch): + out, cache = dirs + monkeypatch.setenv("MODELSCOPE_DOWNLOAD_INTRA_CLOUD", "true") + monkeypatch.setenv("MODELSCOPE_DOWNLOAD_INTRA_CLOUD_REGION", "cn-test") + content = hub.add("r.bin", os.urandom(8192), served_with_region=b"") + + path = api.download_file( + REPO, + "dataset", + "r.bin", + local_dir=out, + cache_dir=cache, + file_size=len(content), + expected_sha256=_sha256(content), + ) + + assert path.read_bytes() == content + assert [r.region for r in hub.seen("r.bin", "api")] == ["cn-test", None] + + def test_forced_refresh_keeps_previous_copy_on_wrong_size(self, hub, api, dirs): + out, cache = dirs + content = hub.add("keep.bin", os.urandom(2048), served=b"") + out.mkdir(parents=True) + (out / "keep.bin").write_bytes(content) + + with pytest.raises(FileSizeMismatchError): + api.download_file( + REPO, "dataset", "keep.bin", local_dir=out, cache_dir=cache, force=True, file_size=len(content) + ) + + assert (out / "keep.bin").read_bytes() == content + + +class TestHashValidation: + def test_cdn_hash_detects_empty_body_without_caller_metadata(self, hub, api, dirs): + hub.add("empty.bin", b"original", served=b"") + with pytest.raises(FileIntegrityError): + api.download_file(REPO, "dataset", "empty.bin", local_dir=dirs[0], cache_dir=dirs[1]) + assert not (dirs[0] / "empty.bin").exists() + assert len(hub.seen("empty.bin", "cdn")) == 3 + + def test_hash_mismatch_is_retried_then_leaves_nothing(self, hub, api, dirs): + out, cache = dirs + content = hub.add("w.bin", os.urandom(10_000), served=os.urandom(10_000)) + + with pytest.raises(FileIntegrityError) as excinfo: + api.download_file( + REPO, + "dataset", + "w.bin", + local_dir=out, + cache_dir=cache, + file_size=len(content), + expected_sha256=_sha256(content), + ) + + assert not isinstance(excinfo.value, FileSizeMismatchError) + assert len(hub.seen("w.bin", "cdn")) == _download.DOWNLOAD_HASH_RETRY_TIMES + assert not (out / "w.bin").exists() + assert not (out / "w.bin.incomplete").exists() + + +class TestResume: + def test_truncated_transfer_resumes_from_received_bytes(self, hub, api, dirs, monkeypatch): + out, cache = dirs + # read1 preserves partial chunks so interrupted transfers resume from the exact byte offset. + monkeypatch.setattr(_download, "DOWNLOAD_CHUNK_SIZE", 8192) + content = hub.add("big.bin", os.urandom(200_000), truncate_first=70_000) + + path = api.download_file( + REPO, + "dataset", + "big.bin", + local_dir=out, + cache_dir=cache, + file_size=len(content), + expected_sha256=_sha256(content), + ) + + assert path.read_bytes() == content + assert [r.range for r in hub.seen("big.bin", "cdn")] == [None, "bytes=70000-"] + + def test_server_ignoring_range_rewrites_instead_of_splicing(self, hub, api, dirs): + out, cache = dirs + content = hub.add("plain.txt", os.urandom(50_000), mode="direct") + out.mkdir(parents=True) + (out / "plain.txt.incomplete").write_bytes(content[:20_000]) + + path = api.download_file( + REPO, + "dataset", + "plain.txt", + local_dir=out, + cache_dir=cache, + file_size=len(content), + expected_sha256=_sha256(content), + ) + + assert path.read_bytes() == content + assert [r.range for r in hub.seen("plain.txt", "api")] == ["bytes=20000-"] + + def test_complete_partial_is_finalised_without_a_request(self, hub, api, dirs): + out, cache = dirs + content = hub.add("done.bin", os.urandom(30_000)) + out.mkdir(parents=True) + (out / "done.bin.incomplete").write_bytes(content) + + path = api.download_file( + REPO, + "dataset", + "done.bin", + local_dir=out, + cache_dir=cache, + file_size=len(content), + expected_sha256=_sha256(content), + ) + + assert path.read_bytes() == content + assert hub.requests == [] + + def test_unresumable_partial_restarts_after_416(self, hub, api, dirs): + out, cache = dirs + content = hub.add("obj.bin", os.urandom(30_000)) + out.mkdir(parents=True) + (out / "obj.bin.incomplete").write_bytes(b"\0" * len(content)) + + # Without a listed size the stale partial file can only be detected by the 416. + path = api.download_file(REPO, "dataset", "obj.bin", local_dir=out, cache_dir=cache) + + assert path.read_bytes() == content + assert [r.range for r in hub.seen("obj.bin", "cdn")] == ["bytes=30000-", None] + + +class TestParallelDownload: + @pytest.fixture(autouse=True) + def _parallel(self, monkeypatch): + monkeypatch.setattr(_download, "DOWNLOAD_PARALLELS", 4) + monkeypatch.setattr(_download, "DOWNLOAD_PARALLEL_THRESHOLD", 1024) + monkeypatch.setattr(_download, "DOWNLOAD_PART_SIZE", 64 * 1024) + + def _fetch(self, api, dirs, path: str, content: bytes) -> Path: + out, cache = dirs + return api.download_file( + REPO, + "dataset", + path, + local_dir=out, + cache_dir=cache, + file_size=len(content), + expected_sha256=_sha256(content), + ) + + def test_parts_are_fetched_by_range_and_merged(self, hub, api, dirs): + content = hub.add("p.bin", os.urandom(300_000)) + + assert self._fetch(api, dirs, "p.bin", content).read_bytes() == content + assert len(hub.seen("p.bin", "cdn")) == 5 + assert all(r.range for r in hub.seen("p.bin", "cdn")) + assert list(dirs[0].glob("p.bin_*")) == [] + + def test_truncated_part_resumes_within_its_range(self, hub, api, dirs, monkeypatch): + monkeypatch.setattr(_download, "DOWNLOAD_CHUNK_SIZE", 8192) + content = hub.add("p.bin", os.urandom(300_000), truncate_first=10_000) + + assert self._fetch(api, dirs, "p.bin", content).read_bytes() == content + starts = [int(r.range.removeprefix("bytes=").split("-")[0]) for r in hub.seen("p.bin", "cdn")] + assert any(start % (64 * 1024) == 10_000 for start in starts) + + def test_server_ignoring_range_falls_back_to_single_stream(self, hub, api, dirs): + content = hub.add("d.bin", os.urandom(300_000), mode="direct") + + assert self._fetch(api, dirs, "d.bin", content).read_bytes() == content + assert hub.seen("d.bin", "api")[-1].range is None + assert list(dirs[0].glob("d.bin_*")) == [] + + def test_part_from_object_of_another_size_is_rejected(self, hub, api, dirs): + content = hub.add("s.bin", os.urandom(300_000), served=os.urandom(100_000)) + + with pytest.raises(FileSizeMismatchError, match="reports 100000 bytes"): + self._fetch(api, dirs, "s.bin", content) + assert list(dirs[0].glob("s.bin*")) == [] + + +class TestFileDownloadError: + def test_lists_failures_and_derives_retryability(self): + transient = FileDownloadError("2 of 5 file(s) failed", failed_files={"a": NetworkError("reset")}, total=5) + mixed = FileDownloadError( + "2 of 5 file(s) failed", + failed_files={"a": NetworkError("reset"), "b": FileSizeMismatchError("short")}, + total=5, + ) + + assert str(mixed).splitlines() == [ + "[E1003] 2 of 5 file(s) failed", + " - a: [E1020] reset", + " - b: [E2020] short", + ] + assert transient.retryable is True + assert mixed.retryable is False + assert FileDownloadError("none").retryable is False + + def test_long_failure_lists_are_truncated(self): + failures = {f"f{i}": NetworkError("x") for i in range(12)} + + assert str(FileDownloadError("12 failed", failed_files=failures)).endswith(" ... and 2 more") + + +class TestBoundedDownloads: + def test_header_stall_has_one_retry_layer(self, hub, api, dirs, monkeypatch): + import modelscope_hub._legacy_api as legacy + + monkeypatch.setattr(legacy, "DOWNLOAD_TIMEOUT", 0.15) + monkeypatch.setattr(_download, "DOWNLOAD_RETRY_TIMES", 1) + hub.add("stall.bin", b"x" * 2048, header_delay=2) + started = time.monotonic() + with pytest.raises(NetworkError): + api.download_file(REPO, "dataset", "stall.bin", local_dir=dirs[0], cache_dir=dirs[1], file_size=2048) + assert len(hub.seen("stall.bin", "cdn")) == 2 + assert time.monotonic() - started < 1 + assert api.legacy._read_session.get_adapter(hub.url).max_retries.total == 0 + assert api.legacy._session.get_adapter(hub.url).max_retries.total == 5 + + @pytest.mark.parametrize("status", [429, 503]) + def test_retry_after_is_capped(self, hub, api, dirs, monkeypatch, status): + monkeypatch.setattr(_download, "DOWNLOAD_RETRY_TIMES", 1) + monkeypatch.setattr(_download, "DOWNLOAD_RETRY_MAX_DELAY_SECONDS", 0.15) + hub.add("busy.bin", b"x", status=status, retry_after="999999") + started = time.monotonic() + with pytest.raises(Exception): + api.download_file(REPO, "dataset", "busy.bin", local_dir=dirs[0], cache_dir=dirs[1]) + assert len(hub.seen("busy.bin", "cdn")) == 2 + assert time.monotonic() - started < 1 + + @pytest.mark.parametrize("headers", [False, True]) + def test_deadline_stops_trickling_headers_and_body(self, hub, api, dirs, monkeypatch, headers): + from modelscope_hub._download_control import _DownloadDeadlineError + + monkeypatch.setattr(_download, "DOWNLOAD_TOTAL_TIMEOUT", 0.35) + hub.add("slow.bin", b"x" * 100, header_trickle=headers, trickle_interval=0 if headers else 0.05) + started = time.monotonic() + with pytest.raises(_DownloadDeadlineError): + api.download_file(REPO, "dataset", "slow.bin", local_dir=dirs[0], cache_dir=dirs[1], file_size=100) + assert time.monotonic() - started < 1.2 + assert len(hub.seen("slow.bin", "cdn")) == 1 + assert not (dirs[0] / "slow.bin").exists() + + def test_minimum_rate_interrupts_a_connection_that_never_goes_idle(self, hub, api, dirs, monkeypatch): + from modelscope_hub._download_control import _DownloadTooSlowError + + monkeypatch.setattr(_download, "DOWNLOAD_LOW_SPEED_WINDOW_SECONDS", 0.25) + monkeypatch.setattr(_download, "DOWNLOAD_MIN_BYTES_PER_SECOND", 1000) + monkeypatch.setattr(_download, "DOWNLOAD_RETRY_TIMES", 0) + hub.add("slow.bin", b"x" * 100, trickle_interval=0.05) + with pytest.raises(_DownloadTooSlowError): + api.download_file(REPO, "dataset", "slow.bin", local_dir=dirs[0], cache_dir=dirs[1], file_size=100) + assert len(hub.seen("slow.bin", "cdn")) == 1 + + def test_ctrl_c_cancels_queued_and_active_downloads(self, hub, api, dirs, monkeypatch): + for i in range(10): + hub.add(f"f{i}.bin", b"x", header_delay=2) + + def interrupt(futures, **kwargs): + deadline = time.monotonic() + 1 + while len([r for r in hub.requests if r.kind == "cdn"]) < 2 and time.monotonic() < deadline: + time.sleep(0.01) + raise KeyboardInterrupt + + monkeypatch.setattr(_download, "wait", interrupt) + started = time.monotonic() + with pytest.raises(KeyboardInterrupt): + api.download_repo(REPO, "dataset", local_dir=dirs[0], cache_dir=dirs[1], max_workers=2) + assert time.monotonic() - started < 1.5 + assert len({r.path for r in hub.requests if r.kind == "api"}) == 2 + + def test_read_pool_does_not_change_global_or_proxy_connection_types(self): + from urllib3.poolmanager import pool_classes_by_scheme + + from modelscope_hub._download_control import _ReadHTTPAdapter + + original = dict(pool_classes_by_scheme) + adapter = _ReadHTTPAdapter(max_retries=0) + proxy = adapter.proxy_manager_for("http://127.0.0.1:9999") + assert pool_classes_by_scheme == original + for scheme, parent in original.items(): + pool = proxy.pool_classes_by_scheme[scheme] + assert issubclass(pool, parent) + assert issubclass(pool.ConnectionCls, parent.ConnectionCls) + adapter.close() + + def test_lock_wait_is_cancellable(self, hub, api, dirs): + from filelock import FileLock + + from modelscope_hub._download_control import _DownloadCancelledError + + stop = threading.Event() + lock_path = api.downloader._lock_path(REPO, "dataset", cache_dir=dirs[1], file_path="f.bin") + lock_path.parent.mkdir(parents=True) + with FileLock(str(lock_path)): + timer = threading.Timer(0.1, stop.set) + timer.start() + with pytest.raises(_DownloadCancelledError): + api.downloader.download_file( + REPO, + "dataset", + "f.bin", + local_dir=dirs[0], + cache_dir=dirs[1], + _cancel_event=stop, + ) + timer.join() + assert hub.requests == [] + + def test_corrupt_forced_refresh_does_not_replace_previous_valid_file(self, hub, api, dirs): + content = hub.add("f.bin", b"good", served=b"evil") + dirs[0].mkdir() + target = dirs[0] / "f.bin" + target.write_bytes(content) + with pytest.raises(FileIntegrityError): + api.download_file( + REPO, + "dataset", + "f.bin", + local_dir=dirs[0], + cache_dir=dirs[1], + force=True, + file_size=4, + expected_sha256=_sha256(content), + ) + assert target.read_bytes() == content + assert not (dirs[0] / "f.bin.incomplete").exists() + + +class TestListingOptimisation: + @pytest.mark.parametrize( + "patterns, expected", + [ + (["a/*.bin", "a/b/*.txt"], ["a"]), + (["a/*.bin", "b/f.json"], ["a", "b"]), + (["*.bin", "a/f.bin"], [None]), + ([r"a/.*\\.bin"], [None]), + (["../a/*.bin"], [None]), + ], + ) + def test_safe_roots(self, patterns, expected): + assert _download._listing_roots(patterns) == expected + + def test_scoped_listing_preserves_glob_selection(self, hub, api, dirs): + hub.add("a/good.bin", b"yes") + hub.add("a/no.txt", b"no") + hub.add("b/other.bin", b"no") + assert ( + api.download_repo( + REPO, + "dataset", + local_dir=dirs[0], + cache_dir=dirs[1], + allow_patterns=["a/*.bin"], + ) + == dirs[0] + ) + assert (dirs[0] / "a/good.bin").read_bytes() == b"yes" + assert not (dirs[0] / "b/other.bin").exists() + assert not (dirs[0] / "a/no.txt").exists() + + def test_file_list_reuse_skips_the_listing_endpoint(self, hub, api, dirs, monkeypatch): + content = hub.add("a/good.bin", b"yes") + manifest = [{"Path": "a/good.bin", "Size": 3, "Sha256": _sha256(content)}] + monkeypatch.setattr( + api.legacy, "list_dataset_files_paginated", lambda *args, **kwargs: pytest.fail("列举了清单") + ) + api.download_repo(REPO, "dataset", local_dir=dirs[0], cache_dir=dirs[1], file_list=manifest) + assert (dirs[0] / "a/good.bin").read_bytes() == content + + def test_endpoint_probe_obeys_its_total_budget(self, hub, monkeypatch): + import modelscope_hub.constants as constants + from modelscope_hub._download_control import _DownloadDeadlineError + + monkeypatch.setattr(constants, "REPO_TREE_PAGE_TIMEOUT_SECONDS", 0.2) + hub.add("metadata.bin", b"x", header_delay=2) + + def exists(probe, repo_id, repo_type): + probe.openapi.request( + "GET", + url=f"{hub.url}/cdn/metadata.bin", + require_token=False, + unwrap=False, + ) + return True + + monkeypatch.setattr(HubApi, "repo_exists", exists) + caller = HubApi(token="") + caller._config._endpoint_overridden = False + started = time.monotonic() + with pytest.raises(_DownloadDeadlineError): + caller.resolve_endpoint_for_read(REPO, repo_type="dataset") + assert time.monotonic() - started < 1 + + def test_region_detection_is_single_flight(self, api, monkeypatch): + calls = [] + + def probe(): + calls.append(1) + time.sleep(0.03) + return "cn-test" + + monkeypatch.setattr(api.downloader, "_detect_region_uncached", probe) + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor(max_workers=8) as pool: + assert list(pool.map(lambda _: api.downloader._detect_region(), range(16))) == ["cn-test"] * 16 + assert calls == [1] + + +class TestDownloadConfiguration: + def test_defaults_exports_registry_and_readme_are_aligned(self): + import modelscope_hub.constants as constants + + expected = { + "DOWNLOAD_TOTAL_TIMEOUT": 3600, + "DOWNLOAD_MIN_BYTES_PER_SECOND": 1024, + "DOWNLOAD_LOW_SPEED_WINDOW_SECONDS": 60, + "DOWNLOAD_RETRY_MAX_DELAY_SECONDS": 60, + "DOWNLOAD_FILE_LOCK_TIMEOUT_SECONDS": 600, + "REPO_TREE_PAGE_TIMEOUT_SECONDS": 15, + "REPO_TREE_LIST_TIMEOUT_SECONDS": 3600, + } + registry = {entry.name: entry.default for entry in constants.ENV_REGISTRY} + readme = (Path(__file__).parents[1] / "README.md").read_text() + for name, value in expected.items(): + assert getattr(constants, name) == value + assert name in constants.__all__ + env_name = f"MODELSCOPE_{name}" + assert registry[env_name] == str(value) + assert f"| `{env_name}` | `{value}` |" in readme + + def test_zero_can_disable_transport_retries(self, monkeypatch): + import subprocess + import sys + + monkeypatch.setenv("MODELSCOPE_DOWNLOAD_MAX_RETRIES", "0") + result = subprocess.run( + [ + sys.executable, + "-c", + "from modelscope_hub.constants import DOWNLOAD_RETRY_TIMES; print(DOWNLOAD_RETRY_TIMES)", + ], + capture_output=True, + text=True, + check=True, + ) + assert result.stdout.strip() == "0" + + +class TestDownloadCallbacks: + def test_resume_reports_each_byte_once_and_closes_callback(self, hub, api, dirs): + reports = [] + + class Callback(_download.ProgressCallback): + def __init__(self, filename, file_size): + super().__init__(filename, file_size) + self.size = 0 + self.closed = False + reports.append(self) + + def update(self, size): + self.size += size + + def end(self): + self.closed = True + + content = hub.add("resume.bin", b"x" * 5000, truncate_first=2000) + api.downloader.download_file( + REPO, + "dataset", + "resume.bin", + local_dir=dirs[0], + cache_dir=dirs[1], + file_size=5000, + expected_sha256=_sha256(content), + progress_callbacks=[Callback], + ) + assert len(reports) == 1 + assert reports[0].size == 5000 + assert reports[0].closed diff --git a/tests/test_user_agent.py b/tests/test_user_agent.py new file mode 100644 index 0000000..0e01658 --- /dev/null +++ b/tests/test_user_agent.py @@ -0,0 +1,331 @@ +"""Verify SDK UA, stable sessions and channel propagation using PreparedRequest, without remote access.""" + +from __future__ import annotations + +import importlib +import io +from argparse import Namespace +from unittest.mock import Mock + +import pytest +import requests +import responses + +from modelscope_hub import HubApi, __version__ +from modelscope_hub import config as config_module +from modelscope_hub._download_control import _DownloadControl +from modelscope_hub._legacy_api import LegacyClient +from modelscope_hub._openapi import OpenAPIClient +from modelscope_hub.agent import AgentApi +from modelscope_hub.cli.base import make_api +from modelscope_hub.cli.main import run_cmd +from modelscope_hub.compat import ( + dataset_file_download, + dataset_snapshot_download, + model_file_download, + snapshot_download, +) +from modelscope_hub.compat.hub_api import LegacyHubApi +from modelscope_hub.config import HubConfig +from modelscope_hub.constants import DEFAULT_ENDPOINT as CN +from modelscope_hub.constants import DEFAULT_INTL_ENDPOINT as AI +from modelscope_hub.errors import CacheError +from modelscope_hub.utils import append_user_agent, build_user_agent + +TOKEN = "ms-ua-test-dummy" +USER = {"success": True, "data": {"username": "alice"}} + + +@pytest.fixture(autouse=True) +def isolated_config(tmp_path, monkeypatch): + for name in ("MODELSCOPE_ENDPOINT", "MODELSCOPE_DOMAIN", "MODELSCOPE_API_TOKEN", "MODELSCOPE_PREFER_AI_SITE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("MODELSCOPE_HOME", str(tmp_path / "home")) + monkeypatch.setenv("MODELSCOPE_CACHE", str(tmp_path / "cache")) + monkeypatch.setattr(config_module, "_default_config", None) + + +@pytest.fixture +def config(): + return HubConfig(token=TOKEN) + + +def _assert_sdk(ua): + assert ua.startswith(f"modelscope_hub/{__version__}; ") + assert "session_id/" in ua + assert "python/" in ua + assert TOKEN not in ua + return ua + + +def _session_id(ua): + return next(part.removeprefix("session_id/") for part in ua.split("; ") if part.startswith("session_id/")) + + +@pytest.mark.parametrize("extra", [None, {"application": "example"}, "example-client/1"]) +def test_build_user_agent_config_and_extra(config, extra): + first = _assert_sdk(build_user_agent(config=config, extra=extra)) + assert first == build_user_agent(config=config, extra=extra) + assert _session_id(first) == config.get_session_id() + if extra: + assert ("application/example" if isinstance(extra, dict) else extra) in first + assert append_user_agent(first, extra) == first + + +@pytest.mark.parametrize("failure", [CacheError("read-only"), PermissionError("read-only"), OSError("read-only")]) +@responses.activate +def test_readonly_home_keeps_requests_and_session_stable(config, monkeypatch, failure): + monkeypatch.setattr(HubConfig, "get_session_id", Mock(side_effect=failure)) + responses.add(responses.GET, f"{CN}/openapi/v1/users/me", json=USER) + api = HubApi(config=config) + api.whoami() + api.whoami() + uas = [call.request.headers["User-Agent"] for call in responses.calls] + uas.extend([api.legacy._user_agent, api.downloader._build_user_agent(), OpenAPIClient(config)._user_agent]) + assert len({_session_id(_assert_sdk(ua)) for ua in uas}) == 1 + assert not config.credentials_dir.exists() + + +def test_hubapi_constructor_remains_lazy(config, monkeypatch): + lookup = Mock(side_effect=AssertionError("不应提前访问用户目录")) + monkeypatch.setattr(HubConfig, "get_session_id", lookup) + HubApi(config=config, user_agent={"application": "example"}) + lookup.assert_not_called() + + +@pytest.mark.parametrize("read_session", [False, True]) +@pytest.mark.parametrize("method", ["GET", "POST", "PUT", "PATCH", "DELETE"]) +@responses.activate +def test_all_http_methods_and_both_sessions(config, method, read_session): + url = f"{CN}/openapi/v1/ua-test" + responses.add(method, url, json={"success": True, "data": {}}) + with OpenAPIClient(config, user_agent={"application": "example"}) as client: + if read_session: + with _DownloadControl(total_timeout=5) as control, client.read_context(control): + client.request(method, "/ua-test") + else: + client.request(method, "/ua-test") + ua = _assert_sdk(responses.calls[0].request.headers["User-Agent"]) + assert "application/example" in ua + assert "python-requests/" not in ua + + +@pytest.mark.parametrize("name", ["User-Agent", "user-agent", "USER-AGENT"]) +@pytest.mark.parametrize("value", [None, "", "other-client/2"]) +@responses.activate +def test_custom_header_is_appended_case_insensitively(config, name, value): + responses.add(responses.GET, f"{CN}/openapi/v1/users/me", json=USER) + headers = {name: value, "X-Example": "preserved"} + OpenAPIClient(config).request("GET", "/users/me", headers=headers) + prepared = responses.calls[0].request.headers + ua = _assert_sdk(prepared["User-Agent"]) + assert sum(key.lower() == "user-agent" for key in prepared) == 1 + assert prepared["X-Example"] == "preserved" + assert headers == {name: value, "X-Example": "preserved"} + if value: + assert value in ua + + +@pytest.mark.parametrize("session_ua", [None, "", "custom-session/1", requests.utils.default_user_agent()]) +@responses.activate +def test_injected_session_has_sdk_ua_without_mutating_headers(config, session_ua): + session = requests.Session() + if session_ua is None: + session.headers.pop("User-Agent", None) + else: + session.headers["User-Agent"] = session_ua + original = dict(session.headers) + responses.add(responses.GET, f"{CN}/openapi/v1/users/me", json=USER) + with OpenAPIClient(config, session=session) as client: + client.get_current_user() + with _DownloadControl(total_timeout=5) as control, client.read_context(control): + client.get_current_user() + uas = [_assert_sdk(call.request.headers["User-Agent"]) for call in responses.calls] + assert uas[0] == uas[1] + assert dict(session.headers) == original + if session_ua == "custom-session/1": + assert session_ua in uas[0] + + +@pytest.mark.parametrize("status", [429, 503]) +@responses.activate +def test_retry_ua_is_stable(config, monkeypatch, status): + sleeps = Mock() + monkeypatch.setattr("modelscope_hub._openapi.time.sleep", sleeps) + url = f"{CN}/openapi/v1/users/me" + responses.add(responses.GET, url, json={"message": "retry"}, status=status, headers={"Retry-After": "2"}) + responses.add(responses.GET, url, json=USER) + with OpenAPIClient(config, max_retries=2) as client: + client.get_current_user() + assert len(responses.calls) == 2 + assert ( + _assert_sdk(responses.calls[0].request.headers["User-Agent"]) + == responses.calls[1].request.headers["User-Agent"] + ) + sleeps.assert_called_once() + if status == 429: + sleeps.assert_called_with(2.0) + + +@responses.activate +def test_anonymous_fallback_keeps_ua_and_drops_auth(config): + url = f"{CN}/openapi/v1/studios" + responses.add(responses.GET, url, json={"message": "forbidden"}, status=403) + responses.add(responses.GET, url, json={"studios": []}) + OpenAPIClient(config).list_studios() + first, second = [call.request.headers for call in responses.calls] + assert _assert_sdk(first["User-Agent"]) == second["User-Agent"] + assert "Authorization" in first and "Cookie" in first + assert "Authorization" not in second and "Cookie" not in second + + +@pytest.mark.parametrize("path", ["/agent_id/.well-known/agentid-configuration", "/agent_id/.well-known/agentid-jwks"]) +@responses.activate +def test_anonymous_idp_metadata_has_ua(config, path): + responses.add(responses.GET, f"{CN}/openapi/v1{path}", json={"data": {}}) + client = OpenAPIClient(config) + if path.endswith("agentid-jwks"): + client.get_agent_id_jwks() + else: + client.get_agent_id_configuration() + headers = responses.calls[0].request.headers + _assert_sdk(headers["User-Agent"]) + assert "Authorization" not in headers and "Cookie" not in headers + + +@responses.activate +def test_multipart_absolute_and_raw_response_have_ua(config): + responses.add(responses.POST, f"{CN}/openapi/v1/files/upload", json={"data": {"id": "f"}}) + responses.add(responses.PUT, "https://storage.example.test/blob", status=204) + with OpenAPIClient(config) as client: + client.upload_file(file=io.BytesIO(b"x")) + response = client.request("PUT", url="https://storage.example.test/blob", data=b"x", unwrap=False) + assert response.status_code == 204 + assert "multipart/form-data" in responses.calls[0].request.headers["Content-Type"] + for call in responses.calls: + _assert_sdk(call.request.headers["User-Agent"]) + assert "Authorization" not in responses.calls[1].request.headers + assert "Cookie" not in responses.calls[1].request.headers + + +@responses.activate +def test_login_and_scoped_probe_share_channel_and_session(config): + responses.add(responses.POST, f"{CN}/api/v1/login", json={"message": "forbidden"}, status=403) + responses.add(responses.GET, f"{CN}/openapi/v1/users/me", json=USER) + assert HubApi(config=config, user_agent={"invoked_by": "cli"}).login(TOKEN).username == "alice" + assert len(responses.calls) == 2 + uas = [_assert_sdk(call.request.headers["User-Agent"]) for call in responses.calls] + assert uas[0] == uas[1] + assert "invoked_by/cli" in uas[0] + + +@responses.activate +def test_peer_login_probe_has_sdk_ua(): + responses.add(responses.POST, f"{AI}/api/v1/login", json={"Code": 200, "Data": {"Username": "alice"}}) + assert HubApi._token_valid_on(TOKEN, AI) + _assert_sdk(responses.calls[0].request.headers["User-Agent"]) + + +@responses.activate +def test_read_site_probes_keep_custom_config_channel_and_session(config): + responses.add(responses.GET, f"{CN}/openapi/v1/models/alice/m", status=404) + responses.add(responses.GET, f"{CN}/api/v1/models/alice/m", status=404) + responses.add(responses.GET, f"{AI}/openapi/v1/models/alice/m", json={"data": {"name": "m"}}) + api = HubApi(config=config, user_agent={"invoked_by": "compat", "application": "example"}) + assert api.resolve_endpoint_for_read("alice/m") == AI + assert len(responses.calls) == 3 + uas = [_assert_sdk(call.request.headers["User-Agent"]) for call in responses.calls] + assert len(set(uas)) == 1 + assert _session_id(uas[0]) == config.get_session_id() + assert "invoked_by/compat" in uas[0] and "application/example" in uas[0] + + +@responses.activate +def test_agent_openapi_legacy_and_blob_have_sdk_ua(config): + responses.add(responses.GET, f"{CN}/openapi/v1/agents/alice/a", json={"data": {}}) + responses.add(responses.GET, f"{CN}/api/v1/agents/alice/a/repo/files", json={"data": {"Trees": []}}) + responses.add(responses.PUT, "https://storage.example.test/blob", status=204) + api = AgentApi(config=config, user_agent={"invoked_by": "cli"}) + api.repo_info("alice", "a") + api.list_repo_files("alice", "a") + api.lfs_upload_blob("https://storage.example.test/blob", b"x") + for call in responses.calls: + assert "invoked_by/cli" in _assert_sdk(call.request.headers["User-Agent"]) + + +@pytest.mark.parametrize("read_session", [False, True]) +@responses.activate +def test_legacy_default_custom_headers_and_blob(config, read_session): + responses.add(responses.GET, f"{CN}/api/v1/models/alice/m", json={"Data": {}}) + responses.add(responses.PUT, f"{CN}/blob", json={}) + client = LegacyClient(token=TOKEN, endpoint=CN) + if read_session: + with _DownloadControl(total_timeout=5) as control, client.read_context(control): + client._request("GET", "models/alice/m", headers={"user-agent": "custom-read/1"}) + else: + client._request("GET", "models/alice/m", headers={"USER-AGENT": "custom-read/1"}) + client.upload_blob(f"{CN}/blob", b"x", size=1, headers={"user-agent": "custom-blob/1"}) + assert "custom-read/1" in _assert_sdk(responses.calls[0].request.headers["User-Agent"]) + assert "custom-blob/1" in _assert_sdk(responses.calls[1].request.headers["User-Agent"]) + + +@responses.activate +def test_cli_whoami_and_explicit_download_endpoint_have_channel(): + responses.add(responses.GET, f"{CN}/openapi/v1/users/me", json=USER) + assert run_cmd(["--token", TOKEN, "whoami"]) == 0 + assert "invoked_by/cli" in _assert_sdk(responses.calls[0].request.headers["User-Agent"]) + from modelscope_hub.cli.download import DownloadCommand + + command = DownloadCommand(Namespace(endpoint=CN, token=TOKEN)) + assert "invoked_by/cli" in _assert_sdk(command._make_api_with_endpoint().openapi._user_agent) + assert "invoked_by/cli" in make_api(Namespace(token=TOKEN)).legacy._user_agent + + +@responses.activate +def test_cli_agent_username_and_repo_calls_have_channel(monkeypatch): + cli_agent = importlib.import_module("modelscope_hub.cli.agent") + responses.add(responses.GET, f"{CN}/openapi/v1/users/me", json=USER) + responses.add(responses.PUT, f"{CN}/api/v1/dolphin/agents", json={"Data": {"AgentList": []}}) + + def list_after_username(**kwargs): + return cli_agent._cmd_list(owner="alice", page_number=1, page_size=10, endpoint=CN, token=TOKEN) + + monkeypatch.setattr(cli_agent, "_cmd_upload", list_after_username) + assert run_cmd(["--token", TOKEN, "agent", "upload", "--repo", "alice/a", "--local-dir", "."]) == 0 + assert responses.calls + for call in responses.calls: + assert "invoked_by/cli" in _assert_sdk(call.request.headers["User-Agent"]) + + +@responses.activate +def test_compat_default_and_override_clients_have_channel(): + responses.add(responses.GET, f"{CN}/openapi/v1/models/alice/m", json={"data": {"name": "m"}}) + api = LegacyHubApi(endpoint=CN, token=TOKEN) + api.repo_exists("alice/m", repo_type="model") + api.repo_exists("alice/m", repo_type="model", token="other-test-token") + for call in responses.calls: + assert "invoked_by/compat" in _assert_sdk(call.request.headers["User-Agent"]) + + +@pytest.mark.parametrize( + "download", [model_file_download, dataset_file_download, snapshot_download, dataset_snapshot_download] +) +def test_compat_download_user_agent_applies_to_all_transports(download, monkeypatch): + seen = [] + + def fake_download(self, *args, **kwargs): + ua = self.openapi._user_agent + assert "invoked_by/compat" in _assert_sdk(ua) and "application/example" in ua + assert self.legacy._user_agent == ua + assert self.downloader._build_user_agent(kwargs["user_agent"]) == ua + seen.append(ua) + return "local-path" + + monkeypatch.setattr(HubApi, "download_file", fake_download) + monkeypatch.setattr(HubApi, "download_repo", fake_download) + kwargs = {"local_files_only": True, "user_agent": {"application": "example"}} + if download in (model_file_download, dataset_file_download): + assert download("alice/m", "x", **kwargs) == "local-path" + else: + assert download("alice/m", **kwargs) == "local-path" + assert len(seen) == 1