356 lines
11 KiB
Python
356 lines
11 KiB
Python
"""
|
||
任务管理器
|
||
用于管理检索任务的生命周期，支持任务取消
|
||
"""
|
||
|
||
import uuid
|
||
import time
|
||
import threading
|
||
from typing import Dict, Any, Optional
|
||
from datetime import datetime, timedelta
|
||
from enum import Enum
|
||
|
||
|
||
class TaskStatus(Enum):
|
||
"""任务状态枚举"""
|
||
PENDING = "pending"
|
||
RUNNING = "running"
|
||
COMPLETED = "completed"
|
||
CANCELLED = "cancelled"
|
||
FAILED = "failed"
|
||
|
||
|
||
class TaskInfo:
|
||
"""任务信息类"""
|
||
|
||
def __init__(self, task_id: str, query: str):
|
||
self.task_id = task_id
|
||
self.query = query
|
||
self.status = TaskStatus.PENDING
|
||
self.created_at = time.time()
|
||
self.started_at: Optional[float] = None
|
||
self.completed_at: Optional[float] = None
|
||
self.result: Optional[Dict[str, Any]] = None
|
||
self.error: Optional[str] = None
|
||
self.progress: int = 0
|
||
self.current_step: str = ""
|
||
|
||
def to_dict(self) -> Dict[str, Any]:
|
||
"""转换为字典"""
|
||
return {
|
||
"task_id": self.task_id,
|
||
"query": self.query,
|
||
"status": self.status.value,
|
||
"created_at": self.created_at,
|
||
"started_at": self.started_at,
|
||
"completed_at": self.completed_at,
|
||
"progress": self.progress,
|
||
"current_step": self.current_step,
|
||
"error": self.error,
|
||
"elapsed_time": self.get_elapsed_time()
|
||
}
|
||
|
||
def get_elapsed_time(self) -> float:
|
||
"""获取已用时间"""
|
||
if self.started_at:
|
||
end_time = self.completed_at or time.time()
|
||
return end_time - self.started_at
|
||
return 0
|
||
|
||
|
||
class TaskManager:
|
||
"""任务管理器"""
|
||
|
||
def __init__(self, max_tasks: int = 1000, task_timeout: float = 1800):
|
||
"""
|
||
初始化任务管理器
|
||
|
||
Args:
|
||
max_tasks: 最大任务数量
|
||
task_timeout: 任务超时时间（秒），默认30分钟
|
||
"""
|
||
self.tasks: Dict[str, TaskInfo] = {}
|
||
self.stop_events: Dict[str, threading.Event] = {}
|
||
self.max_tasks = max_tasks
|
||
self.task_timeout = task_timeout
|
||
self._lock = threading.RLock()
|
||
|
||
# 启动清理线程
|
||
self._cleanup_thread = threading.Thread(target=self._cleanup_expired_tasks, daemon=True)
|
||
self._cleanup_thread.start()
|
||
|
||
def create_task(self, query: str) -> str:
|
||
"""
|
||
创建新任务
|
||
|
||
Args:
|
||
query: 查询内容
|
||
|
||
Returns:
|
||
任务ID
|
||
"""
|
||
with self._lock:
|
||
# 检查任务数量限制
|
||
if len(self.tasks) >= self.max_tasks:
|
||
# 清理已完成的任务
|
||
self._cleanup_completed_tasks()
|
||
|
||
# 如果仍然超限，拒绝创建
|
||
if len(self.tasks) >= self.max_tasks:
|
||
raise RuntimeError(f"任务数量已达上限 ({self.max_tasks})")
|
||
|
||
# 生成任务ID
|
||
task_id = str(uuid.uuid4())
|
||
|
||
# 创建任务信息
|
||
task_info = TaskInfo(task_id, query)
|
||
self.tasks[task_id] = task_info
|
||
|
||
# 创建停止事件
|
||
self.stop_events[task_id] = threading.Event()
|
||
|
||
print(f"[TaskManager] 创建任务: {task_id[:8]}... for query: {query[:50]}...")
|
||
return task_id
|
||
|
||
def start_task(self, task_id: str) -> bool:
|
||
"""
|
||
开始执行任务
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
|
||
Returns:
|
||
是否成功
|
||
"""
|
||
with self._lock:
|
||
if task_id in self.tasks:
|
||
task = self.tasks[task_id]
|
||
if task.status == TaskStatus.PENDING:
|
||
task.status = TaskStatus.RUNNING
|
||
task.started_at = time.time()
|
||
print(f"[TaskManager] 开始任务: {task_id[:8]}...")
|
||
return True
|
||
return False
|
||
|
||
def complete_task(self, task_id: str, result: Dict[str, Any] = None) -> bool:
|
||
"""
|
||
完成任务
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
result: 任务结果
|
||
|
||
Returns:
|
||
是否成功
|
||
"""
|
||
with self._lock:
|
||
if task_id in self.tasks:
|
||
task = self.tasks[task_id]
|
||
task.status = TaskStatus.COMPLETED
|
||
task.completed_at = time.time()
|
||
task.result = result
|
||
task.progress = 100
|
||
print(f"[TaskManager] 完成任务: {task_id[:8]}... 耗时: {task.get_elapsed_time():.2f}s")
|
||
return True
|
||
return False
|
||
|
||
def fail_task(self, task_id: str, error: str) -> bool:
|
||
"""
|
||
标记任务失败
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
error: 错误信息
|
||
|
||
Returns:
|
||
是否成功
|
||
"""
|
||
with self._lock:
|
||
if task_id in self.tasks:
|
||
task = self.tasks[task_id]
|
||
task.status = TaskStatus.FAILED
|
||
task.completed_at = time.time()
|
||
task.error = error
|
||
print(f"[TaskManager] 任务失败: {task_id[:8]}... 错误: {error}")
|
||
return True
|
||
return False
|
||
|
||
def cancel_task(self, task_id: str) -> bool:
|
||
"""
|
||
取消任务
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
|
||
Returns:
|
||
是否成功
|
||
"""
|
||
with self._lock:
|
||
if task_id in self.stop_events:
|
||
# 设置停止信号
|
||
self.stop_events[task_id].set()
|
||
|
||
# 更新任务状态
|
||
if task_id in self.tasks:
|
||
task = self.tasks[task_id]
|
||
if task.status == TaskStatus.RUNNING:
|
||
task.status = TaskStatus.CANCELLED
|
||
task.completed_at = time.time()
|
||
print(f"[TaskManager] 取消任务: {task_id[:8]}...")
|
||
return True
|
||
return False
|
||
|
||
def should_stop(self, task_id: str) -> bool:
|
||
"""
|
||
检查任务是否应该停止
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
|
||
Returns:
|
||
是否应该停止
|
||
"""
|
||
if task_id in self.stop_events:
|
||
return self.stop_events[task_id].is_set()
|
||
return False
|
||
|
||
def update_progress(self, task_id: str, progress: int, current_step: str = "") -> bool:
|
||
"""
|
||
更新任务进度
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
progress: 进度（0-100）
|
||
current_step: 当前步骤描述
|
||
|
||
Returns:
|
||
是否成功
|
||
"""
|
||
with self._lock:
|
||
if task_id in self.tasks:
|
||
task = self.tasks[task_id]
|
||
task.progress = max(0, min(100, progress))
|
||
if current_step:
|
||
task.current_step = current_step
|
||
return True
|
||
return False
|
||
|
||
def get_task_info(self, task_id: str) -> Optional[Dict[str, Any]]:
|
||
"""
|
||
获取任务信息
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
|
||
Returns:
|
||
任务信息字典
|
||
"""
|
||
with self._lock:
|
||
if task_id in self.tasks:
|
||
return self.tasks[task_id].to_dict()
|
||
return None
|
||
|
||
def get_all_tasks(self) -> Dict[str, Dict[str, Any]]:
|
||
"""
|
||
获取所有任务信息
|
||
|
||
Returns:
|
||
所有任务信息字典
|
||
"""
|
||
with self._lock:
|
||
return {tid: task.to_dict() for tid, task in self.tasks.items()}
|
||
|
||
def cleanup_task(self, task_id: str) -> bool:
|
||
"""
|
||
清理任务资源
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
|
||
Returns:
|
||
是否成功
|
||
"""
|
||
with self._lock:
|
||
removed = False
|
||
|
||
# 移除停止事件
|
||
if task_id in self.stop_events:
|
||
del self.stop_events[task_id]
|
||
removed = True
|
||
|
||
# 注意：任务信息暂时保留，供查询使用
|
||
# 将由清理线程定期清理
|
||
|
||
if removed:
|
||
print(f"[TaskManager] 清理任务资源: {task_id[:8]}...")
|
||
|
||
return removed
|
||
|
||
def _cleanup_completed_tasks(self):
|
||
"""清理已完成的任务（内部方法）"""
|
||
with self._lock:
|
||
current_time = time.time()
|
||
to_remove = []
|
||
|
||
for task_id, task in self.tasks.items():
|
||
# 清理已完成/取消/失败且超过5分钟的任务
|
||
if task.status in [TaskStatus.COMPLETED, TaskStatus.CANCELLED, TaskStatus.FAILED]:
|
||
if task.completed_at and (current_time - task.completed_at) > 300:
|
||
to_remove.append(task_id)
|
||
|
||
for task_id in to_remove:
|
||
del self.tasks[task_id]
|
||
if task_id in self.stop_events:
|
||
del self.stop_events[task_id]
|
||
|
||
if to_remove:
|
||
print(f"[TaskManager] 清理了 {len(to_remove)} 个已完成的任务")
|
||
|
||
def _cleanup_expired_tasks(self):
|
||
"""清理过期任务（后台线程）"""
|
||
while True:
|
||
try:
|
||
time.sleep(60) # 每分钟检查一次
|
||
|
||
with self._lock:
|
||
current_time = time.time()
|
||
to_remove = []
|
||
|
||
for task_id, task in self.tasks.items():
|
||
# 超时的运行中任务
|
||
if task.status == TaskStatus.RUNNING:
|
||
if task.started_at and (current_time - task.started_at) > self.task_timeout:
|
||
print(f"[TaskManager] 任务超时: {task_id[:8]}...")
|
||
task.status = TaskStatus.FAILED
|
||
task.error = "Task timeout"
|
||
task.completed_at = current_time
|
||
# 设置停止信号
|
||
if task_id in self.stop_events:
|
||
self.stop_events[task_id].set()
|
||
|
||
# 清理长时间完成的任务（超过30分钟）
|
||
if task.status in [TaskStatus.COMPLETED, TaskStatus.CANCELLED, TaskStatus.FAILED]:
|
||
if task.completed_at and (current_time - task.completed_at) > 1800:
|
||
to_remove.append(task_id)
|
||
|
||
# 移除过期任务
|
||
for task_id in to_remove:
|
||
del self.tasks[task_id]
|
||
if task_id in self.stop_events:
|
||
del self.stop_events[task_id]
|
||
|
||
if to_remove:
|
||
print(f"[TaskManager] 清理了 {len(to_remove)} 个过期任务")
|
||
|
||
except Exception as e:
|
||
print(f"[TaskManager] 清理线程错误: {e}")
|
||
time.sleep(10) # 出错后等待10秒
|
||
|
||
|
||
# 全局任务管理器实例
|
||
task_manager = TaskManager()
|
||
|
||
|
||
def get_task_manager() -> TaskManager:
|
||
"""获取全局任务管理器实例"""
|
||
return task_manager |