117 lines
3.7 KiB
Python
117 lines
3.7 KiB
Python
import logging
|
||
import sys
|
||
from typing import Optional
|
||
from fastapi import APIRouter, HTTPException, Query
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
import threading
|
||
|
||
from app.schemas.response_template import ResponseModel
|
||
from app.service.recommendation_system.import_sys_sketch_to_milvus import main as import_main
|
||
|
||
logger = logging.getLogger()
|
||
router = APIRouter()
|
||
|
||
# 使用线程池执行器来运行长时间任务
|
||
executor = ThreadPoolExecutor(max_workers=1)
|
||
# 用于跟踪任务状态
|
||
task_status = {"running": False}
|
||
|
||
|
||
def run_import_task(batch_size: int, retry_times: int, limit: Optional[int], offset: int, skip_create_collection: bool):
|
||
"""在后台线程中运行导入任务"""
|
||
original_argv = None
|
||
try:
|
||
task_status["running"] = True
|
||
# 保存原始 sys.argv
|
||
original_argv = sys.argv.copy()
|
||
|
||
# 模拟命令行参数
|
||
sys.argv = [
|
||
"import_sys_sketch_to_milvus.py",
|
||
"--batch-size", str(batch_size),
|
||
"--retry-times", str(retry_times),
|
||
]
|
||
if limit is not None:
|
||
sys.argv.extend(["--limit", str(limit)])
|
||
if offset > 0:
|
||
sys.argv.extend(["--offset", str(offset)])
|
||
if skip_create_collection:
|
||
sys.argv.append("--skip-create-collection")
|
||
|
||
import_main()
|
||
task_status["running"] = False
|
||
logger.info("导入任务完成")
|
||
except Exception as e:
|
||
task_status["running"] = False
|
||
logger.error(f"导入任务失败: {e}", exc_info=True)
|
||
raise
|
||
finally:
|
||
# 恢复原始 sys.argv
|
||
if original_argv is not None:
|
||
sys.argv = original_argv
|
||
|
||
|
||
@router.post("/import-sys-sketch", response_model=ResponseModel)
|
||
async def import_sys_sketch(
|
||
batch_size: int = Query(1000, description="批量处理大小（默认：1000）"),
|
||
retry_times: int = Query(3, description="失败重试次数（默认：3）"),
|
||
limit: Optional[int] = Query(None, description="限制处理数量（用于测试，默认：不限制）"),
|
||
offset: int = Query(0, description="起始偏移量（默认：0）"),
|
||
skip_create_collection: bool = Query(False, description="跳过创建集合（如果集合已存在）"),
|
||
):
|
||
"""
|
||
从 t_sys_file 导入系统图向量到 Milvus
|
||
|
||
该接口会异步执行导入任务，任务在后台运行。
|
||
"""
|
||
try:
|
||
# 检查是否有任务正在运行
|
||
if task_status["running"]:
|
||
raise HTTPException(
|
||
status_code=409,
|
||
detail="已有导入任务正在运行，请等待完成后再试"
|
||
)
|
||
|
||
# 在后台线程中执行任务
|
||
executor.submit(
|
||
run_import_task,
|
||
batch_size,
|
||
retry_times,
|
||
limit,
|
||
offset,
|
||
skip_create_collection
|
||
)
|
||
|
||
return ResponseModel(
|
||
code=200,
|
||
msg="导入任务已启动，正在后台执行",
|
||
data={
|
||
"status": "started",
|
||
"batch_size": batch_size,
|
||
"retry_times": retry_times,
|
||
"limit": limit,
|
||
"offset": offset,
|
||
"skip_create_collection": skip_create_collection
|
||
}
|
||
)
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"启动导入任务失败: {e}", exc_info=True)
|
||
raise HTTPException(status_code=500, detail=f"启动导入任务失败: {str(e)}")
|
||
|
||
|
||
@router.get("/import-sys-sketch/status", response_model=ResponseModel)
|
||
async def get_import_status():
|
||
"""
|
||
获取导入任务状态
|
||
"""
|
||
return ResponseModel(
|
||
code=200,
|
||
msg="OK",
|
||
data={
|
||
"running": task_status["running"]
|
||
}
|
||
)
|
||
|