From e0248e810f2756ac819526a65de02ae071ca28db Mon Sep 17 00:00:00 2001 From: Jack Date: Fri, 23 Oct 2026 18:09:59 +0800 Subject: [PATCH] first commit --- .gitignore | 7 ++ Dockerfile | 37 +++++++++ README.md | 177 ++++++++++++++++++++++++++++++++++++++++ app.py | 181 +++++++++++++++++++++++++++++++++++++++++ docker-compose.yml | 62 ++++++++++++++ download_model.py | 199 +++++++++++++++++++++++++++++++++++++++++++++ entrypoint.sh | 17 ++++ requirements.txt | 5 ++ 8 files changed, 685 insertions(+) create mode 100644 .gitignore create mode 100644 Dockerfile create mode 100644 README.md create mode 100644 app.py create mode 100644 docker-compose.yml create mode 100644 download_model.py create mode 100755 entrypoint.sh create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c67cd5d --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +# 模型权重目录(约 290MB,不进版本库) +models/ +*.onnx +*.part + +__pycache__/ +*.pyc \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..ac6fd7c --- /dev/null +++ b/Dockerfile @@ -0,0 +1,37 @@ +# 本地重排服务镜像。 +# +# 模型不在构建期打包 —— 体积大(约 290MB)且会频繁升级,改为容器启动时按需下载, +# 已有则复用。这样换模型只需改环境变量后重启,不用重新 build。 +FROM python:3.11-slim + +# pip 源做成构建参数:国内机器可用清华源显著提速 +ARG PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple +ARG PIP_TRUSTED_HOST=pypi.tuna.tsinghua.edu.cn + +ENV PYTHONUNBUFFERED=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + PIP_NO_CACHE_DIR=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 + +WORKDIR /app + +COPY requirements.txt . +RUN pip install --no-cache-dir \ + -i "${PIP_INDEX_URL}" \ + --trusted-host "${PIP_TRUSTED_HOST}" \ + -r requirements.txt + +COPY app.py download_model.py entrypoint.sh ./ +RUN chmod +x entrypoint.sh + +# 模型目录。挂卷到这里,重启容器时模型不会丢。 +ENV MODEL_DIR=/models +RUN mkdir -p /models + +EXPOSE 8000 + +# slim 镜像里没有 curl,直接用 python 做健康检查 +HEALTHCHECK --interval=30s --timeout=5s --start-period=600s --retries=3 \ + CMD python -c "import urllib.request,sys; sys.exit(0 if urllib.request.urlopen('http://127.0.0.1:8000/health',timeout=4).status==200 else 1)" + +CMD ["./entrypoint.sh"] \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..cb21194 --- /dev/null +++ b/README.md @@ -0,0 +1,177 @@ +# local-remarker + +本地重排(rerank)服务。零 API Key、文档不出内网,对外提供 SiliconFlow 兼容接口。 + +起因是在给 [PandaWiki](https://github.com/chaitin/PandaWiki) 做本地化部署时, +「重排序模型」被卡住了,一共有四个独立的阻塞点: + +| # | 阻塞点 | 实测结论 | +|---|---|---| +| 1 | Ollama 不提供重排 | Ollama 0.40.8 的 `/api/rerank`、`/v1/rerank` 都是纯 404 | +| 2 | 调用方把路径写死 | PandaWiki 的 raglite 只认 `/v1/services/rerank/text`(SilhouFlow 格式) | +| 3 | 装不上 torch | `download.pytorch.org/whl/cpu` 与阿里云 pytorch-wheels 下载均 20 秒零进度 | +| 4 | 官方模型站不可达 | `registry.ollama.ai` 的 blob 下载完全卡死,`hf-mirror.com` 连不上 | + +最终方案是**用 ONNX Runtime,绕开 torch**。ModelScope 上 `Xenova/bge-reranker-base` +自带量化 ONNX,于是第 3、4 两个阻塞点同时消失 —— 整条链路只依赖 +**ModelScope 一个国内源**,实测 15~25 MB/s。 + +## 快速开始 + +```bash +docker compose up -d +``` + +首次启动会下载约 290MB 模型到 `./models`,日志形如: + +``` +[local-remarker] 模型目录: /models + [model_int8.onnx] 需要下载 ... + [model_int8.onnx] 266.0MB / 266.0MB (100%) + [model_int8.onnx] 校验通过 +[local-remarker] 模型准备完成 +``` + +之后无论重启还是重建容器,都不会再产生下载流量。 + +服务监听 **63001**: + +```bash +curl -s localhost:63001/health +# {"status":"ok","model":"bge-reranker-base"} + +curl -s -X POST localhost:63001/v1/services/rerank/text \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "bge-reranker-base", + "query": "如何部署 PandaWiki 知识库", + "documents": [ + "PandaWiki 支持通过 Docker 部署,数据保存在 postgres 与 minio", + "今天股市行情如何,上证指数涨跌", + "部署时需要先配置 embedding 模型,否则无法创建知识库" + ], + "top_n": 3, + "return_documents": true + }' +``` + +## 设计要点 + +### 模型不在镜像里,启动时按需下载 + +`download_model.py` 在 entrypoint 里跑,做三件事: + +1. **检查** —— 每个文件比对 `大小 + sha256`,全部通过就立刻返回,零网络请求 +2. **续传** —— 下载落到 `.part` 文件,重试时带 `Range` 头从断点继续 +3. **原子替换** —— 下载完先校验 sha256,通过才 `os.replace` 成正式文件, + 避免中断留下的半个文件被当成有效模型加载 + +模型仓库、版本、下载源都是环境变量,换模型不用重新 build。 + +### 路径兜底 + +调用方填的 base_url 可能是 `http://host:8000`,也可能是 `http://host:8000/v1`, +框架在其后拼接路径后可能出现 `/v1/v1/services/rerank/text` 之类的变体。所以除了 +显式注册常见路径,还加了一个兜底路由:路径中含 `rerank` 或以 `/text` 结尾的 +POST 一律按重排处理;无关路径照常返回 404,不会误吞流量。 + +### 打分逻辑 + +遵循 bge-reranker 系列的惯例:把 `(query, document)` 拼成句对批量推理, +取 `logits[:, 0]` 作为相关性分数 —— 用的是**原始 logit,不过 sigmoid**, +然后降序取 top_n。 + +### 为什么选 int8 + +宿主机是 i7-1165G7,带 **AVX512-VNNI**,int8 推理能直接吃到该指令集加速。 +266MB 的 int8 vs 1061MB 的 fp32,精度损失很小但速度和体积都明显更优。 + +同仓库可选:`model_quantized.onnx` (266MB)、`model_fp16.onnx` (531MB)、 +`model.onnx` (1061MB)。换法见 `docker-compose.yml` 里的注释。 + +## 接口 + +### `POST /v1/services/rerank/text` + +兼容 SiliconFlow,请求: + +```json +{ + "model": "bge-reranker-base", + "query": "查询词", + "documents": ["片段1", "片段2"], + "top_n": 5, + "return_documents": false +} +``` + +响应: + +```json +{ + "id": "…", + "object": "rerank.result", + "model": "bge-reranker-base", + "results": [ + {"index": 0, "relevance_score": 6.8775} + ], + "usage": {"total_tokens": 44}, + "meta": {"latency_ms": 18} +} +``` + +`relevance_score` 是原始 logit,**无固定取值范围**,只保证相对大小关系正确。 +需要 0~1 区间的话调用方自行过 sigmoid。 + +### 其它 + +- `GET /health` —— 存活探测 +- `GET /info` —— 当前模型、路径、输入张量名 +- 同样接受 `POST /services/rerank/text`、`/v1/rerank`、`/rerank`、`/rank` + +## 配置 + +全部写在 `docker-compose.yml` 的 `environment` 里,改完 `docker compose up -d` 生效。常用的几个: + +| 变量 | 默认 | 说明 | +|---|---|---| +| `PORT`(宿主机侧) | `63001` | compose 的 `ports`,映射到容器 8000 | +| `MODEL_DIR` | `/models` | 容器内模型目录 | +| `MODEL_REPO` | `Xenova/bge-reranker-base` | 换模型改这里 | +| `MODEL_BASE` | ModelScope | 换下载源改这里 | +| `ONNX_FILE` | `model_int8.onnx` | 换量化档位改这里 | +| `MAX_LEN` | `512` | 截断长度,长文档可调大(吃内存) | +| `BATCH_SIZE` | `16` | 每次推理的句对数 | +| `WORKERS` | `1` | 调大前先看内存,每个 worker 各加载一份模型 | + +## 实测性能 + +i7-1165G7 / 8 逻辑核 / 15GB 内存: + +| 场景 | 耗时 | +|---|---| +| 3 个文档重排 | 14~18 ms | +| 镜像大小 | 401 MB | +| 模型大小 | 290 MB | + +## 已知限制 + +- **单 worker**:默认 `--workers 1`,每个 worker 会各自加载一份模型(约 400MB 内存)。 + 高并发场景应调 `WORKERS`,同时注意内存翻倍。 +- **不鉴权**:不校验 `api_key`。放到非受信网络需要自行加反向代理。 +- **仅 CPU**:未做 GPU 加速,也没接 TensorRT/ONNXRuntime-GPU。 +- **只支持 ONNX**:换 safetensors 格式的模型需要调整 `download_model.py` 的 + `MANIFEST` 并改用 torch 运行时。 + +## 文件 + +``` +local-remarker/ +├── app.py FastAPI 服务本体 +├── download_model.py 模型检查/下载/校验(幂等,支持续传) +├── entrypoint.sh 先下模型,再起 uvicorn +├── requirements.txt 运行时依赖(无 torch) +├── Dockerfile +├── docker-compose.yml 参数全部写在这里 +└── .gitignore +``` \ No newline at end of file diff --git a/app.py b/app.py new file mode 100644 index 0000000..c2c55f8 --- /dev/null +++ b/app.py @@ -0,0 +1,181 @@ +"""本地重排(rerank)服务。 + +对外暴露 SiliconFlow 兼容的 `POST /v1/services/rerank/text`,方便直接对接 +各类 RAG 框架。之所以自己实现而不是用现成推理框架,主要是因为: +很多推理运行时(如 Ollama)并不提供 rerank 端点。 + +打分逻辑遵循 bge-reranker 系列的惯例:把 (query, document) 拼成句对批量推理, +取 logits[:, 0] 作为相关性分数(原始 logit,**不过 sigmoid**),降序后取 top_n。 +""" + +import os +import time +import uuid +from typing import Any + +import numpy as np +import onnxruntime as ort +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse +from tokenizers import Tokenizer + +# ---------------------------------------------------------------- 配置 +MODEL_DIR = os.environ.get("MODEL_DIR", "/models") +ONNX_FILE = os.environ.get("ONNX_FILE", "model_int8.onnx") +TOKENIZER_FILE = os.environ.get("TOKENIZER_FILE", "tokenizer.json") +MODEL_NAME = os.environ.get("MODEL_NAME", "bge-reranker-base") +MAX_LEN = int(os.environ.get("MAX_LEN", "512")) +BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "16")) +NUM_THREADS = os.environ.get("NUM_THREADS", "") or None + +app = FastAPI(title="local-remarker", version="1.0") + + +def _build_session() -> "tuple[ort.InferenceSession, Tokenizer]": + onnx_path = os.path.join(MODEL_DIR, ONNX_FILE) + tok_path = os.path.join(MODEL_DIR, TOKENIZER_FILE) + for path in (onnx_path, tok_path): + if not os.path.isfile(path): + raise FileNotFoundError( + f"缺少模型文件: {path}(entrypoint 应先执行 download_model.py)" + ) + + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + if NUM_THREADS: + options.intra_op_num_threads = int(NUM_THREADS) + else: + # 默认吃满可用核;AVX512-VNNI 的机器上 int8 推理收益明显 + options.intra_op_num_threads = os.cpu_count() or 4 + + session = ort.InferenceSession( + onnx_path, sess_options=options, providers=["CPUExecutionProvider"] + ) + + tokenizer = Tokenizer.from_file(tok_path) + tokenizer.enable_truncation(max_length=MAX_LEN) + tokenizer.enable_padding(pad_id=1, pad_token="[PAD]") + return session, tokenizer + + +SESSION, TOKENIZER = _build_session() +INPUT_NAMES = {i.name for i in SESSION.get_inputs()} + + +def _score(query: str, documents: list[str]) -> list[float]: + """对 (query, doc) 句对批量打分,返回与 documents 等长的分数列表。""" + encodings = TOKENIZER.encode_batch([[query, doc] for doc in documents]) + input_ids = np.array([e.ids for e in encodings], dtype=np.int64) + attention_mask = np.array([e.attention_mask for e in encodings], dtype=np.int64) + + feeds = {"input_ids": input_ids, "attention_mask": attention_mask} + if "token_type_ids" in INPUT_NAMES: + feeds["token_type_ids"] = np.array([e.type_ids for e in encodings], dtype=np.int64) + + logits = SESSION.run(None, feeds)[0] + # bge-reranker 只有 1 个 logit,取第 0 列即相关性分数(原始 logit) + return [float(v) for v in np.asarray(logits).reshape(len(documents), -1)[:, 0]] + + +def _rerank(query: str, documents: list[str], top_n: int, + return_documents: bool) -> dict[str, Any]: + started = time.time() + + scores: list[float] = [] + for start in range(0, len(documents), BATCH_SIZE): + scores.extend(_score(query, documents[start:start + BATCH_SIZE])) + + order = sorted(range(len(documents)), key=lambda i: scores[i], + reverse=True)[:top_n] + + results = [] + for index in order: + item: dict[str, Any] = { + "index": index, + "relevance_score": scores[index], + } + if return_documents: + item["document"] = {"text": documents[index]} + results.append(item) + + total_tokens = sum(len(TOKENIZER.encode(query, doc).ids) for doc in documents) + return { + "id": uuid.uuid4().hex, + "object": "rerank.result", + "model": MODEL_NAME, + "results": results, + "usage": {"total_tokens": total_tokens}, + "meta": {"latency_ms": int((time.time() - started) * 1000)}, + } + + +async def _handle(request: Request) -> JSONResponse: + try: + body = await request.json() + except Exception: # noqa: BLE001 + return JSONResponse( + {"error": {"message": "invalid json body"}}, status_code=400 + ) + + query = body.get("query") or "" + documents = body.get("documents") or [] + top_n = int(body.get("top_n") or len(documents) or 1) + return_documents = bool(body.get("return_documents", False)) + + if not query or not documents: + return JSONResponse( + {"error": {"message": "query and documents are required"}}, + status_code=400, + ) + + try: + result = _rerank(query, documents, top_n, return_documents) + except Exception as exc: # noqa: BLE001 + return JSONResponse({"error": {"message": str(exc)}}, status_code=500) + + return JSONResponse(result) + + +# 显式注册常见路径,避免依赖调用方的拼接方式 +for _path in ( + "/v1/services/rerank/text", + "/services/rerank/text", + "/v1/rerank", + "/rerank", + "/rank", + "/v1/rank", +): + app.post(_path)(_handle) + + +@app.api_route("/{full_path:path}", methods=["POST"]) +async def _catch_all(full_path: str, request: Request) -> JSONResponse: + """兜底路由。 + + 调用方填写的 base_url 可能是 `http://host:8000` 或 `http://host:8000/v1`, + 框架在其后拼接路径后可能出现 `/v1/v1/services/rerank/text` 之类的变体。 + 凡是路径中含 rerank、或以 /text 结尾的 POST 一律按重排请求处理; + 无关路径仍返回 404,不会误吞流量。 + """ + if "rerank" not in full_path and not full_path.endswith("/text"): + return JSONResponse( + {"error": {"message": f"not found: /{full_path}"}}, status_code=404 + ) + return await _handle(request) + + +@app.get("/health") +def health() -> dict[str, Any]: + return {"status": "ok", "model": MODEL_NAME} + + +@app.get("/info") +def info() -> dict[str, Any]: + return { + "model": MODEL_NAME, + "onnx_file": ONNX_FILE, + "model_dir": MODEL_DIR, + "max_len": MAX_LEN, + "batch_size": BATCH_SIZE, + "input_names": sorted(INPUT_NAMES), + } \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..3cc1fdc --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,62 @@ +# 本地重排服务。 +# +# 参数全部直接写在本文件里,不再走 .env —— 需要调整时改这里,然后 +# docker compose up -d 重建即可。 +# +# 首次启动会下载约 290MB 模型权重到 ./models; +# 之后重启/重建都会先校验 sha256,通过就跳过下载,零网络流量。 + +services: + local-remarker: + build: . + image: local-remarker:latest + container_name: local-remarker + restart: unless-stopped + + # 宿主机 63001 -> 容器 8000 + ports: + - "63001:8000" + + environment: + # ---- 模型存放位置(容器内)---- + MODEL_DIR: /models + + # ---- 模型来源与选择 ---- + MODEL_REPO: Xenova/bge-reranker-base + MODEL_REV: master + MODEL_BASE: https://modelscope.cn/models/Xenova/bge-reranker-base/resolve/master + MODEL_NAME: bge-reranker-base + # 同仓库可选量化档位: + # model_int8.onnx 266MB 默认,带 AVX512-VNNI 的 CPU 上最快 + # model_quantized.onnx 266MB + # model_fp16.onnx 531MB + # model.onnx 1061MB 原始精度,CPU 上明显变慢 + # 更大的模型可改用 BAAI/bge-reranker-v2-m3(2.2GB,质量最好但重排较慢) + ONNX_FILE: model_int8.onnx + + # ---- 推理参数 ---- + MAX_LEN: "512" # 截断长度,长文档可调大(吃内存) + BATCH_SIZE: "16" # 每次推理的句对数 + + # ---- 服务 ---- + WORKERS: "1" # 调大前先看内存:每个 worker 会各加载一份模型 + PORT: "8000" + LOG_LEVEL: info + TZ: Asia/Shanghai + + # 宿主机 ./models -> 容器 /models。指向大盘路径可避免占满系统盘, + # 例如改成 /data/models/remarker:/models + volumes: + - ./models:/models + + healthcheck: + test: + - CMD + - python + - -c + - "import urllib.request,sys; sys.exit(0 if urllib.request.urlopen('http://127.0.0.1:8000/health',timeout=4).status==200 else 1)" + interval: 30s + timeout: 5s + # 首次启动要下模型,给足预热时间 + start_period: 600s + retries: 3 \ No newline at end of file diff --git a/download_model.py b/download_model.py new file mode 100644 index 0000000..c3431f7 --- /dev/null +++ b/download_model.py @@ -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()) \ No newline at end of file diff --git a/entrypoint.sh b/entrypoint.sh new file mode 100755 index 0000000..c609dbc --- /dev/null +++ b/entrypoint.sh @@ -0,0 +1,17 @@ +#!/bin/sh +# 容器入口:先确保模型就位,再启动服务。 +# +# download_model.py 是幂等的 —— 模型已存在且 sha256 校验通过时会直接返回, +# 因此正常的重启/重建容器不会产生任何下载流量;首次启动才会拉取约 290MB 权重。 + +set -e + +echo "[entrypoint] 检查模型 ..." +python /app/download_model.py + +echo "[entrypoint] 启动 rerank 服务 ..." +exec uvicorn app:app \ + --host "${HOST:-0.0.0.0}" \ + --port "${PORT:-8000}" \ + --workers "${WORKERS:-1}" \ + --log-level "${LOG_LEVEL:-info}" \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..5f7d3d7 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +fastapi==0.115.6 +uvicorn[standard]==0.34.0 +onnxruntime==1.20.1 +tokenizers==0.21.0 +numpy==2.2.1 \ No newline at end of file