ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter2/local_llm_serving/benchmark.py

@@ -0,0 +1,487 @@
+#!/usr/bin/env python3
+"""
+本地 LLM 服务性能基准（实验 2-1 配套）
+
+本脚本通过 OpenAI 兼容接口（vLLM 或 Ollama 均可）测量本地部署的小模型在
+「服务（serving）」层面的三个核心指标，帮助读者建立对吞吐 / 延迟 / 批处理 /
+KV Cache 的直觉：
+
+  1. throughput —— 单流解码吞吐（tokens/s）与首 token 延迟（TTFT）
+  2. kv-cache  —— 前缀缓存命中 vs 未命中的 TTFT 对比
+                （对应实验 2-1 第 5 点：系统提示词不变时缓存命中更快，
+                 修改系统提示词开头几个字符导致缓存失效、需重算整个前缀）
+  3. batching  —— 不同并发度下的聚合吞吐，直观展示批处理带来的吞吐提升
+
+所有数字均来自真实服务端的实测，脚本本身不产生任何合成数据。
+如果尚未启动服务端，可用 --dry-run 离线查看每个场景将要发出的请求配置。
+
+示例：
+    # 先启动服务端（二选一）
+    python server.py                    # vLLM（需要 NVIDIA GPU）
+    ollama serve && ollama pull qwen3:0.6b   # Ollama（Mac / 无 GPU）
+
+    # 跑全部场景并保存结果
+    python benchmark.py --scenario all --output results.json
+
+    # 只看 KV Cache 命中 / 未命中的 TTFT 对比
+    python benchmark.py --scenario kv-cache --backend ollama
+
+    # 批处理吞吐扫描
+    python benchmark.py --scenario batching --concurrency 1,2,4,8
+"""
+
+import argparse
+import json
+import logging
+import statistics
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, List, Optional
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("benchmark")
+
+# 各后端的默认 OpenAI 兼容地址
+BACKEND_DEFAULTS = {
+    "vllm": {"base_url": "http://localhost:8000/v1", "model": "Qwen3-0.6B"},
+    "ollama": {"base_url": "http://localhost:11434/v1", "model": "qwen3:0.6b"},
+}
+
+# 一段确定性的填充文本，用于把共享前缀撑长，让 KV Cache 的效果更明显
+_FILLER_SENTENCE = (
+    "You are a meticulous assistant that follows the operating manual precisely. "
+)
+
+
+def build_padded_system_prompt(target_tokens: int) -> str:
+    """构造一个约含 target_tokens 个 token 的系统提示词（用重复句子填充）。
+
+    这里用「4 字符 ≈ 1 token」的粗略估计来控制长度，只需保证前缀足够长、
+    可复现即可，不追求精确的 token 数。
+    """
+    header = (
+        "# Operating Manual\n"
+        "You are a helpful local assistant deployed for the AI Agent book experiment.\n\n"
+    )
+    approx_chars = max(0, target_tokens * 4 - len(header))
+    repeats = approx_chars // len(_FILLER_SENTENCE) + 1
+    body = _FILLER_SENTENCE * repeats
+    return header + body
+
+
+def make_client(base_url: str, api_key: str):
+    """创建 OpenAI 兼容客户端。"""
+    try:
+        from openai import OpenAI
+    except ImportError:
+        logger.error("缺少依赖 openai，请先执行：pip install openai")
+        sys.exit(1)
+    return OpenAI(base_url=base_url, api_key=api_key)
+
+
+def stream_once(
+    client,
+    model: str,
+    messages: List[Dict[str, str]],
+    max_tokens: int,
+    temperature: float,
+) -> Dict[str, float]:
+    """发起一次流式请求，返回 TTFT、总时长、输出 token 数与解码吞吐。
+
+    - ttft：从发起请求到收到第一个内容或推理分片的时间（秒）
+    - total：整个响应的墙钟时间（秒）
+    - output_tokens：优先取服务端返回的 usage.completion_tokens，
+      否则用收到的内容分片数量作为近似
+    - decode_tps：解码阶段吞吐 = 输出 token / (总时长 - TTFT)
+    """
+    start = time.perf_counter()
+    ttft: Optional[float] = None
+    chunk_count = 0
+    usage_tokens: Optional[int] = None
+
+    stream = client.chat.completions.create(
+        model=model,
+        messages=messages,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        stream=True,
+        stream_options={"include_usage": True},
+    )
+
+    for chunk in stream:
+        # 最后一个分片可能只携带 usage 而没有 choices
+        if getattr(chunk, "usage", None) is not None:
+            try:
+                usage_tokens = chunk.usage.completion_tokens
+            except AttributeError:
+                pass
+        if not chunk.choices:
+            continue
+        delta = chunk.choices[0].delta
+        text = (
+            getattr(delta, "content", None)
+            or getattr(delta, "reasoning_content", None)
+            or getattr(delta, "reasoning", None)
+        )
+        if text:
+            if ttft is None:
+                ttft = time.perf_counter() - start
+            chunk_count += 1
+
+    total = time.perf_counter() - start
+    if ttft is None:
+        ttft = total
+    output_tokens = usage_tokens if usage_tokens is not None else chunk_count
+    decode_time = max(total - ttft, 1e-6)
+    decode_tps = output_tokens / decode_time if output_tokens else 0.0
+
+    return {
+        "ttft": ttft,
+        "total": total,
+        "output_tokens": float(output_tokens),
+        "decode_tps": decode_tps,
+    }
+
+
+# --------------------------------------------------------------------------- #
+# 场景实现
+# --------------------------------------------------------------------------- #
+def scenario_throughput(client, model, args) -> Dict[str, Any]:
+    """单流吞吐 + TTFT：连续发起若干次解码密集的请求并汇总统计。"""
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {
+            "role": "user",
+            "content": "Write a detailed explanation of how KV Cache works in transformer inference.",
+        },
+    ]
+    runs = []
+    for i in range(args.repeats):
+        r = stream_once(client, model, messages, args.max_tokens, args.temperature)
+        runs.append(r)
+        logger.info(
+            "throughput 第 %d/%d 次: TTFT=%.3fs, 解码=%.1f tok/s, 输出=%d tok",
+            i + 1, args.repeats, r["ttft"], r["decode_tps"], int(r["output_tokens"]),
+        )
+    return {
+        "scenario": "throughput",
+        "repeats": args.repeats,
+        "ttft_mean_s": statistics.fmean(x["ttft"] for x in runs),
+        "decode_tps_mean": statistics.fmean(x["decode_tps"] for x in runs),
+        "output_tokens_mean": statistics.fmean(x["output_tokens"] for x in runs),
+        "runs": runs,
+    }
+
+
+def scenario_kv_cache(client, model, args) -> Dict[str, Any]:
+    """KV Cache 命中 vs 未命中的 TTFT 对比（实验 2-1 第 5 点）。
+
+    - 命中组：系统提示词逐字节不变，重复发送同一请求，服务端前缀缓存命中，
+      prefill 几乎可以跳过 → TTFT 明显更低。
+    - 未命中组：每次只在系统提示词「开头」插入一个不同的计数串，前缀被改写，
+      缓存全部失效，服务端必须重算整个前缀 → TTFT 明显更高。
+    两组的提示词长度基本一致，因此差异主要来自前缀缓存是否命中。
+    """
+    base_prompt = build_padded_system_prompt(args.prefix_tokens)
+    user_msg = {"role": "user", "content": "In one short sentence, say hello."}
+
+    # 预热：先发一次把缓存写入（这一次一定是冷启动，不计入统计）
+    warm_msgs = [{"role": "system", "content": base_prompt}, user_msg]
+    stream_once(client, model, warm_msgs, args.max_tokens, args.temperature)
+
+    hit_ttfts, miss_ttfts = [], []
+    for i in range(args.repeats):
+        # 命中：完全相同的前缀
+        hit = stream_once(client, model, warm_msgs, args.max_tokens, args.temperature)
+        hit_ttfts.append(hit["ttft"])
+
+        # 未命中：在开头插入唯一前缀，使缓存失效
+        mutated = f"[req-{i}-{time.time_ns()}] " + base_prompt
+        miss_msgs = [{"role": "system", "content": mutated}, user_msg]
+        miss = stream_once(client, model, miss_msgs, args.max_tokens, args.temperature)
+        miss_ttfts.append(miss["ttft"])
+
+        logger.info(
+            "kv-cache 第 %d/%d 次: 命中 TTFT=%.3fs, 未命中 TTFT=%.3fs",
+            i + 1, args.repeats, hit["ttft"], miss["ttft"],
+        )
+
+    hit_mean = statistics.fmean(hit_ttfts)
+    miss_mean = statistics.fmean(miss_ttfts)
+    return {
+        "scenario": "kv-cache",
+        "prefix_tokens_approx": args.prefix_tokens,
+        "repeats": args.repeats,
+        "ttft_hit_mean_s": hit_mean,
+        "ttft_miss_mean_s": miss_mean,
+        "speedup": (miss_mean / hit_mean) if hit_mean > 0 else None,
+        "ttft_hit_s": hit_ttfts,
+        "ttft_miss_s": miss_ttfts,
+    }
+
+
+def scenario_batching(client, model, args) -> Dict[str, Any]:
+    """批处理：在不同并发度下并发发起请求，测量聚合吞吐。
+
+    连续批处理（continuous batching）是本地 serving 的核心优化：并发越高，
+    GPU 利用率越充分，系统聚合吞吐（所有请求合计 tok/s）通常显著上升，
+    但单个请求的延迟可能上升。此场景把这个权衡直接量化出来。
+    """
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Explain what a large language model is."},
+    ]
+
+    levels = args.concurrency
+    rows = []
+    for level in levels:
+        start = time.perf_counter()
+        with ThreadPoolExecutor(max_workers=level) as pool:
+            futures = [
+                pool.submit(
+                    stream_once, client, model, messages, args.max_tokens, args.temperature
+                )
+                for _ in range(level)
+            ]
+            results = [f.result() for f in futures]
+        wall = time.perf_counter() - start
+        total_tokens = sum(r["output_tokens"] for r in results)
+        agg_tps = total_tokens / wall if wall > 0 else 0.0
+        per_req_tps = agg_tps / level if level else 0.0
+        rows.append(
+            {
+                "concurrency": level,
+                "wall_s": wall,
+                "total_output_tokens": total_tokens,
+                "aggregate_tps": agg_tps,
+                "per_request_tps": per_req_tps,
+                "ttft_mean_s": statistics.fmean(r["ttft"] for r in results),
+            }
+        )
+        logger.info(
+            "batching 并发=%d: 聚合吞吐=%.1f tok/s, 单请求=%.1f tok/s, 墙钟=%.2fs",
+            level, agg_tps, per_req_tps, wall,
+        )
+    return {"scenario": "batching", "levels": rows}
+
+
+# --------------------------------------------------------------------------- #
+# 结果表格
+# --------------------------------------------------------------------------- #
+def print_report(results: List[Dict[str, Any]]) -> None:
+    print("\n" + "=" * 68)
+    print("本地 LLM 服务基准结果")
+    print("=" * 68)
+    for res in results:
+        s = res["scenario"]
+        if s == "throughput":
+            print("\n[throughput] 单流吞吐 / 首 token 延迟")
+            print(f"  次数           : {res['repeats']}")
+            print(f"  平均 TTFT      : {res['ttft_mean_s']:.3f} s")
+            print(f"  平均解码吞吐   : {res['decode_tps_mean']:.1f} tok/s")
+            print(f"  平均输出长度   : {res['output_tokens_mean']:.0f} tok")
+        elif s == "kv-cache":
+            print("\n[kv-cache] 前缀缓存命中 vs 未命中（TTFT）")
+            print(f"  前缀长度(约)   : {res['prefix_tokens_approx']} tok")
+            print(f"  命中平均 TTFT  : {res['ttft_hit_mean_s']:.3f} s")
+            print(f"  未命中平均TTFT : {res['ttft_miss_mean_s']:.3f} s")
+            if res.get("speedup"):
+                print(f"  缓存加速比     : {res['speedup']:.2f}x")
+        elif s == "batching":
+            print("\n[batching] 并发度对聚合吞吐的影响")
+            print(f"  {'并发':>4} | {'聚合tok/s':>10} | {'单请求tok/s':>12} | {'平均TTFT(s)':>11} | {'墙钟(s)':>8}")
+            print(f"  {'-'*4}-+-{'-'*10}-+-{'-'*12}-+-{'-'*11}-+-{'-'*8}")
+            for row in res["levels"]:
+                print(
+                    f"  {row['concurrency']:>4} | {row['aggregate_tps']:>10.1f} | "
+                    f"{row['per_request_tps']:>12.1f} | {row['ttft_mean_s']:>11.3f} | {row['wall_s']:>8.2f}"
+                )
+    print("\n" + "=" * 68)
+
+
+def describe_dry_run(args) -> None:
+    """离线打印将要执行的场景配置，不访问服务端。"""
+    print("=" * 68)
+    print("DRY RUN —— 仅打印计划，不访问服务端")
+    print("=" * 68)
+    print(f"后端         : {args.backend}")
+    print(f"base_url     : {args.base_url}")
+    print(f"模型         : {args.model}")
+    print(f"重复次数     : {args.repeats}")
+    print(f"max_tokens   : {args.max_tokens}")
+    print(f"temperature  : {args.temperature}")
+    scenarios = ["throughput", "kv-cache", "batching"] if args.scenario == "all" else [args.scenario]
+    print(f"待运行场景   : {', '.join(scenarios)}")
+    if "kv-cache" in scenarios:
+        prompt = build_padded_system_prompt(args.prefix_tokens)
+        print(f"  kv-cache   : 填充前缀约 {args.prefix_tokens} tok（实际 {len(prompt)} 字符）")
+    if "batching" in scenarios:
+        print(f"  batching   : 并发扫描 {args.concurrency}")
+    print("=" * 68)
+
+
+def parse_concurrency(value: str) -> List[int]:
+    try:
+        levels = [int(x) for x in value.split(",") if x.strip()]
+    except ValueError:
+        raise argparse.ArgumentTypeError("--concurrency 需为逗号分隔的正整数，例如 1,2,4,8")
+    if not levels or any(x <= 0 for x in levels):
+        raise argparse.ArgumentTypeError("--concurrency 中的并发度必须为正整数")
+    return levels
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="本地 LLM 服务性能基准：吞吐 / 延迟 / KV Cache / 批处理（实验 2-1 配套）",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=(
+            "场景说明:\n"
+            "  throughput  单流解码吞吐(tok/s)与首 token 延迟(TTFT)\n"
+            "  kv-cache    前缀缓存命中 vs 未命中的 TTFT 对比\n"
+            "  batching    不同并发度下的聚合吞吐（批处理权衡）\n"
+            "  all         依次运行以上全部场景\n"
+        ),
+    )
+    parser.add_argument(
+        "--scenario",
+        choices=["throughput", "kv-cache", "batching", "all"],
+        default="all",
+        help="要运行的基准场景（默认: all）",
+    )
+    parser.add_argument(
+        "--backend",
+        choices=["vllm", "ollama"],
+        default="vllm",
+        help="服务端类型，用于推断默认地址与模型名（默认: vllm）",
+    )
+    parser.add_argument(
+        "--base-url",
+        type=str,
+        default=None,
+        help="OpenAI 兼容接口地址，覆盖后端默认值（如 http://localhost:8000/v1）",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        help="模型名，覆盖后端默认值（vLLM 默认 Qwen3-0.6B，Ollama 默认 qwen3:0.6b）",
+    )
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default="EMPTY",
+        help="API Key，本地服务端一般无需真实值（默认: EMPTY）",
+    )
+    parser.add_argument(
+        "--repeats",
+        type=int,
+        default=5,
+        help="throughput / kv-cache 场景的重复次数（默认: 5）",
+    )
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=256,
+        help="每次请求的最大生成 token 数（默认: 256）",
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=0.7,
+        help="采样温度（默认: 0.7）",
+    )
+    parser.add_argument(
+        "--prefix-tokens",
+        type=int,
+        default=1024,
+        help="kv-cache 场景中共享前缀的近似 token 长度，越长缓存效果越明显（默认: 1024）",
+    )
+    parser.add_argument(
+        "--concurrency",
+        type=parse_concurrency,
+        default=[1, 2, 4, 8],
+        help="batching 场景的并发度列表，逗号分隔（默认: 1,2,4,8）",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default=None,
+        help="将结果以 JSON 写入指定文件",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="离线打印计划而不访问服务端，用于验证配置",
+    )
+    return parser
+
+
+def main() -> int:
+    parser = build_parser()
+    args = parser.parse_args()
+
+    # 用后端默认值补全 base_url / model
+    defaults = BACKEND_DEFAULTS[args.backend]
+    if args.base_url is None:
+        args.base_url = defaults["base_url"]
+    if args.model is None:
+        args.model = defaults["model"]
+
+    print("=" * 68)
+    print("🚀 本地 LLM 服务性能基准（实验 2-1）")
+    print("=" * 68)
+
+    if args.dry_run:
+        describe_dry_run(args)
+        return 0
+
+    client = make_client(args.base_url, args.api_key)
+    logger.info("连接服务端: %s（模型: %s）", args.base_url, args.model)
+
+    scenarios = (
+        ["throughput", "kv-cache", "batching"]
+        if args.scenario == "all"
+        else [args.scenario]
+    )
+    dispatch = {
+        "throughput": scenario_throughput,
+        "kv-cache": scenario_kv_cache,
+        "batching": scenario_batching,
+    }
+
+    results: List[Dict[str, Any]] = []
+    try:
+        for name in scenarios:
+            logger.info("开始场景: %s", name)
+            results.append(dispatch[name](client, args.model, args))
+    except Exception as e:  # noqa: BLE001
+        logger.error("基准执行失败: %s", e)
+        logger.info(
+            "请确认服务端已启动：vLLM 用 `python server.py`，"
+            "Ollama 用 `ollama serve` 并已 `ollama pull %s`",
+            args.model,
+        )
+        return 1
+
+    print_report(results)
+
+    if args.output:
+        payload = {
+            "backend": args.backend,
+            "base_url": args.base_url,
+            "model": args.model,
+            "results": results,
+        }
+        with open(args.output, "w", encoding="utf-8") as f:
+            json.dump(payload, f, ensure_ascii=False, indent=2)
+        logger.info("结果已写入: %s", args.output)
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())