from __future__ import annotations import hashlib import json import socket from dataclasses import replace from datetime import timedelta from typing import Any, Mapping from django.conf import settings from django.core.files.base import ContentFile from django.db import IntegrityError, connection, transaction from django.utils import timezone from rest_framework import status from apps.billing.models import CallRecord from apps.billing.services import InvalidCallStateError, refund_call_points from .generation import ( ApiRequestError, GenerationInput, ImageInput, PrechargedGeneration, execute_precharged_generation, prepare_generation, precharge_generation, ) 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( *, user, api_key, request_data: Mapping[str, Any], idempotency_key: str = "", ) -> tuple[ImageGenerationTask, bool]: normalized_key = normalize_idempotency_key(idempotency_key) idempotency_key_hash = hash_text(normalized_key) if normalized_key else None request_hash = request_hash_for_image_request(request_data) existing = find_idempotent_task(api_key, idempotency_key_hash) if existing is not None: return idempotent_task_or_raise(existing, request_hash), False prepared = prepare_generation( GenerationInput( user=user, api_key=api_key, operation_type=CallRecord.OperationType.IMAGE, prompt=request_data["prompt"], alias=request_data.get("model") or None, resolution=request_data.get("resolution") or "1K", parameters=dict(request_data.get("parameters") or {}), image_url=str(request_data.get("image_url") or ""), image_base64=str(request_data.get("image_base64") or ""), aspect_ratio=request_data.get("aspect_ratio") or "1:1", ) ) try: with transaction.atomic(): existing = find_idempotent_task( api_key, idempotency_key_hash, for_update=True, ) if existing is not None: return idempotent_task_or_raise(existing, request_hash), False precharged = precharge_generation(prepared) now = timezone.now() task = ImageGenerationTask.objects.create( user=user, api_key=api_key, call_record=precharged.call_record, idempotency_key=normalized_key, idempotency_key_hash=idempotency_key_hash, request_hash=request_hash, request_payload=task_request_payload( request_data=request_data, image_input=prepared.image_input, ), points_balance_after_charge=precharged.points_balance_after_charge, expires_at=now + timedelta(hours=image_task_retention_hours()), ) store_task_input_image(task, prepared.image_input) return task, True except IntegrityError: if idempotency_key_hash: existing = find_idempotent_task(api_key, idempotency_key_hash) if existing is not None: return idempotent_task_or_raise(existing, request_hash), False raise def find_idempotent_task(api_key, idempotency_key_hash: str | None, *, for_update=False): if not idempotency_key_hash: return None queryset = ImageGenerationTask.objects.select_related("call_record").filter( api_key=api_key, idempotency_key_hash=idempotency_key_hash, ) if for_update: queryset = queryset.select_for_update() return queryset.order_by("id").first() def idempotent_task_or_raise( task: ImageGenerationTask, request_hash: str, ) -> ImageGenerationTask: if task.request_hash != request_hash: raise ApiRequestError( "idempotency_conflict", "Idempotency-Key 已用于不同请求", status.HTTP_409_CONFLICT, ) return task def normalize_idempotency_key(value: str) -> str: normalized = str(value or "").strip() if len(normalized) > IDEMPOTENCY_KEY_MAX_LENGTH: raise ApiRequestError( "bad_request", "Idempotency-Key 过长", status.HTTP_400_BAD_REQUEST, ) return normalized def request_hash_for_image_request(request_data: Mapping[str, Any]) -> str: image_base64 = str(request_data.get("image_base64") or "") payload = { "prompt": str(request_data.get("prompt") or ""), "model": str(request_data.get("model") or ""), "resolution": str(request_data.get("resolution") or "1K"), "aspect_ratio": str(request_data.get("aspect_ratio") or "1:1"), "parameters": dict(request_data.get("parameters") or {}), "image_url": str(request_data.get("image_url") or ""), "image_base64_sha256": hash_text(image_base64) if image_base64 else "", } return hash_json(payload) def task_request_payload( *, request_data: Mapping[str, Any], image_input: ImageInput | None, ) -> dict[str, Any]: image_source = "none" if request_data.get("image_base64"): image_source = "base64_stored" elif request_data.get("image_url"): image_source = "url_stored" payload = { "prompt": str(request_data.get("prompt") or ""), "model": str(request_data.get("model") or ""), "resolution": str(request_data.get("resolution") or "1K"), "aspect_ratio": str(request_data.get("aspect_ratio") or "1:1"), "parameters": dict(request_data.get("parameters") or {}), "image_url": str(request_data.get("image_url") or ""), "image_input": None, } if image_input is not None: payload["image_input"] = { "source": image_source, "mime_type": image_input.mime_type, "filename": image_input.filename, "storage_path": "", } return payload def store_task_input_image( task: ImageGenerationTask, image_input: ImageInput | None, ) -> None: if image_input is None: return task.input_image.save( image_input.filename, ContentFile(image_input.data), save=True, ) payload = dict(task.request_payload or {}) image_info = dict(payload.get("image_input") or {}) image_info["storage_path"] = task.input_image.name payload["image_input"] = image_info task.request_payload = payload task.save(update_fields=("request_payload", "updated_at")) def claim_next_image_task(worker_id: str | None = None) -> ImageGenerationTask | None: normalized_worker_id = normalize_worker_id(worker_id) now = timezone.now() lease_expires_at = now + timedelta(seconds=image_task_lease_seconds()) with transaction.atomic(): 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: queryset = queryset.select_for_update(skip_locked=True) else: queryset = queryset.select_for_update() task = queryset.first() if task is None: return None task.status = ImageGenerationTask.Status.RUNNING task.worker_id = normalized_worker_id task.locked_at = now 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=( "status", "worker_id", "locked_at", "lease_expires_at", "heartbeat_at", "started_at", "next_attempt_at", "attempt_count", "updated_at", ) ) return task def run_one_image_task(worker_id: str | None = None) -> ImageGenerationTask | None: task = claim_next_image_task(worker_id) if task is None: return None return run_image_generation_task(task, worker_id=worker_id) def run_image_generation_task( task: ImageGenerationTask, *, worker_id: str | None = None, ) -> ImageGenerationTask: task = refresh_task(task) if task.status != ImageGenerationTask.Status.RUNNING: return task if worker_id: heartbeat_image_task(task.pk, worker_id) try: precharged = precharged_generation_for_task(task) 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: 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( GenerationInput( user=task.user, api_key=task.api_key, operation_type=CallRecord.OperationType.IMAGE, prompt=str(payload.get("prompt") or ""), alias=str(payload.get("model") or "") or None, resolution=str(payload.get("resolution") or "1K"), parameters=dict(payload.get("parameters") or {}), aspect_ratio=str(payload.get("aspect_ratio") or "1:1"), ) ) image_input = stored_image_input_for_task(task) if image_input is not None: prepared = replace(prepared, image_input=image_input) return PrechargedGeneration( prepared=prepared, call_record=task.call_record, points_cost=task.call_record.points_cost, points_balance_after_charge=task.points_balance_after_charge, ) def stored_image_input_for_task(task: ImageGenerationTask) -> ImageInput | None: if not task.input_image: return None payload = dict(task.request_payload or {}) image_info = dict(payload.get("image_input") or {}) with task.input_image.open("rb") as image_file: image = image_file.read() return ImageInput( data=image, mime_type=str(image_info.get("mime_type") or "image/png"), filename=str(image_info.get("filename") or "input.png"), ) def heartbeat_image_task(task_pk: int, worker_id: str | None = None) -> None: now = timezone.now() ImageGenerationTask.objects.filter( pk=task_pk, status=ImageGenerationTask.Status.RUNNING, ).update( heartbeat_at=now, lease_expires_at=now + timedelta(seconds=image_task_lease_seconds()), worker_id=normalize_worker_id(worker_id), ) def mark_task_succeeded_if_running( task_pk: int, result_url: str, ) -> ImageGenerationTask: now = timezone.now() 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.SUCCEEDED 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( update_fields=( "status", "result_url", "error_code", "error_message", "next_attempt_at", "finished_at", "heartbeat_at", "updated_at", ) ) 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, error_message: str, ) -> ImageGenerationTask: now = timezone.now() 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.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( update_fields=( "status", "error_code", "error_message", "next_attempt_at", "finished_at", "heartbeat_at", "updated_at", ) ) return task def reap_stale_image_tasks( *, now=None, limit: int = 100, ) -> int: current_time = now or timezone.now() stale_ids = list( ImageGenerationTask.objects.filter(status=ImageGenerationTask.Status.RUNNING) .filter(stale_task_filter(current_time)) .order_by("lease_expires_at", "id") .values_list("id", flat=True)[:limit] ) reaped = 0 for task_pk in stale_ids: if reap_stale_image_task(task_pk, current_time): reaped += 1 return reaped def reap_stale_image_task(task_pk: int, now) -> bool: with transaction.atomic(): task = ( ImageGenerationTask.objects.select_for_update() .select_related("call_record") .get(pk=task_pk) ) if task.status != ImageGenerationTask.Status.RUNNING or not task_is_stale(task, now): return False if task.call_record.status == CallRecord.Status.SUCCESS: task.status = ImageGenerationTask.Status.SUCCEEDED 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=( "status", "result_url", "error_code", "error_message", "next_attempt_at", "finished_at", "updated_at", ) ) return True try: refund_call_points( task.call_record, error_message=TASK_TIMEOUT_MESSAGE, reason="Image generation task lease expired; refund precharged points.", ) except InvalidCallStateError: task.call_record.refresh_from_db() 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", "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( update_fields=( "status", "error_code", "error_message", "next_attempt_at", "finished_at", "heartbeat_at", "updated_at", ) ) return True def stale_task_filter(now): stale_heartbeat_before = now - timedelta(seconds=image_task_lease_seconds()) return ( models_q("lease_expires_at__lte", now) | models_q("heartbeat_at__lte", stale_heartbeat_before) ) def task_is_stale(task: ImageGenerationTask, now) -> bool: stale_heartbeat_before = now - timedelta(seconds=image_task_lease_seconds()) if task.lease_expires_at and task.lease_expires_at <= now: return True return bool(task.heartbeat_at and task.heartbeat_at <= stale_heartbeat_before) def refund_task_call(task: ImageGenerationTask, error_message: str) -> None: try: refund_call_points( task.call_record, error_message=error_message, reason="Image generation task failed; refund precharged points.", ) except InvalidCallStateError: return def refresh_task(task: ImageGenerationTask) -> ImageGenerationTask: return ( ImageGenerationTask.objects.select_related("user", "api_key", "call_record") .get(pk=task.pk) ) def task_submit_response(task: ImageGenerationTask) -> dict[str, Any]: call_record = task.call_record return { "task_id": str(task.task_id), "status": task.status, "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, } def task_detail_response(task: ImageGenerationTask) -> dict[str, Any]: call_record = task.call_record data: dict[str, Any] = { "task_id": str(task.task_id), "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, } if task.status == ImageGenerationTask.Status.SUCCEEDED: data["result"] = {"image_url": task.result_url} elif task.status in {ImageGenerationTask.Status.FAILED, ImageGenerationTask.Status.EXPIRED}: data["error"] = { "code": task.error_code or "upstream_error", "message": task.error_message or TASK_FAILURE_MESSAGE, } return data def media_public_url_builder(url: str) -> str: if url.startswith(("http://", "https://")): return url base_url = str(getattr(settings, "MEDIA_PUBLIC_BASE_URL", "") or "").rstrip("/") if base_url and url.startswith("/"): return f"{base_url}{url}" return url def image_task_retention_hours() -> int: return max(1, int(getattr(settings, "IMAGE_TASK_RETENTION_HOURS", 24))) 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: return normalized[:128] return f"{socket.gethostname()}:{hash_text(str(timezone.now().timestamp()))[:8]}" def hash_text(value: str) -> str: return hashlib.sha256(value.encode("utf-8")).hexdigest() def hash_json(value: Mapping[str, Any]) -> str: payload = json.dumps( value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str, ) return hash_text(payload) def models_q(key: str, value): from django.db.models import Q return Q(**{key: value})