451 lines
17 KiB
Python
451 lines
17 KiB
Python
from __future__ import annotations
|
||
|
||
from collections import Counter, defaultdict
|
||
from datetime import datetime, timezone
|
||
from typing import Iterable, Sequence
|
||
|
||
from fastapi import HTTPException
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.enums.generation_task import (
|
||
ChatGenerationPipelineStage,
|
||
ChatGenerationTaskStatus,
|
||
GenerationMode,
|
||
)
|
||
from app.models.chat_generation_task import ChatGenerationTask
|
||
from app.services.generation.pipeline.db_lock_service import execute_with_lock_timeout
|
||
from app.services.operation_log_service import log_operation_event
|
||
from app.services.video_upscale.guard_service import assert_no_recoverable_failed_upscale_tasks
|
||
from app.services.resource_accounting_service import (
|
||
SOURCE_MODEL_CHAT_TASK,
|
||
soft_delete_resources_by_source,
|
||
)
|
||
|
||
|
||
ACTIVE_STAGES = {
|
||
ChatGenerationPipelineStage.QUEUED.value,
|
||
ChatGenerationPipelineStage.PREPARING.value,
|
||
ChatGenerationPipelineStage.CREATING_PROVIDER_TASK.value,
|
||
ChatGenerationPipelineStage.WAITING_REMOTE.value,
|
||
ChatGenerationPipelineStage.POLLING.value,
|
||
ChatGenerationPipelineStage.RESULT_READY.value,
|
||
ChatGenerationPipelineStage.DOWNLOAD_QUEUED.value,
|
||
ChatGenerationPipelineStage.DOWNLOADING.value,
|
||
ChatGenerationPipelineStage.RETRY_WAITING.value,
|
||
ChatGenerationPipelineStage.UPSCALE_QUEUED.value,
|
||
ChatGenerationPipelineStage.UPSCALE_PROCESSING.value,
|
||
ChatGenerationPipelineStage.UPSCALE_POLLING.value,
|
||
ChatGenerationPipelineStage.UPSCALE_DOWNLOADING.value,
|
||
ChatGenerationPipelineStage.UPSCALE_FINALIZING.value,
|
||
ChatGenerationPipelineStage.UPSCALE_RETRY_WAITING.value,
|
||
}
|
||
|
||
|
||
def is_task_active(task: ChatGenerationTask) -> bool:
|
||
return task.deleted_at is None and (
|
||
task.status == ChatGenerationTaskStatus.GENERATING.value
|
||
or (task.pipeline_stage or "") in ACTIVE_STAGES
|
||
)
|
||
|
||
|
||
def get_display_status(task: ChatGenerationTask) -> str:
|
||
if task.deleted_at is not None:
|
||
return "deleted"
|
||
if task.pipeline_stage == ChatGenerationPipelineStage.DOWNLOAD_FAILED.value:
|
||
return "download_failed"
|
||
if task.pipeline_stage == ChatGenerationPipelineStage.UPSCALE_FAILED.value:
|
||
return "failed"
|
||
return task.status or ChatGenerationTaskStatus.PENDING.value
|
||
|
||
|
||
async def load_children_map(
|
||
db: AsyncSession,
|
||
parent_ids: Sequence[str] | Iterable[str],
|
||
*,
|
||
include_deleted: bool = True,
|
||
) -> dict[str, list[ChatGenerationTask]]:
|
||
ids = list(dict.fromkeys(str(item) for item in parent_ids if item))
|
||
if not ids:
|
||
return {}
|
||
query = select(ChatGenerationTask).where(
|
||
ChatGenerationTask.parent_task_id.in_(ids),
|
||
ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_CHILD.value,
|
||
)
|
||
if not include_deleted:
|
||
query = query.where(ChatGenerationTask.deleted_at.is_(None))
|
||
result = await db.execute(
|
||
query.order_by(
|
||
ChatGenerationTask.parent_task_id.asc(),
|
||
ChatGenerationTask.generation_index.asc(),
|
||
ChatGenerationTask.created_at.asc(),
|
||
)
|
||
)
|
||
grouped: dict[str, list[ChatGenerationTask]] = defaultdict(list)
|
||
for task in result.scalars().all():
|
||
if task.parent_task_id:
|
||
grouped[task.parent_task_id].append(task)
|
||
return dict(grouped)
|
||
|
||
|
||
async def load_task_and_children(
|
||
db: AsyncSession,
|
||
*,
|
||
task_id: str,
|
||
user_id: str | None = None,
|
||
include_deleted_children: bool = True,
|
||
) -> tuple[ChatGenerationTask | None, list[ChatGenerationTask]]:
|
||
query = select(ChatGenerationTask).where(ChatGenerationTask.id == task_id)
|
||
if user_id:
|
||
query = query.where(ChatGenerationTask.user_id == user_id)
|
||
result = await db.execute(query.limit(1))
|
||
task = result.scalar_one_or_none()
|
||
if not task:
|
||
return None, []
|
||
if task.generation_mode == GenerationMode.CHATAPI_CHILD.value and task.parent_task_id:
|
||
parent_result = await db.execute(
|
||
select(ChatGenerationTask).where(ChatGenerationTask.id == task.parent_task_id).limit(1)
|
||
)
|
||
parent = parent_result.scalar_one_or_none()
|
||
return parent or task, [task]
|
||
if task.generation_mode != GenerationMode.CHATAPI_MAIN.value:
|
||
return task, []
|
||
children_map = await load_children_map(
|
||
db,
|
||
[task.id],
|
||
include_deleted=include_deleted_children,
|
||
)
|
||
return task, children_map.get(task.id, [])
|
||
|
||
|
||
def _generation_result_status(task: ChatGenerationTask) -> str:
|
||
"""返回任务真实生成结果，不受资源软删除影响。"""
|
||
if task.pipeline_stage == ChatGenerationPipelineStage.DOWNLOAD_FAILED.value:
|
||
return "download_failed"
|
||
if task.status == ChatGenerationTaskStatus.FAILED.value or (task.pipeline_stage or "") in {
|
||
ChatGenerationPipelineStage.FAILED.value,
|
||
ChatGenerationPipelineStage.TIMEOUT.value,
|
||
ChatGenerationPipelineStage.UPSCALE_FAILED.value,
|
||
}:
|
||
return "failed"
|
||
if is_task_active(task):
|
||
return "generating"
|
||
if task.status == ChatGenerationTaskStatus.COMPLETED.value:
|
||
return "completed"
|
||
return task.status or "pending"
|
||
|
||
|
||
def _build_summary(children: list[ChatGenerationTask]) -> str | None:
|
||
if not children:
|
||
return None
|
||
result_counters: Counter[str] = Counter(_generation_result_status(child) for child in children)
|
||
labels = {
|
||
"completed": "完成",
|
||
"failed": "生成失败",
|
||
"download_failed": "下载失败",
|
||
"generating": "生成中",
|
||
"pending": "待处理",
|
||
}
|
||
parts = [f"{count}项{labels.get(status, status)}" for status, count in result_counters.items() if count]
|
||
deleted_count = sum(1 for child in children if child.deleted_at is not None)
|
||
if deleted_count:
|
||
parts.append(f"{deleted_count}项资源已删除")
|
||
return f"{len(children)}项中" + "，".join(parts)
|
||
|
||
|
||
async def aggregate_main_task_status(
|
||
db: AsyncSession,
|
||
*,
|
||
parent_task_id: str,
|
||
) -> ChatGenerationTask | None:
|
||
result = await execute_with_lock_timeout(
|
||
db,
|
||
|
||
select(ChatGenerationTask)
|
||
.where(
|
||
ChatGenerationTask.id == parent_task_id,
|
||
ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_MAIN.value,
|
||
)
|
||
.with_for_update()
|
||
.limit(1)
|
||
)
|
||
main = result.scalar_one_or_none()
|
||
if not main or main.deleted_at is not None:
|
||
return main
|
||
|
||
children_map = await load_children_map(db, [parent_task_id], include_deleted=True)
|
||
children = children_map.get(parent_task_id, [])
|
||
if not children:
|
||
return main
|
||
|
||
previous_status = main.status
|
||
previous_stage = main.pipeline_stage
|
||
active_children = [child for child in children if is_task_active(child)]
|
||
failed_children = [
|
||
child
|
||
for child in children
|
||
if (
|
||
child.status == ChatGenerationTaskStatus.FAILED.value
|
||
or (child.pipeline_stage or "") in {
|
||
ChatGenerationPipelineStage.FAILED.value,
|
||
ChatGenerationPipelineStage.TIMEOUT.value,
|
||
ChatGenerationPipelineStage.DOWNLOAD_FAILED.value,
|
||
ChatGenerationPipelineStage.UPSCALE_FAILED.value,
|
||
}
|
||
)
|
||
]
|
||
completed_children = [
|
||
child for child in children if child.status == ChatGenerationTaskStatus.COMPLETED.value
|
||
]
|
||
|
||
if active_children:
|
||
main.status = ChatGenerationTaskStatus.GENERATING.value
|
||
main.pipeline_stage = active_children[0].pipeline_stage or ChatGenerationPipelineStage.QUEUED.value
|
||
main.generated_at = None
|
||
main.error_message = _build_summary(children)
|
||
elif failed_children:
|
||
main.status = ChatGenerationTaskStatus.FAILED.value
|
||
if any(child.pipeline_stage == ChatGenerationPipelineStage.DOWNLOAD_FAILED.value for child in failed_children):
|
||
main.pipeline_stage = ChatGenerationPipelineStage.DOWNLOAD_FAILED.value
|
||
elif any(child.pipeline_stage == ChatGenerationPipelineStage.UPSCALE_FAILED.value for child in failed_children):
|
||
main.pipeline_stage = ChatGenerationPipelineStage.UPSCALE_FAILED.value
|
||
else:
|
||
main.pipeline_stage = ChatGenerationPipelineStage.FAILED.value
|
||
main.generated_at = max(
|
||
(child.generated_at for child in completed_children if child.generated_at),
|
||
default=datetime.now(timezone.utc),
|
||
)
|
||
main.error_message = _build_summary(children)
|
||
else:
|
||
# 所有子任务真实生成结果均成功；资源是否软删除不改变生成历史终态。
|
||
main.status = ChatGenerationTaskStatus.COMPLETED.value
|
||
main.pipeline_stage = ChatGenerationPipelineStage.DONE.value
|
||
main.generated_at = max(
|
||
(child.generated_at for child in children if child.generated_at),
|
||
default=main.generated_at or datetime.now(timezone.utc),
|
||
)
|
||
main.error_message = None
|
||
|
||
if main.gen_type == "video":
|
||
main.credits_cost = round(sum(float(child.credits_cost or 0) for child in children), 2)
|
||
main.text_credits_cost = round(sum(float(child.text_credits_cost or 0) for child in children), 2)
|
||
main.text_tokens_used = sum(int(child.text_tokens_used or 0) for child in children)
|
||
main.image_tokens_used = sum(int(child.image_tokens_used or 0) for child in children)
|
||
main.video_tokens_used = sum(int(child.video_tokens_used or 0) for child in children)
|
||
# main 的手动重试次数只代表 main 自身，不能累加 child 的轮询/重试次数。
|
||
main.retry_count = int(main.manual_retry_count or 0)
|
||
main.poll_count = sum(int(child.poll_count or 0) for child in children)
|
||
main.poll_error_count = sum(int(child.poll_error_count or 0) for child in children)
|
||
|
||
await db.flush()
|
||
log_operation_event(
|
||
domain="generation_ai_batch",
|
||
event_type="MAIN_STATUS_AGGREGATED",
|
||
event_status="success",
|
||
source="service",
|
||
user_id=main.user_id,
|
||
group_id=main.id,
|
||
task_id=main.id,
|
||
detail={
|
||
"before_status": previous_status,
|
||
"before_stage": previous_stage,
|
||
"after_status": main.status,
|
||
"after_stage": main.pipeline_stage,
|
||
"summary": _build_summary(children),
|
||
},
|
||
)
|
||
return main
|
||
|
||
|
||
async def aggregate_parent_for_child(db: AsyncSession, child: ChatGenerationTask | None) -> ChatGenerationTask | None:
|
||
if not child or child.generation_mode != GenerationMode.CHATAPI_CHILD.value or not child.parent_task_id:
|
||
return None
|
||
# 项目关闭了 autoflush，先显式 flush 子任务的终态，确保聚合查询读取到本事务最新状态。
|
||
await db.flush()
|
||
return await aggregate_main_task_status(db, parent_task_id=str(child.parent_task_id))
|
||
|
||
|
||
async def soft_delete_child_tasks_batch(
|
||
db: AsyncSession,
|
||
*,
|
||
child_task_ids: Sequence[str] | Iterable[str],
|
||
user_id: str,
|
||
deleted_at: datetime | None = None,
|
||
require_completed: bool = False,
|
||
) -> int:
|
||
ids = list(dict.fromkeys(str(item) for item in child_task_ids if item))
|
||
if not ids:
|
||
return 0
|
||
deleted_at = deleted_at or datetime.now(timezone.utc)
|
||
result = await execute_with_lock_timeout(
|
||
db,
|
||
|
||
select(ChatGenerationTask)
|
||
.where(
|
||
ChatGenerationTask.id.in_(ids),
|
||
ChatGenerationTask.user_id == user_id,
|
||
ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_CHILD.value,
|
||
)
|
||
.with_for_update()
|
||
)
|
||
children = list(result.scalars().all())
|
||
found_ids = {str(child.id) for child in children}
|
||
missing_ids = [item for item in ids if item not in found_ids]
|
||
if missing_ids:
|
||
raise HTTPException(status_code=404, detail=f"子任务不存在: {','.join(missing_ids)}")
|
||
|
||
active_children = [child for child in children if child.deleted_at is None]
|
||
running_ids = [child.id for child in active_children if is_task_active(child)]
|
||
if running_ids:
|
||
raise HTTPException(status_code=400, detail=f"仍有 {len(running_ids)} 个子任务生成中，暂不能删除")
|
||
await assert_no_recoverable_failed_upscale_tasks(db, [str(child.id) for child in active_children])
|
||
if require_completed:
|
||
invalid_ids = [
|
||
child.id for child in active_children
|
||
if child.status != ChatGenerationTaskStatus.COMPLETED.value or child.generated_at is None
|
||
]
|
||
if invalid_ids:
|
||
raise HTTPException(status_code=409, detail=f"只有生成完成的资源才能从素材云删除: {','.join(invalid_ids)}")
|
||
|
||
source_ids = [str(child.id) for child in active_children]
|
||
freed_size = await soft_delete_resources_by_source(
|
||
db,
|
||
source_model=SOURCE_MODEL_CHAT_TASK,
|
||
source_ids=source_ids,
|
||
deleted_at=deleted_at,
|
||
)
|
||
parent_ids = list(dict.fromkeys(str(child.parent_task_id) for child in active_children if child.parent_task_id))
|
||
for child in active_children:
|
||
child.deleted_at = deleted_at
|
||
await db.flush()
|
||
for parent_id in parent_ids:
|
||
await aggregate_main_task_status(db, parent_task_id=parent_id)
|
||
return int(freed_size or 0)
|
||
|
||
|
||
async def soft_delete_child_task(
|
||
db: AsyncSession,
|
||
*,
|
||
child_task_id: str,
|
||
user_id: str,
|
||
deleted_at: datetime | None = None,
|
||
) -> int:
|
||
deleted_at = deleted_at or datetime.now(timezone.utc)
|
||
detail_result = await db.execute(
|
||
select(
|
||
ChatGenerationTask.parent_task_id,
|
||
ChatGenerationTask.generation_index,
|
||
).where(
|
||
ChatGenerationTask.id == child_task_id,
|
||
ChatGenerationTask.user_id == user_id,
|
||
ChatGenerationTask.generation_mode == GenerationMode.CHATAPI_CHILD.value,
|
||
).limit(1)
|
||
)
|
||
detail = detail_result.one_or_none()
|
||
if not detail:
|
||
raise HTTPException(status_code=404, detail="子任务不存在")
|
||
log_operation_event(
|
||
domain="generation_ai_batch",
|
||
event_type="CHILD_RESOURCE_DELETE_START",
|
||
event_status="started",
|
||
source="service",
|
||
user_id=user_id,
|
||
group_id=detail.parent_task_id,
|
||
task_id=child_task_id,
|
||
detail={"generation_index": detail.generation_index},
|
||
)
|
||
freed_size = await soft_delete_child_tasks_batch(
|
||
db,
|
||
child_task_ids=[child_task_id],
|
||
user_id=user_id,
|
||
deleted_at=deleted_at,
|
||
)
|
||
log_operation_event(
|
||
domain="generation_ai_batch",
|
||
event_type="CHILD_RESOURCE_DELETE_SUCCESS",
|
||
event_status="success",
|
||
source="service",
|
||
user_id=user_id,
|
||
group_id=detail.parent_task_id,
|
||
task_id=child_task_id,
|
||
detail={"generation_index": detail.generation_index, "freed_size_bytes": freed_size},
|
||
)
|
||
return freed_size
|
||
|
||
|
||
async def soft_delete_top_level_task_group(
|
||
db: AsyncSession,
|
||
*,
|
||
task_id: str,
|
||
user_id: str,
|
||
deleted_at: datetime | None = None,
|
||
) -> int:
|
||
deleted_at = deleted_at or datetime.now(timezone.utc)
|
||
result = await execute_with_lock_timeout(
|
||
db,
|
||
|
||
select(ChatGenerationTask)
|
||
.where(
|
||
ChatGenerationTask.id == task_id,
|
||
ChatGenerationTask.user_id == user_id,
|
||
ChatGenerationTask.generation_mode.in_(
|
||
[GenerationMode.CHATAPI_ASYNC.value, GenerationMode.CHATAPI_MAIN.value]
|
||
),
|
||
ChatGenerationTask.deleted_at.is_(None),
|
||
)
|
||
.with_for_update()
|
||
.limit(1)
|
||
)
|
||
task = result.scalar_one_or_none()
|
||
if not task:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
|
||
if task.generation_mode == GenerationMode.CHATAPI_ASYNC.value:
|
||
if is_task_active(task):
|
||
raise HTTPException(status_code=400, detail="当前任务正在生成中，暂不能删除")
|
||
await assert_no_recoverable_failed_upscale_tasks(db, [str(task.id)])
|
||
freed_size = await soft_delete_resources_by_source(
|
||
db,
|
||
source_model=SOURCE_MODEL_CHAT_TASK,
|
||
source_ids=[task.id],
|
||
deleted_at=deleted_at,
|
||
)
|
||
task.deleted_at = deleted_at
|
||
await db.flush()
|
||
return int(freed_size or 0)
|
||
|
||
children_map = await load_children_map(db, [task.id], include_deleted=True)
|
||
children = children_map.get(task.id, [])
|
||
active_ids = [child.id for child in children if is_task_active(child)]
|
||
if active_ids:
|
||
raise HTTPException(status_code=400, detail=f"任务组仍有 {len(active_ids)} 个子任务生成中，暂不能删除")
|
||
|
||
active_children = [child for child in children if child.deleted_at is None]
|
||
child_ids = [child.id for child in active_children]
|
||
await assert_no_recoverable_failed_upscale_tasks(db, [str(child_id) for child_id in child_ids])
|
||
freed_size = await soft_delete_resources_by_source(
|
||
db,
|
||
source_model=SOURCE_MODEL_CHAT_TASK,
|
||
source_ids=child_ids,
|
||
deleted_at=deleted_at,
|
||
)
|
||
for child in active_children:
|
||
child.deleted_at = deleted_at
|
||
task.deleted_at = deleted_at
|
||
await db.flush()
|
||
log_operation_event(
|
||
domain="generation_ai_batch",
|
||
event_type="BATCH_GROUP_DELETE_SUCCESS",
|
||
event_status="success",
|
||
source="service",
|
||
user_id=user_id,
|
||
group_id=task.id,
|
||
task_id=task.id,
|
||
detail={
|
||
"child_task_ids": child_ids,
|
||
"freed_size_bytes": int(freed_size or 0),
|
||
"physical_files_deleted": False,
|
||
},
|
||
)
|
||
return int(freed_size or 0)
|