from datetime import datetime
|
import logging
|
import os
|
from concurrent.futures import ThreadPoolExecutor
|
import threading
|
from logging.handlers import RotatingFileHandler
|
from langchain_openai import ChatOpenAI
|
from pymilvus import connections, Collection
|
import qwen_vllm
|
|
|
def process_task(detect,image_id,ragurl,ragmode,max_tokens):
|
image_id = detect.tark_do(image_id,ragurl,ragmode,max_tokens)
|
print(f"完成任务 {image_id} (耗时: {10:.2f}s)")
|
return image_id
|
|
|
class qwen_vllm_thread:
|
def __init__(self, max_workers,camera_id,config):
|
self.executor = ThreadPoolExecutor(max_workers=max_workers)
|
self.semaphore = threading.Semaphore(max_workers)
|
self.max_workers = max_workers
|
# 初始化Milvus集合
|
connections.connect("default", host=config.get("milvusurl"), port=config.get("milvusport"))
|
# 加载集合
|
self.collection = Collection(name="smartobject")
|
self.collection.load()
|
|
# 初始化LangChain客户端
|
self.llm = ChatOpenAI(
|
model=config.get("vllmmode"), # 与--served-model-name一致
|
temperature=0.7,
|
max_tokens=200,
|
base_url=config.get("vllmurl"),
|
api_key="EMPTY"
|
)
|
self.config = config
|
# 初始化
|
self.detect = qwen_vllm.detect_tasks()
|
self.detect._thread_name = f"{camera_id}" # 设置名称
|
self.detect.init_logging(f"{camera_id}")
|
# 初始化Milvus集合
|
self.detect.collection = self.collection
|
self.detect.llm = self.llm
|
# 创建实例专属logger
|
self.logger = logging.getLogger(f"{self.__class__}_{id(self)}")
|
self.logger.setLevel(logging.INFO)
|
# 避免重复添加handler
|
if not self.logger.handlers:
|
handler = RotatingFileHandler(
|
filename=os.path.join("logs", 'thread_log.log'),
|
maxBytes=10 * 1024 * 1024,
|
backupCount=3,
|
encoding='utf-8'
|
)
|
formatter = logging.Formatter(
|
'%(asctime)s - %(filename)s:%(lineno)d - %(funcName)s() - %(levelname)s: %(message)s'
|
)
|
handler.setFormatter(formatter)
|
self.logger.addHandler(handler)
|
|
def submit(self,res_a):
|
# 尝试获取信号量（非阻塞）
|
acquired = self.semaphore.acquire(blocking=False)
|
|
if not acquired:
|
self.logger.info(f"线程池已满，等待空闲线程... (当前活跃: {self.max_workers - self.semaphore._value}/{self.max_workers})")
|
# 阻塞等待直到有可用线程
|
self.semaphore.acquire(blocking=True)
|
|
future = self.executor.submit(self._wrap_task, res_a)
|
future.add_done_callback(self._release_semaphore)
|
return future
|
|
def _wrap_task(self, res):
|
try:
|
#self.logger.info(f"处理: { res['id']}开始")
|
current_time = datetime.now()
|
image_id = self.detect.tark_do(res, self.config.get("ragurl"), self.config.get("ragmode"), self.config.get("max_tokens"))
|
self.logger.info(f"处理: { res['id']}完毕{image_id}:{datetime.now() - current_time}")
|
return image_id
|
except Exception as e:
|
self.logger.info(f"任务 { res['id']} 处理出错: {e}")
|
raise
|
|
def _release_semaphore(self, future):
|
self.semaphore.release()
|
self.logger.info(f"释放线程 (剩余空闲: {self.semaphore._value}/{self.max_workers})")
|
|
def shutdown(self):
|
self.executor.shutdown()
|
