chengma/cmhub
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: retry async image tasks

Browse Source
This commit is contained in:
QiuSW
2026-07-09 14:46:09 +08:00
parent e3598bfe8b
commit 8878769ec5
16 changed files with 550 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/image_tasks.py
+149 -10
View File
@@ -31,6 +31,7 @@ from .models import ImageGenerationTask
IDEMPOTENCY_KEY_MAX_LENGTH = 128
TASK_FAILURE_MESSAGE = "图片生成失败，已退回点数"
TASK_TIMEOUT_MESSAGE = "图片生成任务超时，已退回点数"
RETRYABLE_TASK_ERROR_CODES = {"upstream_timeout", "upstream_error"}
def create_image_generation_task(
@@ -207,6 +208,10 @@ def claim_next_image_task(worker_id: str | None = None) -> ImageGenerationTask |
queryset = (
ImageGenerationTask.objects.select_related("user", "api_key", "call_record")
.filter(status=ImageGenerationTask.Status.QUEUED)
.filter(
models_q("next_attempt_at__isnull", True)
| models_q("next_attempt_at__lte", now)
)
.order_by("created_at", "id")
)
if connection.features.has_select_for_update_skip_locked:
@@ -224,6 +229,7 @@ def claim_next_image_task(worker_id: str | None = None) -> ImageGenerationTask |
task.lease_expires_at = lease_expires_at
task.heartbeat_at = now
task.started_at = task.started_at or now
task.next_attempt_at = None
task.attempt_count += 1
task.save(
update_fields=(
@@ -233,6 +239,7 @@ def claim_next_image_task(worker_id: str | None = None) -> ImageGenerationTask |
"lease_expires_at",
"heartbeat_at",
"started_at",
"next_attempt_at",
"attempt_count",
"updated_at",
)
@@ -261,24 +268,56 @@ def run_image_generation_task(
try:
precharged = precharged_generation_for_task(task)
result = execute_precharged_generation(
precharged,
image_url_builder=media_public_url_builder,
)
except ApiRequestError as exc:
refund_task_call(task, exc.message)
return mark_task_failed_if_running(task.pk, exc.code, exc.message)
try:
result = execute_precharged_generation(
precharged,
image_url_builder=media_public_url_builder,
refund_on_failure=False,
)
except ApiRequestError as exc:
return handle_task_generation_failure(
task,
error_code=exc.code,
error_message=exc.message,
retryable=is_retryable_task_error(exc.code),
)
except Exception as exc:
refund_task_call(task, str(exc))
return mark_task_failed_if_running(
task.pk,
"upstream_error",
TASK_FAILURE_MESSAGE,
return handle_task_generation_failure(
task,
error_code="upstream_error",
error_message=str(exc) or TASK_FAILURE_MESSAGE,
retryable=True,
)
return mark_task_succeeded_if_running(task.pk, result.image_url)
def handle_task_generation_failure(
task: ImageGenerationTask,
*,
error_code: str,
error_message: str,
retryable: bool,
) -> ImageGenerationTask:
normalized_code = str(error_code or "upstream_error")
normalized_message = str(error_message or TASK_FAILURE_MESSAGE)
if retryable and task.attempt_count < image_task_max_attempts():
return mark_task_retry_if_running(
task.pk,
normalized_code,
retry_error_message(normalized_code, normalized_message),
next_attempt_at=timezone.now()
+ timedelta(seconds=image_task_retry_backoff_seconds(task.attempt_count)),
)
refund_task_call(task, normalized_message)
return mark_task_failed_if_running(task.pk, normalized_code, normalized_message)
def precharged_generation_for_task(task: ImageGenerationTask) -> PrechargedGeneration:
payload = dict(task.request_payload or {})
prepared = prepare_generation(
@@ -344,6 +383,7 @@ def mark_task_succeeded_if_running(
task.result_url = str(result_url or "")
task.error_code = ""
task.error_message = ""
task.next_attempt_at = None
task.finished_at = now
task.heartbeat_at = now
task.save(
@@ -352,6 +392,7 @@ def mark_task_succeeded_if_running(
"result_url",
"error_code",
"error_message",
"next_attempt_at",
"finished_at",
"heartbeat_at",
"updated_at",
@@ -360,6 +401,41 @@ def mark_task_succeeded_if_running(
return task
def mark_task_retry_if_running(
task_pk: int,
error_code: str,
error_message: str,
*,
next_attempt_at,
) -> ImageGenerationTask:
with transaction.atomic():
task = ImageGenerationTask.objects.select_for_update().get(pk=task_pk)
if task.status != ImageGenerationTask.Status.RUNNING:
return task
task.status = ImageGenerationTask.Status.QUEUED
task.error_code = str(error_code or "upstream_error")
task.error_message = str(error_message or TASK_FAILURE_MESSAGE)
task.next_attempt_at = next_attempt_at
task.worker_id = ""
task.locked_at = None
task.lease_expires_at = None
task.heartbeat_at = None
task.save(
update_fields=(
"status",
"error_code",
"error_message",
"next_attempt_at",
"worker_id",
"locked_at",
"lease_expires_at",
"heartbeat_at",
"updated_at",
)
)
return task
def mark_task_failed_if_running(
task_pk: int,
error_code: str,
@@ -373,6 +449,7 @@ def mark_task_failed_if_running(
task.status = ImageGenerationTask.Status.FAILED
task.error_code = str(error_code or "upstream_error")
task.error_message = str(error_message or TASK_FAILURE_MESSAGE)
task.next_attempt_at = None
task.finished_at = now
task.heartbeat_at = now
task.save(
@@ -380,6 +457,7 @@ def mark_task_failed_if_running(
"status",
"error_code",
"error_message",
"next_attempt_at",
"finished_at",
"heartbeat_at",
"updated_at",
@@ -422,6 +500,7 @@ def reap_stale_image_task(task_pk: int, now) -> bool:
task.result_url = task.call_record.result_ref
task.error_code = ""
task.error_message = ""
task.next_attempt_at = None
task.finished_at = now
task.save(
update_fields=(
@@ -429,6 +508,7 @@ def reap_stale_image_task(task_pk: int, now) -> bool:
"result_url",
"error_code",
"error_message",
"next_attempt_at",
"finished_at",
"updated_at",
)
@@ -446,14 +526,24 @@ def reap_stale_image_task(task_pk: int, now) -> bool:
if task.call_record.status == CallRecord.Status.SUCCESS:
task.status = ImageGenerationTask.Status.SUCCEEDED
task.result_url = task.call_record.result_ref
task.next_attempt_at = None
task.finished_at = now
task.save(update_fields=("status", "result_url", "finished_at", "updated_at"))
task.save(
update_fields=(
"status",
"result_url",
"next_attempt_at",
"finished_at",
"updated_at",
)
)
return True
raise
task.status = ImageGenerationTask.Status.FAILED
task.error_code = "task_timeout"
task.error_message = TASK_TIMEOUT_MESSAGE
task.next_attempt_at = None
task.finished_at = now
task.heartbeat_at = now
task.save(
@@ -461,6 +551,7 @@ def reap_stale_image_task(task_pk: int, now) -> bool:
"status",
"error_code",
"error_message",
"next_attempt_at",
"finished_at",
"heartbeat_at",
"updated_at",
@@ -510,6 +601,9 @@ def task_submit_response(task: ImageGenerationTask) -> dict[str, Any]:
"call_id": call_record.id,
"points_cost": call_record.points_cost,
"points_balance": task.points_balance_after_charge,
"attempt_count": task.attempt_count,
"max_attempts": image_task_max_attempts(),
"next_attempt_at": task.next_attempt_at.isoformat() if task.next_attempt_at else None,
"created_at": task.created_at.isoformat(),
"expires_at": task.expires_at.isoformat() if task.expires_at else None,
}
@@ -522,6 +616,9 @@ def task_detail_response(task: ImageGenerationTask) -> dict[str, Any]:
"status": task.status,
"call_id": call_record.id,
"points_cost": call_record.points_cost,
"attempt_count": task.attempt_count,
"max_attempts": image_task_max_attempts(),
"next_attempt_at": task.next_attempt_at.isoformat() if task.next_attempt_at else None,
"created_at": task.created_at.isoformat(),
"updated_at": task.updated_at.isoformat(),
"expires_at": task.expires_at.isoformat() if task.expires_at else None,
@@ -553,6 +650,48 @@ def image_task_lease_seconds() -> int:
return max(1, int(getattr(settings, "IMAGE_TASK_LEASE_SECONDS", 600)))
def image_task_max_retries() -> int:
return max(0, int(getattr(settings, "IMAGE_TASK_MAX_RETRIES", 2)))
def image_task_max_attempts() -> int:
return 1 + image_task_max_retries()
def image_task_retry_backoff_seconds(attempt_count: int) -> int:
values = image_task_retry_backoff_values()
if not values:
return 0
index = max(0, int(attempt_count or 1) - 1)
return values[min(index, len(values) - 1)]
def image_task_retry_backoff_values() -> list[int]:
raw = str(getattr(settings, "IMAGE_TASK_RETRY_BACKOFF_SECONDS", "10,30") or "")
values: list[int] = []
for part in raw.split(","):
item = part.strip()
if not item:
continue
try:
values.append(max(0, int(item)))
except ValueError:
continue
return values
def is_retryable_task_error(error_code: str) -> bool:
return str(error_code or "") in RETRYABLE_TASK_ERROR_CODES
def retry_error_message(error_code: str, fallback: str) -> str:
if error_code == "upstream_timeout":
return "上游 AI 调用超时，稍后自动重试"
if error_code == "upstream_error":
return "上游 AI 调用失败，稍后自动重试"
return fallback
def normalize_worker_id(worker_id: str | None) -> str:
normalized = str(worker_id or "").strip()
if normalized: