[perf] Optimize model loading flow

Author: LittleMouse

api_server.py

@@ -57,25 +57,36 @@ async def auth_middleware(request: Request, call_next):
 class ModelDispatcher:
     def __init__(self):
         self.backends = {}
-        self.load_models()
+        self.llm_models = []
+        self.lock = asyncio.Lock()
 
-    def load_models(self):
-        for model_name, model_config in config.data["models"].items():
-            if model_config["type"] == "openai_proxy":
-                self.backends[model_name] = OpenAIProxyBackend(model_config)
-            elif model_config["type"] == "tcp_client":
-                self.backends[model_name] = LlmClientBackend(model_config)
-            elif model_config["type"] == "llama.cpp":
-                self.backends[model_name] = TestBackend(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)
-
-    def get_backend(self, model_name):
-        return self.backends.get(model_name)
+    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"):
+                    while len(self.llm_models) >= 2:
+                        oldest_model = self.llm_models.pop(0)
+                        old_instance = self.backends.pop(oldest_model, None)
+                        if old_instance:
+                            await old_instance.close()
+                    self.backends[model_name] = LlmClientBackend(model_config)
+                    self.llm_models.append(model_name)
+                elif model_config["type"] == "llama.cpp":
+                    self.backends[model_name] = TestBackend(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():
     global config
@@ -92,7 +103,7 @@ async def initialize():
 
 @app.post("/v1/chat/completions")
 async def chat_completions(request: Request, body: ChatCompletionRequest):
-    backend = _dispatcher.get_backend(body.model)
+    backend = await _dispatcher.get_backend(body.model)
     if not backend:
         raise HTTPException(
             status_code=400, 
@@ -156,7 +167,7 @@ async def create_completion(request: Request, body: CompletionRequest):
         stream=body.stream
     )
 
-    backend = _dispatcher.get_backend(chat_request.model)
+    backend = await _dispatcher.get_backend(chat_request.model)
     if not backend:
         raise HTTPException(status_code=400, detail=f"Unsupported model: {chat_request.model}")
 
@@ -215,7 +226,7 @@ async def convert_stream():
 async def create_speech(request: Request):
     try:
         request_data = await request.json()
-        backend = _dispatcher.get_backend(request_data.get("model"))
+        backend = await _dispatcher.get_backend(request_data.get("model"))
         if not backend:
             raise HTTPException(status_code=400, detail="Unsupported model")
 
@@ -243,7 +254,7 @@ async def create_transcription(
     response_format: str = Form("json")
 ):
     try:
-        backend = _dispatcher.get_backend(model)
+        backend = await _dispatcher.get_backend(model)
         if not backend:
             raise HTTPException(status_code=400, detail="Unsupported model")
 
@@ -273,7 +284,7 @@ async def create_translation(
     response_format: str = Form("json")
 ):
     try:
-        backend = _dispatcher.get_backend(model)
+        backend = await _dispatcher.get_backend(model)
         if not backend:
             raise HTTPException(status_code=400, detail="Unsupported model")
 

backend/llm_client_backend.py

@@ -72,9 +72,7 @@ async def _get_client(self, request):
                 await asyncio.wait_for(self._pool_lock.acquire(), timeout=timeout - (time.time() - start_time))
 
             if "memory_required" in self.config:
-                await self.memory_checker.check_memory(
-                    self.config["memory_required"]
-                )
+                await self.memory_checker.check_memory(self.config["memory_required"])
 
             self.logger.debug("Creating new LLM client")
             client = LLMClient(
@@ -117,6 +115,14 @@ async def _release_client(self, client):
             self._client_pool.append(client)
             self.logger.debug(f"Returned client to pool | ID:{id(client)}")
 
+    async def close(self):
+        async with self._pool_lock:
+            for client in self._client_pool:
+                client.exit()
+            self._client_pool.clear()
+            self._active_clients.clear()
+            self._inference_executor.shutdown(wait=True)
+
     async def inference_stream(self, query: str, base64_images: list, request: ChatCompletionRequest):
         client = await self._get_client(request)
         task = asyncio.current_task()

services/model_list.py

@@ -46,26 +46,38 @@ async def get_model_list(self, required_mem: int) -> None:
                             "object": f"{model_type}.setup",
                             "system_prompt": "You are a helpful assistant."
                         })
+                        if '-1.5B-' in mode:
+                            new_entry['memory_required'] = 1782579
+                            new_entry['pool_size'] = 1
+                        elif '-1B-' in mode:
+                            new_entry['memory_required'] = 1363148
+                            new_entry['pool_size'] = 2
+                        elif '-0.5B-' in mode:
+                            new_entry['memory_required'] = 560460
+                            new_entry['pool_size'] = 2
+
                     elif model_type == 'tts':
                         if 'melotts' in mode.lower():
                             obj = 'melotts.setup'
+                            new_entry['memory_required'] = 59764
                         else:
                             obj = 'tts.setup'
 
                         new_entry.update({
                             "response_format": "wav.base64",
-                            "object": "melotts.setup",
                             "object": obj
                         })
                     elif model_type == 'asr':
                         if 'whisper' in mode.lower():
                             obj = 'whisper.setup'
+                            if 'tiny' in mode:
+                                new_entry['memory_required'] = 289132
                         else:
                             obj = 'asr.setup'
                         new_entry.update({
                             "input": "pcm.base64",
                             "response_format": "asr.utf-8",
-                            "object": "whisper.setup"
+                            "object": obj
                         })
                     else:
                         continue
@@ -84,4 +96,4 @@ async def _get_model_list(self):
         return await loop.run_in_executor(
             None, 
             self._sys_client.model_list
-        ) 
+        )