first commit
This commit is contained in:
@@ -0,0 +1,199 @@
|
||||
"""启动前检查并下载 rerank 模型。
|
||||
|
||||
设计目标:
|
||||
* 幂等 —— 模型已存在且校验通过就立刻返回,不产生任何网络请求;
|
||||
* 可续传 —— 弱网下载中断后重启容器能从断点继续;
|
||||
* 可校验 —— 每个文件都带 sha256,下载后校验,不匹配则丢弃重下。
|
||||
|
||||
可配置项(均可用环境变量覆盖):
|
||||
MODEL_DIR 模型存放目录,默认 /models
|
||||
MODEL_REPO 仓库名,默认 Xenova/bge-reranker-base
|
||||
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-reranker-base")
|
||||
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, 字节数)
|
||||
MANIFEST = [
|
||||
(
|
||||
"model_int8.onnx",
|
||||
"onnx/model_int8.onnx",
|
||||
"2059d8ef0b6e935b4845e11b38c9af9e9e2e7b91f69fc99efe03254e0a7da8d3",
|
||||
278825308,
|
||||
),
|
||||
(
|
||||
"tokenizer.json",
|
||||
"tokenizer.json",
|
||||
"48564c5c7d3fa64d85d95e65414a542385f88b0f128fd8d4163fd7a57f2be05c",
|
||||
17098079,
|
||||
),
|
||||
(
|
||||
"sentencepiece.bpe.model",
|
||||
"sentencepiece.bpe.model",
|
||||
"cfc8146abe2a0488e9e2a0c56de7952f7c11ab059eca145a0a727afce0db2865",
|
||||
5069051,
|
||||
),
|
||||
(
|
||||
"special_tokens_map.json",
|
||||
"special_tokens_map.json",
|
||||
"d5469a60db23249c7f8945013d78df30b44b6bf686c6bb4740f4223f77b1b535",
|
||||
279,
|
||||
),
|
||||
(
|
||||
"tokenizer_config.json",
|
||||
"tokenizer_config.json",
|
||||
"a1d6bc8734a6f635dc158508bef000f8e2e5a759c7d92f984b2c86e5ff53425b",
|
||||
443,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
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-remarker/1.0")
|
||||
if have:
|
||||
request.add_header("Range", f"bytes={have}-")
|
||||
|
||||
with urllib.request.urlopen(request, timeout=TIMEOUT) as resp:
|
||||
# 服务端不支持 Range 时会返回 200 全量,此时从零重写
|
||||
resuming = resp.status == 206
|
||||
mode = "ab" if (have and resuming) 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-remarker] 模型目录: {MODEL_DIR}")
|
||||
print(f"[local-remarker] 下载源 : {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-remarker] 模型已就绪,跳过下载")
|
||||
return 0
|
||||
|
||||
total_bytes = sum(item[3] for item in pending)
|
||||
print(f"[local-remarker] 需要下载 {len(pending)} 个文件,共 {human(total_bytes)}")
|
||||
|
||||
# 容器内磁盘常较小,给个明确提示而不是让 curl 报一堆 no space
|
||||
free = shutil.disk_usage(MODEL_DIR).free
|
||||
if free < total_bytes * 1.1:
|
||||
raise SystemExit(
|
||||
f"磁盘空间不足:可用 {human(free)},需要约 {human(total_bytes)}。"
|
||||
"请挂载更大的卷或改用 fp16/其它量化版本。"
|
||||
)
|
||||
|
||||
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-remarker] 模型准备完成")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user