308 lines
8.4 KiB
Python
308 lines
8.4 KiB
Python
"""
|
||
异步任务管理器
|
||
|
||
负责管理异步 AI 生成任务的创建、执行和状态跟踪
|
||
"""
|
||
import asyncio
|
||
import uuid
|
||
from typing import Dict, Optional, List, Callable, Any
|
||
from datetime import datetime
|
||
from collections import defaultdict
|
||
|
||
from app.models.task import AsyncTask, TaskStatus, TaskType, TaskProgress
|
||
from app.utils.logger import get_logger
|
||
from app.core.task_persistence import get_task_persistence
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
class TaskManager:
|
||
"""
|
||
异步任务管理器
|
||
|
||
功能：
|
||
1. 创建异步任务
|
||
2. 执行任务（在后台）
|
||
3. 跟踪任务状态和进度
|
||
4. 提供任务查询接口
|
||
"""
|
||
|
||
def __init__(self):
|
||
# 内存存储（生产环境应使用 Redis 或数据库）
|
||
self._tasks: Dict[str, AsyncTask] = {}
|
||
# 按项目ID索引的任务
|
||
self._project_tasks: Dict[str, List[str]] = defaultdict(list)
|
||
# 按类型索引的任务
|
||
self._type_tasks: Dict[TaskType, List[str]] = defaultdict(list)
|
||
# 持久化存储
|
||
self._persistence = get_task_persistence()
|
||
|
||
# 启动时恢复持久化的任务
|
||
self._restore_tasks()
|
||
|
||
def create_task(
|
||
self,
|
||
task_type: TaskType,
|
||
params: Dict[str, Any],
|
||
project_id: Optional[str] = None
|
||
) -> AsyncTask:
|
||
"""
|
||
创建新任务
|
||
|
||
Args:
|
||
task_type: 任务类型
|
||
params: 任务参数
|
||
project_id: 关联的项目ID
|
||
|
||
Returns:
|
||
创建的任务
|
||
"""
|
||
task_id = str(uuid.uuid4())
|
||
|
||
task = AsyncTask(
|
||
id=task_id,
|
||
type=task_type,
|
||
params=params,
|
||
project_id=project_id,
|
||
status=TaskStatus.PENDING
|
||
)
|
||
|
||
self._tasks[task_id] = task
|
||
|
||
if project_id:
|
||
self._project_tasks[project_id].append(task_id)
|
||
|
||
self._type_tasks[task_type].append(task_id)
|
||
|
||
# 保存到持久化存储
|
||
self._persistence.save_task(task)
|
||
|
||
logger.info(f"创建任务: {task_id} ({task_type.value})")
|
||
|
||
return task
|
||
|
||
def get_task(self, task_id: str) -> Optional[AsyncTask]:
|
||
"""获取任务"""
|
||
return self._tasks.get(task_id)
|
||
|
||
def get_tasks_by_project(self, project_id: str) -> List[AsyncTask]:
|
||
"""获取项目的所有任务"""
|
||
task_ids = self._project_tasks.get(project_id, [])
|
||
return [self._tasks[tid] for tid in task_ids if tid in self._tasks]
|
||
|
||
def get_tasks_by_type(self, task_type: TaskType) -> List[AsyncTask]:
|
||
"""获取指定类型的所有任务"""
|
||
task_ids = self._type_tasks.get(task_type, [])
|
||
return [self._tasks[tid] for tid in task_ids if tid in self._tasks]
|
||
|
||
def update_task_progress(
|
||
self,
|
||
task_id: str,
|
||
current: int,
|
||
total: int = 100,
|
||
message: str = "",
|
||
stage: Optional[str] = None
|
||
) -> bool:
|
||
"""
|
||
更新任务进度
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
current: 当前进度
|
||
total: 总进度
|
||
message: 进度消息
|
||
stage: 当前阶段
|
||
|
||
Returns:
|
||
是否更新成功
|
||
"""
|
||
task = self._tasks.get(task_id)
|
||
if not task:
|
||
return False
|
||
|
||
task.progress = TaskProgress(
|
||
current=current,
|
||
total=total,
|
||
message=message,
|
||
stage=stage
|
||
)
|
||
task.updated_at = datetime.now()
|
||
|
||
logger.debug(f"任务进度更新: {task_id} - {current}/{total} - {message}")
|
||
|
||
return True
|
||
|
||
def update_task_status(
|
||
self,
|
||
task_id: str,
|
||
status: TaskStatus,
|
||
result: Optional[Dict[str, Any]] = None,
|
||
error: Optional[str] = None
|
||
) -> bool:
|
||
"""
|
||
更新任务状态
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
status: 新状态
|
||
result: 任务结果
|
||
error: 错误信息
|
||
|
||
Returns:
|
||
是否更新成功
|
||
"""
|
||
task = self._tasks.get(task_id)
|
||
if not task:
|
||
return False
|
||
|
||
old_status = task.status
|
||
task.status = status
|
||
task.updated_at = datetime.now()
|
||
|
||
if result:
|
||
task.result = result
|
||
|
||
if error:
|
||
task.error = error
|
||
|
||
# 状态变更时间戳
|
||
if status == TaskStatus.RUNNING and old_status == TaskStatus.PENDING:
|
||
task.started_at = datetime.now()
|
||
elif status in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED]:
|
||
task.completed_at = datetime.now()
|
||
|
||
# 更新持久化存储
|
||
self._persistence.update_task_status(task_id, status, result, error)
|
||
|
||
logger.info(f"任务状态更新: {task_id} - {old_status.value} -> {status.value}")
|
||
|
||
return True
|
||
|
||
def cancel_task(self, task_id: str) -> bool:
|
||
"""取消任务"""
|
||
return self.update_task_status(task_id, TaskStatus.CANCELLED)
|
||
|
||
def delete_task(self, task_id: str) -> bool:
|
||
"""删除任务"""
|
||
task = self._tasks.get(task_id)
|
||
if not task:
|
||
return False
|
||
|
||
# 从索引中移除
|
||
if task.project_id:
|
||
self._project_tasks[task.project_id] = [
|
||
tid for tid in self._project_tasks[task.project_id] if tid != task_id
|
||
]
|
||
|
||
self._type_tasks[task.type] = [
|
||
tid for tid in self._type_tasks[task.type] if tid != task_id
|
||
]
|
||
|
||
del self._tasks[task_id]
|
||
|
||
# 从持久化存储中删除
|
||
self._persistence.delete_task(task_id)
|
||
|
||
logger.info(f"删除任务: {task_id}")
|
||
|
||
return True
|
||
|
||
async def execute_task_async(
|
||
self,
|
||
task_id: str,
|
||
executor: Callable
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
异步执行任务
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
executor: 执行器函数，接收 params，返回 result
|
||
|
||
Returns:
|
||
任务执行结果
|
||
"""
|
||
task = self._tasks.get(task_id)
|
||
if not task:
|
||
raise ValueError(f"任务不存在: {task_id}")
|
||
|
||
# 更新状态为运行中
|
||
self.update_task_status(task_id, TaskStatus.RUNNING)
|
||
|
||
try:
|
||
# 执行任务
|
||
result = await executor(task.params)
|
||
|
||
# 更新状态为完成
|
||
self.update_task_status(task_id, TaskStatus.COMPLETED, result=result)
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"任务执行失败: {task_id} - {str(e)}")
|
||
self.update_task_status(task_id, TaskStatus.FAILED, error=str(e))
|
||
raise
|
||
|
||
def execute_task_in_background(
|
||
self,
|
||
task_id: str,
|
||
executor: Callable
|
||
) -> asyncio.Task:
|
||
"""
|
||
在后台执行任务
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
executor: 执行器函数
|
||
|
||
Returns:
|
||
asyncio Task 对象
|
||
"""
|
||
return asyncio.create_task(self.execute_task_async(task_id, executor))
|
||
|
||
def _restore_tasks(self) -> None:
|
||
"""从持久化存储恢复任务"""
|
||
try:
|
||
saved_tasks = self._persistence.get_all_tasks()
|
||
restored_count = 0
|
||
|
||
for task_id, task_data in saved_tasks.items():
|
||
# 跳过已完成、失败或取消的任务（太久远的不需要恢复）
|
||
if task_data.get('status') in ['completed', 'failed', 'cancelled']:
|
||
continue
|
||
|
||
# 恢复任务到内存
|
||
try:
|
||
task = AsyncTask(**task_data)
|
||
self._tasks[task_id] = task
|
||
|
||
# 恢复索引
|
||
if task.project_id:
|
||
self._project_tasks[task.project_id].append(task_id)
|
||
self._type_tasks[task.type].append(task_id)
|
||
|
||
restored_count += 1
|
||
except Exception as e:
|
||
logger.warning(f"恢复任务失败: {task_id} - {e}")
|
||
|
||
if restored_count > 0:
|
||
logger.info(f"从持久化存储恢复了 {restored_count} 个任务")
|
||
|
||
# 清理旧任务（超过24小时的已完成任务）
|
||
self._persistence.cleanup_old_tasks(max_age_hours=24)
|
||
|
||
except Exception as e:
|
||
logger.error(f"恢复任务失败: {e}")
|
||
|
||
|
||
# 全局单例
|
||
_task_manager: Optional[TaskManager] = None
|
||
|
||
|
||
def get_task_manager() -> TaskManager:
|
||
"""获取任务管理器单例"""
|
||
global _task_manager
|
||
if _task_manager is None:
|
||
_task_manager = TaskManager()
|
||
return _task_manager
|