1、增加调用 AI 视频生成能力和虚拟素材库管理的对外api
2、增加后台apikkey管理 3、增加apikey单独的模型定价 4、增加apikey调用情况 5、完善所有数据的注释增加
This commit is contained in:
@@ -0,0 +1,10 @@
|
||||
"""API 专用 Celery 任务注册模块。
|
||||
|
||||
复用共享的 celery_app 实例(同一个 broker、同一个 Redis),
|
||||
使 API 任务注册到同一个 Celery 应用上。
|
||||
|
||||
Worker 启动时需要 --include=app.tasks.api_generation_tasks 来加载 API 任务。
|
||||
"""
|
||||
|
||||
from app.tasks.celery_app import celery_app # noqa: F401
|
||||
from app.tasks.celery_app import run_async # noqa: F401
|
||||
@@ -0,0 +1,841 @@
|
||||
"""API 对外开放接口的 Celery 任务。
|
||||
|
||||
处理视频和图片的异步生成流程:
|
||||
- api_create_generation_task: 创建供应商任务(调用 Volcano Ark SDK)
|
||||
- api_poll_generation_task: 轮询视频任务状态
|
||||
- api_download_generation_result_task: 下载生成结果
|
||||
- api_upscale_finalize_task: 超分完成后更新 API 任务
|
||||
|
||||
Worker 启动命令示例:
|
||||
celery -A app.tasks.celery_app worker \\
|
||||
--include=app.tasks.api_generation_tasks \\
|
||||
--queue=gen_api_create,gen_api_poll,gen_api_download \\
|
||||
--concurrency=4
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any
|
||||
|
||||
from app.config import settings
|
||||
from app.enums.celery_queue import CeleryQueue
|
||||
from app.models.api.api_generation_task import ApiGenerationTask
|
||||
from app.models.api.api_key import ApiKey
|
||||
from app.models.base import async_session
|
||||
from sqlalchemy import select
|
||||
from app.models.image_engine import ImageEngine
|
||||
from app.models.video_engine import VideoEngine
|
||||
from app.services.api_v3 import upscale_service, usage_log_service
|
||||
from app.services.api_v3.logging_service import log_model_response, log_upscale_poll, log_upscale_poll_start, log_upscale_poll_end, log_error
|
||||
from app.services.api_v3.quota_service import get_queued_video_tasks, can_start_video_task
|
||||
from app.services.redis_registry_service import redis_acquire_lock
|
||||
from app.services.generation.poll_schedule_service import (
|
||||
build_video_pending_poll_schedule,
|
||||
ensure_video_poll_fields,
|
||||
is_poll_not_due,
|
||||
)
|
||||
from app.services.video_gen import poll_task_status, submit_video_task
|
||||
from app.tasks.api_celery_app import celery_app, run_async
|
||||
|
||||
logger = logging.getLogger("videogen")
|
||||
|
||||
|
||||
# === 队列名称常量 ===
|
||||
QUEUE_CREATE = "gen_api_create"
|
||||
QUEUE_POLL = "gen_api_poll"
|
||||
QUEUE_DOWNLOAD = "gen_api_download"
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
async def _get_quota_info(db, api_key_id: str) -> tuple[float | None, float | None]:
|
||||
"""获取当前配额信息。
|
||||
|
||||
Returns:
|
||||
(quota_before, quota_after) - 当前余额作为 before,after 需要计算
|
||||
"""
|
||||
from app.models.api.api_key import ApiKey
|
||||
key = await db.get(ApiKey, api_key_id)
|
||||
if key:
|
||||
return key.quota_used, key.quota_used
|
||||
return None, None
|
||||
|
||||
|
||||
async def _refund_quota(db, task: ApiGenerationTask):
|
||||
"""退回预扣配额。"""
|
||||
from app.models.api.api_key import ApiKey
|
||||
|
||||
pre_deducted = task.credits_cost or 0.0
|
||||
if pre_deducted <= 0:
|
||||
return
|
||||
|
||||
key_result = await db.execute(
|
||||
select(ApiKey).where(ApiKey.id == task.api_key_id).with_for_update()
|
||||
)
|
||||
key = key_result.scalar_one_or_none()
|
||||
if key:
|
||||
key.quota_used = round(max(0, (key.quota_used or 0.0) - pre_deducted), 2)
|
||||
task.credits_cost = 0 # 标记已退回
|
||||
|
||||
|
||||
async def _start_next_queued_task(db, api_key_id: str):
|
||||
"""检查并启动下一个排队的视频任务。
|
||||
|
||||
当一个任务完成/失败时调用,检查是否有排队的任务可以启动。
|
||||
"""
|
||||
from app.models.api.api_key import ApiKey
|
||||
|
||||
# 加载 API Key
|
||||
key = await db.get(ApiKey, api_key_id)
|
||||
if not key:
|
||||
return
|
||||
|
||||
# 检查是否可以启动新任务
|
||||
if not await can_start_video_task(key, db):
|
||||
return
|
||||
|
||||
# 获取最早的排队任务
|
||||
queued_tasks = await get_queued_video_tasks(key, db, limit=1)
|
||||
if not queued_tasks:
|
||||
return
|
||||
|
||||
next_task = queued_tasks[0]
|
||||
|
||||
# 更新状态并启动
|
||||
next_task.status = "pending"
|
||||
next_task.pipeline_stage = "queued"
|
||||
await db.commit()
|
||||
|
||||
# 入队 Celery 创建任务
|
||||
api_create_generation_task.apply_async(
|
||||
args=[next_task.id],
|
||||
queue=QUEUE_CREATE,
|
||||
)
|
||||
|
||||
logger.info("Started queued API task: %s (key=%s)", next_task.id, api_key_id)
|
||||
|
||||
|
||||
def _build_optimized_prompt(task: ApiGenerationTask) -> str:
|
||||
"""构建优化后的提示词(追加参数信息)。"""
|
||||
base = (task.original_prompt or "").strip().rstrip(",,。;; \n\t")
|
||||
if not base:
|
||||
return task.original_prompt or ""
|
||||
|
||||
parts = []
|
||||
if task.gen_type == "video":
|
||||
parts = [
|
||||
f"时长:{task.duration or 4}秒",
|
||||
f"画面比例:{task.aspect_ratio or '16:9'}",
|
||||
f"分辨率:{task.provider_generation_resolution or task.resolution or '480p'}",
|
||||
]
|
||||
suffix = ",".join(parts)
|
||||
return f"{base},{suffix}" if base and suffix else base
|
||||
|
||||
|
||||
# === 任务 1: 创建供应商任务 ===
|
||||
|
||||
@celery_app.task(
|
||||
name="api.create_generation_task",
|
||||
bind=True,
|
||||
max_retries=3,
|
||||
soft_time_limit=300,
|
||||
time_limit=600,
|
||||
acks_late=True,
|
||||
)
|
||||
def api_create_generation_task(self, task_id: str) -> dict[str, Any]:
|
||||
"""创建视频生成任务并提交到 Volcano Ark SDK。"""
|
||||
return run_async(_create_generation_task(self, task_id))
|
||||
|
||||
|
||||
async def _create_generation_task(self, task_id: str) -> dict[str, Any]:
|
||||
lock_key = f"vg:lock:api_generation:create:{task_id}:attempt:{1}"
|
||||
|
||||
token = await redis_acquire_lock(lock_key=lock_key, ttl_seconds=120)
|
||||
if not token:
|
||||
logger.warning("API create task lock not acquired: %s", task_id)
|
||||
return {"status": "lock_not_acquired", "task_id": task_id}
|
||||
|
||||
try:
|
||||
async with async_session() as db:
|
||||
# 加载任务(带行锁)
|
||||
task = await db.get(ApiGenerationTask, task_id, with_for_update=True)
|
||||
if not task or task.deleted_at:
|
||||
return {"status": "not_found", "task_id": task_id}
|
||||
|
||||
if task.status not in ("pending", "generating"):
|
||||
return {"status": "skipped", "task_id": task_id, "current_status": task.status}
|
||||
|
||||
# 解析引擎配置
|
||||
try:
|
||||
engine_snapshot = json.loads(task.engine_snapshot_json) if task.engine_snapshot_json else {}
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
engine_snapshot = {}
|
||||
|
||||
engine_id = task.engine_id or engine_snapshot.get("id")
|
||||
if not engine_id:
|
||||
# 失败:退回预扣配额
|
||||
await _refund_quota(db, task)
|
||||
task.status = "failed"
|
||||
task.error_message = "无法解析引擎配置"
|
||||
await db.commit()
|
||||
return {"status": "failed", "task_id": task_id, "error": "no_engine"}
|
||||
|
||||
# 加载引擎
|
||||
engine = await db.get(VideoEngine, engine_id) or await db.get(ImageEngine, engine_id)
|
||||
if not engine:
|
||||
# 失败:退回预扣配额
|
||||
await _refund_quota(db, task)
|
||||
task.status = "failed"
|
||||
task.error_message = f"引擎 {engine_id} 不存在"
|
||||
await db.commit()
|
||||
return {"status": "failed", "task_id": task_id, "error": "engine_not_found"}
|
||||
|
||||
# 设置优化提示词
|
||||
task.optimized_prompt = _build_optimized_prompt(task)
|
||||
task.pipeline_stage = "creating_provider_task"
|
||||
await db.flush()
|
||||
|
||||
try:
|
||||
# 调用 Volcano Ark SDK
|
||||
provider_task_id = await submit_video_task(
|
||||
db=db,
|
||||
engine=engine,
|
||||
record=task,
|
||||
include_media_references=True,
|
||||
)
|
||||
|
||||
# 更新任务状态
|
||||
task.provider_task_id = provider_task_id
|
||||
task.provider_response_json = json.dumps({"task_id": provider_task_id}, ensure_ascii=False)
|
||||
task.pipeline_stage = "waiting_remote"
|
||||
task.status = "generating"
|
||||
task.resource_generation_started_at = _now()
|
||||
|
||||
# 设置轮询字段
|
||||
ensure_video_poll_fields(task)
|
||||
task.poll_started_at = _now()
|
||||
task.poll_interval_seconds = 30
|
||||
task.next_poll_at = _now() + timedelta(seconds=30)
|
||||
|
||||
await db.commit()
|
||||
|
||||
logger.info("API video submitted: task_id=%s provider_task_id=%s", task_id, provider_task_id)
|
||||
|
||||
# 记录模型调用成功
|
||||
log_model_response(
|
||||
engine_id=engine_id,
|
||||
model_name=task.model_name,
|
||||
task_id=task_id,
|
||||
success=True,
|
||||
result={"provider_task_id": provider_task_id},
|
||||
)
|
||||
|
||||
# 入队轮询任务
|
||||
api_poll_generation_task.apply_async(
|
||||
args=[task_id],
|
||||
countdown=30,
|
||||
queue=QUEUE_POLL,
|
||||
)
|
||||
|
||||
return {"status": "submitted", "task_id": task_id, "provider_task_id": provider_task_id}
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception("API video submit failed: task_id=%s", task_id)
|
||||
|
||||
# 记录模型调用失败
|
||||
log_model_response(
|
||||
engine_id=engine_id or "unknown",
|
||||
model_name=task.model_name,
|
||||
task_id=task_id,
|
||||
success=False,
|
||||
error=str(exc)[:500],
|
||||
)
|
||||
|
||||
# 失败:退回预扣配额
|
||||
pre_deducted = task.credits_cost or 0.0
|
||||
await _refund_quota(db, task)
|
||||
task.status = "failed"
|
||||
task.error_message = f"提交失败: {str(exc)[:500]}"
|
||||
task.pipeline_stage = "failed"
|
||||
await db.commit()
|
||||
|
||||
# 记录失败日志
|
||||
try:
|
||||
quota_before, _ = await _get_quota_info(db, task.api_key_id)
|
||||
await usage_log_service.record_usage(
|
||||
db=db,
|
||||
api_key_id=task.api_key_id,
|
||||
request_type="video_create",
|
||||
model_name=task.model_name,
|
||||
gen_type="video",
|
||||
status="failed",
|
||||
task_id=task.id,
|
||||
credits_cost=pre_deducted,
|
||||
refund_amount=pre_deducted,
|
||||
price_action="refund",
|
||||
error_message=str(exc)[:500],
|
||||
error_code="submit_failed",
|
||||
quota_before=quota_before,
|
||||
quota_after=quota_before + pre_deducted if quota_before else None,
|
||||
)
|
||||
await db.commit()
|
||||
except Exception as log_exc:
|
||||
logger.error("Failed to record usage log: %s", log_exc)
|
||||
|
||||
# 失败释放并发槽位,检查排队任务
|
||||
await _start_next_queued_task(db, task.api_key_id)
|
||||
|
||||
return {"status": "failed", "task_id": task_id, "error": str(exc)}
|
||||
|
||||
finally:
|
||||
# 释放锁
|
||||
from app.services.redis_registry_service import redis_release_lock
|
||||
await redis_release_lock(lock_key=lock_key, token=token)
|
||||
|
||||
|
||||
# === 任务 2: 轮询任务状态 ===
|
||||
|
||||
@celery_app.task(
|
||||
name="api.poll_generation_task",
|
||||
bind=True,
|
||||
max_retries=3,
|
||||
soft_time_limit=120,
|
||||
time_limit=300,
|
||||
acks_late=True,
|
||||
)
|
||||
def api_poll_generation_task(self, task_id: str) -> dict[str, Any]:
|
||||
"""轮询视频任务状态。"""
|
||||
return run_async(_poll_generation_task(self, task_id))
|
||||
|
||||
|
||||
async def _poll_generation_task(self, task_id: str) -> dict[str, Any]:
|
||||
lock_key = f"vg:lock:api_generation:poll:{task_id}"
|
||||
|
||||
token = await redis_acquire_lock(lock_key=lock_key, ttl_seconds=60)
|
||||
if not token:
|
||||
return {"status": "lock_not_acquired", "task_id": task_id}
|
||||
|
||||
try:
|
||||
async with async_session() as db:
|
||||
task = await db.get(ApiGenerationTask, task_id, with_for_update=True)
|
||||
if not task or task.deleted_at:
|
||||
return {"status": "not_found", "task_id": task_id}
|
||||
|
||||
if task.status != "generating" or not task.provider_task_id:
|
||||
return {"status": "skipped", "task_id": task_id}
|
||||
|
||||
# 检查是否到轮询时间
|
||||
if is_poll_not_due(task):
|
||||
# 重新调度
|
||||
schedule = build_video_pending_poll_schedule(task)
|
||||
task.next_poll_at = schedule.next_poll_at
|
||||
task.poll_interval_seconds = schedule.poll_interval_seconds
|
||||
await db.commit()
|
||||
|
||||
api_poll_generation_task.apply_async(
|
||||
args=[task_id],
|
||||
countdown=schedule.delay_seconds,
|
||||
queue=QUEUE_POLL,
|
||||
)
|
||||
return {"status": "rescheduled", "task_id": task_id, "delay": schedule.delay_seconds}
|
||||
|
||||
# 检查截止时间
|
||||
if task.deadline_at and task.deadline_at <= _now():
|
||||
# 超时:退回预扣配额
|
||||
await _refund_quota(db, task)
|
||||
task.status = "failed"
|
||||
task.error_message = "任务超时(24小时)"
|
||||
task.pipeline_stage = "timeout"
|
||||
await db.commit()
|
||||
return {"status": "timeout", "task_id": task_id}
|
||||
|
||||
# 解析引擎
|
||||
try:
|
||||
engine_snapshot = json.loads(task.engine_snapshot_json) if task.engine_snapshot_json else {}
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
engine_snapshot = {}
|
||||
|
||||
engine_id = task.engine_id or engine_snapshot.get("id")
|
||||
engine = await db.get(VideoEngine, engine_id) if engine_id else None
|
||||
if not engine:
|
||||
task.status = "failed"
|
||||
task.error_message = f"引擎 {engine_id} 不存在"
|
||||
await db.commit()
|
||||
return {"status": "failed", "task_id": task_id, "error": "engine_not_found"}
|
||||
|
||||
# 轮询状态
|
||||
task.last_poll_at = _now()
|
||||
task.poll_count = (task.poll_count or 0) + 1
|
||||
await db.flush()
|
||||
|
||||
try:
|
||||
poll_result = await poll_task_status(engine, task.provider_task_id)
|
||||
except Exception as exc:
|
||||
logger.warning("API poll failed: task_id=%s error=%s", task_id, exc)
|
||||
task.poll_error_count = (task.poll_error_count or 0) + 1
|
||||
await db.commit()
|
||||
|
||||
# 重新调度
|
||||
schedule = build_video_pending_poll_schedule(task)
|
||||
task.next_poll_at = schedule.next_poll_at
|
||||
task.poll_interval_seconds = schedule.poll_interval_seconds
|
||||
await db.commit()
|
||||
|
||||
api_poll_generation_task.apply_async(
|
||||
args=[task_id],
|
||||
countdown=schedule.delay_seconds,
|
||||
queue=QUEUE_POLL,
|
||||
)
|
||||
return {"status": "poll_error", "task_id": task_id}
|
||||
|
||||
status = poll_result.get("status")
|
||||
|
||||
if status == "succeeded":
|
||||
# 成功:入队下载
|
||||
task.remote_result_url = poll_result.get("video_url")
|
||||
task.provider_response_json = poll_result.get("response_data", "")
|
||||
task.pipeline_stage = "result_ready"
|
||||
task.video_tokens_used = poll_result.get("video_tokens", 0)
|
||||
await db.commit()
|
||||
|
||||
api_download_generation_result_task.apply_async(
|
||||
args=[task_id],
|
||||
queue=QUEUE_DOWNLOAD,
|
||||
)
|
||||
return {"status": "succeeded", "task_id": task_id}
|
||||
|
||||
elif status == "failed":
|
||||
# 失败:退回预扣配额
|
||||
await _refund_quota(db, task)
|
||||
task.status = "failed"
|
||||
task.error_message = poll_result.get("error", "视频生成失败")
|
||||
task.pipeline_stage = "failed"
|
||||
task.provider_response_json = poll_result.get("response_data", "")
|
||||
await db.commit()
|
||||
|
||||
# 获取预扣金额(退回前)
|
||||
pre_deducted = task.credits_cost or 0.0
|
||||
await usage_log_service.record_usage(
|
||||
db=db,
|
||||
api_key_id=task.api_key_id,
|
||||
request_type="video_create",
|
||||
model_name=task.model_name,
|
||||
gen_type="video",
|
||||
status="failed",
|
||||
task_id=task.id,
|
||||
credits_cost=pre_deducted,
|
||||
refund_amount=pre_deducted,
|
||||
price_action="refund",
|
||||
error_message=task.error_message,
|
||||
error_code="generation_failed",
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
# 失败释放并发槽位,检查排队任务
|
||||
await _start_next_queued_task(db, task.api_key_id)
|
||||
|
||||
return {"status": "failed", "task_id": task_id}
|
||||
|
||||
else:
|
||||
# 仍在处理中:重新调度
|
||||
schedule = build_video_pending_poll_schedule(task)
|
||||
task.next_poll_at = schedule.next_poll_at
|
||||
task.poll_interval_seconds = schedule.poll_interval_seconds
|
||||
await db.commit()
|
||||
|
||||
api_poll_generation_task.apply_async(
|
||||
args=[task_id],
|
||||
countdown=schedule.delay_seconds,
|
||||
queue=QUEUE_POLL,
|
||||
)
|
||||
return {"status": "pending", "task_id": task_id, "delay": schedule.delay_seconds}
|
||||
|
||||
finally:
|
||||
from app.services.redis_registry_service import redis_release_lock
|
||||
await redis_release_lock(lock_key=lock_key, token=token)
|
||||
|
||||
|
||||
# === 任务 3: 下载生成结果 ===
|
||||
|
||||
@celery_app.task(
|
||||
name="api.download_generation_result_task",
|
||||
bind=True,
|
||||
max_retries=3,
|
||||
soft_time_limit=600,
|
||||
time_limit=900,
|
||||
acks_late=True,
|
||||
)
|
||||
def api_download_generation_result_task(self, task_id: str) -> dict[str, Any]:
|
||||
"""下载视频结果并触发超分(如启用)。"""
|
||||
return run_async(_download_generation_result(self, task_id))
|
||||
|
||||
|
||||
async def _download_generation_result(self, task_id: str) -> dict[str, Any]:
|
||||
lock_key = f"vg:lock:api_generation:download:{task_id}"
|
||||
|
||||
token = await redis_acquire_lock(lock_key=lock_key, ttl_seconds=300)
|
||||
if not token:
|
||||
return {"status": "lock_not_acquired", "task_id": task_id}
|
||||
|
||||
try:
|
||||
|
||||
async with async_session() as db:
|
||||
task = await db.get(ApiGenerationTask, task_id, with_for_update=True)
|
||||
if not task or task.deleted_at:
|
||||
return {"status": "not_found", "task_id": task_id}
|
||||
|
||||
if not task.remote_result_url:
|
||||
task.status = "failed"
|
||||
task.error_message = "无远程结果URL"
|
||||
await db.commit()
|
||||
return {"status": "failed", "task_id": task_id}
|
||||
|
||||
# 检查是否需要超分
|
||||
use_upscale = bool(
|
||||
task.gen_type == "video"
|
||||
and task.video_upscale_enabled_snapshot
|
||||
and task.video_upscale_snapshot_json
|
||||
)
|
||||
|
||||
if use_upscale:
|
||||
# 下载视频到 upscaled 目录(作为超分源)
|
||||
import os
|
||||
from app.services.video_gen import download_video
|
||||
|
||||
date_dir = datetime.now().strftime("%Y%m%d")
|
||||
# 源文件存储路径
|
||||
dest_path = f"./storage/generate/api/upscaled/{date_dir}/{task.id}_source.mp4"
|
||||
abs_dest_path = os.path.abspath(dest_path)
|
||||
os.makedirs(os.path.dirname(abs_dest_path), exist_ok=True)
|
||||
|
||||
logger.info("Downloading source video to: %s", abs_dest_path)
|
||||
try:
|
||||
await download_video(task.remote_result_url, abs_dest_path)
|
||||
# 注意:超分时不设置 video_url,等超分完成后再设置
|
||||
task.local_path = dest_path # 源文件路径
|
||||
task.download_storage_date_dir = date_dir
|
||||
logger.info("Source video downloaded successfully: %s", dest_path)
|
||||
|
||||
# 获取视频信息(尺寸、时长、文件大小)
|
||||
try:
|
||||
abs_path = os.path.abspath(dest_path)
|
||||
if os.path.exists(abs_path):
|
||||
task.source_file_size_bytes = os.path.getsize(abs_path)
|
||||
# 使用 probe_video 探测实际视频尺寸和时长
|
||||
from app.services.video_upscale.media_service import probe_video
|
||||
source_info = await probe_video(abs_path)
|
||||
if source_info:
|
||||
task.source_width = source_info.width
|
||||
task.source_height = source_info.height
|
||||
task.source_duration_seconds = round(source_info.duration_seconds, 2)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to probe video info: %s", exc)
|
||||
|
||||
await db.flush()
|
||||
|
||||
# 创建超分任务(与状态更新在同一事务中)
|
||||
try:
|
||||
upscale_task = await upscale_service.prepare_api_upscale_task(
|
||||
db=db,
|
||||
api_task=task,
|
||||
source_local_path=dest_path,
|
||||
)
|
||||
|
||||
# 如果已存在超分任务(重复调用),检查超分状态
|
||||
if upscale_task is None:
|
||||
logger.info("Upscale task already exists for %s, checking status", task_id)
|
||||
# 重新查询超分任务状态
|
||||
from app.models.video_upscale_task import VideoUpscaleTask
|
||||
from sqlalchemy import select
|
||||
upscale_result = await db.execute(
|
||||
select(VideoUpscaleTask).where(
|
||||
VideoUpscaleTask.api_generation_task_id == task.id
|
||||
).limit(1)
|
||||
)
|
||||
existing_upscale = upscale_result.scalar_one_or_none()
|
||||
|
||||
if existing_upscale and existing_upscale.status == "completed":
|
||||
# 超分已完成
|
||||
task.video_url = existing_upscale.final_resource_url or task.remote_result_url
|
||||
task.status = "completed"
|
||||
task.pipeline_stage = "done"
|
||||
task.generated_at = _now()
|
||||
else:
|
||||
# 超分仍在进行中,保持 generating 状态
|
||||
task.status = "generating"
|
||||
task.pipeline_stage = "upscale_processing"
|
||||
await db.commit()
|
||||
return {"status": task.status, "task_id": task_id, "note": "upscale_already_exists"}
|
||||
|
||||
# 记录超分开始
|
||||
log_upscale_poll_start(task_id=upscale_task.id, api_task_id=task_id)
|
||||
|
||||
# 入队超分任务(使用简化版 API v3 专用任务)
|
||||
from app.tasks.api_upscale_tasks import api_upscale_execute_local_simple, api_upscale_submit_remote_simple
|
||||
|
||||
processor_key = upscale_task.processor_key
|
||||
|
||||
if processor_key in ("local_ffmpeg_crop_v1",):
|
||||
api_upscale_execute_local_simple.apply_async(
|
||||
args=[upscale_task.id],
|
||||
queue="gen_api_upscale",
|
||||
)
|
||||
else:
|
||||
# 远程超分(火山 MediaKit)
|
||||
api_upscale_submit_remote_simple.apply_async(
|
||||
args=[upscale_task.id],
|
||||
queue="gen_api_upscale",
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
return {"status": "upscale_queued", "task_id": task_id, "upscale_task_id": upscale_task.id}
|
||||
|
||||
except Exception as upscale_exc:
|
||||
# 超分创建失败:记录错误,但视频已下载成功
|
||||
# 将任务标记为 completed(有视频但无超分)
|
||||
log_error(
|
||||
"UPSCALE_CREATE_ERROR",
|
||||
f"超分任务创建失败: {str(upscale_exc)[:500]}",
|
||||
{"task_id": task_id, "video_url": task.remote_result_url}
|
||||
)
|
||||
task.video_url = task.remote_result_url
|
||||
task.status = "completed"
|
||||
task.pipeline_stage = "done"
|
||||
task.generated_at = _now()
|
||||
task.error_message = f"超分创建失败,返回原始视频: {str(upscale_exc)[:200]}"
|
||||
await db.commit()
|
||||
|
||||
# 检查并启动下一个排队任务
|
||||
await _start_next_queued_task(db, task.api_key_id)
|
||||
|
||||
return {"status": "completed_without_upscale", "task_id": task_id, "error": str(upscale_exc)[:500]}
|
||||
|
||||
except Exception as exc:
|
||||
# 下载失败:退回预扣配额
|
||||
from app.models.api.api_key import ApiKey
|
||||
pre_deducted = task.credits_cost or 0.0
|
||||
if pre_deducted > 0:
|
||||
key_result = await db.execute(
|
||||
select(ApiKey).where(ApiKey.id == task.api_key_id).with_for_update()
|
||||
)
|
||||
key = key_result.scalar_one_or_none()
|
||||
if key:
|
||||
key.quota_used = round(max(0, (key.quota_used or 0.0) - pre_deducted), 2)
|
||||
|
||||
log_error(
|
||||
"DOWNLOAD_ERROR",
|
||||
f"视频下载失败: {str(exc)[:500]}",
|
||||
{"task_id": task_id}
|
||||
)
|
||||
logger.exception("API video download failed: task_id=%s", task_id)
|
||||
task.status = "failed"
|
||||
task.error_message = f"下载失败: {str(exc)[:500]}"
|
||||
task.credits_cost = 0
|
||||
await db.commit()
|
||||
|
||||
# 记录失败日志
|
||||
await usage_log_service.record_usage(
|
||||
db=db,
|
||||
api_key_id=task.api_key_id,
|
||||
request_type="video_create",
|
||||
model_name=task.model_name,
|
||||
gen_type="video",
|
||||
status="failed",
|
||||
task_id=task.id,
|
||||
credits_cost=pre_deducted,
|
||||
refund_amount=pre_deducted,
|
||||
price_action="refund",
|
||||
resolution=task.resolution,
|
||||
duration=task.duration,
|
||||
error_message=str(exc)[:500],
|
||||
error_code="download_failed",
|
||||
)
|
||||
|
||||
# 失败释放并发槽位,检查排队任务
|
||||
await _start_next_queued_task(db, task.api_key_id)
|
||||
|
||||
return {"status": "failed", "task_id": task_id, "error": str(exc)}
|
||||
|
||||
else:
|
||||
# 直接下载最终结果
|
||||
import os
|
||||
from app.services.video_gen import download_video
|
||||
|
||||
date_dir = datetime.now().strftime("%Y%m%d")
|
||||
# 统一路径格式
|
||||
dest_path = f"./storage/generate/api/videos/{date_dir}/{task.id}.mp4"
|
||||
abs_dest_path = os.path.abspath(dest_path)
|
||||
os.makedirs(os.path.dirname(abs_dest_path), exist_ok=True)
|
||||
|
||||
logger.info("Downloading final video to: %s", abs_dest_path)
|
||||
try:
|
||||
await download_video(task.remote_result_url, abs_dest_path)
|
||||
logger.info("Final video downloaded successfully: %s", dest_path)
|
||||
|
||||
# 配额已在创建时预扣,此处不再重复扣减
|
||||
task.video_url = dest_path # 使用相对路径
|
||||
task.local_path = dest_path
|
||||
task.download_storage_date_dir = date_dir
|
||||
task.status = "completed"
|
||||
task.pipeline_stage = "done"
|
||||
task.generated_at = _now()
|
||||
await db.commit()
|
||||
|
||||
# 记录成功日志(配额已在创建时预扣)
|
||||
try:
|
||||
await usage_log_service.record_usage(
|
||||
db=db,
|
||||
api_key_id=task.api_key_id,
|
||||
request_type="video_create",
|
||||
model_name=task.model_name,
|
||||
gen_type="video",
|
||||
status="success",
|
||||
task_id=task.id,
|
||||
credits_cost=task.credits_cost,
|
||||
tokens_used=task.video_tokens_used,
|
||||
price_action="deduct",
|
||||
resolution=task.resolution,
|
||||
duration=task.duration,
|
||||
)
|
||||
await db.commit()
|
||||
except Exception as log_exc:
|
||||
logger.error("Failed to record usage log: %s", log_exc)
|
||||
# 日志记录失败不应影响任务完成
|
||||
|
||||
# 检查并启动下一个排队任务
|
||||
await _start_next_queued_task(db, task.api_key_id)
|
||||
|
||||
return {"status": "completed", "task_id": task_id}
|
||||
|
||||
except Exception as exc:
|
||||
# 下载失败:退回预扣配额
|
||||
from app.models.api.api_key import ApiKey
|
||||
pre_deducted = task.credits_cost or 0.0
|
||||
if pre_deducted > 0:
|
||||
key_result = await db.execute(
|
||||
select(ApiKey).where(ApiKey.id == task.api_key_id).with_for_update()
|
||||
)
|
||||
key = key_result.scalar_one_or_none()
|
||||
if key:
|
||||
key.quota_used = round(max(0, (key.quota_used or 0.0) - pre_deducted), 2)
|
||||
|
||||
logger.exception("API video download failed: task_id=%s", task_id)
|
||||
task.status = "failed"
|
||||
task.error_message = f"下载失败: {str(exc)[:500]}"
|
||||
task.credits_cost = 0
|
||||
await db.commit()
|
||||
|
||||
# 失败释放并发槽位,检查排队任务
|
||||
await _start_next_queued_task(db, task.api_key_id)
|
||||
|
||||
return {"status": "failed", "task_id": task_id, "error": str(exc)}
|
||||
|
||||
finally:
|
||||
from app.services.redis_registry_service import redis_release_lock
|
||||
await redis_release_lock(lock_key=lock_key, token=token)
|
||||
|
||||
|
||||
# === 任务 4: 超分完成后更新 API 任务 ===
|
||||
|
||||
@celery_app.task(
|
||||
name="api.upscale_finalize_task",
|
||||
bind=True,
|
||||
max_retries=3,
|
||||
soft_time_limit=120,
|
||||
time_limit=300,
|
||||
acks_late=True,
|
||||
)
|
||||
def api_upscale_finalize_task(self, api_task_id: str, upscale_task_id: str) -> dict[str, Any]:
|
||||
"""超分完成后更新 API 任务状态。"""
|
||||
return run_async(_upscale_finalize(self, api_task_id, upscale_task_id))
|
||||
|
||||
|
||||
async def _upscale_finalize(self, api_task_id: str, upscale_task_id: str) -> dict[str, Any]:
|
||||
async with async_session() as db:
|
||||
from app.models.video_upscale_task import VideoUpscaleTask
|
||||
|
||||
task = await db.get(ApiGenerationTask, api_task_id)
|
||||
upscale_task = await db.get(VideoUpscaleTask, upscale_task_id)
|
||||
|
||||
if not task or task.deleted_at:
|
||||
return {"status": "not_found", "api_task_id": api_task_id}
|
||||
|
||||
if upscale_task and upscale_task.status == "completed":
|
||||
# 超分成功:更新视频URL(配额已在创建时预扣)
|
||||
task.video_url = upscale_task.provider_output_url or task.remote_result_url
|
||||
task.status = "completed"
|
||||
task.pipeline_stage = "done"
|
||||
task.generated_at = _now()
|
||||
await db.commit()
|
||||
|
||||
# 记录超分完成日志
|
||||
log_upscale_poll_end(
|
||||
task_id=upscale_task_id,
|
||||
api_task_id=api_task_id,
|
||||
success=True,
|
||||
final_status="completed",
|
||||
total_attempts=upscale_task.attempt_count or 1,
|
||||
)
|
||||
|
||||
await usage_log_service.record_usage(
|
||||
db=db,
|
||||
api_key_id=task.api_key_id,
|
||||
request_type="video_create",
|
||||
model_name=task.model_name,
|
||||
gen_type="video",
|
||||
status="success",
|
||||
task_id=task.id,
|
||||
credits_cost=task.credits_cost,
|
||||
tokens_used=task.video_tokens_used,
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
# 检查并启动下一个排队任务
|
||||
await _start_next_queued_task(db, task.api_key_id)
|
||||
|
||||
return {"status": "completed", "api_task_id": api_task_id}
|
||||
|
||||
elif upscale_task and upscale_task.status == "failed":
|
||||
# 超分失败:回退到原始视频
|
||||
task.video_url = task.remote_result_url
|
||||
task.status = "completed"
|
||||
task.pipeline_stage = "done"
|
||||
task.generated_at = _now()
|
||||
task.error_message = "超分失败,返回原始视频"
|
||||
await db.commit()
|
||||
|
||||
# 记录超分失败日志
|
||||
log_upscale_poll_end(
|
||||
task_id=upscale_task_id,
|
||||
api_task_id=api_task_id,
|
||||
success=False,
|
||||
final_status="failed",
|
||||
total_attempts=upscale_task.attempt_count or 1,
|
||||
)
|
||||
|
||||
# 检查并启动下一个排队任务
|
||||
await _start_next_queued_task(db, task.api_key_id)
|
||||
|
||||
return {"status": "completed_with_fallback", "api_task_id": api_task_id}
|
||||
|
||||
else:
|
||||
# 超分仍在处理中:重新调度
|
||||
log_upscale_poll(
|
||||
task_id=upscale_task_id,
|
||||
api_task_id=api_task_id,
|
||||
status=upscale_task.status or "unknown",
|
||||
attempt=upscale_task.attempt_count or 0,
|
||||
)
|
||||
api_upscale_finalize_task.apply_async(
|
||||
args=[api_task_id, upscale_task_id],
|
||||
countdown=60,
|
||||
queue=QUEUE_DOWNLOAD,
|
||||
)
|
||||
return {"status": "waiting_upscale", "api_task_id": api_task_id}
|
||||
@@ -0,0 +1,146 @@
|
||||
"""API v3 容灾恢复任务。
|
||||
|
||||
处理服务重启后的任务恢复:
|
||||
- 扫描处于中间状态的 ApiGenerationTask
|
||||
- 重新入队未完成的 Celery 任务
|
||||
- 处理租约过期的任务
|
||||
|
||||
Worker 启动时会自动触发恢复扫描。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from app.config import settings
|
||||
from app.enums.celery_queue import CeleryQueue
|
||||
from app.models.api.api_generation_task import ApiGenerationTask
|
||||
from app.models.base import async_session
|
||||
from app.tasks.api_generation_tasks import (
|
||||
QUEUE_CREATE,
|
||||
QUEUE_DOWNLOAD,
|
||||
QUEUE_POLL,
|
||||
api_create_generation_task,
|
||||
api_download_generation_result_task,
|
||||
api_poll_generation_task,
|
||||
)
|
||||
from app.tasks.async_runner import run_async
|
||||
from app.tasks.celery_app import celery_app
|
||||
|
||||
logger = logging.getLogger("videogen")
|
||||
|
||||
|
||||
async def recover_api_generation_tasks_once():
|
||||
"""扫描并恢复未完成的 API v3 生成任务。
|
||||
|
||||
恢复场景:
|
||||
1. status=pending 且未入队 -> 重新入队创建任务
|
||||
2. status=generating 且 provider_task_id 为空 -> 重新入队创建任务
|
||||
3. status=generating 且 provider_task_id 存在 -> 重新入队轮询任务
|
||||
4. pipeline_stage=result_ready -> 重新入队下载任务
|
||||
5. 租约过期但任务未完成 -> 重新入队对应阶段任务
|
||||
"""
|
||||
now = datetime.now(timezone.utc)
|
||||
recovered = 0
|
||||
|
||||
async with async_session() as db:
|
||||
# 1. 恢复 pending/generating 任务(未开始或中断)
|
||||
result = await db.execute(
|
||||
__import__("sqlalchemy").select(ApiGenerationTask).where(
|
||||
ApiGenerationTask.status.in_(["pending", "generating"]),
|
||||
ApiGenerationTask.deleted_at.is_(None),
|
||||
ApiGenerationTask.created_at > now - timedelta(hours=48),
|
||||
)
|
||||
)
|
||||
tasks = list(result.scalars().all())
|
||||
|
||||
for task in tasks:
|
||||
try:
|
||||
if task.status == "pending" or not task.provider_task_id:
|
||||
# 重新入队创建任务
|
||||
api_create_generation_task.apply_async(
|
||||
args=[task.id],
|
||||
queue=QUEUE_CREATE,
|
||||
)
|
||||
logger.info("API recovery: re-enqueued create task %s", task.id)
|
||||
recovered += 1
|
||||
|
||||
elif task.status == "generating" and task.provider_task_id:
|
||||
# 检查是否需要轮询
|
||||
next_poll_at = task.next_poll_at
|
||||
if next_poll_at is None or next_poll_at <= now:
|
||||
# 重新入队轮询任务
|
||||
api_poll_generation_task.apply_async(
|
||||
args=[task.id],
|
||||
queue=QUEUE_POLL,
|
||||
)
|
||||
logger.info("API recovery: re-enqueued poll task %s (provider_task_id=%s)", task.id, task.provider_task_id)
|
||||
recovered += 1
|
||||
|
||||
# 检查下载阶段
|
||||
if task.pipeline_stage == "result_ready" and not task.video_url and not task.image_url:
|
||||
api_download_generation_result_task.apply_async(
|
||||
args=[task.id],
|
||||
queue=QUEUE_DOWNLOAD,
|
||||
)
|
||||
logger.info("API recovery: re-enqueued download task %s", task.id)
|
||||
recovered += 1
|
||||
|
||||
except Exception as exc:
|
||||
logger.warning("API recovery: failed to recover task %s: %s", task.id, exc)
|
||||
|
||||
# 2. 恢复排队任务(服务重启后,排队任务需要重新检查并发)
|
||||
from app.services.api_v3.quota_service import can_start_video_task, get_queued_video_tasks
|
||||
from app.models.api.api_key import ApiKey
|
||||
|
||||
# 获取所有有排队任务的 API Key
|
||||
queued_result = await db.execute(
|
||||
__import__("sqlalchemy").select(ApiGenerationTask.api_key_id).where(
|
||||
ApiGenerationTask.status == "queued",
|
||||
ApiGenerationTask.deleted_at.is_(None),
|
||||
ApiGenerationTask.created_at > now - timedelta(hours=48),
|
||||
).distinct()
|
||||
)
|
||||
api_key_ids = [row[0] for row in queued_result.all()]
|
||||
|
||||
for api_key_id in api_key_ids:
|
||||
try:
|
||||
key = await db.get(ApiKey, api_key_id)
|
||||
if not key:
|
||||
continue
|
||||
|
||||
# 检查是否可以启动排队任务
|
||||
if await can_start_video_task(key, db):
|
||||
queued_tasks = await get_queued_video_tasks(key, db, limit=1)
|
||||
if queued_tasks:
|
||||
next_task = queued_tasks[0]
|
||||
next_task.status = "pending"
|
||||
next_task.pipeline_stage = "queued"
|
||||
await db.commit()
|
||||
|
||||
api_create_generation_task.apply_async(
|
||||
args=[next_task.id],
|
||||
queue=QUEUE_CREATE,
|
||||
)
|
||||
logger.info("API recovery: started queued task %s for key %s", next_task.id, api_key_id)
|
||||
recovered += 1
|
||||
|
||||
except Exception as exc:
|
||||
logger.warning("API recovery: failed to recover queued task for key %s: %s", api_key_id, exc)
|
||||
|
||||
if recovered:
|
||||
logger.info("API recovery: recovered %d tasks", recovered)
|
||||
return recovered
|
||||
|
||||
|
||||
@celery_app.task(
|
||||
name="api_generation.recover_tasks_once",
|
||||
bind=True,
|
||||
max_retries=0,
|
||||
soft_time_limit=300,
|
||||
time_limit=600,
|
||||
)
|
||||
def api_generation_recover_tasks_once(self):
|
||||
"""API v3 任务恢复扫描(Celery Beat 定时触发)。"""
|
||||
return run_async(recover_api_generation_tasks_once())
|
||||
@@ -0,0 +1,364 @@
|
||||
"""API v3 简化超分任务。
|
||||
|
||||
不使用复杂的 CeleryRuntimeLease 锁机制,直接执行超分流程。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from app.config import settings
|
||||
from app.enums.celery_queue import CeleryQueue
|
||||
from app.services.video_upscale.volc_service import VolcSubmitResult, VolcQueryResult
|
||||
from app.tasks.api_celery_app import celery_app
|
||||
|
||||
logger = logging.getLogger("videogen")
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
"""获取当前时间(UTC)。"""
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def _now_str() -> str:
|
||||
"""获取当前日期字符串。"""
|
||||
return _now().strftime("%Y%m%d")
|
||||
|
||||
|
||||
def _format_time(dt: datetime | None) -> str | None:
|
||||
"""格式化时间为字符串(北京时间)。"""
|
||||
if dt is None:
|
||||
return None
|
||||
from datetime import timedelta
|
||||
beijing_time = dt + timedelta(hours=8)
|
||||
return beijing_time.strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
|
||||
async def _maybe_delete_source_file(db, upscale) -> None:
|
||||
"""根据超分配置决定是否删除源文件。"""
|
||||
if not upscale.source_local_path or not upscale.api_generation_task_id:
|
||||
return
|
||||
|
||||
from app.models.api.api_key_upscale_config import ApiKeyUpscaleConfig
|
||||
from app.models.api.api_generation_task import ApiGenerationTask
|
||||
from sqlalchemy import select
|
||||
|
||||
# 获取 API Key ID
|
||||
api_task_result = await db.execute(
|
||||
select(ApiGenerationTask.api_key_id).where(
|
||||
ApiGenerationTask.id == upscale.api_generation_task_id
|
||||
).limit(1)
|
||||
)
|
||||
api_key_id = api_task_result.scalar_one_or_none()
|
||||
if not api_key_id:
|
||||
return
|
||||
|
||||
# 查询超分配置
|
||||
config_result = await db.execute(
|
||||
select(ApiKeyUpscaleConfig).where(
|
||||
ApiKeyUpscaleConfig.api_key_id == api_key_id
|
||||
).limit(1)
|
||||
)
|
||||
upscale_config = config_result.scalar_one_or_none()
|
||||
|
||||
# 只有配置了"成功后删除源文件"才删除
|
||||
if upscale_config and upscale_config.delete_source_after_success:
|
||||
source_abs = os.path.abspath(upscale.source_local_path)
|
||||
if os.path.exists(source_abs):
|
||||
os.remove(source_abs)
|
||||
logger.info("Deleted source file after upscale: %s", source_abs)
|
||||
|
||||
|
||||
@celery_app.task(
|
||||
name="api_upscale.execute_local_simple",
|
||||
bind=True,
|
||||
max_retries=2,
|
||||
soft_time_limit=600,
|
||||
time_limit=900,
|
||||
)
|
||||
def api_upscale_execute_local_simple(self, upscale_task_id: str) -> dict[str, Any]:
|
||||
"""简化版本地超分执行(API v3 专用)。"""
|
||||
from app.models.base import async_session
|
||||
from app.models.video_upscale_task import VideoUpscaleTask
|
||||
from app.services.video_upscale.local_ffmpeg_service import run_local_ffmpeg_upscale
|
||||
from sqlalchemy import select
|
||||
|
||||
async def _execute():
|
||||
async with async_session() as db:
|
||||
result = await db.execute(
|
||||
select(VideoUpscaleTask).where(VideoUpscaleTask.id == upscale_task_id).limit(1)
|
||||
)
|
||||
upscale = result.scalar_one_or_none()
|
||||
if not upscale:
|
||||
return {"status": "not_found"}
|
||||
|
||||
if upscale.status == "completed":
|
||||
return {"status": "already_completed"}
|
||||
|
||||
# 更新状态为处理中
|
||||
upscale.status = "processing"
|
||||
upscale.stage = "local_processing"
|
||||
upscale.attempt_count = int(upscale.attempt_count or 0) + 1
|
||||
upscale.started_at = upscale.started_at or __import__("datetime").datetime.now(__import__("datetime").timezone.utc)
|
||||
await db.commit()
|
||||
|
||||
try:
|
||||
# 执行本地 FFmpeg 超分
|
||||
output_path = await run_local_ffmpeg_upscale(upscale)
|
||||
|
||||
# 计算相对 URL 路径
|
||||
# output_path 是绝对路径,需要转换为 /generate/api/videos/... 格式
|
||||
date_dir = _now_str()
|
||||
url_path = f"/generate/api/videos/{date_dir}/{upscale.api_generation_task_id}.mp4"
|
||||
|
||||
# 更新成功状态
|
||||
upscale.status = "completed"
|
||||
upscale.stage = "upscale_completed"
|
||||
upscale.final_local_path = f"./storage/generate/api/videos/{_now_str()}/{upscale.api_generation_task_id or "unknown"}.mp4"
|
||||
upscale.final_resource_url = url_path # 相对 URL
|
||||
upscale.completed_at = _now()
|
||||
|
||||
# 更新 API 任务的 video_url(使用相对 URL)
|
||||
if upscale.api_generation_task_id:
|
||||
from app.models.api.api_generation_task import ApiGenerationTask
|
||||
api_result = await db.execute(
|
||||
select(ApiGenerationTask).where(ApiGenerationTask.id == upscale.api_generation_task_id).limit(1)
|
||||
)
|
||||
api_task = api_result.scalar_one_or_none()
|
||||
if api_task:
|
||||
api_task.video_url = url_path # 相对 URL
|
||||
api_task.status = "completed"
|
||||
api_task.pipeline_stage = "done"
|
||||
api_task.generated_at = upscale.completed_at
|
||||
|
||||
await db.commit()
|
||||
return {"status": "completed", "output_path": output_path}
|
||||
|
||||
except Exception as exc:
|
||||
upscale.status = "failed"
|
||||
upscale.stage = "failed"
|
||||
upscale.failure_count = int(upscale.failure_count or 0) + 1
|
||||
upscale.last_error = str(exc)[:500]
|
||||
upscale.failed_at = __import__("datetime").datetime.now(__import__("datetime").timezone.utc)
|
||||
await db.commit()
|
||||
raise
|
||||
|
||||
from app.tasks.async_runner import run_async
|
||||
return run_async(_execute())
|
||||
|
||||
|
||||
@celery_app.task(
|
||||
name="api_upscale.submit_remote_simple",
|
||||
bind=True,
|
||||
max_retries=2,
|
||||
soft_time_limit=600,
|
||||
time_limit=900,
|
||||
)
|
||||
def api_upscale_submit_remote_simple(self, upscale_task_id: str) -> dict[str, Any]:
|
||||
"""简化版远程超分提交(API v3 专用)。"""
|
||||
from app.models.base import async_session
|
||||
from app.models.video_upscale_task import VideoUpscaleTask
|
||||
from app.services.video_upscale.volc_service import submit_video_enhance, VolcSubmitResult
|
||||
from sqlalchemy import select
|
||||
|
||||
async def _execute():
|
||||
async with async_session() as db:
|
||||
result = await db.execute(
|
||||
select(VideoUpscaleTask).where(VideoUpscaleTask.id == upscale_task_id).limit(1)
|
||||
)
|
||||
upscale = result.scalar_one_or_none()
|
||||
if not upscale:
|
||||
return {"status": "not_found"}
|
||||
|
||||
if upscale.status == "completed":
|
||||
return {"status": "already_completed"}
|
||||
|
||||
# 更新状态
|
||||
upscale.status = "processing"
|
||||
upscale.stage = "remote_submitting"
|
||||
upscale.attempt_count = int(upscale.attempt_count or 0) + 1
|
||||
upscale.started_at = upscale.started_at or __import__("datetime").datetime.now(__import__("datetime").timezone.utc)
|
||||
await db.commit()
|
||||
|
||||
try:
|
||||
# 提交到火山 MediaKit
|
||||
from app.utils.id_gen import generate_id
|
||||
# 从 API 任务获取超分快照
|
||||
import json
|
||||
api_snapshot = {}
|
||||
if upscale.api_generation_task_id:
|
||||
from app.models.api.api_generation_task import ApiGenerationTask
|
||||
api_result = await db.execute(
|
||||
__import__("sqlalchemy").select(ApiGenerationTask).where(ApiGenerationTask.id == upscale.api_generation_task_id).limit(1)
|
||||
)
|
||||
api_task = api_result.scalar_one_or_none()
|
||||
if api_task and api_task.video_upscale_snapshot_json:
|
||||
try:
|
||||
api_snapshot = json.loads(api_task.video_upscale_snapshot_json)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
|
||||
submit_result: VolcSubmitResult = await submit_video_enhance(
|
||||
processor_key=upscale.processor_key,
|
||||
video_url=upscale.source_remote_url or "",
|
||||
target_resolution=api_snapshot.get("target_resolution", "1080p"),
|
||||
target_width=int(upscale.target_width or 0),
|
||||
target_height=int(upscale.target_height or 0),
|
||||
processor=api_snapshot.get("processor", {}),
|
||||
client_token=generate_id(),
|
||||
)
|
||||
|
||||
# 更新成功状态
|
||||
upscale.provider_task_id = submit_result.task_id
|
||||
upscale.provider_submitted_at = __import__("datetime").datetime.now(__import__("datetime").timezone.utc)
|
||||
upscale.provider_request_json = json.dumps(submit_result.request_payload, ensure_ascii=False) if submit_result.request_payload else None
|
||||
upscale.provider_response_json = json.dumps(submit_result.response_payload, ensure_ascii=False) if submit_result.response_payload else None
|
||||
upscale.celery_task_id = self.request.id if hasattr(self, 'request') else None
|
||||
upscale.status = "processing"
|
||||
upscale.stage = "remote_polling"
|
||||
await db.commit()
|
||||
|
||||
# 立即触发第一次轮询
|
||||
api_upscale_poll_remote_simple.apply_async(
|
||||
args=[upscale_task_id],
|
||||
countdown=30,
|
||||
queue="gen_api_upscale",
|
||||
)
|
||||
|
||||
return {"status": "submitted", "provider_task_id": submit_result.task_id}
|
||||
|
||||
except Exception as exc:
|
||||
upscale.status = "failed"
|
||||
upscale.stage = "failed"
|
||||
upscale.failure_count = int(upscale.failure_count or 0) + 1
|
||||
upscale.last_error = str(exc)[:500]
|
||||
upscale.failed_at = __import__("datetime").datetime.now(__import__("datetime").timezone.utc)
|
||||
await db.commit()
|
||||
raise
|
||||
|
||||
from app.tasks.async_runner import run_async
|
||||
return run_async(_execute())
|
||||
|
||||
|
||||
@celery_app.task(
|
||||
name="api_upscale.poll_remote_simple",
|
||||
bind=True,
|
||||
max_retries=10,
|
||||
soft_time_limit=120,
|
||||
time_limit=300,
|
||||
)
|
||||
def api_upscale_poll_remote_simple(self, upscale_task_id: str) -> dict[str, Any]:
|
||||
"""简化版远程超分轮询(API v3 专用)。"""
|
||||
from app.models.base import async_session
|
||||
from app.models.video_upscale_task import VideoUpscaleTask
|
||||
from app.services.video_upscale.volc_service import query_task
|
||||
from sqlalchemy import select
|
||||
|
||||
async def _execute():
|
||||
async with async_session() as db:
|
||||
result = await db.execute(
|
||||
select(VideoUpscaleTask).where(VideoUpscaleTask.id == upscale_task_id).limit(1)
|
||||
)
|
||||
upscale = result.scalar_one_or_none()
|
||||
if not upscale or upscale.status == "completed":
|
||||
return {"status": "not_found_or_completed"}
|
||||
|
||||
try:
|
||||
# 查询火山 MediaKit 状态
|
||||
query_result = await query_task(upscale.provider_task_id)
|
||||
|
||||
status = query_result.status
|
||||
|
||||
if status == "completed":
|
||||
# 超分完成
|
||||
output_url = query_result.result.get("video_url", "") if query_result.result else ""
|
||||
upscale.provider_output_url = output_url
|
||||
upscale.provider_output_url_expires_at = __import__("datetime").datetime.fromtimestamp(query_result.expires_at, tz=__import__("datetime").timezone.utc) if query_result.expires_at else None
|
||||
|
||||
import os
|
||||
from app.services.video_gen import download_video
|
||||
|
||||
api_task_id = upscale.api_generation_task_id
|
||||
final_path = None
|
||||
|
||||
try:
|
||||
# 直接下载到最终路径
|
||||
date_dir = _now_str()
|
||||
# 相对 URL 路径
|
||||
url_path = f"/generate/api/videos/{date_dir}/{api_task_id}.mp4"
|
||||
# 绝对文件路径
|
||||
z_url_path = f"./storage/generate/api/videos/{date_dir}/{api_task_id}.mp4"
|
||||
abs_path = os.path.abspath(z_url_path)
|
||||
os.makedirs(os.path.dirname(abs_path), exist_ok=True)
|
||||
await download_video(output_url, abs_path)
|
||||
|
||||
upscale.final_local_path = z_url_path
|
||||
upscale.final_resource_url = url_path # 相对 URL
|
||||
final_path = url_path
|
||||
|
||||
# 检查 API Key 超分配置中的"成功后删除源文件"设置
|
||||
_maybe_delete_source_file(db, upscale)
|
||||
except Exception:
|
||||
upscale.final_local_path = output_url
|
||||
upscale.final_resource_url = output_url
|
||||
final_path = output_url
|
||||
|
||||
upscale.status = "completed"
|
||||
upscale.stage = "upscale_completed"
|
||||
upscale.completed_at = _now()
|
||||
# 获取文件大小
|
||||
try:
|
||||
import os
|
||||
abs_path = os.path.abspath(final_path) if final_path and final_path.startswith(".") else None
|
||||
if abs_path and os.path.exists(abs_path):
|
||||
upscale.final_file_size_bytes = os.path.getsize(abs_path)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 更新 API 任务
|
||||
if api_task_id:
|
||||
from app.models.api.api_generation_task import ApiGenerationTask
|
||||
api_result = await db.execute(
|
||||
select(ApiGenerationTask).where(ApiGenerationTask.id == api_task_id).limit(1)
|
||||
)
|
||||
api_task = api_result.scalar_one_or_none()
|
||||
if api_task:
|
||||
api_task.video_url = final_path or output_url
|
||||
api_task.status = "completed"
|
||||
api_task.pipeline_stage = "done"
|
||||
api_task.generated_at = upscale.completed_at
|
||||
|
||||
await db.commit()
|
||||
return {"status": "completed", "output_url": final_path or output_url}
|
||||
|
||||
elif status == "failed":
|
||||
upscale.status = "failed"
|
||||
upscale.stage = "failed"
|
||||
upscale.failure_count = int(upscale.failure_count or 0) + 1
|
||||
upscale.last_error = str(query_result.error) if query_result.error else "超分失败"
|
||||
upscale.failed_at = __import__("datetime").datetime.now(__import__("datetime").timezone.utc)
|
||||
await db.commit()
|
||||
return {"status": "failed"}
|
||||
|
||||
else:
|
||||
# 仍在处理中,继续轮询
|
||||
upscale.stage = "remote_polling"
|
||||
await db.commit()
|
||||
# 重新调度下一次轮询
|
||||
api_upscale_poll_remote_simple.apply_async(
|
||||
args=[upscale_task_id],
|
||||
countdown=30,
|
||||
queue="gen_api_upscale",
|
||||
)
|
||||
return {"status": "polling"}
|
||||
|
||||
except Exception as exc:
|
||||
upscale.failure_count = int(upscale.failure_count or 0) + 1
|
||||
upscale.last_error = str(exc)[:500]
|
||||
await db.commit()
|
||||
raise
|
||||
|
||||
from app.tasks.async_runner import run_async
|
||||
return run_async(_execute())
|
||||
@@ -33,7 +33,11 @@ CELERY_TASK_IMPORTS = (
|
||||
"app.tasks.module_async_recovery_tasks",
|
||||
"app.tasks.module_generation_v2_tasks",
|
||||
"app.tasks.private_portrait_asset_tasks",
|
||||
"app.tasks.vp_v3_asset_tasks",
|
||||
"app.tasks.celery_runtime_tasks",
|
||||
"app.tasks.api_generation_tasks",
|
||||
"app.tasks.api_recovery_tasks",
|
||||
"app.tasks.api_upscale_tasks",
|
||||
)
|
||||
|
||||
|
||||
@@ -85,6 +89,14 @@ def _beat_schedule() -> dict:
|
||||
"priority": settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
|
||||
},
|
||||
}
|
||||
schedule["api-generation-recovery-every-minute"] = {
|
||||
"task": "api_generation.recover_tasks_once",
|
||||
"schedule": 60,
|
||||
"options": {
|
||||
"queue": RECOVERY_QUEUE,
|
||||
"priority": settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
|
||||
},
|
||||
}
|
||||
schedule["generation-download-recovery"] = {
|
||||
"task": CeleryTaskName.RECOVER_DOWNLOAD.value,
|
||||
"schedule": max(1, int(settings.DOWNLOAD_RECOVERY_INTERVAL_SECONDS or 60)),
|
||||
@@ -120,6 +132,16 @@ def _beat_schedule() -> dict:
|
||||
"schedule": 300,
|
||||
"options": {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value},
|
||||
}
|
||||
schedule["vp-v3-sync-due-assets-every-minute"] = {
|
||||
"task": CeleryTaskName.VP_V3_SYNC_DUE_ASSETS.value,
|
||||
"schedule": 60,
|
||||
"options": {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value},
|
||||
}
|
||||
schedule["vp-v3-recover-remote-deletes-every-5-minutes"] = {
|
||||
"task": CeleryTaskName.VP_V3_RECOVER_REMOTE_DELETES.value,
|
||||
"schedule": 300,
|
||||
"options": {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value},
|
||||
}
|
||||
return schedule
|
||||
|
||||
|
||||
@@ -240,6 +262,11 @@ if broker_url:
|
||||
CeleryTaskName.PRIVATE_PORTRAIT_DELETE_GROUP.value: {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value},
|
||||
CeleryTaskName.PRIVATE_PORTRAIT_DELETE_PROJECT.value: {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value},
|
||||
CeleryTaskName.PRIVATE_PORTRAIT_RECOVER_REMOTE_DELETES.value: {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value},
|
||||
CeleryTaskName.VP_V3_POLL_ASSET.value: {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value},
|
||||
CeleryTaskName.VP_V3_SYNC_DUE_ASSETS.value: {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value},
|
||||
CeleryTaskName.VP_V3_DELETE_ASSET.value: {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value},
|
||||
CeleryTaskName.VP_V3_DELETE_PROJECT.value: {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value},
|
||||
CeleryTaskName.VP_V3_RECOVER_REMOTE_DELETES.value: {"queue": CeleryQueue.GEN_PRIVATE_PORTRAIT.value},
|
||||
},
|
||||
)
|
||||
else:
|
||||
@@ -404,6 +431,7 @@ def on_worker_ready(sender=None, **kwargs):
|
||||
try:
|
||||
from app.services.celery_runtime.recovery_service import set_startup_barrier
|
||||
from app.tasks.generation_recovery_tasks import startup_recovery_once
|
||||
from app.tasks.api_recovery_tasks import api_generation_recover_tasks_once
|
||||
|
||||
run_async(set_startup_barrier())
|
||||
countdown = max(0, int(settings.CELERY_STARTUP_RECOVERY_DELAY_SECONDS or 30))
|
||||
@@ -412,8 +440,14 @@ def on_worker_ready(sender=None, **kwargs):
|
||||
queue=RECOVERY_QUEUE,
|
||||
priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
|
||||
)
|
||||
# API v3 任务恢复(延迟 35 秒执行,避免与其他恢复任务冲突)
|
||||
api_generation_recover_tasks_once.apply_async(
|
||||
countdown=countdown + 5,
|
||||
queue=RECOVERY_QUEUE,
|
||||
priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
|
||||
)
|
||||
logger.info(
|
||||
"启动容灾恢复协调任务已投递。queue=%s countdown=%s",
|
||||
"启动容灾恢复协调任务已投递(含 API v3)。queue=%s countdown=%s",
|
||||
RECOVERY_QUEUE,
|
||||
countdown,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,449 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import select
|
||||
|
||||
from app.config import settings
|
||||
from app.enums.celery_queue import CeleryQueue, CeleryTaskName
|
||||
from app.enums.celery_runtime import CeleryRuntimeDomain
|
||||
from app.enums.private_portrait import (
|
||||
PrivatePortraitAssetStatus,
|
||||
PrivatePortraitEventSource,
|
||||
PrivatePortraitEventType,
|
||||
PrivatePortraitRemoteDeleteStatus,
|
||||
)
|
||||
from app.models import async_session
|
||||
from app.models.virtual_portrait_v3 import VpV3Asset, VpV3Project
|
||||
from app.services.celery_runtime.recovery_service import guard_periodic_recovery
|
||||
from app.services.celery_runtime.runtime_service import CeleryRuntimeLease, RuntimeIdentity
|
||||
from app.services.operation_log_service import log_operation_error, log_operation_event
|
||||
from app.services.redis_registry_service import RedisExecutionLockLease
|
||||
from app.services.virtual_portrait_v3.asset_service import (
|
||||
V3_DOMAIN,
|
||||
delete_v3_asset_remote,
|
||||
sync_asset_status,
|
||||
)
|
||||
from app.services.virtual_portrait_v3.project_service import (
|
||||
V3_DOMAIN as V3_PROJECT_DOMAIN,
|
||||
delete_v3_project_remote,
|
||||
)
|
||||
from app.tasks.async_runner import run_async
|
||||
from app.tasks.celery_app import celery_app
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
QUEUE = CeleryQueue.GEN_PRIVATE_PORTRAIT.value
|
||||
|
||||
|
||||
# 北京时间(UTC+8)统一基准:V3 业务所有时间写入 / 时间比较唯一参考
|
||||
_BJ_TZ = timezone(timedelta(hours=8))
|
||||
|
||||
|
||||
def _bj_now() -> datetime:
|
||||
"""返回当前北京时间(UTC+8)naive datetime(去掉 tzinfo,与 DB naive 存储一致)。"""
|
||||
return datetime.now(_BJ_TZ).replace(tzinfo=None)
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
"""统一使用北京时间基准,与业务写入保持一致。"""
|
||||
return _bj_now()
|
||||
|
||||
|
||||
def _naive(dt: datetime | None) -> datetime | None:
|
||||
"""把 datetime 统一成 naive 北京时间(去掉 tzinfo),避免 offset-aware vs naive 比较报错。
|
||||
|
||||
DB 列是 DateTime(timezone=True) 但业务写入都是北京时间(naive),
|
||||
读回时根据方言可能变成 aware 或仍为 naive,比较前统一去掉 tzinfo。
|
||||
"""
|
||||
if dt is None:
|
||||
return None
|
||||
return dt.replace(tzinfo=None) if dt.tzinfo is not None else dt
|
||||
|
||||
|
||||
def _retry_countdown(retries: int) -> int:
|
||||
return min(300, 30 * (2 ** max(0, retries)))
|
||||
|
||||
|
||||
async def _rollback_and_reraise(
|
||||
db,
|
||||
*,
|
||||
event_type: str,
|
||||
exc: BaseException,
|
||||
detail: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
await db.rollback()
|
||||
log_operation_error(
|
||||
domain=V3_DOMAIN,
|
||||
event_type=event_type,
|
||||
source=PrivatePortraitEventSource.CELERY.value,
|
||||
exc=exc,
|
||||
detail=detail,
|
||||
**kwargs,
|
||||
)
|
||||
raise exc
|
||||
|
||||
|
||||
async def _acquire_v3_runtime(
|
||||
*,
|
||||
domain: str,
|
||||
owner_type: str,
|
||||
owner_id: str,
|
||||
task_name: str,
|
||||
hash_key: str,
|
||||
zset_key: str,
|
||||
lock_prefix: str,
|
||||
) -> CeleryRuntimeLease | None:
|
||||
token = uuid.uuid4().hex
|
||||
return await CeleryRuntimeLease.acquire(
|
||||
identity=RuntimeIdentity(
|
||||
domain=domain,
|
||||
owner_type=owner_type,
|
||||
owner_id=owner_id,
|
||||
attempt_no=1,
|
||||
task_name=task_name,
|
||||
queue=QUEUE,
|
||||
),
|
||||
lock_key=f"{lock_prefix}:{owner_type}:{owner_id}:attempt:1",
|
||||
hash_key=hash_key,
|
||||
zset_key=zset_key,
|
||||
token=token,
|
||||
ttl_seconds=max(60, int(settings.VP_V3_RUNTIME_LOCK_TTL_SECONDS or 180)),
|
||||
heartbeat_interval_seconds=max(10, int(settings.VP_V3_RUNTIME_HEARTBEAT_SECONDS or 30)),
|
||||
pipeline_stage="processing",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 轮询:单条素材
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _run_poll_v3_asset(asset_id: str) -> None:
|
||||
lease = await _acquire_v3_runtime(
|
||||
domain=CeleryRuntimeDomain.PRIVATE_PORTRAIT_POLL.value,
|
||||
owner_type="asset",
|
||||
owner_id=asset_id,
|
||||
task_name=CeleryTaskName.VP_V3_POLL_ASSET.value,
|
||||
hash_key=settings.VP_V3_POLL_ACTIVE_REDIS_HASH_KEY,
|
||||
zset_key=settings.VP_V3_POLL_ACTIVE_REDIS_ZSET_KEY,
|
||||
lock_prefix=settings.VP_V3_POLL_LOCK_KEY_PREFIX,
|
||||
)
|
||||
if lease is None:
|
||||
logger.info("vp_v3 poll asset skip: runtime lease not acquired (asset_id=%s)", asset_id)
|
||||
return
|
||||
try:
|
||||
async with async_session() as db:
|
||||
try:
|
||||
row = (await db.execute(
|
||||
select(VpV3Asset.api_key_id).where(VpV3Asset.remote_asset_id == asset_id).limit(1)
|
||||
)).scalar_one_or_none()
|
||||
if not row:
|
||||
return
|
||||
asset = await sync_asset_status(
|
||||
db,
|
||||
api_key_id=str(row),
|
||||
asset_id=asset_id,
|
||||
execution_guard=lease.ensure_owned,
|
||||
)
|
||||
await lease.ensure_owned()
|
||||
await db.commit()
|
||||
logger.info(
|
||||
"vp_v3 poll asset synced: asset_id=%s status=%s poll_count=%s",
|
||||
asset_id, asset.status, asset.poll_count,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.exception("vp_v3 poll asset failed: %s", asset_id)
|
||||
await _rollback_and_reraise(
|
||||
db,
|
||||
event_type=PrivatePortraitEventType.ASSET_POLL_FAILED.value,
|
||||
exc=exc,
|
||||
asset_id=asset_id,
|
||||
detail={"celery_task": CeleryTaskName.VP_V3_POLL_ASSET.value},
|
||||
)
|
||||
finally:
|
||||
await lease.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 轮询:每分钟批量扫描到期素材并分发轮询任务
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _dispatch_v3_due_assets() -> int:
|
||||
barrier = await guard_periodic_recovery()
|
||||
if barrier is not None:
|
||||
logger.info("vp_v3 dispatch due assets skip: periodic recovery barrier active")
|
||||
return 0
|
||||
lock = await RedisExecutionLockLease.acquire(
|
||||
lock_key=settings.VP_V3_DISPATCH_LOCK_KEY,
|
||||
ttl_seconds=55,
|
||||
renew_interval_seconds=20,
|
||||
log_context="vp_v3_poll_dispatch",
|
||||
)
|
||||
if lock is None:
|
||||
logger.info("vp_v3 dispatch due assets skip: dispatch lock not acquired")
|
||||
return 0
|
||||
async with lock:
|
||||
async with async_session() as db:
|
||||
now_naive = _naive(_now())
|
||||
rows = await db.execute(
|
||||
select(VpV3Asset)
|
||||
.where(
|
||||
VpV3Asset.deleted_at.is_(None),
|
||||
VpV3Asset.status == PrivatePortraitAssetStatus.CREATING.value,
|
||||
VpV3Asset.next_poll_at.is_not(None),
|
||||
)
|
||||
.order_by(VpV3Asset.next_poll_at.asc(), VpV3Asset.id.asc())
|
||||
.limit(settings.VP_V3_ASSET_POLL_BATCH_SIZE or 50)
|
||||
.with_for_update(skip_locked=True)
|
||||
)
|
||||
assets = list(rows.scalars().all())
|
||||
# next_poll_at <= now 在内存里过滤(统一 naive 比较,避免 aware vs naive 报错)
|
||||
assets = [a for a in assets if _naive(a.next_poll_at) is not None and _naive(a.next_poll_at) <= now_naive]
|
||||
dispatches: list[tuple[str, int]] = []
|
||||
queue_hold_until_naive = now_naive + timedelta(seconds=120)
|
||||
for asset in assets:
|
||||
poll_no = int(asset.poll_count or 0) + 1
|
||||
dispatches.append((str(asset.remote_asset_id), poll_no))
|
||||
asset.next_poll_at = queue_hold_until_naive
|
||||
await db.commit()
|
||||
|
||||
for asset_id, poll_no in dispatches:
|
||||
poll_v3_asset_status.apply_async(
|
||||
args=[asset_id],
|
||||
queue=QUEUE,
|
||||
countdown=0,
|
||||
task_id=f"vp-v3-poll:{asset_id}:attempt:{poll_no}",
|
||||
)
|
||||
log_operation_event(
|
||||
domain=V3_DOMAIN,
|
||||
event_type=PrivatePortraitEventType.SYNC_DUE_ASSETS_DONE.value,
|
||||
event_status="success",
|
||||
source=PrivatePortraitEventSource.CELERY.value,
|
||||
detail={"matched_count": len(dispatches), "dispatched_count": len(dispatches)},
|
||||
)
|
||||
return len(dispatches)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 删除:素材 / 项目远端删除(已存在)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _run_delete_v3_asset(asset_id: str) -> None:
|
||||
"""执行 V3 素材远端删除。"""
|
||||
lease = await _acquire_v3_runtime(
|
||||
domain=CeleryRuntimeDomain.PRIVATE_PORTRAIT_DELETE.value,
|
||||
owner_type="asset",
|
||||
owner_id=asset_id,
|
||||
task_name=CeleryTaskName.VP_V3_DELETE_ASSET.value,
|
||||
hash_key=settings.VP_V3_DELETE_ACTIVE_REDIS_HASH_KEY,
|
||||
zset_key=settings.VP_V3_DELETE_ACTIVE_REDIS_ZSET_KEY,
|
||||
lock_prefix=settings.VP_V3_DELETE_LOCK_KEY_PREFIX,
|
||||
)
|
||||
if lease is None:
|
||||
logger.info("vp_v3 delete asset skip: runtime lease not acquired (asset_id=%s)", asset_id)
|
||||
return
|
||||
try:
|
||||
async with async_session() as db:
|
||||
try:
|
||||
await delete_v3_asset_remote(
|
||||
db,
|
||||
asset_id=asset_id,
|
||||
execution_guard=lease.ensure_owned,
|
||||
)
|
||||
await lease.ensure_owned()
|
||||
await db.commit()
|
||||
except Exception as exc:
|
||||
await _rollback_and_reraise(
|
||||
db,
|
||||
event_type=PrivatePortraitEventType.ASSET_DELETE_REMOTE_FAILED.value,
|
||||
exc=exc,
|
||||
detail={"asset_id": asset_id},
|
||||
)
|
||||
finally:
|
||||
await lease.close()
|
||||
|
||||
|
||||
async def _run_delete_v3_project(project_id: str) -> int:
|
||||
"""执行 V3 项目远端删除(级联删除素材 + 项目)。"""
|
||||
lease = await _acquire_v3_runtime(
|
||||
domain=CeleryRuntimeDomain.PRIVATE_PORTRAIT_DELETE.value,
|
||||
owner_type="project",
|
||||
owner_id=project_id,
|
||||
task_name=CeleryTaskName.VP_V3_DELETE_PROJECT.value,
|
||||
hash_key=settings.VP_V3_DELETE_ACTIVE_REDIS_HASH_KEY,
|
||||
zset_key=settings.VP_V3_DELETE_ACTIVE_REDIS_ZSET_KEY,
|
||||
lock_prefix=settings.VP_V3_DELETE_LOCK_KEY_PREFIX,
|
||||
)
|
||||
if lease is None:
|
||||
logger.info("vp_v3 delete project skip: runtime lease not acquired (project_id=%s)", project_id)
|
||||
return 0
|
||||
try:
|
||||
async with async_session() as db:
|
||||
try:
|
||||
await delete_v3_project_remote(
|
||||
db,
|
||||
project_id=project_id,
|
||||
)
|
||||
await lease.ensure_owned()
|
||||
await db.commit()
|
||||
except Exception as exc:
|
||||
await _rollback_and_reraise(
|
||||
db,
|
||||
event_type=PrivatePortraitEventType.PROJECT_DELETE_REMOTE_FAILED.value,
|
||||
exc=exc,
|
||||
detail={"project_id": project_id},
|
||||
)
|
||||
finally:
|
||||
await lease.close()
|
||||
return 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 删除恢复:每 5 分钟扫描 pending/failed 的 project/asset 再投递
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _dispatch_v3_remote_delete_recovery() -> dict[str, int]:
|
||||
barrier = await guard_periodic_recovery()
|
||||
if barrier is not None:
|
||||
logger.info("vp_v3 delete recovery skip: periodic recovery barrier active")
|
||||
return {"asset_count": 0, "project_count": 0, "total_count": 0}
|
||||
lock = await RedisExecutionLockLease.acquire(
|
||||
lock_key=settings.VP_V3_DELETE_RECOVERY_LOCK_KEY,
|
||||
ttl_seconds=240,
|
||||
renew_interval_seconds=30,
|
||||
log_context="vp_v3_delete_recovery",
|
||||
)
|
||||
if lock is None:
|
||||
logger.info("vp_v3 delete recovery skip: recovery lock not acquired")
|
||||
return {"asset_count": 0, "project_count": 0, "total_count": 0}
|
||||
async with lock:
|
||||
statuses = [
|
||||
PrivatePortraitRemoteDeleteStatus.PENDING.value,
|
||||
PrivatePortraitRemoteDeleteStatus.FAILED.value,
|
||||
]
|
||||
batch_size = max(1, int(settings.VP_V3_REMOTE_DELETE_RECOVERY_BATCH_SIZE or 50))
|
||||
async with async_session() as db:
|
||||
asset_rows = await db.execute(
|
||||
select(VpV3Asset.id)
|
||||
.where(VpV3Asset.remote_delete_status.in_(statuses))
|
||||
.order_by(VpV3Asset.updated_at.asc(), VpV3Asset.id.asc())
|
||||
.limit(batch_size)
|
||||
)
|
||||
asset_ids = [str(value) for value in asset_rows.scalars().all()]
|
||||
remaining = max(0, batch_size - len(asset_ids))
|
||||
project_ids: list[str] = []
|
||||
if remaining:
|
||||
project_rows = await db.execute(
|
||||
select(VpV3Project.id)
|
||||
.where(VpV3Project.remote_delete_status.in_(statuses))
|
||||
.order_by(VpV3Project.updated_at.asc(), VpV3Project.id.asc())
|
||||
.limit(remaining)
|
||||
)
|
||||
project_ids = [str(value) for value in project_rows.scalars().all()]
|
||||
await db.rollback()
|
||||
|
||||
for asset_id in asset_ids:
|
||||
delete_v3_asset_remote_task.apply_async(
|
||||
args=[asset_id], queue=QUEUE, task_id=f"vp-v3-delete-asset:{asset_id}"
|
||||
)
|
||||
for project_id in project_ids:
|
||||
delete_v3_project_remote_task.apply_async(
|
||||
args=[project_id], queue=QUEUE, task_id=f"vp-v3-delete-project:{project_id}"
|
||||
)
|
||||
return {
|
||||
"asset_count": len(asset_ids),
|
||||
"project_count": len(project_ids),
|
||||
"total_count": len(asset_ids) + len(project_ids),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Celery 任务注册
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@celery_app.task(
|
||||
name=CeleryTaskName.VP_V3_POLL_ASSET.value,
|
||||
queue=QUEUE,
|
||||
bind=True,
|
||||
max_retries=5,
|
||||
default_retry_delay=30,
|
||||
)
|
||||
def poll_v3_asset_status(self, asset_id: str) -> None:
|
||||
"""V3 素材单条状态轮询(Celery 任务)。"""
|
||||
logger.info("vp_v3 poll task START: asset_id=%s task_id=%s", asset_id, self.request.id)
|
||||
try:
|
||||
return run_async(_run_poll_v3_asset(asset_id))
|
||||
except Exception as exc:
|
||||
raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))
|
||||
|
||||
|
||||
@celery_app.task(
|
||||
name=CeleryTaskName.VP_V3_SYNC_DUE_ASSETS.value,
|
||||
queue=QUEUE,
|
||||
bind=True,
|
||||
max_retries=3,
|
||||
default_retry_delay=60,
|
||||
)
|
||||
def sync_v3_due_assets(self) -> int:
|
||||
"""每分钟扫描 V3 到期素材并分发轮询任务(beat schedule)。"""
|
||||
logger.info("vp_v3 sync_due_assets START: task_id=%s", self.request.id)
|
||||
try:
|
||||
return run_async(_dispatch_v3_due_assets())
|
||||
except Exception as exc:
|
||||
raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))
|
||||
|
||||
|
||||
@celery_app.task(
|
||||
name=CeleryTaskName.VP_V3_DELETE_ASSET.value,
|
||||
queue=QUEUE,
|
||||
bind=True,
|
||||
max_retries=3,
|
||||
default_retry_delay=60,
|
||||
)
|
||||
def delete_v3_asset_remote_task(self, asset_id: str) -> None:
|
||||
"""V3 素材远端删除 Celery 任务。"""
|
||||
logger.info("vp_v3 delete asset START: asset_id=%s task_id=%s", asset_id, self.request.id)
|
||||
try:
|
||||
return run_async(_run_delete_v3_asset(asset_id))
|
||||
except Exception as exc:
|
||||
raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))
|
||||
|
||||
|
||||
@celery_app.task(
|
||||
name=CeleryTaskName.VP_V3_DELETE_PROJECT.value,
|
||||
queue=QUEUE,
|
||||
bind=True,
|
||||
max_retries=3,
|
||||
default_retry_delay=60,
|
||||
)
|
||||
def delete_v3_project_remote_task(self, project_id: str) -> int:
|
||||
"""V3 项目远端删除 Celery 任务。"""
|
||||
logger.info("vp_v3 delete project START: project_id=%s task_id=%s", project_id, self.request.id)
|
||||
try:
|
||||
return run_async(_run_delete_v3_project(project_id))
|
||||
except Exception as exc:
|
||||
raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))
|
||||
|
||||
|
||||
@celery_app.task(
|
||||
name=CeleryTaskName.VP_V3_RECOVER_REMOTE_DELETES.value,
|
||||
queue=QUEUE,
|
||||
bind=True,
|
||||
max_retries=3,
|
||||
default_retry_delay=60,
|
||||
)
|
||||
def recover_v3_remote_deletes(self) -> dict[str, int]:
|
||||
"""每 5 分钟扫描 V3 pending/failed 远端删除记录并重新投递(beat schedule)。"""
|
||||
logger.info("vp_v3 recover_remote_deletes START: task_id=%s", self.request.id)
|
||||
try:
|
||||
return run_async(_dispatch_v3_remote_delete_recovery())
|
||||
except Exception as exc:
|
||||
raise self.retry(exc=exc, countdown=_retry_countdown(self.request.retries))
|
||||
Reference in New Issue
Block a user