Blame for code/scripts/serve_cpu.py

#!/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()