Fix datasets issues (#948)

* update

* user-custom timeout and retry in HubApi

* update

* del download count for snapshot_download

modelscope/hub/api.py

@@ -22,7 +22,8 @@ import requests
 from requests import Session
 from requests.adapters import HTTPAdapter, Retry
 
-from modelscope.hub.constants import (API_HTTP_CLIENT_TIMEOUT,
+from modelscope.hub.constants import (API_HTTP_CLIENT_MAX_RETRIES,
+                                      API_HTTP_CLIENT_TIMEOUT,
                                       API_RESPONSE_FIELD_DATA,
                                       API_RESPONSE_FIELD_EMAIL,
                                       API_RESPONSE_FIELD_GIT_ACCESS_TOKEN,
@@ -61,7 +62,10 @@ logger = get_logger()
 class HubApi:
     """Model hub api interface.
     """
-    def __init__(self, endpoint: Optional[str] = None, timeout=API_HTTP_CLIENT_TIMEOUT):
+    def __init__(self,
+                 endpoint: Optional[str] = None,
+                 timeout=API_HTTP_CLIENT_TIMEOUT,
+                 max_retries=API_HTTP_CLIENT_MAX_RETRIES):
         """The ModelScope HubApi。
 
         Args:
@@ -71,7 +75,7 @@ class HubApi:
         self.headers = {'user-agent': ModelScopeConfig.get_user_agent()}
         self.session = Session()
         retry = Retry(
-            total=2,
+            total=max_retries,
             read=2,
             connect=2,
             backoff_factor=1,

modelscope/hub/constants.py

@@ -17,6 +17,7 @@ LOGGER_NAME = 'ModelScopeHub'
 DEFAULT_CREDENTIALS_PATH = Path.home().joinpath('.modelscope', 'credentials')
 REQUESTS_API_HTTP_METHOD = ['get', 'head', 'post', 'put', 'patch', 'delete']
 API_HTTP_CLIENT_TIMEOUT = 60
+API_HTTP_CLIENT_MAX_RETRIES = 2
 API_RESPONSE_FIELD_DATA = 'Data'
 API_FILE_DOWNLOAD_RETRY_TIMES = 5
 API_FILE_DOWNLOAD_TIMEOUT = 60

modelscope/hub/snapshot_download.py

@@ -211,6 +211,7 @@ def _snapshot_download(
             ModelScopeConfig.get_user_agent(user_agent=user_agent, ),
         }
         if 'CI_TEST' not in os.environ:
+            # To count the download statistics, to add the snapshot-identifier as a header.
             headers['snapshot-identifier'] = str(uuid.uuid4())
         _api = HubApi()
         if cookies is None:
@@ -259,7 +260,6 @@ def _snapshot_download(
             group_or_owner, name = model_id_to_group_owner_name(repo_id)
             if not revision:
                 revision = DEFAULT_DATASET_REVISION
-            _api.dataset_download_statistics(name, group_or_owner)
             revision_detail = revision
             page_number = 1
             page_size = 100
@@ -412,6 +412,10 @@ def _download_file_lists(
                 dataset_name=name,
                 namespace=group_or_owner,
                 revision=revision)
+        else:
+            raise InvalidParameter(
+                f'Invalid repo type: {repo_type}, supported types: {REPO_TYPE_SUPPORT}'
+            )
 
         download_file(url, repo_file, temporary_cache_dir, cache, headers,
                       cookies)

modelscope/msdatasets/utils/hf_datasets_util.py

@@ -48,7 +48,7 @@ from datasets.utils.file_utils import (OfflineModeIsEnabled,
                                        relative_to_absolute_path)
 from datasets.utils.info_utils import is_small_dataset
 from datasets.utils.metadata import MetadataConfigs
-from datasets.utils.py_utils import get_imports
+from datasets.utils.py_utils import get_imports, map_nested
 from datasets.utils.track import tracked_str
 from fsspec import filesystem
 from fsspec.core import _un_chain
@@ -218,7 +218,7 @@ def _list_repo_tree(
     token: Optional[Union[bool, str]] = None,
 ) -> Iterable[Union[RepoFile, RepoFolder]]:
 
-    _api = HubApi()
+    _api = HubApi(timeout=3 * 60, max_retries=3)
 
     if is_relative_path(repo_id) and repo_id.count('/') == 1:
         _namespace, _dataset_name = repo_id.split('/')
@@ -231,7 +231,6 @@ def _list_repo_tree(
 
     page_number = 1
     page_size = 100
-    total_data_list = []
     while True:
         data: dict = _api.list_repo_tree(dataset_name=_dataset_name,
                                          namespace=_namespace,
@@ -247,7 +246,6 @@ def _list_repo_tree(
 
         # Parse data (Type: 'tree' or 'blob')
         data_file_list: list = data['Data']['Files']
-        total_data_list.extend(data_file_list)
 
         for file_info_d in data_file_list:
             path_info = {}
@@ -398,7 +396,10 @@ def _resolve_pattern(
         # 10 times faster glob with detail=True (ignores costly info like lastCommit)
         glob_kwargs['expand_info'] = False
 
-    tmp_file_paths = fs.glob(pattern, detail=True, **glob_kwargs)
+    try:
+        tmp_file_paths = fs.glob(pattern, detail=True, **glob_kwargs)
+    except FileNotFoundError:
+        raise DataFilesNotFoundError(f"Unable to find '{pattern}'")
 
     matched_paths = [
         filepath if filepath.startswith(protocol_prefix) else protocol_prefix
@@ -1300,6 +1301,7 @@ class DatasetsWrapperHF:
                     ).get_module()
             except Exception as e1:
                 # All the attempts failed, before raising the error we should check if the module is already cached
+                logger.error(f'>> Error loading {path}: {e1}')
                 try:
                     return CachedDatasetModuleFactory(
                         path,
@@ -1330,6 +1332,8 @@ class DatasetsWrapperHF:
 
 @contextlib.contextmanager
 def load_dataset_with_ctx(*args, **kwargs):
+
+    # Keep the original functions
     hf_endpoint_origin = config.HF_ENDPOINT
     get_from_cache_origin = file_utils.get_from_cache
 
@@ -1344,15 +1348,14 @@ def load_dataset_with_ctx(*args, **kwargs):
     get_module_without_script_origin = HubDatasetModuleFactoryWithoutScript.get_module
     get_module_with_script_origin = HubDatasetModuleFactoryWithScript.get_module
 
+    # Monkey patching with modelscope functions
     config.HF_ENDPOINT = get_endpoint()
     file_utils.get_from_cache = get_from_cache_ms
-
     # Compatible with datasets 2.18.0
     if hasattr(DownloadManager, '_download'):
         DownloadManager._download = _download_ms
     else:
         DownloadManager._download_single = _download_ms
-
     HfApi.dataset_info = _dataset_info
     HfApi.list_repo_tree = _list_repo_tree
     HfApi.get_paths_info = _get_paths_info
@@ -1366,6 +1369,9 @@ def load_dataset_with_ctx(*args, **kwargs):
         dataset_res = DatasetsWrapperHF.load_dataset(*args, **kwargs)
         yield dataset_res
     finally:
+        # Restore the original functions
+        config.HF_ENDPOINT = hf_endpoint_origin
+        file_utils.get_from_cache = get_from_cache_origin
         # Keep the context during the streaming iteration
         if not streaming:
             config.HF_ENDPOINT = hf_endpoint_origin