397 lines
14 KiB
Python
397 lines
14 KiB
Python
from app.tasks.async_runner import run_async
|
|
from datetime import datetime, timezone, timedelta
|
|
import uuid
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.config import settings
|
|
from app.models.base import async_session
|
|
from app.models.chat_generation_task import ChatGenerationTask
|
|
from app.services.celery_download_recovery_service import (
|
|
build_download_active_payload,
|
|
ensure_aware_utc,
|
|
remove_download_active,
|
|
upsert_download_active,
|
|
)
|
|
from app.services.error_codes import extract_error_message
|
|
from app.services.generation_download_service import download_generation_result
|
|
from app.services.generation_log_service import log_task_event
|
|
from app.services.generation_refund_service import mark_chat_generation_task_failed_and_refund_once
|
|
from app.services.resource_accounting_service import record_chat_task_generated_resource
|
|
from app.tasks.celery_app import celery_app
|
|
|
|
ALLOWED_GENERATION_MODES = {"chatapi_async", "hot_opening_replicate", "shot_replicate"}
|
|
|
|
DOWNLOAD_QUEUE = "gen_result_download"
|
|
DOWNLOAD_STAGE_QUEUED = "download_queued"
|
|
DOWNLOAD_STAGE_DOWNLOADING = "downloading"
|
|
DOWNLOAD_STAGE_RETRY_WAITING = "retry_waiting"
|
|
DOWNLOAD_STAGE_DONE = "done"
|
|
DOWNLOAD_STAGE_FAILED = "download_failed"
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def _task_created_date_dir(task: ChatGenerationTask) -> str:
|
|
created_at = ensure_aware_utc(getattr(task, "created_at", None)) or _now()
|
|
return created_at.strftime("%Y/%m/%d")
|
|
|
|
|
|
def _queue_timeout_at(now: datetime | None = None) -> datetime:
|
|
now = now or _now()
|
|
return now + timedelta(seconds=int(settings.DOWNLOAD_TASK_QUEUE_TIMEOUT_SECONDS or 300))
|
|
|
|
|
|
def _lease_until(now: datetime | None = None) -> datetime:
|
|
now = now or _now()
|
|
return now + timedelta(seconds=int(settings.DOWNLOAD_TASK_LEASE_SECONDS or 600))
|
|
|
|
|
|
def _retry_at(attempt: int, now: datetime | None = None) -> datetime:
|
|
now = now or _now()
|
|
base = int(settings.DOWNLOAD_TASK_RETRY_BACKOFF_SECONDS or 30)
|
|
return now + timedelta(seconds=max(1, base * max(1, attempt)))
|
|
|
|
|
|
def _is_expired(value: datetime | None, now: datetime | None = None) -> bool:
|
|
value = ensure_aware_utc(value)
|
|
if value is None:
|
|
return True
|
|
return value <= (now or _now())
|
|
|
|
|
|
def _is_already_completed(task: ChatGenerationTask) -> bool:
|
|
if task.status == "completed" or task.pipeline_stage == DOWNLOAD_STAGE_DONE:
|
|
if task.gen_type == "image" and task.image_url:
|
|
return True
|
|
if task.gen_type == "video" and task.video_url:
|
|
return True
|
|
return False
|
|
|
|
|
|
def _build_celery_task_id(task_id: str, attempt: int | None = None, reason: str | None = None) -> str:
|
|
safe_reason = (reason or "download").replace(" ", "_")[:32]
|
|
return f"download:{task_id}:{int(attempt or 0)}:{safe_reason}:{uuid.uuid4().hex[:12]}"
|
|
|
|
|
|
async def _register_active_from_task(
|
|
task: ChatGenerationTask,
|
|
*,
|
|
check_at: datetime,
|
|
priority: int,
|
|
reason: str | None = None,
|
|
) -> None:
|
|
payload = build_download_active_payload(
|
|
record_id=task.id,
|
|
celery_task_id=task.download_celery_task_id,
|
|
stage=task.pipeline_stage or "",
|
|
attempt=task.download_attempt_count or 0,
|
|
queue=DOWNLOAD_QUEUE,
|
|
priority=priority,
|
|
enqueue_at=task.download_enqueued_at,
|
|
started_at=task.download_started_at,
|
|
lease_until=task.download_lease_until,
|
|
next_retry_at=task.download_next_retry_at,
|
|
check_at=check_at,
|
|
reason=reason,
|
|
)
|
|
await upsert_download_active(record_id=task.id, payload=payload, check_at=check_at)
|
|
|
|
|
|
async def enqueue_download_task(
|
|
db: AsyncSession,
|
|
task: ChatGenerationTask,
|
|
*,
|
|
recover: bool = False,
|
|
reason: str | None = None,
|
|
countdown: int | None = None,
|
|
) -> str | None:
|
|
"""统一投递图片/视频下载任务，并同步 DB + Redis active 注册表。"""
|
|
if not task or task.generation_mode not in ALLOWED_GENERATION_MODES:
|
|
return None
|
|
if task.status != "generating":
|
|
return None
|
|
if not task.remote_result_url:
|
|
return None
|
|
|
|
now = _now()
|
|
priority = settings.DOWNLOAD_TASK_PRIORITY_RECOVER if recover else settings.DOWNLOAD_TASK_PRIORITY_NORMAL
|
|
celery_task_id = _build_celery_task_id(
|
|
task.id,
|
|
attempt=task.download_attempt_count or task.retry_count or 0,
|
|
reason=reason or ("recover" if recover else "normal"),
|
|
)
|
|
|
|
task.pipeline_stage = DOWNLOAD_STAGE_QUEUED
|
|
task.download_celery_task_id = celery_task_id
|
|
task.download_enqueued_at = now
|
|
task.download_next_retry_at = None
|
|
if not task.download_storage_date_dir:
|
|
task.download_storage_date_dir = _task_created_date_dir(task)
|
|
|
|
await db.commit()
|
|
|
|
check_at = _queue_timeout_at(now)
|
|
await _register_active_from_task(task, check_at=check_at, priority=priority, reason=reason)
|
|
|
|
if celery_app:
|
|
download_generation_result_task.apply_async(
|
|
args=[task.id],
|
|
queue=DOWNLOAD_QUEUE,
|
|
priority=priority,
|
|
countdown=countdown,
|
|
task_id=celery_task_id,
|
|
)
|
|
return celery_task_id
|
|
|
|
|
|
async def _reload_task(db: AsyncSession, task_id: str) -> ChatGenerationTask | None:
|
|
result = await db.execute(
|
|
select(ChatGenerationTask).where(
|
|
ChatGenerationTask.id == task_id,
|
|
ChatGenerationTask.deleted_at.is_(None),
|
|
).with_for_update().limit(1)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
|
|
async def _claim_download_lease(db: AsyncSession, task: ChatGenerationTask) -> bool:
|
|
now = _now()
|
|
|
|
if not task or task.generation_mode not in ALLOWED_GENERATION_MODES:
|
|
return False
|
|
if task.status != "generating":
|
|
return False
|
|
if _is_already_completed(task):
|
|
return False
|
|
if not task.remote_result_url:
|
|
return False
|
|
|
|
stage = task.pipeline_stage
|
|
|
|
if stage == DOWNLOAD_STAGE_DOWNLOADING:
|
|
if not _is_expired(task.download_lease_until, now):
|
|
return False
|
|
await log_task_event(
|
|
task,
|
|
event_type="DOWNLOAD_STUCK_RECOVER",
|
|
message=f"downloading lease 已过期，重新抢占下载。lease_until={task.download_lease_until}",
|
|
)
|
|
elif stage == DOWNLOAD_STAGE_RETRY_WAITING:
|
|
if not _is_expired(task.download_next_retry_at, now):
|
|
return False
|
|
elif stage in (DOWNLOAD_STAGE_QUEUED, "result_ready"):
|
|
pass
|
|
else:
|
|
return False
|
|
|
|
old_stage = stage
|
|
task.pipeline_stage = DOWNLOAD_STAGE_DOWNLOADING
|
|
task.download_started_at = now
|
|
task.download_lease_until = _lease_until(now)
|
|
task.download_next_retry_at = None
|
|
task.download_attempt_count = int(task.download_attempt_count or 0) + 1
|
|
task.retry_count = task.download_attempt_count
|
|
if not task.download_storage_date_dir:
|
|
task.download_storage_date_dir = _task_created_date_dir(task)
|
|
|
|
await db.commit()
|
|
|
|
await _register_active_from_task(
|
|
task,
|
|
check_at=task.download_lease_until,
|
|
priority=settings.DOWNLOAD_TASK_PRIORITY_NORMAL,
|
|
reason="claim_download_lease",
|
|
)
|
|
|
|
await log_task_event(
|
|
task,
|
|
event_type="DOWNLOAD_START",
|
|
from_stage=old_stage,
|
|
to_stage=DOWNLOAD_STAGE_DOWNLOADING,
|
|
detail={
|
|
"attempt": task.download_attempt_count,
|
|
"lease_until": task.download_lease_until,
|
|
"download_celery_task_id": task.download_celery_task_id,
|
|
},
|
|
)
|
|
return True
|
|
|
|
|
|
async def _mark_retry_waiting(db: AsyncSession, task: ChatGenerationTask, exc: Exception) -> datetime:
|
|
now = _now()
|
|
attempt = int(task.download_attempt_count or task.retry_count or 0)
|
|
next_retry_at = _retry_at(attempt, now)
|
|
error_message = extract_error_message(exc, "下载") if callable(extract_error_message) else str(exc)
|
|
|
|
retry_celery_task_id = _build_celery_task_id(task.id, attempt=attempt, reason="retry_waiting")
|
|
|
|
task.pipeline_stage = DOWNLOAD_STAGE_RETRY_WAITING
|
|
task.download_celery_task_id = retry_celery_task_id
|
|
task.download_next_retry_at = next_retry_at
|
|
task.download_lease_until = None
|
|
task.download_last_error = error_message
|
|
task.retry_count = attempt
|
|
await db.commit()
|
|
|
|
await _register_active_from_task(
|
|
task,
|
|
check_at=next_retry_at,
|
|
priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
|
|
reason="download_retry_waiting",
|
|
)
|
|
|
|
await log_task_event(
|
|
task,
|
|
event_type="DOWNLOAD_RETRY_WAITING",
|
|
message=error_message,
|
|
to_stage=DOWNLOAD_STAGE_RETRY_WAITING,
|
|
detail={
|
|
"attempt": attempt,
|
|
"next_retry_at": next_retry_at,
|
|
"download_celery_task_id": retry_celery_task_id,
|
|
},
|
|
)
|
|
return next_retry_at
|
|
|
|
|
|
def _should_final_fail(task: ChatGenerationTask) -> bool:
|
|
return int(task.download_attempt_count or task.retry_count or 0) >= int(settings.DOWNLOAD_TASK_MAX_ATTEMPTS or 3)
|
|
|
|
|
|
async def _run(task_id: str):
|
|
async with async_session() as db:
|
|
result = await db.execute(
|
|
select(ChatGenerationTask).where(
|
|
ChatGenerationTask.id == task_id,
|
|
ChatGenerationTask.deleted_at.is_(None),
|
|
).with_for_update().limit(1)
|
|
)
|
|
task = result.scalar_one_or_none()
|
|
if not task:
|
|
return
|
|
|
|
claimed = await _claim_download_lease(db, task)
|
|
if not claimed:
|
|
return
|
|
|
|
try:
|
|
downloaded = await download_generation_result(task)
|
|
|
|
task = await _reload_task(db, task_id)
|
|
if not task:
|
|
return
|
|
|
|
if task.gen_type == "image":
|
|
task.image_url = downloaded.url
|
|
else:
|
|
task.video_url = downloaded.url
|
|
task.video_cover_url = downloaded.cover_url
|
|
|
|
task.status = "completed"
|
|
task.pipeline_stage = DOWNLOAD_STAGE_DONE
|
|
task.generated_at = _now()
|
|
task.retry_count = 0
|
|
task.download_lease_until = None
|
|
task.download_next_retry_at = None
|
|
task.download_last_error = None
|
|
|
|
await record_chat_task_generated_resource(
|
|
db,
|
|
task,
|
|
resource_url=downloaded.url,
|
|
storage_path=downloaded.storage_path,
|
|
file_size_bytes=downloaded.file_size_bytes,
|
|
remote_url=task.remote_result_url,
|
|
generated_at=task.generated_at,
|
|
)
|
|
|
|
await db.commit()
|
|
await remove_download_active(task.id)
|
|
|
|
from app.services.generation_module_hook_service import notify_chat_generation_task_finished
|
|
await notify_chat_generation_task_finished(db, task)
|
|
await db.commit()
|
|
|
|
await log_task_event(
|
|
task,
|
|
event_type="DOWNLOAD_SUCCESS",
|
|
to_status="completed",
|
|
to_stage=DOWNLOAD_STAGE_DONE,
|
|
detail={
|
|
"resource_url": downloaded.url,
|
|
"video_cover_url": downloaded.cover_url,
|
|
"file_size_bytes": downloaded.file_size_bytes,
|
|
"download_attempt_count": task.download_attempt_count,
|
|
},
|
|
)
|
|
|
|
except Exception as exc:
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
|
|
task = await _reload_task(db, task_id)
|
|
if not task:
|
|
return
|
|
|
|
if _should_final_fail(task):
|
|
error_message = extract_error_message(exc, "下载") if callable(extract_error_message) else str(exc)
|
|
await mark_chat_generation_task_failed_and_refund_once(
|
|
db,
|
|
task=task,
|
|
error_message=error_message,
|
|
pipeline_stage=DOWNLOAD_STAGE_FAILED,
|
|
)
|
|
task.download_last_error = error_message
|
|
task.download_lease_until = None
|
|
task.download_next_retry_at = None
|
|
await db.commit()
|
|
|
|
await remove_download_active(task.id)
|
|
|
|
from app.services.generation_module_hook_service import notify_chat_generation_task_finished
|
|
await notify_chat_generation_task_finished(db, task)
|
|
await db.commit()
|
|
|
|
await log_task_event(
|
|
task,
|
|
event_type="DOWNLOAD_FAILED",
|
|
message=task.error_message,
|
|
detail={
|
|
"download_attempt_count": task.download_attempt_count,
|
|
"max_attempts": settings.DOWNLOAD_TASK_MAX_ATTEMPTS,
|
|
},
|
|
)
|
|
else:
|
|
next_retry_at = await _mark_retry_waiting(db, task, exc)
|
|
|
|
if celery_app:
|
|
delay_seconds = max(1, int((next_retry_at - _now()).total_seconds()))
|
|
download_generation_result_task.apply_async(
|
|
args=[task.id],
|
|
queue=DOWNLOAD_QUEUE,
|
|
priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
|
|
countdown=delay_seconds,
|
|
task_id=task.download_celery_task_id,
|
|
)
|
|
|
|
|
|
if celery_app:
|
|
@celery_app.task(name="generation.download_generation_result_task", bind=True, max_retries=3, default_retry_delay=30)
|
|
def download_generation_result_task(self, task_id: str):
|
|
return run_async(_run(task_id))
|
|
else:
|
|
class _DisabledTask:
|
|
def delay(self, *args, **kwargs):
|
|
raise RuntimeError("Celery is disabled")
|
|
|
|
def apply_async(self, *args, **kwargs):
|
|
raise RuntimeError("Celery is disabled")
|
|
|
|
download_generation_result_task = _DisabledTask()
|