161 lines
5.7 KiB
Python
161 lines
5.7 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
皮肤分析 API 接口服务
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
from fastapi import FastAPI, HTTPException, Query
|
||
from contextlib import asynccontextmanager
|
||
import asyncio
|
||
from pydantic import BaseModel, Field
|
||
import uvicorn
|
||
import logging
|
||
from datetime import datetime
|
||
from typing import Optional
|
||
from pathlib import Path
|
||
from config.settings import (API_HOST, API_CALLBACK, TASK_DB_PATH,
|
||
MAX_CONCURRENT_TASKS, MAX_PENDING_TASKS, MODEL_PATH, IMAGE_WORKERS)
|
||
from task_queue import TaskQueue, QueueFull
|
||
from resident_worker import ResidentPool
|
||
|
||
# ==================== 日志配置 ====================
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s | %(levelname)s | %(message)s'
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ==================== 服务生命周期 ====================
|
||
@asynccontextmanager
|
||
async def lifespan(app):
|
||
await asyncio.to_thread(task_queue.start)
|
||
try:
|
||
yield
|
||
finally:
|
||
await asyncio.to_thread(task_queue.stop)
|
||
|
||
|
||
app = FastAPI(
|
||
title="皮肤分析任务接口",
|
||
description="接收文件夹名称，触发皮肤分析任务",
|
||
version="1.0.0",
|
||
lifespan=lifespan,
|
||
)
|
||
|
||
# ==================== 数据模型 ====================
|
||
class TaskRequest(BaseModel):
|
||
folder_name: str = Field(..., description="文件夹全路径")
|
||
task_id: str = Field(..., description="目标文件夹名称（任务 ID）")
|
||
|
||
class TaskResponse(BaseModel):
|
||
code: int = Field(..., description="状态码")
|
||
message: str = Field(..., description="消息")
|
||
task_id: str = Field(..., description="任务 ID")
|
||
status: str = Field(..., description="任务状态")
|
||
|
||
class TaskStatus(BaseModel):
|
||
task_id: str = Field(..., description="任务 ID")
|
||
status: str = Field(..., description="状态")
|
||
created_at: str = Field(..., description="创建时间")
|
||
completed_at: Optional[str] = Field(default=None, description="完成时间")
|
||
error: Optional[str] = Field(default=None, description="错误信息")
|
||
|
||
class CallbackRequest(BaseModel):
|
||
task_id: str = Field(..., description="任务 ID")
|
||
status: int = Field(..., description="处理状态，1=成功，0=失败")
|
||
|
||
# 每个任务槽位对应一个常驻模型进程；启动预热后才开放 API。
|
||
analysis_pool = ResidentPool(MODEL_PATH, MAX_CONCURRENT_TASKS, IMAGE_WORKERS,
|
||
API_HOST + API_CALLBACK)
|
||
task_queue = TaskQueue(TASK_DB_PATH, MAX_CONCURRENT_TASKS,
|
||
MAX_PENDING_TASKS, analysis_pool)
|
||
|
||
# ==================== API 接口 ====================
|
||
@app.get("/")
|
||
async def root():
|
||
return {
|
||
"service": "皮肤分析任务接口",
|
||
"version": "1.0.0",
|
||
"status": "running"
|
||
}
|
||
|
||
@app.get("/health")
|
||
def health_check():
|
||
return {
|
||
"status": "healthy",
|
||
"timestamp": datetime.now().isoformat(),
|
||
"tasks": task_queue.counts(),
|
||
"queue": {"max_concurrent": task_queue.concurrency,
|
||
"max_pending": task_queue.max_pending}
|
||
}
|
||
|
||
|
||
@app.post("/api/v1/analysis/submit", response_model=TaskResponse)
|
||
def submit_task(request: TaskRequest):
|
||
task_id = request.task_id if request.task_id else Path(request.folder_name).name
|
||
if not task_id or len(task_id) < 3:
|
||
raise HTTPException(status_code=400, detail="任务 ID 格式错误，至少 3 个字符")
|
||
# Resolve the date at submission, not after a queued task starts tomorrow.
|
||
input_dir = Path(request.folder_name)
|
||
if not input_dir.is_absolute():
|
||
input_dir = Path("/data/analysis") / datetime.now().strftime("%Y%m%d") / task_id
|
||
try:
|
||
task, created = task_queue.submit(
|
||
task_id, {"folder_name": request.folder_name, "task_id": task_id},
|
||
str(input_dir.resolve()))
|
||
except QueueFull as exc:
|
||
raise HTTPException(status_code=429, detail=str(exc),
|
||
headers={"Retry-After": "5"}) from exc
|
||
return TaskResponse(code=200,
|
||
message="任务已加入队列" if created else "任务已存在，正在等待或运行中",
|
||
task_id=task_id, status=task["status"])
|
||
|
||
@app.post("/api/v1/analysis/callback")
|
||
def analysis_callback(request: CallbackRequest):
|
||
if request.status not in (0, 1):
|
||
raise HTTPException(status_code=400, detail="状态只允许 0 或 1")
|
||
status = task_queue.callback(request.task_id, request.status)
|
||
if status is None:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
return {"code": 200, "message": "回调接收成功", "task_id": request.task_id, "status": status}
|
||
|
||
@app.get("/api/v1/analysis/status/{task_id}", response_model=TaskStatus)
|
||
def get_task_status(task_id: str):
|
||
task = task_queue.get(task_id)
|
||
if task is None:
|
||
raise HTTPException(status_code=404, detail=f"任务不存在：{task_id}")
|
||
|
||
return TaskStatus(
|
||
task_id=task["task_id"],
|
||
status=task["status"],
|
||
created_at=task["created_at"],
|
||
completed_at=task.get("completed_at"),
|
||
error=task.get("error")
|
||
)
|
||
|
||
@app.get("/api/v1/analysis/list")
|
||
def list_tasks(
|
||
status: Optional[str] = Query(default=None),
|
||
limit: int = Query(default=20, ge=1, le=100)
|
||
):
|
||
tasks = task_queue.list(status, limit)
|
||
return {"total": len(tasks), "tasks": tasks}
|
||
|
||
# ==================== 启动服务 ====================
|
||
if __name__ == "__main__":
|
||
print("="*60)
|
||
print("皮肤分析 API 服务")
|
||
print("="*60)
|
||
print("监听地址：http://0.0.0.0:8000")
|
||
print("提交接口：POST /api/v1/analysis/submit")
|
||
print("状态查询：GET /api/v1/analysis/status/{task_id}")
|
||
print("健康检查：GET /health")
|
||
print("="*60)
|
||
print()
|
||
|
||
uvicorn.run(app, host="0.0.0.0", port=8000)
|