feat: add soak runner with budget guards and invariant scoreboard

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()

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