iomgaa/Video-Tree-TRM5
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(tools): batch tree build orchestration with shared API semaphore

Browse Source
This commit is contained in:
iomgaa
2026-07-11 11:56:36 -04:00
parent 978ddef91b
commit 56fe138a46
2 changed files with 316 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit/test_build_trees.py
+125
View File
@@ -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
tools/build_trees.py
+191 -2
View File
@@ -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()