[to #50607174]feat: support parallel download large model file

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/13113235
* [to #50607174]feat: support parallel download large model file

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'

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,

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:

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