From 352336a8a2beb9f8c64f3a81fea85a8f606f62eb Mon Sep 17 00:00:00 2001 From: Yingda Chen Date: Mon, 24 Feb 2025 22:59:06 +0800 Subject: [PATCH] Add download token (#1233) * update token behavior * fix * fixt * address cr --------- Co-authored-by: Yingda Chen --- modelscope/cli/cli.py | 3 --- modelscope/cli/download.py | 28 ++++++++++++++++++++-------- modelscope/cli/upload.py | 13 ++++++++----- modelscope/hub/api.py | 13 ++++++++++++- 4 files changed, 40 insertions(+), 17 deletions(-) diff --git a/modelscope/cli/cli.py b/modelscope/cli/cli.py index afadfa91..1cb98ea0 100644 --- a/modelscope/cli/cli.py +++ b/modelscope/cli/cli.py @@ -40,9 +40,6 @@ def run_cmd(): if not hasattr(args, 'func'): parser.print_help() exit(1) - if args.token is not None: - api = HubApi() - api.login(args.token) cmd = args.func(args) cmd.execute() diff --git a/modelscope/cli/download.py b/modelscope/cli/download.py index 6b430453..723fdfe9 100644 --- a/modelscope/cli/download.py +++ b/modelscope/cli/download.py @@ -3,6 +3,7 @@ import os from argparse import ArgumentParser from modelscope.cli.base import CLICommand +from modelscope.hub.api import HubApi from modelscope.hub.constants import DEFAULT_MAX_WORKERS from modelscope.hub.file_download import (dataset_file_download, model_file_download) @@ -54,16 +55,21 @@ class DownloadCMD(CLICommand): default='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 model.') + help='Revision of the entity (e.g., model).') parser.add_argument( '--cache_dir', type=str, default=None, - help='Cache directory to save model.') + help='Cache directory to save entity (e.g., model).') parser.add_argument( '--local_dir', type=str, @@ -118,6 +124,10 @@ class DownloadCMD(CLICommand): % self.args.repo_type) if not self.args.model and not self.args.dataset: raise Exception('Model or dataset must be set.') + cookies = None + if self.args.token is not None: + api = HubApi() + 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( @@ -125,7 +135,8 @@ class DownloadCMD(CLICommand): self.args.files[0], cache_dir=self.args.cache_dir, local_dir=self.args.local_dir, - revision=self.args.revision) + revision=self.args.revision, + cookies=cookies) elif len( self.args.files) > 1: # download specified multiple files. snapshot_download( @@ -135,7 +146,7 @@ class DownloadCMD(CLICommand): local_dir=self.args.local_dir, allow_file_pattern=self.args.files, max_workers=self.args.max_workers, - ) + cookies=cookies) else: # download repo snapshot_download( self.args.model, @@ -145,7 +156,7 @@ class DownloadCMD(CLICommand): allow_file_pattern=convert_patterns(self.args.include), ignore_file_pattern=convert_patterns(self.args.exclude), max_workers=self.args.max_workers, - ) + cookies=cookies) 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 @@ -154,7 +165,8 @@ class DownloadCMD(CLICommand): self.args.files[0], cache_dir=self.args.cache_dir, local_dir=self.args.local_dir, - revision=dataset_revision) + revision=dataset_revision, + cookies=cookies) elif len( self.args.files) > 1: # download specified multiple files. dataset_snapshot_download( @@ -164,7 +176,7 @@ class DownloadCMD(CLICommand): local_dir=self.args.local_dir, allow_file_pattern=self.args.files, max_workers=self.args.max_workers, - ) + cookies=cookies) else: # download repo dataset_snapshot_download( self.args.dataset, @@ -174,6 +186,6 @@ class DownloadCMD(CLICommand): allow_file_pattern=convert_patterns(self.args.include), ignore_file_pattern=convert_patterns(self.args.exclude), max_workers=self.args.max_workers, - ) + cookies=cookies) else: pass # noop diff --git a/modelscope/cli/upload.py b/modelscope/cli/upload.py index 453a6314..a050c5b2 100644 --- a/modelscope/cli/upload.py +++ b/modelscope/cli/upload.py @@ -91,7 +91,7 @@ class UploadCMD(CLICommand): '--endpoint', type=str, default=get_endpoint(), - help='Endpoint for Modelscope service.') + help='Endpoint for ModelScope service.') parser.set_defaults(func=subparser_func) @@ -137,14 +137,15 @@ class UploadCMD(CLICommand): # Check token and login # The cookies will be reused if the user has logged in before. + cookies = None api = HubApi(endpoint=self.args.endpoint) - if self.args.token: - api.login(access_token=self.args.token) - cookies = ModelScopeConfig.get_cookies() + cookies = api.get_cookies(access_token=self.args.token) + else: + cookies = ModelScopeConfig.get_cookies() if cookies is None: raise ValueError( - 'The `token` is not provided! ' + 'No credential found for entity upload. ' 'You can pass the `--token` argument, ' 'or use api.login(access_token=`your_sdk_token`). ' 'Your token is available at https://modelscope.cn/my/myaccesstoken' @@ -158,6 +159,7 @@ class UploadCMD(CLICommand): 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( @@ -170,6 +172,7 @@ class UploadCMD(CLICommand): 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') diff --git a/modelscope/hub/api.py b/modelscope/hub/api.py index f5a2f39b..3a249f26 100644 --- a/modelscope/hub/api.py +++ b/modelscope/hub/api.py @@ -34,6 +34,7 @@ from modelscope.hub.constants import (API_HTTP_CLIENT_MAX_RETRIES, API_RESPONSE_FIELD_USERNAME, DEFAULT_CREDENTIALS_PATH, DEFAULT_MAX_WORKERS, + DEFAULT_MODELSCOPE_DOMAIN, MODELSCOPE_CLOUD_ENVIRONMENT, MODELSCOPE_CLOUD_USERNAME, MODELSCOPE_REQUEST_ID, ONE_YEAR_SECONDS, @@ -112,9 +113,19 @@ class HubApi: self.upload_checker = UploadingCheck() + def get_cookies(self, access_token): + from requests.cookies import RequestsCookieJar + jar = RequestsCookieJar() + jar.set('m_session_id', + access_token, + domain=os.getenv('MODELSCOPE_DOMAIN', + DEFAULT_MODELSCOPE_DOMAIN), + path='/') + return jar + def login( self, - access_token: Optional[str] = None, + access_token: Optional[str] = None ): """Login with your SDK access token, which can be obtained from https://www.modelscope.cn user center.