🎈 perf: improve reranker performance

Author: hcd233

config/models.yaml.template

@@ -8,8 +8,12 @@ embedding:
 reranker:
   bge-reranker-v2-m3:
     path: BAAI/bge-reranker-v2-m3
-    batch_size: 4
+    batch_size: 4  # 增加batch_size可以提高GPU利用率，但需要更多内存
     max_seq_len: 512
+    # 性能优化建议：
+    # 1. 根据GPU内存调整batch_size：RTX 4090建议4-8，A100建议8-16
+    # 2. 使用GPU并发：多个GPU可以增加更多worker
+    # 3. 启用缓存以避免重复计算（如果applicable）
 
 llm:
   qwen2-7b-instruct-mlx:

src/api/router/v1/reranker.py

@@ -41,7 +41,8 @@ async def rerank(request: RerankRequest) -> RerankResponse:
             detail=f"Invalid model name: {request.model}",
         )
 
-    results = engine.invoke(request.query, request.documents)
+    # 使用引擎的异步方法进行GPU推理
+    results = await engine.async_invoke(request.query, request.documents)
     logger.debug(f"[Cohere Rerank] result: {results}")
 
     sorted_results = sorted(enumerate(results), key=lambda x: x[1]["relevent_score"], reverse=True)

src/engine/base.py

@@ -16,6 +16,10 @@ def invoke(self, *args, **kwargs) -> Any:
     @abstractmethod
     def stream(self, *args, **kwargs) -> Any:
         raise NotImplementedError(f"{self.__class__.__name__} does not implement `stream` method")
+    
+    @abstractmethod
+    def async_invoke(self, *args, **kwargs) -> Any:
+        raise NotImplementedError(f"{self.__class__.__name__} does not implement `async_invoke` method")
 
 
 class RerankerResult(TypedDict):

src/engine/reranker/sentence_transformer.py

@@ -1,7 +1,8 @@
-import torch
-
+import asyncio
+import concurrent.futures
 from typing import List
 
+import torch
 from sentence_transformers import CrossEncoder
 
 from src.config.arg import RerankerConfig
@@ -13,6 +14,7 @@
 
 class SentenceTransformerRerankerEngine(BaseEngine, RerankerConfig):
     model: CrossEncoder
+    _executor: concurrent.futures.ThreadPoolExecutor = None
 
     @classmethod
     def from_config(cls, config: RerankerConfig) -> "SentenceTransformerRerankerEngine":
@@ -21,6 +23,16 @@ def from_config(cls, config: RerankerConfig) -> "SentenceTransformerRerankerEngi
         logger.success(f"[RerankerEngine] load model from {config.path}")
         return cls(model=model, **config.model_dump())
 
+    @property
+    def executor(self) -> concurrent.futures.ThreadPoolExecutor:
+        """延迟初始化线程池执行器"""
+        if self._executor is None:
+            self._executor = concurrent.futures.ThreadPoolExecutor(
+                max_workers=4,  # 单个模型使用一个worker避免GPU内存冲突
+                thread_name_prefix=f"gpu_reranker_{self.alias}"
+            )
+        return self._executor
+
     def invoke(self, query: str, documents: List[str]) -> List[RerankerResult]:
         with torch.inference_mode():  # 使
             scores = self.model.predict(
@@ -36,5 +48,19 @@ def invoke(self, query: str, documents: List[str]) -> List[RerankerResult]:
 
         return [RerankerResult(index=index, relevent_score=score) for score, index in zip(scores, indexes)]
 
+    async def async_invoke(self, query: str, documents: List[str]) -> List[RerankerResult]:
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            self.executor,
+            self.invoke,
+            query,
+            documents
+        )
+
     def stream(self, query: str, documents: List[str]) -> List[RerankerResult]:
         raise NotImplementedError(f"{self.__class__.__name__} does not implement `stream` method")
+
+    def __del__(self):
+        """清理线程池资源"""
+        if hasattr(self, '_executor') and self._executor is not None:
+            self._executor.shutdown(wait=False)