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

项目/AI生成链路合并

Browse Source
This commit is contained in:
GinHa
2026-07-20 13:48:17 +08:00
parent 34ca98f9eb
commit fe5a59d725
73 changed files with 5819 additions and 2573 deletions
Download Patch File Download Diff File Expand all files Collapse all files
video-gen-api/app/services/video_upscale/task_service.py
+139 -23
View File
@@ -5,6 +5,7 @@ import math
import os
import uuid
from datetime import datetime, timedelta, timezone
from collections.abc import Awaitable, Callable
from typing import Any
from urllib.parse import urlsplit, urlunsplit
@@ -30,8 +31,9 @@ from app.enums.video_upscale import (
from app.models.chat_generation_task import ChatGenerationTask
from app.models.generation_record import GenerationRecord
from app.models.video_upscale_task import VideoUpscaleTask
from app.services.generation.ai.task_group_service import aggregate_parent_for_child
from app.services.generation.module_hook_service import notify_chat_generation_task_finished
from app.services.generation.pipeline.db_lock_service import apply_short_lock_timeout
from app.services.generation.pipeline.lifecycle_service import notify_owner_finished
from app.services.redis_registry_service import RedisExecutionLockError, RedisExecutionLockLost
from app.services.media_token_usage_snapshot_service import (
sync_chat_generation_task_media_token_snapshot,
sync_generation_record_media_token_snapshot,
@@ -73,6 +75,14 @@ from app.services.video_upscale.volc_service import (
from app.utils.id_gen import generate_id
ExecutionGuard = Callable[[], Awaitable[None]]
async def _ensure_guard(execution_guard: ExecutionGuard | None) -> None:
if execution_guard is not None:
await execution_guard()
def _now() -> datetime:
return datetime.now(timezone.utc)
@@ -185,6 +195,7 @@ async def _load_pair(
) -> tuple[VideoUpscaleTask | None, VideoUpscaleOwner | None]:
query = select(VideoUpscaleTask).where(VideoUpscaleTask.id == upscale_task_id)
if for_update:
await apply_short_lock_timeout(db)
query = query.with_for_update()
result = await db.execute(query.limit(1))
upscale = result.scalar_one_or_none()
@@ -217,6 +228,7 @@ async def prepare_video_upscale_task(
source_local_path: str,
source_file_size_bytes: int,
source_remote_url: str | None = None,
source_info: Any | None = None,
) -> VideoUpscaleTask:
owner: VideoUpscaleOwner | None = task or generation_record
if owner is None:
@@ -224,7 +236,9 @@ async def prepare_video_upscale_task(
if owner.gen_type != GenerationType.VIDEO.value or not bool(owner.video_upscale_enabled_snapshot):
raise RuntimeError("当前任务未启用视频超分快照")
snapshot = _snapshot(owner)
source_info = await probe_video(source_local_path)
# ffprobe should run before the owner row is locked by the caller.
# Keep this fallback for non-generation callers that do not pass probe data.
source_info = source_info or await probe_video(source_local_path)
remote_url = str(source_remote_url or getattr(owner, "remote_result_url", None) or "").strip() or None
signed_at, expires_at = parse_tos_signed_url_expiry(remote_url)
@@ -233,6 +247,7 @@ async def prepare_video_upscale_task(
if isinstance(owner, ChatGenerationTask)
else VideoUpscaleTask.generation_record_id == owner.id
)
await apply_short_lock_timeout(db)
result = await db.execute(
select(VideoUpscaleTask).where(owner_filter).with_for_update().limit(1)
)
@@ -333,6 +348,7 @@ async def _claim(
stage: str,
chat_stage: str,
increment_attempt: bool,
execution_token: str,
lease_seconds: int | None = None,
) -> tuple[VideoUpscaleTask, VideoUpscaleOwner, dict[str, Any]] | None:
upscale, task = await _load_pair(db, upscale_task_id, for_update=True)
@@ -349,7 +365,7 @@ async def _claim(
snapshot = _snapshot(task)
upscale.status = VideoUpscaleTaskStatus.PROCESSING.value
upscale.stage = stage
upscale.lease_token = uuid.uuid4().hex
upscale.lease_token = execution_token
upscale.lease_until = _lease_until(lease_seconds)
upscale.started_at = upscale.started_at or _now()
upscale.next_retry_at = None
@@ -360,6 +376,21 @@ async def _claim(
return upscale, task, snapshot
async def _load_owned_pair(
db: AsyncSession,
*,
upscale_task_id: str,
execution_token: str,
for_update: bool = True,
) -> tuple[VideoUpscaleTask, VideoUpscaleOwner]:
upscale, task = await _load_pair(db, upscale_task_id, for_update=for_update)
if not upscale or not task:
raise RedisExecutionLockLost(f"超分任务或所属任务不存在: {upscale_task_id}")
if str(upscale.lease_token or "") != str(execution_token):
raise RedisExecutionLockLost(f"超分数据库执行租约已失效: {upscale_task_id}")
return upscale, task
async def _final_fail(
db: AsyncSession,
*,
@@ -381,6 +412,7 @@ async def _final_fail(
upscale.lease_token = None
await mark_owner_upscale_failed(db, task, error_message=error_message)
await db.commit()
await notify_owner_finished(db, task)
log_video_upscale_event(
event_type="upscale_retry_exhausted",
event_status="failed",
@@ -464,8 +496,13 @@ async def _queue_finalize(
upscale_task_id: str,
final_path: str,
reason: str,
execution_token: str,
execution_guard: ExecutionGuard | None = None,
) -> None:
upscale, task = await _load_pair(db, upscale_task_id, for_update=True)
await _ensure_guard(execution_guard)
upscale, task = await _load_owned_pair(
db, upscale_task_id=upscale_task_id, execution_token=execution_token, for_update=True
)
if not upscale or not task:
return
if not is_valid_file(final_path):
@@ -488,6 +525,7 @@ async def _queue_finalize(
task=task,
upscale_task=upscale,
action="finalize",
countdown=max(1, int(settings.VIDEO_UPSCALE_STAGE_HANDOFF_DELAY_SECONDS or 2)),
)
log_video_upscale_event(
event_type="upscale_finalize_enqueued",
@@ -498,8 +536,18 @@ async def _queue_finalize(
)
async def _finalize_success(db: AsyncSession, *, upscale_task_id: str, final_path: str) -> None:
upscale, task = await _load_pair(db, upscale_task_id, for_update=False)
async def _finalize_success(
db: AsyncSession,
*,
upscale_task_id: str,
final_path: str,
execution_token: str,
execution_guard: ExecutionGuard | None = None,
) -> None:
await _ensure_guard(execution_guard)
upscale, task = await _load_owned_pair(
db, upscale_task_id=upscale_task_id, execution_token=execution_token, for_update=False
)
if not upscale or not task:
return
snapshot = _snapshot(task)
@@ -550,7 +598,10 @@ async def _finalize_success(db: AsyncSession, *, upscale_task_id: str, final_pat
if not cover_url or not cover_path:
raise RuntimeError("超分最终视频封面生成失败")
upscale, task = await _load_pair(db, upscale_task_id, for_update=True)
await _ensure_guard(execution_guard)
upscale, task = await _load_owned_pair(
db, upscale_task_id=upscale_task_id, execution_token=execution_token, for_update=True
)
if not upscale or not task:
return
if owner_is_completed(task) and task.video_url:
@@ -569,7 +620,7 @@ async def _finalize_success(db: AsyncSession, *, upscale_task_id: str, final_pat
task.generated_at = now
task.error_message = None
if isinstance(task, ChatGenerationTask):
task.retry_count = 0
task.retry_count = int(task.manual_retry_count or 0)
upscale.status = VideoUpscaleTaskStatus.COMPLETED.value
upscale.stage = VideoUpscaleStage.COMPLETED.value
@@ -595,8 +646,6 @@ async def _finalize_success(db: AsyncSession, *, upscale_task_id: str, final_pat
generated_at=now,
)
await sync_chat_generation_task_media_token_snapshot(db, task)
await notify_chat_generation_task_finished(db, task)
await aggregate_parent_for_child(db, task)
else:
await record_generation_record_generated_resource(
db,
@@ -614,9 +663,12 @@ async def _finalize_success(db: AsyncSession, *, upscale_task_id: str, final_pat
if not delete_source_after_success:
upscale.source_delete_error = VIDEO_UPSCALE_SOURCE_RETAINED_MARKER
await db.commit()
if isinstance(task, ChatGenerationTask):
await notify_owner_finished(db, task)
if delete_source_after_success and source_path and os.path.abspath(source_path) != os.path.abspath(final_path):
removed = safe_remove(source_path)
cleanup_error = None if removed else f"源视频删除失败: {source_path}"
await apply_short_lock_timeout(db)
cleanup_result = await db.execute(
select(VideoUpscaleTask).where(VideoUpscaleTask.id == upscale_task_id).with_for_update().limit(1)
)
@@ -668,13 +720,16 @@ async def _finalize_success(db: AsyncSession, *, upscale_task_id: str, final_pat
)
async def run_local_upscale(db: AsyncSession, upscale_task_id: str) -> None:
async def run_local_upscale(
db: AsyncSession, upscale_task_id: str, *, execution_token: str, execution_guard: ExecutionGuard | None = None
) -> None:
claimed = await _claim(
db,
upscale_task_id=upscale_task_id,
stage=VideoUpscaleStage.LOCAL_PROCESSING.value,
chat_stage=ChatGenerationPipelineStage.UPSCALE_PROCESSING.value,
increment_attempt=True,
execution_token=execution_token,
lease_seconds=int(settings.VIDEO_UPSCALE_LOCAL_TIMEOUT_SECONDS or 3600) + 300,
)
if not claimed:
@@ -690,15 +745,21 @@ async def run_local_upscale(db: AsyncSession, upscale_task_id: str) -> None:
target_width=int(snapshot.get("target_width") or upscale.target_width),
target_height=int(snapshot.get("target_height") or upscale.target_height),
timeout_seconds=int(processor.get("timeout_seconds") or settings.VIDEO_UPSCALE_LOCAL_TIMEOUT_SECONDS),
execution_guard=execution_guard,
)
await _ensure_guard(execution_guard)
log_video_upscale_event(
event_type="upscale_local_success",
task=task,
upscale_task=upscale,
detail={"final_local_path": final_path},
)
except RedisExecutionLockError:
await db.rollback()
raise
except Exception as exc:
await db.rollback()
await _ensure_guard(execution_guard)
await _schedule_retry(
db,
upscale_task_id=upscale_task_id,
@@ -714,9 +775,15 @@ async def run_local_upscale(db: AsyncSession, upscale_task_id: str) -> None:
upscale_task_id=upscale_task_id,
final_path=final_path,
reason="local_upscale_completed",
execution_token=execution_token,
execution_guard=execution_guard,
)
except RedisExecutionLockError:
await db.rollback()
raise
except Exception as exc:
await db.rollback()
await _ensure_guard(execution_guard)
await _schedule_retry(
db,
upscale_task_id=upscale_task_id,
@@ -770,13 +837,17 @@ async def _select_remote_input(db: AsyncSession, upscale: VideoUpscaleTask, proc
return signed_url, VideoUpscaleInputSourceType.LOCAL_SIGNED.value
async def run_remote_submit(db: AsyncSession, upscale_task_id: str, *, count_attempt: bool = True) -> None:
async def run_remote_submit(
db: AsyncSession, upscale_task_id: str, *, count_attempt: bool = True, execution_token: str,
execution_guard: ExecutionGuard | None = None,
) -> None:
claimed = await _claim(
db,
upscale_task_id=upscale_task_id,
stage=VideoUpscaleStage.REMOTE_SUBMITTING.value,
chat_stage=ChatGenerationPipelineStage.UPSCALE_PROCESSING.value,
increment_attempt=count_attempt,
execution_token=execution_token,
lease_seconds=300,
)
if not claimed:
@@ -831,9 +902,10 @@ async def run_remote_submit(db: AsyncSession, upscale_task_id: str, *, count_att
processor=processor,
client_token=f"{upscale.id}-{int(upscale.attempt_count or 0)}-{int(upscale.input_source_fallback_count or 0)}",
)
upscale, task = await _load_pair(db, upscale_task_id, for_update=True)
if not upscale or not task:
return
await _ensure_guard(execution_guard)
upscale, task = await _load_owned_pair(
db, upscale_task_id=upscale_task_id, execution_token=execution_token, for_update=True
)
upscale.provider_task_id = submit_result.task_id
upscale.provider_submitted_at = _now()
upscale.provider_request_json = json.dumps(_sanitize_provider_payload(submit_result.request_payload), ensure_ascii=False, default=str)
@@ -863,6 +935,9 @@ async def run_remote_submit(db: AsyncSession, upscale_task_id: str, *, count_att
remote_request_id=submit_result.request_id,
detail={"provider_task_id": submit_result.task_id, "input_source_type": source_type},
)
except RedisExecutionLockError:
await db.rollback()
raise
except VolcMediaKitError as exc:
log_video_upscale_event(
event_type="upscale_provider_submit_failed",
@@ -875,6 +950,7 @@ async def run_remote_submit(db: AsyncSession, upscale_task_id: str, *, count_att
error=str(exc),
)
await db.rollback()
await _ensure_guard(execution_guard)
await _persist_provider_error_payload(
db,
upscale_task_id=upscale_task_id,
@@ -897,6 +973,7 @@ async def run_remote_submit(db: AsyncSession, upscale_task_id: str, *, count_att
error=str(exc),
)
await db.rollback()
await _ensure_guard(execution_guard)
await _schedule_retry(
db,
upscale_task_id=upscale_task_id,
@@ -905,13 +982,16 @@ async def run_remote_submit(db: AsyncSession, upscale_task_id: str, *, count_att
)
async def run_remote_poll(db: AsyncSession, upscale_task_id: str) -> None:
async def run_remote_poll(
db: AsyncSession, upscale_task_id: str, *, execution_token: str, execution_guard: ExecutionGuard | None = None
) -> None:
claimed = await _claim(
db,
upscale_task_id=upscale_task_id,
stage=VideoUpscaleStage.REMOTE_POLLING.value,
chat_stage=ChatGenerationPipelineStage.UPSCALE_POLLING.value,
increment_attempt=False,
execution_token=execution_token,
lease_seconds=300,
)
if not claimed:
@@ -939,9 +1019,10 @@ async def run_remote_poll(db: AsyncSession, upscale_task_id: str) -> None:
provider_task_id,
request_timeout_seconds=int(processor.get("request_timeout_seconds") or 30),
)
upscale, task = await _load_pair(db, upscale_task_id, for_update=True)
if not upscale or not task:
return
await _ensure_guard(execution_guard)
upscale, task = await _load_owned_pair(
db, upscale_task_id=upscale_task_id, execution_token=execution_token, for_update=True
)
upscale.provider_response_json = json.dumps(query_result.response_payload, ensure_ascii=False, default=str)
upscale.lease_until = None
upscale.lease_token = None
@@ -1001,6 +1082,7 @@ async def run_remote_poll(db: AsyncSession, upscale_task_id: str) -> None:
task=task,
upscale_task=upscale,
action="submit_local_source_fallback",
countdown=max(1, int(settings.VIDEO_UPSCALE_STAGE_HANDOFF_DELAY_SECONDS or 2)),
)
log_video_upscale_event(
event_type="upscale_source_fallback_local",
@@ -1071,6 +1153,7 @@ async def run_remote_poll(db: AsyncSession, upscale_task_id: str) -> None:
task=task,
upscale_task=upscale,
action="download_remote_result",
countdown=max(1, int(settings.VIDEO_UPSCALE_STAGE_HANDOFF_DELAY_SECONDS or 2)),
)
log_video_upscale_event(
event_type="upscale_provider_poll_success",
@@ -1079,6 +1162,9 @@ async def run_remote_poll(db: AsyncSession, upscale_task_id: str) -> None:
remote_request_id=query_result.request_id,
detail={"provider_output_url_expires_at": upscale.provider_output_url_expires_at},
)
except RedisExecutionLockError:
await db.rollback()
raise
except VolcMediaKitError as exc:
log_video_upscale_event(
event_type="upscale_provider_poll_failed",
@@ -1091,6 +1177,7 @@ async def run_remote_poll(db: AsyncSession, upscale_task_id: str) -> None:
error=str(exc),
)
await db.rollback()
await _ensure_guard(execution_guard)
await _persist_provider_error_payload(
db,
upscale_task_id=upscale_task_id,
@@ -1113,6 +1200,7 @@ async def run_remote_poll(db: AsyncSession, upscale_task_id: str) -> None:
error=str(exc),
)
await db.rollback()
await _ensure_guard(execution_guard)
await _schedule_retry(
db,
upscale_task_id=upscale_task_id,
@@ -1121,13 +1209,16 @@ async def run_remote_poll(db: AsyncSession, upscale_task_id: str) -> None:
)
async def run_finalize_upscale(db: AsyncSession, upscale_task_id: str) -> None:
async def run_finalize_upscale(
db: AsyncSession, upscale_task_id: str, *, execution_token: str, execution_guard: ExecutionGuard | None = None
) -> None:
claimed = await _claim(
db,
upscale_task_id=upscale_task_id,
stage=VideoUpscaleStage.FINALIZING.value,
chat_stage=ChatGenerationPipelineStage.UPSCALE_FINALIZING.value,
increment_attempt=False,
execution_token=execution_token,
lease_seconds=max(300, int(settings.VIDEO_COVER_TIMEOUT_SECONDS or 15) + 300),
)
if not claimed:
@@ -1143,13 +1234,19 @@ async def run_finalize_upscale(db: AsyncSession, upscale_task_id: str) -> None:
detail={"final_local_path": final_path},
)
try:
await _finalize_success(db, upscale_task_id=upscale_task_id, final_path=final_path)
await _finalize_success(
db, upscale_task_id=upscale_task_id, final_path=final_path,
execution_token=execution_token, execution_guard=execution_guard,
)
log_video_upscale_event(
event_type="upscale_finalize_success",
task=task,
upscale_task=upscale,
detail={"final_local_path": final_path},
)
except RedisExecutionLockError:
await db.rollback()
raise
except Exception as exc:
log_video_upscale_event(
event_type="upscale_finalize_failed",
@@ -1161,6 +1258,7 @@ async def run_finalize_upscale(db: AsyncSession, upscale_task_id: str) -> None:
error=str(exc),
)
await db.rollback()
await _ensure_guard(execution_guard)
await _schedule_retry(
db,
upscale_task_id=upscale_task_id,
@@ -1170,13 +1268,16 @@ async def run_finalize_upscale(db: AsyncSession, upscale_task_id: str) -> None:
)
async def run_remote_result_download(db: AsyncSession, upscale_task_id: str) -> None:
async def run_remote_result_download(
db: AsyncSession, upscale_task_id: str, *, execution_token: str, execution_guard: ExecutionGuard | None = None
) -> None:
claimed = await _claim(
db,
upscale_task_id=upscale_task_id,
stage=VideoUpscaleStage.RESULT_DOWNLOADING.value,
chat_stage=ChatGenerationPipelineStage.UPSCALE_DOWNLOADING.value,
increment_attempt=False,
execution_token=execution_token,
lease_seconds=int(settings.VIDEO_UPSCALE_REMOTE_RESULT_DOWNLOAD_TIMEOUT_SECONDS or 600) + 120,
)
if not claimed:
@@ -1198,6 +1299,7 @@ async def run_remote_result_download(db: AsyncSession, upscale_task_id: str) ->
task=task,
upscale_task=upscale,
action="renew_remote_result",
countdown=max(1, int(settings.VIDEO_UPSCALE_STAGE_HANDOFF_DELAY_SECONDS or 2)),
)
log_video_upscale_event(
event_type="upscale_provider_result_renew_query",
@@ -1228,13 +1330,18 @@ async def run_remote_result_download(db: AsyncSession, upscale_task_id: str) ->
output_url,
final_path,
int(settings.VIDEO_UPSCALE_REMOTE_RESULT_DOWNLOAD_TIMEOUT_SECONDS or 600),
execution_guard=execution_guard,
)
await _ensure_guard(execution_guard)
log_video_upscale_event(
event_type="upscale_output_download_success",
task=task,
upscale_task=upscale,
detail={"final_local_path": final_path, "file_size_bytes": safe_file_size(final_path)},
)
except RedisExecutionLockError:
await db.rollback()
raise
except Exception as exc:
log_video_upscale_event(
event_type="upscale_output_download_failed",
@@ -1246,6 +1353,7 @@ async def run_remote_result_download(db: AsyncSession, upscale_task_id: str) ->
error=str(exc),
)
await db.rollback()
await _ensure_guard(execution_guard)
expires_at = _aware(upscale.provider_output_url_expires_at)
action = "submit" if expires_at and expires_at <= _now() else "download"
await _schedule_retry(
@@ -1263,9 +1371,15 @@ async def run_remote_result_download(db: AsyncSession, upscale_task_id: str) ->
upscale_task_id=upscale_task_id,
final_path=final_path,
reason="remote_result_download_completed",
execution_token=execution_token,
execution_guard=execution_guard,
)
except RedisExecutionLockError:
await db.rollback()
raise
except Exception as exc:
await db.rollback()
await _ensure_guard(execution_guard)
await _schedule_retry(
db,
upscale_task_id=upscale_task_id,
@@ -1278,6 +1392,7 @@ async def run_remote_result_download(db: AsyncSession, upscale_task_id: str) ->
async def recover_video_upscale_tasks_once(db: AsyncSession) -> dict[str, Any]:
now = _now()
batch_size = max(1, int(settings.VIDEO_UPSCALE_RECOVERY_BATCH_SIZE or 50))
await apply_short_lock_timeout(db)
result = await db.execute(
select(VideoUpscaleTask)
.where(
@@ -1339,6 +1454,7 @@ async def recover_video_upscale_tasks_once(db: AsyncSession) -> dict[str, Any]:
continue
source_path = str(cleanup_upscale.source_local_path or "")
removed = safe_remove(source_path)
await apply_short_lock_timeout(db)
cleanup_result = await db.execute(
select(VideoUpscaleTask).where(VideoUpscaleTask.id == item.id).with_for_update().limit(1)
)