198 lines
7.5 KiB
Python
198 lines
7.5 KiB
Python
"""Celery tasks for async video generation polling."""
|
||
|
||
import logging
|
||
|
||
from celery import shared_task
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@shared_task(ignore_result=True)
|
||
def poll_video_task(record_id):
|
||
"""Poll Volcano API once for a video generation task.
|
||
|
||
一次性任务：查一次 API，更新 DB，结束。
|
||
由 recover_stuck_tasks（beat 每30秒调度）统一驱动，不再自己 retry。
|
||
"""
|
||
from django.utils import timezone
|
||
from apps.generation.models import GenerationRecord
|
||
from utils.airdrama_client import query_task, map_status
|
||
|
||
try:
|
||
record = GenerationRecord.objects.get(pk=record_id)
|
||
except GenerationRecord.DoesNotExist:
|
||
logger.warning('poll_video_task: record %s not found', record_id)
|
||
return
|
||
|
||
if record.status not in ('queued', 'processing'):
|
||
return
|
||
|
||
ark_task_id = record.ark_task_id
|
||
if not ark_task_id:
|
||
logger.warning('poll_video_task: record %s has no ark_task_id', record_id)
|
||
return
|
||
|
||
# Poll Volcano API
|
||
try:
|
||
ark_resp = query_task(ark_task_id)
|
||
new_status = map_status(ark_resp.get('status', ''))
|
||
except Exception:
|
||
logger.exception('poll_video_task: API query failed for record=%s ark=%s', record_id, ark_task_id)
|
||
return
|
||
|
||
if new_status in ('queued', 'processing'):
|
||
record.status = new_status
|
||
record.save(update_fields=['status', 'updated_at'])
|
||
return
|
||
|
||
# Terminal state reached — process result
|
||
record.status = new_status
|
||
|
||
returned_seed = ark_resp.get('seed')
|
||
if returned_seed is not None:
|
||
record.seed = returned_seed
|
||
|
||
if new_status == 'completed':
|
||
_handle_completed(record, ark_resp)
|
||
elif new_status == 'failed':
|
||
_handle_failed(record, ark_resp)
|
||
|
||
record.completed_at = timezone.now()
|
||
record.save(update_fields=[
|
||
'status', 'result_url', 'thumbnail_url', 'error_message', 'raw_error',
|
||
'seed', 'completed_at',
|
||
])
|
||
|
||
logger.info(
|
||
'poll_video_task: record=%s ark=%s final_status=%s',
|
||
record_id, ark_task_id, new_status,
|
||
)
|
||
|
||
|
||
def _handle_completed(record, ark_resp):
|
||
"""Process a completed task: persist video to TOS, extract thumbnail, settle payment."""
|
||
import os
|
||
from utils.airdrama_client import extract_video_url
|
||
|
||
video_url = extract_video_url(ark_resp)
|
||
if video_url:
|
||
# Download once to temp file, reuse for TOS upload + thumbnail extraction
|
||
tmp_path = None
|
||
try:
|
||
from utils.media_utils import download_to_temp, extract_video_info_from_file
|
||
from utils.tos_client import upload_from_file_path, upload_file
|
||
|
||
tmp_path = download_to_temp(video_url, '.mp4')
|
||
|
||
# Upload video to TOS from file (streaming, no full memory load)
|
||
record.result_url = upload_from_file_path(tmp_path, folder='results', content_type='video/mp4')
|
||
|
||
# Extract thumbnail from the same local file (no second download)
|
||
thumb_file, _ = extract_video_info_from_file(tmp_path)
|
||
if thumb_file:
|
||
record.thumbnail_url = upload_file(thumb_file, folder='thumbnails')
|
||
except Exception:
|
||
logger.exception('poll_video_task: failed to persist video / extract thumbnail')
|
||
if not record.result_url:
|
||
record.result_url = video_url
|
||
record.error_message = '视频保存失败，临时链接将在24小时后过期，请联系管理员'
|
||
finally:
|
||
if tmp_path and os.path.exists(tmp_path):
|
||
os.unlink(tmp_path)
|
||
|
||
# 结算：按实际 tokens 扣费
|
||
usage = ark_resp.get('usage', {})
|
||
total_tokens = usage.get('total_tokens', 0) if isinstance(usage, dict) else 0
|
||
if total_tokens > 0:
|
||
from apps.generation.views import _settle_payment
|
||
_settle_payment(record, total_tokens)
|
||
else:
|
||
from apps.generation.views import _release_freeze
|
||
_release_freeze(record)
|
||
|
||
|
||
@shared_task(ignore_result=True)
|
||
def recover_stuck_tasks():
|
||
"""每30秒扫一次所有进行中的任务，统一派发轮询。
|
||
|
||
poll_video_task 是一次性任务，不再自己 retry，由这里统一驱动。
|
||
"""
|
||
from apps.generation.models import GenerationRecord
|
||
|
||
active_records = GenerationRecord.objects.filter(
|
||
status__in=('queued', 'processing'),
|
||
ark_task_id__isnull=False,
|
||
).exclude(ark_task_id='').values_list('id', flat=True)
|
||
|
||
count = 0
|
||
for record_id in active_records:
|
||
try:
|
||
poll_video_task.delay(record_id)
|
||
count += 1
|
||
except Exception:
|
||
logger.error('recover_stuck_tasks: failed to dispatch record=%s', record_id)
|
||
|
||
if count:
|
||
logger.info('recover_stuck_tasks: dispatched %d active tasks', count)
|
||
|
||
|
||
def _handle_failed(record, ark_resp):
|
||
"""Process a failed task: record error and release frozen amount."""
|
||
from utils.airdrama_client import ERROR_MESSAGES
|
||
|
||
error = ark_resp.get('error', {})
|
||
code = error.get('code', '') if isinstance(error, dict) else ''
|
||
raw_msg = error.get('message', '') if isinstance(error, dict) else str(error)
|
||
record.error_message = ERROR_MESSAGES.get(code, raw_msg)
|
||
record.raw_error = f'{code}: {raw_msg}' if code else raw_msg
|
||
|
||
usage = ark_resp.get('usage', {})
|
||
total_tokens = usage.get('total_tokens', 0) if isinstance(usage, dict) else 0
|
||
if total_tokens > 0:
|
||
from apps.generation.views import _settle_payment
|
||
_settle_payment(record, total_tokens)
|
||
else:
|
||
from apps.generation.views import _release_freeze
|
||
_release_freeze(record)
|
||
|
||
|
||
@shared_task(ignore_result=True)
|
||
def process_asset_media(asset_id):
|
||
"""Extract thumbnail + duration for video/audio assets asynchronously."""
|
||
from apps.generation.models import Asset
|
||
try:
|
||
asset = Asset.objects.select_related('group').get(pk=asset_id)
|
||
except Asset.DoesNotExist:
|
||
logger.warning('process_asset_media: asset %s not found', asset_id)
|
||
return
|
||
|
||
from utils.media_utils import extract_video_info, get_audio_duration
|
||
from utils.tos_client import upload_file
|
||
|
||
if asset.asset_type == 'Video':
|
||
thumb_file, dur = extract_video_info(asset.url)
|
||
if thumb_file:
|
||
try:
|
||
asset.thumbnail_url = upload_file(thumb_file, folder='thumbnails')
|
||
except Exception:
|
||
logger.exception('process_asset_media: thumbnail upload failed for asset %s', asset_id)
|
||
asset.duration = dur if dur > 0 else None # None = ffprobe failed, frontend skips duration check
|
||
asset.save(update_fields=['thumbnail_url', 'duration'])
|
||
# Atomic update: only set group thumbnail if still empty (concurrent-safe)
|
||
from apps.generation.models import AssetGroup
|
||
from django.db import transaction
|
||
try:
|
||
with transaction.atomic():
|
||
group = AssetGroup.objects.select_for_update().get(pk=asset.group_id)
|
||
if not group.thumbnail_url and asset.thumbnail_url:
|
||
group.thumbnail_url = asset.thumbnail_url
|
||
group.save(update_fields=['thumbnail_url'])
|
||
except AssetGroup.DoesNotExist:
|
||
logger.warning('process_asset_media: group %s deleted, skipping thumbnail update', asset.group_id)
|
||
elif asset.asset_type == 'Audio':
|
||
dur = get_audio_duration(asset.url)
|
||
asset.duration = dur if dur > 0 else None
|
||
asset.save(update_fields=['duration'])
|
||
|
||
logger.info('process_asset_media: asset %s done (type=%s, dur=%s)', asset_id, asset.asset_type, asset.duration)
|