From 56fe138a465c261ea6d27e55231a12053c975f5d Mon Sep 17 00:00:00 2001 From: iomgaa Date: Sat, 11 Jul 2026 11:56:36 -0400 Subject: [PATCH] feat(tools): batch tree build orchestration with shared API semaphore --- tests/unit/test_build_trees.py | 125 +++++++++++++++++++++ tools/build_trees.py | 193 ++++++++++++++++++++++++++++++++- 2 files changed, 316 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_build_trees.py b/tests/unit/test_build_trees.py index 7a03fda..662cc6a 100644 --- a/tests/unit/test_build_trees.py +++ b/tests/unit/test_build_trees.py @@ -5,9 +5,14 @@ from __future__ import annotations +import argparse +import asyncio +import json import sys from pathlib import Path +import pytest + PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) @@ -17,6 +22,8 @@ from tools.build_trees import ( # noqa: E402 _discover_pending, _find_srt_entries, _tree_is_complete, + load_progress, + main_async, ) @@ -106,3 +113,121 @@ class TestFindSrtEntries: def test_missing_returns_none(self, tmp_path: Path) -> None: """无同名 .srt 返回 None。""" assert _find_srt_entries(tmp_path / "vid.mp4", tmp_path) is None + + +# ── 编排集成（桩 builder，无真实视频/LLM）────────────────────── + + +class _StubBuilder: + """记录并发与注入信号量的桩 builder。""" + + instances: list[_StubBuilder] = [] + inflight = 0 + max_inflight = 0 + + def __init__(self, vlm, llm, config, *, api_semaphore=None) -> None: + """记录注入的 api_semaphore 并登记实例。""" + self.api_semaphore = api_semaphore + _StubBuilder.instances.append(self) + + async def build_async(self, video_path: str, srt_entries=None) -> TreeIndex: + """模拟建树：短暂 sleep 并统计并发峰值，返回最小合法树。""" + _StubBuilder.inflight += 1 + _StubBuilder.max_inflight = max(_StubBuilder.max_inflight, _StubBuilder.inflight) + await asyncio.sleep(0.02) + _StubBuilder.inflight -= 1 + + l1 = L1Node( + id="x_L1_000", + card=L1Card("s", "室内", [], [], [], [], "线性"), + time_range=(0.0, 1.0), + children=[], + ) + return TreeIndex(metadata=IndexMeta(video_path, "video"), roots=[l1]) + + +@pytest.fixture() +def batch_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> dict: + """5 个假视频 + 桩 builder + 隔离的 progress 路径。""" + import tools.build_trees as bt + + _StubBuilder.instances = [] + _StubBuilder.inflight = 0 + _StubBuilder.max_inflight = 0 + monkeypatch.setattr(bt, "VideoTreeBuilder", _StubBuilder) + monkeypatch.setattr(bt, "_build_clients", lambda api_concurrency: (None, None)) + + videos = tmp_path / "videos" + videos.mkdir() + for i in range(5): + (videos / f"v{i}.mp4").write_bytes(b"") + + return { + "videos": videos, + "out": tmp_path / "out", + "progress": tmp_path / "build_progress.json", + } + + +def _make_args(env: dict, video_concurrency: int = 2, limit: int = 0) -> argparse.Namespace: + """构造 main_async 所需的 CLI 参数命名空间。""" + return argparse.Namespace( + videos_dir=str(env["videos"]), + out_dir=str(env["out"]), + srt_dir=str(env["videos"]), + video_concurrency=video_concurrency, + limit=limit, + progress_path=str(env["progress"]), + ) + + +class TestOrchestration: + """main_async 编排行为（桩 builder）。""" + + @pytest.mark.asyncio + async def test_video_concurrency_capped(self, batch_env: dict) -> None: + """同时在建视频数不得超过 video_concurrency。""" + await main_async(_make_args(batch_env, video_concurrency=2)) + assert _StubBuilder.max_inflight <= 2 + assert len(_StubBuilder.instances) == 5 + + @pytest.mark.asyncio + async def test_shared_api_semaphore(self, batch_env: dict) -> None: + """全部 builder 实例共享同一个 API Semaphore 对象。""" + await main_async(_make_args(batch_env)) + sems = {id(b.api_semaphore) for b in _StubBuilder.instances} + assert len(sems) == 1 + assert _StubBuilder.instances[0].api_semaphore is not None + + @pytest.mark.asyncio + async def test_trees_saved_and_progress_recorded(self, batch_env: dict) -> None: + """每个视频产出 tree.json 且 progress 记录全部完成。""" + await main_async(_make_args(batch_env)) + for i in range(5): + assert (batch_env["out"] / f"v{i}" / "tree.json").exists() + + assert load_progress(batch_env["progress"]) == {f"v{i}" for i in range(5)} + + @pytest.mark.asyncio + async def test_resume_skips_finished(self, batch_env: dict) -> None: + """第二次运行跳过全部已完成视频。""" + await main_async(_make_args(batch_env)) + n_first = len(_StubBuilder.instances) + await main_async(_make_args(batch_env)) + assert len(_StubBuilder.instances) == n_first # 无新建 + + @pytest.mark.asyncio + async def test_partial_completion_resume(self, batch_env: dict) -> None: + """部分完成后重跑只建剩余视频（模拟中断后恢复）。""" + batch_env["progress"].write_text( + json.dumps({"finished_video_ids": ["v0", "v1"]}), + encoding="utf-8", + ) + await main_async(_make_args(batch_env)) + assert len(_StubBuilder.instances) == 3 # 仅 v2/v3/v4 + + @pytest.mark.asyncio + async def test_limit(self, batch_env: dict) -> None: + """--limit 2 只建前两个（烟测入口）。""" + await main_async(_make_args(batch_env, limit=2)) + assert len(_StubBuilder.instances) == 2 diff --git a/tools/build_trees.py b/tools/build_trees.py index 23fbe76..5eeeb56 100644 --- a/tools/build_trees.py +++ b/tools/build_trees.py @@ -17,9 +17,12 @@ app/core/adapters 不 import 此脚本。 from __future__ import annotations +import argparse +import asyncio import json import os import sys +import time from pathlib import Path from typing import TYPE_CHECKING @@ -27,17 +30,20 @@ from typing import TYPE_CHECKING PROJECT_ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(PROJECT_ROOT)) +import yaml # noqa: E402 from dotenv import load_dotenv # noqa: E402 from loguru import logger # noqa: E402 load_dotenv(PROJECT_ROOT / ".env") +from app.tree.config import TreeConfig # noqa: E402 from app.tree.index import TreeIndex # noqa: E402 from app.tree.subtitle import parse_srt # noqa: E402 +from app.tree.video_builder import VideoTreeBuilder # noqa: E402 if TYPE_CHECKING: - import asyncio - + from adapters.llm import GovernedLLMClient + from adapters.vlm import GovernedVLMClient from app.tree.subtitle import SRTEntry # --------------------------------------------------------------------------- @@ -173,3 +179,186 @@ def _find_srt_entries(video_path: Path, srt_dir: Path) -> list[SRTEntry] | None: if not srt_path.exists(): return None return parse_srt(str(srt_path)) + + +# --------------------------------------------------------------------------- +# LLM/VLM 客户端构建（复刻 repair_trees.py 惯例） +# --------------------------------------------------------------------------- + + +def _build_clients(api_concurrency: int) -> tuple[GovernedLLMClient, GovernedVLMClient]: + """构建 GovernedLLMClient（LLM + VLM），熔断阈值随 API 并发缩放。 + + 参数: + api_concurrency: 全局 API 并发上限（熔断阈值取 max(cfg, api_concurrency*2)）。 + + 返回: + (llm_client, vlm_client) 元组。 + """ + from adapters.breaker import CircuitBreaker + from adapters.llm import GovernedLLMClient + from adapters.telemetry import SQLiteTelemetryRecorder + from adapters.vlm import GovernedVLMClient + + (PROJECT_ROOT / "logs").mkdir(exist_ok=True) + telemetry = SQLiteTelemetryRecorder(str(PROJECT_ROOT / "logs" / "build_trees_telemetry.db")) + + breaker_threshold = int(os.getenv("LLM_CIRCUIT_BREAKER_THRESHOLD", "5")) + breaker_threshold = max(breaker_threshold, api_concurrency * 2) + breaker_cooldown = int(os.getenv("LLM_CIRCUIT_BREAKER_COOLDOWN", "60")) + timeout_s = float(os.getenv("LLM_TIMEOUT", "120")) + max_retries = int(os.getenv("LLM_MAX_RETRIES", "3")) + base_delay = float(os.getenv("LLM_RETRY_BASE_DELAY", "2.0")) + max_delay = float(os.getenv("LLM_RETRY_MAX_DELAY", "30.0")) + ttft = float(os.getenv("LLM_TTFT_TIMEOUT", "30")) + inter_token = float(os.getenv("LLM_INTER_TOKEN_TIMEOUT", "15")) + + llm = GovernedLLMClient( + model=os.environ["SEARCH_LLM_MODEL"], + base_url=os.environ["SEARCH_LLM_BASE_URL"], + api_key=os.environ["SEARCH_LLM_API_KEY"], + provider="deepseek", + thinking=False, + breaker=CircuitBreaker(fail_threshold=breaker_threshold, cooldown_s=breaker_cooldown), + cache=None, + telemetry=telemetry, + timeout_s=timeout_s, + ttft_timeout_s=ttft, + inter_token_timeout_s=inter_token, + max_retries=max_retries, + retry_base_delay_s=base_delay, + retry_max_delay_s=max_delay, + ) + vlm_base = GovernedLLMClient( + model=os.environ["VL_LLM_MODEL"], + base_url=os.environ["VL_LLM_BASE_URL"], + api_key=os.environ["VL_LLM_API_KEY"], + provider="qwen", + thinking=False, + breaker=CircuitBreaker(fail_threshold=breaker_threshold, cooldown_s=breaker_cooldown), + cache=None, + telemetry=telemetry, + timeout_s=timeout_s, + ttft_timeout_s=ttft, + inter_token_timeout_s=inter_token, + max_retries=max_retries, + retry_base_delay_s=base_delay, + retry_max_delay_s=max_delay, + ) + return llm, GovernedVLMClient(vlm_base) + + +# --------------------------------------------------------------------------- +# 主编排 +# --------------------------------------------------------------------------- + + +async def main_async(args: argparse.Namespace) -> None: + """异步主流程：视频级并发建树 + 全局共享 API 信号量。 + + 参数: + args: CLI 参数（videos_dir/out_dir/srt_dir/video_concurrency/limit/progress_path）。 + """ + videos_dir = Path(args.videos_dir) + out_dir = Path(args.out_dir) + srt_dir = Path(args.srt_dir) + assert videos_dir.is_dir(), f"视频目录不存在: {videos_dir}" + + api_concurrency = int(os.getenv("TREE_BUILD_API_CONCURRENCY", "16")) + + # Phase 1: 待建发现（progress + 完整性双重跳过） + progress_path = Path(args.progress_path) + finished = load_progress(progress_path) + if finished: + logger.info("progress 已记录 {} 个完成视频", len(finished)) + pending = _discover_pending(videos_dir, out_dir, finished) + if args.limit > 0: + pending = pending[: args.limit] + logger.info( + "待建 {} 个视频, video_concurrency={}, api_concurrency={}", + len(pending), + args.video_concurrency, + api_concurrency, + ) + if not pending: + return + + # Phase 2: 客户端与共享信号量 + llm, vlm = _build_clients(api_concurrency) + with open(PROJECT_ROOT / "config" / "default.yaml", encoding="utf-8") as f: + tree_cfg = TreeConfig.from_dict(yaml.safe_load(f)["tree"]) + api_sem = asyncio.Semaphore(api_concurrency) + video_sem = asyncio.Semaphore(args.video_concurrency) + progress_lock = asyncio.Lock() + start_time = time.time() + completed = 0 + failed: list[str] = [] + + # Phase 3: 视频级并发编排（复刻 repair_trees 模式） + async def _build_one(video_path: Path) -> None: + nonlocal completed + async with video_sem: + vid = video_path.stem + logger.info("开始建树 {}", vid) + builder = VideoTreeBuilder(vlm=vlm, llm=llm, config=tree_cfg, api_semaphore=api_sem) + srt_entries = _find_srt_entries(video_path, srt_dir) + try: + index = await builder.build_async(str(video_path), srt_entries) + except Exception as exc: + # WHY: 单视频错误隔离——一个视频失败不拖垮整批； + # 失败清单汇总上报，vid 不进 progress，下次重跑自动重建。 + logger.error("建树失败 {} ({}): {}", vid, type(exc).__name__, exc) + failed.append(vid) + return + tree_path = out_dir / vid / "tree.json" + tree_path.parent.mkdir(parents=True, exist_ok=True) + index.save_json(str(tree_path)) + await save_progress(progress_path, progress_lock, vid) + completed += 1 + if completed % 5 == 0: + elapsed = time.time() - start_time + rate = completed / elapsed * 60 if elapsed > 0 else 0 + logger.info( + "进度: {}/{}, 已用 {:.0f}s, 速率 {:.2f} 视频/分钟", + completed, + len(pending), + elapsed, + rate, + ) + + await asyncio.gather(*[asyncio.create_task(_build_one(p)) for p in pending]) + + # Phase 4: 汇总 + elapsed = time.time() - start_time + logger.info( + "批量建树完成: 成功 {}, 失败 {}, 总耗时 {:.0f}s{}", + completed, + len(failed), + elapsed, + f", 失败清单: {failed}" if failed else "", + ) + + +def main() -> None: + """同步入口。""" + parser = argparse.ArgumentParser(description="批量并行建树") + parser.add_argument("--videos-dir", type=str, required=True, help="视频文件目录") + parser.add_argument("--out-dir", type=str, default="store/videos", help="树输出根目录") + parser.add_argument("--srt-dir", type=str, default="", help="SRT 目录（默认同 videos-dir）") + parser.add_argument("--video-concurrency", type=int, default=16, help="同时在建视频数") + parser.add_argument("--limit", type=int, default=0, help="只建前 N 个（0=全部，烟测用）") + parser.add_argument( + "--progress-path", + type=str, + default=str(PROJECT_ROOT / "logs" / PROGRESS_FILE), + dest="progress_path", + help="progress 文件路径", + ) + args = parser.parse_args() + if not args.srt_dir: + args.srt_dir = args.videos_dir + asyncio.run(main_async(args)) + + +if __name__ == "__main__": + main()