170 lines
5.2 KiB
Python
170 lines
5.2 KiB
Python
"""
|
|
微调任务管理服务
|
|
"""
|
|
|
|
import asyncio
|
|
from typing import Dict, List, Optional
|
|
from datetime import datetime
|
|
from enum import Enum
|
|
|
|
|
|
class FineTuneStatus(Enum):
|
|
"""微调任务状态"""
|
|
PENDING = "pending"
|
|
RUNNING = "running"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
CANCELED = "canceled"
|
|
|
|
|
|
class FineTuneTask:
|
|
"""微调任务"""
|
|
|
|
def __init__(self, task_id: str, name: str, description: str):
|
|
self.task_id = task_id
|
|
self.name = name
|
|
self.description = description
|
|
self.status = FineTuneStatus.PENDING
|
|
self.progress = 0.0
|
|
self.created_at = datetime.now().isoformat()
|
|
self.updated_at = datetime.now().isoformat()
|
|
self.logs: List[str] = []
|
|
|
|
def update_status(self, status: FineTuneStatus):
|
|
"""更新任务状态"""
|
|
self.status = status
|
|
self.updated_at = datetime.now().isoformat()
|
|
|
|
def update_progress(self, progress: float, log: str = None):
|
|
"""更新任务进度"""
|
|
self.progress = min(max(progress, 0.0), 100.0)
|
|
self.updated_at = datetime.now().isoformat()
|
|
if log:
|
|
self.logs.append(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {log}")
|
|
|
|
def to_dict(self):
|
|
"""转换为字典"""
|
|
return {
|
|
"task_id": self.task_id,
|
|
"name": self.name,
|
|
"description": self.description,
|
|
"status": self.status.value,
|
|
"progress": self.progress,
|
|
"created_at": self.created_at,
|
|
"updated_at": self.updated_at,
|
|
"logs": self.logs
|
|
}
|
|
|
|
|
|
class FineTuneManager:
|
|
"""微调任务管理器"""
|
|
|
|
def __init__(self):
|
|
self.tasks: Dict[str, FineTuneTask] = {}
|
|
self.listeners: Dict[str, List[asyncio.Queue]] = {}
|
|
|
|
def create_task(self, task_id: str, name: str, description: str) -> FineTuneTask:
|
|
"""创建新任务"""
|
|
task = FineTuneTask(task_id, name, description)
|
|
self.tasks[task_id] = task
|
|
self.listeners[task_id] = []
|
|
return task
|
|
|
|
def get_task(self, task_id: str) -> Optional[FineTuneTask]:
|
|
"""获取任务"""
|
|
return self.tasks.get(task_id)
|
|
|
|
def get_all_tasks(self) -> List[FineTuneTask]:
|
|
"""获取所有任务"""
|
|
return list(self.tasks.values())
|
|
|
|
def update_task_status(self, task_id: str, status: FineTuneStatus):
|
|
"""更新任务状态"""
|
|
task = self.tasks.get(task_id)
|
|
if task:
|
|
task.update_status(status)
|
|
self._notify_listeners(task_id)
|
|
|
|
def update_task_progress(self, task_id: str, progress: float, log: str = None):
|
|
"""更新任务进度"""
|
|
task = self.tasks.get(task_id)
|
|
if task:
|
|
task.update_progress(progress, log)
|
|
self._notify_listeners(task_id)
|
|
|
|
def delete_task(self, task_id: str):
|
|
"""删除任务"""
|
|
if task_id in self.tasks:
|
|
del self.tasks[task_id]
|
|
if task_id in self.listeners:
|
|
# 通知所有监听器任务已删除
|
|
for queue in self.listeners[task_id]:
|
|
queue.put_nowait(None) # None表示任务已删除
|
|
del self.listeners[task_id]
|
|
|
|
async def add_listener(self, task_id: str) -> asyncio.Queue:
|
|
"""添加任务监听器"""
|
|
queue = asyncio.Queue()
|
|
if task_id in self.listeners:
|
|
self.listeners[task_id].append(queue)
|
|
else:
|
|
self.listeners[task_id] = [queue]
|
|
return queue
|
|
|
|
def remove_listener(self, task_id: str, queue: asyncio.Queue):
|
|
"""移除任务监听器"""
|
|
if task_id in self.listeners:
|
|
self.listeners[task_id].remove(queue)
|
|
# 如果没有监听器了，清理
|
|
if not self.listeners[task_id]:
|
|
del self.listeners[task_id]
|
|
|
|
def _notify_listeners(self, task_id: str):
|
|
"""通知所有监听器"""
|
|
if task_id in self.listeners:
|
|
task = self.tasks.get(task_id)
|
|
if task:
|
|
for queue in self.listeners[task_id]:
|
|
queue.put_nowait(task.to_dict())
|
|
|
|
def task_exists(self, task_id: str) -> bool:
|
|
"""检查任务是否存在"""
|
|
return task_id in self.tasks
|
|
|
|
|
|
# 全局微调任务管理器实例
|
|
fine_tune_manager = FineTuneManager()
|
|
|
|
|
|
async def simulate_fine_tune(task_id: str):
|
|
"""模拟微调任务执行"""
|
|
manager = fine_tune_manager
|
|
task = manager.get_task(task_id)
|
|
|
|
if not task:
|
|
return
|
|
|
|
# 更新状态为运行中
|
|
manager.update_task_status(task_id, FineTuneStatus.RUNNING)
|
|
|
|
# 模拟训练过程
|
|
steps = [
|
|
("初始化模型", 10),
|
|
("加载数据集", 20),
|
|
("训练第1轮", 35),
|
|
("训练第2轮", 50),
|
|
("训练第3轮", 65),
|
|
("训练第4轮", 80),
|
|
("模型评估", 90),
|
|
("保存模型", 100)
|
|
]
|
|
|
|
for step_name, step_progress in steps:
|
|
# 模拟步骤执行
|
|
await asyncio.sleep(2)
|
|
manager.update_task_progress(task_id, step_progress, f"{step_name}完成")
|
|
|
|
# 更新状态为完成
|
|
manager.update_task_status(task_id, FineTuneStatus.COMPLETED)
|
|
manager.update_task_progress(task_id, 100.0, "微调任务已完成")
|