xxx
This commit is contained in:
@@ -0,0 +1,226 @@
|
||||
"""
|
||||
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,
|
||||
}
|
||||
Reference in New Issue
Block a user