diff --git a/modelscope/hub/constants.py b/modelscope/hub/constants.py index 3cde867d..fb8bd5f5 100644 --- a/modelscope/hub/constants.py +++ b/modelscope/hub/constants.py @@ -1,11 +1,15 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +import os from pathlib import Path MODELSCOPE_URL_SCHEME = 'http://' DEFAULT_MODELSCOPE_DOMAIN = 'www.modelscope.cn' DEFAULT_MODELSCOPE_DATA_ENDPOINT = MODELSCOPE_URL_SCHEME + DEFAULT_MODELSCOPE_DOMAIN - +MODELSCOPE_PARALLEL_DOWNLOAD_THRESHOLD_MB = int( + os.environ.get('MODELSCOPE_PARALLEL_DOWNLOAD_THRESHOLD_MB', 500)) +MODELSCOPE_DOWNLOAD_PARALLELS = int( + os.environ.get('MODELSCOPE_DOWNLOAD_PARALLELS', 4)) DEFAULT_MODELSCOPE_GROUP = 'damo' MODEL_ID_SEPARATOR = '/' FILE_HASH = 'Sha256' @@ -16,7 +20,7 @@ API_HTTP_CLIENT_TIMEOUT = 60 API_RESPONSE_FIELD_DATA = 'Data' API_FILE_DOWNLOAD_RETRY_TIMES = 5 API_FILE_DOWNLOAD_TIMEOUT = 60 * 5 -API_FILE_DOWNLOAD_CHUNK_SIZE = 4096 +API_FILE_DOWNLOAD_CHUNK_SIZE = 1024 * 1024 * 16 API_RESPONSE_FIELD_GIT_ACCESS_TOKEN = 'AccessToken' API_RESPONSE_FIELD_USERNAME = 'Username' API_RESPONSE_FIELD_EMAIL = 'Email' diff --git a/modelscope/hub/file_download.py b/modelscope/hub/file_download.py index 6d3ad63d..cb95e164 100644 --- a/modelscope/hub/file_download.py +++ b/modelscope/hub/file_download.py @@ -3,6 +3,8 @@ import copy import os import tempfile +import threading +from concurrent.futures import ThreadPoolExecutor from functools import partial from http.cookiejar import CookieJar from pathlib import Path @@ -13,9 +15,10 @@ from requests.adapters import Retry from tqdm import tqdm from modelscope.hub.api import HubApi, ModelScopeConfig -from modelscope.hub.constants import (API_FILE_DOWNLOAD_CHUNK_SIZE, - API_FILE_DOWNLOAD_RETRY_TIMES, - API_FILE_DOWNLOAD_TIMEOUT, FILE_HASH) +from modelscope.hub.constants import ( + API_FILE_DOWNLOAD_CHUNK_SIZE, API_FILE_DOWNLOAD_RETRY_TIMES, + API_FILE_DOWNLOAD_TIMEOUT, FILE_HASH, MODELSCOPE_DOWNLOAD_PARALLELS, + MODELSCOPE_PARALLEL_DOWNLOAD_THRESHOLD_MB) from modelscope.utils.constant import DEFAULT_MODEL_REVISION from modelscope.utils.logger import get_logger from .errors import FileDownloadError, NotExistError @@ -134,19 +137,25 @@ def model_file_download( # we need to download again url_to_download = get_file_download_url(model_id, file_path, revision) - file_to_download_info = { - 'Path': file_path, - 'Revision': file_to_download_info['Revision'], - FILE_HASH: file_to_download_info[FILE_HASH] - } - temp_file_name = next(tempfile._get_candidate_names()) - http_get_file( - url_to_download, - temporary_cache_dir, - temp_file_name, - headers=headers, - cookies=None if cookies is None else cookies.get_dict()) + + if MODELSCOPE_PARALLEL_DOWNLOAD_THRESHOLD_MB * 1000 * 1000 < file_to_download_info[ + 'Size'] and MODELSCOPE_DOWNLOAD_PARALLELS > 1: + parallel_download( + url_to_download, + temporary_cache_dir, + temp_file_name, + headers=headers, + cookies=None if cookies is None else cookies.get_dict(), + file_size=file_to_download_info['Size']) + else: + http_get_file( + url_to_download, + temporary_cache_dir, + temp_file_name, + headers=headers, + cookies=None if cookies is None else cookies.get_dict()) + temp_file_path = os.path.join(temporary_cache_dir, temp_file_name) # for download with commit we can't get Sha256 if file_to_download_info[FILE_HASH] is not None: @@ -178,6 +187,66 @@ def get_file_download_url(model_id: str, file_path: str, revision: str): ) +def download_part(params): + # unpack parameters + progress, start, end, url, file_name, cookies, headers = params + get_headers = {} if headers is None else copy.deepcopy(headers) + get_headers['Range'] = 'bytes=%s-%s' % (start, end) + with open(file_name, 'rb+') as f: + f.seek(start) + r = requests.get( + url, + stream=True, + headers=get_headers, + cookies=cookies, + timeout=API_FILE_DOWNLOAD_TIMEOUT) + for chunk in r.iter_content(chunk_size=API_FILE_DOWNLOAD_CHUNK_SIZE): + if chunk: # filter out keep-alive new chunks + f.write(chunk) + progress.update(len(chunk)) + + +def parallel_download( + url: str, + local_dir: str, + file_name: str, + cookies: CookieJar, + headers: Optional[Dict[str, str]] = None, + file_size: int = None, +): + # create temp file + temp_file_manager = partial( + tempfile.NamedTemporaryFile, mode='wb', dir=local_dir, delete=False) + with temp_file_manager() as temp_file: + progress = tqdm( + unit='B', + unit_scale=True, + unit_divisor=1024, + total=file_size, + initial=0, + desc='Downloading', + ) + PART_SIZE = 160 * 1024 * 1012 # every part is 160M + tasks = [] + for idx in range(int(file_size / PART_SIZE)): + start = idx * PART_SIZE + end = (idx + 1) * PART_SIZE - 1 + tasks.append( + (progress, start, end, url, temp_file.name, cookies, headers)) + if end + 1 < file_size: + tasks.append((progress, end + 1, file_size - 1, url, + temp_file.name, cookies, headers)) + parallels = MODELSCOPE_DOWNLOAD_PARALLELS if MODELSCOPE_DOWNLOAD_PARALLELS <= 4 else 4 + with ThreadPoolExecutor( + max_workers=parallels, + thread_name_prefix='download') as executor: + list(executor.map(download_part, tasks)) + + progress.close() + + os.replace(temp_file.name, os.path.join(local_dir, file_name)) + + def http_get_file( url: str, local_dir: str, diff --git a/modelscope/hub/snapshot_download.py b/modelscope/hub/snapshot_download.py index 25a97975..82313b15 100644 --- a/modelscope/hub/snapshot_download.py +++ b/modelscope/hub/snapshot_download.py @@ -10,8 +10,10 @@ from typing import Dict, List, Optional, Union from modelscope.hub.api import HubApi, ModelScopeConfig from modelscope.utils.constant import DEFAULT_MODEL_REVISION from modelscope.utils.logger import get_logger -from .constants import FILE_HASH -from .file_download import get_file_download_url, http_get_file +from .constants import (FILE_HASH, MODELSCOPE_DOWNLOAD_PARALLELS, + MODELSCOPE_PARALLEL_DOWNLOAD_THRESHOLD_MB) +from .file_download import (get_file_download_url, http_get_file, + parallel_download) from .utils.caching import ModelFileSystemCache from .utils.utils import (file_integrity_validation, get_cache_dir, model_id_to_group_owner_name) @@ -133,13 +135,24 @@ def snapshot_download(model_id: str, file_path=model_file['Path'], revision=revision) - # First download to /tmp - http_get_file( - url=url, - local_dir=temp_cache_dir, - file_name=model_file['Name'], - headers=headers, - cookies=cookies) + if MODELSCOPE_PARALLEL_DOWNLOAD_THRESHOLD_MB * 1000 * 1000 < model_file[ + 'Size'] and MODELSCOPE_DOWNLOAD_PARALLELS > 1: + parallel_download( + url, + temp_cache_dir, + model_file['Name'], + headers=headers, + cookies=None + if cookies is None else cookies.get_dict(), + file_size=model_file['Size']) + else: + http_get_file( + url, + temp_cache_dir, + model_file['Name'], + headers=headers, + cookies=cookies) + # check file integrity temp_file = os.path.join(temp_cache_dir, model_file['Name']) if FILE_HASH in model_file: diff --git a/requirements/framework.txt b/requirements/framework.txt index 69fa1f40..d98765bf 100644 --- a/requirements/framework.txt +++ b/requirements/framework.txt @@ -14,10 +14,11 @@ Pillow>=6.2.0 pyarrow>=6.0.0,!=9.0.0 python-dateutil>=2.1 pyyaml -requests +requests>=2.25 scipy setuptools simplejson>=3.3.0 sortedcontainers>=1.5.9 tqdm>=4.64.0 +urllib3>=1.26 yapf