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 []