From 736eb3406fdf210515754d08f14b0691fd625bfc Mon Sep 17 00:00:00 2001 From: LittleMouse Date: Thu, 20 Nov 2025 14:50:10 +0800 Subject: [PATCH] [perf] perf model load process --- api_server.py | 139 +++++++++++++++++++++++---------------- services/memory_check.py | 14 ++-- services/model_list.py | 1 + 3 files changed, 94 insertions(+), 60 deletions(-) diff --git a/api_server.py b/api_server.py index e04d4ef..5dde612 100644 --- a/api_server.py +++ b/api_server.py @@ -6,7 +6,6 @@ import logging import time import json import asyncio - from pydub import AudioSegment from fastapi import FastAPI, Request, HTTPException, File, Form, UploadFile from fastapi.responses import JSONResponse, StreamingResponse @@ -20,7 +19,7 @@ from backend import ( CompletionRequest, Message, ) - +from services.memory_check import MemoryChecker from services.model_list import GetModelList logging.basicConfig( @@ -31,7 +30,6 @@ logging.basicConfig( ] ) logger = logging.getLogger("api") - app = FastAPI(title="OpenAI Compatible API Server") class Config: @@ -60,50 +58,91 @@ async def auth_middleware(request: Request, call_next): class ModelDispatcher: def __init__(self): self.backends = {} - self.llm_models = set() - self.asr_models = set() - self.tts_models = set() + self.memory_checker = MemoryChecker( + host=config.data["server"]["host"], + port=config.data["server"]["port"] + ) self.lock = asyncio.Lock() + async def _ensure_memory_available(self, required_mem: int): + if required_mem <= 0: + return + + try: + cmm_info = await self.memory_checker.get_cmminfo() + remain_mem = cmm_info["data"]["remain"] + + logger.debug(f"Memory Check | Required: {required_mem} | Available: {remain_mem}") + + if remain_mem >= required_mem: + return + + needed_mem = required_mem - remain_mem + reclaimable_mem = 0 + models_to_unload = [] + + for model_name, backend in self.backends.items(): + if reclaimable_mem >= needed_mem: + break + + model_conf = config.data["models"].get(model_name, {}) + mem_used = model_conf.get("memory_required", 0) + + reclaimable_mem += mem_used + models_to_unload.append(model_name) + + if remain_mem + reclaimable_mem < required_mem: + total_reclaimable = sum([config.data["models"].get(m, {}).get("memory_required", 0) for m in self.backends]) + raise HTTPException( + status_code=503, + detail=f"Insufficient Memory Resource. Required: {required_mem}, " + f"Available: {remain_mem}, Total Reclaimable: {total_reclaimable}. " + f"Cannot satisfy request even after unloading." + ) + + for model_name in models_to_unload: + logger.info(f"Unloading model '{model_name}' to free memory...") + backend = self.backends.pop(model_name) + if backend: + await backend.close() + + # await asyncio.sleep(0.1) + + except Exception as e: + if isinstance(e, HTTPException): + raise e + logger.error(f"Memory management error: {str(e)}") + raise HTTPException(status_code=500, detail=f"Memory check failed: {str(e)}") + async def get_backend(self, model_name): async with self.lock: - if model_name not in self.backends: - model_config = config.data["models"].get(model_name) - if model_config is None: - return None - if model_config["type"] == "openai_proxy": - self.backends[model_name] = OpenAIProxyBackend(model_config) - elif model_config["type"] in ("llm", "vlm"): - if model_name not in self.llm_models: - for old_model_name in list(self.llm_models): - old_instance = self.backends.pop(old_model_name, None) - if old_instance: - await old_instance.close() - self.llm_models.clear() - self.backends[model_name] = LlmClientBackend(model_config) - self.llm_models.add(model_name) - elif model_config["type"] == "vision_model": - self.backends[model_name] = VisionModelBackend(model_config) - elif model_config["type"] == "tts": - if model_name not in self.tts_models: - for old_model_name in list(self.tts_models): - old_instance = self.backends.pop(old_model_name, None) - if old_instance: - await old_instance.close() - self.tts_models.clear() - self.backends[model_name] = TtsClientBackend(model_config) - self.tts_models.add(model_name) - elif model_config["type"] == "asr": - if model_name not in self.asr_models: - for old_model_name in list(self.asr_models): - old_instance = self.backends.pop(old_model_name, None) - if old_instance: - await old_instance.close() - self.asr_models.clear() - self.backends[model_name] = ASRClientBackend(model_config) - self.asr_models.add(model_name) - else: - return None + if model_name in self.backends: + backend = self.backends.pop(model_name) + self.backends[model_name] = backend + return backend + + model_config = config.data["models"].get(model_name) + if model_config is None: + return None + + required_mem = model_config.get("memory_required", 0) + await self._ensure_memory_available(required_mem) + + logger.info(f"Loading model: {model_name} (Mem Required: {required_mem})") + + if model_config["type"] == "openai_proxy": + self.backends[model_name] = OpenAIProxyBackend(model_config) + elif model_config["type"] in ("llm", "vlm"): + self.backends[model_name] = LlmClientBackend(model_config) + elif model_config["type"] == "vision_model": + self.backends[model_name] = VisionModelBackend(model_config) + elif model_config["type"] == "tts": + self.backends[model_name] = TtsClientBackend(model_config) + elif model_config["type"] == "asr": + self.backends[model_name] = ASRClientBackend(model_config) + else: + return None + return self.backends.get(model_name) async def initialize(): @@ -156,7 +195,6 @@ async def chat_completions(request: Request, body: ChatCompletionRequest): raise finally: logger.debug("Stream connection closed") - return StreamingResponse( format_stream(), media_type="text/event-stream" @@ -238,41 +276,34 @@ async def create_completion(request: Request, body: CompletionRequest): raise HTTPException(status_code=500, detail=str(e)) @app.post("/v1/audio/speech") -async def create_speech( - request: Request, -): +async def create_speech(request: Request): try: request_data = await request.json() model = request_data.get("model") voice = request_data.get("voice", "prompt_data") response_format = request_data.get("response_format", "mp3") - if not model: raise HTTPException( status_code=400, detail="Model is required for speech generation" ) - backend = await _dispatcher.get_backend(model) if not backend: raise HTTPException( status_code=400, detail=f"Unsupported model: {model}" ) - input_text = request_data.get("input") if not input_text: raise HTTPException( status_code=400, detail="Input text is required for speech generation" ) - audio_stream = backend.generate_speech( input_text=input_text, voice=voice, format=response_format ) - return StreamingResponse( audio_stream, media_type=f"audio/{response_format}", @@ -296,7 +327,6 @@ async def create_transcription( status_code=400, detail=f"Unsupported model: {model}" ) - try: audio_data = await file.read() audio = AudioSegment.from_file(io.BytesIO(audio_data), format=file.filename.split('.')[-1]) @@ -351,14 +381,12 @@ async def create_translation( backend = await _dispatcher.get_backend(model) if not backend: raise HTTPException(status_code=400, detail="Unsupported model") - audio_data = await file.read() translation = await backend.create_translation( audio_data, prompt=prompt ) - return JSONResponse(content={ "text": translation, "task": "translate", @@ -390,5 +418,4 @@ async def list_models(): if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000) - logging.getLogger().handlers[0].flush() \ No newline at end of file diff --git a/services/memory_check.py b/services/memory_check.py index 0c0cfc9..0b16e42 100644 --- a/services/memory_check.py +++ b/services/memory_check.py @@ -9,13 +9,17 @@ class MemoryChecker: self.port = port self.logger = logging.getLogger("memory_check") self._sys_client: Optional[SYSClient] = None + + def _ensure_client(self): + if not self._sys_client: + self.logger.debug("Initializing SYSClient...") + self._sys_client = SYSClient(host=self.host, port=self.port) async def check_memory(self, required_mem: int) -> None: try: - if not self._sys_client: - self._sys_client = SYSClient(host=self.host, port=self.port) + self._ensure_client() - cmm_info = await self._get_cmminfo() + cmm_info = await self.get_cmminfo() remain_mem = cmm_info["data"]["remain"] self.logger.debug(f"Memory check - Required: {required_mem}, Available: {remain_mem}") @@ -29,7 +33,9 @@ class MemoryChecker: self.logger.error(f"Memory check failed: {str(e)}") raise - async def _get_cmminfo(self): + async def get_cmminfo(self): + self._ensure_client() + loop = asyncio.get_event_loop() return await loop.run_in_executor( None, diff --git a/services/model_list.py b/services/model_list.py index ac86139..c81bfce 100644 --- a/services/model_list.py +++ b/services/model_list.py @@ -110,6 +110,7 @@ class GetModelList: obj = 'melotts.setup' new_entry['memory_required'] = 59764 new_entry['sample_rate'] = 16000 + new_entry['max_context_length'] = 32768 elif 'cosyvoice' in mode.lower(): obj = 'cosy_voice.setup' new_entry['memory_required'] = 1185772