"""File download — delegates to modelscope_hub for Hub downloads, retains http_get_file for direct HTTP. Hub file downloads (model_file_download, dataset_file_download) are delegated to modelscope_hub.compat. Direct HTTP file downloads (http_get_file, http_get_model_file) are retained as they serve non-Hub use cases. """ import copy import hashlib import io import os import tempfile import time import urllib import uuid from functools import partial from http.cookiejar import CookieJar from typing import Dict, List, Optional, Type import requests # --- Hub file downloads (delegated) --- from modelscope_hub.compat.file_download import \ dataset_file_download as _compat_dataset_file_download from modelscope_hub.compat.file_download import \ model_file_download as _compat_model_file_download from requests.adapters import Retry from tqdm.auto import tqdm from modelscope.hub.constants import (API_FILE_DOWNLOAD_CHUNK_SIZE, API_FILE_DOWNLOAD_RETRY_TIMES, API_FILE_DOWNLOAD_TIMEOUT, MODELSCOPE_SDK_DEBUG) from modelscope.utils.logger import get_logger from .callback import ProgressCallback, TqdmCallback from .errors import FileDownloadError from .utils.utils import find_reusable_legacy_repo_dir, get_endpoint logger = get_logger() def _get_release_timestamp(): """Compute the release timestamp for revision resolution. Returns None (dev-mode) when MODELSCOPE_SDK_DEBUG is set. """ if os.environ.get(MODELSCOPE_SDK_DEBUG): return None try: from modelscope import version dt = getattr(version, '__release_datetime__', None) if not dt: return None return int(time.mktime(time.strptime(dt, '%Y-%m-%d %H:%M:%S'))) except Exception: return None def model_file_download( model_id: str, file_path: str, revision: str = None, *, cache_dir: str = None, local_dir: str = None, cookies: dict = None, token: str = None, endpoint: str = None, local_files_only: bool = False, user_agent=None, ) -> str: """Download a single model file with release-mode revision resolution.""" if revision is None: try: from modelscope.hub.api import HubApi api = HubApi() release_ts = _get_release_timestamp() detail = api.get_valid_revision_detail( model_id, revision=None, release_timestamp=release_ts) revision = detail.get('Revision') except Exception: pass if local_dir is None: local_dir = find_reusable_legacy_repo_dir( model_id, repo_type='model', cache_dir=cache_dir) return _compat_model_file_download( model_id, file_path, revision=revision, cache_dir=cache_dir, local_dir=local_dir, cookies=cookies, token=token, endpoint=endpoint, local_files_only=local_files_only, user_agent=user_agent, ) def dataset_file_download( dataset_id: str, file_path: str, *, cache_dir: str = None, local_dir: str = None, revision: str = None, cookies: dict = None, token: str = None, endpoint: str = None, local_files_only: bool = False, user_agent=None, ) -> str: """Download a single dataset file, reusing flat/hub legacy caches when present.""" if local_dir is None: local_dir = find_reusable_legacy_repo_dir( dataset_id, repo_type='dataset', cache_dir=cache_dir) return _compat_dataset_file_download( dataset_id, file_path, cache_dir=cache_dir, local_dir=local_dir, revision=revision, cookies=cookies, token=token, endpoint=endpoint, local_files_only=local_files_only, user_agent=user_agent, ) # --- Direct HTTP downloads (retained - non-Hub API) --- def get_file_download_url(model_id: str, file_path: str, revision: str, endpoint: Optional[str] = None): """Format file download url according to `model_id`, `revision` and `file_path`. e.g., Given `model_id=john/bert`, `revision=master`, `file_path=README.md`, the resulted download url is: https://modelscope.cn/api/v1/models/john/bert/repo?Revision=master&FilePath=README.md Args: model_id (str): The model_id. file_path (str): File path revision (str): File revision. endpoint (str): The remote endpoint Returns: str: The file url. """ file_path = urllib.parse.quote_plus(file_path) revision = urllib.parse.quote_plus(revision) download_url_template = '{endpoint}/api/v1/models/{model_id}/repo?Revision={revision}&FilePath={file_path}' if not endpoint: endpoint = get_endpoint() return download_url_template.format( endpoint=endpoint, model_id=model_id, revision=revision, file_path=file_path, ) def http_get_model_file( url: str, local_dir: str, file_name: str, file_size: int, cookies: CookieJar, headers: Optional[Dict[str, str]] = None, disable_tqdm: bool = False, progress_callbacks: List[Type[ProgressCallback]] = None, ): """Download remote file, will retry 5 times before giving up on errors. Args: url(str): actual download url of the file local_dir(str): local directory where the downloaded file stores file_name(str): name of the file stored in `local_dir` file_size(int): The file size. cookies(CookieJar): cookies used to authentication the user, which is used for downloading private repos headers(Dict[str, str], optional): http headers to carry necessary info when requesting the remote file disable_tqdm(bool, optional): Disable the progress bar with tqdm. progress_callbacks(List[Type[ProgressCallback]], optional): progress callbacks to track the download progress. Raises: FileDownloadError: File download failed. """ progress_callbacks = [] if progress_callbacks is None else progress_callbacks.copy( ) if not disable_tqdm: progress_callbacks.append(TqdmCallback) progress_callbacks = [ callback(file_name, file_size) for callback in progress_callbacks ] get_headers = {} if headers is None else copy.deepcopy(headers) get_headers['X-Request-ID'] = str(uuid.uuid4().hex) temp_file_path = os.path.join(local_dir, file_name) os.makedirs(os.path.dirname(temp_file_path), exist_ok=True) logger.debug('downloading %s to %s', url, temp_file_path) # retry sleep 0.5s, 1s, 2s, 4s has_retry = False hash_sha256 = hashlib.sha256() retry = Retry( total=API_FILE_DOWNLOAD_RETRY_TIMES, backoff_factor=1, allowed_methods=['GET']) while True: try: if file_size == 0: # Avoid empty file server request with open(temp_file_path, 'w+'): for callback in progress_callbacks: callback.update(1) break # Determine the length of any existing partial download partial_length = 0 # download partial, continue download if os.path.exists(temp_file_path): # resuming from interrupted download is also considered as retry has_retry = True with open(temp_file_path, 'rb') as f: partial_length = f.seek(0, io.SEEK_END) for callback in progress_callbacks: callback.update(partial_length) # Check if download is complete if partial_length >= file_size: break # closed range[], from 0. get_headers['Range'] = 'bytes=%s-%s' % (partial_length, file_size - 1) with open(temp_file_path, 'ab+') as f: r = requests.get( url, stream=True, headers=get_headers, cookies=cookies, timeout=API_FILE_DOWNLOAD_TIMEOUT) r.raise_for_status() for chunk in r.iter_content( chunk_size=API_FILE_DOWNLOAD_CHUNK_SIZE): if chunk: # filter out keep-alive new chunks for callback in progress_callbacks: callback.update(len(chunk)) f.write(chunk) # hash would be discarded in retry case anyway if not has_retry: hash_sha256.update(chunk) break except Exception as e: # no matter what happen, we will retry. has_retry = True retry = retry.increment('GET', url, error=e) retry.sleep() for callback in progress_callbacks: callback.end() # if anything went wrong, we would discard the real-time computed hash and return None return None if has_retry else hash_sha256.hexdigest() def http_get_file( url: str, local_dir: str, file_name: str, cookies: CookieJar, headers: Optional[Dict[str, str]] = None, ): """Download remote file, will retry 5 times before giving up on errors. Args: url(str): actual download url of the file local_dir(str): local directory where the downloaded file stores file_name(str): name of the file stored in `local_dir` cookies(CookieJar): cookies used to authentication the user, which is used for downloading private repos headers(Dict[str, str], optional): http headers to carry necessary info when requesting the remote file Raises: FileDownloadError: File download failed. """ total = -1 temp_file_manager = partial( tempfile.NamedTemporaryFile, mode='wb', dir=local_dir, delete=False) get_headers = {} if headers is None else copy.deepcopy(headers) get_headers['X-Request-ID'] = str(uuid.uuid4().hex) with temp_file_manager() as temp_file: logger.debug('downloading %s to %s', url, temp_file.name) # retry sleep 0.5s, 1s, 2s, 4s retry = Retry( total=API_FILE_DOWNLOAD_RETRY_TIMES, backoff_factor=1, allowed_methods=['GET']) while True: try: downloaded_size = temp_file.tell() get_headers['Range'] = 'bytes=%d-' % downloaded_size r = requests.get( url, stream=True, headers=get_headers, cookies=cookies, timeout=API_FILE_DOWNLOAD_TIMEOUT) r.raise_for_status() content_length = r.headers.get('Content-Length') total = int( content_length) if content_length is not None else None progress = tqdm( unit='B', unit_scale=True, unit_divisor=1024, total=total, initial=downloaded_size, desc='Downloading [' + file_name + ']', ) for chunk in r.iter_content( chunk_size=API_FILE_DOWNLOAD_CHUNK_SIZE): if chunk: # filter out keep-alive new chunks progress.update(len(chunk)) temp_file.write(chunk) progress.close() break except (Exception) as e: # no matter what happen, we will retry. retry = retry.increment('GET', url, error=e) retry.sleep() logger.debug('storing %s in cache at %s', url, local_dir) downloaded_length = os.path.getsize(temp_file.name) if total != downloaded_length: os.remove(temp_file.name) msg = 'File %s download incomplete, content_length: %s but the \ file downloaded length: %s, please download again' % ( file_name, total, downloaded_length) logger.error(msg) raise FileDownloadError(msg) os.replace(temp_file.name, os.path.join(local_dir, file_name)) __all__ = [ 'model_file_download', 'dataset_file_download', 'http_get_file', 'http_get_model_file', 'get_file_download_url', ]