feat: 调度器+并发管线+GPU优化+日志分级+前端修复
- scheduler: ffmpeg 异步线程 + GPU 串行调度 + 模型复用（2N→2 次加载） - pipeline: 阶段拆分（extract/asr/translate），中间数据存 Task 字段 - translate_service: 长度排序批处理，padding 浪费减少 91% - model_manager: ASR/翻译不共驻，BatchedInferencePipeline 批量解码 - 日志分级: INFO=任务流转里程碑，DEBUG=进度详情；默认 INFO - 前端: 日志最新在上+滚动感知+退避轮询；24h 时间；上传中状态显示 - /health: 返回完整 Whisper/NLLB 配置 - upload_service: 单事务 complete + 扩展名白名单 - task_router: 合并 UploadSession 虚拟任务到列表 - Dockerfile: CPU/GPU 独立构建链，deps 缓存稳定 - prefetch_models: 安装时预下载模型权重
This commit is contained in:
audio2text dev
137
scripts/prefetch_models.py
Normal file
137
scripts/prefetch_models.py
Normal file
@@ -0,0 +1,137 @@
|
||||
"""预拉模型权重到本地缓存，避免容器启动时才下载（首次启动慢）。
|
||||
|
||||
在容器内执行（docker run + 挂载 ./models volume），复用镜像里的 huggingface_hub。
|
||||
读 config.yaml 拿 asr.model / translation.model，下载到 HF_HOME（= /models/huggingface）。
|
||||
|
||||
幂等：已下过的模型跳过（HF cache 命中检测）。
|
||||
|
||||
用法（经 scripts/prefetch_models.sh 包装）：
|
||||
./scripts/prefetch_models.sh # 读 config.yaml
|
||||
./scripts/prefetch_models.sh config.gpu.yaml # 读指定配置
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def log(msg: str) -> None:
|
||||
"""带时间戳的日志（容器内无项目 logger，直接 print）。"""
|
||||
ts = time.strftime("%H:%M:%S")
|
||||
print(f"[{ts}] {msg}", flush=True)
|
||||
|
||||
|
||||
def load_config(config_path: str) -> dict:
|
||||
"""读 YAML 配置，返回 asr.model / translation.model。"""
|
||||
import yaml
|
||||
with open(config_path, encoding="utf-8") as f:
|
||||
return yaml.safe_load(f)
|
||||
|
||||
|
||||
def resolve_asr_repo(model_name: str) -> str:
|
||||
"""faster-whisper 模型名 -> HF repo。
|
||||
|
||||
预定义名（tiny.en 等）走 _MODELS 映射；已是 repo 路径（org/name）直接用。
|
||||
"""
|
||||
from faster_whisper.utils import _MODELS
|
||||
if "/" in model_name:
|
||||
return model_name # 已是完整 repo 路径
|
||||
repo = _MODELS.get(model_name)
|
||||
if repo is None:
|
||||
raise ValueError(f"未知 faster-whisper 模型名：{model_name}（不在 _MODELS 里）")
|
||||
return repo
|
||||
|
||||
|
||||
def is_cached(repo_id: str, hf_home: str) -> bool:
|
||||
"""检测 HF cache 是否已有该模型（snapshot 目录存在且非空）。"""
|
||||
# HF 缓存布局：<hf_home>/hub/models--<org>--<name>/snapshots/<hash>/
|
||||
cache_dir = Path(hf_home) / "hub" / f"models--{repo_id.replace('/', '--')}"
|
||||
snapshots = cache_dir / "snapshots"
|
||||
if not snapshots.is_dir():
|
||||
return False
|
||||
return any(snap.is_dir() and any(snap.iterdir()) for snap in snapshots.iterdir())
|
||||
|
||||
|
||||
def download_model(repo_id: str, hf_home: str) -> None:
|
||||
"""用 huggingface_hub 下载模型所有文件到 HF cache。"""
|
||||
from huggingface_hub import snapshot_download
|
||||
log(f" 下载 {repo_id} ...")
|
||||
snapshot_download(
|
||||
repo_id=repo_id,
|
||||
local_dir=None, # 走标准 cache
|
||||
cache_dir=Path(hf_home) / "hub",
|
||||
)
|
||||
|
||||
|
||||
def prefetch_asr(model_name: str, hf_home: str) -> None:
|
||||
"""预拉 faster-whisper ASR 模型。"""
|
||||
repo = resolve_asr_repo(model_name)
|
||||
log(f"[ASR] model={model_name} repo={repo}")
|
||||
if is_cached(repo, hf_home):
|
||||
log(f" ✓ 已缓存，跳过")
|
||||
return
|
||||
download_model(repo, hf_home)
|
||||
log(f" ✓ 完成")
|
||||
|
||||
|
||||
def prefetch_translation(model_name: str, hf_home: str) -> None:
|
||||
"""预拉翻译模型（transformers pipeline 用的 HF repo）。"""
|
||||
log(f"[翻译] model={model_name}")
|
||||
if is_cached(model_name, hf_home):
|
||||
log(f" ✓ 已缓存，跳过")
|
||||
return
|
||||
download_model(model_name, hf_home)
|
||||
log(f" ✓ 完成")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
config_path = os.environ.get("CONFIG_PATH", "/app/config.yaml")
|
||||
if len(sys.argv) > 1:
|
||||
config_path = sys.argv[1]
|
||||
hf_home = os.environ.get("HF_HOME", "/models/huggingface")
|
||||
|
||||
log(f"配置文件: {config_path}")
|
||||
log(f"HF 缓存目录: {hf_home}")
|
||||
|
||||
if not Path(config_path).is_file():
|
||||
log(f"✗ 配置文件不存在: {config_path}")
|
||||
return 1
|
||||
|
||||
cfg = load_config(config_path)
|
||||
asr_model = cfg.get("asr", {}).get("model", "tiny.en")
|
||||
tr_model = cfg.get("translation", {}).get("model", "Helsinki-NLP/opus-mt-en-zh")
|
||||
|
||||
log(f"ASR 模型: {asr_model}")
|
||||
log(f"翻译模型: {tr_model}")
|
||||
log("-" * 50)
|
||||
|
||||
# 预拉 ASR
|
||||
try:
|
||||
prefetch_asr(asr_model, hf_home)
|
||||
except Exception as e:
|
||||
log(f"✗ ASR 模型预拉失败: {e}")
|
||||
return 2
|
||||
|
||||
# 预拉翻译
|
||||
try:
|
||||
prefetch_translation(tr_model, hf_home)
|
||||
except Exception as e:
|
||||
log(f"✗ 翻译模型预拉失败: {e}")
|
||||
return 3
|
||||
|
||||
log("-" * 50)
|
||||
log(f"全部完成。缓存位于: {hf_home}")
|
||||
# 打印缓存大小
|
||||
try:
|
||||
total = sum(f.stat().st_size for f in Path(hf_home).rglob("*") if f.is_file())
|
||||
log(f"缓存总大小: {total / 1024 / 1024 / 1024:.2f} GB")
|
||||
except Exception:
|
||||
pass
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
58
scripts/prefetch_models.sh
Normal file
58
scripts/prefetch_models.sh
Normal file
@@ -0,0 +1,58 @@
|
||||
#!/usr/bin/env bash
|
||||
# 预拉模型权重到 ./models volume，避免容器启动时才下载。
|
||||
#
|
||||
# 用法:
|
||||
# ./scripts/prefetch_models.sh # 读 config.yaml（当前激活配置）
|
||||
# ./scripts/prefetch_models.sh config.gpu.yaml # 读指定配置文件
|
||||
#
|
||||
# 原理：用已构建的 audio2text 镜像跑一次性容器，挂载 ./models volume，
|
||||
# 执行 scripts/prefetch_models.py 把模型下到 HF cache。
|
||||
# 容器运行时 HF_HOME=/models/huggingface 命中缓存，秒级加载。
|
||||
#
|
||||
# 幂等：已下过的模型跳过。换 config 的 model 后重跑即可补下新模型。
|
||||
set -euo pipefail
|
||||
|
||||
cd "$(dirname "$0")/.."
|
||||
ROOT="$(pwd)"
|
||||
|
||||
CONFIG="${1:-config.yaml}"
|
||||
VARIANT="${AUDIO2TEXT_VARIANT:-cpu}"
|
||||
|
||||
# 选镜像：优先用 gpu 镜像（GPU 模型大，可能要 cuda 才能 correctly 下载部分文件），
|
||||
# 否则 cpu。两者都能下 HF 模型（下载是纯网络 IO，不依赖 GPU）。
|
||||
IMAGE="audio2text:cpu"
|
||||
if docker image inspect audio2text:gpu >/dev/null 2>&1 && [ "$VARIANT" = "gpu" ]; then
|
||||
IMAGE="audio2text:gpu"
|
||||
fi
|
||||
|
||||
if ! docker image inspect "$IMAGE" >/dev/null 2>&1; then
|
||||
echo "✗ 镜像 $IMAGE 不存在，请先 ./setup.sh 构建。" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [ ! -f "$ROOT/$CONFIG" ]; then
|
||||
echo "✗ 配置文件不存在: $ROOT/$CONFIG" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
mkdir -p "$ROOT/models"
|
||||
|
||||
echo "==> 预拉模型"
|
||||
echo " 镜像: $IMAGE"
|
||||
echo " 配置: $CONFIG"
|
||||
echo " 缓存: $ROOT/models (volume)"
|
||||
echo
|
||||
|
||||
# 注意：CONFIG_PATH 指向容器内路径，挂载配置文件为只读
|
||||
MSYS_NO_PATHCONV=1 docker run --rm \
|
||||
-v "$ROOT/models:/models" \
|
||||
-v "$ROOT/$CONFIG:/app/config.yaml:ro" \
|
||||
-e HF_HOME=/models/huggingface \
|
||||
-e CT2_CACHE=/models/ctranslate2 \
|
||||
-e CONFIG_PATH=/app/config.yaml \
|
||||
"$IMAGE" \
|
||||
python /app/scripts/prefetch_models.py /app/config.yaml
|
||||
|
||||
echo
|
||||
echo "==> 完成。模型已缓存到 $ROOT/models/huggingface"
|
||||
echo " 容器启动时会命中缓存，无需联网下载。"
|
||||
Reference in New Issue
Block a user