mz/video-gen
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

celery 容灾升级

Browse Source
This commit is contained in:
GinHa
2026-07-22 14:48:29 +08:00
parent 3f1c4063b0
commit 69e7dec807
67 changed files with 6161 additions and 1958 deletions
Download Patch File Download Diff File Expand all files Collapse all files
video-gen-api/app/services/generation/ai/image_batch_service.py
+162 -50
View File
@@ -40,6 +40,7 @@ class ImageBatchClaim:
task_snapshot: SimpleNamespace | None = None
runtime_engine: SimpleNamespace | None = None
existing_child_ids: list[str] | None = None
staged_provider_result: dict | None = None
reason: str | None = None
@@ -53,6 +54,23 @@ def _json(value) -> str | None:
return json.dumps(value, ensure_ascii=False, default=str)
def _parse_staged_provider_result(raw: str | None) -> dict | None:
if not raw:
return None
try:
value = json.loads(raw)
except (TypeError, ValueError, json.JSONDecodeError):
return None
if not isinstance(value, dict):
return None
items = value.get("items")
if not isinstance(items, list) or not items:
return None
return value
def _aware(value: datetime | None) -> datetime | None:
if value is None:
return None
@@ -125,6 +143,13 @@ async def _claim_image_main_batch(
return ImageBatchClaim(False, main_task_id, reason=f"status_{status}")
now = _now()
staged_provider_result = None
if main.pipeline_stage == ChatGenerationPipelineStage.PROVIDER_RESULT_STAGED.value:
staged_provider_result = _parse_staged_provider_result(main.provider_response_json)
if staged_provider_result is None:
main.pipeline_stage = ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value
main.provider_response_json = None
if _lease_alive(main, now):
user_id = str(main.user_id)
group_id = str(main.id)
@@ -143,7 +168,7 @@ async def _claim_image_main_batch(
return ImageBatchClaim(False, main_task_id, reason="lease_alive")
deadline = _aware(main.deadline_at)
if deadline and deadline <= now:
if deadline and deadline <= now and staged_provider_result is None:
main.provider_create_claim_token = None
main.provider_create_lease_until = None
await mark_chat_generation_task_failed_and_refund_once(
@@ -159,7 +184,11 @@ async def _claim_image_main_batch(
main.provider_create_claim_token = claim_token
main.provider_create_started_at = now
main.provider_create_lease_until = now + timedelta(seconds=IMAGE_PROVIDER_CLAIM_LEASE_SECONDS)
main.pipeline_stage = ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value
main.pipeline_stage = (
ChatGenerationPipelineStage.PROVIDER_RESULT_STAGED.value
if staged_provider_result is not None
else ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value
)
runtime_engine = await get_runtime_engine(db, main)
snapshot = _task_snapshot(main)
user_id = str(main.user_id)
@@ -187,6 +216,8 @@ async def _claim_image_main_batch(
claim_token=claim_token,
task_snapshot=snapshot,
runtime_engine=runtime_engine,
staged_provider_result=staged_provider_result,
reason="provider_result_staged" if staged_provider_result is not None else None,
)
@@ -299,6 +330,58 @@ async def _fail_claimed_main(
return True
async def _stage_provider_result(
db: AsyncSession,
*,
main_task_id: str,
claim_token: str,
provider_result: dict,
) -> None:
try:
await db.rollback()
except Exception:
pass
result = await execute_with_lock_timeout(
db,
select(ChatGenerationTask)
.where(
ChatGenerationTask.id == main_task_id,
ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_MAIN.value,
ChatGenerationTask.gen_type == GenerationType.IMAGE.value,
ChatGenerationTask.deleted_at.is_(None),
)
.with_for_update()
.limit(1),
)
main = result.scalar_one_or_none()
if not main:
raise RuntimeError("图片主任务不存在或已删除")
if main.provider_create_claim_token != claim_token:
raise RuntimeError("图片主任务执行租约已失效，拒绝暂存供应商结果")
if main.status != ChatGenerationTaskStatus.GENERATING.value:
raise RuntimeError(f"图片主任务当前状态不允许暂存: {main.status}")
main.provider_response_json = _json(provider_result)
main.image_tokens_used = int(provider_result.get("image_tokens") or 0)
main.pipeline_stage = ChatGenerationPipelineStage.PROVIDER_RESULT_STAGED.value
user_id_snapshot = str(main.user_id)
generation_count_snapshot = int(main.generation_count or 1)
image_tokens_snapshot = int(main.image_tokens_used or 0)
await db.commit()
log_operation_event(
domain="generation_ai_batch",
event_type="IMAGE_BATCH_PROVIDER_RESULT_STAGED",
event_status="success",
source="celery",
user_id=user_id_snapshot,
group_id=main_task_id,
task_id=main_task_id,
detail={
"generation_count": generation_count_snapshot,
"image_tokens": image_tokens_snapshot,
},
)
async def _split_children(
db: AsyncSession,
*,
@@ -474,54 +557,75 @@ async def run_image_main_batch(
return []
generation_count = max(1, int(claim.task_snapshot.generation_count or 1))
try:
log_operation_event(
domain="generation_ai_batch",
event_type=ChatGenerationTaskEventType.IMAGE_BATCH_PROVIDER_START.value,
event_status="started",
source="celery",
user_id=claim.task_snapshot.user_id,
group_id=main_task_id,
task_id=main_task_id,
detail={"generation_count": generation_count},
)
provider_result = await create_image_sync_batch_result_with_engine(
claim.task_snapshot,
claim.runtime_engine,
generation_count=generation_count,
)
await execution_guard()
provider_result = claim.staged_provider_result
if provider_result is None:
try:
log_operation_event(
domain="generation_ai_batch",
event_type=ChatGenerationTaskEventType.IMAGE_BATCH_PROVIDER_START.value,
event_status="started",
source="celery",
user_id=claim.task_snapshot.user_id,
group_id=main_task_id,
task_id=main_task_id,
detail={"generation_count": generation_count},
)
provider_result = await create_image_sync_batch_result_with_engine(
claim.task_snapshot,
claim.runtime_engine,
generation_count=generation_count,
)
await execution_guard()
provider_items = _validate_provider_batch(provider_result, generation_count)
await _stage_provider_result(
db,
main_task_id=main_task_id,
claim_token=claim.claim_token,
provider_result=provider_result,
)
log_operation_event(
domain="generation_ai_batch",
event_type=ChatGenerationTaskEventType.IMAGE_BATCH_PROVIDER_SUCCESS.value,
event_status="success",
source="celery",
user_id=claim.task_snapshot.user_id,
group_id=main_task_id,
task_id=main_task_id,
detail={
"generation_count": generation_count,
"result_count": len(provider_items),
"image_tokens": int(provider_result.get("image_tokens") or 0),
"single_provider_request": True,
"fallback_to_single_requests": False,
"provider_result_staged": True,
},
)
except RedisExecutionLockError:
raise
except Exception as exc:
await execution_guard()
message = exc.safe_message if isinstance(exc, ImageProviderError) else str(exc)
await _fail_claimed_main(
db,
main_task_id=main_task_id,
claim_token=claim.claim_token,
error_message=message or "图片批量生成失败",
event_type=ChatGenerationTaskEventType.IMAGE_BATCH_PROVIDER_FAILED,
exception=exc,
)
return []
else:
provider_items = _validate_provider_batch(provider_result, generation_count)
log_operation_event(
domain="generation_ai_batch",
event_type=ChatGenerationTaskEventType.IMAGE_BATCH_PROVIDER_SUCCESS.value,
event_type="IMAGE_BATCH_STAGED_RESULT_RECOVERED",
event_status="success",
source="celery",
source="recovery",
user_id=claim.task_snapshot.user_id,
group_id=main_task_id,
task_id=main_task_id,
detail={
"generation_count": generation_count,
"result_count": len(provider_items),
"image_tokens": int(provider_result.get("image_tokens") or 0),
"single_provider_request": True,
"fallback_to_single_requests": False,
},
detail={"generation_count": generation_count, "provider_regenerated": False},
)
except RedisExecutionLockError:
raise
except Exception as exc:
await execution_guard()
message = exc.safe_message if isinstance(exc, ImageProviderError) else str(exc)
await _fail_claimed_main(
db,
main_task_id=main_task_id,
claim_token=claim.claim_token,
error_message=message or "图片批量生成失败",
event_type=ChatGenerationTaskEventType.IMAGE_BATCH_PROVIDER_FAILED,
exception=exc,
)
return []
try:
await execution_guard()
@@ -535,14 +639,22 @@ async def run_image_main_batch(
except RedisExecutionLockError:
raise
except Exception as exc:
await execution_guard()
await _fail_claimed_main(
db,
main_task_id=main_task_id,
claim_token=claim.claim_token,
error_message=f"图片批量结果拆分失败: {exc}",
event_type=ChatGenerationTaskEventType.IMAGE_BATCH_SPLIT_FAILED,
exception=exc,
# 供应商结果已经落库；拆分失败只记录并等待恢复，绝不退款或重新调用供应商。
try:
await db.rollback()
except Exception:
pass
log_operation_event(
domain="generation_ai_batch",
event_type=ChatGenerationTaskEventType.IMAGE_BATCH_SPLIT_FAILED.value,
event_status="failed",
source="celery",
user_id=claim.task_snapshot.user_id,
group_id=main_task_id,
task_id=main_task_id,
message=f"图片批量结果拆分失败: {exc}",
detail={"provider_result_staged": True, "provider_regenerated": False},
error=str(exc),
)
return []
video-gen-api/app/services/generation/ai/task_group_service.py
+64 -44
View File
@@ -27,6 +27,7 @@ ACTIVE_STAGES = {
ChatGenerationPipelineStage.QUEUED.value,
ChatGenerationPipelineStage.PREPARING.value,
ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value,
ChatGenerationPipelineStage.PROVIDER_RESULT_STAGED.value,
ChatGenerationPipelineStage.WAITING_REMOTE.value,
ChatGenerationPipelineStage.POLLING.value,
ChatGenerationPipelineStage.RESULT_READY.value,
@@ -153,31 +154,7 @@ def _build_summary(children: list[ChatGenerationTask]) -> str | None:
return f"{len(children)}项中" + "，".join(parts)
async def aggregate_main_task_status(
db: AsyncSession,
*,
parent_task_id: str,
) -> ChatGenerationTask | None:
result = await execute_with_lock_timeout(
db,
select(ChatGenerationTask)
.where(
ChatGenerationTask.id == parent_task_id,
ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_MAIN.value,
)
.with_for_update()
.limit(1)
)
main = result.scalar_one_or_none()
if not main or main.deleted_at is not None:
return main
children_map = await load_children_map(db, [parent_task_id], include_deleted=True)
children = children_map.get(parent_task_id, [])
if not children:
return main
def _apply_main_task_status(main: ChatGenerationTask, children: list[ChatGenerationTask]) -> dict[str, object]:
previous_status = main.status
previous_stage = main.pipeline_stage
active_children = [child for child in children if is_task_active(child)]
@@ -217,7 +194,6 @@ async def aggregate_main_task_status(
)
main.error_message = _build_summary(children)
else:
# 所有子任务真实生成结果均成功；资源是否软删除不改变生成历史终态。
main.status = ChatGenerationTaskStatus.COMPLETED.value
main.pipeline_stage = ChatGenerationPipelineStage.DONE.value
main.generated_at = max(
@@ -232,29 +208,73 @@ async def aggregate_main_task_status(
main.text_tokens_used = sum(int(child.text_tokens_used or 0) for child in children)
main.image_tokens_used = sum(int(child.image_tokens_used or 0) for child in children)
main.video_tokens_used = sum(int(child.video_tokens_used or 0) for child in children)
# main 的手动重试次数只代表 main 自身，不能累加 child 的轮询/重试次数。
main.retry_count = int(main.manual_retry_count or 0)
main.poll_count = sum(int(child.poll_count or 0) for child in children)
main.poll_error_count = sum(int(child.poll_error_count or 0) for child in children)
await db.flush()
log_operation_event(
domain="generation_ai_batch",
event_type="MAIN_STATUS_AGGREGATED",
event_status="success",
source="service",
user_id=main.user_id,
group_id=main.id,
task_id=main.id,
detail={
"before_status": previous_status,
"before_stage": previous_stage,
"after_status": main.status,
"after_stage": main.pipeline_stage,
"summary": _build_summary(children),
},
return {
"before_status": previous_status,
"before_stage": previous_stage,
"after_status": main.status,
"after_stage": main.pipeline_stage,
"summary": _build_summary(children),
}
async def aggregate_main_tasks_status_batch(
db: AsyncSession,
*,
parent_task_ids: Sequence[str] | Iterable[str],
) -> dict[str, ChatGenerationTask]:
ids = list(dict.fromkeys(str(item) for item in parent_task_ids if item))
if not ids:
return {}
result = await execute_with_lock_timeout(
db,
select(ChatGenerationTask)
.where(
ChatGenerationTask.id.in_(ids),
ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_MAIN.value,
)
.order_by(ChatGenerationTask.id.asc())
.with_for_update(),
)
return main
mains = list(result.scalars().all())
children_map = await load_children_map(db, ids, include_deleted=True)
log_snapshots: list[tuple[str, str | None, dict[str, object]]] = []
main_map: dict[str, ChatGenerationTask] = {}
for main in mains:
main_id = str(main.id)
main_map[main_id] = main
if main.deleted_at is not None:
continue
children = children_map.get(main_id, [])
if not children:
continue
detail = _apply_main_task_status(main, children)
log_snapshots.append((main_id, str(main.user_id) if main.user_id else None, detail))
await db.flush()
for main_id, user_id, detail in log_snapshots:
log_operation_event(
domain="generation_ai_batch",
event_type="MAIN_STATUS_AGGREGATED",
event_status="success",
source="service",
user_id=user_id,
group_id=main_id,
task_id=main_id,
detail=detail,
)
return main_map
async def aggregate_main_task_status(
db: AsyncSession,
*,
parent_task_id: str,
) -> ChatGenerationTask | None:
main_map = await aggregate_main_tasks_status_batch(db, parent_task_ids=[parent_task_id])
return main_map.get(str(parent_task_id))
async def aggregate_parent_for_child(db: AsyncSession, child: ChatGenerationTask | None) -> ChatGenerationTask | None:
video-gen-api/app/services/generation/billing_service.py
+35 -2
View File
@@ -7,7 +7,7 @@ from typing import Any, Mapping
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.enums.credit_record import CreditRecordBillingScene, CreditRecordChargeKind, CreditRecordOwnerType, CreditRecordSourceModule
from app.enums.credit_record import CreditRecordChargeKind, CreditRecordOwnerType, CreditRecordSourceModule
from app.models.credit_record import CreditRecord
from app.models.generation_record import GenerationRecord
from app.services.generation.media_reference_service import calculate_media_reference_usage
@@ -22,6 +22,7 @@ from app.services.credit_record_meta_service import (
build_shot_video_analysis_meta,
)
from app.services.credits import calc_image_credits, calc_text_credits, calc_video_credits, deduct_credits
from app.utils.id_gen import generate_id
CHARGE_TEXT_PROMPT = CreditRecordChargeKind.TEXT_PROMPT.value
@@ -416,12 +417,44 @@ async def charge_shot_video_analysis_usage(
charge_kind=CHARGE_VIDEO_ANALYSIS,
action="charge",
)
usage_snapshot = dict(usage)
token_usage_result = await db.execute(
select(TokenUsage)
.where(
TokenUsage.owner_type == owner_type,
TokenUsage.owner_id == owner_id,
TokenUsage.biz_key == biz_key,
)
.order_by(TokenUsage.created_at.asc())
.limit(1)
)
token_usage = token_usage_result.scalar_one_or_none()
if token_usage is None:
token_usage = TokenUsage(
id=generate_id(),
model_config_id=usage_snapshot.get("model_config_id"),
user_id=user_id,
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=_safe_int(
usage_snapshot.get("total_tokens"),
input_tokens + output_tokens,
),
owner_type=owner_type,
owner_id=owner_id,
biz_key=biz_key,
source_module=CreditRecordSourceModule.SHOT_REPLICATE.value,
source_step_code="video_analysis",
)
db.add(token_usage)
await db.flush()
usage_snapshot["token_usage_id"] = token_usage.id
record_meta = await build_shot_video_analysis_meta(
db,
owner_type=owner_type,
owner_id=owner_id,
attempt_no=attempt_no,
usage=usage,
usage=usage_snapshot,
billing_scene=billing_scene,
source_project_id=source_project_id,
source_step_id=source_step_id,
video-gen-api/app/services/generation/pipeline/owner_service.py
+33
View File
@@ -212,3 +212,36 @@ def parse_redis_owner_item_id(value: str) -> GenerationOwnerRef:
return GenerationOwnerRef(owner_type, rest, None)
# Historical Redis/Celery identifiers always belonged to ChatGenerationTask.
return GenerationOwnerRef(GenerationOwnerType.CHAT_GENERATION_TASK.value, text, None)
async def renew_generation_owner_claim_lease(
*,
owner_type: str | GenerationOwnerType | None,
owner_id: str,
attempt_no: int,
claim_field: str,
lease_field: str,
token: str,
lease_seconds: int,
) -> bool:
"""CAS 续期生成所有者租约，不加载 ORM 对象，避免 heartbeat 产生懒加载风险。"""
from datetime import datetime, timedelta, timezone
from sqlalchemy import update
from app.models.base import async_session
normalized = normalize_owner_type(owner_type)
model = ChatGenerationTask if normalized == GenerationOwnerType.CHAT_GENERATION_TASK.value else GenerationRecord
claim_column = getattr(model, claim_field)
now = datetime.now(timezone.utc)
async with async_session() as db:
result = await db.execute(
update(model)
.where(
model.id == owner_id,
model.deleted_at.is_(None),
model.generation_attempt_no == int(attempt_no),
claim_column == token,
)
.values({lease_field: now + timedelta(seconds=max(1, int(lease_seconds)))})
)
await db.commit()
return bool(result.rowcount == 1)
video-gen-api/app/services/generation/prompt_service.py
+1 -1
View File
@@ -37,7 +37,7 @@ def _load_refs(record: ChatGenerationTask) -> list[dict]:
async def _build_user_content(record: ChatGenerationTask, db: AsyncSession | None = None) -> list[dict[str, Any]]:
from app.utils.media import media_to_base64
from app.utils.media import get_llm_media_as_base64, media_to_base64
if record.gen_type == "image":
params = f"图片参数：分辨率档位={record.image_size or '2K'}，比例={record.image_proportion or '1:1'}，像素={record.image_px or '2048x2048'}"
video-gen-api/app/services/generation/recovery_service.py
+195 -109
View File
@@ -5,7 +5,7 @@ import logging
from datetime import datetime, timedelta, timezone
from typing import Any
from sqlalchemy import select
from sqlalchemy import or_, select
from sqlalchemy.ext.asyncio import AsyncSession
from app.config import settings
@@ -20,6 +20,7 @@ from app.enums.generation_task import (
GenerationType,
)
from app.models.chat_generation_task import ChatGenerationTask
from app.services.celery_runtime.runtime_service import runtime_lock_exists, runtime_lock_values
from app.services.celery_download_recovery_service import (
ensure_aware_utc,
get_download_active_payloads,
@@ -49,6 +50,31 @@ logger = logging.getLogger("video_gen")
POLL_QUEUE = CeleryQueue.GEN_PROVIDER_POLL.value
def _create_lock_key(task: ChatGenerationTask) -> str:
return (
f"{settings.GENERATION_CREATE_LOCK_KEY_PREFIX}:"
f"{GenerationOwnerType.CHAT_GENERATION_TASK.value}:{task.id}:"
f"attempt:{int(task.generation_attempt_no or 1)}"
)
def _poll_lock_key(task: ChatGenerationTask) -> str:
return (
f"{settings.GENERATION_POLL_LOCK_KEY_PREFIX}:"
f"{GenerationOwnerType.CHAT_GENERATION_TASK.value}:{task.id}:"
f"attempt:{int(task.generation_attempt_no or 1)}"
)
def _download_lock_key(task: ChatGenerationTask) -> str:
return (
f"{settings.GENERATION_DOWNLOAD_LOCK_KEY_PREFIX}:"
f"{GenerationOwnerType.CHAT_GENERATION_TASK.value}:{task.id}:"
f"attempt:{int(task.generation_attempt_no or 1)}"
)
def _chat_registry_id(task: ChatGenerationTask) -> str:
return redis_owner_item_id(
GenerationOwnerType.CHAT_GENERATION_TASK.value,
@@ -211,6 +237,9 @@ async def recover_one_download_task(
)
return "skip_no_remote_result_url"
if await runtime_lock_exists(_download_lock_key(task)):
return "skip_live_download_runtime_lock"
stage = task.pipeline_stage
redis_payload = payload or {}
@@ -478,6 +507,14 @@ async def recover_one_generation_task(
has_provider_task_id = bool(str(task.provider_task_id or "").strip() or str(task.seedance_task_id or "").strip())
is_deadline_expired = bool(task.deadline_at and _is_expired(task.deadline_at, current_time))
runtime_lock_key = (
_download_lock_key(task)
if has_remote_result
else (_poll_lock_key(task) if has_provider_task_id else _create_lock_key(task))
)
if await runtime_lock_exists(runtime_lock_key):
return "skip_live_runtime_lock"
# 最高优先级：只要远程结果 URL 已经落库，说明生成侧已经成功。
# 不管当前 pipeline_stage 是 queued/creating/waiting/result_ready/download_*，恢复时都不能重复 create 或 poll。
if has_remote_result:
@@ -658,6 +695,152 @@ async def recover_one_generation_task(
return f"skip_stage_{task.pipeline_stage}"
async def recover_image_main_create_tasks_once(db: AsyncSession) -> dict[str, int]:
"""批量恢复同步多图主任务；锁活跃或 DB lease 未过期时绝不接管。"""
from app.tasks.generation_create_tasks import chatapi_create_generation_task
results: dict[str, int] = {}
cursor: str | None = None
batch_size = max(1, int(settings.GENERATION_RECOVERY_BATCH_SIZE or 100))
allowed_stages = [
ChatGenerationPipelineStage.QUEUED.value,
ChatGenerationPipelineStage.PREPARING.value,
ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value,
ChatGenerationPipelineStage.PROVIDER_RESULT_STAGED.value,
]
while True:
query = (
select(ChatGenerationTask)
.where(
ChatGenerationTask.deleted_at.is_(None),
ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_MAIN.value,
ChatGenerationTask.gen_type == GenerationType.IMAGE.value,
ChatGenerationTask.status == ChatGenerationTaskStatus.GENERATING.value,
ChatGenerationTask.pipeline_stage.in_(allowed_stages),
or_(
ChatGenerationTask.provider_create_lease_until <= _now(),
(
ChatGenerationTask.provider_create_lease_until.is_(None)
& (
ChatGenerationTask.updated_at
<= _now()
- timedelta(
seconds=max(
1,
int(settings.GENERATION_CREATE_QUEUE_TIMEOUT_SECONDS or 300),
)
)
)
),
),
)
.order_by(ChatGenerationTask.id.asc())
.limit(batch_size)
.with_for_update(skip_locked=True)
)
if cursor:
query = query.where(ChatGenerationTask.id > cursor)
row_result = await db.execute(query)
mains = list(row_result.scalars().all())
if not mains:
break
main_ids = [str(main.id) for main in mains]
cursor = main_ids[-1]
child_rows = await db.execute(
select(ChatGenerationTask.parent_task_id)
.where(
ChatGenerationTask.parent_task_id.in_(main_ids),
ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_CHILD.value,
ChatGenerationTask.deleted_at.is_(None),
)
.distinct()
)
split_parent_ids = {str(value) for value in child_rows.scalars().all() if value}
lock_key_by_id = {str(main.id): _create_lock_key(main) for main in mains}
lock_values = await runtime_lock_values(list(lock_key_by_id.values()))
now = _now()
dispatches: list[tuple[str, int, bool]] = []
expired_claim_logs: list[tuple[str, int, str]] = []
for main in mains:
main_id = str(main.id)
attempt_no = int(main.generation_attempt_no or 1)
if lock_values.get(lock_key_by_id[main_id]):
results["live_lock"] = results.get("live_lock", 0) + 1
continue
if main_id in split_parent_ids:
main.provider_create_claim_token = None
main.provider_create_lease_until = None
results["already_split"] = results.get("already_split", 0) + 1
continue
lease_until = ensure_aware_utc(main.provider_create_lease_until)
if main.provider_create_claim_token and lease_until and lease_until > now:
results["waiting_db_lease"] = results.get("waiting_db_lease", 0) + 1
continue
has_staged_result = bool(
main.pipeline_stage == ChatGenerationPipelineStage.PROVIDER_RESULT_STAGED.value
and main.provider_response_json
)
if _is_expired(main.deadline_at, now) and not has_staged_result:
main.provider_create_claim_token = None
main.provider_create_lease_until = None
await mark_chat_generation_task_failed_and_refund_once(
db,
task=main,
error_message="图片生成任务超时，系统已自动退回本轮媒体生成积分",
pipeline_stage=ChatGenerationPipelineStage.TIMEOUT.value,
)
results["timeout"] = results.get("timeout", 0) + 1
continue
old_claim = str(main.provider_create_claim_token or "")
main.provider_create_claim_token = None
main.provider_create_lease_until = None
if not has_staged_result:
main.pipeline_stage = ChatGenerationPipelineStage.QUEUED.value
if old_claim:
expired_claim_logs.append((main_id, attempt_no, str(main.generation_mode)))
dispatches.append((main_id, attempt_no, has_staged_result))
await db.commit()
for main_id, attempt_no, generation_mode in expired_claim_logs:
await log_task_event(
task_id=main_id,
generation_attempt_no=attempt_no,
generation_mode=generation_mode,
event_type=ChatGenerationTaskEventType.IMAGE_MAIN_CLAIM_EXPIRED.value,
message="图片主任务执行锁已失效且数据库租约已过期，恢复重新投递",
)
for main_id, attempt_no, has_staged_result in dispatches:
try:
chatapi_create_generation_task.apply_async(
args=[main_id],
kwargs={
"owner_type": GenerationOwnerType.CHAT_GENERATION_TASK.value,
"generation_attempt_no": attempt_no,
},
queue=CeleryQueue.GEN_CHATAPI_CREATE.value,
countdown=0,
task_id=(
f"generation-create:{GenerationOwnerType.CHAT_GENERATION_TASK.value}:"
f"{main_id}:attempt:{attempt_no}"
),
)
key = "recover_staged_split" if has_staged_result else "recover_create"
results[key] = results.get(key, 0) + 1
except Exception:
logger.exception("恢复投递图片主任务失败 task_id=%s", main_id)
results["enqueue_failed"] = results.get("enqueue_failed", 0) + 1
if len(mains) < batch_size:
break
return results
async def recover_generation_tasks_once(db: AsyncSession) -> dict[str, Any]:
"""启动时生成链路容灾扫描。
@@ -670,108 +853,9 @@ async def recover_generation_tasks_once(db: AsyncSession) -> dict[str, Any]:
checked_ids: set[str] = set()
results: dict[str, int] = {}
# 图片多份主任务只补投递，不在恢复服务内直接调用供应商。
# 有效 claim 未过期时必须跳过，防止与正在运行的 Worker 重复调用组图 API。
from app.tasks.generation_create_tasks import chatapi_create_generation_task
image_main_cursor: str | None = None
image_main_batch_size = max(1, int(settings.GENERATION_RECOVERY_BATCH_SIZE or 100))
while True:
image_main_query = select(ChatGenerationTask).where(
ChatGenerationTask.deleted_at.is_(None),
ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_MAIN.value,
ChatGenerationTask.gen_type == GenerationType.IMAGE.value,
ChatGenerationTask.status == ChatGenerationTaskStatus.GENERATING.value,
ChatGenerationTask.pipeline_stage.in_([
ChatGenerationPipelineStage.QUEUED.value,
ChatGenerationPipelineStage.PREPARING.value,
ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value,
]),
)
if image_main_cursor:
image_main_query = image_main_query.where(ChatGenerationTask.id > image_main_cursor)
image_main_result = await db.execute(
image_main_query.with_only_columns(ChatGenerationTask.id)
.order_by(ChatGenerationTask.id.asc())
.limit(image_main_batch_size)
)
image_main_ids = [str(value) for value in image_main_result.scalars().all()]
if not image_main_ids:
break
child_parent_result = await db.execute(
select(ChatGenerationTask.parent_task_id)
.where(
ChatGenerationTask.parent_task_id.in_(image_main_ids),
ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_CHILD.value,
)
.distinct()
)
split_parent_ids = {str(value) for value in child_parent_result.scalars().all() if value}
for main_id in image_main_ids:
image_main_cursor = main_id
checked_ids.add(main_id)
main = await _load_chat_task_for_update(db, main_id)
if main is None:
await db.rollback()
continue
if main_id in split_parent_ids:
main.provider_create_claim_token = None
main.provider_create_lease_until = None
await db.commit()
results["image_main_already_split"] = results.get("image_main_already_split", 0) + 1
continue
now = _now()
lease_until = ensure_aware_utc(main.provider_create_lease_until)
lease_alive = bool(main.provider_create_claim_token and lease_until and lease_until > now)
if lease_alive:
await db.rollback()
results["image_main_claim_alive"] = results.get("image_main_claim_alive", 0) + 1
continue
if _is_expired(main.deadline_at, now):
main.provider_create_claim_token = None
main.provider_create_lease_until = None
await mark_chat_generation_task_failed_and_refund_once(
db,
task=main,
error_message="图片批量生成任务超时",
pipeline_stage=ChatGenerationPipelineStage.TIMEOUT.value,
)
await db.commit()
results["image_main_timeout"] = results.get("image_main_timeout", 0) + 1
continue
claim_expired = False
if main.provider_create_claim_token or main.provider_create_lease_until:
main.provider_create_claim_token = None
main.provider_create_lease_until = None
main.pipeline_stage = ChatGenerationPipelineStage.QUEUED.value
claim_expired = True
attempt_no = int(main.generation_attempt_no or 1)
await db.commit()
if claim_expired:
await log_task_event(
main,
event_type=ChatGenerationTaskEventType.IMAGE_MAIN_CLAIM_EXPIRED.value,
message="图片主任务供应商执行租约已过期，恢复重新投递",
)
try:
chatapi_create_generation_task.apply_async(
args=[main_id],
kwargs={"owner_type": GenerationOwnerType.CHAT_GENERATION_TASK.value, "generation_attempt_no": attempt_no},
queue=CeleryQueue.GEN_CHATAPI_CREATE.value,
countdown=0,
)
results["recover_image_main_create"] = results.get("recover_image_main_create", 0) + 1
except Exception as exc:
logger.exception("恢复投递图片主任务失败 task_id=%s: %s", main_id, exc)
results["recover_image_main_enqueue_failed"] = results.get("recover_image_main_enqueue_failed", 0) + 1
if len(image_main_ids) < image_main_batch_size:
break
image_main_results = await recover_image_main_create_tasks_once(db)
for key, value in image_main_results.items():
results[f"image_main_{key}"] = value
due_poll_ids = await redis_get_due_registry_ids(
zset_key=settings.POLL_ACTIVE_REDIS_ZSET_KEY,
@@ -868,7 +952,7 @@ async def recover_generation_tasks_once(db: AsyncSession) -> dict[str, Any]:
break
# 子任务可能在 worker 中断前已进入终态但主任务尚未汇总，按稳定游标完整重算全部主任务。
from app.services.generation.ai.task_group_service import aggregate_main_task_status
from app.services.generation.ai.task_group_service import aggregate_main_tasks_status_batch
reconciled = 0
main_cursor: str | None = None
while True:
@@ -882,11 +966,10 @@ async def recover_generation_tasks_once(db: AsyncSession) -> dict[str, Any]:
parent_ids = list(main_result.scalars().all())
if not parent_ids:
break
for parent_task_id in parent_ids:
main_cursor = str(parent_task_id)
await aggregate_main_task_status(db, parent_task_id=str(parent_task_id))
await db.commit()
reconciled += 1
main_cursor = str(parent_ids[-1])
await aggregate_main_tasks_status_batch(db, parent_task_ids=[str(value) for value in parent_ids])
await db.commit()
reconciled += len(parent_ids)
if len(parent_ids) < batch_size:
break
if reconciled:
@@ -938,6 +1021,9 @@ async def recover_stale_create_tasks_once(db: AsyncSession) -> dict[str, Any]:
)
task_ids = [str(value) for value in result.scalars().all()]
counts: dict[str, int] = {}
image_main_counts = await recover_image_main_create_tasks_once(db)
for key, value in image_main_counts.items():
counts[f"image_main_{key}"] = value
for task_id in task_ids:
task = await _load_chat_task_for_update(db, task_id)
if task is None: