liqiang
b119135836
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
488 lines
18 KiB
Python
488 lines
18 KiB
Python
#!/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())
|