201 lines
6.7 KiB
Python
201 lines
6.7 KiB
Python
"""启动前检查并下载嵌入模型。
|
|
|
|
设计目标:
|
|
* 幂等 —— 模型已存在且 sha256 校验通过就立刻返回,不产生任何网络请求;
|
|
* 可续传 —— 弱网下载中断后重启容器能从断点继续;
|
|
* 可校验 —— 每个文件都带 sha256,下载后校验,不匹配则丢弃重下。
|
|
|
|
可配置项(均可用环境变量覆盖):
|
|
MODEL_DIR 模型存放目录,默认 /models
|
|
MODEL_REPO 仓库名,默认 Xenova/bge-m3
|
|
MODEL_BASE 下载源前缀
|
|
FORCE 设为 1 时忽略本地文件强制重下
|
|
MAX_RETRIES 单个文件最大重试次数
|
|
"""
|
|
|
|
import hashlib
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
MODEL_DIR = os.environ.get("MODEL_DIR", "/models")
|
|
MODEL_REPO = os.environ.get("MODEL_REPO", "Xenova/bge-m3")
|
|
MODEL_REV = os.environ.get("MODEL_REV", "master")
|
|
MODEL_BASE = os.environ.get(
|
|
"MODEL_BASE", f"https://modelscope.cn/models/{MODEL_REPO}/resolve/{MODEL_REV}"
|
|
)
|
|
FORCE = os.environ.get("FORCE", "0") == "1"
|
|
MAX_RETRIES = int(os.environ.get("MAX_RETRIES", "5"))
|
|
TIMEOUT = int(os.environ.get("DOWNLOAD_TIMEOUT", "60"))
|
|
|
|
# (目标文件名, 仓库内路径, sha256, 字节数)
|
|
#
|
|
# 选 model_int8.onnx:568MB 自包含(没有外置的 .onnx_data),且带 AVX512-VNNI
|
|
# 的 CPU 上 int8 推理最快。fp32 版需要额外下载 model.onnx_data 共 2.1GB。
|
|
MANIFEST = [
|
|
(
|
|
"model_int8.onnx",
|
|
"onnx/model_int8.onnx",
|
|
"a206e10e995aa2a833924bcd725ba5dd6c3425cd34bac3cf2b5677cd2a1c51d6",
|
|
568456694,
|
|
),
|
|
(
|
|
"tokenizer.json",
|
|
"tokenizer.json",
|
|
"6710678b12670bc442b99edc952c4d996ae309a7020c1fa0096dd245c2faf790",
|
|
17082821,
|
|
),
|
|
(
|
|
"sentencepiece.bpe.model",
|
|
"sentencepiece.bpe.model",
|
|
"cfc8146abe2a0488e9e2a0c56de7952f7c11ab059eca145a0a727afce0db2865",
|
|
5069051,
|
|
),
|
|
(
|
|
"special_tokens_map.json",
|
|
"special_tokens_map.json",
|
|
"8c785abebea9ae3257b61681b4e6fd8365ceafde980c21970d001e834cf10835",
|
|
964,
|
|
),
|
|
(
|
|
"tokenizer_config.json",
|
|
"tokenizer_config.json",
|
|
"7e4c1cc848840aeccdd763458c18dd525eb0f795c992e00ebe9c28554e7db2d4",
|
|
1173,
|
|
),
|
|
]
|
|
|
|
|
|
def human(size: float) -> str:
|
|
for unit in ("B", "KB", "MB", "GB"):
|
|
if size < 1024 or unit == "GB":
|
|
return f"{size:.1f}{unit}"
|
|
size /= 1024
|
|
return f"{size:.1f}GB"
|
|
|
|
|
|
def sha256_of(path: str, chunk: int = 1 << 20) -> str:
|
|
digest = hashlib.sha256()
|
|
with open(path, "rb") as fh:
|
|
while True:
|
|
block = fh.read(chunk)
|
|
if not block:
|
|
break
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def is_valid(path: str, expect_sha: str, expect_size: int) -> bool:
|
|
if not os.path.isfile(path):
|
|
return False
|
|
if os.path.getsize(path) != expect_size:
|
|
return False
|
|
return sha256_of(path) == expect_sha
|
|
|
|
|
|
def download(name: str, remote_path: str) -> None:
|
|
"""带断点续传的下载,落到 .part 临时文件,成功后原子改名。"""
|
|
target = os.path.join(MODEL_DIR, name)
|
|
part = target + ".part"
|
|
url = f"{MODEL_BASE}/{remote_path}"
|
|
|
|
last_error: Exception | None = None
|
|
for attempt in range(1, MAX_RETRIES + 1):
|
|
try:
|
|
have = os.path.getsize(part) if os.path.isfile(part) else 0
|
|
total = next((m[3] for m in MANIFEST if m[0] == name), 0)
|
|
|
|
request = urllib.request.Request(url)
|
|
request.add_header("User-Agent", "local-embedding/1.0")
|
|
if have:
|
|
request.add_header("Range", f"bytes={have}-")
|
|
|
|
with urllib.request.urlopen(request, timeout=TIMEOUT) as resp:
|
|
# 服务端不支持 Range 时会返回 200 全量,此时从零重写
|
|
mode = "ab" if (have and resp.status == 206) else "wb"
|
|
if mode == "wb":
|
|
have = 0
|
|
downloaded = have
|
|
|
|
print(f" [{name}] {human(downloaded)} / {human(total)}", flush=True)
|
|
with open(part, mode) as out:
|
|
while True:
|
|
block = resp.read(1 << 20)
|
|
if not block:
|
|
break
|
|
out.write(block)
|
|
downloaded += len(block)
|
|
pct = downloaded * 100 // total if total else 0
|
|
print(
|
|
f"\r [{name}] {human(downloaded)} / {human(total)} ({pct}%)",
|
|
end="",
|
|
flush=True,
|
|
)
|
|
print(flush=True)
|
|
|
|
os.replace(part, target)
|
|
return
|
|
|
|
except (urllib.error.URLError, OSError, TimeoutError) as exc:
|
|
last_error = exc
|
|
wait = min(2 ** attempt, 30)
|
|
print(
|
|
f"\n [{name}] 第 {attempt}/{MAX_RETRIES} 次失败: {exc},{wait}s 后重试",
|
|
flush=True,
|
|
)
|
|
if attempt < MAX_RETRIES:
|
|
time.sleep(wait)
|
|
|
|
raise SystemExit(f"下载 {name} 失败: {last_error}")
|
|
|
|
|
|
def main() -> int:
|
|
print(f"[local-embedding] 模型目录: {MODEL_DIR}")
|
|
print(f"[local-embedding] 下载源 : {MODEL_BASE}")
|
|
|
|
os.makedirs(MODEL_DIR, exist_ok=True)
|
|
|
|
pending = []
|
|
for name, remote, sha, size in MANIFEST:
|
|
target = os.path.join(MODEL_DIR, name)
|
|
if FORCE or not is_valid(target, sha, size):
|
|
if not FORCE and os.path.isfile(target):
|
|
print(f" [{name}] 校验不通过,将重新下载")
|
|
pending.append((name, remote, sha, size))
|
|
else:
|
|
print(f" [{name}] 已存在且校验通过 ({human(size)})")
|
|
|
|
if not pending:
|
|
print("[local-embedding] 模型已就绪,跳过下载")
|
|
return 0
|
|
|
|
total_bytes = sum(item[3] for item in pending)
|
|
print(f"[local-embedding] 需要下载 {len(pending)} 个文件,共 {human(total_bytes)}")
|
|
|
|
# 容器内磁盘常较小,给个明确提示而不是让下载器报一堆 no space
|
|
free = shutil.disk_usage(MODEL_DIR).free
|
|
if free < total_bytes * 1.1:
|
|
raise SystemExit(
|
|
f"磁盘空间不足:可用 {human(free)},需要约 {human(total_bytes)}。"
|
|
"请挂载更大的卷,或改用更小的嵌入模型。"
|
|
)
|
|
|
|
for name, remote, sha, size in pending:
|
|
download(name, remote)
|
|
|
|
# 最终校验,任何文件对不上都不允许带着半个模型启动
|
|
for name, _remote, sha, size in MANIFEST:
|
|
target = os.path.join(MODEL_DIR, name)
|
|
if not is_valid(target, sha, size):
|
|
raise SystemExit(f"{name} 校验失败,拒绝启动")
|
|
print(f" [{name}] 校验通过")
|
|
|
|
print("[local-embedding] 模型准备完成")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main()) |