feat(generation): worker孤儿任务自动恢复 + 429限流结构化提示 (#1677) #1710
@@ -14,6 +14,7 @@ from app.core.task_enqueue import (
|
||||
USER_PENDING_LIMIT,
|
||||
GlobalQueueFull,
|
||||
UserPendingLimitExceeded,
|
||||
build_rate_limit_detail,
|
||||
safe_enqueue_generation_task,
|
||||
)
|
||||
from app.dependencies import (
|
||||
@@ -312,12 +313,12 @@ def create_preview_generation_task(
|
||||
except UserPendingLimitExceeded as e:
|
||||
raise HTTPException(
|
||||
status_code=429,
|
||||
detail=f"您的待处理任务过多(当前 {e.pending_count - count}/{e.limit},本次提交 {count} 个),请等待后再提交",
|
||||
detail=build_rate_limit_detail(e, generation_task_repository, scope="user"),
|
||||
) from e
|
||||
except GlobalQueueFull as e:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="系统繁忙,请稍后再试",
|
||||
detail=build_rate_limit_detail(e, generation_task_repository, scope="global"),
|
||||
) from e
|
||||
|
||||
# 确定视频比例:优先前端传入,否则从模板 mode 推断
|
||||
@@ -498,6 +499,7 @@ def create_preview_generation_task(
|
||||
|
||||
# ── 入队 ──
|
||||
responses: list[PreviewGenerationTaskResponse] = []
|
||||
rate_limit_exc: Exception | None = None # 记录首个限流异常,全部失败时返回结构化提示
|
||||
for variant_index, task in enumerate(created_tasks):
|
||||
try:
|
||||
enqueued = safe_enqueue_generation_task(
|
||||
@@ -510,23 +512,29 @@ def create_preview_generation_task(
|
||||
if not enqueued:
|
||||
logger.warning("[预览生成] 任务入队失败: task_id=%s", task.id)
|
||||
_mark_task_failed(generation_task_repository, task, "任务入队失败")
|
||||
except UserPendingLimitExceeded:
|
||||
except UserPendingLimitExceeded as e:
|
||||
_mark_task_failed(generation_task_repository, task, "待处理任务超限")
|
||||
except GlobalQueueFull:
|
||||
rate_limit_exc = rate_limit_exc or e
|
||||
except GlobalQueueFull as e:
|
||||
_mark_task_failed(generation_task_repository, task, "系统队列已满")
|
||||
rate_limit_exc = rate_limit_exc or e
|
||||
except Exception:
|
||||
logger.exception("[预览生成] 入队异常: task_id=%s", task.id)
|
||||
_mark_task_failed(generation_task_repository, task, "任务入队异常")
|
||||
# enqueue 会原地更新 task 状态/进度,直接用 task 构造响应
|
||||
responses.append(_to_preview_response(task))
|
||||
|
||||
# 队列满/限流时若全部失败,返回明确错误码
|
||||
if all(r.status == "failed" for r in responses):
|
||||
first_err = next((r.error_message for r in responses if r.error_message), "")
|
||||
if "待处理任务" in first_err:
|
||||
raise HTTPException(status_code=429, detail=first_err or "待处理任务超限")
|
||||
if "队列" in first_err:
|
||||
raise HTTPException(status_code=503, detail=first_err or "系统繁忙,请稍后再试")
|
||||
# 队列满/限流时若全部失败,返回结构化错误码(前端区分"排队"与"创建失败")
|
||||
if all(r.status == "failed" for r in responses) and rate_limit_exc is not None:
|
||||
if isinstance(rate_limit_exc, UserPendingLimitExceeded):
|
||||
raise HTTPException(
|
||||
status_code=429,
|
||||
detail=build_rate_limit_detail(rate_limit_exc, generation_task_repository, scope="user"),
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail=build_rate_limit_detail(rate_limit_exc, generation_task_repository, scope="global"),
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"[预览生成] 创建完成: %d 个变体任务, task_ids=%s",
|
||||
|
||||
@@ -10,6 +10,7 @@ from app.core.task_enqueue import (
|
||||
USER_PENDING_LIMIT,
|
||||
GlobalQueueFull,
|
||||
UserPendingLimitExceeded,
|
||||
build_rate_limit_detail,
|
||||
safe_enqueue_generation_task,
|
||||
)
|
||||
from app.dependencies import (
|
||||
@@ -417,12 +418,12 @@ def create_generation_task(
|
||||
except UserPendingLimitExceeded as e:
|
||||
raise HTTPException(
|
||||
status_code=429,
|
||||
detail=f"您的待处理任务过多(当前 {e.pending_count - count}/{e.limit},本次提交 {count} 个),请等待完成后再提交",
|
||||
detail=build_rate_limit_detail(e, generation_task_repository, scope="user"),
|
||||
) from e
|
||||
except GlobalQueueFull as e:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="系统繁忙,请稍后再试",
|
||||
detail=build_rate_limit_detail(e, generation_task_repository, scope="global"),
|
||||
) from e
|
||||
|
||||
# 画中画已下线:strategy_id 中的 pip/voice_pip 统一映射为 one_take
|
||||
@@ -579,7 +580,7 @@ def create_generation_task(
|
||||
if not created_tasks:
|
||||
raise HTTPException(
|
||||
status_code=429,
|
||||
detail="您的待处理任务过多,请等待完成后再提交",
|
||||
detail=build_rate_limit_detail(_e, generation_task_repository, scope="user"),
|
||||
) from _e
|
||||
break
|
||||
except GlobalQueueFull as _e:
|
||||
@@ -587,7 +588,7 @@ def create_generation_task(
|
||||
if not created_tasks:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="系统繁忙,请稍后再试",
|
||||
detail=build_rate_limit_detail(_e, generation_task_repository, scope="global"),
|
||||
) from _e
|
||||
break
|
||||
except HTTPException:
|
||||
@@ -713,15 +714,15 @@ def confirm_generation(
|
||||
log_task_status=True,
|
||||
):
|
||||
logger.warning("[确认生成] 入队失败: task_id=%s", new_task.id)
|
||||
except UserPendingLimitExceeded:
|
||||
except UserPendingLimitExceeded as _e:
|
||||
raise HTTPException(
|
||||
status_code=429,
|
||||
detail="您的待处理任务过多,请等待完成后再提交",
|
||||
detail=build_rate_limit_detail(_e, generation_task_repository, scope="user"),
|
||||
) from None
|
||||
except GlobalQueueFull:
|
||||
except GlobalQueueFull as _e:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="系统繁忙,请稍后再试",
|
||||
detail=build_rate_limit_detail(_e, generation_task_repository, scope="global"),
|
||||
) from None
|
||||
|
||||
return BatchGenerationTaskResponse(
|
||||
@@ -803,12 +804,24 @@ def retry_generation_task(
|
||||
if user_pending >= USER_PENDING_LIMIT:
|
||||
raise HTTPException(
|
||||
status_code=429,
|
||||
detail=f"您的待处理任务过多(当前 {user_pending}/{USER_PENDING_LIMIT}),请等待完成后再提交",
|
||||
detail=build_rate_limit_detail(
|
||||
UserPendingLimitExceeded(
|
||||
user_id=user_id,
|
||||
pending_count=user_pending,
|
||||
limit=USER_PENDING_LIMIT,
|
||||
),
|
||||
generation_task_repository,
|
||||
scope="user",
|
||||
),
|
||||
)
|
||||
if global_pending >= GLOBAL_PENDING_LIMIT:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="系统繁忙,请稍后再试",
|
||||
detail=build_rate_limit_detail(
|
||||
GlobalQueueFull(pending_count=global_pending, limit=GLOBAL_PENDING_LIMIT),
|
||||
generation_task_repository,
|
||||
scope="global",
|
||||
),
|
||||
)
|
||||
|
||||
use_case = CreateGenerationTaskUseCase(generation_task_repository)
|
||||
@@ -843,15 +856,15 @@ def retry_generation_task(
|
||||
log_task_status=True,
|
||||
):
|
||||
logger.warning("[生成任务] 重试入队失败: task_id=%s", retried.id)
|
||||
except UserPendingLimitExceeded:
|
||||
except UserPendingLimitExceeded as _e:
|
||||
raise HTTPException(
|
||||
status_code=429,
|
||||
detail="您的待处理任务过多,请等待完成后再提交",
|
||||
detail=build_rate_limit_detail(_e, generation_task_repository, scope="user"),
|
||||
) from None
|
||||
except GlobalQueueFull:
|
||||
except GlobalQueueFull as _e:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="系统繁忙,请稍后再试",
|
||||
detail=build_rate_limit_detail(_e, generation_task_repository, scope="global"),
|
||||
) from None
|
||||
return _to_generation_task_response(retried)
|
||||
|
||||
|
||||
@@ -8,27 +8,145 @@ logger = logging.getLogger(__name__)
|
||||
# ── 限流阈值常量(全系统统一管理,不要在业务代码里硬编码) ──
|
||||
USER_PENDING_LIMIT = 3 # 单用户 pending 上限
|
||||
GLOBAL_PENDING_LIMIT = 20 # 全局 pending 上限
|
||||
WORKER_CONCURRENCY = 4 # worker 渲染并发数(infra/docker/compose.yml WORKER_CONCURRENCY 默认值)
|
||||
|
||||
# 限流错误码:前端据此区分"排队等待"与"创建失败"
|
||||
ERROR_CODE_USER_QUEUE_FULL = "USER_QUEUE_FULL" # 429:用户自己的任务排队中
|
||||
ERROR_CODE_SYSTEM_QUEUE_FULL = "SYSTEM_QUEUE_FULL" # 503:系统整体繁忙
|
||||
|
||||
|
||||
class UserPendingLimitExceeded(Exception):
|
||||
"""用户 pending 任务数超限,返回 429。"""
|
||||
|
||||
def __init__(self, user_id: str, pending_count: int, limit: int):
|
||||
def __init__(
|
||||
self,
|
||||
user_id: str,
|
||||
pending_count: int,
|
||||
limit: int,
|
||||
*,
|
||||
running_count: int = 0,
|
||||
requested_count: int = 1,
|
||||
queue_ahead: int = 0,
|
||||
estimated_wait_seconds: int = 0,
|
||||
):
|
||||
self.user_id = user_id
|
||||
self.pending_count = pending_count
|
||||
self.limit = limit
|
||||
# 排队上下文(用于 429 结构化提示,前端展示"排队中"而非"创建失败")
|
||||
self.running_count = running_count
|
||||
self.requested_count = requested_count
|
||||
self.queue_ahead = queue_ahead
|
||||
self.estimated_wait_seconds = estimated_wait_seconds
|
||||
super().__init__(f"用户 {user_id} pending 任务数 {pending_count} 超过上限 {limit}")
|
||||
|
||||
|
||||
class GlobalQueueFull(Exception):
|
||||
"""全局限流,返回 503。"""
|
||||
|
||||
def __init__(self, pending_count: int, limit: int):
|
||||
def __init__(
|
||||
self,
|
||||
pending_count: int,
|
||||
limit: int,
|
||||
*,
|
||||
running_count: int = 0,
|
||||
queue_ahead: int = 0,
|
||||
estimated_wait_seconds: int = 0,
|
||||
):
|
||||
self.pending_count = pending_count
|
||||
self.limit = limit
|
||||
self.running_count = running_count
|
||||
self.queue_ahead = queue_ahead
|
||||
self.estimated_wait_seconds = estimated_wait_seconds
|
||||
super().__init__(f"系统 pending 任务数 {pending_count} 超过上限 {limit}")
|
||||
|
||||
|
||||
def _estimate_wait_seconds(queue_ahead: int, generation_task_repository: Any) -> int:
|
||||
"""根据排队任务数 + worker 并发数 + 历史平均任务耗时估算等待秒数。
|
||||
|
||||
估算公式:ceil(排队任务数 / 并发数) × 平均单任务耗时。
|
||||
拿不到历史数据时仓储层返回默认 120 秒。
|
||||
"""
|
||||
import math
|
||||
|
||||
if queue_ahead <= 0:
|
||||
return 0
|
||||
try:
|
||||
estimator = getattr(generation_task_repository, "estimate_avg_duration_seconds", None)
|
||||
avg_seconds = estimator() if estimator is not None else 120.0
|
||||
except Exception:
|
||||
avg_seconds = 120.0
|
||||
return int(math.ceil(queue_ahead / WORKER_CONCURRENCY) * avg_seconds)
|
||||
|
||||
|
||||
def build_rate_limit_detail(
|
||||
exc: Exception,
|
||||
generation_task_repository: Any,
|
||||
*,
|
||||
scope: str = "user",
|
||||
) -> dict:
|
||||
"""构造结构化限流响应体(HTTPException 的 detail)。
|
||||
|
||||
前端按 detail.code 判断场景:
|
||||
- USER_QUEUE_FULL (429):用户自己的任务在排队,应提示"等待/继续排队",不是创建失败
|
||||
- SYSTEM_QUEUE_FULL (503):系统繁忙,稍后重试
|
||||
|
||||
detail 字段:
|
||||
- code: 错误码
|
||||
- message: 可读中文提示(可直接展示)
|
||||
- queued_count: 当前排队(pending)任务数
|
||||
- running_count: 当前渲染中(running)任务数
|
||||
- queue_ahead: 前方排队任务数(预计等待批次依据)
|
||||
- estimated_wait_seconds: 预计等待秒数
|
||||
- limit: 对应限流上限
|
||||
"""
|
||||
if scope == "user" and isinstance(exc, UserPendingLimitExceeded):
|
||||
running = exc.running_count
|
||||
if not running:
|
||||
try:
|
||||
counter = getattr(generation_task_repository, "count_running_by_user", None)
|
||||
running = counter(exc.user_id) if counter is not None else 0
|
||||
except Exception:
|
||||
running = 0
|
||||
queue_ahead = exc.queue_ahead or max(exc.pending_count, 0)
|
||||
wait = exc.estimated_wait_seconds or _estimate_wait_seconds(queue_ahead, generation_task_repository)
|
||||
wait_minutes = max(1, round(wait / 60))
|
||||
message = (
|
||||
f"您有 {exc.pending_count} 个任务正在排队、{running} 个正在渲染,"
|
||||
f"同一时间最多提交 {exc.limit} 个任务。请等待约 {wait_minutes} 分钟后再提交"
|
||||
)
|
||||
return {
|
||||
"code": ERROR_CODE_USER_QUEUE_FULL,
|
||||
"message": message,
|
||||
"queued_count": exc.pending_count,
|
||||
"running_count": running,
|
||||
"queue_ahead": queue_ahead,
|
||||
"estimated_wait_seconds": wait,
|
||||
"limit": exc.limit,
|
||||
}
|
||||
|
||||
# 全局繁忙
|
||||
pending = getattr(exc, "pending_count", 0)
|
||||
running = getattr(exc, "running_count", 0)
|
||||
if not running:
|
||||
try:
|
||||
counter = getattr(generation_task_repository, "count_running_total", None)
|
||||
running = counter() if counter is not None else 0
|
||||
except Exception:
|
||||
running = 0
|
||||
queue_ahead = getattr(exc, "queue_ahead", 0) or pending
|
||||
wait = getattr(exc, "estimated_wait_seconds", 0) or _estimate_wait_seconds(queue_ahead, generation_task_repository)
|
||||
wait_minutes = max(1, round(wait / 60))
|
||||
return {
|
||||
"code": ERROR_CODE_SYSTEM_QUEUE_FULL,
|
||||
"message": f"系统繁忙:当前 {pending} 个任务排队中、{running} 个渲染中,预计等待约 {wait_minutes} 分钟,请稍后再试",
|
||||
"queued_count": pending,
|
||||
"running_count": running,
|
||||
"queue_ahead": queue_ahead,
|
||||
"estimated_wait_seconds": wait,
|
||||
"limit": getattr(exc, "limit", GLOBAL_PENDING_LIMIT),
|
||||
}
|
||||
|
||||
|
||||
def check_queue_limits(
|
||||
user_id: str,
|
||||
generation_task_repository: Any,
|
||||
|
||||
@@ -22,10 +22,18 @@ celery_app.conf.imports = (
|
||||
)
|
||||
|
||||
# Celery Beat 定时任务调度
|
||||
# 注:worker 单实例内嵌 beat(entrypoint-worker.sh -B),定时任务不会重复执行
|
||||
celery_app.conf.beat_schedule = {
|
||||
# pending 任务超时清理:worker 停止消费后,卡 pending 的任务 15 分钟内释放限流名额
|
||||
"cleanup-stale-pending-tasks": {
|
||||
"task": "worker.cleanup_stale_pending_tasks",
|
||||
"schedule": 600.0, # 每 10 分钟(秒)
|
||||
"options": {"expires": 300}, # 5 分钟过期,避免堆积
|
||||
"schedule": 300.0, # 每 5 分钟(秒)
|
||||
"options": {"expires": 240}, # 4 分钟过期,避免堆积
|
||||
},
|
||||
# running 孤儿任务巡检:容器重启/进程被杀后卡 running 的任务,20 分钟无更新则判失败
|
||||
"cleanup-stale-running-tasks": {
|
||||
"task": "worker.cleanup_stale_running_tasks",
|
||||
"schedule": 300.0, # 每 5 分钟(秒)
|
||||
"options": {"expires": 240},
|
||||
},
|
||||
}
|
||||
|
||||
@@ -7,11 +7,34 @@ from worker_app.db import SessionLocal
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 孤儿任务超时阈值:渲染任务超过此时间未更新则视为卡死
|
||||
ORPHAN_TASK_TIMEOUT_MINUTES = 10
|
||||
|
||||
# Pending 任务超时阈值:pending 任务在队列中等待超过此时间则自动清理
|
||||
PENDING_TASK_TIMEOUT_MINUTES = 30
|
||||
def cleanup_stale_running_with_session(repo, timeout_minutes: int) -> int:
|
||||
"""清理超时未更新的 running GenerationTask(可注入 repo 的纯核心,便于单测)。
|
||||
|
||||
Returns:
|
||||
清理的任务数量
|
||||
"""
|
||||
return repo.cleanup_stale_running(timeout_minutes)
|
||||
|
||||
|
||||
def cleanup_stale_pending_with_session(repo, timeout_minutes: int) -> int:
|
||||
"""清理超时 pending GenerationTask(可注入 repo 的纯核心,便于单测)。
|
||||
|
||||
Returns:
|
||||
清理的任务数量
|
||||
"""
|
||||
return repo.cleanup_stale_pending(timeout_minutes)
|
||||
|
||||
|
||||
# 孤儿任务超时阈值:running 任务超过此时间无进度更新则视为卡死。
|
||||
# 依据:worker.generate_video 硬超时 time_limit=11 分钟,正常任务不可能超过;
|
||||
# 20 分钟阈值覆盖硬超时 + 重试 + 余量,绝不误杀正常任务。
|
||||
ORPHAN_TASK_TIMEOUT_MINUTES = 20
|
||||
|
||||
# Pending 任务超时阈值:任务创建后超过此时间仍未被 worker 拉取,
|
||||
# 说明 worker 已停止消费(容器异常/卡死),清掉释放限流名额。
|
||||
# 依据:满队列(20 pending)× 平均 2 分钟 / 并发 4 ≈ 10 分钟,15 分钟留余量。
|
||||
PENDING_TASK_TIMEOUT_MINUTES = 15
|
||||
|
||||
|
||||
def cleanup_orphan_tasks(timeout_minutes: int = ORPHAN_TASK_TIMEOUT_MINUTES) -> int: # pragma: no cover
|
||||
@@ -32,9 +55,11 @@ def cleanup_orphan_tasks(timeout_minutes: int = ORPHAN_TASK_TIMEOUT_MINUTES) ->
|
||||
|
||||
try:
|
||||
session = SessionLocal()
|
||||
repo = SQLAlchemyGenerationTaskRepository(session)
|
||||
count = repo.cleanup_stale_running(timeout_minutes)
|
||||
session.close()
|
||||
try:
|
||||
repo = SQLAlchemyGenerationTaskRepository(session)
|
||||
count = cleanup_stale_running_with_session(repo, timeout_minutes)
|
||||
finally:
|
||||
session.close()
|
||||
if count > 0:
|
||||
logger.warning("清理了 %d 个超时的孤儿 GenerationTask(超过 %d 分钟未更新)", count, timeout_minutes)
|
||||
else:
|
||||
@@ -105,7 +130,7 @@ def cleanup_stale_pending_tasks(timeout_minutes: int = PENDING_TASK_TIMEOUT_MINU
|
||||
session = SessionLocal()
|
||||
try:
|
||||
repo = SQLAlchemyGenerationTaskRepository(session)
|
||||
count = repo.cleanup_stale_pending(timeout_minutes)
|
||||
count = cleanup_stale_pending_with_session(repo, timeout_minutes)
|
||||
if count > 0:
|
||||
logger.warning("清理了 %d 个超时的 pending GenerationTask(超过 %d 分钟未处理)", count, timeout_minutes)
|
||||
else:
|
||||
|
||||
@@ -1,14 +1,18 @@
|
||||
"""定期清理任务 — Celery Beat 调度。
|
||||
|
||||
包含:
|
||||
- cleanup_stale_pending_tasks: 定期清理卡在 pending 超时的 generation_tasks
|
||||
- cleanup_stale_pending_tasks: 定期清理卡在 pending 超时的 generation_tasks(worker 停止消费时占位)
|
||||
- cleanup_stale_running_tasks: 定期清理卡在 running 超时的 generation_tasks(容器重启/进程被杀后的孤儿)
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from celery import shared_task
|
||||
from worker_app.tasks._startup import (
|
||||
ORPHAN_TASK_TIMEOUT_MINUTES,
|
||||
PENDING_TASK_TIMEOUT_MINUTES,
|
||||
cleanup_orphan_tasks,
|
||||
cleanup_stale_jobs,
|
||||
cleanup_stale_pending_tasks,
|
||||
)
|
||||
|
||||
@@ -19,12 +23,12 @@ logger = logging.getLogger(__name__)
|
||||
def scheduled_cleanup_stale_pending(timeout_minutes: int = PENDING_TASK_TIMEOUT_MINUTES) -> dict:
|
||||
"""Celery Beat 调度的定期任务:清理超时的 pending 任务。
|
||||
|
||||
每 10 分钟执行一次(由 celery_app.py 的 beat_schedule 配置),
|
||||
每 5 分钟执行一次(由 celery_app.py 的 beat_schedule 配置),
|
||||
查找所有 status='pending' 且 created_at < NOW() - timeout_minutes
|
||||
的 generation_tasks,批量更新为 failed。
|
||||
的 generation_tasks,批量更新为 failed,释放限流名额。
|
||||
|
||||
Args:
|
||||
timeout_minutes: 超时时间(分钟),默认 30 分钟
|
||||
timeout_minutes: 超时时间(分钟),默认 15 分钟
|
||||
|
||||
Returns:
|
||||
{"cleaned": int}
|
||||
@@ -33,3 +37,33 @@ def scheduled_cleanup_stale_pending(timeout_minutes: int = PENDING_TASK_TIMEOUT_
|
||||
if count > 0:
|
||||
logger.info("[Beat] 清理了 %d 个超时 pending 任务(超时阈值 %d 分钟)", count, timeout_minutes)
|
||||
return {"cleaned": count}
|
||||
|
||||
|
||||
@shared_task(name="worker.cleanup_stale_running_tasks")
|
||||
def scheduled_cleanup_stale_running(timeout_minutes: int = ORPHAN_TASK_TIMEOUT_MINUTES) -> dict:
|
||||
"""Celery Beat 调度的定期任务:清理超时的 running 孤儿任务。
|
||||
|
||||
每 5 分钟执行一次。worker_ready 信号只在 worker 启动时清一次,
|
||||
若 worker 没重启但任务卡死(上传挂起、进程 OOM 被内核杀掉等),
|
||||
任务会永久卡在 running 占位。此任务做持续兜底:
|
||||
查找 status='running' 且 updated_at < NOW() - timeout_minutes 的任务,
|
||||
标记为 failed(原因:容器重启/超时中断),同时清理 Job 表孤儿。
|
||||
|
||||
Args:
|
||||
timeout_minutes: 超时时间(分钟),默认 20 分钟
|
||||
(worker.generate_video 硬超时 11 分钟,正常任务不可能超过 20 分钟)
|
||||
|
||||
Returns:
|
||||
{"generation_tasks": int, "jobs": int}
|
||||
"""
|
||||
gen_count = cleanup_orphan_tasks(timeout_minutes)
|
||||
job_count = cleanup_stale_jobs(timeout_minutes)
|
||||
total = gen_count + job_count
|
||||
if total > 0:
|
||||
logger.warning(
|
||||
"[Beat] 清理孤儿任务: running GenerationTask=%d, Job=%d(超时阈值 %d 分钟)",
|
||||
gen_count,
|
||||
job_count,
|
||||
timeout_minutes,
|
||||
)
|
||||
return {"generation_tasks": gen_count, "jobs": job_count}
|
||||
|
||||
@@ -138,6 +138,52 @@ class SQLAlchemyGenerationTaskRepository:
|
||||
.count()
|
||||
)
|
||||
|
||||
def count_running_by_user(self, user_id: str) -> int:
|
||||
"""统计指定用户处于 running 状态的任务数(用于限流提示展示)。"""
|
||||
return (
|
||||
self.session.query(GenerationTaskModel)
|
||||
.filter(
|
||||
GenerationTaskModel.created_by_user_id == user_id,
|
||||
GenerationTaskModel.status == GenerationTaskStatus.RUNNING.value,
|
||||
)
|
||||
.count()
|
||||
)
|
||||
|
||||
def count_running_total(self) -> int:
|
||||
"""统计全局处于 running 状态的任务数(worker 实际在执行的任务数)。"""
|
||||
return (
|
||||
self.session.query(GenerationTaskModel)
|
||||
.filter(GenerationTaskModel.status == GenerationTaskStatus.RUNNING.value)
|
||||
.count()
|
||||
)
|
||||
|
||||
def estimate_avg_duration_seconds(self, limit: int = 20, default_seconds: float = 120.0) -> float:
|
||||
"""估算最近完成任务的平均耗时(秒),用于 429 限流提示的等待预估。
|
||||
|
||||
取最近 N 条 completed 任务的 (completed_at - started_at) 平均值;
|
||||
无足够历史数据时返回 default_seconds。
|
||||
用 Python 侧计算差值,避免 SQLite/PostgreSQL 方言差异。
|
||||
"""
|
||||
rows = (
|
||||
self.session.query(GenerationTaskModel.started_at, GenerationTaskModel.completed_at)
|
||||
.filter(
|
||||
GenerationTaskModel.status == GenerationTaskStatus.COMPLETED.value,
|
||||
GenerationTaskModel.started_at.isnot(None),
|
||||
GenerationTaskModel.completed_at.isnot(None),
|
||||
)
|
||||
.order_by(GenerationTaskModel.completed_at.desc())
|
||||
.limit(limit)
|
||||
.all()
|
||||
)
|
||||
durations = [
|
||||
(completed - started).total_seconds()
|
||||
for started, completed in rows
|
||||
if completed and started and (completed - started).total_seconds() > 0
|
||||
]
|
||||
if not durations:
|
||||
return default_seconds
|
||||
return sum(durations) / len(durations)
|
||||
|
||||
def list_recent_by_user(self, user_id: str, limit: int = 5) -> list[GenerationTask]:
|
||||
models = (
|
||||
self.session.query(GenerationTaskModel)
|
||||
|
||||
@@ -20,6 +20,12 @@ class GenerationTaskRepository(Protocol):
|
||||
|
||||
def count_pending_total(self) -> int: ...
|
||||
|
||||
def count_running_by_user(self, user_id: str) -> int: ...
|
||||
|
||||
def count_running_total(self) -> int: ...
|
||||
|
||||
def estimate_avg_duration_seconds(self, limit: int = 20, default_seconds: float = 120.0) -> float: ...
|
||||
|
||||
def list_recent_by_user(self, user_id: str, limit: int = 5) -> list[GenerationTask]: ...
|
||||
|
||||
def list_by_source_edit_plan(self, plan_id: str) -> list[GenerationTask]: ...
|
||||
|
||||
@@ -0,0 +1,315 @@
|
||||
"""Issue #1709 任务容错:孤儿任务恢复 + 429 限流结构化提示。
|
||||
|
||||
覆盖:
|
||||
1. 仓储层:count_running_by_user/count_running_total 计数正确(预览/正式任务都计入)
|
||||
2. 仓储层:estimate_avg_duration_seconds 耗时估算(有历史/无历史)
|
||||
3. 限流核心:build_rate_limit_detail 返回结构化 code/message/排队数/预计等待
|
||||
4. worker 侧:cleanup_stale_running/pending 核心函数——中断任务被重置为 failed
|
||||
且原因写明(容器重启/超时中断),正常任务不受影响
|
||||
"""
|
||||
|
||||
import sys
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "apps" / "api"))
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "apps" / "worker"))
|
||||
|
||||
# 预注入 mock worker_app.db,防止真实数据库连接初始化(与其他 worker 测试同模式)
|
||||
_mock_db = MagicMock()
|
||||
_mock_db.SessionLocal = MagicMock()
|
||||
sys.modules.setdefault("worker_app.db", _mock_db)
|
||||
|
||||
from app.core import task_enqueue # noqa: E402
|
||||
from sqlalchemy import create_engine, text # noqa: E402
|
||||
from sqlalchemy.orm import sessionmaker # noqa: E402
|
||||
from worker_app.tasks import _startup # noqa: E402
|
||||
|
||||
from packages.adapters.sqlalchemy_impl.generation_task_repository import ( # noqa: E402
|
||||
SQLAlchemyGenerationTaskRepository,
|
||||
)
|
||||
from packages.adapters.sqlalchemy_impl.models import Base # noqa: E402
|
||||
from packages.domain import GenerationTask, GenerationTaskStatus # noqa: E402
|
||||
|
||||
|
||||
def _repository():
|
||||
engine = create_engine("sqlite:///:memory:", connect_args={"check_same_thread": False})
|
||||
Base.metadata.create_all(engine)
|
||||
session = sessionmaker(bind=engine)()
|
||||
return SQLAlchemyGenerationTaskRepository(session), session, engine
|
||||
|
||||
|
||||
def _make_task(**kwargs) -> GenerationTask:
|
||||
defaults = dict(
|
||||
project_id="proj-1",
|
||||
asset_library_id="lib-1",
|
||||
created_by_user_id="user-1",
|
||||
)
|
||||
defaults.update(kwargs)
|
||||
return GenerationTask.create(**defaults)
|
||||
|
||||
|
||||
def _age_task(engine, task_id, *, updated_minutes=None, created_minutes=None):
|
||||
"""用 SQL 直接把 updated_at/created_at 改到过去(模拟孤儿任务)。"""
|
||||
sets, params = [], {"id": task_id}
|
||||
if updated_minutes is not None:
|
||||
sets.append("updated_at = :uts")
|
||||
params["uts"] = datetime.now(timezone.utc) - timedelta(minutes=updated_minutes)
|
||||
if created_minutes is not None:
|
||||
sets.append("created_at = :cts")
|
||||
params["cts"] = datetime.now(timezone.utc) - timedelta(minutes=created_minutes)
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text(f"UPDATE generation_tasks SET {', '.join(sets)} WHERE id = :id"), params)
|
||||
conn.commit()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. running 计数(限流"渲染中"数量)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_count_running_by_user_mix_statuses():
|
||||
"""count_running_by_user 只统计该用户 running,不含 pending/completed/failed。"""
|
||||
repo, _, _ = _repository()
|
||||
t1 = _make_task(project_id="p1")
|
||||
repo.create(t1) # pending
|
||||
t2 = _make_task(project_id="p2")
|
||||
repo.create(t2)
|
||||
t2.mark_processing()
|
||||
repo.update(t2)
|
||||
t3 = _make_task(project_id="p3")
|
||||
repo.create(t3)
|
||||
t3.mark_processing()
|
||||
repo.update(t3)
|
||||
t4 = _make_task(project_id="p4")
|
||||
repo.create(t4)
|
||||
t4.mark_processing()
|
||||
repo.update(t4)
|
||||
t4.mark_completed()
|
||||
repo.update(t4)
|
||||
t5 = _make_task(project_id="p5", created_by_user_id="user-2")
|
||||
repo.create(t5)
|
||||
t5.mark_processing()
|
||||
repo.update(t5)
|
||||
|
||||
assert repo.count_running_by_user("user-1") == 2
|
||||
assert repo.count_running_by_user("user-2") == 1
|
||||
assert repo.count_running_total() == 3
|
||||
|
||||
|
||||
def test_count_running_total_empty():
|
||||
repo, _, _ = _repository()
|
||||
assert repo.count_running_total() == 0
|
||||
assert repo.count_running_by_user("nobody") == 0
|
||||
|
||||
|
||||
def test_preview_tasks_counted_in_running():
|
||||
"""预览任务(is_preview=True,工单实测卡 80% 的那种)同样计入 running。"""
|
||||
repo, _, _ = _repository()
|
||||
t = _make_task(is_preview=True)
|
||||
repo.create(t)
|
||||
t.mark_processing()
|
||||
repo.update(t)
|
||||
assert repo.count_running_by_user("user-1") == 1
|
||||
assert repo.count_running_total() == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. 平均耗时估算(429 等待预估依据)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _complete_task(repo, engine, task, duration_seconds: float):
|
||||
repo.create(task)
|
||||
task.mark_processing()
|
||||
repo.update(task)
|
||||
task.mark_completed()
|
||||
repo.update(task)
|
||||
now = datetime.now(timezone.utc)
|
||||
with engine.connect() as conn:
|
||||
conn.execute(
|
||||
text("UPDATE generation_tasks SET started_at = :s, completed_at = :c WHERE id = :id"),
|
||||
{"s": now - timedelta(seconds=duration_seconds), "c": now, "id": task.id},
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
def test_estimate_avg_duration_with_history():
|
||||
"""有历史完成任务时返回平均耗时(秒)。"""
|
||||
repo, _, engine = _repository()
|
||||
_complete_task(repo, engine, _make_task(project_id="p1"), 60.0)
|
||||
_complete_task(repo, engine, _make_task(project_id="p2"), 180.0)
|
||||
|
||||
avg = repo.estimate_avg_duration_seconds(default_seconds=120.0)
|
||||
assert 119.0 < avg < 121.0 # (60+180)/2 = 120
|
||||
|
||||
|
||||
def test_estimate_avg_duration_no_history_returns_default():
|
||||
"""无历史数据时返回默认值。"""
|
||||
repo, _, _ = _repository()
|
||||
assert repo.estimate_avg_duration_seconds(default_seconds=90.0) == 90.0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. build_rate_limit_detail 结构化提示(前端区分"排队"与"创建失败")
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_user_rate_limit_detail_structure():
|
||||
"""429 用户限流:返回 USER_QUEUE_FULL + 排队/渲染数 + 预计等待。"""
|
||||
repo, _, _ = _repository()
|
||||
for i in range(2): # 2 个渲染中
|
||||
t = _make_task(project_id=f"rp{i}")
|
||||
repo.create(t)
|
||||
t.mark_processing()
|
||||
repo.update(t)
|
||||
|
||||
exc = task_enqueue.UserPendingLimitExceeded(user_id="user-1", pending_count=3, limit=3)
|
||||
detail = task_enqueue.build_rate_limit_detail(exc, repo, scope="user")
|
||||
|
||||
assert detail["code"] == task_enqueue.ERROR_CODE_USER_QUEUE_FULL
|
||||
assert detail["queued_count"] == 3
|
||||
assert detail["running_count"] == 2
|
||||
assert detail["limit"] == 3
|
||||
assert detail["estimated_wait_seconds"] > 0
|
||||
assert "排队" in detail["message"]
|
||||
assert "user-1" not in detail["message"] # 不泄露内部 ID
|
||||
|
||||
|
||||
def test_global_rate_limit_detail_structure():
|
||||
"""503 全局繁忙:返回 SYSTEM_QUEUE_FULL。"""
|
||||
repo, _, _ = _repository()
|
||||
exc = task_enqueue.GlobalQueueFull(pending_count=20, limit=20)
|
||||
detail = task_enqueue.build_rate_limit_detail(exc, repo, scope="global")
|
||||
|
||||
assert detail["code"] == task_enqueue.ERROR_CODE_SYSTEM_QUEUE_FULL
|
||||
assert detail["queued_count"] == 20
|
||||
assert detail["limit"] == 20
|
||||
assert detail["estimated_wait_seconds"] > 0
|
||||
assert "系统繁忙" in detail["message"]
|
||||
|
||||
|
||||
def test_wait_estimate_uses_concurrency():
|
||||
"""等待预估:排队 8 个 / 并发 4 = 2 批 × 平均耗时。"""
|
||||
|
||||
class FakeRepo:
|
||||
def estimate_avg_duration_seconds(self, limit=20, default_seconds=120.0):
|
||||
return 100.0
|
||||
|
||||
wait = task_enqueue._estimate_wait_seconds(8, FakeRepo())
|
||||
assert wait == 200 # ceil(8/4)=2 批 × 100 秒
|
||||
|
||||
|
||||
def test_wait_estimate_repo_without_methods_uses_default():
|
||||
"""仓储没有新方法(旧 mock/鸭子类型)时用默认 120 秒兜底,不抛错。"""
|
||||
|
||||
class LegacyRepo:
|
||||
"""只实现旧接口的仓储(模拟未升级的调用方)。"""
|
||||
|
||||
def count_pending_total(self):
|
||||
return 0
|
||||
|
||||
wait = task_enqueue._estimate_wait_seconds(4, LegacyRepo())
|
||||
assert wait == 120 # ceil(4/4)=1 批 × 120 默认
|
||||
|
||||
|
||||
def test_rate_limit_detail_running_count_falls_back_to_zero():
|
||||
"""仓储不支持 running 计数时,running_count 优雅降级为 0。"""
|
||||
|
||||
class LegacyRepo:
|
||||
def count_pending_total(self):
|
||||
return 0
|
||||
|
||||
exc = task_enqueue.GlobalQueueFull(pending_count=20, limit=20)
|
||||
detail = task_enqueue.build_rate_limit_detail(exc, LegacyRepo(), scope="global")
|
||||
assert detail["running_count"] == 0
|
||||
assert detail["code"] == task_enqueue.ERROR_CODE_SYSTEM_QUEUE_FULL
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. worker 清理核心:中断任务被重置(worker 重启/超时恢复)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_worker_cleanup_resets_interrupted_running_task():
|
||||
"""模拟 worker 重启:running 超 20 分钟无更新的任务被重置为 failed,原因写明。"""
|
||||
repo, _, engine = _repository()
|
||||
|
||||
t = _make_task(is_preview=True) # 预览任务
|
||||
repo.create(t)
|
||||
t.mark_processing() # running
|
||||
repo.update(t)
|
||||
_age_task(engine, t.id, updated_minutes=25) # 25 分钟无进度更新
|
||||
|
||||
cleaned = _startup.cleanup_stale_running_with_session(repo, 20)
|
||||
assert cleaned == 1
|
||||
|
||||
saved = repo.get(t.id)
|
||||
assert saved.status == GenerationTaskStatus.FAILED
|
||||
assert "中断" in saved.error_message
|
||||
assert saved.error_info.get("error_type") == "WorkerInterrupted"
|
||||
assert saved.completed_at is not None
|
||||
|
||||
|
||||
def test_worker_cleanup_keeps_healthy_running_task():
|
||||
"""正常运行中(5 分钟前有更新)的任务不被误杀。"""
|
||||
repo, _, engine = _repository()
|
||||
|
||||
t = _make_task()
|
||||
repo.create(t)
|
||||
t.mark_processing()
|
||||
repo.update(t)
|
||||
_age_task(engine, t.id, updated_minutes=5)
|
||||
|
||||
assert _startup.cleanup_stale_running_with_session(repo, 20) == 0
|
||||
assert repo.get(t.id).status == GenerationTaskStatus.RUNNING
|
||||
|
||||
|
||||
def test_worker_cleanup_resets_stale_pending_task():
|
||||
"""卡 pending 超 15 分钟(worker 停止消费)的任务被重置,释放限流名额。"""
|
||||
repo, _, engine = _repository()
|
||||
|
||||
t = _make_task(is_preview=True)
|
||||
repo.create(t) # 一直 pending
|
||||
_age_task(engine, t.id, created_minutes=20)
|
||||
|
||||
cleaned = _startup.cleanup_stale_pending_with_session(repo, 15)
|
||||
assert cleaned == 1
|
||||
|
||||
saved = repo.get(t.id)
|
||||
assert saved.status == GenerationTaskStatus.FAILED
|
||||
assert saved.error_info.get("error_type") == "PendingTimeout"
|
||||
# 释放名额后 pending 计数归零,新请求不再被 429 误伤
|
||||
assert repo.count_pending_total() == 0
|
||||
|
||||
|
||||
def test_worker_cleanup_pending_keeps_recent():
|
||||
"""刚创建 3 分钟的 pending 任务不清理。"""
|
||||
repo, _, engine = _repository()
|
||||
|
||||
t = _make_task()
|
||||
repo.create(t)
|
||||
_age_task(engine, t.id, created_minutes=3)
|
||||
|
||||
assert _startup.cleanup_stale_pending_with_session(repo, 15) == 0
|
||||
assert repo.get(t.id).status == GenerationTaskStatus.PENDING
|
||||
|
||||
|
||||
def test_worker_cleanup_multiple_orphans_all_reset():
|
||||
"""3 个卡死 running 任务(工单实测:3 个预览卡 80% 超 10 小时)全部恢复。"""
|
||||
repo, _, engine = _repository()
|
||||
|
||||
ids = []
|
||||
for i in range(3):
|
||||
t = _make_task(project_id=f"p{i}", is_preview=True)
|
||||
repo.create(t)
|
||||
t.mark_processing()
|
||||
repo.update(t)
|
||||
_age_task(engine, t.id, updated_minutes=600) # 10 小时
|
||||
ids.append(t.id)
|
||||
|
||||
cleaned = _startup.cleanup_stale_running_with_session(repo, 20)
|
||||
assert cleaned == 3
|
||||
for tid in ids:
|
||||
assert repo.get(tid).status == GenerationTaskStatus.FAILED
|
||||
Reference in New Issue
Block a user