Преглед изворни кода

io推理上传进行性能限制

AnlaAnla пре 1 месец
родитељ
комит
d51fe87d9c
1 измењених фајлова са 19 додато и 4 уклоњено
  1. 19 4
      app/api/auto_import.py

+ 19 - 4
app/api/auto_import.py

@@ -12,6 +12,21 @@ from app.utils.scheme import CardType
 logger = get_logger(__name__)
 router = APIRouter()
 
+INFERENCE_LIMIT = 2
+UPLOAD_LIMIT = 4
+inference_semaphore = asyncio.Semaphore(INFERENCE_LIMIT)
+upload_semaphore = asyncio.Semaphore(UPLOAD_LIMIT)
+
+# 推理调用加 semaphore
+async def safe_process_main_image(session, file_bytes, filename, score_type, is_reflect_str):
+    async with inference_semaphore:
+        return await process_main_image(session, file_bytes, filename, score_type, is_reflect_str)
+
+# 上传调用加 semaphore
+async def safe_upload_task(task_func, *args, **kwargs):
+    async with upload_semaphore:
+        return await task_func(*args, **kwargs)
+
 
 # --- 内部辅助函数 ---
 async def call_api_with_bytes(
@@ -201,7 +216,7 @@ async def auto_import_script_api(
                 if len(f_bytes) == 0:
                     raise HTTPException(status_code=400, detail=f"图片文件 {f_name} 内容为空")
 
-                res = await process_main_image(session, f_bytes, f_name, img_type, is_reflect_str)
+                res = await safe_process_main_image(session, f_bytes, f_name, img_type, is_reflect_str)
                 processed_results.append(res)
 
             # Step 2: 在自身数据库创建卡片记录
@@ -213,7 +228,7 @@ async def auto_import_script_api(
             # Step 3: 并发调用自身的图片保存接口
             upload_tasks = []
             for res in processed_results:
-                upload_tasks.append(upload_main_image(session, local_base_url, card_id, res))
+                upload_tasks.append(safe_upload_task(session, local_base_url, card_id, res))
 
             for img_type, (f_bytes, f_name) in gray_bytes_data.items():
                 upload_tasks.append(upload_gray_image(session, local_base_url, card_id, img_type, f_bytes, f_name))
@@ -315,7 +330,7 @@ async def auto_import_url_script_api(
             for img_type, (f_bytes, f_name) in main_bytes_data.items():
                 if len(f_bytes) == 0:
                     raise HTTPException(status_code=400, detail=f"图片文件 {f_name} 内容为空")
-                res = await process_main_image(session, f_bytes, f_name, img_type, is_reflect_str)
+                res = await safe_process_main_image(session, f_bytes, f_name, img_type, is_reflect_str)
                 processed_results.append(res)
 
             # 3. 在自身数据库创建卡片记录
@@ -325,7 +340,7 @@ async def auto_import_url_script_api(
             # 4. 并发调用自身的图片保存接口
             upload_tasks = []
             for res in processed_results:
-                upload_tasks.append(upload_main_image(session, local_base_url, card_id, res))
+                upload_tasks.append(safe_upload_task(session, local_base_url, card_id, res))
             for img_type, (f_bytes, f_name) in gray_bytes_data.items():
                 upload_tasks.append(upload_gray_image(session, local_base_url, card_id, img_type, f_bytes, f_name))