From 9aa97a61ae6b6f04693bf390a28a32dbda7d19a3 Mon Sep 17 00:00:00 2001 From: iomgaa Date: Tue, 21 Jul 2026 01:23:40 -0400 Subject: [PATCH] feat: add soak runner with budget guards and invariant scoreboard --- tests/unit/test_soak_corpus.py | 63 ++++++++++ tools/soak/run_soak.py | 219 +++++++++++++++++++++++++++++++++ tools/soak/scoreboard.py | 134 ++++++++++++++++++++ 3 files changed, 416 insertions(+) create mode 100644 tools/soak/run_soak.py create mode 100644 tools/soak/scoreboard.py diff --git a/tests/unit/test_soak_corpus.py b/tests/unit/test_soak_corpus.py index 64a8c1a..ec9f792 100644 --- a/tests/unit/test_soak_corpus.py +++ b/tests/unit/test_soak_corpus.py @@ -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) diff --git a/tools/soak/run_soak.py b/tools/soak/run_soak.py new file mode 100644 index 0000000..308c57c --- /dev/null +++ b/tools/soak/run_soak.py @@ -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() diff --git a/tools/soak/scoreboard.py b/tools/soak/scoreboard.py new file mode 100644 index 0000000..14e4ae4 --- /dev/null +++ b/tools/soak/scoreboard.py @@ -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