#!/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 [--out-dir store/videos] [--srt-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//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))