176 lines
5.4 KiB
Python
176 lines
5.4 KiB
Python
#!/usr/bin/env python3
|
||
"""批量并行建树入口：多视频并发构建三层 TreeIndex。
|
||
|
||
并发模型（Spec-2）:
|
||
视频级 Semaphore(video_concurrency) + gather —— 复刻 repair_trees.py 惯例；
|
||
全局共享一个 API Semaphore(api_concurrency) 注入所有 VideoTreeBuilder，
|
||
端点压力与单视频建树完全一致，吞吐提升来自非 API 阶段跨视频重叠。
|
||
|
||
用法:
|
||
conda activate Video-Tree-TRM
|
||
python tools/build_trees.py --videos-dir <dir> [--out-dir store/videos]
|
||
[--srt-dir <dir>] [--video-concurrency 16] [--limit 0]
|
||
|
||
api_concurrency 为工程配置，从 .env 读取 TREE_BUILD_API_CONCURRENCY（默认 16）。
|
||
app/core/adapters 不 import 此脚本。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
from typing import TYPE_CHECKING
|
||
|
||
# 确保项目根目录在 sys.path 中
|
||
PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||
sys.path.insert(0, str(PROJECT_ROOT))
|
||
|
||
from dotenv import load_dotenv # noqa: E402
|
||
from loguru import logger # noqa: E402
|
||
|
||
load_dotenv(PROJECT_ROOT / ".env")
|
||
|
||
from app.tree.index import TreeIndex # noqa: E402
|
||
from app.tree.subtitle import parse_srt # noqa: E402
|
||
|
||
if TYPE_CHECKING:
|
||
import asyncio
|
||
|
||
from app.tree.subtitle import SRTEntry
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 日志配置：不缓存，立即输出
|
||
# ---------------------------------------------------------------------------
|
||
|
||
logger.remove()
|
||
logger.add(
|
||
sys.stderr,
|
||
format="{time:HH:mm:ss} | {level:<7} | {message}",
|
||
level="DEBUG",
|
||
colorize=True,
|
||
)
|
||
logger.add(
|
||
PROJECT_ROOT / "logs" / "build_trees.log",
|
||
format="{time:YYYY-MM-DD HH:mm:ss} | {level:<7} | {message}",
|
||
level="DEBUG",
|
||
rotation="50 MB",
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 断点续跑 — progress 文件管理（复刻 repair_trees.py 惯例）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
PROGRESS_FILE = "build_progress.json"
|
||
|
||
_VIDEO_SUFFIXES = frozenset({".mp4", ".mkv", ".avi", ".webm"})
|
||
|
||
|
||
def load_progress(path: Path) -> set[str]:
|
||
"""读取 progress 文件，返回已完成视频 ID 集合。
|
||
|
||
参数:
|
||
path: progress JSON 文件路径。
|
||
|
||
返回:
|
||
已完成视频 ID 集合。文件不存在或损坏时返回空集。
|
||
"""
|
||
if not path.exists():
|
||
return set()
|
||
try:
|
||
data = json.loads(path.read_text(encoding="utf-8"))
|
||
return set(data.get("finished_video_ids", []))
|
||
except (json.JSONDecodeError, KeyError, TypeError, AttributeError):
|
||
logger.warning("progress 文件损坏，忽略: {}", path)
|
||
return set()
|
||
|
||
|
||
async def save_progress(path: Path, lock: asyncio.Lock, vid: str) -> None:
|
||
"""原子追加一个视频 ID 到 progress 文件。
|
||
|
||
参数:
|
||
path: progress JSON 文件路径。
|
||
lock: asyncio.Lock，防并发读改写丢更新。
|
||
vid: 要追加的视频 ID。
|
||
"""
|
||
async with lock:
|
||
finished = load_progress(path)
|
||
finished.add(vid)
|
||
tmp = path.with_suffix(".tmp")
|
||
tmp.write_text(
|
||
json.dumps({"finished_video_ids": sorted(finished)}, ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
os.replace(str(tmp), str(path))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 待建发现与完整性校验
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _tree_is_complete(tree_path: Path) -> bool:
|
||
"""判断 tree.json 是否存在且可加载为非空树。
|
||
|
||
参数:
|
||
tree_path: tree.json 路径。
|
||
|
||
返回:
|
||
True 表示完整（跳过重建）；文件缺失/损坏/空树返回 False。
|
||
"""
|
||
if not tree_path.exists():
|
||
return False
|
||
try:
|
||
index = TreeIndex.load_json(str(tree_path))
|
||
except (json.JSONDecodeError, KeyError, TypeError, ValueError, AssertionError) as exc:
|
||
logger.warning("tree.json 无法加载，视为不完整: {} ({})", tree_path, exc)
|
||
return False
|
||
return len(index.roots) > 0
|
||
|
||
|
||
def _discover_pending(
|
||
videos_dir: Path,
|
||
out_dir: Path,
|
||
finished: set[str],
|
||
) -> list[Path]:
|
||
"""扫描视频目录，返回待建视频文件列表（按文件名排序）。
|
||
|
||
跳过条件：video_id 在 progress 中，或 out_dir/<video_id>/tree.json 完整。
|
||
|
||
参数:
|
||
videos_dir: 视频文件目录。
|
||
out_dir: 树输出根目录。
|
||
finished: progress 中已完成的视频 ID 集合。
|
||
|
||
返回:
|
||
待建视频文件路径列表。
|
||
"""
|
||
pending: list[Path] = []
|
||
for f in sorted(videos_dir.iterdir()):
|
||
if not f.is_file() or f.suffix.lower() not in _VIDEO_SUFFIXES:
|
||
continue
|
||
vid = f.stem
|
||
if vid in finished:
|
||
continue
|
||
if _tree_is_complete(out_dir / vid / "tree.json"):
|
||
continue
|
||
pending.append(f)
|
||
return pending
|
||
|
||
|
||
def _find_srt_entries(video_path: Path, srt_dir: Path) -> list[SRTEntry] | None:
|
||
"""按视频同名规则查找并解析 SRT 字幕。
|
||
|
||
参数:
|
||
video_path: 视频文件路径。
|
||
srt_dir: SRT 目录。
|
||
|
||
返回:
|
||
SRTEntry 列表；无同名 .srt 时返回 None。
|
||
"""
|
||
srt_path = srt_dir / f"{video_path.stem}.srt"
|
||
if not srt_path.exists():
|
||
return None
|
||
return parse_srt(str(srt_path))
|