update - embedding add exception and metrics

Author: rtp-llm

maga_transformer/access_logger/access_logger.py

@@ -1,5 +1,6 @@
 import os
 import logging
+from pydantic import BaseModel
 from typing import Any, Union, Dict
 
 from maga_transformer.access_logger.json_util import dump_json
@@ -57,14 +58,18 @@ def log_query_access(self, request: Union[Dict[str, Any], str], id: int) -> None
             access_log = PyAccessLog(request = request_log, response = response_log, id = id)
             self.query_logger.info(dump_json(access_log))
 
-    def log_success_access(self, request: Union[Dict[str, Any], str], response: Any, id: int) -> None:
+    def log_success_access(self, request: Union[Dict[str, Any], str, BaseModel], response: Any, id: int) -> None:
+        if isinstance(request, BaseModel):
+            request = request.model_dump()
         if not self.is_private_request(request):
             response_log = ResponseLog()
             if LOG_RESPONSE:
                 response_log.add_response(response)
             self.log_access(request, response_log, id)
 
-    def log_exception_access(self, request: Union[Dict[str, Any], str], exception: BaseException, id: int) -> None:
+    def log_exception_access(self, request: Union[Dict[str, Any], str, BaseModel], exception: BaseException, id: int) -> None:
+        if isinstance(request, BaseModel):
+            request = request.model_dump()
         response_log = ResponseLog()
         response_log.add_exception(exception)
         if not self.is_private_request(request):

maga_transformer/embedding/api_datatype.py

@@ -5,7 +5,7 @@
 
 class OpenAIEmbeddingRequest(PyDanticModelBase):
     input: Union[str, List[str]]
-    model: str
+    model: str = ""
     encoding_format: str = 'float'
     user: str = ""
     embedding_config: EmbeddingGenerateConfig = EmbeddingGenerateConfig()

maga_transformer/server/inference_server.py

@@ -1,6 +1,7 @@
 import os
 import json
 import time
+import copy
 import logging
 import logging.config
 import traceback
@@ -60,7 +61,6 @@ def start(self):
             self._openai_endpoint = None
             self._embedding_endpoint = None
             if self._inference_worker.model is not None and self._inference_worker.model.model_type == ModelType.EMBEDDING:
-                assert isinstance(self._inference_worker.model, AsyncModel), "only support embedding model in async mode"
                 self._embedding_endpoint = EmbeddingEndpoint(self._inference_worker.model)
             else:
                 self._openai_endpoint = OpenaiEndopoint(self._inference_worker.model)
