feat(tools): batch tree build orchestration with shared API semaphore
This commit is contained in:
+191
-2
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user