Fix the log level and return value of push_to_hub

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/12338218
This commit is contained in:
yuze.zyz
2023-04-15 16:42:35 +08:00
parent a4bafd8b25
commit 2650d37094
5 changed files with 70 additions and 42 deletions

View File

@@ -19,20 +19,27 @@ def _api_push_to_hub(repo_name,
private=True,
commit_message='',
source_repo=''):
api = HubApi()
api.login(token)
api.push_model(
repo_name,
output_dir,
visibility=ModelVisibility.PUBLIC
if not private else ModelVisibility.PRIVATE,
chinese_name=repo_name,
commit_message=commit_message,
original_model_id=source_repo)
commit_message = commit_message or 'No commit message'
logger.info(
f'Successfully upload the model to {repo_name} with message: {commit_message}'
)
try:
api = HubApi()
api.login(token)
api.push_model(
repo_name,
output_dir,
visibility=ModelVisibility.PUBLIC
if not private else ModelVisibility.PRIVATE,
chinese_name=repo_name,
commit_message=commit_message,
original_model_id=source_repo)
commit_message = commit_message or 'No commit message'
logger.info(
f'Successfully upload the model to {repo_name} with message: {commit_message}'
)
return True
except Exception as e:
logger.error(
f'Error happens when uploading model {repo_name} with message: {commit_message}: {e}'
)
return False
def push_to_hub(repo_name,
@@ -41,8 +48,7 @@ def push_to_hub(repo_name,
private=True,
retry=3,
commit_message='',
source_repo='',
async_upload=False):
source_repo=''):
"""
Args:
repo_name: The repo name for the modelhub repo
@@ -52,10 +58,9 @@ def push_to_hub(repo_name,
retry: Retry times if something error in uploading, default 3
commit_message: The commit message
source_repo: The source repo (model id) which this model comes from
async_upload: Upload async, if True, the `retry` parameter will have no affection
Returns:
A handler if async_upload is True, else None
The boolean value to represent whether the model is uploaded.
"""
if token is None:
token = os.environ.get('MODELSCOPE_API_TOKEN')
@@ -66,15 +71,39 @@ def push_to_hub(repo_name,
logger.info(
f'Uploading {output_dir} to {repo_name} with message {commit_message}')
if not async_upload:
for i in range(retry):
try:
_api_push_to_hub(repo_name, output_dir, token, private,
commit_message, source_repo)
except Exception as e:
logger.info(f'Error happens when uploading model: {e}')
continue
break
else:
return _executor.submit(_api_push_to_hub, repo_name, output_dir, token,
private, commit_message, source_repo)
for i in range(retry):
if _api_push_to_hub(repo_name, output_dir, token, private,
commit_message, source_repo):
return True
return False
def push_to_hub_async(repo_name,
output_dir,
token=None,
private=True,
commit_message='',
source_repo=''):
"""
Args:
repo_name: The repo name for the modelhub repo
output_dir: The local output_dir for the checkpoint
token: The user api token, function will check the `MODELSCOPE_API_TOKEN` variable if this argument is None
private: If is a private repo, default True
commit_message: The commit message
source_repo: The source repo (model id) which this model comes from
Returns:
A handler to check the result and the status
"""
if token is None:
token = os.environ.get('MODELSCOPE_API_TOKEN')
assert token is not None, 'Either pass in a token or to set `MODELSCOPE_API_TOKEN` in the environment variables.'
assert os.path.isdir(output_dir)
assert 'configuration.json' in os.listdir(output_dir) or 'configuration.yaml' in os.listdir(output_dir) \
or 'configuration.yml' in os.listdir(output_dir)
logger.info(
f'Uploading {output_dir} to {repo_name} with message {commit_message}')
return _executor.submit(_api_push_to_hub, repo_name, output_dir, token,
private, commit_message, source_repo)

View File

@@ -9,7 +9,7 @@ import torch
from packaging import version
from modelscope.hub.check_model import check_model_is_id
from modelscope.hub.push_to_hub import push_to_hub
from modelscope.hub.push_to_hub import push_to_hub_async
from modelscope.metainfo import Hooks, Pipelines
from modelscope.utils.checkpoint import (load_checkpoint, save_checkpoint,
save_configuration)
@@ -156,14 +156,13 @@ class CheckpointHook(Hook):
self.is_model_id = check_model_is_id(trainer.input_model_id,
self.hub_token)
return push_to_hub(
return push_to_hub_async(
self.model_id_with_org,
os.path.join(self.save_dir, self.output_sub_dir),
token=self.hub_token,
private=self.private_hub,
commit_message=prefix,
source_repo=trainer.input_model_id if self.is_model_id else '',
async_upload=True)
source_repo=trainer.input_model_id if self.is_model_id else '')
def _save_checkpoint(self, trainer, prefix):
"""Save checkpoint files and remove obsolete ones

View File

@@ -13,7 +13,7 @@ python-dateutil>=2.1
pyyaml
requests
scipy
setuptools
setuptools==59.8.0
simplejson>=3.3.0
sortedcontainers>=1.5.9
tqdm>=4.64.0

View File

@@ -1,5 +1,5 @@
accelerate
diffusers>=0.13.1
diffusers>=0.13.1,<0.15.0
ftfy>=6.0.3
librosa<=0.9.2
opencv-python

View File

@@ -9,7 +9,7 @@ import uuid
from modelscope.hub.api import HubApi
from modelscope.hub.constants import Licenses, ModelVisibility
from modelscope.hub.errors import GitError, HTTPError, NotLoginException
from modelscope.hub.push_to_hub import push_to_hub
from modelscope.hub.push_to_hub import push_to_hub, push_to_hub_async
from modelscope.hub.repository import Repository
from modelscope.utils.constant import ModelFile
from modelscope.utils.logger import get_logger
@@ -156,21 +156,21 @@ class HubUploadTest(unittest.TestCase):
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
def test_push_to_hub(self):
future = push_to_hub(
ret = push_to_hub(
repo_name=self.create_model_name,
output_dir=self.finetune_path,
token=TEST_ACCESS_TOKEN1)
self.assertTrue(future is None)
self.assertTrue(ret is True)
@unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
def test_push_to_hub_async(self):
future = push_to_hub(
future = push_to_hub_async(
repo_name=self.create_model_name,
output_dir=self.finetune_path,
token=TEST_ACCESS_TOKEN1,
async_upload=True)
token=TEST_ACCESS_TOKEN1)
while not future.done():
time.sleep(1)
self.assertTrue(future.result())
if __name__ == '__main__':