import asyncio
import os
import uuid
import zipfile
from datetime import datetime
from pathlib import Path
from typing import List, Sequence, Optional, Tuple, Type
import aiofiles
from fastapi import UploadFile
from sqlalchemy import func
from sqlmodel import select, delete
from sqlalchemy.ext.asyncio import AsyncSession
from common.biz_exception import BizException
from common.string_utils import snake_to_camel
from entity.device_model_relation import DeviceModelRelation
from entity.model import AlgoModel, AlgoModelCreate, AlgoModelUpdate, AlgoModelInfo
from common.global_thread_pool import GlobalThreadPool
from common.consts import NotifyChangeType
class ModelService:
def __init__(self, db: AsyncSession):
self.db = db
self.__model_change_callbacks = [] # 用于存储回调函数
self.thread_pool = GlobalThreadPool()
def register_change_callback(self, callback):
"""注册设备变化回调函数"""
self.__model_change_callbacks.append(callback)
def notify_change(self, algo_model_id, change_type):
"""当设备发生变化时，调用回调通知变化"""
for callback in self.__model_change_callbacks:
self.thread_pool.executor.submit(callback, algo_model_id, change_type)
async def get_model_list(self,
name: Optional[str] = None,
remark: Optional[str] = None,
) -> Sequence[AlgoModel]:
statement = self.model_query(name, remark)
results = await self.db.execute(statement)
return results.scalars().all()
async def get_model_page(self,
name: Optional[str] = None,
remark: Optional[str] = None,
offset: int = 0,
limit: int = 10
) -> Tuple[Sequence[AlgoModelInfo], int]:
statement = self.model_query(name, remark)
# 查询总记录数
total_statement = select(func.count()).select_from(statement.subquery())
total_result = await self.db.execute(total_statement)
total = total_result.scalar_one()
# 添加分页限制
statement = statement.offset(offset).limit(limit)
# 执行查询并返回结果
model_list = await self.db.execute(statement)
rows = model_list.scalars().all()
model_info_list: List[AlgoModelInfo] = []
if rows:
for model in rows:
model_info_list.append(AlgoModelInfo(
**model.dict(),
usage_status="使用中" if await self.get_model_usage(model.id) else "未使用"
))
return model_info_list, total # 返回分页数据和总数
def model_query(self, name, remark):
# 构建查询语句
statement = select(AlgoModel)
if name:
statement = statement.where(AlgoModel.name.like(f"%{name}%"))
if remark:
statement = statement.where(AlgoModel.remark.like(f"%{remark}%"))
return statement
async def process_zip(self, file: UploadFile):
model_dir = Path('weights/')
model_handle_dir = Path('model_handler/')
model_dir.mkdir(parents=True, exist_ok=True)
model_handle_dir.mkdir(parents=True, exist_ok=True)
# 支持的模型文件扩展名
SUPPORTED_MODEL_EXTENSIONS = {".pt", ".onnx", ".engine"}
# 临时保存上传文件
temp_path = Path(f"temp_upload_{uuid.uuid4()}.zip")
async with aiofiles.open(temp_path, "wb") as temp_file:
content = await file.read() # 异步读取上传文件内容
await temp_file.write(content)
model_file_path = None
handle_file_path = None
try:
# 使用异步方法读取压缩文件的内容
loop = asyncio.get_event_loop()
with zipfile.ZipFile(temp_path, 'r') as zip_ref:
# 获取压缩包文件列表
file_list = zip_ref.namelist()
# 检查是否有模型权重文件
model_file = next((f for f in file_list if Path(f).suffix in SUPPORTED_MODEL_EXTENSIONS), None)
if not model_file:
raise BizException(
status_code=400,
message=f"Model weight file ({', '.join(SUPPORTED_MODEL_EXTENSIONS)}) is required in the zip."
)
# 异步解压模型文件到模型目录
model_file_path = model_dir / model_file
await loop.run_in_executor(None, zip_ref.extract, model_file, model_dir)
# 检查是否有可选的 Python 脚本
handle_file = next((f for f in file_list if f.endswith(".py")), None)
if handle_file:
handle_file_path = model_handle_dir / handle_file
await loop.run_in_executor(None, zip_ref.extract, handle_file, model_handle_dir)
except zipfile.BadZipFile:
raise BizException(status_code=400, message="Invalid zip file.")
finally:
# 删除临时文件
loop = asyncio.get_event_loop()
await loop.run_in_executor(None, temp_path.unlink)
return str(model_file_path), str(handle_file_path) if handle_file_path else None
async def create_model(self, model_data: AlgoModelCreate, file: UploadFile):
await self.process_model_file(file, model_data)
model = AlgoModel.model_validate(model_data)
model.create_time = datetime.now()
model.update_time = datetime.now()
self.db.add(model)
await self.db.commit()
await self.db.refresh(model)
return model
async def process_model_file(self, file, model):
model_file_path, handle_file_path = await self.process_zip(file)
model.path = model_file_path
if handle_file_path:
model.handle_task = snake_to_camel(os.path.splitext(os.path.basename(handle_file_path))[0])
else:
model.handle_task = 'BaseModelHandler'
async def update_model(self, model_data: AlgoModelUpdate, file: UploadFile):
model = await self.get_model_by_id(model_data.id)
if not model:
return None
update_data = model_data.dict(exclude_unset=True)
for key, value in update_data.items():
setattr(model, key, value)
model.update_time = datetime.now()
if file:
await self.process_model_file(file, model)
self.db.add(model)
await self.db.commit()
await self.db.refresh(model)
self.notify_change(model.id, NotifyChangeType.MODEL_UPDATE)
return model
async def delete_model(self, model_id: int):
model = await self.get_model_by_id(model_id)
if not model:
return None
# 查询 device_model_relation 中是否存在启用的绑定关系
statement = (
select(DeviceModelRelation)
.where(DeviceModelRelation.algo_model_id == model_id)
.where(DeviceModelRelation.is_use == 1)
)
relation_in_use_exec = await self.db.execute(statement)
relation_in_use = relation_in_use_exec.scalars().first()
# 如果存在启用的绑定关系，提示无法删除
if relation_in_use:
raise BizException(message=f"模型 {model.name} 正在被设备使用，无法删除")
statement = delete(AlgoModel).where(AlgoModel.id == model_id)
await self.db.execute(statement)
await self.db.commit()
return model
async def get_models_in_use(self) -> Sequence[AlgoModel]:
"""获取所有在 device_model_relation 表里有启用绑定关系的模型信息"""
statement = (
select(AlgoModel)
.join(DeviceModelRelation, DeviceModelRelation.algo_model_id == AlgoModel.id)
.where(DeviceModelRelation.is_use == 1)
.group_by(AlgoModel.id)
)
results = await self.db.execute(statement)
return results.scalars().all()
async def get_model_usage(self, algo_model_id) -> bool:
statement = (
select(DeviceModelRelation)
.where(
DeviceModelRelation.is_use == 1,
DeviceModelRelation.algo_model_id == algo_model_id,
)
)
result = await self.db.execute(statement)
rows = result.all()
return len(rows) > 0
async def get_model_by_id(self, model_id):
result = await self.db.execute(select(AlgoModel).where(AlgoModel.id == model_id))
model = result.scalar_one_or_none()
return model
