| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586 |
- import argparse
- import traceback
- from typing import List, Union
- from modelscope.hub.api import HubApi
- from modelscope.hub.file_download import model_file_download
- from modelscope.pipelines import pipeline
- from modelscope.utils.config import Config
- from modelscope.utils.constant import ModelFile
- from modelscope.utils.input_output import (
- call_pipeline_with_json, get_pipeline_information_by_pipeline,
- get_task_input_examples, pipeline_output_to_service_base64_output)
- from modelscope.utils.logger import get_logger
- logger = get_logger()
- class DeployChecker:
- def __init__(self):
- self.api = HubApi()
- def check_model(self, model_id: str, model_revision=None):
- # get model_revision & task info
- if not model_revision:
- model_revisions = self.api.list_model_revisions(model_id)
- logger.info(
- f'All model_revisions of `{model_id}`: {model_revisions}')
- if len(model_revisions):
- model_revision = model_revisions[0]
- else:
- logger.error(f'{model_id} has no revision.')
- configuration_file = model_file_download(
- model_id=model_id,
- file_path=ModelFile.CONFIGURATION,
- revision=model_revision)
- cfg = Config.from_file(configuration_file)
- task = cfg.safe_get('task')
- # init pipeline
- ppl = pipeline(
- task=task,
- model=model_id,
- model_revision=model_revision,
- external_engine_for_llm=True)
- pipeline_info = get_pipeline_information_by_pipeline(ppl)
- # call pipeline
- data = get_task_input_examples(task)
- infer_result = call_pipeline_with_json(pipeline_info, ppl, data)
- result = pipeline_output_to_service_base64_output(task, infer_result)
- return result
- def check_deploy(models: Union[str, List], revisions: Union[str, List] = None):
- if not isinstance(models, list):
- models = [models]
- if not isinstance(revisions, list):
- revisions = [revisions] * (1 if revisions else len(models))
- if len(models) != len(revisions):
- logger.error(
- f'The number of models and revisions need to be equal: The number of models'
- f' is {len(model)} while the number of revisions is {len(revision)}.'
- )
- checker = DeployChecker()
- for model, revision in zip(models, revisions):
- try:
- res = checker.check_model(model, revision)
- logger.info(f'{model} {revision}: Deploy pre-check pass. {res}\n')
- except BaseException as e:
- logger.info(
- f'{model} {revision}: Deploy pre-check failed: {e}. {traceback.print_exc()}\n'
- )
- if __name__ == '__main__':
- parser = argparse.ArgumentParser()
- parser.add_argument('--model_id', type=str)
- parser.add_argument('--revision', type=str, default=None)
- args = parser.parse_args()
- check_deploy(args.model_id, args.revision)
|
