254 lines
9.4 KiB
Python
254 lines
9.4 KiB
Python
import asyncio
|
||
import os
|
||
import shutil
|
||
import uuid
|
||
import logging
|
||
from datetime import datetime, timedelta, timezone
|
||
from enum import Enum
|
||
from typing import Dict, Optional, Callable
|
||
|
||
from sqlalchemy import (
|
||
create_engine, Column, String, Enum as SAEnum, Float,
|
||
Integer, Text, DateTime, Boolean
|
||
)
|
||
from sqlalchemy.ext.declarative import declarative_base
|
||
from sqlalchemy.orm import sessionmaker, Session
|
||
|
||
from fastapi import BackgroundTasks
|
||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||
|
||
from configs.translate_config import *
|
||
|
||
# 初始化数据库引擎和会话
|
||
engine = create_engine(
|
||
SQLALCHEMY_DATABASE_URI,
|
||
connect_args={"check_same_thread": False}
|
||
)
|
||
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)
|
||
Base = declarative_base()
|
||
|
||
# 日志设置
|
||
logger = logging.getLogger(__name__)
|
||
logger.setLevel(logging.INFO)
|
||
|
||
# 时区设置（北京时间 UTC+8）
|
||
LOCAL_TZ = timezone(timedelta(hours=8))
|
||
|
||
# ----------------------
|
||
# 任务状态枚举
|
||
# ----------------------
|
||
class TaskStatusEnum(str, Enum):
|
||
"""
|
||
任务状态枚举：
|
||
QUEUED - 已入队，等待处理
|
||
PROCESSING - 正在处理
|
||
COMPLETED - 已完成
|
||
FAILED - 失败
|
||
CANCELLED - 已取消
|
||
"""
|
||
QUEUED = "queued"
|
||
PROCESSING = "processing"
|
||
COMPLETED = "completed"
|
||
FAILED = "failed"
|
||
CANCELLED = "cancelled"
|
||
|
||
# ----------------------
|
||
# TranslationTask ORM 模型
|
||
# ----------------------
|
||
class TranslationTask(Base):
|
||
"""数据库表 file_translate_tasks 对应的 ORM 模型"""
|
||
__tablename__ = "file_translate_tasks"
|
||
|
||
id = Column(String, primary_key=True, index=True) # 任务 ID
|
||
filename = Column(String, nullable=False) # 原始文件名
|
||
src_lang = Column(String, nullable=False) # 源语言
|
||
dst_lang = Column(String, nullable=False) # 目标语言
|
||
is_dual = Column(Boolean, default=True) # 是否双语模式
|
||
file_path = Column(String, nullable=False) # 原文文件路径
|
||
output_path = Column(String, nullable=True) # 翻译后文件路径
|
||
status = Column(SAEnum(TaskStatusEnum), default=TaskStatusEnum.QUEUED) # 当前状态
|
||
progress = Column(Float, default=0.0) # 进度百分比
|
||
retry_count = Column(Integer, default=0) # 已重试次数
|
||
error_msg = Column(Text, nullable=True) # 错误信息
|
||
created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(LOCAL_TZ)) # 创建时间
|
||
updated_at = Column(DateTime(timezone=True), default=lambda: datetime.now(LOCAL_TZ), onupdate=lambda: datetime.now(LOCAL_TZ)) # 最后更新时间
|
||
|
||
# 创建数据库表
|
||
Base.metadata.create_all(bind=engine)
|
||
|
||
# ----------------------
|
||
# TaskManager 定义
|
||
# ----------------------
|
||
class TaskManager:
|
||
def __init__(self, translate_fn: Callable):
|
||
self.translate_fn = translate_fn # 翻译函数引用
|
||
# 设置调度器使用本地时区，避免时间偏移
|
||
self.scheduler = AsyncIOScheduler(timezone=LOCAL_TZ)
|
||
# 定时扫描并恢复卡住任务
|
||
self.scheduler.add_job(
|
||
self.recover_stuck_tasks,
|
||
'interval', minutes=RECOVERY_INTERVAL,
|
||
id='recover_jobs', replace_existing=True
|
||
)
|
||
# 存储运行中任务的取消事件映射
|
||
self._cancel_events: Dict[str, asyncio.Event] = {}
|
||
# self.scheduler.start()
|
||
logger.info("任务管理器已初始化，恢复调度器已启动，扫描间隔：%s 分钟", RECOVERY_INTERVAL)
|
||
|
||
def start(self):
|
||
self.scheduler.start()
|
||
logger.info("恢复调度器已启动")
|
||
|
||
def shutdown(self):
|
||
self.scheduler.shutdown()
|
||
logger.info("恢复调度器已关闭")
|
||
|
||
def db(self) -> Session:
|
||
"""获取数据库会话"""
|
||
return SessionLocal()
|
||
|
||
def generate_task_id(self) -> str:
|
||
"""生成唯一任务 ID"""
|
||
return str(uuid.uuid4())
|
||
|
||
def add_task(
|
||
self,
|
||
filename: str,
|
||
file_path: str,
|
||
src_lang: str,
|
||
dst_lang: str,
|
||
is_dual: bool,
|
||
background_tasks: Optional[BackgroundTasks],
|
||
task_id: Optional[str] = None,
|
||
) -> str:
|
||
"""添加新翻译任务并入队"""
|
||
db = self.db()
|
||
if not task_id:
|
||
task_id = self.generate_task_id() # 如果外部未传入，则生成新ID
|
||
task = TranslationTask(
|
||
id=task_id,
|
||
filename=filename,
|
||
file_path=file_path,
|
||
src_lang=src_lang,
|
||
dst_lang=dst_lang,
|
||
is_dual=is_dual,
|
||
status=TaskStatusEnum.QUEUED,
|
||
)
|
||
db.add(task)
|
||
db.commit()
|
||
db.close()
|
||
|
||
if background_tasks:
|
||
background_tasks.add_task(self.start_task, task_id) # 加入后台执行队列
|
||
return task_id
|
||
|
||
async def start_task(self, task_id: str):
|
||
"""执行翻译任务的后台入口，支持取消"""
|
||
# 创建并注册取消事件
|
||
cancel_event = asyncio.Event()
|
||
self._cancel_events[task_id] = cancel_event
|
||
|
||
db = self.db()
|
||
task = db.query(TranslationTask).get(task_id)
|
||
# 只有 QUEUED 状态可执行
|
||
if not task or task.status != TaskStatusEnum.QUEUED:
|
||
db.close()
|
||
self._cancel_events.pop(task_id, None)
|
||
return
|
||
|
||
# 标记为处理中并更新更新时间
|
||
task.status = TaskStatusEnum.PROCESSING
|
||
task.updated_at = datetime.now(LOCAL_TZ)
|
||
db.commit()
|
||
|
||
try:
|
||
# 调用翻译函数，需支持 cancel_event
|
||
await self.translate_fn(task_id, task, cancel_event)
|
||
task.status = TaskStatusEnum.COMPLETED
|
||
task.progress = 100.0 # 完成进度
|
||
except asyncio.CancelledError:
|
||
# 如果收到取消信号，标记取消
|
||
task.status = TaskStatusEnum.CANCELLED
|
||
task.error_msg = "用户已取消"
|
||
except Exception as e:
|
||
# 其他异常按照重试逻辑处理
|
||
logger.error("任务 %s 翻译失败：%s", task_id, e)
|
||
task.retry_count += 1
|
||
if task.retry_count <= MAX_RETRIES:
|
||
task.status = TaskStatusEnum.QUEUED
|
||
task.error_msg = "意外中断，正在重试"
|
||
run_date = datetime.now(LOCAL_TZ) + timedelta(seconds=RETRY_DELAY)
|
||
self.scheduler.add_job(
|
||
self.start_task, 'date', run_date=run_date, args=[task_id]
|
||
)
|
||
else:
|
||
task.status = TaskStatusEnum.FAILED
|
||
task.error_msg = "意外中断，重试次数已达上限"
|
||
finally:
|
||
# 无论完成、取消或失败，都清理取消事件并更新数据库
|
||
task.updated_at = datetime.now(LOCAL_TZ)
|
||
self._cancel_events.pop(task_id, None)
|
||
db.commit()
|
||
db.close()
|
||
|
||
def get_task(self, task_id: str) -> Optional[TranslationTask]:
|
||
"""获取任务详情"""
|
||
db = self.db()
|
||
task = db.query(TranslationTask).get(task_id)
|
||
db.close()
|
||
return task
|
||
|
||
def cancel_task(self, task_id: str) -> bool:
|
||
"""取消正在 QUEUED 或 PROCESSING 状态的任务并触发取消事件"""
|
||
db = self.db()
|
||
task = db.query(TranslationTask).get(task_id)
|
||
if not task:
|
||
db.close()
|
||
return False
|
||
# 标记为取消
|
||
if task.status in {TaskStatusEnum.QUEUED, TaskStatusEnum.PROCESSING}:
|
||
task.status = TaskStatusEnum.CANCELLED
|
||
task.updated_at = datetime.now(LOCAL_TZ)
|
||
db.commit()
|
||
# 同时清理文件目录
|
||
if os.path.exists(os.path.dirname(task.file_path)):
|
||
shutil.rmtree(os.path.dirname(task.file_path), ignore_errors=True)
|
||
db.close()
|
||
# 触发翻译协程中的取消事件
|
||
if task_id in self._cancel_events:
|
||
self._cancel_events[task_id].set()
|
||
return True
|
||
|
||
def recover_stuck_tasks(self):
|
||
"""扫描并恢复处理超时或卡住的任务"""
|
||
logger.info(">>> 正在扫描翻译任务。。。")
|
||
db = self.db()
|
||
cutoff = datetime.now(LOCAL_TZ) - PROCESSING_TIMEOUT
|
||
logger.info(f"cutoff 时间为：{cutoff}")
|
||
stuck_tasks = db.query(TranslationTask).filter(
|
||
TranslationTask.status.in_([
|
||
TaskStatusEnum.PROCESSING,
|
||
TaskStatusEnum.QUEUED,
|
||
TaskStatusEnum.FAILED
|
||
]),
|
||
TranslationTask.updated_at > cutoff
|
||
).all()
|
||
logger.info(">>> 找到 %s 条卡住的翻译任务", len(stuck_tasks))
|
||
for task in stuck_tasks:
|
||
logger.info("正在恢复卡住的任务 %s", task.id)
|
||
if task.retry_count < MAX_RETRIES:
|
||
task.retry_count += 1
|
||
task.status = TaskStatusEnum.QUEUED
|
||
task.error_msg = "意外中断，正在重试"
|
||
run_date = datetime.now(LOCAL_TZ) + timedelta(seconds=RETRY_DELAY)
|
||
self.scheduler.add_job(
|
||
self.start_task,
|
||
trigger='date',
|
||
run_date=run_date,
|
||
args=[task.id]
|
||
)
|
||
else:
|
||
task.status = TaskStatusEnum.FAILED
|
||
task.error_msg = "意外中断，重试次数已达上限"
|
||
db.commit()
|
||
db.close() |