diff --git a/backend/app/api/v1/__init__.py b/backend/app/api/v1/__init__.py index e69de29..c505308 100644 --- a/backend/app/api/v1/__init__.py +++ b/backend/app/api/v1/__init__.py @@ -0,0 +1,4 @@ +"""API v1 路由""" +from . import jobs, upload, results, tasks + +__all__ = ["jobs", "upload", "results", "tasks"] diff --git a/backend/app/api/v1/jobs.py b/backend/app/api/v1/jobs.py index 778d8a8..6076e1f 100644 --- a/backend/app/api/v1/jobs.py +++ b/backend/app/api/v1/jobs.py @@ -1,7 +1,7 @@ """任务管理 API""" -from fastapi import APIRouter, Depends, HTTPException, UploadFile, File +from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form from sqlalchemy.orm import Session -from typing import List +from typing import List, Optional from pathlib import Path import uuid import shutil @@ -9,7 +9,7 @@ import shutil from ...database import get_db from ...models.job import Job, JobStatus from ...schemas.job import JobResponse -from ...workers.tasks import run_bttoxin_analysis +from ...workers.tasks import run_bttoxin_analysis, update_queue_positions from ...config import settings router = APIRouter() @@ -19,21 +19,55 @@ RESULTS_DIR = Path(settings.RESULTS_DIR) UPLOAD_DIR.mkdir(exist_ok=True) RESULTS_DIR.mkdir(exist_ok=True) + @router.post("/create", response_model=JobResponse) async def create_job( files: List[UploadFile] = File(...), - sequence_type: str = "nucl", - scaf_suffix: str = ".fna", - threads: int = 4, + sequence_type: str = Form("nucl"), + scaf_suffix: str = Form(".fna"), + threads: int = Form(4), + min_identity: float = Form(0.8), + min_coverage: float = Form(0.6), + allow_unknown_families: bool = Form(False), + require_index_hit: bool = Form(True), db: Session = Depends(get_db) ): - """创建新任务""" + """ + 创建新分析任务 + + Args: + files: 上传的文件列表（单文件） + sequence_type: 序列类型 (nucl=核酸, prot=蛋白) + scaf_suffix: 文件后缀 + threads: 线程数 + min_identity: 最小相似度 (0-1) + min_coverage: 最小覆盖度 (0-1) + allow_unknown_families: 是否允许未知家族 + require_index_hit: 是否要求索引命中 + """ + # 验证文件类型 + allowed_extensions = {".fna", ".fa", ".fasta", ".faa"} + for file in files: + ext = Path(file.filename).suffix.lower() + if ext not in allowed_extensions: + raise HTTPException( + status_code=400, + detail=f"Invalid file extension: {ext}. Allowed: {', '.join(allowed_extensions)}" + ) + + # 限制单文件上传 + if len(files) != 1: + raise HTTPException( + status_code=400, + detail="Only one file allowed per task" + ) + job_id = str(uuid.uuid4()) job_input_dir = UPLOAD_DIR / job_id job_output_dir = RESULTS_DIR / job_id - job_input_dir.mkdir(parents=True) - job_output_dir.mkdir(parents=True) + job_input_dir.mkdir(parents=True, exist_ok=True) + job_output_dir.mkdir(parents=True, exist_ok=True) uploaded_files = [] for file in files: @@ -48,20 +82,29 @@ async def create_job( input_files=uploaded_files, sequence_type=sequence_type, scaf_suffix=scaf_suffix, - threads=threads + threads=threads, + min_identity=int(min_identity * 100), # 存储为百分比 + min_coverage=int(min_coverage * 100), + allow_unknown_families=int(allow_unknown_families), + require_index_hit=int(require_index_hit), ) db.add(job) db.commit() db.refresh(job) + # 启动 Celery 任务 task = run_bttoxin_analysis.delay( job_id=job_id, input_dir=str(job_input_dir), output_dir=str(job_output_dir), sequence_type=sequence_type, scaf_suffix=scaf_suffix, - threads=threads + threads=threads, + min_identity=min_identity, + min_coverage=min_coverage, + allow_unknown_families=allow_unknown_families, + require_index_hit=require_index_hit, ) job.celery_task_id = task.id @@ -69,6 +112,7 @@ async def create_job( return job + @router.get("/{job_id}", response_model=JobResponse) async def get_job(job_id: str, db: Session = Depends(get_db)): """获取任务详情""" @@ -77,6 +121,7 @@ async def get_job(job_id: str, db: Session = Depends(get_db)): raise HTTPException(status_code=404, detail="Job not found") return job + @router.get("/{job_id}/progress") async def get_job_progress(job_id: str, db: Session = Depends(get_db)): """获取任务进度""" @@ -84,15 +129,26 @@ async def get_job_progress(job_id: str, db: Session = Depends(get_db)): if not job: raise HTTPException(status_code=404, detail="Job not found") + result = { + 'job_id': job_id, + 'status': job.status.value if isinstance(job.status, JobStatus) else job.status, + 'current_stage': job.current_stage, + 'progress_percent': job.progress_percent, + 'queue_position': job.queue_position, + } + if job.celery_task_id: from ...core.celery_app import celery_app task = celery_app.AsyncResult(job.celery_task_id) + result['celery_state'] = task.state + if task.state == 'PROGRESS': + result['celery_info'] = task.info - return { - 'job_id': job_id, - 'status': job.status, - 'celery_state': task.state, - 'progress': task.info if task.state == 'PROGRESS' else None - } + return result - return {'job_id': job_id, 'status': job.status} + +@router.post("/update-queue-positions") +async def trigger_queue_update(db: Session = Depends(get_db)): + """手动触发队列位置更新""" + task = update_queue_positions.delay() + return {"message": "Queue update triggered", "task_id": task.id} diff --git a/backend/app/api/v1/results.py b/backend/app/api/v1/results.py index dfde9f3..8a43fe6 100644 --- a/backend/app/api/v1/results.py +++ b/backend/app/api/v1/results.py @@ -1,8 +1,76 @@ -"""结果查询 API""" -from fastapi import APIRouter +"""结果下载 API""" +from fastapi import APIRouter, HTTPException, Response +from fastapi.responses import FileResponse +from sqlalchemy.orm import Session +from pathlib import Path +import tarfile +import io +import shutil + +from ...database import get_db +from ...models.job import Job, JobStatus +from ...config import settings router = APIRouter() -@router.get("/") -async def results_info(): - return {"message": "Results endpoint"} +RESULTS_DIR = Path(settings.RESULTS_DIR) + + +@router.get("/{job_id}/download") +async def download_results(job_id: str, db: Session = Depends(get_db)): + """下载任务结果（打包为 .tar.gz）""" + job = db.query(Job).filter(Job.id == job_id).first() + if not job: + raise HTTPException(status_code=404, detail="Job not found") + + if job.status != JobStatus.COMPLETED: + raise HTTPException( + status_code=400, + detail=f"Job not completed. Current status: {job.status}" + ) + + job_output_dir = RESULTS_DIR / job_id + if not job_output_dir.exists(): + raise HTTPException(status_code=404, detail="Results not found on disk") + + # 创建 tar.gz 文件到内存 + tar_buffer = io.BytesIO() + with tarfile.open(fileobj=tar_buffer, mode="w:gz") as tar: + for file_path in job_output_dir.rglob("*"): + if file_path.is_file(): + arcname = file_path.relative_to(job_output_dir) + tar.add(file_path, arcname=arcname) + + tar_buffer.seek(0) + + return Response( + content=tar_buffer.read(), + media_type="application/gzip", + headers={ + "Content-Disposition": f"attachment; filename=bttoxin_{job_id}.tar.gz" + } + ) + + +@router.delete("/{job_id}") +async def delete_job(job_id: str, db: Session = Depends(get_db)): + """删除任务及其结果""" + job = db.query(Job).filter(Job.id == job_id).first() + if not job: + raise HTTPException(status_code=404, detail="Job not found") + + # 删除磁盘上的文件 + job_input_dir = Path(settings.UPLOAD_DIR) / job_id + job_output_dir = RESULTS_DIR / job_id + + if job_input_dir.exists(): + shutil.rmtree(job_input_dir) + + if job_output_dir.exists(): + shutil.rmtree(job_output_dir) + + # 删除数据库记录 + db.delete(job) + db.commit() + + return {"message": f"Job {job_id} deleted successfully"} diff --git a/backend/app/api/v1/tasks.py b/backend/app/api/v1/tasks.py new file mode 100644 index 0000000..954fecc --- /dev/null +++ b/backend/app/api/v1/tasks.py @@ -0,0 +1,70 @@ +"""任务管理 API - 兼容 /api/v1/tasks 路径""" +from fastapi import APIRouter, Depends, HTTPException +from sqlalchemy.orm import Session +from typing import List +from pathlib import Path +from pydantic import BaseModel + +from ...database import get_db +from ...models.job import Job, JobStatus +from ...schemas.job import JobResponse +from ...config import settings + +router = APIRouter() + + +class TaskCreateRequest(BaseModel): + """任务创建请求""" + files: List[str] # 文件名列表 + sequence_type: str = "nucl" + min_identity: float = 0.8 + min_coverage: float = 0.6 + allow_unknown_families: bool = False + require_index_hit: bool = True + + +class QueuePosition(BaseModel): + """队列位置信息""" + position: int + estimated_wait_minutes: int = None + + +@router.post("/", response_model=JobResponse) +async def create_task(request: TaskCreateRequest, db: Session = Depends(get_db)): + """创建新任务（兼容前端）""" + # 暂时复用 jobs 逻辑 + # TODO: 实现完整的文件上传和处理 + raise HTTPException(status_code=501, detail="Use POST /api/v1/jobs/create for now") + + +@router.get("/{task_id}", response_model=JobResponse) +async def get_task(task_id: str, db: Session = Depends(get_db)): + """获取任务状态""" + job = db.query(Job).filter(Job.id == task_id).first() + if not job: + raise HTTPException(status_code=404, detail="Task not found") + return job + + +@router.get("/{task_id}/queue") +async def get_queue_position(task_id: str, db: Session = Depends(get_db)): + """获取排队位置""" + job = db.query(Job).filter(Job.id == task_id).first() + if not job: + raise HTTPException(status_code=404, detail="Task not found") + + if job.status not in [JobStatus.PENDING, JobStatus.QUEUED]: + return {"position": 0, "message": "Task is not in queue"} + + # 计算排队位置 + ahead_jobs = db.query(Job).filter( + Job.id != task_id, + Job.status.in_([JobStatus.PENDING, JobStatus.QUEUED]), + Job.created_at < job.created_at + ).count() + + position = ahead_jobs + 1 + # 假设每个任务约5分钟 + estimated_wait = position * 5 + + return QueuePosition(position=position, estimated_wait_minutes=estimated_wait) diff --git a/backend/app/main.py b/backend/app/main.py index 38da7ed..201cb1e 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -1,51 +1,11 @@ -from fastapi import FastAPI -from fastapi.middleware.cors import CORSMiddleware -from .core.config import settings -from .core.database import init_db -from .api.routes import jobs as jobs_routes -from .core.logging import setup_logging - - -def create_app() -> FastAPI: - app = FastAPI( - title=settings.APP_NAME, - version=settings.APP_VERSION, - docs_url="/docs", - redoc_url="/redoc", - openapi_url=f"{settings.API_V1_PREFIX}/openapi.json", - ) - - app.add_middleware( - CORSMiddleware, - allow_origins=settings.CORS_ORIGINS, - allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], - ) - - app.include_router(jobs_routes.router, prefix=f"{settings.API_V1_PREFIX}/jobs", tags=["jobs"]) - - @app.on_event("startup") - def _on_startup() -> None: - setup_logging(settings.LOG_LEVEL, settings.LOG_FORMAT) - init_db() - - @app.get("/healthz") - def healthz() -> dict: - return {"status": "ok"} - - return app - - -app = create_app() - """FastAPI 主应用""" from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from contextlib import asynccontextmanager from .config import settings -from .api.v1 import jobs, upload, results +from .api.v1 import jobs, upload, results, tasks + @asynccontextmanager async def lifespan(app: FastAPI): @@ -56,6 +16,7 @@ async def lifespan(app: FastAPI): # 关闭时 print("👋 Shutting down BtToxin Pipeline API...") + app = FastAPI( title=settings.APP_NAME, version=settings.APP_VERSION, @@ -74,9 +35,11 @@ app.add_middleware( # 路由 app.include_router(jobs.router, prefix=f"{settings.API_V1_STR}/jobs", tags=["jobs"]) +app.include_router(tasks.router, prefix=f"{settings.API_V1_STR}/tasks", tags=["tasks"]) app.include_router(upload.router, prefix=f"{settings.API_V1_STR}/upload", tags=["upload"]) app.include_router(results.router, prefix=f"{settings.API_V1_STR}/results", tags=["results"]) + @app.get("/") async def root(): return { @@ -85,6 +48,7 @@ async def root(): "status": "healthy" } + @app.get("/health") async def health(): return {"status": "ok"} diff --git a/backend/app/models/job.py b/backend/app/models/job.py index 7422635..b81ca91 100644 --- a/backend/app/models/job.py +++ b/backend/app/models/job.py @@ -1,187 +1,3 @@ -"""任务模型（使用 SQLModel）""" -from typing import Optional, List -from datetime import datetime -from enum import Enum -from sqlmodel import SQLModel, Field, Relationship, Column -from sqlalchemy import JSON -from .base import TimestampModel, generate_uuid - - -class JobStatus(str, Enum): - """任务状态""" - PENDING = "pending" - RUNNING = "running" - COMPLETED = "completed" - FAILED = "failed" - CANCELLED = "cancelled" - - -class StepStatus(str, Enum): - """步骤状态""" - PENDING = "pending" - RUNNING = "running" - COMPLETED = "completed" - FAILED = "failed" - SKIPPED = "skipped" - - -class JobBase(SQLModel): - """Job 基础字段""" - name: str = Field(max_length=255) - description: Optional[str] = None - - sequence_type: str = Field(default="nucl", max_length=20) - scaf_suffix: str = Field(default=".fna", max_length=50) - threads: int = Field(default=4, ge=1, le=32) - update_db: bool = Field(default=False) - - -class Job(JobBase, TimestampModel, table=True): - """Job 数据库模型""" - __tablename__ = "jobs" - - id: str = Field( - default_factory=generate_uuid, - primary_key=True, - index=True, - ) - - user_id: Optional[str] = Field(default=None, index=True) - - status: JobStatus = Field( - default=JobStatus.PENDING, - sa_column_kwargs={"index": True}, - ) - - input_files: List[dict] = Field(default_factory=list, sa_column=Column(JSON)) - workspace_path: Optional[str] = Field(default=None, max_length=500) - result_url: Optional[str] = Field(default=None, max_length=1000) - - celery_task_id: Optional[str] = Field(default=None, max_length=100, index=True) - - current_step: Optional[str] = Field(default=None, max_length=100) - progress: int = Field(default=0, ge=0, le=100) - - error_message: Optional[str] = None - - started_at: Optional[datetime] = None - completed_at: Optional[datetime] = None - - steps: List["Step"] = Relationship( - back_populates="job", - sa_relationship_kwargs={"cascade": "all, delete-orphan"}, - ) - logs: List["JobLog"] = Relationship( - back_populates="job", - sa_relationship_kwargs={"cascade": "all, delete-orphan"}, - ) - - -class JobCreate(JobBase): - """创建 Job 时的请求模型""" - pass - - -class JobRead(JobBase): - """读取 Job 时的响应模型""" - id: str - user_id: Optional[str] - status: JobStatus - workspace_path: Optional[str] - result_url: Optional[str] - celery_task_id: Optional[str] - current_step: Optional[str] - progress: int - error_message: Optional[str] - started_at: Optional[datetime] - completed_at: Optional[datetime] - created_at: datetime - updated_at: datetime - - -class JobUpdate(SQLModel): - """更新 Job 时的请求模型""" - status: Optional[JobStatus] = None - current_step: Optional[str] = None - progress: Optional[int] = None - error_message: Optional[str] = None - started_at: Optional[datetime] = None - completed_at: Optional[datetime] = None - - -class StepBase(SQLModel): - """Step 基础字段""" - step_name: str = Field(max_length=100) - step_order: int - - -class Step(StepBase, table=True): - """Step 数据库模型""" - __tablename__ = "steps" - - id: Optional[int] = Field(default=None, primary_key=True) - job_id: str = Field(foreign_key="jobs.id", index=True) - - status: StepStatus = Field(default=StepStatus.PENDING) - - celery_task_id: Optional[str] = Field(default=None, max_length=100) - log_file: Optional[str] = Field(default=None, max_length=500) - - result_data: Optional[dict] = Field(default=None, sa_column=Column(JSON)) - error_message: Optional[str] = None - - started_at: Optional[datetime] = None - completed_at: Optional[datetime] = None - duration_seconds: Optional[int] = None - - job: "Job" = Relationship(back_populates="steps") - - -class StepRead(StepBase): - """读取 Step 时的响应模型""" - id: int - job_id: str - status: StepStatus - celery_task_id: Optional[str] - log_file: Optional[str] - result_data: Optional[dict] - error_message: Optional[str] - started_at: Optional[datetime] - completed_at: Optional[datetime] - duration_seconds: Optional[int] - - -class JobLogBase(SQLModel): - """JobLog 基础字段""" - level: str = Field(max_length=20) - message: str - step_name: Optional[str] = Field(default=None, max_length=100) - - -class JobLog(JobLogBase, table=True): - """JobLog 数据库模型""" - __tablename__ = "job_logs" - - id: Optional[int] = Field(default=None, primary_key=True) - job_id: str = Field(foreign_key="jobs.id", index=True) - - metadata: Optional[dict] = Field(default=None, sa_column=Column(JSON)) - - timestamp: datetime = Field( - default_factory=datetime.utcnow, - sa_column_kwargs={"index": True}, - ) - - job: "Job" = Relationship(back_populates="logs") - - -class JobLogRead(JobLogBase): - """读取 JobLog 时的响应模型""" - id: int - job_id: str - metadata: Optional[dict] - timestamp: datetime - """任务模型""" from sqlalchemy import Column, String, Integer, DateTime, JSON, Enum, Text from sqlalchemy.sql import func @@ -189,28 +5,46 @@ import enum from ..database import Base + class JobStatus(str, enum.Enum): - PENDING = "pending" - RUNNING = "running" - COMPLETED = "completed" - FAILED = "failed" + """任务状态""" + PENDING = "pending" # 等待进入队列 + QUEUED = "queued" # 已排队，等待执行 + RUNNING = "running" # 正在执行 + COMPLETED = "completed" # 执行完成 + FAILED = "failed" # 执行失败 + class Job(Base): __tablename__ = "jobs" id = Column(String, primary_key=True, index=True) - celery_task_id = Column(String, nullable=True) - status = Column(Enum(JobStatus), default=JobStatus.PENDING) + celery_task_id = Column(String, nullable=True, index=True) + status = Column(Enum(JobStatus), default=JobStatus.PENDING, index=True) input_files = Column(JSON) sequence_type = Column(String, default="nucl") scaf_suffix = Column(String, default=".fna") threads = Column(Integer, default=4) + # 分析参数 + min_identity = Column(Integer, default=80) # 存储为百分比 (0-100) + min_coverage = Column(Integer, default=60) + allow_unknown_families = Column(Integer, default=0) # 0 = False, 1 = True + require_index_hit = Column(Integer, default=1) + result_url = Column(String, nullable=True) logs = Column(Text, nullable=True) error_message = Column(Text, nullable=True) - created_at = Column(DateTime(timezone=True), server_default=func.now()) - updated_at = Column(DateTime(timezone=True), onupdate=func.now()) + # 队列位置 + queue_position = Column(Integer, nullable=True) + + # 进度信息 + current_stage = Column(String, nullable=True) # digger, shoter, plots, bundle + progress_percent = Column(Integer, default=0) + + created_at = Column(DateTime(timezone=True), server_default=func.now(), index=True) + started_at = Column(DateTime(timezone=True), nullable=True) completed_at = Column(DateTime(timezone=True), nullable=True) + updated_at = Column(DateTime(timezone=True), onupdate=func.now()) diff --git a/backend/app/services/concurrency_manager.py b/backend/app/services/concurrency_manager.py new file mode 100644 index 0000000..7924759 --- /dev/null +++ b/backend/app/services/concurrency_manager.py @@ -0,0 +1,114 @@ +"""并发控制服务 - 使用 Redis 实现任务并发限制""" +import asyncio +from typing import Optional +import redis.asyncio as aioredis +from ..config import settings +import logging + +logger = logging.getLogger(__name__) + + +class ConcurrencyManager: + """并发控制管理器 - 使用 Redis Semaphore 实现 16 并发限制""" + + MAX_CONCURRENT_TASKS = 16 + REDIS_SEMAPHORE_KEY = "bttoxin:concurrency:semaphore" + REDIS_QUEUE_KEY = "bttoxin:queue:pending" + REDIS_RUNNING_KEY = "bttoxin:running:tasks" + + def __init__(self): + self._redis: Optional[aioredis.Redis] = None + self._lock = asyncio.Lock() + + async def get_redis(self) -> aioredis.Redis: + """获取 Redis 连接""" + if self._redis is None: + async with self._lock: + if self._redis is None: + self._redis = await aioredis.from_url( + settings.REDIS_URL, + encoding="utf-8", + decode_responses=True + ) + return self._redis + + async def acquire_slot(self, job_id: str) -> bool: + """ + 尝试获取执行槽位 + + Returns: + bool: 是否成功获取槽位 + """ + redis = await self.get_redis() + + # 使用 Redis Sorted Set 实现信号量 + # 检查当前运行的任务数 + running_count = await redis.zcard(self.REDIS_RUNNING_KEY) + + if running_count < self.MAX_CONCURRENT_TASKS: + # 有可用槽位，加入运行队列 + await redis.zadd( + self.REDIS_RUNNING_KEY, + {job_id: asyncio.get_event_loop().time()} + ) + logger.info(f"Job {job_id} acquired slot. Running: {running_count + 1}/{self.MAX_CONCURRENT_TASKS}") + return True + else: + # 没有可用槽位，加入等待队列 + await redis.rpush(self.REDIS_QUEUE_KEY, job_id) + logger.info(f"Job {job_id} queued. Position: {await self.get_queue_position(job_id)}") + return False + + async def release_slot(self, job_id: str): + """释放执行槽位""" + redis = await self.get_redis() + + # 从运行队列移除 + await redis.zrem(self.REDIS_RUNNING_KEY, job_id) + + # 检查是否有等待的任务 + next_job = await redis.lpop(self.REDIS_QUEUE_KEY) + if next_job: + # 将下一个任务加入运行队列 + await redis.zadd( + self.REDIS_RUNNING_KEY, + {next_job: asyncio.get_event_loop().time()} + ) + logger.info(f"Job {next_job} promoted from queue. Job {job_id} released.") + + async def get_queue_position(self, job_id: str) -> int: + """获取任务在队列中的位置""" + redis = await self.get_redis() + position = await redis.lpos(self.REDIS_QUEUE_KEY, job_id) + return position + 1 if position is not None else 0 + + async def get_running_count(self) -> int: + """获取当前运行的任务数""" + redis = await self.get_redis() + return await redis.zcard(self.REDIS_RUNNING_KEY) + + async def get_queue_length(self) -> int: + """获取等待队列长度""" + redis = await self.get_redis() + return await redis.llen(self.REDIS_QUEUE_KEY) + + async def close(self): + """关闭 Redis 连接""" + if self._redis: + await self._redis.close() + self._redis = None + + +# 全局单例 +_concurrency_manager: Optional[ConcurrencyManager] = None +_manager_lock = asyncio.Lock() + + +async def get_concurrency_manager() -> ConcurrencyManager: + """获取并发管理器单例""" + global _concurrency_manager + if _concurrency_manager is None: + async with _manager_lock: + if _concurrency_manager is None: + _concurrency_manager = ConcurrencyManager() + return _concurrency_manager diff --git a/backend/app/workers/tasks.py b/backend/app/workers/tasks.py index bf2394e..d376afb 100644 --- a/backend/app/workers/tasks.py +++ b/backend/app/workers/tasks.py @@ -1,17 +1,23 @@ -"""Celery 任务""" +"""Celery 任务 - 支持并发控制和多阶段 pipeline""" from celery import Task from pathlib import Path import shutil import logging +import asyncio from ..core.celery_app import celery_app from ..core.docker_client import DockerManager from ..database import SessionLocal from ..models.job import Job, JobStatus +from ..services.concurrency_manager import get_concurrency_manager logger = logging.getLogger(__name__) -@celery_app.task(bind=True) +# Pipeline 阶段定义 +PIPELINE_STAGES = ["digger", "shoter", "plots", "bundle"] + + +@celery_app.task(bind=True, max_retries=3) def run_bttoxin_analysis( self, job_id: str, @@ -19,23 +25,52 @@ def run_bttoxin_analysis( output_dir: str, sequence_type: str = "nucl", scaf_suffix: str = ".fna", - threads: int = 4 + threads: int = 4, + min_identity: float = 0.8, + min_coverage: float = 0.6, + allow_unknown_families: bool = False, + require_index_hit: bool = True, ): - """执行分析任务""" + """ + 执行分析任务 - 完整的 4 阶段 pipeline + + Stages: + 1. digger - BtToxin_Digger 识别毒素基因 + 2. shoter - BtToxin_Shoter 评估毒性活性 + 3. plots - 生成热力图 + 4. bundle - 打包结果 + """ db = SessionLocal() try: job = db.query(Job).filter(Job.id == job_id).first() - job.status = JobStatus.RUNNING + if not job: + logger.error(f"Job {job_id} not found") + return {'job_id': job_id, 'status': 'error', 'error': 'Job not found'} + + # 更新状态为 QUEUED + job.status = JobStatus.QUEUED db.commit() + # 尝试获取执行槽位（使用同步 Redis，因为 Celery 是同步的） + # 注意：这里简化处理，实际应该用异步 + # 暂时直接执行，稍后集成真正的并发控制 + + # 更新状态为 RUNNING + job.status = JobStatus.RUNNING + job.current_stage = "digger" + job.progress_percent = 0 + db.commit() + + # 阶段 1: Digger - 识别毒素基因 + logger.info(f"Job {job_id}: Starting Digger stage") self.update_state( state='PROGRESS', - meta={'current': 20, 'total': 100, 'status': 'Running analysis...'} + meta={'stage': 'digger', 'progress': 10, 'status': 'Running BtToxin_Digger...'} ) docker_manager = DockerManager() - result = docker_manager.run_bttoxin_digger( + digger_result = docker_manager.run_bttoxin_digger( input_dir=Path(input_dir), output_dir=Path(output_dir), sequence_type=sequence_type, @@ -43,22 +78,122 @@ def run_bttoxin_analysis( threads=threads ) - if result['success']: - job.status = JobStatus.COMPLETED - job.logs = result.get('logs', '') - else: - job.status = JobStatus.FAILED - job.error_message = result.get('error', 'Analysis failed') + if not digger_result['success']: + raise Exception(f"Digger stage failed: {digger_result.get('error', 'Unknown error')}") + job.progress_percent = 40 db.commit() - return {'job_id': job_id, 'status': job.status} + # 阶段 2: Shoter - 评估毒性活性 + logger.info(f"Job {job_id}: Starting Shoter stage") + job.current_stage = "shoter" + db.commit() + self.update_state( + state='PROGRESS', + meta={'stage': 'shoter', 'progress': 50, 'status': 'Running BtToxin_Shoter...'} + ) + + # TODO: 实现 Shoter 调用 + # shoter_result = run_shoter_pipeline(...) + # 暂时跳过 + logger.info(f"Job {job_id}: Shoter stage not implemented yet, skipping") + + job.progress_percent = 70 + db.commit() + + # 阶段 3: Plots - 生成热力图 + logger.info(f"Job {job_id}: Starting Plots stage") + job.current_stage = "plots" + db.commit() + self.update_state( + state='PROGRESS', + meta={'stage': 'plots', 'progress': 80, 'status': 'Generating plots...'} + ) + + # TODO: 实现 Plots 生成 + logger.info(f"Job {job_id}: Plots stage not implemented yet, skipping") + + job.progress_percent = 90 + db.commit() + + # 阶段 4: Bundle - 打包结果 + logger.info(f"Job {job_id}: Starting Bundle stage") + job.current_stage = "bundle" + db.commit() + self.update_state( + state='PROGRESS', + meta={'stage': 'bundle', 'progress': 95, 'status': 'Bundling results...'} + ) + + # 创建 manifest.json + import json + manifest = { + "job_id": job_id, + "stages_completed": ["digger"], + "stages_skipped": ["shoter", "plots", "bundle"], + "output_files": list(Path(output_dir).rglob("*")), + "parameters": { + "sequence_type": sequence_type, + "min_identity": min_identity, + "min_coverage": min_coverage, + "allow_unknown_families": allow_unknown_families, + "require_index_hit": require_index_hit, + } + } + + manifest_path = Path(output_dir) / "manifest.json" + with open(manifest_path, "w") as f: + json.dump(manifest, f, indent=2, default=str) + + # 完成 + job.status = JobStatus.COMPLETED + job.progress_percent = 100 + job.current_stage = "completed" + job.logs = json.dumps({"stages": ["digger"], "output": str(output_dir)}) + db.commit() + + logger.info(f"Job {job_id}: Completed successfully") + + return { + 'job_id': job_id, + 'status': 'completed', + 'stages': ['digger'], + 'output_dir': str(output_dir) + } except Exception as e: - logger.error(f"Task failed: {e}") + logger.error(f"Job {job_id} failed: {e}") job.status = JobStatus.FAILED job.error_message = str(e) + job.current_stage = "failed" db.commit() raise + + finally: + db.close() + + +@celery_app.task +def update_queue_positions(): + """ + 定期更新排队任务的位置 + 可以通过 Celery Beat 定期调用 + """ + db = SessionLocal() + try: + # 获取所有 QUEUED 状态的任务 + queued_jobs = db.query(Job).filter( + Job.status == JobStatus.QUEUED + ).order_by(Job.created_at).all() + + for idx, job in enumerate(queued_jobs, start=1): + job.queue_position = idx + + db.commit() + logger.info(f"Updated queue positions for {len(queued_jobs)} jobs") + + except Exception as e: + logger.error(f"Failed to update queue positions: {e}") + db.rollback() finally: db.close()