319 lines
9.6 KiB
Python
319 lines
9.6 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, BackgroundTasks, Query
|
||
from pydantic import BaseModel, Field
|
||
import uvicorn
|
||
import logging
|
||
from datetime import datetime
|
||
from typing import Optional
|
||
import subprocess
|
||
from pathlib import Path
|
||
from config.settings import API_HOST, API_CALLBACK
|
||
|
||
# ==================== 日志配置 ====================
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s | %(levelname)s | %(message)s'
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ==================== FastAPI 应用 ====================
|
||
app = FastAPI(
|
||
title="皮肤分析任务接口",
|
||
description="接收文件夹名称，触发皮肤分析任务",
|
||
version="1.0.0"
|
||
)
|
||
|
||
# ==================== 数据存储 ====================
|
||
task_store = {}
|
||
|
||
# ==================== 数据模型 ====================
|
||
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=失败")
|
||
|
||
# ==================== 后台任务 ====================
|
||
def run_analysis_task(task_id: str, request: TaskRequest):
|
||
"""后台运行分析任务"""
|
||
|
||
logger.info(f"🚀 开始执行任务：{task_id}")
|
||
|
||
try:
|
||
task_store[task_id]["status"] = "running"
|
||
|
||
# ==================== 输入目录 ====================
|
||
|
||
|
||
if Path(request.folder_name).is_absolute():
|
||
# 绝对路径
|
||
input_dir = Path(request.folder_name)
|
||
else:
|
||
# 相对路径
|
||
base_dir = "/data/analysis"
|
||
input_dir = (
|
||
Path(base_dir)
|
||
/ datetime.now().strftime("%Y%m%d")
|
||
/ task_id
|
||
)
|
||
|
||
# 检查目录是否存在
|
||
if not input_dir.exists():
|
||
raise Exception(f"输入目录不存在: {input_dir}")
|
||
|
||
# ==================== 脚本路径 ====================
|
||
|
||
script_path = Path(__file__).parent / "wanzheng.py"
|
||
|
||
if not script_path.exists():
|
||
raise Exception(f"分析脚本不存在: {script_path}")
|
||
|
||
# ==================== 模型路径 ====================
|
||
|
||
model_path = "/app/best.pt"
|
||
|
||
if not Path(model_path).exists():
|
||
raise Exception(f"模型文件不存在: {model_path}")
|
||
|
||
# ==================== 构建命令 ====================
|
||
|
||
cmd = [
|
||
sys.executable,
|
||
|
||
str(script_path),
|
||
|
||
str(input_dir),
|
||
|
||
"--model", model_path,
|
||
|
||
"--workers", "4",
|
||
|
||
"--api-url", API_HOST + API_CALLBACK,
|
||
|
||
"--task-id", task_id,
|
||
]
|
||
|
||
logger.info(f"🔧 命令参数：{cmd}")
|
||
logger.info(f"📋 执行命令：{' '.join(cmd)}")
|
||
logger.info(f"📁 输入目录：{input_dir}")
|
||
|
||
# ==================== 执行脚本 ====================
|
||
|
||
result = subprocess.run(
|
||
cmd,
|
||
capture_output=True,
|
||
text=True,
|
||
encoding="utf-8",
|
||
errors="replace",
|
||
timeout=3600
|
||
)
|
||
|
||
# ==================== 输出日志 ====================
|
||
|
||
if result.stdout:
|
||
logger.info(f"📋 标准输出：\n{result.stdout}")
|
||
|
||
if result.stderr:
|
||
logger.error(f"📋 错误输出：\n{result.stderr}")
|
||
|
||
# ==================== 判断执行结果 ====================
|
||
|
||
if result.returncode == 0:
|
||
|
||
task_store[task_id]["status"] = "completed"
|
||
|
||
logger.info(f"✅ 任务完成：{task_id}")
|
||
|
||
else:
|
||
|
||
task_store[task_id]["status"] = "failed"
|
||
|
||
task_store[task_id]["error"] = result.stderr
|
||
|
||
logger.error(
|
||
f"❌ 任务失败：{task_id} "
|
||
f"(code={result.returncode})"
|
||
)
|
||
|
||
task_store[task_id]["completed_at"] = (
|
||
datetime.now().isoformat()
|
||
)
|
||
|
||
except subprocess.TimeoutExpired:
|
||
|
||
task_store[task_id]["status"] = "failed"
|
||
|
||
task_store[task_id]["error"] = "任务超时（超过1小时）"
|
||
|
||
task_store[task_id]["completed_at"] = (
|
||
datetime.now().isoformat()
|
||
)
|
||
|
||
logger.error(f"⏰ 任务超时：{task_id}")
|
||
|
||
except Exception as e:
|
||
|
||
task_store[task_id]["status"] = "failed"
|
||
|
||
task_store[task_id]["error"] = str(e)
|
||
|
||
task_store[task_id]["completed_at"] = (
|
||
datetime.now().isoformat()
|
||
)
|
||
|
||
logger.exception(f"❌ 任务异常：{task_id}")
|
||
|
||
# ==================== API 接口 ====================
|
||
@app.get("/")
|
||
async def root():
|
||
return {
|
||
"service": "皮肤分析任务接口",
|
||
"version": "1.0.0",
|
||
"status": "running"
|
||
}
|
||
|
||
@app.get("/health")
|
||
async def health_check():
|
||
return {
|
||
"status": "healthy",
|
||
"timestamp": datetime.now().isoformat(),
|
||
"tasks": {
|
||
"total": len(task_store),
|
||
"running": sum(1 for t in task_store.values() if t["status"] == "running"),
|
||
"completed": sum(1 for t in task_store.values() if t["status"] == "completed"),
|
||
"failed": sum(1 for t in task_store.values() if t["status"] == "failed")
|
||
}
|
||
}
|
||
|
||
@app.post("/api/v1/analysis/submit", response_model=TaskResponse)
|
||
async def submit_task(request: TaskRequest, background_tasks: BackgroundTasks):
|
||
# 使用传入的 task_id，如果没有则从 folder_name 提取
|
||
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 个字符")
|
||
|
||
if task_id in task_store:
|
||
existing_task = task_store[task_id]
|
||
if existing_task["status"] in ["pending", "running"]:
|
||
return TaskResponse(
|
||
code=200,
|
||
message="任务已存在，正在运行中",
|
||
task_id=task_id,
|
||
status="pending"
|
||
)
|
||
|
||
task_store[task_id] = {
|
||
"task_id": task_id,
|
||
"status": "pending",
|
||
"created_at": datetime.now().isoformat(),
|
||
"completed_at": None,
|
||
"error": None,
|
||
"request": {
|
||
"folder_name": request.folder_name,
|
||
"task_id": request.task_id
|
||
}
|
||
}
|
||
|
||
background_tasks.add_task(run_analysis_task, task_id, request)
|
||
|
||
logger.info(f"📝 任务已提交：{task_id}")
|
||
|
||
return TaskResponse(
|
||
code=200,
|
||
message="任务提交成功",
|
||
task_id=task_id,
|
||
status="pending"
|
||
)
|
||
|
||
@app.post("/api/v1/analysis/callback")
|
||
async def analysis_callback(request: CallbackRequest):
|
||
task_id = request.task_id
|
||
status = "completed" if request.status == 1 else "failed"
|
||
|
||
if task_id not in task_store:
|
||
task_store[task_id] = {
|
||
"task_id": task_id,
|
||
"status": status,
|
||
"created_at": datetime.now().isoformat(),
|
||
"completed_at": datetime.now().isoformat(),
|
||
"error": None if request.status == 1 else "分析任务回调失败",
|
||
"request": None
|
||
}
|
||
else:
|
||
task_store[task_id]["status"] = status
|
||
task_store[task_id]["completed_at"] = datetime.now().isoformat()
|
||
task_store[task_id]["error"] = None if request.status == 1 else "分析任务回调失败"
|
||
|
||
logger.info(f"📩 收到分析回调：{task_id} status={request.status}")
|
||
return {"code": 200, "message": "回调接收成功", "task_id": task_id, "status": status}
|
||
|
||
@app.get("/api/v1/analysis/status/{task_id}", response_model=TaskStatus)
|
||
async def get_task_status(task_id: str):
|
||
if task_id not in task_store:
|
||
raise HTTPException(status_code=404, detail=f"任务不存在：{task_id}")
|
||
|
||
task = task_store[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")
|
||
async def list_tasks(
|
||
status: Optional[str] = Query(default=None),
|
||
limit: int = Query(default=20, ge=1, le=100)
|
||
):
|
||
tasks = list(task_store.values())
|
||
|
||
if status:
|
||
tasks = [t for t in tasks if t["status"] == status]
|
||
|
||
tasks = sorted(tasks, key=lambda x: x["created_at"], reverse=True)[: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)
|