963 lines
33 KiB
Python
963 lines
33 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
import asyncio
|
|
import errno
|
|
import math
|
|
import uuid
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.config import settings
|
|
from app.enums.celery_queue import CeleryQueue, CeleryTaskName
|
|
from app.enums.celery_runtime import CeleryRuntimeDomain
|
|
from app.enums.generation_status import GenerationRecordPipelineStage
|
|
from app.enums.generation_task import (
|
|
ALLOWED_GENERATION_MODES,
|
|
ChatGenerationPipelineStage,
|
|
ChatGenerationTaskEventType,
|
|
GenerationMode,
|
|
GenerationOwnerType,
|
|
GenerationType,
|
|
)
|
|
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,
|
|
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,
|
|
download_video_upscale_source,
|
|
)
|
|
from app.services.generation.log_service import log_task_event
|
|
from app.services.generation.pipeline.db_lock_service import DatabaseRowLockBusy
|
|
from app.services.generation.pipeline.lifecycle_service import (
|
|
mark_owner_failed_and_refund_once,
|
|
notify_owner_finished,
|
|
)
|
|
from app.services.generation.pipeline.owner_service import (
|
|
GenerationOwner,
|
|
is_attempt_current,
|
|
load_generation_owner,
|
|
load_generation_owner_for_update_retry,
|
|
normalize_owner_type,
|
|
owner_is_completed,
|
|
owner_is_generating,
|
|
owner_type_of,
|
|
redis_owner_item_id,
|
|
set_owner_completed,
|
|
renew_generation_owner_claim_lease,
|
|
)
|
|
from app.services.media_token_usage_snapshot_service import (
|
|
sync_chat_generation_task_media_token_snapshot,
|
|
sync_generation_record_media_token_snapshot,
|
|
)
|
|
from app.services.resource_accounting_service import (
|
|
record_chat_task_generated_resource,
|
|
record_generation_record_generated_resource,
|
|
)
|
|
from app.services.redis_registry_service import (
|
|
RedisExecutionLockError,
|
|
ensure_aware_utc,
|
|
)
|
|
from app.services.video_upscale.media_service import probe_video
|
|
from app.services.celery_runtime.runtime_service import CeleryRuntimeLease, RuntimeIdentity
|
|
from app.tasks.async_runner import run_async
|
|
from app.tasks.celery_app import celery_app
|
|
|
|
logger = logging.getLogger("video_gen")
|
|
|
|
DOWNLOAD_QUEUE = CeleryQueue.GEN_RESULT_DOWNLOAD.value
|
|
|
|
DOWNLOAD_STAGE_QUEUED = ChatGenerationPipelineStage.DOWNLOAD_QUEUED.value
|
|
DOWNLOAD_STAGE_DOWNLOADING = ChatGenerationPipelineStage.DOWNLOADING.value
|
|
DOWNLOAD_STAGE_RETRY_WAITING = ChatGenerationPipelineStage.RETRY_WAITING.value
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def _stage(owner: GenerationOwner, chat_stage: ChatGenerationPipelineStage) -> str:
|
|
if isinstance(owner, ChatGenerationTask):
|
|
return chat_stage.value
|
|
try:
|
|
return GenerationRecordPipelineStage(chat_stage.value).value
|
|
except ValueError:
|
|
return chat_stage.value
|
|
|
|
|
|
def _registry_id(owner: GenerationOwner) -> str:
|
|
return redis_owner_item_id(
|
|
owner_type_of(owner), owner.id, int(owner.generation_attempt_no or 1)
|
|
)
|
|
|
|
|
|
def _lock_key(owner_type: str, owner_id: str, attempt_no: int) -> str:
|
|
return (
|
|
f"{settings.GENERATION_DOWNLOAD_LOCK_KEY_PREFIX}:"
|
|
f"{owner_type}:{owner_id}:attempt:{attempt_no}"
|
|
)
|
|
|
|
|
|
def _queue_timeout_at() -> datetime:
|
|
return _now() + timedelta(
|
|
seconds=int(settings.DOWNLOAD_TASK_QUEUE_TIMEOUT_SECONDS or 300)
|
|
)
|
|
|
|
|
|
def _lease_until() -> datetime:
|
|
return _now() + timedelta(
|
|
seconds=int(settings.DOWNLOAD_TASK_LEASE_SECONDS or 600)
|
|
)
|
|
|
|
|
|
def _retry_at(attempt: int) -> datetime:
|
|
return _now() + timedelta(
|
|
seconds=max(
|
|
1,
|
|
int(settings.DOWNLOAD_TASK_RETRY_BACKOFF_SECONDS or 30)
|
|
* max(1, attempt),
|
|
)
|
|
)
|
|
|
|
|
|
def _countdown_until(value: datetime | None) -> int:
|
|
target = ensure_aware_utc(value)
|
|
if target is None:
|
|
return 1
|
|
return max(
|
|
1,
|
|
math.ceil((target - _now()).total_seconds())
|
|
+ int(settings.DOWNLOAD_RETRY_COUNTDOWN_EXTRA_SECONDS or 1),
|
|
)
|
|
|
|
|
|
def _build_celery_task_id(owner: GenerationOwner, *, reason: str) -> str:
|
|
return (
|
|
f"download:{owner_type_of(owner)}:{owner.id}:"
|
|
f"attempt:{int(owner.generation_attempt_no or 1)}:"
|
|
f"{reason[:24]}:{uuid.uuid4().hex[:10]}"
|
|
)
|
|
|
|
|
|
def _is_non_retryable_download_error(exc: Exception) -> bool:
|
|
if not bool(getattr(settings, "DOWNLOAD_NON_RETRYABLE_LOCAL_ERRORS", True)):
|
|
return False
|
|
if isinstance(exc, PermissionError):
|
|
return True
|
|
if isinstance(exc, OSError) and getattr(exc, "errno", None) in {
|
|
errno.EACCES,
|
|
errno.EPERM,
|
|
errno.ENOSPC,
|
|
errno.EROFS,
|
|
errno.ENAMETOOLONG,
|
|
}:
|
|
return True
|
|
message = str(exc).lower()
|
|
return any(
|
|
value in message
|
|
for value in (
|
|
"permission denied",
|
|
"no space left on device",
|
|
"read-only file system",
|
|
"file name too long",
|
|
)
|
|
)
|
|
|
|
|
|
async def _register_active(
|
|
owner: GenerationOwner,
|
|
*,
|
|
check_at: datetime,
|
|
priority: int,
|
|
reason: str,
|
|
) -> None:
|
|
item_id = _registry_id(owner)
|
|
payload = build_download_active_payload(
|
|
record_id=item_id,
|
|
celery_task_id=owner.download_celery_task_id,
|
|
stage=owner.pipeline_stage or "",
|
|
attempt=owner.download_attempt_count or 0,
|
|
queue=DOWNLOAD_QUEUE,
|
|
priority=priority,
|
|
enqueue_at=owner.download_enqueued_at,
|
|
started_at=owner.download_started_at,
|
|
lease_until=owner.download_lease_until,
|
|
next_retry_at=owner.download_next_retry_at,
|
|
check_at=check_at,
|
|
reason=reason,
|
|
)
|
|
payload.update(
|
|
{
|
|
"owner_type": owner_type_of(owner),
|
|
"owner_id": owner.id,
|
|
"generation_attempt_no": int(owner.generation_attempt_no or 1),
|
|
"download_claim_token_suffix": (
|
|
str(owner.download_claim_token)[-8:]
|
|
if owner.download_claim_token
|
|
else None
|
|
),
|
|
}
|
|
)
|
|
await upsert_download_active(
|
|
record_id=item_id, payload=payload, check_at=check_at
|
|
)
|
|
|
|
|
|
async def _remove_active(owner: GenerationOwner) -> None:
|
|
await remove_download_active(_registry_id(owner))
|
|
|
|
|
|
async def _reload_owner_after_commit(
|
|
db: AsyncSession,
|
|
*,
|
|
owner_type: str,
|
|
owner_id: str,
|
|
attempt_no: int,
|
|
) -> GenerationOwner | None:
|
|
owner = await load_generation_owner(
|
|
db,
|
|
owner_type=owner_type,
|
|
owner_id=owner_id,
|
|
for_update=False,
|
|
)
|
|
if owner is None or not is_attempt_current(owner, attempt_no):
|
|
return None
|
|
return owner
|
|
|
|
|
|
async def _apply(
|
|
owner: GenerationOwner,
|
|
*,
|
|
priority: int,
|
|
countdown: int | None,
|
|
reason: str,
|
|
) -> None:
|
|
download_generation_result_task.apply_async(
|
|
args=[owner.id],
|
|
kwargs={
|
|
"owner_type": owner_type_of(owner),
|
|
"generation_attempt_no": int(owner.generation_attempt_no or 1),
|
|
},
|
|
queue=DOWNLOAD_QUEUE,
|
|
priority=priority,
|
|
countdown=countdown,
|
|
task_id=owner.download_celery_task_id,
|
|
)
|
|
await log_task_event(
|
|
owner,
|
|
event_type=ChatGenerationTaskEventType.DOWNLOAD_ENQUEUE.value,
|
|
to_stage=owner.pipeline_stage,
|
|
detail={
|
|
"reason": reason,
|
|
"celery_task_id": owner.download_celery_task_id,
|
|
},
|
|
)
|
|
|
|
|
|
def _recover_enqueue_due(owner: GenerationOwner, now: datetime) -> bool:
|
|
result_ready = _stage(owner, ChatGenerationPipelineStage.RESULT_READY)
|
|
queued = _stage(owner, ChatGenerationPipelineStage.DOWNLOAD_QUEUED)
|
|
downloading = _stage(owner, ChatGenerationPipelineStage.DOWNLOADING)
|
|
retry_waiting = _stage(owner, ChatGenerationPipelineStage.RETRY_WAITING)
|
|
|
|
if owner.pipeline_stage == result_ready:
|
|
return True
|
|
if owner.pipeline_stage == queued:
|
|
enqueued_at = ensure_aware_utc(owner.download_enqueued_at)
|
|
if enqueued_at is None:
|
|
return True
|
|
return enqueued_at + timedelta(
|
|
seconds=int(settings.DOWNLOAD_TASK_QUEUE_TIMEOUT_SECONDS or 300)
|
|
) <= now
|
|
if owner.pipeline_stage == downloading:
|
|
lease_until = ensure_aware_utc(owner.download_lease_until)
|
|
return lease_until is None or lease_until <= now
|
|
if owner.pipeline_stage == retry_waiting:
|
|
next_retry_at = ensure_aware_utc(owner.download_next_retry_at)
|
|
return next_retry_at is None or next_retry_at <= now
|
|
return False
|
|
|
|
|
|
async def enqueue_download_task(
|
|
db: AsyncSession,
|
|
owner: GenerationOwner,
|
|
*,
|
|
recover: bool = False,
|
|
reason: str | None = None,
|
|
countdown: int | None = None,
|
|
) -> str | None:
|
|
reason = reason or ("recover" if recover else "normal")
|
|
if (
|
|
isinstance(owner, ChatGenerationTask)
|
|
and owner.generation_mode not in ALLOWED_GENERATION_MODES
|
|
):
|
|
return None
|
|
if not owner_is_generating(owner) or not owner.remote_result_url:
|
|
return None
|
|
if recover and not _recover_enqueue_due(owner, _now()):
|
|
return owner.download_celery_task_id
|
|
|
|
owner.pipeline_stage = _stage(
|
|
owner, ChatGenerationPipelineStage.DOWNLOAD_QUEUED
|
|
)
|
|
owner.download_celery_task_id = _build_celery_task_id(owner, reason=reason)
|
|
owner.download_enqueued_at = _now()
|
|
owner.download_started_at = None
|
|
owner.download_claim_token = None
|
|
owner.download_lease_until = None
|
|
owner.download_next_retry_at = None
|
|
if not owner.download_storage_date_dir:
|
|
created_at = ensure_aware_utc(owner.created_at) or _now()
|
|
owner.download_storage_date_dir = created_at.strftime("%Y/%m/%d")
|
|
owner_type_snapshot = owner_type_of(owner)
|
|
owner_id_snapshot = str(owner.id)
|
|
attempt_snapshot = int(owner.generation_attempt_no or 1)
|
|
celery_task_id_snapshot = str(owner.download_celery_task_id or "") or None
|
|
await db.commit()
|
|
owner = await load_generation_owner(
|
|
db,
|
|
owner_type=owner_type_snapshot,
|
|
owner_id=owner_id_snapshot,
|
|
for_update=False,
|
|
)
|
|
if owner is None or not is_attempt_current(owner, attempt_snapshot):
|
|
return None
|
|
|
|
priority = int(
|
|
settings.DOWNLOAD_TASK_PRIORITY_RECOVER
|
|
if recover
|
|
else settings.DOWNLOAD_TASK_PRIORITY_NORMAL
|
|
)
|
|
await _register_active(
|
|
owner,
|
|
check_at=_queue_timeout_at(),
|
|
priority=priority,
|
|
reason=reason,
|
|
)
|
|
try:
|
|
await _apply(
|
|
owner,
|
|
priority=priority,
|
|
countdown=countdown,
|
|
reason=reason,
|
|
)
|
|
except Exception as exc:
|
|
await log_task_event(
|
|
owner,
|
|
event_type=ChatGenerationTaskEventType.DOWNLOAD_ENQUEUE_FAILED.value,
|
|
message="下载任务投递失败，等待恢复扫描",
|
|
detail={"reason": reason, "error": str(exc)},
|
|
)
|
|
return None
|
|
return str(owner.download_celery_task_id or celery_task_id_snapshot or "") or None
|
|
|
|
|
|
async def _claim(
|
|
db: AsyncSession,
|
|
owner: GenerationOwner,
|
|
*,
|
|
claim_token: str,
|
|
) -> GenerationOwner | None:
|
|
if not owner_is_generating(owner) or owner_is_completed(owner):
|
|
return None
|
|
allowed = {
|
|
_stage(owner, ChatGenerationPipelineStage.RESULT_READY),
|
|
_stage(owner, ChatGenerationPipelineStage.DOWNLOAD_QUEUED),
|
|
_stage(owner, ChatGenerationPipelineStage.DOWNLOADING),
|
|
_stage(owner, ChatGenerationPipelineStage.RETRY_WAITING),
|
|
}
|
|
if owner.pipeline_stage not in allowed:
|
|
return None
|
|
|
|
now = _now()
|
|
# Redis execution lock is authoritative. A database lease left by a
|
|
# crashed worker is audit/recovery metadata and must not block the worker
|
|
# that owns the current Redis download lock.
|
|
next_retry = ensure_aware_utc(owner.download_next_retry_at)
|
|
if (
|
|
owner.pipeline_stage
|
|
== _stage(owner, ChatGenerationPipelineStage.RETRY_WAITING)
|
|
and next_retry
|
|
and next_retry > now
|
|
):
|
|
return None
|
|
|
|
owner.pipeline_stage = _stage(
|
|
owner, ChatGenerationPipelineStage.DOWNLOADING
|
|
)
|
|
owner.download_claim_token = claim_token
|
|
owner.download_started_at = now
|
|
owner.download_lease_until = _lease_until()
|
|
owner.download_attempt_count = int(owner.download_attempt_count or 0) + 1
|
|
owner.download_last_error = None
|
|
owner_type_snapshot = owner_type_of(owner)
|
|
owner_id_snapshot = str(owner.id)
|
|
attempt_snapshot = int(owner.generation_attempt_no or 1)
|
|
await db.commit()
|
|
owner = await _reload_owner_after_commit(
|
|
db,
|
|
owner_type=owner_type_snapshot,
|
|
owner_id=owner_id_snapshot,
|
|
attempt_no=attempt_snapshot,
|
|
)
|
|
if owner is None:
|
|
return None
|
|
await _register_active(
|
|
owner,
|
|
check_at=owner.download_lease_until,
|
|
priority=int(settings.DOWNLOAD_TASK_PRIORITY_NORMAL),
|
|
reason="downloading_lease",
|
|
)
|
|
await log_task_event(
|
|
owner,
|
|
event_type=ChatGenerationTaskEventType.DOWNLOAD_START.value,
|
|
to_stage=owner.pipeline_stage,
|
|
detail={
|
|
"download_attempt_count": owner.download_attempt_count,
|
|
"claim_token_suffix": claim_token[-8:],
|
|
},
|
|
)
|
|
return owner
|
|
|
|
|
|
async def _sync_snapshot(db: AsyncSession, owner: GenerationOwner) -> None:
|
|
if isinstance(owner, ChatGenerationTask):
|
|
await sync_chat_generation_task_media_token_snapshot(db, owner)
|
|
else:
|
|
await sync_generation_record_media_token_snapshot(
|
|
db, owner, provider_response=owner.provider_response_json
|
|
)
|
|
|
|
|
|
async def _record_resource(
|
|
db: AsyncSession, owner: GenerationOwner, downloaded
|
|
) -> None:
|
|
kwargs = dict(
|
|
resource_url=downloaded.url,
|
|
storage_path=downloaded.storage_path,
|
|
file_size_bytes=downloaded.file_size_bytes,
|
|
remote_url=owner.remote_result_url,
|
|
generated_at=owner.generated_at,
|
|
)
|
|
if isinstance(owner, ChatGenerationTask):
|
|
await record_chat_task_generated_resource(db, owner, **kwargs)
|
|
else:
|
|
await record_generation_record_generated_resource(db, owner, **kwargs)
|
|
|
|
|
|
async def _mark_failed(
|
|
db: AsyncSession,
|
|
owner: GenerationOwner,
|
|
exc: Exception,
|
|
*,
|
|
non_retryable: bool,
|
|
) -> None:
|
|
error_message = (
|
|
extract_error_message(exc, "下载")
|
|
if callable(extract_error_message)
|
|
else str(exc)
|
|
)
|
|
image_child = bool(
|
|
isinstance(owner, ChatGenerationTask)
|
|
and owner.gen_type == GenerationType.IMAGE.value
|
|
and owner.generation_mode == GenerationMode.CHATAPI_CHILD.value
|
|
)
|
|
if image_child:
|
|
owner.status = "failed"
|
|
owner.pipeline_stage = _stage(
|
|
owner, ChatGenerationPipelineStage.DOWNLOAD_FAILED
|
|
)
|
|
owner.error_message = error_message
|
|
else:
|
|
await mark_owner_failed_and_refund_once(
|
|
db,
|
|
owner,
|
|
error_message=error_message,
|
|
pipeline_stage=_stage(
|
|
owner, ChatGenerationPipelineStage.DOWNLOAD_FAILED
|
|
),
|
|
)
|
|
owner.download_claim_token = None
|
|
owner.download_last_error = error_message
|
|
owner.download_lease_until = None
|
|
owner.download_next_retry_at = None
|
|
owner_type_snapshot = owner_type_of(owner)
|
|
owner_id_snapshot = str(owner.id)
|
|
attempt_snapshot = int(owner.generation_attempt_no or 1)
|
|
generation_mode_snapshot = str(getattr(owner, "generation_mode", "") or "") or None
|
|
await db.commit()
|
|
owner = await _reload_owner_after_commit(
|
|
db,
|
|
owner_type=owner_type_snapshot,
|
|
owner_id=owner_id_snapshot,
|
|
attempt_no=attempt_snapshot,
|
|
)
|
|
if owner is not None:
|
|
await notify_owner_finished(db, owner)
|
|
await db.commit()
|
|
owner = await _reload_owner_after_commit(
|
|
db,
|
|
owner_type=owner_type_snapshot,
|
|
owner_id=owner_id_snapshot,
|
|
attempt_no=attempt_snapshot,
|
|
)
|
|
if owner is not None:
|
|
await _remove_active(owner)
|
|
await log_task_event(
|
|
owner_type=owner_type_snapshot,
|
|
owner_id=owner_id_snapshot,
|
|
generation_attempt_no=attempt_snapshot,
|
|
generation_mode=generation_mode_snapshot,
|
|
event_type=(
|
|
ChatGenerationTaskEventType.DOWNLOAD_FAILED_NON_RETRYABLE.value
|
|
if non_retryable
|
|
else ChatGenerationTaskEventType.DOWNLOAD_FAILED.value
|
|
),
|
|
message=error_message,
|
|
)
|
|
|
|
|
|
async def _schedule_retry(
|
|
db: AsyncSession, owner: GenerationOwner, exc: Exception
|
|
) -> None:
|
|
error_message = (
|
|
extract_error_message(exc, "下载")
|
|
if callable(extract_error_message)
|
|
else str(exc)
|
|
)
|
|
owner.pipeline_stage = _stage(
|
|
owner, ChatGenerationPipelineStage.RETRY_WAITING
|
|
)
|
|
owner.download_next_retry_at = _retry_at(
|
|
int(owner.download_attempt_count or 1)
|
|
)
|
|
owner.download_claim_token = None
|
|
owner.download_lease_until = None
|
|
owner.download_last_error = error_message
|
|
owner.download_celery_task_id = _build_celery_task_id(
|
|
owner, reason="retry"
|
|
)
|
|
owner_type_snapshot = owner_type_of(owner)
|
|
owner_id_snapshot = str(owner.id)
|
|
attempt_snapshot = int(owner.generation_attempt_no or 1)
|
|
await db.commit()
|
|
owner = await _reload_owner_after_commit(
|
|
db,
|
|
owner_type=owner_type_snapshot,
|
|
owner_id=owner_id_snapshot,
|
|
attempt_no=attempt_snapshot,
|
|
)
|
|
if owner is None:
|
|
return
|
|
await _register_active(
|
|
owner,
|
|
check_at=owner.download_next_retry_at,
|
|
priority=int(settings.DOWNLOAD_TASK_PRIORITY_RECOVER),
|
|
reason="download_retry_waiting",
|
|
)
|
|
await log_task_event(
|
|
owner,
|
|
event_type=ChatGenerationTaskEventType.DOWNLOAD_RETRY_WAITING.value,
|
|
message=error_message,
|
|
detail={"next_retry_at": owner.download_next_retry_at},
|
|
)
|
|
await _apply(
|
|
owner,
|
|
priority=int(settings.DOWNLOAD_TASK_PRIORITY_RECOVER),
|
|
countdown=_countdown_until(owner.download_next_retry_at),
|
|
reason="download_exception_retry",
|
|
)
|
|
|
|
|
|
async def _resolve_attempt(
|
|
task_id: str, *, owner_type: str, message_attempt: int | None
|
|
) -> int | None:
|
|
async with async_session() as db:
|
|
owner = await load_generation_owner(
|
|
db, owner_type=owner_type, owner_id=task_id, for_update=False
|
|
)
|
|
if not owner:
|
|
await remove_download_active(
|
|
redis_owner_item_id(owner_type, task_id, message_attempt)
|
|
)
|
|
return None
|
|
if not is_attempt_current(owner, message_attempt):
|
|
await remove_download_active(
|
|
redis_owner_item_id(owner_type, task_id, message_attempt)
|
|
)
|
|
await log_task_event(
|
|
owner,
|
|
event_type=ChatGenerationTaskEventType.STALE_ATTEMPT_MESSAGE_SKIPPED.value,
|
|
message="下载消息属于旧生成轮次，已跳过",
|
|
)
|
|
return None
|
|
return int(owner.generation_attempt_no or 1)
|
|
|
|
|
|
async def _restore_after_lock_error(
|
|
*,
|
|
owner_type: str,
|
|
owner_id: str,
|
|
attempt_no: int,
|
|
claim_token: str,
|
|
) -> None:
|
|
async with async_session() as db:
|
|
owner = await load_generation_owner(
|
|
db,
|
|
owner_type=owner_type,
|
|
owner_id=owner_id,
|
|
for_update=True,
|
|
)
|
|
if not owner or not is_attempt_current(owner, attempt_no):
|
|
return
|
|
if owner.download_claim_token != claim_token:
|
|
return
|
|
owner.pipeline_stage = _stage(
|
|
owner, ChatGenerationPipelineStage.DOWNLOAD_QUEUED
|
|
)
|
|
owner.download_claim_token = None
|
|
owner.download_lease_until = None
|
|
owner.download_started_at = None
|
|
owner.download_attempt_count = max(
|
|
0, int(owner.download_attempt_count or 0) - 1
|
|
)
|
|
owner.download_enqueued_at = _now()
|
|
owner_type_snapshot = owner_type_of(owner)
|
|
owner_id_snapshot = str(owner.id)
|
|
attempt_snapshot = int(owner.generation_attempt_no or 1)
|
|
await db.commit()
|
|
owner = await _reload_owner_after_commit(
|
|
db,
|
|
owner_type=owner_type_snapshot,
|
|
owner_id=owner_id_snapshot,
|
|
attempt_no=attempt_snapshot,
|
|
)
|
|
if owner is None:
|
|
return
|
|
await _register_active(
|
|
owner,
|
|
check_at=_queue_timeout_at(),
|
|
priority=int(settings.DOWNLOAD_TASK_PRIORITY_RECOVER),
|
|
reason="download_execution_lock_error",
|
|
)
|
|
|
|
|
|
async def _run(
|
|
task_id: str,
|
|
*,
|
|
owner_type: str = GenerationOwnerType.CHAT_GENERATION_TASK.value,
|
|
generation_attempt_no: int | None = None,
|
|
):
|
|
normalized_owner_type = normalize_owner_type(owner_type)
|
|
effective_attempt = await _resolve_attempt(
|
|
task_id,
|
|
owner_type=normalized_owner_type,
|
|
message_attempt=generation_attempt_no,
|
|
)
|
|
if effective_attempt is None:
|
|
return
|
|
|
|
token = uuid.uuid4().hex
|
|
lease = await CeleryRuntimeLease.acquire(
|
|
identity=RuntimeIdentity(
|
|
domain=CeleryRuntimeDomain.GENERATION_DOWNLOAD.value,
|
|
owner_type=normalized_owner_type,
|
|
owner_id=task_id,
|
|
attempt_no=effective_attempt,
|
|
task_name=CeleryTaskName.DOWNLOAD_GENERATION_RESULT.value,
|
|
queue=DOWNLOAD_QUEUE,
|
|
registry_item_id=redis_owner_item_id(
|
|
normalized_owner_type, task_id, effective_attempt
|
|
),
|
|
),
|
|
lock_key=_lock_key(normalized_owner_type, task_id, effective_attempt),
|
|
hash_key=settings.DOWNLOAD_ACTIVE_REDIS_HASH_KEY,
|
|
zset_key=settings.DOWNLOAD_ACTIVE_REDIS_ZSET_KEY,
|
|
token=token,
|
|
ttl_seconds=int(settings.GENERATION_DOWNLOAD_LOCK_TTL_SECONDS or 600),
|
|
heartbeat_interval_seconds=int(settings.REDIS_EXECUTION_LOCK_RENEW_INTERVAL_SECONDS or 30),
|
|
pipeline_stage=ChatGenerationPipelineStage.DOWNLOADING.value,
|
|
db_heartbeat=lambda owned_token: renew_generation_owner_claim_lease(
|
|
owner_type=normalized_owner_type,
|
|
owner_id=task_id,
|
|
attempt_no=effective_attempt,
|
|
claim_field="download_claim_token",
|
|
lease_field="download_lease_until",
|
|
token=owned_token,
|
|
lease_seconds=int(settings.DOWNLOAD_TASK_LEASE_SECONDS or 600),
|
|
),
|
|
)
|
|
if lease is None:
|
|
return
|
|
|
|
claimed = False
|
|
try:
|
|
async with async_session() as db:
|
|
owner = await load_generation_owner(
|
|
db,
|
|
owner_type=normalized_owner_type,
|
|
owner_id=task_id,
|
|
for_update=True,
|
|
)
|
|
if not owner:
|
|
await remove_download_active(
|
|
redis_owner_item_id(
|
|
normalized_owner_type, task_id, effective_attempt
|
|
)
|
|
)
|
|
return
|
|
if not is_attempt_current(owner, effective_attempt):
|
|
await _remove_active(owner)
|
|
return
|
|
owner = await _claim(db, owner, claim_token=lease.token)
|
|
if owner is None:
|
|
return
|
|
claimed = True
|
|
|
|
try:
|
|
use_upscale = bool(
|
|
owner.gen_type == GenerationType.VIDEO.value
|
|
and owner.video_upscale_enabled_snapshot
|
|
and owner.video_upscale_snapshot_json
|
|
)
|
|
try:
|
|
async with asyncio.timeout(
|
|
max(1, int(settings.GENERATION_DOWNLOAD_TOTAL_TIMEOUT_SECONDS or 480))
|
|
):
|
|
downloaded = (
|
|
await download_video_upscale_source(
|
|
owner, execution_guard=lease.ensure_owned
|
|
)
|
|
if use_upscale
|
|
else await download_generation_result(
|
|
owner, execution_guard=lease.ensure_owned
|
|
)
|
|
)
|
|
except TimeoutError as exc:
|
|
raise TimeoutError("下载总耗时超过限制") from exc
|
|
|
|
# ffprobe is an external process and must not run while holding
|
|
# the generation owner row lock.
|
|
upscale_source_info = None
|
|
if use_upscale:
|
|
await lease.ensure_owned()
|
|
upscale_source_info = await probe_video(
|
|
str(downloaded.storage_path or "")
|
|
)
|
|
|
|
await lease.ensure_owned()
|
|
owner = await load_generation_owner_for_update_retry(
|
|
db,
|
|
owner_type=normalized_owner_type,
|
|
owner_id=task_id,
|
|
)
|
|
if not owner or not is_attempt_current(owner, effective_attempt):
|
|
await remove_download_active(
|
|
redis_owner_item_id(
|
|
normalized_owner_type, task_id, effective_attempt
|
|
)
|
|
)
|
|
return
|
|
if owner.download_claim_token != lease.token:
|
|
return
|
|
|
|
if use_upscale:
|
|
from app.services.video_upscale.task_service import (
|
|
enqueue_upscale_task,
|
|
prepare_video_upscale_task,
|
|
)
|
|
|
|
kwargs = (
|
|
{"task": owner}
|
|
if isinstance(owner, ChatGenerationTask)
|
|
else {"generation_record": owner}
|
|
)
|
|
upscale = await prepare_video_upscale_task(
|
|
db,
|
|
**kwargs,
|
|
source_local_path=str(downloaded.storage_path or ""),
|
|
source_file_size_bytes=downloaded.file_size_bytes,
|
|
source_remote_url=owner.remote_result_url,
|
|
source_info=upscale_source_info,
|
|
)
|
|
owner.download_claim_token = None
|
|
owner.download_lease_until = None
|
|
owner.download_next_retry_at = None
|
|
owner.download_last_error = None
|
|
upscale_mode = str(getattr(owner, "generation_mode", "") or "") or None
|
|
upscale_stage = str(owner.pipeline_stage or "") or None
|
|
await enqueue_upscale_task(
|
|
db, upscale=upscale, reason="source_download_completed"
|
|
)
|
|
owner = await _reload_owner_after_commit(
|
|
db,
|
|
owner_type=normalized_owner_type,
|
|
owner_id=task_id,
|
|
attempt_no=effective_attempt,
|
|
)
|
|
if owner is not None:
|
|
await _remove_active(owner)
|
|
await log_task_event(
|
|
owner_type=normalized_owner_type,
|
|
owner_id=task_id,
|
|
generation_attempt_no=effective_attempt,
|
|
generation_mode=upscale_mode,
|
|
event_type=ChatGenerationTaskEventType.DOWNLOAD_SUCCESS.value,
|
|
to_stage=(str(owner.pipeline_stage or "") if owner is not None else upscale_stage),
|
|
detail={
|
|
"upscale_source_path": downloaded.storage_path
|
|
},
|
|
)
|
|
return
|
|
|
|
if owner.gen_type == GenerationType.IMAGE.value:
|
|
owner.image_url = downloaded.url
|
|
else:
|
|
owner.video_url = downloaded.url
|
|
owner.video_cover_url = downloaded.cover_url
|
|
set_owner_completed(owner)
|
|
owner.pipeline_stage = _stage(
|
|
owner, ChatGenerationPipelineStage.DONE
|
|
)
|
|
owner.generated_at = _now()
|
|
owner.retry_count = int(
|
|
getattr(owner, "manual_retry_count", 0) or 0
|
|
)
|
|
owner.download_claim_token = None
|
|
owner.download_lease_until = None
|
|
owner.download_next_retry_at = None
|
|
owner.download_last_error = None
|
|
await _record_resource(db, owner, downloaded)
|
|
await _sync_snapshot(db, owner)
|
|
completion_mode = str(getattr(owner, "generation_mode", "") or "") or None
|
|
completion_stage = str(owner.pipeline_stage or "") or None
|
|
await db.commit()
|
|
owner = await _reload_owner_after_commit(
|
|
db,
|
|
owner_type=normalized_owner_type,
|
|
owner_id=task_id,
|
|
attempt_no=effective_attempt,
|
|
)
|
|
if owner is not None:
|
|
await notify_owner_finished(db, owner)
|
|
await db.commit()
|
|
owner = await _reload_owner_after_commit(
|
|
db,
|
|
owner_type=normalized_owner_type,
|
|
owner_id=task_id,
|
|
attempt_no=effective_attempt,
|
|
)
|
|
if owner is not None:
|
|
await _remove_active(owner)
|
|
await log_task_event(
|
|
owner_type=normalized_owner_type,
|
|
owner_id=task_id,
|
|
generation_attempt_no=effective_attempt,
|
|
generation_mode=completion_mode,
|
|
event_type=ChatGenerationTaskEventType.DOWNLOAD_SUCCESS.value,
|
|
to_stage=(str(owner.pipeline_stage or "") if owner is not None else completion_stage),
|
|
detail={
|
|
"resource_url": downloaded.url,
|
|
"file_size_bytes": downloaded.file_size_bytes,
|
|
},
|
|
)
|
|
except (RedisExecutionLockError, DatabaseRowLockBusy):
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
raise
|
|
except Exception as exc:
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
pass
|
|
await lease.ensure_owned()
|
|
owner = await load_generation_owner_for_update_retry(
|
|
db,
|
|
owner_type=normalized_owner_type,
|
|
owner_id=task_id,
|
|
)
|
|
if not owner or not is_attempt_current(owner, effective_attempt):
|
|
await remove_download_active(
|
|
redis_owner_item_id(
|
|
normalized_owner_type, task_id, effective_attempt
|
|
)
|
|
)
|
|
return
|
|
if owner.download_claim_token != lease.token:
|
|
return
|
|
non_retryable = _is_non_retryable_download_error(exc)
|
|
max_attempts = int(settings.DOWNLOAD_TASK_MAX_ATTEMPTS or 3)
|
|
if non_retryable or int(owner.download_attempt_count or 0) >= max_attempts:
|
|
await _mark_failed(
|
|
db, owner, exc, non_retryable=non_retryable
|
|
)
|
|
else:
|
|
await _schedule_retry(db, owner, exc)
|
|
except (RedisExecutionLockError, DatabaseRowLockBusy):
|
|
if claimed:
|
|
try:
|
|
await _restore_after_lock_error(
|
|
owner_type=normalized_owner_type,
|
|
owner_id=task_id,
|
|
attempt_no=effective_attempt,
|
|
claim_token=lease.token,
|
|
)
|
|
except Exception:
|
|
logger.exception(
|
|
"下载执行锁异常后的状态恢复失败 owner_type=%s owner_id=%s attempt=%s",
|
|
normalized_owner_type,
|
|
task_id,
|
|
effective_attempt,
|
|
)
|
|
raise
|
|
finally:
|
|
await lease.close()
|
|
|
|
|
|
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,
|
|
owner_type: str = GenerationOwnerType.CHAT_GENERATION_TASK.value,
|
|
generation_attempt_no: int | None = None,
|
|
):
|
|
try:
|
|
return run_async(
|
|
_run(
|
|
task_id,
|
|
owner_type=owner_type,
|
|
generation_attempt_no=generation_attempt_no,
|
|
)
|
|
)
|
|
except Exception as exc:
|
|
retries = int(getattr(self.request, "retries", 0) or 0) + 1
|
|
countdown = int(
|
|
settings.DOWNLOAD_TASK_RETRY_BACKOFF_SECONDS or 30
|
|
) * max(1, retries)
|
|
raise self.retry(exc=exc, countdown=countdown)
|
|
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()
|