471 lines
17 KiB
Python
471 lines
17 KiB
Python
"""数据库驱动的后台任务仓库与调度状态机。
|
||
|
||
已承载图片生成与批量水印；表结构刻意通用，后续可复用到 ZIP 导出等耗时操作。
|
||
Worker 不依赖 FastAPI 进程；第三方 provider task id 保存在子项中，重启后继续轮询。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import uuid
|
||
from dataclasses import dataclass
|
||
from datetime import datetime, timedelta, timezone
|
||
from typing import Any
|
||
|
||
from sqlalchemy import case, func, or_, select, update
|
||
|
||
from db import get_session_factory
|
||
from models import BackgroundJob, BackgroundJobItem
|
||
|
||
JOB_QUEUED = "queued"
|
||
JOB_RUNNING = "running"
|
||
JOB_DONE = "done"
|
||
JOB_PARTIAL = "partial"
|
||
JOB_FAILED = "failed"
|
||
JOB_CANCELLED = "cancelled"
|
||
|
||
ITEM_QUEUED = "queued"
|
||
ITEM_RUNNING = "running"
|
||
ITEM_WAITING_PROVIDER = "waiting_provider"
|
||
ITEM_DONE = "done"
|
||
ITEM_FAILED = "failed"
|
||
ITEM_CANCELLED = "cancelled"
|
||
|
||
TERMINAL_ITEM_STATUSES = (ITEM_DONE, ITEM_FAILED, ITEM_CANCELLED)
|
||
TERMINAL_JOB_STATUSES = (JOB_DONE, JOB_PARTIAL, JOB_FAILED, JOB_CANCELLED)
|
||
|
||
|
||
def utcnow() -> datetime:
|
||
return datetime.now(timezone.utc)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ClaimedItem:
|
||
job_id: str
|
||
item_id: str
|
||
job_kind: str
|
||
product_id: str | None
|
||
provider: str
|
||
model: str | None
|
||
platform: str
|
||
lang: str
|
||
ratio: str
|
||
style_set: int
|
||
request_payload: dict[str, Any]
|
||
item_payload: dict[str, Any]
|
||
provider_task_id: str | None
|
||
provider_task_kind: str | None
|
||
|
||
|
||
async def create_image_job(
|
||
*,
|
||
product_id: str | None,
|
||
kind: str,
|
||
priority: int,
|
||
provider: str,
|
||
model: str | None,
|
||
platform: str,
|
||
lang: str,
|
||
ratio: str,
|
||
style_set: int,
|
||
request_payload: dict[str, Any],
|
||
item_payloads: list[dict[str, Any]],
|
||
) -> BackgroundJob:
|
||
"""一次事务写入父任务和逐图子项，API 成功返回即意味着任务可由 Worker 恢复。"""
|
||
parsed_product_id = uuid.UUID(product_id) if product_id else None
|
||
job = BackgroundJob(
|
||
product_id=parsed_product_id,
|
||
kind=kind,
|
||
status=JOB_QUEUED,
|
||
priority=priority,
|
||
provider=provider,
|
||
model=model,
|
||
platform=platform,
|
||
lang=lang,
|
||
ratio=ratio,
|
||
style_set=style_set,
|
||
request_payload=request_payload,
|
||
total=len(item_payloads),
|
||
)
|
||
async with get_session_factory()() as db:
|
||
db.add(job)
|
||
await db.flush()
|
||
db.add_all([
|
||
BackgroundJobItem(job_id=job.id, position=index, payload=payload, status=ITEM_QUEUED)
|
||
for index, payload in enumerate(item_payloads)
|
||
])
|
||
await db.commit()
|
||
await db.refresh(job)
|
||
return job
|
||
|
||
|
||
async def _refresh_job_status(db, job: BackgroundJob) -> None:
|
||
"""从子项计算父任务进度与终态；调用方已在同一事务中修改子项。"""
|
||
rows = (await db.execute(
|
||
select(BackgroundJobItem.status, func.count(BackgroundJobItem.id))
|
||
.where(BackgroundJobItem.job_id == job.id)
|
||
.group_by(BackgroundJobItem.status)
|
||
)).all()
|
||
counts = {status: int(count) for status, count in rows}
|
||
done = counts.get(ITEM_DONE, 0)
|
||
failed = counts.get(ITEM_FAILED, 0)
|
||
cancelled = counts.get(ITEM_CANCELLED, 0)
|
||
terminal = done + failed + cancelled
|
||
job.completed = done
|
||
job.failed = failed
|
||
|
||
if terminal < job.total:
|
||
job.status = JOB_RUNNING
|
||
return
|
||
job.completed_at = utcnow()
|
||
if done == 0 and cancelled == job.total:
|
||
job.status = JOB_CANCELLED
|
||
elif failed == 0 and cancelled == 0:
|
||
job.status = JOB_DONE
|
||
elif done > 0:
|
||
job.status = JOB_PARTIAL
|
||
elif job.cancel_requested and cancelled > 0:
|
||
job.status = JOB_CANCELLED
|
||
else:
|
||
job.status = JOB_FAILED
|
||
|
||
errors = (await db.scalars(
|
||
select(BackgroundJobItem.error)
|
||
.where(BackgroundJobItem.job_id == job.id, BackgroundJobItem.status == ITEM_FAILED)
|
||
.order_by(BackgroundJobItem.position)
|
||
)).all()
|
||
if errors:
|
||
brief = [str(error)[:200] for error in errors if error]
|
||
job.error = "；".join(brief[:6]) or "任务失败"
|
||
|
||
|
||
async def recover_expired_leases() -> int:
|
||
"""把上个 Worker 遗留的工作单元重新放回队列；provider id 不会被清空。"""
|
||
now = utcnow()
|
||
async with get_session_factory()() as db:
|
||
result = await db.execute(
|
||
update(BackgroundJobItem)
|
||
.where(
|
||
BackgroundJobItem.status == ITEM_RUNNING,
|
||
BackgroundJobItem.lease_expires_at.is_not(None),
|
||
BackgroundJobItem.lease_expires_at < now,
|
||
)
|
||
.values(status=ITEM_QUEUED, lease_owner=None, lease_expires_at=None, next_run_at=now)
|
||
)
|
||
await db.commit()
|
||
return int(result.rowcount or 0)
|
||
|
||
|
||
async def claim_next_item(owner: str, lease_seconds: int = 90) -> ClaimedItem | None:
|
||
"""领取一个到期子项。单机 SQLite Worker 使用租约，未来迁移 PostgreSQL 后可直接强化为 SKIP LOCKED。"""
|
||
now = utcnow()
|
||
async with get_session_factory()() as db:
|
||
# 先回收过期租约；即便 Worker 非正常退出，也不会永久卡住。
|
||
await db.execute(
|
||
update(BackgroundJobItem)
|
||
.where(
|
||
BackgroundJobItem.status == ITEM_RUNNING,
|
||
BackgroundJobItem.lease_expires_at.is_not(None),
|
||
BackgroundJobItem.lease_expires_at < now,
|
||
)
|
||
.values(status=ITEM_QUEUED, lease_owner=None, lease_expires_at=None, next_run_at=now)
|
||
)
|
||
query = (
|
||
select(BackgroundJobItem, BackgroundJob)
|
||
.join(BackgroundJob, BackgroundJob.id == BackgroundJobItem.job_id)
|
||
.where(
|
||
BackgroundJob.cancel_requested.is_(False),
|
||
BackgroundJob.status.not_in(TERMINAL_JOB_STATUSES),
|
||
BackgroundJobItem.status.in_((ITEM_QUEUED, ITEM_WAITING_PROVIDER)),
|
||
or_(BackgroundJobItem.next_run_at.is_(None), BackgroundJobItem.next_run_at <= now),
|
||
)
|
||
.order_by(BackgroundJob.priority.desc(), BackgroundJob.created_at.asc(), BackgroundJobItem.position.asc())
|
||
.limit(1)
|
||
)
|
||
row = (await db.execute(query)).first()
|
||
if row is None:
|
||
await db.commit()
|
||
return None
|
||
item, job = row
|
||
item.status = ITEM_RUNNING
|
||
if not item.provider_task_id:
|
||
item.attempts += 1
|
||
item.lease_owner = owner
|
||
item.lease_expires_at = now + timedelta(seconds=lease_seconds)
|
||
item.next_run_at = None
|
||
item.started_at = item.started_at or now
|
||
job.status = JOB_RUNNING
|
||
job.started_at = job.started_at or now
|
||
await db.commit()
|
||
return ClaimedItem(
|
||
job_id=str(job.id), item_id=str(item.id), job_kind=job.kind,
|
||
product_id=str(job.product_id) if job.product_id else None,
|
||
provider=job.provider, model=job.model, platform=job.platform,
|
||
lang=job.lang, ratio=job.ratio, style_set=job.style_set,
|
||
request_payload=dict(job.request_payload or {}), item_payload=dict(item.payload or {}),
|
||
provider_task_id=item.provider_task_id, provider_task_kind=item.provider_task_kind,
|
||
)
|
||
|
||
|
||
async def mark_waiting_provider(item_id: str, provider_task_id: str, provider_task_kind: str, delay_seconds: int = 3) -> None:
|
||
"""先持久化外部任务 ID，再安排下一次查询，是避免重启重复扣费的关键提交点。"""
|
||
async with get_session_factory()() as db:
|
||
item = await db.get(BackgroundJobItem, uuid.UUID(item_id))
|
||
if item is None:
|
||
return
|
||
item.provider_task_id = provider_task_id
|
||
item.provider_task_kind = provider_task_kind
|
||
item.status = ITEM_WAITING_PROVIDER
|
||
item.lease_owner = None
|
||
item.lease_expires_at = None
|
||
item.next_run_at = utcnow() + timedelta(seconds=delay_seconds)
|
||
await db.commit()
|
||
|
||
|
||
async def mark_item_done(
|
||
item_id: str,
|
||
*,
|
||
result_url: str,
|
||
source_url: str | None,
|
||
watermark_config: dict | None,
|
||
generated_asset_id: str | None,
|
||
) -> None:
|
||
async with get_session_factory()() as db:
|
||
item = await db.get(BackgroundJobItem, uuid.UUID(item_id))
|
||
if item is None:
|
||
return
|
||
job = await db.get(BackgroundJob, item.job_id)
|
||
item.status = ITEM_DONE
|
||
item.result_url = result_url
|
||
item.source_url = source_url
|
||
item.watermark_config = watermark_config
|
||
item.generated_asset_id = uuid.UUID(generated_asset_id) if generated_asset_id else None
|
||
item.error = None
|
||
item.lease_owner = None
|
||
item.lease_expires_at = None
|
||
item.completed_at = utcnow()
|
||
if job is not None:
|
||
await _refresh_job_status(db, job)
|
||
await db.commit()
|
||
|
||
|
||
async def mark_item_failed(item_id: str, error: str) -> None:
|
||
async with get_session_factory()() as db:
|
||
item = await db.get(BackgroundJobItem, uuid.UUID(item_id))
|
||
if item is None:
|
||
return
|
||
job = await db.get(BackgroundJob, item.job_id)
|
||
item.status = ITEM_FAILED
|
||
item.error = error[:500]
|
||
item.lease_owner = None
|
||
item.lease_expires_at = None
|
||
item.completed_at = utcnow()
|
||
if job is not None:
|
||
await _refresh_job_status(db, job)
|
||
await db.commit()
|
||
|
||
|
||
async def job_can_write(job_id: str) -> bool:
|
||
"""结果落盘前的最后一道确认：取消或商品被硬删除时禁止回写素材。"""
|
||
async with get_session_factory()() as db:
|
||
job = await db.get(BackgroundJob, uuid.UUID(job_id))
|
||
if job is None or job.cancel_requested:
|
||
return False
|
||
if job.product_id is None:
|
||
return True
|
||
# 仅用主键查询避免导入 Product 形成循环；不存在即不允许生成结果复活商品。
|
||
from models import Product
|
||
return await db.get(Product, job.product_id) is not None
|
||
|
||
|
||
async def mark_item_cancelled(item_id: str) -> None:
|
||
async with get_session_factory()() as db:
|
||
item = await db.get(BackgroundJobItem, uuid.UUID(item_id))
|
||
if item is None:
|
||
return
|
||
job = await db.get(BackgroundJob, item.job_id)
|
||
item.status = ITEM_CANCELLED
|
||
item.lease_owner = None
|
||
item.lease_expires_at = None
|
||
item.completed_at = utcnow()
|
||
if job is not None:
|
||
await _refresh_job_status(db, job)
|
||
await db.commit()
|
||
|
||
|
||
async def cancel_job(job_id: str) -> BackgroundJob | None:
|
||
async with get_session_factory()() as db:
|
||
job = await db.get(BackgroundJob, uuid.UUID(job_id))
|
||
if job is None:
|
||
return None
|
||
if job.status in TERMINAL_JOB_STATUSES:
|
||
return job
|
||
job.cancel_requested = True
|
||
await db.execute(
|
||
update(BackgroundJobItem)
|
||
.where(BackgroundJobItem.job_id == job.id, BackgroundJobItem.status.in_((ITEM_QUEUED, ITEM_WAITING_PROVIDER)))
|
||
.values(status=ITEM_CANCELLED, completed_at=utcnow(), lease_owner=None, lease_expires_at=None)
|
||
)
|
||
await _refresh_job_status(db, job)
|
||
await db.commit()
|
||
await db.refresh(job)
|
||
return job
|
||
|
||
|
||
async def cancel_product_jobs(product_id: str) -> int:
|
||
"""硬删除商品前取消其未完成任务，Worker 随后只会丢弃结果而不会回写素材。"""
|
||
try:
|
||
parsed = uuid.UUID(product_id)
|
||
except ValueError:
|
||
return 0
|
||
async with get_session_factory()() as db:
|
||
jobs = (await db.scalars(
|
||
select(BackgroundJob).where(
|
||
BackgroundJob.product_id == parsed,
|
||
BackgroundJob.status.not_in(TERMINAL_JOB_STATUSES),
|
||
)
|
||
)).all()
|
||
now = utcnow()
|
||
for job in jobs:
|
||
job.cancel_requested = True
|
||
await db.execute(
|
||
update(BackgroundJobItem)
|
||
.where(BackgroundJobItem.job_id == job.id, BackgroundJobItem.status.in_((ITEM_QUEUED, ITEM_WAITING_PROVIDER)))
|
||
.values(status=ITEM_CANCELLED, completed_at=now, lease_owner=None, lease_expires_at=None)
|
||
)
|
||
await _refresh_job_status(db, job)
|
||
await db.commit()
|
||
return len(jobs)
|
||
|
||
|
||
async def retry_failed_items(job_id: str) -> BackgroundJob | None:
|
||
"""仅把失败子项重新排队；外部失败任务不可复用，清空 provider id 后重新提交。"""
|
||
async with get_session_factory()() as db:
|
||
job = await db.get(BackgroundJob, uuid.UUID(job_id))
|
||
if job is None:
|
||
return None
|
||
now = utcnow()
|
||
result = await db.execute(
|
||
update(BackgroundJobItem)
|
||
.where(BackgroundJobItem.job_id == job.id, BackgroundJobItem.status == ITEM_FAILED)
|
||
.values(
|
||
status=ITEM_QUEUED, error=None, provider_task_id=None, provider_task_kind=None,
|
||
result_url=None, source_url=None, watermark_config=None, generated_asset_id=None,
|
||
next_run_at=now, lease_owner=None, lease_expires_at=None, completed_at=None,
|
||
)
|
||
)
|
||
if not result.rowcount:
|
||
return job
|
||
job.cancel_requested = False
|
||
job.status = JOB_QUEUED
|
||
job.completed_at = None
|
||
job.error = None
|
||
await _refresh_job_status(db, job)
|
||
await db.commit()
|
||
await db.refresh(job)
|
||
return job
|
||
|
||
|
||
async def retry_failed_item(job_id: str, item_id: str) -> BackgroundJob | None:
|
||
"""只重试一个失败子项，供批量水印等任务逐图补偿。"""
|
||
try:
|
||
parsed_job = uuid.UUID(job_id)
|
||
parsed_item = uuid.UUID(item_id)
|
||
except ValueError:
|
||
return None
|
||
async with get_session_factory()() as db:
|
||
job = await db.get(BackgroundJob, parsed_job)
|
||
item = await db.get(BackgroundJobItem, parsed_item)
|
||
if job is None or item is None or item.job_id != job.id:
|
||
return None
|
||
if item.status != ITEM_FAILED:
|
||
return job
|
||
# 水印任务等待期间，素材可能被用户再次处理过。逐项重试时改为以
|
||
# 最新 current 版本为基线，避免把已经更新的图片误判为永久失败。
|
||
if job.kind == "watermark_batch":
|
||
raw_asset_id = (item.payload or {}).get("asset_id")
|
||
try:
|
||
from models import ProductAsset
|
||
asset = await db.get(ProductAsset, uuid.UUID(str(raw_asset_id)))
|
||
except (TypeError, ValueError):
|
||
asset = None
|
||
if asset is not None:
|
||
payload = dict(item.payload or {})
|
||
payload["expected_current_storage_key"] = asset.current_storage_key
|
||
payload["expected_current_url"] = asset.stored_url
|
||
item.payload = payload
|
||
item.status = ITEM_QUEUED
|
||
item.error = None
|
||
item.provider_task_id = None
|
||
item.provider_task_kind = None
|
||
item.result_url = None
|
||
item.source_url = None
|
||
item.watermark_config = None
|
||
item.generated_asset_id = None
|
||
item.next_run_at = utcnow()
|
||
item.lease_owner = None
|
||
item.lease_expires_at = None
|
||
item.completed_at = None
|
||
job.cancel_requested = False
|
||
job.status = JOB_QUEUED
|
||
job.completed_at = None
|
||
job.error = None
|
||
await _refresh_job_status(db, job)
|
||
await db.commit()
|
||
await db.refresh(job)
|
||
return job
|
||
|
||
|
||
async def get_job(job_id: str) -> tuple[BackgroundJob, list[BackgroundJobItem]] | None:
|
||
try:
|
||
parsed = uuid.UUID(job_id)
|
||
except ValueError:
|
||
return None
|
||
async with get_session_factory()() as db:
|
||
job = await db.get(BackgroundJob, parsed)
|
||
if job is None:
|
||
return None
|
||
items = (await db.scalars(
|
||
select(BackgroundJobItem).where(BackgroundJobItem.job_id == job.id).order_by(BackgroundJobItem.position)
|
||
)).all()
|
||
return job, items
|
||
|
||
|
||
async def list_product_jobs(
|
||
product_id: str,
|
||
active_only: bool = False,
|
||
limit: int = 20,
|
||
kinds: tuple[str, ...] | None = None,
|
||
) -> list[BackgroundJob]:
|
||
try:
|
||
parsed = uuid.UUID(product_id)
|
||
except ValueError:
|
||
return []
|
||
async with get_session_factory()() as db:
|
||
query = select(BackgroundJob).where(BackgroundJob.product_id == parsed)
|
||
if kinds:
|
||
query = query.where(BackgroundJob.kind.in_(kinds))
|
||
if active_only:
|
||
query = query.where(BackgroundJob.status.not_in(TERMINAL_JOB_STATUSES))
|
||
return (await db.scalars(query.order_by(BackgroundJob.created_at.desc()).limit(limit))).all()
|
||
|
||
|
||
def suite_view(job: BackgroundJob, items: list[BackgroundJobItem]) -> dict[str, Any]:
|
||
"""兼容现有 /suites/{id} 契约，同时增加任务类型、取消状态与子项状态。"""
|
||
return {
|
||
"id": str(job.id), "status": job.status, "kind": job.kind,
|
||
"style_set": job.style_set, "platform": job.platform, "lang": job.lang,
|
||
"ratio": job.ratio, "provider": job.provider, "model": job.model,
|
||
"total": job.total, "completed": job.completed, "failed": job.failed,
|
||
"cancel_requested": job.cancel_requested, "error": job.error,
|
||
"created_at": job.created_at, "started_at": job.started_at, "completed_at": job.completed_at,
|
||
"images": [
|
||
{
|
||
"type_id": (item.payload or {}).get("kind", "custom"),
|
||
"item_id": str(item.id),
|
||
"name": (item.payload or {}).get("title", "AI 生图"),
|
||
"url": item.result_url or "", "status": item.status,
|
||
"error": item.error,
|
||
}
|
||
for item in items
|
||
],
|
||
}
|