1142 lines
43 KiB
Python
1142 lines
43 KiB
Python
from __future__ import annotations
|
||
|
||
import argparse
|
||
import asyncio
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import re
|
||
import sys
|
||
import time
|
||
import uuid
|
||
from dataclasses import dataclass
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
from typing import Any, Literal
|
||
from urllib.parse import unquote, urlsplit
|
||
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.config import settings
|
||
from app.enums.video_upscale import (
|
||
REMOTE_PROCESSOR_KEYS,
|
||
VIDEO_UPSCALE_RESOLUTIONS,
|
||
VideoUpscaleProcessorKey,
|
||
normalize_video_upscale_resolution,
|
||
)
|
||
from app.models.base import async_session
|
||
from app.models.chat_generation_task import ChatGenerationTask
|
||
from app.models.generated_resource import GeneratedResource
|
||
from app.models.generation_record import GenerationRecord
|
||
from app.services.operation_log_service import (
|
||
log_operation_error,
|
||
log_operation_event,
|
||
log_remote_api_event,
|
||
)
|
||
from app.services.resource_accounting_service import (
|
||
SOURCE_MODEL_CHAT_TASK,
|
||
SOURCE_MODEL_GENERATION_RECORD,
|
||
)
|
||
from app.services.video_upscale.media_service import (
|
||
build_local_source_signed_url,
|
||
download_video_to_path,
|
||
is_valid_file,
|
||
)
|
||
from app.services.video_upscale.snapshot_service import calculate_target_dimensions
|
||
from app.services.video_upscale.volc_service import (
|
||
VolcMediaKitError,
|
||
query_task,
|
||
submit_video_enhance,
|
||
)
|
||
|
||
|
||
OwnerType = Literal["chat_generation_task", "generation_record"]
|
||
|
||
LOG_DOMAIN = "video_upscale_cli"
|
||
LOG_MODULE = "batch_video_upscale"
|
||
LOG_SOURCE = "cli"
|
||
STATE_VERSION = 1
|
||
|
||
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
||
|
||
DEFAULT_OUTPUT_DIR = str(
|
||
PROJECT_ROOT / "storage" / "uploads" / "videos"
|
||
)
|
||
|
||
VERSION_PATTERN = re.compile(r"^[A-Za-z0-9_-]+$")
|
||
ID_PATTERN = re.compile(r"^[A-Za-z0-9_-]+$")
|
||
|
||
OWNER_TYPE_CHAT: OwnerType = "chat_generation_task"
|
||
OWNER_TYPE_RECORD: OwnerType = "generation_record"
|
||
OWNER_TYPES: tuple[OwnerType, ...] = (OWNER_TYPE_CHAT, OWNER_TYPE_RECORD)
|
||
|
||
REMOTE_PROCESSOR_CHOICES = tuple(sorted(REMOTE_PROCESSOR_KEYS))
|
||
|
||
|
||
class RemoteTaskFailedError(RuntimeError):
|
||
pass
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class OwnerRef:
|
||
owner_type: OwnerType
|
||
owner_id: str
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class ResolvedItem:
|
||
owner_type: OwnerType
|
||
owner_id: str
|
||
user_id: str | None
|
||
project_id: str | None
|
||
generation_mode: str | None
|
||
source_local_path: str | None
|
||
source_remote_url: str | None
|
||
source_resource_id: str | None
|
||
source_resolution: str | None
|
||
source_aspect_ratio: str | None
|
||
output_path: str
|
||
|
||
|
||
@dataclass(slots=True)
|
||
class ItemResult:
|
||
owner_type: OwnerType
|
||
owner_id: str
|
||
status: str
|
||
output_path: str | None = None
|
||
provider_task_id: str | None = None
|
||
message: str | None = None
|
||
|
||
def as_dict(self) -> dict[str, Any]:
|
||
return {
|
||
"owner_type": self.owner_type,
|
||
"owner_id": self.owner_id,
|
||
"status": self.status,
|
||
"output_path": self.output_path,
|
||
"provider_task_id": self.provider_task_id,
|
||
"message": self.message,
|
||
}
|
||
|
||
|
||
class StateStore:
|
||
"""进程内并发安全、落盘原子替换的 CLI 断点状态。"""
|
||
|
||
def __init__(self, path: Path):
|
||
self.path = path
|
||
self._lock = asyncio.Lock()
|
||
self._data = self._load()
|
||
|
||
def _load(self) -> dict[str, Any]:
|
||
if not self.path.exists():
|
||
return {"version": STATE_VERSION, "items": {}}
|
||
try:
|
||
raw = json.loads(self.path.read_text(encoding="utf-8"))
|
||
except Exception as exc:
|
||
raise RuntimeError(f"状态文件读取失败: {self.path}: {exc}") from exc
|
||
if not isinstance(raw, dict) or not isinstance(raw.get("items"), dict):
|
||
raise RuntimeError(f"状态文件格式不合法: {self.path}")
|
||
raw.setdefault("version", STATE_VERSION)
|
||
return raw
|
||
|
||
async def get(self, key: str) -> dict[str, Any]:
|
||
async with self._lock:
|
||
value = self._data.get("items", {}).get(key)
|
||
return dict(value) if isinstance(value, dict) else {}
|
||
|
||
async def update(self, key: str, **values: Any) -> dict[str, Any]:
|
||
async with self._lock:
|
||
items = self._data.setdefault("items", {})
|
||
current = dict(items.get(key) or {})
|
||
current.update(values)
|
||
current["updated_at"] = datetime.now(timezone.utc).isoformat()
|
||
items[key] = current
|
||
self._write_unlocked()
|
||
return dict(current)
|
||
|
||
def _write_unlocked(self) -> None:
|
||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||
temp_path = self.path.with_name(f".{self.path.name}.{os.getpid()}.tmp")
|
||
temp_path.write_text(
|
||
json.dumps(self._data, ensure_ascii=False, indent=2, default=str),
|
||
encoding="utf-8",
|
||
)
|
||
os.replace(temp_path, self.path)
|
||
|
||
|
||
def _parser() -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser(
|
||
description=(
|
||
"按 ChatGenerationTask / GenerationRecord ID 调用现有火山视频超分 API，"
|
||
"轮询并下载结果；不修改数据库、原视频或正式 VideoUpscaleTask。"
|
||
)
|
||
)
|
||
parser.add_argument(
|
||
"--item",
|
||
action="append",
|
||
default=[],
|
||
metavar="OWNER_TYPE:ID",
|
||
help=(
|
||
"单条记录，可重复传入；例如 "
|
||
"--item chat_generation_task:001xxx 或 --item generation_record:001xxx"
|
||
),
|
||
)
|
||
parser.add_argument(
|
||
"--owner-type",
|
||
choices=OWNER_TYPES,
|
||
help="同类型批量传参时使用，需配合 --id/--ids",
|
||
)
|
||
parser.add_argument("--id", action="append", default=[], help="记录 ID，可重复传入")
|
||
parser.add_argument("--ids", default="", help="逗号分隔的记录 ID，需配合 --owner-type")
|
||
parser.add_argument(
|
||
"--processor-key",
|
||
required=True,
|
||
choices=REMOTE_PROCESSOR_CHOICES,
|
||
help="仅允许火山远程处理器，不允许 local_ffmpeg_crop_v1",
|
||
)
|
||
parser.add_argument("--output-version", required=True, help="输出版本，例如 v1、v2、standard-test")
|
||
parser.add_argument("--target-aspect-ratio", required=True, help="目标比例，例如 9:16、16:9、1:1")
|
||
parser.add_argument(
|
||
"--target-resolution",
|
||
required=True,
|
||
help=f"目标分辨率，仅支持 {', '.join(VIDEO_UPSCALE_RESOLUTIONS)}",
|
||
)
|
||
parser.add_argument("--output-dir", default=DEFAULT_OUTPUT_DIR, help="结果保存目录")
|
||
parser.add_argument(
|
||
"--max-concurrency",
|
||
type=int,
|
||
default=1,
|
||
help="并行处理数量，默认 1；建议根据火山账号并发限制谨慎调整",
|
||
)
|
||
parser.add_argument(
|
||
"--poll-interval",
|
||
type=int,
|
||
default=max(5, int(settings.VIDEO_UPSCALE_REMOTE_POLL_INTERVAL_SECONDS or 30)),
|
||
help="轮询间隔秒数",
|
||
)
|
||
parser.add_argument(
|
||
"--poll-timeout",
|
||
type=int,
|
||
default=max(60, int(settings.VIDEO_UPSCALE_REMOTE_POLL_TIMEOUT_SECONDS or 7200)),
|
||
help="单条任务最长轮询秒数",
|
||
)
|
||
parser.add_argument(
|
||
"--request-timeout",
|
||
type=int,
|
||
default=max(3, int(settings.VIDEO_UPSCALE_REMOTE_REQUEST_TIMEOUT_SECONDS or 30)),
|
||
help="火山提交/查询请求超时秒数",
|
||
)
|
||
parser.add_argument(
|
||
"--download-timeout",
|
||
type=int,
|
||
default=max(30, int(settings.VIDEO_UPSCALE_REMOTE_RESULT_DOWNLOAD_TIMEOUT_SECONDS or 600)),
|
||
help="结果下载超时秒数",
|
||
)
|
||
parser.add_argument(
|
||
"--source-url-expire",
|
||
type=int,
|
||
default=max(600, int(settings.VIDEO_UPSCALE_LOCAL_SIGNED_URL_EXPIRE_SECONDS or 7200)),
|
||
help="本地原视频签名 URL 有效期秒数",
|
||
)
|
||
parser.add_argument("--state-file", default="", help="自定义断点状态文件路径")
|
||
parser.add_argument("--dry-run", action="store_true", help="仅查询记录和展示计划，不调用 API")
|
||
parser.add_argument(
|
||
"--force-resubmit",
|
||
action="store_true",
|
||
help="忽略已有状态和结果文件，重新提交；不会修改原视频",
|
||
)
|
||
return parser
|
||
|
||
|
||
def _normalize_owner_type(value: str) -> OwnerType:
|
||
normalized = str(value or "").strip().lower()
|
||
if normalized not in OWNER_TYPES:
|
||
raise ValueError(f"不支持的 owner_type: {value}")
|
||
return normalized # type: ignore[return-value]
|
||
|
||
|
||
def _validate_id(value: str) -> str:
|
||
text = str(value or "").strip()
|
||
if not text:
|
||
raise ValueError("记录 ID 不能为空")
|
||
if not ID_PATTERN.fullmatch(text):
|
||
raise ValueError(f"记录 ID 只能包含字母、数字、下划线和短横线: {text}")
|
||
return text
|
||
|
||
|
||
def _collect_owner_refs(args: argparse.Namespace) -> list[OwnerRef]:
|
||
refs: list[OwnerRef] = []
|
||
|
||
for raw in args.item:
|
||
text = str(raw or "").strip()
|
||
if ":" not in text:
|
||
raise ValueError(f"--item 格式错误，应为 OWNER_TYPE:ID: {text}")
|
||
owner_type_text, owner_id_text = text.split(":", 1)
|
||
refs.append(OwnerRef(_normalize_owner_type(owner_type_text), _validate_id(owner_id_text)))
|
||
|
||
ids = [str(value).strip() for value in args.id if str(value).strip()]
|
||
ids.extend(value.strip() for value in str(args.ids or "").split(",") if value.strip())
|
||
if ids:
|
||
if not args.owner_type:
|
||
raise ValueError("使用 --id/--ids 时必须同时传入 --owner-type")
|
||
owner_type = _normalize_owner_type(args.owner_type)
|
||
refs.extend(OwnerRef(owner_type, _validate_id(owner_id)) for owner_id in ids)
|
||
elif args.owner_type:
|
||
raise ValueError("传入 --owner-type 后必须至少传入一个 --id 或 --ids")
|
||
|
||
unique: list[OwnerRef] = []
|
||
seen: set[tuple[str, str]] = set()
|
||
for ref in refs:
|
||
key = (ref.owner_type, ref.owner_id)
|
||
if key not in seen:
|
||
seen.add(key)
|
||
unique.append(ref)
|
||
if not unique:
|
||
raise ValueError("至少传入一个 --item，或使用 --owner-type 配合 --id/--ids")
|
||
return unique
|
||
|
||
|
||
def _validate_args(args: argparse.Namespace) -> None:
|
||
if not args.dry_run and not str(settings.VOLC_API_KEY or "").strip():
|
||
raise RuntimeError("VOLC_API_KEY 未配置，请先在 config.py 对应环境配置中设置")
|
||
if args.processor_key not in REMOTE_PROCESSOR_KEYS:
|
||
raise ValueError(f"CLI 只允许远程处理器: {args.processor_key}")
|
||
if args.processor_key == VideoUpscaleProcessorKey.LOCAL_FFMPEG_CROP_V1.value:
|
||
raise ValueError("该 CLI 不允许使用 FFmpeg 本地处理器")
|
||
if not VERSION_PATTERN.fullmatch(str(args.output_version or "")):
|
||
raise ValueError("--output-version 只能包含字母、数字、下划线和短横线")
|
||
if int(args.max_concurrency) < 1 or int(args.max_concurrency) > 20:
|
||
raise ValueError("--max-concurrency 必须在 1 到 20 之间")
|
||
if int(args.poll_interval) < 5:
|
||
raise ValueError("--poll-interval 不能小于 5 秒")
|
||
if int(args.poll_timeout) < 60:
|
||
raise ValueError("--poll-timeout 不能小于 60 秒")
|
||
if int(args.request_timeout) < 3:
|
||
raise ValueError("--request-timeout 不能小于 3 秒")
|
||
if int(args.download_timeout) < 30:
|
||
raise ValueError("--download-timeout 不能小于 30 秒")
|
||
if int(args.source_url_expire) < 600:
|
||
raise ValueError("--source-url-expire 不能小于 600 秒")
|
||
|
||
|
||
def _normalize_target_resolution(value: str) -> str:
|
||
normalized = normalize_video_upscale_resolution(value)
|
||
if normalized not in VIDEO_UPSCALE_RESOLUTIONS:
|
||
raise ValueError(
|
||
f"不支持的目标分辨率: {value}，仅支持 {', '.join(VIDEO_UPSCALE_RESOLUTIONS)}"
|
||
)
|
||
return normalized
|
||
|
||
|
||
def _source_model(owner_type: OwnerType) -> str:
|
||
return SOURCE_MODEL_CHAT_TASK if owner_type == OWNER_TYPE_CHAT else SOURCE_MODEL_GENERATION_RECORD
|
||
|
||
|
||
def _owner_model(owner_type: OwnerType):
|
||
return ChatGenerationTask if owner_type == OWNER_TYPE_CHAT else GenerationRecord
|
||
|
||
|
||
def _build_output_path(output_dir: Path, output_version: str, owner_id: str, target_resolution: str) -> Path:
|
||
return output_dir / f"{output_version}-{owner_id}_{target_resolution}.mp4"
|
||
|
||
|
||
def _state_key(
|
||
*,
|
||
ref: OwnerRef,
|
||
processor_key: str,
|
||
output_version: str,
|
||
target_aspect_ratio: str,
|
||
target_resolution: str,
|
||
) -> str:
|
||
return ":".join(
|
||
(
|
||
ref.owner_type,
|
||
ref.owner_id,
|
||
processor_key,
|
||
output_version,
|
||
target_aspect_ratio.replace(":", "x"),
|
||
target_resolution,
|
||
)
|
||
)
|
||
|
||
|
||
def _client_token(
|
||
*,
|
||
state_key: str,
|
||
force_resubmit: bool,
|
||
) -> str:
|
||
token_source = state_key
|
||
if force_resubmit:
|
||
token_source += f":force:{time.time_ns()}"
|
||
digest = hashlib.sha256(token_source.encode("utf-8")).hexdigest()
|
||
return f"vg-cli-{digest[:48]}"
|
||
|
||
|
||
def _url_to_local_path(value: str | None) -> str | None:
|
||
text = str(value or "").strip()
|
||
if not text:
|
||
return None
|
||
|
||
direct = Path(text.split("?", 1)[0])
|
||
if direct.is_absolute() and is_valid_file(str(direct)):
|
||
return str(direct)
|
||
|
||
parts = urlsplit(text)
|
||
path_text = unquote(parts.path if parts.scheme or parts.netloc else text.split("?", 1)[0])
|
||
normalized = "/" + path_text.lstrip("/")
|
||
prefix = "/generate/videos/"
|
||
if normalized.startswith(prefix):
|
||
relative = normalized[len(prefix):]
|
||
candidate = Path(settings.STORAGE_LOCAL_PATH) / relative
|
||
if is_valid_file(str(candidate)):
|
||
return str(candidate)
|
||
return None
|
||
|
||
|
||
def _http_url(value: str | None) -> str | None:
|
||
text = str(value or "").strip()
|
||
if text.startswith("http://") or text.startswith("https://"):
|
||
return text
|
||
return None
|
||
|
||
|
||
def _infer_local_path(owner: Any) -> str | None:
|
||
created_at = getattr(owner, "created_at", None)
|
||
if created_at is not None:
|
||
date_dir = created_at.strftime("%Y/%m/%d")
|
||
candidate = Path(settings.STORAGE_LOCAL_PATH) / date_dir / f"{owner.id}.mp4"
|
||
if is_valid_file(str(candidate)):
|
||
return str(candidate)
|
||
candidate = Path(settings.STORAGE_LOCAL_PATH) / f"{owner.id}.mp4"
|
||
if is_valid_file(str(candidate)):
|
||
return str(candidate)
|
||
return None
|
||
|
||
|
||
def _resolve_source(owner: Any, resource: GeneratedResource | None) -> tuple[str | None, str | None]:
|
||
local_candidates = [
|
||
getattr(resource, "storage_path", None) if resource else None,
|
||
_url_to_local_path(getattr(resource, "resource_url", None) if resource else None),
|
||
_url_to_local_path(getattr(owner, "video_url", None)),
|
||
_infer_local_path(owner),
|
||
]
|
||
for candidate in local_candidates:
|
||
if candidate and is_valid_file(str(candidate)):
|
||
return str(candidate), None
|
||
|
||
remote_candidates = [
|
||
getattr(resource, "remote_url", None) if resource else None,
|
||
getattr(resource, "resource_url", None) if resource else None,
|
||
getattr(owner, "video_url", None),
|
||
getattr(owner, "remote_result_url", None),
|
||
]
|
||
for candidate in remote_candidates:
|
||
remote_url = _http_url(candidate)
|
||
if remote_url:
|
||
return None, remote_url
|
||
return None, None
|
||
|
||
|
||
def _safe_log_event(
|
||
*,
|
||
event_type: str,
|
||
event_status: str = "success",
|
||
item: ResolvedItem | None = None,
|
||
task_id: str | None = None,
|
||
message: str | None = None,
|
||
detail: dict[str, Any] | None = None,
|
||
error: str | None = None,
|
||
) -> None:
|
||
final_detail = dict(detail or {})
|
||
if item is not None:
|
||
final_detail.update(
|
||
{
|
||
"owner_type": item.owner_type,
|
||
"owner_id": item.owner_id,
|
||
"source_resource_id": item.source_resource_id,
|
||
"source_local_path": item.source_local_path,
|
||
"source_remote_url": item.source_remote_url,
|
||
"output_path": item.output_path,
|
||
}
|
||
)
|
||
log_operation_event(
|
||
domain=LOG_DOMAIN,
|
||
module=LOG_MODULE,
|
||
source=LOG_SOURCE,
|
||
event_type=event_type,
|
||
event_status=event_status,
|
||
user_id=item.user_id if item else None,
|
||
project_id=item.project_id if item else None,
|
||
task_id=task_id or (item.owner_id if item else None),
|
||
message=message,
|
||
detail=final_detail,
|
||
error=error,
|
||
)
|
||
|
||
|
||
def _safe_log_error(
|
||
*,
|
||
event_type: str,
|
||
exc: BaseException,
|
||
item: ResolvedItem | None = None,
|
||
task_id: str | None = None,
|
||
message: str | None = None,
|
||
detail: dict[str, Any] | None = None,
|
||
) -> None:
|
||
final_detail = dict(detail or {})
|
||
if item is not None:
|
||
final_detail.update(
|
||
{
|
||
"owner_type": item.owner_type,
|
||
"owner_id": item.owner_id,
|
||
"source_resource_id": item.source_resource_id,
|
||
"source_local_path": item.source_local_path,
|
||
"source_remote_url": item.source_remote_url,
|
||
"output_path": item.output_path,
|
||
}
|
||
)
|
||
log_operation_error(
|
||
domain=LOG_DOMAIN,
|
||
module=LOG_MODULE,
|
||
source=LOG_SOURCE,
|
||
event_type=event_type,
|
||
user_id=item.user_id if item else None,
|
||
project_id=item.project_id if item else None,
|
||
task_id=task_id or (item.owner_id if item else None),
|
||
message=message,
|
||
detail=final_detail,
|
||
exc=exc,
|
||
)
|
||
|
||
|
||
async def _load_resources(
|
||
db: AsyncSession,
|
||
*,
|
||
owner_type: OwnerType,
|
||
owner_ids: list[str],
|
||
) -> dict[str, GeneratedResource]:
|
||
if not owner_ids:
|
||
return {}
|
||
result = await db.execute(
|
||
select(GeneratedResource)
|
||
.where(
|
||
GeneratedResource.source_model == _source_model(owner_type),
|
||
GeneratedResource.source_id.in_(owner_ids),
|
||
GeneratedResource.resource_type == "video",
|
||
GeneratedResource.deleted_at.is_(None),
|
||
)
|
||
.order_by(GeneratedResource.source_id.asc(), GeneratedResource.created_at.desc())
|
||
)
|
||
output: dict[str, GeneratedResource] = {}
|
||
for resource in result.scalars().all():
|
||
output.setdefault(str(resource.source_id), resource)
|
||
return output
|
||
|
||
|
||
async def _resolve_items(
|
||
db: AsyncSession,
|
||
*,
|
||
refs: list[OwnerRef],
|
||
output_dir: Path,
|
||
output_version: str,
|
||
target_resolution: str,
|
||
) -> tuple[list[ResolvedItem], list[ItemResult]]:
|
||
resolved_by_key: dict[tuple[str, str], ResolvedItem] = {}
|
||
errors: list[ItemResult] = []
|
||
|
||
for owner_type in OWNER_TYPES:
|
||
type_refs = [ref for ref in refs if ref.owner_type == owner_type]
|
||
if not type_refs:
|
||
continue
|
||
owner_ids = [ref.owner_id for ref in type_refs]
|
||
model = _owner_model(owner_type)
|
||
result = await db.execute(
|
||
select(model).where(
|
||
model.id.in_(owner_ids),
|
||
model.deleted_at.is_(None),
|
||
)
|
||
)
|
||
owners = {str(owner.id): owner for owner in result.scalars().all()}
|
||
resources = await _load_resources(db, owner_type=owner_type, owner_ids=owner_ids)
|
||
|
||
for ref in type_refs:
|
||
owner = owners.get(ref.owner_id)
|
||
output_path = str(_build_output_path(output_dir, output_version, ref.owner_id, target_resolution))
|
||
if owner is None:
|
||
errors.append(
|
||
ItemResult(
|
||
owner_type=ref.owner_type,
|
||
owner_id=ref.owner_id,
|
||
status="owner_missing",
|
||
output_path=output_path,
|
||
message="记录不存在或已软删除",
|
||
)
|
||
)
|
||
continue
|
||
if str(getattr(owner, "gen_type", "") or "").lower() != "video":
|
||
errors.append(
|
||
ItemResult(
|
||
owner_type=ref.owner_type,
|
||
owner_id=ref.owner_id,
|
||
status="not_video",
|
||
output_path=output_path,
|
||
message=f"记录 gen_type={getattr(owner, 'gen_type', None)}，不是视频",
|
||
)
|
||
)
|
||
continue
|
||
|
||
resource = resources.get(ref.owner_id)
|
||
source_local_path, source_remote_url = _resolve_source(owner, resource)
|
||
if not source_local_path and not source_remote_url:
|
||
errors.append(
|
||
ItemResult(
|
||
owner_type=ref.owner_type,
|
||
owner_id=ref.owner_id,
|
||
status="source_missing",
|
||
output_path=output_path,
|
||
message="未找到可用的本地原视频或远程视频 URL",
|
||
)
|
||
)
|
||
continue
|
||
|
||
resolved_by_key[(ref.owner_type, ref.owner_id)] = ResolvedItem(
|
||
owner_type=ref.owner_type,
|
||
owner_id=ref.owner_id,
|
||
user_id=str(getattr(owner, "user_id", "") or "") or None,
|
||
project_id=str(getattr(owner, "project_id", "") or "") or None,
|
||
generation_mode=str(getattr(owner, "generation_mode", "") or "") or None,
|
||
source_local_path=source_local_path,
|
||
source_remote_url=source_remote_url,
|
||
source_resource_id=str(resource.id) if resource else None,
|
||
source_resolution=str(getattr(owner, "resolution", "") or "") or None,
|
||
source_aspect_ratio=str(getattr(owner, "aspect_ratio", "") or "") or None,
|
||
output_path=output_path,
|
||
)
|
||
|
||
ordered = [resolved_by_key[(ref.owner_type, ref.owner_id)] for ref in refs if (ref.owner_type, ref.owner_id) in resolved_by_key]
|
||
return ordered, errors
|
||
|
||
|
||
def _processor_config(args: argparse.Namespace) -> dict[str, Any]:
|
||
return {
|
||
"request_timeout_seconds": int(args.request_timeout),
|
||
"poll_timeout_seconds": int(args.poll_timeout),
|
||
"poll_interval_seconds": int(args.poll_interval),
|
||
"source_url_expire_seconds": int(args.source_url_expire),
|
||
"bitrate_level": "medium",
|
||
"scene": "aigc",
|
||
}
|
||
|
||
|
||
async def _process_item(
|
||
*,
|
||
item: ResolvedItem,
|
||
args: argparse.Namespace,
|
||
target_resolution: str,
|
||
target_width: int,
|
||
target_height: int,
|
||
state_store: StateStore,
|
||
semaphore: asyncio.Semaphore,
|
||
index: int,
|
||
total: int,
|
||
) -> ItemResult:
|
||
async with semaphore:
|
||
ref = OwnerRef(item.owner_type, item.owner_id)
|
||
state_key = _state_key(
|
||
ref=ref,
|
||
processor_key=args.processor_key,
|
||
output_version=args.output_version,
|
||
target_aspect_ratio=args.target_aspect_ratio,
|
||
target_resolution=target_resolution,
|
||
)
|
||
prefix = f"[{index}/{total}] {item.owner_type}:{item.owner_id}"
|
||
output_path = item.output_path
|
||
existing_state = await state_store.get(state_key)
|
||
|
||
if is_valid_file(output_path) and not args.force_resubmit:
|
||
message = f"目标文件已存在，跳过: {output_path}"
|
||
print(f"{prefix} {message}")
|
||
await state_store.update(
|
||
state_key,
|
||
status="completed",
|
||
output_path=output_path,
|
||
owner_type=item.owner_type,
|
||
owner_id=item.owner_id,
|
||
)
|
||
_safe_log_event(
|
||
event_type="item_skipped_existing",
|
||
event_status="skipped",
|
||
item=item,
|
||
message=message,
|
||
)
|
||
return ItemResult(item.owner_type, item.owner_id, "skipped_existing", output_path=output_path, message=message)
|
||
|
||
if existing_state.get("status") == "failed" and not args.force_resubmit:
|
||
message = "该参数组合已有失败状态；如需重新提交请增加 --force-resubmit"
|
||
print(f"{prefix} {message}")
|
||
_safe_log_event(
|
||
event_type="item_skipped_failed_state",
|
||
event_status="skipped",
|
||
item=item,
|
||
task_id=str(existing_state.get("provider_task_id") or "") or None,
|
||
message=message,
|
||
detail={"state": existing_state},
|
||
)
|
||
return ItemResult(
|
||
item.owner_type,
|
||
item.owner_id,
|
||
"skipped_failed_state",
|
||
output_path=output_path,
|
||
provider_task_id=str(existing_state.get("provider_task_id") or "") or None,
|
||
message=message,
|
||
)
|
||
|
||
provider_task_id = None if args.force_resubmit else str(existing_state.get("provider_task_id") or "").strip() or None
|
||
result_url = None if args.force_resubmit else str(existing_state.get("result_url") or "").strip() or None
|
||
|
||
try:
|
||
if result_url:
|
||
print(f"{prefix} 恢复结果下载")
|
||
elif provider_task_id:
|
||
print(f"{prefix} 恢复轮询 provider_task_id={provider_task_id}")
|
||
else:
|
||
source_url = item.source_remote_url
|
||
input_source_type = "remote_url"
|
||
if item.source_local_path:
|
||
source_url = build_local_source_signed_url(
|
||
item.source_local_path,
|
||
expire_seconds=int(args.source_url_expire),
|
||
)
|
||
input_source_type = "local_signed_url"
|
||
if not source_url:
|
||
raise RuntimeError("无法构造火山可访问的原视频 URL")
|
||
|
||
client_token = _client_token(state_key=state_key, force_resubmit=bool(args.force_resubmit))
|
||
print(
|
||
f"{prefix} 提交 {args.processor_key}，目标 "
|
||
f"{target_resolution} {args.target_aspect_ratio} ({target_width}x{target_height})"
|
||
)
|
||
_safe_log_event(
|
||
event_type="remote_submit_started",
|
||
event_status="processing",
|
||
item=item,
|
||
detail={
|
||
"processor_key": args.processor_key,
|
||
"target_resolution": target_resolution,
|
||
"target_aspect_ratio": args.target_aspect_ratio,
|
||
"target_width": target_width,
|
||
"target_height": target_height,
|
||
"input_source_type": input_source_type,
|
||
"client_token": client_token,
|
||
},
|
||
)
|
||
submit_result = await submit_video_enhance(
|
||
processor_key=args.processor_key,
|
||
video_url=source_url,
|
||
target_resolution=target_resolution,
|
||
target_width=target_width,
|
||
target_height=target_height,
|
||
processor=_processor_config(args),
|
||
client_token=client_token,
|
||
)
|
||
provider_task_id = submit_result.task_id
|
||
await state_store.update(
|
||
state_key,
|
||
status="submitted",
|
||
owner_type=item.owner_type,
|
||
owner_id=item.owner_id,
|
||
provider_task_id=provider_task_id,
|
||
output_path=output_path,
|
||
processor_key=args.processor_key,
|
||
output_version=args.output_version,
|
||
target_aspect_ratio=args.target_aspect_ratio,
|
||
target_resolution=target_resolution,
|
||
target_width=target_width,
|
||
target_height=target_height,
|
||
source_local_path=item.source_local_path,
|
||
source_remote_url=item.source_remote_url,
|
||
client_token=client_token,
|
||
)
|
||
log_remote_api_event(
|
||
domain=LOG_DOMAIN,
|
||
module=LOG_MODULE,
|
||
source=LOG_SOURCE,
|
||
event_type="remote_submit_succeeded",
|
||
event_status="success",
|
||
remote_action="submit_video_enhance",
|
||
remote_request_id=submit_result.request_id or provider_task_id,
|
||
task_id=item.owner_id,
|
||
user_id=item.user_id,
|
||
project_id=item.project_id,
|
||
request=submit_result.request_payload,
|
||
response=submit_result.response_payload,
|
||
detail={
|
||
"owner_type": item.owner_type,
|
||
"owner_id": item.owner_id,
|
||
"provider_task_id": provider_task_id,
|
||
},
|
||
)
|
||
print(f"{prefix} 提交成功 provider_task_id={provider_task_id}")
|
||
|
||
if not result_url:
|
||
if not provider_task_id:
|
||
raise RuntimeError("缺少 provider_task_id")
|
||
deadline = time.monotonic() + int(args.poll_timeout)
|
||
last_status: str | None = None
|
||
while True:
|
||
if time.monotonic() >= deadline:
|
||
raise TimeoutError(f"火山任务轮询超过 {args.poll_timeout} 秒")
|
||
query_result = await query_task(
|
||
provider_task_id,
|
||
request_timeout_seconds=int(args.request_timeout),
|
||
)
|
||
if query_result.status != last_status:
|
||
print(f"{prefix} 火山状态={query_result.status}")
|
||
log_remote_api_event(
|
||
domain=LOG_DOMAIN,
|
||
module=LOG_MODULE,
|
||
source=LOG_SOURCE,
|
||
event_type="remote_poll_status_changed",
|
||
event_status=query_result.status,
|
||
remote_action="query_task",
|
||
remote_request_id=query_result.request_id or provider_task_id,
|
||
task_id=item.owner_id,
|
||
user_id=item.user_id,
|
||
project_id=item.project_id,
|
||
response=query_result.response_payload,
|
||
detail={
|
||
"owner_type": item.owner_type,
|
||
"owner_id": item.owner_id,
|
||
"provider_task_id": provider_task_id,
|
||
"previous_status": last_status,
|
||
"current_status": query_result.status,
|
||
},
|
||
)
|
||
last_status = query_result.status
|
||
|
||
if query_result.status == "running":
|
||
await state_store.update(
|
||
state_key,
|
||
status="processing",
|
||
provider_task_id=provider_task_id,
|
||
last_remote_status="running",
|
||
)
|
||
await asyncio.sleep(int(args.poll_interval))
|
||
continue
|
||
if query_result.status == "failed":
|
||
remote_error = query_result.error or {}
|
||
raise RemoteTaskFailedError(
|
||
"火山超分任务失败: "
|
||
+ json.dumps(remote_error, ensure_ascii=False, default=str)
|
||
)
|
||
|
||
result_payload = query_result.result or {}
|
||
result_url = str(result_payload.get("video_url") or "").strip()
|
||
if not result_url:
|
||
raise RuntimeError("火山任务已完成但未返回 result.video_url")
|
||
await state_store.update(
|
||
state_key,
|
||
status="result_ready",
|
||
provider_task_id=provider_task_id,
|
||
result_url=result_url,
|
||
result_expires_at=query_result.expires_at,
|
||
)
|
||
break
|
||
|
||
print(f"{prefix} 下载结果到 {output_path}")
|
||
_safe_log_event(
|
||
event_type="download_started",
|
||
event_status="processing",
|
||
item=item,
|
||
task_id=provider_task_id,
|
||
detail={"provider_task_id": provider_task_id},
|
||
)
|
||
download_target = output_path
|
||
replacement_path: str | None = None
|
||
if args.force_resubmit and is_valid_file(output_path):
|
||
output_obj = Path(output_path)
|
||
replacement_path = str(
|
||
output_obj.with_name(f"{output_obj.stem}.{uuid.uuid4().hex}.replacement.mp4")
|
||
)
|
||
download_target = replacement_path
|
||
try:
|
||
downloaded_path = await download_video_to_path(
|
||
result_url,
|
||
download_target,
|
||
timeout_seconds=int(args.download_timeout),
|
||
)
|
||
if not is_valid_file(downloaded_path):
|
||
raise RuntimeError(f"结果下载后文件不存在或为空: {downloaded_path}")
|
||
if replacement_path:
|
||
os.replace(replacement_path, output_path)
|
||
downloaded_path = output_path
|
||
finally:
|
||
if replacement_path and os.path.exists(replacement_path):
|
||
try:
|
||
os.remove(replacement_path)
|
||
except OSError:
|
||
pass
|
||
await state_store.update(
|
||
state_key,
|
||
status="completed",
|
||
provider_task_id=provider_task_id,
|
||
result_url=result_url,
|
||
output_path=downloaded_path,
|
||
file_size_bytes=os.path.getsize(downloaded_path),
|
||
completed_at=datetime.now(timezone.utc).isoformat(),
|
||
error=None,
|
||
)
|
||
_safe_log_event(
|
||
event_type="item_completed",
|
||
event_status="success",
|
||
item=item,
|
||
task_id=provider_task_id,
|
||
message="火山超分结果下载完成",
|
||
detail={
|
||
"provider_task_id": provider_task_id,
|
||
"downloaded_path": downloaded_path,
|
||
"file_size_bytes": os.path.getsize(downloaded_path),
|
||
"target_resolution": target_resolution,
|
||
"target_aspect_ratio": args.target_aspect_ratio,
|
||
"target_width": target_width,
|
||
"target_height": target_height,
|
||
"note": "CLI 未进行任何 FFmpeg 裁剪、缩放或转码",
|
||
},
|
||
)
|
||
print(f"{prefix} 完成: {downloaded_path}")
|
||
return ItemResult(
|
||
item.owner_type,
|
||
item.owner_id,
|
||
"completed",
|
||
output_path=downloaded_path,
|
||
provider_task_id=provider_task_id,
|
||
)
|
||
except Exception as exc:
|
||
error_detail: dict[str, Any] = {
|
||
"processor_key": args.processor_key,
|
||
"target_resolution": target_resolution,
|
||
"target_aspect_ratio": args.target_aspect_ratio,
|
||
"target_width": target_width,
|
||
"target_height": target_height,
|
||
"provider_task_id": provider_task_id,
|
||
}
|
||
if isinstance(exc, VolcMediaKitError):
|
||
error_detail.update(exc.log_detail())
|
||
error_detail["provider_response"] = exc.response_payload
|
||
state_status = "failed"
|
||
if isinstance(exc, TimeoutError):
|
||
state_status = "interrupted"
|
||
elif isinstance(exc, VolcMediaKitError) and exc.retryable:
|
||
state_status = "interrupted"
|
||
elif result_url and not isinstance(exc, RemoteTaskFailedError):
|
||
state_status = "interrupted"
|
||
await state_store.update(
|
||
state_key,
|
||
status=state_status,
|
||
provider_task_id=provider_task_id,
|
||
result_url=result_url,
|
||
output_path=output_path,
|
||
error=str(exc),
|
||
failed_at=datetime.now(timezone.utc).isoformat(),
|
||
)
|
||
_safe_log_error(
|
||
event_type="item_failed",
|
||
exc=exc,
|
||
item=item,
|
||
task_id=provider_task_id,
|
||
message="视频超分 CLI 处理失败",
|
||
detail=error_detail,
|
||
)
|
||
print(f"{prefix} 失败: {exc}", file=sys.stderr)
|
||
return ItemResult(
|
||
item.owner_type,
|
||
item.owner_id,
|
||
"failed",
|
||
output_path=output_path,
|
||
provider_task_id=provider_task_id,
|
||
message=str(exc),
|
||
)
|
||
|
||
|
||
async def _run(args: argparse.Namespace) -> dict[str, Any]:
|
||
_validate_args(args)
|
||
refs = _collect_owner_refs(args)
|
||
target_resolution = _normalize_target_resolution(args.target_resolution)
|
||
target_width, target_height = calculate_target_dimensions(
|
||
aspect_ratio=args.target_aspect_ratio,
|
||
target_resolution=target_resolution,
|
||
)
|
||
|
||
output_dir = Path(args.output_dir).expanduser().resolve()
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
state_path = (
|
||
Path(args.state_file).expanduser().resolve()
|
||
if str(args.state_file or "").strip()
|
||
else output_dir
|
||
/ f"video_upscale_batch_{args.output_version}_{args.processor_key}_{target_resolution}.state.json"
|
||
)
|
||
state_store = StateStore(state_path)
|
||
|
||
batch_detail = {
|
||
"item_count": len(refs),
|
||
"processor_key": args.processor_key,
|
||
"output_version": args.output_version,
|
||
"target_aspect_ratio": args.target_aspect_ratio,
|
||
"target_resolution": target_resolution,
|
||
"target_width": target_width,
|
||
"target_height": target_height,
|
||
"output_dir": str(output_dir),
|
||
"state_file": str(state_path),
|
||
"max_concurrency": int(args.max_concurrency),
|
||
"dry_run": bool(args.dry_run),
|
||
"force_resubmit": bool(args.force_resubmit),
|
||
"note": "目标比例用于计算请求目标宽高；当前服务不进行 FFmpeg 裁剪、缩放或转码",
|
||
}
|
||
log_operation_event(
|
||
domain=LOG_DOMAIN,
|
||
module=LOG_MODULE,
|
||
source=LOG_SOURCE,
|
||
event_type="batch_started",
|
||
event_status="processing",
|
||
message="视频超分 CLI 批次开始",
|
||
detail=batch_detail,
|
||
)
|
||
|
||
async with async_session() as db:
|
||
items, precheck_errors = await _resolve_items(
|
||
db,
|
||
refs=refs,
|
||
output_dir=output_dir,
|
||
output_version=args.output_version,
|
||
target_resolution=target_resolution,
|
||
)
|
||
|
||
results: list[ItemResult] = list(precheck_errors)
|
||
for error in precheck_errors:
|
||
print(f"{error.owner_type}:{error.owner_id} 预检查失败: {error.message}", file=sys.stderr)
|
||
log_operation_event(
|
||
domain=LOG_DOMAIN,
|
||
module=LOG_MODULE,
|
||
source=LOG_SOURCE,
|
||
event_type="item_precheck_failed",
|
||
event_status="failed",
|
||
task_id=error.owner_id,
|
||
message=error.message,
|
||
detail={
|
||
"owner_type": error.owner_type,
|
||
"owner_id": error.owner_id,
|
||
"output_path": error.output_path,
|
||
"status": error.status,
|
||
},
|
||
error=error.message,
|
||
)
|
||
|
||
if args.dry_run:
|
||
for index, item in enumerate(items, start=1):
|
||
source_type = "local" if item.source_local_path else "remote"
|
||
print(
|
||
json.dumps(
|
||
{
|
||
"index": index,
|
||
"owner_type": item.owner_type,
|
||
"owner_id": item.owner_id,
|
||
"source_type": source_type,
|
||
"source_local_path": item.source_local_path,
|
||
"source_remote_url": item.source_remote_url,
|
||
"source_resolution": item.source_resolution,
|
||
"source_aspect_ratio": item.source_aspect_ratio,
|
||
"processor_key": args.processor_key,
|
||
"target_resolution": target_resolution,
|
||
"target_aspect_ratio": args.target_aspect_ratio,
|
||
"target_width": target_width,
|
||
"target_height": target_height,
|
||
"output_path": item.output_path,
|
||
},
|
||
ensure_ascii=False,
|
||
)
|
||
)
|
||
results.append(
|
||
ItemResult(
|
||
item.owner_type,
|
||
item.owner_id,
|
||
"dry_run",
|
||
output_path=item.output_path,
|
||
)
|
||
)
|
||
else:
|
||
semaphore = asyncio.Semaphore(int(args.max_concurrency))
|
||
tasks = [
|
||
asyncio.create_task(
|
||
_process_item(
|
||
item=item,
|
||
args=args,
|
||
target_resolution=target_resolution,
|
||
target_width=target_width,
|
||
target_height=target_height,
|
||
state_store=state_store,
|
||
semaphore=semaphore,
|
||
index=index,
|
||
total=len(items),
|
||
)
|
||
)
|
||
for index, item in enumerate(items, start=1)
|
||
]
|
||
if tasks:
|
||
results.extend(await asyncio.gather(*tasks))
|
||
|
||
summary = {
|
||
"total": len(refs),
|
||
"resolved": len(items),
|
||
"completed": sum(result.status == "completed" for result in results),
|
||
"skipped": sum(result.status.startswith("skipped") for result in results),
|
||
"dry_run": sum(result.status == "dry_run" for result in results),
|
||
"failed": sum(
|
||
result.status
|
||
not in {"completed", "skipped_existing", "skipped_failed_state", "dry_run"}
|
||
for result in results
|
||
),
|
||
"state_file": str(state_path),
|
||
"items": [result.as_dict() for result in results],
|
||
}
|
||
event_status = "success" if summary["failed"] == 0 else "partial_failed"
|
||
log_operation_event(
|
||
domain=LOG_DOMAIN,
|
||
module=LOG_MODULE,
|
||
source=LOG_SOURCE,
|
||
event_type="batch_completed",
|
||
event_status=event_status,
|
||
message="视频超分 CLI 批次结束",
|
||
detail={**batch_detail, **summary},
|
||
)
|
||
return summary
|
||
|
||
|
||
def main() -> None:
|
||
parser = _parser()
|
||
args = parser.parse_args()
|
||
try:
|
||
summary = asyncio.run(_run(args))
|
||
print(json.dumps(summary, ensure_ascii=False, indent=2, default=str))
|
||
if int(summary.get("failed") or 0) > 0:
|
||
raise SystemExit(1)
|
||
except KeyboardInterrupt:
|
||
print("执行已被用户中断；已提交任务可通过同一命令和状态文件继续轮询。", file=sys.stderr)
|
||
raise SystemExit(130)
|
||
except SystemExit:
|
||
raise
|
||
except Exception as exc:
|
||
log_operation_error(
|
||
domain=LOG_DOMAIN,
|
||
module=LOG_MODULE,
|
||
source=LOG_SOURCE,
|
||
event_type="batch_crashed",
|
||
message="视频超分 CLI 批次异常退出",
|
||
exc=exc,
|
||
)
|
||
print(f"执行失败: {exc}", file=sys.stderr)
|
||
raise SystemExit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|