@@ -111,25 +111,25 @@ async def stream_response(
             self._access_logger.log_exception_access(request, e, id)
             kmonitor.report(AccMetrics.ERROR_QPS_METRIC, 1)
             yield response_data_prefix + \
-                json.dumps(InferenceServer.handler_exceptions(e), ensure_ascii=False) + "\r\n\r\n"
+                json.dumps(InferenceServer.format_exception(e), ensure_ascii=False) + "\r\n\r\n"
 
     @staticmethod
-    def format_exception(errcode: int, message: str) -> Dict[str, Any]:
-        return {'error_code': errcode, "message": message}
+    def format_exception(e: Exception):
+        @staticmethod
+        def _format(errcode: int, message: str) -> Dict[str, Any]:
+            return {'error_code': errcode, "message": message}
 
-    @staticmethod
-    def handler_exceptions(e: Exception):
         if isinstance(e, FtRuntimeException):
-            return InferenceServer.format_exception(e.expcetion_type, e.message)
+            return _format(e.expcetion_type, e.message)
         elif isinstance(e, ConcurrencyException):
-            return InferenceServer.format_exception(ExceptionType.CONCURRENCY_LIMIT_ERROR, str(e))
+            return _format(ExceptionType.CONCURRENCY_LIMIT_ERROR, str(e))
         elif isinstance(e, LoraCountException) or isinstance(e, LoraPathException):
-            return InferenceServer.format_exception(ExceptionType.UPDATE_ERROR, str(e))
+            return _format(ExceptionType.UPDATE_ERROR, str(e))
         elif isinstance(e, Exception):
             error_msg = f'ErrorMsg: {str(e)} \n Traceback: {traceback.format_exc()}'
-            return InferenceServer.format_exception(ExceptionType.UNKNOWN_ERROR, error_msg)
+            return _format(ExceptionType.UNKNOWN_ERROR, error_msg)
         else:
-            return InferenceServer.format_exception(ExceptionType.UNKNOWN_ERROR, str(e))
+            return _format(ExceptionType.UNKNOWN_ERROR, str(e))
 
     def update(self, version_info: VersionInfo):
         id = self._atomic_count.increment()
@@ -146,7 +146,7 @@ def update(self, version_info: VersionInfo):
             self._access_logger.log_exception_access(version_info.__dict__, e, id)
             kmonitor.report(AccMetrics.ERROR_UPDATE_QPS_METRIC, 1)
             error_code = 500
-            rep = JSONResponse(self.handler_exceptions(e), status_code=error_code)
+            rep = JSONResponse(self.format_exception(e), status_code=error_code)
         return rep
 
     async def inference(self, req: Union[str,Dict[Any, Any]], raw_request: RawRequest):
@@ -165,17 +165,7 @@ async def _infer_wrap(self, req: Dict[Any, Any], raw_request: RawRequest, genera
         try:
             rep = await self._infer_impl(req, id, raw_request, generate_call)
         except Exception as e:
-            self._access_logger.log_exception_access(req, e, id)
-            if isinstance(e, ConcurrencyException):
-                kmonitor.report(AccMetrics.CONFLICT_QPS_METRIC)
-                error_code = 409
-            elif isinstance(e, asyncio.CancelledError):
-                kmonitor.report(AccMetrics.CANCAL_QPS_METRIC, 1)
-                error_code = 499
-            else:
-                error_code = 500
-                kmonitor.report(AccMetrics.ERROR_QPS_METRIC, 1)
-            rep = JSONResponse(self.handler_exceptions(e), status_code=error_code)
+            rep = self._handle_exception(req, e, id)
         return rep
 
     async def chat_completion(self, request: ChatCompletionRequest, raw_request: Request):
@@ -187,26 +177,45 @@ def generate_call():
         return await self._infer_wrap(request.model_dump(), raw_request, generate_call)
 
     async def embedding(self, request: Union[Dict[str, Any], str, OpenAIEmbeddingRequest], raw_request: Request):
+        id = self._atomic_count.increment()
+        kmonitor.report(AccMetrics.QPS_METRIC, 1)
         with self._controller:
             try:
                 assert self._embedding_endpoint is not None, "embedding pipeline should not be None"
                 result = await self._embedding_endpoint.embedding(request)
+                log_result = copy.copy(result)
+                # do not log result since too big
+                log_result.data = []
+                self._access_logger.log_success_access(request, log_result, id)
                 return JSONResponse(result.model_dump(exclude_none=True))
-            except FtRuntimeException:
-                raise
             except Exception as e:
-                raise FtRuntimeException(ExceptionType.UNKNOWN_ERROR, str(e))
-            
+                self._handle_exception(request, e, id)
+
     async def similarity(self, request: Union[Dict[str, Any], str, SimilarityRequest], raw_request: Request):
+        id = self._atomic_count.increment()
+        kmonitor.report(AccMetrics.QPS_METRIC, 1)
         with self._controller:
             try:
                 assert self._embedding_endpoint is not None, "embedding pipeline should not be None"
                 result = await self._embedding_endpoint.similarity(request)
+                self._access_logger.log_success_access(request, result.model_dump(exclude_none=True), id)
                 return JSONResponse(result.model_dump(exclude_none=True))
-            except FtRuntimeException:
-                raise
             except Exception as e:
-                raise FtRuntimeException(ExceptionType.UNKNOWN_ERROR, str(e))
+                self._handle_exception(request, e, id)
+
+    def _handle_exception(self, request: Union[Dict[str, Any], str, BaseModel], e: Exception, id: int):
+        self._access_logger.log_exception_access(request, e, id)
+        if isinstance(e, ConcurrencyException):
+            kmonitor.report(AccMetrics.CONFLICT_QPS_METRIC)
+            error_code = 409
+        elif isinstance(e, asyncio.CancelledError):
+            kmonitor.report(AccMetrics.CANCAL_QPS_METRIC, 1)
+            error_code = 499
+        else:
+            error_code = 500
+            kmonitor.report(AccMetrics.ERROR_QPS_METRIC, 1)
+        rep = JSONResponse(self.format_exception(e), status_code=error_code)
+        return rep
 
     async def _call_generate_with_report(self, generate_call: Callable[[], CompleteResponseAsyncGenerator]):
         async def __gen_response_with_report(start_time, response_generator):
@@ -283,4 +292,4 @@ def tokenizer_encode(self, req: Union[str,Dict[Any, Any]]):
             response = TokenizerEncodeResponse(token_ids=token_ids, tokens=tokens)
             return JSONResponse(content=response.model_dump(exclude_none=True))
         except Exception as e:
-            return JSONResponse(self.handler_exceptions(e), status_code=500)
+            return JSONResponse(self.format_exception(e), status_code=500)

maga_transformer/server/inference_worker.py

@@ -18,7 +18,7 @@
 from maga_transformer.config.exceptions import FtRuntimeException, ExceptionType
 from maga_transformer.models.base_model import GenerateResponse, GenerateConfig
 from maga_transformer.model_factory import ModelFactory, AsyncModel
-from maga_transformer.structure.request_extractor import RequestExtractor
+from maga_transformer.structure.request_extractor import RequestExtractor, Request
 
 from pydantic import BaseModel
 
@@ -52,7 +52,7 @@ def __init__(self) -> None:
         if not torch.cuda.is_available():
             raise Exception("GPU not found")
 
-        self.model = ModelFactory.create_from_env()
+        self.model: AsyncModel = ModelFactory.create_from_env()
         self.pipeline = Pipeline(self.model, self.model.tokenizer)
         logging.info("Load model done.")
 
@@ -78,7 +78,7 @@ def inference(self, **kwargs: Any) -> CompleteResponseAsyncGenerator:
         return CompleteResponseAsyncGenerator(response_generator, complete_response_collect_func)
 
 
-    def _inference(self, request, **kwargs):
+    def _inference(self, request: Request, **kwargs: Any):
         if len(request.input_texts) > 1 or request.batch_infer or request.num_return_sequences > 0:
             num_return_sequences = request.generate_configs[0].num_return_sequences
             generators = [self._yield_generate(text, images, generate_config=generate_config, **kwargs)
@@ -127,7 +127,7 @@ def is_streaming(self, req: Dict[str, Any]):
         return RequestExtractor.is_streaming(req) or req.get('stream', False)
 
     def update(self, version_info: VersionInfo):
-        lora_infos = dict()
+        lora_infos: Dict[str, Any] = dict()
         if version_info.peft_info != None:
             lora_infos = version_info.peft_info.get("lora_info", {})
         return self.model.update(lora_infos)