Files
video-gen/video-gen-api/app/services/generation/recovery_service.py
T
2026-07-21 19:22:50 +08:00

1110 lines
45 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
from __future__ import annotations
import json
import logging
from datetime import datetime, timedelta, timezone
from typing import Any
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.config import settings
from app.enums.celery_queue import CeleryQueue
from app.enums.generation_task import (
ALLOWED_GENERATION_MODES,
ChatGenerationPipelineStage,
ChatGenerationTaskEventType,
ChatGenerationTaskStatus,
GenerationMode,
GenerationOwnerType,
GenerationType,
)
from app.models.chat_generation_task import ChatGenerationTask
from app.services.celery_download_recovery_service import (
ensure_aware_utc,
get_download_active_payloads,
get_due_download_record_ids,
postpone_download_active_check,
remove_download_active,
)
from app.services.generation.log_service import log_task_event
from app.services.generation.pipeline.db_lock_service import apply_short_lock_timeout
from app.services.generation.pipeline.lifecycle_service import notify_owner_finished
from app.services.generation.poll_schedule_service import ensure_video_poll_fields, is_poll_not_due, is_video_generation_task
from app.services.generation.refund_service import mark_chat_generation_task_failed_and_refund_once
from app.services.generation.pipeline.owner_service import (
load_generation_owner,
parse_redis_owner_item_id,
redis_owner_item_id,
)
from app.services.redis_registry_service import (
redis_get_due_registry_ids,
redis_get_registry_payloads,
redis_postpone_registry_item,
redis_remove_registry_item,
)
logger = logging.getLogger("video_gen")
POLL_QUEUE = CeleryQueue.GEN_PROVIDER_POLL.value
def _chat_registry_id(task: ChatGenerationTask) -> str:
return redis_owner_item_id(
GenerationOwnerType.CHAT_GENERATION_TASK.value,
str(task.id),
int(getattr(task, "generation_attempt_no", 1) or 1),
)
async def _load_chat_task_for_update(
db: AsyncSession, task_id: str
) -> ChatGenerationTask | None:
owner = await load_generation_owner(
db,
owner_type=GenerationOwnerType.CHAT_GENERATION_TASK.value,
owner_id=str(task_id),
for_update=True,
)
return owner if isinstance(owner, ChatGenerationTask) else None
def _now() -> datetime:
return datetime.now(timezone.utc)
def _is_expired(value: datetime | None, now: datetime | None = None) -> bool:
checked = ensure_aware_utc(value)
if checked is None:
return True
return checked <= (now or _now())
def _queue_timeout_at(task: ChatGenerationTask, now: datetime | None = None) -> datetime:
current_time = now or _now()
enqueued_at = ensure_aware_utc(task.download_enqueued_at)
if enqueued_at is None:
return current_time
return enqueued_at + timedelta(seconds=int(settings.DOWNLOAD_TASK_QUEUE_TIMEOUT_SECONDS or 300))
def _is_queue_timeout(task: ChatGenerationTask, now: datetime | None = None) -> bool:
current_time = now or _now()
return _queue_timeout_at(task, current_time) <= current_time
def _is_final_task_state(task: ChatGenerationTask) -> bool:
return task.status in (ChatGenerationTaskStatus.COMPLETED.value, ChatGenerationTaskStatus.FAILED.value) or task.pipeline_stage in (
ChatGenerationPipelineStage.DONE.value,
ChatGenerationPipelineStage.FAILED.value,
ChatGenerationPipelineStage.TIMEOUT.value,
ChatGenerationPipelineStage.DOWNLOAD_FAILED.value,
)
def _is_success(status: str | None) -> bool:
return str(status or "").lower() in ("succeeded", "success", "completed", "done")
def _is_failed(status: str | None) -> bool:
return str(status or "").lower() in ("failed", "error", "canceled", "cancelled")
def _engine_snapshot(task: ChatGenerationTask) -> dict[str, Any]:
try:
value = json.loads(task.engine_snapshot_json or "{}")
return value if isinstance(value, dict) else {}
except Exception:
return {}
def _poll_queue_timeout_at(now: datetime | None = None) -> datetime:
current_time = now or _now()
return current_time + timedelta(seconds=int(settings.POLL_TASK_QUEUE_TIMEOUT_SECONDS or 120))
async def _remove_poll_active(task_id: str) -> None:
await redis_remove_registry_item(
hash_key=settings.POLL_ACTIVE_REDIS_HASH_KEY,
zset_key=settings.POLL_ACTIVE_REDIS_ZSET_KEY,
item_id=task_id,
log_context="poll_active",
)
async def _postpone_poll_active(
*,
task_id: str,
payload: dict[str, Any] | None = None,
check_at: datetime | int | float | None = None,
) -> None:
await redis_postpone_registry_item(
hash_key=settings.POLL_ACTIVE_REDIS_HASH_KEY,
zset_key=settings.POLL_ACTIVE_REDIS_ZSET_KEY,
item_id=task_id,
payload=payload,
check_at=check_at or _poll_queue_timeout_at(),
log_context="poll_active",
)
async def recover_one_download_task(
db: AsyncSession,
task: ChatGenerationTask,
*,
payload: dict[str, Any] | None = None,
source: str = "startup_db",
) -> str:
from app.tasks.generation_download_tasks import (
DOWNLOAD_STAGE_DOWNLOADING,
DOWNLOAD_STAGE_QUEUED,
DOWNLOAD_STAGE_RETRY_WAITING,
enqueue_download_task,
)
current_time = _now()
if not task:
return "skip_missing_task"
if task.generation_mode not in ALLOWED_GENERATION_MODES:
await remove_download_active(_chat_registry_id(task))
return "clean_invalid_mode"
if _is_final_task_state(task):
await remove_download_active(_chat_registry_id(task))
return "clean_final_state"
if task.status != ChatGenerationTaskStatus.GENERATING.value:
task_id = str(task.id)
generation_attempt_no = int(task.generation_attempt_no or 1)
generation_mode = str(task.generation_mode or "")
status = task.status
stage = task.pipeline_stage
registry_id = _chat_registry_id(task)
await db.rollback()
await remove_download_active(registry_id)
await log_task_event(
owner_type=GenerationOwnerType.CHAT_GENERATION_TASK.value,
owner_id=task_id,
task_id=task_id,
generation_attempt_no=generation_attempt_no,
generation_mode=generation_mode,
event_type=ChatGenerationTaskEventType.DOWNLOAD_SKIP_NOT_GENERATING.value,
message=f"{source} 下载恢复跳过:任务不是 generating",
detail={"status": status, "stage": stage},
)
return "clean_not_generating"
if not task.remote_result_url:
task_id = str(task.id)
generation_attempt_no = int(task.generation_attempt_no or 1)
generation_mode = str(task.generation_mode or "")
status = task.status
stage = task.pipeline_stage
await db.rollback()
await log_task_event(
owner_type=GenerationOwnerType.CHAT_GENERATION_TASK.value,
owner_id=task_id,
task_id=task_id,
generation_attempt_no=generation_attempt_no,
generation_mode=generation_mode,
event_type=ChatGenerationTaskEventType.DOWNLOAD_SKIP_NO_REMOTE_RESULT_URL.value,
message=f"{source} 下载恢复跳过:缺少 remote_result_url",
detail={"status": status, "stage": stage},
)
return "skip_no_remote_result_url"
stage = task.pipeline_stage
redis_payload = payload or {}
if stage == ChatGenerationPipelineStage.RESULT_READY.value:
await enqueue_download_task(
db,
task,
recover=True,
reason=f"{source}_result_ready",
)
await log_task_event(
task,
event_type=ChatGenerationTaskEventType.DOWNLOAD_RECOVERY_ENQUEUE.value,
message=f"{source} 发现 result_ready 未完成下载,启动时恢复投递下载任务",
detail={"payload": redis_payload},
)
return "recover_result_ready"
if stage == DOWNLOAD_STAGE_QUEUED:
if _is_queue_timeout(task, current_time):
await enqueue_download_task(
db,
task,
recover=True,
reason=f"{source}_download_queued_timeout",
)
await log_task_event(
task,
event_type=ChatGenerationTaskEventType.DOWNLOAD_RECOVERY_ENQUEUE.value,
message=f"{source} 发现 download_queued 长时间未消费,启动时恢复投递下载任务",
detail={"payload": redis_payload},
)
return "recover_queued_timeout"
await postpone_download_active_check(
record_id=_chat_registry_id(task),
payload=payload,
check_at=_queue_timeout_at(task, current_time),
)
return "skip_queued_not_timeout"
if stage == DOWNLOAD_STAGE_DOWNLOADING:
if _is_expired(task.download_lease_until, current_time):
await enqueue_download_task(
db,
task,
recover=True,
reason=f"{source}_downloading_lease_expired",
)
await log_task_event(
task,
event_type=ChatGenerationTaskEventType.DOWNLOAD_RECOVERY_ENQUEUE.value,
message=f"{source} 发现 downloading lease 过期,启动时恢复投递下载任务",
detail={"payload": redis_payload},
)
return "recover_downloading_expired"
await postpone_download_active_check(
record_id=_chat_registry_id(task),
payload=payload,
check_at=task.download_lease_until,
)
return "skip_downloading_alive"
if stage == DOWNLOAD_STAGE_RETRY_WAITING:
if _is_expired(task.download_next_retry_at, current_time):
await enqueue_download_task(
db,
task,
recover=True,
reason=f"{source}_retry_waiting_due",
)
await log_task_event(
task,
event_type=ChatGenerationTaskEventType.DOWNLOAD_RECOVERY_ENQUEUE.value,
message=f"{source} 发现 retry_waiting 到期,启动时恢复投递下载任务",
detail={"payload": redis_payload},
)
return "recover_retry_due"
await postpone_download_active_check(
record_id=_chat_registry_id(task),
payload=payload,
check_at=task.download_next_retry_at,
)
return "skip_retry_waiting_not_due"
return f"skip_stage_{stage}"
async def recover_download_tasks_once(db: AsyncSession) -> dict[str, Any]:
"""启动时下载容灾扫描。
先按 Redis active_index 找到到期下载任务;Redis 不可用或索引丢失时,
再通过 DB fallback 扫描 result_ready/download_* 状态,避免任务永久卡住。
"""
checked_ids: set[str] = set()
results: dict[str, int] = {}
due_ids = await get_due_download_record_ids(
limit=settings.DOWNLOAD_RECOVERY_BATCH_SIZE,
)
payloads = await get_download_active_payloads(due_ids)
download_refs = {
registry_item_id: parse_redis_owner_item_id(registry_item_id)
for registry_item_id in due_ids
}
for registry_item_id, ref in download_refs.items():
if ref.owner_type != GenerationOwnerType.CHAT_GENERATION_TASK.value:
continue
task = await _load_chat_task_for_update(db, ref.owner_id)
if task is None:
await db.rollback()
await remove_download_active(registry_item_id)
action = "clean_missing_task"
elif (
ref.generation_attempt_no is not None
and int(task.generation_attempt_no or 1)
!= int(ref.generation_attempt_no)
):
await db.rollback()
await remove_download_active(registry_item_id)
action = "clean_stale_download_attempt"
else:
current_registry_id = _chat_registry_id(task)
if registry_item_id != current_registry_id:
await remove_download_active(registry_item_id)
checked_ids.add(task.id)
action = await recover_one_download_task(
db,
task,
payload=payloads.get(registry_item_id),
source="startup_redis",
)
results[action] = results.get(action, 0) + 1
# DB fallback:不依赖 Redis active 注册表。
fallback_result = await db.execute(
select(ChatGenerationTask.id)
.where(
ChatGenerationTask.deleted_at.is_(None),
ChatGenerationTask.generation_mode.in_(list(ALLOWED_GENERATION_MODES)),
ChatGenerationTask.status == ChatGenerationTaskStatus.GENERATING.value,
ChatGenerationTask.remote_result_url.is_not(None),
ChatGenerationTask.pipeline_stage.in_(
[
ChatGenerationPipelineStage.RESULT_READY.value,
ChatGenerationPipelineStage.DOWNLOAD_QUEUED.value,
ChatGenerationPipelineStage.DOWNLOADING.value,
ChatGenerationPipelineStage.RETRY_WAITING.value,
]
),
)
.order_by(ChatGenerationTask.updated_at.asc())
.limit(int(settings.DOWNLOAD_RECOVERY_BATCH_SIZE or 100))
)
fallback_ids = [str(value) for value in fallback_result.scalars().all()]
for task_id in fallback_ids:
if task_id in checked_ids:
continue
task = await _load_chat_task_for_update(db, task_id)
if task is None:
await db.rollback()
continue
action = await recover_one_download_task(
db,
task,
payload=None,
source="startup_db",
)
results[action] = results.get(action, 0) + 1
checked_ids.add(task_id)
return {"checked": len(checked_ids), "results": results}
async def _mark_timeout(
db: AsyncSession,
task: ChatGenerationTask,
*,
error_message: str = "任务超时",
) -> str:
await mark_chat_generation_task_failed_and_refund_once(
db,
task=task,
error_message=error_message,
pipeline_stage=ChatGenerationPipelineStage.TIMEOUT.value,
)
await db.commit()
await notify_owner_finished(db, task)
await _remove_poll_active(_chat_registry_id(task))
await log_task_event(
task,
event_type=ChatGenerationTaskEventType.TASK_TIMEOUT.value,
to_status="failed",
to_stage=ChatGenerationPipelineStage.TIMEOUT.value,
)
return "mark_timeout"
async def _mark_failed(
db: AsyncSession,
task: ChatGenerationTask,
*,
error_message: str,
event_type: str = "POLL_FAILED",
detail: Any = None,
) -> str:
await mark_chat_generation_task_failed_and_refund_once(
db,
task=task,
error_message=error_message,
pipeline_stage=ChatGenerationPipelineStage.FAILED.value,
)
await db.commit()
await notify_owner_finished(db, task)
await _remove_poll_active(_chat_registry_id(task))
await log_task_event(task, event_type=event_type, message=task.error_message, detail=detail)
return "mark_failed"
async def recover_one_generation_task(
db: AsyncSession,
task: ChatGenerationTask,
*,
payload: dict[str, Any] | None = None,
source: str = "startup_db",
) -> str:
"""恢复单个生成任务。
分流原则:
1. 已有 remote_result_url:只恢复下载,不 poll,不重新 create。
2. 已有 provider_task_id/seedance_task_id:恢复 poll。
3. 无结果 URL、无供应商任务 IDdeadline 未过才恢复 create。
4. 无结果 URL、无供应商任务 ID:deadline 已过直接超时失败,不再补救生成。
"""
from app.tasks.generation_create_tasks import chatapi_create_generation_task
from app.tasks.generation_download_tasks import enqueue_download_task
from app.tasks.generation_poll_tasks import poll_generation_task, register_poll_active
current_time = _now()
redis_payload = payload or {}
if not task:
return "skip_missing_task"
if task.generation_mode not in ALLOWED_GENERATION_MODES:
await _remove_poll_active(_chat_registry_id(task))
return "clean_invalid_mode"
if _is_final_task_state(task):
await _remove_poll_active(_chat_registry_id(task))
return "clean_final_state"
if task.status != ChatGenerationTaskStatus.GENERATING.value:
await _remove_poll_active(_chat_registry_id(task))
return "clean_not_generating"
if bool(getattr(task, "video_upscale_enabled_snapshot", False)) and str(task.pipeline_stage or "").startswith("upscale_"):
# 原视频已经进入超分流水线,后续由 video_upscale 恢复扫描处理。
# 这里禁止再次投递原结果下载,避免覆盖保留的 source.mp4 或提前生成用户资源。
await _remove_poll_active(_chat_registry_id(task))
return "delegate_video_upscale_recovery"
has_remote_result = bool(str(task.remote_result_url or "").strip())
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))
# 最高优先级:只要远程结果 URL 已经落库,说明生成侧已经成功。
# 不管当前 pipeline_stage 是 queued/creating/waiting/result_ready/download_*,恢复时都不能重复 create 或 poll。
if has_remote_result:
await _remove_poll_active(_chat_registry_id(task))
await enqueue_download_task(
db,
task,
recover=True,
reason=f"{source}_has_remote_result_url",
)
await log_task_event(
task,
event_type=ChatGenerationTaskEventType.GENERATION_RECOVERY_ENQUEUE.value,
message=f"{source} 发现任务已存在 remote_result_url,恢复投递下载队列",
detail={
"pipeline_stage": task.pipeline_stage,
"payload": redis_payload,
"deadline_expired": is_deadline_expired,
},
)
return "recover_download_has_remote_result"
# 已经过 deadline 且没有结果 URL
# - 有供应商任务 ID:交给 poll worker 做最后一次状态确认;
# - 没有供应商任务 ID:说明没有可查询的远程任务,直接按超时失败处理,不再重新 create。
if is_deadline_expired:
if has_provider_task_id:
task.pipeline_stage = ChatGenerationPipelineStage.WAITING_REMOTE.value
await db.commit()
await log_task_event(
task,
event_type=ChatGenerationTaskEventType.GENERATION_RECOVERY_ENQUEUE.value,
message=f"{source} 发现任务已到 deadline 且存在供应商任务ID,投递 poll 队列做最终查询",
detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
)
poll_generation_task.apply_async(
args=[task.id],
kwargs={"force_due": True, "owner_type": GenerationOwnerType.CHAT_GENERATION_TASK.value, "generation_attempt_no": int(task.generation_attempt_no or 1)},
queue=POLL_QUEUE,
countdown=0,
)
await register_poll_active(
task,
check_at=_poll_queue_timeout_at(),
reason=f"{source}_deadline_final_poll",
)
return "recover_deadline_final_poll"
return await _mark_timeout(
db,
task,
error_message=f"{source} 发现任务已到 deadline,且没有远程结果或供应商任务ID",
)
# 未过 deadline:有供应商任务 ID 才允许恢复到 poll 队列。
# 视频任务如果 next_poll_at 未到期,不提前 poll,只刷新 active 注册表等待 Beat dispatcher 到期投递。
if has_provider_task_id:
if is_video_generation_task(task):
ensure_video_poll_fields(task, now=current_time)
if is_poll_not_due(task, now=current_time):
await db.commit()
await register_poll_active(
task,
check_at=task.next_poll_at,
next_poll_at=task.next_poll_at,
reason=f"{source}_video_poll_not_due",
)
await log_task_event(
task,
event_type=ChatGenerationTaskEventType.POLL_SKIP_NOT_DUE.value,
message=f"{source} 发现视频任务尚未到下一次轮询时间,启动容灾不提前投递 poll",
detail={
"pipeline_stage": task.pipeline_stage,
"payload": redis_payload,
"next_poll_at": task.next_poll_at,
},
)
return "skip_video_poll_not_due"
original_next_poll_at = ensure_aware_utc(task.next_poll_at)
queue_hold_until = _poll_queue_timeout_at(current_time)
task.pipeline_stage = ChatGenerationPipelineStage.WAITING_REMOTE.value
# 这里仍复用 next_poll_at 做短暂队列保护,避免启动容灾重复投递。
# 真正消费时通过 force_due=True 跳过“未到期”校验,避免保护时间反向阻塞本次 poll。
task.next_poll_at = queue_hold_until
await db.commit()
await log_task_event(
task,
event_type=ChatGenerationTaskEventType.GENERATION_RECOVERY_ENQUEUE.value,
message=f"{source} 发现任务存在供应商任务ID,恢复投递轮询队列",
detail={
"pipeline_stage": task.pipeline_stage,
"payload": redis_payload,
"due_next_poll_at": original_next_poll_at,
"queue_hold_until": queue_hold_until,
},
)
poll_generation_task.apply_async(
args=[task.id],
kwargs={"force_due": True, "owner_type": GenerationOwnerType.CHAT_GENERATION_TASK.value, "generation_attempt_no": int(task.generation_attempt_no or 1)},
queue=POLL_QUEUE,
countdown=0,
)
await register_poll_active(
task,
check_at=task.next_poll_at,
next_poll_at=task.next_poll_at,
reason=f"{source}_has_provider_task_id",
)
return "recover_poll_has_provider_id"
# 未过 deadline,且没有结果 URL / 供应商任务 ID:
# 图片同步任务会重新进入 submit_image_task;视频/其它任务会重新创建供应商任务。
# 这里不能投 poll,因为没有 provider_task_id/seedance_task_id 可查询。
recoverable_create_stages = {
ChatGenerationPipelineStage.QUEUED.value,
ChatGenerationPipelineStage.PREPARING.value,
ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value,
ChatGenerationPipelineStage.WAITING_REMOTE.value,
ChatGenerationPipelineStage.POLLING.value,
}
if task.pipeline_stage in recoverable_create_stages:
if task.pipeline_stage not in (
ChatGenerationPipelineStage.QUEUED.value,
ChatGenerationPipelineStage.PREPARING.value,
ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value,
):
task.pipeline_stage = ChatGenerationPipelineStage.QUEUED.value
# 刷新更新时间形成创建队列保护窗口,避免 Beat 在任务尚未消费时每轮重复补投。
task.updated_at = current_time
# Release the recovery row lock before writing an event through the
# independent logging session or talking to the broker.
await db.commit()
await _remove_poll_active(_chat_registry_id(task))
await log_task_event(
task,
event_type=ChatGenerationTaskEventType.GENERATION_RECOVERY_ENQUEUE.value,
message=f"{source} 发现任务未超时且缺少 remote_result_url/供应商任务ID,恢复投递创建队列",
detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
)
chatapi_create_generation_task.apply_async(
args=[task.id],
kwargs={"owner_type": GenerationOwnerType.CHAT_GENERATION_TASK.value, "generation_attempt_no": int(task.generation_attempt_no or 1)},
queue=CeleryQueue.GEN_CHATAPI_CREATE.value,
countdown=0,
task_id=(
f"generation-create:{GenerationOwnerType.CHAT_GENERATION_TASK.value}:"
f"{task.id}:attempt:{int(task.generation_attempt_no or 1)}"
),
)
return "recover_create_no_remote_no_provider_before_deadline"
# result_ready 但没有 URL 是脏状态;未过 deadline 时回创建队列重新处理,过期上面已标记超时。
if task.pipeline_stage == ChatGenerationPipelineStage.RESULT_READY.value:
task.pipeline_stage = ChatGenerationPipelineStage.QUEUED.value
task.updated_at = current_time
await db.commit()
await _remove_poll_active(_chat_registry_id(task))
await log_task_event(
task,
event_type=ChatGenerationTaskEventType.GENERATION_RECOVERY_ENQUEUE.value,
message=f"{source} 发现 result_ready 但缺少 remote_result_url,未超时,恢复投递创建队列",
detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
)
chatapi_create_generation_task.apply_async(
args=[task.id],
kwargs={"owner_type": GenerationOwnerType.CHAT_GENERATION_TASK.value, "generation_attempt_no": int(task.generation_attempt_no or 1)},
queue=CeleryQueue.GEN_CHATAPI_CREATE.value,
countdown=0,
task_id=(
f"generation-create:{GenerationOwnerType.CHAT_GENERATION_TASK.value}:"
f"{task.id}:attempt:{int(task.generation_attempt_no or 1)}"
),
)
return "recover_create_result_ready_no_url_before_deadline"
return f"skip_stage_{task.pipeline_stage}"
async def recover_generation_tasks_once(db: AsyncSession) -> dict[str, Any]:
"""启动时生成链路容灾扫描。
启动容灾由 worker_ready 触发,只跑一次完整恢复;周期性视频到期轮询由 Celery Beat 调度 dispatch_due_poll_tasks。
恢复顺序:
1. Redis poll active_index 到期任务;
2. DB fallback 扫描 queued/creating/waiting_remote/polling/result_ready
3. 下载阶段仍由 recover_download_tasks_once 兜底。
"""
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
due_poll_ids = await redis_get_due_registry_ids(
zset_key=settings.POLL_ACTIVE_REDIS_ZSET_KEY,
limit=int(settings.POLL_RECOVERY_BATCH_SIZE or settings.GENERATION_RECOVERY_BATCH_SIZE or 100),
log_context="poll_active",
)
poll_payloads = await redis_get_registry_payloads(
hash_key=settings.POLL_ACTIVE_REDIS_HASH_KEY,
item_ids=due_poll_ids,
log_context="poll_active",
)
poll_refs = {
registry_item_id: parse_redis_owner_item_id(registry_item_id)
for registry_item_id in due_poll_ids
}
for registry_item_id, ref in poll_refs.items():
if ref.owner_type != GenerationOwnerType.CHAT_GENERATION_TASK.value:
continue
task = await _load_chat_task_for_update(db, ref.owner_id)
if task is None:
await db.rollback()
await _remove_poll_active(registry_item_id)
action = "clean_missing_poll_task"
elif (
ref.generation_attempt_no is not None
and int(task.generation_attempt_no or 1) != int(ref.generation_attempt_no)
):
await db.rollback()
await _remove_poll_active(registry_item_id)
action = "clean_stale_poll_attempt"
else:
current_registry_id = _chat_registry_id(task)
if registry_item_id != current_registry_id:
await _remove_poll_active(registry_item_id)
checked_ids.add(task.id)
action = await recover_one_generation_task(
db,
task,
payload=poll_payloads.get(registry_item_id),
source="startup_poll_redis",
)
results[action] = results.get(action, 0) + 1
batch_size = int(settings.GENERATION_RECOVERY_BATCH_SIZE or settings.DOWNLOAD_RECOVERY_BATCH_SIZE or 100)
max_rounds = max(1, int(settings.GENERATION_RECOVERY_MAX_ROUNDS or 1))
total_db_checked = 0
for _round in range(max_rounds):
query_result = await db.execute(
select(ChatGenerationTask.id)
.where(
ChatGenerationTask.deleted_at.is_(None),
ChatGenerationTask.generation_mode.in_(list(ALLOWED_GENERATION_MODES)),
ChatGenerationTask.status == ChatGenerationTaskStatus.GENERATING.value,
ChatGenerationTask.pipeline_stage.in_(
[
"queued",
"preparing",
"creating_provider_task",
"waiting_remote",
"polling",
"result_ready",
]
),
)
.order_by(ChatGenerationTask.updated_at.asc())
.limit(batch_size)
)
task_ids = [str(value) for value in query_result.scalars().all()]
if not task_ids:
break
progressed_this_round = 0
for task_id in task_ids:
if task_id in checked_ids:
continue
task = await _load_chat_task_for_update(db, task_id)
if task is None:
await db.rollback()
continue
action = await recover_one_generation_task(
db,
task,
payload=None,
source="startup_db",
)
results[action] = results.get(action, 0) + 1
checked_ids.add(task_id)
total_db_checked += 1
progressed_this_round += 1
if len(task_ids) < batch_size or progressed_this_round <= 0:
break
# 子任务可能在 worker 中断前已进入终态但主任务尚未汇总,按稳定游标完整重算全部主任务。
from app.services.generation.ai.task_group_service import aggregate_main_task_status
reconciled = 0
main_cursor: str | None = None
while True:
main_query = select(ChatGenerationTask.id).where(
ChatGenerationTask.deleted_at.is_(None),
ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_MAIN.value,
)
if main_cursor:
main_query = main_query.where(ChatGenerationTask.id > main_cursor)
main_result = await db.execute(main_query.order_by(ChatGenerationTask.id.asc()).limit(batch_size))
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
if len(parent_ids) < batch_size:
break
if reconciled:
results["reconcile_main"] = reconciled
return {
"checked": len(checked_ids),
"db_checked": total_db_checked,
"results": results,
}
async def recover_stale_create_tasks_once(db: AsyncSession) -> dict[str, Any]:
"""轻量恢复长时间未消费的创建阶段 ChatGenerationTask。
只扫描 queued/preparing/creating_provider_task,避免周期任务重复执行完整
provider poll、下载和主任务汇总逻辑。
"""
current_time = _now()
cutoff = current_time - timedelta(
seconds=max(1, int(settings.GENERATION_CREATE_QUEUE_TIMEOUT_SECONDS or 300))
)
lease_expired_at = current_time
batch_size = max(1, int(settings.GENERATION_RECOVERY_BATCH_SIZE or 20))
result = await db.execute(
select(ChatGenerationTask.id)
.where(
ChatGenerationTask.deleted_at.is_(None),
ChatGenerationTask.generation_mode.in_(list(ALLOWED_GENERATION_MODES)),
ChatGenerationTask.status == ChatGenerationTaskStatus.GENERATING.value,
ChatGenerationTask.pipeline_stage.in_(
[
ChatGenerationPipelineStage.QUEUED.value,
ChatGenerationPipelineStage.PREPARING.value,
ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value,
]
),
ChatGenerationTask.remote_result_url.is_(None),
ChatGenerationTask.provider_task_id.is_(None),
ChatGenerationTask.seedance_task_id.is_(None),
ChatGenerationTask.updated_at <= cutoff,
(
ChatGenerationTask.provider_create_lease_until.is_(None)
| (ChatGenerationTask.provider_create_lease_until <= lease_expired_at)
),
)
.order_by(ChatGenerationTask.updated_at.asc(), ChatGenerationTask.id.asc())
.limit(batch_size)
)
task_ids = [str(value) for value in result.scalars().all()]
counts: dict[str, int] = {}
for task_id in task_ids:
task = await _load_chat_task_for_update(db, task_id)
if task is None:
await db.rollback()
action = "skip_missing_task"
else:
action = await recover_one_generation_task(
db,
task,
payload=None,
source="periodic_create_recovery",
)
counts[action] = counts.get(action, 0) + 1
return {"checked": len(task_ids), "results": counts}
async def dispatch_due_poll_tasks_once(db: AsyncSession) -> dict[str, Any]:
"""周期性轻量到期轮询调度。
只处理视频任务的 next_poll_at 到期记录,不替代启动容灾 recover_generation_tasks_once。
Beat 每分钟触发本任务,本任务把真实供应商轮询投递到 gen_provider_poll 队列。
"""
from app.tasks.generation_poll_tasks import poll_generation_task, register_poll_active
current_time = _now()
batch_size = max(1, int(settings.POLL_DUE_DISPATCH_BATCH_SIZE or 100))
poll_lease_expired_at = current_time - timedelta(seconds=int(settings.POLL_TASK_LEASE_SECONDS or 300))
await apply_short_lock_timeout(db)
query_result = await db.execute(
select(ChatGenerationTask)
.where(
ChatGenerationTask.deleted_at.is_(None),
ChatGenerationTask.generation_mode.in_(list(ALLOWED_GENERATION_MODES)),
ChatGenerationTask.status == ChatGenerationTaskStatus.GENERATING.value,
ChatGenerationTask.gen_type == GenerationType.VIDEO.value,
ChatGenerationTask.next_poll_at.is_not(None),
ChatGenerationTask.next_poll_at <= current_time,
ChatGenerationTask.pipeline_stage.in_(
[
ChatGenerationPipelineStage.WAITING_REMOTE.value,
ChatGenerationPipelineStage.POLLING.value,
]
),
)
.order_by(ChatGenerationTask.next_poll_at.asc(), ChatGenerationTask.updated_at.asc())
.limit(batch_size)
.with_for_update(skip_locked=True)
)
tasks = query_result.scalars().all()
results: dict[str, int] = {}
dispatched_task_ids: list[str] = []
dispatched_due_next_poll_at_by_id: dict[str, datetime | None] = {}
dispatched_queue_hold_until_by_id: dict[str, datetime] = {}
post_commit_logs: list[dict[str, Any]] = []
for task in tasks:
action = "skip_unknown"
try:
if task.pipeline_stage == ChatGenerationPipelineStage.POLLING.value:
last_poll_at = ensure_aware_utc(task.last_poll_at)
if last_poll_at and last_poll_at > poll_lease_expired_at:
await register_poll_active(
task,
check_at=_poll_queue_timeout_at(last_poll_at),
next_poll_at=task.next_poll_at,
reason="due_dispatch_polling_lease_alive",
)
action = "skip_polling_lease_alive"
continue
if not (task.seedance_task_id or task.provider_task_id):
# Defer FK-backed event logging until the batch row locks are released.
post_commit_logs.append({
"task_id": str(task.id),
"generation_attempt_no": int(task.generation_attempt_no or 1),
"event_type": ChatGenerationTaskEventType.POLL_DISPATCH_SKIP.value,
"message": "视频到期轮询调度跳过:缺少外部任务ID",
"detail": {"pipeline_stage": task.pipeline_stage, "next_poll_at": task.next_poll_at},
})
action = "skip_no_provider_task_id"
continue
original_next_poll_at = ensure_aware_utc(task.next_poll_at)
queue_hold_until = _poll_queue_timeout_at(current_time)
task.pipeline_stage = ChatGenerationPipelineStage.WAITING_REMOTE.value
# 设置一个队列消费保护时间,避免 Beat 下一分钟看到旧 next_poll_at 又重复投递。
# poll worker 会通过 force_due=True 消费本次到期任务,避免该保护时间被误判为业务未到期。
task.next_poll_at = queue_hold_until
dispatched_due_next_poll_at_by_id[task.id] = original_next_poll_at
dispatched_queue_hold_until_by_id[task.id] = queue_hold_until
dispatched_task_ids.append(task.id)
action = "dispatch_poll"
except Exception as exc:
logger.exception("视频到期轮询调度单条处理失败。task_id=%s", getattr(task, "id", None))
action = "error"
post_commit_logs.append({
"task_id": str(getattr(task, "id", "") or ""),
"generation_attempt_no": int(getattr(task, "generation_attempt_no", 1) or 1),
"event_type": ChatGenerationTaskEventType.POLL_DISPATCH_SKIP.value,
"message": f"视频到期轮询调度单条处理失败:{exc}",
"detail": None,
})
finally:
results[action] = results.get(action, 0) + 1
await db.commit()
for item in post_commit_logs:
if item["task_id"]:
await log_task_event(
task_id=item["task_id"],
generation_attempt_no=item["generation_attempt_no"],
event_type=item["event_type"],
message=item["message"],
detail=item["detail"],
)
fresh_tasks = []
if dispatched_task_ids:
fresh_result = await db.execute(
select(ChatGenerationTask)
.where(
ChatGenerationTask.id.in_(dispatched_task_ids),
ChatGenerationTask.deleted_at.is_(None),
)
.execution_options(populate_existing=True)
)
fresh_tasks = fresh_result.scalars().all()
enqueued_count = 0
for task in fresh_tasks:
queue_hold_until = dispatched_queue_hold_until_by_id.get(task.id) or ensure_aware_utc(task.next_poll_at) or _poll_queue_timeout_at(current_time)
due_next_poll_at = dispatched_due_next_poll_at_by_id.get(task.id)
await register_poll_active(
task,
check_at=queue_hold_until,
next_poll_at=queue_hold_until,
reason="due_dispatch_poll_queued",
)
await log_task_event(
task,
event_type=ChatGenerationTaskEventType.POLL_DISPATCH_DUE.value,
message="视频 next_poll_at 到期,已投递 provider poll 队列",
detail={
"due_next_poll_at": due_next_poll_at,
"queue_hold_until": queue_hold_until,
"queue": POLL_QUEUE,
},
)
poll_generation_task.apply_async(
args=[task.id],
kwargs={"force_due": True, "owner_type": GenerationOwnerType.CHAT_GENERATION_TASK.value, "generation_attempt_no": int(task.generation_attempt_no or 1)},
queue=POLL_QUEUE,
countdown=0,
)
enqueued_count += 1
# 如果 log_task_event 内部不 commit,这里要提交一次
if fresh_tasks:
await db.commit()
return {
"checked": len(tasks),
"dispatched": len(dispatched_task_ids),
"enqueued": enqueued_count,
"results": results,
}