[perf] perf model load process

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()

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, 

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