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

拆镜复刻开发完成

Browse Source
This commit is contained in:
GinHa
2026-06-11 17:54:40 +08:00
parent a6b6e5f822
commit 74266a126c
34 changed files with 7105 additions and 470 deletions
Download Patch File Download Diff File Expand all files Collapse all files
video-gen-api/app/tasks/generation_poll_tasks.py
+170 -57
View File
@@ -1,6 +1,7 @@
from app.tasks.async_runner import run_async
import json
from datetime import datetime, timezone
from datetime import datetime, timedelta, timezone
from typing import Any
from sqlalchemy import select
@@ -11,9 +12,21 @@ from app.services.error_codes import extract_error_message
from app.services.generation_log_service import log_task_event, log_provider_call
from app.services.generation_refund_service import mark_chat_generation_task_failed_and_refund_once
from app.services.generation_provider_service import poll_provider_task
from app.services.redis_registry_service import (
datetime_to_epoch,
ensure_aware_utc,
redis_remove_registry_item,
redis_upsert_registry_item,
utc_now,
)
from app.tasks.celery_app import celery_app
ALLOWED_GENERATION_MODES = {"chatapi_async", "hot_opening_replicate"}
ALLOWED_GENERATION_MODES = {"chatapi_async", "hot_opening_replicate", "shot_replicate"}
POLL_QUEUE = "gen_provider_poll"
def _now() -> datetime:
return datetime.now(timezone.utc)
def _is_success(status: str) -> bool:
@@ -31,6 +44,86 @@ def _engine_snapshot(task: ChatGenerationTask) -> dict:
return {}
def _deadline_expired(task: ChatGenerationTask, now: datetime | None = None) -> bool:
deadline_at = ensure_aware_utc(task.deadline_at)
return bool(deadline_at and deadline_at <= (now or _now()))
def _poll_check_at(*, delay_seconds: int | float | None = None, now: datetime | None = None) -> datetime:
current_time = now or _now()
delay = int(delay_seconds or settings.CHATAPI_ASYNC_POLL_INTERVAL_SECONDS or 30)
grace = int(settings.POLL_TASK_QUEUE_TIMEOUT_SECONDS or 120)
return current_time + timedelta(seconds=max(1, delay) + max(0, grace))
def _poll_lease_until(now: datetime | None = None) -> datetime:
current_time = now or _now()
return current_time + timedelta(seconds=int(settings.POLL_TASK_LEASE_SECONDS or 300))
def _build_poll_active_payload(
task: ChatGenerationTask,
*,
stage: str,
reason: str,
next_poll_at: datetime | None = None,
check_at: datetime | None = None,
) -> dict[str, Any]:
current_time = utc_now()
checked_next_poll_at = ensure_aware_utc(next_poll_at)
checked_check_at = ensure_aware_utc(check_at)
return {
"task_id": task.id,
"provider_task_id": task.provider_task_id,
"seedance_task_id": task.seedance_task_id,
"generation_mode": task.generation_mode,
"gen_type": task.gen_type,
"stage": stage,
"queue": POLL_QUEUE,
"poll_count": int(task.poll_count or 0),
"retry_count": int(task.retry_count or 0),
"last_poll_at": datetime_to_epoch(task.last_poll_at) if task.last_poll_at else None,
"next_poll_at": datetime_to_epoch(checked_next_poll_at) if checked_next_poll_at else None,
"deadline_at": datetime_to_epoch(task.deadline_at) if task.deadline_at else None,
"check_at": datetime_to_epoch(checked_check_at) if checked_check_at else None,
"updated_at": datetime_to_epoch(current_time),
"reason": reason,
}
async def register_poll_active(
task: ChatGenerationTask,
*,
check_at: datetime,
reason: str,
next_poll_at: datetime | None = None,
) -> None:
payload = _build_poll_active_payload(
task,
stage=task.pipeline_stage or "",
reason=reason,
next_poll_at=next_poll_at,
check_at=check_at,
)
await redis_upsert_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,
log_context="poll_active",
)
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 _notify_finished(db, task: ChatGenerationTask) -> None:
from app.services.generation_module_hook_service import notify_chat_generation_task_finished
@@ -55,6 +148,32 @@ async def _reload_task(db, task_id: str) -> ChatGenerationTask | None:
return result.scalar_one_or_none()
async def _mark_timeout(db, task: ChatGenerationTask, *, message: str = "任务轮询超时") -> None:
await mark_chat_generation_task_failed_and_refund_once(
db,
task=task,
error_message=message,
pipeline_stage="timeout",
)
await _notify_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="timeout")
async def _mark_failed(db, task: ChatGenerationTask, *, message: str, detail: Any = None) -> None:
await mark_chat_generation_task_failed_and_refund_once(
db,
task=task,
error_message=message,
pipeline_stage="failed",
)
await _notify_finished(db, task)
await db.commit()
await remove_poll_active(task.id)
await log_task_event(task, event_type="POLL_FAILED", message=task.error_message, detail=detail)
async def _run(task_id: str):
async with async_session() as db:
result = await db.execute(select(ChatGenerationTask).where(
@@ -62,43 +181,37 @@ async def _run(task_id: str):
ChatGenerationTask.deleted_at.is_(None),
).with_for_update().limit(1))
task = result.scalar_one_or_none()
if not task or task.generation_mode not in ALLOWED_GENERATION_MODES:
if not task:
await remove_poll_active(task_id)
return
if task.generation_mode not in ALLOWED_GENERATION_MODES:
await remove_poll_active(task.id)
return
# 只处理正在生成，且处于远程等待/轮询中的任务。
if task.status != "generating" or task.pipeline_stage not in ("waiting_remote", "polling"):
await remove_poll_active(task.id)
return
if task.deadline_at and datetime.now(timezone.utc) > task.deadline_at:
await mark_chat_generation_task_failed_and_refund_once(
db,
task=task,
error_message="任务轮询超时",
pipeline_stage="timeout",
)
await _notify_finished(db, task)
await db.commit()
await log_task_event(task, event_type="TASK_TIMEOUT", to_status="failed", to_stage="timeout")
if _deadline_expired(task):
await _mark_timeout(db, task, message="任务轮询超时")
return
if not (task.seedance_task_id or task.provider_task_id):
await mark_chat_generation_task_failed_and_refund_once(
db,
task=task,
error_message="缺少外部任务ID",
pipeline_stage="failed",
)
await _notify_finished(db, task)
await db.commit()
await log_task_event(task, event_type="POLL_FAILED", message=task.error_message)
await _mark_failed(db, task, message="缺少外部任务ID")
return
# 标记本次正在轮询。
# 注意：pending 后会再改回 waiting_remote，避免任务长期卡在 polling。
# 标记本次正在轮询，并登记 poll lease。
# 如果 worker 在供应商接口调用过程中退出，启动恢复会在 lease 过期后重新投递。
task.pipeline_stage = "polling"
task.poll_count = (task.poll_count or 0) + 1
task.last_poll_at = datetime.now(timezone.utc)
task.last_poll_at = _now()
await db.commit()
await register_poll_active(
task,
check_at=_poll_lease_until(task.last_poll_at),
reason="polling_lease",
)
try:
poll_result = await poll_provider_task(db, task)
@@ -134,20 +247,13 @@ async def _run(task_id: str):
task.provider_response_json = response_data
if not task.remote_result_url:
await mark_chat_generation_task_failed_and_refund_once(
db,
task=task,
error_message="供应商任务成功但未返回结果URL",
pipeline_stage="failed",
)
await _notify_finished(db, task)
await db.commit()
await log_task_event(task, event_type="POLL_FAILED", message=task.error_message)
await _mark_failed(db, task, message="供应商任务成功但未返回结果URL", detail=poll_result)
return
task.pipeline_stage = "result_ready"
task.retry_count = 0
await db.commit()
await remove_poll_active(task.id)
await log_task_event(task, event_type="POLL_SUCCESS", to_stage="result_ready")
@@ -158,34 +264,38 @@ async def _run(task_id: str):
if _is_failed(status):
task.provider_response_json = response_data
await mark_chat_generation_task_failed_and_refund_once(
await _mark_failed(
db,
task=task,
error_message=poll_result.get("error") or f"供应商任务失败: {status}",
pipeline_stage="failed",
task,
message=poll_result.get("error") or f"供应商任务失败: {status}",
detail=poll_result,
)
await _notify_finished(db, task)
await db.commit()
await log_task_event(task, event_type="POLL_FAILED", message=task.error_message, detail=poll_result)
return
# 关键修改 1：
# 供应商仍在 pending / running 时，把阶段从 polling 改回 waiting_remote。
# 这样数据库状态表示“等待下一次轮询”，不会长期停在 polling。
# 同时可以降低重复 Celery 消息形成多条轮询链的概率。
# 同时登记下一次 poll active，Celery countdown 丢失时可由恢复任务拉起。
task.pipeline_stage = "waiting_remote"
task.retry_count = 0
await db.commit()
await log_task_event(task, event_type="POLL_PENDING", message=f"status={status}")
delay_seconds = int(settings.CHATAPI_ASYNC_POLL_INTERVAL_SECONDS or 30)
next_poll_at = _now() + timedelta(seconds=max(1, delay_seconds))
await register_poll_active(
task,
check_at=_poll_check_at(delay_seconds=delay_seconds),
next_poll_at=next_poll_at,
reason="poll_pending_next",
)
poll_generation_task.apply_async(
args=[task.id],
countdown=settings.CHATAPI_ASYNC_POLL_INTERVAL_SECONDS,
queue=POLL_QUEUE,
countdown=delay_seconds,
)
except Exception as exc:
# 关键修改 2：
# 异常后先 rollback，再重新查询 task，不继续使用 rollback 前的旧 ORM 对象。
try:
await db.rollback()
@@ -194,30 +304,33 @@ async def _run(task_id: str):
task = await _reload_task(db, task_id)
if not task:
await remove_poll_active(task_id)
return
task.retry_count = (task.retry_count or 0) + 1
if task.retry_count > settings.CHATAPI_ASYNC_MAX_RETRIES:
error_message = extract_error_message(exc, "轮询") if callable(extract_error_message) else str(exc)
await mark_chat_generation_task_failed_and_refund_once(
db,
task=task,
error_message=error_message,
pipeline_stage="failed",
)
await _notify_finished(db, task)
await db.commit()
await log_task_event(task, event_type="POLL_FAILED", message=task.error_message)
await _mark_failed(db, task, message=error_message)
else:
# 临时轮询异常时，不让任务停在 polling。
# 回到 waiting_remote，等待下一次重试轮询。
task.pipeline_stage = "waiting_remote"
await db.commit()
delay_seconds = int(settings.CHATAPI_ASYNC_RETRY_BACKOFF_SECONDS or 30) * int(task.retry_count or 1)
next_poll_at = _now() + timedelta(seconds=max(1, delay_seconds))
await register_poll_active(
task,
check_at=_poll_check_at(delay_seconds=delay_seconds),
next_poll_at=next_poll_at,
reason="poll_exception_retry",
)
poll_generation_task.apply_async(
args=[task.id],
countdown=settings.CHATAPI_ASYNC_RETRY_BACKOFF_SECONDS * task.retry_count,
queue=POLL_QUEUE,
countdown=delay_seconds,
)
@@ -233,4 +346,4 @@ else:
def apply_async(self, *args, **kwargs):
raise RuntimeError("Celery is disabled")
poll_generation_task = _DisabledTask()
poll_generation_task = _DisabledTask()