633 lines
23 KiB
Python
633 lines
23 KiB
Python
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.generation_task import (
|
||
ALLOWED_GENERATION_MODES,
|
||
ChatGenerationPipelineStage,
|
||
ChatGenerationTaskEventType,
|
||
ChatGenerationTaskStatus,
|
||
)
|
||
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_module_hook_service import notify_chat_generation_task_finished
|
||
from app.services.generation_refund_service import mark_chat_generation_task_failed_and_refund_once
|
||
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 = "gen_provider_poll"
|
||
|
||
|
||
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(task.id)
|
||
return "clean_invalid_mode"
|
||
if _is_final_task_state(task):
|
||
await remove_download_active(task.id)
|
||
return "clean_final_state"
|
||
if task.status != ChatGenerationTaskStatus.GENERATING.value:
|
||
await remove_download_active(task.id)
|
||
await log_task_event(
|
||
task,
|
||
event_type=ChatGenerationTaskEventType.DOWNLOAD_SKIP_NOT_GENERATING.value,
|
||
message=f"{source} 下载恢复跳过：任务不是 generating",
|
||
detail={"status": task.status, "stage": task.pipeline_stage},
|
||
)
|
||
return "clean_not_generating"
|
||
if not task.remote_result_url:
|
||
await log_task_event(
|
||
task,
|
||
event_type=ChatGenerationTaskEventType.DOWNLOAD_SKIP_NO_REMOTE_RESULT_URL.value,
|
||
message=f"{source} 下载恢复跳过：缺少 remote_result_url",
|
||
detail={"status": task.status, "stage": task.pipeline_stage},
|
||
)
|
||
return "skip_no_remote_result_url"
|
||
|
||
stage = task.pipeline_stage
|
||
redis_payload = payload or {}
|
||
|
||
if stage == ChatGenerationPipelineStage.RESULT_READY.value:
|
||
await log_task_event(
|
||
task,
|
||
event_type=ChatGenerationTaskEventType.DOWNLOAD_RECOVERY_ENQUEUE.value,
|
||
message=f"{source} 发现 result_ready 未完成下载，启动时恢复投递下载任务",
|
||
detail={"payload": redis_payload},
|
||
)
|
||
await enqueue_download_task(
|
||
db,
|
||
task,
|
||
recover=True,
|
||
reason=f"{source}_result_ready",
|
||
)
|
||
return "recover_result_ready"
|
||
|
||
if stage == DOWNLOAD_STAGE_QUEUED:
|
||
if _is_queue_timeout(task, current_time):
|
||
await log_task_event(
|
||
task,
|
||
event_type=ChatGenerationTaskEventType.DOWNLOAD_RECOVERY_ENQUEUE.value,
|
||
message=f"{source} 发现 download_queued 长时间未消费，启动时恢复投递下载任务",
|
||
detail={"payload": redis_payload},
|
||
)
|
||
await enqueue_download_task(
|
||
db,
|
||
task,
|
||
recover=True,
|
||
reason=f"{source}_download_queued_timeout",
|
||
)
|
||
return "recover_queued_timeout"
|
||
|
||
await postpone_download_active_check(
|
||
record_id=task.id,
|
||
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 log_task_event(
|
||
task,
|
||
event_type=ChatGenerationTaskEventType.DOWNLOAD_RECOVERY_ENQUEUE.value,
|
||
message=f"{source} 发现 downloading lease 过期，启动时恢复投递下载任务",
|
||
detail={"payload": redis_payload},
|
||
)
|
||
await enqueue_download_task(
|
||
db,
|
||
task,
|
||
recover=True,
|
||
reason=f"{source}_downloading_lease_expired",
|
||
)
|
||
return "recover_downloading_expired"
|
||
|
||
await postpone_download_active_check(
|
||
record_id=task.id,
|
||
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 log_task_event(
|
||
task,
|
||
event_type=ChatGenerationTaskEventType.DOWNLOAD_RECOVERY_ENQUEUE.value,
|
||
message=f"{source} 发现 retry_waiting 到期，启动时恢复投递下载任务",
|
||
detail={"payload": redis_payload},
|
||
)
|
||
await enqueue_download_task(
|
||
db,
|
||
task,
|
||
recover=True,
|
||
reason=f"{source}_retry_waiting_due",
|
||
)
|
||
return "recover_retry_due"
|
||
|
||
await postpone_download_active_check(
|
||
record_id=task.id,
|
||
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)
|
||
|
||
for task_id in due_ids:
|
||
result = await db.execute(
|
||
select(ChatGenerationTask)
|
||
.where(
|
||
ChatGenerationTask.id == task_id,
|
||
ChatGenerationTask.deleted_at.is_(None),
|
||
)
|
||
.with_for_update()
|
||
.limit(1)
|
||
)
|
||
task = result.scalar_one_or_none()
|
||
if task is None:
|
||
await remove_download_active(task_id)
|
||
action = "clean_missing_task"
|
||
else:
|
||
checked_ids.add(task.id)
|
||
action = await recover_one_download_task(
|
||
db,
|
||
task,
|
||
payload=payloads.get(task_id),
|
||
source="startup_redis",
|
||
)
|
||
results[action] = results.get(action, 0) + 1
|
||
|
||
# DB fallback：不依赖 Redis active 注册表。
|
||
fallback_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.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))
|
||
.with_for_update(skip_locked=True)
|
||
)
|
||
fallback_tasks = fallback_result.scalars().all()
|
||
|
||
for task in fallback_tasks:
|
||
if task.id in checked_ids:
|
||
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 notify_chat_generation_task_finished(db, task)
|
||
await db.commit()
|
||
await _remove_poll_active(task.id)
|
||
await log_task_event(
|
||
task,
|
||
event_type="TASK_TIMEOUT",
|
||
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 notify_chat_generation_task_finished(db, task)
|
||
await db.commit()
|
||
await _remove_poll_active(task.id)
|
||
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、无供应商任务 ID：deadline 未过才恢复 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(task.id)
|
||
return "clean_invalid_mode"
|
||
if _is_final_task_state(task):
|
||
await _remove_poll_active(task.id)
|
||
return "clean_final_state"
|
||
if task.status != ChatGenerationTaskStatus.GENERATING.value:
|
||
await _remove_poll_active(task.id)
|
||
return "clean_not_generating"
|
||
|
||
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(task.id)
|
||
await log_task_event(
|
||
task,
|
||
event_type="GENERATION_RECOVERY_ENQUEUE",
|
||
message=f"{source} 发现任务已存在 remote_result_url，恢复投递下载队列",
|
||
detail={
|
||
"pipeline_stage": task.pipeline_stage,
|
||
"payload": redis_payload,
|
||
"deadline_expired": is_deadline_expired,
|
||
},
|
||
)
|
||
await enqueue_download_task(
|
||
db,
|
||
task,
|
||
recover=True,
|
||
reason=f"{source}_has_remote_result_url",
|
||
)
|
||
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="GENERATION_RECOVERY_ENQUEUE",
|
||
message=f"{source} 发现任务已到 deadline 且存在供应商任务ID，投递 poll 队列做最终查询",
|
||
detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
|
||
)
|
||
poll_generation_task.apply_async(args=[task.id], 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"
|
||
|
||
await log_task_event(
|
||
task,
|
||
event_type="GENERATION_RECOVERY_TIMEOUT",
|
||
message=f"{source} 发现任务已到 deadline，且没有 remote_result_url/供应商任务ID，按超时失败处理",
|
||
detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
|
||
)
|
||
return await _mark_timeout(db, task)
|
||
|
||
# 未过 deadline：有供应商任务 ID 才允许恢复到 poll 队列。
|
||
if has_provider_task_id:
|
||
task.pipeline_stage = ChatGenerationPipelineStage.WAITING_REMOTE.value
|
||
await db.commit()
|
||
await log_task_event(
|
||
task,
|
||
event_type="GENERATION_RECOVERY_ENQUEUE",
|
||
message=f"{source} 发现任务存在供应商任务ID，恢复投递轮询队列",
|
||
detail={"pipeline_stage": task.pipeline_stage, "payload": redis_payload},
|
||
)
|
||
poll_generation_task.apply_async(args=[task.id], queue=POLL_QUEUE, countdown=0)
|
||
await register_poll_active(
|
||
task,
|
||
check_at=_poll_queue_timeout_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
|
||
await db.commit()
|
||
|
||
await _remove_poll_active(task.id)
|
||
await log_task_event(
|
||
task,
|
||
event_type="GENERATION_RECOVERY_ENQUEUE",
|
||
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],
|
||
queue="gen_chatapi_create",
|
||
countdown=0,
|
||
)
|
||
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
|
||
await db.commit()
|
||
await _remove_poll_active(task.id)
|
||
await log_task_event(
|
||
task,
|
||
event_type="GENERATION_RECOVERY_ENQUEUE",
|
||
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],
|
||
queue="gen_chatapi_create",
|
||
countdown=0,
|
||
)
|
||
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]:
|
||
"""启动时生成链路容灾扫描。
|
||
|
||
不新增 Celery beat，不新增 worker 命令；worker 启动时由 Redis 锁保证只投递一次。
|
||
恢复顺序：
|
||
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] = {}
|
||
|
||
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",
|
||
)
|
||
|
||
for task_id in due_poll_ids:
|
||
result = await db.execute(
|
||
select(ChatGenerationTask)
|
||
.where(
|
||
ChatGenerationTask.id == task_id,
|
||
ChatGenerationTask.deleted_at.is_(None),
|
||
)
|
||
.with_for_update()
|
||
.limit(1)
|
||
)
|
||
task = result.scalar_one_or_none()
|
||
if task is None:
|
||
await _remove_poll_active(task_id)
|
||
action = "clean_missing_poll_task"
|
||
else:
|
||
checked_ids.add(task.id)
|
||
action = await recover_one_generation_task(
|
||
db,
|
||
task,
|
||
payload=poll_payloads.get(task_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)
|
||
.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)
|
||
.with_for_update(skip_locked=True)
|
||
)
|
||
tasks = query_result.scalars().all()
|
||
if not tasks:
|
||
break
|
||
|
||
progressed_this_round = 0
|
||
for task in tasks:
|
||
if task.id in checked_ids:
|
||
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(tasks) < batch_size or progressed_this_round <= 0:
|
||
break
|
||
|
||
return {
|
||
"checked": len(checked_ids),
|
||
"db_checked": total_db_checked,
|
||
"results": results,
|
||
} |