diff --git a/core/backend/apps/ai/services.py b/core/backend/apps/ai/services.py index 12292f2..3cc9170 100644 --- a/core/backend/apps/ai/services.py +++ b/core/backend/apps/ai/services.py @@ -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")