165 lines
5.7 KiB
Python
165 lines
5.7 KiB
Python
"""
|
||
异步任务管理器
|
||
负责任务的提交、调度、状态管理
|
||
"""
|
||
import asyncio
|
||
import uuid
|
||
import time
|
||
import os
|
||
from typing import Dict, List, Optional, Callable, Awaitable
|
||
from .task_store import TaskStore
|
||
|
||
|
||
# 任务状态
|
||
STATUS_PENDING = "pending"
|
||
STATUS_RUNNING = "running"
|
||
STATUS_COMPLETED = "completed"
|
||
STATUS_FAILED = "failed"
|
||
STATUS_CANCELLED = "cancelled"
|
||
|
||
|
||
class TaskManager:
|
||
def __init__(self, store_dir: str = "./data/tasks", max_concurrent: int = 3):
|
||
self.store = TaskStore(store_dir)
|
||
self.max_concurrent = max_concurrent
|
||
self._tasks: Dict[str, dict] = {}
|
||
self._running = 0
|
||
self._semaphore = asyncio.Semaphore(max_concurrent)
|
||
self._worker_task: Optional[asyncio.Task] = None
|
||
self._queue: asyncio.Queue = asyncio.Queue()
|
||
self._executor: Optional[Callable[[dict], Awaitable[dict]]] = None
|
||
|
||
async def start(self, executor: Callable[[dict], Awaitable[dict]]):
|
||
"""启动任务管理器，executor 是实际执行任务的异步函数"""
|
||
self._executor = executor
|
||
self._worker_task = asyncio.create_task(self._worker_loop())
|
||
print(f"[TaskManager] 已启动，最大并发: {self.max_concurrent}")
|
||
|
||
async def stop(self):
|
||
"""停止任务管理器"""
|
||
if self._worker_task:
|
||
self._worker_task.cancel()
|
||
try:
|
||
await self._worker_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
print("[TaskManager] 已停止")
|
||
|
||
async def submit_task(self, description: str, project_path: str,
|
||
requirements: str = "", extra: dict = None) -> dict:
|
||
"""提交一个新任务"""
|
||
task_id = str(uuid.uuid4())[:8]
|
||
now = time.time()
|
||
task = {
|
||
"id": task_id,
|
||
"description": description,
|
||
"project_path": project_path,
|
||
"requirements": requirements,
|
||
"status": STATUS_PENDING,
|
||
"created_at": now,
|
||
"updated_at": now,
|
||
"result": None,
|
||
"logs": [],
|
||
"extra": extra or {},
|
||
}
|
||
self._tasks[task_id] = task
|
||
self.store.save(task)
|
||
await self._queue.put(task_id)
|
||
print(f"[TaskManager] 任务已提交: {task_id} - {description[:50]}")
|
||
return task
|
||
|
||
def get_task(self, task_id: str) -> Optional[dict]:
|
||
"""获取任务详情"""
|
||
# 优先从内存取，没有再从文件读
|
||
if task_id in self._tasks:
|
||
return self._tasks[task_id]
|
||
return self.store.load(task_id)
|
||
|
||
def get_task_status(self, task_id: str) -> Optional[str]:
|
||
task = self.get_task(task_id)
|
||
return task["status"] if task else None
|
||
|
||
def list_tasks(self, status: str = None) -> List[dict]:
|
||
"""列出所有任务，可按状态过滤"""
|
||
tasks = list(self._tasks.values())
|
||
# 加上磁盘上的任务
|
||
disk_tasks = self.store.list_all()
|
||
disk_ids = {t["id"] for t in tasks}
|
||
for t in disk_tasks:
|
||
if t["id"] not in disk_ids:
|
||
tasks.append(t)
|
||
if status:
|
||
tasks = [t for t in tasks if t["status"] == status]
|
||
tasks.sort(key=lambda t: t.get("created_at", 0), reverse=True)
|
||
return tasks
|
||
|
||
def append_log(self, task_id: str, message: str):
|
||
"""追加任务日志"""
|
||
task = self._tasks.get(task_id)
|
||
if not task:
|
||
return
|
||
if "logs" not in task:
|
||
task["logs"] = []
|
||
task["logs"].append({
|
||
"time": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||
"message": message,
|
||
})
|
||
if len(task["logs"]) > 500:
|
||
task["logs"] = task["logs"][-500:]
|
||
# 异步持久化（这里直接同步写，简单起见）
|
||
self.store.save(task)
|
||
|
||
async def cancel_task(self, task_id: str) -> bool:
|
||
"""取消任务"""
|
||
task = self._tasks.get(task_id)
|
||
if not task:
|
||
return False
|
||
if task["status"] in (STATUS_COMPLETED, STATUS_FAILED, STATUS_CANCELLED):
|
||
return False
|
||
task["status"] = STATUS_CANCELLED
|
||
task["updated_at"] = time.time()
|
||
self.store.save(task)
|
||
print(f"[TaskManager] 任务已取消: {task_id}")
|
||
return True
|
||
|
||
async def _worker_loop(self):
|
||
"""后台 worker，从队列取任务执行"""
|
||
while True:
|
||
try:
|
||
task_id = await self._queue.get()
|
||
async with self._semaphore:
|
||
await self._execute_task(task_id)
|
||
except asyncio.CancelledError:
|
||
break
|
||
except Exception as e:
|
||
print(f"[TaskManager] Worker 异常: {e}")
|
||
await asyncio.sleep(1)
|
||
|
||
async def _execute_task(self, task_id: str):
|
||
"""执行单个任务"""
|
||
task = self._tasks.get(task_id)
|
||
if not task or task["status"] == STATUS_CANCELLED:
|
||
return
|
||
|
||
task["status"] = STATUS_RUNNING
|
||
task["updated_at"] = time.time()
|
||
self.store.save(task)
|
||
print(f"[TaskManager] 开始执行: {task_id}")
|
||
|
||
try:
|
||
result = await self._executor(task)
|
||
task["result"] = result
|
||
# 检查是否已被取消
|
||
if task["status"] == STATUS_CANCELLED:
|
||
return
|
||
task["status"] = STATUS_COMPLETED
|
||
print(f"[TaskManager] 任务完成: {task_id}")
|
||
except Exception as e:
|
||
task["status"] = STATUS_FAILED
|
||
task["result"] = {"error": str(e)}
|
||
self.append_log(task_id, f"执行失败: {e}")
|
||
print(f"[TaskManager] 任务失败: {task_id} - {e}")
|
||
finally:
|
||
task["updated_at"] = time.time()
|
||
self.store.save(task)
|