优化时间

api_server.py

@@ -6,21 +6,41 @@
 
 import sys
 import os
+<<<<<<< HEAD
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
 from fastapi import FastAPI, HTTPException, Query
 from contextlib import asynccontextmanager
 import asyncio
+=======
+import threading
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from fastapi import FastAPI, HTTPException, BackgroundTasks, Query
+>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c
 from pydantic import BaseModel, Field
 import uvicorn
 import logging
 from datetime import datetime
 from typing import Optional
+<<<<<<< HEAD
 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
+=======
+import subprocess
+from pathlib import Path
+from config.settings import (
+    ANALYSIS_ROOT,
+    API_HOST,
+    API_CALLBACK,
+    MODEL_PATH,
+    PORE_ENABLED,
+    PORE_PREDICTOR_PATH,
+)
+>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c
 
 # ==================== 日志配置 ====================
 logging.basicConfig(
@@ -29,6 +49,7 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+<<<<<<< HEAD
 # ==================== 服务生命周期 ====================
 @asynccontextmanager
 async def lifespan(app):
@@ -46,10 +67,31 @@ app = FastAPI(
     lifespan=lifespan,
 )
 
+=======
+# ==================== FastAPI 应用 ====================
+app = FastAPI(
+    title="皮肤分析任务接口",
+    description="接收文件夹名称，触发皮肤分析任务",
+    version="1.0.0"
+)
+
+# ==================== 数据存储 ====================
+task_store = {}
+active_tasks = set()
+task_store_lock = threading.RLock()
+
+>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c
 # ==================== 数据模型 ====================
 class TaskRequest(BaseModel):
     folder_name: str = Field(..., description="文件夹全路径")
     task_id: str = Field(..., description="目标文件夹名称（任务 ID）")
+<<<<<<< HEAD
+=======
+    task_type: Optional[str] = Field(
+        default=None,
+        description="任务类型，pore 表示仅执行大毛孔检测",
+    )
+>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c
 
 class TaskResponse(BaseModel):
     code: int = Field(..., description="状态码")
@@ -66,6 +108,7 @@ class TaskStatus(BaseModel):
 
 class CallbackRequest(BaseModel):
     task_id: str = Field(..., description="任务 ID")
+<<<<<<< HEAD
     status: int = Field(..., description="处理状态，1=成功，0=失败")
 
 # 每个任务槽位对应一个常驻模型进程；启动预热后才开放 API。
@@ -73,6 +116,203 @@ 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)
+=======
+    status: int = Field(..., ge=0, le=1, description="处理状态，1=成功，0=失败")
+
+
+def validate_task_id(task_id: str) -> str:
+    """校验会进入文件路径、日志和命令参数的任务 ID。"""
+    if not task_id or len(task_id) < 3 or len(task_id) > 128:
+        raise ValueError("任务 ID 格式错误，长度必须为 3 到 128 个字符")
+    if task_id in {".", ".."} or task_id.startswith("-"):
+        raise ValueError("任务 ID 格式错误")
+    if any(char in task_id for char in ("/", "\\")):
+        raise ValueError("任务 ID 不能包含路径分隔符")
+    if any(ord(char) < 32 or ord(char) == 127 for char in task_id):
+        raise ValueError("任务 ID 不能包含控制字符")
+    return task_id
+
+
+def normalize_task_type(task_type: Optional[str]) -> Optional[str]:
+    """规范化可选任务类型；空值保持现有完整分析流程。"""
+    normalized = (task_type or "").strip().lower()
+    if not normalized:
+        return None
+    if normalized != "pore":
+        raise ValueError(f"不支持的任务类型：{task_type}")
+    return normalized
+
+
+def resolve_input_dir(folder_name: str, task_id: str, now: Optional[datetime] = None) -> Path:
+    """解析输入目录，并确保目录始终位于分析数据根目录内。"""
+    root = Path(ANALYSIS_ROOT).resolve()
+    candidate = Path(folder_name)
+    if not candidate.is_absolute():
+        candidate = root / (now or datetime.now()).strftime("%Y%m%d") / task_id
+
+    candidate = candidate.resolve()
+    try:
+        candidate.relative_to(root)
+    except ValueError as exc:
+        raise ValueError(f"输入目录必须位于 {root} 内") from exc
+    if candidate == root:
+        raise ValueError("输入目录不能是分析数据根目录本身")
+    return candidate
+
+
+def update_task(task_id: str, **updates) -> None:
+    """线程安全地更新仍存在的任务。"""
+    with task_store_lock:
+        task = task_store.get(task_id)
+        if task is not None:
+            task.update(updates)
+
+
+def finish_task(task_id: str, status: str, error: Optional[str]) -> None:
+    """原子写入最终状态并解除运行标记，避免迟到回调覆盖结果。"""
+    with task_store_lock:
+        task = task_store.get(task_id)
+        if task is not None:
+            task.update(
+                status=status,
+                error=error,
+                completed_at=datetime.now().isoformat(),
+            )
+        active_tasks.discard(task_id)
+
+# ==================== 后台任务 ====================
+def run_analysis_task(
+    task_id: str,
+    request: TaskRequest,
+    input_dir: Path,
+    task_type: Optional[str] = None,
+):
+    """后台运行分析任务"""
+    
+    logger.info(f"🚀 开始执行任务：{task_id}")
+
+    try:
+        update_task(task_id, status="running")
+
+        # 再次解析路径，防止任务排队期间目录被替换为越界符号链接。
+        revalidated_input_dir = resolve_input_dir(str(input_dir), task_id)
+        if revalidated_input_dir != input_dir:
+            raise Exception("输入目录在任务执行前发生变化")
+        input_dir = revalidated_input_dir
+
+        # 检查目录是否存在
+        if not input_dir.is_dir():
+            raise Exception(f"输入目录不存在或不是目录: {input_dir}")
+
+        # ==================== 脚本路径 ====================
+
+        script_path = Path(__file__).parent / "wanzheng.py"
+
+        if not script_path.exists():
+            raise Exception(f"分析脚本不存在: {script_path}")
+
+        # ==================== 模型路径 ====================
+
+        pore_only = task_type == "pore"
+        model_path = Path(MODEL_PATH)
+
+        if not pore_only and not model_path.is_file():
+            raise Exception(f"模型文件不存在: {model_path}")
+
+        # ==================== 构建命令 ====================
+
+        cmd = [
+            sys.executable,
+
+            str(script_path),
+
+            str(input_dir),
+
+            "--workers", "4",
+
+            "--analysis-root", str(Path(ANALYSIS_ROOT).resolve()),
+
+            "--api-url", f"{API_HOST.rstrip('/')}/{API_CALLBACK.lstrip('/')}",
+
+            "--task-id", task_id,
+        ]
+
+        if pore_only:
+            cmd.extend([
+                "--task-type", "pore",
+                "--pore-predictor", PORE_PREDICTOR_PATH,
+            ])
+        else:
+            cmd.extend(["--model", str(model_path)])
+            if PORE_ENABLED:
+                cmd.extend(["--pore-predictor", PORE_PREDICTOR_PATH])
+            else:
+                cmd.append("--disable-pores")
+
+        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:
+            if result.returncode == 0:
+                logger.info(f"📋 进程日志：\n{result.stderr}")
+            else:
+                logger.error(f"📋 错误输出：\n{result.stderr}")
+
+        # ==================== 判断执行结果 ====================
+
+        if result.returncode == 0:
+            finish_task(task_id, status="completed", error=None)
+
+            logger.info(f"✅ 任务完成：{task_id}")
+
+        else:
+
+            finish_task(
+                task_id,
+                status="failed",
+                error=result.stderr or f"分析进程退出码：{result.returncode}"
+            )
+
+            logger.error(
+                f"❌ 任务失败：{task_id} "
+                f"(code={result.returncode})"
+            )
+
+    except subprocess.TimeoutExpired:
+        finish_task(
+            task_id,
+            status="failed",
+            error="任务超时（超过1小时）"
+        )
+
+        logger.error(f"⏰ 任务超时：{task_id}")
+
+    except Exception as e:
+        finish_task(
+            task_id,
+            status="failed",
+            error=str(e)
+        )
+
+        logger.exception(f"❌ 任务异常：{task_id}")
+>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c
 
 # ==================== API 接口 ====================
 @app.get("/")
