[refactor] Refactor asr_client_backend

Author: LittleMouse

api_server.py

@@ -106,7 +106,7 @@ async def chat_completions(request: Request, body: ChatCompletionRequest):
     backend = await _dispatcher.get_backend(body.model)
     if not backend:
         raise HTTPException(
-            status_code=400, 
+            status_code=400,
             detail=f"Unsupported model: {body.model}"
         )
 
@@ -250,19 +250,20 @@ async def create_transcription(
     prompt: str = Form(""),
     response_format: str = Form("json")
 ):
-    try:
-        backend = await _dispatcher.get_backend(model)
-        if not backend:
-            raise HTTPException(status_code=400, detail="Unsupported model")
+    backend = await _dispatcher.get_backend(model)
+    if not backend:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Unsupported model: {model}"
+        )
 
+    try:
         audio_data = await file.read()
-        
         transcription = await backend.create_transcription(
             audio_data,
             language=language,
             prompt=prompt
         )
-
         return JSONResponse(content={
             "text": transcription,
             "task": "transcribe",

backend/asr_client_backend.py

@@ -1,69 +1,106 @@
-from .base_model_backend import BaseModelBackend
-from client.asr_client import ASRClient
+import time
 import asyncio
+import weakref
 import base64
 import logging
+from .base_model_backend import BaseModelBackend
+from client.asr_client import ASRClient
 from concurrent.futures import ThreadPoolExecutor
-
-logger = logging.getLogger("api.asr")
+from services.memory_check import MemoryChecker
 
 class ASRClientBackend(BaseModelBackend):
-    POOL_SIZE = 1
-    SUPPORTED_FORMATS = ["json", "text", "srt", "verbose_json"]
-    
     def __init__(self, model_config):
         super().__init__(model_config)
-        self._executor = ThreadPoolExecutor(max_workers=self.POOL_SIZE)
-        self.clients = []
-        self._lock = asyncio.Lock()
+        self._client_pool = []
+        self._active_clients = {}
+        self._pool_lock = asyncio.Lock()
+        self.logger = logging.getLogger("api.asr")
+        self.POOL_SIZE = 1
+        self._inference_executor = ThreadPoolExecutor(max_workers=self.POOL_SIZE)
+        self._active_tasks = weakref.WeakSet()
+        self.memory_checker = MemoryChecker(
+            host=self.config["host"],
+            port=self.config["port"]
+        )
 
-    async def create_transcription(self, audio_data: bytes, language: str = "zh", prompt: str = "") -> str:
-        client = await self._get_client()
+    async def _get_client(self):
         try:
-            audio_b64 = base64.b64encode(audio_data).decode('utf-8')
-            return await self._inference_stream(client, audio_b64)
+            await asyncio.wait_for(self._pool_lock.acquire(), timeout=30.0)
+            
+            start_time = time.time()
+            timeout = 30.0
+            retry_interval = 3
+
+            while True:
+                if self._client_pool:
+                    client = self._client_pool.pop()
+                    return client
+                
+                for task in self._active_tasks:
+                    task.cancel()
+                
+                
+                self._pool_lock.release()
+                await asyncio.sleep(retry_interval)
+                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"])
+                client = ASRClient(
+                    host=self.config["host"],
+                    port=self.config["port"]
+                )
+                self._active_clients[id(client)] = client
+
+                loop = asyncio.get_event_loop()
+                await loop.run_in_executor(
+                    None,
+                    client.setup,
+                    "whisper.setup",
+                    {
+                        "model": self.config["model_name"],
+                        "response_format": "asr.utf-8",
+                        "input": "whisper.base64",
+                        "language": "zh",
+                        "enoutput": True
+                    }
+                )
+                return client
+        except asyncio.TimeoutError:
+            raise RuntimeError("Server busy, please try again later.")
         finally:
-            await self._release_client(client)
+            if self._pool_lock.locked():
+                self._pool_lock.release()
 
-    async def _inference_stream(self, client, audio_b64: str) -> str:
+    async def _release_client(self, client):
+        async with self._pool_lock:
+            self._client_pool.append(client)
+ 
+    async def _inference(self, client, audio_b64: str):
         loop = asyncio.get_event_loop()
-        full_text = ""
         for chunk in await loop.run_in_executor(
-            self._executor,
-            client.inference_stream,
+            self._inference_executor,
+            client.inference,
             audio_b64,
             "asr.base64"
         ):
-            full_text += chunk
-        return full_text
-
-    async def _get_client(self):
-        async with self._lock:
-            if self.clients:
-                return self.clients.pop()
-            
-            if len(self.clients) >= self.POOL_SIZE:
-                raise RuntimeError("ASR connection pool exhausted")
-                
-            client = ASRClient(
-                host=self.config["host"],
-                port=self.config["port"]
-            )
-            
-            await asyncio.get_event_loop().run_in_executor(
-                self._executor,
-                client.setup,
-                "whisper.setup",
-                {
-                    "model": self.config["model_name"],
-                    "response_format": "asr.utf-8",
-                    "input": "whisper.base64",
-                    "language": "zh",
-                    "enoutput": True
-                }
-            )
-            return client
+            full_result = chunk
+        return full_result
 
-    async def _release_client(self, client):
-        async with self._lock:
-            self.clients.append(client)
+    async def create_transcription(self, audio_data: bytes, language: str = "zh", prompt: str = "") -> str:
+        client = await self._get_client()
+        task = asyncio.current_task()
+        self._active_tasks.add(task)
+        try:
+            audio_b64 = base64.b64encode(audio_data).decode('utf-8')
+            return await self._inference(client, audio_b64)
+        except asyncio.CancelledError:
+            self.logger.warning("Inference task cancelled, stopping...")
+            client.stop_inference()
+            raise
+        except Exception as e:
+            self.logger.error(f"Inference error: {str(e)}")
+            raise RuntimeError(f"[ERROR: {str(e)}")
+        finally:
+            self._active_tasks.discard(task)
+            await self._release_client(client)

backend/llm_client_backend.py

@@ -95,7 +95,7 @@ async def _get_client(self, request):
 
             loop = asyncio.get_event_loop()
             await loop.run_in_executor(
-                None, 
+                None,
                 lambda: client.setup(
                     self.config["object"],
                     {

client/asr_client.py

@@ -5,7 +5,6 @@
 from typing import Generator
 import logging
 import threading
-import base64
 
 logger = logging.getLogger("asr_client")
 logger.setLevel(logging.DEBUG)
@@ -65,7 +64,7 @@ def setup(self, object: str, model_config: dict) -> dict:
         request_id = self._send_request("setup", object, model_config)
         return self._wait_response(request_id)
 
-    def inference_stream(self, query: str, object_type: str = "asr.base64") -> Generator[str, None, None]:
+    def inference(self, query: str, object_type: str = "asr.base64") -> Generator[str, None, None]:
         request_id = self._send_request("inference", object_type, query)
 
         while True:
@@ -100,21 +99,4 @@ def _wait_response(self, request_id: str) -> dict:
     def connect(self):
         with self._lock:
             if not self.sock:
-                self._connect()
-
-    def create_transcription(self, audio_data: bytes, language: str = "zh") -> str:
-        audio_b64 = base64.b64encode(audio_data).decode('utf-8')
-        
-        self.setup("whisper.setup", {
-            "model": "whisper-tiny",
-            "response_format": "asr.utf-8",
-            "input": "whisper.base64",
-            "language": language,
-            "enoutput": True,
-        })
-        
-        full_text = ""
-        for chunk in self.inference_stream(audio_b64, object_type="asr.base64"):
-            full_text += chunk
-            
-        return full_text
+                self._connect()