feat: add soak runner with budget guards and invariant scoreboard
This commit is contained in:
@@ -129,3 +129,66 @@ class TestWeightedMix:
|
||||
def test_invalid_weights_rejected(self):
|
||||
with pytest.raises(ValueError):
|
||||
weighted_mix({"P1": 0.0}, lambda: 0.5)
|
||||
|
||||
|
||||
# ═══════════ T11: 记分板硬不变量(纯函数) ═══════════
|
||||
|
||||
from tools.soak.scoreboard import ( # noqa: E402
|
||||
inv_call_ids_unique,
|
||||
inv_rows_match_calls,
|
||||
inv_rpm_never_exceeded,
|
||||
inv_rss_stable,
|
||||
structured_success_rate,
|
||||
)
|
||||
|
||||
|
||||
def _row(call_id, *, source="s1", created_at="2026-07-20T10:00:00", error=None, session="r-p3"):
|
||||
return {
|
||||
"call_id": call_id,
|
||||
"source_name": source,
|
||||
"created_at": created_at,
|
||||
"error": error,
|
||||
"session_id": session,
|
||||
}
|
||||
|
||||
|
||||
class TestInvariants:
|
||||
def test_rows_match_calls(self):
|
||||
rows = [_row("a"), _row("b")]
|
||||
inv_rows_match_calls(rows, expected_calls=2)
|
||||
with pytest.raises(AssertionError):
|
||||
inv_rows_match_calls(rows, expected_calls=3)
|
||||
|
||||
def test_call_ids_unique(self):
|
||||
inv_call_ids_unique([_row("a"), _row("b")])
|
||||
with pytest.raises(AssertionError):
|
||||
inv_call_ids_unique([_row("a"), _row("a")])
|
||||
|
||||
def test_rpm_minute_bucket(self):
|
||||
# 固定分钟窗口口径(与限流器 int(sec/60) 同源;滑动窗会对合法的
|
||||
# 跨窗背靠背流量误报,不采用)
|
||||
rows = [
|
||||
_row("a", created_at="2026-07-20T10:00:01"),
|
||||
_row("b", created_at="2026-07-20T10:00:59"),
|
||||
_row("c", created_at="2026-07-20T10:01:01"),
|
||||
]
|
||||
inv_rpm_never_exceeded(rows, {"s1": 2})
|
||||
rows.append(_row("d", created_at="2026-07-20T10:00:30"))
|
||||
with pytest.raises(AssertionError):
|
||||
inv_rpm_never_exceeded(rows, {"s1": 2})
|
||||
|
||||
def test_rpm_ignores_unlimited_sources(self):
|
||||
inv_rpm_never_exceeded([_row(str(i)) for i in range(100)], {"s1": 0})
|
||||
|
||||
def test_rss_stable(self):
|
||||
inv_rss_stable([100.0, 105.0, 110.0], max_growth_mb=50.0)
|
||||
with pytest.raises(AssertionError):
|
||||
inv_rss_stable([100.0, 400.0], max_growth_mb=50.0)
|
||||
|
||||
def test_structured_success_rate(self):
|
||||
rows = [
|
||||
_row("a"),
|
||||
_row("b", error="ResultInvalidError: x"),
|
||||
_row("c", session="r-p1"), # 非 P3 剔除
|
||||
]
|
||||
assert structured_success_rate(rows, session_suffix="-p3") == pytest.approx(0.5)
|
||||
|
||||
@@ -0,0 +1,219 @@
|
||||
"""压测入口(M2 设计 §8): 预算硬顶、隔离守卫、多进程 worker、跑后记分。
|
||||
|
||||
用法(conda 环境内,项目根目录):
|
||||
python tools/soak/run_soak.py --scenario P3 --budget-calls 6 \\
|
||||
--budget-tokens 50000 --workers 2 [--scope LLM] [--concurrency 8]
|
||||
|
||||
守卫(违反即拒跑): REDIS_URL 必须指向 db3(实验室 db0 有在用键);
|
||||
PGW_TELEMETRY_PG_DSN 若配必须指向 polygateway 专用库;--workers>1 时
|
||||
限流/熔断后端必须是 redis(内存后端不跨进程,多 worker 无共享闸即超发)。
|
||||
预算双上限(请求数/token)任一命中优雅停;签字硬顶(设计 §8.1):
|
||||
全程 token ≤ 2 亿。soak 与 pytest 不并跑(FLUSHDB 清 db3 测试状态)。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import resource
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
_ROOT = Path(__file__).resolve().parents[2]
|
||||
sys.path.insert(0, str(_ROOT))
|
||||
sys.path.insert(0, str(_ROOT / "src"))
|
||||
|
||||
TOKEN_HARD_CAP = 200_000_000 # 2 亿,2026-07-20 人类签字(设计 §8.1)
|
||||
|
||||
|
||||
def _merged_env() -> dict[str, str]:
|
||||
from dotenv import dotenv_values
|
||||
|
||||
return {k: v for k, v in {**dotenv_values(_ROOT / ".env"), **os.environ}.items() if v}
|
||||
|
||||
|
||||
def _guard(env: dict[str, str], workers: int) -> None:
|
||||
redis_url = env.get("REDIS_URL", "")
|
||||
if not redis_url.rstrip("/").endswith("/3"):
|
||||
raise SystemExit(f"拒跑: REDIS_URL 必须指向专用 db3,当前 {redis_url!r}")
|
||||
pg = env.get("PGW_TELEMETRY_PG_DSN", "")
|
||||
if pg and not pg.rstrip("/").endswith("/polygateway"):
|
||||
raise SystemExit(f"拒跑: PG DSN 必须指向 polygateway 专用库,当前库名不符")
|
||||
if workers > 1 and (
|
||||
env.get("PGW_LIMITER_BACKEND") != "redis" or env.get("PGW_BREAKER_BACKEND") != "redis"
|
||||
):
|
||||
raise SystemExit("拒跑: --workers>1 需要 PGW_LIMITER_BACKEND/PGW_BREAKER_BACKEND=redis")
|
||||
|
||||
|
||||
async def _flush_db3(redis_url: str) -> None:
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
client = aioredis.from_url(redis_url)
|
||||
try:
|
||||
await client.flushdb()
|
||||
finally:
|
||||
await client.aclose()
|
||||
|
||||
|
||||
def _rss_mb() -> float:
|
||||
peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
|
||||
return peak / 1e6 if sys.platform == "darwin" else peak / 1024
|
||||
|
||||
|
||||
async def _worker_async(args: argparse.Namespace, worker_idx: int) -> None:
|
||||
"""单 worker: 消费场景生成器,信号量限并发,双预算任一命中即停。"""
|
||||
from polygateway import GatewayClient
|
||||
from polygateway.telemetry.sqlite import SQLiteRecorder
|
||||
|
||||
from tools.soak.scenarios import SCENARIOS, SoakCorpus
|
||||
|
||||
env = _merged_env()
|
||||
run_id = args.run_id
|
||||
telemetry_path = _ROOT / f"data/soak/telemetry_{run_id}_{worker_idx}.db"
|
||||
recorder = SQLiteRecorder(telemetry_path)
|
||||
client = GatewayClient.from_env(args.scope, telemetry=recorder, env=env)
|
||||
corpus = SoakCorpus(
|
||||
harness_db=_ROOT / "data/soak/harness.db",
|
||||
telemetry_db=_ROOT / "data/soak/generate_questions_telemetry.db",
|
||||
frames_root=_ROOT / "data/soak/vt_frames",
|
||||
images_root=_ROOT / "data/soak/chs_images",
|
||||
)
|
||||
generator = SCENARIOS[args.scenario](corpus, f"{run_id}-w{worker_idx}")
|
||||
sem = asyncio.Semaphore(args.concurrency)
|
||||
stats = {"calls": 0, "ok": 0, "failed": 0, "cancelled": 0, "tokens": 0}
|
||||
rss_samples = [_rss_mb()]
|
||||
deadline = time.monotonic() + args.max_hours * 3600
|
||||
budget_calls = args.budget_calls // args.workers
|
||||
budget_tokens = min(args.budget_tokens, TOKEN_HARD_CAP) // args.workers
|
||||
inflight: set[asyncio.Task] = set()
|
||||
|
||||
async def _one(kind: str, kwargs: dict) -> None:
|
||||
async with sem:
|
||||
try:
|
||||
# run 级双层隔离: FLUSHDB(入口)+ 每 run 独立 namespace(此处)
|
||||
kwargs.setdefault("cache_namespace", run_id)
|
||||
resp = await client.chat(**kwargs)
|
||||
stats["ok"] += 1
|
||||
stats["tokens"] += resp.prompt_tokens + resp.completion_tokens
|
||||
except asyncio.CancelledError:
|
||||
stats["cancelled"] += 1
|
||||
raise
|
||||
except Exception as exc: # 记分板从遥测读错误分布;这里只计数
|
||||
stats["failed"] += 1
|
||||
print(f"[w{worker_idx}] 调用失败: {type(exc).__name__}: {exc}", flush=True)
|
||||
|
||||
async for kind, kwargs in generator:
|
||||
if (
|
||||
stats["calls"] >= budget_calls
|
||||
or stats["tokens"] >= budget_tokens
|
||||
or time.monotonic() > deadline
|
||||
):
|
||||
break
|
||||
stats["calls"] += 1
|
||||
task = asyncio.create_task(_one(kind, kwargs))
|
||||
inflight.add(task)
|
||||
task.add_done_callback(inflight.discard)
|
||||
if stats["calls"] % 20 == 0:
|
||||
rss_samples.append(_rss_mb())
|
||||
print(f"[w{worker_idx}] {stats}", flush=True)
|
||||
if inflight:
|
||||
await asyncio.gather(*inflight, return_exceptions=True) # 优雅收尾 in-flight
|
||||
rss_samples.append(_rss_mb())
|
||||
await client.aclose()
|
||||
result = {"stats": stats, "rss_mb": rss_samples, "telemetry": str(telemetry_path)}
|
||||
(_ROOT / f"data/soak/result_{run_id}_{worker_idx}.json").write_text(
|
||||
json.dumps(result), encoding="utf-8"
|
||||
)
|
||||
print(f"[w{worker_idx}] 完成: {stats}", flush=True)
|
||||
|
||||
|
||||
def _worker_main(args: argparse.Namespace, worker_idx: int) -> None:
|
||||
asyncio.run(_worker_async(args, worker_idx))
|
||||
|
||||
|
||||
def _scoreboard(args: argparse.Namespace, env: dict[str, str]) -> None:
|
||||
"""跑后记分: 合并 worker 遥测,断言硬不变量,产出报告。"""
|
||||
from tools.soak import scoreboard as sb
|
||||
|
||||
results = []
|
||||
for w in range(args.workers):
|
||||
path = _ROOT / f"data/soak/result_{args.run_id}_{w}.json"
|
||||
results.append(json.loads(path.read_text(encoding="utf-8")))
|
||||
rows = sb.load_rows(*(r["telemetry"] for r in results))
|
||||
calls = sum(r["stats"]["calls"] for r in results)
|
||||
max_attempts = int(env.get(f"{args.scope}__RETRY__MAX_ATTEMPTS", env.get("LLM_MAX_RETRIES", "3")))
|
||||
verdicts: list[tuple[str, str]] = []
|
||||
|
||||
def _check(name: str, fn, *fargs, **fkwargs) -> None:
|
||||
try:
|
||||
fn(*fargs, **fkwargs)
|
||||
verdicts.append((name, "PASS"))
|
||||
except AssertionError as exc:
|
||||
verdicts.append((name, f"FAIL — {exc}"))
|
||||
|
||||
# 行数下界=请求数(每请求至少 1 行),上界容重试放大(每请求 ≤ max_attempts 行)
|
||||
_check(
|
||||
"遥测完备(行数≥请求数,重试容差内)",
|
||||
sb.inv_rows_match_calls,
|
||||
rows,
|
||||
expected_calls=calls,
|
||||
tolerance=calls * max(max_attempts - 1, 0) + sum(r["stats"]["cancelled"] for r in results),
|
||||
)
|
||||
_check("call_id 唯一", sb.inv_call_ids_unique, rows)
|
||||
rpm_conf = {}
|
||||
for key, val in env.items():
|
||||
parts = key.split("__")
|
||||
# 仅四段源键 {SCOPE}__{PROVIDER}__{N}__RPM(排除 {SCOPE}__GLOBAL__RPM)
|
||||
if len(parts) == 4 and parts[0] == args.scope and parts[3] == "RPM":
|
||||
rpm_conf[f"{parts[1].lower()}_{parts[2]}"] = int(val)
|
||||
_check("RPM 从未击穿(分钟桶)", sb.inv_rpm_never_exceeded, rows, rpm_conf)
|
||||
for r in results:
|
||||
_check(f"RSS 平稳(w)", sb.inv_rss_stable, r["rss_mb"], max_growth_mb=args.max_rss_growth_mb)
|
||||
rate = sb.structured_success_rate(rows, session_suffix="-p3")
|
||||
verdicts.append(("P3 结构化成功率", f"{rate:.3f}(基线首跑建立)"))
|
||||
report = sb.render_report(args.run_id, rows, verdicts)
|
||||
path = sb.write_report(args.run_id, report)
|
||||
print(report)
|
||||
print(f"报告: {path}")
|
||||
if any(v.startswith("FAIL") for _, v in verdicts):
|
||||
raise SystemExit("硬不变量被击穿,见报告")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="PolyGateway 真实数据压测")
|
||||
parser.add_argument("--scenario", required=True, choices=["P1", "P2", "P3", "P4", "P5", "P6"])
|
||||
parser.add_argument("--budget-calls", type=int, required=True)
|
||||
parser.add_argument("--budget-tokens", type=int, required=True)
|
||||
parser.add_argument("--workers", type=int, default=1)
|
||||
parser.add_argument("--concurrency", type=int, default=8, help="单 worker 并发上限")
|
||||
parser.add_argument("--scope", default="SOAK")
|
||||
parser.add_argument("--run-id", default=None)
|
||||
parser.add_argument("--max-hours", type=float, default=3.0)
|
||||
parser.add_argument("--max-rss-growth-mb", type=float, default=500.0)
|
||||
args = parser.parse_args()
|
||||
if args.run_id is None:
|
||||
args.run_id = time.strftime("soak_%Y%m%d_%H%M%S")
|
||||
env = _merged_env()
|
||||
_guard(env, args.workers)
|
||||
print(f"run_id={args.run_id}: FLUSHDB db3 + namespace 隔离")
|
||||
asyncio.run(_flush_db3(env["REDIS_URL"]))
|
||||
if args.workers == 1:
|
||||
_worker_main(args, 0)
|
||||
else:
|
||||
ctx = mp.get_context("spawn")
|
||||
procs = [ctx.Process(target=_worker_main, args=(args, w)) for w in range(args.workers)]
|
||||
for p in procs:
|
||||
p.start()
|
||||
for p in procs:
|
||||
p.join()
|
||||
if any(p.exitcode != 0 for p in procs):
|
||||
raise SystemExit(f"worker 退出码异常: {[p.exitcode for p in procs]}")
|
||||
_scoreboard(args, env)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,134 @@
|
||||
"""压测记分板: 六条硬不变量断言 + 报告产出(findings §4;吃自己遥测的狗粮)。
|
||||
|
||||
输入口径三类(M2 计划 T11): 不变量 2/3/5/6 = 纯函数(遥测行迭代器);
|
||||
不变量 1(记账归零)= 活 Redis/后端检查(async);不变量 4(RSS)= 纯函数
|
||||
(run_soak 周期采样 json)。断言函数抛 AssertionError 即不变量被击穿。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
from collections import Counter
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
Row = dict[str, Any]
|
||||
|
||||
|
||||
def load_rows(*db_paths: Path | str) -> list[Row]:
|
||||
"""合并全部 worker 遥测库的 llm_calls 行(dict 形态)。"""
|
||||
rows: list[Row] = []
|
||||
for path in db_paths:
|
||||
conn = sqlite3.connect(path)
|
||||
conn.row_factory = sqlite3.Row
|
||||
try:
|
||||
rows.extend(dict(r) for r in conn.execute("SELECT * FROM llm_calls"))
|
||||
finally:
|
||||
conn.close()
|
||||
return rows
|
||||
|
||||
|
||||
# —— 硬不变量(纯函数) ——
|
||||
|
||||
|
||||
def inv_rows_match_calls(rows: list[Row], *, expected_calls: int, tolerance: int = 0) -> None:
|
||||
"""不变量 2a: 遥测行数 == 发出的请求数(±取消双记的已知语义容差)。"""
|
||||
assert abs(len(rows) - expected_calls) <= tolerance, (
|
||||
f"遥测行数 {len(rows)} ≠ 请求数 {expected_calls}(容差 {tolerance})"
|
||||
)
|
||||
|
||||
|
||||
def inv_call_ids_unique(rows: list[Row]) -> None:
|
||||
"""不变量 2b: call_id 无重复。"""
|
||||
dupes = [cid for cid, n in Counter(r["call_id"] for r in rows).items() if n > 1]
|
||||
assert not dupes, f"call_id 重复: {dupes[:5]}"
|
||||
|
||||
|
||||
def _minute_bucket(created_at: str) -> int:
|
||||
return int(datetime.fromisoformat(str(created_at)).timestamp()) // 60
|
||||
|
||||
def inv_rpm_never_exceeded(rows: list[Row], per_source_rpm: dict[str, int]) -> None:
|
||||
"""不变量 3: 按遥测时间戳重算,任一分钟桶内单源请求数 ≤ RPM 配置。
|
||||
|
||||
口径 = 固定分钟窗口(与限流器 `int(sec/60)` 同源);滑动 60s 窗会对
|
||||
"窗尾+窗头"的合法背靠背流量误报,不采用(findings §4 条 3 的执行口径)。
|
||||
"""
|
||||
buckets: Counter[tuple[str, int]] = Counter(
|
||||
(r["source_name"], _minute_bucket(r["created_at"]))
|
||||
for r in rows
|
||||
if per_source_rpm.get(r["source_name"], 0) > 0
|
||||
)
|
||||
breaches = {
|
||||
key: n for key, n in buckets.items() if n > per_source_rpm[key[0]]
|
||||
}
|
||||
assert not breaches, f"RPM 击穿: {dict(list(breaches.items())[:5])}"
|
||||
|
||||
|
||||
def inv_rss_stable(samples_mb: list[float], *, max_growth_mb: float) -> None:
|
||||
"""不变量 4: 进程 RSS 首末差 < 阈值(窗口计数器/租约表无无界增长)。"""
|
||||
assert len(samples_mb) >= 2, "RSS 采样不足"
|
||||
growth = samples_mb[-1] - samples_mb[0]
|
||||
assert growth < max_growth_mb, f"RSS 增长 {growth:.1f}MB ≥ 阈值 {max_growth_mb}MB"
|
||||
|
||||
|
||||
def structured_success_rate(rows: list[Row], *, session_suffix: str) -> float:
|
||||
"""不变量 5 的量: 指定场景行的最终成功率(error 为空即成功)。"""
|
||||
scoped = [r for r in rows if str(r.get("session_id") or "").endswith(session_suffix)]
|
||||
if not scoped:
|
||||
return 0.0
|
||||
ok = sum(1 for r in scoped if not r.get("error"))
|
||||
return ok / len(scoped)
|
||||
|
||||
|
||||
async def inv_accounting_zeroed(limiter, sources: list[str]) -> None:
|
||||
"""不变量 1: 结束后全部源 inflight == 0(无泄漏租约)。活后端检查。"""
|
||||
for name in sources:
|
||||
stats = await limiter.source_stats(name)
|
||||
assert stats.inflight == 0, f"源 {name} 结束后 inflight={stats.inflight}(租约泄漏)"
|
||||
|
||||
|
||||
# —— 报告 ——
|
||||
|
||||
|
||||
def _percentile(values: list[float], q: float) -> float:
|
||||
if not values:
|
||||
return 0.0
|
||||
ordered = sorted(values)
|
||||
return ordered[min(int(len(ordered) * q), len(ordered) - 1)]
|
||||
|
||||
|
||||
def render_report(run_id: str, rows: list[Row], verdicts: list[tuple[str, str]]) -> str:
|
||||
"""结构化 Markdown 报告(不变量 6: 成本/延迟基线产出)。"""
|
||||
errors = Counter(str(r["error"]).split(":")[0] for r in rows if r.get("error"))
|
||||
latencies = [float(r["latency_ms"]) for r in rows if not r.get("error")]
|
||||
ttfts = [float(r["ttft_ms"]) for r in rows if r.get("ttft_ms") is not None]
|
||||
costs = [float(r["cost"]) for r in rows if r.get("cost") is not None]
|
||||
tokens = sum(int(r["prompt_tokens"]) + int(r["completion_tokens"]) for r in rows)
|
||||
cache_hits = sum(1 for r in rows if r.get("cache_hit"))
|
||||
lines = [
|
||||
f"# Soak 报告: {run_id}",
|
||||
"",
|
||||
"## 不变量裁决",
|
||||
*(f"- {name}: {verdict}" for name, verdict in verdicts),
|
||||
"",
|
||||
"## 规模",
|
||||
f"- 遥测行数: {len(rows)};缓存命中: {cache_hits};总 token: {tokens}",
|
||||
f"- 错误分布: {dict(errors) or '无'}",
|
||||
"",
|
||||
"## 延迟基线(成功行)",
|
||||
f"- 总时长 p50/p95: {_percentile(latencies, 0.5):.0f} / {_percentile(latencies, 0.95):.0f} ms",
|
||||
f"- TTFT p50/p95: {_percentile(ttfts, 0.5):.0f} / {_percentile(ttfts, 0.95):.0f} ms",
|
||||
"",
|
||||
"## 成本",
|
||||
f"- 累计 cost: {sum(costs):.4f}(有单价行 {len(costs)}/{len(rows)})",
|
||||
]
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def write_report(run_id: str, content: str, out_dir: Path | str = "tests/outputs/soak") -> Path:
|
||||
out = Path(out_dir)
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
path = out / f"{run_id}.md"
|
||||
path.write_text(content, encoding="utf-8")
|
||||
return path
|
||||
Reference in New Issue
Block a user