laya-ad / code /scripts /serve_cpu.py
x6nux's picture
laya-ad v1.1:六精度权重 + 模型卡 + 项目代码快照
71dabc7 verified
Raw History Blame Contribute Delete
14.2 kB
#!/usr/bin/env python3
"""mmBERT 广告识别 CPU 推理服务(仅标准库 HTTP + 请求队列 + 微批 worker)。
模型:outputs/ad-clf-int8/quantized.pt(int8 动态量化;采纳结论 outputs/cpu_quant_bench.json:
int8 acc=0.966 vs fp32 0.968、speedup=4.13x、accepted=true)。
加载方法沿 scripts/bench_cpu.py:22-26 注释:必须先 import transformers 让 ModernBert* 类
可反序列化,再 torch.load(..., weights_only=False) 加载完整 pickle 的 nn.Module。
接口(默认 127.0.0.1:8765):
GET /health -> {"status": "ok", ...} 服务与模型就绪状态
POST /predict body {"texts": ["..."]} 一次可传多条文本(一次性多判断),
-> {"results": [{"label": "ad"|"normal", "probs": [p_ad, p_normal]}, ...]}
结果按 texts 输入顺序返回;单条文本在服务内部合并进微批前向。
微批设计:
- HTTP 线程只做 tokenize(不 padding),把"单条文本任务"推入内部队列后等待结果;
- 微批 worker 凑批:拿到第一个任务后开 ≤20ms 凑批窗(--batch-wait-ms),
窗内凑满 max_batch(--max-batch,默认 32)或窗关即执行;
- 长度分桶(--short-bucket-max,默认 64 token):短桶/长桶各一个微批 worker——
动态 padding 的批长由批内最长文本决定,混桶时长文本会把整批拖长
(实测 batch=32: seq≈26 -> 154ms vs seq≈114 -> 867ms);
- 批内动态 padding 到批内最长(截断上限 --max-length,默认 128 = bench_cpu.py 量化
正确率口径 acc=0.966;val n=2077 分布 median=21/p90=78/p95=111),不固定 512 ——
广告文本平均很短(median=21),长截断会白付数倍 seq 算力;
- --no-microbatch:关闭合并,每条文本单独 batch=1 前向(压测对照模式)。
线程:torch.set_num_threads(物理核数);OMP/MKL 等 env 必须在 import torch 之前设置
(同 scripts/bench_cpu.py:47-58)。依据 scripts/bench_threads_scaling.py 实测
(AMD 9700X 8 物理核,int8 形态、seq128/batch8,3 轮中位):
1/2/4/8 线程 = 14.2/23.1/37.1/58.7 样本/s,满物理核 8 比半核 4 快 +58%,
半核档差 36.8%;12/16(SMT 逻辑核)超过物理核数,不在有效档。
注意:端到端 serve 压测(bench_serve.py,100 并发持续)4 与 8 线程持平(118 vs 119 RPS)——
持续高并发下 3 个微批 worker 已饱和 CPU;8 线程收益在单 worker 前向
(seq128/batch8 +58%,低并发/串行大文本受益),代价是串行小请求 p50 约 +3.5ms
(小 GEMM 的线程同步开销)。两者 SLO 均达标。
"""
import argparse
import json
import os
import queue
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
def _physical_cores() -> int:
try:
import psutil
n = psutil.cpu_count(logical=False)
if n:
return n
except Exception:
pass
return max(1, (os.cpu_count() or 2) // 2) # 无 psutil 时按超线程折半兜底
def parse_args():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--host", default="127.0.0.1")
ap.add_argument("--port", type=int, default=8765)
ap.add_argument("--model", default="outputs/ad-clf-int8/quantized.pt")
ap.add_argument("--tokenizer", default="outputs/ad-clf/best")
ap.add_argument("--max-batch", type=int, default=48, help="微批最大条数")
ap.add_argument("--batch-wait-ms", type=float, default=20.0, help="凑批等待窗(ms)")
ap.add_argument("--max-length", type=int, default=128,
help="token 截断上限(= bench_cpu.py 量化正确率口径 acc=0.966;"
"val n=2077: median=21/p90=78/p95=111,128 保住 96% 文本完整;"
"padding 始终为批内最长,见 --short-bucket-max 与排序出批)")
ap.add_argument("--max-chars", type=int, default=20000, help="单条文本字符数上限(超限 400)")
ap.add_argument("--max-texts", type=int, default=128, help="单请求 texts 条数上限(超限 400)")
ap.add_argument("--short-bucket-max", type=int, default=64,
help="长度分桶阈值(token):≤该值的进短桶、>进长桶,长桶 1 个微批 worker,"
"避免长文本把整批 padding 拖长;0 = 不分桶(单队列)")
ap.add_argument("--short-workers", type=int, default=2,
help="短桶微批 worker 数(默认 2:与长桶合计 3 个 worker,"
"burst 瞬时并行换整批 wall time;每个 worker 的前向用满 "
"THREADS=物理核数个 intra-op 线程,持续吞吐由 CPU 总量决定)")
ap.add_argument("--no-microbatch", action="store_true",
help="逐条前向对照模式:不合并批,每条文本单独前向")
return ap.parse_args()
ARGS = parse_args()
# ---- 线程环境变量:必须先于 import torch / transformers ----
# 实测(scripts/bench_threads_scaling.py,int8 seq128/batch8):8 物理核满开 58.7 sps
# 优于半核 4 线程 37.1 sps(+58%),scaling 在物理核内单调上升,无拐点回落。
PHYS_CORES = _physical_cores()
THREADS = PHYS_CORES
for _var in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS",
"NUMEXPR_NUM_THREADS", "VECLIB_MAXIMUM_THREADS"):
os.environ[_var] = str(THREADS)
os.environ["TOKENIZERS_PARALLELISM"] = "false"
import torch # noqa: E402
from transformers import AutoTokenizer # noqa: E402
torch.set_num_threads(THREADS)
if hasattr(sys.stdout, "reconfigure"):
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
def log(msg):
print(f"[serve] {time.strftime('%H:%M:%S')} {msg}", flush=True)
# ---- 模型加载 ----
t0 = time.perf_counter()
import transformers # noqa: E402 必须先于 torch.load(ModernBert* 反序列化注册)
tok = AutoTokenizer.from_pretrained(ARGS.tokenizer)
model = torch.load(ARGS.model, map_location="cpu", weights_only=False)
model.eval()
ID2LABEL = {int(k): v for k, v in model.config.id2label.items()}
NUM_LABELS = model.config.num_labels
PAD_ID = tok.pad_token_id if tok.pad_token_id is not None else 0
log(f"模型加载完成 {time.perf_counter() - t0:.1f}s model={ARGS.model} "
f"labels={ID2LABEL} threads={torch.get_num_threads()} "
f"max_batch={ARGS.max_batch} batch_wait_ms={ARGS.batch_wait_ms} "
f"max_length={ARGS.max_length} microbatch={not ARGS.no_microbatch}")
# ---- 微批核心 ----
# 任务: (input_ids: list[int], slot: dict) ;slot["done"]=Event, slot["probs"]=list[float]
# 分桶:短桶(≤SHORT_BUCKET_MAX token)/ 长桶两个队列,各一个微批 worker——
# 动态 padding 的批长由批内最长文本决定,混桶时长文本会把短文本的整批拖长,
# 实测 batch=32: seq≈26 -> 154ms vs seq≈114 -> 867ms,分桶后短批吞吐差 5 倍以上。
TASKQ_SHORT: "queue.Queue[tuple[list[int], dict]]" = queue.Queue()
TASKQ_LONG: "queue.Queue[tuple[list[int], dict]]" = queue.Queue()
SHORT_BUCKET_MAX = ARGS.short_bucket_max
def _forward_batch(tasks):
"""一批任务合并前向,批内动态 padding 到最长,就地写回 slot。"""
ids_list = [t[0] for t in tasks]
maxl = max(len(x) for x in ids_list)
input_ids = torch.tensor([x + [PAD_ID] * (maxl - len(x)) for x in ids_list],
dtype=torch.long)
attention_mask = torch.tensor([[1] * len(x) + [0] * (maxl - len(x)) for x in ids_list],
dtype=torch.long)
logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
probs = torch.softmax(logits.float(), dim=-1)
for (_, slot), p in zip(tasks, probs):
slot["probs"] = p.tolist()
slot["done"].set()
def _forward_single(task):
"""逐条前向(--no-microbatch 对照)。"""
ids, slot = task
input_ids = torch.tensor([ids], dtype=torch.long)
attention_mask = torch.ones_like(input_ids)
logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
slot["probs"] = torch.softmax(logits.float(), dim=-1)[0].tolist()
slot["done"].set()
def worker(q):
"""微批 worker:凑批窗 ≤batch_wait_ms;缓冲池按 token 长度排序出批——
动态 padding 的批长由批内最长文本决定,排序出批让同批长度相近,padding 损耗最小;
2s 老化优先防长任务被短流饿死。"""
# inference_mode 是线程本地的 RAII,包住 worker 整个生命周期
with torch.inference_mode():
buf = []
while True:
if ARGS.no_microbatch:
_forward_single(q.get())
continue
if not buf:
buf.append(q.get())
deadline = time.perf_counter() + ARGS.batch_wait_ms / 1000.0
while len(buf) < ARGS.max_batch:
remain = deadline - time.perf_counter()
if remain <= 0:
break
try:
buf.append(q.get(timeout=remain))
except queue.Empty:
break
oldest = min(buf, key=lambda t: t[1]["arrival"])
if time.perf_counter() - oldest[1]["arrival"] > 2.0:
# 老化:以最老任务长度为中心取最接近的 max_batch 条(必含最老)
anchor = len(oldest[0])
buf.sort(key=lambda t: abs(len(t[0]) - anchor))
take = buf[:ARGS.max_batch]
for t in take:
buf.remove(t)
else:
buf.sort(key=lambda t: len(t[0]))
take = buf[:ARGS.max_batch]
del buf[:ARGS.max_batch]
_forward_batch(take)
def start_workers():
if ARGS.no_microbatch or SHORT_BUCKET_MAX <= 0:
# 不分桶/逐条:单队列单 worker
threading.Thread(target=worker, args=(TASKQ_SHORT,), daemon=True,
name="mbatch-worker").start()
return
for i in range(max(1, ARGS.short_workers)):
threading.Thread(target=worker, args=(TASKQ_SHORT,), daemon=True,
name=f"mbatch-worker-short{i}").start()
threading.Thread(target=worker, args=(TASKQ_LONG,), daemon=True,
name="mbatch-worker-long").start()
start_workers()
# ---- HTTP 层 ----
class Handler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1" # keep-alive,压测可复用连接
server_version = "laya-ad-serve/1.0"
def log_message(self, fmt, *args): # 静默逐请求日志,压测不刷屏
pass
def _send_json(self, code, obj):
body = json.dumps(obj, ensure_ascii=False).encode("utf-8")
self.send_response(code)
self.send_header("Content-Type", "application/json; charset=utf-8")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def do_GET(self):
if self.path == "/health":
self._send_json(200, {
"status": "ok",
"pid": os.getpid(), # 真实解释器 pid(venv python.exe 是 trampoline,Popen.pid 不是它)
"model": ARGS.model,
"labels": [ID2LABEL[i] for i in range(NUM_LABELS)],
"torch_threads": torch.get_num_threads(),
"physical_cores": PHYS_CORES,
"max_batch": ARGS.max_batch,
"batch_wait_ms": ARGS.batch_wait_ms,
"max_length": ARGS.max_length,
"short_bucket_max": SHORT_BUCKET_MAX,
"microbatch": not ARGS.no_microbatch,
})
else:
self._send_json(404, {"error": "not found"})
def do_POST(self):
if self.path != "/predict":
self._send_json(404, {"error": "not found"})
return
try:
n = int(self.headers.get("Content-Length") or 0)
req = json.loads(self.rfile.read(n).decode("utf-8"))
texts = req["texts"]
if not isinstance(texts, list):
raise ValueError("texts 必须是列表")
texts = [str(t) for t in texts]
if not texts:
raise ValueError("texts 不能为空")
if len(texts) > ARGS.max_texts:
raise ValueError(f"texts 超过单请求上限 {ARGS.max_texts}")
if any(len(t) > ARGS.max_chars for t in texts):
raise ValueError(f"单条文本超过 {ARGS.max_chars} 字符上限")
except Exception as e:
self._send_json(400, {"error": str(e)})
return
# 请求线程 tokenize 后逐条入队(按长度分桶),等各自 slot 完成,按原顺序组装
bucketed = (not ARGS.no_microbatch) and SHORT_BUCKET_MAX > 0
slots = []
for t in texts:
ids = tok(t, truncation=True, max_length=ARGS.max_length)["input_ids"]
slot = {"done": threading.Event(), "probs": None,
"arrival": time.perf_counter()}
if bucketed and len(ids) > SHORT_BUCKET_MAX:
TASKQ_LONG.put((ids, slot))
else:
TASKQ_SHORT.put((ids, slot))
slots.append(slot)
for s in slots:
s["done"].wait()
results = []
for s in slots:
p = s["probs"]
results.append({"label": ID2LABEL[max(range(len(p)), key=p.__getitem__)],
"probs": p})
self._send_json(200, {"results": results})
class Server(ThreadingHTTPServer):
daemon_threads = True
request_queue_size = 256 # socketserver 默认 backlog=5,100 并发瞬时连接会被拒
def main():
srv = Server((ARGS.host, ARGS.port), Handler)
log(f"监听 http://{ARGS.host}:{ARGS.port} (/health, /predict)")
try:
srv.serve_forever()
except KeyboardInterrupt:
pass
finally:
srv.server_close()
if __name__ == "__main__":
main()