223 lines
8.4 KiB
Python
223 lines
8.4 KiB
Python
"""
|
||
API 路由集合
|
||
============
|
||
|
||
把业务接口统一收口到 APIRouter，再由 main.py 通过 include_router 挂载。
|
||
|
||
当前包含的接口：
|
||
GET /api/algorithms 列出已注册的所有去耀斑算法（供前端下拉框）
|
||
POST /api/process/deglint 提交去耀斑处理任务，立即返回 task_id
|
||
GET /api/tasks/{task_id} 查询指定任务的状态与结果
|
||
|
||
派发链：
|
||
POST /api/process/deglint
|
||
└─ BackgroundTasks.add_task(execute_glint_removal_task, ...)
|
||
└─ get_remover(method) 从注册表拿到算法类
|
||
└─ remover.process(input_zarr, output_zarr, **params)
|
||
"""
|
||
|
||
import traceback
|
||
import uuid
|
||
from datetime import datetime
|
||
from typing import Any, Dict
|
||
|
||
from fastapi import APIRouter, BackgroundTasks, HTTPException
|
||
from pydantic import BaseModel, Field
|
||
|
||
# 并发安全的任务状态存储（替代旧版的 MOCK_TASK_DB）
|
||
from app.core.task_store import get_task, set_task, update_task
|
||
|
||
# 算法注册表 API
|
||
from app.core.algorithms import get_remover, list_removers
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 路由实例
|
||
# ---------------------------------------------------------------------------
|
||
# prefix 不在此处设置，统一在 main.py 挂载时给定，便于将来按版本拆分
|
||
# （例如 /api/v1、/api/v2 共存时复用同一个 router 对象）。
|
||
# ---------------------------------------------------------------------------
|
||
router = APIRouter(tags=["deglint"])
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 请求 / 响应数据模型
|
||
# ---------------------------------------------------------------------------
|
||
class DeglintRequest(BaseModel):
|
||
"""POST /api/process/deglint 的请求体"""
|
||
|
||
method: str = Field(
|
||
...,
|
||
description="去耀斑方法名称，必须是已注册算法，例如 'kutser' / 'goodman'",
|
||
examples=["kutser"],
|
||
)
|
||
params: Dict[str, Any] = Field(
|
||
default_factory=dict,
|
||
description=(
|
||
"传递给算法 process() 的超参数字典，例如 "
|
||
"Kutser: {'band_lower': 773, 'band_oxy': 845, 'band_upper': 893}; "
|
||
"Goodman: {'band_ref': 750, 'band_diff': 640, 'A': 0.0, 'B': 0.0}"
|
||
),
|
||
examples=[{"band_lower": 773, "band_oxy": 845, "band_upper": 893}],
|
||
)
|
||
|
||
|
||
class TaskAcceptedResponse(BaseModel):
|
||
"""提交任务成功后立即返回的响应"""
|
||
|
||
task_id: str
|
||
status: str # 一定是 PENDING
|
||
|
||
|
||
class AlgorithmListResponse(BaseModel):
|
||
"""GET /api/algorithms 的响应"""
|
||
|
||
algorithms: list # 已注册算法名列表
|
||
count: int # 算法总数
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 后台任务执行器（真实派发链）
|
||
# ---------------------------------------------------------------------------
|
||
# 注意：这里使用 async def。
|
||
# FastAPI / Starlette 的 BackgroundTasks 支持 async function，
|
||
# 会在响应返回后自动 await 它，不影响主请求链路。
|
||
# ---------------------------------------------------------------------------
|
||
async def execute_glint_removal_task(
|
||
task_id: str,
|
||
method: str,
|
||
params: Dict[str, Any],
|
||
) -> None:
|
||
"""
|
||
后台异步执行器：按 method 名字从注册表取出算法类，实例化并运行 process()。
|
||
|
||
状态机：
|
||
PENDING -> PROCESSING -> SUCCESS
|
||
└──> FAILED（含 error / traceback）
|
||
"""
|
||
# 0. 安全检查：任务记录必须已存在（POST 阶段已写入）
|
||
record = await get_task(task_id)
|
||
if record is None:
|
||
print(f"[{task_id}] 任务不存在, 跳过")
|
||
return
|
||
|
||
# 1. 状态推进到 PROCESSING
|
||
await update_task(
|
||
task_id,
|
||
status="PROCESSING",
|
||
updated_at=datetime.now().isoformat(),
|
||
)
|
||
print(f"[{task_id}] 开始处理 method={method} params={params}")
|
||
|
||
# 2. 临时硬编码 IO 路径（未来由数据管理层提供）
|
||
# TODO: 替换为真实的数据管理服务返回的 zarr 路径
|
||
input_zarr_path = "./data/temp_in.zarr"
|
||
output_zarr_path = f"./data/{task_id}_out.zarr"
|
||
|
||
try:
|
||
# 3. 按 method 名字从注册表取算法类并实例化
|
||
# get_remover 找不到时会抛 KeyError，下面的 except 会兜住
|
||
algorithm_cls = get_remover(method)
|
||
remover = algorithm_cls()
|
||
|
||
# 4. 调用算法（注意 await，因为 BaseGlintRemover.process 是 async）
|
||
await remover.process(input_zarr_path, output_zarr_path, **params)
|
||
|
||
# 5. 成功：写回结果路径与状态
|
||
await update_task(
|
||
task_id,
|
||
status="SUCCESS",
|
||
output_zarr_path=output_zarr_path,
|
||
error=None,
|
||
updated_at=datetime.now().isoformat(),
|
||
)
|
||
print(f"[{task_id}] 处理完成 -> SUCCESS, output={output_zarr_path}")
|
||
|
||
except Exception as exc: # noqa: BLE001 顶层兜底，绝不让后台任务静默失败
|
||
# 6. 失败：记录错误信息与堆栈，便于前端排查
|
||
await update_task(
|
||
task_id,
|
||
status="FAILED",
|
||
output_zarr_path=None,
|
||
error=f"{type(exc).__name__}: {exc}",
|
||
traceback=traceback.format_exc(),
|
||
updated_at=datetime.now().isoformat(),
|
||
)
|
||
print(f"[{task_id}] 处理失败 -> {type(exc).__name__}: {exc}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# GET /algorithms
|
||
# ---------------------------------------------------------------------------
|
||
# 返回当前已注册的所有算法名，供前端动态渲染下拉框 / 选择器。
|
||
# ---------------------------------------------------------------------------
|
||
@router.get("/algorithms", response_model=AlgorithmListResponse)
|
||
async def list_registered_algorithms() -> Dict[str, Any]:
|
||
"""列出已注册的去耀斑算法。"""
|
||
names = list(list_removers().keys())
|
||
return {"algorithms": names, "count": len(names)}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# POST /process/deglint
|
||
# ---------------------------------------------------------------------------
|
||
# 提交去耀斑处理任务。FastAPI 在函数返回后才会把响应发给前端，
|
||
# 因此通过 BackgroundTasks 把耗时操作丢到后台，接口本身立刻返回 task_id。
|
||
# ---------------------------------------------------------------------------
|
||
@router.post("/process/deglint", response_model=TaskAcceptedResponse)
|
||
async def submit_deglint(
|
||
payload: DeglintRequest,
|
||
background_tasks: BackgroundTasks,
|
||
) -> Dict[str, Any]:
|
||
"""提交一个去耀斑处理任务，并立即返回 task_id。"""
|
||
|
||
# 1. 生成唯一任务 ID（UUID4 足以保证全局唯一性）
|
||
task_id = str(uuid.uuid4())
|
||
|
||
# 2. 在任务库中登记一条 PENDING 记录（并发安全）
|
||
# 注意：output_zarr_path / error / traceback 字段在执行过程中被填充
|
||
await set_task(
|
||
task_id,
|
||
{
|
||
"task_id": task_id,
|
||
"method": payload.method,
|
||
"params": payload.params,
|
||
"status": "PENDING",
|
||
"output_zarr_path": None,
|
||
"error": None,
|
||
"traceback": None,
|
||
"created_at": datetime.now().isoformat(),
|
||
"updated_at": datetime.now().isoformat(),
|
||
},
|
||
)
|
||
|
||
# 3. 把真实执行器丢到后台
|
||
background_tasks.add_task(
|
||
execute_glint_removal_task,
|
||
task_id,
|
||
payload.method,
|
||
payload.params,
|
||
)
|
||
|
||
# 4. 立即返回 task_id 与 PENDING 状态
|
||
return {"task_id": task_id, "status": "PENDING"}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# GET /tasks/{task_id}
|
||
# ---------------------------------------------------------------------------
|
||
# 前端轮询此接口获取任务状态。PENDING / PROCESSING 表示仍在跑，
|
||
# SUCCESS 表示成功（含 output_zarr_path），FAILED 表示失败（含 error / traceback）。
|
||
# ---------------------------------------------------------------------------
|
||
@router.get("/tasks/{task_id}")
|
||
async def get_task_status(task_id: str) -> Dict[str, Any]:
|
||
"""查询指定任务的当前状态与结果。"""
|
||
|
||
record = await get_task(task_id)
|
||
if record is None:
|
||
# 找不到 task_id 通常意味着客户端拼错了 ID，或者记录已被清理
|
||
raise HTTPException(status_code=404, detail=f"task_id 不存在: {task_id}")
|
||
|
||
# 直接返回字典，FastAPI 会自动 JSON 序列化
|
||
return record
|