Fine-tuning/backend/app/core/job_queue.py

import asyncio
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Callable, Coroutine, Optional

from pydantic import BaseModel, Field

from app.core.logging import logger


class JobStatus(str, Enum):
    PENDING = "pending"
    QUEUED = "queued"
    PREPROCESSING = "preprocessing"
    TRAINING = "training"
    COMPLETED = "completed"
    EVALUATING = "evaluating"
    EVALUATION_DONE = "evaluation_done"
    FAILED = "failed"
    CANCELLED = "cancelled"

    @property
    def is_terminal(self) -> bool:
        return self in (self.COMPLETED, self.FAILED, self.CANCELLED, self.EVALUATION_DONE)


class TrainingJob(BaseModel):
    id: str
    model_id: str
    model_type: str
    peft_method: str
    dataset_id: str
    config: dict = Field(default_factory=dict)
    status: JobStatus = JobStatus.PENDING
    progress: float = 0.0
    current_epoch: int = 0
    current_step: int = 0
    total_steps: int = 0
    loss: float | None = None
    adapter_path: str | None = None
    error_message: str | None = None
    created_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
    started_at: str | None = None
    finished_at: str | None = None


class JobQueue:
    """异步任务队列，支持取消和并发控制。"""

    def __init__(self, max_concurrent: int = 2):
        self._queue: asyncio.Queue[str] = asyncio.Queue()
        self._jobs: dict[str, TrainingJob] = {}
        self._cancel_events: dict[str, asyncio.Event] = {}
        self._callbacks: list[Callable[[TrainingJob], Coroutine[Any, Any, None]]] = []
        self._max_concurrent = max_concurrent
        self._workers: list[asyncio.Task] = []
        self._running = False

    async def start(self):
        """启动后台 worker。"""
        if self._running:
            return
        self._running = True
        for _ in range(self._max_concurrent):
            worker = asyncio.create_task(self._worker_loop())
            self._workers.append(worker)
        logger.info(f"JobQueue started with {self._max_concurrent} workers")

    async def stop(self):
        """停止所有 worker。"""
        self._running = False
        for event in self._cancel_events.values():
            event.set()
        for worker in self._workers:
            worker.cancel()
        self._workers.clear()
        logger.info("JobQueue stopped")

    async def enqueue(self, job_id: str, job: TrainingJob):
        """将任务加入队列。"""
        self._jobs[job_id] = job
        self._cancel_events[job_id] = asyncio.Event()
        await self._queue.put(job_id)
        logger.info(f"Job {job_id} enqueued")

    async def dequeue(self) -> str:
        """从队列中取出任务 ID。"""
        return await self._queue.get()

    def mark_done(self, job_id: str):
        """标记任务完成。"""
        self._queue.task_done()
        self._cancel_events.pop(job_id, None)

    def get_job(self, job_id: str) -> Optional[TrainingJob]:
        return self._jobs.get(job_id)

    def update_job(self, job_id: str, **kwargs):
        if job_id in self._jobs:
            job = self._jobs[job_id]
            for key, val in kwargs.items():
                if hasattr(job, key):
                    setattr(job, key, val)

    def is_cancelled(self, job_id: str) -> bool:
        event = self._cancel_events.get(job_id)
        return event is not None and event.is_set()

    async def cancel(self, job_id: str):
        """取消任务。"""
        if job_id in self._cancel_events:
            self._cancel_events[job_id].set()
            self.update_job(job_id, status=JobStatus.CANCELLED)
            await self._notify_callbacks()
            logger.info(f"Job {job_id} cancelled")

    def register_callback(self, callback: Callable[[TrainingJob], Coroutine[Any, Any, None]]):
        """注册状态变更回调（用于更新数据库等）。"""
        self._callbacks.append(callback)

    async def _notify_callbacks(self):
        for cb in self._callbacks:
            try:
                for job in self._jobs.values():
                    await cb(job)
            except Exception as e:
                logger.error(f"JobQueue callback error: {e}")

    async def _worker_loop(self):
        """worker 循环：不断从队列取任务并执行。"""
        while self._running:
            try:
                job_id = await asyncio.wait_for(self._queue.get(), timeout=1.0)
            except asyncio.TimeoutError:
                continue

            try:
                await self._run_job(job_id)
            except Exception as e:
                logger.error(f"Job {job_id} failed: {e}")
                self.update_job(job_id, status=JobStatus.FAILED, error_message=str(e))
            finally:
                self._queue.task_done()

    async def _run_job(self, job_id: str):
        """执行单个任务：预处理 → 训练 → 完成。"""
        job = self._jobs.get(job_id)
        if not job:
            return

        self.update_job(job_id, status=JobStatus.QUEUED)
        await self._notify_callbacks()

        if self.is_cancelled(job_id):
            return

        self.update_job(job_id, status=JobStatus.PREPROCESSING, started_at=datetime.now(timezone.utc).isoformat())
        await self._notify_callbacks()

        if self.is_cancelled(job_id):
            return

        try:
            config = job.config
            model_id = job.model_id
            model_type = job.model_type
            peft_method = job.peft_method
            dataset_id = config.get("dataset_id", job.dataset_id)

            # 获取数据集文件路径
            from app.config import get_settings
            settings = get_settings()

            # 查找数据集文件
            dataset_path = self._find_dataset_path(dataset_id)
            if not dataset_path:
                raise FileNotFoundError(f"Dataset not found: {dataset_id}")

            # 预处理
            processed_path = str(settings.processed_dir / f"{job_id}_processed.jsonl")
            task_type = config.get("task_type", "sft")
            template = config.get("dataset_template", "alpaca")

            # 选择引擎
            engine = self._get_engine(model_type)

            # 预处理
            await engine.preprocess_dataset(dataset_path, processed_path, task_type=task_type, template=template)
            self.update_job(job_id, status=JobStatus.TRAINING)
            await self._notify_callbacks()

            # 加载模型
            await engine.load_model(model_id, quantization="4bit" if peft_method == "qlora" else None)

            # 构建 PEFT 配置
            peft_config = engine.get_peft_config(peft_method, config)

            # 训练
            adapter_path = await engine.train(
                job_id=job_id,
                dataset_path=processed_path,
                peft_config=peft_config,
                training_args=config,
            )

            self.update_job(job_id, status=JobStatus.COMPLETED, adapter_path=adapter_path)
            await self._notify_callbacks()
            logger.info(f"Job {job_id} completed successfully")

        except asyncio.CancelledError:
            self.update_job(job_id, status=JobStatus.CANCELLED)
            await self._notify_callbacks()
        except Exception as e:
            logger.error(f"Job {job_id} failed: {e}")
            self.update_job(job_id, status=JobStatus.FAILED, error_message=str(e))
            await self._notify_callbacks()

    def _find_dataset_path(self, dataset_id: str) -> str | None:
        """根据 dataset_id 查找文件路径（数据库或 uploads 目录）。"""
        import asyncio
        from app.core.db import async_session, DatasetRecord
        from sqlalchemy import select
        from app.config import get_settings
        from pathlib import Path

        settings = get_settings()

        # 尝试从数据库查找
        try:
            loop = asyncio.get_event_loop()
            task = loop.create_task(self._lookup_dataset_db(dataset_id))
            path = loop.run_until_complete(task)
            if path:
                return path
        except Exception:
            pass

        # 尝试从 uploads 目录查找
        upload_path = settings.uploads_dir / dataset_id
        if upload_path.exists():
            return str(upload_path)

        # 如果 dataset_id 本身是路径
        if Path(dataset_id).exists():
            return dataset_id

        return None

    async def _lookup_dataset_db(self, dataset_id: str) -> str | None:
        """从数据库查找数据集路径。"""
        from app.core.db import async_session, DatasetRecord
        from sqlalchemy import select

        async with async_session() as session:
            result = await session.execute(select(DatasetRecord).where(
                (DatasetRecord.id == dataset_id) | (DatasetRecord.name == dataset_id)
            ))
            record = result.scalar_one_or_none()
            if record:
                return record.file_path
        return None

    def _get_engine(self, model_type: str):
        """根据模型类型选择训练引擎。"""
        if model_type == "vision":
            from app.engines.vision_engine import vision_engine
            return vision_engine
        elif model_type == "multimodal":
            from app.engines.multimodal_engine import multimodal_engine
            return multimodal_engine
        else:
            from app.engines.text_engine import text_engine
            return text_engine

    @property
    def jobs(self) -> dict[str, TrainingJob]:
        return dict(self._jobs)


# 全局单例
job_queue = JobQueue(max_concurrent=2)