wehub-resource-sync
290 lines
12 KiB
Python
290 lines
12 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import asyncio
|
|
import inspect
|
|
import json
|
|
import multiprocessing
|
|
import time
|
|
import uvicorn
|
|
from aiohttp import ClientConnectorError
|
|
from contextlib import contextmanager
|
|
from dataclasses import asdict
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
|
from http import HTTPStatus
|
|
from threading import Thread
|
|
from typing import List, Optional, Union
|
|
|
|
from swift.arguments import DeployArguments, InferArguments
|
|
from swift.infer_engine import AdapterRequest, InferClient, RequestConfig
|
|
from swift.infer_engine.protocol import (ChatCompletionRequest, CompletionRequest, EmbeddingRequest, Model, ModelList,
|
|
MultiModalRequestMixin, RolloutInferRequest)
|
|
from swift.metrics import InferStats
|
|
from swift.utils import JsonlWriter, get_logger
|
|
from .infer import SwiftInfer
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
class SwiftDeploy(SwiftInfer):
|
|
args_class = DeployArguments
|
|
args: args_class
|
|
|
|
@staticmethod
|
|
def get_infer_engine(args: InferArguments, template=None, **kwargs):
|
|
if isinstance(args, DeployArguments) and args.infer_backend == 'vllm':
|
|
engine_kwargs = (kwargs.get('engine_kwargs') or {}).copy()
|
|
if args.vllm_data_parallel_size > 1:
|
|
if not args.vllm_use_async_engine:
|
|
raise ValueError('vLLM data parallel requires `vllm_use_async_engine=True` in deploy mode.')
|
|
engine_kwargs.setdefault('data_parallel_size', args.vllm_data_parallel_size)
|
|
logger.info(f'Enable vLLM data parallel with size {args.vllm_data_parallel_size}.')
|
|
if args.max_logprobs is not None:
|
|
engine_kwargs['max_logprobs'] = args.max_logprobs
|
|
kwargs['engine_kwargs'] = engine_kwargs
|
|
return SwiftInfer.get_infer_engine(args, template, **kwargs)
|
|
|
|
def _register_app(self):
|
|
self.app.get('/health')(self.health)
|
|
self.app.get('/ping')(self.ping)
|
|
self.app.post('/ping')(self.ping)
|
|
self.app.get('/v1/models')(self.get_available_models)
|
|
self.app.post('/v1/chat/completions')(self.create_chat_completion)
|
|
self.app.post('/v1/completions')(self.create_completion)
|
|
self.app.post('/v1/embeddings')(self.create_embedding)
|
|
self.app.post('/infer/')(self.infer_handler)
|
|
|
|
def __init__(self, args: Optional[Union[List[str], DeployArguments]] = None) -> None:
|
|
super().__init__(args)
|
|
|
|
self.infer_engine.strict = True
|
|
self.infer_stats = InferStats()
|
|
self.app = FastAPI(lifespan=self.lifespan)
|
|
self._register_app()
|
|
|
|
async def _log_stats_hook(self):
|
|
while True:
|
|
await asyncio.sleep(self.args.log_interval)
|
|
self._compute_infer_stats()
|
|
self.infer_stats.reset()
|
|
|
|
def _compute_infer_stats(self):
|
|
global_stats = self.infer_stats.compute()
|
|
for k, v in global_stats.items():
|
|
global_stats[k] = round(v, 8)
|
|
logger.info(global_stats)
|
|
|
|
def lifespan(self, app: FastAPI):
|
|
args = self.args
|
|
if args.log_interval > 0:
|
|
thread = Thread(target=lambda: asyncio.run(self._log_stats_hook()), daemon=True)
|
|
thread.start()
|
|
try:
|
|
yield
|
|
finally:
|
|
if args.log_interval > 0:
|
|
self._compute_infer_stats()
|
|
|
|
def _get_model_list(self):
|
|
args = self.args
|
|
model_list = [args.served_model_name or args.model_suffix]
|
|
if args.adapter_mapping:
|
|
model_list += [name for name in args.adapter_mapping.keys()]
|
|
return model_list
|
|
|
|
async def health(self) -> Response:
|
|
"""Health check endpoint."""
|
|
if self.infer_engine is not None:
|
|
return Response(status_code=200)
|
|
else:
|
|
return Response(status_code=503)
|
|
|
|
async def ping(self) -> Response:
|
|
"""Ping check endpoint. Required for SageMaker compatibility."""
|
|
return await self.health()
|
|
|
|
async def get_available_models(self):
|
|
model_list = self._get_model_list()
|
|
data = [Model(id=model_id, owned_by=self.args.owned_by) for model_id in model_list]
|
|
return ModelList(data=data)
|
|
|
|
async def _check_model(self, request: ChatCompletionRequest) -> Optional[str]:
|
|
available_models = await self.get_available_models()
|
|
model_list = [model.id for model in available_models.data]
|
|
if request.model not in model_list:
|
|
return f'`{request.model}` is not in the model_list: `{model_list}`.'
|
|
|
|
def _check_api_key(self, raw_request: Request) -> Optional[str]:
|
|
api_key = self.args.api_key
|
|
if api_key is None:
|
|
return
|
|
authorization = dict(raw_request.headers).get('authorization')
|
|
error_msg = 'API key error'
|
|
if authorization is None or not authorization.startswith('Bearer '):
|
|
return error_msg
|
|
request_api_key = authorization[7:]
|
|
if request_api_key != api_key:
|
|
return error_msg
|
|
|
|
def _check_max_logprobs(self, request):
|
|
args = self.args
|
|
if isinstance(request.top_logprobs, int) and request.top_logprobs > args.max_logprobs:
|
|
return (f'The value of top_logprobs({request.top_logprobs}) is greater than '
|
|
f'the server\'s max_logprobs({args.max_logprobs}).')
|
|
|
|
@staticmethod
|
|
def create_error_response(status_code: Union[int, str, HTTPStatus], message: str) -> JSONResponse:
|
|
status_code = int(status_code)
|
|
return JSONResponse({'message': message, 'object': 'error'}, status_code)
|
|
|
|
def _post_process(self, request_info, response, return_cmpl_response: bool = False):
|
|
args = self.args
|
|
|
|
for i in range(len(response.choices)):
|
|
if not hasattr(response.choices[i], 'message') or not isinstance(response.choices[i].message.content,
|
|
(tuple, list)):
|
|
continue
|
|
for j, content in enumerate(response.choices[i].message.content):
|
|
if isinstance(content, dict) and content['type'] == 'image':
|
|
b64_image = MultiModalRequestMixin.to_base64(content['image'])
|
|
response.choices[i].message.content[j]['image'] = f'data:image/jpg;base64,{b64_image}'
|
|
|
|
is_finished = all(response.choices[i].finish_reason for i in range(len(response.choices)))
|
|
if 'stream' in response.__class__.__name__.lower():
|
|
request_info['response'] += response.choices[0].delta.content
|
|
else:
|
|
request_info['response'] = response.choices[0].message.content
|
|
if return_cmpl_response:
|
|
response = response.to_cmpl_response()
|
|
if is_finished:
|
|
if args.log_interval > 0:
|
|
self.infer_stats.update(response)
|
|
if self.jsonl_writer:
|
|
self.jsonl_writer.append(request_info)
|
|
if self.args.verbose:
|
|
logger.info(request_info)
|
|
return response
|
|
|
|
def _set_request_config(self, request_config) -> None:
|
|
default_request_config = self.args.get_request_config()
|
|
if default_request_config is None:
|
|
return
|
|
for key, val in asdict(request_config).items():
|
|
default_val = getattr(default_request_config, key)
|
|
if default_val is not None and (val is None or isinstance(val, (list, tuple)) and len(val) == 0):
|
|
setattr(request_config, key, default_val)
|
|
|
|
async def create_chat_completion(self,
|
|
request: ChatCompletionRequest,
|
|
raw_request: Request,
|
|
*,
|
|
return_cmpl_response: bool = False):
|
|
args = self.args
|
|
error_msg = (await self._check_model(request) or self._check_api_key(raw_request)
|
|
or self._check_max_logprobs(request))
|
|
if error_msg:
|
|
return self.create_error_response(HTTPStatus.BAD_REQUEST, error_msg)
|
|
infer_kwargs = self.infer_kwargs.copy()
|
|
adapter_path = args.adapter_mapping.get(request.model)
|
|
if adapter_path:
|
|
infer_kwargs['adapter_request'] = AdapterRequest(request.model, adapter_path)
|
|
|
|
infer_request, request_config = request.parse()
|
|
self._set_request_config(request_config)
|
|
request_info = {'response': '', 'infer_request': infer_request.to_printable()}
|
|
|
|
def pre_infer_hook(kwargs):
|
|
request_info['generation_config'] = kwargs['generation_config']
|
|
return kwargs
|
|
|
|
infer_kwargs['pre_infer_hook'] = pre_infer_hook
|
|
try:
|
|
res_or_gen = await self.infer_async(infer_request, request_config, **infer_kwargs)
|
|
except Exception as e:
|
|
import traceback
|
|
logger.info(traceback.format_exc())
|
|
return self.create_error_response(HTTPStatus.BAD_REQUEST, str(e))
|
|
if request_config.stream:
|
|
|
|
async def _gen_wrapper():
|
|
async for res in res_or_gen:
|
|
res = self._post_process(request_info, res, return_cmpl_response)
|
|
yield f'data: {json.dumps(asdict(res), ensure_ascii=False)}\n\n'
|
|
yield 'data: [DONE]\n\n'
|
|
|
|
return StreamingResponse(_gen_wrapper(), media_type='text/event-stream')
|
|
elif hasattr(res_or_gen, 'choices'):
|
|
# instance of ChatCompletionResponse
|
|
return self._post_process(request_info, res_or_gen, return_cmpl_response)
|
|
else:
|
|
return res_or_gen
|
|
|
|
async def create_completion(self, request: CompletionRequest, raw_request: Request):
|
|
chat_request = ChatCompletionRequest.from_cmpl_request(request)
|
|
return await self.create_chat_completion(chat_request, raw_request, return_cmpl_response=True)
|
|
|
|
async def create_embedding(self, request: EmbeddingRequest, raw_request: Request):
|
|
chat_request = ChatCompletionRequest.from_cmpl_request(request)
|
|
return await self.create_chat_completion(chat_request, raw_request, return_cmpl_response=True)
|
|
|
|
async def infer_handler(self, raw_request: Request):
|
|
body = await raw_request.json()
|
|
infer_requests = [RolloutInferRequest(**r) for r in body.get('infer_requests', [])]
|
|
rc_data = body.get('request_config')
|
|
request_config = RequestConfig(**rc_data) if rc_data else RequestConfig()
|
|
results = await asyncio.gather(*[self.infer_async(req, request_config) for req in infer_requests])
|
|
return results
|
|
|
|
def run(self):
|
|
args = self.args
|
|
self.jsonl_writer = JsonlWriter(args.result_path) if args.result_path else None
|
|
logger.info(f'model_list: {self._get_model_list()}')
|
|
uvicorn.run(
|
|
self.app,
|
|
host=args.host,
|
|
port=args.port,
|
|
ssl_keyfile=args.ssl_keyfile,
|
|
ssl_certfile=args.ssl_certfile,
|
|
log_level=args.log_level)
|
|
|
|
|
|
def deploy_main(args: Optional[Union[List[str], DeployArguments]] = None) -> None:
|
|
SwiftDeploy(args).main()
|
|
|
|
|
|
def is_accessible(port: int):
|
|
infer_client = InferClient(port=port)
|
|
try:
|
|
infer_client.get_model_list()
|
|
except ClientConnectorError:
|
|
return False
|
|
return True
|
|
|
|
|
|
def _deploy_main(args):
|
|
args._import_external_plugins()
|
|
return deploy_main(args)
|
|
|
|
|
|
@contextmanager
|
|
def run_deploy(args: DeployArguments, return_url: bool = False):
|
|
if isinstance(args, DeployArguments) and args.__class__.__name__ == 'DeployArguments':
|
|
deploy_args = args
|
|
else:
|
|
args_dict = asdict(args)
|
|
parameters = inspect.signature(DeployArguments).parameters
|
|
for k in list(args_dict.keys()):
|
|
if k not in parameters or args_dict[k] is None:
|
|
args_dict.pop(k)
|
|
deploy_args = DeployArguments(**args_dict)
|
|
|
|
mp = multiprocessing.get_context('spawn')
|
|
process = mp.Process(target=_deploy_main, args=(deploy_args, ))
|
|
process.start()
|
|
try:
|
|
while not is_accessible(deploy_args.port):
|
|
time.sleep(1)
|
|
yield f'http://127.0.0.1:{deploy_args.port}/v1' if return_url else deploy_args.port
|
|
finally:
|
|
process.terminate()
|
|
logger.info('The deployment process has been terminated.')
|