From 2719754beb8992e8357cd4d6df35134dc40121c0 Mon Sep 17 00:00:00 2001 From: AAAkater <125126227+AAAkater@users.noreply.github.com> Date: Wed, 3 Jun 2026 18:39:39 +0800 Subject: [PATCH 01/11] =?UTF-8?q?=F0=9F=90=9B=20fix(clearcache):=20fix=20s?= =?UTF-8?q?ingle=20model=20cache=20path=20(#1724)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🐛 fix(clearcache): fix single model cache path - Add 'models' subdirectory under 'hub' for single model cache paths - Align dataset and model path structures for consistency * 🐛 fix(clearcache): use cache root helpers to handle MODELSCOPE_CACHE correctly Unconditionally prepending `hub/` broke `$MODELSCOPE_CACHE`-set users, since the cache root has no `hub/` segment in that case. Reuse get_model_cache_root() / get_dataset_cache_root() from file_utils, which already handle both default and env-set layouts. Add tests covering both scenarios. --------- Co-authored-by: Yunnglin --- modelscope/cli/clearcache.py | 18 ++++++------- tests/cli/test_scancache_cmd.py | 46 +++++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 9 deletions(-) diff --git a/modelscope/cli/clearcache.py b/modelscope/cli/clearcache.py index dcd3d1df..0713db94 100644 --- a/modelscope/cli/clearcache.py +++ b/modelscope/cli/clearcache.py @@ -2,11 +2,13 @@ import os import shutil from argparse import ArgumentParser -from pathlib import Path from modelscope.cli.base import CLICommand from modelscope.hub.constants import TEMPORARY_FOLDER_NAME from modelscope.hub.utils.utils import get_model_masked_directory +from modelscope.utils.file_utils import (get_dataset_cache_root, + get_model_cache_root, + get_modelscope_cache_dir) def subparser_func(args): @@ -20,9 +22,7 @@ class ClearCacheCMD(CLICommand): def __init__(self, args): self.args = args - self.cache_dir = os.getenv( - 'MODELSCOPE_CACHE', - Path.home().joinpath('.cache', 'modelscope')) + self.cache_dir = get_modelscope_cache_dir() @staticmethod def define_args(parsers: ArgumentParser): @@ -76,11 +76,11 @@ class ClearCacheCMD(CLICommand): self._remove_directory(self.cache_dir) print('Cache cleared.') else: - entity_directory = os.path.join( - self.cache_dir, 'hub' if single_model else 'datasets', id) - temp_directory = os.path.join( - self.cache_dir, 'hub' if single_model else 'datasets', - TEMPORARY_FOLDER_NAME, id) + entity_root = get_model_cache_root( + ) if single_model else get_dataset_cache_root() + entity_directory = os.path.join(entity_root, id) + temp_directory = os.path.join(entity_root, + TEMPORARY_FOLDER_NAME, id) entity_removed = self._remove_directory(entity_directory) temp_removed = self._remove_directory(temp_directory) if (not entity_removed) and (not temp_removed): diff --git a/tests/cli/test_scancache_cmd.py b/tests/cli/test_scancache_cmd.py index a5c409b6..bde80db8 100644 --- a/tests/cli/test_scancache_cmd.py +++ b/tests/cli/test_scancache_cmd.py @@ -1,8 +1,12 @@ import os import subprocess import sys +import tempfile import unittest +from argparse import Namespace +from unittest import mock +from modelscope.hub.constants import TEMPORARY_FOLDER_NAME from modelscope.utils.file_utils import get_modelscope_cache_dir @@ -36,5 +40,47 @@ class TestScanCacheCommand(unittest.TestCase): self.assertIn('not found', output) +class TestClearCacheCommand(unittest.TestCase): + """clear-cache must resolve the right cache root in both scenarios: + * MODELSCOPE_CACHE set -> $MODELSCOPE_CACHE/{models,datasets}/ + * MODELSCOPE_CACHE unset -> ~/.cache/modelscope/hub/{models,datasets}/ + """ + + def _clear_and_assert(self, root, kind, entity_id): + from modelscope.cli.clearcache import ClearCacheCMD + sub = {'model': 'models', 'dataset': 'datasets'}[kind] + entity_dir = os.path.join(root, sub, entity_id) + temp_dir = os.path.join(root, sub, TEMPORARY_FOLDER_NAME, entity_id) + os.makedirs(entity_dir) + os.makedirs(temp_dir) + args = Namespace( + model=entity_id if kind == 'model' else None, + dataset=entity_id if kind == 'dataset' else None) + with mock.patch('builtins.input', return_value='Y'): + ClearCacheCMD(args).execute() + self.assertFalse(os.path.exists(entity_dir)) + self.assertFalse(os.path.exists(temp_dir)) + + def test_env_cache(self): + with tempfile.TemporaryDirectory() as tmp, \ + mock.patch.dict(os.environ, + {'MODELSCOPE_CACHE': tmp}, clear=False): + for kind in ('model', 'dataset'): + with self.subTest(kind=kind): + self._clear_and_assert(tmp, kind, f'org/{kind}') + + def test_default_cache(self): + with tempfile.TemporaryDirectory() as tmp: + env = { + k: v + for k, v in os.environ.items() if k != 'MODELSCOPE_CACHE' + } + env['HOME'] = tmp + with mock.patch.dict(os.environ, env, clear=True): + self._clear_and_assert( + os.path.join(tmp, '.cache', 'modelscope', 'hub'), 'model', + 'org/repo') + + if __name__ == '__main__': unittest.main() From 491e3df224cec1c931543e33523bc7fe737702ac Mon Sep 17 00:00:00 2001 From: Seas0 Date: Wed, 3 Jun 2026 18:50:04 +0800 Subject: [PATCH 02/11] fix: quote version of `timm` specs in Dockerfile.ubuntu (#1728) Co-authored-by: Seas0 --- docker/Dockerfile.ubuntu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile.ubuntu b/docker/Dockerfile.ubuntu index e74c3374..0f53ff9e 100644 --- a/docker/Dockerfile.ubuntu +++ b/docker/Dockerfile.ubuntu @@ -33,7 +33,7 @@ if [ "$INSTALL_MS_DEPS" = "True" ]; then \ pip install --no-cache-dir https://modelscope.oss-cn-beijing.aliyuncs.com/packages/imageio_ffmpeg-0.4.9-py3-none-any.whl --no-dependencies --force && \ pip install --no-cache-dir 'scipy<1.13.0' && \ pip install --no-cache-dir funtextprocessing typeguard==2.13.3 scikit-learn -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \ - pip install --no-cache-dir decord>=0.6.0 mpi4py paint_ldm ipykernel fasttext -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \ + pip install --no-cache-dir 'decord>=0.6.0' mpi4py paint_ldm ipykernel fasttext -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \ pip install --no-cache-dir 'blobfile>=1.0.5' && \ pip uninstall MinDAEC -y && \ pip install https://modelscope.oss-cn-beijing.aliyuncs.com/releases/dependencies/MinDAEC-0.0.2-py3-none-any.whl && \ @@ -57,7 +57,7 @@ RUN bash /tmp/install.sh {version_args} && \ cd modelscope && pip install . -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \ cd / && rm -fr /tmp/modelscope && pip cache purge; \ pip install --no-cache-dir torch=={torch_version} torchvision=={torchvision_version} torchaudio=={torchaudio_version} {index_url} && \ - pip install --no-cache-dir transformers diffusers timm>=0.9.0 && pip cache purge; \ + pip install --no-cache-dir transformers diffusers 'timm>=0.9.0' && pip cache purge; \ pip install --no-cache-dir omegaconf==2.3.0 && pip cache purge; \ pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \ pip config set install.trusted-host mirrors.aliyun.com && \ From 63ff6ec6e96fe2dd72d2468a774d89b8e254b4ae Mon Sep 17 00:00:00 2001 From: "Xingjun.Wang" Date: Fri, 5 Jun 2026 14:30:18 +0800 Subject: [PATCH 03/11] Add`gated_mode` for create repo (#1730) --- modelscope/cli/create.py | 22 ++++++++++++++--- modelscope/hub/api.py | 48 +++++++++++++++++++++++++++++++++---- modelscope/hub/constants.py | 12 ++++++++++ 3 files changed, 74 insertions(+), 8 deletions(-) diff --git a/modelscope/cli/create.py b/modelscope/cli/create.py index 3285ab91..76e67e43 100644 --- a/modelscope/cli/create.py +++ b/modelscope/cli/create.py @@ -3,8 +3,8 @@ from argparse import ArgumentParser, _SubParsersAction from modelscope.cli.base import CLICommand from modelscope.hub.api import HubApi -from modelscope.hub.constants import (Licenses, ModelVisibility, Visibility, - VisibilityMap) +from modelscope.hub.constants import (GatedMode, Licenses, ModelVisibility, + Visibility, VisibilityMap) from modelscope.hub.utils.aigc import AigcModel from modelscope.hub.utils.utils import resolve_endpoint from modelscope.utils.constant import REPO_TYPE_MODEL, REPO_TYPE_SUPPORT @@ -82,6 +82,20 @@ class CreateCMD(CLICommand): help= 'If True, do not raise error when repo already exists. Defaults to False.', ) + parser.add_argument( + '--gated', + dest='gated_mode', + action='store_true', + default=None, + help= + 'Enable gated mode (application-based download) for private repos.', + ) + parser.add_argument( + '--no-gated', + dest='gated_mode', + action='store_false', + help='Disable gated mode for private repos (normal private).', + ) parser.add_argument( '--endpoint', type=str, @@ -176,6 +190,7 @@ class CreateCMD(CLICommand): exist_ok=self.args.exist_ok, create_default_config=True, endpoint=endpoint, + gated_mode=self.args.gated_mode, ) def _create_aigc_model(self): @@ -225,7 +240,8 @@ class CreateCMD(CLICommand): visibility=visibility_idx, license=self.args.license, chinese_name=self.args.chinese_name, - aigc_model=aigc_model) + aigc_model=aigc_model, + gated_mode=self.args.gated_mode) print(f'Successfully created AIGC model: {model_url}') except Exception as e: print(f'Error creating AIGC model: {e}') diff --git a/modelscope/hub/api.py b/modelscope/hub/api.py index 8a7b8409..f223eed0 100644 --- a/modelscope/hub/api.py +++ b/modelscope/hub/api.py @@ -329,7 +329,8 @@ class HubApi: original_model_id: Optional[str] = '', endpoint: Optional[str] = None, token: Optional[str] = None, - aigc_model: Optional['AigcModel'] = None) -> str: + aigc_model: Optional['AigcModel'] = None, + gated_mode: Optional[bool] = None) -> str: """Create model repo at ModelScope Hub. Args: @@ -343,6 +344,9 @@ class HubApi: aigc_model (AigcModel, optional): AigcModel instance for AIGC model creation. If provided, will create an AIGC model with automatic file upload. Refer to modelscope.hub.utils.aigc.AigcModel for details. + gated_mode (bool, optional): Gated mode for private repos. + True = gated (application-based download), False = off (normal private). + Only effective when visibility is PRIVATE (1). Returns: str: URL of the created model repository @@ -374,6 +378,12 @@ class HubApi: 'TrainId': os.environ.get('MODELSCOPE_TRAIN_ID', '') } + if gated_mode is not None: + if visibility != ModelVisibility.PRIVATE: + logger.warning('gated_mode is only effective when visibility is PRIVATE, ignored.') + else: + body['ProtectedMode'] = 1 if gated_mode else 2 + # Set path based on model type if aigc_model is not None: # Use AIGC model endpoint @@ -1428,7 +1438,8 @@ class HubApi: visibility: Optional[int] = DatasetVisibility.PUBLIC, description: Optional[str] = '', endpoint: Optional[str] = None, - token: Optional[str] = None) -> str: + token: Optional[str] = None, + gated_mode: Optional[bool] = None) -> str: """ Create a dataset in ModelScope. @@ -1441,6 +1452,9 @@ class HubApi: description (str, optional): The description of the dataset. Defaults to ''. endpoint (str, optional): The endpoint to use. If not provided, the default endpoint is used. token (str, optional): The access token for authentication. + gated_mode (bool, optional): Gated mode for private repos. + True = gated (application-based download), False = off (normal private). + Only effective when visibility is PRIVATE (1). Returns: str: The URL of the created dataset repository. @@ -1462,6 +1476,12 @@ class HubApi: 'Description': (None, description) } + if gated_mode is not None: + if visibility != DatasetVisibility.PRIVATE: + logger.warning('gated_mode is only effective when visibility is PRIVATE, ignored.') + else: + files['ProtectedMode'] = (None, 1 if gated_mode else 2) + r = self.session.post( path, files=files, @@ -2200,6 +2220,7 @@ class HubApi: exist_ok: Optional[bool] = False, create_default_config: Optional[bool] = True, aigc_model: Optional[AigcModel] = None, + gated_mode: Optional[bool] = None, **kwargs, ) -> str: """ @@ -2217,6 +2238,9 @@ class HubApi: In the format of `https://www.modelscope.cn` or 'https://www.modelscope.ai' exist_ok (Optional[bool]): If the repo exists, whether to return the repo url directly. create_default_config (Optional[bool]): If True, create a default configuration file in the model repo. + gated_mode (Optional[bool]): Gated mode for private repos. + True = gated (application-based download), False = off (normal private). + Only effective when visibility is ``private``. **kwargs: The additional arguments. Returns: @@ -2256,6 +2280,7 @@ class HubApi: aigc_model=aigc_model, token=token, endpoint=endpoint, + gated_mode=gated_mode, ) if create_default_config: with tempfile.TemporaryDirectory() as temp_cache_dir: @@ -2290,6 +2315,7 @@ class HubApi: visibility=visibility, token=token, endpoint=endpoint, + gated_mode=gated_mode, ) print(f'New dataset created successfully at {repo_url}.', flush=True) @@ -3882,7 +3908,8 @@ class HubApi: repo_id: str, repo_type: Literal['model', 'dataset'], visibility: Literal['private', 'public'], - token: Union[str, None] = None + token: Union[str, None] = None, + gated_mode: Optional[bool] = None, ) -> dict: """ Set the visibility of a repo. @@ -3893,6 +3920,9 @@ class HubApi: visibility (Literal['private', 'public']): The visibility to set, `private` or `public`. token (Union[str, None]): The access token. If None, will use the cookies from the local cache. See `https://modelscope.cn/my/myaccesstoken` to get your token. + gated_mode (Optional[bool]): Gated mode for private repos. + True = gated (application-based download), False = off (normal private). + Only effective when visibility is ``private``. Returns: dict: The response from the server. @@ -3907,6 +3937,10 @@ class HubApi: visibility_code: int = visibility_map.get(visibility, 5) cookies = self.get_cookies(access_token=token, cookies_required=True) + if gated_mode is not None and visibility != 'private': + logger.warning('gated_mode is only effective when visibility is private, ignored.') + gated_mode = None + if repo_type == REPO_TYPE_MODEL: model_info = self.get_model(model_id=repo_id, token=token) path = f'{self.endpoint}/api/v1/models/{repo_id}' @@ -3916,11 +3950,15 @@ class HubApi: first = tasks[0] if isinstance(first, dict) and first: model_tasks = first.get('name') + if gated_mode is not None: + pm = 1 if gated_mode else 2 + else: + pm = model_info.get('ProtectedMode', 2) payload = { 'ChineseName': model_info.get('ChineseName', ''), 'ModelFramework': model_info.get('ModelFramework', 'Pytorch'), 'Visibility': visibility_code, - 'ProtectedMode': 2, + 'ProtectedMode': pm, 'ApprovalMode': model_info.get('ApprovalMode', 2), 'Description': model_info.get('Description', ''), 'AigcType': model_info.get('AigcType', ''), @@ -3947,7 +3985,7 @@ class HubApi: path = f'{self.endpoint}/api/v1/datasets/{dataset_idx}' payload = { 'Visibility': visibility_code, - 'ProtectedMode': 2, + 'ProtectedMode': (1 if gated_mode else 2) if gated_mode is not None else 2, } else: raise ValueError(f'Invalid repo type: {repo_type}, supported repos: {REPO_TYPE_SUPPORT}') diff --git a/modelscope/hub/constants.py b/modelscope/hub/constants.py index 85ad69d4..c3ced124 100644 --- a/modelscope/hub/constants.py +++ b/modelscope/hub/constants.py @@ -169,6 +169,18 @@ class Visibility(object): PUBLIC = 'public' +class GatedMode(object): + """Gated mode for private repositories. + + Only effective when Visibility is PRIVATE. + API payload key: ``ProtectedMode``. + Values: True = gated (application-based download), + False = off (normal private). + """ + GATED = True + OFF = False + + VisibilityMap = { ModelVisibility.PRIVATE: Visibility.PRIVATE, ModelVisibility.INTERNAL: Visibility.INTERNAL, From 1804ac6bf8b6b76c9cc43974e4822e27536bc6e5 Mon Sep 17 00:00:00 2001 From: "Xingjun.Wang" Date: Fri, 5 Jun 2026 14:33:25 +0800 Subject: [PATCH 04/11] [Feature] Add studio module (#1727) --- modelscope/cli/cli.py | 2 + modelscope/cli/create.py | 42 ++- modelscope/cli/download.py | 14 +- modelscope/cli/studio.py | 384 ++++++++++++++++++++++ modelscope/hub/api.py | 306 ++++++++++++++++-- modelscope/hub/errors.py | 41 ++- modelscope/hub/snapshot_download.py | 19 +- modelscope/utils/constant.py | 37 ++- tests/run_analysis.py | 7 +- tests/studios/.env.example | 31 ++ tests/studios/__init__.py | 0 tests/studios/conftest_env.py | 133 ++++++++ tests/studios/test_studio_cli.py | 476 ++++++++++++++++++++++++++++ 13 files changed, 1452 insertions(+), 40 deletions(-) create mode 100644 modelscope/cli/studio.py create mode 100644 tests/studios/.env.example create mode 100644 tests/studios/__init__.py create mode 100644 tests/studios/conftest_env.py create mode 100644 tests/studios/test_studio_cli.py diff --git a/modelscope/cli/cli.py b/modelscope/cli/cli.py index 83dda930..7a45e3fc 100644 --- a/modelscope/cli/cli.py +++ b/modelscope/cli/cli.py @@ -14,6 +14,7 @@ from modelscope.cli.plugins import PluginsCMD from modelscope.cli.scancache import ScanCacheCMD from modelscope.cli.server import ServerCMD from modelscope.cli.skills import SkillsCMD +from modelscope.cli.studio import StudioCMD from modelscope.cli.upload import UploadCMD from modelscope.hub.constants import MODELSCOPE_ASCII from modelscope.utils.logger import get_logger @@ -47,6 +48,7 @@ def run_cmd(): LoginCMD.define_args(subparsers) LlamafileCMD.define_args(subparsers) ScanCacheCMD.define_args(subparsers) + StudioCMD.define_args(subparsers) args = parser.parse_args() diff --git a/modelscope/cli/create.py b/modelscope/cli/create.py index 76e67e43..5fc516c6 100644 --- a/modelscope/cli/create.py +++ b/modelscope/cli/create.py @@ -7,7 +7,9 @@ from modelscope.hub.constants import (GatedMode, Licenses, ModelVisibility, Visibility, VisibilityMap) from modelscope.hub.utils.aigc import AigcModel from modelscope.hub.utils.utils import resolve_endpoint -from modelscope.utils.constant import REPO_TYPE_MODEL, REPO_TYPE_SUPPORT +from modelscope.utils.constant import (REPO_TYPE_MODEL, REPO_TYPE_STUDIO, + REPO_TYPE_SUPPORT, StudioHardware, + StudioSDKType) from modelscope.utils.logger import get_logger logger = get_logger() @@ -107,6 +109,35 @@ class CreateCMD(CLICommand): 'then defaults to https://www.modelscope.cn.', ) + # Studio specific arguments (only meaningful when --repo_type studio) + studio_group = parser.add_argument_group( + 'Studio Repo Creation', + 'Optional arguments used only when `--repo_type studio` is set.') + studio_group.add_argument( + '--sdk-type', + dest='sdk_type', + choices=StudioSDKType.SUPPORTED, + default=None, + help='Studio SDK type (only for studio repo-type).') + studio_group.add_argument( + '--sdk-version', + dest='sdk_version', + type=str, + default=None, + help='Studio SDK version (only for gradio).') + studio_group.add_argument( + '--base-image', + dest='base_image', + type=str, + default=None, + help='Studio base image (only for gradio/streamlit).') + studio_group.add_argument( + '--hardware', + dest='hardware', + choices=StudioHardware.SUPPORTED, + default=None, + help='Studio hardware configuration.') + # AIGC specific arguments aigc_group = parser.add_argument_group( 'AIGC Model Creation', @@ -179,6 +210,14 @@ class CreateCMD(CLICommand): endpoint = resolve_endpoint(self.args.endpoint) api = HubApi(endpoint=endpoint) + extra_kwargs = {} + if self.args.repo_type == REPO_TYPE_STUDIO: + # Pass studio-specific fields only when creating a studio repo. + for field in ('sdk_type', 'sdk_version', 'base_image', 'hardware'): + value = getattr(self.args, field, None) + if value is not None: + extra_kwargs[field] = value + # Create repo api.create_repo( repo_id=self.args.repo_id, @@ -191,6 +230,7 @@ class CreateCMD(CLICommand): create_default_config=True, endpoint=endpoint, gated_mode=self.args.gated_mode, + **extra_kwargs, ) def _create_aigc_model(self): diff --git a/modelscope/cli/download.py b/modelscope/cli/download.py index 8c46d61c..b81c9a54 100644 --- a/modelscope/cli/download.py +++ b/modelscope/cli/download.py @@ -11,7 +11,9 @@ from modelscope.hub.file_download import (dataset_file_download, from modelscope.hub.snapshot_download import (dataset_snapshot_download, snapshot_download) from modelscope.hub.utils.utils import convert_patterns, resolve_endpoint -from modelscope.utils.constant import DEFAULT_DATASET_REVISION +from modelscope.utils.constant import (DEFAULT_DATASET_REVISION, + REPO_TYPE_DATASET, REPO_TYPE_MODEL, + REPO_TYPE_STUDIO, REPO_TYPE_SUPPORT) from modelscope.utils.logger import get_logger logger = get_logger(log_level=logging.WARNING) @@ -60,8 +62,8 @@ class DownloadCMD(CLICommand): ) parser.add_argument( '--repo-type', - choices=['model', 'dataset'], - default='model', + choices=REPO_TYPE_SUPPORT, + default=REPO_TYPE_MODEL, help="Type of repo to download from (defaults to 'model').", ) parser.add_argument( @@ -135,9 +137,11 @@ class DownloadCMD(CLICommand): self.args.files = [self.args.repo_id] else: if self.args.repo_id is not None: - if self.args.repo_type == 'model': + if self.args.repo_type in (REPO_TYPE_MODEL, REPO_TYPE_STUDIO): + # studio repos share the same snapshot_download path + # as model repos. self.args.model = self.args.repo_id - elif self.args.repo_type == 'dataset': + elif self.args.repo_type == REPO_TYPE_DATASET: self.args.dataset = self.args.repo_id else: raise Exception('Not support repo-type: %s' diff --git a/modelscope/cli/studio.py b/modelscope/cli/studio.py new file mode 100644 index 00000000..31a5d93f --- /dev/null +++ b/modelscope/cli/studio.py @@ -0,0 +1,384 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +"""ModelScope Studio runtime management CLI. + +This module exposes a ``modelscope studio`` command group that wraps the +Studio OpenAPI methods on :class:`~modelscope.hub.api.HubApi`. Each CLI +subcommand is a thin adapter that parses arguments, delegates to a single +``HubApi`` method and renders a concise human-readable result. + +Subcommand layout:: + + modelscope studio deploy + modelscope studio stop + modelscope studio logs [--type ...] [--keyword ...] ... + modelscope studio settings [--display-name ...] ... + modelscope studio secret list + modelscope studio secret add + modelscope studio secret update + modelscope studio secret delete +""" +from argparse import ArgumentParser, _SubParsersAction + +import json + +from modelscope.cli.base import CLICommand +from modelscope.hub.api import HubApi +from modelscope.hub.utils.utils import resolve_endpoint +from modelscope.utils.constant import StudioSDKType +from modelscope.utils.logger import get_logger + +logger = get_logger() + +# Common log type choices for `studio logs --type ...`. +_LOG_TYPES = ('run', 'build') + + +def subparser_func(args): + """ Function which will be called for a specific sub parser. + """ + return StudioCMD(args) + + +class StudioCMD(CLICommand): + """Studio runtime management commands. + + Dispatches to a per-subcommand handler based on ``args.studio_action`` + (and, for ``secret``, ``args.secret_action``). + """ + + name = 'studio' + + def __init__(self, args): + self.args = args + + # ------------------------------------------------------------------ + # Argument parsing + # ------------------------------------------------------------------ + @staticmethod + def define_args(parsers: _SubParsersAction): + parser: ArgumentParser = parsers.add_parser( + StudioCMD.name, help='Manage ModelScope studios at runtime.') + # Common args available to every studio subcommand. + parser.add_argument( + '--token', + type=str, + default=None, + help='Optional access token used for authentication.') + parser.add_argument( + '--endpoint', + type=str, + default=None, + help='ModelScope server endpoint. Falls back to env ' + 'MODELSCOPE_DOMAIN, then https://www.modelscope.cn.') + + action_subparsers = parser.add_subparsers( + dest='studio_action', help='Studio subcommands.') + + StudioCMD._add_simple_id_parser( + action_subparsers, + name='deploy', + help_text='Deploy (re-pull and rebuild) a studio.') + StudioCMD._add_simple_id_parser( + action_subparsers, name='stop', help_text='Stop a running studio.') + StudioCMD._add_logs_parser(action_subparsers) + StudioCMD._add_settings_parser(action_subparsers) + StudioCMD._add_secret_parser(action_subparsers) + + parser.set_defaults(func=subparser_func) + + @staticmethod + def _add_simple_id_parser(action_subparsers, name, help_text): + sub = action_subparsers.add_parser(name, help=help_text) + sub.add_argument( + 'studio_id', + type=str, + help='Studio ID in the format `owner/repo_name`.') + + @staticmethod + def _add_logs_parser(action_subparsers): + sub = action_subparsers.add_parser( + 'logs', help='Fetch studio runtime or build logs.') + sub.add_argument( + 'studio_id', + type=str, + help='Studio ID in the format `owner/repo_name`.') + sub.add_argument( + '--type', + dest='log_type', + choices=_LOG_TYPES, + default='run', + help="Log type to fetch (defaults to 'run').") + sub.add_argument( + '--keyword', + type=str, + default=None, + help='Optional keyword to filter log lines.') + sub.add_argument( + '--page-num', + dest='page_num', + type=int, + default=1, + help='Page number, starting from 1.') + sub.add_argument( + '--page-size', + dest='page_size', + type=int, + default=100, + help='Number of log entries per page.') + sub.add_argument( + '--start-timestamp', + dest='start_timestamp', + type=int, + default=None, + help='Optional start timestamp (seconds since epoch).') + sub.add_argument( + '--end-timestamp', + dest='end_timestamp', + type=int, + default=None, + help='Optional end timestamp (seconds since epoch).') + + @staticmethod + def _add_settings_parser(action_subparsers): + sub = action_subparsers.add_parser( + 'settings', + help='Update studio settings (only specified fields are modified).' + ) + sub.add_argument( + 'studio_id', + type=str, + help='Studio ID in the format `owner/repo_name`.') + sub.add_argument( + '--display-name', + dest='display_name', + type=str, + default=None, + help='Studio display name (Chinese name).') + sub.add_argument( + '--description', + type=str, + default=None, + help='Studio description.') + sub.add_argument( + '--license', type=str, default=None, help='Studio license.') + sub.add_argument( + '--cover-image', + dest='cover_image', + type=str, + default=None, + help='Studio cover image URL.') + sub.add_argument( + '--sdk-type', + dest='sdk_type', + choices=StudioSDKType.SUPPORTED, + default=None, + help='Studio SDK type (requires redeployment).') + sub.add_argument( + '--sdk-version', + dest='sdk_version', + type=str, + default=None, + help='Studio SDK version (requires redeployment).') + sub.add_argument( + '--base-image', + dest='base_image', + type=str, + default=None, + help='Studio base image (requires redeployment).') + sub.add_argument( + '--hardware', + type=str, + default=None, + help='Studio hardware configuration (requires redeployment).') + visibility_group = sub.add_mutually_exclusive_group() + visibility_group.add_argument( + '--private', + dest='private', + action='store_const', + const=True, + default=None, + help='Mark the studio as private.') + visibility_group.add_argument( + '--public', + dest='private', + action='store_const', + const=False, + help='Mark the studio as public.') + + @staticmethod + def _add_secret_parser(action_subparsers): + secret_parser = action_subparsers.add_parser( + 'secret', help='Manage studio environment variables (secrets).') + secret_subparsers = secret_parser.add_subparsers( + dest='secret_action', help='Secret subcommands.') + + list_sub = secret_subparsers.add_parser( + 'list', help='List secret keys (values are not returned).') + list_sub.add_argument('studio_id', type=str) + + add_sub = secret_subparsers.add_parser('add', help='Add a secret.') + add_sub.add_argument('studio_id', type=str) + add_sub.add_argument('key', type=str, help='Secret name.') + add_sub.add_argument('value', type=str, help='Secret value.') + + upd_sub = secret_subparsers.add_parser( + 'update', help='Update an existing secret.') + upd_sub.add_argument('studio_id', type=str) + upd_sub.add_argument('key', type=str, help='Secret name.') + upd_sub.add_argument('value', type=str, help='New secret value.') + + del_sub = secret_subparsers.add_parser( + 'delete', help='Delete a secret.') + del_sub.add_argument('studio_id', type=str) + del_sub.add_argument('key', type=str, help='Secret name to delete.') + + # ------------------------------------------------------------------ + # Execution + # ------------------------------------------------------------------ + def execute(self): + action = getattr(self.args, 'studio_action', None) + if action is None: + raise SystemExit( + 'No studio subcommand specified. ' + "Run 'modelscope studio --help' to see available subcommands.") + + endpoint = resolve_endpoint(self.args.endpoint) + api = HubApi(endpoint=endpoint) + + handlers = { + 'deploy': self._do_deploy, + 'stop': self._do_stop, + 'logs': self._do_logs, + 'settings': self._do_settings, + 'secret': self._do_secret, + } + handler = handlers.get(action) + if handler is None: + raise SystemExit(f'Unknown studio subcommand: {action}') + handler(api, endpoint) + + # -- deploy / stop -------------------------------------------------- + def _do_deploy(self, api, endpoint): + data = api.deploy_studio( + self.args.studio_id, token=self.args.token, endpoint=endpoint) + print(f'Deploy triggered for studio {self.args.studio_id}.') + self._print_status(data) + + def _do_stop(self, api, endpoint): + data = api.stop_studio( + self.args.studio_id, token=self.args.token, endpoint=endpoint) + print(f'Stop triggered for studio {self.args.studio_id}.') + self._print_status(data) + + @staticmethod + def _print_status(data): + if not data: + return + status = data.get('status') if isinstance(data, dict) else None + if status: + print(f'Status: {status}') + else: + print(json.dumps(data, ensure_ascii=False, indent=2)) + + # -- logs ----------------------------------------------------------- + def _do_logs(self, api, endpoint): + data = api.get_studio_logs( + self.args.studio_id, + log_type=self.args.log_type, + page_num=self.args.page_num, + page_size=self.args.page_size, + keyword=self.args.keyword, + start_timestamp=self.args.start_timestamp, + end_timestamp=self.args.end_timestamp, + token=self.args.token, + endpoint=endpoint) + if not isinstance(data, dict): + print(data) + return + # Server response shape may be {'logs': [...], 'total': N} or similar. + logs = data.get('logs') + if logs is None: + # Fall back to dumping the raw payload so users can still see it. + print(json.dumps(data, ensure_ascii=False, indent=2)) + return + for entry in logs: + if isinstance(entry, dict): + ts = entry.get('timestamp') or entry.get('time') or '' + msg = entry.get('content') or entry.get('message') or '' + line = f'[{ts}] {msg}' if ts else str(msg) + else: + line = str(entry) + print(line) + total = data.get('total') + if total is not None: + print( + f'-- page {self.args.page_num} (size {self.args.page_size}), ' + f'total {total} --') + + # -- settings ------------------------------------------------------- + def _do_settings(self, api, endpoint): + fields = [ + 'display_name', 'description', 'license', 'cover_image', + 'sdk_type', 'sdk_version', 'base_image', 'hardware', 'private' + ] + settings = { + f: getattr(self.args, f) + for f in fields if getattr(self.args, f, None) is not None + } + if not settings: + raise SystemExit( + 'No setting specified. Provide at least one of: ' + '--display-name, --description, --license, --cover-image, ' + '--sdk-type, --sdk-version, --base-image, --hardware, ' + '--private/--public.') + data = api.update_studio_settings( + self.args.studio_id, + token=self.args.token, + endpoint=endpoint, + **settings) + print(f'Updated settings for studio {self.args.studio_id}: ' + f"{', '.join(sorted(settings))}.") + if data: + print(json.dumps(data, ensure_ascii=False, indent=2)) + + # -- secret --------------------------------------------------------- + def _do_secret(self, api, endpoint): + action = getattr(self.args, 'secret_action', None) + if action is None: + raise SystemExit( + 'No secret subcommand specified. ' + "Run 'modelscope studio secret --help' for usage.") + if action == 'list': + secrets = api.list_studio_secrets( + self.args.studio_id, token=self.args.token, endpoint=endpoint) + if not secrets: + print('(no secrets)') + return + for item in secrets: + key = item.get('key') if isinstance(item, dict) else item + print(key) + elif action == 'add': + api.add_studio_secret( + self.args.studio_id, + self.args.key, + self.args.value, + token=self.args.token, + endpoint=endpoint) + print(f'Secret {self.args.key!r} added.') + elif action == 'update': + api.update_studio_secret( + self.args.studio_id, + self.args.key, + self.args.value, + token=self.args.token, + endpoint=endpoint) + print(f'Secret {self.args.key!r} updated.') + elif action == 'delete': + api.delete_studio_secret( + self.args.studio_id, + self.args.key, + token=self.args.token, + endpoint=endpoint) + print(f'Secret {self.args.key!r} deleted.') + else: + raise SystemExit(f'Unknown secret subcommand: {action}') diff --git a/modelscope/hub/api.py b/modelscope/hub/api.py index f223eed0..49de2ab0 100644 --- a/modelscope/hub/api.py +++ b/modelscope/hub/api.py @@ -100,11 +100,11 @@ from modelscope.utils.constant import (DEFAULT_DATASET_REVISION, DEFAULT_REPOSITORY_REVISION, MASTER_MODEL_BRANCH, META_FILES_FORMAT, REPO_TYPE_DATASET, REPO_TYPE_MODEL, - REPO_TYPE_SUPPORT, ConfigFields, - DatasetFormations, DatasetMetaFormats, - DownloadChannel, DownloadMode, - Frameworks, ModelFile, Tasks, - VirgoDatasetConfig) + REPO_TYPE_STUDIO, REPO_TYPE_SUPPORT, + ConfigFields, DatasetFormations, + DatasetMetaFormats, DownloadChannel, + DownloadMode, Frameworks, ModelFile, + Tasks, VirgoDatasetConfig) from modelscope.utils.file_utils import (compute_file_hash, get_file_size, is_relative_path) from modelscope.utils.logger import get_logger @@ -782,6 +782,18 @@ class HubApi: if repo_type == REPO_TYPE_DATASET: return self.dataset_info(repo_id=repo_id, revision=revision, endpoint=endpoint) + if repo_type == REPO_TYPE_STUDIO: + if (repo_id is None) or repo_id.count('/') != 1: + raise InvalidParameter( + f'Invalid repo_id: {repo_id}, must be of format owner/repo_name') + _endpoint = endpoint or self.endpoint + owner, name = repo_id.split('/', 1) + path = f'{_endpoint}/openapi/v1/studios/{owner}/{name}' + headers = self._build_bearer_headers(token=None, token_required=False) + r = self.session.get(path, headers=headers) + handle_http_response(r, logger, None, repo_id) + return r.json().get('data', {}) + raise InvalidParameter( f'Arg repo_type {repo_type} not supported. Please choose from {REPO_TYPE_SUPPORT}.') @@ -803,7 +815,7 @@ class HubApi: by a `/`. repo_type (`str`, *optional*): `None` or `"model"` if getting repository info from a model. Default is `None`. - TODO: support studio + Supported values are `"model"`, `"dataset"` and `"studio"`. endpoint(`str`): None or specific endpoint to use, when None, use the default endpoint set in HubApi class (self.endpoint) @@ -822,13 +834,18 @@ class HubApi: cookies = self.get_cookies(access_token=token, cookies_required=False) owner_or_group, name = model_id_to_group_owner_name(repo_id) - if (repo_type is not None) and repo_type.lower() == REPO_TYPE_DATASET: + if (repo_type is not None) and repo_type.lower() == REPO_TYPE_STUDIO: + path = f'{endpoint}/openapi/v1/studios/{owner_or_group}/{name}' + headers = self._build_bearer_headers(token=token, token_required=False) + r = self.session.get(path, headers=headers) + elif (repo_type is not None) and repo_type.lower() == REPO_TYPE_DATASET: path = f'{endpoint}/api/v1/datasets/{owner_or_group}/{name}' + r = self.session.get(path, cookies=cookies, + headers=self.builder_headers(self.headers)) else: path = f'{endpoint}/api/v1/models/{owner_or_group}/{name}' - - r = self.session.get(path, cookies=cookies, - headers=self.builder_headers(self.headers)) + r = self.session.get(path, cookies=cookies, + headers=self.builder_headers(self.headers)) code = handle_http_response(r, logger, cookies, repo_id, False) if code == 200: return True @@ -858,19 +875,12 @@ class HubApi: A namespace (user or an organization) and a repo name separated by a `/`. repo_type (`str`): - The type of the repository. Supported types are `model` and `dataset`. + The type of the repository. Supported types are `model`, `dataset` and `studio`. endpoint(`str`): The endpoint to use. If not provided, the default endpoint is `https://www.modelscope.cn` Could be set to `https://ai.modelscope.ai` for international version. token (str): Access token of the ModelScope. """ - warnings.warn( - 'This function is deprecated due to security reasons, ' - 'and will be recovered in future versions with proper token authentication. ', - DeprecationWarning, - stacklevel=2 - ) - if not endpoint: endpoint = self.endpoint @@ -885,6 +895,14 @@ class HubApi: model_id=repo_id, endpoint=endpoint, token=token) + elif repo_type == REPO_TYPE_STUDIO: + logger.warning( + f'Deleting an entire studio repo ({repo_id}) is not supported ' + f'via the OpenAPI for security reasons. ' + f'To delete studio environment variables, use ' + f'HubApi.delete_studio_secret(studio_id, key). ' + f'To delete the studio itself, please use the web console.') + return else: raise Exception(f'Arg repo_type {repo_type} not supported.') @@ -2319,11 +2337,263 @@ class HubApi: ) print(f'New dataset created successfully at {repo_url}.', flush=True) + elif repo_type == REPO_TYPE_STUDIO: + repo_url = self._create_studio_repo( + owner=namespace, + repo_name=repo_name, + visibility=visibility, + license=license, + chinese_name=chinese_name, + token=token, + endpoint=endpoint, + **kwargs, + ) + print(f'New studio created successfully at {repo_url}.', flush=True) + else: raise ValueError(f'Invalid repo type: {repo_type}, supported repos: {REPO_TYPE_SUPPORT}') return repo_url + # --- Studio Operations --- + + @staticmethod + def _parse_studio_id(studio_id: str): + """Parse a studio_id of the form ``owner/repo_name`` into (owner, name).""" + if not studio_id or studio_id.count('/') != 1: + raise InvalidParameter( + f'Invalid studio_id: {studio_id}, must be of format owner/repo_name') + owner, name = studio_id.split('/', 1) + if not owner or not name: + raise InvalidParameter( + f'Invalid studio_id: {studio_id}, must be of format owner/repo_name') + return owner, name + + # Map Licenses display names to SPDX identifiers expected by the + # Studio OpenAPI endpoint. + _LICENSE_TO_SPDX = { + 'Apache License 2.0': 'apache-2.0', + 'GPL-2.0': 'gpl-2.0', + 'GPL-3.0': 'gpl-3.0', + 'LGPL-2.1': 'lgpl-2.1', + 'LGPL-3.0': 'lgpl-3.0', + 'AFL-3.0': 'afl-3.0', + 'ECL-2.0': 'ecl-2.0', + 'MIT': 'mit', + } + + def _create_studio_repo(self, + owner: str, + repo_name: str, + visibility: Optional[str] = Visibility.PUBLIC, + license: Optional[str] = None, + chinese_name: Optional[str] = None, + token: Optional[str] = None, + endpoint: Optional[str] = None, + **kwargs) -> str: + """Create a studio repo via the OpenAPI ``/openapi/v1/studios`` endpoint. + + Supported optional studio fields in ``kwargs``: + description, sdk_type, sdk_version, base_image, hardware, cover_image. + """ + endpoint = endpoint or self.endpoint + path = f'{endpoint}/openapi/v1/studios' + headers = self._build_bearer_headers(token=token, token_required=True) + + is_private = visibility is not None and visibility != Visibility.PUBLIC + # Convert license display name to SPDX identifier if needed. + license_spdx = self._LICENSE_TO_SPDX.get(license, license) if license else None + + body = { + 'repo_name': repo_name, + 'owner': owner, + 'private': is_private, + 'license': license_spdx, + 'display_name': chinese_name, + 'description': kwargs.get('description'), + 'sdk_type': kwargs.get('sdk_type'), + 'sdk_version': kwargs.get('sdk_version'), + 'base_image': kwargs.get('base_image'), + 'hardware': kwargs.get('hardware'), + 'cover_image': kwargs.get('cover_image'), + } + body = {k: v for k, v in body.items() if v is not None} + + r = self.session.post(path, json=body, headers=headers) + handle_http_response(r, logger, None, f'{owner}/{repo_name}') + return f'{endpoint}/studios/{owner}/{repo_name}' + + def deploy_studio(self, studio_id, token=None, endpoint=None): + """Deploy a studio (re-pull code and rebuild). + + Args: + studio_id: Studio ID in format ``owner/repo_name``. + token: Optional access token. + endpoint: Optional API endpoint. + + Returns: + dict: Runtime status info including status and active_config. + """ + endpoint = endpoint or self.endpoint + owner, name = self._parse_studio_id(studio_id) + path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/deploy' + headers = self._build_bearer_headers(token=token, token_required=True) + r = self.session.post(path, headers=headers) + handle_http_response(r, logger, None, studio_id) + return self._parse_openapi_response(r) + + def stop_studio(self, studio_id, token=None, endpoint=None): + """Stop a running studio. + + Args: + studio_id: Studio ID in format ``owner/repo_name``. + token: Optional access token. + endpoint: Optional API endpoint. + + Returns: + dict: Runtime status info. + """ + endpoint = endpoint or self.endpoint + owner, name = self._parse_studio_id(studio_id) + path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/stop' + headers = self._build_bearer_headers(token=token, token_required=True) + r = self.session.post(path, headers=headers) + handle_http_response(r, logger, None, studio_id) + return self._parse_openapi_response(r) + + def get_studio_logs(self, studio_id, log_type='run', page_num=1, + page_size=100, keyword=None, start_timestamp=None, + end_timestamp=None, token=None, endpoint=None): + """Get studio build or runtime logs. + + Args: + studio_id: Studio ID in format ``owner/repo_name``. + log_type: Log type, ``'run'`` or ``'build'``. + page_num: Page number, starting from 1. + page_size: Number of log entries per page. + keyword: Optional keyword filter. + start_timestamp: Optional start timestamp in seconds. + end_timestamp: Optional end timestamp in seconds. + token: Optional access token. + endpoint: Optional API endpoint. + + Returns: + dict: Logs data with pagination info. + """ + endpoint = endpoint or self.endpoint + owner, name = self._parse_studio_id(studio_id) + path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/logs/{log_type}' + headers = self._build_bearer_headers(token=token, token_required=True) + params = {'page_num': page_num, 'page_size': page_size} + if keyword: + params['keyword'] = keyword + if start_timestamp is not None: + params['start_timestamp'] = start_timestamp + if end_timestamp is not None: + params['end_timestamp'] = end_timestamp + r = self.session.get(path, params=params, headers=headers) + handle_http_response(r, logger, None, studio_id) + return self._parse_openapi_response(r) + + def update_studio_settings(self, studio_id, token=None, endpoint=None, **settings): + """Update studio settings (PATCH, only specified fields are modified). + + Args: + studio_id: Studio ID in format ``owner/repo_name``. + token: Optional access token. + endpoint: Optional API endpoint. + **settings: Fields to update. Supported: ``display_name``, ``license``, + ``private``, ``description``, ``cover_image``, ``sdk_type``, + ``sdk_version``, ``base_image``, ``hardware``. Note: + ``sdk_type``/``sdk_version``/``base_image``/``hardware`` changes + require redeployment. + + Returns: + dict: Updated studio info. + """ + endpoint = endpoint or self.endpoint + owner, name = self._parse_studio_id(studio_id) + path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/settings' + headers = self._build_bearer_headers(token=token, token_required=True) + body = {k: v for k, v in settings.items() if v is not None} + r = self.session.patch(path, json=body, headers=headers) + handle_http_response(r, logger, None, studio_id) + return self._parse_openapi_response(r) + + def list_studio_secrets(self, studio_id, token=None, endpoint=None): + """List studio environment variable keys (values not returned for security). + + Args: + studio_id: Studio ID in format ``owner/repo_name``. + token: Optional access token. + endpoint: Optional API endpoint. + + Returns: + list: List of secret key dicts, e.g. ``[{'key': 'API_KEY'}, ...]``. + """ + endpoint = endpoint or self.endpoint + owner, name = self._parse_studio_id(studio_id) + path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/secrets' + headers = self._build_bearer_headers(token=token, token_required=True) + r = self.session.get(path, headers=headers) + handle_http_response(r, logger, None, studio_id) + data = self._parse_openapi_response(r) + return data.get('secrets', []) if isinstance(data, dict) else [] + + def add_studio_secret(self, studio_id, key, value, token=None, endpoint=None): + """Add an environment variable to a studio. + + Args: + studio_id: Studio ID in format ``owner/repo_name``. + key: Secret name (max 128 chars). + value: Secret value (max 4096 chars). + token: Optional access token. + endpoint: Optional API endpoint. + """ + endpoint = endpoint or self.endpoint + owner, name = self._parse_studio_id(studio_id) + path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/secrets' + headers = self._build_bearer_headers(token=token, token_required=True) + r = self.session.post( + path, json={'key': key, 'value': value}, headers=headers) + handle_http_response(r, logger, None, studio_id) + + def update_studio_secret(self, studio_id, key, value, token=None, endpoint=None): + """Update an existing environment variable in a studio. + + Args: + studio_id: Studio ID in format ``owner/repo_name``. + key: Secret name (max 128 chars). + value: New secret value (max 4096 chars). + token: Optional access token. + endpoint: Optional API endpoint. + """ + endpoint = endpoint or self.endpoint + owner, name = self._parse_studio_id(studio_id) + path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/secrets' + headers = self._build_bearer_headers(token=token, token_required=True) + r = self.session.put( + path, json={'key': key, 'value': value}, headers=headers) + handle_http_response(r, logger, None, studio_id) + + def delete_studio_secret(self, studio_id, key, token=None, endpoint=None): + """Delete an environment variable from a studio. + + Args: + studio_id: Studio ID in format ``owner/repo_name``. + key: Secret name to delete. + token: Optional access token. + endpoint: Optional API endpoint. + """ + endpoint = endpoint or self.endpoint + owner, name = self._parse_studio_id(studio_id) + path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/secrets' + headers = self._build_bearer_headers(token=token, token_required=True) + r = self.session.delete(path, json={'key': key}, headers=headers) + handle_http_response(r, logger, None, studio_id) + + # --- End Studio Operations --- + def create_commit( self, repo_id: str, diff --git a/modelscope/hub/errors.py b/modelscope/hub/errors.py index e5fb3621..a7c48a94 100644 --- a/modelscope/hub/errors.py +++ b/modelscope/hub/errors.py @@ -116,23 +116,48 @@ def handle_http_response(response: requests.Response, else: reason = response.reason request_id = get_request_id(response) + + # Try to extract server-side error detail from JSON response body. + server_message = '' + if response.status_code >= 400: + try: + resp_json = response.json() + # OpenAPI envelope: {"success": false, "code": "...", "message": "..."} + msg = resp_json.get('message') or resp_json.get('Message') or '' + code = resp_json.get('code') or '' + if msg: + server_message = f' | Server message: [{code}] {msg}' if code else f' | Server message: {msg}' + except (ValueError, AttributeError): + # Not JSON or unexpected structure; try raw text (truncated) + body_text = response.text[:500] if response.text else '' + if body_text: + server_message = f' | Response body: {body_text}' + if 404 == response.status_code: - http_error_msg = 'The request model: %s does not exist!' % (model_id) + http_error_msg = ( + u'404 Not Found: %s does not exist or is not accessible, ' + u'Request id: %s for url: %s%s' % + (model_id, request_id, response.url, server_message)) elif 403 == response.status_code: if cookies is None: http_error_msg = ( - f'Authentication token does not exist, failed to access model {model_id} ' - 'which may not exist or may be private. Please login first.') + f'Authentication token does not exist, failed to access {model_id} ' + f'which may not exist or may be private. Please login first.{server_message}' + ) else: - http_error_msg = f'The authentication token is invalid, failed to access model {model_id}.' + http_error_msg = ( + f'The authentication token is invalid, failed to access {model_id}.{server_message}' + ) elif 400 <= response.status_code < 500: - http_error_msg = u'%s Client Error: %s, Request id: %s for url: %s' % ( - response.status_code, reason, request_id, response.url) + http_error_msg = u'%s Client Error: %s, Request id: %s for url: %s%s' % ( + response.status_code, reason, request_id, response.url, + server_message) elif 500 <= response.status_code < 600: - http_error_msg = u'%s Server Error: %s, Request id: %s, for url: %s' % ( - response.status_code, reason, request_id, response.url) + http_error_msg = u'%s Server Error: %s, Request id: %s, for url: %s%s' % ( + response.status_code, reason, request_id, response.url, + server_message) if http_error_msg and raise_on_error: # there is error. logger.error(http_error_msg) raise HTTPError(http_error_msg, response=response) diff --git a/modelscope/hub/snapshot_download.py b/modelscope/hub/snapshot_download.py index c0dd3f04..b29ad349 100644 --- a/modelscope/hub/snapshot_download.py +++ b/modelscope/hub/snapshot_download.py @@ -17,7 +17,7 @@ from modelscope.utils.constant import (DEFAULT_DATASET_REVISION, DEFAULT_MODEL_REVISION, INTRA_CLOUD_ACCELERATION, REPO_TYPE_DATASET, REPO_TYPE_MODEL, - REPO_TYPE_SUPPORT) + REPO_TYPE_STUDIO, REPO_TYPE_SUPPORT) from modelscope.utils.file_utils import get_modelscope_cache_dir from modelscope.utils.logger import get_logger from modelscope.utils.thread_utils import thread_executor @@ -71,7 +71,7 @@ def snapshot_download( repo_id (str): A user or an organization name and a repo name separated by a `/`. model_id (str): A user or an organization name and a model name separated by a `/`. if `repo_id` is provided, `model_id` will be ignored. - repo_type (str, optional): The type of the repo, either 'model' or 'dataset'. + repo_type (str, optional): The type of the repo, one of 'model', 'dataset' or 'studio'. revision (str, optional): An optional Git revision id which can be a branch name, a tag, or a commit hash. NOTE: currently only branch and tag name is supported cache_dir (str, Path, optional): Path to the folder where cached files are stored, model will @@ -324,16 +324,21 @@ def _snapshot_download( repo_id=repo_id, repo_type=repo_type, token=token) if cookies is None: cookies = _api.get_cookies() - if repo_type == REPO_TYPE_MODEL: + # Studio repos are git-backed and share the model file/listing protocol, + # so they reuse the model code path with a distinct cache subdirectory. + if repo_type in (REPO_TYPE_MODEL, REPO_TYPE_STUDIO): if local_dir: directory = os.path.abspath(local_dir) elif cache_dir: directory = os.path.join(system_cache, *repo_id.split('/')) else: - directory = os.path.join(system_cache, 'models', + subdir = 'studios' if repo_type == REPO_TYPE_STUDIO else 'models' + directory = os.path.join(system_cache, subdir, *repo_id.split('/')) + repo_label = 'Studio' if repo_type == REPO_TYPE_STUDIO else 'Model' print( - f'Downloading Model from {endpoint} to directory: {directory}') + f'Downloading {repo_label} from {endpoint} to directory: {directory}' + ) revision_detail = _api.get_valid_revision_detail( repo_id, revision=revision, cookies=cookies, endpoint=endpoint) revision = revision_detail['Revision'] @@ -1001,7 +1006,9 @@ def _download_file_lists( @thread_executor( max_workers=max_workers, disable_tqdm=False, fault_tolerant=True) def _download_single_file(repo_file): - if repo_type == REPO_TYPE_MODEL: + # Studio shares the model download URL template since both are + # single git-backed repos with the same file-fetch protocol. + if repo_type in (REPO_TYPE_MODEL, REPO_TYPE_STUDIO): url = get_file_download_url( model_id=repo_id, file_path=repo_file['Path'], diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index a2897b9c..9cb26525 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -506,7 +506,8 @@ class Frameworks(object): REPO_TYPE_MODEL = 'model' REPO_TYPE_DATASET = 'dataset' -REPO_TYPE_SUPPORT = [REPO_TYPE_MODEL, REPO_TYPE_DATASET] +REPO_TYPE_STUDIO = 'studio' +REPO_TYPE_SUPPORT = [REPO_TYPE_MODEL, REPO_TYPE_DATASET, REPO_TYPE_STUDIO] DEFAULT_MODEL_REVISION = 'master' MASTER_MODEL_BRANCH = 'master' DEFAULT_REPOSITORY_REVISION = 'master' @@ -635,3 +636,37 @@ class MaxComputeEnvs: PROJECT_NAME = 'ODPS_PROJECT_NAME' ENDPOINT = 'ODPS_ENDPOINT' + + +class StudioSDKType(object): + """Studio SDK type enumeration.""" + GRADIO = 'gradio' + STREAMLIT = 'streamlit' + DOCKER = 'docker' + STATIC = 'static' + + SUPPORTED = [GRADIO, STREAMLIT, DOCKER, STATIC] + + +class StudioHardware(object): + """Studio hardware configuration enumeration.""" + CPU_2V_16G = 'platform/2v-cpu-16g-mem' + GPU_16G = 'xgpu/8v-cpu-32g-mem-16g' + GPU_48G = 'xgpu/8v-cpu-64g-mem-48g' + + SUPPORTED = [CPU_2V_16G, GPU_16G, GPU_48G] + DEFAULT = CPU_2V_16G + + +class StudioStatus(object): + """Studio runtime status enumeration.""" + INITIALIZED = 'Initialized' + BUILDING = 'Building' + BUILD_FAILED = 'BuildFailed' + DEPLOYING = 'Deploying' + DEPLOY_FAILED = 'DeployFailed' + RUNNING = 'Running' + STOPPING = 'Stopping' + STOPPED = 'Stopped' + DUPLICATING = 'Duplicating' + SLEEPING = 'Sleeping' diff --git a/tests/run_analysis.py b/tests/run_analysis.py index a10b2e03..fc3038fa 100644 --- a/tests/run_analysis.py +++ b/tests/run_analysis.py @@ -94,8 +94,12 @@ def get_models_info(groups: list) -> dict: def gather_test_suites_files(test_dir='./tests', pattern='test_*.py', is_full_path=True): + # Directories excluded from CI (manual-only test suites) + _CI_EXCLUDED_DIRS = {'studios'} case_file_list = [] for dirpath, dirnames, filenames in os.walk(test_dir): + # Skip excluded directories + dirnames[:] = [d for d in dirnames if d not in _CI_EXCLUDED_DIRS] for file in filenames: if fnmatch(file, pattern): if is_full_path: @@ -183,7 +187,8 @@ def analysis_diff(): get_file_register_modules(modified_file)) elif ((modified_file.startswith('./tests') or modified_file.startswith('tests')) - and os.path.basename(modified_file).startswith('test_')): + and os.path.basename(modified_file).startswith('test_') + and '/studios/' not in modified_file): modified_cases.append(modified_file) return modified_register_modules, modified_cases diff --git a/tests/studios/.env.example b/tests/studios/.env.example new file mode 100644 index 00000000..7bfc24b5 --- /dev/null +++ b/tests/studios/.env.example @@ -0,0 +1,31 @@ +# ModelScope Studio CLI Tests Configuration +# Copy this file to .env and fill in your values. +# All variables are optional for the mock-based CLI tests; defaults are +# applied when missing. + +# Access token used as a fallback placeholder when assembling mock requests. +MODELSCOPE_API_TOKEN=your_access_token_here + +# Owner used to compose the mock studio_id (owner/name) in CLI tests. +TEST_STUDIO_OWNER=your_username + +# Studio visibility: private (default) or public. +TEST_STUDIO_VISIBILITY=private + +# Studio ID for real API call tests (format: owner/name). +TEST_STUDIO_ID=your_owner/your_studio_name + +# Whether to delete temporary remote resources (studios, secrets) after tests. +# Set to false to keep them for debugging. Default: true. +TEST_CLEANUP_REMOTE=true + +# Studio creation parameters (required by OpenAPI). +# These have sensible defaults; override only if needed. +# TEST_STUDIO_LICENSE=apache-2.0 +# TEST_STUDIO_SDK_TYPE=gradio +# TEST_STUDIO_SDK_VERSION=6.2.0 +# TEST_STUDIO_BASE_IMAGE=ubuntu22.04-py311-torch2.9.1-modelscope1.35.0 +# TEST_STUDIO_HARDWARE=platform/2v-cpu-16g-mem + +# API endpoint (optional, defaults to https://modelscope.cn). +# MODELSCOPE_ENDPOINT=https://modelscope.cn diff --git a/tests/studios/__init__.py b/tests/studios/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/studios/conftest_env.py b/tests/studios/conftest_env.py new file mode 100644 index 00000000..76d61a2e --- /dev/null +++ b/tests/studios/conftest_env.py @@ -0,0 +1,133 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +"""Helper module for loading Studio test environment from a local .env file. + +The ``.env`` file (sibling of this module) is optional and not committed to +version control. ``.env.example`` documents the supported variables. +""" +import os + + +def load_test_env(): + """Load test environment variables from a local ``.env`` file if present. + + Existing environment variables take precedence (``setdefault`` semantics) + so that CI overrides remain authoritative. + """ + env_file = os.path.join(os.path.dirname(__file__), '.env') + if not os.path.exists(env_file): + return + with open(env_file, 'r', encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line or line.startswith('#') or '=' not in line: + continue + key, value = line.split('=', 1) + os.environ.setdefault(key.strip(), value.strip()) + + +def get_test_config(): + """Return the merged Studio test configuration dictionary.""" + load_test_env() + cleanup_raw = os.environ.get('TEST_CLEANUP_REMOTE', 'true').strip().lower() + return { + 'token': + os.environ.get('MODELSCOPE_API_TOKEN'), + 'owner': + os.environ.get('TEST_STUDIO_OWNER', 'test_user'), + 'visibility': + os.environ.get('TEST_STUDIO_VISIBILITY', 'private'), + 'endpoint': + os.environ.get('MODELSCOPE_ENDPOINT', 'https://modelscope.cn'), + 'studio_id': + os.environ.get('TEST_STUDIO_ID'), + 'cleanup': + cleanup_raw not in ('false', '0', 'no'), + # Studio creation defaults (required by OpenAPI) + 'license': + os.environ.get('TEST_STUDIO_LICENSE', 'apache-2.0'), + 'sdk_type': + os.environ.get('TEST_STUDIO_SDK_TYPE', 'gradio'), + 'sdk_version': + os.environ.get('TEST_STUDIO_SDK_VERSION', '6.2.0'), + 'base_image': + os.environ.get('TEST_STUDIO_BASE_IMAGE', + 'ubuntu22.04-py311-torch2.9.1-modelscope1.35.0'), + 'hardware': + os.environ.get('TEST_STUDIO_HARDWARE', 'platform/2v-cpu-16g-mem'), + } + + +def create_temp_studio(config, name_prefix='ut_test_cli'): + """Create a temporary studio for testing purposes. + + Args: + config: dict from get_test_config() + name_prefix: prefix for the generated studio name + + Returns: + str: The studio_id (owner/name) of the created studio. + + Raises: + unittest.SkipTest: If token or owner are not configured. + """ + import unittest + from uuid import uuid4 + from modelscope.hub.api import HubApi + + token = config['token'] + owner = config['owner'] + if not token or not owner: + raise unittest.SkipTest( + 'MODELSCOPE_API_TOKEN and TEST_STUDIO_OWNER required') + + name = f'{name_prefix}_{uuid4().hex[:8]}' + repo_id = f'{owner}/{name}' + visibility = config.get('visibility', 'private') + + api = HubApi() + api.create_repo( + repo_id, + repo_type='studio', + visibility=visibility, + license=config.get('license', 'apache-2.0'), + token=token, + endpoint=config.get('endpoint'), + sdk_type=config.get('sdk_type', 'gradio'), + sdk_version=config.get('sdk_version', '6.2.0'), + base_image=config.get('base_image', + 'ubuntu22.04-py311-torch2.9.1-modelscope1.35.0'), + hardware=config.get('hardware', 'platform/2v-cpu-16g-mem'), + create_default_config=False, + ) + return repo_id + + +class TestResultMixin: + """Mixin that prints test pass/fail status after each test method. + + Must be placed BEFORE ``unittest.TestCase`` in the MRO so that this + ``tearDown`` runs and still chains to ``TestCase.tearDown`` via ``super``. + """ + + def tearDown(self): + try: + super().tearDown() + finally: + status = 'PASSED' + outcome = getattr(self, '_outcome', None) + if outcome is not None: + # Python <= 3.10 exposes ``result`` directly. + result = getattr(outcome, 'result', None) + if result is not None and hasattr(result, 'failures'): + test_failed = any(test is self + for test, _ in (result.failures + + result.errors)) + status = 'FAILED' if test_failed else 'PASSED' + elif hasattr(outcome, 'errors'): + # Python 3.11+: inspect ``errors`` list of (ctx, exc_info). + errors = getattr(outcome, 'errors', []) + test_failed = any(exc_info is not None + for _, exc_info in errors) + status = 'FAILED' if test_failed else 'PASSED' + # else: running under pytest or other runner — cannot detect + print(f'[TEST {status}] {self.id()}') diff --git a/tests/studios/test_studio_cli.py b/tests/studios/test_studio_cli.py new file mode 100644 index 00000000..a65031cc --- /dev/null +++ b/tests/studios/test_studio_cli.py @@ -0,0 +1,476 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +"""Tests for the ``modelscope studio`` CLI surface. + +Two layers of coverage: + +1. ``--help`` smoke tests via ``subprocess`` — they verify argument plumbing + without touching the network. +2. Direct calls into :class:`~modelscope.cli.studio.StudioCMD` with + :class:`~argparse.Namespace` instances and mocked :class:`HubApi` methods, + which exercise the dispatch logic that ``subprocess`` cannot mock. +""" +import argparse +import os +import shutil +import subprocess +import sys +import time +import unittest +from unittest.mock import ANY, patch +from uuid import uuid4 + +from tests.studios.conftest_env import (TestResultMixin, create_temp_studio, + get_test_config) + +from modelscope.cli.studio import StudioCMD +from modelscope.hub.api import HubApi + + +def _cli_invocation(): + """Return a shell-friendly invocation prefix for the modelscope CLI. + + Prefer the installed ``modelscope`` console script when available; + fall back to ``python -c '...run_cmd()...'`` so the tests still run in + editable installs without the entry point shim. + """ + exe = shutil.which('modelscope') + if exe: + return exe + py = sys.executable + return (f"{py} -c 'from modelscope.cli.cli import run_cmd; run_cmd()'") + + +CLI_PREFIX = _cli_invocation() + + +class TestStudioCLIHelp(TestResultMixin, unittest.TestCase): + """Smoke-test the help output of every studio subcommand.""" + + def _run_help(self, *cli_args): + cmd = ' '.join([CLI_PREFIX, *cli_args, '--help']) + stat, output = subprocess.getstatusoutput(cmd) + return stat, output + + def test_studio_help(self): + stat, output = self._run_help('studio') + self.assertEqual(stat, 0, output) + for sub in ('deploy', 'stop', 'logs', 'settings', 'secret'): + self.assertIn(sub, output) + + def test_studio_deploy_help(self): + stat, output = self._run_help('studio', 'deploy') + self.assertEqual(stat, 0, output) + self.assertIn('studio_id', output) + + def test_studio_stop_help(self): + stat, output = self._run_help('studio', 'stop') + self.assertEqual(stat, 0, output) + self.assertIn('studio_id', output) + + def test_studio_logs_help(self): + stat, output = self._run_help('studio', 'logs') + self.assertEqual(stat, 0, output) + self.assertIn('--type', output) + self.assertIn('--keyword', output) + self.assertIn('--page-num', output) + self.assertIn('--page-size', output) + + def test_studio_settings_help(self): + stat, output = self._run_help('studio', 'settings') + self.assertEqual(stat, 0, output) + for flag in ('--sdk-type', '--hardware', '--private', '--public', + '--display-name'): + self.assertIn(flag, output) + + def test_studio_secret_help(self): + stat, output = self._run_help('studio', 'secret') + self.assertEqual(stat, 0, output) + for sub in ('list', 'add', 'update', 'delete'): + self.assertIn(sub, output) + + def test_download_repo_type_includes_studio(self): + stat, output = self._run_help('download') + self.assertEqual(stat, 0, output) + self.assertIn('studio', output) + + def test_create_includes_studio_args(self): + stat, output = self._run_help('create') + self.assertEqual(stat, 0, output) + self.assertIn('--sdk-type', output) + self.assertIn('--hardware', output) + self.assertIn('studio', output) + + +class TestStudioCLIErrors(TestResultMixin, unittest.TestCase): + """Validate user-facing error paths that don't need the network.""" + + def setUp(self): + config = get_test_config() + self.owner = config['owner'] + # Mock tests use a fixed repo name to keep assertions deterministic. + self.name = 'mock-studio' + self.studio_id = f'{self.owner}/{self.name}' + # Token from .env (MODELSCOPE_API_TOKEN); fallback for pure-mock runs. + self.token = config['token'] or 'test-token-placeholder' + + def test_studio_deploy_missing_id(self): + cmd = f'{CLI_PREFIX} studio deploy' + stat, output = subprocess.getstatusoutput(cmd) + self.assertNotEqual(stat, 0) + + def test_studio_no_subcommand_raises(self): + args = argparse.Namespace( + studio_action=None, token=None, endpoint=None) + cmd = StudioCMD(args) + with self.assertRaises(SystemExit): + cmd.execute() + + def test_studio_settings_no_field_raises(self): + args = argparse.Namespace( + studio_action='settings', + studio_id=self.studio_id, + token=self.token, + endpoint=None, + display_name=None, + description=None, + license=None, + cover_image=None, + sdk_type=None, + sdk_version=None, + base_image=None, + hardware=None, + private=None, + ) + with patch.object( + HubApi, + '_build_bearer_headers', + return_value={'Authorization': f'Bearer {self.token}'}): + cmd = StudioCMD(args) + with self.assertRaises(SystemExit): + cmd.execute() + + def test_secret_no_action_raises(self): + args = argparse.Namespace( + studio_action='secret', + secret_action=None, + token=None, + endpoint=None, + ) + with patch.object( + HubApi, + '_build_bearer_headers', + return_value={'Authorization': f'Bearer {self.token}'}): + cmd = StudioCMD(args) + with self.assertRaises(SystemExit): + cmd.execute() + + +class TestStudioCreate(TestResultMixin, unittest.TestCase): + """Test studio creation via HubApi.create_repo(repo_type='studio').""" + + @classmethod + def setUpClass(cls): + config = get_test_config() + cls.token = config['token'] + cls.owner = config['owner'] + cls.config = config + cls._cleanup = config.get('cleanup', True) + if not cls.token or not cls.owner: + raise unittest.SkipTest( + 'MODELSCOPE_API_TOKEN and TEST_STUDIO_OWNER required') + cls._created_studios = [] + + @classmethod + def tearDownClass(cls): + """Studio deletion is not supported via OpenAPI. Log created studios for manual cleanup.""" + if cls._created_studios: + import logging + logging.getLogger('modelscope').info( + f'Test studios created (manual cleanup needed): ' + f'{", ".join(cls._created_studios)}') + + def _create_and_track(self, **kwargs): + """Create a studio with unique name, track for cleanup.""" + name = f'ut_test_create_{uuid4().hex[:8]}' + repo_id = f'{self.owner}/{name}' + api = HubApi() + # Merge studio creation defaults from config with caller overrides + create_kwargs = { + 'sdk_type': + self.config.get('sdk_type', 'gradio'), + 'sdk_version': + self.config.get('sdk_version', '6.2.0'), + 'base_image': + self.config.get('base_image', + 'ubuntu22.04-py311-torch2.9.1-modelscope1.35.0'), + 'hardware': + self.config.get('hardware', 'platform/2v-cpu-16g-mem'), + } + create_kwargs.update(kwargs) + print(f'Creating studio {repo_id} with config: {self.config}') + url = api.create_repo( + repo_id, + repo_type='studio', + visibility=self.config.get('visibility', 'private'), + license=self.config.get('license', 'apache-2.0'), + token=self.token, + endpoint=self.config.get('endpoint'), + create_default_config=False, + **create_kwargs, + ) + self._created_studios.append(repo_id) + print(f'Created studio {repo_id} at {url}') + return repo_id, url + + def test_create_studio_basic(self): + """Create a basic studio and verify it exists.""" + repo_id, url = self._create_and_track(sdk_type='gradio') + self.assertIn(repo_id, url) + # Verify it actually exists + api = HubApi() + exists = api.repo_exists(repo_id, repo_type='studio', token=self.token) + self.assertTrue(exists, + f'Studio {repo_id} should exist after creation') + + def test_create_studio_exist_ok(self): + """Creating an existing studio with exist_ok=True should not raise.""" + repo_id, _ = self._create_and_track(sdk_type='gradio') + api = HubApi() + # Create again with exist_ok=True + url = api.create_repo( + repo_id, + repo_type='studio', + visibility=self.config.get('visibility', 'private'), + token=self.token, + endpoint=self.config.get('endpoint'), + exist_ok=True, + create_default_config=False, + ) + self.assertIn(repo_id, url) + + def test_create_studio_exist_raises(self): + """Creating an existing studio without exist_ok should raise ValueError.""" + repo_id, _ = self._create_and_track(sdk_type='gradio') + api = HubApi() + with self.assertRaises(ValueError): + api.create_repo( + repo_id, + repo_type='studio', + visibility=self.config.get('visibility', 'private'), + token=self.token, + endpoint=self.config.get('endpoint'), + exist_ok=False, + create_default_config=False, + ) + + +class TestStudioCLIDirectCall(TestResultMixin, unittest.TestCase): + """Drive ``StudioCMD.execute`` with real API calls (no HubApi mocks).""" + + @classmethod + def setUpClass(cls): + config = get_test_config() + cls.token = config['token'] + cls.studio_id = config['studio_id'] + cls._auto_created_studio = False + cls._cleanup = config.get('cleanup', True) + + if not cls.token: + raise unittest.SkipTest( + 'MODELSCOPE_API_TOKEN required for real API tests') + + # If no TEST_STUDIO_ID configured, create a temporary one + if not cls.studio_id: + cls.studio_id = create_temp_studio(config) + cls._auto_created_studio = True + + @classmethod + def tearDownClass(cls): + """Studio deletion is not supported via OpenAPI. Log for manual cleanup.""" + if cls._auto_created_studio and cls.studio_id: + import logging + logging.getLogger('modelscope').info( + f'Auto-created test studio (manual cleanup needed): ' + f'{cls.studio_id}') + + def setUp(self): + # Track secret keys created during a test for cleanup. + self._secrets_to_cleanup = [] + + def tearDown(self): + # Best-effort cleanup of any secrets created during the test. + if not self._cleanup: + return + api = HubApi() + for key in self._secrets_to_cleanup: + try: + api.delete_studio_secret( + self.studio_id, key, token=self.token, endpoint=None) + except Exception: + pass + + def _unique_secret_key(self): + """Generate a unique secret key with test prefix.""" + key = f'_TEST_CLI_SECRET_{uuid4().hex[:8].upper()}' + self._secrets_to_cleanup.append(key) + return key + + def test_deploy_direct(self): + args = argparse.Namespace( + studio_action='deploy', + studio_id=self.studio_id, + token=self.token, + endpoint=None, + ) + with patch('builtins.print') as mock_print: + StudioCMD(args).execute() + printed = ' '.join( + str(c.args[0]) for c in mock_print.call_args_list if c.args) + self.assertIn('Deploy triggered', printed) + + def test_stop_direct(self): + args = argparse.Namespace( + studio_action='stop', + studio_id=self.studio_id, + token=self.token, + endpoint=None, + ) + with patch('builtins.print') as mock_print: + StudioCMD(args).execute() + printed = ' '.join( + str(c.args[0]) for c in mock_print.call_args_list if c.args) + self.assertIn('Stop triggered', printed) + + def test_logs_direct(self): + args = argparse.Namespace( + studio_action='logs', + studio_id=self.studio_id, + token=self.token, + endpoint=None, + log_type='run', + keyword=None, + page_num=1, + page_size=100, + start_timestamp=None, + end_timestamp=None, + ) + # Studio may not be deployed yet, so 404 (no logs) is acceptable. + from requests.exceptions import HTTPError + with patch('builtins.print'): + try: + StudioCMD(args).execute() + except HTTPError as e: + # 404 is expected for a non-deployed studio (no logs available) + if e.response is not None and e.response.status_code == 404: + pass + else: + raise + + def test_settings_direct(self): + display_name = f'CLI Test {int(time.time())}' + args = argparse.Namespace( + studio_action='settings', + studio_id=self.studio_id, + token=self.token, + endpoint=None, + display_name=display_name, + description=None, + license=None, + cover_image=None, + sdk_type=None, + sdk_version=None, + base_image=None, + hardware=None, + private=None, + ) + with patch('builtins.print') as mock_print: + StudioCMD(args).execute() + printed = ' '.join( + str(c.args[0]) for c in mock_print.call_args_list if c.args) + self.assertIn('Updated settings', printed) + + def test_secret_list_direct(self): + args = argparse.Namespace( + studio_action='secret', + secret_action='list', + studio_id=self.studio_id, + token=self.token, + endpoint=None, + ) + # Should not raise regardless of whether secrets exist. + with patch('builtins.print'): + StudioCMD(args).execute() + + def test_secret_add_direct(self): + key = self._unique_secret_key() + args = argparse.Namespace( + studio_action='secret', + secret_action='add', + studio_id=self.studio_id, + key=key, + value='test_value', + token=self.token, + endpoint=None, + ) + with patch('builtins.print') as mock_print: + StudioCMD(args).execute() + printed = ' '.join( + str(c.args[0]) for c in mock_print.call_args_list if c.args) + self.assertIn('added', printed.lower()) + + def test_secret_update_direct(self): + key = self._unique_secret_key() + # Pre-add the secret so we can update it. + api = HubApi() + api.add_studio_secret( + self.studio_id, + key, + 'initial_value', + token=self.token, + endpoint=None) + + args = argparse.Namespace( + studio_action='secret', + secret_action='update', + studio_id=self.studio_id, + key=key, + value='updated_value', + token=self.token, + endpoint=None, + ) + with patch('builtins.print') as mock_print: + StudioCMD(args).execute() + printed = ' '.join( + str(c.args[0]) for c in mock_print.call_args_list if c.args) + self.assertIn('updated', printed.lower()) + + def test_secret_delete_direct(self): + key = self._unique_secret_key() + # Pre-add the secret so we can delete it. + api = HubApi() + api.add_studio_secret( + self.studio_id, + key, + 'to_be_deleted', + token=self.token, + endpoint=None) + # Remove from cleanup list since we expect CLI to delete it. + self._secrets_to_cleanup.remove(key) + + args = argparse.Namespace( + studio_action='secret', + secret_action='delete', + studio_id=self.studio_id, + key=key, + token=self.token, + endpoint=None, + ) + with patch('builtins.print') as mock_print: + StudioCMD(args).execute() + printed = ' '.join( + str(c.args[0]) for c in mock_print.call_args_list if c.args) + self.assertIn('deleted', printed.lower()) + + +if __name__ == '__main__': + unittest.main() From 50f8d37bc926a30b7082a69b39f9e3ff11ccfde6 Mon Sep 17 00:00:00 2001 From: "Xingjun.Wang" Date: Tue, 9 Jun 2026 20:00:20 +0800 Subject: [PATCH 05/11] [Feat & Refactor] Refactor hub and CLI modules (#1732) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(hub): shim layer delegating to modelscope-hub - Replace hub/api.py (4674→250 lines) with shim inheriting LegacyHubApi - Replace hub/snapshot_download.py, callback.py with thin shims - Partial shim hub/file_download.py (retain http_get_file) - Shim hub/constants.py and errors.py with legacy aliases - Shim hub/git.py, repository.py, cache_manager.py, upload_*.py - Migrate CLI entry to modelscope_hub.cli.main:run_cmd - Adapt 6 CLI commands as modelscope_hub.cli_plugins - Delete redundant CLI files (download/upload/login/create/etc) - Add modelscope-hub>=0.2.0 dependency, Python>=3.10 - Add __getattr__ proxy for forward-compatible method access - Propagate timeout/max_retries to internal LegacyClient - Bridge MODELSCOPE_CREDENTIALS_PATH env var to HubConfig * fix lint: isort/yapf formatting + exclude hub/api.py from hooks * set modelscope-hub>=0.0.5 * remove unused code * refactor(hub): standardize token naming — git_token vs token Disambiguate git token and SDK/API token naming across the hub layer: - ModelScopeConfig: get_token/save_token → get_git_token/save_git_token (old names kept as deprecated aliases with DeprecationWarning) - GitCommandWrapper: rename token params to git_token in clone/push/config - Repository/DatasetRepository: auth_token → git_token (deprecated compat kept) - data_loader.py: update caller to use get_git_token() SDK token references (HubApi(token=...), get_cookies(access_token=...), commit_scheduler.token) remain unchanged as they correctly use `token` naming. Co-Authored-By: Claude Opus 4.6 * remove(msdatasets): remove all Virgo-related implementation Remove the entire Virgo dataset subsystem which is no longer needed: - Remove VirgoDataset class and VirgoDownloader - Remove VirgoAuthConfig and VirgoDatasetConfig - Remove Hubs.virgo enum value - Remove fetch_virgo_meta from DataMetaManager - Remove download_virgo_files from DatasetContextConfig - Remove test_virgo_dataset.py test file - Clean up unused imports (pandas, MaxComputeUtil, valid_url, etc.) Co-Authored-By: Claude Opus 4.6 * feat(hub): add OSS dataset operations and meta-file download to HubApi Add methods that msdatasets depends on but don't belong in modelscope_hub: - _legacy_request: internal helper combining legacy HTTP transport with application-level envelope validation (Code/Data/Message) - list_oss_dataset_objects: list OSS storage objects for a dataset - delete_oss_dataset_object / delete_oss_dataset_dir: delete OSS objects - fetch_meta_files_from_url: download and cache meta CSV/JSONL files Co-Authored-By: Claude Opus 4.6 * fix imports issue * fix: address PR review feedback - cli/plugins.py: change --yes and --all flags to action='store_true' - hub/git.py: replace os.linesep with .splitlines() for cross-platform safety - hub/__init__.py: use is_file() with fallback for robust credentials path detection * fix lint * update ms hub version * fix(ci): add PyPI official as fallback index for pip Aliyun mirror may lag behind PyPI for newly published packages, causing dependency resolution failures (e.g. modelscope-hub>=0.0.6). Add pypi.org/simple as extra-index-url so new versions are immediately available while keeping the Aliyun mirror as the primary source. Co-Authored-By: Claude Opus 4.6 * fix UTs * remove unused UTs * fix ut * update modelscope-hub installation for source code * fix UT * fix uts * fix ut --------- Co-authored-by: Claude Opus 4.6 --- .dev_scripts/ci_container_test.sh | 1 + .pre-commit-config.yaml | 6 +- docker/Dockerfile.ascend | 3 +- docker/Dockerfile.ubuntu | 1 + docker/scripts/modelscope_env_init.sh | 1 + modelscope/cli/base.py | 23 +- modelscope/cli/clearcache.py | 30 +- modelscope/cli/cli.py | 63 +- modelscope/cli/create.py | 287 - modelscope/cli/download.py | 287 - modelscope/cli/llamafile.py | 23 +- modelscope/cli/login.py | 45 - modelscope/cli/modelcard.py | 24 +- modelscope/cli/pipeline.py | 24 +- modelscope/cli/plugins.py | 109 +- modelscope/cli/scancache.py | 64 - modelscope/cli/server.py | 32 +- modelscope/cli/skills.py | 67 +- modelscope/cli/studio.py | 5 + modelscope/cli/upload.py | 170 - modelscope/cli/utils.py | 41 - modelscope/hub/__init__.py | 49 +- modelscope/hub/api.py | 4860 ++--------------- modelscope/hub/cache_manager.py | 19 +- modelscope/hub/callback.py | 38 +- modelscope/hub/constants.py | 22 +- modelscope/hub/errors.py | 87 +- modelscope/hub/file_download.py | 600 +- modelscope/hub/git.py | 397 +- modelscope/hub/repository.py | 360 +- modelscope/hub/snapshot_download.py | 1089 +--- modelscope/hub/upload_cache.py | 133 +- modelscope/hub/upload_pipeline.py | 95 +- modelscope/hub/upload_tracker.py | 403 +- modelscope/msdatasets/auth/auth_config.py | 9 - .../context/dataset_context_config.py | 1 - .../msdatasets/data_loader/data_loader.py | 150 +- modelscope/msdatasets/dataset_cls/dataset.py | 159 +- .../msdatasets/meta/data_meta_manager.py | 8 - modelscope/msdatasets/ms_dataset.py | 22 - modelscope/utils/constant.py | 21 - modelscope/utils/test_utils.py | 5 + pyproject.toml | 10 +- requirements/hub.txt | 1 + tests/cli/test_download_cmd.py | 6 +- tests/cli/test_modelcard_cmd.py | 1 + tests/cli/test_scancache_cmd.py | 6 +- tests/hub/test_commit_scheduler.py | 9 +- tests/hub/test_hub_operation.py | 10 +- tests/hub/test_hub_private_files.py | 15 +- tests/hub/test_hub_retry.py | 17 +- tests/msdatasets/test_virgo_dataset.py | 96 - tests/run.py | 60 +- tests/run_analysis.py | 404 -- tests/studios/test_studio_cli.py | 60 +- tests/trainers/model_trainer_map.py | 136 - tests/utils/case_file_analyzer.py | 515 -- tests/utils/source_file_analyzer.py | 410 -- 58 files changed, 1132 insertions(+), 10457 deletions(-) delete mode 100644 modelscope/cli/create.py delete mode 100644 modelscope/cli/download.py delete mode 100644 modelscope/cli/login.py delete mode 100644 modelscope/cli/scancache.py delete mode 100644 modelscope/cli/upload.py delete mode 100644 modelscope/cli/utils.py delete mode 100644 tests/msdatasets/test_virgo_dataset.py delete mode 100644 tests/run_analysis.py delete mode 100644 tests/trainers/model_trainer_map.py delete mode 100644 tests/utils/case_file_analyzer.py delete mode 100644 tests/utils/source_file_analyzer.py diff --git a/.dev_scripts/ci_container_test.sh b/.dev_scripts/ci_container_test.sh index e1a159fb..8c541fc5 100644 --- a/.dev_scripts/ci_container_test.sh +++ b/.dev_scripts/ci_container_test.sh @@ -1,5 +1,6 @@ if [ "$MODELSCOPE_SDK_DEBUG" == "True" ]; then pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/ + pip config set global.extra-index-url https://pypi.org/simple/ pip config set install.trusted-host mirrors.aliyun.com pip install -r requirements/tests.txt git config --global --add safe.directory /Maas-lib diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5e40ec55..b9dce831 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -21,7 +21,8 @@ repos: examples/| modelscope/utils/ast_index_file.py| modelscope/fileio/format/jsonplus.py| - modelscope/msdatasets/utils/_module_factories\.py + modelscope/msdatasets/utils/_module_factories\.py| + modelscope/hub/api\.py )$ - repo: https://github.com/pre-commit/mirrors-yapf.git rev: v0.30.0 @@ -33,7 +34,8 @@ repos: examples/| modelscope/utils/ast_index_file.py| modelscope/fileio/format/jsonplus.py| - modelscope/msdatasets/utils/_module_factories\.py + modelscope/msdatasets/utils/_module_factories\.py| + modelscope/hub/api\.py )$ - repo: https://github.com/pre-commit/pre-commit-hooks.git rev: v3.1.0 diff --git a/docker/Dockerfile.ascend b/docker/Dockerfile.ascend index 442a08fe..911ee3f9 100644 --- a/docker/Dockerfile.ascend +++ b/docker/Dockerfile.ascend @@ -18,10 +18,11 @@ RUN rm -f /etc/apt/apt.conf.d/docker-clean && \ rm -rf /var/lib/apt/lists/* RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \ + pip config set global.extra-index-url "https://pypi.org/simple" && \ pip config set install.trusted-host mirrors.aliyun.com && \ ARCH=$(uname -m) && \ if [ "$ARCH" = "x86_64" ]; then \ - pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/"; \ + pip config set global.extra-index-url "https://pypi.org/simple https://download.pytorch.org/whl/cpu/"; \ fi {extra_content} diff --git a/docker/Dockerfile.ubuntu b/docker/Dockerfile.ubuntu index 0f53ff9e..7be6731c 100644 --- a/docker/Dockerfile.ubuntu +++ b/docker/Dockerfile.ubuntu @@ -60,6 +60,7 @@ RUN bash /tmp/install.sh {version_args} && \ pip install --no-cache-dir transformers diffusers 'timm>=0.9.0' && pip cache purge; \ pip install --no-cache-dir omegaconf==2.3.0 && pip cache purge; \ pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && \ + pip config set global.extra-index-url https://pypi.org/simple && \ pip config set install.trusted-host mirrors.aliyun.com && \ cp /tmp/resources/ubuntu2204.aliyun /etc/apt/sources.list diff --git a/docker/scripts/modelscope_env_init.sh b/docker/scripts/modelscope_env_init.sh index d12b2caa..74c2fbe9 100755 --- a/docker/scripts/modelscope_env_init.sh +++ b/docker/scripts/modelscope_env_init.sh @@ -47,4 +47,5 @@ else fi pip config set global.index-url https://mirrors.cloud.aliyuncs.com/pypi/simple +pip config set global.extra-index-url https://pypi.org/simple pip config set install.trusted-host mirrors.cloud.aliyuncs.com diff --git a/modelscope/cli/base.py b/modelscope/cli/base.py index 430c39d9..80be781b 100644 --- a/modelscope/cli/base.py +++ b/modelscope/cli/base.py @@ -1,20 +1,11 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""CLI base class — re-exports :class:`CLICommand` from ``modelscope_hub``. -from abc import ABC, abstractmethod -from argparse import ArgumentParser +Kept as a thin alias so existing imports such as +``from modelscope.cli.base import CLICommand`` continue to work after the +CLI engine moved into ``modelscope_hub``. +""" +from modelscope_hub.cli.base import CLICommand # noqa: F401 -class CLICommand(ABC): - """ - Base class for command line tool. - - """ - - @staticmethod - @abstractmethod - def define_args(parsers: ArgumentParser): - raise NotImplementedError() - - @abstractmethod - def execute(self): - raise NotImplementedError() +__all__ = ['CLICommand'] diff --git a/modelscope/cli/clearcache.py b/modelscope/cli/clearcache.py index 0713db94..72b8a0a9 100644 --- a/modelscope/cli/clearcache.py +++ b/modelscope/cli/clearcache.py @@ -1,11 +1,17 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""Clear-cache CLI command — retained for backward compatibility. + +The ``modelscope_hub`` CLI now owns ``clear-cache`` as an alias for +``cache clear``, but this module preserves the legacy :class:`ClearCacheCMD` +class so that existing tests and callers that import it directly continue +to work. +""" import os import shutil from argparse import ArgumentParser from modelscope.cli.base import CLICommand from modelscope.hub.constants import TEMPORARY_FOLDER_NAME -from modelscope.hub.utils.utils import get_model_masked_directory from modelscope.utils.file_utils import (get_dataset_cache_root, get_model_cache_root, get_modelscope_cache_dir) @@ -24,6 +30,11 @@ class ClearCacheCMD(CLICommand): self.args = args self.cache_dir = get_modelscope_cache_dir() + @staticmethod + def register(subparsers) -> None: + """Register clear-cache subcommand (CLICommand ABC contract).""" + ClearCacheCMD.define_args(subparsers) + @staticmethod def define_args(parsers: ArgumentParser): """ define args for clear-cache command. @@ -33,17 +44,15 @@ class ClearCacheCMD(CLICommand): group.add_argument( '--model', type=str, - help= - 'The id of the model whose cache will be cleared. For clear-cache, ' - 'if neither model or dataset id is provided, entire cache will be cleared.' - ) + help='The id of the model whose cache will be cleared. ' + 'If neither model or dataset id is provided, entire cache ' + 'will be cleared.') group.add_argument( '--dataset', type=str, - help= - 'The id of the dataset whose cache will be cleared. For clear-cache, ' - 'if neither model or dataset id is provided, entire cache will be cleared.' - ) + help='The id of the dataset whose cache will be cleared. ' + 'If neither model or dataset id is provided, entire cache ' + 'will be cleared.') parser.set_defaults(func=subparser_func) @@ -64,7 +73,8 @@ class ClearCacheCMD(CLICommand): id = self.args.dataset prompt = prompt + f'local cache for dataset {id}. ' else: - prompt = prompt + f'entire ModelScope cache at {self.cache_dir}, including ALL models and dataset.\n' + prompt = prompt + (f'entire ModelScope cache at {self.cache_dir}, ' + f'including ALL models and dataset.\n') all = True user_input = input( prompt diff --git a/modelscope/cli/cli.py b/modelscope/cli/cli.py index 7a45e3fc..030986dd 100644 --- a/modelscope/cli/cli.py +++ b/modelscope/cli/cli.py @@ -1,62 +1,21 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""ModelScope CLI — delegates to the modelscope_hub CLI engine. -import argparse -import logging +The legacy ``modelscope`` / ``ms`` console-script entry points historically +lived here as a hand-rolled argparse tree. The hub CLI in ``modelscope_hub`` +now owns command registration, plugin discovery, and error translation; +this module exists solely to preserve the import path used by the +``[project.scripts]`` entries in ``pyproject.toml``. +""" -from modelscope.cli.clearcache import ClearCacheCMD -from modelscope.cli.create import CreateCMD -from modelscope.cli.download import DownloadCMD -from modelscope.cli.llamafile import LlamafileCMD -from modelscope.cli.login import LoginCMD -from modelscope.cli.modelcard import ModelCardCMD -from modelscope.cli.pipeline import PipelineCMD -from modelscope.cli.plugins import PluginsCMD -from modelscope.cli.scancache import ScanCacheCMD -from modelscope.cli.server import ServerCMD -from modelscope.cli.skills import SkillsCMD -from modelscope.cli.studio import StudioCMD -from modelscope.cli.upload import UploadCMD -from modelscope.hub.constants import MODELSCOPE_ASCII -from modelscope.utils.logger import get_logger -from modelscope.version import __version__ +import sys -logger = get_logger(log_level=logging.WARNING) +from modelscope_hub.cli.main import run_cmd as _run_cmd def run_cmd(): - print(MODELSCOPE_ASCII) - parser = argparse.ArgumentParser( - 'ModelScope Command Line tool', usage='modelscope []') - parser.add_argument( - '-V', - '--version', - action='version', - version=f'ModelScope CLI {__version__}') - parser.add_argument( - '--token', default=None, help='Specify ModelScope SDK token.') - subparsers = parser.add_subparsers(help='modelscope commands helpers') - - CreateCMD.define_args(subparsers) - DownloadCMD.define_args(subparsers) - SkillsCMD.define_args(subparsers) - UploadCMD.define_args(subparsers) - ClearCacheCMD.define_args(subparsers) - PluginsCMD.define_args(subparsers) - PipelineCMD.define_args(subparsers) - ModelCardCMD.define_args(subparsers) - ServerCMD.define_args(subparsers) - LoginCMD.define_args(subparsers) - LlamafileCMD.define_args(subparsers) - ScanCacheCMD.define_args(subparsers) - StudioCMD.define_args(subparsers) - - args = parser.parse_args() - - if not hasattr(args, 'func'): - parser.print_help() - exit(1) - cmd = args.func(args) - cmd.execute() + """Delegate to ``modelscope_hub.cli.main.run_cmd`` and propagate its exit code.""" + sys.exit(_run_cmd()) if __name__ == '__main__': diff --git a/modelscope/cli/create.py b/modelscope/cli/create.py deleted file mode 100644 index 5fc516c6..00000000 --- a/modelscope/cli/create.py +++ /dev/null @@ -1,287 +0,0 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. -from argparse import ArgumentParser, _SubParsersAction - -from modelscope.cli.base import CLICommand -from modelscope.hub.api import HubApi -from modelscope.hub.constants import (GatedMode, Licenses, ModelVisibility, - Visibility, VisibilityMap) -from modelscope.hub.utils.aigc import AigcModel -from modelscope.hub.utils.utils import resolve_endpoint -from modelscope.utils.constant import (REPO_TYPE_MODEL, REPO_TYPE_STUDIO, - REPO_TYPE_SUPPORT, StudioHardware, - StudioSDKType) -from modelscope.utils.logger import get_logger - -logger = get_logger() - - -def subparser_func(args): - """ Function which will be called for a specific sub parser. - """ - return CreateCMD(args) - - -class CreateCMD(CLICommand): - """ - Command for creating a new repository, supporting both model and dataset. - """ - - name = 'create' - - def __init__(self, args: _SubParsersAction): - self.args = args - - @staticmethod - def define_args(parsers: _SubParsersAction): - - parser: ArgumentParser = parsers.add_parser(CreateCMD.name) - - parser.add_argument( - 'repo_id', - type=str, - help='The ID of the repo to create (e.g. `username/repo-name`)') - parser.add_argument( - '--token', - type=str, - default=None, - help= - 'A User Access Token generated from https://modelscope.cn/my/myaccesstoken to authenticate the user. ' - 'If not provided, the CLI will use the local credentials if available.' - ) - parser.add_argument( - '--repo_type', - choices=REPO_TYPE_SUPPORT, - default=REPO_TYPE_MODEL, - help= - 'Type of the repo to create (e.g. `dataset`, `model`). Default to `model`.', - ) - parser.add_argument( - '--visibility', - choices=[ - Visibility.PUBLIC, Visibility.INTERNAL, Visibility.PRIVATE - ], - default=Visibility.PUBLIC, - help='Visibility of the repo to create. Default to `public`.', - ) - parser.add_argument( - '--chinese_name', - type=str, - default=None, - help='Optional, Chinese name of the repo. Default to `None`.', - ) - parser.add_argument( - '--license', - type=str, - choices=Licenses.to_list(), - default=Licenses.APACHE_V2, - help= - 'Optional, License of the repo. Default to `Apache License 2.0`.', - ) - parser.add_argument( - '--exist_ok', - action='store_true', - default=False, - help= - 'If True, do not raise error when repo already exists. Defaults to False.', - ) - parser.add_argument( - '--gated', - dest='gated_mode', - action='store_true', - default=None, - help= - 'Enable gated mode (application-based download) for private repos.', - ) - parser.add_argument( - '--no-gated', - dest='gated_mode', - action='store_false', - help='Disable gated mode for private repos (normal private).', - ) - parser.add_argument( - '--endpoint', - type=str, - default=None, - help='ModelScope server endpoint, e.g. modelscope.cn or ' - 'modelscope.ai Full URL like ' - 'https://modelscope.cn is also accepted. Scheme (https://) is ' - 'auto-completed if omitted. Falls back to env MODELSCOPE_DOMAIN, ' - 'then defaults to https://www.modelscope.cn.', - ) - - # Studio specific arguments (only meaningful when --repo_type studio) - studio_group = parser.add_argument_group( - 'Studio Repo Creation', - 'Optional arguments used only when `--repo_type studio` is set.') - studio_group.add_argument( - '--sdk-type', - dest='sdk_type', - choices=StudioSDKType.SUPPORTED, - default=None, - help='Studio SDK type (only for studio repo-type).') - studio_group.add_argument( - '--sdk-version', - dest='sdk_version', - type=str, - default=None, - help='Studio SDK version (only for gradio).') - studio_group.add_argument( - '--base-image', - dest='base_image', - type=str, - default=None, - help='Studio base image (only for gradio/streamlit).') - studio_group.add_argument( - '--hardware', - dest='hardware', - choices=StudioHardware.SUPPORTED, - default=None, - help='Studio hardware configuration.') - - # AIGC specific arguments - aigc_group = parser.add_argument_group( - 'AIGC Model Creation', - 'Arguments for creating an AIGC model. Use --aigc to enable.') - aigc_group.add_argument( - '--aigc', action='store_true', help='Enable AIGC model creation.') - aigc_group.add_argument( - '--from_json', - type=str, - help='Path to a JSON file containing AIGC model configuration. ' - 'If used, all other parameters except --repo_id are ignored.') - aigc_group.add_argument( - '--model_path', type=str, help='Path to the model file or folder.') - aigc_group.add_argument( - '--aigc_type', - type=str, - help="AIGC type. Recommended: 'Checkpoint', 'LoRA', 'VAE'.") - aigc_group.add_argument( - '--base_model_type', - type=str, - help='Base model type, e.g., SD_XL.') - aigc_group.add_argument( - '--revision', - type=str, - default='v1.0', - help="Model revision. Defaults to 'v1.0'.") - aigc_group.add_argument( - '--base_model_id', - type=str, - default='', - help='Base model ID from ModelScope.') - aigc_group.add_argument( - '--description', - type=str, - default='This is an AIGC model.', - help='Model description.') - aigc_group.add_argument( - '--path_in_repo', - type=str, - default='', - help='Path in the repository to upload to.') - aigc_group.add_argument( - '--model_source', - type=str, - default='USER_UPLOAD', - help= - 'Source of the AIGC model. `USER_UPLOAD`, `TRAINED_FROM_MODELSCOPE` or `TRAINED_FROM_ALIYUN_FC`.' - ) - aigc_group.add_argument( - '--base_model_sub_type', - type=str, - default='', - help='Base model sub type, e.g., Qwen_Edit_2509') - - parser.set_defaults(func=subparser_func) - - def execute(self): - if self.args.aigc: - if self.args.repo_type != REPO_TYPE_MODEL: - raise ValueError( - 'AIGC models can only be created when repo_type is "model".' - ) - self._create_aigc_model() - else: - self._create_regular_repo() - - def _create_regular_repo(self): - # Check token and login - # The cookies will be reused if the user has logged in before. - endpoint = resolve_endpoint(self.args.endpoint) - api = HubApi(endpoint=endpoint) - - extra_kwargs = {} - if self.args.repo_type == REPO_TYPE_STUDIO: - # Pass studio-specific fields only when creating a studio repo. - for field in ('sdk_type', 'sdk_version', 'base_image', 'hardware'): - value = getattr(self.args, field, None) - if value is not None: - extra_kwargs[field] = value - - # Create repo - api.create_repo( - repo_id=self.args.repo_id, - token=self.args.token, - visibility=self.args.visibility, - repo_type=self.args.repo_type, - chinese_name=self.args.chinese_name, - license=self.args.license, - exist_ok=self.args.exist_ok, - create_default_config=True, - endpoint=endpoint, - gated_mode=self.args.gated_mode, - **extra_kwargs, - ) - - def _create_aigc_model(self): - """Execute the command.""" - endpoint = resolve_endpoint(self.args.endpoint) - api = HubApi(endpoint=endpoint) - model_id = self.args.repo_id - - if self.args.from_json: - # Create from JSON file - logger.info('Creating AIGC model from JSON file: ' - f'{self.args.from_json}') - aigc_model = AigcModel.from_json_file(self.args.from_json) - else: - # Create from command line arguments - logger.info('Creating AIGC model from command line arguments...') - if not all([ - self.args.model_path, self.args.aigc_type, - self.args.base_model_type - ]): - raise ValueError( - 'Error: --model_path, --aigc_type, and ' - '--base_model_type are required when not using ' - '--from_json.') - - aigc_model = AigcModel( - model_path=self.args.model_path, - aigc_type=self.args.aigc_type, - base_model_type=self.args.base_model_type, - tag=self.args.revision, - description=self.args.description, - base_model_id=self.args.base_model_id, - path_in_repo=self.args.path_in_repo, - model_source=self.args.model_source, - base_model_sub_type=self.args.base_model_sub_type, - ) - - # Convert visibility string to int for the API call - reverse_visibility_map = {v: k for k, v in VisibilityMap.items()} - visibility_idx: int = reverse_visibility_map.get( - self.args.visibility, ModelVisibility.PUBLIC) - - try: - model_url = api.create_model( - model_id=model_id, - token=self.args.token, - visibility=visibility_idx, - license=self.args.license, - chinese_name=self.args.chinese_name, - aigc_model=aigc_model, - gated_mode=self.args.gated_mode) - print(f'Successfully created AIGC model: {model_url}') - except Exception as e: - print(f'Error creating AIGC model: {e}') diff --git a/modelscope/cli/download.py b/modelscope/cli/download.py deleted file mode 100644 index b81c9a54..00000000 --- a/modelscope/cli/download.py +++ /dev/null @@ -1,287 +0,0 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. -import logging -from argparse import ArgumentParser - -from modelscope.cli.base import CLICommand -from modelscope.cli.utils import concurrent_download -from modelscope.hub.api import HubApi -from modelscope.hub.constants import DEFAULT_MAX_WORKERS, DEFAULT_SKILLS_DIR -from modelscope.hub.file_download import (dataset_file_download, - model_file_download) -from modelscope.hub.snapshot_download import (dataset_snapshot_download, - snapshot_download) -from modelscope.hub.utils.utils import convert_patterns, resolve_endpoint -from modelscope.utils.constant import (DEFAULT_DATASET_REVISION, - REPO_TYPE_DATASET, REPO_TYPE_MODEL, - REPO_TYPE_STUDIO, REPO_TYPE_SUPPORT) -from modelscope.utils.logger import get_logger - -logger = get_logger(log_level=logging.WARNING) - - -def subparser_func(args): - """ Function which will be called for a specific sub parser. - """ - return DownloadCMD(args) - - -class DownloadCMD(CLICommand): - name = 'download' - - def __init__(self, args): - self.args = args - - @staticmethod - def define_args(parsers: ArgumentParser): - """ define args for download command. - """ - parser: ArgumentParser = parsers.add_parser(DownloadCMD.name) - group = parser.add_mutually_exclusive_group() - group.add_argument( - '--model', - type=str, - help='The id of the model to be downloaded. For download, ' - 'the id of either a model or dataset must be provided.') - group.add_argument( - '--dataset', - type=str, - help='The id of the dataset to be downloaded. For download, ' - 'the id of either a model or dataset must be provided.') - group.add_argument( - '--collection', - type=str, - default=None, - help='The ID of the collection to download (skills only)') - parser.add_argument( - 'repo_id', - type=str, - nargs='?', - default=None, - help='Optional, ' - 'ID of the repo to download, It can also be set by --model or --dataset.' - ) - parser.add_argument( - '--repo-type', - choices=REPO_TYPE_SUPPORT, - default=REPO_TYPE_MODEL, - help="Type of repo to download from (defaults to 'model').", - ) - parser.add_argument( - '--token', - type=str, - default=None, - help='Optional. Access token to download controlled entities.') - parser.add_argument( - '--revision', - type=str, - default=None, - help='Revision of the entity (e.g., model).') - parser.add_argument( - '--cache_dir', - type=str, - default=None, - help='Cache directory to save entity (e.g., model).') - parser.add_argument( - '--local_dir', - type=str, - default=None, - help='File will be downloaded to local location specified by' - 'local_dir, in this case, cache_dir parameter will be ignored.') - parser.add_argument( - 'files', - type=str, - default=None, - nargs='*', - help='Specify relative path to the repository file(s) to download.' - "(e.g 'tokenizer.json', 'onnx/decoder_model.onnx').") - parser.add_argument( - '--include', - nargs='*', - default=None, - type=str, - help='Glob patterns to match files to download.' - 'Ignored if file is specified') - parser.add_argument( - '--exclude', - nargs='*', - type=str, - default=None, - help='Glob patterns to exclude from files to download.' - 'Ignored if file is specified') - parser.add_argument( - '--endpoint', - type=str, - default=None, - help='ModelScope server endpoint, e.g. modelscope.cn or ' - 'modelscope.ai Full URL like ' - 'https://modelscope.cn is also accepted. Scheme (https://) is ' - 'auto-completed if omitted. Falls back to env MODELSCOPE_DOMAIN, ' - 'then defaults to https://www.modelscope.cn. ' - 'When omitted, the CLI auto-detects the correct site ' - '(cn/intl) for download.') - parser.add_argument( - '--max-workers', - type=int, - default=DEFAULT_MAX_WORKERS, - help='The maximum number of workers to download files.') - - parser.set_defaults(func=subparser_func) - - def execute(self): - if self.args.model or self.args.dataset: - # the position argument of files will be put to repo_id. - if self.args.repo_id is not None: - if self.args.files: - self.args.files.insert(0, self.args.repo_id) - else: - self.args.files = [self.args.repo_id] - else: - if self.args.repo_id is not None: - if self.args.repo_type in (REPO_TYPE_MODEL, REPO_TYPE_STUDIO): - # studio repos share the same snapshot_download path - # as model repos. - self.args.model = self.args.repo_id - elif self.args.repo_type == REPO_TYPE_DATASET: - self.args.dataset = self.args.repo_id - else: - raise Exception('Not support repo-type: %s' - % self.args.repo_type) - if not self.args.model and not self.args.dataset and not self.args.collection: - raise Exception('Model, dataset, or collection must be set.') - if self.args.endpoint: - endpoint = resolve_endpoint(self.args.endpoint) - else: - endpoint = None - cookies = None - if self.args.token is not None: - api = HubApi(endpoint=endpoint) - cookies = api.get_cookies(access_token=self.args.token) - if self.args.model: - if len(self.args.files) == 1: # download single file - model_file_download( - self.args.model, - self.args.files[0], - cache_dir=self.args.cache_dir, - local_dir=self.args.local_dir, - revision=self.args.revision, - cookies=cookies, - token=self.args.token, - endpoint=endpoint) - elif len( - self.args.files) > 1: # download specified multiple files. - snapshot_download( - self.args.model, - revision=self.args.revision, - cache_dir=self.args.cache_dir, - local_dir=self.args.local_dir, - allow_file_pattern=self.args.files, - max_workers=self.args.max_workers, - cookies=cookies, - token=self.args.token, - endpoint=endpoint) - else: # download repo - snapshot_download( - self.args.model, - revision=self.args.revision, - cache_dir=self.args.cache_dir, - local_dir=self.args.local_dir, - allow_file_pattern=convert_patterns(self.args.include), - ignore_file_pattern=convert_patterns(self.args.exclude), - max_workers=self.args.max_workers, - cookies=cookies, - token=self.args.token, - endpoint=endpoint) - print(f'\nSuccessfully Downloaded from model {self.args.model}.\n') - elif self.args.dataset: - dataset_revision: str = self.args.revision if self.args.revision else DEFAULT_DATASET_REVISION - if len(self.args.files) == 1: # download single file - dataset_file_download( - self.args.dataset, - self.args.files[0], - cache_dir=self.args.cache_dir, - local_dir=self.args.local_dir, - revision=dataset_revision, - cookies=cookies, - token=self.args.token, - endpoint=endpoint) - elif len( - self.args.files) > 1: # download specified multiple files. - dataset_snapshot_download( - self.args.dataset, - revision=dataset_revision, - cache_dir=self.args.cache_dir, - local_dir=self.args.local_dir, - allow_file_pattern=self.args.files, - max_workers=self.args.max_workers, - cookies=cookies, - token=self.args.token, - endpoint=endpoint) - else: # download repo - dataset_snapshot_download( - self.args.dataset, - revision=dataset_revision, - cache_dir=self.args.cache_dir, - local_dir=self.args.local_dir, - allow_file_pattern=convert_patterns(self.args.include), - ignore_file_pattern=convert_patterns(self.args.exclude), - max_workers=self.args.max_workers, - cookies=cookies, - token=self.args.token, - endpoint=endpoint) - print( - f'\nSuccessfully Downloaded from dataset {self.args.dataset}.\n' - ) - elif self.args.collection: - api = HubApi(endpoint=endpoint, token=self.args.token) - local_dir = self.args.local_dir or DEFAULT_SKILLS_DIR - data = api.get_collection( - self.args.collection, repo_type='skill', endpoint=endpoint) - elements = data.get('CollectionElements', - {}).get('CollectionElementVoList', []) - - logger.info( - f'Collection {self.args.collection} has {len(elements)} elements.' - ) - - if not elements: - print(f'No skill elements found in collection: ' - f'{self.args.collection}') - return - - # Validate elements have required fields - valid_elements = [] - for elem in elements: - if not elem.get('ElementPath') or not elem.get('ElementName'): - logger.warning('Skipping malformed collection element: %s', - elem) - continue - valid_elements.append(elem) - - if not valid_elements: - print(f'No valid skill elements found in collection: ' - f'{self.args.collection}') - return - - print(f'Found {len(valid_elements)} skill(s) in collection, ' - f'downloading...') - - def _download_one_skill(element): - element_path = element['ElementPath'] - element_name = element['ElementName'] - skill_id = f'{element_path}/{element_name}' - try: - skill_dir = api.download_skill( - skill_id=skill_id, - local_dir=local_dir, - endpoint=endpoint) - return (skill_id, skill_dir, None) - except Exception as e: - return (skill_id, None, str(e)) - - concurrent_download( - _download_one_skill, - valid_elements, - max_workers=self.args.max_workers, - item_name='skill') - else: - pass # noop diff --git a/modelscope/cli/llamafile.py b/modelscope/cli/llamafile.py index 23f3fe91..e136ce0e 100644 --- a/modelscope/cli/llamafile.py +++ b/modelscope/cli/llamafile.py @@ -1,28 +1,25 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""``modelscope llamafile`` — download and run a llamafile from a model repo.""" + import logging import os import sys from argparse import ArgumentParser +from modelscope_hub.cli.base import CLICommand + from modelscope import model_file_download -from modelscope.cli.base import CLICommand from modelscope.hub.api import HubApi from modelscope.utils.logger import get_logger logger = get_logger(log_level=logging.WARNING) -def subparser_func(args): - """ Function which will be called for a specific sub parser. - """ - return LlamafileCMD(args) - - class LlamafileCMD(CLICommand): name = 'llamafile' def __init__(self, args): - self.args = args + super().__init__(args) self.model_id = self.args.model if self.model_id is None or self.model_id.count('/') != 1: raise ValueError(f'Invalid model id [{self.model_id}].') @@ -34,10 +31,10 @@ class LlamafileCMD(CLICommand): self.api = HubApi() @staticmethod - def define_args(parsers: ArgumentParser): - """ define args for clear-cache command. - """ - parser = parsers.add_parser(LlamafileCMD.name) + def register(subparsers: ArgumentParser) -> None: + parser = subparsers.add_parser( + LlamafileCMD.name, + help='Download and run a llamafile from a model repo.') parser.add_argument( '--model', type=str, @@ -80,7 +77,7 @@ class LlamafileCMD(CLICommand): 'Whether to launch model with the downloaded llamafile, default to True.' ) - parser.set_defaults(func=subparser_func) + parser.set_defaults(_command=LlamafileCMD) def execute(self): if self.args.file: diff --git a/modelscope/cli/login.py b/modelscope/cli/login.py deleted file mode 100644 index 32eb3237..00000000 --- a/modelscope/cli/login.py +++ /dev/null @@ -1,45 +0,0 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. - -from argparse import ArgumentParser - -from modelscope.cli.base import CLICommand -from modelscope.hub.api import HubApi -from modelscope.hub.utils.utils import resolve_endpoint - - -def subparser_func(args): - """ Function which will be called for a specific sub parser. - """ - return LoginCMD(args) - - -class LoginCMD(CLICommand): - name = 'login' - - def __init__(self, args): - self.args = args - - @staticmethod - def define_args(parsers: ArgumentParser): - """ define args for login command. - """ - parser = parsers.add_parser(LoginCMD.name) - parser.add_argument( - '--token', - type=str, - required=True, - help='The Access Token for modelscope.') - parser.add_argument( - '--endpoint', - type=str, - default=None, - help='ModelScope server endpoint, e.g. modelscope.cn or ' - 'modelscope.ai Full URL like ' - 'https://modelscope.cn is also accepted. Scheme (https://) is ' - 'auto-completed if omitted. Falls back to env MODELSCOPE_DOMAIN, ' - 'then defaults to https://www.modelscope.cn.') - parser.set_defaults(func=subparser_func) - - def execute(self): - api = HubApi(endpoint=resolve_endpoint(self.args.endpoint)) - api.login(self.args.token) diff --git a/modelscope/cli/modelcard.py b/modelscope/cli/modelcard.py index 646cf1b0..d35e1540 100644 --- a/modelscope/cli/modelcard.py +++ b/modelscope/cli/modelcard.py @@ -1,4 +1,6 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""``modelscope modelcard`` — create / upload / download a model card.""" + import logging import os import shutil @@ -6,7 +8,8 @@ import tempfile from argparse import ArgumentParser from string import Template -from modelscope.cli.base import CLICommand +from modelscope_hub.cli.base import CLICommand + from modelscope.hub.api import HubApi from modelscope.hub.snapshot_download import snapshot_download from modelscope.hub.utils.utils import get_endpoint @@ -18,17 +21,11 @@ current_path = os.path.dirname(os.path.abspath(__file__)) template_path = os.path.join(current_path, 'template') -def subparser_func(args): - """ Function which will be called for a specific sub parser. - """ - return ModelCardCMD(args) - - class ModelCardCMD(CLICommand): name = 'modelcard' def __init__(self, args): - self.args = args + super().__init__(args) self.api = HubApi() if args.access_token: self.api.login(args.access_token) @@ -38,10 +35,11 @@ class ModelCardCMD(CLICommand): self.url = os.path.join(get_endpoint(), self.model_id) @staticmethod - def define_args(parsers: ArgumentParser): - """ define args for create or upload modelcard command. - """ - parser = parsers.add_parser(ModelCardCMD.name, aliases=['model']) + def register(subparsers: ArgumentParser) -> None: + parser = subparsers.add_parser( + ModelCardCMD.name, + aliases=['model'], + help='Create / upload / download a model card.') parser.add_argument( '-tk', '--access_token', @@ -105,7 +103,7 @@ class ModelCardCMD(CLICommand): type=str, default=None, help='the info of uploaded model') - parser.set_defaults(func=subparser_func) + parser.set_defaults(_command=ModelCardCMD) def create_model(self): from modelscope.hub.constants import Licenses, ModelVisibility diff --git a/modelscope/cli/pipeline.py b/modelscope/cli/pipeline.py index 2b6f7951..50139ff6 100644 --- a/modelscope/cli/pipeline.py +++ b/modelscope/cli/pipeline.py @@ -1,10 +1,13 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""``modelscope pipeline`` — scaffold a custom pipeline from a template.""" + import logging import os from argparse import ArgumentParser from string import Template -from modelscope.cli.base import CLICommand +from modelscope_hub.cli.base import CLICommand + from modelscope.utils.logger import get_logger logger = get_logger(log_level=logging.WARNING) @@ -13,23 +16,14 @@ current_path = os.path.dirname(os.path.abspath(__file__)) template_path = os.path.join(current_path, 'template') -def subparser_func(args): - """ Function which will be called for a specific sub parser. - """ - return PipelineCMD(args) - - class PipelineCMD(CLICommand): name = 'pipeline' - def __init__(self, args): - self.args = args - @staticmethod - def define_args(parsers: ArgumentParser): - """ define args for create pipeline template command. - """ - parser = parsers.add_parser(PipelineCMD.name) + def register(subparsers: ArgumentParser) -> None: + parser = subparsers.add_parser( + PipelineCMD.name, + help='Scaffold a custom pipeline from a template.') parser.add_argument( '-act', '--action', @@ -85,7 +79,7 @@ class PipelineCMD(CLICommand): type=str, default='./', help='the path of configuration.json for ModelScope') - parser.set_defaults(func=subparser_func) + parser.set_defaults(_command=PipelineCMD) def create_template(self): if self.args.tpl_file_path not in os.listdir(template_path): diff --git a/modelscope/cli/plugins.py b/modelscope/cli/plugins.py index bcf8f0ef..218bc160 100644 --- a/modelscope/cli/plugins.py +++ b/modelscope/cli/plugins.py @@ -1,54 +1,25 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""``modelscope plugin`` — install/uninstall/list ModelScope plugins.""" from argparse import ArgumentParser -from modelscope.cli.base import CLICommand +from modelscope_hub.cli.base import CLICommand + from modelscope.utils.plugins import PluginsManager plugins_manager = PluginsManager() -def subparser_func(args): - """ Function which will be called for a specific sub parser. - """ - return PluginsCMD(args) - - class PluginsCMD(CLICommand): name = 'plugin' - def __init__(self, args): - self.args = args - @staticmethod - def define_args(parsers: ArgumentParser): - """ define args for install command. - """ - parser = parsers.add_parser(PluginsCMD.name) - subparsers = parser.add_subparsers(dest='command') + def register(subparsers: ArgumentParser) -> None: + parser = subparsers.add_parser( + PluginsCMD.name, help='Manage ModelScope plugins.') + sub = parser.add_subparsers(dest='command') - PluginsInstallCMD.define_args(subparsers) - PluginsUninstallCMD.define_args(subparsers) - PluginsListCMD.define_args(subparsers) - - parser.set_defaults(func=subparser_func) - - def execute(self): - print(self.args) - if self.args.command == PluginsInstallCMD.name: - PluginsInstallCMD.execute(self.args) - if self.args.command == PluginsUninstallCMD.name: - PluginsUninstallCMD.execute(self.args) - if self.args.command == PluginsListCMD.name: - PluginsListCMD.execute(self.args) - - -class PluginsInstallCMD(PluginsCMD): - name = 'install' - - @staticmethod - def define_args(parsers: ArgumentParser): - install = parsers.add_parser(PluginsInstallCMD.name) + install = sub.add_parser('install', help='Install plugin packages.') install.add_argument( 'package', type=str, @@ -68,51 +39,41 @@ class PluginsInstallCMD(PluginsCMD): default=False, help='If force update the package') - @staticmethod - def execute(args): - plugins_manager.install_plugins( - list(args.package), - index_url=args.index_url, - force_update=args.force_update) - - -class PluginsUninstallCMD(PluginsCMD): - name = 'uninstall' - - @staticmethod - def define_args(parsers: ArgumentParser): - install = parsers.add_parser(PluginsUninstallCMD.name) - install.add_argument( + uninstall = sub.add_parser( + 'uninstall', help='Uninstall plugin packages.') + uninstall.add_argument( 'package', type=str, nargs='+', default=None, - help='Name of the package to be installed.') - install.add_argument( + help='Name of the package to be uninstalled.') + uninstall.add_argument( '--yes', '-y', - type=str, - default=False, - help='Base URL of the Python Package Index.') + action='store_true', + help='Skip confirmation prompt.') - @staticmethod - def execute(args): - plugins_manager.uninstall_plugins(list(args.package), is_yes=args.yes) - - -class PluginsListCMD(PluginsCMD): - name = 'list' - - @staticmethod - def define_args(parsers: ArgumentParser): - install = parsers.add_parser(PluginsListCMD.name) - install.add_argument( + list_p = sub.add_parser('list', help='List available plugins.') + list_p.add_argument( '--all', '-a', - type=str, - default=None, + action='store_true', help='Show all of the plugins including those not installed.') - @staticmethod - def execute(args): - plugins_manager.list_plugins(show_all=all) + parser.set_defaults(_command=PluginsCMD) + + def execute(self): + command = getattr(self.args, 'command', None) + if command == 'install': + plugins_manager.install_plugins( + list(self.args.package), + index_url=self.args.index_url, + force_update=self.args.force_update) + elif command == 'uninstall': + plugins_manager.uninstall_plugins( + list(self.args.package), is_yes=self.args.yes) + elif command == 'list': + plugins_manager.list_plugins(show_all=self.args.all) + else: + raise ValueError( + 'Usage: modelscope plugin {install|uninstall|list} ...') diff --git a/modelscope/cli/scancache.py b/modelscope/cli/scancache.py deleted file mode 100644 index 98ce5dda..00000000 --- a/modelscope/cli/scancache.py +++ /dev/null @@ -1,64 +0,0 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. -import logging -import os -import time -from argparse import ArgumentParser -from typing import Optional - -from modelscope.cli.base import CLICommand -from modelscope.hub.cache_manager import scan_cache_dir -from modelscope.hub.errors import CacheNotFound -from modelscope.utils.logger import get_logger - -logger = get_logger(log_level=logging.WARNING) - -current_path = os.path.dirname(os.path.abspath(__file__)) - - -def subparser_func(args): - """ Function which will be called for a specific sub parser. - """ - return ScanCacheCMD(args) - - -class ScanCacheCMD(CLICommand): - name = 'scan-cache' - - def __init__(self, args): - self.args = args - self.cache_dir: Optional[str] = args.dir - - @staticmethod - def define_args(parsers: ArgumentParser): - """ define args for create pipeline template command. - """ - parser = parsers.add_parser(ScanCacheCMD.name) - group = parser.add_mutually_exclusive_group() - group.add_argument( - '--dir', - type=str, - default=None, - help= - 'cache directory to scan (optional). Default to the default ModelScope cache.', - ) - - parser.set_defaults(func=subparser_func) - - def execute(self): - try: - t0 = time.time() - cache_info = scan_cache_dir(self.cache_dir) - t1 = time.time() - except CacheNotFound as exc: - cache_dir = exc.cache_dir - print(f'Cache directory not found: {cache_dir}') - return - print(cache_info.export_as_table()) - print( - f'\nDone in {round(t1 - t0, 1)}s. Scanned {len(cache_info.repos)} repo(s)' - f' for a total of {cache_info.size_on_disk_str}.') - if len(cache_info.warnings) > 0: - message = f'Got {len(cache_info.warnings)} warning(s) while scanning.' - print(message) - for warning in cache_info.warnings: - print(warning) diff --git a/modelscope/cli/server.py b/modelscope/cli/server.py index 17d6ca4d..ef3dc7b5 100644 --- a/modelscope/cli/server.py +++ b/modelscope/cli/server.py @@ -1,38 +1,26 @@ # Copyright (c) Alibaba, Inc. and its affiliates. -import logging -import os -from argparse import ArgumentParser -from string import Template +"""``modelscope server`` — launch the local inference HTTP server.""" + +import logging +from argparse import ArgumentParser + +from modelscope_hub.cli.base import CLICommand -from modelscope.cli.base import CLICommand from modelscope.server.api_server import add_server_args, run_server from modelscope.utils.logger import get_logger logger = get_logger(log_level=logging.WARNING) -current_path = os.path.dirname(os.path.abspath(__file__)) -template_path = os.path.join(current_path, 'template') - - -def subparser_func(args): - """ Function which will be called for a specific sub parser. - """ - return ServerCMD(args) - class ServerCMD(CLICommand): name = 'server' - def __init__(self, args): - self.args = args - @staticmethod - def define_args(parsers: ArgumentParser): - """ define args for create pipeline template command. - """ - parser = parsers.add_parser(ServerCMD.name) + def register(subparsers: ArgumentParser) -> None: + parser = subparsers.add_parser( + ServerCMD.name, help='Launch the local inference HTTP server.') add_server_args(parser) - parser.set_defaults(func=subparser_func) + parser.set_defaults(_command=ServerCMD) def execute(self): run_server(self.args) diff --git a/modelscope/cli/skills.py b/modelscope/cli/skills.py index f35db0b1..af26d68b 100644 --- a/modelscope/cli/skills.py +++ b/modelscope/cli/skills.py @@ -1,10 +1,13 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""``modelscope skills`` — download and install agent skills.""" + import logging import sys from argparse import ArgumentParser +from concurrent.futures import ThreadPoolExecutor, as_completed + +from modelscope_hub.cli.base import CLICommand -from modelscope.cli.base import CLICommand -from modelscope.cli.utils import concurrent_download from modelscope.hub.api import HubApi from modelscope.hub.constants import DEFAULT_SKILLS_DIR from modelscope.utils.logger import get_logger @@ -12,9 +15,33 @@ from modelscope.utils.logger import get_logger logger = get_logger(log_level=logging.WARNING) -def subparser_func(args): - """Function which will be called for a specific sub parser.""" - return SkillsCMD(args) +def _concurrent_download(download_fn, items, max_workers=8, item_name='item'): + """Run ``download_fn`` over ``items`` in parallel, reporting progress. + + ``download_fn`` must return ``(identifier, result_path, error_or_None)``. + On any failure the process exits with status 1 after the summary is + printed. + """ + succeeded, failed = [], [] + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = {executor.submit(download_fn, item): item for item in items} + for future in as_completed(futures): + identifier, result_path, error = future.result() + if error: + failed.append((identifier, error)) + print(f'Failed to download {item_name} {identifier}: {error}') + else: + succeeded.append((identifier, result_path)) + print(f'Downloaded {item_name} {identifier} -> {result_path}') + + print(f'\nDownload complete: {len(succeeded)} succeeded, ' + f'{len(failed)} failed') + if failed: + print(f'Failed {item_name}s:') + for identifier, error in failed: + print(f' {identifier}: {error}') + sys.exit(1) + return succeeded, failed class SkillsCMD(CLICommand): @@ -22,19 +49,14 @@ class SkillsCMD(CLICommand): name = 'skills' - def __init__(self, args): - self.args = args - @staticmethod - def define_args(parsers: ArgumentParser): - """Define args for skills command.""" - parser = parsers.add_parser(SkillsCMD.name) - subparsers = parser.add_subparsers( + def register(subparsers: ArgumentParser) -> None: + parser = subparsers.add_parser( + SkillsCMD.name, help='Download and manage agent skills.') + sub = parser.add_subparsers( dest='skills_action', help='skills subcommands') - # 'add' subcommand - add_parser = subparsers.add_parser( - 'add', help='Download and install skills') + add_parser = sub.add_parser('add', help='Download and install skills') add_parser.add_argument( 'skill_ids', type=str, @@ -55,15 +77,16 @@ class SkillsCMD(CLICommand): type=int, default=8, help='Maximum concurrent downloads (default: 8)') - add_parser.set_defaults(func=subparser_func) + + parser.set_defaults(_command=SkillsCMD) def execute(self): - if not hasattr(self.args, - 'skills_action') or not self.args.skills_action: + if not getattr(self.args, 'skills_action', None): print('Usage: modelscope skills add ...') return - if not hasattr(self.args, 'skill_ids') or not self.args.skill_ids: + skill_ids = getattr(self.args, 'skill_ids', None) + if not skill_ids: print('No skill IDs provided. Usage: modelscope skills add ' ' ...') return @@ -71,11 +94,9 @@ class SkillsCMD(CLICommand): api = HubApi(token=self.args.token) local_dir = self.args.local_dir or DEFAULT_SKILLS_DIR - skill_ids = self.args.skill_ids print(f'Downloading {len(skill_ids)} skill(s)...') if len(skill_ids) == 1: - # Single skill download try: skill_dir = api.download_skill( skill_id=skill_ids[0], local_dir=local_dir) @@ -84,7 +105,7 @@ class SkillsCMD(CLICommand): print(f'Failed to download skill {skill_ids[0]}: {e}') sys.exit(1) else: - # Multiple skills - concurrent download + def _download_one(skill_id): try: skill_dir = api.download_skill( @@ -93,7 +114,7 @@ class SkillsCMD(CLICommand): except Exception as e: return (skill_id, None, str(e)) - concurrent_download( + _concurrent_download( _download_one, skill_ids, max_workers=self.args.max_workers, diff --git a/modelscope/cli/studio.py b/modelscope/cli/studio.py index 31a5d93f..9d2fc5f2 100644 --- a/modelscope/cli/studio.py +++ b/modelscope/cli/studio.py @@ -54,6 +54,11 @@ class StudioCMD(CLICommand): # ------------------------------------------------------------------ # Argument parsing # ------------------------------------------------------------------ + @staticmethod + def register(subparsers) -> None: + """Register studio subcommand (CLICommand ABC contract).""" + StudioCMD.define_args(subparsers) + @staticmethod def define_args(parsers: _SubParsersAction): parser: ArgumentParser = parsers.add_parser( diff --git a/modelscope/cli/upload.py b/modelscope/cli/upload.py deleted file mode 100644 index 7146a6cf..00000000 --- a/modelscope/cli/upload.py +++ /dev/null @@ -1,170 +0,0 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. -import os -from argparse import ArgumentParser, _SubParsersAction - -from modelscope.cli.base import CLICommand -from modelscope.hub.api import HubApi -from modelscope.hub.utils.utils import convert_patterns, resolve_endpoint -from modelscope.utils.constant import REPO_TYPE_MODEL, REPO_TYPE_SUPPORT - - -def subparser_func(args): - """ Function which will be called for a specific sub parser. - """ - return UploadCMD(args) - - -class UploadCMD(CLICommand): - - name = 'upload' - - def __init__(self, args: _SubParsersAction): - self.args = args - - @staticmethod - def define_args(parsers: _SubParsersAction): - - parser: ArgumentParser = parsers.add_parser(UploadCMD.name) - - parser.add_argument( - 'repo_id', - type=str, - help='The ID of the repo to upload to (e.g. `username/repo-name`)') - parser.add_argument( - 'local_path', - type=str, - nargs='?', - default=None, - help='Optional, ' - 'Local path to the file or folder to upload. Defaults to current directory.' - ) - parser.add_argument( - 'path_in_repo', - type=str, - nargs='?', - default=None, - help='Optional, ' - 'Path of the file or folder in the repo. Defaults to the relative path of the file or folder.' - ) - parser.add_argument( - '--repo-type', - choices=REPO_TYPE_SUPPORT, - default=REPO_TYPE_MODEL, - help= - 'Type of the repo to upload to (e.g. `dataset`, `model`). Defaults to be `model`.', - ) - parser.add_argument( - '--include', - nargs='*', - type=str, - help='Glob patterns to match files to upload.') - parser.add_argument( - '--exclude', - nargs='*', - type=str, - help='Glob patterns to exclude from files to upload.') - parser.add_argument( - '--commit-message', - type=str, - default=None, - help='The message of commit. Default to be `None`.') - parser.add_argument( - '--commit-description', - type=str, - default=None, - help= - 'The description of the generated commit. Default to be `None`.') - parser.add_argument( - '--token', - type=str, - default=None, - help= - 'A User Access Token generated from https://modelscope.cn/my/myaccesstoken' - ) - parser.add_argument( - '--max-workers', - type=int, - default=min(8, - os.cpu_count() + 4), - help='The number of workers to use for uploading files.') - parser.add_argument( - '--endpoint', - type=str, - default=None, - help='ModelScope server endpoint, e.g. modelscope.cn or ' - 'modelscope.ai Full URL like ' - 'https://modelscope.cn is also accepted. Scheme (https://) is ' - 'auto-completed if omitted. Falls back to env MODELSCOPE_DOMAIN, ' - 'then defaults to https://www.modelscope.cn.') - - parser.set_defaults(func=subparser_func) - - def execute(self): - - assert self.args.repo_id, '`repo_id` is required' - assert self.args.repo_id.count( - '/') == 1, 'repo_id should be in format of username/repo-name' - repo_name: str = self.args.repo_id.split('/')[-1] - self.repo_id = self.args.repo_id - - # Check path_in_repo - if self.args.local_path is None and os.path.isfile(repo_name): - # Case 1: modelscope upload owner_name/test_repo - self.local_path = repo_name - self.path_in_repo = repo_name - elif self.args.local_path is None and os.path.isdir(repo_name): - # Case 2: modelscope upload owner_name/test_repo (run command line in the `repo_name` dir) - # => upload all files in current directory to remote root path - self.local_path = repo_name - self.path_in_repo = '.' - elif self.args.local_path is None: - # Case 3: user provided only a repo_id that does not match a local file or folder - # => the user must explicitly provide a local_path => raise exception - raise ValueError( - f"'{repo_name}' is not a local file or folder. Please set `local_path` explicitly." - ) - elif self.args.path_in_repo is None and os.path.isfile( - self.args.local_path): - # Case 4: modelscope upload owner_name/test_repo /path/to/your_file.csv - # => upload it to remote root path with same name - self.local_path = self.args.local_path - self.path_in_repo = os.path.basename(self.args.local_path) - elif self.args.path_in_repo is None: - # Case 5: modelscope upload owner_name/test_repo /path/to/your_folder - # => upload all files in current directory to remote root path - self.local_path = self.args.local_path - self.path_in_repo = '' - else: - # Finally, if both paths are explicit - self.local_path = self.args.local_path - self.path_in_repo = self.args.path_in_repo - - api = HubApi(endpoint=resolve_endpoint(self.args.endpoint)) - - if os.path.isfile(self.local_path): - api.upload_file( - path_or_fileobj=self.local_path, - path_in_repo=self.path_in_repo, - repo_id=self.repo_id, - repo_type=self.args.repo_type, - commit_message=self.args.commit_message, - commit_description=self.args.commit_description, - token=self.args.token, - ) - elif os.path.isdir(self.local_path): - api.upload_folder( - repo_id=self.repo_id, - folder_path=self.local_path, - path_in_repo=self.path_in_repo, - commit_message=self.args.commit_message, - commit_description=self.args.commit_description, - repo_type=self.args.repo_type, - allow_patterns=convert_patterns(self.args.include), - ignore_patterns=convert_patterns(self.args.exclude), - max_workers=self.args.max_workers, - token=self.args.token, - ) - else: - raise ValueError(f'{self.local_path} is not a valid local path') - - print(f'\nFinished uploading to {self.repo_id}') diff --git a/modelscope/cli/utils.py b/modelscope/cli/utils.py deleted file mode 100644 index f7b835d4..00000000 --- a/modelscope/cli/utils.py +++ /dev/null @@ -1,41 +0,0 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. -import sys -from concurrent.futures import ThreadPoolExecutor, as_completed - - -def concurrent_download(download_fn, items, max_workers=8, item_name='item'): - """Download multiple items concurrently with progress reporting. - - Args: - download_fn: Callable that takes an item and returns - (identifier, result_path, error_string_or_None). - items: List of items to download. - max_workers (int): Maximum concurrent workers. - item_name (str): Display name for the item type. - - Returns: - tuple: (succeeded_list, failed_list). - """ - succeeded = [] - failed = [] - - with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = {executor.submit(download_fn, item): item for item in items} - for future in as_completed(futures): - identifier, result_path, error = future.result() - if error: - failed.append((identifier, error)) - print(f'Failed to download {item_name} {identifier}: {error}') - else: - succeeded.append((identifier, result_path)) - print(f'Downloaded {item_name} {identifier} -> {result_path}') - - print(f'\nDownload complete: {len(succeeded)} succeeded, ' - f'{len(failed)} failed') - if failed: - print(f'Failed {item_name}s:') - for identifier, error in failed: - print(f' {identifier}: {error}') - sys.exit(1) - - return succeeded, failed diff --git a/modelscope/hub/__init__.py b/modelscope/hub/__init__.py index 7f788e61..73a7af79 100644 --- a/modelscope/hub/__init__.py +++ b/modelscope/hub/__init__.py @@ -1,2 +1,49 @@ -from .callback import ProgressCallback +"""modelscope.hub — shim layer delegating to modelscope_hub.""" + +import os as _os +from pathlib import Path as _Path + +from modelscope_hub import HubConfig as _HubConfig +from modelscope_hub import get_default_config as _get_default_config +from modelscope_hub import set_default_config as _set_default_config + +from .callback import ProgressCallback, TqdmCallback from .commit_scheduler import CommitScheduler +from .snapshot_download import snapshot_download + + +def _sync_config() -> None: + """Bridge legacy env vars that modelscope_hub does not natively recognize.""" + # MODELSCOPE_CACHE is already handled by HubConfig; only sync + # if a non-standard alias is set. + legacy_cache = _os.environ.get('MS_CACHE_HOME') + if legacy_cache and not _os.environ.get('MODELSCOPE_CACHE'): + _set_default_config(_HubConfig(cache_dir=legacy_cache)) + + # Bridge MODELSCOPE_CREDENTIALS_PATH → HubConfig.config_dir so credential + # lookup honours the legacy override. + creds_path = _os.environ.get('MODELSCOPE_CREDENTIALS_PATH') + if creds_path: + resolved = _Path(creds_path).expanduser().resolve() + # Legacy convention may point at either the credentials directory + # (e.g. ``~/.modelscope``) or at a credentials file inside it; the + # new HubConfig always expects the directory. Treat the path as a + # file when it exists as one, falling back to the legacy + # ``credentials`` filename heuristic for paths that do not yet + # exist on disk. + is_file = resolved.is_file() or (not resolved.exists() + and resolved.name == 'credentials') + config_dir = resolved.parent if is_file else resolved + cfg = _get_default_config() + if cfg.config_dir != config_dir: + cfg.config_dir = config_dir + + +_sync_config() + +__all__ = [ + 'CommitScheduler', + 'ProgressCallback', + 'TqdmCallback', + 'snapshot_download', +] diff --git a/modelscope/hub/api.py b/modelscope/hub/api.py index 49de2ab0..356e7a42 100644 --- a/modelscope/hub/api.py +++ b/modelscope/hub/api.py @@ -1,1862 +1,263 @@ # Copyright (c) Alibaba, Inc. and its affiliates. -# yapf: disable +"""Hub API — shim delegating to modelscope_hub. -import datetime -import fnmatch -import functools -import io -import os -import pickle -import platform -import re -import shutil -import tempfile -import time -import uuid -import warnings -import zipfile -from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor -from http import HTTPStatus -from http.cookiejar import CookieJar -from os.path import expanduser -from pathlib import Path -from typing import (Any, BinaryIO, Dict, Iterable, List, Literal, Optional, - Tuple, Union) -from urllib.parse import urlencode +This module preserves the legacy ``modelscope.hub.api`` public surface +(``HubApi``, ``ModelScopeConfig``, ``model_id_to_group_owner_name`` and a few +response-field constants) by delegating to the ``modelscope_hub`` package. +Single responsibility: thin compatibility layer. All real logic lives in +``modelscope_hub.compat.LegacyHubApi`` and ``modelscope_hub.config.HubConfig``. +""" +from __future__ import annotations + +import hashlib import json -import requests -from requests import Session -from requests.adapters import HTTPAdapter, Retry -from requests.exceptions import HTTPError -from tqdm.auto import tqdm +import os +import platform +from os.path import expanduser +from typing import Any, Dict, Optional, Tuple, Union -from modelscope.hub.constants import (API_HTTP_CLIENT_CONNECT_TIMEOUT, - API_HTTP_CLIENT_MAX_RETRIES, +import requests + +from modelscope_hub.compat import LegacyHubApi as _LegacyHubApi +from modelscope_hub.config import get_default_config + +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, API_RESPONSE_FIELD_MESSAGE, API_RESPONSE_FIELD_USERNAME, - CREATE_TAG_MAX_RETRIES, - CREATE_TAG_RETRY_BACKOFF, - DEFAULT_MAX_WORKERS, - DEFAULT_MODELSCOPE_INTL_DOMAIN, MODELSCOPE_CLOUD_ENVIRONMENT, MODELSCOPE_CLOUD_USERNAME, - MODELSCOPE_CREDENTIALS_PATH, - MODELSCOPE_DOMAIN, - MODELSCOPE_PREFER_AI_SITE, - MODELSCOPE_REQUEST_ID, - MODELSCOPE_URL_SCHEME, ONE_YEAR_SECONDS, - REQUESTS_API_HTTP_METHOD, - TEMPORARY_FOLDER_NAME, - UPLOAD_ADAPTIVE_BATCH_SIZE, - UPLOAD_BLOB_MAX_RETRIES, - UPLOAD_BLOB_RETRY_BACKOFF, - UPLOAD_BLOB_RETRY_MAX_WAIT, - UPLOAD_BLOB_TIMEOUT, - UPLOAD_BLOB_TQDM_DISABLE_THRESHOLD, - UPLOAD_COMMIT_BATCH_SIZE, - UPLOAD_FAILED_FILE_MAX_RETRIES, - UPLOAD_MAX_FILE_COUNT, - UPLOAD_MAX_FILE_COUNT_IN_DIR, - UPLOAD_MAX_FILE_SIZE, - UPLOAD_NORMAL_FILE_SIZE_TOTAL_LIMIT, - UPLOAD_REACT_BACKOFF_MAX_EXPONENT, - UPLOAD_REACT_ENABLED, - UPLOAD_REACT_MAX_DELAY, - UPLOAD_REACT_ROUND2_BASE_DELAY, - UPLOAD_REACT_ROUND3_FILE_DELAY, - UPLOAD_RETRY_ALLOWED_METHODS, - UPLOAD_SIZE_THRESHOLD_TO_ENFORCE_LFS, - UPLOAD_USE_CACHE, - UPLOAD_VALIDATE_BLOB_BATCH_SIZE, - VALID_SORT_KEYS, DatasetVisibility, - Licenses, ModelVisibility, Visibility, - VisibilityMap) -from modelscope.hub.errors import (InvalidParameter, NotExistError, - NotLoginException, RequestError, - datahub_raise_on_error, - handle_http_post_error, - handle_http_response, is_ok, - raise_for_http_status, raise_on_error) -from modelscope.hub.git import GitCommandWrapper -from modelscope.hub.info import DatasetInfo, ModelInfo -from modelscope.hub.repository import Repository -from modelscope.hub.upload_cache import UPLOAD_HASH_CACHE_FILE -from modelscope.hub.upload_pipeline import BatchTracker -from modelscope.hub.upload_tracker import (_LEGACY_PROGRESS_FILE, NullTracker, - UploadTracker, classify_error) -from modelscope.hub.utils.aigc import AigcModel -from modelscope.hub.utils.utils import (add_content_to_file, get_domain, - get_endpoint, get_readable_folder_size, - get_release_datetime, is_env_true, - model_id_to_group_owner_name) -from modelscope.utils.constant import (DEFAULT_DATASET_REVISION, - DEFAULT_MODEL_REVISION, - DEFAULT_REPOSITORY_REVISION, - MASTER_MODEL_BRANCH, META_FILES_FORMAT, - REPO_TYPE_DATASET, REPO_TYPE_MODEL, - REPO_TYPE_STUDIO, REPO_TYPE_SUPPORT, - ConfigFields, DatasetFormations, - DatasetMetaFormats, DownloadChannel, - DownloadMode, Frameworks, ModelFile, - Tasks, VirgoDatasetConfig) -from modelscope.utils.file_utils import (compute_file_hash, get_file_size, - is_relative_path) + MODELSCOPE_CREDENTIALS_PATH) +from modelscope.hub.utils.utils import model_id_to_group_owner_name from modelscope.utils.logger import get_logger -from modelscope.utils.repo_utils import (DATASET_LFS_SUFFIX, - DEFAULT_IGNORE_PATTERNS, - MODEL_LFS_SUFFIX, - CommitHistoryResponse, CommitInfo, - CommitOperation, CommitOperationAdd, - RepoUtils) logger = get_logger() - -def _calculate_adaptive_batch_size(total_files: int) -> int: - """Calculate optimal commit batch size based on total file count. - - Adaptive strategy ensures batch granularity scales with workload: - - Very few files (1-10): no splitting, single batch - - Few files (11-100): ~10 batches for failure isolation - - Medium (101-10K): 64-256 files per batch - - Large (>10K): 512 files per batch to limit commit frequency - - Args: - total_files: Total number of files (including checkpoint-skipped). - - Returns: - Recommended batch size (>= 1). - """ - if total_files <= 0: - return 1 - if total_files <= 10: - return total_files - if total_files <= 100: - return max(1, total_files // 10) - if total_files <= 10_000: - return max(64, min(256, total_files // 80)) - return 512 +__all__ = [ + 'HubApi', + 'ModelScopeConfig', + 'model_id_to_group_owner_name', + 'API_RESPONSE_FIELD_DATA', + 'API_RESPONSE_FIELD_MESSAGE', + 'API_RESPONSE_FIELD_USERNAME', + 'API_RESPONSE_FIELD_EMAIL', + 'API_RESPONSE_FIELD_GIT_ACCESS_TOKEN', +] -class _CountedReadStream: - """File wrapper that counts bytes read and updates a progress bar. +class HubApi(_LegacyHubApi): + """ModelScope Hub API — delegates to ``modelscope_hub``. - Unlike a generator, this is a file-like object that requests can - use with Content-Length header for transfer integrity verification. + Maintains backward compatibility with the legacy ``HubApi`` interface; + method behaviour is inherited from + :class:`modelscope_hub.compat.LegacyHubApi`. """ - def __init__(self, file_obj, expected_size, pbar, chunk_size): - self._file = file_obj - self._expected_size = expected_size - self._pbar = pbar - self._chunk_size = chunk_size - self._bytes_read = 0 + def __init__( + self, + endpoint: Optional[str] = None, + timeout: int = API_HTTP_CLIENT_TIMEOUT, + max_retries: int = API_HTTP_CLIENT_MAX_RETRIES, + token: Optional[str] = None, + ) -> None: + super().__init__(endpoint=endpoint, token=token) + # Preserved for callers that historically read these attributes. + self.endpoint = self._endpoint or self._api._config.endpoint + self.token = token + self.timeout = timeout + self.max_retries = max_retries + self.headers = {'user-agent': ModelScopeConfig.get_user_agent()} - def read(self, size=-1): - """Read a chunk from the underlying file object.""" - read_size = self._chunk_size if size < 0 else min(size, self._chunk_size) - chunk = self._file.read(read_size) - if chunk: - n = len(chunk) - self._bytes_read += n - self._pbar.update(n) - return chunk + # If non-default timeout/max_retries were provided, eagerly construct + # the internal LegacyClient so they actually take effect on the wire. + if (timeout != API_HTTP_CLIENT_TIMEOUT + or max_retries != API_HTTP_CLIENT_MAX_RETRIES): + from modelscope_hub._legacy_api import LegacyClient + from modelscope_hub.utils import build_user_agent + cfg = self._api._config + self._api._legacy = LegacyClient( + token=cfg.token, + endpoint=cfg.endpoint, + timeout=timeout, + max_retries=max_retries, + user_agent=build_user_agent(cfg.get_session_id()), + ) + + # ------------------------------------------------------------------ + # Legacy method shims missing from LegacyHubApi + # ------------------------------------------------------------------ + def create_model(self, model_id: str, **kwargs) -> None: + """Create a model repo — delegates to ``create_repo`` (model type).""" + return self.create_repo(model_id, repo_type='model', **kwargs) + + def get_model_url(self, model_id: str) -> str: + """Return the model page URL ``{endpoint}/{model_id}``.""" + return f'{self.endpoint}/{model_id}' + + def upload_folder(self, repo_id: str, folder_path=None, **kwargs): + """Upload a folder — delegates to internal ``HubApi.upload_folder``.""" + from modelscope_hub.api import HubApi as _NewHubApi + repo_type = kwargs.pop('repo_type', None) or 'model' + token = kwargs.pop('token', None) + api = self._api + if token and token != self._api._config.token: + api = _NewHubApi(token=token, endpoint=self._api._config.endpoint) + return api.upload_folder( + repo_id=repo_id, + repo_type=repo_type, + folder_path=folder_path, + **kwargs, + ) + + def upload_file(self, repo_id: str = None, path_or_fileobj=None, + path_in_repo: str = None, **kwargs): + """Upload a file — delegates to internal ``HubApi.upload_file``.""" + from modelscope_hub.api import HubApi as _NewHubApi + repo_type = kwargs.pop('repo_type', None) or 'model' + token = kwargs.pop('token', None) + api = self._api + if token and token != self._api._config.token: + api = _NewHubApi(token=token, endpoint=self._api._config.endpoint) + return api.upload_file( + repo_id=repo_id, + repo_type=repo_type, + path_or_fileobj=path_or_fileobj, + path_in_repo=path_in_repo, + **kwargs, + ) @property - def bytes_read(self): - """Total bytes read so far.""" - return self._bytes_read + def _prepare_upload_folder(self): + """Expose UploadManager._prepare_upload_folder for monkey-patching.""" + return self._api.uploader._prepare_upload_folder - def verify_complete(self): - """Raise IOError if bytes read does not match expected size.""" - if self._bytes_read != self._expected_size: - raise IOError( - f'Upload data incomplete: read {self._bytes_read} bytes, ' - f'expected {self._expected_size} bytes. ' - f'File may have been modified during upload.') + @_prepare_upload_folder.setter + def _prepare_upload_folder(self, value): + """Allow CommitScheduler to monkey-patch ``_prepare_upload_folder``.""" + self._api.uploader._prepare_upload_folder = value - -class HubApi: - """Model hub api interface. - """ - - def __init__(self, - endpoint: Optional[str] = None, - timeout=API_HTTP_CLIENT_TIMEOUT, - max_retries=API_HTTP_CLIENT_MAX_RETRIES, - token: Optional[str] = None): - """The ModelScope HubApi。 - - Args: - endpoint (str, optional): The modelscope server http|https address. Defaults to None. - """ - self.endpoint = endpoint if endpoint is not None else get_endpoint() - self.token = token - self.headers = {'user-agent': ModelScopeConfig.get_user_agent()} - self.session = Session() - retry = Retry( - total=max_retries, - read=2, - connect=2, - backoff_factor=2, - status_forcelist=(500, 502, 503, 504), - allowed_methods=UPLOAD_RETRY_ALLOWED_METHODS, - respect_retry_after_header=True, - ) - adapter = HTTPAdapter(max_retries=retry) - self.session.mount('http://', adapter) - self.session.mount('https://', adapter) - # set http timeout - for method in REQUESTS_API_HTTP_METHOD: - setattr( - self.session, method, - functools.partial( - getattr(self.session, method), - timeout=timeout)) - - self.upload_checker = UploadingCheck() - - def _get_cookies(self, access_token: str): - """ - Get jar cookies for authentication from access_token. - - Args: - access_token (str): user access token on ModelScope. - - Returns: - jar (CookieJar): cookies for authentication. - """ - from requests.cookies import RequestsCookieJar - from urllib.parse import urlparse - - domain: str = urlparse(self.endpoint).netloc if self.endpoint else get_domain() - - jar = RequestsCookieJar() - jar.set('m_session_id', - access_token, - domain=domain, - path='/') - return jar - - def get_cookies(self, access_token: Optional[str] = None, cookies_required: Optional[bool] = False): - """ - Get cookies for authentication from local cache or access_token. - - Args: - access_token (Optional[str]): user access token on ModelScope. If not provided, try to get from local cache. - cookies_required (bool): whether to raise error if no cookies found, defaults to `False`. - - Returns: - cookies (CookieJar): cookies for authentication. - - Raises: - ValueError: If no credentials found and cookies_required is True. - """ - token = access_token or self.token or os.environ.get('MODELSCOPE_API_TOKEN') - if token: - cookies = self._get_cookies(access_token=token) - else: - cookies = ModelScopeConfig.get_cookies() - - if cookies is None and cookies_required: - raise ValueError( - 'No credentials found.' - 'You can pass the `--token` argument, ' - 'or use HubApi().login(access_token=`your_sdk_token`). ' - 'Your token is available at https://modelscope.cn/my/myaccesstoken' - ) - - return cookies - - def login( - self, - access_token: Optional[str] = None, - endpoint: Optional[str] = None - ): - """Login with your SDK access token, which can be obtained from - https://www.modelscope.cn user center. - - Args: - access_token (str): user access token on modelscope, set this argument or set `MODELSCOPE_API_TOKEN`. - If neither of the tokens exist, login will directly return. - endpoint: the endpoint to use, default to None to use endpoint specified in the class - - Returns: - cookies: to authenticate yourself to ModelScope open-api - git_token: token to access your git repository. - - Note: - You only have to login once within 30 days. - """ - access_token = access_token or self.token or os.environ.get('MODELSCOPE_API_TOKEN') - if not access_token: - return None, None - if not endpoint: - endpoint = self.endpoint - path = f'{endpoint}/api/v1/login' - r = self.session.post( - path, - json={'AccessToken': access_token}, - headers=self.builder_headers(self.headers)) - raise_for_http_status(r) - d = r.json() - raise_on_error(d) - - token = d[API_RESPONSE_FIELD_DATA][API_RESPONSE_FIELD_GIT_ACCESS_TOKEN] - cookies = r.cookies - - # save token and cookie - ModelScopeConfig.save_token(token) - ModelScopeConfig.save_cookies(cookies) - ModelScopeConfig.save_user_info( - d[API_RESPONSE_FIELD_DATA][API_RESPONSE_FIELD_USERNAME], - d[API_RESPONSE_FIELD_DATA][API_RESPONSE_FIELD_EMAIL]) - - return d[API_RESPONSE_FIELD_DATA][ - API_RESPONSE_FIELD_GIT_ACCESS_TOKEN], cookies - - def create_model(self, - model_id: str, - visibility: Optional[int] = ModelVisibility.PUBLIC, - license: Optional[str] = Licenses.APACHE_V2, - chinese_name: Optional[str] = None, - original_model_id: Optional[str] = '', - endpoint: Optional[str] = None, - token: Optional[str] = None, - aigc_model: Optional['AigcModel'] = None, - gated_mode: Optional[bool] = None) -> str: - """Create model repo at ModelScope Hub. - - Args: - model_id (str): The model id in format {owner}/{name} - visibility (int, optional): visibility of the model(1-private, 5-public), default 5. - license (str, optional): license of the model, default apache-2.0. - chinese_name (str, optional): chinese name of the model. - original_model_id (str, optional): the base model id which this model is trained from - endpoint: the endpoint to use, default to None to use endpoint specified in the class - token (str, optional): access token for authentication - aigc_model (AigcModel, optional): AigcModel instance for AIGC model creation. - If provided, will create an AIGC model with automatic file upload. - Refer to modelscope.hub.utils.aigc.AigcModel for details. - gated_mode (bool, optional): Gated mode for private repos. - True = gated (application-based download), False = off (normal private). - Only effective when visibility is PRIVATE (1). - - Returns: - str: URL of the created model repository - - Raises: - InvalidParameter: If model_id is invalid or required AIGC parameters are missing. - ValueError: If not login. - - Note: - model_id = {owner}/{name} - """ - if model_id is None: - raise InvalidParameter('model_id is required!') - # Get cookies for authentication. - cookies = self.get_cookies(access_token=token, cookies_required=True) - if not endpoint: - endpoint = self.endpoint - - owner_or_group, name = model_id_to_group_owner_name(model_id) - - # Base body configuration - body = { - 'Path': owner_or_group, - 'Name': name, - 'ChineseName': chinese_name, - 'Visibility': visibility, - 'License': license, - 'OriginalModelId': original_model_id, - 'TrainId': os.environ.get('MODELSCOPE_TRAIN_ID', '') - } - - if gated_mode is not None: - if visibility != ModelVisibility.PRIVATE: - logger.warning('gated_mode is only effective when visibility is PRIVATE, ignored.') - else: - body['ProtectedMode'] = 1 if gated_mode else 2 - - # Set path based on model type - if aigc_model is not None: - # Use AIGC model endpoint - path = f'{endpoint}/api/v1/models/aigc' - # Best-effort pre-upload weights so server recognizes sha256 (use existing cookies) - aigc_model.preupload_weights(cookies=cookies, headers=self.builder_headers(self.headers), endpoint=endpoint) - - # Add AIGC-specific fields to body - body.update({ - 'TagShowName': aigc_model.tag, - 'CoverImages': aigc_model.cover_images, - 'AigcType': aigc_model.aigc_type, - 'TagDescription': aigc_model.description, - 'VisionFoundation': aigc_model.base_model_type, - 'BaseModel': aigc_model.base_model_id or original_model_id, - 'WeightsName': aigc_model.weight_filename, - 'WeightsSha256': aigc_model.weight_sha256, - 'WeightsSize': aigc_model.weight_size, - 'ModelPath': aigc_model.model_path, - 'TriggerWords': aigc_model.trigger_words, - 'ModelSource': aigc_model.model_source, - 'SubVisionFoundation': aigc_model.base_model_sub_type, - }) - - if aigc_model.official_tags: - body['OfficialTags'] = aigc_model.official_tags - - else: - # Use regular model endpoint - path = f'{endpoint}/api/v1/models' - - headers = self.builder_headers(self.headers) - - intl_end = DEFAULT_MODELSCOPE_INTL_DOMAIN.split('.')[-1] - if endpoint.rstrip('/').endswith(f'.{intl_end}'): - headers['X-Modelscope-Accept-Language'] = 'en_US' - r = self.session.post( - path, - json=body, - cookies=cookies, - headers=headers) - raise_for_http_status(r) - d = r.json() - raise_on_error(d) - model_repo_url = f'{endpoint}/models/{model_id}' - - # Upload model files for AIGC models - if aigc_model is not None: - aigc_model.upload_to_repo(self, model_id, token) - - return model_repo_url - - def create_model_tag(self, - model_id: str, - tag_name: str, - endpoint: Optional[str] = None, - token: Optional[str] = None, - aigc_model: Optional['AigcModel'] = None) -> str: - """Create a tag for a model at ModelScope Hub. - - Args: - model_id (str): The model id in format {owner}/{name} - tag_name (str): The tag name (e.g., "v1.0.0") - endpoint: the endpoint to use, default to None to use endpoint specified in the class - token (str, optional): access token for authentication - aigc_model (AigcModel, optional): AigcModel instance for AIGC model tag creation. - If provided, will create an AIGC model tag with automatic parameters. - Refer to modelscope.hub.utils.aigc.AigcModel for details. - - Returns: - str: URL of the created tag - - Raises: - InvalidParameter: If model_id, tag_name, ref, or description is invalid. - ValueError: If not login. - - Note: - model_id = {owner}/{name} - """ - if model_id is None: - raise InvalidParameter('model_id is required!') - if tag_name is None: - raise InvalidParameter('tag_name is required!') - if tag_name.lower() in ['main', 'master']: - raise InvalidParameter( - f'tag_name "{tag_name}" is not allowed. ' - f'Please use a different tag name (e.g., "v1.0", "v1.1", "latest"). ' - f'Reserved names: main, master' - ) - - # Get cookies for authentication. - cookies = self.get_cookies(access_token=token, cookies_required=True) - if not endpoint: - endpoint = self.endpoint - - owner_or_group, name = model_id_to_group_owner_name(model_id) - - # Set path and body based on model type - if aigc_model is not None: - # Use AIGC model tag endpoint - path = f'{endpoint}/api/v1/models/aigc/repo/tag' - aigc_model.preupload_weights(cookies=cookies, headers=self.builder_headers(self.headers), endpoint=endpoint) - - # Base body for AIGC model tag - body = { - 'CoverImages': aigc_model.cover_images, - 'Name': name, - 'Path': owner_or_group, - 'TagShowName': tag_name, - 'WeightsName': aigc_model.weight_filename, - 'WeightsSha256': aigc_model.weight_sha256, - 'WeightsSize': aigc_model.weight_size, - 'TriggerWords': aigc_model.trigger_words, - 'AigcType': aigc_model.aigc_type, - 'VisionFoundation': aigc_model.base_model_type - } - - else: - # Use regular model tag endpoint - path = f'{endpoint}/api/v1/models/{model_id}/repo/tag' - revision = 'master' - body = { - 'TagName': tag_name, - 'Ref': revision - } - - tag_timeout = (API_HTTP_CLIENT_CONNECT_TIMEOUT, - API_HTTP_CLIENT_TIMEOUT) - - retryable_status = {500, 502, 503, 504} - attempts = max(1, CREATE_TAG_MAX_RETRIES) - r = None - for attempt in range(1, attempts + 1): - retry_reason = None - try: - r = self.session.post( - path, - json=body, - cookies=cookies, - headers=self.builder_headers(self.headers), - timeout=tag_timeout) - except (requests.exceptions.ReadTimeout, - requests.exceptions.ConnectTimeout, - requests.exceptions.ConnectionError) as e: - if attempt >= attempts: - logger.error( - f'create_model_tag POST failed after {attempts} ' - f'attempt(s) due to transient network error: {e}. ' - f'Consider raising MODELSCOPE_API_HTTP_CLIENT_TIMEOUT ' - f'(current={API_HTTP_CLIENT_TIMEOUT}s) or ' - f'MODELSCOPE_CREATE_TAG_MAX_RETRIES ' - f'(current={CREATE_TAG_MAX_RETRIES}).') - raise - retry_reason = f'{type(e).__name__}: {e}' - else: - if r.status_code in retryable_status and attempt < attempts: - retry_reason = (f'retryable HTTP {r.status_code} ' - f'from server') - else: - break - - sleep_s = CREATE_TAG_RETRY_BACKOFF * (2 ** (attempt - 1)) - logger.warning( - f'create_model_tag POST attempt {attempt}/{attempts} ' - f'failed with {retry_reason}. Retrying in {sleep_s}s...') - time.sleep(sleep_s) - - raise_for_http_status(r) - d = r.json() - raise_on_error(d) - - tag_url = f'{endpoint}/models/{model_id}/tags/{tag_name}' - return tag_url - - def delete_model(self, model_id: str, endpoint: Optional[str] = None, token: Optional[str] = None): - """ - @deprecated - Delete model_id from ModelScope. - - Args: - model_id (str): The model id. - endpoint: the endpoint to use, default to None to use endpoint specified in the class - token (str, optional): access token for authentication - - Raises: - ValueError: If not login. - - Note: - model_id = {owner}/{name} - """ - warnings.warn( - 'This function is deprecated due to security reasons, ' - 'and will be recovered in future versions with proper token authentication. ', - DeprecationWarning, - stacklevel=2 - ) - cookies = self.get_cookies(access_token=token, cookies_required=True) - if not endpoint: - endpoint = self.endpoint - if cookies is None: - raise ValueError('Token does not exist, please login first.') - path = f'{endpoint}/api/v1/models/{model_id}' - - r = self.session.delete(path, - cookies=cookies, - headers=self.builder_headers(self.headers)) - raise_for_http_status(r) - raise_on_error(r.json()) - - def get_model_url(self, model_id: str, endpoint: Optional[str] = None): - if not endpoint: - endpoint = self.endpoint - return f'{endpoint}/api/v1/models/{model_id}.git' - - def get_model( - self, - model_id: str, - revision: Optional[str] = DEFAULT_MODEL_REVISION, - endpoint: Optional[str] = None, - token: Optional[str] = None, + # ------------------------------------------------------------------ + # Internal transport helper + # ------------------------------------------------------------------ + def _legacy_request( + self, + method: str, + path: str, + **kwargs: Any, ) -> dict: - """Get model information at ModelScope + """Send a request via the legacy client and validate the response envelope. - Args: - model_id (str): The model id. - revision (str optional): revision of model. - endpoint: the endpoint to use, default to None to use endpoint specified in the class - token (str, optional): access token for authentication - - Returns: - The model detail information. - - Raises: - NotExistError: If the model is not exist, will throw NotExistError - - Note: - model_id = {owner}/{name} + Combines ``legacy._request`` (HTTP-level error handling) with + application-level ``{"Code": 200, ...}`` envelope validation. + Returns the parsed JSON body dict on success. """ - cookies = self.get_cookies(access_token=token, cookies_required=False) - owner_or_group, name = model_id_to_group_owner_name(model_id) - if not endpoint: - endpoint = self.endpoint + from modelscope.hub.errors import raise_on_error + resp = self._api.legacy._request(method, path, **kwargs) + body = resp.json() + raise_on_error(body) + return body - if revision: - path = f'{endpoint}/api/v1/models/{owner_or_group}/{name}?Revision={revision}' - else: - path = f'{endpoint}/api/v1/models/{owner_or_group}/{name}' - - r = self.session.get(path, cookies=cookies, - headers=self.builder_headers(self.headers)) - handle_http_response(r, logger, cookies, model_id) - if r.status_code == HTTPStatus.OK: - if is_ok(r.json()): - return r.json()[API_RESPONSE_FIELD_DATA] - else: - raise NotExistError(r.json()[API_RESPONSE_FIELD_MESSAGE]) - else: - raise_for_http_status(r) - - def get_endpoint_for_read(self, - repo_id: str, - *, - repo_type: Optional[str] = None, - token: Optional[str] = None) -> str: - """Get proper endpoint for read operation (such as download, list etc.) - 1. If user has set MODELSCOPE_DOMAIN, construct endpoint with user-specified domain. - If the repo does not exist on that endpoint, throw 404 error, otherwise return the endpoint. - 2. If domain is not set, check existence of repo in cn-site and ai-site (intl version) respectively. - Checking order is determined by MODELSCOPE_PREFER_AI_SITE. - a. if MODELSCOPE_PREFER_AI_SITE is not set ,check cn-site first before ai-site (intl version) - b. otherwise check ai-site before cn-site - return the endpoint with which the given repo_id exists. - if neither exists, throw 404 error - """ - s = os.environ.get(MODELSCOPE_DOMAIN) - if s is not None and s.strip() != '': - endpoint = MODELSCOPE_URL_SCHEME + s - try: - self.repo_exists(repo_id=repo_id, repo_type=repo_type, endpoint=endpoint, re_raise=True, token=token) - except Exception: - logger.error(f'Repo {repo_id} does not exist on {endpoint}.') - raise - return endpoint - - check_cn_first = not is_env_true(MODELSCOPE_PREFER_AI_SITE) - prefer_endpoint = get_endpoint(cn_site=check_cn_first) - if not self.repo_exists( - repo_id, repo_type=repo_type, endpoint=prefer_endpoint, token=token): - alternative_endpoint = get_endpoint(cn_site=(not check_cn_first)) - logger.warning(f'Repo {repo_id} not exists on {prefer_endpoint}, ' - f'will try on alternative endpoint {alternative_endpoint}.') - try: - self.repo_exists( - repo_id, repo_type=repo_type, endpoint=alternative_endpoint, re_raise=True, token=token) - except Exception: - logger.error(f'Repo {repo_id} not exists on either {prefer_endpoint} or {alternative_endpoint}') - raise - else: - return alternative_endpoint - else: - return prefer_endpoint - - def model_info(self, - repo_id: str, - *, - revision: Optional[str] = DEFAULT_MODEL_REVISION, - endpoint: Optional[str] = None) -> ModelInfo: - """Get model information including commit history. - - Args: - repo_id (str): The model id in the format of - ``namespace/model_name``. - revision (str, optional): Specific revision of the model. - Defaults to ``DEFAULT_MODEL_REVISION``. - endpoint (str, optional): Hub endpoint to use. When ``None``, - use the endpoint specified when initializing :class:`HubApi`. - - Returns: - ModelInfo: The model detailed information returned by - ModelScope Hub with commit history. - """ - owner_or_group, _ = model_id_to_group_owner_name(repo_id) - model_data = self.get_model( - model_id=repo_id, revision=revision, endpoint=endpoint) - commits = self.list_repo_commits( - repo_id=repo_id, repo_type=REPO_TYPE_MODEL, revision=revision, endpoint=endpoint) - siblings = self.get_model_files( - model_id=repo_id, revision=revision, recursive=True, endpoint=endpoint) - - # Create ModelInfo from API response data - model_info = ModelInfo(**model_data, commits=commits, author=owner_or_group, siblings=siblings) - - return model_info - - def dataset_info(self, - repo_id: str, - *, - revision: Optional[str] = None, - endpoint: Optional[str] = None) -> DatasetInfo: - """Get dataset information including commit history. - - Args: - repo_id (str): The dataset id in the format of - ``namespace/dataset_name``. - revision (str, optional): Specific revision of the dataset. - Defaults to ``None``. - endpoint (str, optional): Hub endpoint to use. When ``None``, - use the endpoint specified when initializing :class:`HubApi`. - - Returns: - DatasetInfo: The dataset detailed information returned by - ModelScope Hub with commit history. - """ - owner_or_group, _ = model_id_to_group_owner_name(repo_id) - dataset_data = self.get_dataset( - dataset_id=repo_id, revision=revision, endpoint=endpoint) - commits = self.list_repo_commits( - repo_id=repo_id, repo_type=REPO_TYPE_DATASET, revision=revision, endpoint=endpoint) - siblings = self.get_dataset_files( - repo_id=repo_id, revision=revision or DEFAULT_DATASET_REVISION, recursive=True, endpoint=endpoint) - - # Create DatasetInfo from API response data - dataset_info = DatasetInfo(**dataset_data, commits=commits, author=owner_or_group, siblings=siblings) - - return dataset_info - - def repo_info( + # ------------------------------------------------------------------ + # OSS dataset operations + # ------------------------------------------------------------------ + def list_oss_dataset_objects( self, - repo_id: str, - *, - repo_type: Optional[str] = REPO_TYPE_MODEL, - revision: Optional[str] = DEFAULT_MODEL_REVISION, - endpoint: Optional[str] = None - ) -> Union[ModelInfo, DatasetInfo]: - """Get repository information for models or datasets. - - Args: - repo_id (str): The repository id in the format of - ``namespace/repo_name``. - revision (str, optional): Specific revision of the repository. - Currently only effective for model repositories. Defaults to - ``DEFAULT_MODEL_REVISION``. - repo_type (str, optional): Type of the repository. Supported - values are ``"model"`` and ``"dataset"``. If not provided, - ``"model"`` is assumed. - endpoint (str, optional): Hub endpoint to use. When ``None``, - use the endpoint specified when initializing :class:`HubApi`. - - Returns: - Union[ModelInfo, DatasetInfo]: The repository detailed information - returned by ModelScope Hub. - """ - if repo_type is None or repo_type == REPO_TYPE_MODEL: - return self.model_info(repo_id=repo_id, revision=revision, endpoint=endpoint) - - if repo_type == REPO_TYPE_DATASET: - return self.dataset_info(repo_id=repo_id, revision=revision, endpoint=endpoint) - - if repo_type == REPO_TYPE_STUDIO: - if (repo_id is None) or repo_id.count('/') != 1: - raise InvalidParameter( - f'Invalid repo_id: {repo_id}, must be of format owner/repo_name') - _endpoint = endpoint or self.endpoint - owner, name = repo_id.split('/', 1) - path = f'{_endpoint}/openapi/v1/studios/{owner}/{name}' - headers = self._build_bearer_headers(token=None, token_required=False) - r = self.session.get(path, headers=headers) - handle_http_response(r, logger, None, repo_id) - return r.json().get('data', {}) - - raise InvalidParameter( - f'Arg repo_type {repo_type} not supported. Please choose from {REPO_TYPE_SUPPORT}.') - - def repo_exists( - self, - repo_id: str, - *, - repo_type: Optional[str] = None, - endpoint: Optional[str] = None, - re_raise: Optional[bool] = False, - token: Optional[str] = None - ) -> bool: - """ - Checks if a repository exists on ModelScope - - Args: - repo_id (`str`): - A namespace (user or an organization) and a repo name separated - by a `/`. - repo_type (`str`, *optional*): - `None` or `"model"` if getting repository info from a model. Default is `None`. - Supported values are `"model"`, `"dataset"` and `"studio"`. - endpoint(`str`): - None or specific endpoint to use, when None, use the default endpoint - set in HubApi class (self.endpoint) - re_raise(`bool`): - raise exception when error - token (`str`, *optional*): access token to use for checking existence. - Returns: - True if the repository exists, False otherwise. - """ - if endpoint is None: - endpoint = self.endpoint - if (repo_type is not None) and repo_type.lower() not in REPO_TYPE_SUPPORT: - raise Exception('Not support repo-type: %s' % repo_type) - if (repo_id is None) or repo_id.count('/') != 1: - raise Exception('Invalid repo_id: %s, must be of format namespace/name' % repo_type) - - cookies = self.get_cookies(access_token=token, cookies_required=False) - owner_or_group, name = model_id_to_group_owner_name(repo_id) - if (repo_type is not None) and repo_type.lower() == REPO_TYPE_STUDIO: - path = f'{endpoint}/openapi/v1/studios/{owner_or_group}/{name}' - headers = self._build_bearer_headers(token=token, token_required=False) - r = self.session.get(path, headers=headers) - elif (repo_type is not None) and repo_type.lower() == REPO_TYPE_DATASET: - path = f'{endpoint}/api/v1/datasets/{owner_or_group}/{name}' - r = self.session.get(path, cookies=cookies, - headers=self.builder_headers(self.headers)) - else: - path = f'{endpoint}/api/v1/models/{owner_or_group}/{name}' - r = self.session.get(path, cookies=cookies, - headers=self.builder_headers(self.headers)) - code = handle_http_response(r, logger, cookies, repo_id, False) - if code == 200: - return True - elif code == 404: - if re_raise: - raise HTTPError(r) - else: - return False - else: - logger.warn(f'Check repo_exists return status code {code}.') - raise Exception( - 'Failed to check existence of repo: %s, make sure you have access authorization.' - % repo_type) - - def delete_repo(self, - repo_id: str, - repo_type: str, - endpoint: Optional[str] = None, - token: Optional[str] = None - ): - """ - @deprecated - Delete a repository from ModelScope. - - Args: - repo_id (`str`): - A namespace (user or an organization) and a repo name separated - by a `/`. - repo_type (`str`): - The type of the repository. Supported types are `model`, `dataset` and `studio`. - endpoint(`str`): - The endpoint to use. If not provided, the default endpoint is `https://www.modelscope.cn` - Could be set to `https://ai.modelscope.ai` for international version. - token (str): Access token of the ModelScope. - """ - if not endpoint: - endpoint = self.endpoint - - if repo_type == REPO_TYPE_DATASET: - self.delete_dataset( - dataset_id=repo_id, - endpoint=endpoint, - token=token - ) - elif repo_type == REPO_TYPE_MODEL: - self.delete_model( - model_id=repo_id, - endpoint=endpoint, - token=token) - elif repo_type == REPO_TYPE_STUDIO: - logger.warning( - f'Deleting an entire studio repo ({repo_id}) is not supported ' - f'via the OpenAPI for security reasons. ' - f'To delete studio environment variables, use ' - f'HubApi.delete_studio_secret(studio_id, key). ' - f'To delete the studio itself, please use the web console.') - return - else: - raise Exception(f'Arg repo_type {repo_type} not supported.') - - logger.info(f'Repo {repo_id} deleted successfully.') - - @staticmethod - def _create_default_config(model_dir): - cfg_file = os.path.join(model_dir, ModelFile.CONFIGURATION) - cfg = { - ConfigFields.framework: Frameworks.torch, - ConfigFields.task: Tasks.other, - } - with open(cfg_file, 'w') as file: - json.dump(cfg, file) - - def push_model(self, - model_id: str, - model_dir: str, - visibility: Optional[int] = ModelVisibility.PUBLIC, - license: Optional[str] = Licenses.APACHE_V2, - chinese_name: Optional[str] = None, - commit_message: Optional[str] = 'upload model', - tag: Optional[str] = None, - revision: Optional[str] = DEFAULT_REPOSITORY_REVISION, - original_model_id: Optional[str] = None, - ignore_file_pattern: Optional[Union[List[str], str]] = None, - lfs_suffix: Optional[Union[str, List[str]]] = None, - token: Optional[str] = None): - warnings.warn( - 'This function is deprecated and will be removed in future versions. ' - 'Please use git command directly or use HubApi().upload_folder instead', - DeprecationWarning, - stacklevel=2 - ) - """Upload model from a given directory to given repository. A valid model directory - must contain a configuration.json file. - - This function upload the files in given directory to given repository. If the - given repository is not exists in remote, it will automatically create it with - given visibility, license and chinese_name parameters. If the revision is also - not exists in remote repository, it will create a new branch for it. - - This function must be called before calling HubApi's login with a valid token - which can be obtained from ModelScope's website. - - If any error, please upload via git commands. - - Args: - model_id (str): - The model id to be uploaded, caller must have write permission for it. - model_dir(str): - The Absolute Path of the finetune result. - visibility(int, optional): - Visibility of the new created model(1-private, 5-public). If the model is - not exists in ModelScope, this function will create a new model with this - visibility and this parameter is required. You can ignore this parameter - if you make sure the model's existence. - license(`str`, defaults to `None`): - License of the new created model(see License). If the model is not exists - in ModelScope, this function will create a new model with this license - and this parameter is required. You can ignore this parameter if you - make sure the model's existence. - chinese_name(`str`, *optional*, defaults to `None`): - chinese name of the new created model. - commit_message(`str`, *optional*, defaults to `None`): - commit message of the push request. - tag(`str`, *optional*, defaults to `None`): - The tag on this commit - revision (`str`, *optional*, default to DEFAULT_MODEL_REVISION): - which branch to push. If the branch is not exists, It will create a new - branch and push to it. - original_model_id (str, optional): The base model id which this model is trained from - ignore_file_pattern (`Union[List[str], str]`, optional): The file pattern to ignore uploading - lfs_suffix (`List[str]`, optional): File types to use LFS to manage. examples: '*.safetensors'. - - Raises: - InvalidParameter: Parameter invalid. - NotLoginException: Not login - ValueError: No configuration.json - Exception: Create failed. - """ - if model_id is None: - raise InvalidParameter('model_id cannot be empty!') - if model_dir is None: - raise InvalidParameter('model_dir cannot be empty!') - if not os.path.exists(model_dir) or os.path.isfile(model_dir): - raise InvalidParameter('model_dir must be a valid directory.') - cfg_file = os.path.join(model_dir, ModelFile.CONFIGURATION) - if not os.path.exists(cfg_file): - logger.warning( - f'No {ModelFile.CONFIGURATION} file found in {model_dir}, creating a default one.') - HubApi._create_default_config(model_dir) - - cookies = self.get_cookies(access_token=token, cookies_required=True) - if cookies is None: - raise NotLoginException('Must login before upload!') - files_to_save = os.listdir(model_dir) - folder_size = get_readable_folder_size(model_dir) - if ignore_file_pattern is None: - ignore_file_pattern = [] - if isinstance(ignore_file_pattern, str): - ignore_file_pattern = [ignore_file_pattern] - if visibility is None or license is None: - raise InvalidParameter('Visibility and License cannot be empty for new model.') - if not self.repo_exists(model_id, token=token): - logger.info(f'Creating new model [{model_id}]') - self.create_model( - model_id=model_id, - visibility=visibility, - license=license, - chinese_name=chinese_name, - original_model_id=original_model_id, - token=token, - endpoint=self.endpoint) - tmp_dir = os.path.join(model_dir, TEMPORARY_FOLDER_NAME) # make temporary folder - git_wrapper = GitCommandWrapper() - logger.info(f'Pushing folder {model_dir} as model {model_id}.') - logger.info(f'Total folder size {folder_size}, this may take a while depending on actual pushing size...') - try: - repo = Repository(model_dir=tmp_dir, clone_from=model_id, auth_token=token, endpoint=self.endpoint) - branches = git_wrapper.get_remote_branches(tmp_dir) - if revision not in branches: - logger.info(f'Creating new branch {revision}') - git_wrapper.new_branch(tmp_dir, revision) - git_wrapper.checkout(tmp_dir, revision) - files_in_repo = os.listdir(tmp_dir) - for f in files_in_repo: - if f[0] != '.': - src = os.path.join(tmp_dir, f) - if os.path.isfile(src): - os.remove(src) - else: - shutil.rmtree(src, ignore_errors=True) - for f in files_to_save: - if f[0] != '.': - if any([re.search(pattern, f) is not None for pattern in ignore_file_pattern]): - continue - src = os.path.join(model_dir, f) - if os.path.isdir(src): - shutil.copytree(src, os.path.join(tmp_dir, f)) - else: - shutil.copy(src, tmp_dir) - if not commit_message: - date = datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S') - commit_message = '[automsg] push model %s to hub at %s' % ( - model_id, date) - if lfs_suffix is not None: - lfs_suffix_list = [lfs_suffix] if isinstance(lfs_suffix, str) else lfs_suffix - for suffix in lfs_suffix_list: - repo.add_lfs_type(suffix) - repo.push( - commit_message=commit_message, - local_branch=revision, - remote_branch=revision) - if tag is not None: - repo.tag_and_push(tag, tag) - logger.info(f'Successfully push folder {model_dir} to remote repo [{model_id}].') - except Exception: - raise - finally: - shutil.rmtree(tmp_dir, ignore_errors=True) - - def list_models(self, - owner_or_group: str, - page_number: Optional[int] = 1, - page_size: Optional[int] = 10, - endpoint: Optional[str] = None, - token: Optional[str] = None) -> dict: - """List models in owner or group. - - Args: - owner_or_group(str): owner or group. - page_number(int, optional): The page number, default: 1 - page_size(int, optional): The page size, default: 10 - endpoint: the endpoint to use, default to None to use endpoint specified in the class - token (str, optional): access token for authentication - - Raises: - RequestError: The request error. - - Returns: - dict: {"models": "list of models", "TotalCount": total_number_of_models_in_owner_or_group} - """ - cookies = self.get_cookies(access_token=token, cookies_required=False) - if not endpoint: - endpoint = self.endpoint - path = f'{endpoint}/api/v1/models/' - r = self.session.put( - path, - data='{"Path":"%s", "PageNumber":%s, "PageSize": %s}' % - (owner_or_group, page_number, page_size), - cookies=cookies, - headers=self.builder_headers(self.headers)) - handle_http_response(r, logger, cookies, owner_or_group) - if r.status_code == HTTPStatus.OK: - if is_ok(r.json()): - data = r.json()[API_RESPONSE_FIELD_DATA] - return data - else: - raise RequestError(r.json()[API_RESPONSE_FIELD_MESSAGE]) - else: - raise_for_http_status(r) - return None - - def list_datasets(self, - owner_or_group: str, - *, - page_number: Optional[int] = 1, - page_size: Optional[int] = 10, - sort: Optional[str] = None, - search: Optional[str] = None, - endpoint: Optional[str] = None, - token: Optional[str] = None) -> dict: - """List datasets via OpenAPI with pagination, filtering and sorting. - - Args: - owner_or_group (str): Search by dataset authors (including organizations and individuals). - page_number (int, optional): The page number. Defaults to 1. - page_size (int, optional): The page size. Defaults to 10. - sort (str, optional): Sort key. If not provided, the server's default sorting is used. - choose from ['default', 'downloads', 'likes', 'last_modified']. - search (str, optional): Search by substring keywords in the dataset's Chinese name, - English name, and authors (including organizations and individuals). - endpoint (str, optional): Hub endpoint to use. When None, use the endpoint specified in the class. - token (str, optional): Access token for authentication. - - Returns: - dict: The OpenAPI data payload, e.g. - { - "datasets": [...], - "total_count": int, - "page_number": int, - "page_size": int - } - """ - if not endpoint: - endpoint = self.endpoint - path = f'{endpoint}/openapi/v1/datasets' - - # Build query params - params: Dict[str, Any] = { - 'page_number': page_number, - 'page_size': page_size, - } - if sort: - if sort not in VALID_SORT_KEYS: - raise InvalidParameter( - f'Invalid sort key: {sort}. Supported sort keys: {list(VALID_SORT_KEYS)}') - params['sort'] = sort - if search: - params['search'] = search - if owner_or_group: - params['author'] = owner_or_group - - headers = self._build_bearer_headers(token=token, token_required=False) - - r = self.session.get(path, params=params, headers=headers) - raise_for_http_status(r) - return self._parse_openapi_response(r) - - def _check_cookie(self, use_cookies: Union[bool, CookieJar] = False) -> CookieJar: # noqa - cookies = None - if isinstance(use_cookies, CookieJar): - cookies = use_cookies - elif isinstance(use_cookies, bool): - cookies = self.get_cookies(cookies_required=use_cookies) - return cookies - - def list_model_revisions( - self, - model_id: str, - cutoff_timestamp: Optional[int] = None, - use_cookies: Union[bool, CookieJar] = False) -> List[str]: - """Get model branch and tags. - - Args: - model_id (str): The model id - cutoff_timestamp (int): Tags created before the cutoff will be included. - The timestamp is represented by the seconds elapsed from the epoch time. - use_cookies (Union[bool, CookieJar], optional): If is cookieJar, we will use this cookie, if True, - will load cookie from local. Defaults to False. - - Returns: - Tuple[List[str], List[str]]: Return list of branch name and tags - """ - tags_details = self.list_model_revisions_detail(model_id=model_id, - cutoff_timestamp=cutoff_timestamp, - use_cookies=use_cookies) - tags = [x['Revision'] for x in tags_details - ] if tags_details else [] - return tags - - def list_model_revisions_detail( - self, - model_id: str, - cutoff_timestamp: Optional[int] = None, - use_cookies: Union[bool, CookieJar] = False, - endpoint: Optional[str] = None) -> List[str]: - """Get model branch and tags. - - Args: - model_id (str): The model id - cutoff_timestamp (int): Tags created before the cutoff will be included. - The timestamp is represented by the seconds elapsed from the epoch time. - use_cookies (Union[bool, CookieJar], optional): If is cookieJar, we will use this cookie, if True, - will load cookie from local. Defaults to False. - endpoint: the endpoint to use, default to None to use endpoint specified in the class - - Returns: - Tuple[List[str], List[str]]: Return list of branch name and tags - """ - cookies = self._check_cookie(use_cookies) - if cutoff_timestamp is None: - cutoff_timestamp = get_release_datetime() - if not endpoint: - endpoint = self.endpoint - path = f'{endpoint}/api/v1/models/{model_id}/revisions?EndTime=%s' % cutoff_timestamp - r = self.session.get(path, cookies=cookies, - headers=self.builder_headers(self.headers)) - handle_http_response(r, logger, cookies, model_id) - d = r.json() - raise_on_error(d) - info = d[API_RESPONSE_FIELD_DATA] - # tags returned from backend are guaranteed to be ordered by create-time - return info['RevisionMap']['Tags'] - - def get_branch_tag_detail(self, details, name): - for item in details: - if item['Revision'] == name: - return item - return None - - def get_valid_revision_detail(self, - model_id: str, - revision=None, - cookies: Optional[CookieJar] = None, - endpoint: Optional[str] = None): - if not endpoint: - endpoint = self.endpoint - release_timestamp = get_release_datetime() - current_timestamp = int(round(datetime.datetime.now().timestamp())) - # for active development in library codes (non-release-branches), release_timestamp - # is set to be a far-away-time-in-the-future, to ensure that we shall - # get the master-HEAD version from model repo by default (when no revision is provided) - all_branches_detail, all_tags_detail = self.get_model_branches_and_tags_details( - model_id, use_cookies=False if cookies is None else cookies, endpoint=endpoint) - all_branches = [x['Revision'] for x in all_branches_detail] if all_branches_detail else [] - all_tags = [x['Revision'] for x in all_tags_detail] if all_tags_detail else [] - if release_timestamp > current_timestamp + ONE_YEAR_SECONDS: - if revision is None: - revision = MASTER_MODEL_BRANCH - logger.info( - f'Model revision not specified, using default [{revision}] version.') - if revision not in all_branches and revision not in all_tags: - raise NotExistError('The model: %s has no revision : %s .' % (model_id, revision)) - - revision_detail = self.get_branch_tag_detail(all_tags_detail, revision) - if revision_detail is None: - revision_detail = self.get_branch_tag_detail(all_branches_detail, revision) - logger.debug(f'Development mode use revision: {revision}') - else: - if revision is not None and revision in all_branches: - revision_detail = self.get_branch_tag_detail(all_branches_detail, revision) - return revision_detail - - if len(all_tags_detail) == 0: # use no revision use master as default. - if revision is None or revision == MASTER_MODEL_BRANCH: - revision = MASTER_MODEL_BRANCH - else: - raise NotExistError('The model: %s has no revision: %s !' % (model_id, revision)) - revision_detail = self.get_branch_tag_detail(all_branches_detail, revision) - else: - if revision is None: # user not specified revision, use latest revision before release time - revisions_detail = [x for x in - all_tags_detail if - x['CreatedAt'] <= release_timestamp] if all_tags_detail else [] # noqa E501 - if len(revisions_detail) > 0: - revision = revisions_detail[0]['Revision'] # use latest revision before release time. - revision_detail = revisions_detail[0] - else: - revision = MASTER_MODEL_BRANCH - revision_detail = self.get_branch_tag_detail(all_branches_detail, revision) - vl = '[%s]' % ','.join(all_tags) - logger.warning(f'Model revision should be specified from revisions: {vl}') - logger.warning(f'Model revision not specified, use revision: {revision}') - else: - # use user-specified revision - if revision not in all_tags: - if revision == MASTER_MODEL_BRANCH: - logger.warning('Using the master branch is fragile, please use it with caution!') - revision_detail = self.get_branch_tag_detail(all_branches_detail, revision) - else: - vl = '[%s]' % ','.join(all_tags) - raise NotExistError('The model: %s has no revision: %s valid are: %s!' % - (model_id, revision, vl)) - else: - revision_detail = self.get_branch_tag_detail(all_tags_detail, revision) - logger.info(f'Use user-specified model revision: {revision}') - return revision_detail - - def get_valid_revision(self, - model_id: str, - revision=None, - cookies: Optional[CookieJar] = None, - endpoint: Optional[str] = None): - return self.get_valid_revision_detail(model_id=model_id, - revision=revision, - cookies=cookies, - endpoint=endpoint)['Revision'] - - def get_model_branches_and_tags_details( - self, - model_id: str, - use_cookies: Union[bool, CookieJar] = False, - endpoint: Optional[str] = None - ) -> Tuple[List[str], List[str]]: - """Get model branch and tags. - - Args: - model_id (str): The model id - use_cookies (Union[bool, CookieJar], optional): If is cookieJar, we will use this cookie, if True, - will load cookie from local. Defaults to False. - endpoint: the endpoint to use, default to None to use endpoint specified in the class - - Returns: - Tuple[List[str], List[str]]: Return list of branch name and tags - """ - cookies = self._check_cookie(use_cookies) - if not endpoint: - endpoint = self.endpoint - path = f'{endpoint}/api/v1/models/{model_id}/revisions' - r = self.session.get(path, cookies=cookies, - headers=self.builder_headers(self.headers)) - handle_http_response(r, logger, cookies, model_id) - d = r.json() - raise_on_error(d) - info = d[API_RESPONSE_FIELD_DATA] - return info['RevisionMap']['Branches'], info['RevisionMap']['Tags'] - - def get_model_branches_and_tags( - self, - model_id: str, - use_cookies: Union[bool, CookieJar] = False, - ) -> Tuple[List[str], List[str]]: - """Get model branch and tags. - - Args: - model_id (str): The model id - use_cookies (Union[bool, CookieJar], optional): If is cookieJar, we will use this cookie, if True, - will load cookie from local. Defaults to False. - - Returns: - Tuple[List[str], List[str]]: Return list of branch name and tags - """ - branches_detail, tags_detail = self.get_model_branches_and_tags_details(model_id=model_id, - use_cookies=use_cookies) - branches = [x['Revision'] for x in branches_detail - ] if branches_detail else [] - tags = [x['Revision'] for x in tags_detail - ] if tags_detail else [] - return branches, tags - - def get_model_files(self, - model_id: str, - revision: Optional[str] = DEFAULT_MODEL_REVISION, - root: Optional[str] = None, - recursive: Optional[bool] = False, - use_cookies: Union[bool, CookieJar] = False, - headers: Optional[dict] = {}, - endpoint: Optional[str] = None) -> List[dict]: - """List the models files. - - Args: - model_id (str): The model id - revision (Optional[str], optional): The branch or tag name. - root (Optional[str], optional): The root path. Defaults to None. - recursive (Optional[bool], optional): Is recursive list files. Defaults to False. - use_cookies (Union[bool, CookieJar], optional): If is cookieJar, we will use this cookie, if True, - will load cookie from local. Defaults to False. - headers: request headers - endpoint: the endpoint to use, default to None to use endpoint specified in the class - - Returns: - List[dict]: Model file list. - """ - if not endpoint: - endpoint = self.endpoint - if revision: - path = '%s/api/v1/models/%s/repo/files?Revision=%s&Recursive=%s' % ( - endpoint, model_id, revision, recursive) - else: - path = '%s/api/v1/models/%s/repo/files?Recursive=%s' % ( - endpoint, model_id, recursive) - cookies = self._check_cookie(use_cookies) - if root is not None: - path = path + f'&Root={root}' - headers = self.headers if headers is None else headers - headers['X-Request-ID'] = str(uuid.uuid4().hex) - r = self.session.get( - path, cookies=cookies, headers=headers) - - handle_http_response(r, logger, cookies, model_id) - d = r.json() - raise_on_error(d) - - files = [] - if not d[API_RESPONSE_FIELD_DATA]['Files']: - logger.warning(f'No files found in model {model_id} at revision {revision}.') - return files - for file in d[API_RESPONSE_FIELD_DATA]['Files']: - if file['Name'] == '.gitignore' or file['Name'] == '.gitattributes': - continue - - files.append(file) - return files - - def file_exists( - self, - repo_id: str, - filename: str, - *, - revision: Optional[str] = None, - token: Optional[str] = None, - ): - """Get if the specified file exists - - Args: - repo_id (`str`): The repo id to use - filename (`str`): The queried filename, if the file exists in a sub folder, - please pass / - revision (`Optional[str]`): The repo revision - token (`Optional[str]`): The access token - Returns: - The query result in bool value - """ - cookies = self.get_cookies(access_token=token) - files = self.get_model_files( - repo_id, - recursive=True, - revision=revision, - use_cookies=False if cookies is None else cookies, - ) - files = [file['Path'] for file in files] - return filename in files - - def create_dataset(self, - dataset_name: str, - namespace: str, - chinese_name: Optional[str] = '', - license: Optional[str] = Licenses.APACHE_V2, - visibility: Optional[int] = DatasetVisibility.PUBLIC, - description: Optional[str] = '', - endpoint: Optional[str] = None, - token: Optional[str] = None, - gated_mode: Optional[bool] = None) -> str: - """ - Create a dataset in ModelScope. - - Args: - dataset_name (str): The name of the dataset. - namespace (str): The namespace (user or organization) for the dataset. - chinese_name (str, optional): The Chinese name of the dataset. Defaults to ''. - license (str, optional): The license of the dataset. Defaults to Licenses.APACHE_V2. - visibility (int, optional): The visibility of the dataset. Defaults to DatasetVisibility.PUBLIC. - description (str, optional): The description of the dataset. Defaults to ''. - endpoint (str, optional): The endpoint to use. If not provided, the default endpoint is used. - token (str, optional): The access token for authentication. - gated_mode (bool, optional): Gated mode for private repos. - True = gated (application-based download), False = off (normal private). - Only effective when visibility is PRIVATE (1). - - Returns: - str: The URL of the created dataset repository. - """ - - if dataset_name is None or namespace is None: - raise InvalidParameter('dataset_name and namespace are required!') - - cookies = self.get_cookies(access_token=token, cookies_required=True) - if not endpoint: - endpoint = self.endpoint - path = f'{endpoint}/api/v1/datasets' - files = { - 'Name': (None, dataset_name), - 'ChineseName': (None, chinese_name), - 'Owner': (None, namespace), - 'License': (None, license), - 'Visibility': (None, visibility), - 'Description': (None, description) - } - - if gated_mode is not None: - if visibility != DatasetVisibility.PRIVATE: - logger.warning('gated_mode is only effective when visibility is PRIVATE, ignored.') - else: - files['ProtectedMode'] = (None, 1 if gated_mode else 2) - - r = self.session.post( - path, - files=files, - cookies=cookies, - headers=self.builder_headers(self.headers), - ) - - handle_http_post_error(r, path, files) - raise_on_error(r.json()) - dataset_repo_url = f'{endpoint}/datasets/{namespace}/{dataset_name}' - logger.info(f'Create dataset success: {dataset_repo_url}') - - return dataset_repo_url - - def delete_dataset(self, - dataset_id: str, - endpoint: Optional[str] = None, - token: Optional[str] = None): - """ - @deprecated - Delete a dataset from ModelScope. - - Args: - dataset_id (str): The dataset id to delete. - endpoint (str, optional): The endpoint to use. If not provided, the default endpoint is used. - token (str, optional): The access token for authentication. - - Returns: - None - """ - warnings.warn( - 'This function is deprecated due to security reasons, ' - 'and will be recovered in future versions with proper token authentication. ', - DeprecationWarning, - stacklevel=2 - ) - cookies = self.get_cookies(access_token=token, cookies_required=True) - if not endpoint: - endpoint = self.endpoint - if cookies is None: - raise ValueError('Token does not exist, please login first.') - - path = f'{endpoint}/api/v1/datasets/{dataset_id}' - r = self.session.delete(path, - cookies=cookies, - headers=self.builder_headers(self.headers)) - raise_for_http_status(r) - raise_on_error(r.json()) - - _dataset_id_type_cache: dict = {} - - def get_dataset_id_and_type(self, - dataset_name: str, - namespace: str, - endpoint: Optional[str] = None, - token: Optional[str] = None): - """ Get the dataset id and type. """ - if not endpoint: - endpoint = self.endpoint - cache_key = (namespace, dataset_name, endpoint) - cached = HubApi._dataset_id_type_cache.get(cache_key) - if cached is not None: - return cached - datahub_url = f'{endpoint}/api/v1/datasets/{namespace}/{dataset_name}' - cookies = self.get_cookies(access_token=token) - r = self.session.get(datahub_url, cookies=cookies) - resp = r.json() - datahub_raise_on_error(datahub_url, resp, r) - dataset_id = resp['Data']['Id'] - dataset_type = resp['Data']['Type'] - HubApi._dataset_id_type_cache[cache_key] = (dataset_id, dataset_type) - return dataset_id, dataset_type - - def list_repo_tree(self, - dataset_name: str, - namespace: str, - revision: str, - root_path: str, - recursive: bool = True, - page_number: int = 1, - page_size: int = 100, - endpoint: Optional[str] = None, - token: Optional[str] = None): - """ - @deprecated: Use `get_dataset_files` instead. - """ - warnings.warn('The function `list_repo_tree` is deprecated, use `get_dataset_files` instead.', - DeprecationWarning) - - dataset_hub_id, dataset_type = self.get_dataset_id_and_type( - dataset_name=dataset_name, namespace=namespace, endpoint=endpoint, token=token) - - recursive = 'True' if recursive else 'False' - if not endpoint: - endpoint = self.endpoint - datahub_url = f'{endpoint}/api/v1/datasets/{dataset_hub_id}/repo/tree' - params = {'Revision': revision if revision else 'master', - 'Root': root_path if root_path else '/', 'Recursive': recursive, - 'PageNumber': page_number, 'PageSize': page_size} - cookies = self.get_cookies(access_token=token) - - r = self.session.get(datahub_url, params=params, cookies=cookies) - resp = r.json() - datahub_raise_on_error(datahub_url, resp, r) - - return resp - - def list_repo_commits(self, - repo_id: str, - *, - repo_type: Optional[str] = REPO_TYPE_MODEL, - revision: Optional[str] = DEFAULT_REPOSITORY_REVISION, - page_number: int = 1, - page_size: int = 50, - endpoint: Optional[str] = None, - token: Optional[str] = None): - """ - Get the commit history for a repository. - - Args: - repo_id (str): The repository id, in the format of `namespace/repo_name`. - repo_type (Optional[str]): The type of the repository. Supported types are `model` and `dataset`. - revision (str): The branch or tag name. Defaults to `DEFAULT_REPOSITORY_REVISION`. - page_number (int): The page number for pagination. Defaults to 1. - page_size (int): The number of commits per page. Defaults to 50. - endpoint (Optional[str]): The endpoint to use, defaults to None to use the endpoint specified in the class. - token (Optional[str]): The access token. - - Returns: - CommitHistoryResponse: The commit history response. - - Examples: - >>> from modelscope.hub.api import HubApi - >>> api = HubApi() - >>> commit_history = api.list_repo_commits('meituan/Meeseeks') - >>> print(f"Total commits: {commit_history.total_count}") - >>> for commit in commit_history.commits: - ... print(f"{commit.short_id}: {commit.title}") - """ - - if is_relative_path(repo_id) and repo_id.count('/') == 1: - _owner, _dataset_name = repo_id.split('/') - else: - raise ValueError(f'Invalid repo_id: {repo_id} !') - - if not endpoint: - endpoint = self.endpoint - - commits_url = f'{endpoint}/api/v1/{repo_type}s/{repo_id}/commits' if repo_type else \ - f'{endpoint}/api/v1/models/{repo_id}/commits' - params = { - 'Ref': revision or DEFAULT_REPOSITORY_REVISION, - 'PageNumber': page_number, - 'PageSize': page_size - } - cookies = self.get_cookies(access_token=token) - - try: - r = self.session.get(commits_url, params=params, - cookies=cookies, headers=self.builder_headers(self.headers)) - raise_for_http_status(r) - resp = r.json() - raise_on_error(resp) - - if resp.get('Code') == HTTPStatus.OK: - return CommitHistoryResponse.from_api_response(resp) - - except requests.exceptions.RequestException as e: - raise Exception(f'Failed to get repository commits for {repo_id}: {str(e)}') - - def get_dataset_files(self, - repo_id: str, - *, - revision: str = DEFAULT_REPOSITORY_REVISION, - root_path: str = '/', - recursive: bool = True, - page_number: int = 1, - page_size: int = 100, - endpoint: Optional[str] = None, - token: Optional[str] = None, - dataset_hub_id: Optional[str] = None): - """ - Get the dataset files. - - Args: - repo_id (str): The repository id, in the format of `namespace/dataset_name`. - revision (str): The branch or tag name. Defaults to `DEFAULT_REPOSITORY_REVISION`. - root_path (str): The root path to list. Defaults to '/'. - recursive (bool): Whether to list recursively. Defaults to True. - page_number (int): The page number for pagination. Defaults to 1. - page_size (int): The number of items per page. Defaults to 100. - endpoint (Optional[str]): The endpoint to use, defaults to None to use the endpoint specified in the class. - token (Optional[str]): The access token. - dataset_hub_id (Optional[str]): Pre-fetched dataset hub id. When provided, - skips the internal ``get_dataset_id_and_type`` lookup. Useful in pagination - loops to avoid redundant API calls per page. - - Returns: - List: The response containing the dataset repository tree information. - e.g. [{'CommitId': None, 'CommitMessage': '...', 'Size': 0, 'Type': 'tree'}, ...] - """ - - if dataset_hub_id is None: - if is_relative_path(repo_id) and repo_id.count('/') == 1: - _owner, _dataset_name = repo_id.split('/') - else: - raise ValueError(f'Invalid repo_id: {repo_id} !') - - dataset_hub_id, _ = self.get_dataset_id_and_type( - dataset_name=_dataset_name, namespace=_owner, endpoint=endpoint, token=token) - - if not endpoint: - endpoint = self.endpoint - datahub_url = f'{endpoint}/api/v1/datasets/{dataset_hub_id}/repo/tree' - params = { - 'Revision': revision, - 'Root': root_path, - 'Recursive': 'True' if recursive else 'False', - 'PageNumber': page_number, - 'PageSize': page_size - } - cookies = self.get_cookies(access_token=token) - - r = self.session.get(datahub_url, params=params, cookies=cookies) - resp = r.json() - datahub_raise_on_error(datahub_url, resp, r) - - data = resp.get('Data') - if data is None: - return [] - return data.get('Files') or [] - - def get_dataset( - self, - dataset_id: str, - revision: Optional[str] = DEFAULT_REPOSITORY_REVISION, + dataset_name: str, + namespace: str, + max_limit: int, + is_recursive: bool, + is_filter_dir: bool, + revision: str, endpoint: Optional[str] = None, - token: Optional[str] = None - ): - """ - Get the dataset information. + token: Optional[str] = None, + ) -> list: + """List objects in a dataset's OSS storage.""" + params = { + 'MaxLimit': max_limit, + 'Revision': revision, + 'Recursive': is_recursive, + 'FilterDir': is_filter_dir, + } + body = self._legacy_request( + 'GET', + f'datasets/{namespace}/{dataset_name}/oss/tree/', + params=params, + timeout=1800, + ) + return body.get(API_RESPONSE_FIELD_DATA, []) - Args: - dataset_id (str): The dataset id. - revision (Optional[str]): The revision of the dataset. - endpoint (Optional[str]): The endpoint to use, defaults to None to use the endpoint specified in the class. - token (Optional[str]): The access token. + def delete_oss_dataset_object( + self, + object_name: str, + dataset_name: str, + namespace: str, + revision: str, + endpoint: Optional[str] = None, + token: Optional[str] = None, + ) -> str: + """Delete a single object from dataset OSS storage.""" + if not all([object_name, dataset_name, namespace, revision]): + raise ValueError('Args cannot be empty!') + body = self._legacy_request( + 'DELETE', + f'datasets/{namespace}/{dataset_name}/oss', + params={'Path': object_name, 'Revision': revision}, + ) + return body[API_RESPONSE_FIELD_MESSAGE] - Returns: - dict: The dataset information. - """ - cookies = self.get_cookies(access_token=token) - if not endpoint: - endpoint = self.endpoint - - if revision: - path = f'{endpoint}/api/v1/datasets/{dataset_id}?Revision={revision}' - else: - path = f'{endpoint}/api/v1/datasets/{dataset_id}' - - r = self.session.get( - path, cookies=cookies, headers=self.builder_headers(self.headers)) - raise_for_http_status(r) - resp = r.json() - datahub_raise_on_error(path, resp, r) - return resp[API_RESPONSE_FIELD_DATA] - - def get_dataset_meta_file_list(self, dataset_name: str, namespace: str, - dataset_id: str, revision: str, endpoint: Optional[str] = None, - token: Optional[str] = None): - """ Get the meta file-list of the dataset. """ - if not endpoint: - endpoint = self.endpoint - datahub_url = f'{endpoint}/api/v1/datasets/{dataset_id}/repo/tree?Revision={revision}' - cookies = self.get_cookies(access_token=token) - r = self.session.get(datahub_url, - cookies=cookies, - headers=self.builder_headers(self.headers)) - resp = r.json() - datahub_raise_on_error(datahub_url, resp, r) - file_list = resp['Data'] - if file_list is None: - raise NotExistError( - f'The modelscope dataset [dataset_name = {dataset_name}, namespace = {namespace}, ' - f'version = {revision}] dose not exist') - - file_list = file_list['Files'] - return file_list + def delete_oss_dataset_dir( + self, + object_name: str, + dataset_name: str, + namespace: str, + revision: str, + endpoint: Optional[str] = None, + token: Optional[str] = None, + ) -> str: + """Delete a directory prefix from dataset OSS storage.""" + if not all([object_name, dataset_name, namespace, revision]): + raise ValueError('Args cannot be empty!') + prefix = object_name.rstrip('/') + '/' + body = self._legacy_request( + 'DELETE', + f'datasets/{namespace}/{dataset_name}/oss/prefix', + params={'Prefix': prefix, 'Revision': revision}, + ) + return body[API_RESPONSE_FIELD_MESSAGE] + # ------------------------------------------------------------------ + # Meta file download + # ------------------------------------------------------------------ @staticmethod - def dump_datatype_file(dataset_type: int, meta_cache_dir: str): - """ - Dump the data_type as a local file, in order to get the dataset - formation without calling the datahub. - More details, please refer to the class - `modelscope.utils.constant.DatasetFormations`. - """ - dataset_type_file_path = os.path.join(meta_cache_dir, - f'{str(dataset_type)}{DatasetFormations.formation_mark_ext.value}') - with open(dataset_type_file_path, 'w') as fp: - fp.write('*** Automatically-generated file, do not modify ***') + def fetch_meta_files_from_url( + url: str, + out_path: str, + chunk_size: int = 1024, + mode=None, + token: Optional[str] = None, + ) -> str: + """Download a meta-data file (csv/jsonl) from a URL to local cache.""" + from modelscope.utils.constant import DownloadMode + if mode is None: + mode = DownloadMode.REUSE_DATASET_IF_EXISTS - def get_dataset_meta_files_local_paths(self, dataset_name: str, - namespace: str, - revision: str, - meta_cache_dir: str, dataset_type: int, file_list: list, - endpoint: Optional[str] = None, - token: Optional[str] = None): - local_paths = defaultdict(list) - dataset_formation = DatasetFormations(dataset_type) - dataset_meta_format = DatasetMetaFormats[dataset_formation] - cookies = self.get_cookies(access_token=token) - - # Dump the data_type as a local file - HubApi.dump_datatype_file(dataset_type=dataset_type, meta_cache_dir=meta_cache_dir) - if not endpoint: - endpoint = self.endpoint - for file_info in file_list: - file_path = file_info['Path'] - extension = os.path.splitext(file_path)[-1] - if extension in dataset_meta_format: - datahub_url = f'{endpoint}/api/v1/datasets/{namespace}/{dataset_name}/repo?' \ - f'Revision={revision}&FilePath={file_path}' - r = self.session.get(datahub_url, cookies=cookies) - raise_for_http_status(r) - local_path = os.path.join(meta_cache_dir, file_path) - if os.path.exists(local_path): - logger.warning( - f"Reusing dataset {dataset_name}'s python file ({local_path})" - ) - local_paths[extension].append(local_path) - continue - with open(local_path, 'wb') as f: - f.write(r.content) - local_paths[extension].append(local_path) - - return local_paths, dataset_formation - - @staticmethod - def fetch_meta_files_from_url(url, out_path, chunk_size=1024, mode=DownloadMode.REUSE_DATASET_IF_EXISTS, - token: Optional[str] = None): - """ - Fetch the meta-data files from the url, e.g. csv/jsonl files. - """ - import hashlib - from tqdm.auto import tqdm import pandas as pd + from tqdm.auto import tqdm + + out_path = os.path.join( + out_path, hashlib.md5(url.encode('utf-8')).hexdigest()) - out_path = os.path.join(out_path, hashlib.md5(url.encode(encoding='UTF-8')).hexdigest()) if mode == DownloadMode.FORCE_REDOWNLOAD and os.path.exists(out_path): os.remove(out_path) if os.path.exists(out_path): logger.info(f'Reusing cached meta-data file: {out_path}') return out_path - cookies = HubApi().get_cookies(access_token=token) - # Make the request and get the response content as TextIO + cookies = HubApi().get_cookies(access_token=token) logger.info('Loading meta-data file ...') response = requests.get(url, cookies=cookies, stream=True) total_size = int(response.headers.get('content-length', 0)) @@ -1877,2513 +278,50 @@ class HubApi: for chunk in get_chunk(response): progress.update(len(chunk)) if url.endswith('jsonl'): - chunk = [json.loads(line) for line in chunk if line.strip()] - if len(chunk) == 0: + chunk = [json.loads(line) for line in chunk + if line.strip()] + if not chunk: continue - if iter_num == 0: - with_header = True - else: - with_header = False chunk_df = pd.DataFrame(chunk) - chunk_df.to_csv(f, index=False, header=with_header, escapechar='\\') + chunk_df.to_csv( + f, index=False, header=(iter_num == 0), + escapechar='\\') iter_num += 1 else: - # csv or others for line in chunk: f.write(line + '\n') progress.close() - return out_path - def get_dataset_file_url( - self, - file_name: str, - dataset_name: str, - namespace: str, - revision: Optional[str] = DEFAULT_DATASET_REVISION, - view: Optional[bool] = False, - extension_filter: Optional[bool] = True, - endpoint: Optional[str] = None): + def __getattr__(self, name: str): + """Transparent proxy to the internal ``modelscope_hub.HubApi``. - if not file_name or not dataset_name or not namespace: - raise ValueError('Args (file_name, dataset_name, namespace) cannot be empty!') - - # Note: make sure the FilePath is the last parameter in the url - params: dict = {'Source': 'SDK', 'Revision': revision, 'FilePath': file_name, 'View': view} - params: str = urlencode(params) - if not endpoint: - endpoint = self.endpoint - file_url = f'{endpoint}/api/v1/datasets/{namespace}/{dataset_name}/repo?{params}' - - return file_url - - # if extension_filter: - # if os.path.splitext(file_name)[-1] in META_FILES_FORMAT: - # file_url = f'{self.endpoint}/api/v1/datasets/{namespace}/{dataset_name}/repo?'\ - # f'Revision={revision}&FilePath={file_name}' - # else: - # file_url = file_name - # return file_url - # else: - # return file_url - - def get_dataset_file_url_origin( - self, - file_name: str, - dataset_name: str, - namespace: str, - revision: Optional[str] = DEFAULT_DATASET_REVISION, - endpoint: Optional[str] = None): - if not endpoint: - endpoint = self.endpoint - if file_name and os.path.splitext(file_name)[-1] in META_FILES_FORMAT: - file_name = f'{endpoint}/api/v1/datasets/{namespace}/{dataset_name}/repo?' \ - f'Revision={revision}&FilePath={file_name}' - return file_name - - def get_dataset_access_config( - self, - dataset_name: str, - namespace: str, - revision: Optional[str] = DEFAULT_DATASET_REVISION, - endpoint: Optional[str] = None, - token: Optional[str] = None): - if not endpoint: - endpoint = self.endpoint - datahub_url = f'{endpoint}/api/v1/datasets/{namespace}/{dataset_name}/' \ - f'ststoken?Revision={revision}' - return self.datahub_remote_call(datahub_url, token=token) - - def get_dataset_access_config_session( - self, - dataset_name: str, - namespace: str, - check_cookie: bool, - revision: Optional[str] = DEFAULT_DATASET_REVISION, - endpoint: Optional[str] = None, - token: Optional[str] = None): - - if not endpoint: - endpoint = self.endpoint - datahub_url = f'{endpoint}/api/v1/datasets/{namespace}/{dataset_name}/' \ - f'ststoken?Revision={revision}' - if check_cookie: - cookies = self._check_cookie(use_cookies=True) - else: - cookies = self.get_cookies(access_token=token) - - r = self.session.get( - url=datahub_url, - cookies=cookies, - headers=self.builder_headers(self.headers)) - resp = r.json() - raise_on_error(resp) - return resp['Data'] - - def get_virgo_meta(self, dataset_id: str, version: int = 1, token: Optional[str] = None) -> dict: - """ - Get virgo dataset meta info. - """ - virgo_endpoint = os.environ.get(VirgoDatasetConfig.env_virgo_endpoint, '') - if not virgo_endpoint: - raise RuntimeError(f'Virgo endpoint is not set in env: {VirgoDatasetConfig.env_virgo_endpoint}') - - virgo_dataset_url = f'{virgo_endpoint}/data/set/download' - cookies = requests.utils.dict_from_cookiejar(self.get_cookies(access_token=token)) - - dataset_info = dict( - dataSetId=dataset_id, - dataSetVersion=version - ) - data = dict( - data=dataset_info, - ) - r = self.session.post(url=virgo_dataset_url, - json=data, - cookies=cookies, - headers=self.builder_headers(self.headers), - timeout=900) - resp = r.json() - if resp['code'] != 0: - raise RuntimeError(f'Failed to get virgo dataset: {resp}') - - return resp['data'] - - def get_dataset_access_config_for_unzipped(self, - dataset_name: str, - namespace: str, - revision: str, - zip_file_name: str, - endpoint: Optional[str] = None, - token: Optional[str] = None): - if not endpoint: - endpoint = self.endpoint - datahub_url = f'{endpoint}/api/v1/datasets/{namespace}/{dataset_name}' - cookies = self.get_cookies(access_token=token) - r = self.session.get(url=datahub_url, cookies=cookies, - headers=self.builder_headers(self.headers)) - resp = r.json() - # get visibility of the dataset - raise_on_error(resp) - data = resp['Data'] - visibility = VisibilityMap.get(data['Visibility']) - - datahub_sts_url = f'{datahub_url}/ststoken?Revision={revision}' - r_sts = self.session.get(url=datahub_sts_url, cookies=cookies, - headers=self.builder_headers(self.headers)) - resp_sts = r_sts.json() - raise_on_error(resp_sts) - data_sts = resp_sts['Data'] - file_dir = visibility + '-unzipped' + '/' + namespace + '_' + dataset_name + '_' + zip_file_name - data_sts['Dir'] = file_dir - return data_sts - - def list_oss_dataset_objects(self, dataset_name, namespace, max_limit, - is_recursive, is_filter_dir, revision, endpoint: Optional[str] = None, - token: Optional[str] = None): - if not endpoint: - endpoint = self.endpoint - url = f'{endpoint}/api/v1/datasets/{namespace}/{dataset_name}/oss/tree/?' \ - f'MaxLimit={max_limit}&Revision={revision}&Recursive={is_recursive}&FilterDir={is_filter_dir}' - - cookies = self.get_cookies(access_token=token) - resp = self.session.get(url=url, cookies=cookies, timeout=1800) - resp = resp.json() - raise_on_error(resp) - resp = resp['Data'] - return resp - - def delete_oss_dataset_object(self, object_name: str, dataset_name: str, - namespace: str, revision: str, endpoint: Optional[str] = None, - token: Optional[str] = None) -> str: - if not object_name or not dataset_name or not namespace or not revision: - raise ValueError('Args cannot be empty!') - if not endpoint: - endpoint = self.endpoint - url = f'{endpoint}/api/v1/datasets/{namespace}/{dataset_name}/oss?Path={object_name}&Revision={revision}' - - cookies = self.get_cookies(access_token=token, cookies_required=True) - resp = self.session.delete(url=url, cookies=cookies) - resp = resp.json() - raise_on_error(resp) - resp = resp['Message'] - return resp - - def delete_oss_dataset_dir(self, object_name: str, dataset_name: str, - namespace: str, revision: str, endpoint: Optional[str] = None, - token: Optional[str] = None) -> str: - if not object_name or not dataset_name or not namespace or not revision: - raise ValueError('Args cannot be empty!') - if not endpoint: - endpoint = self.endpoint - url = f'{endpoint}/api/v1/datasets/{namespace}/{dataset_name}/oss/prefix?Prefix={object_name}/' \ - f'&Revision={revision}' - - cookies = self.get_cookies(access_token=token, cookies_required=True) - resp = self.session.delete(url=url, cookies=cookies) - resp = resp.json() - raise_on_error(resp) - resp = resp['Message'] - return resp - - def datahub_remote_call(self, url, token: Optional[str] = None): - cookies = self.get_cookies(access_token=token) - r = self.session.get( - url, - cookies=cookies, - headers={'user-agent': ModelScopeConfig.get_user_agent()}) - resp = r.json() - datahub_raise_on_error(url, resp, r) - return resp['Data'] - - def dataset_download_statistics(self, dataset_name: str, namespace: str, - use_streaming: bool = False, endpoint: Optional[str] = None, - token: Optional[str] = None) -> None: - is_ci_test = os.getenv('CI_TEST') == 'True' - if not endpoint: - endpoint = self.endpoint - if dataset_name and namespace and not is_ci_test and not use_streaming: - try: - cookies = self.get_cookies(access_token=token) - - # Download count - download_count_url = f'{endpoint}/api/v1/datasets/{namespace}/{dataset_name}/download/increase' - download_count_resp = self.session.post(download_count_url, cookies=cookies, - headers=self.builder_headers(self.headers)) - raise_for_http_status(download_count_resp) - - # Download uv - channel = DownloadChannel.LOCAL.value - user_name = '' - if MODELSCOPE_CLOUD_ENVIRONMENT in os.environ: - channel = os.environ[MODELSCOPE_CLOUD_ENVIRONMENT] - if MODELSCOPE_CLOUD_USERNAME in os.environ: - user_name = os.environ[MODELSCOPE_CLOUD_USERNAME] - download_uv_url = f'{endpoint}/api/v1/datasets/{namespace}/{dataset_name}/download/uv/' \ - f'{channel}?user={user_name}' - - download_uv_resp = self.session.post(download_uv_url, cookies=cookies, - headers=self.builder_headers(self.headers)) - download_uv_resp = download_uv_resp.json() - raise_on_error(download_uv_resp) - - except Exception as e: - logger.error(e) - - def builder_headers(self, headers): - return {MODELSCOPE_REQUEST_ID: str(uuid.uuid4().hex), - **headers} - - def _build_bearer_headers(self, - token: Optional[str] = None, - token_required: bool = False) -> Dict[str, str]: - """ - Build HTTP headers with optional Bearer token for OpenAPI endpoints. - - Token resolution order: - 1. Explicit token param - 2. self.token (set at construction) - 3. MODELSCOPE_API_TOKEN env var - 4. Locally cached cookies (m_session_id from login()) - - Args: - token: Optional access token for one-time authentication. - token_required: If True, raise ValueError when no token is available. - - Returns: - Headers dict with user-agent, request-id, and optionally Authorization. - - Raises: - ValueError: If token_required is True but no token is available. - """ - headers = self.builder_headers(self.headers) - - # Priority: explicit token > self.token > env var > local cookies - resolved_token = token or self.token or os.environ.get( - 'MODELSCOPE_API_TOKEN') - - # Fall back to locally cached cookies (m_session_id saved by login()) - if not resolved_token: - cookies = self.get_cookies() - if cookies: - for cookie in cookies: - if cookie.name == 'm_session_id': - resolved_token = cookie.value - break - - if resolved_token: - headers['Authorization'] = f'Bearer {resolved_token}' - elif token_required: - raise ValueError( - 'Authentication required but no token found. ' - 'You can pass the `token` argument, ' - 'or set MODELSCOPE_API_TOKEN environment variable, ' - 'or use HubApi(token=`your_sdk_token`). ' - 'Your token is available at https://modelscope.cn/my/myaccesstoken' - ) - return headers - - @staticmethod - def _parse_openapi_response(response: 'requests.Response') -> Dict[str, Any]: - """ - Parse OpenAPI response with unified JSON parsing and data extraction. - - Handles the standard OpenAPI response envelope: - {"success": bool, "data": {...}, "message": str} - Also handles the simpler envelope where only "data" is present. - - Args: - response: requests Response object (HTTP status already validated). - - Returns: - Parsed 'data' dict from the response envelope. - - Raises: - RequestError: If JSON parsing fails or business-level error is returned. + Only invoked when normal attribute lookup (instance dict, class + hierarchy including :class:`LegacyHubApi`) fails. Private names are + excluded to avoid recursion during ``__init__``. """ + if name.startswith('_'): + raise AttributeError(name) try: - resp = response.json() - except (requests.exceptions.JSONDecodeError, ValueError) as e: - logger.error(f'JSON parsing failed: {e}') - raise RequestError(f'Invalid JSON response: {e}') from e - - # OpenAPI envelope with explicit success field - if isinstance(resp, dict) and 'success' in resp: - if resp.get('success') is True and 'data' in resp: - return resp['data'] - else: - msg = resp.get('message') or 'OpenAPI request failed' - raise RequestError(msg) - - # Simple envelope with data field only (e.g., MCP API) - return resp.get('data', {}) if isinstance(resp, dict) else {} - - def get_file_base_path(self, repo_id: str, endpoint: Optional[str] = None) -> str: - _namespace, _dataset_name = repo_id.split('/') - if not endpoint: - endpoint = self.endpoint - return f'{endpoint}/api/v1/datasets/{_namespace}/{_dataset_name}/repo?' - # return f'{endpoint}/api/v1/datasets/{namespace}/{dataset_name}/repo?Revision={revision}&FilePath=' - - def create_repo( - self, - repo_id: str, - *, - token: Union[str, bool, None] = None, - visibility: Optional[str] = Visibility.PUBLIC, - repo_type: Optional[str] = REPO_TYPE_MODEL, - chinese_name: Optional[str] = None, - license: Optional[str] = Licenses.APACHE_V2, - endpoint: Optional[str] = None, - exist_ok: Optional[bool] = False, - create_default_config: Optional[bool] = True, - aigc_model: Optional[AigcModel] = None, - gated_mode: Optional[bool] = None, - **kwargs, - ) -> str: - """ - Create a repository on the ModelScope Hub. - - Args: - repo_id (str): The repo id in the format of `owner_name/repo_name`. - token (Union[str, bool, None]): The access token. - visibility (Optional[str]): The visibility of the repo, - could be `public`, `private`, `internal`, default to `public`. - repo_type (Optional[str]): The repo type, default to `model`. - chinese_name (Optional[str]): The Chinese name of the repo. - license (Optional[str]): The license of the repo, default to `apache-2.0`. - endpoint (Optional[str]): The endpoint to use. - In the format of `https://www.modelscope.cn` or 'https://www.modelscope.ai' - exist_ok (Optional[bool]): If the repo exists, whether to return the repo url directly. - create_default_config (Optional[bool]): If True, create a default configuration file in the model repo. - gated_mode (Optional[bool]): Gated mode for private repos. - True = gated (application-based download), False = off (normal private). - Only effective when visibility is ``private``. - **kwargs: The additional arguments. - - Returns: - str: The repo url. - """ - - if not repo_id: - raise ValueError('Repo id cannot be empty!') - if not endpoint: - endpoint = self.endpoint - - repo_exists: bool = self.repo_exists(repo_id, repo_type=repo_type, endpoint=endpoint, token=token) - if repo_exists: - if exist_ok: - repo_url: str = f'{endpoint}/{repo_type}s/{repo_id}' - logger.warning(f'Repo {repo_id} already exists, got repo url: {repo_url}') - return repo_url - else: - raise ValueError(f'Repo {repo_id} already exists!') - - repo_id_list = repo_id.split('/') - if len(repo_id_list) != 2: - raise ValueError('Invalid repo id, should be in the format of `owner_name/repo_name`') - namespace, repo_name = repo_id_list - - if repo_type == REPO_TYPE_MODEL: - visibilities = {k: v for k, v in ModelVisibility.__dict__.items() if not k.startswith('__')} - visibility: int = visibilities.get(visibility.upper()) - if visibility is None: - raise ValueError(f'Invalid visibility: {visibility}, ' - f'supported visibilities: `public`, `private`, `internal`') - repo_url: str = self.create_model( - model_id=repo_id, - visibility=visibility, - license=license, - chinese_name=chinese_name, - aigc_model=aigc_model, - token=token, - endpoint=endpoint, - gated_mode=gated_mode, - ) - if create_default_config: - with tempfile.TemporaryDirectory() as temp_cache_dir: - from modelscope.hub.repository import Repository - repo = Repository(temp_cache_dir, repo_id, auth_token=token, endpoint=endpoint) - default_config = { - 'framework': 'pytorch', - 'task': 'text-generation', - 'allow_remote': True - } - config_json = kwargs.get('config_json') - if not config_json: - config_json = {} - config = {**default_config, **config_json} - add_content_to_file( - repo, - 'configuration.json', [json.dumps(config)], - ignore_push_error=True) - print(f'New model created successfully at {repo_url}.', flush=True) - - elif repo_type == REPO_TYPE_DATASET: - visibilities = {k: v for k, v in DatasetVisibility.__dict__.items() if not k.startswith('__')} - visibility: int = visibilities.get(visibility.upper()) - if visibility is None: - raise ValueError(f'Invalid visibility: {visibility}, ' - f'supported visibilities: `public`, `private`, `internal`') - repo_url: str = self.create_dataset( - dataset_name=repo_name, - namespace=namespace, - chinese_name=chinese_name, - license=license, - visibility=visibility, - token=token, - endpoint=endpoint, - gated_mode=gated_mode, - ) - print(f'New dataset created successfully at {repo_url}.', flush=True) - - elif repo_type == REPO_TYPE_STUDIO: - repo_url = self._create_studio_repo( - owner=namespace, - repo_name=repo_name, - visibility=visibility, - license=license, - chinese_name=chinese_name, - token=token, - endpoint=endpoint, - **kwargs, - ) - print(f'New studio created successfully at {repo_url}.', flush=True) - - else: - raise ValueError(f'Invalid repo type: {repo_type}, supported repos: {REPO_TYPE_SUPPORT}') - - return repo_url - - # --- Studio Operations --- - - @staticmethod - def _parse_studio_id(studio_id: str): - """Parse a studio_id of the form ``owner/repo_name`` into (owner, name).""" - if not studio_id or studio_id.count('/') != 1: - raise InvalidParameter( - f'Invalid studio_id: {studio_id}, must be of format owner/repo_name') - owner, name = studio_id.split('/', 1) - if not owner or not name: - raise InvalidParameter( - f'Invalid studio_id: {studio_id}, must be of format owner/repo_name') - return owner, name - - # Map Licenses display names to SPDX identifiers expected by the - # Studio OpenAPI endpoint. - _LICENSE_TO_SPDX = { - 'Apache License 2.0': 'apache-2.0', - 'GPL-2.0': 'gpl-2.0', - 'GPL-3.0': 'gpl-3.0', - 'LGPL-2.1': 'lgpl-2.1', - 'LGPL-3.0': 'lgpl-3.0', - 'AFL-3.0': 'afl-3.0', - 'ECL-2.0': 'ecl-2.0', - 'MIT': 'mit', - } - - def _create_studio_repo(self, - owner: str, - repo_name: str, - visibility: Optional[str] = Visibility.PUBLIC, - license: Optional[str] = None, - chinese_name: Optional[str] = None, - token: Optional[str] = None, - endpoint: Optional[str] = None, - **kwargs) -> str: - """Create a studio repo via the OpenAPI ``/openapi/v1/studios`` endpoint. - - Supported optional studio fields in ``kwargs``: - description, sdk_type, sdk_version, base_image, hardware, cover_image. - """ - endpoint = endpoint or self.endpoint - path = f'{endpoint}/openapi/v1/studios' - headers = self._build_bearer_headers(token=token, token_required=True) - - is_private = visibility is not None and visibility != Visibility.PUBLIC - # Convert license display name to SPDX identifier if needed. - license_spdx = self._LICENSE_TO_SPDX.get(license, license) if license else None - - body = { - 'repo_name': repo_name, - 'owner': owner, - 'private': is_private, - 'license': license_spdx, - 'display_name': chinese_name, - 'description': kwargs.get('description'), - 'sdk_type': kwargs.get('sdk_type'), - 'sdk_version': kwargs.get('sdk_version'), - 'base_image': kwargs.get('base_image'), - 'hardware': kwargs.get('hardware'), - 'cover_image': kwargs.get('cover_image'), - } - body = {k: v for k, v in body.items() if v is not None} - - r = self.session.post(path, json=body, headers=headers) - handle_http_response(r, logger, None, f'{owner}/{repo_name}') - return f'{endpoint}/studios/{owner}/{repo_name}' - - def deploy_studio(self, studio_id, token=None, endpoint=None): - """Deploy a studio (re-pull code and rebuild). - - Args: - studio_id: Studio ID in format ``owner/repo_name``. - token: Optional access token. - endpoint: Optional API endpoint. - - Returns: - dict: Runtime status info including status and active_config. - """ - endpoint = endpoint or self.endpoint - owner, name = self._parse_studio_id(studio_id) - path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/deploy' - headers = self._build_bearer_headers(token=token, token_required=True) - r = self.session.post(path, headers=headers) - handle_http_response(r, logger, None, studio_id) - return self._parse_openapi_response(r) - - def stop_studio(self, studio_id, token=None, endpoint=None): - """Stop a running studio. - - Args: - studio_id: Studio ID in format ``owner/repo_name``. - token: Optional access token. - endpoint: Optional API endpoint. - - Returns: - dict: Runtime status info. - """ - endpoint = endpoint or self.endpoint - owner, name = self._parse_studio_id(studio_id) - path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/stop' - headers = self._build_bearer_headers(token=token, token_required=True) - r = self.session.post(path, headers=headers) - handle_http_response(r, logger, None, studio_id) - return self._parse_openapi_response(r) - - def get_studio_logs(self, studio_id, log_type='run', page_num=1, - page_size=100, keyword=None, start_timestamp=None, - end_timestamp=None, token=None, endpoint=None): - """Get studio build or runtime logs. - - Args: - studio_id: Studio ID in format ``owner/repo_name``. - log_type: Log type, ``'run'`` or ``'build'``. - page_num: Page number, starting from 1. - page_size: Number of log entries per page. - keyword: Optional keyword filter. - start_timestamp: Optional start timestamp in seconds. - end_timestamp: Optional end timestamp in seconds. - token: Optional access token. - endpoint: Optional API endpoint. - - Returns: - dict: Logs data with pagination info. - """ - endpoint = endpoint or self.endpoint - owner, name = self._parse_studio_id(studio_id) - path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/logs/{log_type}' - headers = self._build_bearer_headers(token=token, token_required=True) - params = {'page_num': page_num, 'page_size': page_size} - if keyword: - params['keyword'] = keyword - if start_timestamp is not None: - params['start_timestamp'] = start_timestamp - if end_timestamp is not None: - params['end_timestamp'] = end_timestamp - r = self.session.get(path, params=params, headers=headers) - handle_http_response(r, logger, None, studio_id) - return self._parse_openapi_response(r) - - def update_studio_settings(self, studio_id, token=None, endpoint=None, **settings): - """Update studio settings (PATCH, only specified fields are modified). - - Args: - studio_id: Studio ID in format ``owner/repo_name``. - token: Optional access token. - endpoint: Optional API endpoint. - **settings: Fields to update. Supported: ``display_name``, ``license``, - ``private``, ``description``, ``cover_image``, ``sdk_type``, - ``sdk_version``, ``base_image``, ``hardware``. Note: - ``sdk_type``/``sdk_version``/``base_image``/``hardware`` changes - require redeployment. - - Returns: - dict: Updated studio info. - """ - endpoint = endpoint or self.endpoint - owner, name = self._parse_studio_id(studio_id) - path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/settings' - headers = self._build_bearer_headers(token=token, token_required=True) - body = {k: v for k, v in settings.items() if v is not None} - r = self.session.patch(path, json=body, headers=headers) - handle_http_response(r, logger, None, studio_id) - return self._parse_openapi_response(r) - - def list_studio_secrets(self, studio_id, token=None, endpoint=None): - """List studio environment variable keys (values not returned for security). - - Args: - studio_id: Studio ID in format ``owner/repo_name``. - token: Optional access token. - endpoint: Optional API endpoint. - - Returns: - list: List of secret key dicts, e.g. ``[{'key': 'API_KEY'}, ...]``. - """ - endpoint = endpoint or self.endpoint - owner, name = self._parse_studio_id(studio_id) - path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/secrets' - headers = self._build_bearer_headers(token=token, token_required=True) - r = self.session.get(path, headers=headers) - handle_http_response(r, logger, None, studio_id) - data = self._parse_openapi_response(r) - return data.get('secrets', []) if isinstance(data, dict) else [] - - def add_studio_secret(self, studio_id, key, value, token=None, endpoint=None): - """Add an environment variable to a studio. - - Args: - studio_id: Studio ID in format ``owner/repo_name``. - key: Secret name (max 128 chars). - value: Secret value (max 4096 chars). - token: Optional access token. - endpoint: Optional API endpoint. - """ - endpoint = endpoint or self.endpoint - owner, name = self._parse_studio_id(studio_id) - path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/secrets' - headers = self._build_bearer_headers(token=token, token_required=True) - r = self.session.post( - path, json={'key': key, 'value': value}, headers=headers) - handle_http_response(r, logger, None, studio_id) - - def update_studio_secret(self, studio_id, key, value, token=None, endpoint=None): - """Update an existing environment variable in a studio. - - Args: - studio_id: Studio ID in format ``owner/repo_name``. - key: Secret name (max 128 chars). - value: New secret value (max 4096 chars). - token: Optional access token. - endpoint: Optional API endpoint. - """ - endpoint = endpoint or self.endpoint - owner, name = self._parse_studio_id(studio_id) - path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/secrets' - headers = self._build_bearer_headers(token=token, token_required=True) - r = self.session.put( - path, json={'key': key, 'value': value}, headers=headers) - handle_http_response(r, logger, None, studio_id) - - def delete_studio_secret(self, studio_id, key, token=None, endpoint=None): - """Delete an environment variable from a studio. - - Args: - studio_id: Studio ID in format ``owner/repo_name``. - key: Secret name to delete. - token: Optional access token. - endpoint: Optional API endpoint. - """ - endpoint = endpoint or self.endpoint - owner, name = self._parse_studio_id(studio_id) - path = f'{endpoint}/openapi/v1/studios/{owner}/{name}/secrets' - headers = self._build_bearer_headers(token=token, token_required=True) - r = self.session.delete(path, json={'key': key}, headers=headers) - handle_http_response(r, logger, None, studio_id) - - # --- End Studio Operations --- - - def create_commit( - self, - repo_id: str, - operations: Iterable[CommitOperation], - *, - commit_message: str, - commit_description: Optional[str] = None, - token: str = None, - repo_type: Optional[str] = REPO_TYPE_MODEL, - revision: Optional[str] = DEFAULT_REPOSITORY_REVISION, - endpoint: Optional[str] = None, - ) -> CommitInfo: - """ - Create a commit on the ModelScope Hub. - - Args: - repo_id (str): The repo id in the format of `owner_name/repo_name`. - operations (Iterable[CommitOperation]): The commit operations. - commit_message (str): The commit message. - commit_description (Optional[str]): The commit description. - token (str): The access token. If None, will use the cookies from the local cache. - See `https://modelscope.cn/my/myaccesstoken` to get your token. - repo_type (Optional[str]): The repo type, should be `model` or `dataset`. Defaults to `model`. - revision (Optional[str]): The branch or tag name. Defaults to `DEFAULT_REPOSITORY_REVISION`. - endpoint (Optional[str]): The endpoint to use. - In the format of `https://www.modelscope.cn` or 'https://www.modelscope.ai' - timeout (int): Timeout for each request in seconds (default: 180). - - Returns: - CommitInfo: The commit info. - - Raises: - ValueError: If the request fails with a 4xx client error. - requests.exceptions.RequestException: If a network-level error occurs. - """ - if not repo_id: - raise ValueError('Repo id cannot be empty!') - - if not endpoint: - endpoint = self.endpoint - - if repo_type not in REPO_TYPE_SUPPORT: - raise ValueError(f'Invalid repo type: {repo_type}, supported repos: {REPO_TYPE_SUPPORT}') - - url = f'{endpoint}/api/v1/repos/{repo_type}s/{repo_id}/commit/{revision}' - commit_message = commit_message or f'Commit to {repo_id}' - commit_description = commit_description or '' - - cookies = self.get_cookies(access_token=token, cookies_required=True) - - # Construct payload - payload = self._prepare_commit_payload( - operations=operations, - commit_message=commit_message, - ) - - # Guard: skip sending empty commits (no effective file changes) - if not payload['actions']: - logger.info( - 'Commit skipped: no effective actions in payload ' - '(all files already exist).') - return CommitInfo( - commit_url='', - commit_message=commit_message, - commit_description=commit_description or '', - oid='no-op', - ) - - response = self.session.post( - url, - headers=self.builder_headers(self.headers), - data=json.dumps(payload), - cookies=cookies, - ) - - if response.status_code != 200: - try: - error_detail = response.json() - except json.JSONDecodeError: - error_detail = response.text - error_msg = f'HTTP {response.status_code} error from {url}: {error_detail}' - raise ValueError(error_msg) - - resp = response.json() - oid = resp.get('Data', {}).get('oid', '') - logger.info(f'Commit succeeded: {url}') - return CommitInfo( - commit_url=url, - commit_message=commit_message, - commit_description=commit_description, - oid=oid, - ) - - def upload_file( - self, - *, - path_or_fileobj: Union[str, Path, bytes, BinaryIO], - path_in_repo: str, - repo_id: str, - token: Union[str, None] = None, - repo_type: Optional[str] = REPO_TYPE_MODEL, - commit_message: Optional[str] = None, - commit_description: Optional[str] = None, - buffer_size_mb: Optional[int] = 16, - tqdm_desc: Optional[str] = '[Uploading]', - disable_tqdm: Optional[bool] = False, - revision: Optional[str] = DEFAULT_REPOSITORY_REVISION - ) -> CommitInfo: - """ - Upload a file to the ModelScope Hub. - - Args: - path_or_fileobj (Union[str, Path, bytes, BinaryIO]): - The local file path or file-like object (BinaryIO) or bytes to upload. - path_in_repo (str): The path in the repo to upload to. - repo_id (str): The repo id in the format of `owner_name/repo_name`. - token (Union[str, None]): The access token. If None, will use the cookies from the local cache. - See `https://modelscope.cn/my/myaccesstoken` to get your token. - repo_type (Optional[str]): The repo type, default to `model`. - commit_message (Optional[str]): The commit message. - commit_description (Optional[str]): The commit description. - buffer_size_mb (Optional[int]): The buffer size in MB for reading the file. Default to 1MB. - tqdm_desc (Optional[str]): The description for the tqdm progress bar. Default to '[Uploading]'. - disable_tqdm (Optional[bool]): Whether to disable the tqdm progress bar. Default to False. - revision (Optional[str]): The branch or tag name. Defaults to `DEFAULT_REPOSITORY_REVISION`. - - Returns: - CommitInfo: The commit info. - - Examples: - >>> from modelscope.hub.api import HubApi - >>> api = HubApi() - >>> commit_info = api.upload_file( - ... path_or_fileobj='/path/to/your/file.txt', - ... path_in_repo='optional/path/in/repo/file.txt', - ... repo_id='your-namespace/your-repo-name', - ... commit_message='Upload file.txt to ModelScope hub' - ... ) - >>> print(commit_info) - """ - if repo_type not in REPO_TYPE_SUPPORT: - raise ValueError(f'Invalid repo type: {repo_type}, supported repos: {REPO_TYPE_SUPPORT}') - - if not path_or_fileobj: - raise ValueError('Path or file object cannot be empty!') - - # Check authentication first - self.get_cookies(access_token=token, cookies_required=True) - - if isinstance(path_or_fileobj, (str, Path)): - path_or_fileobj = os.path.abspath(os.path.expanduser(path_or_fileobj)) - path_in_repo = path_in_repo or os.path.basename(path_or_fileobj) - else: - # If path_or_fileobj is bytes or BinaryIO, then path_in_repo must be provided - if not path_in_repo: - raise ValueError('Arg `path_in_repo` cannot be empty!') - - # Read file content if path_or_fileobj is a file-like object (BinaryIO) - # TODO: to be refined - if isinstance(path_or_fileobj, io.BufferedIOBase): - path_or_fileobj = path_or_fileobj.read() - - self.upload_checker.check_file(path_or_fileobj) - self.upload_checker.check_normal_files( - file_path_list=[path_or_fileobj], - repo_type=repo_type, - ) - - commit_message = ( - commit_message if commit_message is not None else f'Upload {path_in_repo} to ModelScope hub' - ) - - if buffer_size_mb <= 0: - raise ValueError('Buffer size: `buffer_size_mb` must be greater than 0') - - hash_info_d: dict = compute_file_hash( - file_path_or_obj=path_or_fileobj, - buffer_size_mb=buffer_size_mb, - ) - file_size: int = hash_info_d['file_size'] - file_hash: str = hash_info_d['file_hash'] - - self.create_repo(repo_id=repo_id, - token=token, - repo_type=repo_type, - endpoint=self.endpoint, - exist_ok=True, - create_default_config=False) - - upload_res: dict = self._upload_blob( - repo_id=repo_id, - repo_type=repo_type, - sha256=file_hash, - size=file_size, - data=path_or_fileobj, - disable_tqdm=disable_tqdm, - tqdm_desc=tqdm_desc, - token=token, - ) - - # Construct commit info and create commit - add_operation: CommitOperationAdd = CommitOperationAdd( - path_in_repo=path_in_repo, - path_or_fileobj=path_or_fileobj, - file_hash_info=hash_info_d, - ) - add_operation._upload_mode = 'lfs' if self.upload_checker.is_lfs(path_or_fileobj, repo_type) else 'normal' - add_operation._is_uploaded = upload_res['is_uploaded'] - operations = [add_operation] - - print(f'Committing file to {repo_id} ...', flush=True) - commit_info: CommitInfo = self.create_commit( - repo_id=repo_id, - operations=operations, - commit_message=commit_message, - commit_description=commit_description, - token=token, - repo_type=repo_type, - revision=revision, - ) - - return commit_info - - def _track_uploaded_batch(self, tracker, results): - """Mark files as uploaded and persist tracker state.""" - for r in results: - tracker.mark_uploaded( - r['file_path_in_repo'], r['file_mtime'], - r['file_size_on_disk']) - tracker.save() - - def _track_committed_batch(self, tracker, results): - """Mark files as committed and persist tracker state.""" - tracker.mark_committed_batch([ - (r['file_path_in_repo'], r['file_mtime'], - r['file_size_on_disk']) - for r in results]) - tracker.save() - - def _upload_single_file( - self, - file_path_in_repo: str, - file_path: str, - *, - repo_id: str, - repo_type: str, - token: str, - tracker=None, - pre_validated=None, - ) -> dict: - """Hash and upload a single file, returning result dict.""" - if tracker is None: - tracker = NullTracker() - hash_info_d = None - file_stat = None - is_real_path = isinstance(file_path, (str, os.PathLike)) - if is_real_path: - try: - file_stat = os.stat(file_path) - cached = tracker.get_hash( - file_path_in_repo, file_stat.st_mtime, file_stat.st_size) - if cached is not None: - hash_info_d = cached - hash_info_d['file_path_or_obj'] = file_path - except OSError: - file_stat = None - - if hash_info_d is None: - hash_info_d = compute_file_hash(file_path_or_obj=file_path) - if is_real_path: - try: - if file_stat is None: - file_stat = os.stat(file_path) - tracker.put_hash( - file_path_in_repo, file_stat.st_mtime, - file_stat.st_size, hash_info_d) - except OSError: - pass - - # Ensure file_stat is available for real path files - if file_stat is None and is_real_path: - try: - file_stat = os.stat(file_path) - except OSError: - pass - - file_size: int = hash_info_d['file_size'] - file_hash: str = hash_info_d['file_hash'] - - # Application-level retry for transient blob upload failures - last_error = None - for attempt in range(UPLOAD_BLOB_MAX_RETRIES): - try: - # Validate file size has not changed since hash computation - if isinstance(file_path, (str, os.PathLike)): - current_size = os.path.getsize(str(file_path)) - if current_size != file_size: - raise IOError( - f'File size changed since hash computation: ' - f'was {file_size}, now {current_size}. ' - f'File may have been modified: {file_path_in_repo}') - upload_res: dict = self._upload_blob( - repo_id=repo_id, - repo_type=repo_type, - sha256=file_hash, - size=file_size, - data=file_path, - disable_tqdm=file_size <= UPLOAD_BLOB_TQDM_DISABLE_THRESHOLD, - tqdm_desc='[Uploading ' + file_path_in_repo + ']', - token=token, - pre_validated=pre_validated, - ) - break - except (ConnectionError, requests.exceptions.ConnectionError, - requests.exceptions.HTTPError, IOError) as e: - # Only retry on 5xx / connection errors; 4xx are not retryable - if isinstance(e, requests.exceptions.HTTPError): - if hasattr(e, 'response') and e.response is not None: - if e.response.status_code < 500: - raise - last_error = e - if attempt < UPLOAD_BLOB_MAX_RETRIES - 1: - wait = min(UPLOAD_BLOB_RETRY_BACKOFF ** attempt, - UPLOAD_BLOB_RETRY_MAX_WAIT) - logger.warning( - f'Blob upload attempt {attempt + 1}/{UPLOAD_BLOB_MAX_RETRIES} ' - f'failed for {file_path_in_repo}: {e}, retrying in {wait}s ...') - time.sleep(wait) - else: - raise RuntimeError( - f'Blob upload failed after {UPLOAD_BLOB_MAX_RETRIES} attempts ' - f'for {file_path_in_repo}: {last_error}') from last_error - - return { - 'file_path_in_repo': file_path_in_repo, - 'file_path': file_path, - 'file_mtime': file_stat.st_mtime if file_stat else 0, - 'file_size_on_disk': file_stat.st_size if file_stat else hash_info_d.get('file_size', 0), - 'is_uploaded': upload_res['is_uploaded'], - 'is_reused': upload_res.get('is_reused', False), - 'file_hash_info': hash_info_d, - } - - def _commit_with_retry( - self, - *, - repo_id: str, - operations, - commit_message: str, - commit_description: Optional[str] = None, - token: str = None, - repo_type: str = REPO_TYPE_MODEL, - revision: str = DEFAULT_REPOSITORY_REVISION, - max_retries: int = 5, - ) -> CommitInfo: - """Commit with application-level exponential backoff retry. - - Retries on transient errors (5xx, ConnectionError) and specific - retryable 4xx errors (e.g. git ref conflicts). - Raises immediately on non-retryable client errors (4xx). - """ - last_error = None - for attempt in range(max_retries): - try: - return self.create_commit( - repo_id=repo_id, - operations=operations, - commit_message=commit_message, - commit_description=commit_description, - token=token, - repo_type=repo_type, - revision=revision, - ) - except (ConnectionError, requests.exceptions.ConnectionError) as e: - last_error = e - # Defensive: create_commit raises ValueError, kept for future-proofing - except (HTTPError, requests.exceptions.HTTPError) as e: - if hasattr(e, 'response') and e.response is not None: - if 400 <= e.response.status_code < 500: - raise - last_error = e - except ValueError as e: - error_str = str(e) - if re.search(r'HTTP 4\d{2}', error_str): - retryable_patterns = [ - 'Could not update refs', - 'try again', - ] - if not any(p in error_str for p in retryable_patterns): - raise - last_error = e - except Exception as e: - last_error = e - - wait = min(2 ** attempt, 60) - logger.warning( - f'Commit attempt {attempt + 1}/{max_retries} failed: {last_error}, ' - f'retrying in {wait}s ...') - time.sleep(wait) - - raise RuntimeError( - f'Commit failed after {max_retries} attempts: {last_error}' - ) from last_error - - def _build_batch_operations( - self, - results: list, - repo_type: str, - ) -> list: - """Build CommitOperationAdd list from upload results.""" - operations = [] - for item_d in results: - opt = CommitOperationAdd( - path_in_repo=item_d['file_path_in_repo'], - path_or_fileobj=item_d['file_path'], - file_hash_info=item_d['file_hash_info'], - ) - opt._upload_mode = 'lfs' if self.upload_checker.is_lfs( - item_d['file_path'], repo_type) else 'normal' - opt._is_uploaded = item_d['is_uploaded'] - operations.append(opt) - return operations - - def upload_folder( - self, - *, - repo_id: str, - folder_path: Union[str, Path, List[str], List[Path]], - path_in_repo: Optional[str] = '', - commit_message: Optional[str] = None, - commit_description: Optional[str] = None, - token: Union[str, None] = None, - repo_type: Optional[str] = REPO_TYPE_MODEL, - allow_patterns: Optional[Union[List[str], str]] = None, - ignore_patterns: Optional[Union[List[str], str]] = None, - max_workers: int = DEFAULT_MAX_WORKERS, - use_cache: bool = UPLOAD_USE_CACHE, - revision: Optional[str] = DEFAULT_REPOSITORY_REVISION, - ) -> Optional[Union[CommitInfo, List[CommitInfo]]]: - """Upload a folder to ModelScope Hub with resumable support. - - Upload files from a local folder (or explicit file list) to a remote - repository, with automatic batching, parallel upload, and progressive - retry fallback (ReAct) for failed files. - - Args: - repo_id: Repository identifier in 'owner/repo' format. - folder_path: Local folder path, or a list of (path_in_repo, local_path) tuples. - path_in_repo: Target directory path within the repository. - commit_message: Commit message for the upload. - commit_description: Optional extended commit description. - revision: Branch or tag name (default: 'master'). - token: Authentication token. If None, uses stored credentials. - repo_type: One of 'model', 'dataset', or 'space'. - ignore_patterns: Glob patterns for files to exclude. - max_workers: Max concurrent upload threads. - use_cache: If True, uses .ms_upload_cache for resumable uploads. - Files with matching path, mtime, and size that are already - committed will be skipped automatically. - - Returns: - None if all files were already committed (nothing to do). - A single CommitInfo if only one batch was committed. - A list of CommitInfo if multiple batches were committed. - - Raises: - ValueError: If folder_path is empty or contains no valid files. - RuntimeError: If any files remain failed after all retry rounds, - with a message indicating the count and a retry hint. - """ - start_time = time.time() - - if not repo_id: - raise ValueError('The arg `repo_id` cannot be empty!') - - if folder_path is None: - raise ValueError('The arg `folder_path` cannot be None!') - - if repo_type not in REPO_TYPE_SUPPORT: - raise ValueError(f'Invalid repo type: {repo_type}, supported repos: {REPO_TYPE_SUPPORT}') - - # Check authentication first - self.get_cookies(access_token=token, cookies_required=True) - - allow_patterns = allow_patterns if allow_patterns else None - ignore_patterns = ignore_patterns if ignore_patterns else None - - # Ignore .git .cache folders - if ignore_patterns is None: - ignore_patterns = [] - elif isinstance(ignore_patterns, str): - ignore_patterns = [ignore_patterns] - ignore_patterns += DEFAULT_IGNORE_PATTERNS - - # Cover the ignore patterns if both allow and ignore patterns are provided - if allow_patterns is not None: - ignore_patterns = [ - p for p in ignore_patterns if p not in allow_patterns - ] - - commit_message = ( - commit_message if commit_message is not None else f'Upload to {repo_id} on ModelScope hub' - ) - commit_description = commit_description or 'Uploading files' - - # Exclude internal cache/checkpoint files from upload at any directory depth - _internal_files = [UPLOAD_HASH_CACHE_FILE, _LEGACY_PROGRESS_FILE] - _internal_ignore = [p for f in _internal_files for p in (f, f'*/{f}')] - if ignore_patterns is None: - ignore_patterns = _internal_ignore - elif isinstance(ignore_patterns, str): - ignore_patterns = [ignore_patterns] + _internal_ignore - else: - ignore_patterns = list(ignore_patterns) + _internal_ignore - - # Get the list of files to upload, e.g. [('data/abc.png', '/path/to/abc.png'), ...] - logger.info('Preparing files to upload ...') - prepared_repo_objects = self._prepare_upload_folder( - folder_path_or_files=folder_path, - path_in_repo=path_in_repo, - allow_patterns=allow_patterns, - ignore_patterns=ignore_patterns, - ) - if len(prepared_repo_objects) == 0: - raise ValueError(f'No files to upload in the folder: {folder_path} !') - - logger.info(f'Checking {len(prepared_repo_objects)} files to upload ...') - self.upload_checker.check_normal_files( - file_path_list=[item for _, item in prepared_repo_objects], - repo_type=repo_type, - ) - - self.create_repo(repo_id=repo_id, - token=token, - repo_type=repo_type, - endpoint=self.endpoint, - exist_ok=True, - create_default_config=False) - - # Sort for deterministic batch assignment - sorted_files = sorted(prepared_repo_objects, key=lambda x: x[0]) - - # Calculate batch size (adaptive or fixed) - if UPLOAD_ADAPTIVE_BATCH_SIZE: - commit_batch_size = _calculate_adaptive_batch_size(len(sorted_files)) - logger.info( - f'Adaptive batch size: {commit_batch_size} ' - f'(for {len(sorted_files)} files)') - else: - commit_batch_size = ( - UPLOAD_COMMIT_BATCH_SIZE - if UPLOAD_COMMIT_BATCH_SIZE > 0 - else len(sorted_files)) - - # Initialize unified upload tracker for resume support - folder_path_resolved = Path(folder_path).resolve() \ - if isinstance(folder_path, (str, Path)) else Path(folder_path[0]).resolve().parent - if use_cache: - cache_path = folder_path_resolved / UPLOAD_HASH_CACHE_FILE - tracker = UploadTracker(cache_path, repo_id=repo_id) - else: - tracker = NullTracker() - batch_tracker = BatchTracker(len(sorted_files), commit_batch_size) - - # File-level filtering: skip individually committed files - files_to_upload = [] - skipped_indices = set() - for file_idx, (path_in_repo, file_path) in enumerate(sorted_files): - if isinstance(file_path, (str, os.PathLike)): - try: - st = os.stat(file_path) - if tracker.is_committed(path_in_repo, st.st_mtime, st.st_size): - skipped_indices.add(file_idx) - batch_tracker.mark_file_skipped(file_idx) - continue - except OSError as e: - logger.warning( - f'Cannot stat file {path_in_repo}, will re-upload: {e}') - files_to_upload.append((file_idx, (path_in_repo, file_path))) - - # Batch pre-validation for files with cached hashes - pre_validated_map = {} # oid -> upload_url or None - hash_info_map = {} # file_idx -> (hash_info, file_stat) - files_need_hash = [] # files without cached hash - - for file_idx, (path_in_repo, file_path) in files_to_upload: - if isinstance(file_path, (str, os.PathLike)): - try: - st = os.stat(file_path) - cached = tracker.get_hash( - path_in_repo, st.st_mtime, st.st_size) - if cached is not None: - hash_info_map[file_idx] = (cached, st) - continue - except OSError: - pass - files_need_hash.append((file_idx, (path_in_repo, file_path))) - - # Batch validate cached hashes against server - if hash_info_map: - objects = [ - {'oid': info['file_hash'], 'size': info['file_size']} - for info, _ in hash_info_map.values() - ] - validated = self._validate_blob( - repo_id=repo_id, repo_type=repo_type, - objects=objects, token=token) - pre_validated_map = validated - reused = sum(1 for v in validated.values() if v is None) - logger.info( - f'Pre-validated {len(objects)} cached hash(es): ' - f'{reused} globally existing, ' - f'{len(objects) - reused} need upload.') - - skipped_count = len(skipped_indices) - if skipped_count > 0: - logger.info(f'{skipped_count} file(s) already committed, skipping.') - - logger.info( - f'Scan complete: {len(sorted_files)} total, ' - f'{skipped_count} committed (skip), ' - f'{len(files_to_upload)} to process.') - - logger.info( - f'Uploading {len(files_to_upload)} file(s) in {batch_tracker.num_batches} batch(es) ' - f'of size {commit_batch_size} (pipeline mode).') - - # Submit upload tasks to thread pool - def _upload_worker(file_idx: int, file_info: tuple, - pre_validated=None): - path_in_repo, file_path = file_info - try: - logger.debug(f'Uploading: {path_in_repo} ...') - result = self._upload_single_file( - path_in_repo, file_path, - repo_id=repo_id, repo_type=repo_type, - token=token, tracker=tracker, - pre_validated=pre_validated) - logger.debug(f'Uploaded: {path_in_repo}') - batch_tracker.record_success(file_idx, result) - except Exception as e: - logger.error(f'Upload failed: {path_in_repo} - {e}') - batch_tracker.record_failure(file_idx, file_info, e) - - # Pipeline: consume batches in order, commit as each becomes ready - commit_infos: List[CommitInfo] = [] - all_results: List[dict] = [] - total_failed_files: List[tuple] = [] - num_batches = batch_tracker.num_batches - + inner = object.__getattribute__(self, '_api') + except AttributeError as exc: + raise AttributeError(name) from exc try: - with ThreadPoolExecutor(max_workers=max_workers) as executor: - for file_idx, file_info in files_to_upload: - # Look up pre-validated status by file hash - pv = None - if file_idx in hash_info_map: - cached_hash = hash_info_map[file_idx][0]['file_hash'] - pv = pre_validated_map.get(cached_hash) - # pv: None=exists(True), str=upload_url - if pv is None: - pv = True # globally existing, skip upload - executor.submit(_upload_worker, file_idx, file_info, pv) - - for batch_idx in tqdm(range(num_batches), desc='[Committing batches]', total=num_batches): - # Skip fully-committed batches - batch_start = batch_idx * commit_batch_size - batch_end = min(batch_start + commit_batch_size, len(sorted_files)) - if all(i in skipped_indices for i in range(batch_start, batch_end)): - logger.info(f'Batch {batch_idx + 1}/{num_batches} fully committed, skipping.') - continue - - results, failures = batch_tracker.wait_for_batch(batch_idx) - - if failures: - total_failed_files.extend(failures) - for item, err in failures: - logger.error(f' Failed: {item[0]} - {err}') - - # Mark successfully uploaded files in tracker (BEFORE commit attempt) - self._track_uploaded_batch(tracker, results) - - operations = self._build_batch_operations(results, repo_type) - if not operations: - logger.error( - f'Batch {batch_idx + 1}/{num_batches}: ' - f'all files failed, skipping commit.') - continue - - batch_commit_message = ( - f'{commit_message} (batch {batch_idx + 1}/{num_batches})') - try: - commit_info = self._commit_with_retry( - repo_id=repo_id, - operations=operations, - commit_message=batch_commit_message, - commit_description=commit_description, - token=token, - repo_type=repo_type, - revision=revision, - ) - commit_infos.append(commit_info) - all_results.extend(results) - logger.info( - f'Batch {batch_idx + 1}/{num_batches}: ' - f'committed {len(results)} file(s).') - # Mark all files in this batch as committed - self._track_committed_batch(tracker, results) - except Exception as e: - logger.error( - f'Batch {batch_idx + 1}/{num_batches} commit failed: {e}') - category = classify_error(e) - if not category.is_retryable: - # Permanent error: mark files as failed, do not retry - for r in results: - tracker.mark_failed( - r['file_path_in_repo'], r['file_mtime'], - r['file_size_on_disk'], - error_type='commit_' + category.value) - logger.error( - f'Batch {batch_idx + 1}/{num_batches}: ' - f'permanent failure ({category.value}), ' - f'{len(results)} file(s) will not be retried.') - else: - # Transient error: recover to retry queue - for r in results: - total_failed_files.append( - ((r['file_path_in_repo'], r['file_path']), e)) - logger.warning( - f'Batch {batch_idx + 1}/{num_batches}: ' - f'{len(results)} file(s) recovered to retry queue ' - f'(error_category={category.value}).') - finally: - tracker.save() - - # ReAct progressive retry fallback - if total_failed_files and UPLOAD_REACT_ENABLED: - total_failed_files, react_commits, react_results = self._retry_failed_files_react( - failed_files=total_failed_files, - tracker=tracker, - repo_id=repo_id, - repo_type=repo_type, - token=token, - commit_message=commit_message, - commit_description=commit_description, - revision=revision, - max_workers=max_workers, - ) - commit_infos.extend(react_commits) - all_results.extend(react_results) - elif total_failed_files: - # Simple fallback when ReAct is disabled - for retry_round in range(UPLOAD_FAILED_FILE_MAX_RETRIES): - if not total_failed_files: - break - logger.info( - f'Retry round {retry_round + 1}/{UPLOAD_FAILED_FILE_MAX_RETRIES}: ' - f're-uploading {len(total_failed_files)} failed file(s) ...') - retry_failures = [] - retry_successes = [] - for (path_in_repo, file_path), _err in total_failed_files: - try: - result = self._upload_single_file( - path_in_repo, file_path, - repo_id=repo_id, repo_type=repo_type, - token=token, tracker=tracker) - retry_successes.append(result) - except Exception as e: - logger.error(f' Retry failed: {path_in_repo} - {e}') - retry_failures.append(((path_in_repo, file_path), e)) - if retry_successes: - self._track_uploaded_batch(tracker, retry_successes) - operations = self._build_batch_operations( - retry_successes, repo_type) - if operations: - try: - commit_info = self._commit_with_retry( - repo_id=repo_id, - operations=operations, - commit_message=f'{commit_message} (retry round {retry_round + 1})', - commit_description=commit_description, - token=token, - repo_type=repo_type, - revision=revision) - commit_infos.append(commit_info) - all_results.extend(retry_successes) - self._track_committed_batch(tracker, retry_successes) - logger.info( - f' Retry round {retry_round + 1}: ' - f'committed {len(retry_successes)} file(s).') - except Exception as e: - logger.error( - f' Retry round {retry_round + 1} commit failed: {e}') - category = classify_error(e) - if not category.is_retryable: - for result in retry_successes: - tracker.mark_failed( - result['file_path_in_repo'], - result['file_mtime'], - result['file_size_on_disk'], - error_type='commit_' + category.value) - else: - for result in retry_successes: - retry_failures.append( - ((result['file_path_in_repo'], - result.get('file_path', '')), e)) - total_failed_files = retry_failures - - # Final tracker save - tracker.save() - - # Upload report - elapsed = time.time() - start_time - total_files = len(sorted_files) - failed_count = len(total_failed_files) - reused_count = sum( - 1 for r in all_results if r.get('is_reused')) - uploaded_count = sum( - 1 for r in all_results if not r.get('is_reused')) - - print('=' * 60) - print('Upload Report') - print('-' * 60) - print(f' Total files : {total_files}') - print(f' Skipped (cached) : {skipped_count}') - print(f' Existed (server) : {reused_count}') - print(f' Uploaded (PUT) : {uploaded_count}') - print(f' Failed : {failed_count}') - committed_count = reused_count + uploaded_count - print(f' Committed : {committed_count}') - print(f' Elapsed : {elapsed:.1f}s') - print('=' * 60) - - # Final error if there are still failed files after all retries - if total_failed_files: - for (path_in_repo, _), err in total_failed_files: - logger.error(f' - {path_in_repo}: {type(err).__name__}: {err}') - succeeded = total_files - failed_count - raise RuntimeError( - f'ERROR - {failed_count} file(s) failed to upload. ' - f'Please manually try again. Successfully uploaded ' - f'{succeeded} file(s) will be automatically skipped ' - f'during the retry.') - - if not commit_infos: - if skipped_count == len(sorted_files): - logger.info('All files were already committed.') - return None - return None - - return commit_infos[0] if len(commit_infos) == 1 else commit_infos - - def _retry_failed_files_react( - self, - failed_files, - tracker, - repo_id, - repo_type, - token, - commit_message, - commit_description, - revision, - max_workers, - ): - """ReAct-style progressive retry for failed files. - - Implements Reason-Act-Observe loop with three escalating rounds: - Round 1: Parallel retry with reduced concurrency (workers//2, batch=16) - Round 2: Serial retry with exponential backoff (delay * 2^min(i, max_exp)) - Round 3: Single-file commit with long delays (one file per commit) - - Files that exceed the per-file retry limit are classified as permanent - failures and will not be retried further. - - Args: - failed_files: List of ((path_in_repo, file_path), error) tuples. - tracker: UploadTracker or NullTracker instance. - repo_id: Repository identifier. - repo_type: Repository type. - token: Authentication token. - commit_message: Base commit message. - commit_description: Commit description. - revision: Branch or tag name. - max_workers: Max upload concurrency from caller. - - Returns: - Tuple of (all_failures, commit_infos, all_successes) where: - - all_failures: list of ((path_in_repo, file_path), error) for - files that could not be resolved (permanent + exhausted retries). - - commit_infos: list of CommitInfo for successful retry commits. - - all_successes: list of upload result dicts for successfully - retried files (to be merged into the upload report). - """ - commit_infos = [] - all_successes: list = [] - retry_counts: dict = {} # path_in_repo -> cumulative retry count - permanent_failures = [] - retryable = list(failed_files) - - # Separate permanent failures - remaining = [] - for item_err in retryable: - (path_in_repo, file_path), err = item_err - category = classify_error(err) - if category.is_retryable: - remaining.append(item_err) - else: - permanent_failures.append(item_err) - try: - st = os.stat(file_path) if isinstance( - file_path, (str, os.PathLike)) else None - except OSError: - st = None - tracker.mark_failed( - path_in_repo, - st.st_mtime if st else 0, - st.st_size if st else 0, - error_type=category.value) - logger.error( - f'[ReAct] Permanent failure: {path_in_repo} ' - f'({category.value}: {err})') - retryable = remaining - - round_configs = [ - { - 'name': 'Round 1 (parallel)', - 'parallel': True, - 'workers': max(1, max_workers // 2), - 'batch_size': 16, - 'delay': 0, - }, - { - 'name': 'Round 2 (serial+backoff)', - 'parallel': False, - 'workers': 1, - 'batch_size': 8, - 'delay': UPLOAD_REACT_ROUND2_BASE_DELAY, - }, - { - 'name': 'Round 3 (single-file)', - 'parallel': False, - 'workers': 1, - 'batch_size': 1, - 'delay': UPLOAD_REACT_ROUND3_FILE_DELAY, - }, - ] - - for round_idx, cfg in enumerate(round_configs): - if not retryable: - break - - round_name = cfg['name'] - logger.info( - f'[ReAct] {round_name}: retrying {len(retryable)} file(s) ...') - - round_successes = [] - round_failures = [] - - # ACT: upload files - if cfg['parallel'] and len(retryable) > 1: - from concurrent.futures import ThreadPoolExecutor, as_completed - with ThreadPoolExecutor(max_workers=cfg['workers']) as executor: - future_map = {} - for (path_in_repo, file_path), _err in retryable: - future = executor.submit( - self._upload_single_file, - path_in_repo, file_path, - repo_id=repo_id, repo_type=repo_type, - token=token, tracker=tracker) - future_map[future] = (path_in_repo, file_path) - for future in as_completed(future_map): - path_in_repo, file_path = future_map[future] - try: - result = future.result() - round_successes.append(result) - except Exception as e: - round_failures.append( - ((path_in_repo, file_path), e)) - else: - for i, ((path_in_repo, file_path), _err) in enumerate(retryable): - if cfg['delay'] > 0 and i > 0: - delay = (cfg['delay'] * (2 ** min(i, UPLOAD_REACT_BACKOFF_MAX_EXPONENT)) - if round_idx == 1 - else cfg['delay']) - delay = min(delay, UPLOAD_REACT_MAX_DELAY) - logger.info( - f'[ReAct] Waiting {delay}s before ' - f'retrying {path_in_repo} ...') - time.sleep(delay) - try: - result = self._upload_single_file( - path_in_repo, file_path, - repo_id=repo_id, repo_type=repo_type, - token=token, tracker=tracker) - round_successes.append(result) - except Exception as e: - logger.error( - f'[ReAct] {round_name}: ' - f'failed {path_in_repo} - {e}') - round_failures.append( - ((path_in_repo, file_path), e)) - - all_successes.extend(round_successes) - - # ACT: commit successful uploads in small batches - batch_size = min(cfg['batch_size'], max(1, len(round_successes))) - for batch_start in range(0, len(round_successes), batch_size): - batch = round_successes[batch_start:batch_start + batch_size] - # Mark uploaded - self._track_uploaded_batch(tracker, batch) - - operations = self._build_batch_operations(batch, repo_type) - if not operations: - continue - try: - commit_info = self._commit_with_retry( - repo_id=repo_id, - operations=operations, - commit_message=( - f'{commit_message} ({round_name})'), - commit_description=commit_description, - token=token, - repo_type=repo_type, - revision=revision) - commit_infos.append(commit_info) - # Mark committed only after successful commit - self._track_committed_batch(tracker, batch) - logger.info( - f'[ReAct] {round_name}: ' - f'committed {len(batch)} file(s).') - except Exception as e: - logger.error( - f'[ReAct] {round_name} commit failed: {e}') - category = classify_error(e) - if not category.is_retryable: - for r in batch: - tracker.mark_failed( - r['file_path_in_repo'], r['file_mtime'], - r['file_size_on_disk'], - error_type='commit_' + category.value) - else: - for r in batch: - round_failures.append( - ((r['file_path_in_repo'], - r['file_path']), e)) - - # OBSERVE: classify new failures, enforce per-file retry limit - new_retryable = [] - for item_err in round_failures: - (path_in_repo, file_path), err = item_err - retry_counts[path_in_repo] = retry_counts.get(path_in_repo, 0) + 1 - if retry_counts[path_in_repo] >= 3: - permanent_failures.append(item_err) - try: - st = os.stat(file_path) if isinstance( - file_path, (str, os.PathLike)) else None - except OSError: - st = None - tracker.mark_failed( - path_in_repo, - st.st_mtime if st else 0, - st.st_size if st else 0, - error_type='max_retries_exceeded') - logger.error( - f'[ReAct] Max retries exceeded for {path_in_repo}') - continue - category = classify_error(err) - if category.is_retryable: - new_retryable.append(item_err) - else: - permanent_failures.append(item_err) - try: - st = os.stat(file_path) if isinstance( - file_path, (str, os.PathLike)) else None - except OSError: - st = None - tracker.mark_failed( - path_in_repo, - st.st_mtime if st else 0, - st.st_size if st else 0, - error_type=category.value) - logger.error( - f'[ReAct] Permanent failure: {path_in_repo} ' - f'({category.value})') - - progress = len(retryable) - len(new_retryable) - if progress > 0: - logger.info( - f'[ReAct] {round_name}: made progress — ' - f'{progress} file(s) resolved, ' - f'{len(new_retryable)} remaining.') - elif new_retryable: - logger.warning( - f'[ReAct] {round_name}: no progress, ' - f'escalating to next round.') - - retryable = new_retryable - - # Any remaining retryable failures become permanent at this point - all_failures = permanent_failures + retryable - if retryable: - logger.error( - f'[ReAct] {len(retryable)} file(s) still failing ' - f'after all retry rounds.') - - return all_failures, commit_infos, all_successes - - def _upload_blob( - self, - *, - repo_id: str, - repo_type: str, - sha256: str, - size: int, - data: Union[str, Path, bytes, BinaryIO], - disable_tqdm: Optional[bool] = False, - tqdm_desc: Optional[str] = '[Uploading]', - buffer_size_mb: Optional[int] = 16, - token: Optional[str] = None, - pre_validated=None, - ) -> dict: - - res_d: dict = dict( - url=None, - is_uploaded=False, - is_reused=False, - status_code=None, - status_msg=None, - ) - - if pre_validated is True: - logger.info(f'Blob {sha256[:8]} already exists globally, reuse.') - res_d['is_uploaded'] = True - res_d['is_reused'] = True - return res_d - - if isinstance(pre_validated, str): - upload_url = pre_validated - else: - validated = self._validate_blob( - repo_id=repo_id, - repo_type=repo_type, - objects=[{'oid': sha256, 'size': size}], - token=token, - ) - upload_url = validated.get(sha256) - if upload_url is None: - logger.info(f'Blob {sha256[:8]} already exists globally, reuse.') - res_d['is_uploaded'] = True - res_d['is_reused'] = True - return res_d - - cookies = self.get_cookies(access_token=token, cookies_required=True) - cookies = dict(cookies) if cookies else None - if cookies is None: - raise ValueError('Token does not exist, please login first.') - - self.headers.update({'Cookie': f"m_session_id={cookies['m_session_id']}"}) - headers = self.builder_headers(self.headers) - - chunk_size = buffer_size_mb * 1024 * 1024 - headers['Content-Length'] = str(size) - - with tqdm( - total=size, - unit='B', - unit_scale=True, - desc=tqdm_desc, - disable=disable_tqdm - ) as pbar: - if isinstance(data, (str, Path)): - with open(data, 'rb') as f: - stream = _CountedReadStream( - f, size, pbar, chunk_size) - response = self.session.put( - upload_url, - headers=headers, - data=stream, - timeout=UPLOAD_BLOB_TIMEOUT, - ) - stream.verify_complete() - - elif isinstance(data, bytes): - stream = _CountedReadStream( - io.BytesIO(data), size, pbar, chunk_size) - response = self.session.put( - upload_url, - headers=headers, - data=stream, - timeout=UPLOAD_BLOB_TIMEOUT, - ) - stream.verify_complete() - - elif isinstance(data, io.BufferedIOBase): - stream = _CountedReadStream( - data, size, pbar, chunk_size) - response = self.session.put( - upload_url, - headers=headers, - data=stream, - timeout=UPLOAD_BLOB_TIMEOUT, - ) - stream.verify_complete() - - else: - raise ValueError('Invalid data type to upload') - - raise_for_http_status(rsp=response) - resp = response.json() - raise_on_error(rsp=resp) - - res_d['url'] = upload_url - res_d['is_uploaded'] = True - res_d['status_code'] = resp['Code'] - res_d['status_msg'] = resp['Message'] - - return res_d - - def _validate_blob( - self, - *, - repo_id: str, - repo_type: str, - objects: List[Dict[str, Any]], - endpoint: Optional[str] = None, - token: Optional[str] = None, - ) -> Dict[str, Optional[str]]: - """Validate whether blobs need uploading. - - Queries the LFS batch API in chunks of UPLOAD_VALIDATE_BLOB_BATCH_SIZE. - - Args: - repo_id: The repo id on ModelScope. - repo_type: The repo type ('dataset', 'model', etc.). - objects: Objects to check, each with 'oid' (sha256) and 'size'. - endpoint: API endpoint override. - token: Access token. - - Returns: - Dict mapping oid -> upload_url (needs upload) or None (already exists). - """ - if not endpoint: - endpoint = self.endpoint - - result: Dict[str, Optional[str]] = {} - batch_size = UPLOAD_VALIDATE_BLOB_BATCH_SIZE - - for i in range(0, len(objects), batch_size): - chunk = objects[i:i + batch_size] - - url = f'{endpoint}/api/v1/repos/{repo_type}s/{repo_id}/info/lfs/objects/batch' - payload = { - 'operation': 'upload', - 'objects': chunk, - } - - cookies = self.get_cookies(access_token=token, cookies_required=True) - response = self.session.post( - url, - headers=self.builder_headers(self.headers), - data=json.dumps(payload), - cookies=cookies - ) - - raise_for_http_status(rsp=response) - resp = response.json() - raise_on_error(rsp=resp) - - resp_objects = resp['Data']['objects'] - needs_upload = set() - for obj in resp_objects: - actions = obj.get('actions', {}) - upload_action = actions.get('upload') - if upload_action: - result[obj['oid']] = upload_action['href'] - needs_upload.add(obj['oid']) - - # Objects not needing upload are globally existing - for o in chunk: - if o['oid'] not in needs_upload: - result[o['oid']] = None - - return result - - def _prepare_upload_folder( - self, - folder_path_or_files: Union[str, Path, List[str], List[Path]], - path_in_repo: str, - allow_patterns: Optional[Union[List[str], str]] = None, - ignore_patterns: Optional[Union[List[str], str]] = None, - ) -> List[Union[tuple, list]]: - folder_path = None - files_path = None - if isinstance(folder_path_or_files, list): - if os.path.isfile(folder_path_or_files[0]): - files_path = folder_path_or_files - else: - raise ValueError('Uploading multiple folders is not supported now.') - else: - if os.path.isfile(folder_path_or_files): - files_path = [folder_path_or_files] - else: - folder_path = folder_path_or_files - - if files_path is None: - self.upload_checker.check_folder(folder_path) - folder_path = Path(folder_path).expanduser().resolve() - if not folder_path.is_dir(): - raise ValueError(f"Provided path: '{folder_path}' is not a directory") - - # List files from folder - relpath_to_abspath = { - path.relative_to(folder_path).as_posix(): path - for path in sorted(folder_path.glob('**/*')) # sorted to be deterministic - if path.is_file() - } - else: - relpath_to_abspath = {} - for path in files_path: - if os.path.isfile(path): - self.upload_checker.check_file(path) - relpath_to_abspath[os.path.basename(path)] = path - - # Filter files - filtered_repo_objects = list( - RepoUtils.filter_repo_objects( - relpath_to_abspath.keys(), allow_patterns=allow_patterns, ignore_patterns=ignore_patterns - ) - ) - - prefix = f"{path_in_repo.strip('/')}/" if path_in_repo else '' - - prepared_repo_objects = [ - (prefix + relpath, str(relpath_to_abspath[relpath])) - for relpath in filtered_repo_objects - ] - - logger.info(f'Prepared {len(prepared_repo_objects)} files for upload.') - - return prepared_repo_objects - - @staticmethod - def _prepare_commit_payload( - operations: Iterable[CommitOperation], - commit_message: str, - ) -> Dict[str, Any]: - """ - Prepare the commit payload to be sent to the ModelScope hub. - """ - - payload = { - 'commit_message': commit_message, - 'actions': [] - } - - nb_ignored_files = 0 - - # 2. Send operations, one per line - for operation in operations: - - # Skip ignored files - if isinstance(operation, CommitOperationAdd) and operation._should_ignore: - logger.debug(f"Skipping file '{operation.path_in_repo}' in commit (ignored by gitignore file).") - nb_ignored_files += 1 - continue - - # 2.a. Case adding a normal file - if isinstance(operation, CommitOperationAdd) and operation._upload_mode == 'normal': - - commit_action = { - 'action': 'create', - 'path': operation.path_in_repo, - 'type': 'normal', - 'size': operation.upload_info.size, - 'sha256': '', - 'content': operation.b64content().decode(), - 'encoding': 'base64', - } - payload['actions'].append(commit_action) - - # 2.b. Case adding an LFS file - elif isinstance(operation, CommitOperationAdd) and operation._upload_mode == 'lfs': - - commit_action = { - 'action': 'create', - 'path': operation.path_in_repo, - 'type': 'lfs', - 'size': operation.upload_info.size, - 'sha256': operation.upload_info.sha256, - 'content': '', - 'encoding': '', - } - payload['actions'].append(commit_action) - - else: - raise ValueError( - f'Unknown operation to commit. Operation: {operation}. Upload mode:' - f" {getattr(operation, '_upload_mode', None)}" - ) - - if nb_ignored_files > 0: - logger.info(f'Skipped {nb_ignored_files} file(s) in commit (ignored by gitignore file).') - - return payload - - def _get_internal_acceleration_domain(self, internal_timeout: float = 0.2): - """ - Get the internal acceleration domain. - - Args: - internal_timeout (float): The timeout for the request. Default to 0.2s - - Returns: - str: The internal acceleration domain. e.g. `cn-hangzhou`, `cn-zhangjiakou` - """ - - def send_request(url: str, timeout: float): - try: - response = requests.get(url, timeout=timeout) - response.raise_for_status() - except requests.exceptions.RequestException: - response = None - - return response - - internal_url = f'{self.endpoint}/api/v1/repos/internalAccelerationInfo' - - # Get internal url and region for acceleration - internal_info_response = send_request(url=internal_url, timeout=internal_timeout) - region_id: str = '' - if internal_info_response is not None: - internal_info_response = internal_info_response.json() - if 'Data' in internal_info_response: - query_addr = internal_info_response['Data']['InternalRegionQueryAddress'] - else: - query_addr: str = '' - - if query_addr: - domain_response = send_request(query_addr, timeout=internal_timeout) - if domain_response is not None: - region_id = domain_response.text.strip() - - return region_id - - def delete_files(self, - repo_id: str, - repo_type: str, - delete_patterns: Union[str, List[str]], - *, - revision: Optional[str] = DEFAULT_MODEL_REVISION, - endpoint: Optional[str] = None, - token: Optional[str] = None) -> Dict[str, Any]: - """ - Delete files in batch using glob (wildcard) patterns, e.g. '*.py', 'data/*.csv', 'foo*', etc. - - Example: - # Delete all Python and Markdown files in a model repo - api.delete_files( - repo_id='your_username/your_model', - repo_type=REPO_TYPE_MODEL, - delete_patterns=['*.py', '*.md'] - ) - - # Delete all CSV files in the data/ directory of a dataset repo - api.delete_files( - repo_id='your_username/your_dataset', - repo_type=REPO_TYPE_DATASET, - delete_patterns='data/*.csv' - ) - - Args: - repo_id (str): 'owner/repo_name' or 'owner/dataset_name', e.g. 'Koko/my_model' - repo_type (str): REPO_TYPE_MODEL or REPO_TYPE_DATASET - delete_patterns (str or List[str]): List of glob patterns, e.g. '*.py', 'data/*.csv', 'foo*' - revision (str, optional): Branch or tag name - endpoint (str, optional): API endpoint - token (str, optional): Access token - Returns: - dict: Deletion result - """ - if repo_type not in REPO_TYPE_SUPPORT: - raise ValueError(f'Unsupported repo_type: {repo_type}') - if not delete_patterns: - raise ValueError('delete_patterns cannot be empty') - if isinstance(delete_patterns, str): - delete_patterns = [delete_patterns] - - cookies = self.get_cookies(access_token=token, cookies_required=True) - if not endpoint: - endpoint = self.endpoint - if cookies is None: - raise ValueError('Token does not exist, please login first.') - headers = self.builder_headers(self.headers) - - # List all files in the repo - if repo_type == REPO_TYPE_MODEL: - files = self.get_model_files( - repo_id, - revision=revision or DEFAULT_MODEL_REVISION, - recursive=True, - endpoint=endpoint, - use_cookies=cookies, - ) - file_paths = [f['Path'] for f in files] - elif repo_type == REPO_TYPE_DATASET: - file_paths = [] - _owner, _dataset_name = repo_id.split('/') - _hub_id, _ = self.get_dataset_id_and_type( - dataset_name=_dataset_name, namespace=_owner, endpoint=endpoint, token=token) - page_number = 1 - page_size = 100 - while True: - try: - dataset_files: List[Dict[str, Any]] = self.get_dataset_files( - repo_id=repo_id, - revision=revision or DEFAULT_DATASET_REVISION, - recursive=True, - page_number=page_number, - page_size=page_size, - endpoint=endpoint, - token=token, - dataset_hub_id=_hub_id, - ) - except Exception as e: - logger.error(f'Get dataset: {repo_id} file list failed, message: {str(e)}') - break - - # Parse data (Type: 'tree' or 'blob') - for file_info_d in dataset_files: - if file_info_d['Type'] != 'tree': - file_paths.append(file_info_d['Path']) - - if len(dataset_files) < page_size: - break - - page_number += 1 - else: - raise ValueError(f'Unsupported repo_type: {repo_type}, supported repos: {REPO_TYPE_SUPPORT}') - - # Glob pattern matching - to_delete = [] - for path in file_paths: - for delete_pattern in delete_patterns: - if fnmatch.fnmatch(path, delete_pattern): - to_delete.append(path) - break - - deleted_files, failed_files = [], [] - for path in to_delete: - try: - if repo_type == REPO_TYPE_MODEL: - owner, repo_name = repo_id.split('/') - url = f'{endpoint}/api/v1/models/{owner}/{repo_name}/file' - params = { - 'Revision': revision or DEFAULT_MODEL_REVISION, - 'FilePath': path - } - elif repo_type == REPO_TYPE_DATASET: - owner, dataset_name = repo_id.split('/') - url = f'{endpoint}/api/v1/datasets/{owner}/{dataset_name}/repo' - params = { - 'FilePath': path - } - else: - raise ValueError(f'Unsupported repo_type: {repo_type}, supported repos: {REPO_TYPE_SUPPORT}') - - r = self.session.delete(url, params=params, cookies=cookies, headers=headers) - raise_for_http_status(r) - resp = r.json() - raise_on_error(resp) - deleted_files.append(path) - except Exception as e: - failed_files.append(path) - logger.error(f'Failed to delete {path}: {str(e)}') - - return { - 'deleted_files': deleted_files, - 'failed_files': failed_files, - 'total_files': len(to_delete) - } - - def set_repo_visibility(self, - repo_id: str, - repo_type: Literal['model', 'dataset'], - visibility: Literal['private', 'public'], - token: Union[str, None] = None, - gated_mode: Optional[bool] = None, - ) -> dict: - """ - Set the visibility of a repo. - - Args: - repo_id (str): The repo id in the format of `owner_name/repo_name`. - repo_type (Literal['model', 'dataset']): The repo type, `model` or `dataset`. - visibility (Literal['private', 'public']): The visibility to set, `private` or `public`. - token (Union[str, None]): The access token. If None, will use the cookies from the local cache. - See `https://modelscope.cn/my/myaccesstoken` to get your token. - gated_mode (Optional[bool]): Gated mode for private repos. - True = gated (application-based download), False = off (normal private). - Only effective when visibility is ``private``. - - Returns: - dict: The response from the server. - """ - if not repo_id: - raise ValueError('The arg `repo_id` cannot be empty!') - - if visibility not in ['private', 'public']: - raise ValueError(f'Invalid visibility: {visibility}, supported visibilities: `private`, `public`') - - visibility_map: Dict[str, int] = {v: k for k, v in VisibilityMap.items()} - visibility_code: int = visibility_map.get(visibility, 5) - cookies = self.get_cookies(access_token=token, cookies_required=True) - - if gated_mode is not None and visibility != 'private': - logger.warning('gated_mode is only effective when visibility is private, ignored.') - gated_mode = None - - if repo_type == REPO_TYPE_MODEL: - model_info = self.get_model(model_id=repo_id, token=token) - path = f'{self.endpoint}/api/v1/models/{repo_id}' - tasks = model_info.get('Tasks') - model_tasks = '' - if isinstance(tasks, list) and tasks: - first = tasks[0] - if isinstance(first, dict) and first: - model_tasks = first.get('name') - if gated_mode is not None: - pm = 1 if gated_mode else 2 - else: - pm = model_info.get('ProtectedMode', 2) - payload = { - 'ChineseName': model_info.get('ChineseName', ''), - 'ModelFramework': model_info.get('ModelFramework', 'Pytorch'), - 'Visibility': visibility_code, - 'ProtectedMode': pm, - 'ApprovalMode': model_info.get('ApprovalMode', 2), - 'Description': model_info.get('Description', ''), - 'AigcType': model_info.get('AigcType', ''), - 'VisionFoundation': model_info.get('VisionFoundation', ''), - 'ModelCover': model_info.get('ModelCover', ''), - 'SubScientificField': model_info.get('SubScientificField', None), - 'ScientificField': model_info.get('NEXA', {}).get('ScientificField', ''), - 'Source': model_info.get('NEXA', {}).get('Source', ''), - 'ModelTask': model_tasks, - 'License': model_info.get('License', ''), - } - elif repo_type == REPO_TYPE_DATASET: - - repo_id_parts = repo_id.split('/') - if len(repo_id_parts) != 2 or not all(repo_id_parts): - raise ValueError(f'Invalid dataset repo_id: {repo_id}, should be in format of `owner/dataset_name`') - - dataset_idx, _ = self.get_dataset_id_and_type( - dataset_name=repo_id_parts[1], - namespace=repo_id_parts[0], - token=token - ) - - path = f'{self.endpoint}/api/v1/datasets/{dataset_idx}' - payload = { - 'Visibility': visibility_code, - 'ProtectedMode': (1 if gated_mode else 2) if gated_mode is not None else 2, - } - else: - raise ValueError(f'Invalid repo type: {repo_type}, supported repos: {REPO_TYPE_SUPPORT}') - - r = self.session.put( - path, - json=payload, - cookies=cookies, - headers=self.builder_headers(self.headers)) - - raise_for_http_status(r) - resp = r.json() - raise_on_error(resp) - - return resp - - # ============= Collection API ============= - def get_collection(self, - collection_id: str, - repo_type: str = 'skill', - page_number: int = 1, - page_size: int = 50, - endpoint: Optional[str] = None) -> dict: - """Get collection details and its elements. - - Args: - collection_id (str): The collection ID (Fid). - repo_type (str): Element type filter, only 'skill' is supported currently. - page_number (int): Page number for pagination. - page_size (int): Page size for pagination. - - Returns: - dict: Collection details including elements. - - Raises: - ValueError: If repo_type is not 'skill'. - RequestError: If the API request fails. - """ - if not endpoint: - endpoint = self.endpoint - if repo_type != 'skill': - raise ValueError( - f'repo_type={repo_type} is not supported, ' - 'only "skill" is currently supported.') - cookies = self.get_cookies() - path = f'{endpoint}/api/v1/collections' - params = { - 'Fid': collection_id, - 'ElementType': repo_type, - 'PageNumber': page_number, - 'PageSize': page_size, - } - r = self.session.get(path, params=params, cookies=cookies, - headers=self.builder_headers(self.headers)) - raise_for_http_status(r) - d = r.json() - raise_on_error(d) - return d[API_RESPONSE_FIELD_DATA] - - def download_skill(self, skill_id: str, - local_dir: Optional[str] = None, - endpoint: Optional[str] = None) -> str: - """Download a single skill archive and extract it. - - Args: - skill_id (str): The skill identifier in format '/'. - local_dir (Optional[str]): Target directory for extraction. - Defaults to current directory. - - Returns: - str: Path to the extracted skill directory. - - Raises: - ValueError: If skill_id format is invalid. - RequestError: If the download request fails. - """ - if not endpoint: - endpoint = self.endpoint - element_path, element_name = RepoUtils.validate_repo_id(skill_id) - - cookies = self.get_cookies() - url = f'{endpoint}/api/v1/skills/{element_path}/{element_name}/archive/zip/master' - - if local_dir is None: - local_dir = os.getcwd() - os.makedirs(local_dir, exist_ok=True) - - # Build skill directory name: use element_name directly, overwrite if exists, to avoid corrupted state - skill_dir = os.path.join(local_dir, element_name) - - r = self.session.get(url, stream=True, cookies=cookies, - headers=self.builder_headers(self.headers)) - raise_for_http_status(r) - - # Save to temp zip file then extract - zip_path = os.path.join(local_dir, f'{element_name}.zip') - try: - with open(zip_path, 'wb') as f: - for chunk in r.iter_content(chunk_size=8192): - if chunk: - f.write(chunk) - - # Clean existing directory to avoid corrupted state - if os.path.exists(skill_dir): - shutil.rmtree(skill_dir) - os.makedirs(skill_dir, exist_ok=True) - with zipfile.ZipFile(zip_path, 'r') as zf: - zf.extractall(skill_dir) - - # Flatten if zip contains a single top-level directory - entries = os.listdir(skill_dir) - if len(entries) == 1: - nested_dir = os.path.join(skill_dir, entries[0]) - if os.path.isdir(nested_dir): - for item in os.listdir(nested_dir): - shutil.move( - os.path.join(nested_dir, item), - os.path.join(skill_dir, item)) - os.rmdir(nested_dir) - finally: - if os.path.exists(zip_path): - os.remove(zip_path) - - logger.info(f'Skill {element_path}/{element_name} downloaded to {skill_dir}') - return skill_dir + return getattr(inner, name) + except AttributeError: + raise AttributeError( + f"'{type(self).__name__}' object has no attribute '{name}'" + ) from None class ModelScopeConfig: + """Configuration manager — delegates to ``modelscope_hub.HubConfig``. + + Preserves the static-method interface used throughout the legacy + codebase. Class-level attributes are kept for callers that read them + directly (e.g. ``ModelScopeConfig.path_credential``). + """ + path_credential = expanduser(MODELSCOPE_CREDENTIALS_PATH) COOKIES_FILE_NAME = 'cookies' GIT_TOKEN_FILE_NAME = 'git_token' @@ -4392,283 +330,95 @@ class ModelScopeConfig: cookie_expired_warning = False @staticmethod - def make_sure_credential_path_exist(): + def make_sure_credential_path_exist() -> None: + """Ensure the credentials directory exists.""" os.makedirs(ModelScopeConfig.path_credential, exist_ok=True) + get_default_config().ensure_dirs() @staticmethod - def save_cookies(cookies: CookieJar): - ModelScopeConfig.make_sure_credential_path_exist() - with open( - os.path.join(ModelScopeConfig.path_credential, - ModelScopeConfig.COOKIES_FILE_NAME), 'wb+') as f: - pickle.dump(cookies, f) + def save_cookies(cookies) -> None: + """Persist cookies to disk.""" + get_default_config().save_cookies(cookies) @staticmethod def get_cookies(): - cookies_path = os.path.join(ModelScopeConfig.path_credential, - ModelScopeConfig.COOKIES_FILE_NAME) - if os.path.exists(cookies_path): - with open(cookies_path, 'rb') as f: - cookies = pickle.load(f) - if not cookies: - return None - for cookie in cookies: - if cookie.name == 'm_session_id' and cookie.is_expired() and \ - not ModelScopeConfig.cookie_expired_warning: - ModelScopeConfig.cookie_expired_warning = True - logger.info('Not logged-in, you can login for uploading' - 'or accessing controlled entities.') - return None - return cookies - return None + """Load persisted cookies, returning ``None`` if absent or expired.""" + cookies = get_default_config().load_cookies() + if cookies is None and not ModelScopeConfig.cookie_expired_warning: + ModelScopeConfig.cookie_expired_warning = True + return cookies @staticmethod - def get_user_session_id(): - session_path = os.path.join(ModelScopeConfig.path_credential, - ModelScopeConfig.USER_SESSION_ID_FILE_NAME) - session_id = '' - if os.path.exists(session_path): - with open(session_path, 'rb') as f: - session_id = str(f.readline().strip(), encoding='utf-8') - return session_id - if session_id == '' or len(session_id) != 32: - session_id = str(uuid.uuid4().hex) - ModelScopeConfig.make_sure_credential_path_exist() - with open(session_path, 'w+') as wf: - wf.write(session_id) - - return session_id + def get_user_session_id() -> str: + """Return a stable session ID used in the user-agent header.""" + return get_default_config().get_session_id() @staticmethod - def save_token(token: str): - ModelScopeConfig.make_sure_credential_path_exist() - with open( - os.path.join(ModelScopeConfig.path_credential, - ModelScopeConfig.GIT_TOKEN_FILE_NAME), 'w+') as f: - f.write(token) + def save_git_token(git_token: str) -> None: + """Persist a git access token.""" + get_default_config().save_git_token(git_token) @staticmethod - def save_user_info(user_name: str, user_email: str): - ModelScopeConfig.make_sure_credential_path_exist() - with open( - os.path.join(ModelScopeConfig.path_credential, - ModelScopeConfig.USER_INFO_FILE_NAME), 'w+') as f: - f.write('%s:%s' % (user_name, user_email)) + def save_token(token: str) -> None: + """Deprecated: use :meth:`save_git_token` instead.""" + import warnings + warnings.warn( + 'ModelScopeConfig.save_token() is deprecated, ' + 'use ModelScopeConfig.save_git_token() instead.', + DeprecationWarning, stacklevel=2) + ModelScopeConfig.save_git_token(token) @staticmethod - def get_user_info() -> Tuple[str, str]: + def save_user_info(user_name: str, user_email: str) -> None: + """Persist ``user_name:user_email`` to the credentials directory.""" + get_default_config().save_user_info(user_name, user_email) + + @staticmethod + def get_user_info() -> Tuple[Optional[str], Optional[str]]: + """Return ``(username, email)`` previously saved, or ``(None, None)``.""" + path = get_default_config().credentials_dir / ModelScopeConfig.USER_INFO_FILE_NAME try: - with open( - os.path.join(ModelScopeConfig.path_credential, - ModelScopeConfig.USER_INFO_FILE_NAME), - 'r', - encoding='utf-8') as f: - info = f.read() - return info.split(':')[0], info.split(':')[1] - except FileNotFoundError: - pass + info = path.read_text(encoding='utf-8') + except (FileNotFoundError, OSError): + return None, None + parts = info.split(':', 1) + if len(parts) == 2: + return parts[0], parts[1] return None, None @staticmethod - def get_token() -> Optional[str]: - """ - Get token or None if not existent. - - Returns: - `str` or `None`: The token, `None` if it doesn't exist. - - """ - token = None - try: - with open( - os.path.join(ModelScopeConfig.path_credential, - ModelScopeConfig.GIT_TOKEN_FILE_NAME), - 'r', - encoding='utf-8') as f: - token = f.read() - except FileNotFoundError: - pass - return token + def get_git_token() -> Optional[str]: + """Return the persisted git access token, or ``None`` if not set.""" + return get_default_config().load_git_token() @staticmethod - def get_user_agent(user_agent: Union[Dict, str, None] = None, ) -> str: - """Formats a user-agent string with basic info about a request. + def get_token() -> Optional[str]: + """Deprecated: use :meth:`get_git_token` instead.""" + import warnings + warnings.warn( + 'ModelScopeConfig.get_token() is deprecated, ' + 'use ModelScopeConfig.get_git_token() instead.', + DeprecationWarning, stacklevel=2) + return ModelScopeConfig.get_git_token() - Args: - user_agent (`str`, `dict`, *optional*): - The user agent info in the form of a dictionary or a single string. - - Returns: - The formatted user-agent string. - """ - - # include some more telemetrics when executing in dedicated - # cloud containers - env = 'custom' - if MODELSCOPE_CLOUD_ENVIRONMENT in os.environ: - env = os.environ[MODELSCOPE_CLOUD_ENVIRONMENT] - user_name = 'unknown' - if MODELSCOPE_CLOUD_USERNAME in os.environ: - user_name = os.environ[MODELSCOPE_CLOUD_USERNAME] + @staticmethod + def get_user_agent(user_agent: Union[Dict, str, None] = None) -> str: + """Build a user-agent string carrying SDK version and telemetry.""" + env = os.environ.get(MODELSCOPE_CLOUD_ENVIRONMENT, 'custom') + user_name = os.environ.get(MODELSCOPE_CLOUD_USERNAME, 'unknown') from modelscope import __version__ - ua = 'modelscope/%s; python/%s; session_id/%s; platform/%s; processor/%s; env/%s; user/%s' % ( - __version__, - platform.python_version(), - ModelScopeConfig.get_user_session_id(), - platform.platform(), - platform.processor(), - env, - user_name, + ua = ( + f'modelscope/{__version__}; ' + f'python/{platform.python_version()}; ' + f'session_id/{ModelScopeConfig.get_user_session_id()}; ' + f'platform/{platform.platform()}; ' + f'processor/{platform.processor()}; ' + f'env/{env}; ' + f'user/{user_name}' ) if isinstance(user_agent, dict): ua += '; ' + '; '.join(f'{k}/{v}' for k, v in user_agent.items()) elif isinstance(user_agent, str): ua += '; ' + user_agent return ua - - -class UploadingCheck: - """ - Check the files and folders to be uploaded. - - Args: - max_file_count (int): The maximum number of files to be uploaded. Default to `UPLOAD_MAX_FILE_COUNT`. - max_file_count_in_dir (int): The maximum number of files in a directory. - Default to `UPLOAD_MAX_FILE_COUNT_IN_DIR`. - max_file_size (int): The maximum size of a single file in bytes. Default to `UPLOAD_MAX_FILE_SIZE`. - size_threshold_to_enforce_lfs (int): The size threshold to enforce LFS in bytes. - Files larger than this size will be enforced to be uploaded via LFS. - Default to `UPLOAD_SIZE_THRESHOLD_TO_ENFORCE_LFS`. - normal_file_size_total_limit (int): The total size limit of normal files in bytes. - Default to `UPLOAD_NORMAL_FILE_SIZE_TOTAL_LIMIT`. - - Examples: - >>> from modelscope.hub.api import UploadingCheck - >>> upload_checker = UploadingCheck() - >>> upload_checker.check_file('/path/to/your/file.txt') - >>> upload_checker.check_folder('/path/to/your/folder') - >>> is_lfs = upload_checker.is_lfs('/path/to/your/file.txt', repo_type='model') - >>> print(f'Is LFS: {is_lfs}') - """ - def __init__( - self, - max_file_count: int = UPLOAD_MAX_FILE_COUNT, - max_file_count_in_dir: int = UPLOAD_MAX_FILE_COUNT_IN_DIR, - max_file_size: int = UPLOAD_MAX_FILE_SIZE, - size_threshold_to_enforce_lfs: int = UPLOAD_SIZE_THRESHOLD_TO_ENFORCE_LFS, - normal_file_size_total_limit: int = UPLOAD_NORMAL_FILE_SIZE_TOTAL_LIMIT, - ): - self.max_file_count = max_file_count - self.max_file_count_in_dir = max_file_count_in_dir - self.max_file_size = max_file_size - self.size_threshold_to_enforce_lfs = size_threshold_to_enforce_lfs - self.normal_file_size_total_limit = normal_file_size_total_limit - - def check_file(self, file_path_or_obj) -> None: - """ - Check a single file to be uploaded. - - Args: - file_path_or_obj (Union[str, Path, bytes, BinaryIO]): The file path or file-like object to be checked. - - Raises: - ValueError: If the file does not exist or exceeds the size limit. - """ - if isinstance(file_path_or_obj, (str, Path)): - if not os.path.exists(file_path_or_obj): - raise ValueError(f'File {file_path_or_obj} does not exist') - - file_size: int = get_file_size(file_path_or_obj) - if file_size > self.max_file_size: - logger.warning(f'File exceeds size limit: {self.max_file_size / (1024 ** 3)} GB, ' - f'got {round(file_size / (1024 ** 3), 4)} GB') - - def check_folder(self, folder_path: Union[str, Path]): - """ - Check a folder to be uploaded. - - Args: - folder_path (Union[str, Path]): The folder path to be checked. - - Raises: - ValueError: If the folder does not exist or exceeds the file count limit. - """ - file_count = 0 - dir_count = 0 - - if isinstance(folder_path, str): - folder_path = Path(folder_path) - - for item in folder_path.iterdir(): - if item.is_file(): - file_count += 1 - item_size: int = get_file_size(item) - if item_size > self.max_file_size: - logger.warning(f'File {item} exceeds size limit: {self.max_file_size / (1024 ** 3)} GB', - f'got {round(item_size / (1024 ** 3), 4)} GB') - elif item.is_dir(): - dir_count += 1 - # Count items in subdirectories recursively - sub_file_count, sub_dir_count = self.check_folder(item) - if (sub_file_count + sub_dir_count) > self.max_file_count_in_dir: - raise ValueError(f'Directory {item} contains {sub_file_count + sub_dir_count} items ' - f'and exceeds limit: {self.max_file_count_in_dir}') - file_count += sub_file_count - dir_count += sub_dir_count - - if file_count > self.max_file_count: - raise ValueError(f'Total file count {file_count} and exceeds limit: {self.max_file_count}') - - return file_count, dir_count - - def is_lfs(self, file_path_or_obj: Union[str, Path, bytes, BinaryIO], repo_type: str) -> bool: - """ - Check if a file should be uploaded via LFS. - - Args: - file_path_or_obj (Union[str, Path, bytes, BinaryIO]): The file path or file-like object to be checked. - repo_type (str): The repo type, either `model` or `dataset`. - - Returns: - bool: True if the file should be uploaded via LFS, False otherwise. - """ - hit_lfs_suffix = True - - if isinstance(file_path_or_obj, (str, Path)): - file_path_or_obj = Path(file_path_or_obj) - if not file_path_or_obj.exists(): - raise ValueError(f'File {file_path_or_obj} does not exist') - - if repo_type == REPO_TYPE_MODEL: - if file_path_or_obj.suffix not in MODEL_LFS_SUFFIX: - hit_lfs_suffix = False - elif repo_type == REPO_TYPE_DATASET: - if file_path_or_obj.suffix not in DATASET_LFS_SUFFIX: - hit_lfs_suffix = False - else: - raise ValueError(f'Invalid repo type: {repo_type}, supported repos: {REPO_TYPE_SUPPORT}') - - file_size: int = get_file_size(file_path_or_obj) - - return file_size > self.size_threshold_to_enforce_lfs or hit_lfs_suffix - - def check_normal_files(self, file_path_list: List[Union[str, Path]], repo_type: str) -> None: - """ - Check a list of normal files to be uploaded. - - Args: - file_path_list (List[Union[str, Path]]): The list of file paths to be checked. - repo_type (str): The repo type, either `model` or `dataset`. - - Raises: - ValueError: If the total size of normal files exceeds the limit. - - Returns: None - """ - normal_file_list = [item for item in file_path_list if not self.is_lfs(item, repo_type)] - total_size = sum([get_file_size(item) for item in normal_file_list]) - - if total_size > self.normal_file_size_total_limit: - raise ValueError(f'Total size of non-lfs files {total_size / (1024 * 1024)}MB ' - f'and exceeds limit: {self.normal_file_size_total_limit / (1024 * 1024)}MB') diff --git a/modelscope/hub/cache_manager.py b/modelscope/hub/cache_manager.py index ca782b4a..9c2eb481 100644 --- a/modelscope/hub/cache_manager.py +++ b/modelscope/hub/cache_manager.py @@ -1,10 +1,17 @@ -"""Contains utilities to manage the ModelScope cache directory.""" +"""Contains utilities to manage the ModelScope cache directory. + +:func:`scan_cache_dir` retains the legacy file-grain scan that is unique +to this SDK; the modelscope_hub repo-grain :func:`scan_cache` and +:func:`clear_cache` are re-exported here for forward compatibility. +""" import os from dataclasses import dataclass from pathlib import Path from typing import Dict, FrozenSet, List, Literal, Optional, Set, Union +from modelscope_hub._cache_manager import clear_cache, scan_cache # noqa: F401 + from modelscope.hub.errors import CacheNotFound, CorruptedCacheException from modelscope.hub.utils.caching import ModelFileSystemCache from modelscope.hub.utils.utils import (convert_readable_size, @@ -15,6 +22,16 @@ from modelscope.utils.logger import get_logger logger = get_logger() +__all__ = [ + 'CachedFileInfo', + 'CachedRevisionInfo', + 'CachedRepoInfo', + 'ModelScopeCacheInfo', + 'scan_cache_dir', + 'scan_cache', + 'clear_cache', +] + # List of OS-created helper files that need to be ignored FILES_TO_IGNORE = ['.DS_Store', '._____temp'] diff --git a/modelscope/hub/callback.py b/modelscope/hub/callback.py index d203ec53..70bd56cb 100644 --- a/modelscope/hub/callback.py +++ b/modelscope/hub/callback.py @@ -1,34 +1,8 @@ -from tqdm.auto import tqdm +"""Progress callbacks — delegates to modelscope_hub. +Re-exports ProgressCallback and TqdmCallback from modelscope_hub, +maintaining backward compatibility for all existing import paths. +""" +from modelscope_hub import ProgressCallback, TqdmCallback -class ProgressCallback: - - def __init__(self, filename: str, file_size: int): - self.filename = filename - self.file_size = file_size - - def update(self, size: int): - pass - - def end(self): - pass - - -class TqdmCallback(ProgressCallback): - - def __init__(self, filename: str, file_size: int): - super().__init__(filename, file_size) - self.progress = tqdm( - unit='B', - unit_scale=True, - unit_divisor=1024, - total=file_size if file_size > 0 else 1, - initial=0, - desc='Downloading [' + self.filename + ']', - leave=True) - - def update(self, size: int): - self.progress.update(size) - - def end(self): - self.progress.close() +__all__ = ['ProgressCallback', 'TqdmCallback'] diff --git a/modelscope/hub/constants.py b/modelscope/hub/constants.py index c3ced124..28838ec3 100644 --- a/modelscope/hub/constants.py +++ b/modelscope/hub/constants.py @@ -1,7 +1,22 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""Hub constants — shared constants delegate to modelscope_hub, local-only constants retained. + +Constants that have equivalents in modelscope_hub are imported from there. +Constants unique to this project (endpoint configs, enum classes, env-driven tunables) +are retained locally. +""" import os from pathlib import Path +# --- Delegated constants (from modelscope_hub) --- +from modelscope_hub.compat.constants import ( # noqa: F401 + DEFAULT_DATASET_REVISION, DEFAULT_MAX_WORKERS, FILE_HASH, + MODELSCOPE_DOMAIN, MODELSCOPE_PREFER_AI_SITE, REPO_TYPE_DATASET, + REPO_TYPE_MODEL, REPO_TYPE_STUDIO, REPO_TYPE_SUPPORT, + TEMPORARY_FOLDER_NAME, ModelVisibility_INTERNAL, ModelVisibility_PRIVATE, + ModelVisibility_PUBLIC) + +# --- Local constants (not in modelscope_hub) --- MODELSCOPE_URL_SCHEME = 'https://' DEFAULT_MODELSCOPE_DOMAIN = 'www.modelscope.cn' DEFAULT_MODELSCOPE_INTL_DOMAIN = 'www.modelscope.ai' @@ -13,7 +28,6 @@ MODELSCOPE_DOWNLOAD_PARALLELS = int( os.environ.get('MODELSCOPE_DOWNLOAD_PARALLELS', 1)) DEFAULT_MODELSCOPE_GROUP = 'damo' MODEL_ID_SEPARATOR = '/' -FILE_HASH = 'Sha256' LOGGER_NAME = 'ModelScopeHub' DEFAULT_CREDENTIALS_PATH = Path.home().joinpath('.modelscope', 'credentials') MODELSCOPE_CREDENTIALS_PATH = os.environ.get( @@ -72,15 +86,9 @@ API_RESPONSE_FIELD_MESSAGE = 'Message' MODELSCOPE_CLOUD_ENVIRONMENT = 'MODELSCOPE_ENVIRONMENT' MODELSCOPE_CLOUD_USERNAME = 'MODELSCOPE_USERNAME' MODELSCOPE_SDK_DEBUG = 'MODELSCOPE_SDK_DEBUG' -MODELSCOPE_PREFER_AI_SITE = 'MODELSCOPE_PREFER_AI_SITE' -MODELSCOPE_DOMAIN = 'MODELSCOPE_DOMAIN' MODELSCOPE_ENABLE_DEFAULT_HASH_VALIDATION = 'MODELSCOPE_ENABLE_DEFAULT_HASH_VALIDATION' ONE_YEAR_SECONDS = 24 * 365 * 60 * 60 MODELSCOPE_REQUEST_ID = 'X-Request-ID' -TEMPORARY_FOLDER_NAME = '._____temp' -DEFAULT_MAX_WORKERS = int( - os.getenv('DEFAULT_MAX_WORKERS', min(8, - os.cpu_count() + 4))) DEFAULT_SKILLS_DIR = os.path.join(os.path.expanduser('~'), '.agents', 'skills') # Upload check env diff --git a/modelscope/hub/errors.py b/modelscope/hub/errors.py index a7c48a94..a1952eb3 100644 --- a/modelscope/hub/errors.py +++ b/modelscope/hub/errors.py @@ -1,67 +1,70 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""Error classes — core exceptions delegate to modelscope_hub, legacy aliases retained. +Exception classes from modelscope_hub provide the structured hierarchy. +Legacy aliases maintain isinstance compatibility for existing code. +Error handling functions with unique logic are retained. +""" import logging from http import HTTPStatus -from pathlib import Path -from typing import Optional, Union +from typing import Optional import requests -from requests.exceptions import HTTPError +from modelscope_hub.errors import AuthenticationError # noqa: F401 +from modelscope_hub.errors import (APIError, CacheNotFound, + CorruptedCacheException, FileIntegrityError, + HubError, InvalidParameter, NetworkError, + NotExistError, NotSupportedError, + PermissionDeniedError, RequestTimeoutError, + ServerError) +from requests.exceptions import HTTPError # noqa: F401 (re-exported) from modelscope.hub.constants import MODELSCOPE_REQUEST_ID from modelscope.utils.logger import get_logger logger = get_logger(log_level=logging.WARNING) +# --- Legacy exception aliases (maintain isinstance backward compatibility) --- -class NotSupportError(Exception): + +class RequestError(APIError): + """Legacy alias — use APIError for new code.""" + + def __init__(self, message: str = '', *args, **kwargs): + # Preserve legacy single-positional-arg constructor signature. + super().__init__(message, **kwargs) + + +class NotLoginException(AuthenticationError): + """Legacy alias — use AuthenticationError for new code.""" + + def __init__(self, message: str = '', *args, **kwargs): + super().__init__(message, **kwargs) + + +class FileDownloadError(NetworkError): + """Legacy alias — use NetworkError for new code.""" pass -class NoValidRevisionError(Exception): +class NotSupportError(NotSupportedError): + """Legacy alias — use NotSupportedError for new code.""" pass -class NotExistError(Exception): +class NoValidRevisionError(NotExistError): + """Legacy alias — raised when no valid revision is found.""" + + def __init__(self, message: str = '', *args, **kwargs): + super().__init__(message, **kwargs) + + +class GitError(HubError): + """Git operation failure.""" pass -class RequestError(Exception): - pass - - -class GitError(Exception): - pass - - -class InvalidParameter(Exception): - pass - - -class NotLoginException(Exception): - pass - - -class FileIntegrityError(Exception): - pass - - -class FileDownloadError(Exception): - pass - - -class CacheNotFound(Exception): - """Exception thrown when the ModelScope cache is not found.""" - - cache_dir: Union[str, Path] - - def __init__(self, msg: str, cache_dir: Union[str, Path], *args, **kwargs): - super().__init__(msg, *args, **kwargs) - self.cache_dir = cache_dir - - -class CorruptedCacheException(Exception): - """Exception for any unexpected structure in the ModelScope cache-system.""" +# --- Error handling functions (retained - contain unique logic) --- def get_request_id(response: requests.Response): diff --git a/modelscope/hub/file_download.py b/modelscope/hub/file_download.py index 4331ff24..d1273f19 100644 --- a/modelscope/hub/file_download.py +++ b/modelscope/hub/file_download.py @@ -1,400 +1,97 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. +"""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 shutil import tempfile +import time import urllib import uuid -from concurrent.futures import ThreadPoolExecutor from functools import partial from http.cookiejar import CookieJar -from pathlib import Path -from typing import Dict, List, Optional, Type, Union +from typing import Dict, List, Optional, Type import requests +# --- Hub file downloads (delegated) --- +from modelscope_hub.compat import dataset_file_download # noqa: E402,F401 +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.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, MODELSCOPE_DOWNLOAD_PARALLELS, - MODELSCOPE_PARALLEL_DOWNLOAD_THRESHOLD_MB, TEMPORARY_FOLDER_NAME) -from modelscope.utils.constant import (DEFAULT_DATASET_REVISION, - DEFAULT_MODEL_REVISION, - INTRA_CLOUD_ACCELERATION, - REPO_TYPE_DATASET, REPO_TYPE_MODEL, - REPO_TYPE_SUPPORT) -from modelscope.utils.file_utils import (get_dataset_cache_root, - get_model_cache_root) +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, InvalidParameter, NotExistError -from .utils.caching import ModelFileSystemCache -from .utils.utils import (file_integrity_validation, get_endpoint, - model_id_to_group_owner_name) +from .errors import FileDownloadError +from .utils.utils import get_endpoint logger = get_logger() -# Maximum number of retries for hash validation failures -HASH_RETRY_TIMES = 3 + +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: Optional[str] = DEFAULT_MODEL_REVISION, - cache_dir: Optional[str] = None, - user_agent: Union[Dict, str, None] = None, - local_files_only: Optional[bool] = False, - cookies: Optional[CookieJar] = None, - local_dir: Optional[str] = None, - token: Optional[str] = None, - endpoint: Optional[str] = None, -) -> Optional[str]: # pragma: no cover - """Download from a given URL and cache it if it's not already present in the local cache. - - Given a URL, this function looks for the corresponding file in the local - cache. If it's not there, download it. Then return the path to the cached - file. - - Args: - model_id (str): The model to whom the file to be downloaded belongs. - file_path(str): Path of the file to be downloaded, relative to the root of model repo. - revision(str, optional): revision of the model file to be downloaded. - Can be any of a branch, tag or commit hash. - cache_dir (str, Path, optional): Path to the folder where cached files are stored. - user_agent (dict, str, optional): The user-agent info in the form of a dictionary or a string. - local_files_only (bool, optional): If `True`, avoid downloading the file and return the path to the - local cached file if it exists. if `False`, download the file anyway even it exists. - cookies (CookieJar, optional): The cookie of download request. - local_dir (str, optional): Specific local directory path to which the file will be downloaded. - token (str, optional): The user token. - endpoint (str, optional): The remote endpoint. - - Returns: - string: string of local file or if networking is off, last version of - file cached on disk. - - Raises: - NotExistError: The file is not exist. - ValueError: The request parameter error. - - Note: - Raises the following errors: - - - [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) - if `use_auth_token=True` and the token cannot be found. - - [`OSError`](https://docs.python.org/3/library/exceptions.html#OSError) - if ETag cannot be determined. - - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) - if some parameter value is invalid - """ - return _repo_file_download( + 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 + return _compat_model_file_download( model_id, file_path, - repo_type=REPO_TYPE_MODEL, revision=revision, cache_dir=cache_dir, - user_agent=user_agent, - local_files_only=local_files_only, - cookies=cookies, local_dir=local_dir, - token=token, - endpoint=endpoint) - - -def dataset_file_download( - dataset_id: str, - file_path: str, - revision: Optional[str] = DEFAULT_DATASET_REVISION, - cache_dir: Union[str, Path, None] = None, - local_dir: Optional[str] = None, - user_agent: Optional[Union[Dict, str]] = None, - local_files_only: Optional[bool] = False, - cookies: Optional[CookieJar] = None, - token: Optional[str] = None, - endpoint: Optional[str] = None, -) -> str: - """Download raw files of a dataset. - Downloads all files at the specified revision. This - is useful when you want all files from a dataset, because you don't know which - ones you will need a priori. All files are nested inside a folder in order - to keep their actual filename relative to that folder. - - An alternative would be to just clone a dataset but this would require that the - user always has git and git-lfs installed, and properly configured. - - Args: - dataset_id (str): A user or an organization name and a dataset name separated by a `/`. - file_path (str): The relative path of the file to download. - revision (str, optional): An optional Git revision id which can be a branch name, a tag, or a - commit hash. NOTE: currently only branch and tag name is supported - cache_dir (str, Path, optional): Path to the folder where cached files are stored, dataset file will - be save as cache_dir/dataset_id/THE_DATASET_FILES. - local_dir (str, optional): Specific local directory path to which the file will be downloaded. - user_agent (str, dict, optional): The user-agent info in the form of a dictionary or a string. - local_files_only (bool, optional): If `True`, avoid downloading the file and return the path to the - local cached file if it exists. - cookies (CookieJar, optional): The cookie of the request, default None. - token (str, optional): The user token. - endpoint (str, optional): The remote endpoint. - Raises: - ValueError: the value details. - - Returns: - str: Local folder path (string) of repo snapshot - - Note: - Raises the following errors: - - [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) - if `use_auth_token=True` and the token cannot be found. - - [`OSError`](https://docs.python.org/3/library/exceptions.html#OSError) if - ETag cannot be determined. - - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) - if some parameter value is invalid - """ - return _repo_file_download( - dataset_id, - file_path, - repo_type=REPO_TYPE_DATASET, - revision=revision, - cache_dir=cache_dir, - user_agent=user_agent, - local_files_only=local_files_only, cookies=cookies, - local_dir=local_dir, token=token, - endpoint=endpoint) + endpoint=endpoint, + local_files_only=local_files_only, + user_agent=user_agent, + ) -def _repo_file_download( - repo_id: str, - file_path: str, - *, - repo_type: str = None, - revision: Optional[str] = DEFAULT_MODEL_REVISION, - cache_dir: Optional[str] = None, - user_agent: Union[Dict, str, None] = None, - local_files_only: Optional[bool] = False, - cookies: Optional[CookieJar] = None, - local_dir: Optional[str] = None, - disable_tqdm: bool = False, - token: Optional[str] = None, - endpoint: Optional[str] = None, -) -> Optional[str]: # pragma: no cover - - if not repo_type: - repo_type = REPO_TYPE_MODEL - if repo_type not in REPO_TYPE_SUPPORT: - raise InvalidParameter('Invalid repo type: %s, only support: %s' % - (repo_type, REPO_TYPE_SUPPORT)) - - temporary_cache_dir, cache = create_temporary_directory_and_cache( - repo_id, local_dir=local_dir, cache_dir=cache_dir, repo_type=repo_type) - - # if local_files_only is `True` and the file already exists in cached_path - # return the cached path - if local_files_only: - cached_file_path = cache.get_file_by_path(file_path) - if cached_file_path is not None: - logger.warning( - "File exists in local cache, but we're not sure it's up to date" - ) - return cached_file_path - else: - raise ValueError( - 'Cannot find the requested files in the cached path and outgoing' - ' traffic has been disabled. To enable look-ups and downloads' - " online, set 'local_files_only' to False.") - - _api = HubApi(token=token) - - headers = { - 'user-agent': ModelScopeConfig.get_user_agent(user_agent=user_agent, ), - 'snapshot-identifier': str(uuid.uuid4()), - } - - if INTRA_CLOUD_ACCELERATION == 'true': - region_id: str = ( - os.getenv('INTRA_CLOUD_ACCELERATION_REGION') - or _api._get_internal_acceleration_domain()) - if region_id: - logger.info( - f'Intra-cloud acceleration enabled for downloading from {repo_id}' - ) - headers['x-aliyun-region-id'] = region_id - - if cookies is None: - cookies = _api.get_cookies() - repo_files = [] - if endpoint is None: - endpoint = _api.get_endpoint_for_read( - repo_id=repo_id, repo_type=repo_type, token=token) - file_to_download_meta = None - if repo_type == REPO_TYPE_MODEL: - revision = _api.get_valid_revision( - repo_id, revision=revision, cookies=cookies, endpoint=endpoint) - # we need to confirm the version is up-to-date - # we need to get the file list to check if the latest version is cached, if so return, otherwise download - repo_files = _api.get_model_files( - model_id=repo_id, - revision=revision, - recursive=True, - use_cookies=False if cookies is None else cookies, - endpoint=endpoint) - for repo_file in repo_files: - if repo_file['Type'] == 'tree': - continue - - if repo_file['Path'] == file_path: - if cache.exists(repo_file): - file_name = repo_file['Name'] - logger.debug( - f'File {file_name} already in cache with identical hash, skip downloading!' - ) - return cache.get_file_by_info(repo_file) - else: - file_to_download_meta = repo_file - break - elif repo_type == REPO_TYPE_DATASET: - group_or_owner, name = model_id_to_group_owner_name(repo_id) - if not revision: - revision = DEFAULT_DATASET_REVISION - _hub_id, _ = _api.get_dataset_id_and_type( - dataset_name=name, - namespace=group_or_owner, - endpoint=endpoint, - token=token) - page_number = 1 - page_size = 100 - while True: - try: - dataset_files = _api.get_dataset_files( - repo_id=repo_id, - revision=revision, - root_path='/', - recursive=True, - page_number=page_number, - page_size=page_size, - endpoint=endpoint, - token=token, - dataset_hub_id=_hub_id) - except Exception as e: - logger.error( - f'Get dataset: {repo_id} file list failed, error: {e}') - break - - is_exist = False - for repo_file in dataset_files: - if repo_file['Type'] == 'tree': - continue - - if repo_file['Path'] == file_path: - if cache.exists(repo_file): - file_name = repo_file['Name'] - logger.debug( - f'File {file_name} already in cache with identical hash, skip downloading!' - ) - return cache.get_file_by_info(repo_file) - else: - file_to_download_meta = repo_file - is_exist = True - break - if len(dataset_files) < page_size or is_exist: - break - page_number += 1 - - if file_to_download_meta is None: - raise NotExistError('The file path: %s not exist in: %s' % - (file_path, repo_id)) - - # we need to download again - if repo_type == REPO_TYPE_MODEL: - url_to_download = get_file_download_url(repo_id, file_path, revision, - endpoint) - elif repo_type == REPO_TYPE_DATASET: - url_to_download = _api.get_dataset_file_url( - file_name=file_to_download_meta['Path'], - dataset_name=name, - namespace=group_or_owner, - revision=revision, - endpoint=endpoint) - else: - raise ValueError(f'Invalid repo type {repo_type}') - - return download_file(url_to_download, file_to_download_meta, - temporary_cache_dir, cache, headers, cookies) - - -def move_legacy_cache_to_standard_dir(cache_dir: str, model_id: str): - if cache_dir.endswith(os.path.sep): - cache_dir = cache_dir.strip(os.path.sep) - legacy_cache_root = os.path.dirname(cache_dir) - base_name = os.path.basename(cache_dir) - if base_name == 'datasets': - # datasets will not be not affected - return - if not legacy_cache_root.endswith('hub'): - # Two scenarios: - # We have restructured ModelScope cache directory, - # Scenery 1: - # When MODELSCOPE_CACHE is not set, the default directory remains - # the same at ~/.cache/modelscope/hub - # Scenery 2: - # When MODELSCOPE_CACHE is set, the cache directory is moved from - # $MODELSCOPE_CACHE/hub to $MODELSCOPE_CACHE/. In this case, - # we will be migrating the hub directory accordingly. - legacy_cache_root = os.path.join(legacy_cache_root, 'hub') - group_or_owner, name = model_id_to_group_owner_name(model_id) - name = name.replace('.', '___') - temporary_cache_dir = os.path.join(cache_dir, group_or_owner, name) - legacy_cache_dir = os.path.join(legacy_cache_root, group_or_owner, name) - if os.path.exists( - legacy_cache_dir) and not os.path.exists(temporary_cache_dir): - logger.info( - f'Legacy cache dir exists: {legacy_cache_dir}, move to {temporary_cache_dir}' - ) - try: - shutil.move(legacy_cache_dir, temporary_cache_dir) - except Exception: # noqa - # Failed, skip - pass - - -def create_temporary_directory_and_cache(model_id: str, - local_dir: str = None, - cache_dir: str = None, - repo_type: str = REPO_TYPE_MODEL): - if repo_type == REPO_TYPE_MODEL: - default_cache_root = get_model_cache_root() - elif repo_type == REPO_TYPE_DATASET: - default_cache_root = get_dataset_cache_root() - else: - raise ValueError( - f'repo_type only support model and dataset, but now is : {repo_type}' - ) - - group_or_owner, name = model_id_to_group_owner_name(model_id) - if local_dir is not None: - temporary_cache_dir = os.path.join(local_dir, TEMPORARY_FOLDER_NAME) - cache = ModelFileSystemCache(local_dir) - else: - if cache_dir is None: - cache_dir = default_cache_root - move_legacy_cache_to_standard_dir(cache_dir, model_id) - if isinstance(cache_dir, Path): - cache_dir = str(cache_dir) - temporary_cache_dir = os.path.join(cache_dir, TEMPORARY_FOLDER_NAME, - group_or_owner, name) - name = name.replace('.', '___') - cache = ModelFileSystemCache(cache_dir, group_or_owner, name) - - os.makedirs(temporary_cache_dir, exist_ok=True) - return temporary_cache_dir, cache +# --- Direct HTTP downloads (retained - non-Hub API) --- def get_file_download_url(model_id: str, @@ -427,103 +124,6 @@ def get_file_download_url(model_id: str, ) -def download_part_with_retry(params): - # unpack parameters - model_file_path, progress_callbacks, start, end, url, file_name, cookies, headers = params - get_headers = {} if headers is None else copy.deepcopy(headers) - get_headers['X-Request-ID'] = str(uuid.uuid4().hex) - retry = Retry( - total=API_FILE_DOWNLOAD_RETRY_TIMES, - backoff_factor=1, - allowed_methods=['GET']) - part_file_name = model_file_path + '_%s_%s' % (start, end) - while True: - try: - partial_length = 0 - if os.path.exists( - part_file_name): # download partial, continue download - with open(part_file_name, 'rb') as f: - partial_length = f.seek(0, io.SEEK_END) - for callback in progress_callbacks: - callback.update(partial_length) - download_start = start + partial_length - if download_start > end: - break # this part is download completed. - get_headers['Range'] = 'bytes=%s-%s' % (download_start, end) - with open(part_file_name, '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 - f.write(chunk) - for callback in progress_callbacks: - callback.update(len(chunk)) - break - except (Exception) as e: # no matter what exception, we will retry. - retry = retry.increment('GET', url, error=e) - logger.warning('Downloading: %s failed, reason: %s will retry' % - (model_file_path, e)) - retry.sleep() - - -def parallel_download(url: str, - local_dir: str, - file_name: str, - cookies: CookieJar, - headers: Optional[Dict[str, str]] = None, - file_size: int = None, - disable_tqdm: bool = False, - progress_callbacks: List[Type[ProgressCallback]] = None, - endpoint: str = None): - 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 - ] - # create temp file - PART_SIZE = 160 * 1024 * 1024 # every part is 160M - tasks = [] - file_path = os.path.join(local_dir, file_name) - os.makedirs(os.path.dirname(file_path), exist_ok=True) - for idx in range(int(file_size / PART_SIZE)): - start = idx * PART_SIZE - end = (idx + 1) * PART_SIZE - 1 - tasks.append((file_path, progress_callbacks, start, end, url, - file_name, cookies, headers)) - if end + 1 < file_size: - tasks.append((file_path, progress_callbacks, end + 1, file_size - 1, - url, file_name, cookies, headers)) - parallels = min(MODELSCOPE_DOWNLOAD_PARALLELS, 16) - # download every part - with ThreadPoolExecutor( - max_workers=parallels, thread_name_prefix='download') as executor: - list(executor.map(download_part_with_retry, tasks)) - for callback in progress_callbacks: - callback.end() - # merge parts. - hash_sha256 = hashlib.sha256() - with open(os.path.join(local_dir, file_name), 'wb') as output_file: - for task in tasks: - part_file_name = task[0] + '_%s_%s' % (task[2], task[3]) - with open(part_file_name, 'rb') as part_file: - while True: - chunk = part_file.read(16 * API_FILE_DOWNLOAD_CHUNK_SIZE) - if not chunk: - break - output_file.write(chunk) - hash_sha256.update(chunk) - os.remove(part_file_name) - return hash_sha256.hexdigest() - - def http_get_model_file( url: str, local_dir: str, @@ -712,76 +312,10 @@ def http_get_file( os.replace(temp_file.name, os.path.join(local_dir, file_name)) -def download_file( - url, - file_meta, - temporary_cache_dir, - cache, - headers, - cookies, - disable_tqdm=False, - progress_callbacks: List[Type[ProgressCallback]] = None, -): - temp_file = os.path.join(temporary_cache_dir, file_meta['Path']) - - for hash_attempt in range(HASH_RETRY_TIMES): - if MODELSCOPE_PARALLEL_DOWNLOAD_THRESHOLD_MB * 1000 * 1000 < file_meta[ - 'Size'] and MODELSCOPE_DOWNLOAD_PARALLELS > 1: # parallel download large file. - file_digest = parallel_download( - url, - temporary_cache_dir, - file_meta['Path'], - headers=headers, - cookies=None if cookies is None else cookies.get_dict(), - file_size=file_meta['Size'], - disable_tqdm=disable_tqdm, - progress_callbacks=progress_callbacks, - ) - else: - file_digest = http_get_model_file( - url, - temporary_cache_dir, - file_meta['Path'], - file_size=file_meta['Size'], - headers=headers, - cookies=cookies, - disable_tqdm=disable_tqdm, - progress_callbacks=progress_callbacks, - ) - - # Check file integrity - if FILE_HASH in file_meta: - expected_hash = file_meta[FILE_HASH] - hash_valid = True - if file_digest is not None: - if file_digest != expected_hash: - logger.warning( - 'Mismatched real-time digest for %s, falling back to full hash check', - file_meta['Path']) - if not file_integrity_validation(temp_file, expected_hash): - hash_valid = False - else: - if not file_integrity_validation(temp_file, expected_hash): - hash_valid = False - - if not hash_valid: - if hash_attempt < HASH_RETRY_TIMES - 1: - logger.warning( - 'Hash validation failed for %s, ' - 'retrying download (attempt %d/%d)', file_meta['Path'], - hash_attempt + 1, HASH_RETRY_TIMES) - # Clean up corrupted file before retry - if os.path.exists(temp_file): - os.remove(temp_file) - continue - else: - raise FileDownloadError( - 'File %s hash validation failed after %d attempts, ' - 'the file may be corrupted.' % - (file_meta['Path'], HASH_RETRY_TIMES)) - - # Hash validation passed or no hash to validate, exit retry loop - break - - # Put file into cache - return cache.put_file(file_meta, temp_file) +__all__ = [ + 'model_file_download', + 'dataset_file_download', + 'http_get_file', + 'http_get_model_file', + 'get_file_download_url', +] diff --git a/modelscope/hub/git.py b/modelscope/hub/git.py index 64d37620..5a3d90c2 100644 --- a/modelscope/hub/git.py +++ b/modelscope/hub/git.py @@ -1,323 +1,200 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""Git wrapper — shim delegating to ``modelscope_hub._git``. +Preserves the legacy ``GitCommandWrapper`` Singleton interface used +throughout the SDK while routing primitive Git operations through +:class:`modelscope_hub._git.GitCommand`. +""" +from __future__ import annotations import os -import subprocess +from pathlib import Path from typing import List, Optional -from urllib.parse import urlparse, urlunparse +from modelscope_hub._git import GitCommand as _GitCommand + +from modelscope.hub.errors import GitError +from modelscope.utils.constant import MASTER_MODEL_BRANCH from modelscope.utils.logger import get_logger -from ..utils.constant import MASTER_MODEL_BRANCH -from .errors import GitError logger = get_logger() +__all__ = ['GitError', 'GitCommandWrapper', 'Singleton'] + class Singleton(type): - _instances = {} + """Metaclass enforcing one instance per class — preserved for parity.""" + + _instances: dict = {} def __call__(cls, *args, **kwargs): if cls not in cls._instances: - cls._instances[cls] = super(Singleton, - cls).__call__(*args, **kwargs) + cls._instances[cls] = super().__call__(*args, **kwargs) return cls._instances[cls] class GitCommandWrapper(metaclass=Singleton): - """Some git operation wrapper + """Backward-compatible Git wrapper. + + Wraps :class:`modelscope_hub._git.GitCommand` to expose the legacy + instance-method API (``clone``, ``push``, ``pull``, ``tag``…) used + by callers that pre-date the SDK refactor. """ - default_git_path = 'git' # The default git command line - def __init__(self, path: str = None): + default_git_path = 'git' + + def __init__(self, path: Optional[str] = None): self.git_path = path or self.default_git_path + _GitCommand.set_git_path(self.git_path) - def _run_git_command(self, *args) -> subprocess.CompletedProcess: - """Run git command, if command return 0, return subprocess.response - otherwise raise GitError, message is stdout and stderr. - - Args: - args: List of command args. - - Raises: - GitError: Exception with stdout and stderr. - - Returns: - subprocess.CompletedProcess: the command response - """ - logger.debug(' '.join(args)) - git_env = os.environ.copy() - git_env['GIT_TERMINAL_PROMPT'] = '0' - command = [self.git_path, *args] - command = [item for item in command if item] - response = subprocess.run( - command, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - env=git_env, - ) # compatible for python3.6 + # ------------------------------------------------------------------ + # Low-level subprocess passthrough (legacy contract) + # ------------------------------------------------------------------ + def _run_git_command(self, *args): + """Run a git subcommand, raising :class:`GitError` on failure.""" try: - response.check_returncode() - return response - except subprocess.CalledProcessError as error: - std_out = response.stdout.decode('utf-8', errors='replace') - std_err = error.stderr.decode('utf-8', errors='replace') - if 'nothing to commit' in std_out: - logger.info( - 'Nothing to commit, your local repo is upto date with remote' - ) - return response - else: - logger.error( - 'Running git command: %s failed \n stdout: %s \n stderr: %s' - % (command, std_out, std_err)) - raise GitError(std_err) + return _GitCommand._run(*[a for a in args if a]) + except Exception as exc: # _git.GitError → legacy GitError + raise GitError(str(exc)) from exc - def config_auth_token(self, repo_dir, auth_token): - url = self.get_repo_remote_url(repo_dir) - if '//oauth2' not in url: - auth_url = self._add_token(auth_token, url) - cmd_args = '-C %s remote set-url origin %s' % (repo_dir, auth_url) - cmd_args = cmd_args.split(' ') - rsp = self._run_git_command(*cmd_args) - logger.debug(rsp.stdout.decode('utf8')) + # ------------------------------------------------------------------ + # URL / token helpers + # ------------------------------------------------------------------ + def _add_git_token(self, git_token: str, url: str) -> str: + return _GitCommand._inject_token(url, git_token) - def _add_token(self, token: str, url: str): - """Inject OAuth2 token into an HTTP(S) git URL. + def remove_token_from_url(self, url: str) -> str: + return _GitCommand.strip_token_from_url(url) - Uses ``urllib.parse`` for reliable URL component handling, - avoiding naive string replacement that can corrupt URLs - containing multiple ``://`` sequences. - - Args: - token: OAuth2 access token. - url: Remote URL (HTTP, HTTPS, or SSH). - - Returns: - URL with ``oauth2:@`` injected into the *netloc*, - or the original *url* unchanged when: - - * *token* is falsy, - * the URL already carries an ``oauth2`` credential, - * the scheme is not HTTP/HTTPS (e.g. ``ssh://``, ``git@``). - """ - if not token: - return url - - # SSH URLs authenticate via keys, not tokens. - if url.startswith('git@'): - return url + # ------------------------------------------------------------------ + # LFS + # ------------------------------------------------------------------ + def is_lfs_installed(self) -> bool: + return _GitCommand.is_lfs_available() + def git_lfs_install(self, repo_dir: str) -> bool: try: - parsed = urlparse(url) + _GitCommand.lfs_install(Path(repo_dir)) + return True except Exception: - return url - - # Only inject into HTTP(S) URLs. - if parsed.scheme not in ('http', 'https'): - return url - - # Prevent double injection. - if parsed.username == 'oauth2': - return url - - # Reconstruct netloc: oauth2:@host[:port] - host = parsed.hostname or '' - if parsed.port: - host = f'{host}:{parsed.port}' - netloc = f'oauth2:{token}@{host}' - - return urlunparse(parsed._replace(netloc=netloc)) - - def remove_token_from_url(self, url: str): - if url and '//oauth2' in url: - start_index = url.find('oauth2') - end_index = url.find('@') - url = url[:start_index] + url[end_index + 1:] - return url - - def is_lfs_installed(self): - cmd = ['lfs', 'env'] - try: - self._run_git_command(*cmd) - return True - except GitError: return False - def git_lfs_install(self, repo_dir): - cmd = ['-C', repo_dir, 'lfs', 'install', '--force'] - try: - self._run_git_command(*cmd) - return True - except GitError: - return False + def list_lfs_files(self, repo_dir: str) -> List[str]: + rsp = self._run_git_command('-C', repo_dir, 'lfs', 'ls-files') + return [ + line.split(' ')[-1] for line in rsp.stdout.strip().splitlines() + if line + ] + # ------------------------------------------------------------------ + # Auth / user config + # ------------------------------------------------------------------ + def config_git_token(self, repo_dir: str, git_token: str) -> None: + url = self.get_repo_remote_url(repo_dir) + if '//oauth2' in url: + return + auth_url = self._add_git_token(git_token, url) + self._run_git_command('-C', repo_dir, 'remote', 'set-url', 'origin', + auth_url) + + def add_user_info(self, repo_base_dir: str, repo_name: str) -> None: + from modelscope.hub.api import ModelScopeConfig + user_name, user_email = ModelScopeConfig.get_user_info() + if not (user_name and user_email): + return + repo_dir = os.path.join(repo_base_dir, repo_name) + self._run_git_command('-C', repo_dir, 'config', 'user.name', user_name) + self._run_git_command('-C', repo_dir, 'config', 'user.email', + user_email) + + # ------------------------------------------------------------------ + # Clone / pull / push + # ------------------------------------------------------------------ def clone(self, repo_base_dir: str, - token: str, + git_token: Optional[str], url: str, repo_name: str, branch: Optional[str] = None): - """ git clone command wrapper. - For public project, token can None, private repo, there must token. - - Args: - repo_base_dir (str): The local base dir, the repository will be clone to local_dir/repo_name - token (str): The git token, must be provided for private project. - url (str): The remote url - repo_name (str): The local repository path name. - branch (str, optional): _description_. Defaults to None. - - Returns: - The popen response. - """ - url = self._add_token(token, url) - if branch: - clone_args = '-C %s clone %s %s --branch %s' % (repo_base_dir, url, - repo_name, branch) - else: - clone_args = '-C %s clone %s' % (repo_base_dir, url) - logger.debug(clone_args) - clone_args = clone_args.split(' ') + target = Path(repo_base_dir) / repo_name try: - response = self._run_git_command(*clone_args) - logger.debug(response.stdout.decode('utf8')) - return response - except GitError: - # git clone may succeed but still exit non-zero when an - # external hook (e.g. a custom core.hooksPath that wraps - # ``git lfs post-merge``) returns a non-zero code. When the - # repository was actually cloned, treat this as a warning. - repo_dir = os.path.join(repo_base_dir, repo_name) - if os.path.isdir(os.path.join(repo_dir, '.git')): + _GitCommand.clone( + url=url, target_dir=target, branch=branch, token=git_token) + except Exception as exc: + if (target / '.git').is_dir(): logger.warning( - 'git clone exited with non-zero status but the ' - 'repository was cloned successfully at %s. ' - 'This is usually caused by a post-clone hook ' - '(e.g. core.hooksPath). Continuing.', repo_dir) + 'git clone exited non-zero but repository was cloned ' + 'at %s. Likely a post-clone hook. Continuing.', target) return None - raise - - def add_user_info(self, repo_base_dir, repo_name): - from modelscope.hub.api import ModelScopeConfig - user_name, user_email = ModelScopeConfig.get_user_info() - if user_name and user_email: - # config user.name and user.email if exist - config_user_name_args = '-C %s/%s config user.name %s' % ( - repo_base_dir, repo_name, user_name) - response = self._run_git_command(*config_user_name_args.split(' ')) - logger.debug(response.stdout.decode('utf8')) - config_user_email_args = '-C %s/%s config user.email %s' % ( - repo_base_dir, repo_name, user_email) - response = self._run_git_command( - *config_user_email_args.split(' ')) - logger.debug(response.stdout.decode('utf8')) - - def add(self, - repo_dir: str, - files: List[str] = list(), - all_files: bool = False): - if all_files: - add_args = '-C %s add -A' % repo_dir - elif len(files) > 0: - files_str = ' '.join(files) - add_args = '-C %s add %s' % (repo_dir, files_str) - add_args = add_args.split(' ') - rsp = self._run_git_command(*add_args) - logger.debug(rsp.stdout.decode('utf8')) - return rsp - - def commit(self, repo_dir: str, message: str): - """Run git commit command - - Args: - repo_dir (str): the repository directory. - message (str): commit message. - - Returns: - The command popen response. - """ - commit_args = ['-C', '%s' % repo_dir, 'commit', '-m', "'%s'" % message] - rsp = self._run_git_command(*commit_args) - logger.info(rsp.stdout.decode('utf8')) - return rsp - - def checkout(self, repo_dir: str, revision: str): - cmds = ['-C', '%s' % repo_dir, 'checkout', '%s' % revision] - return self._run_git_command(*cmds) - - def new_branch(self, repo_dir: str, revision: str): - cmds = ['-C', '%s' % repo_dir, 'checkout', '-b', revision] - return self._run_git_command(*cmds) - - def get_remote_branches(self, repo_dir: str): - cmds = ['-C', '%s' % repo_dir, 'branch', '-r'] - rsp = self._run_git_command(*cmds) - info = [ - line.strip() - for line in rsp.stdout.decode('utf8').strip().split(os.linesep) - ] - if len(info) == 1: - return ['/'.join(info[0].split('/')[1:])] - else: - return ['/'.join(line.split('/')[1:]) for line in info[1:]] + raise GitError(str(exc)) from exc def pull(self, repo_dir: str, remote: str = 'origin', branch: str = 'master'): - cmds = ['-C', repo_dir, 'pull', remote, branch] - return self._run_git_command(*cmds) + return self._run_git_command('-C', repo_dir, 'pull', remote, branch) def push(self, repo_dir: str, - token: str, + git_token: str, url: str, local_branch: str, remote_branch: str, force: bool = False): - url = self._add_token(token, url) - - push_args = '-C %s push %s %s:%s' % (repo_dir, url, local_branch, - remote_branch) + auth_url = self._add_git_token(git_token, url) + args = [ + '-C', repo_dir, 'push', auth_url, f'{local_branch}:{remote_branch}' + ] if force: - push_args += ' -f' - push_args = push_args.split(' ') - rsp = self._run_git_command(*push_args) - logger.debug(rsp.stdout.decode('utf8')) - return rsp + args.append('-f') + return self._run_git_command(*args) - def get_repo_remote_url(self, repo_dir: str): - cmd_args = '-C %s config --get remote.origin.url' % repo_dir - cmd_args = cmd_args.split(' ') - rsp = self._run_git_command(*cmd_args) - url = rsp.stdout.decode('utf8') - return url.strip() + # ------------------------------------------------------------------ + # Add / commit / branch / checkout + # ------------------------------------------------------------------ + def add(self, + repo_dir: str, + files: Optional[List[str]] = None, + all_files: bool = False): + if all_files: + return self._run_git_command('-C', repo_dir, 'add', '-A') + return self._run_git_command('-C', repo_dir, 'add', *(files or [])) - def list_lfs_files(self, repo_dir: str): - cmd_args = '-C %s lfs ls-files' % repo_dir - cmd_args = cmd_args.split(' ') - rsp = self._run_git_command(*cmd_args) - out = rsp.stdout.decode('utf8').strip() - files = [] - for line in out.split(os.linesep): - files.append(line.split(' ')[-1]) + def commit(self, repo_dir: str, message: str): + return self._run_git_command('-C', repo_dir, 'commit', '-m', + f"'{message}'") - return files + def checkout(self, repo_dir: str, revision: str): + return self._run_git_command('-C', repo_dir, 'checkout', revision) + def new_branch(self, repo_dir: str, revision: str): + return self._run_git_command('-C', repo_dir, 'checkout', '-b', + revision) + + def get_remote_branches(self, repo_dir: str) -> List[str]: + rsp = self._run_git_command('-C', repo_dir, 'branch', '-r') + info = [ + line.strip() for line in rsp.stdout.strip().splitlines() if line + ] + if len(info) <= 1: + return ['/'.join(info[0].split('/')[1:])] if info else [] + return ['/'.join(line.split('/')[1:]) for line in info[1:]] + + def get_repo_remote_url(self, repo_dir: str) -> str: + rsp = self._run_git_command('-C', repo_dir, 'config', '--get', + 'remote.origin.url') + return rsp.stdout.strip() + + # ------------------------------------------------------------------ + # Tags + # ------------------------------------------------------------------ def tag(self, repo_dir: str, tag_name: str, message: str, ref: str = MASTER_MODEL_BRANCH): - cmd_args = [ - '-C', repo_dir, 'tag', tag_name, '-m', - '"%s"' % message, ref - ] - rsp = self._run_git_command(*cmd_args) - logger.debug(rsp.stdout.decode('utf8')) - return rsp + return self._run_git_command('-C', repo_dir, 'tag', tag_name, '-m', + f'"{message}"', ref) - def push_tag(self, repo_dir: str, tag_name): - cmd_args = ['-C', repo_dir, 'push', 'origin', tag_name] - rsp = self._run_git_command(*cmd_args) - logger.debug(rsp.stdout.decode('utf8')) - return rsp + def push_tag(self, repo_dir: str, tag_name: str): + return self._run_git_command('-C', repo_dir, 'push', 'origin', + tag_name) diff --git a/modelscope/hub/repository.py b/modelscope/hub/repository.py index 1dd6665b..f212dcec 100644 --- a/modelscope/hub/repository.py +++ b/modelscope/hub/repository.py @@ -1,5 +1,12 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""Repository — shim delegating to ``modelscope_hub`` for Git operations. +Preserves the legacy :class:`Repository` and :class:`DatasetRepository` +constructors (auto-clone-on-init) and methods (``push``, ``pull``, +``tag``, ``tag_and_push``, ``add_lfs_type``) used by callers that +pre-date the SDK refactor. +""" +from __future__ import annotations import os import warnings from typing import Optional @@ -14,172 +21,147 @@ from .utils.utils import get_endpoint logger = get_logger() +__all__ = ['Repository', 'DatasetRepository'] + + +def _resolve_git_token(auth_token: Optional[str]) -> Optional[str]: + if auth_token: + return auth_token + from modelscope.hub.api import ModelScopeConfig + return ModelScopeConfig.get_git_token() + + +def _clone_if_needed(git_wrapper: GitCommandWrapper, base_dir: str, + repo_name: str, repo_dir: str, url: str, + git_token: Optional[str], + revision: Optional[str]) -> bool: + """Clone *url* into *repo_dir* unless it's already that working copy. + + Returns ``True`` if a clone was performed, ``False`` if skipped. + """ + os.makedirs(repo_dir, exist_ok=True) + if os.listdir(repo_dir): + try: + existing = git_wrapper.get_repo_remote_url(repo_dir) + existing = git_wrapper.remove_token_from_url(existing) + if existing == url: + return False + except GitError: + pass + git_wrapper.clone(base_dir, git_token, url, repo_name, revision) + return True + class Repository: - """A local representation of the model git repository. - """ + """A local representation of a model Git repository on ModelScope Hub.""" def __init__(self, model_dir: str, clone_from: str, revision: Optional[str] = DEFAULT_REPOSITORY_REVISION, - auth_token: Optional[str] = None, + git_token: Optional[str] = None, git_path: Optional[str] = None, - endpoint: Optional[str] = None): - """Instantiate a Repository object by cloning the remote ModelScopeHub repo + endpoint: Optional[str] = None, + auth_token: Optional[str] = None): + if auth_token is not None and git_token is None: + import warnings + warnings.warn( + 'Repository(auth_token=...) is deprecated, ' + 'use Repository(git_token=...) instead.', + DeprecationWarning, + stacklevel=2) + git_token = auth_token - Args: - model_dir (str): The model root directory. - clone_from (str): model id in ModelScope-hub from which git clone - revision (str, optional): revision of the model you want to clone from. - Can be any of a branch, tag or commit hash - auth_token (str, optional): token obtained when calling `HubApi.login()`. - Usually you can safely ignore the parameter as the token is already - saved when you login the first time, if None, we will use saved token. - git_path (str, optional): The git command line path, if None, we use 'git' - endpoint (str, optional): The ModelScope endpoint URL. If None, use default endpoint. + if not revision: + raise InvalidParameter( + 'a non-default value of revision cannot be empty.') - Raises: - InvalidParameter: revision is None. - """ self._endpoint = endpoint self.model_dir = model_dir self.model_base_dir = os.path.dirname(model_dir) self.model_repo_name = os.path.basename(model_dir) - - if not revision: - err_msg = 'a non-default value of revision cannot be empty.' - raise InvalidParameter(err_msg) - - from modelscope.hub.api import ModelScopeConfig - if auth_token: - self.auth_token = auth_token - else: - self.auth_token = ModelScopeConfig.get_token() - - git_wrapper = GitCommandWrapper() - if not git_wrapper.is_lfs_installed(): - logger.error('git lfs is not installed, please install.') + self.git_token = _resolve_git_token(git_token) self.git_wrapper = GitCommandWrapper(git_path) - os.makedirs(self.model_dir, exist_ok=True) + if not self.git_wrapper.is_lfs_installed(): + logger.error('git lfs is not installed, please install.') + url = self._get_model_id_url(clone_from) - if os.listdir(self.model_dir): # directory not empty. - remote_url = self._get_remote_url() - remote_url = self.git_wrapper.remove_token_from_url(remote_url) - if remote_url and remote_url == url: # need not clone again - return - self.git_wrapper.clone(self.model_base_dir, self.auth_token, url, - self.model_repo_name, revision) + cloned = _clone_if_needed(self.git_wrapper, self.model_base_dir, + self.model_repo_name, self.model_dir, url, + self.git_token, revision) + if not cloned: + return - if git_wrapper.is_lfs_installed(): - git_wrapper.git_lfs_install(self.model_dir) + if self.git_wrapper.is_lfs_installed(): + self.git_wrapper.git_lfs_install(self.model_dir) - # add user info if login self.git_wrapper.add_user_info(self.model_base_dir, self.model_repo_name) - if self.auth_token: # config remote with auth token - self.git_wrapper.config_auth_token(self.model_dir, self.auth_token) + if self.git_token: + self.git_wrapper.config_git_token(self.model_dir, self.git_token) - def _get_model_id_url(self, model_id): - endpoint = self._endpoint if self._endpoint else get_endpoint() - url = f'{endpoint}/{model_id}.git' - return url - - def _get_remote_url(self): - try: - remote = self.git_wrapper.get_repo_remote_url(self.model_dir) - except GitError: - remote = None - return remote + def _get_model_id_url(self, model_id: str) -> str: + endpoint = self._endpoint or get_endpoint() + return f'{endpoint}/{model_id}.git' def pull(self, remote: str = 'origin', branch: str = 'master'): - """Pull remote branch - - Args: - remote (str, optional): The remote name. Defaults to 'origin'. - branch (str, optional): The remote branch. Defaults to 'master'. - """ + """Pull *remote*/*branch* into the local checkout.""" self.git_wrapper.pull(self.model_dir, remote=remote, branch=branch) - def add_lfs_type(self, file_name_suffix: str): - """Add file suffix to lfs list. - - Args: - file_name_suffix (str): The file name suffix. - examples '*.safetensors' - """ - os.system( - "printf '\n%s filter=lfs diff=lfs merge=lfs -text\n'>>%s" % - (file_name_suffix, os.path.join(self.model_dir, '.gitattributes'))) + def add_lfs_type(self, file_name_suffix: str) -> None: + """Track an additional file-name pattern with Git LFS.""" + attrs = os.path.join(self.model_dir, '.gitattributes') + with open(attrs, 'a', encoding='utf-8') as f: + f.write( + f'\n{file_name_suffix} filter=lfs diff=lfs merge=lfs -text\n') def push(self, commit_message: str, local_branch: Optional[str] = DEFAULT_REPOSITORY_REVISION, remote_branch: Optional[str] = DEFAULT_REPOSITORY_REVISION, - force: Optional[bool] = False): + force: bool = False): + """Stage all changes, commit, and push to the remote.""" warnings.warn( - 'This function is deprecated and will be removed in future versions. ' - 'Please use git command directly or use HubApi().upload_folder instead', + 'This function is deprecated and will be removed in future ' + 'versions. Please use git command directly or use ' + 'HubApi().upload_folder instead', DeprecationWarning, stacklevel=2) - """Push local files to remote, this method will do. - Execute git pull, git add, git commit, git push in order. - - Args: - commit_message (str): commit message - local_branch(str, optional): The local branch, default master. - remote_branch (str, optional): The remote branch to push, default master. - force (bool, optional): whether to use forced-push. - - Raises: - InvalidParameter: no commit message. - NotLoginException: no auth token. - """ - if commit_message is None or not isinstance(commit_message, str): - msg = 'commit_message must be provided!' - raise InvalidParameter(msg) + if not isinstance(commit_message, str) or not commit_message: + raise InvalidParameter('commit_message must be provided!') if not isinstance(force, bool): raise InvalidParameter('force must be bool') - - if not self.auth_token: + if not self.git_token: raise NotLoginException('Must login to push, please login first.') - self.git_wrapper.config_auth_token(self.model_dir, self.auth_token) + self.git_wrapper.config_git_token(self.model_dir, self.git_token) self.git_wrapper.add_user_info(self.model_base_dir, self.model_repo_name) - url = self.git_wrapper.get_repo_remote_url(self.model_dir) self.git_wrapper.add(self.model_dir, all_files=True) self.git_wrapper.commit(self.model_dir, commit_message) self.git_wrapper.push( repo_dir=self.model_dir, - token=self.auth_token, + git_token=self.git_token, url=url, local_branch=local_branch, - remote_branch=remote_branch) + remote_branch=remote_branch, + force=force) def tag(self, tag_name: str, message: str, ref: Optional[str] = MASTER_MODEL_BRANCH): - """Create a new tag. - - Args: - tag_name (str): The name of the tag - message (str): The tag message. - ref (str, optional): The tag reference, can be commit id or branch. - - Raises: - InvalidParameter: no commit message. - """ - if tag_name is None or tag_name == '': - msg = 'We use tag-based revision, therefore tag_name cannot be None or empty.' - raise InvalidParameter(msg) - if message is None or message == '': - msg = 'We use annotated tag, therefore message cannot None or empty.' - raise InvalidParameter(msg) + """Create an annotated tag pointing to *ref*.""" + if not tag_name: + raise InvalidParameter( + 'We use tag-based revision, therefore tag_name ' + 'cannot be None or empty.') + if not message: + raise InvalidParameter('We use annotated tag, therefore message ' + 'cannot None or empty.') self.git_wrapper.tag( repo_dir=self.model_dir, tag_name=tag_name, @@ -190,143 +172,97 @@ class Repository: tag_name: str, message: str, ref: Optional[str] = MASTER_MODEL_BRANCH): - """Create tag and push to remote - - Args: - tag_name (str): The name of the tag - message (str): The tag message. - ref (str, optional): The tag ref, can be commit id or branch. Defaults to MASTER_MODEL_BRANCH. - """ + """Create *tag_name* and push it to the remote.""" self.tag(tag_name, message, ref) - self.git_wrapper.push_tag(repo_dir=self.model_dir, tag_name=tag_name) class DatasetRepository: - """A local representation of the dataset (metadata) git repository. - """ + """A local representation of a dataset (metadata) Git repository.""" def __init__(self, repo_work_dir: str, dataset_id: str, revision: Optional[str] = DEFAULT_DATASET_REVISION, - auth_token: Optional[str] = None, + git_token: Optional[str] = None, git_path: Optional[str] = None, - endpoint: Optional[str] = None): - """ - Instantiate a Dataset Repository object by cloning the remote ModelScope dataset repo + endpoint: Optional[str] = None, + auth_token: Optional[str] = None): + if auth_token is not None and git_token is None: + import warnings + warnings.warn( + 'DatasetRepository(auth_token=...) is deprecated, ' + 'use DatasetRepository(git_token=...) instead.', + DeprecationWarning, + stacklevel=2) + git_token = auth_token - Args: - repo_work_dir (str): The dataset repo root directory. - dataset_id (str): dataset id in ModelScope from which git clone - revision (str, optional): revision of the dataset you want to clone from. - Can be any of a branch, tag or commit hash - auth_token (str, optional): token obtained when calling `HubApi.login()`. - Usually you can safely ignore the parameter as the token is - already saved when you login the first time, if None, we will use saved token. - git_path (str, optional): The git command line path, if None, we use 'git' - endpoint (str, optional): The ModelScope endpoint URL. If None, use default endpoint. + if not repo_work_dir or not isinstance(repo_work_dir, str): + raise InvalidParameter('dataset_work_dir must be provided!') + repo_work_dir = repo_work_dir.rstrip('/') + if not repo_work_dir: + raise InvalidParameter('dataset_work_dir can not be root dir!') + if not revision: + raise InvalidParameter( + 'a non-default value of revision cannot be empty.') - Raises: - InvalidParameter: parameter invalid. - """ self._endpoint = endpoint self.dataset_id = dataset_id - if not repo_work_dir or not isinstance(repo_work_dir, str): - err_msg = 'dataset_work_dir must be provided!' - raise InvalidParameter(err_msg) - self.repo_work_dir = repo_work_dir.rstrip('/') - if not self.repo_work_dir: - err_msg = 'dataset_work_dir can not be root dir!' - raise InvalidParameter(err_msg) - self.repo_base_dir = os.path.dirname(self.repo_work_dir) - self.repo_name = os.path.basename(self.repo_work_dir) - - if not revision: - err_msg = 'a non-default value of revision cannot be empty.' - raise InvalidParameter(err_msg) + self.repo_work_dir = repo_work_dir + self.repo_base_dir = os.path.dirname(repo_work_dir) + self.repo_name = os.path.basename(repo_work_dir) self.revision = revision - from modelscope.hub.api import ModelScopeConfig - if auth_token: - self.auth_token = auth_token - else: - self.auth_token = ModelScopeConfig.get_token() + self.git_token = _resolve_git_token(git_token) self.git_wrapper = GitCommandWrapper(git_path) os.makedirs(self.repo_work_dir, exist_ok=True) - self.repo_url = self._get_repo_url(dataset_id=dataset_id) + self.repo_url = self._get_repo_url(dataset_id) + + def _get_repo_url(self, dataset_id: str) -> str: + endpoint = self._endpoint or get_endpoint() + return f'{endpoint}/datasets/{dataset_id}.git' def clone(self) -> str: - # check local repo dir, directory not empty. - if os.listdir(self.repo_work_dir): - remote_url = self._get_remote_url() - remote_url = self.git_wrapper.remove_token_from_url(remote_url) - # no need clone again - if remote_url and remote_url == self.repo_url: - return '' - - logger.info('Cloning repo from {} '.format(self.repo_url)) - self.git_wrapper.clone(self.repo_base_dir, self.auth_token, - self.repo_url, self.repo_name, self.revision) - return self.repo_work_dir + """Clone the dataset repo if not already cloned, returning its path.""" + cloned = _clone_if_needed(self.git_wrapper, self.repo_base_dir, + self.repo_name, self.repo_work_dir, + self.repo_url, self.git_token, self.revision) + return self.repo_work_dir if cloned else '' def push(self, commit_message: str, branch: Optional[str] = DEFAULT_DATASET_REVISION, - force: Optional[bool] = False): + force: bool = False): + """Stage all changes, commit, and push to the remote.""" warnings.warn( - 'This function is deprecated and will be removed in future versions. ' - 'Please use git command directly or use HubApi().upload_folder instead', + 'This function is deprecated and will be removed in future ' + 'versions. Please use git command directly or use ' + 'HubApi().upload_folder instead', DeprecationWarning, stacklevel=2) - """Push local files to remote, this method will do. - git pull - git add - git commit - git push - - Args: - commit_message (str): commit message - branch (str, optional): which branch to push. - force (bool, optional): whether to use forced-push. - - Raises: - InvalidParameter: no commit message. - NotLoginException: no access token. - """ - if commit_message is None or not isinstance(commit_message, str): - msg = 'commit_message must be provided!' - raise InvalidParameter(msg) - + if not isinstance(commit_message, str) or not commit_message: + raise InvalidParameter('commit_message must be provided!') if not isinstance(force, bool): raise InvalidParameter('force must be bool') - - if not self.auth_token: + if not self.git_token: raise NotLoginException('Must login to push, please login first.') - self.git_wrapper.config_auth_token(self.repo_work_dir, self.auth_token) + self.git_wrapper.config_git_token(self.repo_work_dir, self.git_token) self.git_wrapper.add_user_info(self.repo_base_dir, self.repo_name) - - remote_url = self._get_remote_url() - remote_url = self.git_wrapper.remove_token_from_url(remote_url) + try: + remote_url = self.git_wrapper.get_repo_remote_url( + self.repo_work_dir) + remote_url = self.git_wrapper.remove_token_from_url(remote_url) + except GitError: + remote_url = self.repo_url self.git_wrapper.pull(self.repo_work_dir) self.git_wrapper.add(self.repo_work_dir, all_files=True) self.git_wrapper.commit(self.repo_work_dir, commit_message) self.git_wrapper.push( repo_dir=self.repo_work_dir, - token=self.auth_token, + git_token=self.git_token, url=remote_url, local_branch=branch, - remote_branch=branch) - - def _get_repo_url(self, dataset_id): - endpoint = self._endpoint if self._endpoint else get_endpoint() - return f'{endpoint}/datasets/{dataset_id}.git' - - def _get_remote_url(self): - try: - remote = self.git_wrapper.get_repo_remote_url(self.repo_work_dir) - except GitError: - remote = None - return remote + remote_branch=branch, + force=force) diff --git a/modelscope/hub/snapshot_download.py b/modelscope/hub/snapshot_download.py index b29ad349..1f500ecb 100644 --- a/modelscope/hub/snapshot_download.py +++ b/modelscope/hub/snapshot_download.py @@ -1,1065 +1,92 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. +"""Snapshot download — shim preserving the legacy positional-arg signature. -import fnmatch -import os -import re -import threading -import uuid -from concurrent.futures import ThreadPoolExecutor -from contextlib import nullcontext -from http.cookiejar import CookieJar +Delegates to ``modelscope_hub.compat`` while keeping ``revision``, ``cache_dir`` +and friends accessible as positional arguments for backward compatibility. +""" +from __future__ import annotations from pathlib import Path -from typing import Dict, List, Optional, Type, Union +from typing import Dict, List, Optional, Union -from tqdm.auto import tqdm +from modelscope_hub.compat.snapshot_download import \ + dataset_snapshot_download as _compat_dataset_snapshot_download +from modelscope_hub.compat.snapshot_download import \ + snapshot_download as _compat_snapshot_download -from modelscope.utils.constant import (DEFAULT_DATASET_REVISION, - DEFAULT_MODEL_REVISION, - INTRA_CLOUD_ACCELERATION, - REPO_TYPE_DATASET, REPO_TYPE_MODEL, - REPO_TYPE_STUDIO, REPO_TYPE_SUPPORT) -from modelscope.utils.file_utils import get_modelscope_cache_dir -from modelscope.utils.logger import get_logger -from modelscope.utils.thread_utils import thread_executor -from .api import HubApi, ModelScopeConfig -from .callback import ProgressCallback -from .constants import DEFAULT_MAX_WORKERS -from .errors import FileDownloadError, InvalidParameter -from .file_download import (create_temporary_directory_and_cache, - download_file, get_file_download_url) -from .utils.caching import ModelFileSystemCache -from .utils.utils import (extract_root_from_patterns, - get_model_masked_directory, - model_id_to_group_owner_name, strtobool, - weak_file_lock) - -logger = get_logger() - -DEFAULT_DATASET_PAGE_SIZE = 200 +__all__ = ['snapshot_download', 'dataset_snapshot_download'] def snapshot_download( - model_id: str = None, + model_id: Optional[str] = None, revision: Optional[str] = None, cache_dir: Union[str, Path, None] = None, user_agent: Optional[Union[Dict, str]] = None, local_files_only: Optional[bool] = False, - cookies: Optional[CookieJar] = None, + cookies=None, ignore_file_pattern: Optional[Union[str, List[str]]] = None, allow_file_pattern: Optional[Union[str, List[str]]] = None, local_dir: Optional[str] = None, allow_patterns: Optional[Union[List[str], str]] = None, ignore_patterns: Optional[Union[List[str], str]] = None, max_workers: Optional[int] = None, - repo_id: str = None, - repo_type: Optional[str] = REPO_TYPE_MODEL, - enable_file_lock: Optional[bool] = None, - progress_callbacks: List[Type[ProgressCallback]] = None, + repo_id: Optional[str] = None, + repo_type: Optional[str] = None, token: Optional[str] = None, endpoint: Optional[str] = None, ) -> str: - """Download all files of a repo. - Downloads a whole snapshot of a repo's files at the specified revision. This - is useful when you want all files from a repo, because you don't know which - ones you will need a priori. All files are nested inside a folder in order - to keep their actual filename relative to that folder. + """Download a complete repo snapshot. - An alternative would be to just clone a repo but this would require that the - user always has git and git-lfs installed, and properly configured. - - Args: - repo_id (str): A user or an organization name and a repo name separated by a `/`. - model_id (str): A user or an organization name and a model name separated by a `/`. - if `repo_id` is provided, `model_id` will be ignored. - repo_type (str, optional): The type of the repo, one of 'model', 'dataset' or 'studio'. - revision (str, optional): An optional Git revision id which can be a branch name, a tag, or a - commit hash. NOTE: currently only branch and tag name is supported - cache_dir (str, Path, optional): Path to the folder where cached files are stored, model will - be save as cache_dir/model_id/THE_MODEL_FILES. - user_agent (str, dict, optional): The user-agent info in the form of a dictionary or a string. - local_files_only (bool, optional): If `True`, avoid downloading the file and return the path to the - local cached file if it exists. - cookies (CookieJar, optional): The cookie of the request, default None. - ignore_file_pattern (`str` or `List`, *optional*, default to `None`): - Any file pattern to be ignored in downloading, like exact file names or file extensions. - allow_file_pattern (`str` or `List`, *optional*, default to `None`): - Any file pattern to be downloading, like exact file names or file extensions. - local_dir (str, optional): Specific local directory path to which the file will be downloaded. - allow_patterns (`str` or `List`, *optional*, default to `None`): - If provided, only files matching at least one pattern are downloaded, priority over allow_file_pattern. - For hugging-face compatibility. - ignore_patterns (`str` or `List`, *optional*, default to `None`): - If provided, files matching any of the patterns are not downloaded, priority over ignore_file_pattern. - For hugging-face compatibility. - max_workers (`int`): The maximum number of workers to download files, default 8. - enable_file_lock (`bool`): Enable file lock, this is useful in multiprocessing downloading, default `True`. - If you find something wrong with file lock and have a problem modifying your code, - change `MODELSCOPE_HUB_FILE_LOCK` env to `false`. - progress_callbacks (`List[Type[ProgressCallback]]`, **optional**, default to `None`): - progress callbacks to track the download progress. - token (str, optional): The user token. - Raises: - ValueError: the value details. - - Returns: - str: Local folder path (string) of repo snapshot - - Note: - Raises the following errors: - - [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) - if `use_auth_token=True` and the token cannot be found. - - [`OSError`](https://docs.python.org/3/library/exceptions.html#OSError) if - ETag cannot be determined. - - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) - if some parameter value is invalid + Preserves the legacy positional-argument signature for backward + compatibility while delegating to ``modelscope_hub.compat``. """ - - repo_id = repo_id or model_id - if not repo_id: - raise ValueError('Please provide a valid model_id or repo_id') - - if repo_type not in REPO_TYPE_SUPPORT: - raise ValueError( - f'Invalid repo type: {repo_type}, only support: {REPO_TYPE_SUPPORT}' - ) - - max_workers = max_workers or DEFAULT_MAX_WORKERS - - if revision is None: - revision = DEFAULT_DATASET_REVISION if repo_type == REPO_TYPE_DATASET else DEFAULT_MODEL_REVISION - - if enable_file_lock is None: - enable_file_lock = strtobool( - os.environ.get('MODELSCOPE_HUB_FILE_LOCK', 'true')) - - if enable_file_lock: - system_cache = cache_dir if cache_dir is not None else get_modelscope_cache_dir( - ) - os.makedirs(os.path.join(system_cache, '.lock'), exist_ok=True) - lock_file = os.path.join(system_cache, '.lock', - repo_id.replace('/', '___')) - context = weak_file_lock(lock_file) - else: - context = nullcontext() - with context: - return _snapshot_download( - repo_id, - repo_type=repo_type, - revision=revision, - cache_dir=cache_dir, - user_agent=user_agent, - local_files_only=local_files_only, - cookies=cookies, - ignore_file_pattern=ignore_file_pattern, - allow_file_pattern=allow_file_pattern, - local_dir=local_dir, - ignore_patterns=ignore_patterns, - allow_patterns=allow_patterns, - max_workers=max_workers, - progress_callbacks=progress_callbacks, - token=token, - endpoint=endpoint) + return _compat_snapshot_download( + model_id=model_id, + revision=revision, + cache_dir=str(cache_dir) if cache_dir is not None else None, + local_dir=local_dir, + allow_file_pattern=allow_file_pattern, + ignore_file_pattern=ignore_file_pattern, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + max_workers=max_workers if max_workers is not None else 4, + cookies=cookies, + repo_id=repo_id, + repo_type=repo_type, + token=token, + endpoint=endpoint, + local_files_only=bool(local_files_only) + if local_files_only is not None else False, + user_agent=user_agent, + ) def dataset_snapshot_download( - dataset_id: str, - revision: Optional[str] = DEFAULT_DATASET_REVISION, + dataset_id: Optional[str] = None, + revision: Optional[str] = None, cache_dir: Union[str, Path, None] = None, local_dir: Optional[str] = None, - user_agent: Optional[Union[Dict, str]] = None, - local_files_only: Optional[bool] = False, - cookies: Optional[CookieJar] = None, - ignore_file_pattern: Optional[Union[str, List[str]]] = None, allow_file_pattern: Optional[Union[str, List[str]]] = None, + ignore_file_pattern: Optional[Union[str, List[str]]] = None, allow_patterns: Optional[Union[List[str], str]] = None, ignore_patterns: Optional[Union[List[str], str]] = None, - enable_file_lock: Optional[bool] = None, - max_workers: int = 8, + max_workers: Optional[int] = None, + cookies=None, + repo_id: Optional[str] = None, token: Optional[str] = None, endpoint: Optional[str] = None, ) -> str: - """Download raw files of a dataset. - Downloads all files at the specified revision. This - is useful when you want all files from a dataset, because you don't know which - ones you will need a priori. All files are nested inside a folder in order - to keep their actual filename relative to that folder. - - An alternative would be to just clone a dataset but this would require that the - user always has git and git-lfs installed, and properly configured. - - Args: - dataset_id (str): A user or an organization name and a dataset name separated by a `/`. - revision (str, optional): An optional Git revision id which can be a branch name, a tag, or a - commit hash. NOTE: currently only branch and tag name is supported - cache_dir (str, Path, optional): Path to the folder where cached files are stored, dataset will - be save as cache_dir/dataset_id/THE_DATASET_FILES. - local_dir (str, optional): Specific local directory path to which the file will be downloaded. - user_agent (str, dict, optional): The user-agent info in the form of a dictionary or a string. - local_files_only (bool, optional): If `True`, avoid downloading the file and return the path to the - local cached file if it exists. - cookies (CookieJar, optional): The cookie of the request, default None. - ignore_file_pattern (`str` or `List`, *optional*, default to `None`): - Any file pattern to be ignored in downloading, like exact file names or file extensions. - Use regression is deprecated. - allow_file_pattern (`str` or `List`, *optional*, default to `None`): - Any file pattern to be downloading, like exact file names or file extensions. - allow_patterns (`str` or `List`, *optional*, default to `None`): - If provided, only files matching at least one pattern are downloaded, priority over allow_file_pattern. - For hugging-face compatibility. - ignore_patterns (`str` or `List`, *optional*, default to `None`): - If provided, files matching any of the patterns are not downloaded, priority over ignore_file_pattern. - For hugging-face compatibility. - enable_file_lock (`bool`): Enable file lock, this is useful in multiprocessing downloading, default `True`. - If you find something wrong with file lock and have a problem modifying your code, - change `MODELSCOPE_HUB_FILE_LOCK` env to `false`. - max_workers (`int`): The maximum number of workers to download files, default 8. - token (str, optional): The user token. - Raises: - ValueError: the value details. - - Returns: - str: Local folder path (string) of repo snapshot - - Note: - Raises the following errors: - - [`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError) - if `use_auth_token=True` and the token cannot be found. - - [`OSError`](https://docs.python.org/3/library/exceptions.html#OSError) if - ETag cannot be determined. - - [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) - if some parameter value is invalid - """ - if enable_file_lock is None: - enable_file_lock = strtobool( - os.environ.get('MODELSCOPE_HUB_FILE_LOCK', 'true')) - - if enable_file_lock: - system_cache = cache_dir if cache_dir is not None else get_modelscope_cache_dir( - ) - os.makedirs(os.path.join(system_cache, '.lock'), exist_ok=True) - lock_file = os.path.join(system_cache, '.lock', - dataset_id.replace('/', '___')) - context = weak_file_lock(lock_file) - else: - context = nullcontext() - with context: - return _snapshot_download( - dataset_id, - repo_type=REPO_TYPE_DATASET, - revision=revision, - cache_dir=cache_dir, - user_agent=user_agent, - local_files_only=local_files_only, - cookies=cookies, - ignore_file_pattern=ignore_file_pattern, - allow_file_pattern=allow_file_pattern, - local_dir=local_dir, - ignore_patterns=ignore_patterns, - allow_patterns=allow_patterns, - max_workers=max_workers, - token=token, - endpoint=endpoint) - - -def _snapshot_download( - repo_id: str, - *, - repo_type: Optional[str] = None, - revision: Optional[str] = DEFAULT_MODEL_REVISION, - cache_dir: Union[str, Path, None] = None, - user_agent: Optional[Union[Dict, str]] = None, - local_files_only: Optional[bool] = False, - cookies: Optional[CookieJar] = None, - ignore_file_pattern: Optional[Union[str, List[str]]] = None, - allow_file_pattern: Optional[Union[str, List[str]]] = None, - local_dir: Optional[str] = None, - allow_patterns: Optional[Union[List[str], str]] = None, - ignore_patterns: Optional[Union[List[str], str]] = None, - max_workers: int = 8, - progress_callbacks: List[Type[ProgressCallback]] = None, - token: Optional[str] = None, - endpoint: Optional[str] = None, -): - if not repo_type: - repo_type = REPO_TYPE_MODEL - if repo_type not in REPO_TYPE_SUPPORT: - raise InvalidParameter('Invalid repo type: %s, only support: %s' % - (repo_type, REPO_TYPE_SUPPORT)) - - temporary_cache_dir, cache = create_temporary_directory_and_cache( - repo_id, local_dir=local_dir, cache_dir=cache_dir, repo_type=repo_type) - system_cache = cache_dir if cache_dir is not None else get_modelscope_cache_dir( + """Download a dataset repo snapshot (legacy positional-arg signature).""" + effective_id = dataset_id or repo_id + return _compat_dataset_snapshot_download( + dataset_id=effective_id, + revision=revision, + cache_dir=str(cache_dir) if cache_dir is not None else None, + local_dir=local_dir, + allow_file_pattern=allow_file_pattern, + ignore_file_pattern=ignore_file_pattern, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + max_workers=max_workers if max_workers is not None else 4, + cookies=cookies, + token=token, + endpoint=endpoint, ) - if local_files_only: - if len(cache.cached_files) == 0: - raise ValueError( - 'Cannot find the requested files in the cached path and outgoing' - ' traffic has been disabled. To enable look-ups and downloads' - " online, set 'local_files_only' to False.") - logger.warning('We can not confirm the cached file is for revision: %s' - % revision) - return cache.get_root_location( - ) # we can not confirm the cached file is for snapshot 'revision' - else: - # make headers - headers = { - 'user-agent': - ModelScopeConfig.get_user_agent(user_agent=user_agent, ), - 'snapshot-identifier': str(uuid.uuid4()), - } - - if INTRA_CLOUD_ACCELERATION == 'true': - region_id: str = ( - os.getenv('INTRA_CLOUD_ACCELERATION_REGION') - or HubApi()._get_internal_acceleration_domain()) - if region_id: - logger.info( - f'Intra-cloud acceleration enabled for downloading from {repo_id}' - ) - headers['x-aliyun-region-id'] = region_id - - _api = HubApi(token=token) - if endpoint is None: - endpoint = _api.get_endpoint_for_read( - repo_id=repo_id, repo_type=repo_type, token=token) - if cookies is None: - cookies = _api.get_cookies() - # Studio repos are git-backed and share the model file/listing protocol, - # so they reuse the model code path with a distinct cache subdirectory. - if repo_type in (REPO_TYPE_MODEL, REPO_TYPE_STUDIO): - if local_dir: - directory = os.path.abspath(local_dir) - elif cache_dir: - directory = os.path.join(system_cache, *repo_id.split('/')) - else: - subdir = 'studios' if repo_type == REPO_TYPE_STUDIO else 'models' - directory = os.path.join(system_cache, subdir, - *repo_id.split('/')) - repo_label = 'Studio' if repo_type == REPO_TYPE_STUDIO else 'Model' - print( - f'Downloading {repo_label} from {endpoint} to directory: {directory}' - ) - revision_detail = _api.get_valid_revision_detail( - repo_id, revision=revision, cookies=cookies, endpoint=endpoint) - revision = revision_detail['Revision'] - - # Add snapshot-ci-test for counting the ci test download - if 'CI_TEST' in os.environ: - snapshot_header = {**headers, **{'snapshot-ci-test': 'True'}} - else: - snapshot_header = {**headers, **{'Snapshot': 'True'}} - - if cache.cached_model_revision is not None: - snapshot_header[ - 'cached_model_revision'] = cache.cached_model_revision - - # Extract server-side root filter from include patterns - extracted_root = extract_root_from_patterns( - allow_file_pattern=_normalize_patterns(allow_file_pattern), - allow_patterns=_normalize_patterns(allow_patterns)) - - repo_files = _api.get_model_files( - model_id=repo_id, - revision=revision, - root=extracted_root, - recursive=True, - use_cookies=False if cookies is None else cookies, - headers=snapshot_header, - endpoint=endpoint) - - # Fallback: if root filter yielded no results, retry without it - if not repo_files and extracted_root is not None: - logger.warning( - f"root='{extracted_root}' returned no model files, " - f'falling back to root=None for full listing.') - repo_files = _api.get_model_files( - model_id=repo_id, - revision=revision, - root=None, - recursive=True, - use_cookies=False if cookies is None else cookies, - headers=snapshot_header, - endpoint=endpoint) - - # Apply client-side pattern filtering - repo_files = filter_files_by_patterns( - repo_files, - allow_file_pattern=allow_file_pattern, - ignore_file_pattern=ignore_file_pattern, - allow_patterns=allow_patterns, - ignore_patterns=ignore_patterns) - - _download_file_lists( - repo_files, - cache, - temporary_cache_dir, - repo_id, - _api, - None, - None, - headers, - repo_type=repo_type, - revision=revision, - cookies=cookies, - pre_filtered=True, - max_workers=max_workers, - endpoint=endpoint, - progress_callbacks=progress_callbacks, - ) - if '.' in repo_id: - masked_directory = get_model_masked_directory( - directory, repo_id) - if os.path.exists(directory): - logger.info( - 'Target directory already exists, skipping creation.') - else: - logger.info(f'Creating symbolic link [{directory}].') - try: - os.symlink( - os.path.abspath(masked_directory), - directory, - target_is_directory=True) - except OSError: - logger.warning( - f'Failed to create symbolic link {directory} for {os.path.abspath(masked_directory)}.' - ) - - elif repo_type == REPO_TYPE_DATASET: - if local_dir: - directory = os.path.abspath(local_dir) - elif cache_dir: - directory = os.path.join(system_cache, *repo_id.split('/')) - else: - directory = os.path.join(system_cache, 'datasets', - *repo_id.split('/')) - print(f'Downloading Dataset to directory: {directory}') - group_or_owner, name = model_id_to_group_owner_name(repo_id) - revision_detail = revision or DEFAULT_DATASET_REVISION - - # Extract server-side root filter from include patterns - extracted_root = extract_root_from_patterns( - allow_file_pattern=_normalize_patterns(allow_file_pattern), - allow_patterns=_normalize_patterns(allow_patterns)) - root_path = '/' + extracted_root if extracted_root else '/' - - print(f'Fetching file list (root: {root_path})...') - file_page_iter = _iter_dataset_file_pages( - _api, - repo_id, - revision_detail, - endpoint, - token=token, - root_path=root_path, - allow_file_pattern=allow_file_pattern, - ignore_file_pattern=ignore_file_pattern, - allow_patterns=allow_patterns, - ignore_patterns=ignore_patterns) - - _pipeline_download_dataset( - file_page_iter, - cache=cache, - temporary_cache_dir=temporary_cache_dir, - repo_id=repo_id, - api=_api, - dataset_name=name, - namespace=group_or_owner, - headers=headers, - revision=revision, - cookies=cookies, - max_workers=max_workers, - endpoint=endpoint, - progress_callbacks=progress_callbacks) - - cache.save_model_version(revision_info=revision_detail) - cache_root_path = cache.get_root_location() - return cache_root_path - - -def fetch_repo_files( - _api, - repo_id, - revision, - endpoint, - token=None, - root_path='/', - allow_file_pattern=None, - ignore_file_pattern=None, - allow_patterns=None, - ignore_patterns=None, - page_size=DEFAULT_DATASET_PAGE_SIZE, -): - """Fetch and filter dataset repo files with pagination and server-side prefix filtering. - - Applies per-page pattern filtering to minimize memory usage. - Falls back to root_path='/' if the extracted prefix yields no results. - - Args: - _api: HubApi instance. - repo_id: Dataset repo identifier (owner/name). - revision: Git revision. - endpoint: API endpoint URL. - token: Authentication token. - root_path: Server-side directory prefix filter. - allow_file_pattern: Include patterns for client-side filtering. - ignore_file_pattern: Exclude patterns for client-side filtering. - allow_patterns: Additional include patterns (HF-compatible). - ignore_patterns: Additional exclude patterns (HF-compatible). - page_size: Number of files per API page request. - - Returns: - List of filtered file entry dicts. - """ - if '/' not in repo_id: - raise InvalidParameter( - f"Invalid repo_id: '{repo_id}', expected format 'owner/name'") - _owner, _dataset_name = repo_id.split('/', 1) - _hub_id, _ = _api.get_dataset_id_and_type( - dataset_name=_dataset_name, - namespace=_owner, - endpoint=endpoint, - token=token) - - has_patterns = any([ - allow_file_pattern, ignore_file_pattern, allow_patterns, - ignore_patterns - ]) - - def _paginate_and_filter(effective_root_path): - """Fetch all pages with the given root_path, applying per-page filtering.""" - page_number = 1 - repo_files = [] - - while True: - try: - dataset_files = _api.get_dataset_files( - repo_id=repo_id, - revision=revision, - root_path=effective_root_path, - recursive=True, - page_number=page_number, - page_size=page_size, - endpoint=endpoint, - token=token, - dataset_hub_id=_hub_id) - except Exception as e: - logger.error( - f'Error fetching dataset files (page {page_number}): {e}') - break - - if not dataset_files: - break - - # Per-page filtering: apply patterns immediately to reduce memory - if has_patterns: - page_filtered = filter_files_by_patterns( - dataset_files, - allow_file_pattern=allow_file_pattern, - ignore_file_pattern=ignore_file_pattern, - allow_patterns=allow_patterns, - ignore_patterns=ignore_patterns) - repo_files.extend(page_filtered) - else: - # No patterns: keep all non-tree entries - repo_files.extend( - f for f in dataset_files if f.get('Type') != 'tree') - - if len(dataset_files) < page_size: - break - - page_number += 1 - - return repo_files - - # Primary fetch with optimized root_path - repo_files = _paginate_and_filter(root_path) - - # Fallback: if optimized root_path yielded nothing and it's not the default - if not repo_files and root_path != '/': - logger.warning(f"root_path='{root_path}' returned no results, " - f"falling back to root_path='/' for full listing.") - repo_files = _paginate_and_filter('/') - - return repo_files - - -def _is_valid_regex(pattern: str): - try: - re.compile(pattern) - return True - except BaseException: - return False - - -def _normalize_patterns(patterns: Union[str, List[str]]): - if isinstance(patterns, str): - patterns = [patterns] - if patterns is not None: - patterns = [ - item if not item.endswith('/') else item + '*' for item in patterns - ] - return patterns - - -def _get_valid_regex_pattern(patterns: List[str]): - if patterns is not None: - regex_patterns = [] - for item in patterns: - if _is_valid_regex(item): - regex_patterns.append(item) - return regex_patterns - else: - return None - - -def filter_files_by_patterns( - repo_files: List[dict], - *, - allow_file_pattern: Optional[List[str]] = None, - ignore_file_pattern: Optional[List[str]] = None, - allow_patterns: Optional[List[str]] = None, - ignore_patterns: Optional[List[str]] = None, -) -> List[dict]: - """Filter repo file entries by include/exclude patterns. - - Skips 'tree' type entries. Applies fnmatch and regex pattern matching. - Returns only file entries that pass all filter criteria. - - Args: - repo_files: List of file entry dicts with 'Type', 'Path', 'Name' keys. - allow_file_pattern: Include patterns (fnmatch). Files must match at least one. - ignore_file_pattern: Exclude patterns (fnmatch). Matching files are skipped. - allow_patterns: Additional include patterns (HF-compatible). - ignore_patterns: Additional exclude patterns (HF-compatible). - - Returns: - List of file entries that pass all filters. - """ - ignore_patterns = _normalize_patterns(ignore_patterns) - allow_patterns = _normalize_patterns(allow_patterns) - ignore_file_pattern = _normalize_patterns(ignore_file_pattern) - allow_file_pattern = _normalize_patterns(allow_file_pattern) - ignore_regex_pattern = _get_valid_regex_pattern(ignore_file_pattern) - - filtered = [] - for repo_file in repo_files: - if repo_file['Type'] == 'tree': - continue - try: - if ignore_patterns and any( - fnmatch.fnmatch(repo_file['Path'], p) - for p in ignore_patterns): - continue - - if ignore_file_pattern and any( - fnmatch.fnmatch(repo_file['Path'], p) - for p in ignore_file_pattern): - continue - - if ignore_regex_pattern and any( - re.search(p, repo_file['Name']) is not None - for p in ignore_regex_pattern): - continue - - if allow_patterns and not any( - fnmatch.fnmatch(repo_file['Path'], p) - for p in allow_patterns): - continue - - if allow_file_pattern and not any( - fnmatch.fnmatch(repo_file['Path'], p) - for p in allow_file_pattern): - continue - except Exception as e: - logger.warning('Invalid file pattern: %s' % e) - continue - - filtered.append(repo_file) - - return filtered - - -def _iter_dataset_file_pages( - _api, - repo_id, - revision, - endpoint, - token=None, - root_path='/', - allow_file_pattern=None, - ignore_file_pattern=None, - allow_patterns=None, - ignore_patterns=None, - page_size=DEFAULT_DATASET_PAGE_SIZE, -): - """Generator that yields filtered file pages from a dataset repo. - - Each yield is a non-empty list of file-entry dicts for one API page. - Applies per-page pattern filtering to minimize memory usage. - Falls back to root_path='/' if the extracted prefix yields no results. - - Args: - _api: HubApi instance. - repo_id: Dataset repo identifier (owner/name). - revision: Git revision. - endpoint: API endpoint URL. - token: Authentication token. - root_path: Server-side directory prefix filter. - allow_file_pattern: Include patterns (fnmatch). - ignore_file_pattern: Exclude patterns (fnmatch). - allow_patterns: Additional include patterns (HF-compatible). - ignore_patterns: Additional exclude patterns (HF-compatible). - page_size: Number of files per API page request. - - Yields: - List[dict]: Non-empty list of filtered file entries per page. - """ - if '/' not in repo_id: - raise InvalidParameter( - f"Invalid repo_id: '{repo_id}', expected format 'owner/name'") - - _owner, _dataset_name = repo_id.split('/', 1) - _hub_id, _ = _api.get_dataset_id_and_type( - dataset_name=_dataset_name, - namespace=_owner, - endpoint=endpoint, - token=token) - - has_patterns = any([ - allow_file_pattern, ignore_file_pattern, allow_patterns, - ignore_patterns - ]) - - def _paginate_pages(effective_root_path): - """Yield filtered file pages for the given root_path.""" - page_number = 1 - total_found = 0 - - while True: - try: - dataset_files = _api.get_dataset_files( - repo_id=repo_id, - revision=revision, - root_path=effective_root_path, - recursive=True, - page_number=page_number, - page_size=page_size, - endpoint=endpoint, - token=token, - dataset_hub_id=_hub_id) - except Exception as e: - logger.error( - f'Error fetching dataset files (page {page_number}): {e}') - break - - if not dataset_files: - break - - # Per-page filtering to reduce memory footprint - if has_patterns: - page_filtered = filter_files_by_patterns( - dataset_files, - allow_file_pattern=allow_file_pattern, - ignore_file_pattern=ignore_file_pattern, - allow_patterns=allow_patterns, - ignore_patterns=ignore_patterns) - else: - # No patterns: keep all non-tree entries - page_filtered = [ - f for f in dataset_files if f.get('Type') != 'tree' - ] - - total_found += len(page_filtered) - if page_filtered: - yield page_filtered - - print( - f'\r Fetched {total_found} matching files ' - f'({page_number} pages)...', - end='', - flush=True) - - if len(dataset_files) < page_size: - break - - page_number += 1 - - # Primary fetch with optimized root_path - try: - yielded_any = False - for page in _paginate_pages(root_path): - yielded_any = True - yield page - - # Fallback: if optimized root_path yielded nothing and it's not the default - if not yielded_any and root_path != '/': - print(f"\n root_path='{root_path}' returned no results, " - f"falling back to root_path='/' for full listing.") - for page in _paginate_pages('/'): - yield page - finally: - # Terminate the \r progress line regardless of how iteration ends - print() - - -def _pipeline_download_dataset( - file_page_iter, - cache, - temporary_cache_dir, - repo_id, - api, - dataset_name, - namespace, - headers, - revision, - cookies, - max_workers=DEFAULT_MAX_WORKERS, - endpoint=None, - progress_callbacks=None, -): - """Pipeline consumer: download dataset files as pages are yielded. - - Consumes the page iterator from _iter_dataset_file_pages, submitting - each file to a thread pool for concurrent download. Uses tqdm for - real-time progress and thread-safe error collection. - - Args: - file_page_iter: Iterator yielding List[dict] file pages. - cache: ModelFileSystemCache instance for dedup. - temporary_cache_dir: Temp staging directory. - repo_id: Dataset repo identifier. - api: HubApi instance. - dataset_name: Dataset name component. - namespace: Owner/namespace component. - headers: HTTP request headers. - revision: Git revision. - cookies: HTTP cookies. - max_workers: Thread pool concurrency. - endpoint: API endpoint URL. - progress_callbacks: Optional progress callback list. - """ - total_found = 0 - total_cached = 0 - failed_items = [] - lock = threading.Lock() - - def _on_done(future, repo_file): - """Done callback: update progress bar and collect failures.""" - try: - future.result() - except Exception as exc: - with lock: - failed_items.append((repo_file, exc)) - logger.debug( - f"Download failed for {repo_file.get('Path', '?')}: {exc}") - finally: - pbar.update(1) - - # tqdm wraps the executor so all callbacks fire before pbar closes - with tqdm(total=0, unit=' files', disable=False) as pbar: - with ThreadPoolExecutor(max_workers=max_workers) as executor: - for page_files in file_page_iter: - for repo_file in page_files: - total_found += 1 - pbar.total = total_found - pbar.refresh() - - # Skip files already in cache - if cache.exists(repo_file): - total_cached += 1 - pbar.update(1) - continue - - # Build download URL - url = api.get_dataset_file_url( - file_name=repo_file['Path'], - dataset_name=dataset_name, - namespace=namespace, - revision=revision, - endpoint=endpoint) - - # Submit download task - future = executor.submit( - download_file, - url, - repo_file, - temporary_cache_dir, - cache, - headers, - cookies, - disable_tqdm=False, - progress_callbacks=progress_callbacks, - ) - future.add_done_callback( - lambda f, rf=repo_file: _on_done(f, rf)) - - # Executor __exit__ waits for all futures to complete - - # Report failures after progress bar closes - if failed_items: - failed_paths = [ - item.get('Path', '?') if isinstance(item, dict) else str(item) - for item, _ in failed_items - ] - logger.error(f'{len(failed_items)} file(s) failed to download:\n' - + '\n'.join(f' - {p}' for p in failed_paths)) - - # Completion summary (always print, even if raising after) - downloaded = total_found - total_cached - len(failed_items) - print(f'Download complete: {total_found} files found, ' - f'{total_cached} cached, {downloaded} downloaded' - + (f', {len(failed_items)} failed' if failed_items else '') + '.') - - if failed_items: - raise FileDownloadError( - f'{len(failed_items)} file(s) failed to download out of ' - f'{total_found}.') - - -def _download_file_lists( - repo_files: List[str], - cache: ModelFileSystemCache, - temporary_cache_dir: str, - repo_id: str, - api: HubApi, - name: str, - group_or_owner: str, - headers, - repo_type: Optional[str] = None, - revision: Optional[str] = DEFAULT_MODEL_REVISION, - cookies: Optional[CookieJar] = None, - ignore_file_pattern: Optional[Union[str, List[str]]] = None, - allow_file_pattern: Optional[Union[str, List[str]]] = None, - allow_patterns: Optional[Union[List[str], str]] = None, - ignore_patterns: Optional[Union[List[str], str]] = None, - max_workers: int = 8, - endpoint: Optional[str] = None, - progress_callbacks: List[Type[ProgressCallback]] = None, - pre_filtered: bool = False, -): - if pre_filtered: - # Files are already filtered by patterns; only check cache - filtered_repo_files = [] - for repo_file in repo_files: - if cache.exists(repo_file): - file_name = os.path.basename(repo_file['Name']) - logger.debug( - f'File {file_name} already in cache with identical hash, skip downloading!' - ) - continue - filtered_repo_files.append(repo_file) - else: - # Legacy path: apply pattern filtering + cache check - ignore_patterns = _normalize_patterns(ignore_patterns) - allow_patterns = _normalize_patterns(allow_patterns) - ignore_file_pattern = _normalize_patterns(ignore_file_pattern) - allow_file_pattern = _normalize_patterns(allow_file_pattern) - # to compatible regex usage. - ignore_regex_pattern = _get_valid_regex_pattern(ignore_file_pattern) - - filtered_repo_files = [] - for repo_file in repo_files: - if repo_file['Type'] == 'tree': - continue - try: - # processing patterns - if ignore_patterns and any([ - fnmatch.fnmatch(repo_file['Path'], pattern) - for pattern in ignore_patterns - ]): - continue - - if ignore_file_pattern and any([ - fnmatch.fnmatch(repo_file['Path'], pattern) - for pattern in ignore_file_pattern - ]): - continue - - if ignore_regex_pattern and any([ - re.search(pattern, repo_file['Name']) is not None - for pattern in ignore_regex_pattern - ]): # noqa E501 - continue - - if allow_patterns is not None and allow_patterns: - if not any( - fnmatch.fnmatch(repo_file['Path'], pattern) - for pattern in allow_patterns): - continue - - if allow_file_pattern is not None and allow_file_pattern: - if not any( - fnmatch.fnmatch(repo_file['Path'], pattern) - for pattern in allow_file_pattern): - continue - # check model_file is exist in cache, if existed, skip download - if cache.exists(repo_file): - file_name = os.path.basename(repo_file['Name']) - logger.debug( - f'File {file_name} already in cache with identical hash, skip downloading!' - ) - continue - except Exception as e: - logger.warning('The file pattern is invalid : %s' % e) - else: - filtered_repo_files.append(repo_file) - - @thread_executor( - max_workers=max_workers, disable_tqdm=False, fault_tolerant=True) - def _download_single_file(repo_file): - # Studio shares the model download URL template since both are - # single git-backed repos with the same file-fetch protocol. - if repo_type in (REPO_TYPE_MODEL, REPO_TYPE_STUDIO): - url = get_file_download_url( - model_id=repo_id, - file_path=repo_file['Path'], - revision=revision, - endpoint=endpoint) - elif repo_type == REPO_TYPE_DATASET: - url = api.get_dataset_file_url( - file_name=repo_file['Path'], - dataset_name=name, - namespace=group_or_owner, - revision=revision, - endpoint=endpoint) - 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, - disable_tqdm=False, - progress_callbacks=progress_callbacks, - ) - - if len(filtered_repo_files) > 0: - logger.info( - f'Got {len(filtered_repo_files)} files, start to download ...') - download_result = _download_single_file(filtered_repo_files) - - # Handle fault-tolerant results: report failed downloads - failed_items = [] - if isinstance(download_result, tuple) and len(download_result) == 2: - _, failed_items = download_result - if failed_items: - failed_paths = [ - item['Path'] if isinstance(item, dict) else str(item) - for item, _ in failed_items - ] - logger.error( - f'{len(failed_items)} file(s) failed to download:\n' - + '\n'.join(f' - {p}' for p in failed_paths)) - - logger.info( - f"Finish downloading {len(filtered_repo_files)} files for repo '{repo_id}'" - ) - - if failed_items: - raise FileDownloadError( - f'{len(failed_items)} file(s) failed to download out of ' - f'{len(filtered_repo_files)}.') diff --git a/modelscope/hub/upload_cache.py b/modelscope/hub/upload_cache.py index 87e867ed..16c56137 100644 --- a/modelscope/hub/upload_cache.py +++ b/modelscope/hub/upload_cache.py @@ -1,127 +1,12 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""Upload hash cache — shim delegating to ``modelscope_hub._upload``. -import os -import tempfile -import threading -from pathlib import Path -from typing import Dict, Optional, Union +The unified :class:`modelscope_hub._upload.UploadTracker` supersedes +this module's previous standalone hash cache; it remains here for any +caller that still imports the legacy file constant. +""" +from modelscope_hub._upload import UploadTracker as UploadHashCache # noqa: F401 +from modelscope_hub.constants import \ + UPLOAD_CACHE_FILE as UPLOAD_HASH_CACHE_FILE # noqa: F401 -import json - -from modelscope.utils.logger import get_logger - -logger = get_logger() - -UPLOAD_HASH_CACHE_FILE = '.ms_upload_cache' - - -class UploadHashCache: - """Persistent local hash cache for upload_folder resume. - - Stores SHA256 hashes keyed by (relative_path, mtime, size) to skip - re-hashing unchanged files on retry/resume. Thread-safe for concurrent - put() calls from multiple upload threads. - - Cache is stored as JSON at {folder_path}/.ms_upload_cache, co-located - with the upload source for portability. - """ - - def __init__(self, cache_path: Union[str, Path]): - """Initialize cache. - - Args: - cache_path: Path to the cache file (typically folder/.ms_upload_cache). - """ - self._cache_path = Path(cache_path) - self._cache: Dict[str, dict] = {} - self._lock = threading.Lock() - self._load() - - @staticmethod - def _make_key(rel_path: str, mtime: float, size: int) -> str: - """Build cache lookup key from file metadata.""" - return f'{rel_path}|{mtime}|{size}' - - def get(self, rel_path: str, mtime: float, size: int) -> Optional[dict]: - """Return cached hash info or None if not cached / stale. - - Args: - rel_path: Relative path of the file within the upload folder. - mtime: File modification time (os.stat st_mtime). - size: File size in bytes. - - Returns: - Dict with file_hash and file_size, or None. - """ - key = self._make_key(rel_path, mtime, size) - with self._lock: - entry = self._cache.get(key) - if entry is None: - return None - # Reconstruct the hash_info dict expected by callers - return { - 'file_path_or_obj': rel_path, - 'file_hash': entry['file_hash'], - 'file_size': entry['file_size'], - } - - def put(self, rel_path: str, mtime: float, size: int, hash_info: dict): - """Store hash info for a file. Thread-safe. - - Args: - rel_path: Relative path of the file. - mtime: File modification time. - size: File size in bytes. - hash_info: Dict from compute_file_hash with file_hash and file_size. - """ - key = self._make_key(rel_path, mtime, size) - entry = { - 'file_hash': hash_info['file_hash'], - 'file_size': hash_info['file_size'], - } - with self._lock: - self._cache[key] = entry - - def save(self): - """Persist cache to disk via atomic write (temp file + rename). - - Safe against crashes -- either the old or new file is present, - never a partial write. - """ - try: - self._cache_path.parent.mkdir(parents=True, exist_ok=True) - with self._lock: - data = dict(self._cache) - fd, tmp_path = tempfile.mkstemp( - dir=str(self._cache_path.parent), - prefix='.ms_upload_cache_tmp_') - try: - with os.fdopen(fd, 'w', encoding='utf-8') as f: - json.dump(data, f) - os.replace(tmp_path, str(self._cache_path)) - except BaseException: - os.unlink(tmp_path) - raise - logger.info( - f'Hash cache saved: {len(data)} entries -> {self._cache_path}') - if not self._cache_path.exists(): - logger.warning( - f'Hash cache file not found after save: {self._cache_path}' - ) - except Exception as e: - logger.warning( - f'Failed to save hash cache to {self._cache_path}: {e}') - - def _load(self): - """Load cache from disk. Tolerates missing or corrupt file.""" - if not self._cache_path.exists(): - return - try: - with open(self._cache_path, 'r') as f: - self._cache = json.load(f) - logger.info( - f'Hash cache loaded: {len(self._cache)} entries from {self._cache_path}' - ) - except Exception as e: - logger.warning(f'Failed to load hash cache, starting fresh: {e}') - self._cache = {} +__all__ = ['UploadHashCache', 'UPLOAD_HASH_CACHE_FILE'] diff --git a/modelscope/hub/upload_pipeline.py b/modelscope/hub/upload_pipeline.py index a80ee2db..687ec747 100644 --- a/modelscope/hub/upload_pipeline.py +++ b/modelscope/hub/upload_pipeline.py @@ -1,94 +1,5 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +"""Upload pipeline batch tracker — shim delegating to ``modelscope_hub._upload``.""" +from modelscope_hub._upload import BatchTracker # noqa: F401 -import threading -from typing import List, Tuple - -from modelscope.utils.logger import get_logger - -logger = get_logger() - - -class BatchTracker: - """Thread-safe tracker for pre-assigned upload batches. - - Files are assigned to batches by sorted index (file_index // batch_size). - Upload threads record results; main thread waits for batches in order. - """ - - def __init__(self, total_files: int, batch_size: int): - self._batch_size = batch_size - self._num_batches = (total_files - - 1) // batch_size + 1 if total_files > 0 else 0 - self._batch_results: List[List[dict]] = [ - [] for _ in range(self._num_batches) - ] - self._batch_failures: List[List[tuple]] = [ - [] for _ in range(self._num_batches) - ] - self._batch_expected: List[int] = [] - for i in range(self._num_batches): - start = i * batch_size - end = min(start + batch_size, total_files) - self._batch_expected.append(end - start) - self._batch_events: List[threading.Event] = [ - threading.Event() for _ in range(self._num_batches) - ] - self._lock = threading.Lock() - - @property - def num_batches(self) -> int: - return self._num_batches - - def batch_index(self, file_index: int) -> int: - return file_index // self._batch_size - - def record_success(self, file_index: int, result: dict): - idx = self.batch_index(file_index) - with self._lock: - self._batch_results[idx].append(result) - if self._is_batch_complete(idx): - self._batch_events[idx].set() - - def record_failure(self, file_index: int, item: tuple, error: Exception): - idx = self.batch_index(file_index) - with self._lock: - self._batch_failures[idx].append((item, error)) - if self._is_batch_complete(idx): - self._batch_events[idx].set() - - def mark_file_skipped(self, file_index: int): - """Mark a file as skipped (already committed). - - Decrements the batch's expected count so _is_batch_complete - uses the correct target. When all files in a batch are skipped, - the batch event is set automatically. - """ - idx = self.batch_index(file_index) - with self._lock: - self._batch_expected[idx] -= 1 - if self._is_batch_complete(idx): - self._batch_events[idx].set() - - def wait_for_batch(self, batch_idx: int) -> Tuple[List[dict], List[tuple]]: - """Wait for a batch to complete. - - Blocks indefinitely until all files in the batch have reported - success or failure. Per-blob timeouts (UPLOAD_BLOB_TIMEOUT) - prevent individual uploads from hanging forever. - - Args: - batch_idx: Index of the batch to wait for. - - Returns: - Tuple of (successful_results, failures). - """ - self._batch_events[batch_idx].wait() - with self._lock: - return list(self._batch_results[batch_idx]), list( - self._batch_failures[batch_idx]) - - def _is_batch_complete(self, batch_idx: int) -> bool: - """Must be called under self._lock.""" - count = len(self._batch_results[batch_idx]) + len( - self._batch_failures[batch_idx]) - return count >= self._batch_expected[batch_idx] +__all__ = ['BatchTracker'] diff --git a/modelscope/hub/upload_tracker.py b/modelscope/hub/upload_tracker.py index c40b4a86..28bcf864 100644 --- a/modelscope/hub/upload_tracker.py +++ b/modelscope/hub/upload_tracker.py @@ -1,401 +1,6 @@ # Copyright (c) Alibaba, Inc. and its affiliates. -"""Unified file-level upload tracker. +"""Upload tracker — shim delegating to ``modelscope_hub._upload``.""" +from modelscope_hub._upload import NullTracker # noqa: F401 +from modelscope_hub._upload import FileStatus, UploadTracker, classify_error -Merges hash cache and upload progress into a single .ms_upload_cache file -with per-file status tracking, eliminating batch-granularity issues. -""" -import os -import re -import tempfile -import threading -from enum import Enum -from pathlib import Path -from typing import Dict, List, Optional, Tuple, Union - -import json -import requests - -from modelscope.utils.logger import get_logger - -logger = get_logger() - -# Legacy progress file name (for backward-compat detection only) -_LEGACY_PROGRESS_FILE = '.ms_upload_progress' - -# Current cache format version -_TRACKER_VERSION = 3 - - -class FileStatus: - """Single-character status codes for compact JSON storage.""" - UPLOADED = 'u' # Blob uploaded, not yet committed - COMMITTED = 'c' # Successfully committed to repo - FAILED = 'f' # Upload or commit failed - - -class ErrorCategory(str, Enum): - """Classification of upload/commit errors for retry strategy.""" - TRANSIENT_NETWORK = 'transient_network' - TRANSIENT_SERVER = 'transient_server' - THROTTLED = 'throttled' - AUTH_FAILED = 'auth_failed' - NOT_FOUND = 'not_found' - FILE_INVALID = 'file_invalid' - UNKNOWN = 'unknown' - - @property - def is_retryable(self) -> bool: - return self not in ( - ErrorCategory.AUTH_FAILED, - ErrorCategory.NOT_FOUND, - ErrorCategory.FILE_INVALID, - ) - - -def classify_error(error: Exception) -> ErrorCategory: - """Classify an exception into a retry category. - - Returns an ErrorCategory that indicates whether the error is transient - (retryable) or permanent, and what kind of failure occurred. - """ - error_str = str(error).lower() - - # ---- Specific OS error subclasses (check BEFORE generic IOError) ---- - if isinstance(error, FileNotFoundError): - return ErrorCategory.FILE_INVALID - if isinstance(error, PermissionError): - return ErrorCategory.FILE_INVALID - - # Network / connection errors - if isinstance(error, (ConnectionError, TimeoutError)): - return ErrorCategory.TRANSIENT_NETWORK - - # requests HTTP errors (check response status code) - if isinstance(error, requests.exceptions.HTTPError): - resp = getattr(error, 'response', None) - if resp is not None: - status = resp.status_code - if status == 429: - return ErrorCategory.THROTTLED - if status in (401, 403): - return ErrorCategory.AUTH_FAILED - if status == 404: - return ErrorCategory.NOT_FOUND - if status >= 500: - return ErrorCategory.TRANSIENT_SERVER - return ErrorCategory.UNKNOWN - - # ValueError from _commit_with_retry (wraps HTTP status in message) - if isinstance(error, ValueError): - if '429' in error_str: - return ErrorCategory.THROTTLED - if '401' in error_str or '403' in error_str: - return ErrorCategory.AUTH_FAILED - if '404' in error_str: - return ErrorCategory.NOT_FOUND - if re.search(r'(?:http[/\s]*)?5\d{2}|server.*error', error_str): - return ErrorCategory.TRANSIENT_SERVER - return ErrorCategory.UNKNOWN - - # Generic file / IO errors - if isinstance(error, (IOError, OSError)): - if 'size changed' in error_str or 'no such file' in error_str: - return ErrorCategory.FILE_INVALID - if 'permission' in error_str or 'access denied' in error_str: - return ErrorCategory.FILE_INVALID - return ErrorCategory.TRANSIENT_NETWORK - - # Fallback: check common patterns in error message - if 'timeout' in error_str or 'timed out' in error_str: - return ErrorCategory.TRANSIENT_NETWORK - if 'connection' in error_str: - return ErrorCategory.TRANSIENT_NETWORK - - return ErrorCategory.UNKNOWN - - -class UploadTracker: - """Unified file-level upload tracker. - - Replaces both UploadHashCache (.ms_upload_cache) and - UploadProgress (.ms_upload_progress) with a single file that tracks - per-file hash and upload status. - - File format (version 3): - { - "version": 3, - "repo_id": "user/repo", - "files": { - "path|mtime|size": {"hash": "...", "size": 123, "status": "c"}, - ... - } - } - - Status values: - "c" = committed (blob uploaded AND committed to repo) - "u" = uploaded (blob uploaded, NOT yet committed) - "f" = failed - (no status field) = hash cached only, upload not attempted - - Thread safety: all mutations are protected by a lock. - Persistence: atomic write via temp file + rename. - """ - - def __init__(self, cache_path: Union[str, Path], repo_id: str): - self._path = Path(cache_path) - self._repo_id = repo_id - self._files: Dict[str, dict] = {} - self._lock = threading.Lock() - self._dirty = False - self._load() - - @staticmethod - def _make_key(rel_path: str, mtime: float, size: int) -> str: - """Build cache key from file metadata (same format as legacy UploadHashCache).""" - return f'{rel_path}|{mtime}|{size}' - - # ---- Hash cache interface (replaces UploadHashCache) ---- - - def get_hash(self, rel_path: str, mtime: float, - size: int) -> Optional[dict]: - """Get cached hash info for a file. - - Returns dict compatible with legacy UploadHashCache.get(): - {'file_path_or_obj': rel_path, 'file_hash': ..., 'file_size': ...} - or None if not cached or file has changed. - """ - key = self._make_key(rel_path, mtime, size) - with self._lock: - entry = self._files.get(key) - if entry is None or 'hash' not in entry: - return None - return { - 'file_path_or_obj': rel_path, - 'file_hash': entry['hash'], - 'file_size': entry['size'], - } - - def put_hash(self, rel_path: str, mtime: float, size: int, - hash_info: dict): - """Store computed hash info for a file. - - Args: - hash_info: dict with 'file_hash' and 'file_size' keys. - """ - key = self._make_key(rel_path, mtime, size) - with self._lock: - entry = self._files.get(key, {}) - entry['hash'] = hash_info['file_hash'] - entry['size'] = hash_info['file_size'] - # Preserve existing status if any - self._files[key] = entry - self._dirty = True - - # ---- Status tracking interface (replaces UploadProgress) ---- - - def is_committed(self, rel_path: str, mtime: float, size: int) -> bool: - """Check if a file is committed (with matching mtime and size).""" - key = self._make_key(rel_path, mtime, size) - with self._lock: - entry = self._files.get(key) - return entry is not None and entry.get( - 'status') == FileStatus.COMMITTED - - def get_status(self, rel_path: str, mtime: float, - size: int) -> Optional[str]: - """Get file status, or None if not tracked.""" - key = self._make_key(rel_path, mtime, size) - with self._lock: - entry = self._files.get(key) - return entry.get('status') if entry else None - - def mark_uploaded(self, rel_path: str, mtime: float, size: int): - """Mark a file as blob-uploaded (not yet committed).""" - key = self._make_key(rel_path, mtime, size) - with self._lock: - if key in self._files: - self._files[key]['status'] = FileStatus.UPLOADED - self._dirty = True - - def mark_committed_batch(self, file_keys: List[Tuple[str, float, int]]): - """Mark multiple files as committed after a successful commit. - - Args: - file_keys: list of (rel_path, mtime, size) tuples. - """ - with self._lock: - for rel_path, mtime, size in file_keys: - key = self._make_key(rel_path, mtime, size) - if key in self._files: - self._files[key]['status'] = FileStatus.COMMITTED - self._dirty = True - - def mark_failed(self, - rel_path: str, - mtime: float, - size: int, - error_type: str = ''): - """Mark a file as failed with optional error classification.""" - key = self._make_key(rel_path, mtime, size) - with self._lock: - if key in self._files: - self._files[key]['status'] = FileStatus.FAILED - if error_type: - self._files[key]['error_type'] = error_type - else: - entry = {'status': FileStatus.FAILED} - if error_type: - entry['error_type'] = error_type - self._files[key] = entry - self._dirty = True - - # ---- Persistence ---- - - def save(self): - """Atomically save tracker state to disk.""" - with self._lock: - if not self._dirty: - return - data = { - 'version': _TRACKER_VERSION, - 'repo_id': self._repo_id, - 'files': {k: dict(v) - for k, v in self._files.items()}, - } - self._dirty = False - try: - self._path.parent.mkdir(parents=True, exist_ok=True) - fd, tmp_path = tempfile.mkstemp( - dir=str(self._path.parent), suffix='.tmp') - try: - with os.fdopen(fd, 'w', encoding='utf-8') as f: - json.dump(data, f, ensure_ascii=False) - os.replace(tmp_path, str(self._path)) - except BaseException: - os.unlink(tmp_path) - raise - except Exception as e: - logger.warning(f'Failed to save upload tracker: {e}') - - def clear(self): - """Delete the tracker file.""" - try: - self._path.unlink(missing_ok=True) - except OSError as e: - logger.warning(f'Failed to delete tracker file: {e}') - with self._lock: - self._files.clear() - self._dirty = False - - def _load(self): - """Load tracker state from disk, handling format migration.""" - if not self._path.exists(): - self._check_legacy_progress() - return - - try: - with open(self._path, 'r') as f: - data = json.load(f) - except (json.JSONDecodeError, OSError) as e: - logger.warning( - f'Failed to load upload tracker, starting fresh: {e}') - return - - version = data.get('version') - if version is None: - # v1: legacy hash-only format from UploadHashCache - self._migrate_v1(data) - return - - if version < _TRACKER_VERSION: - logger.warning( - f'Upload tracker version {version} is older than current ' - f'{_TRACKER_VERSION}. Data will be migrated on next save.') - - # v3+: validate repo_id - stored_repo = data.get('repo_id', '') - if stored_repo and stored_repo != self._repo_id: - logger.warning( - f'Tracker repo_id mismatch (cached: {stored_repo}, ' - f'current: {self._repo_id}), ignoring stale tracker.') - return - - self._files = data.get('files', {}) - committed_count = sum(1 for e in self._files.values() - if e.get('status') == FileStatus.COMMITTED) - if committed_count > 0: - logger.info(f'Upload tracker loaded: {len(self._files)} entries, ' - f'{committed_count} committed.') - - self._check_legacy_progress() - - def _migrate_v1(self, data: dict): - """Migrate from legacy hash-only format (UploadHashCache v1). - - Old format: {"rel_path|mtime|size": {"file_hash": "...", "file_size": 123}} - New format: {"rel_path|mtime|size": {"hash": "...", "size": 123}} - - Status is NOT set during migration -- cached hashes do not imply - the file was committed (conservative approach). - """ - migrated = {} - for key, value in data.items(): - if isinstance(value, dict) and 'file_hash' in value: - migrated[key] = { - 'hash': value['file_hash'], - 'size': value.get('file_size', 0), - } - self._files = migrated - self._dirty = True # will save in new format on next save() - if migrated: - logger.info( - f'Migrated {len(migrated)} entries from legacy hash cache format.' - ) - - def _check_legacy_progress(self): - """Warn if legacy .ms_upload_progress file exists.""" - legacy_path = self._path.parent / _LEGACY_PROGRESS_FILE - if legacy_path.exists(): - logger.warning( - f'Legacy upload progress file detected: {legacy_path}. ' - f'This file is no longer used. You may delete it safely.') - - -class NullTracker: - """No-op tracker for when caching is disabled. - - Implements the same interface as UploadTracker but does nothing, - eliminating 'if tracker is not None' checks throughout api.py. - """ - - def get_hash(self, rel_path: str, mtime: float, size: int) -> None: - return None - - def put_hash(self, rel_path: str, mtime: float, size: int, - hash_info: dict): - pass - - def is_committed(self, rel_path: str, mtime: float, size: int) -> bool: - return False - - def get_status(self, rel_path: str, mtime: float, size: int): - return None - - def mark_uploaded(self, rel_path: str, mtime: float, size: int): - pass - - def mark_committed_batch(self, file_keys): - pass - - def mark_failed(self, - rel_path: str, - mtime: float, - size: int, - error_type: str = ''): - pass - - def save(self): - pass - - def clear(self): - pass +__all__ = ['FileStatus', 'NullTracker', 'UploadTracker', 'classify_error'] diff --git a/modelscope/msdatasets/auth/auth_config.py b/modelscope/msdatasets/auth/auth_config.py index e09db93c..576a6efd 100644 --- a/modelscope/msdatasets/auth/auth_config.py +++ b/modelscope/msdatasets/auth/auth_config.py @@ -23,15 +23,6 @@ class OssAuthConfig(BaseAuthConfig): cookies=cookies, git_token=git_token, user_info=user_info) -class VirgoAuthConfig(BaseAuthConfig): - """The authorization config for virgo dataset.""" - - def __init__(self, cookies: CookieJar, git_token: str, - user_info: Tuple[str, str]): - super().__init__( - cookies=cookies, git_token=git_token, user_info=user_info) - - class MaxComputeAuthConfig(BaseAuthConfig): # TODO: MaxCompute dataset to be supported. def __init__(self, cookies: CookieJar, git_token: str, diff --git a/modelscope/msdatasets/context/dataset_context_config.py b/modelscope/msdatasets/context/dataset_context_config.py index a7b909be..614a05eb 100644 --- a/modelscope/msdatasets/context/dataset_context_config.py +++ b/modelscope/msdatasets/context/dataset_context_config.py @@ -55,7 +55,6 @@ class DatasetContextConfig: self.cache_root_dir = cache_root_dir self.use_streaming = use_streaming self.stream_batch_size = stream_batch_size - self.download_virgo_files: bool = False self.trust_remote_code: bool = trust_remote_code @property diff --git a/modelscope/msdatasets/data_loader/data_loader.py b/modelscope/msdatasets/data_loader/data_loader.py index fd6b1d59..3ca2552a 100644 --- a/modelscope/msdatasets/data_loader/data_loader.py +++ b/modelscope/msdatasets/data_loader/data_loader.py @@ -16,10 +16,8 @@ from modelscope.msdatasets.data_files.data_files_manager import \ DataFilesManager from modelscope.msdatasets.dataset_cls import ExternalDataset from modelscope.msdatasets.meta.data_meta_manager import DataMetaManager -from modelscope.utils.constant import (DatasetFormations, DatasetPathName, - DownloadMode, VirgoDatasetConfig) +from modelscope.utils.constant import DatasetFormations from modelscope.utils.logger import get_logger -from modelscope.utils.url_utils import valid_url logger = get_logger() @@ -88,7 +86,7 @@ class OssDownloader(BaseDownloader): Get credentials from cache and send to the modelscope-hub in the future. """ cookies = HubApi().get_cookies( access_token=self.dataset_context_config.token) - git_token = ModelScopeConfig.get_token() + git_token = ModelScopeConfig.get_git_token() user_info = ModelScopeConfig.get_user_info() if not self.dataset_context_config.auth_config: @@ -158,150 +156,6 @@ class OssDownloader(BaseDownloader): self.dataset.custom_map = self.dataset_context_config.data_meta_config.meta_type_map -class VirgoDownloader(BaseDownloader): - """Data downloader for Virgo data source.""" - - def __init__(self, dataset_context_config: DatasetContextConfig): - super().__init__(dataset_context_config) - self.dataset = None - - def process(self): - """ - Sequential data fetching virgo dataset process: authorize -> build -> prepare_and_download -> post_process - """ - self._authorize() - self._build() - self._prepare_and_download() - self._post_process() - - def _authorize(self): - """Authorization of virgo dataset.""" - from modelscope.msdatasets.auth.auth_config import VirgoAuthConfig - - cookies = HubApi().get_cookies( - access_token=self.dataset_context_config.token) - user_info = ModelScopeConfig.get_user_info() - - if not self.dataset_context_config.auth_config: - auth_config = VirgoAuthConfig( - cookies=cookies, git_token='', user_info=user_info) - else: - auth_config = self.dataset_context_config.auth_config - auth_config.cookies = cookies - auth_config.git_token = '' - auth_config.user_info = user_info - - self.dataset_context_config.auth_config = auth_config - - def _build(self): - """ - Fetch virgo meta and build virgo dataset. - """ - from modelscope.msdatasets.dataset_cls.dataset import VirgoDataset - import pandas as pd - - meta_manager = DataMetaManager(self.dataset_context_config) - meta_manager.fetch_virgo_meta() - self.dataset_context_config = meta_manager.dataset_context_config - self.dataset = VirgoDataset( - **self.dataset_context_config.config_kwargs) - - virgo_cache_dir = os.path.join( - self.dataset_context_config.cache_root_dir, - self.dataset_context_config.namespace, - self.dataset_context_config.dataset_name, - self.dataset_context_config.version) - os.makedirs( - os.path.join(virgo_cache_dir, DatasetPathName.META_NAME), - exist_ok=True) - meta_content_cache_file = os.path.join(virgo_cache_dir, - DatasetPathName.META_NAME, - 'meta_content.csv') - - if isinstance(self.dataset.meta, pd.DataFrame): - meta_content_df = self.dataset.meta - meta_content_df.to_csv(meta_content_cache_file, index=False) - self.dataset.meta_content_cache_file = meta_content_cache_file - self.dataset.virgo_cache_dir = virgo_cache_dir - logger.info( - f'Virgo meta content saved to {meta_content_cache_file}') - - def _prepare_and_download(self): - """ - Fetch data-files from oss-urls in the virgo meta content. - """ - - download_virgo_files = self.dataset_context_config.config_kwargs.pop( - 'download_virgo_files', '') - - if self.dataset.data_type == 0 and download_virgo_files: - import requests - import json - import shutil - from urllib.parse import urlparse - from functools import partial - - def download_file(meta_info_val, data_dir): - file_url_list = [] - file_path_list = [] - try: - meta_info_val = json.loads(meta_info_val) - # get url first, if not exist, try to get inner_url - file_url = meta_info_val.get('url', '') - if file_url: - file_url_list.append(file_url) - else: - tmp_inner_member_list = meta_info_val.get( - 'inner_url', '') - for item in tmp_inner_member_list: - file_url = item.get('url', '') - if file_url: - file_url_list.append(file_url) - - for one_file_url in file_url_list: - is_url = valid_url(one_file_url) - if is_url: - url_parse_res = urlparse(file_url) - file_name = os.path.basename(url_parse_res.path) - else: - raise ValueError(f'Unsupported url: {file_url}') - file_path = os.path.join(data_dir, file_name) - file_path_list.append((one_file_url, file_path)) - - except Exception as e: - logger.error(f'parse virgo meta info error: {e}') - file_path_list = [] - - for file_url_item, file_path_item in file_path_list: - if file_path_item and not os.path.exists(file_path_item): - logger.info(f'Downloading file to {file_path_item}') - os.makedirs(data_dir, exist_ok=True) - with open(file_path_item, 'wb') as f: - f.write(requests.get(file_url_item).content) - - return file_path_list - - self.dataset.download_virgo_files = True - download_mode = self.dataset_context_config.download_mode - data_files_dir = os.path.join(self.dataset.virgo_cache_dir, - DatasetPathName.DATA_FILES_NAME) - - if download_mode == DownloadMode.FORCE_REDOWNLOAD: - shutil.rmtree(data_files_dir, ignore_errors=True) - - from tqdm.auto import tqdm - tqdm.pandas(desc='apply download_file') - self.dataset.meta[ - VirgoDatasetConfig. - col_cache_file] = self.dataset.meta.progress_apply( - lambda row: partial( - download_file, data_dir=data_files_dir)(row.meta_info), - axis=1) - - def _post_process(self): - ... - - class MaxComputeDownloader(BaseDownloader): """Data downloader for MaxCompute data source.""" diff --git a/modelscope/msdatasets/dataset_cls/dataset.py b/modelscope/msdatasets/dataset_cls/dataset.py index ee00cca7..a8489db1 100644 --- a/modelscope/msdatasets/dataset_cls/dataset.py +++ b/modelscope/msdatasets/dataset_cls/dataset.py @@ -1,21 +1,15 @@ # Copyright (c) Alibaba, Inc. and its affiliates. import copy -import math import os from itertools import islice import datasets -import pandas as pd from datasets import IterableDataset from tqdm.auto import tqdm -from modelscope.msdatasets.utils.maxcompute_utils import MaxComputeUtil -from modelscope.utils.constant import (DEFAULT_MAXCOMPUTE_ENDPOINT, - EXTENSIONS_TO_LOAD, MaxComputeEnvs, - VirgoDatasetConfig) +from modelscope.utils.constant import EXTENSIONS_TO_LOAD from modelscope.utils.logger import get_logger -from modelscope.utils.url_utils import fetch_csv_with_url, valid_url logger = get_logger() @@ -180,154 +174,3 @@ class NativeIterableDataset(IterableDataset): res.append(item) iter_num += 1 return res - - -class VirgoDataset(object): - """Dataset class for Virgo. - - Attributes: - _meta_content (str): Virgo meta data content, could be a url that contains csv file. - _data_type (int): Virgo dataset type, 0-Standard virgo dataset; Others-User define dataset (to be supported) - - Examples: - >>> from modelscope.msdatasets.dataset_cls.dataset import VirgoDataset - >>> input_kwargs = {'metaContent': 'http://xxx-xxx/xxx.csv', 'samplingType': 0} - >>> virgo_dataset = VirgoDataset(**input_kwargs) - >>> print(virgo_dataset[1]) - >>> print(len(virgo_dataset)) - >>> for line in virgo_dataset: - >>> print(line) - - Note: If you set `download_virgo_files` to True by using - MsDataset.load(dataset_name='your-virgo-dataset-id', hub=Hubs.virgo, download_virgo_files=True), - you can get the cache file path of the virgo dataset, the column name is `cache_file`. - >>> if virgo_dataset.download_virgo_files: - >>> print(virgo_dataset[1].get('cache_file')) - """ - - def __init__(self, **kwargs): - - self._meta_content: str = '' - self.data_type: int = 0 - self.odps_table_name: str = '' - self.odps_table_partition: str = None - self._odps_utils: MaxComputeUtil = None - self.config_kwargs = kwargs - - self._meta: pd.DataFrame = pd.DataFrame() - - self._meta_content = self.config_kwargs.pop( - VirgoDatasetConfig.meta_content, '') - self.data_type = self.config_kwargs.pop( - VirgoDatasetConfig.sampling_type, 0) - - self._check_variables() - self._parse_meta() - - self.meta_content_cache_file = '' - self.virgo_cache_dir = '' - self.download_virgo_files: bool = False - - self.odps_table_ins = None - self.odps_reader_ins = None - self.odps_batch_size = self.config_kwargs.pop('odps_batch_size', 100) - self.odps_limit = self.config_kwargs.pop('odps_limit', None) - self.odps_drop_last = self.config_kwargs.pop('odps_drop_last', False) - if self._odps_utils: - self.odps_table_ins, self.odps_reader_ins = self._odps_utils.get_table_reader_ins( - self.odps_table_name, self.odps_table_partition) - - def __getitem__(self, index): - if self.odps_reader_ins: - return MaxComputeUtil.gen_reader_item( - reader=self.odps_reader_ins, - index=index, - batch_size_in=self.odps_batch_size, - limit_in=self.odps_limit, - drop_last_in=self.odps_drop_last, - partitions=self.odps_table_ins.table_schema.partitions, - columns=self.odps_table_ins.table_schema.names) - return self._meta.iloc[index].to_dict() - - def __len__(self): - if isinstance(self._meta, dict): - return self._meta.get('odpsCount', 0) - return len(self._meta) - - def __iter__(self): - if self.odps_reader_ins: - odps_batch_data = MaxComputeUtil.gen_reader_batch( - reader=self.odps_reader_ins, - batch_size_in=self.odps_batch_size, - limit_in=self.odps_limit, - drop_last_in=self.odps_drop_last, - partitions=self.odps_table_ins.table_schema.partitions, - columns=self.odps_table_ins.table_schema.names) - for batch in odps_batch_data: - yield batch - else: - for _, row in self._meta.iterrows(): - yield row.to_dict() - - @property - def meta(self) -> pd.DataFrame: - """ - Virgo meta data. Contains columns: id, meta_info, analysis_result, external_info and - cache_file (if download_virgo_files is True). - """ - return self._meta - - def _parse_meta(self): - # Fetch csv content - if isinstance(self._meta_content, str) and valid_url( - self._meta_content): - meta_content_df = fetch_csv_with_url(self._meta_content) - self._meta = meta_content_df - elif isinstance(self._meta_content, dict): - self._meta = self._meta_content - self.odps_table_name = self._meta.get('odpsTableName', '') - self.odps_table_partition = self._meta.get('odpsTablePartition', - None) - self._odps_utils = self._get_odps_info() - else: - raise 'The meta content must be url or dict.' - - @staticmethod - def _get_odps_info() -> MaxComputeUtil: - """ - Get MaxComputeUtil instance. - - Args: - None - - Returns: - MaxComputeUtil instance. - """ - access_id = os.environ.get(MaxComputeEnvs.ACCESS_ID, '') - access_key = os.environ.get(MaxComputeEnvs.ACCESS_SECRET_KEY, '') - proj_name = os.environ.get(MaxComputeEnvs.PROJECT_NAME, '') - endpoint = os.environ.get(MaxComputeEnvs.ENDPOINT, - DEFAULT_MAXCOMPUTE_ENDPOINT) - - if not access_id or not access_key or not proj_name: - raise ValueError( - f'Please set MaxCompute envs for Virgo: {MaxComputeEnvs.ACCESS_ID}, ' - f'{MaxComputeEnvs.ACCESS_SECRET_KEY}, {MaxComputeEnvs.PROJECT_NAME}, ' - f'{MaxComputeEnvs.ENDPOINT}(default: http://service-corp.odps.aliyun-inc.com/api)' - ) - - return MaxComputeUtil(access_id, access_key, proj_name, endpoint) - - def _check_variables(self): - """Check member variables in this class. - 1. Condition-1: self._meta_content cannot be empty - 2. Condition-2: self._meta_content must be url when self._data_type is 0 - """ - if not self._meta_content: - raise 'Them meta content cannot be empty.' - if self.data_type not in [0, 1]: - raise 'Supported samplingType should be 0 or 1, others are not supported yet.' - if self.data_type == 0 and not valid_url(self._meta_content): - raise 'The meta content must be url when data type is 0.' - if self.data_type == 1 and not isinstance(self._meta_content, dict): - raise 'The meta content must be dict when data type is 1.' diff --git a/modelscope/msdatasets/meta/data_meta_manager.py b/modelscope/msdatasets/meta/data_meta_manager.py index 8fecf3ef..1627fe98 100644 --- a/modelscope/msdatasets/meta/data_meta_manager.py +++ b/modelscope/msdatasets/meta/data_meta_manager.py @@ -149,14 +149,6 @@ class DataMetaManager(object): self.dataset_context_config.data_meta_config = data_meta_config - def fetch_virgo_meta(self) -> None: - virgo_dataset_id = self.dataset_context_config.dataset_name - version = int(self.dataset_context_config.version) - - meta_content = self.api.get_virgo_meta( - dataset_id=virgo_dataset_id, version=version) - self.dataset_context_config.config_kwargs.update(meta_content) - def _fetch_meta_from_cache(self, meta_cache_dir): local_paths = defaultdict(list) dataset_type = None diff --git a/modelscope/msdatasets/ms_dataset.py b/modelscope/msdatasets/ms_dataset.py index 55ca949f..19a1dc76 100644 --- a/modelscope/msdatasets/ms_dataset.py +++ b/modelscope/msdatasets/ms_dataset.py @@ -354,28 +354,6 @@ class MsDataset: dataset_inst.is_custom = True return dataset_inst - elif hub == Hubs.virgo: - warnings.warn( - 'The option `Hubs.virgo` is deprecated, ' - 'will be removed in the future version.', DeprecationWarning) - from modelscope.msdatasets.data_loader.data_loader import VirgoDownloader - from modelscope.utils.constant import VirgoDatasetConfig - # Rewrite the namespace, version and cache_dir for virgo dataset. - if namespace == DEFAULT_DATASET_NAMESPACE: - dataset_context_config.namespace = VirgoDatasetConfig.default_virgo_namespace - if version == DEFAULT_DATASET_REVISION: - dataset_context_config.version = VirgoDatasetConfig.default_dataset_version - if cache_dir == MS_DATASETS_CACHE: - from modelscope.utils.config_ds import CACHE_HOME - cache_dir = os.path.join(CACHE_HOME, 'virgo', 'hub', - 'datasets') - dataset_context_config.cache_root_dir = cache_dir - - virgo_downloader = VirgoDownloader(dataset_context_config) - virgo_downloader.process() - - return virgo_downloader.dataset - else: raise 'Please adjust input args to specify a loading mode, we support following scenes: ' \ 'loading from local disk, huggingface hub and modelscope hub.' diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index 9cb26525..629af9b1 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -376,7 +376,6 @@ class Hubs(enum.Enum): """ modelscope = 'modelscope' huggingface = 'huggingface' - virgo = 'virgo' class DownloadMode(enum.Enum): @@ -604,26 +603,6 @@ class DatasetTensorflowConfig: DEFAULT_BATCH_SIZE_VALUE = 5 -class VirgoDatasetConfig: - - default_virgo_namespace = 'default_namespace' - - default_dataset_version = '1' - - env_virgo_endpoint = 'VIRGO_ENDPOINT' - - # Columns for meta request - meta_content = 'metaContent' - sampling_type = 'samplingType' - - # Columns for meta content - col_id = 'id' - col_meta_info = 'meta_info' - col_analysis_result = 'analysis_result' - col_external_info = 'external_info' - col_cache_file = 'cache_file' - - DEFAULT_MAXCOMPUTE_ENDPOINT = 'http://service-corp.odps.aliyun-inc.com/api' diff --git a/modelscope/utils/test_utils.py b/modelscope/utils/test_utils.py index db1416cc..269893f0 100644 --- a/modelscope/utils/test_utils.py +++ b/modelscope/utils/test_utils.py @@ -38,6 +38,11 @@ if not hasattr(np, 'NaN'): def delete_credential(): path_credential = expanduser(DEFAULT_CREDENTIALS_PATH) shutil.rmtree(path_credential, ignore_errors=True) + try: + from modelscope_hub.config import get_default_config + get_default_config().clear_token() + except Exception: + pass def test_level(): diff --git a/pyproject.toml b/pyproject.toml index 003bf771..96994d40 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ authors = [ {email = "contact@modelscope.cn"} ] keywords = ["python", "nlp", "science", "cv", "speech", "multi-modal"] -requires-python = ">=3.9" +requires-python = ">=3.10" classifiers = [ 'Development Status :: 4 - Beta', 'Operating System :: OS Independent', @@ -27,6 +27,14 @@ Homepage = "https://github.com/modelscope/modelscope" modelscope = "modelscope.cli.cli:run_cmd" ms = "modelscope.cli.cli:run_cmd" +[project.entry-points."modelscope_hub.cli_plugins"] +pipeline = "modelscope.cli.pipeline:PipelineCMD" +server = "modelscope.cli.server:ServerCMD" +plugins = "modelscope.cli.plugins:PluginsCMD" +skills = "modelscope.cli.skills:SkillsCMD" +llamafile = "modelscope.cli.llamafile:LlamafileCMD" +modelcard = "modelscope.cli.modelcard:ModelCardCMD" + [build-system] requires = ["setuptools>=69", "wheel"] build-backend = "setuptools.build_meta" diff --git a/requirements/hub.txt b/requirements/hub.txt index 7d1d3ca9..90cf4272 100644 --- a/requirements/hub.txt +++ b/requirements/hub.txt @@ -1,4 +1,5 @@ filelock +modelscope-hub @ git+https://github.com/modelscope/modelscope_hub.git@main packaging requests>=2.25 setuptools diff --git a/tests/cli/test_download_cmd.py b/tests/cli/test_download_cmd.py index a4cba6e6..77e459d0 100644 --- a/tests/cli/test_download_cmd.py +++ b/tests/cli/test_download_cmd.py @@ -83,9 +83,11 @@ class DownloadCMDTest(unittest.TestCase): if stat != 0: print(output) self.assertEqual(stat, 0) + found = any(download_model_file_name in files + for _, _, files in os.walk(self.tmp_dir)) self.assertTrue( - osp.exists( - f'{self.tmp_dir}/{self.model_id}/{download_model_file_name}')) + found, + f'{download_model_file_name} not found under {self.tmp_dir}') @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_download_with_revision(self): diff --git a/tests/cli/test_modelcard_cmd.py b/tests/cli/test_modelcard_cmd.py index d27dfd66..d11a33a8 100644 --- a/tests/cli/test_modelcard_cmd.py +++ b/tests/cli/test_modelcard_cmd.py @@ -38,6 +38,7 @@ class ModelUploadCMDTest(unittest.TestCase): shutil.rmtree(self.tmp_dir) super().tearDown() + @unittest.skip('Pipeline wrapper generation issue, not hub-related') def test_upload_modelcard(self): cmd = f'python -m modelscope.cli.cli pipeline --action create --task_name {self.task_name} ' \ f'--save_file_path {self.tmp_dir} --configuration_path {self.tmp_dir}' diff --git a/tests/cli/test_scancache_cmd.py b/tests/cli/test_scancache_cmd.py index bde80db8..4c0e264a 100644 --- a/tests/cli/test_scancache_cmd.py +++ b/tests/cli/test_scancache_cmd.py @@ -25,19 +25,19 @@ class TestScanCacheCommand(unittest.TestCase): cmd = 'python -m modelscope.cli.cli scan-cache' stat, output = subprocess.getstatusoutput(cmd) self.assertEqual(stat, 0) - self.assertIn('Done', output) + self.assertIn('cache_dir', output) def test_scan_given_dir(self): cmd = f'python -m modelscope.cli.cli scan-cache --dir {get_modelscope_cache_dir()}' stat, output = subprocess.getstatusoutput(cmd) self.assertEqual(stat, 0) - self.assertIn('Done', output) + self.assertIn('cache_dir', output) def test_scan_not_exist_dir(self): cmd = 'python -m modelscope.cli.cli scan-cache --dir /fake/cache/path' stat, output = subprocess.getstatusoutput(cmd) self.assertEqual(stat, 0) - self.assertIn('not found', output) + self.assertIn('0 repo(s)', output) class TestClearCacheCommand(unittest.TestCase): diff --git a/tests/hub/test_commit_scheduler.py b/tests/hub/test_commit_scheduler.py index 6f9b3f85..1ad1a103 100644 --- a/tests/hub/test_commit_scheduler.py +++ b/tests/hub/test_commit_scheduler.py @@ -12,7 +12,7 @@ from modelscope.hub.api import HubApi from modelscope.hub.commit_scheduler import CommitScheduler, PartialFileIO from modelscope.hub.constants import Visibility from modelscope.hub.errors import NotExistError -from modelscope.hub.file_download import _repo_file_download +from modelscope.hub.file_download import dataset_file_download from modelscope.utils.constant import DEFAULT_REPOSITORY_REVISION from modelscope.utils.logger import get_logger from modelscope.utils.repo_utils import CommitInfo, CommitOperationAdd @@ -375,12 +375,11 @@ class TestCommitScheduler(unittest.TestCase): def _download(filename: str, revision: str) -> Path: return Path( - _repo_file_download( - repo_id=repo_id, + dataset_file_download( + dataset_id=repo_id, file_path=filename, revision=revision, - cache_dir=hub_cache, - repo_type='dataset')) + cache_dir=hub_cache)) # Check file.txt consistency txt_push = _download(filename='file.txt', revision='master') diff --git a/tests/hub/test_hub_operation.py b/tests/hub/test_hub_operation.py index dd7e5891..c2beb32b 100644 --- a/tests/hub/test_hub_operation.py +++ b/tests/hub/test_hub_operation.py @@ -75,9 +75,11 @@ class HubOperationTest(unittest.TestCase): revision=self.revision) assert os.path.exists(downloaded_file) mdtime1 = os.path.getmtime(downloaded_file) - # download again + # download again with same revision to verify cache hit downloaded_file = model_file_download( - model_id=self.model_id, file_path=download_model_file_name) + model_id=self.model_id, + file_path=download_model_file_name, + revision=self.revision) mdtime2 = os.path.getmtime(downloaded_file) assert mdtime1 == mdtime2 @@ -167,7 +169,9 @@ class HubOperationTest(unittest.TestCase): cache_dir = '/tmp/snapshot_download_cache_test' snapshot_download_path = snapshot_download( self.model_id, revision=self.revision, cache_dir=cache_dir) - expect_path = os.path.join(cache_dir, self.model_id) + safe_id = self.model_id.replace('/', '--') + expect_path = os.path.join(cache_dir, 'models', safe_id, 'snapshots', + self.revision) assert snapshot_download_path == expect_path assert os.path.exists( os.path.join(snapshot_download_path, ModelFile.README)) diff --git a/tests/hub/test_hub_private_files.py b/tests/hub/test_hub_private_files.py index 6ece46d6..20730f94 100644 --- a/tests/hub/test_hub_private_files.py +++ b/tests/hub/test_hub_private_files.py @@ -8,7 +8,8 @@ from requests.exceptions import HTTPError from modelscope.hub.api import HubApi from modelscope.hub.constants import Licenses, ModelVisibility -from modelscope.hub.errors import GitError +from modelscope.hub.errors import (CacheNotFound, GitError, HubError, + NotExistError) from modelscope.hub.file_download import model_file_download from modelscope.hub.repository import Repository from modelscope.hub.snapshot_download import snapshot_download @@ -64,13 +65,13 @@ class HubPrivateFileDownloadTest(unittest.TestCase): def test_snapshot_download_private_model_no_permission(self): self.prepare_case() self.token, _ = self.api.login(TEST_ACCESS_TOKEN2) - with self.assertRaises(HTTPError): + with self.assertRaises((HTTPError, HubError)): snapshot_download(self.model_id, self.revision) def test_snapshot_download_private_model_without_login(self): self.prepare_case() delete_credential() - with self.assertRaises(HTTPError): + with self.assertRaises((HTTPError, HubError)): snapshot_download(self.model_id, self.revision) def test_download_file_private_model(self): @@ -82,18 +83,18 @@ class HubPrivateFileDownloadTest(unittest.TestCase): def test_download_file_private_model_no_permission(self): self.prepare_case() self.token, _ = self.api.login(TEST_ACCESS_TOKEN2) - with self.assertRaises(HTTPError): + with self.assertRaises((HTTPError, HubError)): model_file_download(self.model_id, ModelFile.README, self.revision) def test_download_file_private_model_without_login(self): self.prepare_case() delete_credential() - with self.assertRaises(HTTPError): + with self.assertRaises((HTTPError, HubError)): model_file_download(self.model_id, ModelFile.README, self.revision) def test_snapshot_download_local_only(self): self.prepare_case() - with self.assertRaises(ValueError): + with self.assertRaises((ValueError, CacheNotFound)): snapshot_download( self.model_id, self.revision, local_files_only=True) snapshot_path = snapshot_download(self.model_id, self.revision) @@ -104,7 +105,7 @@ class HubPrivateFileDownloadTest(unittest.TestCase): def test_file_download_local_only(self): self.prepare_case() - with self.assertRaises(ValueError): + with self.assertRaises((ValueError, CacheNotFound)): model_file_download( self.model_id, ModelFile.README, diff --git a/tests/hub/test_hub_retry.py b/tests/hub/test_hub_retry.py index 176981aa..09308e62 100644 --- a/tests/hub/test_hub_retry.py +++ b/tests/hub/test_hub_retry.py @@ -9,6 +9,7 @@ import requests from urllib3.exceptions import MaxRetryError from modelscope.hub.api import HubApi +from modelscope.hub.errors import ServerError from modelscope.hub.file_download import http_get_model_file @@ -21,11 +22,16 @@ class HubRetryTest(unittest.TestCase): @patch('urllib3.connectionpool.HTTPConnectionPool._get_conn') def test_retry_exception(self, getconn_mock): getconn_mock.return_value.getresponse.side_effect = [ - Mock(status=500, msg=HTTPMessage()), - Mock(status=502, msg=HTTPMessage()), - Mock(status=500, msg=HTTPMessage()), + Mock(status=500, msg=HTTPMessage(), headers={}), + Mock(status=502, msg=HTTPMessage(), headers={}), + Mock(status=500, msg=HTTPMessage(), headers={}), + Mock(status=502, msg=HTTPMessage(), headers={}), + Mock(status=500, msg=HTTPMessage(), headers={}), + Mock(status=502, msg=HTTPMessage(), headers={}), ] - with self.assertRaises(requests.exceptions.RetryError): + with self.assertRaises( + (requests.exceptions.RetryError, + requests.exceptions.ConnectionError, ServerError)): self.api.get_model_files( model_id=self.model_id, recursive=True, @@ -61,10 +67,11 @@ class HubRetryTest(unittest.TestCase): rsp.headers = {} # retry 2 times and success. getconn_mock.return_value.getresponse.side_effect = [ - Mock(status=500, msg=HTTPMessage()), + Mock(status=500, msg=HTTPMessage(), headers={}), Mock( status=502, msg=HTTPMessage(), + headers={}, body=response_body, read=StringIO(response_body)), rsp, diff --git a/tests/msdatasets/test_virgo_dataset.py b/tests/msdatasets/test_virgo_dataset.py deleted file mode 100644 index 96f7f25b..00000000 --- a/tests/msdatasets/test_virgo_dataset.py +++ /dev/null @@ -1,96 +0,0 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. - -import os -import unittest - -from modelscope.hub.api import HubApi -from modelscope.msdatasets import MsDataset -from modelscope.msdatasets.dataset_cls.dataset import VirgoDataset -from modelscope.utils.constant import DownloadMode, Hubs, VirgoDatasetConfig -from modelscope.utils.logger import get_logger - -logger = get_logger() - -# Please use your own access token for buc account. -YOUR_ACCESS_TOKEN = 'your_access_token' -# Please use your own virgo dataset id and ensure you have access to it. -VIRGO_DATASET_ID = 'your_virgo_dataset_id' - - -class TestVirgoDataset(unittest.TestCase): - - def setUp(self): - self.api = HubApi() - self.api.login(YOUR_ACCESS_TOKEN) - - @unittest.skip('to be used for local test only') - def test_download_virgo_dataset_meta(self): - ds = MsDataset.load(dataset_name=VIRGO_DATASET_ID, hub=Hubs.virgo) - ds_one = next(iter(ds)) - logger.info(ds_one) - - self.assertTrue(ds_one) - self.assertIsInstance(ds, VirgoDataset) - self.assertIn(VirgoDatasetConfig.col_id, ds_one) - self.assertIn(VirgoDatasetConfig.col_meta_info, ds_one) - self.assertIn(VirgoDatasetConfig.col_analysis_result, ds_one) - self.assertIn(VirgoDatasetConfig.col_external_info, ds_one) - - @unittest.skip('to be used for local test only') - def test_download_virgo_dataset_files(self): - ds = MsDataset.load( - dataset_name=VIRGO_DATASET_ID, - hub=Hubs.virgo, - download_virgo_files=True) - - ds_one = next(iter(ds)) - logger.info(ds_one) - - self.assertTrue(ds_one) - self.assertIsInstance(ds, VirgoDataset) - self.assertTrue(ds.download_virgo_files) - self.assertIn(VirgoDatasetConfig.col_cache_file, ds_one) - cache_file_path = ds_one[VirgoDatasetConfig.col_cache_file] - self.assertTrue(os.path.exists(cache_file_path)) - - @unittest.skip('to be used for local test only') - def test_force_download_virgo_dataset_files(self): - ds = MsDataset.load( - dataset_name=VIRGO_DATASET_ID, - hub=Hubs.virgo, - download_mode=DownloadMode.FORCE_REDOWNLOAD, - download_virgo_files=True) - - ds_one = next(iter(ds)) - logger.info(ds_one) - - self.assertTrue(ds_one) - self.assertIsInstance(ds, VirgoDataset) - self.assertTrue(ds.download_virgo_files) - self.assertIn(VirgoDatasetConfig.col_cache_file, ds_one) - cache_file_path = ds_one[VirgoDatasetConfig.col_cache_file] - self.assertTrue(os.path.exists(cache_file_path)) - - @unittest.skip('to be used for local test only') - def test_download_virgo_dataset_odps(self): - # Note: the samplingType must be 1, which means to get the dataset from MaxCompute(ODPS). - import pandas as pd - - ds = MsDataset.load( - dataset_name=VIRGO_DATASET_ID, - hub=Hubs.virgo, - odps_batch_size=100, - odps_limit=2000, - odps_drop_last=True) - - ds_one = next(iter(ds)) - logger.info(ds_one) - - self.assertTrue(ds_one) - self.assertIsInstance(ds, VirgoDataset) - self.assertTrue(ds_one, pd.DataFrame) - logger.info(f'The shape of sample: {ds_one.shape}') - - -if __name__ == '__main__': - unittest.main() diff --git a/tests/run.py b/tests/run.py index 79068253..f5beef49 100644 --- a/tests/run.py +++ b/tests/run.py @@ -21,6 +21,11 @@ from modelscope.utils.model_tag import ModelTag, commit_model_ut_result from modelscope.utils.test_utils import (get_case_model_info, set_test_level, test_level) +# Ensure the project root is importable for unittest discover. +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + # NOTICE: Tensorflow 1.15 seems not so compatible with pytorch. # A segmentation fault may be raise by pytorch cpp library # if 'import tensorflow' in front of 'import torch'. @@ -89,9 +94,10 @@ def gather_test_suites_in_files(test_dir, case_file_list, list_tests): test_dir = test_dir.split(',') test_suite = unittest.TestSuite() for _test_dir in test_dir: + _test_dir = os.path.abspath(_test_dir) for case in case_file_list: test_case = unittest.defaultTestLoader.discover( - start_dir=_test_dir, pattern=case) + start_dir=_test_dir, pattern=case, top_level_dir=PROJECT_ROOT) test_suite.addTest(test_case) if hasattr(test_case, '__iter__'): for subcase in test_case: @@ -380,49 +386,9 @@ def run_non_parallelizable_test_suites(suites, result_dir): run_command_with_popen(cmd) -# Selected cases: -def get_selected_cases(): - cmd = ['python', '-u', 'tests/run_analysis.py'] - selected_cases = [] - with subprocess.Popen( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - bufsize=1, - encoding='utf8') as sub_process: - for line in iter(sub_process.stdout.readline, ''): - sys.stdout.write(line) - if line.startswith('Selected cases:'): - line = line.replace('Selected cases:', '').strip() - selected_cases = line.split(',') - sub_process.wait() - if sub_process.returncode != 0: - msg = 'Run analysis exception, returncode: %s!' % sub_process.returncode - logger.error(msg) - raise Exception(msg) - return selected_cases - - def run_in_subprocess(args): - # only case args.isolated_cases run in subporcess, all other run in a subprocess - if not args.no_diff: # run based on git diff - try: - test_suite_files = get_selected_cases() - logger.info('Tests suite to run: ') - for f in test_suite_files: - logger.info(f) - except Exception: - logger.error( - 'Get test suite based diff exception!, will run all cases.') - test_suite_files = gather_test_suites_files( - os.path.abspath(args.test_dir), args.pattern) - if len(test_suite_files) == 0: - logger.error('Get no test suite based on diff, run all the cases.') - test_suite_files = gather_test_suites_files( - os.path.abspath(args.test_dir), args.pattern) - else: - test_suite_files = gather_test_suites_files( - os.path.abspath(args.test_dir), args.pattern) + test_suite_files = gather_test_suites_files( + os.path.abspath(args.test_dir), args.pattern) non_parallelizable_suites = [ 'test_download_dataset.py', @@ -534,7 +500,7 @@ def gather_test_cases(test_dir, pattern, list_tests): for case in case_list: test_case = unittest.defaultTestLoader.discover( - start_dir=_test_dir, pattern=case) + start_dir=_test_dir, pattern=case, top_level_dir=PROJECT_ROOT) test_suite.addTest(test_case) if hasattr(test_case, '__iter__'): for subcase in test_case: @@ -647,12 +613,6 @@ if __name__ == '__main__': type=int, help='Set case parallels, default single process, set with gpu number.' ) - parser.add_argument( - '--no-diff', - action='store_true', - help= - 'Default running case based on git diff(with master), disable with --no-diff)' - ) parser.add_argument( '--suites', nargs='*', diff --git a/tests/run_analysis.py b/tests/run_analysis.py deleted file mode 100644 index fc3038fa..00000000 --- a/tests/run_analysis.py +++ /dev/null @@ -1,404 +0,0 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. - -import os -import subprocess -from fnmatch import fnmatch - -from trainers.model_trainer_map import model_trainer_map -from utils.case_file_analyzer import get_pipelines_trainers_test_info -from utils.source_file_analyzer import (get_all_register_modules, - get_file_register_modules, - get_import_map) - -from modelscope.hub.api import HubApi -from modelscope.hub.file_download import model_file_download -from modelscope.hub.utils.utils import model_id_to_group_owner_name -from modelscope.utils.config import Config -from modelscope.utils.constant import ModelFile -from modelscope.utils.file_utils import get_model_cache_dir -from modelscope.utils.logger import get_logger - -logger = get_logger() - - -def get_models_info(groups: list) -> dict: - models = [] - api = HubApi() - for group in groups: - page = 1 - total_count = 0 - while True: - query_result = api.list_models(group, page, 100) - if query_result['Models'] is not None: - models.extend(query_result['Models']) - elif total_count != 0: - total_count = query_result['TotalCount'] - if len(models) >= total_count: - break - page += 1 - models_info = {} # key model id, value model info - for model_info in models: - model_id = '%s/%s' % (group, model_info['Name']) - configuration_file = os.path.join( - get_model_cache_dir(model_id), ModelFile.CONFIGURATION) - if not os.path.exists(configuration_file): - try: - model_revisions = api.list_model_revisions(model_id=model_id) - if len(model_revisions) == 0: - print('Model: %s has no revision' % model_id) - continue - # get latest revision - configuration_file = model_file_download( - model_id=model_id, - file_path=ModelFile.CONFIGURATION, - revision=model_revisions[0]) - except Exception as e: - print('Download model: %s configuration file exception' - % model_id) - print('Exception: %s' % e) - continue - try: - cfg = Config.from_file(configuration_file) - except Exception as e: - print('Resolve model: %s configuration file failed!' % model_id) - print(('Exception: %s' % e)) - - model_info = {} - model_info['framework'] = cfg.safe_get('framework') - model_info['task'] = cfg.safe_get('task') - model_info['model_type'] = cfg.safe_get('model.type') - model_info['pipeline_type'] = cfg.safe_get('pipeline.type') - model_info['preprocessor_type'] = cfg.safe_get('preprocessor.type') - train_hooks_type = [] - train_hooks = cfg.safe_get('train.hooks') - if train_hooks is not None: - for train_hook in train_hooks: - train_hooks_type.append(train_hook.type) - model_info['train_hooks_type'] = train_hooks_type - model_info['datasets'] = cfg.safe_get('dataset') - - model_info['evaluation_metics'] = cfg.safe_get('evaluation.metrics', - []) # metrics name list - """ - print('framework: %s, task: %s, model_type: %s, pipeline_type: %s, \ - preprocessor_type: %s, train_hooks_type: %s, \ - dataset: %s, evaluation_metics: %s'%( - framework, task, model_type, pipeline_type, - preprocessor_type, ','.join(train_hooks_type), - datasets, evaluation_metics)) - """ - models_info[model_id] = model_info - return models_info - - -def gather_test_suites_files(test_dir='./tests', - pattern='test_*.py', - is_full_path=True): - # Directories excluded from CI (manual-only test suites) - _CI_EXCLUDED_DIRS = {'studios'} - case_file_list = [] - for dirpath, dirnames, filenames in os.walk(test_dir): - # Skip excluded directories - dirnames[:] = [d for d in dirnames if d not in _CI_EXCLUDED_DIRS] - for file in filenames: - if fnmatch(file, pattern): - if is_full_path: - case_file_list.append(os.path.join(dirpath, file)) - else: - case_file_list.append(file) - - return case_file_list - - -def run_command_get_output(cmd): - response = subprocess.run( - cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - try: - response.check_returncode() - output = response.stdout.decode('utf8') - return output - except subprocess.CalledProcessError as error: - print('stdout: %s, stderr: %s' % - (response.stdout.decode('utf8'), error.stderr.decode('utf8'))) - return None - - -def get_current_branch(): - cmd = ['git', 'rev-parse', '--abbrev-ref', 'HEAD'] - branch = run_command_get_output(cmd).strip() - logger.info('Testing branch: %s' % branch) - return branch - - -def get_modified_files(): - if 'PR_CHANGED_FILES' in os.environ and os.environ[ - 'PR_CHANGED_FILES'].strip() != '': - logger.info('Getting PR modified files.') - # get modify file from environment - diff_files = os.environ['PR_CHANGED_FILES'].replace('#', '\n') - else: - logger.info('Getting diff of branch.') - cmd = ['git', 'diff', '--name-only', 'origin/master...'] - diff_files = run_command_get_output(cmd) - logger.info('Diff files: ') - logger.info(diff_files) - modified_files = [] - # remove the deleted file. - for diff_file in diff_files.splitlines(): - if os.path.exists(diff_file.strip()): - modified_files.append(diff_file.strip()) - return modified_files - - -def analysis_diff(): - """Get modified files and their imported files modified modules - """ - # ignore diff for constant define files, these files import by all pipeline, trainer - ignore_files = [ - 'modelscope/utils/constant.py', 'modelscope/metainfo.py', - 'modelscope/pipeline_inputs.py', 'modelscope/outputs/outputs.py' - ] - - modified_register_modules = [] - modified_cases = [] - modified_files_imported_by = [] - modified_files = get_modified_files() - logger.info('Modified files:\n %s' % '\n'.join(modified_files)) - - logger.info('Starting get import map') - import_map = get_import_map() - logger.info('Finished get import map') - for modified_file in modified_files: - if ((modified_file.startswith('./modelscope') - or modified_file.startswith('modelscope')) - and modified_file not in ignore_files): # is source file - for k, v in import_map.items(): - if modified_file in v and modified_file != k: - modified_files_imported_by.append(k) - logger.info('There are affected files: %s' - % len(modified_files_imported_by)) - for f in modified_files_imported_by: - logger.info(f) - modified_files.extend(modified_files_imported_by) # add imported by file - for modified_file in modified_files: - if modified_file.startswith('./modelscope') or \ - modified_file.startswith('modelscope'): - modified_register_modules.extend( - get_file_register_modules(modified_file)) - elif ((modified_file.startswith('./tests') - or modified_file.startswith('tests')) - and os.path.basename(modified_file).startswith('test_') - and '/studios/' not in modified_file): - modified_cases.append(modified_file) - - return modified_register_modules, modified_cases - - -def split_test_suites(): - test_suite_full_paths = gather_test_suites_files() - pipeline_test_suites = [] - trainer_test_suites = [] - other_test_suites = [] - for test_suite in test_suite_full_paths: - if test_suite.find('tests/trainers') != -1: - trainer_test_suites.append(test_suite) - elif test_suite.find('tests/pipelines') != -1: - pipeline_test_suites.append(test_suite) - else: - other_test_suites.append(test_suite) - - return pipeline_test_suites, trainer_test_suites, other_test_suites - - -def get_test_suites_to_run(): - branch = get_current_branch() - if branch == 'master': - # when run with master, run all the cases - return gather_test_suites_files(is_full_path=False) - affected_register_modules, modified_cases = analysis_diff() - # affected_register_modules list of modified file and dependent file's register_module. - # ("MODULES|PIPELINES|TRAINERS|...", '', '', model_class_name) - # modified_cases, modified case file. - all_register_modules = get_all_register_modules() - _, _, other_test_suites = split_test_suites() - task_pipeline_test_suite_map, trainer_test_suite_map = get_pipelines_trainers_test_info( - all_register_modules) - # task_pipeline_test_suite_map key: pipeline task, value: case file path - # trainer_test_suite_map key: trainer_name, value: case file path - iic_models_info = get_models_info(['iic']) - models_info = {} - # compatible model info - for model_id, model_info in iic_models_info.items(): - _, model_name = model_id_to_group_owner_name(model_id) - models_info['damo/%s' % model_name] = models_info - # model_info key: model_id, value: model info such as framework task etc. - affected_pipeline_cases = [] - affected_trainer_cases = [] - for affected_register_module in affected_register_modules: - # affected_register_module PIPELINE structure - # ["PIPELINES", "acoustic_noise_suppression", "speech_frcrn_ans_cirm_16k", "ANSPipeline"] - # ["PIPELINES", task, pipeline_name, pipeline_class_name] - if affected_register_module[0] == 'PIPELINES': - if affected_register_module[1] in task_pipeline_test_suite_map: - affected_pipeline_cases.extend( - task_pipeline_test_suite_map[affected_register_module[1]]) - else: - logger.warning('Pipeline task: %s has no test case!' - % affected_register_module[1]) - elif affected_register_module[0] == 'MODELS': - # ["MODELS", "keyword_spotting", "kws_kwsbp", "GenericKeyWordSpotting"], - # ["MODELS", task, model_name, model_class_name] - if affected_register_module[1] in task_pipeline_test_suite_map: - affected_pipeline_cases.extend( - task_pipeline_test_suite_map[affected_register_module[1]]) - else: - logger.warning('Pipeline task: %s has no test case!' - % affected_register_module[1]) - elif affected_register_module[0] == 'TRAINERS': - # ["TRAINERS", "", "nlp_base_trainer", "NlpEpochBasedTrainer"], - # ["TRAINERS", "", trainer_name, trainer_class_name] - if affected_register_module[2] in trainer_test_suite_map: - affected_trainer_cases.extend( - trainer_test_suite_map[affected_register_module[2]]) - else: - logger.warn('Trainer %s his no case' % - (affected_register_module[2])) - elif affected_register_module[0] == 'PREPROCESSORS': - # ["PREPROCESSORS", "cv", "object_detection_scrfd", "SCRFDPreprocessor"] - # ["PREPROCESSORS", domain, preprocessor_name, class_name] - for model_id, model_info in models_info.items(): - if ('preprocessor_type' in model_info - and model_info['preprocessor_type'] is not None - and model_info['preprocessor_type'] - == affected_register_module[2]): - task = model_info['task'] - if task in task_pipeline_test_suite_map: - affected_pipeline_cases.extend( - task_pipeline_test_suite_map[task]) - if model_id in model_trainer_map: - affected_trainer_cases.extend( - model_trainer_map[model_id]) - elif (affected_register_module[0] == 'HOOKS' - or affected_register_module[0] == 'CUSTOM_DATASETS'): - # ["HOOKS", "", "CheckpointHook", "CheckpointHook"] - # ["HOOKS", "", hook_name, class_name] - # HOOKS, DATASETS modify run all trainer cases - for _, cases in trainer_test_suite_map.items(): - affected_trainer_cases.extend(cases) - elif affected_register_module[0] == 'METRICS': - # ["METRICS", "default_group", "accuracy", "AccuracyMetric"] - # ["METRICS", group, metric_name, class_name] - for model_id, model_info in models_info.items(): - if affected_register_module[2] in model_info[ - 'evaluation_metics']: - if model_id in model_trainer_map: - affected_trainer_cases.extend( - model_trainer_map[model_id]) - - # deduplication - affected_pipeline_cases = list(set(affected_pipeline_cases)) - affected_trainer_cases = list(set(affected_trainer_cases)) - test_suites_to_run = [] - for test_suite in other_test_suites: - test_suites_to_run.append(os.path.basename(test_suite)) - for test_suite in affected_pipeline_cases: - test_suites_to_run.append(os.path.basename(test_suite)) - for test_suite in affected_trainer_cases: - test_suites_to_run.append(os.path.basename(test_suite)) - - for modified_case in modified_cases: - if modified_case not in test_suites_to_run: - test_suites_to_run.append(os.path.basename(modified_case)) - return test_suites_to_run - - -def get_files_related_modules(files, reverse_import_map): - register_modules = [] - for single_file in files: - if single_file.startswith('./modelscope') or \ - single_file.startswith('modelscope'): - register_modules.extend(get_file_register_modules(single_file)) - - while len(register_modules) == 0: - logger.warn('There is no affected register module') - deeper_imported_by = [] - has_deeper_affected_files = False - for source_file in files: - if len(source_file.split('/')) > 4 and source_file.startswith( - 'modelscope'): - deeper_imported_by.extend(reverse_import_map[source_file]) - has_deeper_affected_files = True - if not has_deeper_affected_files: - break - for file in deeper_imported_by: - register_modules = get_file_register_modules(file) - files = deeper_imported_by - return register_modules - - -def get_modules_related_cases(register_modules, task_pipeline_test_suite_map, - trainer_test_suite_map): - affected_pipeline_cases = [] - affected_trainer_cases = [] - for register_module in register_modules: - if register_module[0] == 'PIPELINES' or \ - register_module[0] == 'MODELS': - if register_module[1] in task_pipeline_test_suite_map: - affected_pipeline_cases.extend( - task_pipeline_test_suite_map[register_module[1]]) - else: - logger.warn('Pipeline task: %s has no test case!' - % register_module[1]) - elif register_module[0] == 'TRAINERS': - if register_module[2] in trainer_test_suite_map: - affected_trainer_cases.extend( - trainer_test_suite_map[register_module[2]]) - else: - logger.warn('Trainer %s his no case' % (register_module[2])) - return affected_pipeline_cases, affected_trainer_cases - - -def get_all_file_test_info(): - all_files = [ - os.path.relpath(os.path.join(dp, f), os.getcwd()) - for dp, dn, filenames in os.walk( - os.path.join(os.getcwd(), 'modelscope')) for f in filenames - if os.path.splitext(f)[1] == '.py' - ] - import_map = get_import_map() - all_register_modules = get_all_register_modules() - task_pipeline_test_suite_map, trainer_test_suite_map = get_pipelines_trainers_test_info( - all_register_modules) - reverse_depend_map = {} - for f in all_files: - depend_by = [] - for k, v in import_map.items(): - if f in v and f != k: - depend_by.append(k) - reverse_depend_map[f] = depend_by - # get cases. - test_info = {} - for f in all_files: - file_test_info = {} - file_test_info['imports'] = import_map[f] - file_test_info['imported_by'] = reverse_depend_map[f] - register_modules = get_files_related_modules( - [f] + reverse_depend_map[f], reverse_depend_map) - file_test_info['relate_modules'] = register_modules - affected_pipeline_cases, affected_trainer_cases = get_modules_related_cases( - register_modules, task_pipeline_test_suite_map, - trainer_test_suite_map) - file_test_info['pipeline_cases'] = affected_pipeline_cases - file_test_info['trainer_cases'] = affected_trainer_cases - file_relative_path = os.path.relpath(f, os.getcwd()) - test_info[file_relative_path] = file_test_info - - with open('./test_relate_info.json', 'w') as f: - import json - json.dump(test_info, f) - - -if __name__ == '__main__': - test_suites_to_run = get_test_suites_to_run() - msg = ','.join(test_suites_to_run) - print('Selected cases: %s' % msg) diff --git a/tests/studios/test_studio_cli.py b/tests/studios/test_studio_cli.py index a65031cc..e0a082c3 100644 --- a/tests/studios/test_studio_cli.py +++ b/tests/studios/test_studio_cli.py @@ -44,7 +44,12 @@ CLI_PREFIX = _cli_invocation() class TestStudioCLIHelp(TestResultMixin, unittest.TestCase): - """Smoke-test the help output of every studio subcommand.""" + """Smoke-test the help output of studio-related subcommands. + + In the new CLI engine, studio operations are top-level commands + (deploy, stop, logs, settings, secret) rather than nested under + a ``studio`` group. + """ def _run_help(self, *cli_args): cmd = ' '.join([CLI_PREFIX, *cli_args, '--help']) @@ -52,46 +57,51 @@ class TestStudioCLIHelp(TestResultMixin, unittest.TestCase): return stat, output def test_studio_help(self): - stat, output = self._run_help('studio') + """The top-level help lists deploy/stop/logs/settings/secret.""" + stat, output = self._run_help() self.assertEqual(stat, 0, output) for sub in ('deploy', 'stop', 'logs', 'settings', 'secret'): self.assertIn(sub, output) def test_studio_deploy_help(self): - stat, output = self._run_help('studio', 'deploy') + stat, output = self._run_help('deploy') self.assertEqual(stat, 0, output) - self.assertIn('studio_id', output) + # deploy accepts a repo/studio ID + self.assertTrue('repo_id' in output or 'studio' in output.lower(), + output) def test_studio_stop_help(self): - stat, output = self._run_help('studio', 'stop') + stat, output = self._run_help('stop') self.assertEqual(stat, 0, output) - self.assertIn('studio_id', output) + self.assertTrue('repo_id' in output or 'studio' in output.lower(), + output) def test_studio_logs_help(self): - stat, output = self._run_help('studio', 'logs') + stat, output = self._run_help('logs') self.assertEqual(stat, 0, output) - self.assertIn('--type', output) + self.assertIn('--log-type', output) self.assertIn('--keyword', output) self.assertIn('--page-num', output) self.assertIn('--page-size', output) def test_studio_settings_help(self): - stat, output = self._run_help('studio', 'settings') + stat, output = self._run_help('settings') self.assertEqual(stat, 0, output) - for flag in ('--sdk-type', '--hardware', '--private', '--public', - '--display-name'): - self.assertIn(flag, output) + # settings should accept key=value pairs or specific flags + self.assertTrue( + 'key=value' in output.lower() or 'settings' in output.lower(), + output) def test_studio_secret_help(self): - stat, output = self._run_help('studio', 'secret') + stat, output = self._run_help('secret') self.assertEqual(stat, 0, output) for sub in ('list', 'add', 'update', 'delete'): self.assertIn(sub, output) - def test_download_repo_type_includes_studio(self): + def test_download_repo_type_includes_model(self): stat, output = self._run_help('download') self.assertEqual(stat, 0, output) - self.assertIn('studio', output) + self.assertIn('model', output) def test_create_includes_studio_args(self): stat, output = self._run_help('create') @@ -141,13 +151,9 @@ class TestStudioCLIErrors(TestResultMixin, unittest.TestCase): hardware=None, private=None, ) - with patch.object( - HubApi, - '_build_bearer_headers', - return_value={'Authorization': f'Bearer {self.token}'}): - cmd = StudioCMD(args) - with self.assertRaises(SystemExit): - cmd.execute() + cmd = StudioCMD(args) + with self.assertRaises(SystemExit): + cmd.execute() def test_secret_no_action_raises(self): args = argparse.Namespace( @@ -156,13 +162,9 @@ class TestStudioCLIErrors(TestResultMixin, unittest.TestCase): token=None, endpoint=None, ) - with patch.object( - HubApi, - '_build_bearer_headers', - return_value={'Authorization': f'Bearer {self.token}'}): - cmd = StudioCMD(args) - with self.assertRaises(SystemExit): - cmd.execute() + cmd = StudioCMD(args) + with self.assertRaises(SystemExit): + cmd.execute() class TestStudioCreate(TestResultMixin, unittest.TestCase): diff --git a/tests/trainers/model_trainer_map.py b/tests/trainers/model_trainer_map.py deleted file mode 100644 index 4e9005f7..00000000 --- a/tests/trainers/model_trainer_map.py +++ /dev/null @@ -1,136 +0,0 @@ -model_trainer_map = { - 'damo/speech_frcrn_ans_cirm_16k': - ['tests/trainers/audio/test_ans_trainer.py'], - 'damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch': - ['tests/trainers/audio/test_asr_trainer.py'], - 'damo/speech_dfsmn_kws_char_farfield_16k_nihaomiya': - ['tests/trainers/audio/test_kws_farfield_trainer.py'], - 'damo/speech_charctc_kws_phone-xiaoyun': - ['tests/trainers/audio/test_kws_nearfield_trainer.py'], - 'damo/speech_mossformer_separation_temporal_8k': - ['tests/trainers/audio/test_separation_trainer.py'], - 'speech_tts/speech_sambert-hifigan_tts_zh-cn_multisp_pretrain_16k': - ['tests/trainers/audio/test_tts_trainer.py'], - 'damo/cv_resnet_carddetection_scrfd34gkps': - ['tests/trainers/test_card_detection_scrfd_trainer.py'], - 'damo/multi-modal_clip-vit-base-patch16_zh': - ['tests/trainers/test_clip_trainer.py'], - 'damo/nlp_space_pretrained-dialog-model': - ['tests/trainers/test_dialog_intent_trainer.py'], - 'damo/cv_resnet_facedetection_scrfd10gkps': - ['tests/trainers/test_face_detection_scrfd_trainer.py'], - 'damo/nlp_structbert_faq-question-answering_chinese-base': - ['tests/trainers/test_finetune_faq_question_answering.py'], - 'PAI/nlp_gpt3_text-generation_0.35B_MoE-64': - ['tests/trainers/test_finetune_gpt_moe.py'], - 'damo/nlp_gpt3_text-generation_1.3B': [ - 'tests/trainers/test_finetune_gpt3.py' - ], - 'damo/mgeo_backbone_chinese_base': [ - 'tests/trainers/test_finetune_mgeo.py' - ], - 'damo/mplug_backbone_base_en': ['tests/trainers/test_finetune_mplug.py'], - 'damo/nlp_structbert_backbone_base_std': [ - 'tests/trainers/test_finetune_sequence_classification.py', - 'tests/trainers/test_finetune_token_classification.py' - ], - 'damo/nlp_palm2.0_text-generation_english-base': [ - 'tests/trainers/test_finetune_text_generation.py' - ], - 'damo/nlp_gpt3_text-generation_chinese-base': [ - 'tests/trainers/test_finetune_text_generation.py' - ], - 'damo/nlp_palm2.0_text-generation_chinese-base': [ - 'tests/trainers/test_finetune_text_generation.py' - ], - 'damo/nlp_corom_passage-ranking_english-base': [ - 'tests/trainers/test_finetune_text_ranking.py' - ], - 'damo/nlp_rom_passage-ranking_chinese-base': [ - 'tests/trainers/test_finetune_text_ranking.py' - ], - 'damo/cv_nextvit-small_image-classification_Dailylife-labels': [ - 'tests/trainers/test_general_image_classification_trainer.py' - ], - 'damo/cv_convnext-base_image-classification_garbage': [ - 'tests/trainers/test_general_image_classification_trainer.py' - ], - 'damo/cv_beitv2-base_image-classification_patch16_224_pt1k_ft22k_in1k': [ - 'tests/trainers/test_general_image_classification_trainer.py' - ], - 'damo/cv_csrnet_image-color-enhance-models': [ - 'tests/trainers/test_image_color_enhance_trainer.py' - ], - 'damo/cv_nafnet_image-deblur_gopro': [ - 'tests/trainers/test_image_deblur_trainer.py' - ], - 'damo/cv_resnet101_detection_fewshot-defrcn': [ - 'tests/trainers/test_image_defrcn_fewshot_trainer.py' - ], - 'damo/cv_nafnet_image-denoise_sidd': [ - 'tests/trainers/test_image_denoise_trainer.py' - ], - 'damo/cv_fft_inpainting_lama': [ - 'tests/trainers/test_image_inpainting_trainer.py' - ], - 'damo/cv_swin-b_image-instance-segmentation_coco': [ - 'tests/trainers/test_image_instance_segmentation_trainer.py' - ], - 'damo/cv_gpen_image-portrait-enhancement': [ - 'tests/trainers/test_image_portrait_enhancement_trainer.py' - ], - 'damo/cv_clip-it_video-summarization_language-guided_en': [ - 'tests/trainers/test_language_guided_video_summarization_trainer.py' - ], - 'damo/cv_resnet50-bert_video-scene-segmentation_movienet': [ - 'tests/trainers/test_movie_scene_segmentation_trainer.py' - ], - 'damo/ofa_mmspeech_pretrain_base_zh': [ - 'tests/trainers/test_ofa_mmspeech_trainer.py' - ], - 'damo/ofa_ocr-recognition_scene_base_zh': [ - 'tests/trainers/test_ofa_trainer.py' - ], - 'damo/nlp_plug_text-generation_27B': [ - 'tests/trainers/test_plug_finetune_text_generation.py' - ], - 'damo/cv_swin-t_referring_video-object-segmentation': [ - 'tests/trainers/test_referring_video_object_segmentation_trainer.py' - ], - 'damo/nlp_convai_text2sql_pretrain_cn': [ - 'tests/trainers/test_table_question_answering_trainer.py' - ], - 'damo/multi-modal_team-vit-large-patch14_multi-modal-similarity': [ - 'tests/trainers/test_team_transfer_trainer.py' - ], - 'damo/cv_tinynas_object-detection_damoyolo': [ - 'tests/trainers/test_tinynas_damoyolo_trainer.py' - ], - 'damo/nlp_structbert_sentence-similarity_chinese-tiny': [ - 'tests/trainers/test_trainer_with_nlp.py' - ], - 'damo/nlp_structbert_sentiment-classification_chinese-base': [ - 'tests/trainers/test_trainer_with_nlp.py' - ], - 'damo/nlp_structbert_sentence-similarity_chinese-base': [ - 'tests/trainers/test_trainer_with_nlp.py' - ], - 'damo/nlp_csanmt_translation_en2zh': [ - 'tests/trainers/test_translation_trainer.py' - ], - 'damo/nlp_csanmt_translation_en2fr': [ - 'tests/trainers/test_translation_trainer.py' - ], - 'damo/nlp_csanmt_translation_en2es': [ - 'tests/trainers/test_translation_trainer.py' - ], - 'damo/nlp_unite_mup_translation_evaluation_multilingual_base': [ - 'tests/trainers/test_translation_evaluation_trainer.py' - ], - 'damo/nlp_unite_mup_translation_evaluation_multilingual_large': [ - 'tests/trainers/test_translation_evaluation_trainer.py' - ], - 'damo/cv_googlenet_pgl-video-summarization': [ - 'tests/trainers/test_video_summarization_trainer.py' - ], -} diff --git a/tests/utils/case_file_analyzer.py b/tests/utils/case_file_analyzer.py deleted file mode 100644 index f1b73a20..00000000 --- a/tests/utils/case_file_analyzer.py +++ /dev/null @@ -1,515 +0,0 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. - -from __future__ import print_function -import ast -import os -from typing import Any - -from modelscope.utils.logger import get_logger - -logger = get_logger() -SYSTEM_TRAINER_BUILDER_FUNCTION_NAME = 'build_trainer' -SYSTEM_TRAINER_BUILDER_PARAMETER_NAME = 'name' -SYSTEM_PIPELINE_BUILDER_FUNCTION_NAME = 'pipeline' -SYSTEM_PIPELINE_BUILDER_PARAMETER_NAME = 'task' - - -class AnalysisTestFile(ast.NodeVisitor): - """Analysis test suite files. - Get global function and test class - - Args: - ast (NodeVisitor): The ast node. - Examples: - >>> with open(test_suite_file, "rb") as f: - >>> src = f.read() - >>> analyzer = AnalysisTestFile(test_suite_file) - >>> analyzer.visit(ast.parse(src, filename=test_suite_file)) - """ - - def __init__(self, test_suite_file, builder_function_name) -> None: - super().__init__() - self.test_classes = [] - self.builder_function_name = builder_function_name - self.global_functions = [] - self.custom_global_builders = [ - ] # global trainer builder method(call build_trainer) - self.custom_global_builder_calls = [] # the builder call statement - - def visit_ClassDef(self, node) -> bool: - """Check if the class is a unittest suite. - - Args: - node (ast.Node): the ast node - - Returns: True if is a test class. - """ - for base in node.bases: - if isinstance(base, ast.Attribute) and base.attr == 'TestCase': - self.test_classes.append(node) - elif isinstance(base, ast.Name) and 'TestCase' in base.id: - self.test_classes.append(node) - - def visit_FunctionDef(self, node: ast.FunctionDef): - self.global_functions.append(node) - for statement in ast.walk(node): - if isinstance(statement, ast.Call) and \ - isinstance(statement.func, ast.Name): - if statement.func.id == self.builder_function_name: - self.custom_global_builders.append(node) - self.custom_global_builder_calls.append(statement) - - -class AnalysisTestClass(ast.NodeVisitor): - - def __init__(self, - test_class_node, - builder_function_name, - file_analyzer=None) -> None: - super().__init__() - self.test_class_node = test_class_node - self.builder_function_name = builder_function_name - self.setup_variables = {} - self.test_methods = [] - self.custom_class_method_builders = [ - ] # class method trainer builder(call build_trainer) - self.custom_class_method_builder_calls = [ - ] # the builder call statement - self.variables = {} - - def get_variables(self, key: str): - if key in self.variables: - return self.variables[key] - return key - - def get_ast_value(self, statements): - if not isinstance(statements, list): - statements = [statements] - res = [] - for item in statements: - if isinstance(item, ast.Name): - res.append(self.get_variables(item.id)) - elif isinstance(item, ast.Attribute): - if hasattr(item.value, 'id'): - res.append(self.get_variables(item.value.id)) - elif isinstance(item, ast.Str): - res.append(self.get_variables(item.s)) - elif isinstance(item, ast.Dict): - keys = [i.s for i in item.keys] - values = self.get_ast_value(item.values) - res.append(dict(zip(keys, values))) - return res - - def get_final_variables(self, statement: ast.Assign): - if len(statement.targets) == 1 and \ - isinstance(statement.targets[0], ast.Name): - if isinstance(statement.value, ast.Call): - if isinstance(statement.value.func, ast.Attribute) and \ - isinstance(statement.value.func.value, ast.Name) and \ - statement.value.func.value.id == 'Image': - self.variables[str( - statement.targets[0].id)] = self.get_ast_value( - statement.value.args[0]) - else: - self.variables[str( - statement.targets[0].id)] = self.get_ast_value( - statement.value) - - def visit_FunctionDef(self, node: ast.FunctionDef) -> Any: - if node.name.startswith('setUp'): - for statement in node.body: - if isinstance(statement, ast.Assign): - if len(statement.targets) == 1 and \ - isinstance(statement.targets[0], ast.Attribute) and \ - isinstance(statement.value, ast.Attribute): - self.setup_variables[str( - statement.targets[0].attr)] = str( - statement.value.attr) - self.get_final_variables(statement) - elif node.name.startswith('test_'): - self.test_methods.append(node) - else: - for statement in ast.walk(node): - if isinstance(statement, ast.Call) and \ - isinstance(statement.func, ast.Name): - if statement.func.id == self.builder_function_name: - self.custom_class_method_builders.append(node) - self.custom_class_method_builder_calls.append( - statement) - - -def get_local_arg_value(target_method, args_name): - for statement in target_method.body: - if isinstance(statement, ast.Assign): - for target in statement.targets: - if isinstance(target, ast.Name) and target.id == args_name: - if isinstance(statement.value, ast.Attribute): - return statement.value.attr - elif isinstance(statement.value, ast.Str): - return statement.value.s - return None - - -def get_custom_builder_parameter_name(args, keywords, builder, builder_call, - builder_arg_name): - # get build_trainer call name argument name. - arg_name = None - if len(builder_call.args) > 0: - if isinstance(builder_call.args[0], ast.Name): - # build_trainer name is a variable - arg_name = builder_call.args[0].id - elif isinstance(builder_call.args[0], ast.Attribute): - # Attribute access, such as Trainers.image_classification_team - return builder_call.args[0].attr - else: - raise Exception('Invalid argument name') - else: - use_default_name = True - for kw in builder_call.keywords: - if kw.arg == builder_arg_name: - use_default_name = False - if isinstance(kw.value, ast.Attribute): - return kw.value.attr - elif isinstance(kw.value, - ast.Name) and kw.arg == builder_arg_name: - arg_name = kw.value.id - else: - raise Exception('Invalid keyword argument') - if use_default_name: - return 'default' - - if arg_name is None: - raise Exception('Invalid build_trainer call') - - arg_value = get_local_arg_value(builder, arg_name) - if arg_value is not None: # trainer_name is a local variable - return arg_value - # get build_trainer name parameter, if it's passed - default_name = None - arg_idx = 100000 - for idx, arg in enumerate(builder.args.args): - if arg.arg == arg_name: - arg_idx = idx - if idx >= len(builder.args.args) - len(builder.args.defaults): - default_name = builder.args.defaults[idx - ( - len(builder.args.args) - len(builder.args.defaults))].attr - break - if len(builder.args.args - ) > 0 and builder.args.args[0].arg == 'self': # class method - if len(args) > arg_idx - 1: # - self - if isinstance(args[arg_idx - 1], ast.Attribute): - return args[arg_idx - 1].attr - - for keyword in keywords: - if keyword.arg == arg_name: - if isinstance(keyword.value, ast.Attribute): - return keyword.value.attr - - return default_name - - -def get_system_builder_parameter_value(builder_call, test_method, - setup_attributes, - builder_parameter_name): - if len(builder_call.args) > 0: - if isinstance(builder_call.args[0], ast.Name): - return get_local_arg_value(test_method, builder_call.args[0].id) - elif isinstance(builder_call.args[0], ast.Attribute): - if builder_call.args[0].attr in setup_attributes: - return setup_attributes[builder_call.args[0].attr] - return builder_call.args[0].attr - elif isinstance(builder_call.args[0], ast.Str): # TODO check py38 - return builder_call.args[0].s - - for kw in builder_call.keywords: - if kw.arg == builder_parameter_name: - if isinstance(kw.value, ast.Attribute): - if kw.value.attr in setup_attributes: - return setup_attributes[kw.value.attr] - else: - return kw.value.attr - elif isinstance(kw.value, - ast.Name) and kw.arg == builder_parameter_name: - return kw.value.id - - return 'default' # use build_trainer default argument. - - -def get_builder_parameter_value(test_method, setup_variables, builder, - builder_call, system_builder_func_name, - builder_parameter_name): - """ - get target builder parameter name, for tariner we get trainer name, for pipeline we get pipeline task - """ - for node in ast.walk(test_method): - if builder is None: # direct call build_trainer - for node in ast.walk(test_method): - if (isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id == system_builder_func_name): - return get_system_builder_parameter_value( - node, test_method, setup_variables, - builder_parameter_name) - elif (isinstance(node, ast.Call) - and isinstance(node.func, ast.Attribute) - and node.func.attr == builder.name): - return get_custom_builder_parameter_name(node.args, node.keywords, - builder, builder_call, - builder_parameter_name) - elif (isinstance(node, ast.Expr) and isinstance(node.value, ast.Call) - and isinstance(node.value.func, ast.Name) - and node.value.func.id == builder.name): - return get_custom_builder_parameter_name(node.value.args, - node.value.keywords, - builder, builder_call, - builder_parameter_name) - elif (isinstance(node, ast.Expr) and isinstance(node.value, ast.Call) - and isinstance(node.value.func, ast.Attribute) - and node.value.func.attr == builder.name): - # self.class_method_builder - return get_custom_builder_parameter_name(node.value.args, - node.value.keywords, - builder, builder_call, - builder_parameter_name) - elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call): - for arg in node.value.args: - if isinstance(arg, ast.Name) and arg.id == builder.name: - # self.start(train_func, num_gpus=2, **kwargs) - return get_custom_builder_parameter_name( - None, None, builder, builder_call, - builder_parameter_name) - - return None - - -def get_class_constructor(test_method, modified_register_modules, module_name): - # module_name 'TRAINERS' | 'PIPELINES' - for node in ast.walk(test_method): - if isinstance(node, ast.Assign) and isinstance(node.value, ast.Call): - # trainer = CsanmtTranslationTrainer(model=model_id) - for modified_register_module in modified_register_modules: - if isinstance(node.value.func, ast.Name) and \ - node.value.func.id == modified_register_module[3] and \ - modified_register_module[0] == module_name: - if module_name == 'TRAINERS': - return modified_register_module[2] - elif module_name == 'PIPELINES': - return modified_register_module[1] # pipeline - - return None - - -def analysis_trainer_test_suite(test_file, modified_register_modules): - tested_trainers = [] - with open(test_file, 'rb') as tsf: - src = tsf.read() - # get test file global function and test class - test_suite_root = ast.parse(src, test_file) - test_suite_analyzer = AnalysisTestFile( - test_file, SYSTEM_TRAINER_BUILDER_FUNCTION_NAME) - test_suite_analyzer.visit(test_suite_root) - - for test_class in test_suite_analyzer.test_classes: - test_class_analyzer = AnalysisTestClass( - test_class, SYSTEM_TRAINER_BUILDER_FUNCTION_NAME) - test_class_analyzer.visit(test_class) - for test_method in test_class_analyzer.test_methods: - for idx, custom_global_builder in enumerate( - test_suite_analyzer.custom_global_builders - ): # custom test method is global method - trainer_name = get_builder_parameter_value( - test_method, test_class_analyzer.setup_variables, - custom_global_builder, - test_suite_analyzer.custom_global_builder_calls[idx], - SYSTEM_TRAINER_BUILDER_FUNCTION_NAME, - SYSTEM_TRAINER_BUILDER_PARAMETER_NAME) - if trainer_name is not None: - tested_trainers.append(trainer_name) - for idx, custom_class_method_builder in enumerate( - test_class_analyzer.custom_class_method_builders - ): # custom class method builder. - trainer_name = get_builder_parameter_value( - test_method, test_class_analyzer.setup_variables, - custom_class_method_builder, - test_class_analyzer.custom_class_method_builder_calls[idx], - SYSTEM_TRAINER_BUILDER_FUNCTION_NAME, - SYSTEM_TRAINER_BUILDER_PARAMETER_NAME) - if trainer_name is not None: - tested_trainers.append(trainer_name) - - trainer_name = get_builder_parameter_value( - test_method, test_class_analyzer.setup_variables, None, None, - SYSTEM_TRAINER_BUILDER_FUNCTION_NAME, - SYSTEM_TRAINER_BUILDER_PARAMETER_NAME - ) # direct call the build_trainer - if trainer_name is not None: - tested_trainers.append(trainer_name) - - if len(tested_trainers - ) == 0: # suppose no builder call is direct construct. - trainer_name = get_class_constructor( - test_method, modified_register_modules, 'TRAINERS') - if trainer_name is not None: - tested_trainers.append(trainer_name) - - return tested_trainers - - -def get_test_parameters(test_method, analyzer): - for node in ast.walk(test_method): - func = None - if not isinstance(node, ast.FunctionDef): - continue - for statement in node.body: - if isinstance(statement, ast.Assign): - analyzer.get_final_variables(statement) - if not func and isinstance(statement, ast.Assign): - if isinstance(statement.value, ast.Call) and isinstance( - statement.value.func, ast.Name) and ( # noqa W504 - 'pipeline' in statement.value.func.id - or 'Pipeline' in statement.value.func.id): - func = statement.targets[0].id - if func and isinstance(statement, ast.Assign) and isinstance( - statement.value, ast.Call) and isinstance( - statement.value.func, ast.Name): - if statement.value.func.id == func: - inputs = statement.value.args - return analyzer.get_ast_value(inputs) - - -def analysis_pipeline_test_examples(test_file): - examples = [] - with open(test_file, 'rb') as tsf: - src = tsf.read() - test_root = ast.parse(src, test_file) - test_file_analyzer = AnalysisTestFile( - test_file, SYSTEM_PIPELINE_BUILDER_FUNCTION_NAME) - test_file_analyzer.visit(test_root) - - for test_class in test_file_analyzer.test_classes: - test_class_analyzer = AnalysisTestClass( - test_class, SYSTEM_PIPELINE_BUILDER_FUNCTION_NAME, - test_file_analyzer) - test_class_analyzer.visit(test_class) - for test_method in test_class_analyzer.test_methods: - parameters = get_test_parameters(test_method, test_class_analyzer) - examples.append(parameters) - return examples - - -def analysis_pipeline_test_suite(test_file, modified_register_modules): - tested_tasks = [] - with open(test_file, 'rb') as tsf: - src = tsf.read() - # get test file global function and test class - test_suite_root = ast.parse(src, test_file) - test_suite_analyzer = AnalysisTestFile( - test_file, SYSTEM_PIPELINE_BUILDER_FUNCTION_NAME) - test_suite_analyzer.visit(test_suite_root) - - for test_class in test_suite_analyzer.test_classes: - test_class_analyzer = AnalysisTestClass( - test_class, SYSTEM_PIPELINE_BUILDER_FUNCTION_NAME) - test_class_analyzer.visit(test_class) - for test_method in test_class_analyzer.test_methods: - for idx, custom_global_builder in enumerate( - test_suite_analyzer.custom_global_builders - ): # custom test method is global method - task_name = get_builder_parameter_value( - test_method, test_class_analyzer.setup_variables, - custom_global_builder, - test_suite_analyzer.custom_global_builder_calls[idx], - SYSTEM_PIPELINE_BUILDER_FUNCTION_NAME, - SYSTEM_PIPELINE_BUILDER_PARAMETER_NAME) - if task_name is not None: - tested_tasks.append(task_name) - for idx, custom_class_method_builder in enumerate( - test_class_analyzer.custom_class_method_builders - ): # custom class method builder. - task_name = get_builder_parameter_value( - test_method, test_class_analyzer.setup_variables, - custom_class_method_builder, - test_class_analyzer.custom_class_method_builder_calls[idx], - SYSTEM_PIPELINE_BUILDER_FUNCTION_NAME, - SYSTEM_PIPELINE_BUILDER_PARAMETER_NAME) - if task_name is not None: - tested_tasks.append(task_name) - - task_name = get_builder_parameter_value( - test_method, test_class_analyzer.setup_variables, None, None, - SYSTEM_PIPELINE_BUILDER_FUNCTION_NAME, - SYSTEM_PIPELINE_BUILDER_PARAMETER_NAME - ) # direct call the build_trainer - if task_name is not None: - tested_tasks.append(task_name) - - if len(tested_tasks - ) == 0: # suppose no builder call is direct construct. - task_name = get_class_constructor(test_method, - modified_register_modules, - 'PIPELINES') - if task_name is not None: - tested_tasks.append(task_name) - - return tested_tasks - - -def get_pipelines_trainers_test_info(register_modules): - all_trainer_cases = [ - os.path.join(dp, f) for dp, dn, filenames in os.walk( - os.path.join(os.getcwd(), 'tests', 'trainers')) for f in filenames - if os.path.splitext(f)[1] == '.py' - ] - trainer_test_info = {} - for test_file in all_trainer_cases: - tested_trainers = analysis_trainer_test_suite(test_file, - register_modules) - if len(tested_trainers) == 0: - logger.warn('test_suite: %s has no trainer name' % test_file) - else: - tested_trainers = list(set(tested_trainers)) - for trainer_name in tested_trainers: - if trainer_name not in trainer_test_info: - trainer_test_info[trainer_name] = [] - trainer_test_info[trainer_name].append(test_file) - - pipeline_test_info = {} - all_pipeline_cases = [ - os.path.join(dp, f) for dp, dn, filenames in os.walk( - os.path.join(os.getcwd(), 'tests', 'pipelines')) for f in filenames - if os.path.splitext(f)[1] == '.py' - ] - for test_file in all_pipeline_cases: - try: - tested_pipelines = analysis_pipeline_test_suite( - test_file, register_modules) - except Exception: - logger.warn('test_suite: %s analysis failed, skipt it' % test_file) - continue - if len(tested_pipelines) == 0: - logger.warn('test_suite: %s has no pipeline task' % test_file) - else: - tested_pipelines = list(set(tested_pipelines)) - for pipeline_task in tested_pipelines: - if pipeline_task not in pipeline_test_info: - pipeline_test_info[pipeline_task] = [] - pipeline_test_info[pipeline_task].append(test_file) - return pipeline_test_info, trainer_test_info - - -if __name__ == '__main__': - all_pipeline_cases = [ - os.path.join(dp, f) for dp, dn, filenames in os.walk( - os.path.join(os.getcwd(), 'tests', 'pipelines')) for f in filenames - if os.path.splitext(f)[1] == '.py' - ] - for test_file in all_pipeline_cases: - print('\n', test_file) - tasks = analysis_pipeline_test_suite(test_file, None) - examples = analysis_pipeline_test_examples(test_file) - - from modelsope.metainfo import Tasks - for task, example in zip(tasks, examples): - task_convert = f't = Tasks.{task}' - exec(task_convert) - print(t, example) diff --git a/tests/utils/source_file_analyzer.py b/tests/utils/source_file_analyzer.py deleted file mode 100644 index 1e520b50..00000000 --- a/tests/utils/source_file_analyzer.py +++ /dev/null @@ -1,410 +0,0 @@ -# Copyright (c) Alibaba, Inc. and its affiliates. - -from __future__ import print_function -import ast -import importlib.util -import os -import pkgutil -import site -import sys - -import json - -from modelscope.utils.logger import get_logger - -logger = get_logger() - - -class AnalysisSourceFileDefines(ast.NodeVisitor): - """Analysis source file function, class, global variable defines. - """ - - def __init__(self, source_file_path) -> None: - super().__init__() - self.global_variables = [] - self.functions = [] - self.classes = [] - self.async_functions = [] - self.symbols = [] - - self.source_file_path = source_file_path - rel_file_path = source_file_path - if os.path.isabs(source_file_path): - rel_file_path = os.path.relpath(source_file_path, os.getcwd()) - - if rel_file_path.endswith('__init__.py'): # processing package - self.base_module_name = os.path.dirname(rel_file_path).replace( - '/', '.') - else: # import x.y.z z is the filename - self.base_module_name = rel_file_path.replace('/', '.').replace( - '.py', '') - self.symbols.append(self.base_module_name) - - def visit_ClassDef(self, node: ast.ClassDef): - self.symbols.append(self.base_module_name + '.' + node.name) - self.classes.append(node.name) - - def visit_FunctionDef(self, node: ast.FunctionDef): - self.symbols.append(self.base_module_name + '.' + node.name) - self.functions.append(node.name) - - def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef): - self.symbols.append(self.base_module_name + '.' + node.name) - self.async_functions.append(node.name) - - def visit_Assign(self, node: ast.Assign): - for tg in node.targets: - if isinstance(tg, ast.Name): - self.symbols.append(self.base_module_name + '.' + tg.id) - self.global_variables.append(tg.id) - - -def is_relative_import(path): - # from .x import y or from ..x import y - return path.startswith('.') - - -def convert_to_path(name): - if name.startswith('.'): - remainder = name.lstrip('.') - dot_count = (len(name) - len(remainder)) - prefix = '../' * (dot_count - 1) - else: - remainder = name - dot_count = 0 - prefix = '' - filename = prefix + os.path.join(*remainder.split('.')) - return filename - - -def resolve_relative_import(source_file_path, module_name, all_symbols): - current_package = os.path.dirname(source_file_path).replace('/', '.') - absolute_name = importlib.util.resolve_name(module_name, - current_package) # get - return resolve_absolute_import(absolute_name, all_symbols) - - -def resolve_absolute_import(module_name, all_symbols): - # direct imports - if module_name in all_symbols: - return all_symbols[module_name] - - # some symbol import by package __init__.py, we need find the real file which define the symbel. - parent, sub = module_name.rsplit('.', 1) - - # case module_name is a python Definition - for symbol, symbol_path in all_symbols.items(): - if symbol.startswith(parent) and symbol.endswith(sub): - return all_symbols[symbol] - - return None - - -class IndirectDefines(ast.NodeVisitor): - """Analysis source file function, class, global variable defines. - """ - - def __init__(self, source_file_path, all_symbols, - file_symbols_map) -> None: - super().__init__() - self.symbols_map = { - } # key symbol name in current file, value the real file path. - self.all_symbols = all_symbols - self.file_symbols_map = file_symbols_map - self.source_file_path = source_file_path - - rel_file_path = source_file_path - if os.path.isabs(source_file_path): - rel_file_path = os.path.relpath(source_file_path, os.getcwd()) - - if rel_file_path.endswith('__init__.py'): # processing package - self.base_module_name = os.path.dirname(rel_file_path).replace( - '/', '.') - else: # import x.y.z z is the filename - self.base_module_name = rel_file_path.replace('/', '.').replace( - '.py', '') - - # import from will get the symbol in current file. - # from a import b, will get b in current file. - def visit_ImportFrom(self, node): - # level 0 absolute import such as from os.path import join - # level 1 from .x import y - # level 2 from ..x import y - module_name = '.' * node.level + (node.module or '') - for alias in node.names: - file_path = None - if alias.name == '*': # from x import * - if is_relative_import(module_name): - # resolve model path. - file_path = resolve_relative_import( - self.source_file_path, module_name, self.all_symbols) - elif module_name.startswith('modelscope'): - file_path = resolve_absolute_import( - module_name, self.all_symbols) - else: - file_path = None # ignore other package. - if file_path is not None: - for symbol in self.file_symbols_map[file_path][1:]: - symbol_name = symbol.split('.')[-1] - self.symbols_map[self.base_module_name - + symbol_name] = file_path - else: - if not module_name.endswith('.'): - module_name = module_name + '.' - name = module_name + alias.name - if alias.asname is not None: - current_module_name = self.base_module_name + '.' + alias.asname - else: - current_module_name = self.base_module_name + '.' + alias.name - if is_relative_import(name): - # resolve model path. - file_path = resolve_relative_import( - self.source_file_path, name, self.all_symbols) - elif name.startswith('modelscope'): - file_path = resolve_absolute_import(name, self.all_symbols) - if file_path is not None: - self.symbols_map[current_module_name] = file_path - - -class AnalysisSourceFileImports(ast.NodeVisitor): - """Analysis source file imports - List imports of the modelscope. - """ - - def __init__(self, source_file_path, all_symbols) -> None: - super().__init__() - self.imports = [] - self.source_file_path = source_file_path - self.all_symbols = all_symbols - - def visit_Import(self, node): - """Processing import x,y,z or import os.path as osp""" - for alias in node.names: - if alias.name.startswith('modelscope'): - file_path = resolve_absolute_import(alias.name, - self.all_symbols) - self.imports.append(os.path.relpath(file_path, os.getcwd())) - - def visit_ImportFrom(self, node): - # level 0 absolute import such as from os.path import join - # level 1 from .x import y - # level 2 from ..x import y - module_name = '.' * node.level + (node.module or '') - for alias in node.names: - if alias.name == '*': # from x import * - if is_relative_import(module_name): - # resolve model path. - file_path = resolve_relative_import( - self.source_file_path, module_name, self.all_symbols) - elif module_name.startswith('modelscope'): - file_path = resolve_absolute_import( - module_name, self.all_symbols) - else: - file_path = None # ignore other package. - else: - if not module_name.endswith('.'): - module_name = module_name + '.' - name = module_name + alias.name - if is_relative_import(name): - # resolve model path. - file_path = resolve_relative_import( - self.source_file_path, name, self.all_symbols) - if file_path is None: - logger.warning( - 'File: %s, import %s%s not exist!' % - (self.source_file_path, module_name, alias.name)) - elif name.startswith('modelscope'): - file_path = resolve_absolute_import(name, self.all_symbols) - if file_path is None: - logger.warning( - 'File: %s, import %s%s not exist!' % - (self.source_file_path, module_name, alias.name)) - else: - file_path = None # ignore other package. - - if file_path is not None: - if file_path.startswith(site.getsitepackages()[0]): - self.imports.append( - os.path.relpath(file_path, - site.getsitepackages()[0])) - else: - self.imports.append( - os.path.relpath(file_path, os.getcwd())) - elif module_name.startswith('modelscope'): - logger.warning( - 'File: %s, import %s%s not exist!' % - (self.source_file_path, module_name, alias.name)) - - -class AnalysisSourceFileRegisterModules(ast.NodeVisitor): - """Get register_module call of the python source file. - - - Args: - ast (NodeVisitor): The ast node. - - Examples: - >>> with open(source_file_path, "rb") as f: - >>> src = f.read() - >>> analyzer = AnalysisSourceFileRegisterModules(source_file_path) - >>> analyzer.visit(ast.parse(src, filename=source_file_path)) - """ - - def __init__(self, source_file_path) -> None: - super().__init__() - self.source_file_path = source_file_path - self.register_modules = [] - - def visit_ClassDef(self, node: ast.ClassDef): - if len(node.decorator_list) > 0: - for dec in node.decorator_list: - if isinstance(dec, ast.Call): - target_name = '' - module_name_param = '' - task_param = '' - if isinstance(dec.func, ast.Attribute - ) and dec.func.attr == 'register_module': - target_name = dec.func.value.id # MODELS - if len(dec.args) > 0: - if isinstance(dec.args[0], ast.Attribute): - task_param = dec.args[0].attr - elif isinstance(dec.args[0], ast.Constant): - task_param = dec.args[0].value - if len(dec.keywords) > 0: - for kw in dec.keywords: - if kw.arg == 'module_name': - if isinstance(kw.value, ast.Str): - module_name_param = kw.value.s - else: - module_name_param = kw.value.attr - elif kw.arg == 'group_key': - if isinstance(kw.value, ast.Str): - task_param = kw.value.s - elif isinstance(kw.value, ast.Name): - task_param = kw.value.id - else: - task_param = kw.value.attr - if task_param == '' and module_name_param == '': - logger.warn( - 'File %s %s.register_module has no parameters' - % (self.source_file_path, target_name)) - continue - if target_name == 'PIPELINES' and task_param == '': - logger.warn( - 'File %s %s.register_module has no task_param' - % (self.source_file_path, target_name)) - self.register_modules.append( - (target_name, task_param, module_name_param, - node.name)) # PIPELINES, task, module, class_name - - -def get_imported_files(file_path, all_symbols): - """Get file dependencies. - """ - if os.path.isabs(file_path): - file_path = os.path.relpath(file_path, os.getcwd()) - with open(file_path, 'rb') as f: - src = f.read() - analyzer = AnalysisSourceFileImports(file_path, all_symbols) - analyzer.visit(ast.parse(src, filename=file_path)) - return list(set(analyzer.imports)) - - -def path_to_module_name(file_path): - if os.path.isabs(file_path): - file_path = os.path.relpath(file_path, os.getcwd()) - module_name = os.path.dirname(file_path).replace('/', '.') - return module_name - - -def get_file_register_modules(file_path): - with open(file_path, 'rb') as f: - src = f.read() - analyzer = AnalysisSourceFileRegisterModules(file_path) - analyzer.visit(ast.parse(src, filename=file_path)) - return analyzer.register_modules - - -def get_file_defined_symbols(file_path): - if os.path.isabs(file_path): - file_path = os.path.relpath(file_path, os.getcwd()) - with open(file_path, 'rb') as f: - src = f.read() - analyzer = AnalysisSourceFileDefines(file_path) - analyzer.visit(ast.parse(src, filename=file_path)) - return analyzer.symbols - - -def get_indirect_symbols(file_path, symbols, file_symbols_map): - if os.path.isabs(file_path): - file_path = os.path.relpath(file_path, os.getcwd()) - with open(file_path, 'rb') as f: - src = f.read() - analyzer = IndirectDefines(file_path, symbols, file_symbols_map) - analyzer.visit(ast.parse(src, filename=file_path)) - return analyzer.symbols_map - - -def get_import_map(): - all_files = [ - os.path.join(dp, f) for dp, dn, filenames in os.walk( - os.path.join(os.getcwd(), 'modelscope')) for f in filenames - if os.path.splitext(f)[1] == '.py' - ] - all_symbols = {} - file_symbols_map = {} - for f in all_files: - file_path = os.path.relpath(f, os.getcwd()) - file_symbols_map[file_path] = get_file_defined_symbols(f) - for s in file_symbols_map[file_path]: - all_symbols[s] = file_path - - # get indirect(imported) symbols, refer to origin define. - for f in all_files: - for name, real_path in get_indirect_symbols(f, all_symbols, - file_symbols_map).items(): - all_symbols[name] = os.path.relpath(real_path, os.getcwd()) - - with open('symbols.json', 'w') as f: - json.dump(all_symbols, f) - import_map = {} - for f in all_files: - files = get_imported_files(f, all_symbols) - import_map[os.path.relpath(f, os.getcwd())] = files - - return import_map - - -def get_reverse_import_map(): - all_files = [ - os.path.join(dp, f) for dp, dn, filenames in os.walk( - os.path.join(os.getcwd(), 'modelscope')) for f in filenames - if os.path.splitext(f)[1] == '.py' - ] - import_map = get_import_map() - - reverse_depend_map = {} - for f in all_files: - depend_by = [] - for k, v in import_map.items(): - if f in v and f != k: - depend_by.append(k) - reverse_depend_map[f] = depend_by - - return reverse_depend_map, import_map - - -def get_all_register_modules(): - all_files = [ - os.path.join(dp, f) for dp, dn, filenames in os.walk( - os.path.join(os.getcwd(), 'modelscope')) for f in filenames - if os.path.splitext(f)[1] == '.py' - ] - all_register_modules = [] - for f in all_files: - all_register_modules.extend(get_file_register_modules(f)) - return all_register_modules - - -if __name__ == '__main__': - pass From d3effb7370e88141f03f28879e11232c2a0e8e82 Mon Sep 17 00:00:00 2001 From: "Xingjun.Wang" Date: Tue, 9 Jun 2026 20:08:33 +0800 Subject: [PATCH 06/11] [Update] Update modelscope-hub requirement (#1733) --- requirements/hub.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/hub.txt b/requirements/hub.txt index 90cf4272..4b52c05d 100644 --- a/requirements/hub.txt +++ b/requirements/hub.txt @@ -1,5 +1,5 @@ filelock -modelscope-hub @ git+https://github.com/modelscope/modelscope_hub.git@main +modelscope-hub>=0.0.7 packaging requests>=2.25 setuptools From 37eb369e9bd8490f1cf895ad62695ff920b3c959 Mon Sep 17 00:00:00 2001 From: dwd Date: Mon, 15 Jun 2026 10:04:01 +0800 Subject: [PATCH 07/11] add metax dockerfile and its requirements for ms-swift 4.2.x (#1734) --- docker/Metax/4.2/Dockerfile.metax | 145 ++++++++++++++++++ docker/Metax/4.2/Dockerfile.with_metax_image | 78 ++++++++++ docker/Metax/4.2/build.sh | 12 ++ docker/Metax/4.2/build_from_metax_image.sh | 9 ++ docker/Metax/4.2/override.txt | 6 + docker/Metax/4.2/requirements_extra.txt | 14 ++ .../Metax/4.2/swift_building_instructions.md | 52 +++++++ 7 files changed, 316 insertions(+) create mode 100644 docker/Metax/4.2/Dockerfile.metax create mode 100644 docker/Metax/4.2/Dockerfile.with_metax_image create mode 100644 docker/Metax/4.2/build.sh create mode 100644 docker/Metax/4.2/build_from_metax_image.sh create mode 100644 docker/Metax/4.2/override.txt create mode 100644 docker/Metax/4.2/requirements_extra.txt create mode 100644 docker/Metax/4.2/swift_building_instructions.md diff --git a/docker/Metax/4.2/Dockerfile.metax b/docker/Metax/4.2/Dockerfile.metax new file mode 100644 index 00000000..fae1d85c --- /dev/null +++ b/docker/Metax/4.2/Dockerfile.metax @@ -0,0 +1,145 @@ +ARG BUILD_BASE_IMAGE=registry.access.redhat.com/ubi9/ubi:9.6 +ARG PYTHON_VERSION=3.12 +ARG UV_EXTRA_INDEX_URL=https://repos.metax-tech.com/r/maca-pypi/simple +ARG UV_TRUSTED_HOST=repos.metax-tech.com + +# may need passing a particular vllm version during build +ARG VLLM_VERSION +ARG MACA_VERSION +ARG CU_BRIDGE_VERSION=${MACA_VERSION} + +#################### BASE BUILD IMAGE #################### +FROM ${BUILD_BASE_IMAGE} AS base +ARG UV_TRUSTED_HOST + +# maca environment variables +ENV MACA_PATH=/opt/maca +ENV MACA_CLANG_PATH=/opt/maca/mxgpu_llvm/bin +ENV CUCC_PATH="${MACA_PATH}/tools/cu-bridge" +ENV CUDA_PATH=/root/cu-bridge/CUDA_DIR +ENV CUCC_CMAKE_ENTRY=2 +ENV PATH="/opt/venv/bin:/root/.local/bin:$PATH" +ENV PATH=/opt/mxdriver/bin:${MACA_PATH}/bin:${MACA_PATH}/mxgpu_llvm/bin:${MACA_PATH}/tools/cu-bridge/tools:${MACA_PATH}/tools/cu-bridge/bin:${PATH} +ENV LD_LIBRARY_PATH=/opt/mxdriver/lib:${MACA_PATH}/lib:${MACA_PATH}/mxgpu_llvm/lib:${MACA_PATH}/ompi/lib:${MACA_PATH}/ucx/lib:${LD_LIBRARY_PATH} + +# uv environment variables +ENV VIRTUAL_ENV=/opt/venv +ENV UV_INDEX_STRATEGY="unsafe-best-match" +ENV UV_HTTP_TIMEOUT=6000 +ENV UV_LINK_MODE=copy +ARG UV_EXTRA_INDEX_URL +ENV UV_EXTRA_INDEX_URL=${UV_EXTRA_INDEX_URL} +ARG UV_INDEX_URL +ENV UV_INDEX_URL=https://mirrors.aliyun.com/pypi/simple +ENV UV_TRUSTED_INDEX_HOST=mirrors.aliyun.com +ENV UV_OVERRIDE=/workspace/override.txt + +# vllm compile option +ENV VLLM_INSTALL_PUNICA_KERNELS=1 + +# AI version arguments +ARG PYTHON_VERSION +ARG VLLM_VERSION +ARG VLLM_METAX_VERSION +ARG MACA_VERSION +ARG MEGATRON_VERSION +ARG SWIFT_VERSION +ARG CU_BRIDGE_VERSION +ARG TE_VERSION + +WORKDIR /workspace +COPY override.txt /workspace/override.txt +COPY requirements_extra.txt /workspace/requirements_extra.txt + +RUN printf "[metax-centos]\n\ +name=Maca Driver Yum Repository\n\ +baseurl=https://repos.metax-tech.com/r/metax-driver-centos-$(uname -m)/\n\ +enabled=1\n\ +gpgcheck=0" > /etc/yum.repos.d/metax-driver-centos.repo && \ + printf "[maca-sdk]\n\ +name=Maca Sdk Yum Repository\n\ +baseurl=https://repos.metax-tech.com/r/maca-sdk-rpm-$(uname -m)/\n\ +enabled=1\n\ +gpgcheck=0" > /etc/yum.repos.d/maca-sdk-rpm.repo + +RUN yum install -y \ + python3-pip hostname \ + unzip vim git openblas-devel make cmake \ + ninja-build gcc g++ procps-ng \ + libibverbs librdmacm libibumad \ + binutils numactl-libs \ + metax-driver-${MACA_VERSION}* mxgvm \ + maca_sdk-${MACA_VERSION}* \ + && yum clean all && rm -rf /var/cache/yum /tmp/* + +RUN python3 -m pip install uv -i $UV_INDEX_URL --trusted-host ${UV_TRUSTED_INDEX_HOST} && \ + uv venv /opt/venv --python=${PYTHON_VERSION} + +RUN python3 --version && \ + uv self version + + +RUN git clone --depth 1 --branch ${SWIFT_VERSION} https://github.com/modelscope/ms-swift.git && \ + git clone --depth 1 --branch ${VLLM_METAX_VERSION} https://github.com/MetaX-MACA/vLLM-metax.git && \ + git clone --depth 1 --branch ${VLLM_VERSION} https://github.com/vllm-project/vllm.git && \ + git clone --depth 1 --branch ${MEGATRON_VERSION} https://github.com/NVIDIA/Megatron-LM.git + +# Step 1: install MACA SDK, Metax-Driver and cu-bridge +# Metax-Driver mainly contains vbios and kmd files, which are not needed in a container. +# Here we keep the mx-smi management tool. Kernel version mismatch errors are ignored. +RUN cd /tmp/ && \ + export MACA_PATH=/opt/maca && \ + curl -o ${CU_BRIDGE_VERSION}.zip -LsSf https://gitee.com/metax-maca/cu-bridge/repository/archive/${CU_BRIDGE_VERSION}.zip && \ + unzip ${CU_BRIDGE_VERSION}.zip && \ + mv cu-bridge-${CU_BRIDGE_VERSION} cu-bridge && \ + chmod -R 755 cu-bridge && \ + cd cu-bridge && \ + mkdir build && cd build && \ + cmake -DCMAKE_INSTALL_PREFIX=/opt/maca/tools/cu-bridge ../ && \ + make && make install + +# Step 2: trim unused MACA packages and install build prerequisites +RUN cd vLLM-metax && \ + uv pip install -r requirements/build.txt && \ + uv pip install build + +# Step 3: install Metax python requirements +RUN cd vLLM-metax && \ + UV_HTTP_TIMEOUT=960 uv pip install -r requirements/maca.txt --trusted-host ${UV_TRUSTED_HOST} + +# Step 4: build vLLM with empty device to avoid CUDA dependency +RUN cd vllm && \ + python3 use_existing_torch.py && \ + uv pip install -r requirements/build/cuda.txt && \ + VLLM_TARGET_DEVICE=empty uv pip install -v . --no-build-isolation + +# Step 5: build vLLM-metax +RUN cd vLLM-metax && \ + uv pip install -r requirements/build.txt && \ + python3 -m build -w -n && \ + uv pip install dist/*.whl + +# Step 6: install Megatron-LM +RUN cd /workspace/Megatron-LM && \ + if [ -f megatron/legacy/fused_kernels/__init__.py ]; then \ + sed -i 's/nvcc/cucc/g' megatron/legacy/fused_kernels/__init__.py; \ + fi && \ + uv pip install . + +# Step 7: install transformer-engine +RUN uv pip install transformer_engine==${TE_VERSION} -i https://repos.metax-tech.com/r/maca-pypi/simple --trusted-host ${UV_TRUSTED_HOST} + +# Step 8: patch and install ms-swift v4.2.0 with Megatron extra dependencies +RUN sed -i '0,/^\(from \|import \)/{s//import vllm_metax.patch\n&/}' ms-swift/swift/__init__.py && \ + cd ms-swift && \ + uv pip install '.[megatron]' + +# Step 9: install optional runtime dependencies used by swift 4.2.0 +RUN uv pip install deepspeed -i https://repos.metax-tech.com/r/maca-pypi/simple --trusted-host ${UV_TRUSTED_HOST} +RUN uv pip install pip +RUN uv pip install -r requirements_extra.txt +RUN ln -sf ${CUDA_PATH}/bin/nvcc ${CUDA_PATH}/bin/cucc + +# vllm installation may bring in incompatible CUDA-only wheels. Remove them here. +RUN uv pip uninstall flashinfer-python cupy-cuda12x +#################### FINAL IMAGE #################### diff --git a/docker/Metax/4.2/Dockerfile.with_metax_image b/docker/Metax/4.2/Dockerfile.with_metax_image new file mode 100644 index 00000000..738befa1 --- /dev/null +++ b/docker/Metax/4.2/Dockerfile.with_metax_image @@ -0,0 +1,78 @@ +ARG BUILD_BASE_IMAGE=mx-devops-acr-cn-shanghai.cr.volces.com/opensource/public-ai-release/maca/sglang:0.5.10-maca.ai3.7.1.12-torch2.8-py312-ubuntu22.04-amd64 +ARG PYTHON_VERSION=3.12 + +FROM ${BUILD_BASE_IMAGE} AS base + +# NOTE: +# This fast-build path inherits Python/Torch/TE from a prebuilt Metax release image. +# We keep the verified base image tag here instead of guessing a newer one. +# As a result, this path may lag behind the Megatron-SWIFT Quick Start recommendations. + +# may need passing a particular vllm version during build +ARG VLLM_VERSION +ARG VLLM_METAX_VERSION +ARG MEGATRON_VERSION +ARG SWIFT_VERSION + +ENV MACA_PATH=/opt/maca +ENV CUCC_CMAKE_ENTRY=2 +ENV CUDA_PATH=/root/cu-bridge/CUDA_DIR +ENV CUCC_PATH=${MACA_PATH}/tools/cu-bridge +ENV PATH=/opt/conda/bin:/opt/conda/condabin:${CUDA_PATH}/bin:${CUCC_PATH}/tools:${CUCC_PATH}/bin:${MACA_PATH}/bin:${PATH} +ENV LD_LIBRARY_PATH=${CUDA_PATH}/lib64:${MACA_PATH}/lib:${MACA_PATH}/mxgpu_llvm/lib:${LD_LIBRARY_PATH} + +WORKDIR /workspace +COPY requirements_extra.txt /workspace/requirements_extra.txt + +RUN apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/* + +# Initialize cu-bridge if it is not already prepared in the base image. +RUN if [ ! -d /root/cu-bridge ]; then \ + ${MACA_PATH}/tools/cu-bridge/tools/pre_make; \ + fi + +# Clone all GitHub sources while the external proxy is enabled. +RUN rm -rf /workspace/ms-swift /workspace/vLLM-metax /workspace/vllm /workspace/Megatron-LM && \ + git clone --depth 1 --branch ${SWIFT_VERSION} https://github.com/modelscope/ms-swift.git && \ + git clone --depth 1 --branch ${VLLM_METAX_VERSION} https://github.com/MetaX-MACA/vLLM-metax.git && \ + git clone --depth 1 --branch ${VLLM_VERSION} https://github.com/vllm-project/vllm.git && \ + git clone --depth 1 --branch ${MEGATRON_VERSION} https://github.com/NVIDIA/Megatron-LM.git + +# install cmake +RUN pip install cmake ninja + +# Step 1: build original vLLM for torch setup +RUN cd vllm && \ + python3 use_existing_torch.py && \ + pip install -r requirements/build/cuda.txt && \ + VLLM_TARGET_DEVICE=empty pip install -v . --no-build-isolation + +# Step 3: build vLLM-metax +RUN cd vLLM-metax && \ + python3 use_existing_metax.py && \ + pip install -r requirements/build.txt && \ + python3 -m build -w -n && \ + pip install dist/*.whl + +# Step 4: patch and install Megatron-LM +RUN cd /workspace/Megatron-LM && \ + if [ -f megatron/legacy/fused_kernels/__init__.py ]; then \ + sed -i 's/nvcc/cucc/g' megatron/legacy/fused_kernels/__init__.py; \ + fi && \ + pip install . + +# Step 5: install transformer-engine +ARG TE_VERSION=2.8.0+4a002bf5.maca3.5.3.105.torch2.8 +ARG UV_TRUSTED_HOST=repos.metax-tech.com +RUN pip install transformer_engine==${TE_VERSION} -i https://repos.metax-tech.com/r/maca-pypi/simple --trusted-host ${UV_TRUSTED_HOST} + +# Step 6: patch and install ms-swift v4.2.0 with its Megatron extra +RUN sed -i '0,/^\(from \|import \)/{s//import vllm_metax.patch\n&/}' ms-swift/swift/__init__.py && \ + cd ms-swift && \ + pip install "transformers<5.4.0" && \ + pip install '.[megatron]' && \ + pip install -r /workspace/requirements_extra.txt + +RUN pip uninstall torchao -y + +CMD ["bash"] diff --git a/docker/Metax/4.2/build.sh b/docker/Metax/4.2/build.sh new file mode 100644 index 00000000..88d88102 --- /dev/null +++ b/docker/Metax/4.2/build.sh @@ -0,0 +1,12 @@ +docker build \ + --network host \ + -f Dockerfile.metax \ + -t swift:v4.2.3 \ + --build-arg VLLM_VERSION=v0.20.0 \ + --build-arg VLLM_METAX_VERSION=v0.20.0 \ + --build-arg MACA_VERSION=3.7.0 \ + --build-arg MEGATRON_VERSION=core_v0.16.0 \ + --build-arg SWIFT_VERSION=v4.2.3 \ + --build-arg TE_VERSION=2.8.0+4a002bf5.maca3.5.3.105.torch2.8 \ + --build-arg CU_BRIDGE_VERSION=3.7.0 \ + . diff --git a/docker/Metax/4.2/build_from_metax_image.sh b/docker/Metax/4.2/build_from_metax_image.sh new file mode 100644 index 00000000..e105df92 --- /dev/null +++ b/docker/Metax/4.2/build_from_metax_image.sh @@ -0,0 +1,9 @@ +docker build \ + --network host \ + -f Dockerfile.with_metax_image \ + -t swift:v4.2.3-with-metax-base \ + --build-arg VLLM_VERSION=v0.20.0 \ + --build-arg VLLM_METAX_VERSION=v0.20.0 \ + --build-arg MEGATRON_VERSION=core_v0.16.0 \ + --build-arg SWIFT_VERSION=v4.2.3 \ + . diff --git a/docker/Metax/4.2/override.txt b/docker/Metax/4.2/override.txt new file mode 100644 index 00000000..d317a962 --- /dev/null +++ b/docker/Metax/4.2/override.txt @@ -0,0 +1,6 @@ +setuptools>=77.0.3,<80 +datasets>=3.0,<4.0 +flash-linear-attention +mcoplib +transformers<5.4.0 +numpy<2 diff --git a/docker/Metax/4.2/requirements_extra.txt b/docker/Metax/4.2/requirements_extra.txt new file mode 100644 index 00000000..3480e1d9 --- /dev/null +++ b/docker/Metax/4.2/requirements_extra.txt @@ -0,0 +1,14 @@ +decord +diffusers==0.35.2 +evalscope>=1.0.0 +evalscope[opencompass] +evalscope[vlmeval] +keye_vl_utils>=1.5.2 +librosa +mpi4py +optimum==1.27.0 +pytorchvideo +qwen_omni_utils>=0.0.9 +qwen_vl_utils==0.0.14 +soundfile +timm diff --git a/docker/Metax/4.2/swift_building_instructions.md b/docker/Metax/4.2/swift_building_instructions.md new file mode 100644 index 00000000..90168f04 --- /dev/null +++ b/docker/Metax/4.2/swift_building_instructions.md @@ -0,0 +1,52 @@ +# 1. Build swift 4.2 image from a UBI9 base image + Full build from a minimal base image, using a venv virtual environment. + +## 1.1. Build + ``` bash + bash build.sh + ``` + +## 1.2. Run a container + ``` bash + docker run -d -it --net=host --uts=host --ipc=host --privileged=true --group-add video \ + --shm-size 100gb --ulimit memlock=-1 \ + --security-opt seccomp=unconfined --security-opt apparmor=unconfined \ + --device=/dev/dri --device=/dev/mxcd \ + --name base_image \ + ${IMAGE_ID} bash + ``` + +## 1.3. Activate the venv environment + ``` bash + source /opt/venv/bin/activate + ``` + +## 1.4. Run swift examples + ``` bash + cd /workspace/ms-swift + bash examples/train/full/train.sh + ``` + +# 2. Build swift 4.2 image from a Metax release image + Faster build based on the pre-built Metax release image. + +## 2.1. Build + ``` bash + bash build_from_metax_image.sh + ``` + +## 2.2. Run a container + ``` bash + docker run -d -it --net=host --uts=host --ipc=host --privileged=true --group-add video \ + --shm-size 100gb --ulimit memlock=-1 \ + --security-opt seccomp=unconfined --security-opt apparmor=unconfined \ + --device=/dev/dri --device=/dev/mxcd \ + --name base_image \ + ${IMAGE_ID} bash + ``` + +## 2.3. Run swift examples + ``` bash + cd /workspace/ms-swift + bash examples/train/full/train.sh + ``` From c7dba8de1e4a1f571aa952f5e0db3410eb6d35c9 Mon Sep 17 00:00:00 2001 From: tastelikefeet <58414341+tastelikefeet@users.noreply.github.com> Date: Mon, 15 Jun 2026 10:35:50 +0800 Subject: [PATCH 08/11] Fix some security issues (#1737) --- modelscope/models/base/base_model.py | 4 +- .../abnormal_object_detection/mmdet_model.py | 3 +- .../annotator/annotator.py | 16 +- .../cv/face_detection/scrfd/scrfd_detect.py | 9 +- .../cv/image_classification/mmcls_model.py | 6 + .../panseg_model.py | 6 + .../semantic_seg_model.py | 6 + .../models/cv/object_detection/mmdet_model.py | 3 +- .../cv/salient_detection/salient_model.py | 3 +- modelscope/pipelines/base.py | 4 +- .../cv/human3d_animation_pipeline.py | 10 + .../cv/object_detection_3d_pipeline.py | 3 +- .../cv/tinynas_classification_pipeline.py | 10 +- .../cv/action_detection_mapper.py | 8 +- .../cv/controllable_image_generation.py | 11 +- .../preprocessors/cv/mmcls_preprocessor.py | 8 + modelscope/utils/config.py | 48 +++- .../test_human3d_animation_security.py | 100 +++++++ .../pipelines/test_remote_config_security.py | 261 ++++++++++++++++++ .../test_tinynas_classification_security.py | 95 +++++++ .../test_action_detection_mapper_security.py | 162 +++++++++++ 21 files changed, 752 insertions(+), 24 deletions(-) create mode 100644 tests/pipelines/test_human3d_animation_security.py create mode 100644 tests/pipelines/test_remote_config_security.py create mode 100644 tests/pipelines/test_tinynas_classification_security.py create mode 100644 tests/preprocessors/test_action_detection_mapper_security.py diff --git a/modelscope/models/base/base_model.py b/modelscope/models/base/base_model.py index 4a1b6494..e2a305cf 100644 --- a/modelscope/models/base/base_model.py +++ b/modelscope/models/base/base_model.py @@ -53,7 +53,9 @@ class Model(ABC): 'import extra libs or execute the code in the model repo, setting this to true ' 'means you trust the files in it.') if not check_model_from_owner_group(model_dir=model_dir): - assert self.trust_remote_code, info_str + # `raise` (not `assert`) so the gate also holds under `python -O`. + if not self.trust_remote_code: + raise RuntimeError(info_str) @abstractmethod def forward(self, *args, **kwargs) -> Dict[str, Any]: diff --git a/modelscope/models/cv/abnormal_object_detection/mmdet_model.py b/modelscope/models/cv/abnormal_object_detection/mmdet_model.py index d341b13a..4626022d 100644 --- a/modelscope/models/cv/abnormal_object_detection/mmdet_model.py +++ b/modelscope/models/cv/abnormal_object_detection/mmdet_model.py @@ -27,7 +27,8 @@ class AbnormalDetectionModel(TorchModel): model_path = osp.join(model_dir, ModelFile.TORCH_MODEL_FILE) config_path = osp.join(model_dir, 'mmcv_config.py') - config = Config.from_file(config_path) + config = Config.from_file( + config_path, trust_remote_code=self.trust_remote_code) config.model.pretrained = None self.model = build_detector( config.model, test_cfg=config.get('test_cfg')) diff --git a/modelscope/models/cv/controllable_image_generation/annotator/annotator.py b/modelscope/models/cv/controllable_image_generation/annotator/annotator.py index 12076080..eff53416 100644 --- a/modelscope/models/cv/controllable_image_generation/annotator/annotator.py +++ b/modelscope/models/cv/controllable_image_generation/annotator/annotator.py @@ -361,14 +361,22 @@ def show_result_pyplot(model, class SegformerDetector: - def __init__(self, annotator_ckpts_path, device='cuda'): + def __init__(self, + annotator_ckpts_path, + device='cuda', + trust_remote_code=False): + from modelscope.utils.config import \ + check_trust_remote_code_for_config modelpath = os.path.join( annotator_ckpts_path, 'segformer_mit-b4_512x512_160k_ade20k_20220620_112216-4fa4f58f.pth' ) - config_file = os.path.join( - annotator_ckpts_path.replace('ckpt/annotator/', ''), - 'config/config.py') + annotator_root = annotator_ckpts_path.replace('ckpt/annotator/', '') + config_file = os.path.join(annotator_root, 'config/config.py') + check_trust_remote_code_for_config( + config_file, + trust_remote_code=trust_remote_code, + model_dir=annotator_root) self.model = init_segmentor(config_file, modelpath).to(device) def __call__(self, img): diff --git a/modelscope/models/cv/face_detection/scrfd/scrfd_detect.py b/modelscope/models/cv/face_detection/scrfd/scrfd_detect.py index e382ec9a..fd45045d 100644 --- a/modelscope/models/cv/face_detection/scrfd/scrfd_detect.py +++ b/modelscope/models/cv/face_detection/scrfd/scrfd_detect.py @@ -39,8 +39,15 @@ class ScrfdDetect(TorchModel): from modelscope.models.cv.face_detection.scrfd.mmdet_patch.models.backbones import ResNetV1e from modelscope.models.cv.face_detection.scrfd.mmdet_patch.models.dense_heads import SCRFDHead from modelscope.models.cv.face_detection.scrfd.mmdet_patch.models.detectors import SCRFD + from modelscope.utils.config import \ + check_trust_remote_code_for_config cfg_file = kwargs.get('config_file', 'mmcv_scrfd.py') - cfg = Config.fromfile(osp.join(model_dir, cfg_file)) + cfg_path = osp.join(model_dir, cfg_file) + check_trust_remote_code_for_config( + cfg_path, + trust_remote_code=self.trust_remote_code, + model_dir=model_dir) + cfg = Config.fromfile(cfg_path) model_file = kwargs.get('model_file', ModelFile.TORCH_MODEL_BIN_FILE) ckpt_path = osp.join(model_dir, model_file) cfg.model.test_cfg.score_thr = kwargs.get('score_thr', 0.3) diff --git a/modelscope/models/cv/image_classification/mmcls_model.py b/modelscope/models/cv/image_classification/mmcls_model.py index bd37d3de..a038551d 100644 --- a/modelscope/models/cv/image_classification/mmcls_model.py +++ b/modelscope/models/cv/image_classification/mmcls_model.py @@ -16,12 +16,18 @@ class ClassificationModel(TorchModel): from mmcls.models import build_classifier import modelscope.models.cv.image_classification.backbones from modelscope.utils.hub import read_config + from modelscope.utils.config import \ + check_trust_remote_code_for_config super().__init__(model_dir) self.config_type = 'ms_config' mm_config = os.path.join(model_dir, 'config.py') if os.path.exists(mm_config): + check_trust_remote_code_for_config( + mm_config, + trust_remote_code=self.trust_remote_code, + model_dir=model_dir) cfg = mmcv.Config.fromfile(mm_config) cfg.model.pretrained = None self.cls_model = build_classifier(cfg.model) diff --git a/modelscope/models/cv/image_panoptic_segmentation/panseg_model.py b/modelscope/models/cv/image_panoptic_segmentation/panseg_model.py index f44c01e8..685d6ec1 100644 --- a/modelscope/models/cv/image_panoptic_segmentation/panseg_model.py +++ b/modelscope/models/cv/image_panoptic_segmentation/panseg_model.py @@ -20,9 +20,15 @@ class SwinLPanopticSegmentation(TorchModel): from mmcv.runner import load_checkpoint import mmcv from mmdet.models import build_detector + from modelscope.utils.config import \ + check_trust_remote_code_for_config config = osp.join(model_dir, 'config.py') + check_trust_remote_code_for_config( + config, + trust_remote_code=self.trust_remote_code, + model_dir=model_dir) cfg = mmcv.Config.fromfile(config) if 'pretrained' in cfg.model: cfg.model.pretrained = None diff --git a/modelscope/models/cv/image_semantic_segmentation/semantic_seg_model.py b/modelscope/models/cv/image_semantic_segmentation/semantic_seg_model.py index 455f29fb..6ff088c6 100644 --- a/modelscope/models/cv/image_semantic_segmentation/semantic_seg_model.py +++ b/modelscope/models/cv/image_semantic_segmentation/semantic_seg_model.py @@ -27,8 +27,14 @@ class SemanticSegmentation(TorchModel): from mmcv.runner import load_checkpoint import mmcv from mmdet.models import build_detector + from modelscope.utils.config import \ + check_trust_remote_code_for_config config = osp.join(model_dir, 'mmcv_config.py') + check_trust_remote_code_for_config( + config, + trust_remote_code=self.trust_remote_code, + model_dir=model_dir) cfg = mmcv.Config.fromfile(config) if 'pretrained' in cfg.model: cfg.model.pretrained = None diff --git a/modelscope/models/cv/object_detection/mmdet_model.py b/modelscope/models/cv/object_detection/mmdet_model.py index 485d440a..26b6f57f 100644 --- a/modelscope/models/cv/object_detection/mmdet_model.py +++ b/modelscope/models/cv/object_detection/mmdet_model.py @@ -31,7 +31,8 @@ class DetectionModel(TorchModel): model_path = osp.join(model_dir, ModelFile.TORCH_MODEL_FILE) config_path = osp.join(model_dir, 'mmcv_config.py') - config = Config.from_file(config_path) + config = Config.from_file( + config_path, trust_remote_code=self.trust_remote_code) config.model.pretrained = None self.model = build_detector(config.model) diff --git a/modelscope/models/cv/salient_detection/salient_model.py b/modelscope/models/cv/salient_detection/salient_model.py index e25166c8..5d4c5646 100644 --- a/modelscope/models/cv/salient_detection/salient_model.py +++ b/modelscope/models/cv/salient_detection/salient_model.py @@ -32,7 +32,8 @@ class SalientDetection(TorchModel): self.model = U2NET(3, 1) else: self.model = SENet(backbone_path=None, pretrained=False) - config = Config.from_file(config_path) + config = Config.from_file( + config_path, trust_remote_code=self.trust_remote_code) self.norm_mean = config.norm_mean self.norm_std = config.norm_std self.norm_size = config.norm_size diff --git a/modelscope/pipelines/base.py b/modelscope/pipelines/base.py index 115a3b95..f7d06efc 100644 --- a/modelscope/pipelines/base.py +++ b/modelscope/pipelines/base.py @@ -151,7 +151,9 @@ class Pipeline(ABC): 'import extra libs or execute the code in the model repo, setting this to true ' 'means you trust the files in it.') if not check_model_from_owner_group(model_dir=model_dir): - assert self.trust_remote_code, info_str + # `raise` (not `assert`) so the gate also holds under `python -O`. + if not self.trust_remote_code: + raise RuntimeError(info_str) def prepare_model(self): """ Place model on certain device for pytorch models before first inference diff --git a/modelscope/pipelines/cv/human3d_animation_pipeline.py b/modelscope/pipelines/cv/human3d_animation_pipeline.py index 4e5ab46d..683e499e 100644 --- a/modelscope/pipelines/cv/human3d_animation_pipeline.py +++ b/modelscope/pipelines/cv/human3d_animation_pipeline.py @@ -72,6 +72,16 @@ class Human3DAnimationPipeline(Pipeline): (case_name, action_name)) exec_path = os.path.join(self.model_dir, 'skinning.py') + # `skinning.py` ships inside the model repo and is executed via Blender; + # gate it behind trust_remote_code to prevent RCE from untrusted repos. + self.check_trust_remote_code( + info_str= + ('Human3DAnimationPipeline executes `skinning.py` from the model ' + 'repository via Blender, which can run arbitrary code. ' + 'Pass `trust_remote_code=True` to pipeline() to opt in if you ' + 'trust the model repository.'), + model_dir=self.model_dir) + cmd = f'{self.blender} -b -P {exec_path} -- --input {self.case_dir}' \ f' --gltf_path {gltf_path} --action {self.action}' os.system(cmd) diff --git a/modelscope/pipelines/cv/object_detection_3d_pipeline.py b/modelscope/pipelines/cv/object_detection_3d_pipeline.py index 640d76f6..fc4eab48 100644 --- a/modelscope/pipelines/cv/object_detection_3d_pipeline.py +++ b/modelscope/pipelines/cv/object_detection_3d_pipeline.py @@ -49,7 +49,8 @@ class ObjectDetection3DPipeline(Pipeline): """ super().__init__(model=model, **kwargs) config_path = osp.join(model, 'mmcv_depe.py') - self.cfg = Config.from_file(config_path) + self.cfg = Config.from_file( + config_path, trust_remote_code=self.trust_remote_code) if torch.cuda.is_available(): self.device = torch.device('cuda') else: diff --git a/modelscope/pipelines/cv/tinynas_classification_pipeline.py b/modelscope/pipelines/cv/tinynas_classification_pipeline.py index ca039eca..8190b63d 100644 --- a/modelscope/pipelines/cv/tinynas_classification_pipeline.py +++ b/modelscope/pipelines/cv/tinynas_classification_pipeline.py @@ -1,5 +1,6 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +import ast import math import os.path as osp from typing import Any, Dict @@ -83,10 +84,11 @@ class TinynasClassificationPipeline(Pipeline): def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: label_mapping_path = osp.join(self.path, 'label_map.txt') - f = open(label_mapping_path, encoding='utf-8') - content = f.read() - f.close() - label_dict = eval(content) + with open(label_mapping_path, encoding='utf-8') as f: + content = f.read() + # `label_map.txt` ships from a remote model repo; restrict parsing to + # plain literal containers so a malicious file cannot trigger RCE (#1668). + label_dict = ast.literal_eval(content) output_prob = torch.nn.functional.softmax(inputs['outputs'], dim=-1) score = torch.max(output_prob) diff --git a/modelscope/preprocessors/cv/action_detection_mapper.py b/modelscope/preprocessors/cv/action_detection_mapper.py index 9bb6d422..669e45dc 100644 --- a/modelscope/preprocessors/cv/action_detection_mapper.py +++ b/modelscope/preprocessors/cv/action_detection_mapper.py @@ -1,5 +1,6 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +import ast import copy import random @@ -84,7 +85,12 @@ class VideoDetMapper: def _call(self, data_dict): video_name = data_dict['path:FILE'] if data_dict['actions'] is not None: - data_dict['actions'] = eval(data_dict['actions']) + actions = data_dict['actions'] + if isinstance(actions, bytes): + actions = actions.decode('utf-8') + if isinstance(actions, str): + actions = ast.literal_eval(actions) + data_dict['actions'] = actions else: data_dict['actions'] = [] diff --git a/modelscope/preprocessors/cv/controllable_image_generation.py b/modelscope/preprocessors/cv/controllable_image_generation.py index 054baca3..6fb77872 100644 --- a/modelscope/preprocessors/cv/controllable_image_generation.py +++ b/modelscope/preprocessors/cv/controllable_image_generation.py @@ -58,7 +58,7 @@ def resize_image(input_image, resolution): return img -def build_detector(control_type, model_path, device): +def build_detector(control_type, model_path, device, trust_remote_code=False): if control_type == 'scribble': detector = None elif control_type == 'canny': @@ -74,7 +74,8 @@ def build_detector(control_type, model_path, device): elif control_type == 'pose': detector = OpenposeDetector(model_path, device) elif control_type == 'seg': - detector = SegformerDetector(model_path, device) + detector = SegformerDetector( + model_path, device, trust_remote_code=trust_remote_code) elif control_type == 'fake_scribble': detector = HEDdetector(model_path, device) else: @@ -136,8 +137,10 @@ class ControllableImageGenerationPreprocessor(Preprocessor): def __init__(self, mode=ModeKeys.INFERENCE, *args, **kwargs): super().__init__(mode=ModeKeys.INFERENCE, *args, **kwargs) self.detector = build_detector( - kwargs.get('control_type', 'hed'), kwargs.get('model_path', None), - kwargs.get('device', 'cuda')) + kwargs.get('control_type', 'hed'), + kwargs.get('model_path', None), + kwargs.get('device', 'cuda'), + trust_remote_code=kwargs.get('trust_remote_code', False)) @type_assert(object, object) def __call__(self, data: input, **kwargs) -> Dict[str, Any]: diff --git a/modelscope/preprocessors/cv/mmcls_preprocessor.py b/modelscope/preprocessors/cv/mmcls_preprocessor.py index 36e7ac4d..b549ee09 100644 --- a/modelscope/preprocessors/cv/mmcls_preprocessor.py +++ b/modelscope/preprocessors/cv/mmcls_preprocessor.py @@ -38,11 +38,19 @@ class ImageClassificationMmcvPreprocessor(Preprocessor): import mmcv from mmcls.datasets.pipelines import Compose from modelscope.models.cv.image_classification.utils import preprocess_transform + from modelscope.utils.config import \ + check_trust_remote_code_for_config + # Preprocessor base does not carry trust_remote_code; read it directly. + trust_remote_code = kwargs.get('trust_remote_code', False) super().__init__(**kwargs) self.config_type = 'ms_config' mm_config = os.path.join(model_dir, 'config.py') if os.path.exists(mm_config): + check_trust_remote_code_for_config( + mm_config, + trust_remote_code=trust_remote_code, + model_dir=model_dir) cfg = mmcv.Config.fromfile(mm_config) cfg.model.pretrained = None config_type = 'mmcv_config' diff --git a/modelscope/utils/config.py b/modelscope/utils/config.py index 099dbf11..2b2920cb 100644 --- a/modelscope/utils/config.py +++ b/modelscope/utils/config.py @@ -28,6 +28,40 @@ DEPRECATION_KEY = '_deprecation_' RESERVED_KEYS = ['filename', 'text', 'pretty_text'] +def check_trust_remote_code_for_config(filename, + trust_remote_code: bool = False, + model_dir=None): + """Refuse to exec a `.py` config file that comes from an untrusted source. + + Loading a Python config (via `Config.from_file`, `mmcv.Config.fromfile`, + `mmseg.apis.init_segmentor`, etc.) imports the file as a module, which + runs any top-level code it contains. Anything that ultimately reads a + `.py` config from a remote model repo MUST gate that load with this + helper. JSON / YAML configs are passive data and pass through. + + Args: + filename: Path to the candidate config file. + trust_remote_code: Caller opt-in flag; pass through + ``self.trust_remote_code`` from ``Model`` / ``Pipeline`` / + ``Preprocessor`` callers. + model_dir: Repo root used by the owner-group check. Defaults to the + parent directory of ``filename``. + """ + if not str(filename).endswith('.py'): + return + from modelscope.utils.automodel_utils import check_model_from_owner_group + if model_dir is None: + model_dir = osp.dirname(osp.abspath(osp.expanduser(str(filename)))) + if check_model_from_owner_group(model_dir=model_dir): + return + if trust_remote_code: + return + raise RuntimeError( + f'Refusing to load Python config "{filename}": doing so would execute ' + 'code from the model repository. Pass `trust_remote_code=True` to opt ' + 'in if you trust the source.') + + class ConfigDict(addict.Dict): """ Dict which support get value through getattr @@ -82,7 +116,7 @@ class Config: """ @staticmethod - def _file2dict(filename): + def _file2dict(filename, trust_remote_code: bool = False, model_dir=None): filename = osp.abspath(osp.expanduser(filename)) if not osp.exists(filename): raise ValueError(f'File does not exists {filename}') @@ -90,6 +124,9 @@ class Config: if fileExtname not in ['.py', '.json', '.yaml', '.yml']: raise IOError('Only py/yml/yaml/json type are supported now!') + check_trust_remote_code_for_config( + filename, trust_remote_code=trust_remote_code, model_dir=model_dir) + with tempfile.TemporaryDirectory() as tmp_cfg_dir: tmp_cfg_file = tempfile.NamedTemporaryFile( dir=tmp_cfg_dir, suffix=fileExtname) @@ -126,10 +163,11 @@ class Config: return cfg_dict, cfg_text @staticmethod - def from_file(filename): + def from_file(filename, trust_remote_code: bool = False, model_dir=None): if isinstance(filename, Path): filename = str(filename) - cfg_dict, cfg_text = Config._file2dict(filename) + cfg_dict, cfg_text = Config._file2dict( + filename, trust_remote_code=trust_remote_code, model_dir=model_dir) return Config(cfg_dict, cfg_text=cfg_text, filename=filename) @staticmethod @@ -156,7 +194,9 @@ class Config: temp_file.write(cfg_str) # on windows, previous implementation cause error # see PR 1077 for details - cfg = Config.from_file(temp_file.name) + # `from_string` materializes a caller-provided in-process string into + # a tempfile; the threat model for `trust_remote_code` does not apply. + cfg = Config.from_file(temp_file.name, trust_remote_code=True) os.remove(temp_file.name) return cfg diff --git a/tests/pipelines/test_human3d_animation_security.py b/tests/pipelines/test_human3d_animation_security.py new file mode 100644 index 00000000..ca403a66 --- /dev/null +++ b/tests/pipelines/test_human3d_animation_security.py @@ -0,0 +1,100 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +"""Regression test for Issue #1673. + +Verifies that Human3DAnimationPipeline.gen_weights() refuses to execute the +remote `skinning.py` script unless either (a) the model is owned by a trusted +group (`iic` / `damo`), or (b) the user explicitly opted in with +`trust_remote_code=True`. +""" +import os +import shutil +import sys +import tempfile +import types +import unittest +from unittest import mock + +from modelscope.pipelines.cv.human3d_animation_pipeline import \ + Human3DAnimationPipeline # noqa: E402 + +# `modelscope.models.cv.human3d_animation` lazily resolves to symbols that +# pull in heavy optional deps (e.g. `nvdiffrast`). We never invoke those +# symbols inside `gen_weights`, so substitute a lightweight stub *before* the +# pipeline module is imported. Only stub if the real package cannot be +# resolved, so we don't shadow it for environments that do have the deps. +_PKG = 'modelscope.models.cv.human3d_animation' +if _PKG not in sys.modules or not hasattr(sys.modules[_PKG], + 'gen_skeleton_bvh'): + try: + from modelscope.models.cv.human3d_animation import ( # noqa: F401 + gen_skeleton_bvh, read_obj, write_obj, + ) + except Exception: + _stub = types.ModuleType(_PKG) + _stub.gen_skeleton_bvh = lambda *a, **kw: None + _stub.read_obj = lambda *a, **kw: None + _stub.write_obj = lambda *a, **kw: None + sys.modules[_PKG] = _stub + + +class Human3DAnimationSecurityTest(unittest.TestCase): + """Pin the trust_remote_code gate around remote `skinning.py` execution.""" + + def setUp(self): + self.tmp_root = tempfile.mkdtemp(prefix='ms_human3d_security_') + # `check_model_from_owner_group` extracts the owner from the parent + # directory name, so model dirs must live two levels deep. + self.untrusted_dir = os.path.join(self.tmp_root, 'attacker', 'badrepo') + self.trusted_dir = os.path.join(self.tmp_root, 'damo', 'goodrepo') + for d in (self.untrusted_dir, self.trusted_dir): + os.makedirs(d) + with open(os.path.join(d, 'skinning.py'), 'w') as f: + f.write("raise SystemExit('this script must not run')\n") + + self.case_dir = os.path.join(self.tmp_root, 'case') + os.makedirs(self.case_dir) + self.save_dir = os.path.join(self.tmp_root, 'out') + + def tearDown(self): + shutil.rmtree(self.tmp_root, ignore_errors=True) + + def _make_pipeline(self, model_dir, trust_remote_code): + # Bypass Pipeline.__init__ which would try to download a real model. + p = Human3DAnimationPipeline.__new__(Human3DAnimationPipeline) + p.model_dir = model_dir + p.trust_remote_code = trust_remote_code + p.case_dir = self.case_dir + p.action = 'SwingDancing' + p.blender = 'blender' + return p + + def test_untrusted_repo_without_optin_is_blocked(self): + p = self._make_pipeline(self.untrusted_dir, trust_remote_code=False) + with mock.patch( + 'modelscope.pipelines.cv.human3d_animation_pipeline.os.system' + ) as m_system: + with self.assertRaises(RuntimeError): + p.gen_weights(save_dir=self.save_dir) + m_system.assert_not_called() + + def test_untrusted_repo_with_optin_is_allowed(self): + p = self._make_pipeline(self.untrusted_dir, trust_remote_code=True) + with mock.patch( + 'modelscope.pipelines.cv.human3d_animation_pipeline.os.system', + return_value=0, + ) as m_system: + p.gen_weights(save_dir=self.save_dir) + m_system.assert_called_once() + + def test_trusted_owner_is_allowed_without_optin(self): + p = self._make_pipeline(self.trusted_dir, trust_remote_code=False) + with mock.patch( + 'modelscope.pipelines.cv.human3d_animation_pipeline.os.system', + return_value=0, + ) as m_system: + p.gen_weights(save_dir=self.save_dir) + m_system.assert_called_once() + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/pipelines/test_remote_config_security.py b/tests/pipelines/test_remote_config_security.py new file mode 100644 index 00000000..1ef988a3 --- /dev/null +++ b/tests/pipelines/test_remote_config_security.py @@ -0,0 +1,261 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +"""Regression test for Issue #1672. + +Pins the `trust_remote_code` gate that prevents `Config.from_file` and +`check_trust_remote_code_for_config` from executing a `.py` config from an +untrusted model repository. + +Coverage matrix +--------------- +The helper / choke point exercised here is the only mechanism guarding every +sink listed in the issue: + + * via `Config.from_file('.py')`: + - `SalientDetection` (models/cv/salient_detection) + - `AbnormalDetectionModel` (models/cv/abnormal_object_detection) + - `DetectionModel` (models/cv/object_detection) + - `ObjectDetection3DPipeline` (pipelines/cv/object_detection_3d_pipeline) + + * via `check_trust_remote_code_for_config(...)` invoked before + `mmcv.Config.fromfile` / `mmseg.apis.init_segmentor`: + - `ClassificationModel` (image_classification/mmcls_model) + - `SwinLPanopticSegmentation` (image_panoptic_segmentation) + - `SemanticSegmentation` (image_semantic_segmentation) + - `ScrfdDetect` / `TinyMogDetect` / `DamoFdDetect` (face_detection/scrfd) + - `ImageClassificationMmcvPreprocessor` (preprocessors/cv) + - `SegformerDetector` (controllable_image_generation) +""" +import os +import shutil +import sys +import tempfile +import unittest +from unittest import mock + +from modelscope.utils.config import Config, check_trust_remote_code_for_config + + +def _make_repo(root, owner, name): + """`check_model_from_owner_group` reads the owner from the parent dir.""" + repo = os.path.join(root, owner, name) + os.makedirs(repo) + return repo + + +def _write(path, content): + with open(path, 'w') as f: + f.write(content) + + +# Canary: any caller that fails the gate must NEVER execute this body. +_CANARY = ('raise SystemExit("this script must not run")\n' + 'CANARY = "should not be reachable"\n') + + +class ConfigChokePointSecurityTest(unittest.TestCase): + """End-to-end gating at `Config.from_file` for `.py` configs. + + Covers SalientDetection / AbnormalDetectionModel / DetectionModel / + ObjectDetection3DPipeline, all of which call `Config.from_file(*.py, + trust_remote_code=self.trust_remote_code)`. + """ + + def setUp(self): + self.tmp_root = tempfile.mkdtemp(prefix='ms_cfg_security_') + self.untrusted_dir = _make_repo(self.tmp_root, 'attacker', 'badrepo') + self.trusted_dir = _make_repo(self.tmp_root, 'damo', 'goodrepo') + + self.untrusted_py = os.path.join(self.untrusted_dir, 'mmcv_config.py') + self.trusted_py = os.path.join(self.trusted_dir, 'mmcv_config.py') + for p in (self.untrusted_py, self.trusted_py): + _write(p, _CANARY) + + self.untrusted_json = os.path.join(self.untrusted_dir, + 'configuration.json') + _write(self.untrusted_json, '{"a": 1}') + + def tearDown(self): + shutil.rmtree(self.tmp_root, ignore_errors=True) + + def test_untrusted_py_without_optin_is_blocked(self): + with self.assertRaises(RuntimeError): + Config.from_file(self.untrusted_py) + + def test_untrusted_py_with_optin_is_allowed(self): + # opt-in must let the load proceed, which means the canary will run + # and raise SystemExit -- proving the gate did not block. + with self.assertRaises(SystemExit): + Config.from_file(self.untrusted_py, trust_remote_code=True) + + def test_trusted_owner_py_is_allowed_without_optin(self): + with self.assertRaises(SystemExit): + Config.from_file(self.trusted_py) + + def test_json_passive_load_unaffected(self): + cfg = Config.from_file(self.untrusted_json) + self.assertEqual(cfg.a, 1) + + def test_from_string_in_process_is_not_gated(self): + # `Config.from_string` materializes a caller-supplied in-process + # string and must keep working without `trust_remote_code`. + cfg = Config.from_string('a = 42\n', '.py') + self.assertEqual(cfg.a, 42) + + def test_explicit_model_dir_overrides_default(self): + # Some callers pass `model_dir` explicitly; verify trust derives + # from that path, not from the file's parent directory. + nested = os.path.join(self.untrusted_dir, 'subdir') + os.makedirs(nested) + nested_py = os.path.join(nested, 'mmcv_config.py') + _write(nested_py, _CANARY) + with self.assertRaises(RuntimeError): + Config.from_file(nested_py, model_dir=self.untrusted_dir) + # Same file, but with a trusted model_dir override -> allowed. + with self.assertRaises(SystemExit): + Config.from_file(nested_py, model_dir=self.trusted_dir) + + +class CheckTrustRemoteCodeForConfigTest(unittest.TestCase): + """Unit tests for the helper invoked at every mmcv / init_segmentor sink. + + Covers ClassificationModel, SwinLPanopticSegmentation, SemanticSegmentation, + ScrfdDetect (and subclasses TinyMogDetect / DamoFdDetect), + ImageClassificationMmcvPreprocessor, and SegformerDetector. + """ + + def setUp(self): + self.tmp_root = tempfile.mkdtemp(prefix='ms_cfg_security_helper_') + self.untrusted_dir = _make_repo(self.tmp_root, 'attacker', 'badrepo') + self.trusted_dir = _make_repo(self.tmp_root, 'damo', 'goodrepo') + # Helper only inspects file extensions; no need to create real files. + self.untrusted_py = os.path.join(self.untrusted_dir, 'config.py') + self.trusted_py = os.path.join(self.trusted_dir, 'config.py') + self.untrusted_json = os.path.join(self.untrusted_dir, + 'configuration.json') + + def tearDown(self): + shutil.rmtree(self.tmp_root, ignore_errors=True) + + def test_non_py_path_always_passes(self): + # JSON / YAML configs are passive data; helper must not block them. + check_trust_remote_code_for_config( + self.untrusted_json, + trust_remote_code=False, + model_dir=self.untrusted_dir) + + def test_untrusted_py_without_optin_raises(self): + with self.assertRaises(RuntimeError): + check_trust_remote_code_for_config( + self.untrusted_py, + trust_remote_code=False, + model_dir=self.untrusted_dir) + + def test_untrusted_py_with_optin_passes(self): + check_trust_remote_code_for_config( + self.untrusted_py, + trust_remote_code=True, + model_dir=self.untrusted_dir) + + def test_trusted_owner_py_passes(self): + check_trust_remote_code_for_config( + self.trusted_py, + trust_remote_code=False, + model_dir=self.trusted_dir) + + def test_default_model_dir_inferred_from_filename(self): + # When `model_dir` is omitted, helper must infer it from the file's + # parent directory and still gate untrusted owners correctly. + with self.assertRaises(RuntimeError): + check_trust_remote_code_for_config( + self.untrusted_py, trust_remote_code=False) + # Trusted parent should pass. + check_trust_remote_code_for_config( + self.trusted_py, trust_remote_code=False) + + def test_helper_uses_raise_not_assert(self): + # Security gates must survive `python -O`. The helper raises + # RuntimeError, not AssertionError; verify the failure type so the + # contract is locked in. + try: + check_trust_remote_code_for_config( + self.untrusted_py, + trust_remote_code=False, + model_dir=self.untrusted_dir) + except RuntimeError: + return + except AssertionError: + self.fail( + 'Gate used `assert`; would be a no-op under `python -O`.') + self.fail('Gate did not raise on untrusted `.py` config.') + + +class SinkWiringTest(unittest.TestCase): + """Spot-check that a representative mmcv sink calls the helper before + reaching `mmcv.Config.fromfile`. + + We patch the helper to raise a sentinel; if the sink is wired correctly, + the sentinel propagates and the unsafe `mmcv.Config.fromfile` is never + invoked. + """ + + def setUp(self): + self.tmp_root = tempfile.mkdtemp(prefix='ms_cfg_sink_wiring_') + self.untrusted_dir = _make_repo(self.tmp_root, 'attacker', 'badrepo') + + def tearDown(self): + shutil.rmtree(self.tmp_root, ignore_errors=True) + + def _import_module_text(self, dotted): + # Read the source of a module without importing it (avoids dragging + # in mmcv / mmdet / torch on test machines that lack them). + import importlib.util + spec = importlib.util.find_spec(dotted) + self.assertIsNotNone(spec, f'{dotted} not found') + with open(spec.origin, 'r') as f: + return f.read() + + def _assert_helper_precedes_sink(self, dotted, sink_call): + src = self._import_module_text(dotted) + helper_pos = src.find('check_trust_remote_code_for_config(') + sink_pos = src.find(sink_call) + self.assertGreater( + helper_pos, -1, + f'{dotted} does not call check_trust_remote_code_for_config') + self.assertGreater(sink_pos, -1, + f'{dotted} no longer contains `{sink_call}`') + self.assertLess(helper_pos, sink_pos, + f'{dotted}: helper call must precede `{sink_call}`') + + def test_classification_model_wires_helper(self): + self._assert_helper_precedes_sink( + 'modelscope.models.cv.image_classification.mmcls_model', + 'mmcv.Config.fromfile(') + + def test_panoptic_segmentation_wires_helper(self): + self._assert_helper_precedes_sink( + 'modelscope.models.cv.image_panoptic_segmentation.panseg_model', + 'mmcv.Config.fromfile(') + + def test_semantic_segmentation_wires_helper(self): + self._assert_helper_precedes_sink( + 'modelscope.models.cv.image_semantic_segmentation.semantic_seg_model', + 'mmcv.Config.fromfile(') + + def test_scrfd_detect_wires_helper(self): + self._assert_helper_precedes_sink( + 'modelscope.models.cv.face_detection.scrfd.scrfd_detect', + 'Config.fromfile(') + + def test_mmcls_preprocessor_wires_helper(self): + self._assert_helper_precedes_sink( + 'modelscope.preprocessors.cv.mmcls_preprocessor', + 'mmcv.Config.fromfile(') + + def test_segformer_detector_wires_helper(self): + self._assert_helper_precedes_sink( + 'modelscope.models.cv.controllable_image_generation.annotator.annotator', + 'init_segmentor(') + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/pipelines/test_tinynas_classification_security.py b/tests/pipelines/test_tinynas_classification_security.py new file mode 100644 index 00000000..238aff06 --- /dev/null +++ b/tests/pipelines/test_tinynas_classification_security.py @@ -0,0 +1,95 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +"""Regression test for Issue #1668. + +`TinynasClassificationPipeline.postprocess` previously called `eval()` on +the raw contents of `label_map.txt` shipped from the remote model repo, +giving attackers RCE the moment the pipeline produced a prediction. The fix +swaps `eval` for `ast.literal_eval`. +""" +import os +import shutil +import sys +import tempfile +import types +import unittest + +import torch # noqa: E402 + +from modelscope.pipelines.cv.tinynas_classification_pipeline import \ + TinynasClassificationPipeline # noqa: E402 + +# `tinynas_classification_pipeline` imports torch/torchvision and our own +# `tinynas_classfication` module at top level. Stub the latter only when +# missing so we don't drag heavy custom CUDA ops into the test. +if 'modelscope.models.cv.tinynas_classfication' not in sys.modules: + try: + from modelscope.models.cv.tinynas_classfication import \ + get_zennet # noqa: F401 + except Exception: + stub = types.ModuleType('modelscope.models.cv.tinynas_classfication') + stub.get_zennet = lambda *a, **kw: None + sys.modules['modelscope.models.cv.tinynas_classfication'] = stub + + +class TinynasLabelMapSecurityTest(unittest.TestCase): + """Pin the `ast.literal_eval` gate around remote `label_map.txt`.""" + + def setUp(self): + self.tmp_root = tempfile.mkdtemp(prefix='ms_tinynas_security_') + self.label_path = os.path.join(self.tmp_root, 'label_map.txt') + # Bypass __init__ — it would try to load a real checkpoint. + self.pipe = TinynasClassificationPipeline.__new__( + TinynasClassificationPipeline) + self.pipe.path = self.tmp_root + + def tearDown(self): + shutil.rmtree(self.tmp_root, ignore_errors=True) + + def _write(self, content): + with open(self.label_path, 'w', encoding='utf-8') as f: + f.write(content) + + def _fake_inputs(self, argmax_idx): + # `outputs` must support `.argmax().item()` and softmax. + outputs = torch.zeros(1, 3) + outputs[0, argmax_idx] = 10.0 + return {'outputs': outputs} + + def test_legitimate_label_map_parses(self): + self._write("{0: 'cat', 1: 'dog', 2: 'wolf'}") + out = self.pipe.postprocess(self._fake_inputs(2)) + self.assertEqual(out['labels'], ['wolf']) + self.assertEqual(len(out['scores']), 1) + + def test_malicious_import_payload_blocked(self): + canary = '/tmp/ms_tinynas_rce_canary' + if os.path.exists(canary): + os.remove(canary) + self._write( + "__import__('os').system('echo HACKED && touch {}')".format( + canary)) + with self.assertRaises((ValueError, SyntaxError)): + self.pipe.postprocess(self._fake_inputs(0)) + self.assertFalse( + os.path.exists(canary), + 'os.system was reached -- eval gate is bypassed.') + + def test_malicious_function_call_blocked(self): + self._write("print('pwned')") + with self.assertRaises((ValueError, SyntaxError)): + self.pipe.postprocess(self._fake_inputs(0)) + + def test_malicious_subclass_escape_blocked(self): + self._write('().__class__.__bases__[0].__subclasses__()') + with self.assertRaises((ValueError, SyntaxError)): + self.pipe.postprocess(self._fake_inputs(0)) + + def test_malformed_label_map_raises_cleanly(self): + # Empty file is not a literal container; must not silently no-op. + self._write('') + with self.assertRaises((ValueError, SyntaxError)): + self.pipe.postprocess(self._fake_inputs(0)) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/preprocessors/test_action_detection_mapper_security.py b/tests/preprocessors/test_action_detection_mapper_security.py new file mode 100644 index 00000000..77aa4753 --- /dev/null +++ b/tests/preprocessors/test_action_detection_mapper_security.py @@ -0,0 +1,162 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +"""Regression test for Issue #1667. + +`VideoDetMapper._call` previously fed the dataset's `actions` field straight +to `eval()`, allowing a malicious remote dataset to inject arbitrary Python +expressions (e.g. `__import__('os').system(...)`) and gain RCE during +training. The fix swaps `eval` for `ast.literal_eval`, which only parses +plain literal containers. +""" +import os +import sys +import types +import unittest + +from modelscope.preprocessors.cv.action_detection_mapper import \ + VideoDetMapper # noqa: E402 + + +# `action_detection_mapper` imports `decord`, `detectron2`, `scipy.interpolate` +# at module level. Stub the heavy ones so the test can run on machines that do +# not have them installed; only stub when the real package is missing. +def _stub_module(name, attrs=None): + if name in sys.modules: + return + mod = types.ModuleType(name) + for k, v in (attrs or {}).items(): + setattr(mod, k, v) + sys.modules[name] = mod + + +try: + import decord # noqa: F401 +except Exception: + _stub_module('decord', { + 'cpu': lambda *a, **kw: None, + 'VideoReader': object + }) + +try: + import detectron2 # noqa: F401 + import detectron2.data.transforms # noqa: F401 + import detectron2.structures # noqa: F401 +except Exception: + _stub_module('detectron2') + _stub_module( + 'detectron2.data', + {'transforms': types.ModuleType('detectron2.data.transforms')}) + _stub_module( + 'detectron2.data.transforms', { + 'ExtentTransform': object, + 'RandomBrightness': object, + 'RandomFlip': object, + 'ResizeShortestEdge': object, + }) + _stub_module('detectron2.structures', { + 'Boxes': object, + 'Instances': object + }) + +try: + import scipy.interpolate # noqa: F401 +except Exception: + _stub_module('scipy') + _stub_module('scipy.interpolate', {'interp1d': lambda *a, **kw: None}) + + +class ActionDetectionMapperSecurityTest(unittest.TestCase): + """Pin the `ast.literal_eval` gate around remote `actions` payloads.""" + + def setUp(self): + # Bypass __init__ which constructs heavy detectron2 transforms. + self.mapper = VideoDetMapper.__new__(VideoDetMapper) + + def _parse(self, actions_value): + # Exercise only the literal-eval branch of `_call` without invoking + # the rest of the heavy pipeline. + data_dict = {'path:FILE': 'dummy.mp4', 'actions': actions_value} + if data_dict['actions'] is not None: + actions = data_dict['actions'] + if isinstance(actions, bytes): + actions = actions.decode('utf-8') + if isinstance(actions, str): + import ast + actions = ast.literal_eval(actions) + data_dict['actions'] = actions + else: + data_dict['actions'] = [] + return data_dict['actions'] + + def test_legitimate_python_repr_payload_parses(self): + legit = ("[{'start': 0, 'end': 30, 'label': 'walk'," + " 'boxes': {'0': [1, 2, 3, 4]}}]") + result = self._parse(legit) + self.assertEqual(result[0]['label'], 'walk') + self.assertEqual(result[0]['boxes']['0'], [1, 2, 3, 4]) + + def test_legitimate_json_payload_parses(self): + legit = ('[{"start": 0, "end": 30, "label": "walk",' + ' "boxes": {"0": [1, 2, 3, 4]}}]') + result = self._parse(legit) + self.assertEqual(result[0]['label'], 'walk') + + def test_none_payload_becomes_empty_list(self): + self.assertEqual(self._parse(None), []) + + def test_already_parsed_list_passes_through(self): + # If a downstream caller pre-parses the JSON, we should not try to + # literal_eval a list (which would raise TypeError). + already = [{ + 'start': 0, + 'end': 1, + 'label': 'x', + 'boxes': { + '0': [0, 0, 1, 1] + } + }] + self.assertEqual(self._parse(already), already) + + def test_bytes_payload_is_decoded_then_parsed(self): + legit = ("[{'start': 0, 'end': 1, 'label': 'walk'," + " 'boxes': {'0': [1, 2, 3, 4]}}]").encode('utf-8') + result = self._parse(legit) + self.assertEqual(result[0]['label'], 'walk') + + def test_malicious_import_payload_blocked(self): + # Canary file: literal_eval must NEVER reach `os.system`. + canary = '/tmp/ms_action_det_rce_canary' + if os.path.exists(canary): + os.remove(canary) + mal = ("__import__('os').system(" + "'echo HACKED && touch {}')").format(canary) + with self.assertRaises((ValueError, SyntaxError)): + self._parse(mal) + self.assertFalse( + os.path.exists(canary), + 'os.system was reached -- eval gate is bypassed.') + + def test_malicious_function_call_blocked(self): + # Bare function call: `print('pwned')` is not a literal. + with self.assertRaises((ValueError, SyntaxError)): + self._parse("print('pwned')") + + def test_malicious_attribute_access_blocked(self): + with self.assertRaises((ValueError, SyntaxError)): + self._parse('().__class__.__bases__[0].__subclasses__()') + + def test_mapper_call_swallows_malicious_payload(self): + # The public `__call__` wraps `_call` in try/except and returns None + # on failure. Ensure malicious payloads land there cleanly rather + # than executing. + canary = '/tmp/ms_action_det_rce_canary_call' + if os.path.exists(canary): + os.remove(canary) + mal = ("__import__('os').system(" "'touch {}')").format(canary) + # Use the real public entry point to confirm end-to-end behaviour. + out = self.mapper.__call__({'path:FILE': 'dummy.mp4', 'actions': mal}) + self.assertIsNone(out) + self.assertFalse(os.path.exists(canary)) + + +if __name__ == '__main__': + unittest.main() From c1d88e6ab722c79efef6b3d949f6404c971874a5 Mon Sep 17 00:00:00 2001 From: addsubmuldiv Date: Mon, 15 Jun 2026 11:07:56 +0800 Subject: [PATCH 09/11] Update npu dockerfile (#1736) --- docker/Dockerfile.ascend | 42 ++++++++-- docker/OVERVIEW.ascend.md | 156 +++++++++++++++++++++++++++++++++++ docker/OVERVIEW.ascend.zh.md | 156 +++++++++++++++++++++++++++++++++++ docker/build_image.py | 41 ++++++++- 4 files changed, 389 insertions(+), 6 deletions(-) create mode 100644 docker/OVERVIEW.ascend.md create mode 100644 docker/OVERVIEW.ascend.zh.md diff --git a/docker/Dockerfile.ascend b/docker/Dockerfile.ascend index 911ee3f9..a9524f8f 100644 --- a/docker/Dockerfile.ascend +++ b/docker/Dockerfile.ascend @@ -3,7 +3,8 @@ FROM {base_image} ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \ PIP_DEFAULT_TIMEOUT=300 \ PIP_RETRIES=10 \ - SOC_VERSION={soc_version} + SOC_VERSION={soc_version} \ + CANN_VERSION={cann_version} SHELL ["/bin/bash", "-c"] @@ -36,15 +37,15 @@ RUN ARCH=$(uname -m) && \ source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ source /usr/local/Ascend/nnal/atb/set_env.sh && \ # Install torch & torch_npu & torchvision - pip install torch==2.9.0 torch_npu==2.9.0 torchvision==0.24.0 && \ + pip install torch==2.9.0 torch_npu==2.9.0.post2 torchvision==0.24.0 && \ # Install vllm cd vllm && VLLM_TARGET_DEVICE=empty pip install -v -e . && cd .. && \ # Install vllm-ascend cd vllm-ascend && pip install -v -e . && cd .. # ---------- Clone training-side repositories ---------- -RUN git clone --depth 1 --branch v0.15.3 https://github.com/NVIDIA/Megatron-LM.git /Megatron-LM && \ - git clone --depth 1 --branch core_r0.15.3 https://gitcode.com/Ascend/MindSpeed.git /MindSpeed && \ +RUN git clone --depth 1 --branch {megatron_branch} https://github.com/NVIDIA/Megatron-LM.git /Megatron-LM && \ + git clone --depth 1 --branch {mindspeed_branch} https://gitcode.com/Ascend/MindSpeed.git /MindSpeed && \ GIT_LFS_SKIP_SMUDGE=1 git clone --depth 1 -b {swift_branch} --single-branch https://github.com/modelscope/ms-swift.git /ms-swift && \ git clone --depth 1 https://github.com/modelscope/mcore-bridge.git /mcore-bridge @@ -70,7 +71,7 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ torch==2.9.0 torchvision==0.24.0 torchaudio==2.9.0; \ fi && \ pip install --no-cache-dir --force-reinstall --no-deps \ - torch_npu==2.9.0 && \ + torch_npu==2.9.0.post2 && \ rm -rf /root/.cache/pip # ---------- Remove CUDA-only dependencies pulled in by vllm (they cause missing libtorch_cuda.so errors on NPU) ---------- @@ -122,6 +123,37 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ pip install --no-cache-dir omegaconf==2.3.0 && \ pip cache purge +# ---------- Reinstall triton-ascend for the selected CANN version ---------- +RUN set -eux; \ + pip uninstall -y triton || true; \ + pip uninstall -y triton-ascend || true; \ + case "${CANN_VERSION}" in \ + 8.5.*) \ + pip install --no-cache-dir --force-reinstall triton-ascend==3.2.0; \ + ;; \ + 9.0.0) \ + PY_ABI="cp$(python -c 'import sys; print(f"{sys.version_info.major}{sys.version_info.minor}")')"; \ + case "${PY_ABI}" in \ + cp310|cp311|cp312|cp313) ;; \ + *) echo "Unsupported Python ABI for triton-ascend 3.2.1: ${PY_ABI}" >&2; exit 1 ;; \ + esac; \ + ARCH="$(uname -m)"; \ + case "${ARCH}" in \ + aarch64|x86_64) ;; \ + *) echo "Unsupported architecture for triton-ascend 3.2.1: ${ARCH}" >&2; exit 1 ;; \ + esac; \ + WHEEL_NAME="triton_ascend-3.2.1-${PY_ABI}-${PY_ABI}-manylinux_2_27_${ARCH}.manylinux_2_28_${ARCH}.whl"; \ + WHEEL_PATH="/tmp/${WHEEL_NAME}"; \ + curl -fL "https://gitcode.com/Ascend/triton-ascend/releases/download/v3.2.1/${WHEEL_NAME}" -o "${WHEEL_PATH}"; \ + pip install --no-cache-dir --force-reinstall "${WHEEL_PATH}"; \ + rm -f "${WHEEL_PATH}"; \ + ;; \ + *) \ + echo "Unsupported CANN_VERSION for triton-ascend install: ${CANN_VERSION}" >&2; \ + exit 1; \ + ;; \ + esac + RUN echo 'source /usr/local/Ascend/ascend-toolkit/set_env.sh' >> /root/.bashrc && \ echo '[ -f /usr/local/Ascend/nnal/atb/set_env.sh ] && source /usr/local/Ascend/nnal/atb/set_env.sh' >> /root/.bashrc && \ echo 'set +H' >> /root/.bashrc diff --git a/docker/OVERVIEW.ascend.md b/docker/OVERVIEW.ascend.md new file mode 100644 index 00000000..eb33c451 --- /dev/null +++ b/docker/OVERVIEW.ascend.md @@ -0,0 +1,156 @@ +# ms-swift Ascend + +> English | [中文](./OVERVIEW.ascend.zh.md) + +ms-swift Ascend images provide a ready-to-use ms-swift environment for Huawei Ascend Atlas NPUs. The images are built on top of the Ascend CANN container images and include the Python, CANN, PyTorch NPU, vLLM Ascend, Megatron, MindSpeed, mcore-bridge, ms-swift, and ModelScope runtime components needed for Ascend inference and training workflows. + +## Quick Reference + +- Base image: `quay.io/ascend/cann:---py` +- Build template: `docker/Dockerfile.ascend` +- Build entrypoint: `docker/build_image.py --image_type ascend` +- Default base image: `quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11` +- Default output tag: `${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-` +- Ascend runtime environment is sourced from `/usr/local/Ascend/ascend-toolkit/set_env.sh` +- If available, NNAL/ATB runtime is sourced from `/usr/local/Ascend/nnal/atb/set_env.sh` + +## Image Contents + +The Ascend Dockerfile installs and configures: + +| Component | Version / Source | +| --- | --- | +| CANN | inherited from the selected `quay.io/ascend/cann` base image | +| Python | inherited from the base image tag, for example `py3.11` | +| PyTorch | `torch==2.9.0` | +| torch-npu | `torch_npu==2.9.0.post2` | +| torchvision / torchaudio | `torchvision==0.24.0`, `torchaudio==2.9.0` | +| vLLM | source install from `vllm-project/vllm`, default branch `v0.18.0` | +| vLLM Ascend | source install from `vllm-project/vllm-ascend`, default branch `v0.18.0` | +| Megatron-LM | source checkout, default branch `v0.15.3` | +| MindSpeed | source checkout, default branch `core_r0.15.3` | +| mcore-bridge | source checkout from `modelscope/mcore-bridge` | +| ms-swift | source checkout from `modelscope/ms-swift`, default branch `main` | +| ModelScope | source checkout from `modelscope/modelscope`, default branch `master` | +| triton-ascend | `3.2.0` for CANN `8.5.*`; local wheel install of `3.2.1` for CANN `9.0.0` | + +## Supported Tag Format + +Images built by `docker/build_image.py --image_type ascend` use this tag format: + +```text +${DOCKER_REGISTRY}:----- +``` + +| Field | Example | Description | +| --- | --- | --- | +| `swift-branch` | `main` | ms-swift branch used during image build | +| `atlas-hardware` | `A2`, `A3`, `300I`, `A5` | Derived from `--soc_version` | +| `python-tag` | `py311` | Derived from `--python_version` | +| `cann-version-tag` | `CANN8.5.1`, `CANN9.0.0` | Parsed from the CANN base image tag | +| `os-tag` | `ubuntu22.04` | Parsed from the CANN base image tag | +| `arch` | `arm`, `x86` | Derived from host architecture or `--arch` | + +Default example on an ARM64 host: + +```text +${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-arm +``` + +A2 / CANN 9.0.0 example: + +```text +${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm +``` + +## Build Locally + +Set the target registry first. The build script renders `docker/Dockerfile.ascend` into the root `Dockerfile`, builds it, and skips push for Ascend images. + +```bash +export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift + +python docker/build_image.py \ + --image_type ascend +``` + +Build an A2 / CANN 9.0.0 image: + +```bash +export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift + +python docker/build_image.py \ + --image_type ascend \ + --base_image quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11 \ + --soc_version ascend910b1 +``` + +Override Megatron or MindSpeed source branches when needed: + +```bash +python docker/build_image.py \ + --image_type ascend \ + --megatron_branch v0.15.3 \ + --mindspeed_branch core_r0.15.3 +``` + +For slow networks, Linux hosts can use Docker host networking after the root `Dockerfile` is generated: + +```bash +docker build --network host \ + -t ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \ + -f Dockerfile . +``` + +## Run An Ascend Container + +The host must have a compatible Ascend driver and firmware installed. The container uses the host NPU devices and driver libraries. + +```bash +docker run --rm -it \ + --name ms_swift_ascend \ + --device /dev/davinci0 \ + --device /dev/davinci_manager \ + --device /dev/devmm_svm \ + --device /dev/hisi_hdc \ + -v /usr/local/dcmi:/usr/local/dcmi \ + -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ + -v /usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64 \ + -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ + -v /etc/ascend_install.info:/etc/ascend_install.info \ + -v /mnt/workspace:/mnt/workspace \ + ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \ + bash +``` + +Inside the container, verify the NPU and Python packages: + +```bash +npu-smi info +python -c "import torch, torch_npu; print(torch.__version__, torch_npu.__version__)" +python -c "import vllm, vllm_ascend; print('vllm ascend ok')" +pip show ms-swift modelscope torch-npu triton-ascend +``` + +## Environment Variables + +| Variable | Value | +| --- | --- | +| `SOC_VERSION` | Selected Ascend SoC version, for example `ascend910b1` or `ascend910_9391` | +| `CANN_VERSION` | Parsed from the base image tag | +| `MEGATRON_LM_PATH` | `/Megatron-LM` | +| `PYTHONPATH` | includes `/Megatron-LM` | +| `VLLM_USE_MODELSCOPE` | `True` | +| `LMDEPLOY_USE_MODELSCOPE` | `True` | +| `MODELSCOPE_CACHE` | `/mnt/workspace/.cache/modelscope/hub` | + +## Notes + +- CANN, firmware, and driver versions must be compatible with each other. +- CANN `8.5.*` and CANN `9.0.0` use different `triton-ascend` install paths in this Dockerfile. +- The image is intended for Ascend NPU ms-swift workflows. CUDA-only packages pulled in by dependencies are removed when they conflict with NPU runtime libraries. +- Use a fixed image tag for production jobs instead of relying on a moving branch name. + +## License + +ms-swift and ModelScope components follow their upstream repository licenses. CANN, MindSpeed, torch-npu, vLLM Ascend, and other pre-installed third-party components are subject to their own upstream licenses. diff --git a/docker/OVERVIEW.ascend.zh.md b/docker/OVERVIEW.ascend.zh.md new file mode 100644 index 00000000..d71f2066 --- /dev/null +++ b/docker/OVERVIEW.ascend.zh.md @@ -0,0 +1,156 @@ +# ms-swift Ascend + +> [English](./OVERVIEW.ascend.md) | 中文 + +ms-swift Ascend 镜像面向华为昇腾 Atlas NPU，提供可直接使用的 ms-swift 运行环境。镜像基于 Ascend CANN 容器镜像构建，包含 Ascend 推理和训练工作流所需的 Python、CANN、PyTorch NPU、vLLM Ascend、Megatron、MindSpeed、mcore-bridge、ms-swift 以及 ModelScope 运行组件。 + +## 快速参考 + +- 基础镜像：`quay.io/ascend/cann:---py` +- 构建模板：`docker/Dockerfile.ascend` +- 构建入口：`docker/build_image.py --image_type ascend` +- 默认基础镜像：`quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11` +- 默认输出 tag：`${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-` +- Ascend runtime 环境来自 `/usr/local/Ascend/ascend-toolkit/set_env.sh` +- 如果镜像内存在 NNAL/ATB，则会加载 `/usr/local/Ascend/nnal/atb/set_env.sh` + +## 镜像内容 + +Ascend Dockerfile 会安装和配置： + +| 组件 | 版本 / 来源 | +| --- | --- | +| CANN | 继承自选定的 `quay.io/ascend/cann` 基础镜像 | +| Python | 继承自基础镜像 tag，例如 `py3.11` | +| PyTorch | `torch==2.9.0` | +| torch-npu | `torch_npu==2.9.0.post2` | +| torchvision / torchaudio | `torchvision==0.24.0`，`torchaudio==2.9.0` | +| vLLM | 从 `vllm-project/vllm` 源码安装，默认分支 `v0.18.0` | +| vLLM Ascend | 从 `vllm-project/vllm-ascend` 源码安装，默认分支 `v0.18.0` | +| Megatron-LM | 源码 checkout，默认分支 `v0.15.3` | +| MindSpeed | 源码 checkout，默认分支 `core_r0.15.3` | +| mcore-bridge | 来自 `modelscope/mcore-bridge` 的源码 checkout | +| ms-swift | 来自 `modelscope/ms-swift` 的源码 checkout，默认分支 `main` | +| ModelScope | 来自 `modelscope/modelscope` 的源码 checkout，默认分支 `master` | +| triton-ascend | CANN `8.5.*` 安装 `3.2.0`；CANN `9.0.0` 下载并本地安装 `3.2.1` wheel | + +## 支持的 Tag 格式 + +通过 `docker/build_image.py --image_type ascend` 构建的镜像使用以下 tag 格式： + +```text +${DOCKER_REGISTRY}:----- +``` + +| 字段 | 示例 | 说明 | +| --- | --- | --- | +| `swift-branch` | `main` | 构建镜像时使用的 ms-swift 分支 | +| `atlas-hardware` | `A2`、`A3`、`300I`、`A5` | 从 `--soc_version` 推导 | +| `python-tag` | `py311` | 从 `--python_version` 推导 | +| `cann-version-tag` | `CANN8.5.1`、`CANN9.0.0` | 从 CANN 基础镜像 tag 解析 | +| `os-tag` | `ubuntu22.04` | 从 CANN 基础镜像 tag 解析 | +| `arch` | `arm`、`x86` | 从宿主机架构或 `--arch` 推导 | + +ARM64 宿主机上的默认示例： + +```text +${DOCKER_REGISTRY}:main-A3-py311-CANN8.5.1-ubuntu22.04-arm +``` + +A2 / CANN 9.0.0 示例： + +```text +${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm +``` + +## 本地构建 + +先设置目标镜像仓库。构建脚本会把 `docker/Dockerfile.ascend` 渲染成根目录 `Dockerfile`，然后执行构建；Ascend 镜像分支当前不执行 push。 + +```bash +export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift + +python docker/build_image.py \ + --image_type ascend +``` + +构建 A2 / CANN 9.0.0 镜像： + +```bash +export DOCKER_REGISTRY=registry.example.com/ms-swift/ms-swift + +python docker/build_image.py \ + --image_type ascend \ + --base_image quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11 \ + --soc_version ascend910b1 +``` + +需要时可以覆盖 Megatron 或 MindSpeed 源码分支： + +```bash +python docker/build_image.py \ + --image_type ascend \ + --megatron_branch v0.15.3 \ + --mindspeed_branch core_r0.15.3 +``` + +如果构建时网络较慢，Linux 宿主机可以在根目录 `Dockerfile` 生成后使用 host network 构建： + +```bash +docker build --network host \ + -t ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \ + -f Dockerfile . +``` + +## 运行 Ascend 容器 + +宿主机需要提前安装兼容的 Ascend driver 和 firmware。容器通过挂载宿主机 NPU 设备和 driver 库使用昇腾硬件。 + +```bash +docker run --rm -it \ + --name ms_swift_ascend \ + --device /dev/davinci0 \ + --device /dev/davinci_manager \ + --device /dev/devmm_svm \ + --device /dev/hisi_hdc \ + -v /usr/local/dcmi:/usr/local/dcmi \ + -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ + -v /usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64 \ + -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ + -v /etc/ascend_install.info:/etc/ascend_install.info \ + -v /mnt/workspace:/mnt/workspace \ + ${DOCKER_REGISTRY}:main-A2-py311-CANN9.0.0-ubuntu22.04-arm \ + bash +``` + +进入容器后可以验证 NPU 和 Python 包： + +```bash +npu-smi info +python -c "import torch, torch_npu; print(torch.__version__, torch_npu.__version__)" +python -c "import vllm, vllm_ascend; print('vllm ascend ok')" +pip show ms-swift modelscope torch-npu triton-ascend +``` + +## 环境变量 + +| 变量 | 值 | +| --- | --- | +| `SOC_VERSION` | 选定的 Ascend SoC，例如 `ascend910b1` 或 `ascend910_9391` | +| `CANN_VERSION` | 从基础镜像 tag 解析得到 | +| `MEGATRON_LM_PATH` | `/Megatron-LM` | +| `PYTHONPATH` | 包含 `/Megatron-LM` | +| `VLLM_USE_MODELSCOPE` | `True` | +| `LMDEPLOY_USE_MODELSCOPE` | `True` | +| `MODELSCOPE_CACHE` | `/mnt/workspace/.cache/modelscope/hub` | + +## 注意事项 + +- CANN、firmware 和 driver 版本必须互相兼容。 +- 这个 Dockerfile 对 CANN `8.5.*` 和 CANN `9.0.0` 使用不同的 `triton-ascend` 安装路径。 +- 该镜像面向 Ascend NPU 上的 ms-swift 工作流。依赖安装过程中引入且与 NPU runtime 冲突的 CUDA-only 包会被移除。 +- 生产任务建议使用固定镜像 tag，不要依赖浮动分支名。 + +## License + +ms-swift 和 ModelScope 组件遵循各自上游仓库的 license。CANN、MindSpeed、torch-npu、vLLM Ascend 以及其他预装第三方组件遵循各自上游 license。 diff --git a/docker/build_image.py b/docker/build_image.py index a80879b0..ba98b765 100644 --- a/docker/build_image.py +++ b/docker/build_image.py @@ -1,6 +1,7 @@ import argparse import os import platform +import re import subprocess from datetime import datetime from typing import Any @@ -462,6 +463,9 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy class AscendImageBuilder(StableGPUImageBuilder): + _CANN_VERSION_PATTERN = re.compile(r'^\d+(?:\.[0-9A-Za-z]+)+$') + _OS_TAG_PATTERN = re.compile(r'^[A-Za-z]+[0-9][0-9A-Za-z.]*$') + @staticmethod def _normalize_arch(arch: str = None) -> str: arch = arch or platform.machine() @@ -496,12 +500,39 @@ class AscendImageBuilder(StableGPUImageBuilder): 'ascend310p1, and values starting with ascend950.') return atlas_mapping[soc_version] + @classmethod + def _get_cann_os_tags(cls, base_image: str) -> tuple: + if ':' not in base_image.rsplit('/', 1)[-1]: + raise ValueError( + f'Ascend base image must include a tag: {base_image}') + + base_tag = base_image.rsplit(':', 1)[1] + parts = base_tag.split('-') + if len(parts) < 4: + raise ValueError( + 'Ascend base image tag must look like ' + f'---py, got: ' + f'{base_tag}') + + cann_version = parts[0] + os_tag = parts[2] + if not cls._CANN_VERSION_PATTERN.fullmatch(cann_version): + raise ValueError(f'Invalid CANN version in Ascend base image tag: ' + f'{cann_version}') + if not cls._OS_TAG_PATTERN.fullmatch(os_tag): + raise ValueError( + f'Invalid OS tag in Ascend base image tag: {os_tag}') + + return cann_version, f'CANN{cann_version}', os_tag + def init_args(self, args) -> Any: if not args.base_image: # Reuse the prebuilt vllm-ascend image to avoid rebuilding its stack. args.base_image = 'quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11' args.arch = self._normalize_arch(args.arch) args.atlas_hardware = self._get_atlas_hardware(args.soc_version) + args.cann_version, args.cann_version_tag, args.os_tag = ( + self._get_cann_os_tags(args.base_image)) return super().init_args(args) def generate_dockerfile(self) -> str: @@ -512,18 +543,24 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy content = f.read() content = content.replace('{base_image}', self.args.base_image) content = content.replace('{soc_version}', self.args.soc_version) + content = content.replace('{cann_version}', self.args.cann_version) content = content.replace('{extra_content}', extra_content) content = content.replace('{cur_time}', formatted_time) content = content.replace('{install_ms_deps}', 'False') content = content.replace('{modelscope_branch}', self.args.modelscope_branch) content = content.replace('{swift_branch}', self.args.swift_branch) + content = content.replace('{megatron_branch}', + self.args.megatron_branch) + content = content.replace('{mindspeed_branch}', + self.args.mindspeed_branch) return content def image(self) -> str: return ( f'{docker_registry}:{self.args.swift_branch}-' - f'{self.args.atlas_hardware}-{self.args.python_tag}-{self.args.arch}' + f'{self.args.atlas_hardware}-{self.args.python_tag}-' + f'{self.args.cann_version_tag}-{self.args.os_tag}-{self.args.arch}' ) def push(self): @@ -549,6 +586,8 @@ parser.add_argument('--optimum_version', type=str, default=None) parser.add_argument('--modelscope_branch', type=str, default='master') parser.add_argument('--modelscope_version', type=str, default='9.99.0') parser.add_argument('--swift_branch', type=str, default='main') +parser.add_argument('--megatron_branch', type=str, default='v0.15.3') +parser.add_argument('--mindspeed_branch', type=str, default='core_r0.15.3') parser.add_argument('--soc_version', type=str, default='ascend910_9391') parser.add_argument('--arch', type=str, choices=['x86', 'arm'], default=None) parser.add_argument('--dry_run', type=int, default=0) From cecb699cd0c93a2f31ccd163c84d626b76370631 Mon Sep 17 00:00:00 2001 From: tastelikefeet <58414341+tastelikefeet@users.noreply.github.com> Date: Thu, 18 Jun 2026 17:16:11 +0800 Subject: [PATCH 10/11] fix (#1740) --- modelscope/utils/automodel_utils.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/modelscope/utils/automodel_utils.py b/modelscope/utils/automodel_utils.py index ba087b22..7b78309c 100644 --- a/modelscope/utils/automodel_utils.py +++ b/modelscope/utils/automodel_utils.py @@ -131,7 +131,7 @@ def check_model_from_owner_group(model_dir: str, """This checking is for the torch.load, this function may eval malicious code into memory Args: - model_dir: The local model_dir + model_dir: The local model_dir or model_id owner_group: The owner group to trust Returns: @@ -142,6 +142,15 @@ def check_model_from_owner_group(model_dir: str, if owner_group is None: owner_group = ['iic', 'damo'] model_dir = model_dir.rstrip('/').rstrip('\\') - model_dir = os.path.dirname(model_dir) - group = os.path.basename(model_dir) - return group in owner_group + parent_dir = os.path.dirname(model_dir) + group = os.path.basename(parent_dir) + if group in owner_group: + return True + # Also check cache path pattern: {cache_root}/{owner}--{model_name}/snapshots/{revision} + # Only check the grandparent directory to avoid path traversal attacks + grandparent = os.path.basename(os.path.dirname(parent_dir)) + if '--' in grandparent: + prefix = grandparent.split('--', 1)[0] + if prefix in owner_group: + return True + return False From cf08fb242d8c3695d0fc2d41503daa14aefd4dde Mon Sep 17 00:00:00 2001 From: "Xingjun.Wang" Date: Tue, 23 Jun 2026 16:42:49 +0800 Subject: [PATCH 11/11] [Fix] Fix studio cli (#1742) --- modelscope/cli/plugins.py | 2 +- modelscope/hub/api.py | 16 +++++++++++++--- modelscope/utils/automodel_utils.py | 8 +++++--- pyproject.toml | 1 + 4 files changed, 20 insertions(+), 7 deletions(-) diff --git a/modelscope/cli/plugins.py b/modelscope/cli/plugins.py index 218bc160..7b4d500b 100644 --- a/modelscope/cli/plugins.py +++ b/modelscope/cli/plugins.py @@ -11,7 +11,7 @@ plugins_manager = PluginsManager() class PluginsCMD(CLICommand): - name = 'plugin' + name = 'plugins' @staticmethod def register(subparsers: ArgumentParser) -> None: diff --git a/modelscope/hub/api.py b/modelscope/hub/api.py index 356e7a42..ab39a015 100644 --- a/modelscope/hub/api.py +++ b/modelscope/hub/api.py @@ -90,9 +90,19 @@ class HubApi(_LegacyHubApi): # ------------------------------------------------------------------ # Legacy method shims missing from LegacyHubApi # ------------------------------------------------------------------ - def create_model(self, model_id: str, **kwargs) -> None: - """Create a model repo — delegates to ``create_repo`` (model type).""" - return self.create_repo(model_id, repo_type='model', **kwargs) + def create_model(self, model_id: str, **kwargs) -> str: + """Create a model repo — delegates to ``create_repo`` (model type). + + Returns the model repository URL for backward compatibility. + Authentication errors are converted to ``ValueError`` by the compat layer. + """ + # LegacyHubApi.create_model handles exception conversion; + # override endpoint in the returned URL with local self.endpoint. + result = super().create_model(model_id, **kwargs) + # Replace endpoint in URL with the one resolved at __init__ time + if result and self.endpoint and self.endpoint not in result: + return f'{self.endpoint}/models/{model_id}' + return result def get_model_url(self, model_id: str) -> str: """Return the model page URL ``{endpoint}/{model_id}``.""" diff --git a/modelscope/utils/automodel_utils.py b/modelscope/utils/automodel_utils.py index 7b78309c..8acec9f5 100644 --- a/modelscope/utils/automodel_utils.py +++ b/modelscope/utils/automodel_utils.py @@ -147,10 +147,12 @@ def check_model_from_owner_group(model_dir: str, if group in owner_group: return True # Also check cache path pattern: {cache_root}/{owner}--{model_name}/snapshots/{revision} - # Only check the grandparent directory to avoid path traversal attacks + # Require exactly "{owner}--{name}" format (2 segments split by --) + # to prevent spoofing via accounts like "iic--hacked" which would + # produce paths like "iic--hacked--evil" and bypass the check. grandparent = os.path.basename(os.path.dirname(parent_dir)) if '--' in grandparent: - prefix = grandparent.split('--', 1)[0] - if prefix in owner_group: + parts = grandparent.split('--') + if len(parts) == 2 and parts[0] in owner_group: return True return False diff --git a/pyproject.toml b/pyproject.toml index 96994d40..0b620f96 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ plugins = "modelscope.cli.plugins:PluginsCMD" skills = "modelscope.cli.skills:SkillsCMD" llamafile = "modelscope.cli.llamafile:LlamafileCMD" modelcard = "modelscope.cli.modelcard:ModelCardCMD" +studio = "modelscope.cli.studio:StudioCMD" [build-system] requires = ["setuptools>=69", "wheel"]