@@ -84,6 +324,7 @@ async def root():
     }
 
 @app.get("/health")
+<<<<<<< HEAD
 def health_check():
     return {
         "status": "healthy",
@@ -129,6 +370,118 @@ def get_task_status(task_id: str):
     if task is None:
         raise HTTPException(status_code=404, detail=f"任务不存在：{task_id}")
 
+=======
+async def health_check():
+    with task_store_lock:
+        statuses = [task["status"] for task in task_store.values()]
+    return {
+        "status": "healthy",
+        "timestamp": datetime.now().isoformat(),
+        "tasks": {
+            "total": len(statuses),
+            "pending": statuses.count("pending"),
+            "running": statuses.count("running"),
+            "completed": statuses.count("completed"),
+            "failed": statuses.count("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
+
+    try:
+        validate_task_id(task_id)
+        task_type = normalize_task_type(request.task_type)
+        input_dir = resolve_input_dir(request.folder_name, task_id)
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+    with task_store_lock:
+        if task_id in active_tasks:
+            existing_task = task_store[task_id]
+            return TaskResponse(
+                code=200,
+                message="任务已存在，正在运行中",
+                task_id=task_id,
+                status=existing_task["status"]
+            )
+
+        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,
+                "task_type": task_type,
+            }
+        }
+        active_tasks.add(task_id)
+
+    background_tasks.add_task(
+        run_analysis_task,
+        task_id,
+        request,
+        input_dir,
+        task_type,
+    )
+    
+    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):
+    try:
+        task_id = validate_task_id(request.task_id)
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+    callback_status = "completed" if request.status == 1 else "failed"
+
+    with task_store_lock:
+        if task_id not in task_store:
+            task_store[task_id] = {
+                "task_id": task_id,
+                "status": callback_status,
+                "created_at": datetime.now().isoformat(),
+                "completed_at": datetime.now().isoformat(),
+                "error": None if request.status == 1 else "分析任务回调失败",
+                "request": None
+            }
+        else:
+            current_status = task_store[task_id]["status"]
+            is_finished = current_status in {"completed", "failed"} and task_id not in active_tasks
+            if not is_finished:
+                task_store[task_id]["status"] = callback_status
+                task_store[task_id]["completed_at"] = datetime.now().isoformat()
+                task_store[task_id]["error"] = (
+                    None if request.status == 1 else "分析任务回调失败"
+                )
+
+        status = task_store[task_id]["status"]
+
+    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):
+    with task_store_lock:
+        task = task_store.get(task_id)
+        if task is not None:
+            task = dict(task)
+    if task is None:
+        raise HTTPException(status_code=404, detail=f"任务不存在：{task_id}")
+>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c
     return TaskStatus(
         task_id=task["task_id"],
         status=task["status"],
@@ -138,12 +491,33 @@ def get_task_status(task_id: str):
     )
 
 @app.get("/api/v1/analysis/list")
+<<<<<<< HEAD
 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}
+=======
+async def list_tasks(
+    status: Optional[str] = Query(default=None),
+    limit: int = Query(default=20, ge=1, le=100)
+):
+    with task_store_lock:
+        tasks = [dict(task) for task in 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)
+    total = len(tasks)
+    tasks = tasks[:limit]
+    
+    return {
+        "total": total,
+        "tasks": tasks
+    }
+>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c
 
 # ==================== 启动服务 ====================
 if __name__ == "__main__":