227 lines
7.4 KiB
Python
227 lines
7.4 KiB
Python
|
"""
|
||
|
|
moldinsight/api/batch_router.py — 批量分析端点
|
||
|
|
|
||
|
|
- POST /api/batch-upload 批量上传多文件，返回 batch_id + 各 task_id
|
||
|
|
- GET /api/batch/{batch_id} 聚合查询批量任务进度
|
||
|
|
"""
|
||
|
|
import uuid
|
||
|
|
from datetime import datetime
|
||
|
|
from typing import List, Dict, Any
|
||
|
|
|
||
|
|
from fastapi import APIRouter, UploadFile, File, Form, HTTPException, Depends
|
||
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
|
|
||
|
|
from shared.database.database import get_db_session
|
||
|
|
from shared.services.auth_service import get_current_active_user
|
||
|
|
from shared.models.database import User
|
||
|
|
from shared.models.schemas import ProcessingStatus, create_task_info
|
||
|
|
from shared.services.redis_task_manager import redis_task_manager
|
||
|
|
from shared.utils.file_handler import FileHandler
|
||
|
|
from shared.utils.logger import get_logger
|
||
|
|
from moldinsight.services.storage_integration_rustfs import StorageIntegrationService
|
||
|
|
|
||
|
|
try:
|
||
|
|
from celery_tasks import process_stp_task
|
||
|
|
_use_celery = True
|
||
|
|
except ImportError:
|
||
|
|
process_stp_task = None
|
||
|
|
_use_celery = False
|
||
|
|
|
||
|
|
logger = get_logger(__name__)
|
||
|
|
|
||
|
|
router = APIRouter()
|
||
|
|
|
||
|
|
file_handler = FileHandler()
|
||
|
|
|
||
|
|
# ─── 批量元数据 Redis key 约定 ──────────────────────────────────────
|
||
|
|
_BATCH_KEY_PREFIX = "batch:"
|
||
|
|
_BATCH_TTL = 86400 # 24h
|
||
|
|
|
||
|
|
|
||
|
|
def _batch_redis_key(batch_id: str) -> str:
|
||
|
|
return f"{_BATCH_KEY_PREFIX}{batch_id}"
|
||
|
|
|
||
|
|
|
||
|
|
@router.post("/batch-upload")
|
||
|
|
async def batch_upload(
|
||
|
|
files: List[UploadFile] = File(...),
|
||
|
|
material: str = Form("ABS"),
|
||
|
|
draft_angle: float = Form(2.0),
|
||
|
|
shrinkage_rate: float = Form(0.5),
|
||
|
|
parting_precision: float = Form(0.1),
|
||
|
|
cavity_match: int = Form(95),
|
||
|
|
db_session: AsyncSession = Depends(get_db_session),
|
||
|
|
current_user: User = Depends(get_current_active_user),
|
||
|
|
):
|
||
|
|
"""批量上传多个 STP 文件，每个文件创建独立分析任务，用 batch_id 聚合。"""
|
||
|
|
if not files:
|
||
|
|
raise HTTPException(400, "请至少上传一个文件")
|
||
|
|
if len(files) > 20:
|
||
|
|
raise HTTPException(400, "单次批量上传最多 20 个文件")
|
||
|
|
|
||
|
|
process_params = {
|
||
|
|
"material": material,
|
||
|
|
"draft_angle": float(draft_angle),
|
||
|
|
"shrinkage_rate": float(shrinkage_rate),
|
||
|
|
"parting_precision": float(parting_precision),
|
||
|
|
"cavity_match": int(cavity_match),
|
||
|
|
}
|
||
|
|
|
||
|
|
batch_id = str(uuid.uuid4())
|
||
|
|
tasks: List[Dict[str, Any]] = []
|
||
|
|
storage_service = StorageIntegrationService()
|
||
|
|
|
||
|
|
for file in files:
|
||
|
|
# 文件类型检查
|
||
|
|
if not file.filename.lower().endswith(('.stp', '.step')):
|
||
|
|
tasks.append({
|
||
|
|
"filename": file.filename,
|
||
|
|
"task_id": None,
|
||
|
|
"status": "rejected",
|
||
|
|
"error": "不支持的文件类型",
|
||
|
|
})
|
||
|
|
continue
|
||
|
|
|
||
|
|
task_id = str(uuid.uuid4())
|
||
|
|
try:
|
||
|
|
file_path, file_size, file_meta = await file_handler.save_uploaded_file(file)
|
||
|
|
|
||
|
|
stp_file = await storage_service.save_stp_file(
|
||
|
|
session=db_session,
|
||
|
|
file_path=file_path,
|
||
|
|
original_filename=file_meta["safe_original_name"],
|
||
|
|
user_id=current_user.id,
|
||
|
|
)
|
||
|
|
|
||
|
|
await storage_service.create_processing_task(
|
||
|
|
db_session, task_id, stp_file.id, parameters=process_params,
|
||
|
|
)
|
||
|
|
|
||
|
|
task_info = create_task_info(
|
||
|
|
task_id=task_id,
|
||
|
|
status=ProcessingStatus.PROCESSING,
|
||
|
|
filename=file.filename,
|
||
|
|
file_path=str(file_path),
|
||
|
|
file_size=file_size,
|
||
|
|
upload_time=str(datetime.now()),
|
||
|
|
)
|
||
|
|
task_info["material"] = material
|
||
|
|
task_info["parameters"] = process_params
|
||
|
|
task_info["batch_id"] = batch_id
|
||
|
|
await redis_task_manager.set_task(task_id, task_info)
|
||
|
|
|
||
|
|
# 调度处理
|
||
|
|
if _use_celery:
|
||
|
|
process_stp_task.delay(task_id, str(file_path), stp_file.id, process_params)
|
||
|
|
else:
|
||
|
|
import asyncio
|
||
|
|
from moldinsight.services.processing_service import processing_service
|
||
|
|
asyncio.create_task(processing_service.process_file_with_storage(
|
||
|
|
task_id, str(file_path), stp_file.id, process_params
|
||
|
|
))
|
||
|
|
|
||
|
|
tasks.append({
|
||
|
|
"filename": file.filename,
|
||
|
|
"task_id": task_id,
|
||
|
|
"status": "processing",
|
||
|
|
"stp_file_id": stp_file.id,
|
||
|
|
})
|
||
|
|
logger.info(
|
||
|
|
f"[BATCH] batch_id={batch_id} task_id={task_id} "
|
||
|
|
f"file={file.filename} user={current_user.username}"
|
||
|
|
)
|
||
|
|
except Exception as exc:
|
||
|
|
logger.warning(f"[BATCH] 文件 {file.filename} 上传失败: {exc}")
|
||
|
|
tasks.append({
|
||
|
|
"filename": file.filename,
|
||
|
|
"task_id": task_id,
|
||
|
|
"status": "error",
|
||
|
|
"error": str(exc),
|
||
|
|
})
|
||
|
|
|
||
|
|
# 将 batch 元数据写入 Redis
|
||
|
|
batch_meta = {
|
||
|
|
"batch_id": batch_id,
|
||
|
|
"user_id": current_user.id,
|
||
|
|
"created_at": str(datetime.now()),
|
||
|
|
"task_ids": [t["task_id"] for t in tasks if t.get("task_id")],
|
||
|
|
"total": len(tasks),
|
||
|
|
"params": process_params,
|
||
|
|
}
|
||
|
|
await redis_task_manager.redis_client.set(
|
||
|
|
_batch_redis_key(batch_id),
|
||
|
|
__import__("json").dumps(batch_meta),
|
||
|
|
ex=_BATCH_TTL,
|
||
|
|
)
|
||
|
|
|
||
|
|
return {
|
||
|
|
"batch_id": batch_id,
|
||
|
|
"total": len(tasks),
|
||
|
|
"accepted": sum(1 for t in tasks if t.get("status") != "rejected"),
|
||
|
|
"tasks": tasks,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
@router.get("/batch/{batch_id}")
|
||
|
|
async def get_batch_status(
|
||
|
|
batch_id: str,
|
||
|
|
current_user: User = Depends(get_current_active_user),
|
||
|
|
):
|
||
|
|
"""聚合查询批量任务进度"""
|
||
|
|
import json
|
||
|
|
|
||
|
|
raw = await redis_task_manager.redis_client.get(_batch_redis_key(batch_id))
|
||
|
|
if not raw:
|
||
|
|
raise HTTPException(404, "批量任务不存在或已过期")
|
||
|
|
|
||
|
|
batch_meta = json.loads(raw)
|
||
|
|
|
||
|
|
# 权限检查
|
||
|
|
if batch_meta.get("user_id") and batch_meta["user_id"] != current_user.id:
|
||
|
|
raise HTTPException(403, "无权访问该批量任务")
|
||
|
|
|
||
|
|
task_ids = batch_meta.get("task_ids", [])
|
||
|
|
task_statuses = []
|
||
|
|
completed = 0
|
||
|
|
failed = 0
|
||
|
|
processing = 0
|
||
|
|
|
||
|
|
for tid in task_ids:
|
||
|
|
task_data = await redis_task_manager.get_task(tid)
|
||
|
|
if not task_data:
|
||
|
|
task_statuses.append({"task_id": tid, "status": "unknown"})
|
||
|
|
continue
|
||
|
|
status = task_data.get("status", "unknown")
|
||
|
|
progress = task_data.get("progress", 0)
|
||
|
|
filename = task_data.get("filename", "")
|
||
|
|
error = task_data.get("error", "")
|
||
|
|
html_file = task_data.get("html_file", "")
|
||
|
|
|
||
|
|
if status == ProcessingStatus.COMPLETED:
|
||
|
|
completed += 1
|
||
|
|
elif status == ProcessingStatus.FAILED:
|
||
|
|
failed += 1
|
||
|
|
else:
|
||
|
|
processing += 1
|
||
|
|
|
||
|
|
task_statuses.append({
|
||
|
|
"task_id": tid,
|
||
|
|
"status": status,
|
||
|
|
"progress": progress,
|
||
|
|
"filename": filename,
|
||
|
|
"error": error,
|
||
|
|
"html_file": html_file,
|
||
|
|
})
|
||
|
|
|
||
|
|
total = len(task_ids)
|
||
|
|
return {
|
||
|
|
"batch_id": batch_id,
|
||
|
|
"created_at": batch_meta.get("created_at"),
|
||
|
|
"total": total,
|
||
|
|
"completed": completed,
|
||
|
|
"failed": failed,
|
||
|
|
"processing": processing,
|
||
|
|
"progress_percent": round((completed + failed) / max(total, 1) * 100, 1),
|
||
|
|
"tasks": task_statuses,
|
||
|
|
}
|