异步生图兜底:回收僵尸出图任务(worker 崩溃后退还被占额度)

worker 崩溃/重启可能留下卡在 RESERVED 的出图任务,额度被一直占住、前端也永远等不到。
每次新提交时顺手回收:>10 分钟仍 RESERVED 的标记失败并退还预留额度(与导出僵尸清理同思路,
无需额外定时任务)。

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
zyc
2026-06-15 16:41:23 +08:00
co-authored by Claude Opus 4.8
parent a749be5837
commit 77c037495e
+31
View File
@@ -6,6 +6,7 @@ from datetime import timedelta
from decimal import Decimal
from io import BytesIO
from pathlib import Path
from django.core.exceptions import ObjectDoesNotExist
from django.db import transaction
from django.utils import timezone
@@ -819,6 +820,35 @@ _STANDALONE_TASK_TYPE = {
}
def _reap_stale_standalone_image_tasks(*, team) -> None:
"""兜底:worker 崩溃/重启(OOM、部署)可能留下卡在 RESERVED 的出图任务,额度被一直占住、
前端轮询也永远等不到结果。超过 10 分钟(远大于单张真实出图耗时 ~60s)仍 RESERVED 的判为僵尸:
标记失败并退还预留额度。趁每次新提交时顺手回收,无需额外的定时任务(与导出僵尸清理同思路)。"""
cutoff = timezone.now() - timedelta(minutes=10)
stale = AITask.objects.filter(
team=team,
project__isnull=True,
task_type__in=[AITask.Type.PERSON_IMAGE, AITask.Type.PRODUCT_IMAGE],
status=AITask.Status.RESERVED,
updated_at__lt=cutoff,
)
for task in stale:
try:
with transaction.atomic():
task.status = AITask.Status.FAILED
task.error_message = "worker 未在预期时间内完成(僵尸任务自动回收)"
task.completed_at = timezone.now()
task.save(update_fields=["status", "error_message", "completed_at", "updated_at"])
try:
reservation = task.credit_reservation
except ObjectDoesNotExist:
reservation = None
if reservation is not None:
release_credit(reservation=reservation, reason="僵尸出图任务自动回收")
except Exception: # noqa: BLE001 — 单个回收失败不应阻断新任务提交
continue
def enqueue_standalone_images(*, team, user, prompt: str, mode: str = "image", count: int = 1) -> list[AITask]:
"""独立生图(图片创作 / 模特上身图 / 平台套图)改为**异步**:本函数在 Web 请求里只做「建任务 +
预留额度」这种秒级的活,真正 ~30s 的 ARK 出图交给 Celery worker(generate_standalone_image_task)。
@@ -828,6 +858,7 @@ def enqueue_standalone_images(*, team, user, prompt: str, mode: str = "image", c
worker 内闭环。返回已 RESERVED 的 AITask 列表,前端拿 id 轮询 GET /api/ai/generate-image/?ids=… 取结果。"""
from apps.ai.tasks import generate_standalone_image_task
_reap_stale_standalone_image_tasks(team=team)
model_config = get_default_model(ModelConfig.Capability.IMAGE)
if model_config is None:
raise ValueError("no active image model configured")