846 lines
34 KiB
Python
846 lines
34 KiB
Python
"""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):
|
||
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:
|
||
source_width = source_info.width
|
||
source_height = source_info.height
|
||
source_duration = 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,
|
||
source_width=source_width,
|
||
source_height=source_height,
|
||
source_duration=source_duration,
|
||
source_file_size_bytes=source_file_size_bytes
|
||
)
|
||
|
||
# 如果已存在超分任务（重复调用），检查超分状态
|
||
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=CeleryQueue.GEN_API_UPSCALE.value,
|
||
)
|
||
else:
|
||
# 远程超分（火山 MediaKit）
|
||
api_upscale_submit_remote_simple.apply_async(
|
||
args=[upscale_task.id],
|
||
queue=CeleryQueue.GEN_API_UPSCALE.value,
|
||
)
|
||
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}
|