377 lines
15 KiB
Python
377 lines
15 KiB
Python
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
from datetime import datetime
|
|
from types import SimpleNamespace
|
|
|
|
from sqlalchemy import select
|
|
|
|
from app.models.base import async_session
|
|
from app.models.generation_record import GenerationRecord
|
|
from app.models.image_engine import ImageEngine
|
|
from app.models.video_engine import VideoEngine
|
|
from app.enums.model_pricing import PricingSnapshotStage, ProviderCostStatus
|
|
from app.services.video_gen import poll_task_status, download_video, _log_video_response
|
|
from app.services.image_gen import download_image, is_sync_image_provider_result_uncertain
|
|
from app.services.media_token_usage_snapshot_service import (
|
|
mark_media_provider_cost_status,
|
|
sync_generation_record_media_token_snapshot,
|
|
)
|
|
from app.services.resource_accounting_service import (
|
|
record_generation_record_generated_resource,
|
|
safe_file_size,
|
|
)
|
|
from app.services.video_cover_service import create_video_cover_for_local_video
|
|
from app.config import settings
|
|
from app.services.generation_refund_service import mark_generation_record_failed_and_refund_once
|
|
|
|
logger = logging.getLogger("videogen")
|
|
|
|
POLL_INTERVAL = 30 # seconds between polls
|
|
MAX_POLLS = 60 # max 30 minutes total
|
|
|
|
|
|
def _loads(data: str | None) -> dict:
|
|
if not data:
|
|
return {}
|
|
try:
|
|
value = json.loads(data)
|
|
return value if isinstance(value, dict) else {}
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
async def _get_runtime_engine(db, record: GenerationRecord):
|
|
"""使用生成开始时冻结的引擎快照，数据库行只读取当前密钥。"""
|
|
if not record.engine_id:
|
|
raise ValueError("生成记录缺少锁定的 engine_id")
|
|
snapshot = _loads(record.engine_snapshot_json)
|
|
model = ImageEngine if record.gen_type == "image" else VideoEngine
|
|
engine = (await db.execute(select(model).where(model.id == record.engine_id).limit(1))).scalar_one_or_none()
|
|
if not engine:
|
|
raise ValueError("锁定的生成引擎不存在")
|
|
return SimpleNamespace(
|
|
id=record.engine_id,
|
|
name=snapshot.get("name") or engine.name,
|
|
provider=snapshot.get("provider") or engine.provider,
|
|
api_base=snapshot.get("api_base") or engine.api_base,
|
|
api_key=engine.api_key,
|
|
model_name=snapshot.get("model_name") or engine.model_name,
|
|
generate_url=snapshot.get("generate_url") or getattr(engine, "generate_url", ""),
|
|
query_url=snapshot.get("query_url") or getattr(engine, "query_url", ""),
|
|
default_size=snapshot.get("default_size") or getattr(engine, "default_size", "2K"),
|
|
)
|
|
|
|
|
|
class TaskQueue:
|
|
def __init__(self):
|
|
self.queue: asyncio.Queue[str] = asyncio.Queue()
|
|
self.running = False
|
|
self._active: dict[str, int] = {} # record_id -> poll count
|
|
|
|
async def enqueue(self, record_id: str):
|
|
"""Add a record to the polling queue."""
|
|
await self.queue.put(record_id)
|
|
|
|
async def recover(self):
|
|
"""Recover in-progress tasks from DB on startup."""
|
|
async with async_session() as db:
|
|
result = await db.execute(
|
|
select(GenerationRecord).where(
|
|
GenerationRecord.status == "generating",
|
|
GenerationRecord.seedance_task_id.isnot(None),
|
|
GenerationRecord.deleted_at.is_(None),
|
|
)
|
|
)
|
|
records = result.scalars().all()
|
|
for record in records:
|
|
await self.queue.put(record.id)
|
|
logger.info(f"Recovered task: {record.id} (seedance: {record.seedance_task_id})")
|
|
|
|
async def run(self):
|
|
"""Main polling loop."""
|
|
self.running = True
|
|
logger.info("Video queue started")
|
|
|
|
while self.running:
|
|
try:
|
|
record_id = await asyncio.wait_for(self.queue.get(), timeout=5.0)
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
|
|
try:
|
|
await self._process(record_id)
|
|
except Exception as e:
|
|
logger.error(f"Error processing {record_id}: {e}")
|
|
finally:
|
|
self.queue.task_done()
|
|
|
|
logger.info("Video queue stopped")
|
|
|
|
async def _process(self, record_id: str):
|
|
"""Process a single record: poll status and update DB."""
|
|
async with async_session() as db:
|
|
result = await db.execute(
|
|
select(GenerationRecord).where(
|
|
GenerationRecord.id == record_id,
|
|
GenerationRecord.deleted_at.is_(None),
|
|
)
|
|
.with_for_update()
|
|
.limit(1)
|
|
)
|
|
record = result.scalar_one_or_none()
|
|
if not record or record.status != "generating":
|
|
return
|
|
|
|
if record.gen_type == "video":
|
|
await self._process_video(db, record)
|
|
else:
|
|
await self._process_image(db, record)
|
|
|
|
async def _process_video(self, db, record):
|
|
"""Process video generation task."""
|
|
record_id = record.id
|
|
|
|
if not record.seedance_task_id:
|
|
await mark_generation_record_failed_and_refund_once(
|
|
db,
|
|
record=record,
|
|
error_message="缺少外部任务ID",
|
|
)
|
|
await db.commit()
|
|
return
|
|
|
|
try:
|
|
engine = await _get_runtime_engine(db, record)
|
|
poll_result = await poll_task_status(engine, record.seedance_task_id)
|
|
except Exception as e:
|
|
logger.error(f"Poll error for {record_id}: {e}")
|
|
count = self._active.get(record_id, 0) + 1
|
|
self._active[record_id] = count
|
|
if count >= MAX_POLLS:
|
|
await mark_generation_record_failed_and_refund_once(
|
|
db,
|
|
record=record,
|
|
error_message=f"轮询超时: {e}",
|
|
)
|
|
await db.commit()
|
|
del self._active[record_id]
|
|
else:
|
|
await asyncio.sleep(POLL_INTERVAL)
|
|
await self.queue.put(record_id)
|
|
return
|
|
|
|
status = poll_result["status"]
|
|
try:
|
|
resp_data = json.loads(poll_result.get("response_data", "{}"))
|
|
except (json.JSONDecodeError, TypeError):
|
|
resp_data = {}
|
|
|
|
_log_video_response(record_id, resp_data, poll_result.get("error"))
|
|
|
|
if status == "succeeded":
|
|
file_url = poll_result.get("video_url", "")
|
|
resp_data["video_url"] = file_url
|
|
resp_data["task_id"] = record.seedance_task_id
|
|
record.provider_response_json = json.dumps(resp_data, ensure_ascii=False, default=str)
|
|
record.video_tokens_used = poll_result.get("video_tokens", 0)
|
|
# 视频 Provider 已完成时核算供应商成本；下载失败不影响已发生的供应商费用。
|
|
await sync_generation_record_media_token_snapshot(
|
|
db,
|
|
record,
|
|
provider_response=resp_data,
|
|
stage=PricingSnapshotStage.PROVIDER_ASYNC_COMPLETED.value,
|
|
)
|
|
storage_path = None
|
|
file_size_bytes = 0
|
|
if settings.STORAGE_TYPE == "local" and file_url:
|
|
try:
|
|
date_dir = datetime.now().strftime("%Y/%m/%d")
|
|
dest_dir = os.path.join(settings.STORAGE_LOCAL_PATH, date_dir)
|
|
os.makedirs(dest_dir, exist_ok=True)
|
|
dest = os.path.join(dest_dir, f"{record_id}.mp4")
|
|
await download_video(file_url, dest)
|
|
record.video_url = f"/generate/videos/{date_dir}/{record_id}.mp4"
|
|
cover_url, _cover_storage_path = create_video_cover_for_local_video(
|
|
record_id=record_id,
|
|
video_path=dest,
|
|
date_dir=date_dir,
|
|
log_prefix=f"GenerationRecord视频封面生成 record_id={record_id}",
|
|
)
|
|
record.video_cover_url = cover_url
|
|
storage_path = dest
|
|
file_size_bytes = safe_file_size(dest)
|
|
except Exception as e:
|
|
logger.warning(f"Download failed, using remote URL: {e}")
|
|
record.video_url = file_url
|
|
else:
|
|
record.video_url = file_url
|
|
record.status = "completed"
|
|
record.generated_at = datetime.now()
|
|
if record.video_url:
|
|
await record_generation_record_generated_resource(
|
|
db,
|
|
record,
|
|
resource_url=record.video_url,
|
|
storage_path=storage_path,
|
|
file_size_bytes=file_size_bytes,
|
|
remote_url=file_url,
|
|
generated_at=record.generated_at,
|
|
)
|
|
await sync_generation_record_media_token_snapshot(
|
|
db,
|
|
record,
|
|
provider_response=resp_data,
|
|
stage=PricingSnapshotStage.RESOURCE_DOWNLOAD_COMPLETED.value,
|
|
)
|
|
self._active.pop(record_id, None)
|
|
await db.commit()
|
|
logger.info(f"Video task completed: {record_id}")
|
|
|
|
elif status == "failed":
|
|
await mark_generation_record_failed_and_refund_once(
|
|
db,
|
|
record=record,
|
|
error_message=poll_result.get("error", "视频生成失败"),
|
|
)
|
|
self._active.pop(record_id, None)
|
|
await db.commit()
|
|
logger.info(f"Video task failed: {record_id}")
|
|
|
|
else:
|
|
count = self._active.get(record_id, 0) + 1
|
|
self._active[record_id] = count
|
|
if count >= MAX_POLLS:
|
|
await mark_generation_record_failed_and_refund_once(
|
|
db,
|
|
record=record,
|
|
error_message="视频生成超时",
|
|
)
|
|
self._active.pop(record_id, None)
|
|
await db.commit()
|
|
logger.info(f"Video task timed out: {record_id}")
|
|
else:
|
|
await db.commit()
|
|
await asyncio.sleep(POLL_INTERVAL)
|
|
await self.queue.put(record_id)
|
|
|
|
async def _process_image(self, db, record):
|
|
"""Process image generation task - calls API directly."""
|
|
record_id = record.id
|
|
from app.services.image_gen import submit_image_task, _log_image_response
|
|
|
|
provider_call_started = False
|
|
provider_call_completed = False
|
|
try:
|
|
engine = await _get_runtime_engine(db, record)
|
|
provider_call_started = True
|
|
poll_result = await asyncio.to_thread(
|
|
submit_image_task,
|
|
db,
|
|
engine,
|
|
record,
|
|
include_media_references=False,
|
|
)
|
|
provider_call_completed = True
|
|
|
|
if poll_result["error"] == "":
|
|
remote_url = poll_result.get("image_url")
|
|
try:
|
|
provider_response = json.loads(poll_result.get("response_data") or "{}")
|
|
except (json.JSONDecodeError, TypeError):
|
|
provider_response = {}
|
|
record.provider_response_json = json.dumps(provider_response, ensure_ascii=False, default=str)
|
|
record.image_tokens_used = poll_result.get("image_tokens", 0)
|
|
# 火山图片接口为同步生成：最终响应返回后立即完成供应商成本核算。
|
|
await sync_generation_record_media_token_snapshot(
|
|
db,
|
|
record,
|
|
provider_response=provider_response,
|
|
stage=PricingSnapshotStage.PROVIDER_SYNC_COMPLETED.value,
|
|
)
|
|
storage_path = None
|
|
file_size_bytes = 0
|
|
if settings.STORAGE_TYPE == "local" and remote_url:
|
|
try:
|
|
date_dir = datetime.now().strftime("%Y/%m/%d")
|
|
dest_dir = os.path.join(settings.STORAGE_IMAGE_LOCAL_PATH, date_dir)
|
|
os.makedirs(dest_dir, exist_ok=True)
|
|
dest = os.path.join(dest_dir, f"{record_id}.png")
|
|
await download_image(remote_url, dest)
|
|
record.image_url = f"/generate/images/{date_dir}/{record_id}.png"
|
|
storage_path = dest
|
|
file_size_bytes = safe_file_size(dest)
|
|
except Exception as e:
|
|
logger.warning(f"Download failed, using remote URL: {e}")
|
|
record.image_url = remote_url
|
|
else:
|
|
record.image_url = remote_url
|
|
record.status = "completed"
|
|
record.generated_at = datetime.now()
|
|
if record.image_url:
|
|
await record_generation_record_generated_resource(
|
|
db,
|
|
record,
|
|
resource_url=record.image_url,
|
|
storage_path=storage_path,
|
|
file_size_bytes=file_size_bytes,
|
|
remote_url=remote_url,
|
|
generated_at=record.generated_at,
|
|
)
|
|
await sync_generation_record_media_token_snapshot(
|
|
db,
|
|
record,
|
|
provider_response=provider_response,
|
|
stage=PricingSnapshotStage.RESOURCE_DOWNLOAD_COMPLETED.value,
|
|
)
|
|
await db.commit()
|
|
logger.info(f"Image task completed: {record_id}")
|
|
else:
|
|
error_message = poll_result.get("error", "图片生成失败")
|
|
await mark_media_provider_cost_status(
|
|
db,
|
|
owner=record,
|
|
status=ProviderCostStatus.NOT_INCURRED.value,
|
|
reason=error_message,
|
|
usage_stage="provider_sync_failed",
|
|
)
|
|
await mark_generation_record_failed_and_refund_once(
|
|
db,
|
|
record=record,
|
|
error_message=error_message,
|
|
)
|
|
await db.commit()
|
|
logger.info(f"Image task failed: {record_id}")
|
|
_log_image_response(record_id, poll_result)
|
|
|
|
except Exception as e:
|
|
if provider_call_started:
|
|
uncertain = provider_call_completed or is_sync_image_provider_result_uncertain(e)
|
|
await mark_media_provider_cost_status(
|
|
db,
|
|
owner=record,
|
|
status=(
|
|
ProviderCostStatus.PROVIDER_RESULT_UNCERTAIN.value
|
|
if uncertain
|
|
else ProviderCostStatus.NOT_INCURRED.value
|
|
),
|
|
reason=str(e),
|
|
usage_stage="provider_sync_exception",
|
|
)
|
|
await mark_generation_record_failed_and_refund_once(
|
|
db,
|
|
record=record,
|
|
error_message=str(e),
|
|
)
|
|
_log_image_response(record_id, {}, str(e))
|
|
await db.commit()
|
|
logger.error(f"Image task failed: {record_id}, error: {e}")
|
|
|
|
def stop(self):
|
|
"""Signal the queue to stop."""
|
|
self.running = False
|
|
|
|
|
|
task_queue = TaskQueue()
|